From 0c3626ec06e913e2aa5ee62256d63ed0d95d969e Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 21 Sep 2026 01:00:28 -0700 Subject: [PATCH 01/42] hexagon: overhaul of buffer and DMA handling to support 64bit mappings + improvements (#29197) * hex-dma64: enable support extended buffer mappings and 64bit dma hex-dma64: expand binary ops to support more DMA scenarios hex-dma64: add binary-ops.h hex-dma64: add --hex-dma64 to run.py and fix minor issues hex-dma64: update SSM_CONV to use dma with proper support for 64bit hex-ops: remove obsolete gate for % 128 in binary ops hex-l2: dont check weight tensors against dirty ranges hex-dma64: most binary ops now support dma hex-dma: use dma_addr_t instead of plain uint64_t to avoid overhead on older targets hex-dma: update all dma users to use dma_data (instead of pointers) hex-dma64: simplify lazy buffer mapping and clonning hex-fusion: factor out try_fuse_common that checks for dma64 buffers hex-bufs: minor cleanup for mmaping logic hex-bufs: simplify buffer clonning hex-ssm-conv: tighten gating checks and check vtcm size in kparams hex-binary: fix incorred mod/wrap in scalar ops hex-binary: make sure to call precompute kparams in support checks hex-dma64: update addr handling in mm,concat,binary hex-dma64: fixing up leftover of dma_addr_t conversion hex-binary: redo the kernel selection again and fix regressions in MOEs hex-binary: specialize per-type/per-op hex-binary: vtcm-layout and per-src dma-queue hex-dma64: update dma_push to transparently handle 64bit/extended * hex-cpy: fix improper rebase with the fixes for cont. tensors * hex-dma-cpy: update CPY to use safe dma rows/size limits * hex-mmap: bump number of mmaps to 64 to allow avoid eviction in larger models * hex-dma: add support for the secondary ring as a fallback for too-large transactions * hex-rope: fix freq_factors access with 64bit dma * hex-dma: audit all ops for proper use/gards for 64bit addresses * hex-dma64: uninline glu-compute funcs to avoid register pressure due to 64bit addr math * hex-dma64: refactor binary ops to separate dma loops * hex-devel: add inspect script to help with dbg and analysis * hex-dma: refactor dma-pipelines in unary-ops * hex-dma: rewrite softmax to use dma * hex-dma: rewrite GDN dma loops and improve HVX register usage * hex-gdn: fuse GDN+CPY * hex-mm: factor out HVX solver * hex-mm: remove hvx-flat kernels, the chunked version now handles vtcm limits much better * hex-buffs: reject huge buffer allocations that we cannot memory map * hex-inspect: add logic to look for float promo calls * hex-mm: reduce HVX register spills in HVX prompt kernels * hex-bufs: do not double count buffers from tensors in the same op * hex-roll: fix merge conflict * hex-dma: reroute all matmul ddr kernels to new chunked dma/vtcm kernels * hex-dev: update developer docs to include inspection for register spils and float promos * hex-ops: forgot to add new headers * hex-softmax: fix gpt-oss dims * hex-dma64: cleanup dma_addr_t casts * hex-dma64: add support for dma/vtcm for flash-atten with sinks * hex-mm-add: fix MUL_MAT+ADD fusion with bias.weights in extended bufs * hex-add-id: add support for dma for src1 (exp. table) * hex-dma: imrpove v73 fallback paths * hex-bufs: do not drop extended mappings during va defrag * hex-scripts: fix flake8 warnings * hex-docs: fix editor-config warnings * hex-inspect: fix warnings from ty --- docs/backend/snapdragon/README.md | 3 + docs/backend/snapdragon/developer.md | 22 + ggml/src/ggml-hexagon/ggml-hexagon.cpp | 1052 ++++++++--- ggml/src/ggml-hexagon/htp-opnode.h | 23 +- ggml/src/ggml-hexagon/htp/act-ops.c | 288 ++- ggml/src/ggml-hexagon/htp/allreduce-ops.c | 226 +-- ggml/src/ggml-hexagon/htp/argsort-ops.c | 6 +- ggml/src/ggml-hexagon/htp/binary-ops.c | 1074 +++++++---- ggml/src/ggml-hexagon/htp/binary-ops.h | 111 ++ ggml/src/ggml-hexagon/htp/concat-ops.c | 46 +- ggml/src/ggml-hexagon/htp/cpy-ops.c | 358 ++-- ggml/src/ggml-hexagon/htp/cumsum-ops.c | 47 +- ggml/src/ggml-hexagon/htp/diag-ops.c | 35 +- ggml/src/ggml-hexagon/htp/dma-queue.c | 186 +- ggml/src/ggml-hexagon/htp/dma-queue.h | 278 +-- ggml/src/ggml-hexagon/htp/fill-ops.c | 4 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 276 +-- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 89 +- .../ggml-hexagon/htp/gated-delta-net-ops.c | 903 +++++---- .../ggml-hexagon/htp/gated-delta-net-ops.h | 297 +++ ggml/src/ggml-hexagon/htp/get-rows-ops.c | 262 +-- ggml/src/ggml-hexagon/htp/hex-dma.h | 2 - ggml/src/ggml-hexagon/htp/htp-ctx.h | 12 +- ggml/src/ggml-hexagon/htp/htp-ops.h | 11 +- ggml/src/ggml-hexagon/htp/htp-tensor.c | 11 +- ggml/src/ggml-hexagon/htp/htp-tensor.h | 4 + ggml/src/ggml-hexagon/htp/htp_iface.idl | 2 +- ggml/src/ggml-hexagon/htp/hvx-exp.h | 4 +- .../ggml-hexagon/htp/hvx-mm-kernels-flat.h | 1648 ----------------- .../ggml-hexagon/htp/hvx-mm-kernels-float.h | 382 ++++ .../ggml-hexagon/htp/hvx-mm-kernels-tiled.h | 334 +--- ggml/src/ggml-hexagon/htp/im2col-ops.c | 454 ++--- ggml/src/ggml-hexagon/htp/main.c | 171 +- ggml/src/ggml-hexagon/htp/matmul-ops.c | 1374 ++++++++------ ggml/src/ggml-hexagon/htp/matmul-ops.h | 169 +- ggml/src/ggml-hexagon/htp/pad-ops.c | 107 +- ggml/src/ggml-hexagon/htp/repeat-ops.c | 4 +- ggml/src/ggml-hexagon/htp/roll-ops.c | 48 +- ggml/src/ggml-hexagon/htp/rope-ops.c | 48 +- ggml/src/ggml-hexagon/htp/rope-ops.h | 22 +- ggml/src/ggml-hexagon/htp/set-rows-ops.c | 42 +- ggml/src/ggml-hexagon/htp/softmax-ops.c | 704 ++++--- ggml/src/ggml-hexagon/htp/softmax-ops.h | 106 ++ ggml/src/ggml-hexagon/htp/solve-tri-ops.c | 4 +- ggml/src/ggml-hexagon/htp/ssm-conv.c | 528 +++--- ggml/src/ggml-hexagon/htp/ssm-conv.h | 40 + ggml/src/ggml-hexagon/htp/sum-rows-ops.c | 6 +- ggml/src/ggml-hexagon/htp/unary-ops.c | 1419 ++++++++------ scripts/snapdragon/ggml-hexagon-inspect.py | 1106 +++++++++++ scripts/snapdragon/run.py | 5 +- tests/test-backend-ops.cpp | 2 + 51 files changed, 8234 insertions(+), 6121 deletions(-) create mode 100644 ggml/src/ggml-hexagon/htp/binary-ops.h create mode 100644 ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h delete mode 100644 ggml/src/ggml-hexagon/htp/hex-dma.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h create mode 100644 ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h create mode 100644 ggml/src/ggml-hexagon/htp/softmax-ops.h create mode 100644 ggml/src/ggml-hexagon/htp/ssm-conv.h create mode 100755 scripts/snapdragon/ggml-hexagon-inspect.py diff --git a/docs/backend/snapdragon/README.md b/docs/backend/snapdragon/README.md index 5d32a5877..ca79ca885 100644 --- a/docs/backend/snapdragon/README.md +++ b/docs/backend/snapdragon/README.md @@ -327,6 +327,9 @@ on 4 physical NPUs, or `--devices 'HTP0[0-1:0],HTP1[0-1:1]'` on 2 physical NPUs - `GGML_HEXAGON_HOSTBUF=1` (default: 0, disabled) Enables allocating host buffers for debugging. By default, host buffers are disabled. +- `GGML_HEXAGON_DMA64=0` (default: enabled on v81+) + Disables 64-bit DMA for model weights. Set to `1` to enable it explicitly on a supported architecture. + - `GGML_HEXAGON_VERBOSE=1` Enables verbose logging of Ops from the backend. Example output: diff --git a/docs/backend/snapdragon/developer.md b/docs/backend/snapdragon/developer.md index 633643c16..378d47653 100644 --- a/docs/backend/snapdragon/developer.md +++ b/docs/backend/snapdragon/developer.md @@ -146,6 +146,28 @@ Writing high-performance operators for Hexagon requires following specific guide python3 scripts/snapdragon/ggml-hexagon-align-macros.py --fix ggml/src/ggml-hexagon/htp/ ``` +### Binary Inspection and Spill Analysis + +Use [`scripts/snapdragon/ggml-hexagon-inspect.py`](../../../scripts/snapdragon/ggml-hexagon-inspect.py) to audit Hexagon binaries for register +spills, unexpected float promotions, or disassembly: + +- Always verify that compute kernels have zero in-loop vector spills (`--spills --strict`) and no float promotions (`--promotions`). +- Avoid excessive loop unrolling (`#pragma unroll`), which increases register pressure and causes spills. + +```bash +# Check for vector and scalar register spills +python3 scripts/snapdragon/ggml-hexagon-inspect.py --spills --strict --func "^compute_" + +# Check for float promotions +python3 scripts/snapdragon/ggml-hexagon-inspect.py --promotions --func "^compute_" + +# Disassemble with annotated loops and spill markers +python3 scripts/snapdragon/ggml-hexagon-inspect.py --disasm compute_same_shape_div_f32 + +# Resolve crash addresses to function symbols and lines +python3 scripts/snapdragon/ggml-hexagon-inspect.py --addr2line 0x51a30 0x5ba54 +``` + ## Multi-Device Partitioning (mdev) Multi-device (mdev) mode enables row-level tensor parallel execution across multiple physical NPU cores or virtual NPU diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 352434b6a..ec5a4aeb6 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #ifdef _WIN32 # define WIN32_LEAN_AND_MEAN @@ -53,11 +54,15 @@ #include "htp-opnode.h" #include "htp-ops.h" #include "htp/matmul-ops.h" +#include "htp/binary-ops.h" #include "htp/flash-attn-ops.h" #include "htp/unary-ops.h" #include "htp/get-rows-ops.h" #include "htp/set-rows-ops.h" +#include "htp/softmax-ops.h" #include "htp/rope-ops.h" +#include "htp/ssm-conv.h" +#include "htp/gated-delta-net-ops.h" #include "htp_iface.h" #include "htp-drv.h" @@ -91,8 +96,9 @@ static int opt_etm = 0; static int opt_verbose = 0; static int opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu) static bool opt_hostbuf = false; +static bool opt_dma64 = false; -static int opt_mm_select = 3; // 3 = HMX -> Tiled -> Flat -> CPU, 2 = Tiled -> Flat -> CPU, 1 = Flat -> CPU +static int opt_mm_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported) static int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported) static int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 = unfused ALLREDUCE (DMA), 0 = fallback to CPY+FENCE @@ -113,6 +119,7 @@ enum ggml_hexagon_fusion_flags { GGML_HEXAGON_FUSE_MUL_MAT_ADD = (1 << 3), // 8 GGML_HEXAGON_FUSE_MUL_MAT_NX = (1 << 4), // 16 GGML_HEXAGON_FUSE_MUL_MAT_ID_NX = (1 << 5), // 32 + GGML_HEXAGON_FUSE_GDN_CPY = (1 << 6), // 64 }; static inline bool ggml_hexagon_is_fusion_enabled(int flag) { @@ -299,6 +306,15 @@ static void ggml_hexagon_precompute_unary_params( struct htp_unary_kernel_params * kparams ); +static bool ggml_hexagon_precompute_binary_params( + const struct ggml_hexagon_session * sess, + uint32_t op, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_binary_kernel_params * kparams +); + static void ggml_hexagon_precompute_get_rows_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, @@ -315,12 +331,32 @@ static void ggml_hexagon_precompute_set_rows_params( struct htp_set_rows_kernel_params * kparams ); +static void ggml_hexagon_precompute_softmax_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * op, + struct htp_softmax_kernel_params * kparams +); + static void ggml_hexagon_precompute_rope_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * op, struct htp_rope_kernel_params * kparams ); +static void ggml_hexagon_precompute_ssm_conv_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_ssm_conv_kernel_params * kparams +); + +static void ggml_hexagon_precompute_gated_delta_net_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * op, + struct htp_gdn_kernel_params * kparams +); + static void ggml_hexagon_precompute_fused_mmnx_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, @@ -349,6 +385,7 @@ static bool ggml_hexagon_precompute_allreduce_params( ); static bool mm_is_hmx_eligible(const ggml_tensor * t); +static htp_op_code op_remap_to_htp(const ggml_tensor * t); static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams); static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams); static bool is_mergeable_mul_mat(const ggml_tensor * t); @@ -474,12 +511,14 @@ struct ggml_hexagon_session { const std::vector & sync_tensors, uint32_t rank, uint32_t n_ranks, uint32_t fence_seq_entry = 0, uint32_t fence_seq_exit = 0); + void start_batch(); void flush_sync(bool all = true); void flush_async(); void flush_batch(size_t min_ops = 1); void flush_peers(); void flush_pending(bool all = true); + ggml_hexagon_shared_buffer * mmap_tensor(const ggml_tensor * t); bool clone_buffer(const ggml_hexagon_shared_buffer*); void release_buffer(const ggml_hexagon_shared_buffer*); void unclone_buffer(const ggml_hexagon_shared_buffer*); @@ -515,6 +554,9 @@ struct ggml_backend_hexagon_device_context { ggml_hexagon_session * session() { if (!sess) { sess = std::make_unique(config, dev); + if (max_bufsize > sess->max_vmem) { + max_bufsize = sess->max_vmem; + } } return sess.get(); } @@ -563,15 +605,21 @@ struct ggml_hexagon_shared_buffer { std::vector tensor_extra; bool mapped; bool pinned; + bool extended; const char * c_name() const { return sess->c_name(); } uint8_t * base() const { return mem ? mem->base : nullptr; } size_t size() const { return mem ? mem->size : 0; } int fd() const { return mem ? mem->fd : -1; } - void mmap() { - if (!this->mem) return; - fastrpc_map_flags flags = this->pinned ? FASTRPC_MAP_FD : FASTRPC_MAP_FD_DELAYED; + void mmap(bool extended = false) { + if (!this->mem) return; + if (this->mapped) return; + + GGML_ASSERT(!this->pinned || !extended); + + this->extended = extended; + fastrpc_map_flags flags = this->pinned ? FASTRPC_MAP_FD : (extended ? FASTRPC_MAP_FD_DELAYED_EXTENDED : FASTRPC_MAP_FD_DELAYED); int err = fastrpc_mmap(sess->domain_id, fd(), (void *) base(), 0, size(), flags); if (err != 0) { @@ -580,8 +628,8 @@ struct ggml_hexagon_shared_buffer { throw std::runtime_error("ggml-hex: fastrpc_mmap failed (see log for details)"); } - HEX_VERBOSE("ggml-hex: %s mapped buffer: base %p size %zu fd %d pinned %u\n", - sess->c_name(), (void *) base(), size(), fd(), pinned); + HEX_VERBOSE("ggml-hex: %s mapped buffer: base %p size %zu fd %d pinned %u extended %u\n", + sess->c_name(), (void *) base(), size(), fd(), pinned, extended); this->mapped = true; } @@ -611,7 +659,9 @@ struct ggml_hexagon_shared_buffer { HEX_VERBOSE("ggml-hex: %s allocated buffer: base %p size %zu fd %d pinned %d\n", sess->c_name(), (void *) base(), this->size(), fd(), (int) pinned); - mmap(); + if (this->pinned) { + mmap(); + } } void free() { @@ -623,13 +673,18 @@ struct ggml_hexagon_shared_buffer { } ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, size_t size, bool pinned = false) { - this->sess = sess; - this->mapped = false; - this->pinned = pinned; + this->sess = sess; + this->mapped = false; + this->pinned = pinned; + this->extended = false; // Size adjustment inside the buffer class: 4K aligned data size + 4K guard page size_t guard_offset = (size + 4095) & ~4095; size_t total_size = guard_offset + 4096; + if (!pinned && opt_dma64) { + constexpr size_t extended_align = 2 * 1024 * 1024; + total_size = (total_size + extended_align - 1) & ~(extended_align - 1); + } alloc(total_size); } @@ -640,6 +695,7 @@ struct ggml_hexagon_shared_buffer { this->mem = other.mem; this->mapped = false; this->pinned = other.pinned; + this->extended = other.extended; } ~ggml_hexagon_shared_buffer() { @@ -663,6 +719,7 @@ struct ggml_hexagon_fence_buffer : public ggml_hexagon_shared_buffer { backend_buffer.buft = buft; backend_buffer.context = static_cast(this); backend_buffer.size = size; + mmap(false); } uint8_t * alloc_slot(uint32_t n_slots = 1) { @@ -703,11 +760,6 @@ inline void ggml_hexagon_session::free_fence(void * ptr, uint32_t n_slots) { } } -static ggml_hexagon_session * ggml_backend_hexagon_buffer_get_sess(ggml_backend_buffer_t buffer) { - auto sbuf = static_cast(buffer->context); - return sbuf->sess; -} - static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer) { auto sbuf = static_cast(buffer->context); sbuf->sess->unclone_buffer(sbuf); @@ -2024,7 +2076,17 @@ static const char * ggml_backend_hexagon_buffer_type_name(ggml_backend_buffer_ty static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto dev_ctx = static_cast(buffer_type->context)->dev_ctx; + if (size > dev_ctx->max_bufsize) { + GGML_LOG_ERROR("ggml-hex: %s buffer size %zu exceeds max_bufsize %zu\n", + dev_ctx->c_name(), size, dev_ctx->max_bufsize); + return nullptr; + } auto sess = dev_ctx->session(); + if (sess && sess->max_vmem && size > sess->max_vmem) { + GGML_LOG_ERROR("ggml-hex: %s buffer size %zu exceeds max_vmem %zu\n", + dev_ctx->c_name(), size, sess->max_vmem); + return nullptr; + } try { ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false); return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size); @@ -2037,7 +2099,17 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto dev_ctx = static_cast(buffer_type->context)->dev_ctx; + if (size > dev_ctx->max_bufsize) { + GGML_LOG_ERROR("ggml-hex: %s host buffer size %zu exceeds max_bufsize %zu\n", + dev_ctx->c_name(), size, dev_ctx->max_bufsize); + return nullptr; + } auto sess = dev_ctx->session(); + if (sess && sess->max_vmem && size > sess->max_vmem) { + GGML_LOG_ERROR("ggml-hex: %s host buffer size %zu exceeds max_vmem %zu\n", + dev_ctx->c_name(), size, sess->max_vmem); + return nullptr; + } try { ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false); return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_host_buffer_interface, sbuf, size); @@ -2067,7 +2139,9 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { auto * context = static_cast(buft->context); - return context->dev_ctx->max_bufsize; + auto dev_ctx = context->dev_ctx; + dev_ctx->session(); + return dev_ctx->max_bufsize; } static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) { @@ -2139,7 +2213,7 @@ struct ggml_hexagon_opbatch { unsigned int n_bufs; // num buffers in the batch unsigned int n_tens; // num tensors ... unsigned int n_ops; // num ops ... - size_t b_vmem; // sum of all buffer sizes + size_t b_vmem; // sum of non-extended buffer sizes unsigned int n_bufs_max; unsigned int n_tens_max; @@ -2198,11 +2272,14 @@ struct ggml_hexagon_opbatch { b_map.insert({sbuf->fd(), bi}); htp_buf_desc &b = h_bufs[bi]; - b.base = (uint64_t) sbuf->base(); - b.fd = sbuf->fd(); - b.size = sbuf->size(); + b.base = (uint64_t) sbuf->base(); + b.fd = sbuf->fd(); + b.size = sbuf->size(); + b.flags = sbuf->extended ? HTP_BUF_EXTENDED : 0; - b_vmem += b.size; + if (!sbuf->extended) { + b_vmem += b.size; + } HEX_VERBOSE("ggml-hex: %s add-buffer #%u : fd %d base %p size %zu : vmem %zu\n", sess->c_name(), bi, b.fd, (void*) sbuf->base(), (size_t) b.size, b_vmem); @@ -2298,21 +2375,33 @@ struct ggml_hexagon_opbatch { } bool fit_op(const htp_opnode & node) const { - if (n_ops >= n_ops_max ) return false; + if (n_ops >= n_ops_max) return false; // check how much extras we will need size_t extra_bufs = 0; size_t extra_vmem = 0; size_t extra_tens = 0; + int seen_bufs[HTP_OP_MAX_BUFS]; + int n_seen_bufs = 0; + auto fit_tensor = [&](const ggml_tensor *t) { if (!t) return; if (!t_map.count(t)) { extra_tens++; auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); + int fd = sbuf->fd(); + if (!b_map.count(fd)) { + for (int i = 0; i < n_seen_bufs; i++) { + if (seen_bufs[i] == fd) return; + } + if (n_seen_bufs < HTP_OP_MAX_BUFS) { + seen_bufs[n_seen_bufs++] = fd; + } + if (!sbuf->extended) { + extra_vmem += sbuf->size(); + } extra_bufs += 1; } } @@ -2402,6 +2491,50 @@ struct ggml_hexagon_opbatch { } } + bool try_fuse_common(std::initializer_list tensors) const { + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + + int seen_bufs[HTP_OP_MAX_BUFS]; + int n_seen_bufs = 0; + + for (const auto * t : tensors) { + if (!t || t_map.count(t)) { + continue; + } + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + int fd = sbuf->fd(); + if (!b_map.count(fd)) { + bool found = false; + for (int i = 0; i < n_seen_bufs; i++) { + if (seen_bufs[i] == fd) { + found = true; + break; + } + } + if (!found) { + if (n_seen_bufs < HTP_OP_MAX_BUFS) { + seen_bufs[n_seen_bufs++] = fd; + } + if (!sbuf->extended) { + extra_vmem += sbuf->size(); + } + extra_bufs += 1; + } + } + } + + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + return true; + } + + bool try_fuse_common(const ggml_tensor * t1, const ggml_tensor * t2) const { + return try_fuse_common({t1, t2}); + } + bool try_fuse_allreduce_add(const htp_opnode & node) { if (n_ops == 0 || opt_ar_select != 2) return false; if (node.opcode != HTP_OP_ADD) return false; @@ -2466,20 +2599,7 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } - } - }; - fit_t(res_tensor); - fit_t(add_dst); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + if (!try_fuse_common(res_tensor, add_dst)) { return false; } @@ -2556,20 +2676,7 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } - } - }; - fit_t(weight); - fit_t(node.dst()); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + if (!try_fuse_common(weight, node.dst())) { return false; } @@ -2626,8 +2733,15 @@ struct ggml_hexagon_opbatch { const ggml_tensor * src0 = last_node.src0(); const ggml_tensor * src1 = last_node.src1(); + if (src2->type != GGML_TYPE_F32) return false; + + const struct htp_mm_kernel_params * orig_kparams = (const struct htp_mm_kernel_params *) last_node.kernel_params; struct htp_mm_kernel_params kparams; ggml_hexagon_precompute_fused_matmul_add_params(sess, src0, src1, src2, node.dst(), &kparams); + if (kparams.kernel_type == HTP_MM_KERNEL_UNSUPPORTED) { + return false; + } + const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1); if (!can_fuse) return false; @@ -2638,20 +2752,19 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } + if (kparams.n_hmx > 0 && orig_kparams->n_hmx > 0) { + if (kparams.m_chunk < orig_kparams->m_chunk || + kparams.n_chunk < orig_kparams->n_chunk || + kparams.n_act_threads < orig_kparams->n_act_threads) { + HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, th %d->%d)\n", + sess->c_name(), orig_kparams->m_chunk, kparams.m_chunk, + orig_kparams->n_chunk, kparams.n_chunk, + orig_kparams->n_act_threads, kparams.n_act_threads); + return false; } - }; - fit_t(src2); - fit_t(node.dst()); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + } + + if (!try_fuse_common(src2, node.dst())) { return false; } @@ -2723,20 +2836,7 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } - } - }; - fit_t(w_in); - fit_t(d_in); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + if (!try_fuse_common(w_in, d_in)) { return false; } @@ -2787,20 +2887,7 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } - } - }; - fit_t(w1); - fit_t(node.dst()); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + if (!try_fuse_common(w1, node.dst())) { return false; } @@ -2882,20 +2969,7 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } - } - }; - fit_t(w_in); - fit_t(d_in); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + if (!try_fuse_common(w_in, d_in)) { return false; } @@ -2948,20 +3022,7 @@ struct ggml_hexagon_opbatch { return false; } - size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; - auto fit_t = [&](const ggml_tensor * t) { - if (!t_map.count(t)) { - extra_tens++; - auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd())) { - extra_vmem += sbuf->size(); - extra_bufs += 1; - } - } - }; - fit_t(w1); - fit_t(node.dst()); - if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + if (!try_fuse_common(w1, node.dst())) { return false; } @@ -3005,6 +3066,70 @@ struct ggml_hexagon_opbatch { return false; } + bool try_fuse_gdn_cpy(const htp_opnode & node) { + if (n_ops == 0 || node.opcode != HTP_OP_CPY) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_GATED_DELTA_NET) return false; + if (last_node.outputs.size() != 1) return false; + + const ggml_tensor * gdn_out = last_node.dst(); + const ggml_tensor * cpy_node = node.node; + const ggml_tensor * cpy_src = node.src0(); + const ggml_tensor * cpy_dst = node.dst(); + + if (!cpy_src || !cpy_dst || !cpy_dst->data) return false; + if (gdn_out->type != GGML_TYPE_F32 || cpy_src->type != GGML_TYPE_F32 || cpy_dst->type != GGML_TYPE_F32) return false; + if ((gdn_out->flags & GGML_TENSOR_FLAG_OUTPUT) || (cpy_node->flags & GGML_TENSOR_FLAG_OUTPUT)) return false; + + const ggml_tensor * v = last_node.node->src[2]; + if (!v) return false; + + 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]; + const int64_t K = ggml_get_op_params_i32(last_node.node, 0); + const size_t tail_off = (size_t) S_v * H * n_tokens * n_seqs * sizeof(float); + + const int64_t D = S_v * S_v * H; + const int64_t n_written = std::min(n_tokens, K); + + if (cpy_src->op != GGML_OP_VIEW || (cpy_src->view_src != gdn_out && cpy_src->view_src->data != gdn_out->data) || + cpy_src->view_offs != tail_off || !ggml_is_contiguous(cpy_src)) { + return false; + } + + if (cpy_dst->ne[0] != D || cpy_dst->ne[1] != n_seqs || cpy_dst->nb[0] != sizeof(float)) { + return false; + } + if (n_seqs > 1 && cpy_dst->nb[1] != (size_t) D * sizeof(float)) { + return false; + } + if (n_written > 1) { + if (cpy_dst->ne[2] != n_written || cpy_dst->nb[2] != (size_t) D * n_seqs * sizeof(float)) { + return false; + } + } + + if (!try_fuse_common({cpy_dst})) { + return false; + } + + last_node.name += "+CPY"; + last_node.outputs.push_back(cpy_dst); + last_node.fused.push_back(node.node); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.dst[1] = add_tensor(cpy_dst); + for (uint32_t d = 2; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused GATED_DELTA_NET+CPY (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + bool try_fuse(const htp_opnode & node) { if (!opt_opfusion) return false; if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_ALLREDUCE_ADD) && try_fuse_allreduce_add(node)) return true; @@ -3012,6 +3137,7 @@ struct ggml_hexagon_opbatch { if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ADD) && try_fuse_mul_mat_add(node)) return true; if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_NX) && try_fuse_mul_mat_nx(node)) return true; if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ID_NX) && try_fuse_mul_mat_id_nx(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_GDN_CPY) && try_fuse_gdn_cpy(node)) return true; return false; } }; @@ -3130,8 +3256,8 @@ struct ggml_hexagon_opqueue { } htp_tensor *t = (htp_tensor*) t_ptr; for (unsigned int i=0; i < req.n_tensors; i++) { - GGML_LOG_DEBUG("ggml-hex: %s htp-tensor #%u : bi %u offset %u size %u : %zu:%zu:%zu:%zu\n", - shm_buf->sess->c_name(), i, t[i].bi, t[i].data, t[i].size, + GGML_LOG_DEBUG("ggml-hex: %s htp-tensor #%u : bi %u offset %llu size %u : %zu:%zu:%zu:%zu\n", + shm_buf->sess->c_name(), i, t[i].bi, (unsigned long long) t[i].data, t[i].size, (size_t) t[i].ne[0], (size_t) t[i].ne[1], (size_t) t[i].ne[2], (size_t) t[i].ne[3]); } } @@ -3269,6 +3395,12 @@ void ggml_hexagon_session::flush_sync(bool all) { flush_pending(all); } +void ggml_hexagon_session::start_batch() { + if (this->mdev.count > 1) { + enqueue_mdev_group(); + } +} + void ggml_hexagon_session::flush_batch(size_t min_ops) { if (op_batch->n_ops < min_ops) { return; } @@ -3318,14 +3450,13 @@ void ggml_hexagon_session::flush_batch(size_t min_ops) { void ggml_hexagon_session::enqueue_op(const htp_opnode & node) { auto clone_tensor_buffer = [this](const ggml_tensor * t) { - if (t && t->buffer && ggml_backend_buffer_is_hexagon(t->buffer)) { - auto sbuf = static_cast(t->buffer->context); - if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != this) { - this->clone_buffer(sbuf); - } - for (auto & sub : this->mdev.sessions) { - sub->clone_buffer(sbuf); - } + auto sbuf = this->mmap_tensor(t); + if (!sbuf) return; + if (sbuf->sess != this) { + this->clone_buffer(sbuf); + } + for (auto & sub : this->mdev.sessions) { + sub->clone_buffer(sbuf); } }; @@ -3344,8 +3475,13 @@ void ggml_hexagon_session::enqueue_op(const htp_opnode & node) { flush_async(); } - if (this->mdev.count > 1 && op_batch->n_ops == 0) { - enqueue_mdev_group(); + if (op_batch->empty()) { + start_batch(); + } + + if (!op_batch->fit_op(node)) { + GGML_ABORT("ggml-hex: %s op does not fit into empty batch (vmem/tensor/buffer limit exceeded)\n", + c_name()); } op_batch->add_op(node); @@ -3358,19 +3494,19 @@ void ggml_hexagon_session::enqueue_mdev_group() { static ggml_hexagon_tensor_extra fence_extra { {}, 0, GGML_HEXAGON_TENSOR_FENCE }; ggml_tensor dummy_t {}; - dummy_t.buffer = &this->fence_buf->backend_buffer; - dummy_t.extra = &fence_extra; - dummy_t.data = (void *) fence_slot; - dummy_t.type = GGML_TYPE_I8; - dummy_t.ne[0] = HTP_FENCE_SLOT_SIZE; - dummy_t.ne[1] = (int64_t) this->mdev.count; - dummy_t.ne[2] = 1; - dummy_t.ne[3] = 1; - dummy_t.nb[0] = 1; - dummy_t.nb[1] = HTP_FENCE_SLOT_SIZE; - dummy_t.nb[2] = dummy_t.nb[1] * dummy_t.ne[1]; - dummy_t.nb[3] = dummy_t.nb[2]; - dummy_t.op = GGML_OP_NONE; + dummy_t.buffer = &this->fence_buf->backend_buffer; + dummy_t.extra = &fence_extra; + dummy_t.data = (void *) fence_slot; + dummy_t.type = GGML_TYPE_I8; + dummy_t.ne[0] = HTP_FENCE_SLOT_SIZE; + dummy_t.ne[1] = (int64_t) this->mdev.count; + dummy_t.ne[2] = 1; + dummy_t.ne[3] = 1; + dummy_t.nb[0] = 1; + dummy_t.nb[1] = HTP_FENCE_SLOT_SIZE; + dummy_t.nb[2] = dummy_t.nb[1] * dummy_t.ne[1]; + dummy_t.nb[3] = dummy_t.nb[2]; + dummy_t.op = GGML_OP_NONE; dummy_t.op_params[0] = (int32_t) this->mdev.idx; ggml_tensor * node = group_node.add_dummy(dummy_t); @@ -3379,9 +3515,6 @@ void ggml_hexagon_session::enqueue_mdev_group() { group_node.outputs.clear(); group_node.name = "MDEV_GROUP"; - if (this->fence_buf->sess != this) { - this->clone_buffer(this->fence_buf); - } for (auto & sub : this->mdev.sessions) { sub->clone_buffer(this->fence_buf); } @@ -3588,7 +3721,19 @@ void ggml_hexagon_session::enqueue_allreduce( this->enqueue_op(ar_node); } -bool ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer *sbuf) +ggml_hexagon_shared_buffer * ggml_hexagon_session::mmap_tensor(const ggml_tensor * t) { + if (!t) return nullptr; + + auto sbuf = static_cast(t->buffer->context); + if (!sbuf->mapped) { + const bool is_weight = ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; + const bool extended = opt_dma64 && is_weight; + sbuf->mmap(extended); + } + return sbuf; +} + +bool ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer * sbuf) { GGML_ASSERT(sbuf && sbuf->mem); if (sbuf->sess == this) return true; @@ -3600,12 +3745,14 @@ bool ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer *sbuf) if (this->cloned_buffers.find(fd) != this->cloned_buffers.end()) return true; + GGML_ASSERT(sbuf->mapped); + HEX_VERBOSE("ggml-hex: %s clone-buffer: %s base %p size %zu fd %d\n", this->name.c_str(), sbuf->c_name(), sbuf->base(), sbuf->size(), fd); auto clone = std::make_unique(this, *sbuf); try { - clone->mmap(); + clone->mmap(sbuf->extended); } catch (const std::exception & exc) { GGML_LOG_ERROR("ggml-hex: %s lazy mapping of buffer context failed: %s\n", this->c_name(), exc.what()); return false; @@ -4081,7 +4228,7 @@ static bool ggml_hexagon_precompute_flash_attn_params( const uint32_t DK_pad = hex_round_up(DK, 64); const uint32_t DV_pad = hex_round_up(DV, 64); size_t Br = 0, Bc = 0; - int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0); + int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0, sinks != nullptr, n_head); if (ret == 0) { kparams->kernel_type = HTP_FA_KERNEL_HMX; kparams->Br = Br; @@ -4091,7 +4238,7 @@ static bool ggml_hexagon_precompute_flash_attn_params( kparams->u.hmx.g_br = hex_align_up(G * Br, 32); kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0; - kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK_pad, DV_pad, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0); + kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK_pad, DV_pad, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0, sinks != nullptr, n_head); const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256); kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes @@ -4122,7 +4269,7 @@ static bool ggml_hexagon_precompute_flash_attn_params( const size_t size_k_row_padded = hex_round_up(k->ne[0] * 2, 128); const size_t size_v_row_padded = hex_round_up(v->ne[0] * 2, 128); - kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, kparams->is_q_fp32 != 0, mask != nullptr, sess->n_threads); + kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, kparams->is_q_fp32 != 0, mask != nullptr, sinks != nullptr, n_head, sess->n_threads); kparams->u.hvx.size_q_row_padded = size_q_row_padded; kparams->u.hvx.size_k_row_padded = size_k_row_padded; @@ -4238,9 +4385,15 @@ static bool ggml_hexagon_supported_gated_delta_net(const struct ggml_hexagon_ses return false; } - return true; + const uint32_t total_rows = (uint32_t) (H * n_seqs); + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, total_rows); + struct htp_gdn_vtcm_layout layout; + htp_gdn_vtcm_layout_build(&layout, (uint32_t) S_v, n_threads ? n_threads : 1); + if (layout.total_bytes > sess->vtcm_size) { + return false; + } - GGML_UNUSED(sess); + return true; } static bool ggml_hexagon_matmul_is_hmx_eligible( @@ -4311,6 +4464,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( int ne11_padded, bool is_matmul_id, bool is_batched, + size_t src2_size, size_t vtcm_budget, struct htp_mm_kernel_params * kparams ) { @@ -4331,7 +4485,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) { // Try grouped path first const bool use_dma_activation = (src1->nb[1]/sizeof(float) > (size_t)ne00_padded); - if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, use_dma_activation, n_threads, pipeline, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { + if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, use_dma_activation, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { use_grouped = true; } } @@ -4339,7 +4493,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( if (!use_grouped) { // Fallback to simple 2D path (group_size = 1) const int m_id_rows = (dst && is_matmul_id) ? (int) ((size_t) dst->ne[1] * dst->ne[2]) : 0; - if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { + if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { return false; } } @@ -4358,6 +4512,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( kparams->div_n_act_threads = init_fastdiv_values(act_threads_selected); kparams->div_ne00_padded = init_fastdiv_values(ne00_padded); kparams->vtcm_src1_size = 0; + kparams->vtcm_src2_size = (int32_t) src2_size; kparams->vtcm_dst_size = 0; if (is_batched && !is_matmul_id) { @@ -4387,6 +4542,11 @@ static void ggml_hexagon_precompute_hvx_mm_params( size_t vtcm_budget, struct htp_mm_kernel_params * kparams ) { + if (opt_mm_select < 1) { + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; + } + kparams->n_hmx = 0; kparams->n_threads = sess->n_threads; @@ -4410,29 +4570,30 @@ static void ggml_hexagon_precompute_hvx_mm_params( for (uint32_t d = max_prefetch; d >= 2; d /= 2) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0->nb[1], 0, src2_row_size, d, true, false + 0, src0->nb[1], kparams->src1_row_size, 0, d, true, false ); if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; break; } } - if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { - htp_mm_hvx_vtcm_layout_build( - &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0->nb[1], 0, src2_row_size, 2, true, false - ); + if (L.total_bytes > vtcm_budget) { + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } - kparams->n_prefetch = best_n_prefetch; + kparams->n_prefetch = best_n_prefetch; kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; kparams->vtcm_dst_size = L.dst_bytes; + goto done_quant; } else { - bool try_tiled = (k_align && opt_mm_select >= 2); + bool try_tiled = (k_align && opt_mm_select >= 1); if (try_tiled) { - kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1 || wtype == GGML_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - if (src1_nrows < (int)sess->n_threads) { + kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1 || wtype == GGML_TYPE_Q4_K) + ? htp_mm_q8_1_tiled_row_size(ne10) + : htp_mm_q8_0_tiled_row_size(ne10); + if (src1_nrows < (int) sess->n_threads) { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_BLOCK; } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; @@ -4451,113 +4612,72 @@ static void ggml_hexagon_precompute_hvx_mm_params( break; } } - if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { - htp_mm_hvx_vtcm_layout_build( - &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false - ); - } - kparams->n_prefetch = best_n_prefetch; - - if (L.total_bytes <= vtcm_budget) { - kparams->vtcm_size = L.total_bytes; + uint32_t m_chunk = 0; + if (htp_mm_hvx_solve_vtcm_params( + kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, best_n_prefetch, vtcm_budget, + &L, &m_chunk)) { + kparams->n_prefetch = best_n_prefetch; + kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0; + kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_src2_size = L.src2_bytes; + kparams->vtcm_dst_size = L.dst_bytes; goto done_quant; } - HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), L.total_bytes, vtcm_budget); } - // Flat HVX fallback - { - kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1 || wtype == GGML_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; - - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false - ); - - kparams->n_prefetch = 16; - kparams->vtcm_size = L.total_bytes; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - } + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } done_quant:; } else if (wtype == GGML_TYPE_F16) { // F16 HVX - const bool is_batched = (ne02 > 1) || (ne03 > 1); - const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1); - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false - ); - - if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { + uint32_t m_chunk = 0; + if (htp_mm_hvx_solve_vtcm_params( + HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, vtcm_budget, + &L, &m_chunk)) { kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_VTCM; + kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0; kparams->src1_row_size = hex_round_up(ne10 * 2, 128); kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_src2_size = L.src2_bytes; kparams->vtcm_dst_size = L.dst_bytes; kparams->n_prefetch = 16; - } else { - if (src1->type == GGML_TYPE_F32) { - kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F32_DDR; - } else { - kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_DDR; - } - kparams->src1_row_size = src1->nb[1]; - htp_mm_hvx_vtcm_layout_build( - &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false - ); - kparams->vtcm_size = L.total_bytes; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->n_prefetch = 16; + return; } + + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } else { // F32 HVX - const bool is_batched = (ne02 > 1) || (ne03 > 1); - const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1); - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false - ); - - if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { + uint32_t m_chunk = 0; + if (htp_mm_hvx_solve_vtcm_params( + HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, vtcm_budget, + &L, &m_chunk)) { kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_VTCM; + kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0; kparams->src1_row_size = hex_round_up(ne10 * 4, 128); kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_src2_size = L.src2_bytes; kparams->vtcm_dst_size = L.dst_bytes; kparams->n_prefetch = 16; - } else { - kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_DDR; - kparams->src1_row_size = src1->nb[1]; - htp_mm_hvx_vtcm_layout_build( - &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false - ); - kparams->vtcm_size = L.total_bytes; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->n_prefetch = 16; + return; } + + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } } @@ -4567,6 +4687,7 @@ static void ggml_hexagon_precompute_matmul_params_impl( const struct ggml_tensor * src1, const struct ggml_tensor * dst, const size_t src2_row_size, + const size_t src2_size, struct htp_mm_kernel_params * kparams ) { memset(kparams, 0, sizeof(*kparams)); @@ -4593,9 +4714,9 @@ static void ggml_hexagon_precompute_matmul_params_impl( const size_t vtcm_budget = sess->vtcm_size; // Check HMX eligibility and try precomputing HMX parameters - bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 3); + bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2); if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, dst, ne01_padded, is_matmul_id, is_batched)) { - if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, vtcm_budget, kparams)) { + if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, src2_size, vtcm_budget, kparams)) { goto finalize; } } @@ -4608,7 +4729,7 @@ finalize: kparams->div_ne1 = init_fastdiv_values(ne11); kparams->div_r2 = init_fastdiv_values(ne02 > 0 ? ne12 / ne02 : 1); kparams->div_r3 = init_fastdiv_values(ne03 > 0 ? ne13 / ne03 : 1); - kparams->div_ne11 = init_fastdiv_values(ne11); + kparams->div_ne12 = init_fastdiv_values(ne12); } static void ggml_hexagon_precompute_matmul_params( @@ -4618,7 +4739,7 @@ static void ggml_hexagon_precompute_matmul_params( const struct ggml_tensor * dst, struct htp_mm_kernel_params * kparams ) { - ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams); + ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, 0, kparams); } static void ggml_hexagon_precompute_fused_matmul_add_params( @@ -4629,7 +4750,83 @@ static void ggml_hexagon_precompute_fused_matmul_add_params( const struct ggml_tensor * dst, struct htp_mm_kernel_params * kparams ) { - ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2->nb[1], kparams); + const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0; + ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, kparams); +} + +static bool ggml_hexagon_precompute_binary_params( + const struct ggml_hexagon_session * sess, + uint32_t op, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_binary_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + + const size_t elem_size = ggml_type_size(src0->type); + const size_t src0_row_size = src0->ne[0] * elem_size; + const size_t src1_row_size = src1->ne[0] * elem_size; + const size_t dst_row_size = dst->ne[0] * elem_size; + + const size_t src0_row_size_aligned = hex_round_up(src0_row_size, 128); + const size_t src1_row_size_aligned = hex_round_up(src1_row_size, 128); + const size_t dst_row_size_aligned = hex_round_up(dst_row_size, 128); + + const bool is_add_id = op == HTP_OP_ADD_ID; + const bool is_scalar = !is_add_id && src1->ne[0] == 1; + const bool is_transposed = src0->nb[1] < src0_row_size || src1->nb[1] < src1_row_size || dst->nb[1] < dst_row_size; + const bool is_same_shape = !is_add_id && !is_scalar && !is_transposed && + src1->ne[0] == src0->ne[0] && + (src1->ne[1] == src0->ne[1] || src1->ne[1] == 1) && + (src1->ne[2] == src0->ne[2] || src1->ne[2] == 1) && + (src1->ne[3] == src0->ne[3] || src1->ne[3] == 1); + const bool is_row_bcast = is_same_shape && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1; + const bool is_complex = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] == src0->ne[0]); + + enum htp_binary_kernel_type kernel_type; + size_t src1_size = 0; + + if (is_add_id) { + kernel_type = HTP_BINARY_KERNEL_ADD_ID; + src1_size = hex_round_up(src1->ne[1] * src1_row_size_aligned, 128); + } else if (is_row_bcast) { + kernel_type = HTP_BINARY_KERNEL_ROW_BCAST; + src1_size = src1_row_size_aligned; + } else if (is_scalar) { + const bool is_scalar_static = (src1->ne[2] == 1 && src1->ne[3] == 1) && + (src1->ne[1] == 1 || src1->nb[1] == elem_size); + if (is_scalar_static) { + kernel_type = HTP_BINARY_KERNEL_SCALAR_DMA; + src1_size = hex_round_up(src1->ne[1] * elem_size, 128); + } else { + kernel_type = HTP_BINARY_KERNEL_SCALAR; + } + } else if (is_same_shape) { + kernel_type = HTP_BINARY_KERNEL_SAME_SHAPE; + } else if (is_complex) { + kernel_type = HTP_BINARY_KERNEL_COMPLEX; + } else { + kernel_type = HTP_BINARY_KERNEL_REPEAT; + } + + kparams->kernel_type = kernel_type; + kparams->n_threads = sess->n_threads; + kparams->src0_row_size_aligned = src0_row_size_aligned; + kparams->src1_row_size_aligned = src1_row_size_aligned; + kparams->dst_row_size_aligned = dst_row_size_aligned; + kparams->src1_size = src1_size; + + struct htp_binary_vtcm_layout L; + htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size); + if (L.rows_per_buffer == 0 || L.total_bytes > sess->vtcm_size) { + return false; + } + + kparams->rows_per_buffer = L.rows_per_buffer; + kparams->vtcm_size = L.total_bytes; + + return true; } static void ggml_hexagon_precompute_unary_params( @@ -4790,6 +4987,78 @@ static void ggml_hexagon_precompute_set_rows_params( kparams->vtcm_size = vtcm_layout.total_bytes; } +static void ggml_hexagon_precompute_softmax_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * op, + struct htp_softmax_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + + const struct ggml_tensor * src0 = op->src[0]; + const struct ggml_tensor * src1 = op->src[1]; + + const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, src0_nrows); + + float scale = 1.0f; + float max_bias = 0.0f; + memcpy(&scale, &op->op_params[0], sizeof(float)); + memcpy(&max_bias, &op->op_params[1], sizeof(float)); + + kparams->scale = scale; + kparams->max_bias = max_bias; + + const uint32_t n_head = src0->ne[2]; + const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); + kparams->n_head = n_head; + kparams->n_head_log2 = n_head_log2; + + if (max_bias > 0.0f && n_head_log2 > 0) { + kparams->m0 = powf(2.0f, -(max_bias) / n_head_log2); + kparams->m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + } else { + kparams->m0 = 1.0f; + kparams->m1 = 1.0f; + } + + kparams->use_src1 = (src1 != nullptr) ? 1 : 0; + kparams->use_f16 = (src1 != nullptr && src1->type == GGML_TYPE_F16) ? 1 : 0; + + const uint32_t ne00 = src0->ne[0]; + const uint32_t ne10 = src1 ? src1->ne[0] : 1; + + struct htp_softmax_vtcm_layout layout; + htp_softmax_vtcm_layout_build(&layout, ne00, ne10, kparams->use_src1 != 0, kparams->use_f16 != 0, n_threads); + + kparams->n_threads = n_threads; + kparams->src0_nrows = src0_nrows; + kparams->src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads; + kparams->vtcm_size = (uint32_t) layout.total_bytes; + kparams->vtcm_src0_size_per_thread = (uint32_t) layout.src0_bytes_per_thread; + kparams->vtcm_src1_size_per_thread = (uint32_t) layout.src1_bytes_per_thread; + kparams->vtcm_dst_size_per_thread = (uint32_t) layout.dst_bytes_per_thread; + kparams->src0_row_size_aligned = (uint32_t) layout.src0_spad_half_size; + kparams->src1_row_size_aligned = (uint32_t) layout.src1_spad_half_size; + kparams->dst_row_size_aligned = (uint32_t) layout.dst_spad_half_size; + kparams->src0_spad_half_size = (uint32_t) layout.src0_spad_half_size; + kparams->src1_spad_half_size = (uint32_t) layout.src1_spad_half_size; + kparams->dst_spad_half_size = (uint32_t) layout.dst_spad_half_size; + if (!kparams->use_src1) { + kparams->kernel_id = HTP_SOFTMAX_KERNEL_NOMASK; + } else if (kparams->use_f16) { + kparams->kernel_id = HTP_SOFTMAX_KERNEL_MASK_F16; + } else { + kparams->kernel_id = HTP_SOFTMAX_KERNEL_MASK_F32; + } + + if (src0->ne[1] > 0) kparams->div_ne01 = init_fastdiv_values(src0->ne[1]); + if (src0->ne[2] > 0) kparams->div_ne02 = init_fastdiv_values(src0->ne[2]); + const uint32_t ne12 = src1 ? src1->ne[2] : 1; + const uint32_t ne13 = src1 ? src1->ne[3] : 1; + if (ne12 > 0) kparams->div_ne12 = init_fastdiv_values(ne12); + if (ne13 > 0) kparams->div_ne13 = init_fastdiv_values(ne13); +} + static void ggml_hexagon_precompute_rope_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * op, @@ -4798,13 +5067,15 @@ static void ggml_hexagon_precompute_rope_params( memset(kparams, 0, sizeof(*kparams)); const struct ggml_tensor * src0 = op->src[0]; + const struct ggml_tensor * src2 = op->src[2]; const struct ggml_tensor * dst = op; const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, src0_nrows); + const uint32_t n_freq_factors = src2 ? (uint32_t) src2->ne[0] : 0; struct htp_rope_vtcm_layout layout; - htp_rope_vtcm_layout_build(&layout, src0->ne[0], n_threads); + htp_rope_vtcm_layout_build(&layout, src0->ne[0], n_threads, n_freq_factors); kparams->n_threads = n_threads; kparams->src0_nrows = src0_nrows; @@ -4813,6 +5084,8 @@ static void ggml_hexagon_precompute_rope_params( kparams->spad_per_thread = (uint32_t) layout.bytes_per_thread; kparams->theta_cache_offset = (uint32_t) layout.theta_cache_size_aligned; kparams->src0_row_size_aligned = (uint32_t) layout.src0_row_size_aligned; + kparams->freq_factors_offset = (uint32_t) (layout.bytes_per_thread * n_threads); + kparams->freq_factors_size = (uint32_t) layout.freq_factors_size_aligned; if (src0_nrows > 0) { kparams->div_ne2_ne1 = init_fastdiv_values(dst->ne[2] * dst->ne[1]); @@ -4820,6 +5093,146 @@ static void ggml_hexagon_precompute_rope_params( } } +static void ggml_hexagon_precompute_ssm_conv_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_ssm_conv_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + + const uint32_t d_conv = (uint32_t) src1->ne[0]; + const uint32_t d_inner = (uint32_t) src0->ne[1]; + const uint32_t n_t = (uint32_t) dst->ne[1]; + const uint32_t n_s = (uint32_t) dst->ne[2]; + const uint32_t ncs = (uint32_t) src0->ne[0]; + + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, (d_inner + 31) / 32); + + kparams->n_threads = n_threads; + kparams->d_conv = d_conv; + kparams->d_inner = d_inner; + kparams->n_t = n_t; + kparams->n_s = n_s; + + const uint32_t raw_rpt = (d_inner + n_threads - 1) / n_threads; + const uint32_t d_inner_per_thread = hex_round_up(raw_rpt, 32); + kparams->d_inner_per_thread = d_inner_per_thread; + + kparams->src0_row_size_aligned = hex_round_up(ncs * sizeof(float), 128); + kparams->src1_row_size_aligned = hex_round_up(d_conv * sizeof(float), 128); + kparams->dst_row_size_aligned = hex_round_up(d_inner * sizeof(float), 128); + + if (n_t == 1) { + kparams->d_inner_tile = d_inner_per_thread; + + const uint32_t src1_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128; + const uint32_t src1_T_bytes = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128); + const uint32_t vtcm_src1_per_thread = src1_raw_bytes + src1_T_bytes; + + const uint32_t src0_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128; + const uint32_t src0_T_bytes = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128); + const uint32_t vtcm_src0_per_thread = src0_raw_bytes + src0_T_bytes; + + const uint32_t vtcm_dst_per_thread = hex_round_up(d_inner_per_thread * sizeof(float), 128); + + kparams->vtcm_src0_size_per_thread = vtcm_src0_per_thread; + kparams->vtcm_src1_size_per_thread = vtcm_src1_per_thread; + kparams->vtcm_dst_size_per_thread = vtcm_dst_per_thread; + + kparams->vtcm_src0_size = vtcm_src0_per_thread * n_threads; + kparams->vtcm_src1_size = vtcm_src1_per_thread * n_threads; + kparams->vtcm_dst_size = vtcm_dst_per_thread * n_threads; + kparams->vtcm_size = kparams->vtcm_src0_size + kparams->vtcm_src1_size + kparams->vtcm_dst_size; + } else { + const uint32_t src1_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128; + const uint32_t src1_T_bytes = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128); + const uint32_t vtcm_src1_per_thread = src1_raw_bytes + src1_T_bytes; + + const size_t vtcm_budget = (sess->vtcm_size > 0 ? sess->vtcm_size / n_threads : (1024 * 1024)); + const size_t avail_for_src0 = vtcm_budget > vtcm_src1_per_thread ? vtcm_budget - vtcm_src1_per_thread : (128 * 1024); + + uint32_t d_inner_tile = (uint32_t)((avail_for_src0 / 2) / (ncs * sizeof(float) + n_t * sizeof(float) + 1)); + d_inner_tile = (d_inner_tile / 32) * 32; + if (d_inner_tile == 0) { + d_inner_tile = 32; + } + if (d_inner_tile > d_inner_per_thread) { + d_inner_tile = d_inner_per_thread; + } + kparams->d_inner_tile = d_inner_tile; + + const uint32_t src0_tile_raw = hex_round_up(d_inner_tile * ncs * sizeof(float), 128) + 128; + const uint32_t src0_tile_T = hex_round_up(ncs * d_inner_tile * sizeof(float), 128); + const uint32_t vtcm_src0_per_thread = src0_tile_raw + src0_tile_T; + + const uint32_t vtcm_dst_per_thread = hex_round_up(d_inner_tile * n_t * sizeof(float), 128); + + kparams->vtcm_src0_size_per_thread = vtcm_src0_per_thread; + kparams->vtcm_src1_size_per_thread = vtcm_src1_per_thread; + kparams->vtcm_dst_size_per_thread = vtcm_dst_per_thread; + + kparams->vtcm_src0_size = vtcm_src0_per_thread * n_threads; + kparams->vtcm_src1_size = vtcm_src1_per_thread * n_threads; + kparams->vtcm_dst_size = vtcm_dst_per_thread * n_threads; + kparams->vtcm_size = kparams->vtcm_src0_size + kparams->vtcm_src1_size + kparams->vtcm_dst_size; + } + + kparams->div_n_threads = init_fastdiv_values(n_threads); +} + +static void ggml_hexagon_precompute_gated_delta_net_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * op, + struct htp_gdn_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + + const struct ggml_tensor * q = op->src[0]; + const struct ggml_tensor * k = op->src[1]; + const struct ggml_tensor * v = op->src[2]; + const struct ggml_tensor * g = op->src[3]; + const struct ggml_tensor * state = op->src[5]; + + const uint32_t S_v = (uint32_t) v->ne[0]; + const uint32_t H = (uint32_t) v->ne[1]; + const uint32_t n_tokens = (uint32_t) v->ne[2]; + const uint32_t n_seqs = (uint32_t) v->ne[3]; + const uint32_t K = (uint32_t) ggml_get_op_params_i32(op, 0); + + const uint32_t rq3 = (uint32_t) (n_seqs / q->ne[3]); + const uint32_t rk3 = (uint32_t) (n_seqs / k->ne[3]); + const uint32_t total_rows = H * n_seqs; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, total_rows); + + struct htp_gdn_vtcm_layout layout; + htp_gdn_vtcm_layout_build(&layout, S_v, n_threads ? n_threads : 1); + + kparams->n_threads = n_threads ? n_threads : 1; + kparams->S_v = S_v; + kparams->H = H; + kparams->n_tokens = n_tokens; + kparams->n_seqs = n_seqs; + kparams->K = K; + kparams->total_rows = total_rows; + kparams->rows_per_thread = (total_rows + kparams->n_threads - 1) / kparams->n_threads; + kparams->kda = (g->ne[0] == S_v) ? 1 : 0; + kparams->state_aligned = (uint32_t) layout.state_aligned; + kparams->vtcm_per_thread = (uint32_t) layout.bytes_per_thread; + kparams->vtcm_size = (uint32_t) layout.total_bytes; + kparams->state_seq_stride = (uint32_t) (state->nb[3] / sizeof(float)); + kparams->state_size_per_snap = S_v * S_v * H * n_seqs; + kparams->scale = 1.0f / sqrtf((float) S_v); + + if (H > 0) kparams->div_H = init_fastdiv_values(H); + if (q->ne[1] > 0) kparams->div_q1 = init_fastdiv_values((uint32_t) q->ne[1]); + if (k->ne[1] > 0) kparams->div_k1 = init_fastdiv_values((uint32_t) k->ne[1]); + if (rq3 > 0) kparams->div_rq3 = init_fastdiv_values(rq3); + if (rk3 > 0) kparams->div_rk3 = init_fastdiv_values(rk3); + if (kparams->n_threads > 0) kparams->div_n_threads = init_fastdiv_values(kparams->n_threads); +} + static void ggml_hexagon_precompute_fused_mmnx_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, // W0 @@ -4849,9 +5262,9 @@ static void ggml_hexagon_precompute_fused_mmnx_params( const size_t vtcm_budget = sess->vtcm_size; const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1); - bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 3); + bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2); if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, nullptr, ne01_padded, false, is_batched)) { - if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, nullptr, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, false, is_batched, vtcm_budget, kparams)) { + if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, nullptr, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, false, is_batched, 0, vtcm_budget, kparams)) { kparams->n_weights = n_weights; goto finalize; } @@ -4886,7 +5299,7 @@ static void ggml_hexagon_precompute_fused_mmnx_params( } struct htp_mm_hvx_vtcm_layout L; - bool try_tiled = (opt_mm_select >= 2); + bool try_tiled = (opt_mm_select >= 1); // Test tiled first htp_mm_hvx_vtcm_layout_build( @@ -4903,19 +5316,8 @@ static void ggml_hexagon_precompute_fused_mmnx_params( kparams->n_prefetch = best_n_prefetch; kparams->n_weights = n_weights; } else { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; - size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1 || wtype == GGML_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true - ); - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - kparams->n_weights = n_weights; + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } } @@ -4924,7 +5326,7 @@ finalize: kparams->div_ne1 = init_fastdiv_values(ne11); kparams->div_r2 = init_fastdiv_values(ne02 > 0 ? ne12 / ne02 : 1); kparams->div_r3 = init_fastdiv_values(ne03 > 0 ? ne13 / ne03 : 1); - kparams->div_ne11 = init_fastdiv_values(ne11); + kparams->div_ne12 = init_fastdiv_values(ne12); } static void ggml_hexagon_precompute_fused_mmidnx_params( @@ -4935,7 +5337,7 @@ static void ggml_hexagon_precompute_fused_mmidnx_params( int32_t n_weights, struct htp_mm_kernel_params * kparams ) { - ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams); + ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, 0, kparams); kparams->n_weights = n_weights; } @@ -4969,12 +5371,19 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s case GGML_TYPE_MXFP4: case GGML_TYPE_Q4_K: case GGML_TYPE_Q6_K: + if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { + return false; + } + if (src0->ne[0] % ((src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q4_K) ? QK_K : 32)) { return false; } - if (src1->ne[2] != 1 || src1->ne[3] != 1) { - return false; // no broadcasting (for now) + if (src1->ne[2] < src0->ne[2] || src1->ne[3] < src0->ne[3]) { + return false; + } + if (src1->ne[2] % src0->ne[2] != 0 || src1->ne[3] % src0->ne[3] != 0) { + return false; } if (!src0->buffer) { @@ -4989,6 +5398,9 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s if (src1->ne[2] < src0->ne[2] || src1->ne[3] < src0->ne[3]) { return false; } + if (src1->ne[2] % src0->ne[2] != 0 || src1->ne[3] % src0->ne[3] != 0) { + return false; + } break; case GGML_TYPE_F32: @@ -5001,6 +5413,9 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s if (src1->ne[2] < src0->ne[2] || src1->ne[3] < src0->ne[3]) { return false; } + if (src1->ne[2] % src0->ne[2] != 0 || src1->ne[3] % src0->ne[3] != 0) { + return false; + } break; default: @@ -5009,7 +5424,7 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s struct htp_mm_kernel_params kparams; ggml_hexagon_precompute_matmul_params(sess, src0, src1, dst, &kparams); - if ((size_t)kparams.vtcm_size > sess->vtcm_size) { + if (kparams.kernel_type == HTP_MM_KERNEL_UNSUPPORTED || (size_t) kparams.vtcm_size > sess->vtcm_size) { HEX_VERBOSE("ggml-hex: %s supported MUL_MAT VTCM size needed (%d) > budget (%zu)\n", sess->c_name(), kparams.vtcm_size, sess->vtcm_size); return false; } @@ -5035,6 +5450,10 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session case GGML_TYPE_MXFP4: case GGML_TYPE_Q4_K: case GGML_TYPE_Q6_K: + if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { + return false; + } + if (src0->ne[0] % ((src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q4_K) ? QK_K : 32)) { return false; } @@ -5050,7 +5469,7 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session struct htp_mm_kernel_params kparams; ggml_hexagon_precompute_matmul_params(sess, src0, src1, dst, &kparams); - if ((size_t)kparams.vtcm_size > sess->vtcm_size) { + if (kparams.kernel_type == HTP_MM_KERNEL_UNSUPPORTED || (size_t) kparams.vtcm_size > sess->vtcm_size) { HEX_VERBOSE("ggml-hex: %s supported MUL_MAT_ID VTCM size needed (%d) > budget (%zu)\n", sess->c_name(), kparams.vtcm_size, sess->vtcm_size); return false; } @@ -5093,37 +5512,42 @@ static bool ggml_hexagon_supported_binary(const struct ggml_hexagon_session * se return false; } - return true; - - GGML_UNUSED(sess); + struct htp_binary_kernel_params kparams; + return ggml_hexagon_precompute_binary_params(sess, op_remap_to_htp(op), src0, src1, dst, &kparams); } static bool ggml_hexagon_supported_add_id(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const struct ggml_tensor * src0 = op->src[0]; const struct ggml_tensor * src1 = op->src[1]; + const struct ggml_tensor * src2 = op->src[2]; const struct ggml_tensor * dst = op; - if (src0->type != GGML_TYPE_F32) { + if (!src2) { return false; } - if (src1->type != GGML_TYPE_F32) { - return false; - } - if (dst->type != GGML_TYPE_F32) { + if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || src2->type != GGML_TYPE_I32) { return false; } if (!ggml_are_same_shape(src0, dst)) { return false; } + if (src1->ne[0] != src0->ne[0] || src1->ne[2] != 1 || src1->ne[3] != 1) { + return false; + } + if (src2->ne[0] != src0->ne[1] || src2->ne[1] != src0->ne[2]) { + return false; + } + if (src0->nb[0] != sizeof(float) || src1->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float) || src2->nb[0] != sizeof(int32_t)) { + return false; + } - // REVISIT: add support for non-contigiuos tensors + // REVISIT: add support for non-contiguous tensors if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { return false; } - return true; - - GGML_UNUSED(sess); + struct htp_binary_kernel_params kparams; + return ggml_hexagon_precompute_binary_params(sess, HTP_OP_ADD_ID, src0, src1, dst, &kparams); } static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { @@ -5251,6 +5675,10 @@ static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * s return false; } + if (src0->ne[2] > 512) { + return false; + } + if (src1) { if (src1->type != GGML_TYPE_F32 && src1->type != GGML_TYPE_F16) { return false; @@ -5296,6 +5724,14 @@ static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * s return false; } + const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, src0_nrows); + struct htp_softmax_vtcm_layout layout; + htp_softmax_vtcm_layout_build(&layout, src0->ne[0], src1 ? src1->ne[0] : 1, src1 != nullptr, src1 && src1->type == GGML_TYPE_F16, n_threads); + if (layout.total_bytes > sess->vtcm_size) { + return false; + } + return true; GGML_UNUSED(sess); @@ -5505,9 +5941,10 @@ static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess } const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, src0_nrows); + const uint32_t n_freq_factors = src2 ? (uint32_t) src2->ne[0] : 0; struct htp_rope_vtcm_layout layout; - htp_rope_vtcm_layout_build(&layout, src0->ne[0], n_threads); + htp_rope_vtcm_layout_build(&layout, src0->ne[0], n_threads, n_freq_factors); if (layout.total_bytes > sess->vtcm_size) { return false; } @@ -5530,11 +5967,14 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * return false; // src0 should be effectively 3D } - const int d_conv = src1->ne[0]; + const int d_conv = src1->ne[0]; const int d_inner = src0->ne[1]; const int n_t = dst->ne[1]; const int n_s = dst->ne[2]; + if (d_conv == 0 || d_conv > 32 || d_inner == 0) { + return false; + } if (src0->ne[0] != d_conv - 1 + n_t || src0->ne[1] != d_inner || src0->ne[2] != n_s) { return false; } @@ -5551,9 +5991,13 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * return false; } - return true; + struct htp_ssm_conv_kernel_params kparams; + ggml_hexagon_precompute_ssm_conv_params(sess, src0, src1, dst, &kparams); + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + return false; + } - GGML_UNUSED(sess); + return true; } static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { @@ -5796,7 +6240,7 @@ static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struc return false; // Q6_K has no fused HVX kernel } - return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; + return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW; } static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) { @@ -5926,6 +6370,11 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg node.node->src[0], node.node->src[1], node.node, (struct htp_mm_kernel_params *)node.kernel_params ); + } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID || node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) { + const ggml_tensor * src1 = node.node->src[1]; + GGML_ASSERT(ggml_hexagon_precompute_binary_params(sess, + node.opcode, node.node->src[0], src1, node.node, + (struct htp_binary_kernel_params *) node.kernel_params)); } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) { ggml_hexagon_precompute_flash_attn_params(sess, node.node, @@ -5954,6 +6403,21 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg node.node, (struct htp_rope_kernel_params *)node.kernel_params ); + } else if (node.opcode == HTP_OP_SSM_CONV) { + ggml_hexagon_precompute_ssm_conv_params(sess, + node.node->src[0], node.node->src[1], node.dst(), + (struct htp_ssm_conv_kernel_params *)node.kernel_params + ); + } else if (node.opcode == HTP_OP_SOFTMAX) { + ggml_hexagon_precompute_softmax_params(sess, + node.node, + (struct htp_softmax_kernel_params *)node.kernel_params + ); + } else if (node.opcode == HTP_OP_GATED_DELTA_NET) { + ggml_hexagon_precompute_gated_delta_net_params(sess, + node.node, + (struct htp_gdn_kernel_params *)node.kernel_params + ); } computed_nodes.push_back(std::move(node)); } @@ -6230,7 +6694,9 @@ static uint64_t ggml_hexagon_session_key(const ggml_hexagon_session * sess) { static bool ggml_hexagon_cpy_tensor_async_phys(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { auto sess_src = static_cast(backend_src->context); auto sess_dst = static_cast(backend_dst->context); - auto sbuf_dst = (ggml_hexagon_shared_buffer *) dst->buffer->context; + + sess_src->mmap_tensor(src); + auto sbuf_dst = sess_dst->mmap_tensor(dst); if (!sess_src->clone_buffer(sbuf_dst)) { return false; } @@ -6277,7 +6743,9 @@ static bool ggml_hexagon_cpy_tensor_async_phys(ggml_backend_t backend_src, ggml_ static bool ggml_hexagon_cpy_tensor_async_virt(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { auto sess_src = static_cast(backend_src->context); auto sess_dst = static_cast(backend_dst->context); - auto sbuf_src = (ggml_hexagon_shared_buffer *) src->buffer->context; + + auto sbuf_src = sess_src->mmap_tensor(src); + sess_dst->mmap_tensor(dst); if (!sess_dst->clone_buffer(sbuf_src)) { return false; } @@ -6905,7 +7373,7 @@ static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { GGML_LOG_INFO("ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev %zu\n", opt_ndev); - GGML_LOG_INFO("ggml-hex: Hexagon Arch version v%d\n", opt_arch); + GGML_LOG_INFO("ggml-hex: Hexagon Arch version v%d, DMA64 %s\n", opt_arch, opt_dma64 ? "enabled" : "disabled"); // Create devices for (size_t i = 0; i < opt_ndev; i++) { @@ -7270,6 +7738,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { const char * str_mbuf = getenv("GGML_HEXAGON_MBUF"); const char * str_optrace = getenv("GGML_HEXAGON_OPTRACE"); const char * str_hostbuf = getenv("GGML_HEXAGON_HOSTBUF"); + const char * str_dma64 = getenv("GGML_HEXAGON_DMA64"); // Init Arch first since it affects other defaults if (!str_arch) { @@ -7297,6 +7766,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { // Update vmem default opt_vmem = opt_arch >= 75 ? HTP_OP_MAX_VMEM_DEFAULT : 3000 * MiB; + opt_dma64 = opt_arch > 79 && (!str_dma64 || atoi(str_dma64) != 0); auto RE_ICASE = std::regex_constants::icase; diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index ef7b5184f..0716a8d21 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -14,7 +14,11 @@ #include "htp/matmul-ops.h" #include "htp/flash-attn-ops.h" #include "htp/unary-ops.h" +#include "htp/binary-ops.h" #include "htp/allreduce-ops.h" +#include "htp/ssm-conv.h" +#include "htp/gated-delta-net-ops.h" +#include "htp/softmax-ops.h" struct htp_opnode { ggml_tensor * node { nullptr }; @@ -325,10 +329,6 @@ struct htp_opformat { } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { path = "hvx-tiled"; - } else if (type == HTP_MM_KERNEL_HVX_F16_F16_DDR || type == HTP_MM_KERNEL_HVX_F16_F32_DDR || - type == HTP_MM_KERNEL_HVX_F32_F32_DDR || type == HTP_MM_KERNEL_HVX_F32_F16_DDR || - type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - path = "hvx-flat"; } snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size); } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) { @@ -350,6 +350,21 @@ struct htp_opformat { snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]); } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) { snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]); + } else if (node.opcode == HTP_OP_SSM_CONV) { + const auto * kparams = (const struct htp_ssm_conv_kernel_params *) node.kernel_params; + snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size); + } else if (node.opcode == HTP_OP_SOFTMAX) { + const auto * kparams = (const struct htp_softmax_kernel_params *) node.kernel_params; + snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_id, (int) kparams->n_threads, (int) kparams->vtcm_size); + } else if (node.opcode == HTP_OP_GATED_DELTA_NET) { + const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params; + snprintf(str, max_size, "%s vtcm %u", + kparams->kda ? "kda" : "scalar", + (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads)); + } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID || + node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) { + const auto * kparams = (const struct htp_binary_kernel_params *) node.kernel_params; + snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size); } else { snprintf(str, max_size, "----"); } diff --git a/ggml/src/ggml-hexagon/htp/act-ops.c b/ggml/src/ggml-hexagon/htp/act-ops.c index 5911c0890..d59ac0770 100644 --- a/ggml/src/ggml-hexagon/htp/act-ops.c +++ b/ggml/src/ggml-hexagon/htp/act-ops.c @@ -7,7 +7,7 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hvx-utils.h" #define GGML_COMMON_DECL_C @@ -53,13 +53,24 @@ const uint32_t nb2 = dst->nb[2]; \ const uint32_t nb3 = dst->nb[3]; +struct htp_act_context; + +typedef void (*glu_compute_fn_t)(const float * restrict src0, + const float * restrict src1, + float * restrict dst, + const uint32_t num_rows, + const struct htp_act_context * actx); + struct htp_act_context { struct htp_ops_context * octx; + glu_compute_fn_t compute; + const char * op_str; + // Precomputed values - const uint8_t * data_src0; - const uint8_t * data_src1; - uint8_t * data_dst; + dma_addr_t data_src0; + dma_addr_t data_src1; + dma_addr_t data_dst; size_t src0_row_size; size_t src1_row_size; @@ -134,10 +145,10 @@ static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L, // swiglu(x) = x1 * sigmoid(x0) static void swiglu_f32(const float * restrict src0, - const float * restrict src1, - float * restrict dst, - const uint32_t num_rows, - const struct htp_act_context * actx) { + const float * restrict src1, + float * restrict dst, + const uint32_t num_rows, + const struct htp_act_context * actx) { htp_glu_op_preamble; for (uint32_t ib = 0; ib < num_rows; ib++) { @@ -152,10 +163,10 @@ static void swiglu_f32(const float * restrict src0, // out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f) static void swiglu_oai_f32(const float * restrict src0, - const float * restrict src1, - float * restrict dst, - const uint32_t num_rows, - const struct htp_act_context * actx) { + const float * restrict src1, + float * restrict dst, + const uint32_t num_rows, + const struct htp_act_context * actx) { htp_glu_op_preamble; const float alpha = ((const float *) (actx->octx->op_params))[2]; const float limit = ((const float *) (actx->octx->op_params))[3]; @@ -181,10 +192,10 @@ static void swiglu_oai_f32(const float * restrict src0, } static void swiglu_clamp_f32(const float * restrict src0, - const float * restrict src1, - float * restrict dst, - const uint32_t num_rows, - const struct htp_act_context * actx) { + const float * restrict src1, + float * restrict dst, + const uint32_t num_rows, + const struct htp_act_context * actx) { htp_glu_op_preamble; const float limit = ((const float *) (actx->octx->op_params))[3]; @@ -353,10 +364,10 @@ static inline void hvx_geglu_quick_f32_aa(uint8_t * restrict dst, const uint8_t // geglu(x, g) = gelu(x) * g static void geglu_f32(const float * restrict src0, - const float * restrict src1, - float * restrict dst, - const uint32_t num_rows, - const struct htp_act_context * actx) { + const float * restrict src1, + float * restrict dst, + const uint32_t num_rows, + const struct htp_act_context * actx) { htp_glu_op_preamble; for (uint32_t ib = 0; ib < num_rows; ib++) { @@ -385,111 +396,100 @@ static void geglu_quick_f32(const float * restrict src0, } } -#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR) \ - static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_act_context * actx = (struct htp_act_context *) data; \ - htp_act_preamble; \ - \ - struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL; \ - \ - size_t src0_row_size = actx->src0_row_size; \ - size_t src1_row_size = actx->src1_row_size; \ - size_t dst_row_size = actx->dst_row_size; \ - \ - size_t src0_row_stride = actx->src0_row_stride; \ - size_t src1_row_stride = actx->src1_row_stride; \ - \ - const uint32_t src0_nrows = actx->src0_nrows; \ - const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread; \ - \ - const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows); \ - \ - /* no work for this thread */ \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - const uint8_t * restrict data_src0 = actx->data_src0; \ - const uint8_t * restrict data_src1 = actx->data_src1; \ - uint8_t * restrict data_dst = actx->data_dst; \ - \ - const size_t src0_row_size_aligned = actx->src0_row_size_aligned; \ - const size_t src1_row_size_aligned = actx->src1_row_size_aligned; \ - const size_t dst_row_size_aligned = actx->dst_row_size_aligned; \ - \ - uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread); \ - uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread); \ - uint8_t * restrict dst_spad_data = actx->vtcm_dst + (ith * actx->vtcm_dst_size_per_thread); \ - \ - size_t src0_spad_half_size = actx->src0_spad_half_size; \ - size_t src1_spad_half_size = actx->src1_spad_half_size; \ - size_t dst_spad_half_size = actx->dst_spad_half_size; \ - \ - const int BLOCK = actx->block; \ - if (BLOCK == 0) { \ - FARF(ERROR, \ - OP_STR \ - " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", \ - actx->vtcm_src0_size_per_thread, src0_row_size_aligned); \ - return; \ - } \ - \ - dma_queue * dma_queue = actx->octx->ctx->dma[ith]; \ - \ - /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */ \ - for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \ - const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); \ - \ - /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */ \ - dma_queue_push_vtcm_to_ddr(dma_queue, \ - dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)), \ - dst_row_size, dst_row_size_aligned, 0); \ - \ - dma_queue_push( \ - dma_queue, \ - dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)), \ - src0_row_size_aligned, src0_row_stride, src0_row_size, block_size); \ - dma_queue_push( \ - dma_queue, \ - dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)), \ - src1_row_size_aligned, src1_row_stride, src1_row_size, block_size); \ - } \ - \ - for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) { \ - const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); \ - \ - float * dst_spad = (float *) dma_queue_pop(dma_queue).src; \ - float * src0_spad = (float *) dma_queue_pop(dma_queue).dst; \ - float * src1_spad = (float *) dma_queue_pop(dma_queue).dst; \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); \ - CORE_EXPR; \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); \ - \ - dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad), \ - dst_row_size, dst_row_size_aligned, block_size); \ - \ - /* prefetch N+2 loop iteration if any */ \ - const uint32_t pref_block = (ir + BLOCK * 2); \ - if (pref_block < src0_end_row) { \ - const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block); \ - dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)), \ - src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size); \ - dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)), \ - src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size); \ - } \ - } \ - \ - dma_queue_flush(dma_queue); \ - \ +static void glu_f32_per_thread(unsigned int nth, unsigned int ith, void * data) { + struct htp_act_context * actx = (struct htp_act_context *) data; + htp_act_preamble; + + struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL; + + size_t src0_row_size = actx->src0_row_size; + size_t src1_row_size = actx->src1_row_size; + size_t dst_row_size = actx->dst_row_size; + + size_t src0_row_stride = actx->src0_row_stride; + size_t src1_row_stride = actx->src1_row_stride; + + const uint32_t src0_nrows = actx->src0_nrows; + const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread; + + const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows); + + /* no work for this thread */ + if (src0_start_row >= src0_end_row) { + return; } -DEFINE_GLU_PER_THREAD(swiglu, "swiglu-f32", swiglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx)) -DEFINE_GLU_PER_THREAD(swiglu_oai, "swiglu-oai-f32", swiglu_oai_f32(src0_spad, src1_spad, dst_spad, block_size, actx)) -DEFINE_GLU_PER_THREAD(swiglu_clamp, "swiglu-clamp-f32", swiglu_clamp_f32(src0_spad, src1_spad, dst_spad, block_size, actx)) -DEFINE_GLU_PER_THREAD(geglu, "geglu-f32", geglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx)) -DEFINE_GLU_PER_THREAD(geglu_quick, "geglu-quick-f32", geglu_quick_f32(src0_spad, src1_spad, dst_spad, block_size, actx)) + const dma_addr_t data_src0 = actx->data_src0; + const dma_addr_t data_src1 = actx->data_src1; + const dma_addr_t data_dst = actx->data_dst; + + const size_t src0_row_size_aligned = actx->src0_row_size_aligned; + const size_t src1_row_size_aligned = actx->src1_row_size_aligned; + const size_t dst_row_size_aligned = actx->dst_row_size_aligned; + + uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread); + uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread); + uint8_t * restrict dst_spad_data = actx->vtcm_dst + (ith * actx->vtcm_dst_size_per_thread); + + size_t src0_spad_half_size = actx->src0_spad_half_size; + size_t src1_spad_half_size = actx->src1_spad_half_size; + size_t dst_spad_half_size = actx->dst_spad_half_size; + + const int BLOCK = actx->block; + if (BLOCK == 0) { + FARF(ERROR, "%s : VTCM reservation %zu is too small, needed %zu\n", + actx->op_str, actx->vtcm_src0_size_per_thread, src0_row_size_aligned); + return; + } + + dma_queue * dma_q = actx->octx->ctx->dma[ith]; + glu_compute_fn_t compute = actx->compute; + + for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { + const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); + + /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */ + dma_queue_push(dma_q, + dma_make_data(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)), + dst_row_size, dst_row_size_aligned, dst_row_size, 0); + + dma_queue_push(dma_q, + dma_make_data(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)), + src0_row_size_aligned, src0_row_stride, src0_row_size, block_size); + + dma_queue_push(dma_q, + dma_make_data(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)), + src1_row_size_aligned, src1_row_stride, src1_row_size, block_size); + } + + for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) { + const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); + + float * dst_spad = (float *) dma_queue_pop(dma_q).src; + float * src0_spad = (float *) dma_queue_pop(dma_q).dst; + float * src1_spad = (float *) dma_queue_pop(dma_q).dst; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); + compute(src0_spad, src1_spad, dst_spad, block_size, actx); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); + + dma_queue_push(dma_q, dma_make_data(data_dst + (ir * dst_row_size), dst_spad), + dst_row_size, dst_row_size_aligned, dst_row_size, block_size); + + /* prefetch N+2 loop iteration if any */ + const uint32_t pref_block = (ir + BLOCK * 2); + if (pref_block < src0_end_row) { + const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block); + dma_queue_push(dma_q, dma_make_data(src0_spad, data_src0 + (pref_block * src0_row_stride)), + src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size); + dma_queue_push(dma_q, dma_make_data(src1_spad, data_src1 + (pref_block * src1_row_stride)), + src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size); + } + } + + dma_queue_flush(dma_q); +} static int execute_op_activations_f32(struct htp_ops_context * octx) { const struct htp_tensor * src0 = octx->src[0]; @@ -501,33 +501,33 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - worker_callback_t act_op_func; - const char * op_type = NULL; + glu_compute_fn_t compute_fn = NULL; + const char * op_type = NULL; switch (octx->op) { case HTP_OP_GLU_SWIGLU: - act_op_func = (worker_callback_t)glu_swiglu_f32_per_thread; - op_type = "swiglu-f32"; + compute_fn = swiglu_f32; + op_type = "swiglu-f32"; break; case HTP_OP_GLU_SWIGLU_OAI: - act_op_func = (worker_callback_t)glu_swiglu_oai_f32_per_thread; - op_type = "swiglu-oai-f32"; + compute_fn = swiglu_oai_f32; + op_type = "swiglu-oai-f32"; break; case HTP_OP_GLU_SWIGLU_CLAMP: - act_op_func = (worker_callback_t) glu_swiglu_clamp_f32_per_thread; - op_type = "swiglu-clamp-f32"; + compute_fn = swiglu_clamp_f32; + op_type = "swiglu-clamp-f32"; break; case HTP_OP_GLU_GEGLU: - act_op_func = (worker_callback_t)glu_geglu_f32_per_thread; - op_type = "geglu-f32"; + compute_fn = geglu_f32; + op_type = "geglu-f32"; break; case HTP_OP_GLU_GEGLU_QUICK: - act_op_func = (worker_callback_t)glu_geglu_quick_f32_per_thread; - op_type = "geglu-quick-f32"; + compute_fn = geglu_quick_f32; + op_type = "geglu-quick-f32"; break; default: FARF(ERROR, "Unsupported activations Op %u\n", octx->op); @@ -588,13 +588,11 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) { L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads); } - if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) { - return HTP_STATUS_OK; - } - // Prepare context struct htp_act_context actx; - actx.octx = octx; + actx.octx = octx; + actx.compute = compute_fn; + actx.op_str = op_type; actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); @@ -628,9 +626,9 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) { actx.nc = dst->ne[0]; - // Pointers and GLU logic - const uint8_t * data_src0 = (const uint8_t *) src0->data; - const uint8_t * data_src1 = src1 ? (const uint8_t *) src1->data : NULL; + // Addresses and GLU logic + dma_addr_t data_src0 = src0->data; + dma_addr_t data_src1 = src1 ? src1->data : 0; if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU || octx->op == HTP_OP_GLU_SWIGLU_OAI || @@ -651,9 +649,9 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) { actx.data_src0 = data_src0; actx.data_src1 = data_src1; - actx.data_dst = (uint8_t *) dst->data; + actx.data_dst = dst->data; - work_queue_run(octx->ctx->work_queue, act_op_func, &actx, n_threads); + work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &actx, n_threads); return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/allreduce-ops.c b/ggml/src/ggml-hexagon/htp/allreduce-ops.c index d6e7f0d10..7b577befb 100644 --- a/ggml/src/ggml-hexagon/htp/allreduce-ops.c +++ b/ggml/src/ggml-hexagon/htp/allreduce-ops.c @@ -14,7 +14,7 @@ #include "htp-ops.h" #include "hvx-utils.h" #include "htp-tensor.h" -#include "hex-dma.h" +#include "dma-queue.h" #include "hex-profile.h" #include "allreduce-ops.h" #include "htp-fence.h" @@ -38,97 +38,97 @@ struct htp_allreduce_context { uint8_t * res_spad_base; }; -#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD) \ -static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \ - struct htp_ops_context * octx = actx->octx; \ - \ - const uint32_t n_ranks = actx->n_ranks; \ - const uint32_t n_dsts = actx->n_dsts; \ - const uint32_t block_elems = actx->block_elems; \ - \ - const uint32_t dr = actx->elems_per_thread; \ - const uint32_t ir0 = actx->rank_elem_start + dr * ith; \ - const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem); \ - if (ir0 >= ir1) return; \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - dma_queue * q = octx->ctx->dma[ith]; \ - \ - uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ - for (uint32_t s = 0; s < n_ranks; s++) { \ - src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \ - } \ - uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \ - uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL; \ - \ - const size_t spad_half = actx->vtcm_size_per_thread / 2; \ - uint32_t ir_prefetch = ir0; \ - int spad_idx = 0; \ - \ - for (int k = 0; k < 2 && ir_prefetch < ir1; k++) { \ - uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch); \ - size_t cur_bytes = cur_elems * sizeof(TYPE); \ - uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ - for (uint32_t d = 0; d < n_dsts; d++) { \ - uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE); \ - dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0); \ - } \ - for (uint32_t s = 0; s < n_ranks; s++) { \ - uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ - const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ - dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ - } \ - if (HAS_ADD) { \ - uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ - const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ - dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ - } \ - ir_prefetch += cur_elems; \ - spad_idx ^= 1; \ - } \ - \ - for (uint32_t ir = ir0; ir < ir1; ) { \ - uint32_t cur_elems = MIN(block_elems, ir1 - ir); \ - size_t cur_bytes = cur_elems * sizeof(TYPE); \ - uint8_t * d_spad = NULL; \ - for (uint32_t d = 0; d < n_dsts; d++) { \ - d_spad = (uint8_t *) dma_queue_pop(q).src; \ - } \ - uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ - for (uint32_t s = 0; s < n_ranks; s++) { \ - s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \ - } \ - uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL; \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ - HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems); \ - for (uint32_t s = 2; s < n_ranks; s++) { \ - HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems); \ - } \ - if (HAS_ADD) { \ - HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ - for (uint32_t d = 0; d < n_dsts; d++) { \ - uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE); \ - dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1); \ - } \ - if (ir_prefetch < ir1) { \ - uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch); \ - size_t next_bytes = next_elems * sizeof(TYPE); \ - for (uint32_t s = 0; s < n_ranks; s++) { \ - const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ - dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \ - } \ - if (HAS_ADD) { \ - const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ - dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1); \ - } \ - ir_prefetch += next_elems; \ - } \ - ir += cur_elems; \ - } \ - dma_queue_flush(q); \ +#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD) \ +static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \ + struct htp_ops_context * octx = actx->octx; \ + \ + const uint32_t n_ranks = actx->n_ranks; \ + const uint32_t n_dsts = actx->n_dsts; \ + const uint32_t block_elems = actx->block_elems; \ + \ + const uint32_t dr = actx->elems_per_thread; \ + const uint32_t ir0 = actx->rank_elem_start + dr * ith; \ + const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem); \ + if (ir0 >= ir1) return; \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ + \ + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \ + } \ + uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \ + uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL; \ + \ + const size_t spad_half = actx->vtcm_size_per_thread / 2; \ + uint32_t ir_prefetch = ir0; \ + int spad_idx = 0; \ + \ + for (int k = 0; k < 2 && ir_prefetch < ir1; k++) { \ + uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch); \ + size_t cur_bytes = cur_elems * sizeof(TYPE); \ + uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + dma_addr_t d_ddr = octx->dsts[d]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0); \ + } \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ + const dma_addr_t s_ddr = octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + if (HAS_ADD) { \ + uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ + const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(dma_q, dma_make_data(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + ir_prefetch += cur_elems; \ + spad_idx ^= 1; \ + } \ + \ + for (uint32_t ir = ir0; ir < ir1; ) { \ + uint32_t cur_elems = MIN(block_elems, ir1 - ir); \ + size_t cur_bytes = cur_elems * sizeof(TYPE); \ + uint8_t * d_spad = NULL; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + d_spad = (uint8_t *) dma_queue_pop(dma_q).src; \ + } \ + uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + s_spad[s] = (uint8_t *) dma_queue_pop(dma_q).dst; \ + } \ + uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(dma_q).dst : NULL; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ + HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems); \ + for (uint32_t s = 2; s < n_ranks; s++) { \ + HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems); \ + } \ + if (HAS_ADD) { \ + HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems); \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + dma_addr_t d_ddr = octx->dsts[d]->data + ir * sizeof(TYPE); \ + dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + if (ir_prefetch < ir1) { \ + uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch); \ + size_t next_bytes = next_elems * sizeof(TYPE); \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + const dma_addr_t s_next = octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(dma_q, dma_make_data(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \ + } \ + if (HAS_ADD) { \ + const dma_addr_t r_next = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(dma_q, dma_make_data(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1); \ + } \ + ir_prefetch += next_elems; \ + } \ + ir += cur_elems; \ + } \ + dma_queue_flush(dma_q); \ } DEFINE_ALLREDUCE_THREAD_DMA_1D(f16, __fp16, hvx_add_f16_aaa, 0) @@ -154,7 +154,7 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, if (r0 >= r1) return; \ \ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - dma_queue * q = octx->ctx->dma[ith]; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ \ uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ for (uint32_t s = 0; s < n_ranks; s++) { \ @@ -171,18 +171,18 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch); \ uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ for (uint32_t d = 0; d < n_dsts; d++) { \ - uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1]; \ - dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0); \ + dma_addr_t d_ddr = octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1]; \ + dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0); \ } \ for (uint32_t s = 0; s < n_ranks; s++) { \ uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ - const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ - dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows); \ + const dma_addr_t s_ddr = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ + dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows); \ } \ if (HAS_ADD && !IS_ROW_BCAST) { \ uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ - const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ - dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows); \ + const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ + dma_queue_push(dma_q, dma_make_data(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows); \ } \ r_prefetch += cur_rows; \ spad_idx ^= 1; \ @@ -192,13 +192,13 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, uint32_t cur_rows = MIN(block_rows, r1 - r); \ uint8_t * d_spad = NULL; \ for (uint32_t d = 0; d < n_dsts; d++) { \ - d_spad = (uint8_t *) dma_queue_pop(q).src; \ + d_spad = (uint8_t *) dma_queue_pop(dma_q).src; \ } \ uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ for (uint32_t s = 0; s < n_ranks; s++) { \ - s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \ + s_spad[s] = (uint8_t *) dma_queue_pop(dma_q).dst; \ } \ - uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL; \ + uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(dma_q).dst : NULL; \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \ for (uint32_t row = 0; row < cur_rows; row++) { \ uint8_t * d_row = d_spad + row * row_size_aligned; \ @@ -216,24 +216,24 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \ for (uint32_t d = 0; d < n_dsts; d++) { \ - uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1]; \ - dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows); \ + dma_addr_t d_ddr = octx->dsts[d]->data + r * octx->dsts[d]->nb[1]; \ + dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows); \ } \ if (r_prefetch < r1) { \ uint32_t next_rows = MIN(block_rows, r1 - r_prefetch); \ for (uint32_t s = 0; s < n_ranks; s++) { \ - const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ - dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows); \ + const dma_addr_t s_next = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ + dma_queue_push(dma_q, dma_make_data(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows); \ } \ if (HAS_ADD && !IS_ROW_BCAST) { \ - const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ - dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows); \ + const dma_addr_t r_next = octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ + dma_queue_push(dma_q, dma_make_data(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows); \ } \ r_prefetch += next_rows; \ } \ r += cur_rows; \ } \ - dma_queue_flush(q); \ + dma_queue_flush(dma_q); \ } DEFINE_ALLREDUCE_THREAD_DMA_2D(f16, __fp16, hvx_add_f16_aaa, 0, 0) @@ -406,11 +406,11 @@ int op_allreduce(struct htp_ops_context * octx) { } if (has_add && actx.is_row_bcast) { - const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data; + const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data; const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float)); - dma_queue * q = octx->ctx->dma[0]; - dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1); - dma_queue_pop(q); + dma_queue * dma_q = octx->ctx->dma[0]; + dma_queue_push(dma_q, dma_make_data(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1); + dma_queue_pop(dma_q); } work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads); diff --git a/ggml/src/ggml-hexagon/htp/argsort-ops.c b/ggml/src/ggml-hexagon/htp/argsort-ops.c index 6ee614d3d..9e9e74651 100644 --- a/ggml/src/ggml-hexagon/htp/argsort-ops.c +++ b/ggml/src/ggml-hexagon/htp/argsort-ops.c @@ -9,7 +9,7 @@ #include "ggml.h" #include "hvx-utils.h" -#include "hex-dma.h" +#include "dma-queue.h" #include "hex-common.h" #include "htp-ctx.h" @@ -591,6 +591,10 @@ int op_argsort(struct htp_ops_context * octx) { const struct htp_tensor * src0 = octx->src[0]; const struct htp_tensor * dst = octx->dst; + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } + const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3]; const size_t dst_row_size = dst->ne[0] * sizeof(int32_t); diff --git a/ggml/src/ggml-hexagon/htp/binary-ops.c b/ggml/src/ggml-hexagon/htp/binary-ops.c index bfa849e0e..155e85237 100644 --- a/ggml/src/ggml-hexagon/htp/binary-ops.c +++ b/ggml/src/ggml-hexagon/htp/binary-ops.c @@ -8,13 +8,14 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hvx-utils.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" #include "hex-common.h" #include "hex-profile.h" +#include "binary-ops.h" #include "htp-ctx.h" #include "htp-ops.h" #include "htp-tensor.h" @@ -26,6 +27,8 @@ // Context for binary operations struct htp_binary_context { struct htp_ops_context * octx; + struct htp_binary_vtcm_layout vtcm_layout; + uint8_t * vtcm_base; struct fastdiv_values src0_dim1_div; // ne01 struct fastdiv_values src0_dim2_div; // ne02 @@ -42,9 +45,12 @@ struct htp_binary_context { size_t src0_row_size_aligned; size_t src1_row_size_aligned; size_t dst_row_size_aligned; + size_t row_size_bytes; bool split_at_ne01; bool split_at_ne02; + + void * compute; }; #define htp_binary_preamble \ @@ -95,114 +101,213 @@ static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_ return MIN(bctx->block_max, block_limit); } -// Macro for scalar op switch -#define COMPUTE_SCALAR_OP(DST, SRC, VAL, TYPE, N) \ - if(TYPE == HTP_TYPE_F32) { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \ - case HTP_OP_SUB: hvx_sub_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \ - case HTP_OP_MUL: hvx_mul_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \ - case HTP_OP_DIV: hvx_mul_scalar_f32_aa(DST, SRC, 1.0f / (*(float *)VAL), N); break; \ - default: break; \ - } \ - } \ - else { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \ - case HTP_OP_SUB: hvx_sub_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \ - case HTP_OP_MUL: hvx_mul_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \ - case HTP_OP_DIV: hvx_div_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \ - default: break; \ - } \ - } +// Out-of-line compute micro-kernels -// Macro for vector op switch (All Aligned) -#define COMPUTE_VECTOR_OP_AAA(DST, SRC0, SRC1, TYPE, N) \ - if(TYPE == HTP_TYPE_F32) { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_f32_aaa(DST, SRC0, SRC1, N); break; \ - case HTP_OP_SUB: hvx_sub_f32_aaa(DST, SRC0, SRC1, N); break; \ - case HTP_OP_MUL: hvx_mul_f32_aaa(DST, SRC0, SRC1, N); break; \ - case HTP_OP_DIV: hvx_div_f32_aaa(DST, SRC0, SRC1, N); break; \ - default: break; \ - } \ - } \ - else { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_f16_aaa(DST, SRC0, SRC1, N); break; \ - case HTP_OP_SUB: hvx_sub_f16_aaa(DST, SRC0, SRC1, N); break; \ - case HTP_OP_MUL: hvx_mul_f16_aaa(DST, SRC0, SRC1, N); break; \ - case HTP_OP_DIV: hvx_div_f16_aaa(DST, SRC0, SRC1, N); break; \ - default: break; \ - } \ - } +typedef void (*compute_scalar_dma_t)( + uint8_t * dst, const uint8_t * src0, const void * s1_table, + uint32_t cur_i11, uint32_t ne11, uint32_t n_rows, + size_t dst_stride, size_t src0_stride, uint32_t ne00); -// Macro for vector op switch (Dst Aligned, Src0 Aligned, Src1 Unaligned) -#define COMPUTE_VECTOR_OP_AAU(DST, SRC0, SRC1, TYPE, N) \ - if(TYPE == HTP_TYPE_F32) { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_f32_aau(DST, SRC0, SRC1, N); break; \ - case HTP_OP_SUB: hvx_sub_f32_aau(DST, SRC0, SRC1, N); break; \ - case HTP_OP_MUL: hvx_mul_f32_aau(DST, SRC0, SRC1, N); break; \ - case HTP_OP_DIV: hvx_div_f32_aau(DST, SRC0, SRC1, N); break; \ - default: break; \ - } \ - } \ - else { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_f16_aau(DST, SRC0, SRC1, N); break; \ - case HTP_OP_SUB: hvx_sub_f16_aau(DST, SRC0, SRC1, N); break; \ - case HTP_OP_MUL: hvx_mul_f16_aau(DST, SRC0, SRC1, N); break; \ - case HTP_OP_DIV: hvx_div_f16_aau(DST, SRC0, SRC1, N); break; \ - default: break; \ - } \ - } +#define DEFINE_COMPUTE_SCALAR_DMA(NAME, TYPE, HVX_STMT) \ +static void compute_scalar_dma_##NAME( \ + uint8_t * dst, const uint8_t * src0, const void * s1_table, \ + uint32_t cur_i11, uint32_t ne11, uint32_t n_rows, \ + size_t dst_stride, size_t src0_stride, uint32_t ne00) { \ + const TYPE * table = (const TYPE *) s1_table; \ + for (uint32_t r = 0; r < n_rows; r++) { \ + uint8_t * r_dst = dst + r * dst_stride; \ + const uint8_t * r_src0 = src0 + r * src0_stride; \ + TYPE val = table[cur_i11]; \ + HVX_STMT; \ + if (ne11 > 1 && ++cur_i11 == ne11) { \ + cur_i11 = 0; \ + } \ + } \ +} -// Macro for vector op switch (All Unaligned - generic loop used in element repeat) -#define COMPUTE_VECTOR_OP_UUU(DST, SRC0, SRC1, TYPE, N) \ - if(TYPE == HTP_TYPE_F32) { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_f32_uuu(DST, SRC0, SRC1, N); break; \ - case HTP_OP_SUB: hvx_sub_f32_uuu(DST, SRC0, SRC1, N); break; \ - case HTP_OP_MUL: hvx_mul_f32_uuu(DST, SRC0, SRC1, N); break; \ - case HTP_OP_DIV: hvx_div_f32_uuu(DST, SRC0, SRC1, N); break; \ - default: break; \ - } \ - } \ - else { \ - switch (octx->op) { \ - case HTP_OP_ADD: hvx_add_f16_uuu(DST, SRC0, SRC1, N); break; \ - case HTP_OP_SUB: hvx_sub_f16_uuu(DST, SRC0, SRC1, N); break; \ - case HTP_OP_MUL: hvx_mul_f16_uuu(DST, SRC0, SRC1, N); break; \ - case HTP_OP_DIV: hvx_div_f16_uuu(DST, SRC0, SRC1, N); break; \ - default: break; \ - } \ - } +DEFINE_COMPUTE_SCALAR_DMA(add_f32, float, hvx_add_scalar_f32_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR_DMA(add_f16, _Float16, hvx_add_scalar_f16_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR_DMA(sub_f32, float, hvx_sub_scalar_f32_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR_DMA(sub_f16, _Float16, hvx_sub_scalar_f16_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR_DMA(mul_f32, float, hvx_mul_scalar_f32_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR_DMA(mul_f16, _Float16, hvx_mul_scalar_f16_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR_DMA(div_f32, float, hvx_mul_scalar_f32_aa(r_dst, r_src0, 1.0f / (val), ne00)) +DEFINE_COMPUTE_SCALAR_DMA(div_f16, _Float16, hvx_div_scalar_f16_aa(r_dst, r_src0, val, ne00)) -// 1. Scalar src1 (ne10 == 1) -static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) { +typedef void (*compute_scalar_t)( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_ptr, uint32_t s1_stride, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00); + +#define DEFINE_COMPUTE_SCALAR(NAME, TYPE, HVX_STMT) \ +static void compute_scalar_##NAME( \ + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_ptr, uint32_t s1_stride, \ + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) { \ + for (uint32_t r = 0; r < n_rows; r++) { \ + uint8_t * r_dst = dst + r * dst_stride; \ + const uint8_t * r_src0 = src0 + r * src0_stride; \ + TYPE val = *(const TYPE *)(src1_ptr + r * s1_stride); \ + HVX_STMT; \ + } \ +} + +DEFINE_COMPUTE_SCALAR(add_f32, float, hvx_add_scalar_f32_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR(add_f16, _Float16, hvx_add_scalar_f16_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR(sub_f32, float, hvx_sub_scalar_f32_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR(sub_f16, _Float16, hvx_sub_scalar_f16_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR(mul_f32, float, hvx_mul_scalar_f32_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR(mul_f16, _Float16, hvx_mul_scalar_f16_aa(r_dst, r_src0, val, ne00)) +DEFINE_COMPUTE_SCALAR(div_f32, float, hvx_mul_scalar_f32_aa(r_dst, r_src0, 1.0f / (val), ne00)) +DEFINE_COMPUTE_SCALAR(div_f16, _Float16, hvx_div_scalar_f16_aa(r_dst, r_src0, val, ne00)) + +typedef void (*compute_same_shape_t)( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, size_t src1_stride, uint32_t ne00); + +#define DEFINE_COMPUTE_SAME_SHAPE(NAME, HVX_FN) \ +static void compute_same_shape_##NAME( \ + uint8_t * dst, const uint8_t * src0, const uint8_t * src1, \ + uint32_t n_rows, size_t dst_stride, size_t src0_stride, size_t src1_stride, uint32_t ne00) { \ + for (uint32_t r = 0; r < n_rows; r++) { \ + HVX_FN(dst + r * dst_stride, src0 + r * src0_stride, src1 + r * src1_stride, ne00); \ + } \ +} + +DEFINE_COMPUTE_SAME_SHAPE(add_f32, hvx_add_f32_aaa) +DEFINE_COMPUTE_SAME_SHAPE(add_f16, hvx_add_f16_aaa) +DEFINE_COMPUTE_SAME_SHAPE(sub_f32, hvx_sub_f32_aaa) +DEFINE_COMPUTE_SAME_SHAPE(sub_f16, hvx_sub_f16_aaa) +DEFINE_COMPUTE_SAME_SHAPE(mul_f32, hvx_mul_f32_aaa) +DEFINE_COMPUTE_SAME_SHAPE(mul_f16, hvx_mul_f16_aaa) +DEFINE_COMPUTE_SAME_SHAPE(div_f32, hvx_div_f32_aaa) +DEFINE_COMPUTE_SAME_SHAPE(div_f16, hvx_div_f16_aaa) + +typedef void (*compute_row_bcast_t)( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00); + +#define DEFINE_COMPUTE_ROW_BCAST(NAME, HVX_FN) \ +static void compute_row_bcast_##NAME( \ + uint8_t * dst, const uint8_t * src0, const uint8_t * src1, \ + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) { \ + for (uint32_t r = 0; r < n_rows; r++) { \ + HVX_FN(dst + r * dst_stride, src0 + r * src0_stride, src1, ne00); \ + } \ +} + +DEFINE_COMPUTE_ROW_BCAST(add_f32, hvx_add_f32_aaa) +DEFINE_COMPUTE_ROW_BCAST(add_f16, hvx_add_f16_aaa) +DEFINE_COMPUTE_ROW_BCAST(sub_f32, hvx_sub_f32_aaa) +DEFINE_COMPUTE_ROW_BCAST(sub_f16, hvx_sub_f16_aaa) +DEFINE_COMPUTE_ROW_BCAST(mul_f32, hvx_mul_f32_aaa) +DEFINE_COMPUTE_ROW_BCAST(mul_f16, hvx_mul_f16_aaa) +DEFINE_COMPUTE_ROW_BCAST(div_f32, hvx_div_f32_aaa) +DEFINE_COMPUTE_ROW_BCAST(div_f16, hvx_div_f16_aaa) + +typedef void (*compute_complex_t)( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane, + uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00); + +#define DEFINE_COMPUTE_COMPLEX(NAME, HVX_FN) \ +static void compute_complex_##NAME( \ + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane, \ + uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11, \ + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) { \ + for (uint32_t r = 0; r < n_rows; r++) { \ + uint32_t i11 = fastmodulo(i01 + r, ne11, div11); \ + const uint8_t * r_src1 = src1_plane + i11 * nb11; \ + HVX_FN(dst + r * dst_stride, src0 + r * src0_stride, r_src1, ne00); \ + } \ +} + +DEFINE_COMPUTE_COMPLEX(add_f32, hvx_add_f32_aau) +DEFINE_COMPUTE_COMPLEX(add_f16, hvx_add_f16_aau) +DEFINE_COMPUTE_COMPLEX(sub_f32, hvx_sub_f32_aau) +DEFINE_COMPUTE_COMPLEX(sub_f16, hvx_sub_f16_aau) +DEFINE_COMPUTE_COMPLEX(mul_f32, hvx_mul_f32_aau) +DEFINE_COMPUTE_COMPLEX(mul_f16, hvx_mul_f16_aau) +DEFINE_COMPUTE_COMPLEX(div_f32, hvx_div_f32_aau) +DEFINE_COMPUTE_COMPLEX(div_f16, hvx_div_f16_aau) + +typedef void (*compute_repeat_t)( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane, + uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00, uint32_t ne10); + +#define DEFINE_COMPUTE_REPEAT(NAME, TYPE, HVX_FN) \ +static void compute_repeat_##NAME( \ + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane, \ + uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11, \ + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00, uint32_t ne10) { \ + for (uint32_t r = 0; r < n_rows; r++) { \ + uint32_t i11 = fastmodulo(i01 + r, ne11, div11); \ + const uint8_t * r_src1_row = src1_plane + i11 * nb11; \ + uint8_t * r_dst = dst + r * dst_stride; \ + const uint8_t * r_src0 = src0 + r * src0_stride; \ + for (uint32_t c = 0; c < ne00; c += ne10) { \ + uint32_t len = MIN(ne10, ne00 - c); \ + HVX_FN(r_dst + c * sizeof(TYPE), r_src0 + c * sizeof(TYPE), r_src1_row, len); \ + } \ + } \ +} + +DEFINE_COMPUTE_REPEAT(add_f32, float, hvx_add_f32_uuu) +DEFINE_COMPUTE_REPEAT(add_f16, _Float16, hvx_add_f16_uuu) +DEFINE_COMPUTE_REPEAT(sub_f32, float, hvx_sub_f32_uuu) +DEFINE_COMPUTE_REPEAT(sub_f16, _Float16, hvx_sub_f16_uuu) +DEFINE_COMPUTE_REPEAT(mul_f32, float, hvx_mul_f32_uuu) +DEFINE_COMPUTE_REPEAT(mul_f16, _Float16, hvx_mul_f16_uuu) +DEFINE_COMPUTE_REPEAT(div_f32, float, hvx_div_f32_uuu) +DEFINE_COMPUTE_REPEAT(div_f16, _Float16, hvx_div_f16_uuu) + +typedef void (*compute_add_id_t)( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_data, const char * src2_data, + uint32_t i01, uint32_t i02, uint32_t nb20, uint32_t nb21, uint32_t src1_stride, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00); + +static void compute_add_id_f32( + uint8_t * dst, const uint8_t * src0, const uint8_t * src1_data, const char * src2_data, + uint32_t i01, uint32_t i02, uint32_t nb20, uint32_t nb21, uint32_t src1_stride, + uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) { + for (uint32_t r = 0; r < n_rows; r++) { + uint32_t r_i01 = i01 + r; + const int32_t idx = *(const int32_t *)(src2_data + r_i01 * nb20 + i02 * nb21); + if (idx < 0) { + memcpy(dst + r * dst_stride, src0 + r * src0_stride, ne00 * sizeof(float)); + continue; + } + const uint8_t * r_src1 = src1_data + idx * src1_stride; + const uint8_t * r_src0 = src0 + r * src0_stride; + uint8_t * r_dst = dst + r * dst_stride; + hvx_add_f32_aaa(r_dst, r_src0, r_src1, ne00); + } +} + +// 1a. Scalar src1 in VTCM via DMA (ne10 == 1, ne12 == 1, ne13 == 1) +static void binary_thread_scalar_dma(unsigned int nth, unsigned int ith, void * data) { struct htp_binary_context * bctx = (struct htp_binary_context *) data; struct htp_ops_context * octx = bctx->octx; htp_binary_preamble; - const uint32_t src0_type = octx->src[0]->type; - const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16); + const uint32_t row_size_bytes = bctx->row_size_bytes; const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); if (start_row >= end_row) return; - FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); + FARF(HIGH, "binary-scalar-dma: %d/%d (%u:%u) row-size %u (%u)", + ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); - uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * dst_spad_base = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); - size_t src0_spad_half = octx->src0_spad.size_per_thread / 2; - size_t dst_spad_half = octx->dst_spad.size_per_thread / 2; + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); + size_t src0_spad_half = layout->src0_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; + const void * s1_table = VTCM_LAYOUT_PTR(const void, bctx->vtcm_base, layout->off_src1); - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint32_t ir_prefetch = start_row; int spad_idx = 0; - // Preamble for (int k = 0; k < 2 && ir_prefetch < end_row; k++) { uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); uint32_t i03, i02, i01, rem; @@ -211,26 +316,27 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) { i02 = fastdiv(rem, &bctx->src0_dim1_div); i01 = rem - i02 * ne01; - uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); - dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, + current_block_size); ir_prefetch += current_block_size; spad_idx ^= 1; } - // Main loop struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + compute_scalar_dma_t compute = (compute_scalar_dma_t) bctx->compute; for (uint32_t ir = start_row; ir < end_row; ) { uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); - uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src; - uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst; + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; uint32_t i03, i02, i01, rem; i03 = fastdiv(ir, &bctx->src0_dim12_div); @@ -238,67 +344,150 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) { i02 = fastdiv(rem, &bctx->src0_dim1_div); i01 = rem - i02 * ne01; - // src1 indices (broadcast/repeat) - uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div); - uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div); - uint32_t i11 = fastmodulo(i01, ne11, &bctx->src1_dim1_div); - - uint8_t * src1_ptr = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11; - uint32_t s1_stride = (ne11 == 1) ? 0 : nb11; + uint32_t cur_i11 = fastmodulo(i01, ne11, &bctx->src1_dim1_div); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - for (uint32_t r = 0; r < current_block_size; r++) { - uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned; - uint8_t * r_dst = d_spad + r * bctx->dst_row_size_aligned; - COMPUTE_SCALAR_OP(r_dst, r_src0, src1_ptr, src0_type, ne00); - src1_ptr += s1_stride; - } + compute(d_spad, s0_spad, s1_table, cur_i11, ne11, current_block_size, + bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, + current_block_size); if (ir_prefetch < end_row) { - uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); - uint32_t p03, p02, p01, prem; - p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); - prem = ir_prefetch - p03 * (ne02 * ne01); - p02 = fastdiv(prem, &bctx->src0_dim1_div); - p01 = prem - p02 * ne01; - uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; - - dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); - ir_prefetch += next_block_size; + uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, + next_block_size); + ir_prefetch += next_block_size; } ir += current_block_size; } - dma_queue_flush(q); + dma_queue_flush(dma_q); } -// 2. Vector Same Shape (ne1x == ne0x) or Simple Broadcast -static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, void * data) { +// 1b. Scalar src1 dynamic / pointer (ne10 == 1) +static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data) { struct htp_binary_context * bctx = (struct htp_binary_context *) data; struct htp_ops_context * octx = bctx->octx; htp_binary_preamble; - const uint32_t src0_type = octx->src[0]->type; - const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16); + const uint32_t row_size_bytes = bctx->row_size_bytes; const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); if (start_row >= end_row) return; - FARF(HIGH, "binary-same-shape: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); + FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)", + ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); - uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * src1_spad_base = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread); - uint8_t * dst_spad_base = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); + size_t src0_spad_half = layout->src0_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; - size_t src0_spad_half = octx->src0_spad.size_per_thread / 2; - size_t src1_spad_half = octx->src1_spad.size_per_thread / 2; - size_t dst_spad_half = octx->dst_spad.size_per_thread / 2; + dma_queue * dma_q = octx->ctx->dma[ith]; + uint32_t ir_prefetch = start_row; + int spad_idx = 0; - dma_queue * q = octx->ctx->dma[ith]; + for (int k = 0; k < 2 && ir_prefetch < end_row; k++) { + uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t i03, i02, i01, rem; + i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + rem = ir_prefetch - i03 * (ne02 * ne01); + i02 = fastdiv(rem, &bctx->src0_dim1_div); + i01 = rem - i02 * ne01; + + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + + uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; + uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; + + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); + ir_prefetch += current_block_size; + spad_idx ^= 1; + } + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + compute_scalar_t compute = (compute_scalar_t) bctx->compute; + + for (uint32_t ir = start_row; ir < end_row; ) { + uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); + + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; + + uint32_t i03, i02, i01, rem; + i03 = fastdiv(ir, &bctx->src0_dim12_div); + rem = ir - i03 * (ne02 * ne01); + i02 = fastdiv(rem, &bctx->src0_dim1_div); + i01 = rem - i02 * ne01; + + uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div); + uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div); + uint32_t i11 = fastmodulo(i01, ne11, &bctx->src1_dim1_div); + + const uint8_t * src1_ptr = (const uint8_t *)(uintptr_t) src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11; + uint32_t s1_stride = (ne11 == 1) ? 0 : nb11; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); + compute(d_spad, s0_spad, src1_ptr, s1_stride, current_block_size, + bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); + + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); + + if (ir_prefetch < end_row) { + uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); + ir_prefetch += next_block_size; + } + ir += current_block_size; + } + + dma_queue_flush(dma_q); +} + +// 2. Vector Same Shape (ne1x == ne0x) or Simple Broadcast +static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith, void * data) { + struct htp_binary_context * bctx = (struct htp_binary_context *) data; + struct htp_ops_context * octx = bctx->octx; + htp_binary_preamble; + + const uint32_t row_size_bytes = bctx->row_size_bytes; + const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; + const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); + if (start_row >= end_row) return; + + FARF(HIGH, "binary-same-shape: %d/%d (%u:%u) row-size %u (%u)", + ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); + + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * src1_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); + + size_t src0_spad_half = layout->src0_spad_half_size; + size_t src1_spad_half = layout->src1_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; + + dma_queue * dma_q = octx->ctx->dma[ith]; uint32_t ir_prefetch = start_row; int spad_idx = 0; @@ -314,36 +503,33 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi uint32_t i12 = (ne12 == 1) ? 0 : i02; uint32_t i11 = (ne11 == 1) ? 0 : i01; - uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; - uint8_t * src1_curr = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t src1_curr = src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; uint8_t * s1_spad = src1_spad_base + spad_idx * src1_spad_half; uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); - dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); - dma_queue_push(q, dma_make_ptr(s1_spad, src1_curr), bctx->src1_row_size_aligned, nb11, row_size_bytes, current_block_size); + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); + dma_queue_push(dma_q, dma_make_data(s1_spad, src1_curr), bctx->src1_row_size_aligned, nb11, row_size_bytes, current_block_size); ir_prefetch += current_block_size; spad_idx ^= 1; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + compute_same_shape_t compute = (compute_same_shape_t) bctx->compute; for (uint32_t ir = start_row; ir < end_row; ) { uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); - uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src; - uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst; - uint8_t * s1_spad = (uint8_t *) dma_queue_pop(q).dst; + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; + uint8_t * s1_spad = (uint8_t *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - for (uint32_t r = 0; r < current_block_size; r++) { - uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned; - uint8_t * r_src1 = s1_spad + r * bctx->src1_row_size_aligned; - uint8_t * r_dst = d_spad + r * bctx->dst_row_size_aligned; - COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00); - } + compute(d_spad, s0_spad, s1_spad, current_block_size, + bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, bctx->src1_row_size_aligned, ne00); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); uint32_t i03, i02, i01, rem; @@ -351,61 +537,61 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi rem = ir - i03 * (ne02 * ne01); i02 = fastdiv(rem, &bctx->src0_dim1_div); i01 = rem - i02 * ne01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); if (ir_prefetch < end_row) { - uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); - uint32_t p03, p02, p01, prem; - p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); - prem = ir_prefetch - p03 * (ne02 * ne01); - p02 = fastdiv(prem, &bctx->src0_dim1_div); - p01 = prem - p02 * ne01; + uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; - uint32_t p13 = (ne13 == 1) ? 0 : p03; - uint32_t p12 = (ne12 == 1) ? 0 : p02; - uint32_t p11 = (ne11 == 1) ? 0 : p01; + uint32_t p13 = (ne13 == 1) ? 0 : p03; + uint32_t p12 = (ne12 == 1) ? 0 : p02; + uint32_t p11 = (ne11 == 1) ? 0 : p01; - uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; - uint8_t * s1_next = (uint8_t *)src1->data + p13 * nb13 + p12 * nb12 + p11 * nb11; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_addr_t s1_next = src1->data + p13 * nb13 + p12 * nb12 + p11 * nb11; - dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); - dma_queue_push(q, dma_make_ptr(s1_spad, s1_next), bctx->src1_row_size_aligned, nb11, row_size_bytes, next_block_size); + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); + dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), bctx->src1_row_size_aligned, nb11, row_size_bytes, next_block_size); - ir_prefetch += next_block_size; + ir_prefetch += next_block_size; } ir += current_block_size; } - dma_queue_flush(q); + dma_queue_flush(dma_q); } // 3. Row Broadcast (ne11 == 1, ne12 == 1, single row src1) -static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith, void * data) { +static void binary_thread_vector_row_broadcast(unsigned int nth, unsigned int ith, void * data) { struct htp_binary_context * bctx = (struct htp_binary_context *) data; struct htp_ops_context * octx = bctx->octx; htp_binary_preamble; - const uint32_t src0_type = octx->src[0]->type; - const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16); + const uint32_t row_size_bytes = bctx->row_size_bytes; const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); if (start_row >= end_row) return; - FARF(HIGH, "binary-row-bcast: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); + FARF(HIGH, "binary-row-bcast: %d/%d (%u:%u) row-size %u (%u)", + ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); - uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * src1_spad_base = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread); - uint8_t * dst_spad_base = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); - size_t src0_spad_half = octx->src0_spad.size_per_thread / 2; - size_t dst_spad_half = octx->dst_spad.size_per_thread / 2; + size_t src0_spad_half = layout->src0_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint32_t ir_prefetch = start_row; int spad_idx = 0; - void * s1_ptr = (void *) src1_spad_base; + void * s1_ptr = VTCM_LAYOUT_PTR(void, bctx->vtcm_base, layout->off_src1); for (int k = 0; k < 2 && ir_prefetch < end_row; k++) { uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); @@ -414,77 +600,76 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith, uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; - uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); - dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); ir_prefetch += current_block_size; spad_idx ^= 1; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + compute_row_bcast_t compute = (compute_row_bcast_t) bctx->compute; for (uint32_t ir = start_row; ir < end_row; ) { uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); - uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src; - uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst; + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - for (uint32_t r = 0; r < current_block_size; r++) { - uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned; - uint8_t * r_src1 = (uint8_t *)s1_ptr; // Constant - uint8_t * r_dst = d_spad + r * bctx->dst_row_size_aligned; - COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00); - } + compute(d_spad, s0_spad, (const uint8_t *)s1_ptr, current_block_size, + bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div); uint32_t rem = ir - i03 * (ne02 * ne01); uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); if (ir_prefetch < end_row) { - uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); - uint32_t p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); - uint32_t prem = ir_prefetch - p03 * (ne02 * ne01); - uint32_t p02 = fastdiv(prem, &bctx->src0_dim1_div); - uint32_t p01 = prem - p02 * ne01; - uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; - dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); - ir_prefetch += next_block_size; + uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); + ir_prefetch += next_block_size; } ir += current_block_size; } - dma_queue_flush(q); + dma_queue_flush(dma_q); } // 4. Vector Complex (ne10 == ne00, complex broadcast) -static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void * data) { +static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, void * data) { struct htp_binary_context * bctx = (struct htp_binary_context *) data; struct htp_ops_context * octx = bctx->octx; htp_binary_preamble; - const uint32_t src0_type = octx->src[0]->type; - const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16); + const uint32_t row_size_bytes = bctx->row_size_bytes; const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); if (start_row >= end_row) return; - FARF(HIGH, "binary-complex: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); + FARF(HIGH, "binary-complex: %d/%d (%u:%u) row-size %u (%u)", + ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); - uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * dst_spad_base = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); - size_t src0_spad_half = octx->src0_spad.size_per_thread / 2; - size_t dst_spad_half = octx->dst_spad.size_per_thread / 2; + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); + size_t src0_spad_half = layout->src0_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint32_t ir_prefetch = start_row; int spad_idx = 0; @@ -495,86 +680,81 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void * uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; - uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); - dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); ir_prefetch += current_block_size; spad_idx ^= 1; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + compute_complex_t compute = (compute_complex_t) bctx->compute; for (uint32_t ir = start_row; ir < end_row; ) { uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); - uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src; - uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst; + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div); uint32_t rem = ir - i03 * (ne02 * ne01); uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; + uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div); + uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div); + const uint8_t * src1_plane = (const uint8_t *)(uintptr_t) src1->data + i13 * nb13 + i12 * nb12; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - for (uint32_t r = 0; r < current_block_size; r++) { - uint32_t r_i01 = i01 + r; - uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div); - uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div); - uint32_t i11 = fastmodulo(r_i01, ne11, &bctx->src1_dim1_div); - - uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned; - uint8_t * r_src1 = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11; - uint8_t * r_dst = d_spad + r * bctx->dst_row_size_aligned; - - // Read src1 from DDR (unaligned) - COMPUTE_VECTOR_OP_AAU(r_dst, r_src0, r_src1, src0_type, ne00); - } + compute(d_spad, s0_spad, src1_plane, i01, ne11, &bctx->src1_dim1_div, nb11, + current_block_size, bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); if (ir_prefetch < end_row) { - uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); - uint32_t p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); - uint32_t prem = ir_prefetch - p03 * (ne02 * ne01); - uint32_t p02 = fastdiv(prem, &bctx->src0_dim1_div); - uint32_t p01 = prem - p02 * ne01; - uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; - dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); - ir_prefetch += next_block_size; + uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); + ir_prefetch += next_block_size; } ir += current_block_size; } - dma_queue_flush(q); + dma_queue_flush(dma_q); } // 5. Element Repeat (ne10 != ne00) -static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void * data) { +static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, void * data) { struct htp_binary_context * bctx = (struct htp_binary_context *) data; struct htp_ops_context * octx = bctx->octx; htp_binary_preamble; - const uint32_t src0_type = octx->src[0]->type; - const uint32_t elem_size_bytes = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16); - const uint32_t row_size_bytes = ne00 * elem_size_bytes;; + const uint32_t row_size_bytes = bctx->row_size_bytes; const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); if (start_row >= end_row) return; - uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * dst_spad_base = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); - size_t src0_spad_half = octx->src0_spad.size_per_thread / 2; - size_t dst_spad_half = octx->dst_spad.size_per_thread / 2; + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); + size_t src0_spad_half = layout->src0_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; - FARF(HIGH, "binary-repeat: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); + FARF(HIGH, "binary-repeat: %d/%d (%u:%u) row-size %u (%u)", + ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned); - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint32_t ir_prefetch = start_row; int spad_idx = 0; @@ -585,71 +765,62 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void * uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; - uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); - dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); ir_prefetch += current_block_size; spad_idx ^= 1; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + compute_repeat_t compute = (compute_repeat_t) bctx->compute; for (uint32_t ir = start_row; ir < end_row; ) { uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); - uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src; - uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst; + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div); uint32_t rem = ir - i03 * (ne02 * ne01); uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; + uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div); + uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div); + const uint8_t * src1_plane = (const uint8_t *)(uintptr_t) src1->data + i13 * nb13 + i12 * nb12; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - for (uint32_t r = 0; r < current_block_size; r++) { - uint32_t r_i01 = i01 + r; - uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div); - uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div); - uint32_t i11 = fastmodulo(r_i01, ne11, &bctx->src1_dim1_div); - - uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned; - uint8_t * r_src1_row = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11; - uint8_t * r_dst = d_spad + r * bctx->dst_row_size_aligned; - - // Repeat src1 row - for (uint32_t c = 0; c < ne00; c += ne10) { - uint32_t len = MIN(ne10, ne00 - c); - // Use UUU for speed and simplicity - COMPUTE_VECTOR_OP_UUU(r_dst + c * elem_size_bytes, r_src0 + c * elem_size_bytes, r_src1_row, src0_type, len); - } - } + compute(d_spad, s0_spad, src1_plane, i01, ne11, &bctx->src1_dim1_div, nb11, + current_block_size, bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00, ne10); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); if (ir_prefetch < end_row) { - uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); - uint32_t p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); - uint32_t prem = ir_prefetch - p03 * (ne02 * ne01); - uint32_t p02 = fastdiv(prem, &bctx->src0_dim1_div); - uint32_t p01 = prem - p02 * ne01; - uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; - dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); - ir_prefetch += next_block_size; + uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); + ir_prefetch += next_block_size; } ir += current_block_size; } - dma_queue_flush(q); + dma_queue_flush(dma_q); } // 6. ADD_ID (src1 gathered via src2 indices) -static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) { +static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void * data) { struct htp_binary_context * bctx = (struct htp_binary_context *) data; struct htp_ops_context * octx = bctx->octx; @@ -662,27 +833,29 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) { const uint32_t ne01 = src0->ne[1]; const uint32_t ne02 = src0->ne[2]; const uint32_t ne03 = src0->ne[3]; - const uint32_t ne11 = src1->ne[1]; // for bounds check const uint32_t nb01 = src0->nb[1]; const uint32_t nb02 = src0->nb[2]; const uint32_t nb03 = src0->nb[3]; - const uint32_t nb11 = src1->nb[1]; // src1 row stride + const uint32_t src1_stride = bctx->src1_row_size_aligned; const uint32_t nb1 = dst->nb[1]; const uint32_t nb2 = dst->nb[2]; const uint32_t nb3 = dst->nb[3]; + const uint32_t row_size_bytes = bctx->row_size_bytes; const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith; const uint32_t end_row = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows); if (start_row >= end_row) return; - uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * dst_spad_base = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); - size_t src0_spad_half = octx->src0_spad.size_per_thread / 2; - size_t dst_spad_half = octx->dst_spad.size_per_thread / 2; + const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout; + uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread); + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst) + (ith * layout->dst_bytes_per_thread); + const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1); + size_t src0_spad_half = layout->src0_spad_half_size; + size_t dst_spad_half = layout->dst_spad_half_size; - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint32_t ir_prefetch = start_row; int spad_idx = 0; @@ -693,14 +866,14 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) { uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div); uint32_t i01 = rem - i02 * ne01; - uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01; + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half; uint8_t * d_spad = dst_spad_base + spad_idx * dst_spad_half; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, ne00 * sizeof(float), 0); - dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, ne00 * sizeof(float), current_block_size); + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0); + dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size); ir_prefetch += current_block_size; spad_idx ^= 1; } @@ -709,8 +882,8 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) { for (uint32_t ir = start_row; ir < end_row; ) { uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02); - uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src; - uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst; + uint8_t * d_spad = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst; uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div); uint32_t rem = ir - i03 * (ne02 * ne01); @@ -718,42 +891,37 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) { uint32_t i01 = rem - i02 * ne01; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - for (uint32_t r = 0; r < current_block_size; r++) { - uint32_t r_i01 = i01 + r; // linear within block since we split at ne01 - - const int32_t idx = *(int32_t *)((char *)src2->data + r_i01 * src2->nb[0] + i02 * src2->nb[1]); - - uint8_t * r_src1 = (uint8_t *)src1->data + idx * nb11; - uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned; - uint8_t * r_dst = d_spad + r * bctx->dst_row_size_aligned; - - hvx_add_f32_aau(r_dst, r_src0, r_src1, ne00); - } + compute_add_id_t compute = (compute_add_id_t) bctx->compute; + compute(d_spad, s0_spad, vtcm_src1, (const char *)(uintptr_t)src2->data, + i01, i02, src2->nb[0], src2->nb[1], src1_stride, + current_block_size, bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; - dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, ne00 * sizeof(float), current_block_size); + dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size); if (ir_prefetch < end_row) { uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02); - uint32_t p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); - uint32_t prem = ir_prefetch - p03 * (ne02 * ne01); - uint32_t p02 = fastdiv(prem, &bctx->src0_dim1_div); - uint32_t p01 = prem - p02 * ne01; - uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; - dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, ne00 * sizeof(float), next_block_size); + uint32_t p03, p02, p01, prem; + p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div); + prem = ir_prefetch - p03 * (ne02 * ne01); + p02 = fastdiv(prem, &bctx->src0_dim1_div); + p01 = prem - p02 * ne01; + dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01; + dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size); ir_prefetch += next_block_size; } ir += current_block_size; } - dma_queue_flush(q); + dma_queue_flush(dma_q); } static int execute_op_binary(struct htp_ops_context * octx) { const struct htp_tensor * src0 = octx->src[0]; const struct htp_tensor * src1 = octx->src[1]; const struct htp_tensor * dst = octx->dst; + const struct htp_binary_kernel_params * kparams = (const struct htp_binary_kernel_params *) octx->kernel_params; const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; @@ -770,7 +938,8 @@ static int execute_op_binary(struct htp_ops_context * octx) { if (octx->ctx->mdev.count > 1) { uint32_t rows_per_chunk = 0; htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_row_size, &rows_per_chunk); - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, + octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); row_start = range.start; nrows = range.count; } @@ -779,92 +948,57 @@ static int execute_op_binary(struct htp_ops_context * octx) { return HTP_STATUS_OK; } + if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) { + return HTP_STATUS_INVAL_PARAMS; + } + const uint32_t n_threads = octx->n_threads; + const size_t src0_row_size_aligned = kparams->src0_row_size_aligned; + const size_t src1_row_size_aligned = kparams->src1_row_size_aligned; + const size_t dst_row_size_aligned = kparams->dst_row_size_aligned; - size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN); - size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN); - size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN); - - bool is_add_id = (octx->op == HTP_OP_ADD_ID); - bool is_scalar = !is_add_id && (src1->ne[0] == 1); - - bool is_transposed = (src0->nb[1] < src0_row_size || src1->nb[1] < src1_row_size || dst->nb[1] < dst_row_size); - - bool is_same_shape = !is_add_id && !is_scalar && !is_transposed && - (src1->ne[0] == src0->ne[0] && src0->ne[0] % VLEN == 0) && - (src1->ne[1] == src0->ne[1] || src1->ne[1] == 1) && - (src1->ne[2] == src0->ne[2] || src1->ne[2] == 1) && - (src1->ne[3] == src0->ne[3] || src1->ne[3] == 1); - - bool is_row_bcast = is_same_shape && (src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1); - bool is_complex = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] == src0->ne[0]); - bool is_repeat = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] != src0->ne[0]); - - size_t spad_row_total; - if (is_same_shape) { - spad_row_total = 2 * (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned); - } else { - spad_row_total = 2 * (src0_row_size_aligned + dst_row_size_aligned); + if (htp_tensor_is_extended(src1)) { + if (kparams->kernel_type != HTP_BINARY_KERNEL_SAME_SHAPE && + kparams->kernel_type != HTP_BINARY_KERNEL_ROW_BCAST && + kparams->kernel_type != HTP_BINARY_KERNEL_SCALAR_DMA && + kparams->kernel_type != HTP_BINARY_KERNEL_ADD_ID) { + return HTP_STATUS_NO_SUPPORT; + } } - size_t rows_per_buffer = octx->ctx->vtcm_size / (n_threads * spad_row_total); - - // Adjust for static src1 in row_bcast case - if (is_row_bcast) { - size_t needed_static = src1_row_size_aligned; - if (octx->ctx->vtcm_size < needed_static) return HTP_STATUS_VTCM_TOO_SMALL; - size_t avail = octx->ctx->vtcm_size - needed_static; - rows_per_buffer = avail / (n_threads * spad_row_total); - } - - if (rows_per_buffer < 1) { - FARF(ERROR, "binary: VTCM too small\n"); - return HTP_STATUS_VTCM_TOO_SMALL; - } - - octx->src0_spad.size_per_thread = rows_per_buffer * 2 * src0_row_size_aligned; - octx->dst_spad.size_per_thread = rows_per_buffer * 2 * dst_row_size_aligned; - - if (is_add_id || is_scalar || is_complex || is_repeat || is_row_bcast) { - octx->src1_spad.size_per_thread = 0; - } else { - octx->src1_spad.size_per_thread = rows_per_buffer * 2 * src1_row_size_aligned; - } - - octx->dst_spad.size = n_threads * octx->dst_spad.size_per_thread; - octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread; - if (is_row_bcast) { - octx->src1_spad.size = src1_row_size_aligned; - } else { - octx->src1_spad.size = n_threads * octx->src1_spad.size_per_thread; - } - - if (octx->ctx->vtcm_size < (octx->src0_spad.size + octx->src1_spad.size + octx->dst_spad.size)) { - return HTP_STATUS_VTCM_TOO_SMALL; - } - - octx->src0_spad.data = octx->ctx->vtcm_base; octx->src0_spad.src = NULL; - octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL; - octx->dst_spad.data = octx->src1_spad.data + octx->src1_spad.size; octx->dst_spad.src = NULL; - - if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) { - return HTP_STATUS_OK; - } - - dma_queue * q = octx->ctx->dma[0]; - if (is_row_bcast) { - dma_queue_push(q, dma_make_ptr(octx->src1_spad.data, (const void *) src1->data), src1_row_size_aligned, 0, src1->ne[0] * elem_size, 1); + if (octx->op == HTP_OP_ADD_ID && htp_tensor_is_extended(octx->src[2])) { + return HTP_STATUS_NO_SUPPORT; } struct htp_binary_context bctx; + bctx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base; + htp_binary_vtcm_layout_build(&bctx.vtcm_layout, kparams, octx->ctx->vtcm_size); + + if (bctx.vtcm_layout.rows_per_buffer == 0 || bctx.vtcm_layout.total_bytes > octx->ctx->vtcm_size) { + return HTP_STATUS_VTCM_TOO_SMALL; + } + + dma_queue * dma_q = octx->ctx->dma[0]; + uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1); + if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) { + dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1); + } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR_DMA) { + dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[1] * elem_size, 1); + } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) { + dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), + kparams->src1_row_size_aligned, src1->nb[1], + src1->ne[0] * elem_size, src1->ne[1]); + } + bctx.octx = octx; bctx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); bctx.total_rows = nrows; bctx.row_start = row_start; - bctx.block_max = rows_per_buffer; + bctx.block_max = bctx.vtcm_layout.rows_per_buffer; bctx.src0_row_size_aligned = src0_row_size_aligned; bctx.src1_row_size_aligned = src1_row_size_aligned; bctx.dst_row_size_aligned = dst_row_size_aligned; + bctx.row_size_bytes = src0_row_size; bctx.src0_dim1_div = init_fastdiv_values(src0->ne[1]); bctx.src0_dim2_div = init_fastdiv_values(src0->ne[2]); @@ -880,19 +1014,153 @@ static int execute_op_binary(struct htp_ops_context * octx) { bool src0_contig_dim2 = (src0->nb[3] == src0->ne[2] * src0->nb[2]); bool dst_contig_dim2 = (dst->nb[3] == src0->ne[2] * dst->nb[2]); - bctx.split_at_ne01 = (src0->ne[2] > 1) && ((src1->ne[1] > 1) || (src1->ne[2] > 1) || !src0_contig_dim1 || !dst_contig_dim1); + bctx.split_at_ne01 = (octx->op == HTP_OP_ADD_ID) || + ((src0->ne[2] > 1) && ((src1->ne[1] > 1) || (src1->ne[2] > 1) || !src0_contig_dim1 || !dst_contig_dim1)); bctx.split_at_ne02 = (src0->ne[3] > 1) && ((src1->ne[2] > 1) || (src1->ne[3] > 1) || !src0_contig_dim2 || !dst_contig_dim2); - worker_callback_t worker_func; - if (is_add_id) worker_func = binary_job_add_id; - else if (is_scalar) worker_func = binary_job_scalar; - else if (is_row_bcast) worker_func = binary_job_vector_row_broadcast; - else if (is_same_shape) worker_func = binary_job_vector_same_shape; - else if (is_complex) worker_func = binary_job_vector_complex; - else worker_func = binary_job_element_repeat; + worker_callback_t worker_func = NULL; + void * compute_func = NULL; - if (is_row_bcast) { - dma_queue_pop(q); + switch (kparams->kernel_type) { + case HTP_BINARY_KERNEL_SAME_SHAPE: + worker_func = binary_thread_vector_same_shape; + if (src0_type == HTP_TYPE_F32) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_same_shape_add_f32; break; + case HTP_OP_SUB: compute_func = compute_same_shape_sub_f32; break; + case HTP_OP_MUL: compute_func = compute_same_shape_mul_f32; break; + case HTP_OP_DIV: compute_func = compute_same_shape_div_f32; break; + default: break; + } + } else if (src0_type == HTP_TYPE_F16) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_same_shape_add_f16; break; + case HTP_OP_SUB: compute_func = compute_same_shape_sub_f16; break; + case HTP_OP_MUL: compute_func = compute_same_shape_mul_f16; break; + case HTP_OP_DIV: compute_func = compute_same_shape_div_f16; break; + default: break; + } + } + break; + case HTP_BINARY_KERNEL_ROW_BCAST: + worker_func = binary_thread_vector_row_broadcast; + if (src0_type == HTP_TYPE_F32) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_row_bcast_add_f32; break; + case HTP_OP_SUB: compute_func = compute_row_bcast_sub_f32; break; + case HTP_OP_MUL: compute_func = compute_row_bcast_mul_f32; break; + case HTP_OP_DIV: compute_func = compute_row_bcast_div_f32; break; + default: break; + } + } else if (src0_type == HTP_TYPE_F16) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_row_bcast_add_f16; break; + case HTP_OP_SUB: compute_func = compute_row_bcast_sub_f16; break; + case HTP_OP_MUL: compute_func = compute_row_bcast_mul_f16; break; + case HTP_OP_DIV: compute_func = compute_row_bcast_div_f16; break; + default: break; + } + } + break; + case HTP_BINARY_KERNEL_SCALAR_DMA: + worker_func = binary_thread_scalar_dma; + if (src0_type == HTP_TYPE_F32) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_scalar_dma_add_f32; break; + case HTP_OP_SUB: compute_func = compute_scalar_dma_sub_f32; break; + case HTP_OP_MUL: compute_func = compute_scalar_dma_mul_f32; break; + case HTP_OP_DIV: compute_func = compute_scalar_dma_div_f32; break; + default: break; + } + } else if (src0_type == HTP_TYPE_F16) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_scalar_dma_add_f16; break; + case HTP_OP_SUB: compute_func = compute_scalar_dma_sub_f16; break; + case HTP_OP_MUL: compute_func = compute_scalar_dma_mul_f16; break; + case HTP_OP_DIV: compute_func = compute_scalar_dma_div_f16; break; + default: break; + } + } + break; + case HTP_BINARY_KERNEL_SCALAR: + worker_func = binary_thread_scalar; + if (src0_type == HTP_TYPE_F32) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_scalar_add_f32; break; + case HTP_OP_SUB: compute_func = compute_scalar_sub_f32; break; + case HTP_OP_MUL: compute_func = compute_scalar_mul_f32; break; + case HTP_OP_DIV: compute_func = compute_scalar_div_f32; break; + default: break; + } + } else if (src0_type == HTP_TYPE_F16) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_scalar_add_f16; break; + case HTP_OP_SUB: compute_func = compute_scalar_sub_f16; break; + case HTP_OP_MUL: compute_func = compute_scalar_mul_f16; break; + case HTP_OP_DIV: compute_func = compute_scalar_div_f16; break; + default: break; + } + } + break; + case HTP_BINARY_KERNEL_COMPLEX: + worker_func = binary_thread_vector_complex; + if (src0_type == HTP_TYPE_F32) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_complex_add_f32; break; + case HTP_OP_SUB: compute_func = compute_complex_sub_f32; break; + case HTP_OP_MUL: compute_func = compute_complex_mul_f32; break; + case HTP_OP_DIV: compute_func = compute_complex_div_f32; break; + default: break; + } + } else if (src0_type == HTP_TYPE_F16) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_complex_add_f16; break; + case HTP_OP_SUB: compute_func = compute_complex_sub_f16; break; + case HTP_OP_MUL: compute_func = compute_complex_mul_f16; break; + case HTP_OP_DIV: compute_func = compute_complex_div_f16; break; + default: break; + } + } + break; + case HTP_BINARY_KERNEL_REPEAT: + worker_func = binary_thread_element_repeat; + if (src0_type == HTP_TYPE_F32) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_repeat_add_f32; break; + case HTP_OP_SUB: compute_func = compute_repeat_sub_f32; break; + case HTP_OP_MUL: compute_func = compute_repeat_mul_f32; break; + case HTP_OP_DIV: compute_func = compute_repeat_div_f32; break; + default: break; + } + } else if (src0_type == HTP_TYPE_F16) { + switch (octx->op) { + case HTP_OP_ADD: compute_func = compute_repeat_add_f16; break; + case HTP_OP_SUB: compute_func = compute_repeat_sub_f16; break; + case HTP_OP_MUL: compute_func = compute_repeat_mul_f16; break; + case HTP_OP_DIV: compute_func = compute_repeat_div_f16; break; + default: break; + } + } + break; + case HTP_BINARY_KERNEL_ADD_ID: + if (octx->op == HTP_OP_ADD_ID && src0_type == HTP_TYPE_F32) { + worker_func = binary_thread_add_id_f32; + compute_func = (void *) compute_add_id_f32; + } + break; + default: break; + } + + if (!worker_func || !compute_func) { + return HTP_STATUS_NO_SUPPORT; + } + + bctx.compute = compute_func; + + if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST || + kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR_DMA || + kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) { + dma_queue_pop(dma_q); } work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads); diff --git a/ggml/src/ggml-hexagon/htp/binary-ops.h b/ggml/src/ggml-hexagon/htp/binary-ops.h new file mode 100644 index 000000000..b99f2ad64 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/binary-ops.h @@ -0,0 +1,111 @@ +#ifndef HTP_BINARY_OPS_H +#define HTP_BINARY_OPS_H + +#include +#include +#include + +#include "hex-common.h" +#include "htp-ops.h" +#include "htp-vtcm.h" + +enum htp_binary_kernel_type { + HTP_BINARY_KERNEL_SAME_SHAPE = 0, + HTP_BINARY_KERNEL_ROW_BCAST, + HTP_BINARY_KERNEL_SCALAR_DMA, + HTP_BINARY_KERNEL_SCALAR, + HTP_BINARY_KERNEL_ADD_ID, + HTP_BINARY_KERNEL_COMPLEX, + HTP_BINARY_KERNEL_REPEAT, +}; + +struct htp_binary_kernel_params { + uint32_t kernel_type; + uint32_t n_threads; + uint32_t rows_per_buffer; + + uint32_t src0_row_size_aligned; + uint32_t src1_row_size_aligned; + uint32_t dst_row_size_aligned; + + uint32_t src1_size; + uint32_t vtcm_size; +}; + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_binary_kernel_params) <= 128, "htp_binary_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_binary_kernel_params) <= 128, "htp_binary_kernel_params is too large for kernel_params blob"); +#endif + +struct htp_binary_vtcm_layout { + size_t total_bytes; + size_t off_src0; + size_t off_src1; + size_t off_dst; + + size_t src0_bytes_per_thread; + size_t src1_bytes_per_thread; + size_t dst_bytes_per_thread; + + size_t src0_spad_half_size; + size_t src1_spad_half_size; + size_t dst_spad_half_size; + + size_t src1_size; + uint32_t rows_per_buffer; +}; + +static inline void htp_binary_vtcm_layout_build( + struct htp_binary_vtcm_layout * L, + const struct htp_binary_kernel_params * kparams, + size_t vtcm_size +) { + memset(L, 0, sizeof(*L)); + + const uint32_t n_threads = kparams->n_threads; + if (n_threads == 0) { + return; + } + + const size_t spad_row_total = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) + ? 2 * (kparams->src0_row_size_aligned + kparams->src1_row_size_aligned + kparams->dst_row_size_aligned) + : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned); + + if (spad_row_total == 0 || vtcm_size < kparams->src1_size) { + return; + } + + const size_t rows_per_buffer = (vtcm_size - kparams->src1_size) / (n_threads * spad_row_total); + if (rows_per_buffer == 0) { + return; + } + + L->rows_per_buffer = (uint32_t) rows_per_buffer; + L->src1_size = kparams->src1_size; + + L->src0_bytes_per_thread = rows_per_buffer * 2 * kparams->src0_row_size_aligned; + L->dst_bytes_per_thread = rows_per_buffer * 2 * kparams->dst_row_size_aligned; + L->src1_bytes_per_thread = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) + ? rows_per_buffer * 2 * kparams->src1_row_size_aligned + : 0; + + L->src0_spad_half_size = L->src0_bytes_per_thread / 2; + L->src1_spad_half_size = L->src1_bytes_per_thread / 2; + L->dst_spad_half_size = L->dst_bytes_per_thread / 2; + + const size_t src0_total = n_threads * L->src0_bytes_per_thread; + const size_t src1_total = (kparams->src1_size > 0) + ? kparams->src1_size + : n_threads * L->src1_bytes_per_thread; + const size_t dst_total = n_threads * L->dst_bytes_per_thread; + + size_t off = 0; + VTCM_LAYOUT_ALLOC(off, off_src0, src0_total); + VTCM_LAYOUT_ALLOC(off, off_src1, src1_total); + VTCM_LAYOUT_ALLOC(off, off_dst, dst_total); + + L->total_bytes = off; +} + +#endif diff --git a/ggml/src/ggml-hexagon/htp/concat-ops.c b/ggml/src/ggml-hexagon/htp/concat-ops.c index 966e867b3..1fa6ec1bd 100644 --- a/ggml/src/ggml-hexagon/htp/concat-ops.c +++ b/ggml/src/ggml-hexagon/htp/concat-ops.c @@ -6,7 +6,7 @@ #include "hexagon_types.h" #include "hexagon_protos.h" #include "hvx_hexagon_protos.h" -#include "hex-dma.h" +#include "dma-queue.h" #include "htp-vtcm.h" #include "hvx-utils.h" #include "hex-fastdiv.h" @@ -41,7 +41,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * const uint32_t end_i = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end; if (start_i >= end_i) return; - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint8_t * spad0_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; uint8_t * spad1_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread; @@ -64,14 +64,14 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i; uint32_t src1_width_bytes = current_block_i * sizeof(float); - uint8_t * src1_ptr = (uint8_t *)src1->data + i * src1->nb[1]; - dma_queue_push(q, dma_make_ptr(spad1_base, src1_ptr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); + const dma_addr_t src1_addr = src1->data + i * src1->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); uint32_t src0_row_bytes = src0_ne0 * sizeof(float); - uint8_t * src0_ptr = (uint8_t *)src0->data + i * src0->nb[1]; - dma_queue_push(q, dma_make_ptr(spad0_base, src0_ptr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); + const dma_addr_t src0_addr = src0->data + i * src0->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); - dma_queue_pop(q); // src1 + dma_queue_pop(dma_q); // src1 HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride); @@ -87,12 +87,12 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - dma_queue_pop(q); // src0 + dma_queue_pop(dma_q); // src0 - uint8_t * dst_ptr = (uint8_t *)dst->data + i * dst->nb[1]; - dma_queue_push(q, dma_make_ptr(dst_ptr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i); + const dma_addr_t dst_addr = dst->data + i * dst->nb[1]; + dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i); - dma_queue_pop(q); + dma_queue_pop(dma_q); } } @@ -112,7 +112,7 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void * const uint32_t end_i = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end; if (start_i >= end_i) return; - dma_queue * q = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; uint8_t * spad0_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; uint8_t * spad1_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread; @@ -135,14 +135,14 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void * uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i; uint32_t src1_width_bytes = current_block_i * sizeof(__fp16); - uint8_t * src1_ptr = (uint8_t *)src1->data + i * src1->nb[1]; - dma_queue_push(q, dma_make_ptr(spad1_base, src1_ptr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); + const dma_addr_t src1_addr = src1->data + i * src1->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0); uint32_t src0_row_bytes = src0_ne0 * sizeof(__fp16); - uint8_t * src0_ptr = (uint8_t *)src0->data + i * src0->nb[1]; - dma_queue_push(q, dma_make_ptr(spad0_base, src0_ptr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); + const dma_addr_t src0_addr = src0->data + i * src0->nb[1]; + dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i); - dma_queue_pop(q); // src1 + dma_queue_pop(dma_q); // src1 HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride); @@ -158,12 +158,12 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void * } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - dma_queue_pop(q); // src0 + dma_queue_pop(dma_q); // src0 - uint8_t * dst_ptr = (uint8_t *)dst->data + i * dst->nb[1]; - dma_queue_push(q, dma_make_ptr(dst_ptr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i); + const dma_addr_t dst_addr = dst->data + i * dst->nb[1]; + dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i); - dma_queue_pop(q); + dma_queue_pop(dma_q); } } @@ -304,6 +304,10 @@ int op_concat(struct htp_ops_context * octx) { worker_func = concat_2d_f16_transposed; } } else { + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } + const uint32_t total_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3]; uint32_t elem_start = 0; uint32_t nelems = total_elements; diff --git a/ggml/src/ggml-hexagon/htp/cpy-ops.c b/ggml/src/ggml-hexagon/htp/cpy-ops.c index 490efd687..4453dda3a 100644 --- a/ggml/src/ggml-hexagon/htp/cpy-ops.c +++ b/ggml/src/ggml-hexagon/htp/cpy-ops.c @@ -49,6 +49,30 @@ struct htp_copy_context { struct fastdiv_values div_ne02_ne01_ne00; }; +static inline void cpy_dma_sametype_reshape_contig( + dma_queue * dma_q, + dma_addr_t dst, + dma_addr_t src0, + uint32_t total_bytes +) { + if (total_bytes == 0) { + return; + } + + const uint32_t max_chunk = DMA_SAFE_CHUNK_SIZE; + while (total_bytes > 0) { + const uint32_t chunk = MIN(total_bytes, max_chunk); + if (!dma_queue_push(dma_q, dma_make_data(dst, src0), chunk, chunk, chunk, /*nrows=*/ 1)) { + dma_queue_flush(dma_q); + dma_queue_push(dma_q, dma_make_data(dst, src0), chunk, chunk, chunk, /*nrows=*/ 1); + } + dst += chunk; + src0 += chunk; + total_bytes -= chunk; + } + dma_queue_flush(dma_q); +} + #define cpy_preamble \ const struct htp_tensor *src0 = octx->src[0]; \ const struct htp_tensor *dst = octx->dst; \ @@ -73,129 +97,131 @@ struct htp_copy_context { const uint32_t nb2 = dst->nb[2]; \ const uint32_t nb3 = dst->nb[3]; -#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ -static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_copy_context * ct = (struct htp_copy_context *) data; \ - struct htp_ops_context * octx = ct->octx; \ - cpy_preamble; \ - const uint32_t dr = ct->src0_nrows_per_thread; \ - const uint32_t ir0 = ct->row_start + dr * ith; \ - const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows); \ - if (ir0 >= ir1) return; \ - const bool contiguous = (nb01 == ne00 * ELEM_SIZE) && (nb1 == nb01) && \ - (nb02 == ne01 * nb01) && (nb2 == nb02) && \ - (nb03 == ne02 * nb02) && (nb3 == nb03); \ - const uint32_t ne02_ne01 = ne02 * ne01; \ - uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01); \ - uint32_t rem = ir0 - i03 * ne02_ne01; \ - uint32_t i02 = fastdiv(rem, &ct->div_ne01); \ - uint32_t i01 = rem - i02 * ne01; \ - uint8_t * dst_ptr = (uint8_t *) dst->data + i01*nb1 + i02*nb2 + i03*nb3; \ - uint8_t * src0_ptr = (uint8_t *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; \ - if (contiguous) { \ - hvx_copy_uu(dst_ptr, src0_ptr, (ir1 - ir0) * ne00, ELEM_SIZE); \ - return; \ - } \ - for (uint32_t r = ir0; r < ir1; r++) { \ - hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2); \ - hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE); \ - dst_ptr += nb1; \ - src0_ptr += nb01; \ - if (++i01 == ne01) { \ - i01 = 0; \ - if (++i02 == ne02) { \ - i02 = 0; \ - i03++; \ - } \ - dst_ptr = (uint8_t *) dst->data + i02*nb2 + i03*nb3; \ - src0_ptr = (uint8_t *) src0->data + i02*nb02 + i03*nb03; \ - } \ - } \ +#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ +static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_copy_context * ct = (struct htp_copy_context *) data; \ + struct htp_ops_context * octx = ct->octx; \ + cpy_preamble; \ + const uint32_t dr = ct->src0_nrows_per_thread; \ + const uint32_t ir0 = ct->row_start + dr * ith; \ + const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows); \ + if (ir0 >= ir1) return; \ + const bool contiguous = htp_tensor_is_contiguous(src0, ELEM_SIZE) && htp_tensor_is_contiguous(dst, ELEM_SIZE); \ + if (contiguous) { \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ + dma_addr_t dst_addr = dst->data + ir0 * ne00 * ELEM_SIZE; \ + dma_addr_t src0_addr = src0->data + ir0 * ne00 * ELEM_SIZE; \ + cpy_dma_sametype_reshape_contig(dma_q, dst_addr, src0_addr, (ir1 - ir0) * ne00 * ELEM_SIZE); \ + return; \ + } \ + const uint32_t ne02_ne01 = ne02 * ne01; \ + uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01); \ + uint32_t rem = ir0 - i03 * ne02_ne01; \ + uint32_t i02 = fastdiv(rem, &ct->div_ne01); \ + uint32_t i01 = rem - i02 * ne01; \ + uint8_t * dst_ptr = (uint8_t *) dst->data + i01*nb1 + i02*nb2 + i03*nb3; \ + uint8_t * src0_ptr = (uint8_t *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; \ + for (uint32_t r = ir0; r < ir1; r++) { \ + hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2); \ + hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE); \ + dst_ptr += nb1; \ + src0_ptr += nb01; \ + if (++i01 == ne01) { \ + i01 = 0; \ + if (++i02 == ne02) { \ + i02 = 0; \ + i03++; \ + } \ + dst_ptr = (uint8_t *) dst->data + i02*nb2 + i03*nb3; \ + src0_ptr = (uint8_t *) src0->data + i02*nb02 + i03*nb03; \ + } \ + } \ } DEFINE_CPY_SAMESHAPE(f32, float, 4) DEFINE_CPY_SAMESHAPE(f16, __fp16, 2) -#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ -static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_copy_context * ct = (struct htp_copy_context *) data; \ - struct htp_ops_context * octx = ct->octx; \ - cpy_preamble; \ - const uint32_t th_nelem = ct->elem_per_thread; \ - const uint32_t th_start = ct->elem_start + ith * th_nelem; \ - const uint32_t th_end = MIN(th_start + th_nelem, ct->elem_start + ct->nelem); \ - if (th_start >= th_end) return; \ - \ - if (htp_tensor_is_contiguous(src0, ELEM_SIZE) && htp_tensor_is_contiguous(dst, ELEM_SIZE)) { \ - hvx_copy_uu((uint8_t *) dst->data + (size_t) th_start * ELEM_SIZE, \ - (const uint8_t *) src0->data + (size_t) th_start * ELEM_SIZE, \ - th_end - th_start, ELEM_SIZE); \ - return; \ - } \ - \ - const uint32_t ne01_ne00 = ne01 * ne00; \ - const uint32_t ne02_ne01_ne00 = ne02 * ne01_ne00; \ - const uint32_t ne1_ne0 = ne1 * ne0; \ - const uint32_t ne2_ne1_ne0 = ne2 * ne1_ne0; \ - \ - uint32_t e = th_start; \ - uint32_t i13 = fastdiv(e, &ct->div_ne2_ne1_ne0); \ - uint32_t rem = e - i13 * ne2_ne1_ne0; \ - uint32_t i12 = fastdiv(rem, &ct->div_ne1_ne0); \ - uint32_t rem2 = rem - i12 * ne1_ne0; \ - uint32_t i11 = fastdiv(rem2, &ct->div_ne0); \ - uint32_t i10 = rem2 - i11 * ne0; \ - \ - uint32_t i03 = fastdiv(e, &ct->div_ne02_ne01_ne00); \ - uint32_t rem_s = e - i03 * ne02_ne01_ne00; \ - uint32_t i02 = fastdiv(rem_s, &ct->div_ne01_ne00); \ - uint32_t rem2_s = rem_s - i02 * ne01_ne00; \ - uint32_t i01 = fastdiv(rem2_s, &ct->div_ne00); \ - uint32_t i00 = rem2_s - i01 * ne00; \ - \ - char * dst_ptr = (char *) dst->data + i10*nb0 + i11*nb1 + i12*nb2 + i13*nb3; \ - const char * src0_ptr = (const char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03; \ - \ - const bool rows_contig = (nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE); \ - \ - while (e < th_end) { \ - uint32_t run = 1; \ - if (rows_contig) { \ - run = MIN(MIN(ne00 - i00, ne0 - i10), th_end - e); \ - hvx_copy_uu((uint8_t *) dst_ptr, (const uint8_t *) src0_ptr, run, ELEM_SIZE); \ - } else { \ - *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr); \ - } \ - e += run; \ - \ - dst_ptr += run * nb0; \ - i10 += run; \ - if (i10 == ne0) { \ - i10 = 0; \ - if (++i11 == ne1) { \ - i11 = 0; \ - if (++i12 == ne2) { \ - i12 = 0; \ - i13++; \ - } \ - } \ - dst_ptr = (char *) dst->data + i11*nb1 + i12*nb2 + i13*nb3; \ - } \ - \ - src0_ptr += run * nb00; \ - i00 += run; \ - if (i00 == ne00) { \ - i00 = 0; \ - if (++i01 == ne01) { \ - i01 = 0; \ - if (++i02 == ne02) { \ - i02 = 0; \ - i03++; \ - } \ - } \ - src0_ptr = (const char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; \ - } \ - } \ +#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ +static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_copy_context * ct = (struct htp_copy_context *) data; \ + struct htp_ops_context * octx = ct->octx; \ + cpy_preamble; \ + const uint32_t th_nelem = ct->elem_per_thread; \ + const uint32_t th_start = ct->elem_start + ith * th_nelem; \ + const uint32_t th_end = MIN(th_start + th_nelem, ct->elem_start + ct->nelem); \ + if (th_start >= th_end) return; \ + \ + if (htp_tensor_is_contiguous(src0, ELEM_SIZE) && htp_tensor_is_contiguous(dst, ELEM_SIZE)) { \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ + dma_addr_t dst_addr = dst->data + th_start * ELEM_SIZE; \ + dma_addr_t src0_addr = src0->data + th_start * ELEM_SIZE; \ + cpy_dma_sametype_reshape_contig(dma_q, dst_addr, src0_addr, (th_end - th_start) * ELEM_SIZE); \ + return; \ + } \ + \ + const uint32_t ne01_ne00 = ne01 * ne00; \ + const uint32_t ne02_ne01_ne00 = ne02 * ne01_ne00; \ + const uint32_t ne1_ne0 = ne1 * ne0; \ + const uint32_t ne2_ne1_ne0 = ne2 * ne1_ne0; \ + \ + uint32_t e = th_start; \ + uint32_t i13 = fastdiv(e, &ct->div_ne2_ne1_ne0); \ + uint32_t rem = e - i13 * ne2_ne1_ne0; \ + uint32_t i12 = fastdiv(rem, &ct->div_ne1_ne0); \ + uint32_t rem2 = rem - i12 * ne1_ne0; \ + uint32_t i11 = fastdiv(rem2, &ct->div_ne0); \ + uint32_t i10 = rem2 - i11 * ne0; \ + \ + uint32_t i03 = fastdiv(e, &ct->div_ne02_ne01_ne00); \ + uint32_t rem_s = e - i03 * ne02_ne01_ne00; \ + uint32_t i02 = fastdiv(rem_s, &ct->div_ne01_ne00); \ + uint32_t rem2_s = rem_s - i02 * ne01_ne00; \ + uint32_t i01 = fastdiv(rem2_s, &ct->div_ne00); \ + uint32_t i00 = rem2_s - i01 * ne00; \ + \ + char * dst_ptr = (char *) dst->data + i10*nb0 + i11*nb1 + i12*nb2 + i13*nb3; \ + const char * src0_ptr = (const char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03; \ + \ + const bool rows_contig = (nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE); \ + \ + while (e < th_end) { \ + uint32_t run = 1; \ + if (rows_contig) { \ + run = MIN(MIN(ne00 - i00, ne0 - i10), th_end - e); \ + hvx_copy_uu((uint8_t *) dst_ptr, (const uint8_t *) src0_ptr, run, ELEM_SIZE); \ + } else { \ + *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr); \ + } \ + e += run; \ + \ + dst_ptr += run * nb0; \ + i10 += run; \ + if (i10 == ne0) { \ + i10 = 0; \ + if (++i11 == ne1) { \ + i11 = 0; \ + if (++i12 == ne2) { \ + i12 = 0; \ + i13++; \ + } \ + } \ + dst_ptr = (char *) dst->data + i11*nb1 + i12*nb2 + i13*nb3; \ + } \ + \ + src0_ptr += run * nb00; \ + i00 += run; \ + if (i00 == ne00) { \ + i00 = 0; \ + if (++i01 == ne01) { \ + i01 = 0; \ + if (++i02 == ne02) { \ + i02 = 0; \ + i03++; \ + } \ + } \ + src0_ptr = (const char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; \ + } \ + } \ } DEFINE_CPY_RESHAPE(f32, float, 4) @@ -273,6 +299,27 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi } } +static inline void cpy_dma_push_2d_chunked( + dma_queue * dma_q, + dma_addr_t dst, + dma_addr_t src, + size_t dst_stride, + size_t src_stride, + size_t row_size, + uint32_t nrows +) { + while (nrows > 0) { + const uint32_t cur_rows = MIN(nrows, DMA_MAX_NROWS); + if (!dma_queue_push(dma_q, dma_make_data(dst, src), dst_stride, src_stride, row_size, cur_rows)) { + dma_queue_flush(dma_q); + dma_queue_push(dma_q, dma_make_data(dst, src), dst_stride, src_stride, row_size, cur_rows); + } + dst += cur_rows * dst_stride; + src += cur_rows * src_stride; + nrows -= cur_rows; + } +} + static inline void cpy_dma_sametype_sameshape( struct htp_ops_context * octx, const struct htp_tensor * dst, @@ -282,46 +329,35 @@ static inline void cpy_dma_sametype_sameshape( uint32_t nb01, uint32_t nb02, uint32_t nb03, uint32_t nb1, uint32_t nb2, uint32_t nb3 ) { + const bool contiguous = htp_tensor_is_contiguous(src0, elem_size) && htp_tensor_is_contiguous(dst, elem_size); + + dma_queue * dma_q = octx->ctx->dma[0]; + + if (contiguous) { + cpy_dma_sametype_reshape_contig(dma_q, dst->data, src0->data, ne00 * elem_size * ne01 * ne02 * ne03); + return; + } + const bool contiguous_outer = (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) && (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2)); - dma_queue * q = octx->ctx->dma[0]; - if (contiguous_outer) { - if (!dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03)) { - dma_queue_flush(q); - dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03); - } - dma_queue_flush(q); + uint32_t total_rows = ne01 * ne02 * ne03; + cpy_dma_push_2d_chunked(dma_q, dst->data, src0->data, nb1, nb01, ne00 * elem_size, total_rows); + dma_queue_flush(dma_q); return; } for (uint32_t i03 = 0; i03 < ne03; i03++) { for (uint32_t i02 = 0; i02 < ne02; i02++) { - uint8_t * dst_ptr = (uint8_t *) dst->data + i02 * nb2 + i03 * nb3; - uint8_t * src0_ptr = (uint8_t *) src0->data + i02 * nb02 + i03 * nb03; - - if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) { - dma_queue_flush(q); - dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01); - } + dma_addr_t dst_data = dst->data + i02 * nb2 + i03 * nb3; + dma_addr_t src0_data = src0->data + i02 * nb02 + i03 * nb03; + cpy_dma_push_2d_chunked(dma_q, dst_data, src0_data, nb1, nb01, ne00 * elem_size, ne01); } } - dma_queue_flush(q); -} - -static inline void cpy_dma_sametype_reshape_contig( - struct htp_ops_context * octx, - const struct htp_tensor * dst, - const struct htp_tensor * src0, - uint32_t total_bytes -) { - dma_queue * q = octx->ctx->dma[0]; - dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), - total_bytes, total_bytes, total_bytes, /*nrows=*/ 1); - dma_queue_pop(q); + dma_queue_flush(dma_q); } static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { @@ -345,10 +381,6 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } - const bool sametype = (src0->type == dst->type); const bool transposed = (nb00 > nb01) || (nb0 > nb1) || (nb00 != ct.src0_type_size) || (nb0 != ct.dst_type_size) || @@ -360,6 +392,15 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { const bool src_is_contiguous = htp_tensor_is_contiguous(src0, ct.src0_type_size); const bool dst_is_contiguous = htp_tensor_is_contiguous(dst, ct.dst_type_size); + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + if (!sametype) { + return HTP_STATUS_NO_SUPPORT; + } + if (!sameshape && !(src_is_contiguous && dst_is_contiguous && octx->ctx->mdev.count <= 1)) { + return HTP_STATUS_NO_SUPPORT; + } + } + if (sameshape) { const uint32_t total_rows = ne01 * ne02 * ne03; const uint32_t row_size = ne00 * ct.dst_type_size; @@ -373,7 +414,8 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { if (octx->ctx->mdev.count > 1) { const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1; const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous; - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, can_split ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, can_split ? rows_per_chunk : 0, + octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); row_start = range.start; nrows = range.count; } @@ -386,9 +428,11 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { ct.nrows = nrows; ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); - if (sametype && octx->ctx->mdev.count <= 1) { - *use_dma = true; - cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3); + if (sametype && (octx->ctx->mdev.count <= 1 || htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) { + if (octx->ctx->mdev.idx == 0) { + *use_dma = true; + cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3); + } } else { work_queue_func_t copy_fun = NULL; if (sametype) { @@ -408,7 +452,7 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { if (octx->ctx->mdev.count <= 1 && dst_is_contiguous && src_is_contiguous) { *use_dma = true; - cpy_dma_sametype_reshape_contig(octx, dst, src0, total_elems * ct.dst_type_size); + cpy_dma_sametype_reshape_contig(octx->ctx->dma[0], dst->data, src0->data, total_elems * ct.dst_type_size); return HTP_STATUS_OK; } @@ -424,7 +468,8 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) { if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous; - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elems, can_split ? elems_per_line : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elems, can_split ? elems_per_line : 0, + octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); elem_start = range.start; nelem = range.count; } @@ -461,6 +506,9 @@ int op_cpy(struct htp_ops_context * octx) { if (octx->ctx->mdev.idx == 0) { const struct htp_tensor * sync = octx->src[1]; + if (htp_tensor_is_extended(sync)) { + return HTP_STATUS_NO_SUPPORT; + } const uint32_t seq = (uint32_t) octx->op_params[0]; atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data; htp_fence_write(sync_fence, seq, octx->status); diff --git a/ggml/src/ggml-hexagon/htp/cumsum-ops.c b/ggml/src/ggml-hexagon/htp/cumsum-ops.c index 971fa3bcc..eaab7d7e5 100644 --- a/ggml/src/ggml-hexagon/htp/cumsum-ops.c +++ b/ggml/src/ggml-hexagon/htp/cumsum-ops.c @@ -14,7 +14,7 @@ #include "htp-tensor.h" #include "hvx-types.h" #include "hvx-utils.h" -#include "hex-dma.h" +#include "dma-queue.h" #define htp_cumsum_tensors_preamble \ const struct htp_tensor * restrict src0 = octx->src[0]; \ @@ -55,7 +55,7 @@ struct htp_cumsum_context { struct htp_cumsum_context * cctx = (struct htp_cumsum_context *) data; \ struct htp_ops_context * octx = cctx->octx; \ htp_cumsum_tensors_preamble; \ - dma_queue * dma_queue = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; // --------------------------------------------------------------------------- // HVX prefix scan helpers @@ -131,47 +131,47 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat const size_t src_row_size_aligned = cctx->src_row_size_aligned; const size_t dst_row_size_aligned = cctx->dst_row_size_aligned; - const uint8_t * src_data = (const uint8_t *) src0->data; - uint8_t * dst_data = (uint8_t *) dst->data; + const dma_addr_t src_data = src0->data; + const dma_addr_t dst_data = dst->data; uint8_t * src_spad = octx->src0_spad.data + (ith * src_row_size_aligned * 2); uint8_t * dst_spad = octx->dst_spad.data + (ith * dst_row_size_aligned * 2); for (uint32_t ir = ir0, spad_idx = 0; ir < ir1 && spad_idx < 2; ir++, spad_idx++) { // Dummy dst writeback to establish queue ordering - dma_queue_push_vtcm_to_ddr(dma_queue, - dma_make_ptr(dst_data, dst_spad + (spad_idx * dst_row_size_aligned)), - dst_row_size, dst_row_size_aligned, 0); + dma_queue_push(dma_q, + dma_make_data(dst_data, dst_spad + (spad_idx * dst_row_size_aligned)), + dst_row_size, dst_row_size_aligned, dst_row_size, 0); - dma_queue_push_ddr_to_vtcm(dma_queue, - dma_make_ptr(src_spad + (spad_idx * src_row_size_aligned), - src_data + (ir * src_row_size)), - src_row_size_aligned, src_row_size, 1); + dma_queue_push(dma_q, + dma_make_data(src_spad + (spad_idx * src_row_size_aligned), + src_data + (ir * src_row_size)), + src_row_size_aligned, src_row_size, src_row_size, 1); } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; for (uint32_t ir = ir0; ir < ir1; ir++) { - float * dst_spad_row = (float *) dma_queue_pop(dma_queue).src; - float * src_spad_row = (float *) dma_queue_pop(dma_queue).dst; + float * dst_spad_row = (float *) dma_queue_pop(dma_q).src; + float * src_spad_row = (float *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); hvx_cumsum_row_f32(src_spad_row, dst_spad_row, ne00); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - dma_queue_push_vtcm_to_ddr(dma_queue, - dma_make_ptr(dst_data + (ir * dst_row_size), (uint8_t *) dst_spad_row), - dst_row_size, dst_row_size_aligned, 1); + dma_queue_push(dma_q, + dma_make_data(dst_data + (ir * dst_row_size), dst_spad_row), + dst_row_size, dst_row_size_aligned, dst_row_size, 1); const uint32_t next_row = ir + 2; if (next_row < ir1) { - dma_queue_push_ddr_to_vtcm(dma_queue, - dma_make_ptr((uint8_t *) src_spad_row, src_data + (next_row * src_row_size)), - src_row_size_aligned, src_row_size, 1); + dma_queue_push(dma_q, + dma_make_data(src_spad_row, src_data + (next_row * src_row_size)), + src_row_size_aligned, src_row_size, src_row_size, 1); } } - dma_queue_flush(dma_queue); + dma_queue_flush(dma_q); FARF(HIGH, "cumsum-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n", ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1, @@ -211,10 +211,6 @@ int op_cumsum_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_rows = src0->ne[1] * src0->ne[2] * src0->ne[3]; const size_t dst_data_row_size = dst->ne[0] * sizeof(float); @@ -264,6 +260,9 @@ int op_cumsum_f32(struct htp_ops_context * octx) { }; if (octx->ctx->vtcm_size < spad_per_thread * n_threads) { + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads); } else { work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads); diff --git a/ggml/src/ggml-hexagon/htp/diag-ops.c b/ggml/src/ggml-hexagon/htp/diag-ops.c index a69fd89d3..162214d3e 100644 --- a/ggml/src/ggml-hexagon/htp/diag-ops.c +++ b/ggml/src/ggml-hexagon/htp/diag-ops.c @@ -13,7 +13,7 @@ #include "hvx-types.h" #include "hex-utils.h" #include "hvx-copy.h" -#include "hex-dma.h" +#include "dma-queue.h" #define htp_diag_tensors_preamble \ const struct htp_tensor * restrict src0 = octx->src[0]; \ @@ -59,7 +59,7 @@ static inline void hvx_diag_row_f32(const float * restrict src, float * restrict 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]; + dma_queue * dma_q = octx->ctx->dma[ith]; const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith; const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches); @@ -73,8 +73,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) 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; + const dma_addr_t src_data = src0->data; + const dma_addr_t dst_data = 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); @@ -86,13 +86,13 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) const uint32_t i3 = ib / ne02; const uint32_t i2 = ib % ne02; - const uint8_t * src_batch = src_data + i3 * nb03 + i2 * nb02; + const dma_addr_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); + dma_queue_push(dma_q, + dma_make_data(src_spad, src_batch), + src_batch_size_aligned, src_batch_size, src_batch_size, 1); + dma_queue_flush(dma_q); const float * src_spad_f32 = (const float *) src_spad; float * dst_spad_f32 = (float *) dst_spad; @@ -104,11 +104,11 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1)); // 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); + const dma_addr_t dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1; + dma_queue_push(dma_q, + dma_make_data(dst_row, dst_spad), + dst_row_size, dst_row_size_aligned, dst_row_size, 1); + dma_queue_flush(dma_q); } } @@ -156,10 +156,6 @@ 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 size_t dst_batch_size = dst->ne[1] * dst->nb[1]; @@ -221,6 +217,9 @@ int op_diag_f32(struct htp_ops_context * octx) { }; if (octx->ctx->vtcm_size < spad_per_thread * n_threads) { + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads); } else { work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads); diff --git a/ggml/src/ggml-hexagon/htp/dma-queue.c b/ggml/src/ggml-hexagon/htp/dma-queue.c index 4beded1de..464e4b849 100644 --- a/ggml/src/ggml-hexagon/htp/dma-queue.c +++ b/ggml/src/ggml-hexagon/htp/dma-queue.c @@ -22,58 +22,69 @@ static inline uintptr_t align_up(uintptr_t addr, size_t align) { return (addr + align - 1) & ~(align - 1); } -size_t dma_queue_sizeof(size_t capacity) { +static inline size_t dma_ring_sizeof(size_t capacity) { capacity = pow2_ceil(capacity); - size_t size_q = sizeof(dma_queue); - size_t offset_r = align_up(size_q, HEX_L2_LINE_SIZE); size_t size_r = sizeof(dma_ring); - size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE); + size_t offset_desc = align_up(size_r, HEX_L2_LINE_SIZE); size_t size_desc = capacity * sizeof(dma_descriptor_2d); - size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE); - size_t size_dptr = capacity * sizeof(dma_ptr); + size_t offset_data = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE); + size_t size_data = capacity * sizeof(dma_data); - return offset_dptr + size_dptr; + return offset_data + size_data; +} + +static inline dma_ring * dma_ring_init(void * ptr, size_t capacity, struct htp_thread_trace * trace) { + capacity = pow2_ceil(capacity); + + size_t size_r = sizeof(dma_ring); + size_t offset_desc = align_up(size_r, HEX_L2_LINE_SIZE); + size_t size_desc = capacity * sizeof(dma_descriptor_2d); + size_t offset_data = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE); + + dma_ring * r = (dma_ring *) ptr; + r->trace = trace; + r->capacity = capacity; + r->idx_mask = capacity - 1; + r->push_idx = 0; + r->pop_idx = 0; + r->desc = (dma_descriptor_2d *) ((uintptr_t) ptr + offset_desc); + r->data = (dma_data *) ((uintptr_t) ptr + offset_data); + r->tail = &r->desc[capacity - 1]; + + return r; +} + +size_t dma_queue_sizeof(size_t capacity) { + size_t size_q = sizeof(dma_queue); + size_t offset_r0 = align_up(size_q, HEX_L2_LINE_SIZE); + size_t size_r0 = dma_ring_sizeof(capacity); + size_t offset_r1 = align_up(offset_r0 + size_r0, HEX_L2_LINE_SIZE); + size_t size_r1 = dma_ring_sizeof(DMA_FALLBACK_CAPACITY); + + return offset_r1 + size_r1; } size_t dma_queue_alignof(void) { return HEX_L2_LINE_SIZE; } -dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace) { - capacity = pow2_ceil(capacity); - - size_t size_q = sizeof(dma_queue); - size_t offset_r = align_up(size_q, HEX_L2_LINE_SIZE); - size_t size_r = sizeof(dma_ring); - size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE); - size_t size_desc = capacity * sizeof(dma_descriptor_2d); - size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE); - size_t size_dptr = capacity * sizeof(dma_ptr); - - size_t total_size = offset_dptr + size_dptr; +dma_queue_t dma_queue_init(void * ptr, size_t capacity, struct htp_thread_trace * trace) { + size_t total_size = dma_queue_sizeof(capacity); memset(ptr, 0, total_size); dma_queue * q = (dma_queue *) ptr; - dma_ring * r = (dma_ring *) ((uintptr_t) ptr + offset_r); - q->ring = r; - q->nocache = 0; - q->alias = false; + size_t size_q = sizeof(dma_queue); + size_t offset_r0 = align_up(size_q, HEX_L2_LINE_SIZE); + size_t size_r0 = dma_ring_sizeof(capacity); + size_t offset_r1 = align_up(offset_r0 + size_r0, HEX_L2_LINE_SIZE); - r->trace = trace; - r->vtcm_base = vtcm_base; - r->vtcm_end = vtcm_base + vtcm_size; - r->capacity = capacity; - r->idx_mask = capacity - 1; - r->push_idx = 0; - r->pop_idx = 0; + q->ring0 = dma_ring_init((void *) ((uintptr_t) ptr + offset_r0), capacity, trace); + q->ring1 = dma_ring_init((void *) ((uintptr_t) ptr + offset_r1), DMA_FALLBACK_CAPACITY, trace); + q->alias = false; - r->desc = (dma_descriptor_2d *) ((uintptr_t) ptr + offset_desc); - r->dptr = (dma_ptr *) ((uintptr_t) ptr + offset_dptr); - r->tail = &r->desc[capacity - 1]; - - FARF(HIGH, "dma-queue: capacity %u, unified memory size %zu\n", capacity, total_size); + FARF(HIGH, "dma-queue: capacity %u, unified memory size %zu\n", (unsigned) capacity, total_size); return q; } @@ -86,13 +97,13 @@ size_t dma_queue_alias_sizeof(void) { return sizeof(dma_queue); } -dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache) { +dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q) { dma_queue * q = (dma_queue *) ptr; memset(q, 0, sizeof(dma_queue)); - q->ring = main_q->ring; - q->nocache = nocache; - q->alias = true; + q->ring0 = main_q->ring0; + q->ring1 = main_q->ring1; + q->alias = true; return q; } @@ -101,4 +112,101 @@ void dma_queue_alias_free(dma_queue_t q) { (void) q; } +bool dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { + dma_ring * r0 = q->ring0; + dma_ring * r1 = q->ring1; + if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) { + return false; + } + + r1->tail = r0->tail; + + size_t rem_rows = nrows; + dma_addr_t cur_dst = ddata.dst; + dma_addr_t cur_src = ddata.src; + + while (rem_rows > 0) { + const uint32_t cur_rows = MIN(rem_rows, DMA_MAX_NROWS); + dma_data cur_data = dma_make_data(cur_dst, cur_src); + if (!dma_ring_push_single_2d(r1, cur_data, dst_stride, src_stride, row_size, cur_rows)) { + dma_ring_flush(r1); + dma_ring_push_single_2d(r1, cur_data, dst_stride, src_stride, row_size, cur_rows); + } + cur_dst += cur_rows * dst_stride; + cur_src += cur_rows * src_stride; + rem_rows -= cur_rows; + } + + dma_ring_flush(r1); + r0->tail = r1->tail; + + return dma_ring_push_single_2d(r0, ddata, 0, 0, 0, /*nrows=*/ 0); +} + +bool dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total) { + dma_ring * r0 = q->ring0; + dma_ring * r1 = q->ring1; + + if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) { + return false; + } + + r1->tail = r0->tail; + + size_t rem_bytes = total; + dma_addr_t cur_dst = ddata.dst; + dma_addr_t cur_src = ddata.src; + + while (rem_bytes > 0) { + const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE); + dma_data cur_data = dma_make_data(cur_dst, cur_src); + if (!dma_ring_push_single_1d(r1, cur_data, cur_bytes)) { + dma_ring_flush(r1); + dma_ring_push_single_1d(r1, cur_data, cur_bytes); + } + cur_dst += cur_bytes; + cur_src += cur_bytes; + rem_bytes -= cur_bytes; + } + + dma_ring_flush(r1); + r0->tail = r1->tail; + + return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0); +} + +#if __HVX_ARCH__ < 75 + +bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { + dma_ring * r0 = q->ring0; + dma_ring * r1 = q->ring1; + + if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) { + return false; + } + + r1->tail = r0->tail; + + size_t rem_rows = nrows; + dma_addr_t cur_dst = ddata.dst; + dma_addr_t cur_src = ddata.src; + + while (rem_rows > 0) { + dma_data cur_data = dma_make_data(cur_dst, cur_src); + if (!dma_ring_push_single_1d(r1, cur_data, row_size)) { + dma_ring_flush(r1); + dma_ring_push_single_1d(r1, cur_data, row_size); + } + cur_dst += dst_stride; + cur_src += src_stride; + rem_rows -= 1; + } + + dma_ring_flush(r1); + r0->tail = r1->tail; + + return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0); +} + +#endif diff --git a/ggml/src/ggml-hexagon/htp/dma-queue.h b/ggml/src/ggml-hexagon/htp/dma-queue.h index 190ca3a9b..d256e6bef 100644 --- a/ggml/src/ggml-hexagon/htp/dma-queue.h +++ b/ggml/src/ggml-hexagon/htp/dma-queue.h @@ -3,8 +3,10 @@ #include #include +#include #include #include +#include #include "hex-utils.h" #include "hex-profile.h" @@ -24,8 +26,8 @@ typedef struct dma_descriptor_1d_s { uint32_t src_bypass:1; uint32_t order:1; uint32_t done:1; - void * src; - void * dst; + uint32_t src; + uint32_t dst; } dma_descriptor_1d; #if __HVX_ARCH__ < 75 @@ -40,8 +42,8 @@ typedef struct dma_descriptor_2d_s { uint32_t src_bypass:1; uint32_t order:1; uint32_t done:1; - void * src; - void * dst; + uint32_t src; + uint32_t dst; uint32_t desc_type:8; uint32_t reserved1:24; uint32_t row_size:16; @@ -64,10 +66,18 @@ typedef struct dma_descriptor_2d_s { uint32_t src_bypass:1; uint32_t order:1; uint32_t done:1; - void * src; - void * dst; + uint32_t src; + uint32_t dst; uint32_t desc_type:8; +#if __HVX_ARCH__ > 79 + uint32_t src_upper:8; + uint32_t dst_upper:8; + uint32_t allocation:2; + uint32_t reserved0:2; + uint32_t transform:4; +#else uint32_t reserved0:24; +#endif uint32_t row_size:24; uint32_t nrows_lo:8; uint32_t nrows_hi:8; @@ -78,45 +88,63 @@ typedef struct dma_descriptor_2d_s { #endif +#if __HVX_ARCH__ > 79 +typedef uint64_t dma_addr_t; +#else +typedef uint32_t dma_addr_t; +#endif + typedef struct { - void *dst; - const void *src; -} dma_ptr; + dma_addr_t dst; + dma_addr_t src; +} dma_data; + +// Hardware descriptor field limits +#define DMA_MAX_NROWS 0xFFFFu // 16-bit HW descriptor limit (65535) +#define DMA_MAX_SIZE_16B 0xFFFFu // 16-bit HW descriptor limit for row_size (65535) +#define DMA_MAX_STRIDE_16B 0xFFFFu // 16-bit HW descriptor limit for strides (65535) +#define DMA_MAX_SIZE_24B 0x00FFFFFFu // 24-bit HW descriptor limit for row_size / 1D size (16MB - 1) +#define DMA_MAX_STRIDE_24B 0x00FFFFFFu // 24-bit HW descriptor limit for strides (16MB - 1) +#define DMA_SAFE_CHUNK_SIZE 0x00F00000u // ~15MB safe contiguous chunk size + +#define DMA_FALLBACK_CAPACITY 16u // descriptors in secondary fallback ring typedef struct dma_ring_s dma_ring; struct dma_ring_s { dma_descriptor_2d * desc; // descriptor pointers dma_descriptor_2d * tail; // tail pointer - dma_ptr * dptr; // dst/src pointers + dma_data * data; // dst/src data uint32_t push_idx; uint32_t pop_idx; uint32_t capacity; uint32_t idx_mask; struct htp_thread_trace * trace; - uintptr_t vtcm_base; - uintptr_t vtcm_end; }; typedef struct dma_queue_s dma_queue; typedef dma_queue * dma_queue_t; struct dma_queue_s { - dma_ring * ring; // Points to the descriptor ring state - uint8_t nocache; // Queue-specific bypass flag + dma_ring * ring0; // Main descriptor ring state + dma_ring * ring1; // Secondary fallback descriptor ring state bool alias; // When set, dma_queue_delete will not free the ring }; - - size_t dma_queue_sizeof(size_t capacity); size_t dma_queue_alignof(void); -dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace); +dma_queue_t dma_queue_init(void * ptr, size_t capacity, struct htp_thread_trace * trace); void dma_queue_free(dma_queue_t q); size_t dma_queue_alias_sizeof(void); -dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache); +dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q); void dma_queue_alias_free(dma_queue_t q); +bool dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows); +bool dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total); +#if __HVX_ARCH__ < 75 +bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows); +#endif + // TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead // but those do not seem to always compiler properly. static inline void dmstart(void * next) { @@ -141,36 +169,37 @@ static inline unsigned int dmwait(void) { return ret; } -static inline dma_ptr dma_make_ptr(void *dst, const void *src) +static inline dma_data dma_make_data_impl(dma_addr_t dst, dma_addr_t src) { - dma_ptr p = { dst, src }; - return p; + dma_data d = { dst, src }; + return d; } -static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) { - return (uintptr_t) ptr >= q->ring->vtcm_base && (uintptr_t) ptr < q->ring->vtcm_end; -} +#define dma_make_data(dst, src) dma_make_data_impl((dma_addr_t) (dst), (dma_addr_t) (src)) + +static inline bool dma_ring_push_single_1d(dma_ring * r, dma_data ddata, size_t size) { +#if __HVX_ARCH__ > 79 + assert(!((ddata.src | ddata.dst) >> 32) || size == 0); +#endif -static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) { - dma_ring * r = q->ring; if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) { return false; } dma_descriptor_1d * desc = (dma_descriptor_1d *) &r->desc[r->push_idx]; - desc->src = (void *) dptr.src; - desc->dst = (void *) dptr.dst; + desc->src = (uint32_t) ddata.src; + desc->dst = (uint32_t) ddata.dst; desc->size = size; - r->dptr[r->push_idx] = dptr; + r->data[r->push_idx] = ddata; htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx); if (size) { desc->next = NULL; desc->desc_size = 0; // 1D mode - desc->src_bypass = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache; - desc->dst_bypass = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache; + desc->src_bypass = 1; + desc->dst_bypass = 1; desc->order = 0; desc->done = 0; @@ -185,8 +214,17 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t return true; } -static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { - dma_ring * r = q->ring; +static inline bool dma_ring_push_single_2d(dma_ring * r, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { +#if __HVX_ARCH__ > 79 + const uint32_t src_hi = (uint32_t) (ddata.src >> 32); + const uint32_t dst_hi = (uint32_t) (ddata.dst >> 32); + const bool is_ext = (src_hi | dst_hi) != 0; + + if (is_ext && ((ddata.src >> 40) || (ddata.dst >> 40))) { + return false; + } +#endif + if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) { return false; } @@ -194,34 +232,44 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dma_descriptor_2d * desc = &r->desc[r->push_idx]; desc->next = NULL; - desc->reserved0 = 0; desc->reserved1 = 0; desc->desc_size = 1; // 2d mode - desc->src_bypass = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache; - desc->dst_bypass = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache; + desc->src_bypass = 1; + desc->dst_bypass = 1; desc->src_comp = 0; desc->dst_comp = 0; desc->order = 0; desc->done = 0; desc->src_stride = src_stride; desc->dst_stride = dst_stride; - desc->src = (void *) dptr.src; - desc->dst = (void *) dptr.dst; + desc->src = (uint32_t) ddata.src; + desc->dst = (uint32_t) ddata.dst; desc->row_size = row_size; #if __HVX_ARCH__ < 75 + desc->reserved0 = 0; desc->desc_type = 0; // 2d (16-bit) mode desc->nrows = nrows; desc->src_offset = 0; desc->dst_offset = 0; #else +#if __HVX_ARCH__ > 79 + desc->src_upper = src_hi; + desc->dst_upper = dst_hi; + desc->allocation = 0; + desc->reserved0 = 0; + desc->transform = 0; + desc->desc_type = is_ext ? 10 : 9; // 2d 40-bit or 24-bit mode +#else + desc->reserved0 = 0; desc->desc_type = 9; // 2d (24-bit) mode +#endif desc->nrows_lo = (nrows & 0xff); desc->nrows_hi = (nrows >> 8); desc->offset = 0; #endif - r->dptr[r->push_idx] = dptr; + r->data[r->push_idx] = ddata; htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx); @@ -236,21 +284,20 @@ static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t return true; } -static inline dma_ptr dma_queue_pop(dma_queue * q) { - dma_ring * r = q->ring; - dma_ptr dptr = { NULL }; +static inline dma_data dma_ring_pop(dma_ring * r) { + dma_data ddata = { 0 }; if (r->push_idx == r->pop_idx) { - return dptr; + return ddata; } - dptr = r->dptr[r->pop_idx]; + ddata = r->data[r->pop_idx]; volatile dma_descriptor_2d * desc = &r->desc[r->pop_idx]; // Wait for desc to complete if (!desc->done) { - // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, dptr.dst, dptr.src); + // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, ddata.dst, ddata.src); while (!desc->done) { dmpoll(); } @@ -259,108 +306,133 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) { htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx); r->pop_idx = (r->pop_idx + 1) & r->idx_mask; - return dptr; + return ddata; } -static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) { - dma_ring * r = q->ring; - dma_ptr dptr = { NULL }; +static inline dma_data dma_ring_pop_nowait(dma_ring * r) { + dma_data ddata = { 0 }; if (r->push_idx == r->pop_idx) { - return dptr; + return ddata; } - dptr = r->dptr[r->pop_idx]; + ddata = r->data[r->pop_idx]; htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx); r->pop_idx = (r->pop_idx + 1) & r->idx_mask; - return dptr; + return ddata; +} + +static inline bool dma_ring_empty(dma_ring * r) { + return r->push_idx == r->pop_idx; +} + +static inline void dma_ring_flush(dma_ring * r) { + while (!dma_ring_empty(r)) { + dma_ring_pop(r); + } +} + +static inline uint32_t dma_ring_depth(dma_ring * r) { + return (r->push_idx - r->pop_idx) & r->idx_mask; +} + +static inline uint32_t dma_ring_capacity(dma_ring * r) { + return r->capacity; +} + +static inline bool dma_queue_push_single_1d(dma_queue * q, dma_data ddata, size_t size) { + return dma_ring_push_single_1d(q->ring0, ddata, size); +} + +static inline bool dma_queue_push_single_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { + return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows); +} + +static inline dma_data dma_queue_pop(dma_queue * q) { + return dma_ring_pop(q->ring0); +} + +static inline dma_data dma_queue_pop_nowait(dma_queue * q) { + return dma_ring_pop_nowait(q->ring0); } static inline bool dma_queue_empty(dma_queue * q) { - return q->ring->push_idx == q->ring->pop_idx; + return dma_ring_empty(q->ring0); } static inline void dma_queue_flush(dma_queue * q) { - while (dma_queue_pop(q).dst != NULL) ; + dma_ring_flush(q->ring0); } static inline uint32_t dma_queue_depth(dma_queue * q) { - return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask; + return dma_ring_depth(q->ring0); } static inline uint32_t dma_queue_capacity(dma_queue * q) { - return q->ring->capacity; + return dma_ring_capacity(q->ring0); } #if __HVX_ARCH__ < 75 -// Overflow-safe DMA push: all 2d descriptor fields (row_size, nrows, src_stride, dst_stride) are 16-bit, max 65535. -// This version transparently handles values that exceed the 16-bit limit and submits chained DMA transtions. - -#define DMA_MAX_FIELD_VAL 65535u - -static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { +static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { // Fast path: everything fits in 16 bits if (nrows == 0 || __builtin_expect( - row_size <= DMA_MAX_FIELD_VAL && - nrows <= DMA_MAX_FIELD_VAL && - src_stride <= DMA_MAX_FIELD_VAL && - dst_stride <= DMA_MAX_FIELD_VAL, 1)) { - return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows); + nrows <= DMA_MAX_NROWS && + row_size <= DMA_MAX_SIZE_16B && + src_stride <= DMA_MAX_STRIDE_16B && + dst_stride <= DMA_MAX_STRIDE_16B, 1)) { + return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows); } - // Contiguous block - // Use 1d DMA mode which supports sizes up to 24-bits (16MB) + // Contiguous block: 1D DMA mode supports up to 24-bit size (16MB) if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) { size_t total = row_size * nrows; - return dma_queue_push_single_1d(q, dptr, total); + if (total <= DMA_MAX_SIZE_24B) { + return dma_ring_push_single_1d(q->ring0, ddata, total); + } + return dma_queue_push_fallback_contig(q, ddata, total); } - // Stride overflow - fall back to row-by-row. - { - const uint8_t *src = (const uint8_t *) dptr.src; - uint8_t *dst = (uint8_t *) dptr.dst; - size_t r = 0; - while (r + 1 < nrows) { - dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride); - if (!dma_queue_push_single_1d(q, p, row_size)) { - dma_queue_flush(q); - } else { - r++; - } - } - dma_queue_flush(q); - dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride); - return dma_queue_push_single_1d(q, p, row_size); + // Row count overflow with 16-bit strides: chunk 2D descriptors via fallback ring + if (row_size <= DMA_MAX_SIZE_16B && src_stride <= DMA_MAX_STRIDE_16B && dst_stride <= DMA_MAX_STRIDE_16B) { + return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows); } + + // Stride or row_size overflow: row-by-row 1D via fallback ring + return dma_queue_push_fallback_1d(q, ddata, dst_stride, src_stride, row_size, nrows); } #else // HVX_ARCH >= 75 -static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { - // On v75 and up we always use 2d 24-bit mode - return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows); +static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) { + if (nrows == 0 || __builtin_expect( + nrows <= DMA_MAX_NROWS && + row_size <= DMA_MAX_SIZE_24B && + src_stride <= DMA_MAX_STRIDE_24B && + dst_stride <= DMA_MAX_STRIDE_24B, 1)) { + return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows); + } + + // Contiguous block exceeding 24 bits + if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) { + size_t total = row_size * nrows; + return dma_queue_push_fallback_contig(q, ddata, total); + } + + return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows); } #endif -static inline bool dma_queue_push_ddr_to_vtcm(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) { - return dma_queue_push(q, dptr, dst_row_size, src_row_size, src_row_size, nrows); -} - -static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) { - return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows); -} - #define DMA_CACHE_MAX_SIZE 256U typedef struct { uint8_t *base; uint32_t line_size; uint32_t capacity; - uint32_t src[DMA_CACHE_MAX_SIZE]; + dma_addr_t src[DMA_CACHE_MAX_SIZE]; uint16_t age[DMA_CACHE_MAX_SIZE]; } dma_cache; @@ -376,14 +448,14 @@ static inline void dma_cache_init(dma_cache *c, uint8_t *base, uint32_t line_siz } } -static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * src, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows) +static inline bool dma_cache_push(dma_queue *q, dma_cache *c, dma_addr_t src_addr, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows) { uint32_t o_idx = 0; uint16_t o_age = 0; uint8_t * dst = 0; for (unsigned i=0; i < c->capacity; i++) { - if (c->src[i] == (uint32_t) src) { + if (c->src[i] == src_addr) { c->age[i] = 0; dst = c->base + (i * c->line_size); nrows = 0; // dummy dma } else { @@ -393,12 +465,12 @@ static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * sr } if (!dst) { c->age[o_idx] = 0; - c->src[o_idx] = (uint32_t) src; + c->src[o_idx] = src_addr; dst = c->base + o_idx * c->line_size; // normal nrows dma - return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows); + return dma_queue_push(q, dma_make_data(dst, src_addr), dst_stride, src_stride, row_size, nrows); } - return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0); + return dma_queue_push_single_1d(q, dma_make_data(dst, src_addr), 0); } #ifdef __cplusplus diff --git a/ggml/src/ggml-hexagon/htp/fill-ops.c b/ggml/src/ggml-hexagon/htp/fill-ops.c index 1f6eaafad..212104a23 100644 --- a/ggml/src/ggml-hexagon/htp/fill-ops.c +++ b/ggml/src/ggml-hexagon/htp/fill-ops.c @@ -88,8 +88,8 @@ int op_fill(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + if (htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; } uint32_t row_start = 0; diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 75422f420..988886082 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -12,7 +12,7 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hex-fastdiv.h" #include "hex-profile.h" #include "hmx-queue.h" @@ -86,6 +86,7 @@ struct htp_fa_context { uint8_t * spad_v; uint8_t * spad_m; uint8_t * spad_a; + float * spad_sinks; const struct htp_tensor * k; const struct htp_tensor * v; @@ -149,6 +150,7 @@ struct hmx_fa_context { uint8_t * vtcm_hmx_scales_qk; // HMX output scales (qk_scale) __fp16 * vtcm_mask_buf; // VTCM mask buffer [Br * m_line], DMA'd per KV block __fp16 * vtcm_slopes; // ALiBi slopes [g_br] + float * vtcm_sinks; // Attention sinks size_t row_buf_stride; // HVX vectors per row buffer (Bc/64) size_t mask_buf_row_stride; // elements (__fp16) per row in mask buffer size_t q_tile_bytes; @@ -213,7 +215,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - dma_queue * dma = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; const uint32_t DK = nek0; const uint32_t DV = nev0; @@ -243,27 +245,27 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3); const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2); - const __fp16 * mp_base = NULL; + dma_addr_t mp_base = 0; if (mask) { const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2); const uint32_t im3 = fastmodulo(iq3, mask->ne[3], &factx->src3_div3); - mp_base = (const __fp16 *) ((const uint8_t *) mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]); + mp_base = mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]; } // Precalculate next row variables if there is a next row bool has_next_ir = (ir + 1 < ir1); uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0; - const uint8_t * next_q_row_ptr = NULL; - const __fp16 * next_mp_base = NULL; + dma_addr_t next_q_row_ptr = 0; + dma_addr_t next_mp_base = 0; - const uint8_t * next_k_src0 = NULL; - const uint8_t * next_v_src0 = NULL; - const uint8_t * next_m_src0 = NULL; + dma_addr_t next_k_src0 = 0; + dma_addr_t next_v_src0 = 0; + dma_addr_t next_m_src0 = 0; uint32_t next_block_size0 = 0; - const uint8_t * next_k_src1 = NULL; - const uint8_t * next_v_src1 = NULL; - const uint8_t * next_m_src1 = NULL; + dma_addr_t next_k_src1 = 0; + dma_addr_t next_v_src1 = 0; + dma_addr_t next_m_src1 = 0; uint32_t next_block_size1 = 0; if (has_next_ir) { @@ -278,22 +280,22 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3); next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2); - next_q_row_ptr = (const uint8_t *) q->data + (next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3); + next_q_row_ptr = q->data + next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3; if (mask) { const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2); const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3); - next_mp_base = (const __fp16 *) ((const uint8_t *) mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]); + next_mp_base = mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]; } // Precalculate next K/V block 0 source pointers { const uint32_t ic_start = 0; next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); - next_k_src0 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3); - next_v_src0 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3); + next_k_src0 = k->data + ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3; + next_v_src0 = v->data + ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3; if (mask) { - next_m_src0 = (const uint8_t *) (next_mp_base + ic_start); + next_m_src0 = next_mp_base + ic_start * sizeof(__fp16); } } @@ -301,18 +303,18 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * if (factx->n_blocks > 1) { const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE; next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); - next_k_src1 = (const uint8_t *) k->data + (ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3); - next_v_src1 = (const uint8_t *) v->data + (ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3); + next_k_src1 = k->data + ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3; + next_v_src1 = v->data + ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3; if (mask) { - next_m_src1 = (const uint8_t *) (next_mp_base + ic_start); + next_m_src1 = next_mp_base + ic_start * sizeof(__fp16); } } } if (ir == ir0) { // Fetch Q row - const uint8_t * q_row_ptr = (const uint8_t *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3); - dma_queue_push(dma, dma_make_ptr(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); + const dma_addr_t q_row_ptr = q->data + iq1*nbq1 + iq2*nbq2 + iq3*nbq3; + dma_queue_push(dma_q, dma_make_data(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); // Prefetch first two blocks for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) { @@ -320,20 +322,20 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); // K - const uint8_t * k_src = (const uint8_t *) k->data + (ic_start*nbk1 + ik2*nbk2 + ik3*nbk3); + const dma_addr_t k_src = k->data + ic_start*nbk1 + ik2*nbk2 + ik3*nbk3; uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block; - dma_queue_push(dma, dma_make_ptr(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size); + dma_queue_push(dma_q, dma_make_data(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size); // V - const uint8_t * v_src = (const uint8_t *) v->data + (ic_start*nbv1 + iv2*nbv2 + iv3*nbv3); + const dma_addr_t v_src = v->data + ic_start*nbv1 + iv2*nbv2 + iv3*nbv3; uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block; - dma_queue_push(dma, dma_make_ptr(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size); + dma_queue_push(dma_q, dma_make_data(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size); // Mask if (mask) { - const uint8_t * m_src = (const uint8_t *) (mp_base + ic_start); + const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16); // Mask is 1D contiguous for this row - dma_cache_push(dma, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1); + dma_cache_push(dma_q, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1); } } } @@ -348,7 +350,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * hvx_splat_f32_a(spad_a, 0, DV); float * VKQ32 = (float *) (spad_a + 0); - uint8_t * q_ptr_vtcm = dma_queue_pop(dma).dst; + uint8_t * q_ptr_vtcm = (void *) dma_queue_pop(dma_q).dst; if (factx->is_q_fp32) { hvx_copy_f16_f32_aa(q_ptr_vtcm, q_ptr_vtcm, DK); // inplace convert f32 to f16 } @@ -365,9 +367,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); // Wait for DMA - uint8_t * k_base = dma_queue_pop(dma).dst; // K - uint8_t * v_base = dma_queue_pop(dma).dst; // V - __fp16 * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M + uint8_t * k_base = (void *) dma_queue_pop(dma_q).dst; // K + uint8_t * v_base = (void *) dma_queue_pop(dma_q).dst; // V + __fp16 * m_base = mask ? (__fp16 *) dma_queue_pop(dma_q).dst : NULL; // M if (factx->k->type == HTP_TYPE_Q8_0) { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); @@ -424,7 +426,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * if (ib + 1 == factx->n_blocks && has_next_ir) { // Queue next row's Q row! - dma_queue_push(dma, dma_make_ptr(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); + dma_queue_push(dma_q, dma_make_data(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); if (factx->n_blocks % 2 == 0) { // Queue next row's block 0 (into buffer slot 0) @@ -432,14 +434,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * uint8_t * v_dst = spad_v + 0 * factx->size_v_block; // K (block 0 of next row) - dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0); + dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0); // V (block 0 of next row) - dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0); + dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0); // Mask (block 0 of next row) if (mask) { - dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1); + dma_cache_push(dma_q, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1); } } } @@ -502,17 +504,17 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start); // K - const uint8_t * k_src = (const uint8_t *) k->data + (next_ic_start*nbk1 + ik2*nbk2 + ik3*nbk3); - dma_queue_push(dma, dma_make_ptr(k_base, k_src), factx->size_k_row_padded, nbk1, size_k_row, next_block_size); + const dma_addr_t k_src = k->data + next_ic_start*nbk1 + ik2*nbk2 + ik3*nbk3; + dma_queue_push(dma_q, dma_make_data(k_base, k_src), factx->size_k_row_padded, nbk1, size_k_row, next_block_size); // V - const uint8_t * v_src = (const uint8_t *) v->data + (next_ic_start*nbv1 + iv2*nbv2 + iv3*nbv3); - dma_queue_push(dma, dma_make_ptr(v_base, v_src), factx->size_v_row_padded, nbv1, size_v_row, next_block_size); + const dma_addr_t v_src = v->data + next_ic_start*nbv1 + iv2*nbv2 + iv3*nbv3; + dma_queue_push(dma_q, dma_make_data(v_base, v_src), factx->size_v_row_padded, nbv1, size_v_row, next_block_size); // Mask if (mask) { - const uint8_t * m_src = (const uint8_t *) (mp_base + next_ic_start); - dma_cache_push(dma, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1); + const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16); + dma_cache_push(dma_q, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1); } } } @@ -525,14 +527,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * uint8_t * v_dst = spad_v + 1 * factx->size_v_block; // K (block 1 of next row) - dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1); + dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1); // V (block 1 of next row) - dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1); + dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1); // Mask (block 1 of next row) if (mask) { - dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1); + dma_cache_push(dma_q, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1); } } } else { @@ -542,14 +544,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * uint8_t * v_dst = spad_v + 0 * factx->size_v_block; // K (block 0 of next row) - dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0); + dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0); // V (block 0 of next row) - dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0); + dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0); // Mask (block 0 of next row) if (mask) { - dma_cache_push(dma, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1); + dma_cache_push(dma_q, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1); } } @@ -559,14 +561,14 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * uint8_t * v_dst = spad_v + 1 * factx->size_v_block; // K (block 1 of next row) - dma_queue_push(dma, dma_make_ptr(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1); + dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1); // V (block 1 of next row) - dma_queue_push(dma, dma_make_ptr(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1); + dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1); // Mask (block 1 of next row) if (mask) { - dma_cache_push(dma, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1); + dma_cache_push(dma_q, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1); } } } @@ -578,7 +580,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * float S = hvx_vec_get_f32(S_vec); if (sinks) { - const float s = ((float *)((char *) sinks->data))[h]; + const float s = factx->spad_sinks[h]; float vs = 1.0f; @@ -781,7 +783,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { const size_t m_end = hex_smin(m_start + m_bytes_per_t, col_vec_bytes); if (factx->sinks) { - const float * sinks_data = (const float *) (uintptr_t) factx->sinks->data; + const float * sinks_data = factx->vtcm_sinks; float * m_vec = (float *) factx->vtcm_m_vec; const size_t r_start = l_start / sizeof(float); const size_t r_end = l_end / sizeof(float); @@ -1779,7 +1781,7 @@ static __attribute__((noinline)) void fa_compute_slopes( } static void fa_push_mask_dma_gqa( - dma_queue * dma, + dma_queue * dma_q, const struct htp_tensor * mask, uint32_t q_start, uint32_t im3, @@ -1794,36 +1796,36 @@ static void fa_push_mask_dma_gqa( for (uint32_t g = 0; g < G; ++g) { const uint32_t h_idx = kv_head * G + g; const uint32_t im2 = fastmodulo(h_idx, mask->ne[2], &factx->src3_div2); - const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + - im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16); + const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + + im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16); uint8_t * ms_dst = (uint8_t *) factx->vtcm_mask_buf + g * m_line_bytes; - dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); + dma_queue_push(dma_q, dma_make_data(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); } } -static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) { +static void fa_pop_mask_dma_gqa(dma_queue * dma_q, uint32_t G) { for (uint32_t g = 0; g < G; ++g) { - dma_queue_pop(dma); + dma_queue_pop(dma_q); } } -static inline void fa_prefetch_block(dma_queue * dma, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask, +static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor * k, const struct htp_tensor * v, const struct htp_tensor * mask, uint32_t b, size_t Bc, size_t size_k_row_padded, size_t size_k_row, size_t size_v_row_padded, size_t size_v_row, uint32_t ik2, uint32_t ik3, uint32_t iv2, uint32_t iv3, uint32_t q_start, uint32_t im3, uint32_t kv_head, uint32_t G, size_t m_line_bytes, size_t n_rows_q, size_t nek1, size_t prefetch_buf, struct hmx_fa_context * factx) { const uint32_t prefetch_start = b * Bc; const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start); - const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows); - const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows); + const dma_addr_t k_prefetch_src = k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx->vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows); + const dma_addr_t v_prefetch_src = v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx->vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows); if (mask) { if (__builtin_expect(factx->mask_broadcast, true)) { - const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16); - dma_cache_push(dma, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q); + const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16); + dma_cache_push(dma_q, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q); } else { - fa_push_mask_dma_gqa(dma, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx); + fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx); } } } @@ -1953,7 +1955,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Build the VTCM layout once (shared with the host estimator) and place every // scratch buffer at its computed offset. Padded head dims size the HMX tiles. struct hmx_fa_vtcm_layout L; - hmx_fa_vtcm_layout_build(&L, G, DK_pad, DV_pad, Br, Bc, n_threads, pipeline, factx.is_q_fp32); + hmx_fa_vtcm_layout_build(&L, G, DK_pad, DV_pad, Br, Bc, n_threads, pipeline, factx.is_q_fp32, factx.sinks != NULL, factx.n_heads); if (L.total_bytes > ctx->vtcm_size) { return HTP_STATUS_VTCM_TOO_SMALL; @@ -1995,6 +1997,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { factx.col_vec_bytes = L.col_vec_bytes; factx.d_tile_bytes = L.d_tile_bytes; factx.vtcm_slopes = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes); + factx.vtcm_sinks = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_sinks, factx.sinks != NULL); const size_t m_line_bytes = L.m_line_bytes; // used by the mask DMAs in the KV loop @@ -2022,13 +2025,13 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0 hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale)); - // ======== Skip compute if profiling ======== - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } - // ======== DMA setup ======== - dma_queue * const dma = ctx->dma[0]; + dma_queue * const dma_q = ctx->dma[0]; + + if (factx.sinks) { + dma_queue_push(dma_q, dma_make_data(factx.vtcm_sinks, factx.sinks->data), L.sinks_bytes, 0, factx.sinks->size, 1); + dma_queue_pop(dma_q); + } const size_t n_row_tiles_g_br = g_br / HMX_FP16_TILE_N_ROWS; const size_t n_tiles_per_bc = Bc / HMX_FP16_TILE_N_COLS; @@ -2064,32 +2067,32 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // 1. Push Q and KV DMAs for the very first iteration. // Subsequent iterations are enqueued early at the end of the previous iteration. if (ib3 == 0 && q_start == q_start_min && kv_head == 0) { - const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + + const dma_addr_t q_ptr = q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3]; const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed; const size_t n_rows = q_transposed ? factx.G : n_rows_q; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows); + dma_queue_push(dma_q, dma_make_data(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows); if (factx.n_kv_blocks > 0) { - const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0); + const dma_addr_t k_src = k->data + ik2 * k->nb[2] + ik3 * k->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0); - const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0); + const dma_addr_t v_src = v->data + iv2 * v->nb[2] + iv3 * v->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0); if (factx.pipeline && mask) { if (__builtin_expect(factx.mask_broadcast, true)) { - const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0; - dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q); + const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0; + dma_cache_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q); } else { - fa_push_mask_dma_gqa(dma, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx); + fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx); } } } } // 2. Pop Q DMA (blocks until Q is loaded) - dma_queue_pop(dma); + dma_queue_pop(dma_q); // ---- Load Q block & Initialize per-block state ---- fa_phase_q_load(&factx, q, q_start, kv_head, ib3, n_rows_g); @@ -2116,12 +2119,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Prefetch block 1 early if there are multiple blocks if (factx.n_kv_blocks > 1) { - fa_prefetch_block(dma, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row, + fa_prefetch_block(dma_q, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row, ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx); } // Prep and start QK-dot(0) - void * curr_k0 = dma_queue_pop(dma).dst; + void * curr_k0 = (void *) dma_queue_pop(dma_q).dst; fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0); qk_job[0].q_tiles = factx.vtcm_q_tiles; @@ -2140,16 +2143,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const size_t n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS); // ---- 1. Pop and run V-prep for current block ---- - void * curr_v = dma_queue_pop(dma).dst; + void * curr_v = (void *) dma_queue_pop(dma_q).dst; fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start); // ---- 2. Pop and run mask-prep for current block ---- __fp16 * current_mask_vtcm = NULL; if (mask) { if (__builtin_expect(factx.mask_broadcast, true)) { - current_mask_vtcm = (__fp16 *) dma_queue_pop(dma).dst; + current_mask_vtcm = (__fp16 *) dma_queue_pop(dma_q).dst; } else { - fa_pop_mask_dma_gqa(dma, G); + fa_pop_mask_dma_gqa(dma_q, G); current_mask_vtcm = factx.vtcm_mask_buf; } } @@ -2183,7 +2186,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const uint32_t next_rows = hex_smin(Bc, nek1 - next_start); const size_t next_buf = 1 - buf_idx; - void * next_k = dma_queue_pop(dma).dst; + void * next_k = (void *) dma_queue_pop(dma_q).dst; fa_phase_k_interleave(&factx, next_rows, k_src_stride, next_k, next_start, next_buf); qk_job[next_buf].q_tiles = factx.vtcm_q_tiles; @@ -2234,7 +2237,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Prefetch block kv_blk + 2 if (kv_blk + 2 < factx.n_kv_blocks) { - fa_prefetch_block(dma, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row, + fa_prefetch_block(dma_q, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row, ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx); } @@ -2276,10 +2279,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { if (mask) { if (__builtin_expect(factx.mask_broadcast, true)) { - const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16); - dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); + const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16); + dma_cache_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q); } else { - fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx); + fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx); } } @@ -2287,14 +2290,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { const uint32_t prefetch_start = (kv_blk + 1) * Bc; const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start); const size_t prefetch_buf = 1 - buf_idx; - const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows); - const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows); + const dma_addr_t k_prefetch_src = k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows); + const dma_addr_t v_prefetch_src = v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows); } // Wait for current K DMA and interleave - void * curr_k = dma_queue_pop(dma).dst; + void * curr_k = (void *) dma_queue_pop(dma_q).dst; fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start, 0); { @@ -2312,16 +2315,16 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { } // Wait for current V DMA and interleave - void * curr_v = dma_queue_pop(dma).dst; + void * curr_v = (void *) dma_queue_pop(dma_q).dst; fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[0], n_tiles_per_bc, kv_start); // ---- Phase 3: softmax + build_D ---- __fp16 * current_mask_vtcm = NULL; if (mask) { if (__builtin_expect(factx.mask_broadcast, true)) { - current_mask_vtcm = (__fp16 *) dma_queue_pop(dma).dst; + current_mask_vtcm = (__fp16 *) dma_queue_pop(dma_q).dst; } else { - fa_pop_mask_dma_gqa(dma, G); + fa_pop_mask_dma_gqa(dma_q, G); current_mask_vtcm = factx.vtcm_mask_buf; } } @@ -2393,10 +2396,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { if (has_next) { const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start); - const uint8_t * next_q_ptr = (const uint8_t *) q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3]; + const dma_addr_t next_q_ptr = q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3]; const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed; const size_t next_n_rows = q_transposed ? factx.G : next_n_rows_q; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows); + dma_queue_push(dma_q, dma_make_data(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows); if (factx.n_kv_blocks > 0) { const uint32_t next_ik2 = next_kv_head; @@ -2408,11 +2411,11 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3); } - const uint8_t * next_k_src = (const uint8_t *) k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0); + const dma_addr_t next_k_src = k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0); - const uint8_t * next_v_src = (const uint8_t *) v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3]; - dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0); + const dma_addr_t next_v_src = v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3]; + dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0); if (factx.pipeline && mask) { uint32_t next_im3 = im3; @@ -2420,10 +2423,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3); } if (__builtin_expect(factx.mask_broadcast, true)) { - const uint8_t * ms_src = (const uint8_t *) mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0; - dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q); + const dma_addr_t ms_src = mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0; + dma_cache_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q); } else { - fa_push_mask_dma_gqa(dma, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx); + fa_push_mask_dma_gqa(dma_q, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx); } } } @@ -2465,6 +2468,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } + if (htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } + const struct htp_fa_kernel_params * kparams = (const struct htp_fa_kernel_params *) octx->kernel_params; if (kparams->kernel_type == HTP_FA_KERNEL_UNSUPPORTED) { @@ -2502,11 +2509,6 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { factx.size_k_row_padded = kparams->u.hvx.size_k_row_padded; factx.size_v_row_padded = kparams->u.hvx.size_v_row_padded; - size_t size_q_block = factx.size_q_row_padded * 1; // single row for now - factx.size_k_block = factx.size_k_row_padded * FLASH_ATTN_BLOCK_SIZE; - factx.size_v_block = factx.size_v_row_padded * FLASH_ATTN_BLOCK_SIZE; - factx.size_m_block = hex_round_up(FLASH_ATTN_BLOCK_SIZE * sizeof(__fp16), 128); - factx.n_blocks = kparams->n_kv_blocks; factx.scale = kparams->scale; @@ -2551,26 +2553,38 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { factx.qrow_start = qrow_start; factx.qrows_per_thread = fastdiv(qrows + n_threads - 1, &octx->n_threads_div); - size_t size_vkq_acc = hex_round_up(v->ne[0] * sizeof(float), 128); // VKQ32 + const bool has_mask = (mask != NULL); + const bool has_sinks = (octx->src[4] != NULL); + struct hvx_fa_vtcm_layout L; + hvx_fa_vtcm_layout_build(&L, k->ne[0], v->ne[0], factx.is_q_fp32, has_mask, has_sinks, n_head, n_threads); - factx.size_q_block = size_q_block; - factx.size_vkq_acc = size_vkq_acc; - - uint8_t * vtcm_cur = octx->ctx->vtcm_base; - - factx.spad_q = vtcm_seq_alloc(&vtcm_cur, size_q_block * n_threads); - factx.spad_k = vtcm_seq_alloc(&vtcm_cur, factx.size_k_block * 2 * n_threads); - factx.spad_v = vtcm_seq_alloc(&vtcm_cur, factx.size_v_block * 2 * n_threads); - factx.spad_m = vtcm_seq_alloc(&vtcm_cur, (mask ? factx.size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * n_threads); - factx.spad_a = vtcm_seq_alloc(&vtcm_cur, size_vkq_acc * n_threads); - - if ((size_t) (vtcm_cur - octx->ctx->vtcm_base) > octx->ctx->vtcm_size) { + if (L.total_bytes > octx->ctx->vtcm_size) { return HTP_STATUS_VTCM_TOO_SMALL; } - if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) { - work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, n_threads); + factx.size_q_block = L.size_q_block; + factx.size_k_block = L.size_k_block; + factx.size_v_block = L.size_v_block; + factx.size_m_block = L.size_m_block; + factx.size_vkq_acc = L.size_vkq_acc; + + uint8_t * const base = octx->ctx->vtcm_base; + + factx.spad_q = VTCM_LAYOUT_PTR(uint8_t, base, L.off_q); + factx.spad_k = VTCM_LAYOUT_PTR(uint8_t, base, L.off_k); + factx.spad_v = VTCM_LAYOUT_PTR(uint8_t, base, L.off_v); + factx.spad_m = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_m, has_mask); + factx.spad_a = VTCM_LAYOUT_PTR(uint8_t, base, L.off_a); + factx.spad_sinks = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_sinks, has_sinks); + + if (has_sinks) { + const struct htp_tensor * sinks = octx->src[4]; + dma_queue * dma_q = octx->ctx->dma[0]; + dma_queue_push(dma_q, dma_make_data(factx.spad_sinks, sinks->data), L.size_sinks, 0, sinks->size, 1); + dma_queue_pop(dma_q); } + work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, n_threads); + return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index 027845411..2bd232190 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -121,6 +121,7 @@ struct hmx_fa_vtcm_layout { size_t off_hmx_scales_qk; size_t off_mask_buf; size_t off_slopes; + size_t off_sinks; // Region byte sizes reused by the device at runtime (not just for allocation). size_t q_tile_bytes; @@ -130,6 +131,7 @@ struct hmx_fa_vtcm_layout { size_t m_line_bytes; // one mask row size_t m_buf_slot_bytes; // one dma_cache slot = align_up(Br * m_line_bytes, 4096) size_t col_vec_bytes; + size_t sinks_bytes; // Derived strides. size_t row_buf_stride; // HVX vectors (128B) per row buffer @@ -142,8 +144,10 @@ struct hmx_fa_vtcm_layout { // Build the VTCM layout. static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, - size_t gqa_factor, size_t DK, size_t DV, - size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) { + size_t gqa_factor, size_t DK, size_t DV, + size_t Br, size_t Bc, size_t n_threads, + bool pipeline, bool is_q_fp32, + bool has_sinks, size_t n_heads) { const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS); const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE); @@ -166,6 +170,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, const size_t m_buf_slot = hex_align_up(Br * m_line_size, 256); const size_t m_buf_size = m_buf_slot * HMX_FA_DMA_CACHE_SIZE; const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128); + const size_t sinks_size = hex_round_up(n_heads * sizeof(float), 128); size_t off = 0; @@ -215,6 +220,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256); VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size); VTCM_LAYOUT_ALLOC(off, off_slopes, slopes_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_sinks, sinks_size, has_sinks); L->q_tile_bytes = q_tile_size; L->o_tile_bytes = o_tile_size; @@ -228,20 +234,43 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L, L->m_buf_slot_bytes = m_buf_slot; L->row_buf_stride = row_vec_size / 128; L->mask_buf_row_stride = m_line_size / sizeof(__fp16); + L->sinks_bytes = has_sinks ? sinks_size : 0; L->pipeline = pipeline; L->total_bytes = off; } // Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration. -static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32) { +static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline, bool is_q_fp32, bool has_sinks, size_t n_heads) { struct hmx_fa_vtcm_layout L; - hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32); + hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline, is_q_fp32, has_sinks, n_heads); return L.total_bytes; } #define FA_HVX_BLOCK_SIZE 64 -static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, size_t n_threads) { +struct hvx_fa_vtcm_layout { + size_t off_q; + size_t off_k; + size_t off_v; + size_t off_m; + size_t off_a; + size_t off_sinks; + + size_t size_q_block; + size_t size_k_block; + size_t size_v_block; + size_t size_m_block; + size_t size_vkq_acc; + size_t size_sinks; + + size_t total_bytes; +}; + +static inline void hvx_fa_vtcm_layout_build(struct hvx_fa_vtcm_layout * L, + size_t DK, size_t DV, + bool is_q_fp32, bool has_mask, + bool has_sinks, size_t n_heads, + size_t n_threads) { const size_t size_q_row_padded = hex_round_up(DK * (is_q_fp32 ? 4 : 2), 128); const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128); const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128); @@ -251,29 +280,47 @@ static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_f const size_t size_v_block = size_v_row_padded * FA_HVX_BLOCK_SIZE; const size_t size_m_block = hex_round_up(FA_HVX_BLOCK_SIZE * sizeof(__fp16), 128); const size_t size_vkq_acc = hex_round_up(DV * sizeof(float), 128); + const size_t size_sinks = hex_round_up(n_heads * sizeof(float), 128); - const size_t size_per_thread = size_q_block * 1 - + size_k_block * 2 - + size_v_block * 2 - + (has_mask ? size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) - + size_vkq_acc; + size_t off = 0; - return size_per_thread * n_threads; + VTCM_LAYOUT_ALLOC(off, off_q, size_q_block * n_threads); + VTCM_LAYOUT_ALLOC(off, off_k, size_k_block * 2 * n_threads); + VTCM_LAYOUT_ALLOC(off, off_v, size_v_block * 2 * n_threads); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_m, size_m_block * HVX_FA_DMA_CACHE_SIZE * n_threads, has_mask); + VTCM_LAYOUT_ALLOC(off, off_a, size_vkq_acc * n_threads); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_sinks, size_sinks, has_sinks); + + L->size_q_block = size_q_block; + L->size_k_block = size_k_block; + L->size_v_block = size_v_block; + L->size_m_block = size_m_block; + L->size_vkq_acc = size_vkq_acc; + L->size_sinks = has_sinks ? size_sinks : 0; + L->total_bytes = off; +} + +static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, bool has_sinks, size_t n_heads, size_t n_threads) { + struct hvx_fa_vtcm_layout L; + hvx_fa_vtcm_layout_build(&L, DK, DV, is_q_fp32, has_mask, has_sinks, n_heads, n_threads); + return L.total_bytes; } #define FA_MIN_KV_BLOCKS 3 // Cost-based (Br, Bc) search for flash attention with pipeline constraint. static inline int hmx_fa_find_chunk_size(size_t * Br_out, - size_t * Bc_out, - size_t gqa_factor, - size_t DK, - size_t DV, - size_t qo_len, - size_t kv_len, - size_t vtcm_budget, - size_t n_threads, - bool is_q_fp32) { + size_t * Bc_out, + size_t gqa_factor, + size_t DK, + size_t DV, + size_t qo_len, + size_t kv_len, + size_t vtcm_budget, + size_t n_threads, + bool is_q_fp32, + bool has_sinks, + size_t n_heads) { const size_t T = HMX_FP16_TILE_N_ROWS; // 32 const size_t br_unit = hmx_ceil_div(T, gqa_factor); const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64 @@ -297,7 +344,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out, for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) { // Try all Bc candidates from Bc_limit down to bc_unit for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) { - size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32); + size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline, is_q_fp32, has_sinks, n_heads); if (vtcm_needed <= vtcm_budget) { // This Bc fits for this Br! const size_t q_blocks = (qo_len + Br - 1) / Br; diff --git a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c index 0b6529571..b37313370 100644 --- a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c +++ b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c @@ -1,44 +1,40 @@ -#include #include +#include #include +#include -#include "hvx-utils.h" -#include "hex-fastdiv.h" -#include "hex-common.h" -#include "hex-profile.h" - -#define GGML_COMMON_DECL_C +#include "hvx-base.h" +#include "hvx-copy.h" +#include "hvx-reduce.h" +#include "hvx-exp.h" +#include "dma-queue.h" #include "ggml-common.h" #include "htp-ctx.h" #include "htp-tensor.h" +#include "gated-delta-net-ops.h" #ifndef MIN #define MIN(a, b) ((a) < (b) ? (a) : (b)) #endif -#define HTP_GDN_MAX_SV 128 - - struct htp_gdn_context { struct htp_ops_context * octx; - uint32_t rows_per_thread; - size_t state_bytes; + const struct htp_gdn_kernel_params * kparams; + struct htp_gdn_vtcm_layout layout; uint8_t * vtcm_base; - size_t vtcm_per_thread; uint32_t row_start; uint32_t nrows; }; -static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const float * restrict mul, const float * restrict dot, uint32_t n) { +static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const HVX_Vector * restrict mul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc = Q6_V_vzero(); - - const uint32_t epv = 128 / sizeof(float); + const uint32_t epv = 128 / sizeof(float); const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { HVX_Vector vd = hvx_vmemu(dst + i * epv); - HVX_Vector vm = hvx_vmem(mul + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vm = mul[i]; + HVX_Vector vdot = dot[i]; HVX_Vector out = hvx_vec_mul_f32_f32(vd, vm); hvx_vmemu(dst + i * epv) = out; acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(out, vdot)); @@ -46,29 +42,28 @@ static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const float * res if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vd = hvx_vmemu(dst + off); - HVX_Vector vm = hvx_vmem(mul + off); - HVX_Vector vdot = hvx_vmem(dot + off); - HVX_Vector out = hvx_vec_mul_f32_f32(vd, vm); - hvx_vec_store_u(dst + off, nloe * sizeof(float), out); + HVX_Vector vm = mul[nvec]; + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); - HVX_Vector prod = hvx_vec_mul_f32_f32(out, vdot); - acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, prod, Q6_V_vzero())); + HVX_Vector zero = Q6_V_vzero(); + + HVX_Vector out = hvx_vec_mul_f32_f32(hvx_vmemu(dst + off), vm); + hvx_vec_store_u(dst + off, nloe * sizeof(float), out); + acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out, vdot), zero)); } return hvx_vec_reduce_sum_f32(acc); } -static inline HVX_Vector gdn_mul_scalar_dot_f32(float * restrict dst, float mul, const float * restrict dot, uint32_t n) { +static inline HVX_Vector gdn_mul_scalar_dot_f32(float * restrict dst, float mul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc = Q6_V_vzero(); const HVX_Vector vmul = hvx_vec_splat_f32(mul); - - const uint32_t epv = 128 / sizeof(float); + const uint32_t epv = 128 / sizeof(float); const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { HVX_Vector vd = hvx_vmemu(dst + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vdot = dot[i]; HVX_Vector out = hvx_vec_mul_f32_f32(vd, vmul); hvx_vmemu(dst + i * epv) = out; acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(out, vdot)); @@ -76,29 +71,28 @@ static inline HVX_Vector gdn_mul_scalar_dot_f32(float * restrict dst, float mul, if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vd = hvx_vmemu(dst + off); - HVX_Vector vdot = hvx_vmem(dot + off); - HVX_Vector out = hvx_vec_mul_f32_f32(vd, vmul); - hvx_vec_store_u(dst + off, nloe * sizeof(float), out); + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); - HVX_Vector prod = hvx_vec_mul_f32_f32(out, vdot); - acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, prod, Q6_V_vzero())); + HVX_Vector zero = Q6_V_vzero(); + + HVX_Vector out = hvx_vec_mul_f32_f32(hvx_vmemu(dst + off), vmul); + hvx_vec_store_u(dst + off, nloe * sizeof(float), out); + acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out, vdot), zero)); } return hvx_vec_reduce_sum_f32(acc); } -static inline HVX_Vector gdn_add_scaled_dot_f32(float * restrict dst, const float * restrict src, - HVX_Vector vscale, const float * restrict dot, uint32_t n) { +static inline HVX_Vector gdn_add_scaled_dot_f32(float * restrict dst, const HVX_Vector * restrict src, + HVX_Vector vscale, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc = Q6_V_vzero(); - - const uint32_t epv = 128 / sizeof(float); + const uint32_t epv = 128 / sizeof(float); const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { HVX_Vector vd = hvx_vmemu(dst + i * epv); - HVX_Vector vs = hvx_vmem(src + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vs = src[i]; + HVX_Vector vdot = dot[i]; HVX_Vector out = hvx_vec_add_f32_f32(vd, hvx_vec_mul_f32_f32(vs, vscale)); hvx_vmemu(dst + i * epv) = out; acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(out, vdot)); @@ -106,22 +100,22 @@ static inline HVX_Vector gdn_add_scaled_dot_f32(float * restrict dst, const floa if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vd = hvx_vmemu(dst + off); - HVX_Vector vs = hvx_vmem(src + off); - HVX_Vector vdot = hvx_vmem(dot + off); - HVX_Vector out = hvx_vec_add_f32_f32(vd, hvx_vec_mul_f32_f32(vs, vscale)); - hvx_vec_store_u(dst + off, nloe * sizeof(float), out); + HVX_Vector vs = src[nvec]; + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); - HVX_Vector prod = hvx_vec_mul_f32_f32(out, vdot); - acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, prod, Q6_V_vzero())); + HVX_Vector zero = Q6_V_vzero(); + + HVX_Vector out = hvx_vec_add_f32_f32(hvx_vmemu(dst + off), hvx_vec_mul_f32_f32(vs, vscale)); + hvx_vec_store_u(dst + off, nloe * sizeof(float), out); + acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out, vdot), zero)); } return hvx_vec_reduce_sum_f32(acc); } -static inline void gdn_mul_dot4_f32(float * restrict dst0, float * restrict dst1, - float * restrict dst2, float * restrict dst3, const float * restrict mul, - const float * restrict dot, uint32_t n, float * restrict sums) { +static inline HVX_Vector gdn_mul_dot4_f32(float * restrict dst0, float * restrict dst1, + float * restrict dst2, float * restrict dst3, + const HVX_Vector * restrict mul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc0 = Q6_V_vzero(); HVX_Vector acc1 = Q6_V_vzero(); HVX_Vector acc2 = Q6_V_vzero(); @@ -131,8 +125,8 @@ static inline void gdn_mul_dot4_f32(float * restrict dst0, float * restrict dst1 const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { - HVX_Vector vm = hvx_vmem(mul + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vm = mul[i]; + HVX_Vector vdot = dot[i]; HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vm); HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vm); @@ -152,8 +146,8 @@ static inline void gdn_mul_dot4_f32(float * restrict dst0, float * restrict dst1 if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vm = hvx_vmem(mul + off); - HVX_Vector vdot = hvx_vmem(dot + off); + HVX_Vector vm = mul[nvec]; + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); HVX_Vector zero = Q6_V_vzero(); @@ -174,23 +168,22 @@ static inline void gdn_mul_dot4_f32(float * restrict dst0, float * restrict dst1 } HVX_Vector_x4 acc = { .v = { acc0, acc1, acc2, acc3 } }; - hvx_vec_store_u(sums, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(acc)); + return hvx_vec_reduce_sum_f32x4(acc); } -static inline void gdn_mul_scalar_dot4_f32(float * restrict dst0, float * restrict dst1, - float * restrict dst2, float * restrict dst3, float mul, - const float * restrict dot, uint32_t n, float * restrict sums) { +static inline HVX_Vector gdn_mul_scalar_dot4_f32(float * restrict dst0, float * restrict dst1, + float * restrict dst2, float * restrict dst3, + HVX_Vector vmul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc0 = Q6_V_vzero(); HVX_Vector acc1 = Q6_V_vzero(); HVX_Vector acc2 = Q6_V_vzero(); HVX_Vector acc3 = Q6_V_vzero(); - const HVX_Vector vmul = hvx_vec_splat_f32(mul); const uint32_t epv = 128 / sizeof(float); const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vdot = dot[i]; HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vmul); HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vmul); @@ -210,7 +203,7 @@ static inline void gdn_mul_scalar_dot4_f32(float * restrict dst0, float * restri if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vdot = hvx_vmem(dot + off); + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); HVX_Vector zero = Q6_V_vzero(); @@ -231,13 +224,13 @@ static inline void gdn_mul_scalar_dot4_f32(float * restrict dst0, float * restri } HVX_Vector_x4 acc = { .v = { acc0, acc1, acc2, acc3 } }; - hvx_vec_store_u(sums, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(acc)); + return hvx_vec_reduce_sum_f32x4(acc); } -static inline void gdn_add_scaled_dot4_f32(float * restrict dst0, float * restrict dst1, - float * restrict dst2, float * restrict dst3, const float * restrict src, - const float * restrict scale, const float * restrict dot, uint32_t n, - float * restrict sums) { +static inline HVX_Vector gdn_add_scaled_dot4_f32(float * restrict dst0, float * restrict dst1, + float * restrict dst2, float * restrict dst3, + const HVX_Vector * restrict src, const float * restrict scale, + const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc0 = Q6_V_vzero(); HVX_Vector acc1 = Q6_V_vzero(); HVX_Vector acc2 = Q6_V_vzero(); @@ -251,8 +244,8 @@ static inline void gdn_add_scaled_dot4_f32(float * restrict dst0, float * restri const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { - HVX_Vector vs = hvx_vmem(src + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vs = src[i]; + HVX_Vector vdot = dot[i]; HVX_Vector out0 = hvx_vec_add_f32_f32(hvx_vmemu(dst0 + i * epv), hvx_vec_mul_f32_f32(vs, scale0)); HVX_Vector out1 = hvx_vec_add_f32_f32(hvx_vmemu(dst1 + i * epv), hvx_vec_mul_f32_f32(vs, scale1)); @@ -272,8 +265,8 @@ static inline void gdn_add_scaled_dot4_f32(float * restrict dst0, float * restri if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vs = hvx_vmem(src + off); - HVX_Vector vdot = hvx_vmem(dot + off); + HVX_Vector vs = src[nvec]; + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); HVX_Vector zero = Q6_V_vzero(); @@ -294,14 +287,13 @@ static inline void gdn_add_scaled_dot4_f32(float * restrict dst0, float * restri } HVX_Vector_x4 acc = { .v = { acc0, acc1, acc2, acc3 } }; - hvx_vec_store_u(sums, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(acc)); + return hvx_vec_reduce_sum_f32x4(acc); } -static inline void gdn_mul_dot8_f32(float * restrict dst0, float * restrict dst1, +static inline HVX_Vector gdn_mul_dot8_f32(float * restrict dst0, float * restrict dst1, float * restrict dst2, float * restrict dst3, float * restrict dst4, float * restrict dst5, float * restrict dst6, float * restrict dst7, - const float * restrict mul, const float * restrict dot, uint32_t n, - float * restrict sums) { + const HVX_Vector * restrict mul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc0 = Q6_V_vzero(); HVX_Vector acc1 = Q6_V_vzero(); HVX_Vector acc2 = Q6_V_vzero(); @@ -315,8 +307,8 @@ static inline void gdn_mul_dot8_f32(float * restrict dst0, float * restrict dst1 const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { - HVX_Vector vm = hvx_vmem(mul + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vm = mul[i]; + HVX_Vector vdot = dot[i]; HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vm); HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vm); @@ -348,8 +340,8 @@ static inline void gdn_mul_dot8_f32(float * restrict dst0, float * restrict dst1 if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vm = hvx_vmem(mul + off); - HVX_Vector vdot = hvx_vmem(dot + off); + HVX_Vector vm = mul[nvec]; + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); HVX_Vector zero = Q6_V_vzero(); @@ -383,14 +375,16 @@ static inline void gdn_mul_dot8_f32(float * restrict dst0, float * restrict dst1 HVX_Vector_x4 accA = { .v = { acc0, acc1, acc2, acc3 } }; HVX_Vector_x4 accB = { .v = { acc4, acc5, acc6, acc7 } }; - hvx_vec_store_u(sums + 0, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accA)); - hvx_vec_store_u(sums + 4, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accB)); + HVX_Vector rA = hvx_vec_reduce_sum_f32x4(accA); + HVX_Vector rB = hvx_vec_reduce_sum_f32x4(accB); + HVX_VectorPred q16 = Q6_Q_vsetq2_R(16); + return Q6_V_vmux_QVV(q16, rA, Q6_V_vror_VR(rB, 128 - 16)); } -static inline void gdn_mul_scalar_dot8_f32(float * restrict dst0, float * restrict dst1, +static inline HVX_Vector gdn_mul_scalar_dot8_f32(float * restrict dst0, float * restrict dst1, float * restrict dst2, float * restrict dst3, float * restrict dst4, float * restrict dst5, float * restrict dst6, float * restrict dst7, - float mul, const float * restrict dot, uint32_t n, float * restrict sums) { + HVX_Vector vmul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc0 = Q6_V_vzero(); HVX_Vector acc1 = Q6_V_vzero(); HVX_Vector acc2 = Q6_V_vzero(); @@ -399,13 +393,12 @@ static inline void gdn_mul_scalar_dot8_f32(float * restrict dst0, float * restri HVX_Vector acc5 = Q6_V_vzero(); HVX_Vector acc6 = Q6_V_vzero(); HVX_Vector acc7 = Q6_V_vzero(); - const HVX_Vector vmul = hvx_vec_splat_f32(mul); const uint32_t epv = 128 / sizeof(float); const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vdot = dot[i]; HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vmul); HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vmul); @@ -437,7 +430,7 @@ static inline void gdn_mul_scalar_dot8_f32(float * restrict dst0, float * restri if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vdot = hvx_vmem(dot + off); + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); HVX_Vector zero = Q6_V_vzero(); @@ -471,15 +464,17 @@ static inline void gdn_mul_scalar_dot8_f32(float * restrict dst0, float * restri HVX_Vector_x4 accA = { .v = { acc0, acc1, acc2, acc3 } }; HVX_Vector_x4 accB = { .v = { acc4, acc5, acc6, acc7 } }; - hvx_vec_store_u(sums + 0, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accA)); - hvx_vec_store_u(sums + 4, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accB)); + HVX_Vector rA = hvx_vec_reduce_sum_f32x4(accA); + HVX_Vector rB = hvx_vec_reduce_sum_f32x4(accB); + HVX_VectorPred q16 = Q6_Q_vsetq2_R(16); + return Q6_V_vmux_QVV(q16, rA, Q6_V_vror_VR(rB, 128 - 16)); } -static inline void gdn_add_scaled_dot8_f32(float * restrict dst0, float * restrict dst1, +static inline HVX_Vector gdn_add_scaled_dot8_f32(float * restrict dst0, float * restrict dst1, float * restrict dst2, float * restrict dst3, float * restrict dst4, float * restrict dst5, float * restrict dst6, float * restrict dst7, - const float * restrict src, const float * restrict scale, - const float * restrict dot, uint32_t n, float * restrict sums) { + const HVX_Vector * restrict src, const float * restrict scale, + const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc0 = Q6_V_vzero(); HVX_Vector acc1 = Q6_V_vzero(); HVX_Vector acc2 = Q6_V_vzero(); @@ -501,8 +496,8 @@ static inline void gdn_add_scaled_dot8_f32(float * restrict dst0, float * restri const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; for (uint32_t i = 0; i < nvec; ++i) { - HVX_Vector vs = hvx_vmem(src + i * epv); - HVX_Vector vdot = hvx_vmem(dot + i * epv); + HVX_Vector vs = src[i]; + HVX_Vector vdot = dot[i]; HVX_Vector out0 = hvx_vec_add_f32_f32(hvx_vmemu(dst0 + i * epv), hvx_vec_mul_f32_f32(vs, scale0)); HVX_Vector out1 = hvx_vec_add_f32_f32(hvx_vmemu(dst1 + i * epv), hvx_vec_mul_f32_f32(vs, scale1)); @@ -534,8 +529,8 @@ static inline void gdn_add_scaled_dot8_f32(float * restrict dst0, float * restri if (nloe) { const uint32_t off = nvec * epv; - HVX_Vector vs = hvx_vmem(src + off); - HVX_Vector vdot = hvx_vmem(dot + off); + HVX_Vector vs = src[nvec]; + HVX_Vector vdot = dot[nvec]; HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float)); HVX_Vector zero = Q6_V_vzero(); @@ -569,13 +564,196 @@ static inline void gdn_add_scaled_dot8_f32(float * restrict dst0, float * restri HVX_Vector_x4 accA = { .v = { acc0, acc1, acc2, acc3 } }; HVX_Vector_x4 accB = { .v = { acc4, acc5, acc6, acc7 } }; - hvx_vec_store_u(sums + 0, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accA)); - hvx_vec_store_u(sums + 4, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accB)); + HVX_Vector rA = hvx_vec_reduce_sum_f32x4(accA); + HVX_Vector rB = hvx_vec_reduce_sum_f32x4(accB); + HVX_VectorPred q16 = Q6_Q_vsetq2_R(16); + return Q6_V_vmux_QVV(q16, rA, Q6_V_vror_VR(rB, 128 - 16)); +} + +static inline void gdn_step_kda_f32( + float * restrict s_work, + float * restrict attn_out, + const float * restrict q_t, + const float * restrict k_t, + const float * restrict v_t, + const float * restrict g_t, + float beta_val, + float scale, + uint32_t S_v +) { + const uint32_t epv = 128 / sizeof(float); + const uint32_t nvec = S_v / epv; + const uint32_t nloe = S_v % epv; + + HVX_Vector vq[4]; + HVX_Vector vk[4]; + HVX_Vector vg[4]; + + 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); + + for (uint32_t i = 0; i < nvec; ++i) { + vq[i] = hvx_vmemu(q_t + i * epv); + vk[i] = hvx_vmemu(k_t + i * epv); + vg[i] = hvx_vec_exp_f32_guard(hvx_vmemu(g_t + i * epv), max_exp, inf); + } + if (nloe) { + vq[nvec] = hvx_vmemu(q_t + nvec * epv); + vk[nvec] = hvx_vmemu(k_t + nvec * epv); + vg[nvec] = hvx_vec_exp_f32_guard(hvx_vmemu(g_t + nvec * epv), max_exp, inf); + } + + const HVX_Vector vbeta = hvx_vec_splat_f32(beta_val); + const HVX_Vector vscale = hvx_vec_splat_f32(scale); + + float delta[8] __attribute__((aligned(128))); + + uint32_t j = 0; + for (; j + 8 <= S_v; j += 8) { + float * row0 = s_work + (uint64_t) (j + 0) * S_v; + float * row1 = s_work + (uint64_t) (j + 1) * S_v; + float * row2 = s_work + (uint64_t) (j + 2) * S_v; + float * row3 = s_work + (uint64_t) (j + 3) * S_v; + float * row4 = s_work + (uint64_t) (j + 4) * S_v; + float * row5 = s_work + (uint64_t) (j + 5) * S_v; + float * row6 = s_work + (uint64_t) (j + 6) * S_v; + float * row7 = s_work + (uint64_t) (j + 7) * S_v; + + HVX_Vector vsums = gdn_mul_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, + vg, vk, S_v); + + HVX_Vector vv_t = hvx_vmemu(v_t + j); + HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, vsums); + HVX_Vector vdelta = hvx_vec_mul_f32_f32(diff, vbeta); + hvx_vec_store_u(delta, 8 * sizeof(float), vdelta); + + HVX_Vector vattn = gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, + vk, delta, vq, S_v); + + HVX_Vector res_attn = hvx_vec_mul_f32_f32(vattn, vscale); + hvx_vec_store_u(attn_out + j, 8 * sizeof(float), res_attn); + } + for (; j + 4 <= S_v; j += 4) { + float * row0 = s_work + (uint64_t) (j + 0) * S_v; + float * row1 = s_work + (uint64_t) (j + 1) * S_v; + float * row2 = s_work + (uint64_t) (j + 2) * S_v; + float * row3 = s_work + (uint64_t) (j + 3) * S_v; + + HVX_Vector vsums = gdn_mul_dot4_f32(row0, row1, row2, row3, vg, vk, S_v); + + HVX_Vector vv_t = hvx_vmemu(v_t + j); + HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, vsums); + HVX_Vector vdelta = hvx_vec_mul_f32_f32(diff, vbeta); + hvx_vec_store_u(delta, 4 * sizeof(float), vdelta); + + HVX_Vector vattn = gdn_add_scaled_dot4_f32(row0, row1, row2, row3, vk, delta, vq, S_v); + + HVX_Vector res_attn = hvx_vec_mul_f32_f32(vattn, vscale); + hvx_vec_store_u(attn_out + j, 4 * sizeof(float), res_attn); + } + for (; j < S_v; ++j) { + float * row = s_work + (uint64_t) j * S_v; + HVX_Vector vsum = gdn_mul_dot_f32(row, vg, vk, S_v); + HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); + HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), vbeta); + HVX_Vector vres = gdn_add_scaled_dot_f32(row, vk, vdj, vq, S_v); + attn_out[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale)); + } +} + +static inline void gdn_step_scalar_f32( + float * restrict s_work, + float * restrict attn_out, + const float * restrict q_t, + const float * restrict k_t, + const float * restrict v_t, + const float * restrict g_t, + float beta_val, + float scale, + uint32_t S_v +) { + const uint32_t epv = 128 / sizeof(float); + const uint32_t nvec = S_v / epv; + const uint32_t nloe = S_v % epv; + + HVX_Vector vq[4]; + HVX_Vector vk[4]; + + for (uint32_t i = 0; i < nvec; ++i) { + vq[i] = hvx_vmemu(q_t + i * epv); + vk[i] = hvx_vmemu(k_t + i * epv); + } + if (nloe) { + vq[nvec] = hvx_vmemu(q_t + nvec * epv); + vk[nvec] = hvx_vmemu(k_t + nvec * epv); + } + + const float gate = expf(g_t[0]); + const HVX_Vector vgate = hvx_vec_splat_f32(gate); + const HVX_Vector vbeta = hvx_vec_splat_f32(beta_val); + const HVX_Vector vscale = hvx_vec_splat_f32(scale); + + float delta[8] __attribute__((aligned(128))); + + uint32_t j = 0; + for (; j + 8 <= S_v; j += 8) { + float * row0 = s_work + (uint64_t) (j + 0) * S_v; + float * row1 = s_work + (uint64_t) (j + 1) * S_v; + float * row2 = s_work + (uint64_t) (j + 2) * S_v; + float * row3 = s_work + (uint64_t) (j + 3) * S_v; + float * row4 = s_work + (uint64_t) (j + 4) * S_v; + float * row5 = s_work + (uint64_t) (j + 5) * S_v; + float * row6 = s_work + (uint64_t) (j + 6) * S_v; + float * row7 = s_work + (uint64_t) (j + 7) * S_v; + + HVX_Vector vsums = gdn_mul_scalar_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, + vgate, vk, S_v); + + HVX_Vector vv_t = hvx_vmemu(v_t + j); + HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, vsums); + HVX_Vector vdelta = hvx_vec_mul_f32_f32(diff, vbeta); + hvx_vec_store_u(delta, 8 * sizeof(float), vdelta); + + HVX_Vector vattn = gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, + vk, delta, vq, S_v); + + HVX_Vector res_attn = hvx_vec_mul_f32_f32(vattn, vscale); + hvx_vec_store_u(attn_out + j, 8 * sizeof(float), res_attn); + } + for (; j + 4 <= S_v; j += 4) { + float * row0 = s_work + (uint64_t) (j + 0) * S_v; + float * row1 = s_work + (uint64_t) (j + 1) * S_v; + float * row2 = s_work + (uint64_t) (j + 2) * S_v; + float * row3 = s_work + (uint64_t) (j + 3) * S_v; + + HVX_Vector vsums = gdn_mul_scalar_dot4_f32(row0, row1, row2, row3, vgate, vk, S_v); + + HVX_Vector vv_t = hvx_vmemu(v_t + j); + HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, vsums); + HVX_Vector vdelta = hvx_vec_mul_f32_f32(diff, vbeta); + hvx_vec_store_u(delta, 4 * sizeof(float), vdelta); + + HVX_Vector vattn = gdn_add_scaled_dot4_f32(row0, row1, row2, row3, vk, delta, vq, S_v); + + HVX_Vector res_attn = hvx_vec_mul_f32_f32(vattn, vscale); + hvx_vec_store_u(attn_out + j, 4 * sizeof(float), res_attn); + } + for (; j < S_v; ++j) { + float * row = s_work + (uint64_t) j * S_v; + HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, gate, vk, S_v); + HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); + HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), vbeta); + HVX_Vector vres = gdn_add_scaled_dot_f32(row, vk, vdj, vq, S_v); + attn_out[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale)); + } } static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, void * data) { struct htp_gdn_context * gctx = (struct htp_gdn_context *) data; struct htp_ops_context * octx = gctx->octx; + const struct htp_gdn_kernel_params * kparams = gctx->kparams; const struct htp_tensor * q = octx->src[0]; const struct htp_tensor * k = octx->src[1]; @@ -585,66 +763,55 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo const struct htp_tensor * state = octx->src[5]; const struct htp_tensor * dst = octx->dst; - const uint32_t S_v = v->ne[0]; - const uint32_t H = v->ne[1]; - const uint32_t n_tokens = v->ne[2]; - const uint32_t n_seqs = v->ne[3]; - const uint32_t K = octx->op_params[0]; - - const uint32_t row_end = gctx->row_start + gctx->nrows; + const uint32_t S_v = kparams->S_v; + const uint32_t H = kparams->H; + const uint32_t n_tokens = kparams->n_tokens; + const uint32_t n_seqs = kparams->n_seqs; + const uint32_t K = kparams->K; + const uint32_t row_end = gctx->row_start + gctx->nrows; if (ith >= gctx->nrows) { return; } - const uint32_t rq3 = n_seqs / q->ne[3]; - const uint32_t rk3 = n_seqs / k->ne[3]; - const float scale = 1.0f / sqrtf((float) S_v); - + const struct htp_tensor * dst_cache = octx->dsts[1]; + const float scale = kparams->scale; float * dst_base = (float *) (uintptr_t) dst->data; - float * state_out_base = dst_base + (uint64_t) S_v * H * n_tokens * n_seqs; - const float * state_in_base = (const float *) (uintptr_t) state->data; + float * state_out_base = dst_cache ? (float *) (uintptr_t) dst_cache->data : (dst_base + S_v * H * n_tokens * n_seqs); - const bool kda = (g->ne[0] == S_v); - float local_gate[HTP_GDN_MAX_SV] __attribute__((aligned(128))); - float local_q[HTP_GDN_MAX_SV] __attribute__((aligned(128))); - float local_k[HTP_GDN_MAX_SV] __attribute__((aligned(128))); - float local_sums[32] __attribute__((aligned(128))); - - dma_queue * dma = octx->ctx->dma[ith]; - size_t state_aligned = (size_t) S_v * S_v * sizeof(float); - state_aligned = (state_aligned + 127) & ~(size_t)127; + dma_queue * dma_q = octx->ctx->dma[ith]; + const struct htp_gdn_vtcm_layout * layout = &gctx->layout; float * s_work[2]; - s_work[0] = (float *) (gctx->vtcm_base + gctx->vtcm_per_thread * ith); - s_work[1] = s_work[0] + state_aligned / sizeof(float); + s_work[0] = (float *) (gctx->vtcm_base + layout->bytes_per_thread * ith); + s_work[1] = s_work[0] + layout->state_aligned / sizeof(float); - struct fastdiv_values fd_H = init_fastdiv_values(H); - struct fastdiv_values fd_q1 = init_fastdiv_values(q->ne[1]); - struct fastdiv_values fd_k1 = init_fastdiv_values(k->ne[1]); - struct fastdiv_values fd_rq3 = init_fastdiv_values(rq3); - struct fastdiv_values fd_rk3 = init_fastdiv_values(rk3); + const struct fastdiv_values * fd_H = &kparams->div_H; + const struct fastdiv_values * fd_q1 = &kparams->div_q1; + const struct fastdiv_values * fd_k1 = &kparams->div_k1; + const struct fastdiv_values * fd_rq3 = &kparams->div_rq3; + const struct fastdiv_values * fd_rk3 = &kparams->div_rk3; - const uint64_t state_seq_stride = state->nb[3] / sizeof(float); - const uint64_t state_size_per_snap = (uint64_t) S_v * S_v * H * n_seqs; + const uint32_t state_seq_stride = kparams->state_seq_stride; + const uint64_t state_size_per_snap = (uint64_t) kparams->state_size_per_snap; + const dma_addr_t state_out_dma_base = dst_cache ? dst_cache->data : (dst->data + S_v * H * n_tokens * n_seqs * sizeof(float)); uint32_t ir_prefetch = gctx->row_start + ith; int spad_idx = 0; // Prefetch preamble (up to 2 steps) - for (int k = 0; k < 2 && ir_prefetch < row_end; k++) { - const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H); - const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H); - const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v; - // final state lands in snapshot slot 0 (most-recent-first ordering) - float * ps_out = state_out_base + ((uint64_t) piv3 * H + piv1) * S_v * S_v; + for (int step = 0; step < 2 && ir_prefetch < row_end; step++) { + const uint32_t piv1 = fastmodulo(ir_prefetch, H, fd_H); + const uint32_t piv3 = fastdiv(ir_prefetch, fd_H); + dma_addr_t ps_in = state->data + ((uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v) * sizeof(float); + dma_addr_t ps_out = state_out_dma_base + ((uint64_t) piv3 * H + piv1) * S_v * S_v * sizeof(float); // Push dummy write-back - dma_queue_push(dma, dma_make_ptr(ps_out, s_work[spad_idx]), + dma_queue_push(dma_q, dma_make_data(ps_out, s_work[spad_idx]), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), 0); // Push fetch - dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in), + dma_queue_push(dma_q, dma_make_data(s_work[spad_idx], ps_in), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); @@ -653,28 +820,26 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith)); int curr_spad_idx = 0; for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) { - dma_queue_pop(dma); - dma_queue_pop(dma); + dma_queue_pop(dma_q); + dma_queue_pop(dma_q); float * s_work_curr = s_work[curr_spad_idx]; - const uint32_t iv1 = fastmodulo(ir, H, &fd_H); - const uint32_t iv3 = fastdiv(ir, &fd_H); + const uint32_t iv1 = fastmodulo(ir, H, fd_H); + const uint32_t iv3 = fastdiv(ir, fd_H); - const uint32_t iq1 = fastmodulo(iv1, q->ne[1], &fd_q1); - const uint32_t ik1 = fastmodulo(iv1, k->ne[1], &fd_k1); - const uint32_t iq3 = fastdiv(iv3, &fd_rq3); - const uint32_t ik3 = fastdiv(iv3, &fd_rk3); - - // final state lands in snapshot slot 0 (most-recent-first ordering) - float * s_out = state_out_base + ((uint64_t) iv3 * H + iv1) * S_v * S_v; + const uint32_t iq1 = fastmodulo(iv1, q->ne[1], fd_q1); + const uint32_t ik1 = fastmodulo(iv1, k->ne[1], fd_k1); + const uint32_t iq3 = fastdiv(iv3, fd_rq3); + const uint32_t ik3 = fastdiv(iv3, fd_rk3); + dma_addr_t s_out = state_out_dma_base + ((uint64_t) iv3 * H + iv1) * S_v * S_v * sizeof(float); float * attn_data = dst_base + ((uint64_t) iv3 * n_tokens * H + iv1) * S_v; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); for (uint32_t t = 0; t < n_tokens; ++t) { const float * q_t = (const float *) ((const uint8_t *) (uintptr_t) q->data + (uint64_t) iq3 * q->nb[3] + (uint64_t) t * q->nb[2] + (uint64_t) iq1 * q->nb[1]); @@ -687,146 +852,36 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo const float beta_val = *(const float *) ((const uint8_t *) (uintptr_t) beta->data + (uint64_t) iv3 * beta->nb[3] + (uint64_t) t * beta->nb[2] + (uint64_t) iv1 * beta->nb[1]); - hvx_copy_f32_au((uint8_t *) local_q, (const uint8_t *) q_t, S_v); - hvx_copy_f32_au((uint8_t *) local_k, (const uint8_t *) k_t, S_v); - - if (kda) { - hvx_exp_f32((uint8_t *) local_gate, (const uint8_t *) g_t, S_v, false); - - uint32_t j = 0; - for (; j + 8 <= S_v; j += 8) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v; - float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v; - float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v; - float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v; - gdn_mul_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - local_gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn); - } - for (; j + 4 <= S_v; j += 4) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - gdn_mul_dot4_f32(row0, row1, row2, row3, local_gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn); - } - HVX_Vector vscale_splat = hvx_vec_splat_f32(scale); - for (; j < S_v; ++j) { - float * row = s_work_curr + (uint64_t) j * S_v; - HVX_Vector vsum = gdn_mul_dot_f32(row, local_gate, local_k, S_v); - HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); - HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val)); - HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v); - attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat)); - } + if (kparams->kda) { + gdn_step_kda_f32(s_work_curr, attn_data, q_t, k_t, v_t, g_t, beta_val, scale, S_v); } else { - const float gate = expf(g_t[0]); - uint32_t j = 0; - for (; j + 8 <= S_v; j += 8) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v; - float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v; - float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v; - float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v; - gdn_mul_scalar_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn); - } - for (; j + 4 <= S_v; j += 4) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - gdn_mul_scalar_dot4_f32(row0, row1, row2, row3, gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn); - } - HVX_Vector vscale_splat = hvx_vec_splat_f32(scale); - for (; j < S_v; ++j) { - float * row = s_work_curr + (uint64_t) j * S_v; - HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, gate, local_k, S_v); - HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); - HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val)); - HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v); - attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat)); - } + gdn_step_scalar_f32(s_work_curr, attn_data, q_t, k_t, v_t, g_t, beta_val, scale, S_v); } if (K > 1) { - // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back. const int64_t target_slot = (int64_t) n_tokens - 1 - (int64_t) t; - if (target_slot >= 0 && target_slot < (int64_t) K) { + if (target_slot > 0 && target_slot < (int64_t) K) { float * curr_state_o = state_out_base + (uint64_t) target_slot * state_size_per_snap + ((uint64_t) iv3 * H + iv1) * S_v * S_v; - if (curr_state_o != s_out) { - hvx_copy_f32_uu((uint8_t *) curr_state_o, (const uint8_t *) s_work_curr, S_v * S_v); - } + hvx_copy_f32_uu((uint8_t *) curr_state_o, (const uint8_t *) s_work_curr, S_v * S_v); } } attn_data += (uint64_t) S_v * H; } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); // Push real write-back - dma_queue_push(dma, dma_make_ptr(s_out, s_work_curr), + dma_queue_push(dma_q, dma_make_data(s_out, s_work_curr), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); // Prefetch next block (if any) if (ir_prefetch < row_end) { - const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H); - const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H); - const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v; + const uint32_t piv1 = fastmodulo(ir_prefetch, H, fd_H); + const uint32_t piv3 = fastdiv(ir_prefetch, fd_H); + dma_addr_t ps_in = state->data + ((uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v) * sizeof(float); - dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in), + dma_queue_push(dma_q, dma_make_data(s_work[spad_idx], ps_in), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); @@ -836,14 +891,13 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo curr_spad_idx ^= 1; } - dma_queue_flush(dma); - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end); + dma_queue_flush(dma_q); } - static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, void * data) { struct htp_gdn_context * gctx = (struct htp_gdn_context *) data; struct htp_ops_context * octx = gctx->octx; + const struct htp_gdn_kernel_params * kparams = gctx->kparams; const struct htp_tensor * q = octx->src[0]; const struct htp_tensor * k = octx->src[1]; @@ -853,63 +907,51 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo const struct htp_tensor * state = octx->src[5]; const struct htp_tensor * dst = octx->dst; - const uint32_t S_v = v->ne[0]; - const uint32_t H = v->ne[1]; - const uint32_t n_seqs = v->ne[3]; - - const uint32_t row_end = gctx->row_start + gctx->nrows; + const uint32_t S_v = kparams->S_v; + const uint32_t H = kparams->H; + const uint32_t n_seqs = kparams->n_seqs; + const uint32_t row_end = gctx->row_start + gctx->nrows; if (ith >= gctx->nrows) { return; } - const uint32_t rq3 = n_seqs / q->ne[3]; - const uint32_t rk3 = n_seqs / k->ne[3]; - const float scale = 1.0f / sqrtf((float) S_v); + const struct htp_tensor * dst_cache = octx->dsts[1]; + const float scale = kparams->scale; + float * dst_base = (float *) (uintptr_t) dst->data; - float * dst_base = (float *) (uintptr_t) dst->data; - float * state_out_base = dst_base + (uint64_t) S_v * H * n_seqs; - const float * state_in_base = (const float *) (uintptr_t) state->data; - - const bool kda = (g->ne[0] == S_v); - float local_gate[HTP_GDN_MAX_SV] __attribute__((aligned(128))); - float local_q[HTP_GDN_MAX_SV] __attribute__((aligned(128))); - float local_k[HTP_GDN_MAX_SV] __attribute__((aligned(128))); - float local_sums[32] __attribute__((aligned(128))); - - dma_queue * dma = octx->ctx->dma[ith]; - size_t state_aligned = (size_t) S_v * S_v * sizeof(float); - state_aligned = (state_aligned + 127) & ~(size_t)127; + dma_queue * dma_q = octx->ctx->dma[ith]; + const struct htp_gdn_vtcm_layout * layout = &gctx->layout; float * s_work[2]; - s_work[0] = (float *) (gctx->vtcm_base + gctx->vtcm_per_thread * ith); - s_work[1] = s_work[0] + state_aligned / sizeof(float); + s_work[0] = (float *) (gctx->vtcm_base + layout->bytes_per_thread * ith); + s_work[1] = s_work[0] + layout->state_aligned / sizeof(float); - struct fastdiv_values fd_H = init_fastdiv_values(H); - struct fastdiv_values fd_q1 = init_fastdiv_values(q->ne[1]); - struct fastdiv_values fd_k1 = init_fastdiv_values(k->ne[1]); - struct fastdiv_values fd_rq3 = init_fastdiv_values(rq3); - struct fastdiv_values fd_rk3 = init_fastdiv_values(rk3); + const struct fastdiv_values * fd_H = &kparams->div_H; + const struct fastdiv_values * fd_q1 = &kparams->div_q1; + const struct fastdiv_values * fd_k1 = &kparams->div_k1; + const struct fastdiv_values * fd_rq3 = &kparams->div_rq3; + const struct fastdiv_values * fd_rk3 = &kparams->div_rk3; - const uint64_t state_seq_stride = state->nb[3] / sizeof(float); + const uint32_t state_seq_stride = kparams->state_seq_stride; + const dma_addr_t state_out_dma_base = dst_cache ? dst_cache->data : (dst->data + S_v * H * n_seqs * sizeof(float)); uint32_t ir_prefetch = gctx->row_start + ith; int spad_idx = 0; // Prefetch preamble (up to 2 steps) - for (int k = 0; k < 2 && ir_prefetch < row_end; k++) { - const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H); - const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H); - const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v; - // final state lands in snapshot slot 0 (most-recent-first ordering) - float * ps_out = state_out_base + ((uint64_t) piv3 * H + piv1) * S_v * S_v; + for (int step = 0; step < 2 && ir_prefetch < row_end; step++) { + const uint32_t piv1 = fastmodulo(ir_prefetch, H, fd_H); + const uint32_t piv3 = fastdiv(ir_prefetch, fd_H); + dma_addr_t ps_in = state->data + ((uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v) * sizeof(float); + dma_addr_t ps_out = state_out_dma_base + ((uint64_t) piv3 * H + piv1) * S_v * S_v * sizeof(float); // Push dummy write-back - dma_queue_push(dma, dma_make_ptr(ps_out, s_work[spad_idx]), + dma_queue_push(dma_q, dma_make_data(ps_out, s_work[spad_idx]), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), 0); // Push fetch - dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in), + dma_queue_push(dma_q, dma_make_data(s_work[spad_idx], ps_in), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); @@ -918,26 +960,23 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith)); int curr_spad_idx = 0; for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) { - dma_queue_pop(dma); - dma_queue_pop(dma); + dma_queue_pop(dma_q); + dma_queue_pop(dma_q); float * s_work_curr = s_work[curr_spad_idx]; - const uint32_t iv1 = fastmodulo(ir, H, &fd_H); - const uint32_t iv3 = fastdiv(ir, &fd_H); + const uint32_t iv1 = fastmodulo(ir, H, fd_H); + const uint32_t iv3 = fastdiv(ir, fd_H); - const uint32_t iq1 = fastmodulo(iv1, q->ne[1], &fd_q1); - const uint32_t ik1 = fastmodulo(iv1, k->ne[1], &fd_k1); - const uint32_t iq3 = fastdiv(iv3, &fd_rq3); - const uint32_t ik3 = fastdiv(iv3, &fd_rk3); - - // final state lands in snapshot slot 0 (most-recent-first ordering) - float * s_out = state_out_base + ((uint64_t) iv3 * H + iv1) * S_v * S_v; + const uint32_t iq1 = fastmodulo(iv1, q->ne[1], fd_q1); + const uint32_t ik1 = fastmodulo(iv1, k->ne[1], fd_k1); + const uint32_t iq3 = fastdiv(iv3, fd_rq3); + const uint32_t ik3 = fastdiv(iv3, fd_rk3); + dma_addr_t s_out = state_out_dma_base + ((uint64_t) iv3 * H + iv1) * S_v * S_v * sizeof(float); float * attn_data = dst_base + ((uint64_t) iv3 * H + iv1) * S_v; const float * q_t = (const float *) ((const uint8_t *) (uintptr_t) q->data + @@ -951,132 +990,26 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo const float beta_val = *(const float *) ((const uint8_t *) (uintptr_t) beta->data + (uint64_t) iv3 * beta->nb[3] + (uint64_t) iv1 * beta->nb[1]); - hvx_copy_f32_au((uint8_t *) local_q, (const uint8_t *) q_t, S_v); - hvx_copy_f32_au((uint8_t *) local_k, (const uint8_t *) k_t, S_v); - - if (kda) { - hvx_exp_f32((uint8_t *) local_gate, (const uint8_t *) g_t, S_v, false); - - uint32_t j = 0; - for (; j + 8 <= S_v; j += 8) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v; - float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v; - float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v; - float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v; - gdn_mul_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - local_gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn); - } - for (; j + 4 <= S_v; j += 4) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - gdn_mul_dot4_f32(row0, row1, row2, row3, local_gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn); - } - HVX_Vector vscale_splat = hvx_vec_splat_f32(scale); - for (; j < S_v; ++j) { - float * row = s_work_curr + (uint64_t) j * S_v; - HVX_Vector vsum = gdn_mul_dot_f32(row, local_gate, local_k, S_v); - HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); - HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val)); - HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v); - attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat)); - } + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); + if (kparams->kda) { + gdn_step_kda_f32(s_work_curr, attn_data, q_t, k_t, v_t, g_t, beta_val, scale, S_v); } else { - const float gate = expf(g_t[0]); - uint32_t j = 0; - for (; j + 8 <= S_v; j += 8) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v; - float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v; - float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v; - float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v; - gdn_mul_scalar_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7, - local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn); - } - for (; j + 4 <= S_v; j += 4) { - float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v; - float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v; - float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v; - float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v; - gdn_mul_scalar_dot4_f32(row0, row1, row2, row3, gate, local_k, S_v, local_sums); - - float local_delta_b[32] __attribute__((aligned(128))); - HVX_Vector vv_t = hvx_vmemu(v_t + j); - HVX_Vector v_local_sums = hvx_vmem(local_sums); - HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums); - hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val)); - - gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums); - - HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale)); - hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn); - } - HVX_Vector vscale_splat = hvx_vec_splat_f32(scale); - for (; j < S_v; ++j) { - float * row = s_work_curr + (uint64_t) j * S_v; - HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, gate, local_k, S_v); - HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); - HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val)); - HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v); - attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat)); - } + gdn_step_scalar_f32(s_work_curr, attn_data, q_t, k_t, v_t, g_t, beta_val, scale, S_v); } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); // Push real write-back - dma_queue_push(dma, dma_make_ptr(s_out, s_work_curr), + dma_queue_push(dma_q, dma_make_data(s_out, s_work_curr), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); // Prefetch next block (if any) if (ir_prefetch < row_end) { - const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H); - const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H); - const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v; + const uint32_t piv1 = fastmodulo(ir_prefetch, H, fd_H); + const uint32_t piv3 = fastdiv(ir_prefetch, fd_H); + dma_addr_t ps_in = state->data + ((uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v) * sizeof(float); - dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in), + dma_queue_push(dma_q, dma_make_data(s_work[spad_idx], ps_in), S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); @@ -1086,11 +1019,9 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo curr_spad_idx ^= 1; } - dma_queue_flush(dma); - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end); + dma_queue_flush(dma_q); } - int op_gated_delta_net(struct htp_ops_context * octx) { const struct htp_tensor * q = octx->src[0]; const struct htp_tensor * k = octx->src[1]; @@ -1131,44 +1062,88 @@ int op_gated_delta_net(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + for (int i = 0; i < 5; i++) { + if (htp_tensor_is_extended(octx->src[i])) { + return HTP_STATUS_NO_SUPPORT; + } + } + if (htp_tensor_is_extended(octx->dst)) { + return HTP_STATUS_NO_SUPPORT; + } + if (octx->dsts[1]) { + const struct htp_tensor * dst_cache = octx->dsts[1]; + if (dst_cache->type != HTP_TYPE_F32 || htp_tensor_is_extended(dst_cache)) { + return HTP_STATUS_NO_SUPPORT; + } } - const uint32_t total_rows = H * n_seqs; + const struct htp_gdn_kernel_params * kparams = (const struct htp_gdn_kernel_params *) octx->kernel_params; + struct htp_gdn_kernel_params kparams_local; + if (!kparams || kparams->S_v == 0) { + const uint32_t rq3 = n_seqs / q->ne[3]; + const uint32_t rk3 = n_seqs / k->ne[3]; + const uint32_t total_rows = H * n_seqs; + uint32_t n_threads = (total_rows < octx->n_threads) ? total_rows : octx->n_threads; + if (n_threads == 0) { + n_threads = 1; + } + memset(&kparams_local, 0, sizeof(kparams_local)); + kparams_local.n_threads = n_threads; + kparams_local.S_v = S_v; + kparams_local.H = H; + kparams_local.n_tokens = n_tokens; + kparams_local.n_seqs = n_seqs; + kparams_local.K = K; + kparams_local.total_rows = total_rows; + kparams_local.rows_per_thread = (total_rows + n_threads - 1) / n_threads; + struct htp_gdn_vtcm_layout layout_local; + htp_gdn_vtcm_layout_build(&layout_local, S_v, n_threads); + kparams_local.state_aligned = (uint32_t) layout_local.state_aligned; + kparams_local.vtcm_per_thread = (uint32_t) layout_local.bytes_per_thread; + kparams_local.vtcm_size = (uint32_t) layout_local.total_bytes; + kparams_local.kda = (g->ne[0] == S_v) ? 1 : 0; + kparams_local.scale = 1.0f / sqrtf((float) S_v); + kparams_local.state_seq_stride = (uint32_t) (state->nb[3] / sizeof(float)); + kparams_local.state_size_per_snap = S_v * S_v * H * n_seqs; + + kparams_local.div_H = init_fastdiv_values(H); + kparams_local.div_q1 = init_fastdiv_values(q->ne[1]); + kparams_local.div_k1 = init_fastdiv_values(k->ne[1]); + kparams_local.div_rq3 = init_fastdiv_values(rq3); + kparams_local.div_rk3 = init_fastdiv_values(rk3); + kparams_local.div_n_threads = init_fastdiv_values(n_threads); + + kparams = &kparams_local; + } + + const uint32_t total_rows = kparams->total_rows; uint32_t row_start = 0; uint32_t nrows = total_rows; - if (octx->ctx->mdev.count > 1) { - const uint32_t head_bytes = S_v * sizeof(float); - const uint32_t rows_per_chunk = (head_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(head_bytes, HEX_L2_LINE_SIZE)) : 1; - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, - octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); - row_start = range.start; - nrows = range.count; + if (octx->op_params[1] != 0) { + row_start = octx->op_params[1]; + nrows = octx->op_params[2]; } if (nrows == 0) { return HTP_STATUS_OK; } - const uint32_t n_threads = octx->n_threads; + const uint32_t n_threads = (nrows < kparams->n_threads) ? nrows : kparams->n_threads; struct htp_gdn_context gctx; - gctx.octx = octx; - gctx.row_start = row_start; - gctx.nrows = nrows; - gctx.rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); - gctx.state_bytes = (size_t) S_v * S_v * sizeof(float); - - size_t state_aligned = (size_t) S_v * S_v * sizeof(float); - state_aligned = (state_aligned + 127) & ~(size_t)127; - - assert(octx->ctx->vtcm_size >= 2 * state_aligned * n_threads); - + gctx.octx = octx; + gctx.kparams = kparams; + gctx.row_start = row_start; + gctx.nrows = nrows; gctx.vtcm_base = octx->ctx->vtcm_base; - gctx.vtcm_per_thread = 2 * state_aligned; + + htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_threads); + + if (gctx.layout.total_bytes > octx->ctx->vtcm_size) { + return HTP_STATUS_VTCM_TOO_SMALL; + } FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : " "vtcm-size %zu n_threads %u\n", @@ -1177,7 +1152,7 @@ int op_gated_delta_net(struct htp_ops_context * octx) { v->ne[0], v->ne[1], v->ne[2], v->ne[3], state->ne[0], state->ne[1], state->ne[2], state->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - gctx.vtcm_per_thread * octx->n_threads, octx->n_threads); + gctx.layout.total_bytes, n_threads); if (n_tokens == 1) { work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_tg_thread, &gctx, n_threads); diff --git a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h new file mode 100644 index 000000000..fd703142e --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h @@ -0,0 +1,297 @@ +#ifndef HTP_GATED_DELTA_NET_OPS_H +#define HTP_GATED_DELTA_NET_OPS_H + +#include +#include +#include + +#include "hex-fastdiv.h" +#include "hex-common.h" +#include "htp-vtcm.h" + +#define HTP_GDN_MAX_SV 128 +#define HTP_GDN_CHUNK_SIZE 64 + +#ifndef HMX_FP16_TILE_SIZE +#define HMX_FP16_TILE_SIZE 2048 +#endif + +enum htp_gdn_kernel_type { + HTP_GDN_KERNEL_HVX_RECURRENT = 0, + HTP_GDN_KERNEL_HMX_CHUNKED = 1, +}; + +struct htp_gdn_kernel_params { + uint8_t kernel_type; + uint8_t pipeline; + uint16_t chunk_size; + uint16_t n_chunks; + uint16_t n_heads_batch; + + uint32_t n_threads; + uint32_t S_v; + uint32_t H; + uint32_t n_tokens; + uint32_t n_seqs; + uint32_t K; + + uint32_t total_rows; + uint32_t row_start; + uint32_t nrows; + uint32_t rows_per_thread; + + uint32_t kda; + uint32_t state_aligned; + uint32_t vtcm_per_thread; + uint32_t vtcm_size; + uint32_t state_seq_stride; + uint32_t state_size_per_snap; + + float scale; + + struct fastdiv_values div_H; + struct fastdiv_values div_q1; + struct fastdiv_values div_k1; + struct fastdiv_values div_rq3; + struct fastdiv_values div_rk3; + struct fastdiv_values div_n_threads; +}; + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_params is too large for kernel_params blob"); +#endif + +struct htp_gdn_vtcm_layout { + size_t state_aligned; + size_t bytes_per_thread; + size_t total_bytes; +}; + +static inline void htp_gdn_vtcm_layout_build( + struct htp_gdn_vtcm_layout * layout, + uint32_t S_v, + uint32_t n_threads +) { + size_t state_bytes = (size_t) S_v * S_v * sizeof(float); + layout->state_aligned = hex_round_up(state_bytes, 128); + layout->bytes_per_thread = 2 * layout->state_aligned; + layout->total_bytes = layout->bytes_per_thread * n_threads; +} + +struct htp_gdn_hmx_vtcm_layout { + size_t off_s_state; + size_t off_s_f16; + size_t off_s_col_tiles; + size_t off_s_update_f32; + size_t off_s_update_tiles; + + size_t off_q_f32[2]; + size_t off_k_f32[2]; + size_t off_v_f32[2]; + size_t off_g_f32[2]; + size_t off_b_f32[2]; + size_t off_g_raw[2]; + size_t off_b_raw[2]; + size_t off_o_f32[2]; + + size_t off_v_inter_f32; + size_t off_o_inter_f32; + size_t off_o_intra_f32; + size_t off_k_f16; + size_t off_v_prime_f16; + size_t off_delta_f16; + size_t off_d_f16; + + size_t off_q_row_tiles; + size_t off_q_prime_row_tiles; + size_t off_k_row_tiles; + size_t off_k_col_tiles; + size_t off_k_prime_row_tiles; + size_t off_k_col_tiles_64x128; + size_t off_kk_tiles; + size_t off_qk_tiles; + size_t off_v_inter_tiles; + size_t off_o_inter_tiles; + size_t off_inv_row_tiles; + size_t off_a_row_tiles; + size_t off_v_prime_col_tiles; + size_t off_delta_tiles; + size_t off_delta_col_tiles; + size_t off_o_intra_tiles; + size_t off_d_row_tiles; + + size_t off_gamma; + size_t off_lambda_init; + size_t off_decay_m; + size_t off_decay_a; + size_t off_rows_kk; + size_t off_rows_qk; + size_t off_rows_inv; + size_t off_rows_a; + + size_t off_thread_scratch; + size_t off_attn_rem; + size_t off_scales_1; + + size_t state_f32_bytes; + size_t state_f16_bytes; + size_t state_tiles_bytes; + size_t dma_chunk_bytes; + size_t act_f16_bytes; + size_t tile_64xSv_bytes; + size_t tile_64x64_bytes; + + uint32_t n_heads_batch; + uint32_t n_threads; + bool pipeline; + size_t total_bytes; +}; + +static inline void htp_gdn_hmx_vtcm_layout_build( + struct htp_gdn_hmx_vtcm_layout * L, + uint32_t S_v, + uint32_t chunk_size, + uint32_t n_heads_batch, + uint32_t n_threads, + bool pipeline +) { + memset(L, 0, sizeof(*L)); + L->n_heads_batch = n_heads_batch; + L->n_threads = n_threads; + L->pipeline = pipeline; + + const size_t bh = (size_t) n_heads_batch; + const size_t nth = (size_t) (n_threads > 0 ? n_threads : 1); + + const size_t state_f32_sz = hex_round_up(S_v * S_v * sizeof(float), 2048); + const size_t state_f16_sz = hex_round_up(S_v * S_v * sizeof(__fp16), 2048); + const size_t n_sv_tiles = S_v / 32; + const size_t state_tiles_sz = n_sv_tiles * n_sv_tiles * 2048; + + const size_t dma_chunk_sz = hex_round_up(chunk_size * S_v * sizeof(float), 2048); + const size_t dma_scalar_sz = hex_round_up(chunk_size * sizeof(float), 128); + + const size_t act_f16_sz = hex_round_up(chunk_size * S_v * sizeof(__fp16), 2048); + const size_t tile_64xSv_sz = 2 * n_sv_tiles * 2048; + const size_t tile_64x64_sz = 4 * 2048; + + const size_t decay_sz = 64 * 64 * sizeof(__fp16); + const size_t row_vecs_sz = 64 * 128; + + L->state_f32_bytes = state_f32_sz; + L->state_f16_bytes = state_f16_sz; + L->state_tiles_bytes = state_tiles_sz; + L->dma_chunk_bytes = dma_chunk_sz; + L->act_f16_bytes = act_f16_sz; + L->tile_64xSv_bytes = tile_64xSv_sz; + L->tile_64x64_bytes = tile_64x64_sz; + + size_t off = 0; + + VTCM_LAYOUT_ALLOC(off, off_s_state, bh * state_f32_sz); + VTCM_LAYOUT_ALLOC(off, off_s_f16, bh * state_f16_sz); + VTCM_LAYOUT_ALLOC(off, off_s_col_tiles, bh * state_tiles_sz); + VTCM_LAYOUT_ALLOC(off, off_s_update_f32, bh * state_f32_sz); + VTCM_LAYOUT_ALLOC(off, off_s_update_tiles, bh * state_tiles_sz); + + VTCM_LAYOUT_ALLOC(off, off_q_f32[0], bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_q_f32[1], bh * dma_chunk_sz, pipeline); + VTCM_LAYOUT_ALLOC(off, off_k_f32[0], bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_k_f32[1], bh * dma_chunk_sz, pipeline); + VTCM_LAYOUT_ALLOC(off, off_v_f32[0], bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_f32[1], bh * dma_chunk_sz, pipeline); + VTCM_LAYOUT_ALLOC(off, off_g_f32[0], bh * dma_scalar_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_f32[1], bh * dma_scalar_sz, pipeline); + VTCM_LAYOUT_ALLOC(off, off_b_f32[0], bh * dma_scalar_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_b_f32[1], bh * dma_scalar_sz, pipeline); + const size_t raw_gb_sz = hex_round_up(bh * chunk_size * sizeof(float), 128); + VTCM_LAYOUT_ALLOC(off, off_g_raw[0], raw_gb_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_raw[1], raw_gb_sz, pipeline); + VTCM_LAYOUT_ALLOC(off, off_b_raw[0], raw_gb_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_b_raw[1], raw_gb_sz, pipeline); + VTCM_LAYOUT_ALLOC(off, off_o_f32[0], bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_o_f32[1], bh * dma_chunk_sz, pipeline); + + VTCM_LAYOUT_ALLOC(off, off_v_inter_f32, bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC(off, off_o_inter_f32, bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC(off, off_o_intra_f32, bh * dma_chunk_sz); + VTCM_LAYOUT_ALLOC(off, off_k_f16, bh * act_f16_sz); + VTCM_LAYOUT_ALLOC(off, off_v_prime_f16, bh * act_f16_sz); + VTCM_LAYOUT_ALLOC(off, off_delta_f16, bh * act_f16_sz); + VTCM_LAYOUT_ALLOC(off, off_d_f16, bh * act_f16_sz); + + VTCM_LAYOUT_ALLOC(off, off_q_row_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_q_prime_row_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_k_row_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_k_col_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_k_prime_row_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_k_col_tiles_64x128, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_kk_tiles, bh * tile_64x64_sz); + VTCM_LAYOUT_ALLOC(off, off_qk_tiles, bh * tile_64x64_sz); + VTCM_LAYOUT_ALLOC(off, off_v_inter_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_o_inter_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_inv_row_tiles, bh * tile_64x64_sz); + VTCM_LAYOUT_ALLOC(off, off_a_row_tiles, bh * tile_64x64_sz); + VTCM_LAYOUT_ALLOC(off, off_v_prime_col_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_delta_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_delta_col_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_o_intra_tiles, bh * tile_64xSv_sz); + VTCM_LAYOUT_ALLOC(off, off_d_row_tiles, bh * tile_64xSv_sz); + + VTCM_LAYOUT_ALLOC(off, off_gamma, bh * hex_round_up(chunk_size * sizeof(float), 128)); + VTCM_LAYOUT_ALLOC(off, off_lambda_init, bh * hex_round_up(chunk_size * sizeof(float), 128)); + VTCM_LAYOUT_ALLOC(off, off_decay_m, bh * decay_sz); + VTCM_LAYOUT_ALLOC(off, off_decay_a, bh * decay_sz); + VTCM_LAYOUT_ALLOC(off, off_rows_kk, bh * row_vecs_sz); + VTCM_LAYOUT_ALLOC(off, off_rows_qk, bh * row_vecs_sz); + VTCM_LAYOUT_ALLOC(off, off_rows_inv, bh * row_vecs_sz); + VTCM_LAYOUT_ALLOC(off, off_rows_a, bh * row_vecs_sz); + + const size_t thread_scratch_sz = 64 * 128; + VTCM_LAYOUT_ALLOC(off, off_thread_scratch, nth * thread_scratch_sz); + VTCM_LAYOUT_ALLOC(off, off_attn_rem, nth * (128 * sizeof(float))); + VTCM_LAYOUT_ALLOC(off, off_scales_1, 256); + + L->total_bytes = off; +} + +static inline bool htp_gdn_hmx_solve_layout( + struct htp_gdn_hmx_vtcm_layout * layout_out, + uint32_t S_v, + uint32_t chunk_size, + uint32_t total_rows, + size_t vtcm_budget, + uint32_t n_threads, + bool pipeline, + uint32_t * n_heads_batch_out +) { + uint32_t max_batch = 8; + if (max_batch > total_rows) { + max_batch = total_rows; + } + if (max_batch > n_threads) { + max_batch = n_threads; + } + static const uint32_t candidates[] = { 8, 6, 4, 2, 1 }; + for (size_t i = 0; i < sizeof(candidates) / sizeof(candidates[0]); ++i) { + uint32_t bh = candidates[i]; + if (bh > max_batch) { + continue; + } + struct htp_gdn_hmx_vtcm_layout L; + htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, bh, n_threads, pipeline); + if (L.total_bytes <= vtcm_budget) { + *layout_out = L; + *n_heads_batch_out = bh; + return true; + } + } + if (pipeline) { + return htp_gdn_hmx_solve_layout(layout_out, S_v, chunk_size, total_rows, vtcm_budget, n_threads, false, n_heads_batch_out); + } + return false; +} + +#endif // HTP_GATED_DELTA_NET_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.c b/ggml/src/ggml-hexagon/htp/get-rows-ops.c index 958ecac3f..f51e00c15 100644 --- a/ggml/src/ggml-hexagon/htp/get-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.c @@ -11,6 +11,7 @@ #define GGML_COMMON_DECL_C #include "ggml-common.h" #include "hex-common.h" +#include "dma-queue.h" #include "htp-ctx.h" #include "htp-ops.h" #include "htp-tensor.h" @@ -59,140 +60,139 @@ struct get_rows_context { \ const uint32_t nr = ne10 * ne11 * ne12; -#define GET_ROWS_THREAD_ST_FN(IDX_TYPE) \ -static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ - struct get_rows_context * grctx = (struct get_rows_context *)data; \ - struct htp_ops_context * octx = grctx->octx; \ - const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ - get_rows_preamble; \ - const uint32_t dr = grctx->tasks_per_thread; \ - const uint32_t ir0 = grctx->task_start + dr * ith; \ - if (ir0 >= grctx->task_start + grctx->tasks) { \ - return; \ - } \ - const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); \ - const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00); \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ - for (uint32_t i = ir0; i < ir1; ++i) { \ - const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); \ - const uint32_t rem = i - i12 * ne11 * ne10; \ - const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ - const uint32_t i10 = rem - i11 * ne10; \ - const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ - const uint32_t i01 = (uint32_t)*src1_ptr; \ - assert(i01 < ne01); \ - const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ - const uint32_t i02 = i11 - q02 * ne02; \ - const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ - const uint32_t i03 = i12 - q03 * ne03; \ - const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; \ - const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; \ - while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01, \ - row_size_bytes, 1)) { \ - dma_queue_pop(dma_queue); \ - } \ - } \ - dma_queue_flush(dma_queue); \ +#define GET_ROWS_THREAD_ST_FN(IDX_TYPE) \ +static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct get_rows_context * grctx = (struct get_rows_context *)data; \ + struct htp_ops_context * octx = grctx->octx; \ + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ + get_rows_preamble; \ + const uint32_t dr = grctx->tasks_per_thread; \ + const uint32_t ir0 = grctx->task_start + dr * ith; \ + if (ir0 >= grctx->task_start + grctx->tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); \ + const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00); \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ + for (uint32_t i = ir0; i < ir1; ++i) { \ + const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); \ + const uint32_t rem = i - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ + const uint32_t i01 = (uint32_t)*src1_ptr; \ + assert(i01 < ne01); \ + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ + const uint32_t i02 = i11 - q02 * ne02; \ + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ + const uint32_t i03 = i12 - q03 * ne03; \ + const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; \ + const dma_addr_t dst_data = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; \ + while (!dma_queue_push(dma_q, dma_make_data(dst_data, src0_data), nb1, nb01, \ + row_size_bytes, 1)) { \ + dma_queue_pop(dma_q); \ + } \ + } \ + dma_queue_flush(dma_q); \ } GET_ROWS_THREAD_ST_FN(int32_t) GET_ROWS_THREAD_ST_FN(int64_t) -#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR) \ -static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ - struct get_rows_context * grctx = (struct get_rows_context *)data; \ - struct htp_ops_context * octx = grctx->octx; \ - const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ - get_rows_preamble; \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - const uint32_t dr = grctx->tasks_per_thread; \ - const uint32_t ir0 = grctx->task_start + dr * ith; \ - if (ir0 >= grctx->task_start + grctx->tasks) { \ - return; \ - } \ - const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); \ - const uint32_t chunks_per_row = kparams->chunks_per_row; \ - const uint32_t chunk_size = kparams->chunk_size; \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ - const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; \ - uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ - uint8_t * vtcm_dst = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ - for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) { \ - const uint32_t i = ir0 + step; \ - const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ - const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ - const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ - const uint32_t rem = row_idx - i12 * ne11 * ne10; \ - const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ - const uint32_t i10 = rem - i11 * ne10; \ - const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ - const uint32_t i01 = (uint32_t)*src1_ptr; \ - assert(i01 < ne01); \ - const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ - const uint32_t i02 = i11 - q02 * ne02; \ - const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ - const uint32_t i03 = i12 - q03 * ne03; \ - const uint32_t offset = chunk_idx * chunk_size; \ - const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ - const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems); \ - const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ - const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset); \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)(uintptr_t)octx->dst->data, \ - vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ - cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0); \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size), \ - (const void *)src0_ptr), \ - vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1); \ - } \ - for (uint32_t step = 0; step < ir1 - ir0; ++step) { \ - const uint32_t i = ir0 + step; \ - void * dst_spad = (void *) dma_queue_pop(dma_queue).src; \ - void * src_spad = (void *) dma_queue_pop(dma_queue).dst; \ - const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ - const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ - const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ - const uint32_t rem = row_idx - i12 * ne11 * ne10; \ - const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ - const uint32_t i10 = rem - i11 * ne10; \ - const uint32_t offset = chunk_idx * chunk_size; \ - const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ - const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i); \ - COMPUTE_EXPR; \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i); \ - const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)dst_ptr, (const void *)dst_spad), \ - cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1); \ - const uint32_t next_step = step + 2; \ - if (next_step < ir1 - ir0) { \ - const uint32_t pi = ir0 + next_step; \ - const uint32_t prow_idx = fastdiv(pi, &kparams->div_chunks_per_row); \ - const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row; \ - const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11); \ - const uint32_t prem = prow_idx - pi12 * ne11 * ne10; \ - const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10); \ - const uint32_t pi10 = prem - pi11 * ne10; \ - const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \ - const uint32_t pi01 = (uint32_t)*psrc1_ptr; \ - assert(pi01 < ne01); \ - const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02); \ - const uint32_t pi02 = pi11 - pq02 * ne02; \ - const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03); \ - const uint32_t pi03 = pi12 - pq03 * ne03; \ - const uint32_t poffset = pchunk_idx * chunk_size; \ - const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0; \ - const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems); \ - const uintptr_t psrc0_ptr = \ - octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset); \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr), \ - vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1); \ - } \ - } \ - dma_queue_flush(dma_queue); \ +#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR) \ +static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct get_rows_context * grctx = (struct get_rows_context *)data; \ + struct htp_ops_context * octx = grctx->octx; \ + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ + get_rows_preamble; \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + const uint32_t dr = grctx->tasks_per_thread; \ + const uint32_t ir0 = grctx->task_start + dr * ith; \ + if (ir0 >= grctx->task_start + grctx->tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); \ + const uint32_t chunks_per_row = kparams->chunks_per_row; \ + const uint32_t chunk_size = kparams->chunk_size; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ + const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; \ + uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ + uint8_t * vtcm_dst = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ + for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) { \ + const uint32_t i = ir0 + step; \ + const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ + const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ + const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ + const uint32_t rem = row_idx - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ + const uint32_t i01 = (uint32_t)*src1_ptr; \ + assert(i01 < ne01); \ + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ + const uint32_t i02 = i11 - q02 * ne02; \ + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ + const uint32_t i03 = i12 - q03 * ne03; \ + const uint32_t offset = chunk_idx * chunk_size; \ + const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ + const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems); \ + const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ + const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset); \ + dma_queue_push(dma_q, \ + dma_make_data(octx->dst->data, \ + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ + cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0); \ + dma_queue_push(dma_q, \ + dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, src0_data), \ + vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1); \ + } \ + for (uint32_t step = 0; step < ir1 - ir0; ++step) { \ + const uint32_t i = ir0 + step; \ + void * dst_spad = (void *) dma_queue_pop(dma_q).src; \ + void * src_spad = (void *) dma_queue_pop(dma_q).dst; \ + const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ + const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ + const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ + const uint32_t rem = row_idx - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const uint32_t offset = chunk_idx * chunk_size; \ + const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ + const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i); \ + COMPUTE_EXPR; \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i); \ + const dma_addr_t dst_data = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); \ + dma_queue_push(dma_q, \ + dma_make_data(dst_data, dst_spad), \ + cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1); \ + const uint32_t next_step = step + 2; \ + if (next_step < ir1 - ir0) { \ + const uint32_t pi = ir0 + next_step; \ + const uint32_t prow_idx = fastdiv(pi, &kparams->div_chunks_per_row); \ + const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row; \ + const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11); \ + const uint32_t prem = prow_idx - pi12 * ne11 * ne10; \ + const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10); \ + const uint32_t pi10 = prem - pi11 * ne10; \ + const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \ + const uint32_t pi01 = (uint32_t)*psrc1_ptr; \ + assert(pi01 < ne01); \ + const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02); \ + const uint32_t pi02 = pi11 - pq02 * ne02; \ + const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03); \ + const uint32_t pi03 = pi12 - pq03 * ne03; \ + const uint32_t poffset = pchunk_idx * chunk_size; \ + const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0; \ + const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems); \ + const dma_addr_t psrc0_data = \ + octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset); \ + dma_queue_push(dma_q, \ + dma_make_data(src_spad, psrc0_data), \ + vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1); \ + } \ + } \ + dma_queue_flush(dma_q); \ } #define F32_BYTES(n) ((n) * sizeof(float)) @@ -227,8 +227,8 @@ int op_get_rows(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + if (htp_tensor_is_extended(octx->src[1])) { + return HTP_STATUS_NO_SUPPORT; } const struct htp_tensor * dst = octx->dst; diff --git a/ggml/src/ggml-hexagon/htp/hex-dma.h b/ggml/src/ggml-hexagon/htp/hex-dma.h deleted file mode 100644 index 9e9a5f950..000000000 --- a/ggml/src/ggml-hexagon/htp/hex-dma.h +++ /dev/null @@ -1,2 +0,0 @@ -#pragma once -#include "dma-queue.h" diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index a3d5e8cef..814feef7c 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -1,7 +1,7 @@ #ifndef HTP_CTX_H #define HTP_CTX_H -#include "hex-dma.h" +#include "dma-queue.h" #include "hmx-queue.h" #include "htp-ops.h" #include "hex-profile.h" @@ -17,16 +17,16 @@ #ifndef HTP_MAX_NTHREADS #define HTP_MAX_NTHREADS 10 #endif -#define HTP_MAX_MMAPS 16 -#define HTP_MAX_DIRTY_RANGES 32 +#define HTP_MAX_MMAPS 64 +#define HTP_MAX_DIRTY_RANGES 64 // Memory mapping struct htp_mmap { uint64_t size; uint64_t base; uint32_t fd; - uint32_t reserved; + uint32_t flags; }; struct htp_dirty_range { @@ -68,9 +68,6 @@ struct htp_ops_context { const struct htp_tensor * dsts[HTP_OP_MAX_OUTPUTS]; }; - dma_queue ** src_dma[HTP_OP_MAX_INPUTS]; - dma_queue ** dst_dma[HTP_OP_MAX_OUTPUTS]; - // TODO convert these to an array struct htp_spad src0_spad; struct htp_spad src1_spad; @@ -90,7 +87,6 @@ struct htp_context { struct htp_mmap mmap[HTP_MAX_MMAPS]; dma_queue_t dma[HTP_MAX_NTHREADS]; - dma_queue_t dma_cached[HTP_MAX_NTHREADS]; struct htp_thread_trace trace[HTP_MAX_NTHREADS + 1]; work_queue_t work_queue; hmx_queue_t hmx_queue; diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index faf3118c4..0e63febdd 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -133,10 +133,13 @@ enum htp_tensor_flags { HTP_TENSOR_FENCE = (1U << 2) // Tensor is synchronization fence (explicitly managed) }; +enum htp_buf_flags { + HTP_BUF_EXTENDED = (1U << 0), +}; + // Tensor descriptor struct htp_tensor { - uint32_t data; // Buffer offset in the messages, and data pointer on the NPU - uint32_t reserved; // Reserved for alignment padding (must be multiple of 8) + uint64_t data; // Buffer offset in the messages, and data pointer on the NPU uint32_t size; // Data size in bytes uint32_t flags; // Buffer / tensor flags uint32_t type; // Data type @@ -150,12 +153,12 @@ struct htp_tensor { struct htp_buf_desc { uint64_t base; // base address uint64_t size; // total size - uint32_t flags; // buffer flags (unused) + uint32_t flags; // HTP_BUF_* uint32_t fd; // file descriptor }; enum htp_op_flags { - HTP_OPFLAGS_SKIP_COMPUTE = (1U << 0), // Skip actual computation (used for profiling) + HTP_OPFLAGS_STUB = (1U << 0), }; // Op descriptor diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.c b/ggml/src/ggml-hexagon/htp/htp-tensor.c index 760ccd831..03b0070d7 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.c +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.c @@ -226,6 +226,10 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co } static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) { + if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) { + return; + } + uint32_t t_start = t->data; uint32_t t_end = t_start + t->size; @@ -236,6 +240,7 @@ static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor if (r->start < t_end && t_start < r->end) { if (t_start <= r->start && r->end <= t_end) { r->start = 0; + r->end = 0; } else if (t_start <= r->start) { r->start = t_end; } else if (r->end <= t_end) { @@ -246,6 +251,10 @@ static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor } static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) { + if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) { + return false; + } + uint32_t t_start = t->data; uint32_t t_end = t_start + t->size; @@ -327,7 +336,7 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co for (uint32_t i = 0; i < n; i++) { const struct htp_tensor * t = tensors[i]; - if (t && is_tensor_dirty(ctx, t)) { + if (is_tensor_dirty(ctx, t)) { dirty_tensors[n_dirty++] = t; ranges[n_dirty - 1].start = t->data; ranges[n_dirty - 1].end = t->data + t->size; diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.h b/ggml/src/ggml-hexagon/htp/htp-tensor.h index 1e32bf09f..f7a966836 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.h +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.h @@ -21,6 +21,10 @@ static inline void * htp_tensor_data(const struct htp_tensor * t) { return (void *) (uintptr_t) t->data; } +static inline bool htp_tensor_is_extended(const struct htp_tensor * t) { + return t && (t->data >> 32) != 0; +} + static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) { return (uint32_t *) &t->flags; } diff --git a/ggml/src/ggml-hexagon/htp/htp_iface.idl b/ggml/src/ggml-hexagon/htp/htp_iface.idl index 47693d8b8..b46e25296 100644 --- a/ggml/src/ggml-hexagon/htp/htp_iface.idl +++ b/ggml/src/ggml-hexagon/htp/htp_iface.idl @@ -13,7 +13,7 @@ struct htp_iface_pmu_conf { interface htp_iface : remote_handle64 { AEEResult start(in uint32 sess_id, in uint64 dsp_queue_id, in uint32 n_hvx, in uint32 n_hmx, in uint64 max_vmem); AEEResult stop(); - AEEResult mmap(in uint32 fd, in uint32 size); + AEEResult mmap(in uint32 fd, in uint64 size); AEEResult munmap(in uint32 fd); AEEResult profiler(in uint32 mode, in htp_iface_pmu_conf pmu); AEEResult etm(in uint32 enable); diff --git a/ggml/src/ggml-hexagon/htp/hvx-exp.h b/ggml/src/ggml-hexagon/htp/hvx-exp.h index bcd3d2d32..93ca8cf51 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-exp.h +++ b/ggml/src/ggml-hexagon/htp/hvx-exp.h @@ -173,7 +173,7 @@ static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict HVX_Vector * p_vec_in1 = (HVX_Vector *) src; HVX_Vector * p_vec_out = (HVX_Vector *) dst; - #pragma unroll(4) + #pragma unroll(2) 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++); @@ -183,7 +183,7 @@ static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict } } } else { - #pragma unroll(4) + #pragma unroll(2) for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { HVX_Vector in = *(HVX_UVector *) (src + i * SIZEOF_FP32); diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h deleted file mode 100644 index 5c1372cf1..000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h +++ /dev/null @@ -1,1648 +0,0 @@ -// Dynamic quantizers that produce flat (non-tiled) activations - -static inline void quantize_block_f32_q8_0_flat( - float * restrict x, - uint8_t * restrict y_quants, - __fp16 * restrict y_scales, - uint32_t block_idx -) { - HVX_Vector * vx = (HVX_Vector *) x; - HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); - HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); - HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); - HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); - - HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); - HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); - HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); - HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); - - HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); - HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); - HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); - HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); - - HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); - HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); - - HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); - HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - - HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); - HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); - - HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); - HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); - - HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); - HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); - HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - - * (HVX_Vector *) (y_quants + block_idx * 128) = vx_i8; - - HVX_VectorPair vp1 = Q6_W_vshuff_VVR(vd23_hf, vd01_hf, -2); - HVX_VectorPair vp2 = Q6_W_vshuff_VVR(Q6_V_hi_W(vp1), Q6_V_lo_W(vp1), -2); - HVX_Vector v_scales = Q6_V_lo_W(vp2); - hvx_vec_store_u(y_scales + block_idx * 4, 8, v_scales); -} - -static inline void quantize_block_f32_q8_1_flat( - float * restrict x, - uint8_t * restrict y_quants, - __fp16 * restrict y_scales, - uint32_t block_idx -) { - HVX_Vector * vx = (HVX_Vector *) x; - HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); - HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); - HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); - HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); - - HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); - HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); - HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); - HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); - - HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); - HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); - HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); - HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); - - HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); - HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); - - HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); - HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - - HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); - HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); - - HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); - HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); - - HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); - HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); - HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - - const HVX_Vector ones = Q6_Vb_vsplat_R(1); - HVX_Vector v_sums = Q6_Vw_vrmpy_VbVb(vx_i8, ones); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 4)); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 8)); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 16)); - - * (HVX_Vector *) (y_quants + block_idx * 128) = vx_i8; - - HVX_VectorPair vp1 = Q6_W_vshuff_VVR(vd23_hf, vd01_hf, -2); - HVX_VectorPair vp2 = Q6_W_vshuff_VVR(Q6_V_hi_W(vp1), Q6_V_lo_W(vp1), -2); - HVX_Vector v_scales = Q6_V_lo_W(vp2); - - HVX_VectorPair v_deal1 = Q6_W_vdeal_VVR(v_sums, v_sums, -4); - HVX_Vector v_even1 = Q6_V_lo_W(v_deal1); - HVX_VectorPair v_deal2 = Q6_W_vdeal_VVR(v_even1, v_even1, -4); - HVX_Vector v_even2 = Q6_V_lo_W(v_deal2); - HVX_VectorPair v_deal3 = Q6_W_vdeal_VVR(v_even2, v_even2, -4); - HVX_Vector v_sums_shuffled = Q6_V_lo_W(v_deal3); - - HVX_Vector v_sums_sf = Q6_Vsf_equals_Vw(v_sums_shuffled); - HVX_Vector v_sums_hf = hvx_vec_f32_to_f16(v_sums_sf, Q6_V_vzero()); - - HVX_Vector v_prod = hvx_vec_mul_f16_f16(v_scales, v_sums_hf); - - HVX_VectorPair vp_scales = Q6_W_vshuff_VVR(v_prod, v_scales, -2); - HVX_Vector v_final = Q6_V_lo_W(vp_scales); - - hvx_vec_store_u(y_scales + block_idx * 8, 16, v_final); -} - -static inline void quantize_row_f32_q8_0_flat(float * restrict x, uint8_t * restrict y, uint32_t k) { - assert(k % 32 == 0); - const uint32_t quants_size = hex_round_up(k, 128); - uint8_t * restrict y_quants = y; - __fp16 * restrict y_scales = (__fp16 *) (y + quants_size); - - const uint32_t nb = (k + 127) / 128; - for (uint32_t i = 0; i < nb; i++) { - quantize_block_f32_q8_0_flat(x + i * 128, y_quants, y_scales, i); - } -} - -static inline void quantize_row_f32_q8_1_flat(float * restrict x, uint8_t * restrict y, uint32_t k) { - assert(k % 32 == 0); - const uint32_t quants_size = hex_round_up(k, 128); - uint8_t * restrict y_quants = y; - __fp16 * restrict y_scales = (__fp16 *) (y + quants_size); - - const uint32_t nb = (k + 127) / 128; - for (uint32_t i = 0; i < nb; i++) { - quantize_block_f32_q8_1_flat(x + i * 128, y_quants, y_scales, i); - } -} - -static inline void quantize_f32_q8_0_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_row_size, - size_t dst_row_size -) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_0_flat((float *) tmp_data, dst_data, ne0); - dst_data += dst_row_size; - src_data += src_row_size; - } -} - -static inline void quantize_f32_q8_1_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_row_size, - size_t dst_row_size -) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_1_flat((float *) tmp_data, dst_data, ne0); - dst_data += dst_row_size; - src_data += src_row_size; - } -} - -static inline void quantize_f32_f32_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_stride, - size_t dst_stride -) { - (void) tmp_data; - const size_t src_row_size = ne0 * sizeof(float); - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_stride, 2); - hvx_copy_f32_au(dst_data, src_data, ne0); - - dst_data += dst_stride; - src_data += src_stride; - } -} - -static inline void quantize_f32_f16_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_stride, - size_t dst_stride -) { - (void) tmp_data; - const size_t src_row_size = ne0 * sizeof(float); - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_stride, 2); - hvx_copy_f16_f32_au(dst_data, src_data, ne0); - - dst_data += dst_stride; - src_data += src_stride; - } -} - -static inline void quantize_f16_f16_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_stride, - size_t dst_stride -) { - (void) tmp_data; - const size_t src_row_size = ne0 * sizeof(float); - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_stride, 2); - hvx_copy_f16_au(dst_data, src_data, ne0); - - dst_data += dst_stride; - src_data += src_stride; - } -} - -// Dot kernels that consume flat (non-tiled) activations - -static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act_rep, i8); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0_rep, v_act1_rep, i8); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act_rep, Q6_V_vzero()); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_offset = vptr[4]; - HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); - HVX_Vector v_scale = Q6_V_lo_W(p_deal); - HVX_Vector v_offset = Q6_V_hi_W(p_deal); - - __fp16 scale_a_val = y_scales[kt * 2 + 0]; - __fp16 sum_a_val = y_scales[kt * 2 + 1]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - HVX_Vector v_sum_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a); - HVX_Vector v_offset_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a); - - HVX_Vector v_scaled_dot = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - HVX_Vector v_sum_scaled = hvx_vec_add_f32_f32(v_scaled_dot, v_offset_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0_rep, v_act1_rep, Q6_V_vzero()); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_offset = vptr[4]; - HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); - HVX_Vector v_scale = Q6_V_lo_W(p_deal); - HVX_Vector v_offset = Q6_V_hi_W(p_deal); - - __fp16 scale_a0_val = y0_scales[kt * 2 + 0]; - __fp16 sum_a0_val = y0_scales[kt * 2 + 1]; - __fp16 scale_a1_val = y1_scales[kt * 2 + 0]; - __fp16 sum_a1_val = y1_scales[kt * 2 + 1]; - - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_sum_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - HVX_Vector v_sum_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a0); - HVX_Vector v_offset_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a1); - HVX_Vector v_offset_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a1); - - HVX_Vector v_scaled_dot_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c0 = hvx_vec_add_f32_f32(v_scaled_dot_c0, v_offset_comb_c0); - - HVX_Vector v_scaled_dot_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - HVX_Vector v_sum_scaled_c1 = hvx_vec_add_f32_f32(v_scaled_dot_c1, v_offset_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_q8_0_32x1(vptr, v_act_rep); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[8]; - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_q8_0_32x2(vptr, v_act0_rep, v_act1_rep); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[8]; - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_q6_k_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector i32 = Q6_Vb_vsplat_R(32); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 896); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_q6_k_32x1(vptr, v_act_rep, i32); - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, scale_q6_k_32x1(v_sums, vptr[6], v_scale_a)); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q6_k_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector i32 = Q6_Vb_vsplat_R(32); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 896); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - HVX_Vector v_act1_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums0, v_sums1; - accum_q6_k_32x2(vptr, v_act0_rep, v_act1_rep, i32, &v_sums0, &v_sums1); - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, scale_q6_k_32x1(v_sums0, vptr[6], v_scale_a0)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, scale_q6_k_32x1(v_sums1, vptr[6], v_scale_a1)); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act_rep, mask_h4, lut); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0_rep, v_act1_rep, mask_h4, lut); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut; - HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0; - HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act_rep, mask_h4, lut); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512); - HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand); - r0_d = Q6_V_vand_VV(r0_d, e8m0_mask); - HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23); - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - HVX_VectorPair p_scale_a_f32 = hvx_vec_f16_to_f32(v_scale_a_f16); - HVX_Vector v_scale_a = Q6_V_lo_W(p_scale_a_f32); - - HVX_Vector v_scale_comb = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f)); - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut; - HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0; - HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0_rep, v_act1_rep, mask_h4, lut); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512); - HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand); - r0_d = Q6_V_vand_VV(r0_d, e8m0_mask); - HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23); - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - HVX_VectorPair p_scale_a0_f32 = hvx_vec_f16_to_f32(v_scale_a0_f16); - HVX_VectorPair p_scale_a1_f32 = hvx_vec_f16_to_f32(v_scale_a1_f16); - HVX_Vector v_scale_a0 = Q6_V_lo_W(p_scale_a0_f32); - HVX_Vector v_scale_a1 = Q6_V_lo_W(p_scale_a1_f32); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f)); - v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f)); - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -#if __HVX_ARCH__ < 79 -#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_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_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) -#endif - -static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { - const HVX_Vector * restrict x = (const HVX_Vector *) vx; - const HVX_Vector * restrict y = (const HVX_Vector *) vy; - - uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors - uint32_t nloe = n % VLEN_FP32; // leftover elements - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(4) - for (i = 0; i < nvec; i++) { - HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]); - rsum = HVX_OP_ADD_F32(rsum, prod); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]); - HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]); - HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf); - rsum = HVX_OP_ADD_F32(rsum, prod); - } - - *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum)); -} - -static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0, - const void * restrict vx0, const void * restrict vx1, - const void * restrict vy0) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y = (const HVX_Vector *) vy0; - - uint32_t nvec = n / VLEN_FP32; - uint32_t nloe = n % VLEN_FP32; - - HVX_Vector rsum0 = Q6_V_vzero(); - HVX_Vector rsum1 = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector y_sf = y[i]; - HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf); - HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf); - rsum0 = HVX_OP_ADD_F32(rsum0, prod0); - rsum1 = HVX_OP_ADD_F32(rsum1, prod1); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]); - HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]); - HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf); - HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf); - rsum0 = HVX_OP_ADD_F32(rsum0, prod0); - rsum1 = HVX_OP_ADD_F32(rsum1, prod1); - } - - HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); - hvx_vec_store_u(s0, 8, rsum); -} - -static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, - const void * restrict vx0, const void * restrict vx1, - const void * restrict vy0, const void * restrict vy1) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; - const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; - - uint32_t nvec = n / VLEN_FP32; - uint32_t nloe = n % VLEN_FP32; - - HVX_Vector r0_c0_sum = Q6_V_vzero(); - HVX_Vector r0_c1_sum = Q6_V_vzero(); - HVX_Vector r1_c0_sum = Q6_V_vzero(); - HVX_Vector r1_c1_sum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector r0_sf = x0[i]; - HVX_Vector r1_sf = x1[i]; - HVX_Vector c0_sf = y0[i]; - HVX_Vector c1_sf = y1[i]; - - r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf)); - r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf)); - r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf)); - r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf)); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - - HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]); - HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]); - HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]); - - r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf)); - r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf)); - r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf)); - r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf)); - } - - // Reduce and store results - HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); - HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); - - hvx_vec_store_u(s0, 8, r0_r1_c0_sum); - hvx_vec_store_u(s1, 8, r0_r1_c1_sum); -} - -static inline void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) { - const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x; - const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y; - - uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors - uint32_t nloe = n % VLEN_FP32; // leftover elements - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector x_sf = vx[i]; - HVX_Vector y_sf = vy[i]; - - rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf)); - } - - if (nloe) { - HVX_Vector x_sf = vx[i]; - HVX_Vector y_sf = vy[i]; - - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - x_sf = Q6_V_vand_QV(bmask, x_sf); - y_sf = Q6_V_vand_QV(bmask, y_sf); - - rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf)); - } - - rsum = hvx_vec_reduce_sum_f32(rsum); - hvx_vec_store_u(&s[0], 4, rsum); -} - -#undef HVX_OP_ADD_F32 -#undef HVX_OP_MUL_F32 - -static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { - const HVX_Vector * restrict x = (const HVX_Vector *) vx; - const HVX_Vector * restrict y = (const HVX_Vector *) vy; - - uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors - uint32_t nloe = n % VLEN_FP16; // leftover elements - - HVX_VectorPair rsum_p = Q6_W_vzero(); - - uint32_t i = 0; - - #pragma unroll(4) - for (i = 0; i < nvec; i++) { - rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]); - HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); - rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf); - } - - HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p))); - hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum)); -} - -static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0, - const void * restrict vx0, const void * restrict vx1, - const void * restrict vy0) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y = (const HVX_Vector *) vy0; - - uint32_t nvec = n / VLEN_FP16; - uint32_t nloe = n % VLEN_FP16; - - HVX_VectorPair rsum0_p = Q6_W_vzero(); - HVX_VectorPair rsum1_p = Q6_W_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector y_hf = y[i]; - rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf); - rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); - HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]); - rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf); - rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf); - } - - HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p))); - HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p))); - HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); - hvx_vec_store_u(s0, 8, rsum); -} - -static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, - const void * restrict vx0, const void * restrict vx1, - const void * restrict vy0, const void * restrict vy1) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; - const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; - - uint32_t nvec = n / VLEN_FP16; - uint32_t nloe = n % VLEN_FP16; - - // Row sums (sf) - 4 accumulators for 2x2 tile - HVX_VectorPair r0_c0_sum_p = Q6_W_vzero(); - HVX_VectorPair r0_c1_sum_p = Q6_W_vzero(); - HVX_VectorPair r1_c0_sum_p = Q6_W_vzero(); - HVX_VectorPair r1_c1_sum_p = Q6_W_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector r0_hf = x0[i]; - HVX_Vector r1_hf = x1[i]; - HVX_Vector c0_hf = y0[i]; - HVX_Vector c1_hf = y1[i]; - - // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1 - r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf); - r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf); - r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf); - r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - - HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]); - HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]); - HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]); - - r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf); - r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf); - r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf); - r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf); - } - - HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p))); - HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p))); - HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p))); - HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p))); - - // Reduce and store results - HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); - HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); - - hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum); // row0,col0 row1,col0 - hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum); // row0,col1 row1,col1 -} - -static inline void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { - const HVX_UVector * restrict x = (const HVX_UVector *) vx; - const HVX_UVector * restrict y = (const HVX_UVector *) vy; - - uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors - uint32_t nloe = n % VLEN_FP16; // leftover elements - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(4) - for (i = 0; i < nvec; i++) { - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]); - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]); - HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); - - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf); - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum)); - hvx_vec_store_u(&s[0], 4, rsum); -} - -static inline void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) { - const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x; - const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y; - - uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors - uint32_t nloe = n % VLEN_FP16; // leftover elements - - const HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - // Load y (fp32) and convert into fp16 - HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements - HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements - HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf))); - - // Load x (fp16) - HVX_Vector x_hf = vx[i]; - - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf); - - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - if (nloe) { - // Load y (fp32) and convert into fp16 - HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements - HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements - HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf))); - - // Load x (fp16) - HVX_Vector x_hf = vx[i]; - - // Zero-out unused elements - // Note that we need to clear both x and y because they may contain NANs - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - x_hf = Q6_V_vand_QV(bmask, x_hf); - y_hf = Q6_V_vand_QV(bmask, y_hf); - - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf); - - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - // Convert into fp32 and reduce - rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum)); - hvx_vec_store_u(&s[0], 4, rsum); -} - -static inline void hvx_tensor_add_f32_grid( - const struct htp_tensor * restrict dst, - const struct htp_tensor * restrict src2, - uint32_t start_row, - uint32_t end_row, - uint32_t start_col, - uint32_t end_col, - const struct fastdiv_values * div_ne11_12, - const struct fastdiv_values * div_ne11 -) { - if (start_row >= end_row || start_col >= end_col) return; - const uint32_t nb1 = dst->nb[1]; // row stride in bytes - - const uint32_t ne11 = dst->ne[1]; - const uint32_t ne12 = dst->ne[2]; - const uint32_t ne11_12 = ne11 * ne12; - - const bool is_broadcast1 = (src2->ne[1] == 1); - const bool is_broadcast2 = (src2->ne[2] == 1); - const bool is_broadcast3 = (src2->ne[3] == 1); - - for (uint32_t r = start_row; r < end_row; r++) { - float * dst_row = (float *) ((uint8_t *) dst->data + r * nb1); - - uint32_t i13 = fastdiv(r, div_ne11_12); - uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11); - uint32_t i11 = r - i13 * ne11_12 - i12 * ne11; - - uint32_t i23 = is_broadcast3 ? 0 : i13; - uint32_t i22 = is_broadcast2 ? 0 : i12; - uint32_t i21 = is_broadcast1 ? 0 : i11; - - const float * src2_row = (const float *) ((const uint8_t *) src2->data + - i21 * src2->nb[1] + i22 * src2->nb[2] + i23 * src2->nb[3]); - - float * dst_ptr = &dst_row[start_col]; - const float * src2_ptr = &src2_row[start_col]; - int remaining = end_col - start_col; - while (remaining >= 32) { - HVX_Vector v_out = hvx_vmemu(dst_ptr); - HVX_Vector v_z = hvx_vmemu(src2_ptr); - hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z); - dst_ptr += 32; - src2_ptr += 32; - remaining -= 32; - } - if (remaining > 0) { - HVX_Vector v_out = hvx_vmemu(dst_ptr); - HVX_Vector v_z = hvx_vmemu(src2_ptr); - hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z)); - } - } -} - diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h new file mode 100644 index 000000000..605892aa7 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h @@ -0,0 +1,382 @@ +#ifndef HVX_MM_KERNELS_FLOAT_H +#define HVX_MM_KERNELS_FLOAT_H + +#include "hvx-utils.h" +#include "htp-tensor.h" + +// Float activation copy/quantization kernels (DDR -> VTCM) + +static inline void quantize_f32_f32_kernel( + const uint8_t * restrict src_data, + uint8_t * restrict dst_data, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t nrows, + size_t src_stride, + size_t dst_stride +) { + (void) tmp_data; + const size_t src_row_size = ne0 * sizeof(float); + for (uint32_t i = 0; i < nrows; ++i) { + hex_l2fetch(src_data, src_row_size, src_stride, 2); + hvx_copy_f32_au(dst_data, src_data, ne0); + + dst_data += dst_stride; + src_data += src_stride; + } +} + +static inline void quantize_f32_f16_kernel( + const uint8_t * restrict src_data, + uint8_t * restrict dst_data, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t nrows, + size_t src_stride, + size_t dst_stride +) { + (void) tmp_data; + const size_t src_row_size = ne0 * sizeof(float); + for (uint32_t i = 0; i < nrows; ++i) { + hex_l2fetch(src_data, src_row_size, src_stride, 2); + hvx_copy_f16_f32_au(dst_data, src_data, ne0); + + dst_data += dst_stride; + src_data += src_stride; + } +} + +static inline void quantize_f16_f16_kernel( + const uint8_t * restrict src_data, + uint8_t * restrict dst_data, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t nrows, + size_t src_stride, + size_t dst_stride +) { + (void) tmp_data; + const size_t src_row_size = ne0 * sizeof(float); + for (uint32_t i = 0; i < nrows; ++i) { + hex_l2fetch(src_data, src_row_size, src_stride, 2); + hvx_copy_f16_au(dst_data, src_data, ne0); + + dst_data += dst_stride; + src_data += src_stride; + } +} + +// Float dot product kernels (HVX) + +#if __HVX_ARCH__ < 79 +#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_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_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) +#endif + +static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { + const HVX_Vector * restrict x = (const HVX_Vector *) vx; + const HVX_Vector * restrict y = (const HVX_Vector *) vy; + + uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors + uint32_t nloe = n % VLEN_FP32; // leftover elements + + HVX_Vector rsum = Q6_V_vzero(); + + uint32_t i = 0; + + #pragma unroll(4) + for (i = 0; i < nvec; i++) { + HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]); + rsum = HVX_OP_ADD_F32(rsum, prod); + } + + if (nloe) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); + HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]); + HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]); + HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf); + rsum = HVX_OP_ADD_F32(rsum, prod); + } + + *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum)); +} + +static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0, + const void * restrict vx0, const void * restrict vx1, + const void * restrict vy0) { + const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; + const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; + const HVX_Vector * restrict y = (const HVX_Vector *) vy0; + + uint32_t nvec = n / VLEN_FP32; + uint32_t nloe = n % VLEN_FP32; + + HVX_Vector rsum0 = Q6_V_vzero(); + HVX_Vector rsum1 = Q6_V_vzero(); + + uint32_t i = 0; + + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_Vector y_sf = y[i]; + HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf); + HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf); + rsum0 = HVX_OP_ADD_F32(rsum0, prod0); + rsum1 = HVX_OP_ADD_F32(rsum1, prod1); + } + + if (nloe) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); + HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]); + HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]); + HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]); + HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf); + HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf); + rsum0 = HVX_OP_ADD_F32(rsum0, prod0); + rsum1 = HVX_OP_ADD_F32(rsum1, prod1); + } + + HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); + hvx_vec_store_u(s0, 8, rsum); +} + +static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, + const void * restrict vx0, const void * restrict vx1, + const void * restrict vy0, const void * restrict vy1) { + const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; + const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; + const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; + const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; + + uint32_t nvec = n / VLEN_FP32; + uint32_t nloe = n % VLEN_FP32; + + HVX_Vector r0_c0_sum = Q6_V_vzero(); + HVX_Vector r0_c1_sum = Q6_V_vzero(); + HVX_Vector r1_c0_sum = Q6_V_vzero(); + HVX_Vector r1_c1_sum = Q6_V_vzero(); + + uint32_t i = 0; + + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_Vector r0_sf = x0[i]; + HVX_Vector r1_sf = x1[i]; + HVX_Vector c0_sf = y0[i]; + HVX_Vector c1_sf = y1[i]; + + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf)); + } + + if (nloe) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); + + HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]); + HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]); + HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]); + HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]); + + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf)); + } + + // Reduce and store results + HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); + HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); + + hvx_vec_store_u(s0, 8, r0_r1_c0_sum); + hvx_vec_store_u(s1, 8, r0_r1_c1_sum); +} + +#undef HVX_OP_ADD_F32 +#undef HVX_OP_MUL_F32 + +static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { + const HVX_Vector * restrict x = (const HVX_Vector *) vx; + const HVX_Vector * restrict y = (const HVX_Vector *) vy; + + uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors + uint32_t nloe = n % VLEN_FP16; // leftover elements + + HVX_VectorPair rsum_p = Q6_W_vzero(); + + uint32_t i = 0; + + #pragma unroll(4) + for (i = 0; i < nvec; i++) { + rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]); + } + + if (nloe) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); + HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]); + HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); + rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf); + } + + HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p))); + hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum)); +} + +static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0, + const void * restrict vx0, const void * restrict vx1, + const void * restrict vy0) { + const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; + const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; + const HVX_Vector * restrict y = (const HVX_Vector *) vy0; + + uint32_t nvec = n / VLEN_FP16; + uint32_t nloe = n % VLEN_FP16; + + HVX_VectorPair rsum0_p = Q6_W_vzero(); + HVX_VectorPair rsum1_p = Q6_W_vzero(); + + uint32_t i = 0; + + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_Vector y_hf = y[i]; + rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf); + rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf); + } + + if (nloe) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); + HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); + HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]); + HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]); + rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf); + rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf); + } + + HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p))); + HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p))); + HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); + hvx_vec_store_u(s0, 8, rsum); +} + +static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, + const void * restrict vx0, const void * restrict vx1, + const void * restrict vy0, const void * restrict vy1) { + const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; + const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; + const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; + const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; + + uint32_t nvec = n / VLEN_FP16; + uint32_t nloe = n % VLEN_FP16; + + // Row sums (sf) - 4 accumulators for 2x2 tile + HVX_VectorPair r0_c0_sum_p = Q6_W_vzero(); + HVX_VectorPair r0_c1_sum_p = Q6_W_vzero(); + HVX_VectorPair r1_c0_sum_p = Q6_W_vzero(); + HVX_VectorPair r1_c1_sum_p = Q6_W_vzero(); + + uint32_t i = 0; + + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_Vector r0_hf = x0[i]; + HVX_Vector r1_hf = x1[i]; + HVX_Vector c0_hf = y0[i]; + HVX_Vector c1_hf = y1[i]; + + // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1 + r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf); + r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf); + r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf); + r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf); + } + + if (nloe) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); + + HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]); + HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]); + HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]); + HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]); + + r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf); + r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf); + r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf); + r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf); + } + + HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p))); + HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p))); + HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p))); + HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p))); + + // Reduce and store results + HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); + HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); + + hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum); // row0,col0 row1,col0 + hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum); // row0,col1 row1,col1 +} + + + +static inline void hvx_tensor_add_f32_grid( + const struct htp_tensor * restrict dst, + const struct htp_tensor * restrict src2, + uint32_t start_row, + uint32_t end_row, + uint32_t start_col, + uint32_t end_col, + const struct fastdiv_values * div_ne11_12, + const struct fastdiv_values * div_ne11 +) { + if (start_row >= end_row || start_col >= end_col) return; + const uint32_t nb1 = dst->nb[1]; // row stride in bytes + + const uint32_t ne11 = dst->ne[1]; + const uint32_t ne12 = dst->ne[2]; + const uint32_t ne11_12 = ne11 * ne12; + + const bool is_broadcast1 = (src2->ne[1] == 1); + const bool is_broadcast2 = (src2->ne[2] == 1); + const bool is_broadcast3 = (src2->ne[3] == 1); + + for (uint32_t r = start_row; r < end_row; r++) { + float * dst_row = (float *) ((uint8_t *) dst->data + (size_t) r * nb1); + + uint32_t i13 = fastdiv(r, div_ne11_12); + uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11); + uint32_t i11 = r - i13 * ne11_12 - i12 * ne11; + + uint32_t i23 = is_broadcast3 ? 0 : i13; + uint32_t i22 = is_broadcast2 ? 0 : i12; + uint32_t i21 = is_broadcast1 ? 0 : i11; + + const float * src2_row = (const float *) ((const uint8_t *) src2->data + + (size_t) i21 * src2->nb[1] + (size_t) i22 * src2->nb[2] + (size_t) i23 * src2->nb[3]); + + float * dst_ptr = &dst_row[start_col]; + const float * src2_ptr = &src2_row[start_col]; + int remaining = end_col - start_col; + while (remaining >= 32) { + HVX_Vector v_out = hvx_vmemu(dst_ptr); + HVX_Vector v_z = hvx_vmemu(src2_ptr); + hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z); + dst_ptr += 32; + src2_ptr += 32; + remaining -= 32; + } + if (remaining > 0) { + HVX_Vector v_out = hvx_vmemu(dst_ptr); + HVX_Vector v_z = hvx_vmemu(src2_ptr); + hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z)); + } + } +} + +#endif // HVX_MM_KERNELS_FLOAT_H diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h index c889538ac..4d6110ffa 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h @@ -48,22 +48,33 @@ static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * r v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 8)); v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 16)); - float vmax0[32] __attribute__((aligned(128))); - float vmax1[32] __attribute__((aligned(128))); - float vmax2[32] __attribute__((aligned(128))); - float vmax3[32] __attribute__((aligned(128))); - int32_t sums[32] __attribute__((aligned(128))); + const HVX_Vector v_inv127 = hvx_vec_splat_f32(1.0f / 127.0f); + HVX_Vector vd0_sf = hvx_vec_mul_f32_f32(vmax0_sf, v_inv127); + HVX_Vector vd1_sf = hvx_vec_mul_f32_f32(vmax1_sf, v_inv127); + HVX_Vector vd2_sf = hvx_vec_mul_f32_f32(vmax2_sf, v_inv127); + HVX_Vector vd3_sf = hvx_vec_mul_f32_f32(vmax3_sf, v_inv127); - hvx_vec_store_u(vmax0, 128, vmax0_sf); - hvx_vec_store_u(vmax1, 128, vmax1_sf); - hvx_vec_store_u(vmax2, 128, vmax2_sf); - hvx_vec_store_u(vmax3, 128, vmax3_sf); - hvx_vec_store_u(sums, 128, v_sums); + HVX_Vector v_sums_sf = Q6_Vsf_equals_Vw(v_sums); + HVX_Vector voff0_sf = hvx_vec_mul_f32_f32(vd0_sf, v_sums_sf); + HVX_Vector voff1_sf = hvx_vec_mul_f32_f32(vd1_sf, Q6_V_vror_VR(v_sums_sf, 32)); + HVX_Vector voff2_sf = hvx_vec_mul_f32_f32(vd2_sf, Q6_V_vror_VR(v_sums_sf, 64)); + HVX_Vector voff3_sf = hvx_vec_mul_f32_f32(vd3_sf, Q6_V_vror_VR(v_sums_sf, 96)); - float d0 = vmax0[0] / 127.0f; - float d1 = vmax1[0] / 127.0f; - float d2 = vmax2[0] / 127.0f; - float d3 = vmax3[0] / 127.0f; + HVX_Vector voff01_hf = hvx_vec_f32_to_f16(voff0_sf, voff1_sf); + HVX_Vector voff23_hf = hvx_vec_f32_to_f16(voff2_sf, voff3_sf); + + HVX_Vector r_scale[4] = { + hvx_vec_repl_f16(vd01_hf), + hvx_vec_repl_f16(Q6_V_vror_VR(vd01_hf, 64)), + hvx_vec_repl_f16(vd23_hf), + hvx_vec_repl_f16(Q6_V_vror_VR(vd23_hf, 64)), + }; + HVX_Vector r_offset[4] = { + hvx_vec_repl_f16(voff01_hf), + hvx_vec_repl_f16(Q6_V_vror_VR(voff01_hf, 64)), + hvx_vec_repl_f16(voff23_hf), + hvx_vec_repl_f16(Q6_V_vror_VR(voff23_hf, 64)), + }; static const uint8_t __attribute__((aligned(128))) repl[128] = { 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, @@ -89,24 +100,6 @@ static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * r HVX_Vector r6 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 24), v_repl_ctrl); HVX_Vector r7 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 28), v_repl_ctrl); - __fp16 scale_h, offset_h; - if (b == 0) { - scale_h = (__fp16) d0; - offset_h = (__fp16) (sums[0] * d0); - } else if (b == 1) { - scale_h = (__fp16) d1; - offset_h = (__fp16) (sums[8] * d1); - } else if (b == 2) { - scale_h = (__fp16) d2; - offset_h = (__fp16) (sums[16] * d2); - } else { - scale_h = (__fp16) d3; - offset_h = (__fp16) (sums[24] * d3); - } - - HVX_Vector r_scale = Q6_Vh_vsplat_R(*(int16_t *)&scale_h); - HVX_Vector r_offset = Q6_Vh_vsplat_R(*(int16_t *)&offset_h); - HVX_Vector * restrict dst = (HVX_Vector *) (y_block + b * 1280); dst[0] = r0; dst[1] = r1; @@ -116,8 +109,8 @@ static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * r dst[5] = r5; dst[6] = r6; dst[7] = r7; - dst[8] = r_scale; - dst[9] = r_offset; + dst[8] = r_scale[b]; + dst[9] = r_offset[b]; } } @@ -486,51 +479,7 @@ static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float HVX_Vector i8 = Q6_Vb_vsplat_R(8); uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_4bit_32x2(vptr0, v_act0_0, v_act1_0, i8); - HVX_VectorPair v_sums1 = accum_4bit_32x2(vptr1, v_act0_1, v_act1_1, i8); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = vptr0[4]; - HVX_Vector v_scale_w1 = vptr1[4]; - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); @@ -615,76 +564,7 @@ static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float HVX_Vector v_sum_float_c1 = Q6_V_vzero(); uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1280); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1280); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1280); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1280); - - HVX_VectorPair v_sums0 = accum_4bit_32x2(vptr0, v_act0_0, v_act1_0, Q6_V_vzero()); - HVX_VectorPair v_sums1 = accum_4bit_32x2(vptr1, v_act0_1, v_act1_1, Q6_V_vzero()); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_offset0 = vptr0[4]; - HVX_VectorPair p_deal0 = Q6_W_vdeal_VVR(v_scale_offset0, v_scale_offset0, -2); - HVX_Vector v_scale0 = Q6_V_lo_W(p_deal0); - HVX_Vector v_offset0 = Q6_V_hi_W(p_deal0); - - HVX_Vector v_scale_offset1 = vptr1[4]; - HVX_VectorPair p_deal1 = Q6_W_vdeal_VVR(v_scale_offset1, v_scale_offset1, -2); - HVX_Vector v_scale1 = Q6_V_lo_W(p_deal1); - HVX_Vector v_offset1 = Q6_V_hi_W(p_deal1); - - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_sum_a_c0_0 = v_act0_0[9]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_sum_a_c1_0 = v_act1_0[9]; - - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_sum_a_c0_1 = v_act0_1[9]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - HVX_Vector v_sum_a_c1_1 = v_act1_1[9]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale0, v_scale_a_c0_0); - HVX_Vector v_offset_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset0, v_sum_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale0, v_scale_a_c1_0); - HVX_Vector v_offset_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset0, v_sum_a_c1_0); - - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale1, v_scale_a_c0_1); - HVX_Vector v_offset_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset1, v_sum_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale1, v_scale_a_c1_1); - HVX_Vector v_offset_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset1, v_sum_a_c1_1); - - HVX_Vector v_scaled_dot_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_add_f32_f32(v_scaled_dot_c0_0, v_offset_comb_c0_0); - - HVX_Vector v_scaled_dot_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_add_f32_f32(v_scaled_dot_c1_0, v_offset_comb_c1_0); - - HVX_Vector v_scaled_dot_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_add_f32_f32(v_scaled_dot_c0_1, v_offset_comb_c0_1); - - HVX_Vector v_scaled_dot_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_add_f32_f32(v_scaled_dot_c1_1, v_offset_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1280); const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1280); @@ -771,51 +651,7 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float HVX_Vector v_sum_float_c1 = Q6_V_vzero(); uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 1152); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 1152); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_q8_0_32x2(vptr0, v_act0_0, v_act1_0); - HVX_VectorPair v_sums1 = accum_q8_0_32x2(vptr1, v_act0_1, v_act1_1); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = vptr0[8]; - HVX_Vector v_scale_w1 = vptr1[8]; - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); @@ -952,51 +788,7 @@ static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, floa HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_4bit_32x2_lut(vptr0, v_act0_0, v_act1_0, mask_h4, lut); - HVX_VectorPair v_sums1 = accum_4bit_32x2_lut(vptr1, v_act0_1, v_act1_1, mask_h4, lut); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = vptr0[4]; - HVX_Vector v_scale_w1 = vptr1[4]; - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); @@ -1089,69 +881,7 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_4bit_32x2_lut(vptr0, v_act0_0, v_act1_0, mask_h4, lut); - HVX_VectorPair v_sums1 = accum_4bit_32x2_lut(vptr1, v_act0_1, v_act1_1, mask_h4, lut); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = hvx_vmem(tile_ptr + (kt + 0) * 640 + 512); - HVX_Vector r0_d0 = Q6_V_vdelta_VV(v_scale_w0, expand); - r0_d0 = Q6_V_vand_VV(r0_d0, e8m0_mask); - HVX_Vector v_scale_w_f32_0 = Q6_Vw_vasl_VwR(r0_d0, 23); - - HVX_Vector v_scale_w1 = hvx_vmem(tile_ptr + (kt + 1) * 640 + 512); - HVX_Vector r0_d1 = Q6_V_vdelta_VV(v_scale_w1, expand); - r0_d1 = Q6_V_vand_VV(r0_d1, e8m0_mask); - HVX_Vector v_scale_w_f32_1 = Q6_Vw_vasl_VwR(r0_d1, 23); - - HVX_Vector v_scale_a_c0_f16_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_f16_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_f16_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_f16_1 = v_act1_1[8]; - - HVX_VectorPair p_scale_a_c0_f32_0 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16_0); - HVX_VectorPair p_scale_a_c1_f32_0 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16_0); - HVX_VectorPair p_scale_a_c0_f32_1 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16_1); - HVX_VectorPair p_scale_a_c1_f32_1 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16_1); - - HVX_Vector v_scale_a_c0_0 = Q6_V_lo_W(p_scale_a_c0_f32_0); - HVX_Vector v_scale_a_c1_0 = Q6_V_lo_W(p_scale_a_c1_f32_0); - HVX_Vector v_scale_a_c0_1 = Q6_V_lo_W(p_scale_a_c0_f32_1); - HVX_Vector v_scale_a_c1_1 = Q6_V_lo_W(p_scale_a_c1_f32_1); - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f32_f32(v_scale_w_f32_0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f32_f32(v_scale_w_f32_0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f32_f32(v_scale_w_f32_1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f32_f32(v_scale_w_f32_1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { + for (uint32_t kt = 0; kt < n_k_tiles; kt++) { const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); diff --git a/ggml/src/ggml-hexagon/htp/im2col-ops.c b/ggml/src/ggml-hexagon/htp/im2col-ops.c index 26af14ed5..2e05cf3e1 100644 --- a/ggml/src/ggml-hexagon/htp/im2col-ops.c +++ b/ggml/src/ggml-hexagon/htp/im2col-ops.c @@ -14,7 +14,7 @@ #include "htp-ctx.h" #include "htp-ops.h" #include "hvx-utils.h" -#include "hex-dma.h" +#include "dma-queue.h" #include "hex-profile.h" #include "htp-vtcm.h" #include "htp-tensor.h" @@ -86,8 +86,8 @@ static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * const uint32_t OH = is_2D ? dst->ne[2] : 1; \ const uint32_t OW = dst->ne[1]; \ const uint32_t patch_stride = IC * KH * KW; \ - const float * restrict src_data = (const float *) src1->data; \ - DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + const float * restrict src_data = (const float *) (uintptr_t) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) (uintptr_t) dst->data; \ const uint32_t patch_end = ictx->patch_base + ictx->npatches; \ const uint32_t patch_start = ictx->patch_base + ictx->npatches_per_thread * ith; \ const uint32_t patch_stop = MIN(patch_start + ictx->npatches_per_thread, patch_end); \ @@ -153,155 +153,156 @@ IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx // this block's store-out) waits for both - safe because the ring is strict // FIFO and each buffer slot is only reused after its prior consumer (compute // or store-out) already finished in program order. -#define IM2COL_BLOCKED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ - static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ - struct htp_ops_context * octx = ictx->octx; \ - struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ - const struct htp_tensor * restrict src1 = octx->src[1]; \ - const struct htp_tensor * restrict dst = octx->dst; \ - const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1]; \ - const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; \ - const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; \ - const int32_t is_2D = octx->op_params[6] == 1; \ - const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ - const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ - const uint32_t IH = is_2D ? src1->ne[1] : 1; \ - const uint32_t IW = src1->ne[0]; \ - const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; \ - const uint32_t KW = octx->src[0]->ne[0]; \ - const uint32_t OH = is_2D ? dst->ne[2] : 1; \ - const uint32_t OW = dst->ne[1]; \ - const uint32_t owb = ictx->pe_owb, Wb = ictx->pe_wb; \ - const uint32_t patch_stride = IC * KH * KW; \ - const float * restrict src_data = (const float *) src1->data; \ - DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ - dma_queue * dmaq = octx->ctx->dma[ith]; \ - uint8_t * srcb_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ - uint8_t * dstb_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ - float * srcb2[2] = { (float *) srcb_base, (float *) (srcb_base + ictx->pe_src_row_bytes) }; \ - DST_CTYPE * dstb2[2] = { (DST_CTYPE *) dstb_base, (DST_CTYPE *) (dstb_base + ictx->pe_dst_row_bytes) }; \ - const uint32_t nrows = N * OH; \ - const uint32_t per_thread = ictx->pe_rows_per_thread; \ - const uint32_t row_start = per_thread * ith; \ - const uint32_t row_end = MIN(row_start + per_thread, nrows); \ - if (row_start >= row_end) \ - return; \ - const uint32_t nbpr = (OW + owb - 1) / owb; \ - const uint32_t nrows_local = row_end - row_start; \ - const uint32_t total_blocks = nrows_local * nbpr; \ - for (uint32_t bi = 0; bi < total_blocks; bi++) { \ - const uint32_t buf = bi & 1u; \ - float * srcb = srcb2[buf]; \ - DST_CTYPE * dstb = dstb2[buf]; \ - const uint32_t r = row_start + bi / nbpr; \ - const uint32_t in = r / OH; \ - const uint32_t ioh = r % OH; \ - const uint32_t c0 = (bi % nbpr) * owb; \ - const uint32_t nb = MIN(owb, OW - c0); \ - const int32_t win0 = (int32_t) c0 * s0 - p0; \ - if (bi == 0) { \ - /* prologue: stage block 0 and wait - nothing to overlap with yet */ \ - for (uint32_t ikh = 0; ikh < KH; ikh++) { \ - const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ - if (iih < 0 || iih >= (int32_t) IH) \ - continue; \ - const int32_t lo = win0 < 0 ? -win0 : 0; \ - int32_t hi = (int32_t) IW - win0; \ - if (hi > (int32_t) Wb) \ - hi = (int32_t) Wb; \ - if (hi <= lo) \ - continue; \ - const uint32_t cpw = (uint32_t) (hi - lo); \ - float * vdst = srcb + (uint64_t) ikh * Wb + (uint32_t) lo; \ - const float * vsrc = src_data + ((uint64_t) (in * IC) * IH + iih) * IW + (win0 + lo); \ - while (!dma_queue_push(dmaq, dma_make_ptr((uint8_t *) vdst, (const uint8_t *) vsrc), \ - (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float), \ - cpw * sizeof(float), IC)) { \ - dma_queue_pop(dmaq); \ - } \ - } \ - dma_queue_flush(dmaq); \ - } \ - if (bi + 1 < total_blocks) { \ - /* prefetch: stage block bi+1 into the other slot; overlaps with this block's compute below */ \ - const uint32_t nbuf = 1u - buf; \ - float * nsrcb = srcb2[nbuf]; \ - const uint32_t nr = row_start + (bi + 1) / nbpr; \ - const uint32_t nin = nr / OH; \ - const uint32_t nioh = nr % OH; \ - const uint32_t nc0 = ((bi + 1) % nbpr) * owb; \ - const int32_t nwin0 = (int32_t) nc0 * s0 - p0; \ - for (uint32_t ikh = 0; ikh < KH; ikh++) { \ - const int32_t iih = (int32_t) nioh * s1 + (int32_t) ikh * d1 - p1; \ - if (iih < 0 || iih >= (int32_t) IH) \ - continue; \ - const int32_t lo = nwin0 < 0 ? -nwin0 : 0; \ - int32_t hi = (int32_t) IW - nwin0; \ - if (hi > (int32_t) Wb) \ - hi = (int32_t) Wb; \ - if (hi <= lo) \ - continue; \ - const uint32_t cpw = (uint32_t) (hi - lo); \ - float * vdst = nsrcb + (uint64_t) ikh * Wb + (uint32_t) lo; \ - const float * vsrc = src_data + ((uint64_t) (nin * IC) * IH + iih) * IW + (nwin0 + lo); \ - while (!dma_queue_push(dmaq, dma_make_ptr((uint8_t *) vdst, (const uint8_t *) vsrc), \ - (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float), \ - cpw * sizeof(float), IC)) { \ - dma_queue_pop(dmaq); \ - } \ - } \ - } \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ - for (uint32_t j = 0; j < nb; j++) { \ - const uint32_t iow = c0 + j; \ - DST_CTYPE * dst_patch = dstb + (uint64_t) j * patch_stride; \ - const int32_t iiw0 = (int32_t) iow * s0 - p0; \ - for (uint32_t ikh = 0; ikh < KH; ikh++) { \ - const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ - const int okh = (iih >= 0 && iih < (int32_t) IH); \ - for (uint32_t iic = 0; iic < IC; iic++) { \ - DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ - if (!okh) { \ - SPLAT_FN(out_run, 0.0f, KW); \ - continue; \ - } \ - const float * vrow = srcb + ((uint64_t) (iic * KH + ikh)) * Wb; /* col win0 at idx 0*/ \ - if (d0 == 1) { \ - /* contiguous run within the staged window: [lo,hi) in-bounds, tails zero pad */ \ - const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ - int32_t hi = (int32_t) IW - iiw0; \ - if (hi > (int32_t) KW) { \ - hi = (int32_t) KW; \ - } \ - if (hi <= lo) { \ - SPLAT_FN(out_run, 0.0f, KW); \ - } else { \ - if (lo > 0) { \ - SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ - } \ - COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (vrow + (iiw0 + lo - win0)), \ - (uint32_t) (hi - lo)); \ - if (hi < (int32_t) KW) { \ - SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ - } \ - } \ - continue; \ - } \ - for (uint32_t ikw = 0; ikw < KW; ikw++) { \ - const int32_t iiw = iiw0 + (int32_t) ikw * d0; \ - out_run[ikw] = \ - (iiw < 0 || iiw >= (int32_t) IW) ? (DST_CTYPE) 0.0f : (DST_CTYPE) vrow[iiw - win0]; \ - } \ - } \ - } \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ - DST_CTYPE * ddr = dst_data + ((uint64_t) (in * OH + ioh) * OW + c0) * patch_stride; \ - dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr, (uint8_t *) dstb), \ - nb * patch_stride * (DST_ELEM), nb * patch_stride * (DST_ELEM), 1); \ - dma_queue_flush(dmaq); \ - } \ +#define IM2COL_BLOCKED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1]; \ + const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; \ + const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; \ + const int32_t is_2D = octx->op_params[6] == 1; \ + const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ + const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ + const uint32_t IH = is_2D ? src1->ne[1] : 1; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; \ + const uint32_t KW = octx->src[0]->ne[0]; \ + const uint32_t OH = is_2D ? dst->ne[2] : 1; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t owb = ictx->pe_owb, Wb = ictx->pe_wb; \ + const uint32_t patch_stride = IC * KH * KW; \ + const dma_addr_t src_data = src1->data; \ + const dma_addr_t dst_data = dst->data; \ + dma_queue * dmaq = octx->ctx->dma[ith]; \ + uint8_t * srcb_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ + uint8_t * dstb_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ + float * srcb2[2] = { (float *) srcb_base, (float *) (srcb_base + ictx->pe_src_row_bytes) }; \ + DST_CTYPE * dstb2[2] = { (DST_CTYPE *) dstb_base, (DST_CTYPE *) (dstb_base + ictx->pe_dst_row_bytes) }; \ + const uint32_t nrows = N * OH; \ + const uint32_t per_thread = ictx->pe_rows_per_thread; \ + const uint32_t row_start = per_thread * ith; \ + const uint32_t row_end = MIN(row_start + per_thread, nrows); \ + if (row_start >= row_end) \ + return; \ + const uint32_t nbpr = (OW + owb - 1) / owb; \ + const uint32_t nrows_local = row_end - row_start; \ + const uint32_t total_blocks = nrows_local * nbpr; \ + for (uint32_t bi = 0; bi < total_blocks; bi++) { \ + const uint32_t buf = bi & 1u; \ + float * srcb = srcb2[buf]; \ + DST_CTYPE * dstb = dstb2[buf]; \ + const uint32_t r = row_start + bi / nbpr; \ + const uint32_t in = r / OH; \ + const uint32_t ioh = r % OH; \ + const uint32_t c0 = (bi % nbpr) * owb; \ + const uint32_t nb = MIN(owb, OW - c0); \ + const int32_t win0 = (int32_t) c0 * s0 - p0; \ + if (bi == 0) { \ + /* prologue: stage block 0 and wait - nothing to overlap with yet */ \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + if (iih < 0 || iih >= (int32_t) IH) \ + continue; \ + const int32_t lo = win0 < 0 ? -win0 : 0; \ + int32_t hi = (int32_t) IW - win0; \ + if (hi > (int32_t) Wb) \ + hi = (int32_t) Wb; \ + if (hi <= lo) \ + continue; \ + const uint32_t cpw = (uint32_t) (hi - lo); \ + float * vdst = srcb + (size_t) ikh * Wb + lo; \ + const dma_addr_t vsrc = src_data + (size_t) (((in * IC) * IH + iih) * IW + (win0 + lo)) * sizeof(float); \ + while (!dma_queue_push(dmaq, dma_make_data(vdst, vsrc), \ + (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float), \ + cpw * sizeof(float), IC)) { \ + dma_queue_pop(dmaq); \ + } \ + } \ + dma_queue_flush(dmaq); \ + } \ + if (bi + 1 < total_blocks) { \ + /* prefetch: stage block bi+1 into the other slot; overlaps with this block's compute below */ \ + const uint32_t nbuf = 1u - buf; \ + float * nsrcb = srcb2[nbuf]; \ + const uint32_t nr = row_start + (bi + 1) / nbpr; \ + const uint32_t nin = nr / OH; \ + const uint32_t nioh = nr % OH; \ + const uint32_t nc0 = ((bi + 1) % nbpr) * owb; \ + const int32_t nwin0 = (int32_t) nc0 * s0 - p0; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) nioh * s1 + (int32_t) ikh * d1 - p1; \ + if (iih < 0 || iih >= (int32_t) IH) \ + continue; \ + const int32_t lo = nwin0 < 0 ? -nwin0 : 0; \ + int32_t hi = (int32_t) IW - nwin0; \ + if (hi > (int32_t) Wb) \ + hi = (int32_t) Wb; \ + if (hi <= lo) \ + continue; \ + const uint32_t cpw = (uint32_t) (hi - lo); \ + float * vdst = nsrcb + (size_t) ikh * Wb + lo; \ + const dma_addr_t vsrc = src_data + (size_t) (((nin * IC) * IH + iih) * IW + (nwin0 + lo)) * sizeof(float); \ + while (!dma_queue_push(dmaq, dma_make_data(vdst, vsrc), \ + (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float), \ + cpw * sizeof(float), IC)) { \ + dma_queue_pop(dmaq); \ + } \ + } \ + } \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + for (uint32_t j = 0; j < nb; j++) { \ + const uint32_t iow = c0 + j; \ + DST_CTYPE * dst_patch = dstb + (uint64_t) j * patch_stride; \ + const int32_t iiw0 = (int32_t) iow * s0 - p0; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + const int okh = (iih >= 0 && iih < (int32_t) IH); \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (!okh) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const float * vrow = srcb + ((uint64_t) (iic * KH + ikh)) * Wb; /* col win0 at idx 0*/ \ + if (d0 == 1) { \ + /* contiguous run within the staged window: [lo,hi) in-bounds, tails zero pad */ \ + const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ + int32_t hi = (int32_t) IW - iiw0; \ + if (hi > (int32_t) KW) { \ + hi = (int32_t) KW; \ + } \ + if (hi <= lo) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + } else { \ + if (lo > 0) { \ + SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ + } \ + COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (vrow + (iiw0 + lo - win0)), \ + (uint32_t) (hi - lo)); \ + if (hi < (int32_t) KW) { \ + SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ + } \ + } \ + continue; \ + } \ + for (uint32_t ikw = 0; ikw < KW; ikw++) { \ + const int32_t iiw = iiw0 + (int32_t) ikw * d0; \ + out_run[ikw] = \ + (iiw < 0 || iiw >= (int32_t) IW) ? (DST_CTYPE) 0.0f : (DST_CTYPE) vrow[iiw - win0]; \ + } \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + const dma_addr_t ddr = dst_data + (size_t) ((in * OH + ioh) * OW + c0) * patch_stride * (DST_ELEM); \ + dma_queue_push(dmaq, dma_make_data(ddr, dstb), \ + nb * patch_stride * (DST_ELEM), nb * patch_stride * (DST_ELEM), \ + nb * patch_stride * (DST_ELEM), 1); \ + dma_queue_flush(dmaq); \ + } \ } IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "blk-dma-f16") IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "blk-dma-f32") @@ -309,75 +310,76 @@ IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_f32_thread, float, hvx_copy_f32_uu, // Exact-tiling patch-embed DMA fast path (s0==KW, p0=0, d0=1; and 2D s1==KH, // p1=0, d1=1). Intentionally reads no stride/pad/dilation params so the inner // copy stays tight and fully hoisted - do NOT graft the general gather in here. -#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ - static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ - struct htp_ops_context * octx = ictx->octx; \ - struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ - const struct htp_tensor * restrict src1 = octx->src[1]; \ - const struct htp_tensor * restrict dst = octx->dst; \ - const int32_t is_2D = octx->op_params[6] == 1; \ - const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ - const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ - const uint32_t IH = is_2D ? src1->ne[1] : 1; \ - const uint32_t IW = src1->ne[0]; \ - const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; \ - const uint32_t KW = octx->src[0]->ne[0]; \ - const uint32_t OH = is_2D ? dst->ne[2] : 1; \ - const uint32_t OW = dst->ne[1]; \ - const uint32_t patch_stride = IC * KH * KW; \ - const float * restrict src_data = (const float *) src1->data; \ - DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ - dma_queue * dmaq = octx->ctx->dma[ith]; \ - uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ - uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ - float * srcb = (float *) src_base; \ - DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \ - const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows; \ - const uint32_t per_thread = ictx->pe_rows_per_thread; \ - const uint32_t row_start = ictx->pe_row_base + per_thread * ith; \ - const uint32_t row_end = MIN(row_start + per_thread, row_end_max); \ - if (row_start >= row_end) \ - return; \ - for (uint32_t r = row_start; r < row_end; r++) { \ - const uint32_t in = r / OH; \ - const uint32_t ioh = r % OH; \ - for (uint32_t ikh = 0; ikh < KH; ikh++) { \ - int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ - int ok = (iih >= 0 && iih < (int32_t) IH); \ - for (uint32_t iic = 0; iic < IC; iic++) { \ - float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \ - const float * _vsrc = \ - ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \ - dma_queue_push_ddr_to_vtcm( \ - dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \ - IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \ - } \ - } \ - for (uint32_t i = 0; i < IC * KH; i++) \ - dma_queue_pop(dmaq); \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ - for (uint32_t iow = 0; iow < OW; iow++) { \ - DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \ - for (uint32_t ikh = 0; ikh < KH; ikh++) { \ - int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ - for (uint32_t iic = 0; iic < IC; iic++) { \ - DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ - if (iih < 0 || iih >= (int32_t) IH) { \ - SPLAT_FN(out_run, 0.0f, KW); \ - continue; \ - } \ - const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \ - COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \ - } \ - } \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ - DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \ - dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \ - OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \ - dma_queue_flush(dmaq); \ - } \ +#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const int32_t is_2D = octx->op_params[6] == 1; \ + const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ + const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ + const uint32_t IH = is_2D ? src1->ne[1] : 1; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; \ + const uint32_t KW = octx->src[0]->ne[0]; \ + const uint32_t OH = is_2D ? dst->ne[2] : 1; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const dma_addr_t src_data = src1->data; \ + const dma_addr_t dst_data = dst->data; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ + uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ + uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ + float * srcb = (float *) src_base; \ + DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \ + const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows; \ + const uint32_t per_thread = ictx->pe_rows_per_thread; \ + const uint32_t row_start = ictx->pe_row_base + per_thread * ith; \ + const uint32_t row_end = MIN(row_start + per_thread, row_end_max); \ + if (row_start >= row_end) \ + return; \ + for (uint32_t r = row_start; r < row_end; r++) { \ + const uint32_t in = r / OH; \ + const uint32_t ioh = r % OH; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ + int ok = (iih >= 0 && iih < (int32_t) IH); \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + float * vdst = srcb + (size_t) (iic * KH + ikh) * IW; \ + const dma_addr_t vsrc = ok \ + ? (src_data + (size_t) ((in * IC + iic) * IH + iih) * IW * sizeof(float)) \ + : src_data; \ + dma_queue_push(dma_q, dma_make_data(vdst, vsrc), \ + IW * sizeof(float), IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \ + } \ + } \ + for (uint32_t i = 0; i < IC * KH; i++) \ + dma_queue_pop(dma_q); \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + for (uint32_t iow = 0; iow < OW; iow++) { \ + DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (iih < 0 || iih >= (int32_t) IH) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \ + COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + const dma_addr_t ddr_row = dst_data + (size_t) (in * OH + ioh) * OW * patch_stride * (DST_ELEM); \ + dma_queue_push(dma_q, dma_make_data(ddr_row, dstb), \ + OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), \ + OW * patch_stride * (DST_ELEM), 1); \ + dma_queue_flush(dma_q); \ + } \ } IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16") @@ -496,10 +498,6 @@ int op_im2col(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } - const int32_t is_2D = octx->op_params[6] == 1; const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; const uint32_t OH = is_2D ? dst->ne[2] : 1; @@ -571,12 +569,14 @@ int op_im2col(struct htp_ops_context * octx) { } } // Fall through to pure-DDR. - - if (npatches == 0) { return HTP_STATUS_OK; } + if (htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } + if (dst->type == HTP_TYPE_F16) { work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads); } else { diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index b324cfd3a..653c9a250 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -21,7 +21,7 @@ #include #include "hex-utils.h" -#include "hex-dma.h" +#include "dma-queue.h" #include "hmx-queue.h" #define GGML_COMMON_DECL_C @@ -49,33 +49,74 @@ struct htp_handle { struct htp_context * ctx; }; -static inline void * htp_mmap(uint32_t fd, uint32_t size) { +static inline uint64_t htp_mmap(uint32_t fd, uint64_t size, uint32_t flags) { +#if __HVX_ARCH__ > 79 + if (flags & HTP_BUF_EXTENDED) { + HAP_mem_req_payload_t payload; + memset(&payload, 0, sizeof(payload)); + payload.request_id = HAP_MEM_MAP; + payload.mmap.len = size; + payload.mmap.prot = HAP_MEM_CACHE_NON_SHARED | HAP_PROT_READ; + payload.mmap.flags = HAP_MEM_FLAGS_EXTENDED_MAP; + payload.mmap.fd = fd; + + if (HAP_mem_request(&payload) != 0) { + FARF(ERROR, "extended mmap failed : fd %u size %llu", fd, (unsigned long long) size); + return 0; + } + + return payload.mmap.dsp_va; + } +#else + if (flags & HTP_BUF_EXTENDED) { + FARF(ERROR, "extended mmap is unsupported on v%d", __HVX_ARCH__); + return 0; + } +#endif + + if (size > UINT32_MAX) { + FARF(ERROR, "mmap failed : size %llu exceeds 32-bit limit", (unsigned long long) size); + return 0; + } + void * va = (void *)-1; for (int retry = 0; retry < 2; retry++) { #if __HVX_ARCH__ > 73 - va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); + va = HAP_mmap2(NULL, (size_t) size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); #else if (size > HTP_MMAP_MAX_VMEM) { - FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size); + FARF(ERROR, "mmap failed : size %llu exceeds 2GB limit for HAP_mmap", (unsigned long long) size); abort(); } - va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); + va = HAP_mmap(NULL, (int) size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); #endif if (va != (void *)-1 && va != NULL) { - return va; + return (uint64_t) (uintptr_t) va; } if (retry == 0) { - FARF(HIGH, "mmap failed first try (va %p fd %u size %u), retrying...", va, fd, size); + FARF(HIGH, "mmap failed first try (va %p fd %u size %llu), retrying...", va, fd, (unsigned long long) size); } } - return NULL; + return 0; } -static inline void htp_munmap(void * va, uint32_t size) { +static inline void htp_munmap(uint64_t va, uint64_t size, uint32_t flags) { +#if __HVX_ARCH__ > 79 + if (flags & HTP_BUF_EXTENDED) { + HAP_mem_req_payload_t payload; + memset(&payload, 0, sizeof(payload)); + payload.request_id = HAP_MEM_UNMAP; + payload.munmap.dsp_va = va; + payload.munmap.len = size; + HAP_mem_request(&payload); + return; + } +#endif + #if __HVX_ARCH__ > 73 - HAP_munmap2(va, size); + HAP_munmap2((void *) (uintptr_t) va, (size_t) size); #else - HAP_munmap(va, size); + HAP_munmap((void *) (uintptr_t) va, (int) size); #endif } @@ -160,10 +201,11 @@ AEEResult htp_iface_close(remote_handle64 handle) { // release the mmaps (if any) for (uint32_t i=0; immap[i].size) { - htp_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size); + htp_munmap(ctx->mmap[i].base, ctx->mmap[i].size, ctx->mmap[i].flags); ctx->mmap[i].size = 0; - ctx->mmap[i].base = NULL; + ctx->mmap[i].base = 0; ctx->mmap[i].fd = -1; + ctx->mmap[i].flags = 0; } } @@ -184,7 +226,7 @@ AEEResult htp_iface_close(remote_handle64 handle) { return AEE_SUCCESS; } -AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) { +AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint64_t size) { struct htp_handle * h = (struct htp_handle *) handle; if (!h || !h->ctx) { return AEE_EBADPARM; @@ -203,16 +245,17 @@ AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) { for (uint32_t i=0; immap[i]; if (!m->size) { - FARF(HIGH, "mmap : fd %u size %u", fd, size); - void *va = htp_mmap(fd, size); - if (va == NULL) { - FARF(ERROR, "mmap failed : fd %u size %u", fd, (uint32_t) size); + FARF(HIGH, "mmap : fd %u size %llu", fd, (unsigned long long) size); + uint64_t va = htp_mmap(fd, size, 0); + if (va == 0) { + FARF(ERROR, "mmap failed : fd %u size %llu", fd, (unsigned long long) size); return AEE_EFAILED; } - m->base = (uint64_t) va; + m->base = va; m->fd = fd; m->size = size; + m->flags = 0; return AEE_SUCCESS; } @@ -231,11 +274,12 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) { for (uint32_t i=0; immap[i]; if (fd < 0 || m->fd == fd) { - FARF(HIGH, "unmmap : base %p fd %u size %u", (void*) m->base, m->fd, (uint32_t) m->size); - htp_munmap((void *) m->base, m->size); + FARF(HIGH, "unmmap : base 0x%llx fd %u size %llu", (unsigned long long) m->base, m->fd, (unsigned long long) m->size); + htp_munmap(m->base, m->size, m->flags); m->size = 0; m->base = NULL; m->fd = -1; + m->flags = 0; } } @@ -394,8 +438,6 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp for (uint32_t i = 0; i < n_hvx; i++) { size_dma = hex_align_up(size_dma, dma_queue_alignof()); size_dma += dma_queue_sizeof(256); - size_dma = hex_align_up(size_dma, dma_queue_alignof()); - size_dma += dma_queue_alias_sizeof(); } offset = offset_dma + size_dma; @@ -538,16 +580,11 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp // Initialize DMA queues uint8_t * dma_ptr_curr = (uint8_t *) ((uintptr_t) block + offset_dma); size_t size_dma_q = dma_queue_sizeof(256); - size_t size_dma_alias = dma_queue_alias_sizeof(); for (int i = 0; i < ctx->n_threads; i++) { dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof()); - ctx->dma_cached[i] = dma_queue_init(dma_ptr_curr, 256, (uintptr_t) ctx->vtcm_base, ctx->vtcm_size, &ctx->trace[i]); + ctx->dma[i] = dma_queue_init(dma_ptr_curr, 256, &ctx->trace[i]); dma_ptr_curr += size_dma_q; - - dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof()); - ctx->dma[i] = dma_queue_alias_init(dma_ptr_curr, ctx->dma_cached[i], 1); - dma_ptr_curr += size_dma_alias; } ctx->ddr_spad_size = 512 * 1024; // 512 KB @@ -608,8 +645,7 @@ AEEResult htp_iface_stop(remote_handle64 handle) { work_queue_free(ctx->work_queue); for (int i = 0; i < ctx->n_threads; i++) { - dma_queue_alias_free(ctx->dma[i]); - dma_queue_free(ctx->dma_cached[i]); + dma_queue_free(ctx->dma[i]); } if (ctx->hmx_queue) { @@ -908,7 +944,7 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct for (uint32_t i=0; immap + i; - if (m->size && m->fd == b->fd) { + if (m->size && m->fd == b->fd && m->flags == b->flags) { b->base = m->base; *m_reuse |= (1 << i); return true; @@ -920,11 +956,12 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) { if (m->size) { - FARF(ALWAYS, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); - htp_munmap((void *) m->base, m->size); + FARF(ALWAYS, "unmap : fd %u base 0x%llx size %llu", m->fd, (unsigned long long) m->base, (unsigned long long) m->size); + htp_munmap(m->base, m->size, m->flags); m->size = 0; m->base = 0; m->fd = -1; + m->flags = 0; } } @@ -935,17 +972,18 @@ static inline bool mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { struct htp_mmap *m = &ctx->mmap[i]; if (!m->size) { - void *va = htp_mmap(b->fd, b->size); - if (va == NULL) { - FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %u", b->fd, (uint32_t) b->size); + uint64_t va = htp_mmap(b->fd, b->size, b->flags); + if (va == 0) { + FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %llu", b->fd, (unsigned long long) b->size); return false; } - m->base = b->base = (uint64_t) va; + m->base = b->base = va; m->fd = b->fd; m->size = b->size; + m->flags = b->flags; - FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); + FARF(ALWAYS, "mmap : fd %u base 0x%llx size %llu flags 0x%x", m->fd, (unsigned long long) m->base, (unsigned long long) m->size, m->flags); return true; } } @@ -964,14 +1002,22 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin // See what we can reuse for (uint32_t i=0; i < n_bufs; i++) { struct htp_buf_desc *b = bufs + i; - if (reuse_buf(ctx, &m_reuse, b)) { b_reuse++; } else { e_vmem += b->size; } - FARF(HIGH, "prep-buf #%u : pass0 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags); + if (reuse_buf(ctx, &m_reuse, b)) { + b_reuse++; + } else if (!(b->flags & HTP_BUF_EXTENDED)) { + e_vmem += b->size; + } + FARF(HIGH, "prep-buf #%u : pass0 fd %u base 0x%llx size %llu flags 0x%x", i, b->fd, (unsigned long long) b->base, (unsigned long long) b->size, b->flags); } if (b_reuse == n_bufs) return; // all bufs reuse existing mappings // See how much vmem we have mmaped right now - for (uint32_t i=0; immap[i].size; } + for (uint32_t i=0; immap[i].flags & HTP_BUF_EXTENDED)) { + m_vmem += ctx->mmap[i].size; + } + } FARF(HIGH, "prep-bufs : pass1 mmap-vmem %zu extra-vmem %zu max-vmem %zu : n-bufs %u b-reuse %u", (size_t) m_vmem, (size_t) e_vmem, (size_t) ctx->max_vmem, n_bufs, b_reuse); @@ -980,7 +1026,9 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin // Drop unused mappings for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { bool used = m_reuse & (1<mmap + i); } + if (!used && !(ctx->mmap[i].flags & HTP_BUF_EXTENDED)) { + drop_mmap(ctx, ctx->mmap + i); + } } } @@ -992,35 +1040,40 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin mmap_ok = false; break; } - FARF(HIGH, "prep-buf #%u : pass1 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags); + FARF(HIGH, "prep-buf #%u : pass1 fd %u base 0x%llx size %llu flags 0x%x", i, b->fd, (unsigned long long) b->base, (unsigned long long) b->size, b->flags); } if (!mmap_ok) { - // Attempt clean defragmentation: drop all mappings and remap (pass 2) - FARF(HIGH, "prep-bufs : dropping all mappings to defragment address space"); - for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { drop_mmap(ctx, ctx->mmap + i); } + // Attempt defragmentation: drop 32-bit mappings and remap (pass 2) + FARF(HIGH, "prep-bufs : dropping 32-bit mappings to defragment address space"); + for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { + if (!(ctx->mmap[i].flags & HTP_BUF_EXTENDED)) { + drop_mmap(ctx, ctx->mmap + i); + } + } for (uint32_t i=0; i < n_bufs; i++) { struct htp_buf_desc *b = bufs + i; - b->base = 0; + if (!(b->flags & HTP_BUF_EXTENDED)) { + b->base = 0; + } if (!mmap_buf(ctx, b)) { - FARF(ERROR, "prep-bufs : mmap failed after defragmentation (fd %u size %u)", b->fd, (uint32_t) b->size); + FARF(ERROR, "prep-bufs : mmap failed after defragmentation (fd %u size %llu)", b->fd, (unsigned long long) b->size); abort(); } - FARF(HIGH, "prep-buf #%u : pass2 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags); + FARF(HIGH, "prep-buf #%u : pass2 fd %u base 0x%llx size %llu flags 0x%x", i, b->fd, (unsigned long long) b->base, (unsigned long long) b->size, b->flags); } } } static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) { - uint32_t offset = t->data; - uint32_t size = t->size; + uint64_t offset = t->data; uint32_t bi = t->bi; - t->data = (uint32_t) (bufs[bi].base + offset); // update data to the actual pointer + t->data = bufs[bi].base + offset; // update data to the actual pointer - FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data, - t->ne[0], t->ne[1], t->ne[3], t->ne[3]); + FARF(HIGH, "prep-tensor #%u: bi %u offset %llu size %u data 0x%llx : %u:%u:%u:%u", idx, t->bi, (unsigned long long) offset, t->size, (unsigned long long) t->data, + t->ne[0], t->ne[1], t->ne[2], t->ne[3]); } static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t n_tens) { @@ -1050,15 +1103,13 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs uint16_t src_idx = op->src[i]; if (src_idx == 0xffff) { octx->src[i] = NULL; - octx->src_dma[i] = NULL; continue; } struct htp_tensor *src = tens + src_idx; octx->src[i] = src; - octx->src_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma; - FARF(HIGH, "prep-src #%u: data %p size %u : %u:%u:%u:%u", op->src[i], (void*) src->data, src->size, + FARF(HIGH, "prep-src #%u: data 0x%llx size %u : %u:%u:%u:%u", op->src[i], (unsigned long long) src->data, src->size, src->ne[0], src->ne[1], src->ne[2], src->ne[3]); } @@ -1069,14 +1120,12 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs uint16_t dst_idx = op->dst[i]; if (dst_idx == 0xffff) { octx->dsts[i] = NULL; - octx->dst_dma[i] = NULL; continue; } struct htp_tensor *dst = tens + dst_idx; octx->dsts[i] = dst; - octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma; - FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size, + FARF(HIGH, "prep-dst[%u] #%u: data 0x%llx size %u : %u:%u:%u:%u", i, dst_idx, (unsigned long long) dst->data, dst->size, dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]); } diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index e16cfdcbe..a09bc7a28 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -11,7 +11,7 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hvx-utils.h" #include "hvx-dump.h" #include "hvx-arith.h" @@ -25,22 +25,13 @@ #include "matmul-ops.h" #include "htp-vtcm.h" -static void hvx_tensor_add_f32_grid( - const struct htp_tensor * restrict dst, - const struct htp_tensor * restrict src2, - uint32_t start_row, - uint32_t end_row, - uint32_t start_col, - uint32_t end_col, - const struct fastdiv_values * div_ne11_12, - const struct fastdiv_values * div_ne11 -); - typedef struct { float *dst; - const float *src2; + dma_addr_t src2_addr; + size_t src2_bytes; const float *activation; - const __fp16 *weight; + dma_addr_t weight; + dma_queue * weight_dma; int m; int k; int n; @@ -66,6 +57,15 @@ typedef struct { struct fastdiv_values div_r3; } hmx_mm_f16_f32_batched_params_t; +static bool htp_matmul_has_extended_weight(const struct htp_ops_context * octx, uint32_t n_weights) { + for (uint32_t i = 0; i < n_weights; ++i) { + if (htp_tensor_is_extended(octx->src[i])) { + return true; + } + } + return false; +} + struct htp_mm_context { const char * type; struct htp_ops_context * octx; @@ -94,6 +94,8 @@ struct htp_mm_context { uint32_t src0_row_end; uint32_t src0_row_size_padded; uint32_t src1_nrows; + uint32_t cur_m_start; + uint32_t cur_m_rows; struct fastdiv_values mm_div_ne12_ne1; struct fastdiv_values mm_div_ne1; @@ -228,7 +230,7 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = { #define htp_matmul_preamble \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ - dma_queue *dma_queue = octx->ctx->dma[ith]; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; \ htp_matmul_tensors_preamble; @@ -244,291 +246,213 @@ static inline void hvx_mm_run_quant_task(struct htp_mm_context * mmctx, unsigned } } -// *** matmul with support for 4d tensors and full broadcasting -static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { - htp_matmul_preamble; - - assert(ne12 % ne02 == 0); - assert(ne13 % ne03 == 0); - - // This is the size of the first dimension of the result, so we can iterate that way. (see the ASSERT above, these are the same numbers) - const uint32_t nr0 = ne0; - - // This is the size of the rest of the dimensions of the result - const uint32_t nr1 = ne1 * ne2 * ne3; - - const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; - - // distribute the thread work across the inner or outer loop based on which one is larger - uint32_t dr0, dr1, ith0, ith1; - if (nr0 > nr1) { - dr0 = fastdiv(src0_nrows + nth - 1, &octx->n_threads_div); - dr1 = nr1; - ith0 = ith; - ith1 = 0; - } else { - dr0 = src0_nrows; - dr1 = fastdiv(nr1 + nth - 1, &octx->n_threads_div); - ith0 = 0; - ith1 = ith; - } - - const uint32_t ir0_start = mmctx->src0_row_start + dr0 * ith0; - const uint32_t ir0_end = MIN(ir0_start + dr0, mmctx->src0_row_end); - - const uint32_t ir1_start = dr1 * ith1; - const uint32_t ir1_end = MIN(ir1_start + dr1, nr1); - - // no work for this thread - if (ir0_start >= ir0_end || ir1_start >= ir1_end) { - return; - } - - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start); - - const uint32_t blck_0 = 64; - const uint32_t blck_1 = 64; - - for (uint32_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) { - for (uint32_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) { - for (uint32_t ir1 = iir1; ir1 < MIN(iir1 + blck_1, ir1_end); ir1++) { - const uint32_t i13 = fastdiv(ir1, &mmctx->mm_div_ne12_ne1); - const uint32_t i12 = fastdiv(ir1 - i13 * ne12 * ne1, &mmctx->mm_div_ne1); - const uint32_t i11 = (ir1 - i13 * ne12 * ne1 - i12 * ne1); - - // broadcast src0 into src1 - const uint32_t i03 = fastdiv(i13, &mmctx->mm_div_r3); - const uint32_t i02 = fastdiv(i12, &mmctx->mm_div_r2); - - const uint32_t i1 = i11; - const uint32_t i2 = i12; - const uint32_t i3 = i13; - - const uint8_t * restrict src0_base = (const uint8_t *) src0->data + (0 + i02 * nb02 + i03 * nb03); - const uint8_t * restrict src1_col = (const uint8_t *) src1->data + (i11 * nb11 + i12 * nb12 + i13 * nb13); - float * dst_col = (float *) ((uint8_t * restrict) dst->data + (i1 * nb1 + i2 * nb2 + i3 * nb3)); - - const uint32_t ir0_block_end = MIN(iir0 + blck_0, ir0_end); - for (uint32_t ir0 = iir0; ir0 < ir0_block_end; ir0++) { - const uint8_t * restrict src0_row = src0_base + ir0 * nb01; - mmctx->vec_dot_1x1(ne00, &dst_col[ir0], src0_row, src1_col); - } - } - } - } - - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start); - if (src2) { - hvx_tensor_add_f32_grid(dst, src2, ir1_start, ir1_end, ir0_start, ir0_end, &mmctx->mm_div_ne12_ne1, &mmctx->mm_div_ne1); - } -} // hvx kernels first: the HMX Q6_K dequantizer reuses unpack_q6_k_group from there #include "hvx-mm-kernels-tiled.h" +#include "hvx-mm-kernels-float.h" #include "hmx-mm-kernels-tiled.h" -#include "hvx-mm-kernels-flat.h" // Specialized repacked matmul macros -#define MATMUL_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - htp_matmul_preamble; \ - \ - const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ - const uint32_t src1_nrows = ne11 * ne12 * ne13; \ - \ - const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - const size_t dst_row_size = nb1; \ - const size_t src1_row_size = nb11; \ - const size_t src1_stride = mmctx->vtcm_src1_stride; \ - const size_t src2_stride = src2 ? ((src2->ne[1] == 1) ? 0 : src2->nb[1]) : 0; \ - \ - uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict src1_data = mmctx->vtcm_src1; \ - \ - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - \ - uint32_t ct_start = src0_start_row / 32; \ - uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - if (src0_start_row < src0_end_row) { \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)ne0 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ - float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); \ - float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); \ - \ - float * dst_ptr0 = &dst_row0[ct * 32]; \ - float * dst_ptr1 = &dst_row1[ct * 32]; \ - \ - const float * src2_ptr0 = NULL; \ - const float * src2_ptr1 = NULL; \ - if (src2) { \ - const float * restrict src2_row0 = (const float *) ((const uint8_t *) src2->data + ((ir1+0) * src2_stride)); \ - const float * restrict src2_row1 = (const float *) ((const uint8_t *) src2->data + ((ir1+1) * src2_stride)); \ - src2_ptr0 = &src2_row0[ct * 32]; \ - src2_ptr1 = &src2_row1[ct * 32]; \ - } \ - DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, src2_ptr0, src2_ptr1); \ - } \ - \ - for (; ir1 < src1_nrows; ++ir1) { \ - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \ - float * dst_ptr = &dst_row[ct * 32]; \ - \ - const float * src2_ptr = NULL; \ - if (src2) { \ - const float * restrict src2_row = (const float *) ((const uint8_t *) src2->data + (ir1 * src2_stride)); \ - src2_ptr = &src2_row[ct * 32]; \ - } \ - DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, src2_ptr); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ +#define MATMUL_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ +static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + htp_matmul_preamble; \ + \ + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ + const uint32_t src1_nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); \ + const uint32_t cur_m_start = mmctx->cur_m_start; \ + \ + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + \ + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ + const uint32_t n_prefetch = kparams->n_prefetch; \ + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ + \ + const size_t dst_row_size = nb1; \ + const size_t src1_row_size = nb11; \ + const size_t src1_stride = mmctx->vtcm_src1_stride; \ + const size_t src2_stride = src2 ? ((src2->ne[1] == 1) ? 0 : src2->nb[1]) : 0; \ + \ + uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ + uint8_t * restrict src1_data = mmctx->vtcm_src1; \ + \ + const dma_addr_t src0_row = src0->data; \ + \ + const uint32_t tile_size = TILE_SIZE; \ + const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ + \ + uint32_t n_k_tiles_w = ne00 / 32; \ + uint32_t n_k_tiles_a = ne10 / 32; \ + uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ + uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ + \ + uint32_t ct_start = src0_start_row / 32; \ + uint32_t ct_end = (src0_end_row + 31) / 32; \ + \ + uint32_t push_ct = ct_start; \ + if (src0_start_row < src0_end_row) { \ + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ + src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + } \ + } \ + \ + hvx_mm_run_quant_task(mmctx, ith); \ + \ + if (src0_start_row >= src0_end_row) { \ + return; \ + } \ + \ + for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ + \ + int valid_rows = (int)ne0 - (int)(ct * 32); \ + valid_rows = MIN(32, MAX(0, valid_rows)); \ + \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + uint32_t ir1 = 0; \ + for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ + const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ + const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ + float * restrict dst_row0 = (float *) (dst->data + ((cur_m_start + ir1+0) * dst_row_size)); \ + float * restrict dst_row1 = (float *) (dst->data + ((cur_m_start + ir1+1) * dst_row_size)); \ + \ + float * dst_ptr0 = &dst_row0[ct * 32]; \ + float * dst_ptr1 = &dst_row1[ct * 32]; \ + \ + const float * src2_ptr0 = NULL; \ + const float * src2_ptr1 = NULL; \ + if (src2) { \ + const float * restrict src2_row0 = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1+0) * src2_stride)); \ + const float * restrict src2_row1 = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1+1) * src2_stride)); \ + src2_ptr0 = &src2_row0[ct * 32]; \ + src2_ptr1 = &src2_row1[ct * 32]; \ + } \ + DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, src2_ptr0, src2_ptr1); \ + } \ + \ + for (; ir1 < src1_nrows; ++ir1) { \ + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ + float * restrict dst_row = (float *) (dst->data + ((cur_m_start + ir1) * dst_row_size)); \ + float * dst_ptr = &dst_row[ct * 32]; \ + \ + const float * src2_ptr = NULL; \ + if (src2) { \ + const float * restrict src2_row = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1) * src2_stride)); \ + src2_ptr = &src2_row[ct * 32]; \ + } \ + DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, src2_ptr); \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + \ + if (push_ct < ct_end) { \ + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), \ + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + push_ct++; \ + } \ + } \ } -#define MATVEC_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X1) \ -static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - htp_matmul_preamble; \ - \ - const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ - \ - const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - const size_t dst_row_size = nb1; \ - const size_t src1_row_size = nb11; \ - const size_t src1_stride = mmctx->vtcm_src1_stride; \ - \ - uint8_t * vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ - uint8_t * vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * src1_data = mmctx->vtcm_src1; \ - \ - float * tmp = (float *) vtcm_dst_ptr; \ - \ - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ - \ - const uint8_t * restrict src1_col = (const uint8_t *) src1_data; \ - float * restrict dst_col = (float *) dst->data; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - \ - uint32_t ct_start = src0_start_row / 32; \ - uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - if (src0_start_row < src0_end_row) { \ - if (src2) { \ - float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; \ - const float * src2_ptr = (const float *) src2->data + src0_start_row; \ - int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ - if (slice_size > 0) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), \ - slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \ - dma_queue_pop_nowait(dma_queue); \ - } \ - } \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ - \ - float * dst_ptr = &tmp[ct * 32 - src0_start_row]; \ - int valid_rows = (int)ne0 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ - \ - int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ - if (copy_cnt > 0) { \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ - if (src2) { \ - hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], \ - (const uint8_t *) tmp, \ - (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), \ - copy_cnt); \ - } else { \ - hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ - } \ +#define MATVEC_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X1) \ +static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + htp_matmul_preamble; \ + \ + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ + \ + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + \ + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ + const uint32_t n_prefetch = kparams->n_prefetch; \ + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ + \ + const size_t dst_row_size = nb1; \ + const size_t src1_row_size = nb11; \ + const size_t src1_stride = mmctx->vtcm_src1_stride; \ + \ + uint8_t * vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ + uint8_t * vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ + uint8_t * src1_data = mmctx->vtcm_src1; \ + \ + float * tmp = (float *) vtcm_dst_ptr; \ + \ + const dma_addr_t src0_row = src0->data; \ + \ + const uint8_t * restrict src1_col = (const uint8_t *) src1_data; \ + float * restrict dst_col = (float *) dst->data; \ + \ + const uint32_t tile_size = TILE_SIZE; \ + const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ + \ + uint32_t n_k_tiles_w = ne00 / 32; \ + uint32_t n_k_tiles_a = ne10 / 32; \ + uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ + uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ + \ + uint32_t ct_start = src0_start_row / 32; \ + uint32_t ct_end = (src0_end_row + 31) / 32; \ + \ + uint32_t push_ct = ct_start; \ + if (src0_start_row < src0_end_row) { \ + if (src2) { \ + float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; \ + const dma_addr_t src2_addr = src2->data + src0_start_row * sizeof(float); \ + int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ + if (slice_size > 0) { \ + dma_queue_push(dma_q, dma_make_data(vtcm_src2_ptr, src2_addr), \ + slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \ + dma_queue_pop_nowait(dma_q); \ + } \ + } \ + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ + src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + } \ + } \ + \ + hvx_mm_run_quant_task(mmctx, ith); \ + \ + if (src0_start_row >= src0_end_row) { \ + return; \ + } \ + \ + for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ + \ + float * dst_ptr = &tmp[ct * 32 - src0_start_row]; \ + int valid_rows = (int)ne0 - (int)(ct * 32); \ + valid_rows = MIN(32, MAX(0, valid_rows)); \ + \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + \ + if (push_ct < ct_end) { \ + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), \ + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + push_ct++; \ + } \ + } \ + \ + int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ + if (copy_cnt > 0) { \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ + if (src2) { \ + hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], \ + (const uint8_t *) tmp, \ + (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), \ + copy_cnt); \ + } else { \ + hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ + } \ } #define MATMUL_NX_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ @@ -555,19 +479,18 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v uint32_t n_k_tiles_a = ne10 / 32; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ - \ hvx_mm_run_quant_task(mmctx, ith); \ \ for (uint32_t widx = 0; widx < n_weights; widx++) { \ const struct htp_tensor * restrict src_w = octx->src[widx]; \ const struct htp_tensor * restrict dst = octx->dsts[widx]; \ if (!src_w || !dst) continue; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ \ const uint32_t ne00 = src_w->ne[0]; \ const uint32_t ne01 = src_w->ne[1]; \ const size_t dst_row_size = dst->nb[1]; \ - const uint8_t * restrict src_w_row = (const uint8_t *) src_w->data; \ + const dma_addr_t src_w_row = src_w->data; \ \ uint32_t n_k_tiles_w = ne00 / 32; \ uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ @@ -595,12 +518,12 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v \ uint32_t push_ct = ct_start; \ for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_weight_ptr + d * tile_row_transfer_size_aligned, \ + dma_queue_push(dma_q, dma_make_data(vtcm_weight_ptr + d * tile_row_transfer_size_aligned, \ src_w_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ } \ \ for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ int valid_rows = (int)ne01 - (int)(ct * 32); \ valid_rows = MIN(32, MAX(0, valid_rows)); \ \ @@ -627,7 +550,7 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ if (push_ct < ct_end) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src_w_row + push_ct * tile_row_stride), \ + dma_queue_push(dma_q, dma_make_data(w_tile, src_w_row + push_ct * tile_row_stride), \ aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ push_ct++; \ } \ @@ -642,52 +565,57 @@ MATMUL_2D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x2, tiled_vec_do MATMUL_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) MATMUL_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) -MATMUL_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) -MATMUL_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) -MATMUL_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) -MATMUL_2D_REPACKED_IMPL(q6_k_flat, 896, flat_vec_dot_q6_k_32x2, flat_vec_dot_q6_k_32x1) -MATMUL_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) -MATMUL_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) - -#define QUANTIZE_IMPL(name, log_name, kernel_fn, dst_row_size_expr) \ -static void name(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_mm_context * mmctx = data; \ - struct htp_ops_context * octx = mmctx->octx; \ - const struct htp_tensor * src = mmctx->act; \ - const uint32_t ne0 = src->ne[0]; \ - const uint32_t ne1 = src->ne[1]; \ - const uint32_t ne2 = src->ne[2]; \ - const uint32_t ne3 = src->ne[3]; \ - const uint32_t nrows = ne1 * ne2 * ne3; \ - const uint32_t nrows_per_thread = mmctx->n_quant_rows_per_thread; \ - \ - const uint32_t ir_first = nrows_per_thread * ith; \ - if (ir_first >= nrows) { \ - return; \ - } \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ - \ - uint8_t * restrict dst = mmctx->vtcm_src1; \ - const uint32_t ir_last = MIN(ir_first + nrows_per_thread, nrows); \ - const size_t src_row_size = src->nb[1]; \ - const size_t dst_row_size = (dst_row_size_expr); \ - const uint8_t * restrict src_data = (const uint8_t *) src->data + (src_row_size * ir_first); \ - uint8_t * restrict dst_data = (uint8_t *) dst + (dst_row_size * ir_first); \ - uint8_t * restrict tmp_data = (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith); \ - kernel_fn(src_data, dst_data, tmp_data, ne0, ir_last - ir_first, src_row_size, dst_row_size); \ - \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ +#define QUANTIZE_IMPL(name, log_name, kernel_fn, dst_row_size_expr) \ +static void name(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_mm_context * mmctx = data; \ + struct htp_ops_context * octx = mmctx->octx; \ + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ + const struct htp_tensor * src = mmctx->act; \ + const uint32_t ne0 = src->ne[0]; \ + const uint32_t nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : mmctx->src1_nrows; \ + const uint32_t nrows_per_thread = mmctx->n_quant_rows_per_thread; \ + \ + const uint32_t ir_first = nrows_per_thread * ith; \ + if (ir_first >= nrows) { \ + return; \ + } \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ + \ + uint8_t * restrict dst = mmctx->vtcm_src1; \ + const uint32_t ir_last = MIN(ir_first + nrows_per_thread, nrows); \ + const size_t src_row_size = src->nb[1]; \ + const size_t dst_row_size = (dst_row_size_expr); \ + uint8_t * restrict tmp_data = (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith); \ + \ + const bool is_contiguous = (src->nb[2] == src->ne[1] * src->nb[1]) && (src->nb[3] == src->ne[2] * src->nb[2]); \ + if (is_contiguous) { \ + const uint8_t * restrict src_data = (const uint8_t *) src->data + (src_row_size * (mmctx->cur_m_start + ir_first)); \ + uint8_t * restrict dst_data = (uint8_t *) dst + (dst_row_size * ir_first); \ + kernel_fn(src_data, dst_data, tmp_data, ne0, ir_last - ir_first, src_row_size, dst_row_size); \ + } else { \ + const uint32_t ne12_ne1 = src->ne[2] * src->ne[1]; \ + for (uint32_t ir = ir_first; ir < ir_last; ++ir) { \ + const uint32_t ir1 = mmctx->cur_m_start + ir; \ + const uint32_t i13 = fastdiv(ir1, &kparams->div_ne12_ne1); \ + const uint32_t rem = ir1 - i13 * ne12_ne1; \ + const uint32_t i12 = fastdiv(rem, &kparams->div_ne1); \ + const uint32_t i11 = rem - i12 * src->ne[1]; \ + const uint8_t * restrict row_src = (const uint8_t *) src->data + ((size_t) i11 * src->nb[1] + (size_t) i12 * src->nb[2] + (size_t) i13 * src->nb[3]); \ + uint8_t * restrict row_dst = dst + (dst_row_size * ir); \ + kernel_fn(row_src, row_dst, tmp_data, ne0, 1, src_row_size, dst_row_size); \ + } \ + } \ + \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ } QUANTIZE_IMPL(quantize_f32_q8_0_tiled, "quantize-f32-q8_0_tiled", quantize_f32_q8_0_tiled_kernel, htp_mm_q8_0_tiled_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_q8_1_tiled, "quantize-f32-q8_1_tiled", quantize_f32_q8_1_tiled_kernel, htp_mm_q8_1_tiled_row_size(ne0)) -QUANTIZE_IMPL(quantize_f32_q8_0_flat, "quantize-f32-q8_0_flat", quantize_f32_q8_0_flat_kernel, htp_mm_q8_0_flat_row_size(ne0)) -QUANTIZE_IMPL(quantize_f32_q8_1_flat, "quantize-f32-q8_1_flat", quantize_f32_q8_1_flat_kernel, htp_mm_q8_1_flat_row_size(ne0)) -QUANTIZE_IMPL(quantize_f32_f32_flat, "quantize-f32-f32", quantize_f32_f32_flat_kernel, mmctx->vtcm_src1_stride) -QUANTIZE_IMPL(quantize_f32_f16_flat, "quantize-f32-f16", quantize_f32_f16_flat_kernel, mmctx->vtcm_src1_stride) -QUANTIZE_IMPL(quantize_f16_f16_flat, "quantize-f16-f16", quantize_f16_f16_flat_kernel, mmctx->vtcm_src1_stride) +QUANTIZE_IMPL(quantize_f32_f32, "quantize-f32-f32", quantize_f32_f32_kernel, mmctx->vtcm_src1_stride) +QUANTIZE_IMPL(quantize_f32_f16, "quantize-f32-f16", quantize_f32_f16_kernel, mmctx->vtcm_src1_stride) +QUANTIZE_IMPL(quantize_f16_f16, "quantize-f16-f16", quantize_f16_f16_kernel, mmctx->vtcm_src1_stride) static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = data; @@ -744,25 +672,146 @@ MATVEC_2D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x1) MATVEC_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x1) MATVEC_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x1) -MATVEC_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x1) -MATVEC_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x1) -MATVEC_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x1) -MATVEC_2D_REPACKED_IMPL(q6_k_flat, 896, flat_vec_dot_q6_k_32x1) -MATVEC_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x1) -MATVEC_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x1) - - MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) -MATMUL_NX_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) -MATMUL_NX_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) +#define MATMUL_4D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ +static void hvx_mm_4d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + htp_matmul_preamble; \ + \ + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ + const uint32_t cur_m_rows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); \ + const uint32_t cur_m_start = mmctx->cur_m_start; \ + \ + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + \ + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ + const uint32_t n_prefetch = kparams->n_prefetch; \ + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ + \ + const size_t dst_row_size = nb1; \ + const size_t src1_stride = mmctx->vtcm_src1_stride; \ + \ + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ + uint8_t * restrict src1_data = mmctx->vtcm_src1; \ + \ + const uint32_t tile_size = TILE_SIZE; \ + const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ + \ + const uint32_t n_k_tiles_w = ne00 / 32; \ + const uint32_t n_k_tiles_a = ne10 / 32; \ + const uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ + const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ + const uint32_t src0_slice_stride = ((ne01 + 31) / 32) * tile_row_stride; \ + \ + const uint32_t ct_start = src0_start_row / 32; \ + const uint32_t ct_end = (src0_end_row + 31) / 32; \ + \ + hvx_mm_run_quant_task(mmctx, ith); \ + \ + if (src0_start_row >= src0_end_row || cur_m_rows == 0) { \ + return; \ + } \ + \ + const uint32_t total_batches = ne12 * ne13; \ + const uint32_t b_start = fastdiv(cur_m_start, &kparams->div_ne1); \ + uint32_t b_end = fastdiv(cur_m_start + cur_m_rows + ne11 - 1, &kparams->div_ne1); \ + b_end = MIN(b_end, total_batches); \ + \ + uint32_t b_grp_start = b_start; \ + while (b_grp_start < b_end) { \ + const uint32_t b3 = fastdiv(b_grp_start, &kparams->div_ne12); \ + const uint32_t b2 = b_grp_start - b3 * ne12; \ + const uint32_t i02 = fastdiv(b2, &kparams->div_r2); \ + const uint32_t i03 = fastdiv(b3, &kparams->div_r3); \ + \ + uint32_t b_grp_end = b_grp_start + 1; \ + while (b_grp_end < b_end) { \ + const uint32_t cur_b3 = fastdiv(b_grp_end, &kparams->div_ne12); \ + const uint32_t cur_b2 = b_grp_end - cur_b3 * ne12; \ + const uint32_t cur_i02 = fastdiv(cur_b2, &kparams->div_r2); \ + const uint32_t cur_i03 = fastdiv(cur_b3, &kparams->div_r3); \ + if (cur_i02 != i02 || cur_i03 != i03) { \ + break; \ + } \ + b_grp_end++; \ + } \ + \ + const uint32_t slice_idx = i03 * ne02 + i02; \ + const dma_addr_t src0_slice = src0->data + (size_t) slice_idx * src0_slice_stride; \ + \ + uint32_t push_ct = ct_start; \ + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ + src0_slice + (size_t) push_ct * tile_row_stride), \ + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + } \ + \ + for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ + \ + int valid_rows = (int)ne0 - (int)(ct * 32); \ + valid_rows = MIN(32, MAX(0, valid_rows)); \ + \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + for (uint32_t b = b_grp_start; b < b_grp_end; b++) { \ + const uint32_t b_m_start = b * ne11; \ + const uint32_t m_first = MAX(cur_m_start, b_m_start); \ + const uint32_t m_last = MIN(cur_m_start + cur_m_rows, b_m_start + ne11); \ + if (m_first >= m_last) continue; \ + \ + const uint32_t cur_b3 = fastdiv(b, &kparams->div_ne12); \ + const uint32_t cur_b2 = b - cur_b3 * ne12; \ + uint8_t * dst_batch_base = (uint8_t *) dst->data + (size_t) cur_b2 * nb2 + (size_t) cur_b3 * nb3; \ + \ + const uint32_t chunk_m_offset = m_first - cur_m_start; \ + const uint32_t dst_m_offset = m_first - b_m_start; \ + const uint32_t batch_nrows = m_last - m_first; \ + \ + uint32_t ir1 = 0; \ + for (; ir1 + 1 < batch_nrows; ir1 += 2) { \ + const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (chunk_m_offset + ir1 + 0) * src1_stride); \ + const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (chunk_m_offset + ir1 + 1) * src1_stride); \ + float * restrict dst_row0 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 0) * dst_row_size); \ + float * restrict dst_row1 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 1) * dst_row_size); \ + float * dst_ptr0 = &dst_row0[ct * 32]; \ + float * dst_ptr1 = &dst_row1[ct * 32]; \ + DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, NULL, NULL); \ + } \ + for (; ir1 < batch_nrows; ++ir1) { \ + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (chunk_m_offset + ir1) * src1_stride); \ + float * restrict dst_row = (float *) (dst_batch_base + (dst_m_offset + ir1) * dst_row_size); \ + float * dst_ptr = &dst_row[ct * 32]; \ + DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + \ + if (push_ct < ct_end) { \ + dma_queue_push(dma_q, dma_make_data(w_tile, src0_slice + (size_t) push_ct * tile_row_stride), \ + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + push_ct++; \ + } \ + } \ + b_grp_start = b_grp_end; \ + } \ + if (src2) { \ + hvx_tensor_add_f32_grid(dst, src2, cur_m_start, cur_m_start + cur_m_rows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); \ + } \ +} + +MATMUL_4D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) +MATMUL_4D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) +MATMUL_4D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) +MATMUL_4D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x2, tiled_vec_dot_q6_k_32x1) +MATMUL_4D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) +MATMUL_4D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; @@ -773,7 +822,8 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { const uint32_t prefetch_mask = n_prefetch - 1; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; // src0 rows - const uint32_t src1_nrows = ne11 * ne12 * ne13; // src1 rows + const uint32_t src1_nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : mmctx->src1_nrows; // src1 rows + const uint32_t cur_m_start = mmctx->cur_m_start; const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); @@ -793,7 +843,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; uint8_t * restrict src1_data = mmctx->vtcm_src1; - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; + const dma_addr_t src0_row = src0->data; // Prefill vtcm with src0 rows if (src0_start_row < src0_end_row) { @@ -802,7 +852,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { if (is0 >= (int)n_prefetch) { break; } - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } @@ -815,7 +865,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { // Process src0 rows for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); // Process src1 columns in pairs (2x2 tiling) @@ -823,15 +873,15 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { for (; ir1 + 1 < src1_nrows; ir1 += 2) { const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); - float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); - float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); + float * restrict dst_row0 = (float *) (dst->data + ((cur_m_start + ir1+0) * dst_row_size)); + float * restrict dst_row1 = (float *) (dst->data + ((cur_m_start + ir1+1) * dst_row_size)); mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); } // Handle remaining src1 rows (fallback to 2x1) for (; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + float * restrict dst_row = (float *) (dst->data + ((cur_m_start + ir1) * dst_row_size)); mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, src1_col); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); @@ -840,7 +890,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { const int pr0 = (ir0 + n_prefetch); const int is0 = (pr0 - src0_start_row) & prefetch_mask; if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } @@ -849,21 +899,21 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { if (src0_end_row != src0_end_row_x2) { uint32_t ir0 = src0_end_row_x2; const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 1); - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); #pragma unroll(2) for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + float * restrict dst_row = (float *) (dst->data + ((cur_m_start + ir1) * dst_row_size)); mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, src1_col); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); } if (src2) { - hvx_tensor_add_f32_grid(dst, src2, 0, src1_nrows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); + hvx_tensor_add_f32_grid(dst, src2, cur_m_start, cur_m_start + src1_nrows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); } } @@ -891,7 +941,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { float * tmp = (float *) vtcm_dst_ptr; - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; + const dma_addr_t src0_row = src0->data; const uint8_t * restrict src1_col = (const uint8_t *) src1_data; float * restrict dst_col = (float *) dst->data; @@ -906,12 +956,12 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { if (src0_start_row < src0_end_row) { if (src2) { float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; - const float * src2_ptr = (const float *) src2->data + src0_start_row; + const dma_addr_t src2_addr = src2->data + src0_start_row * sizeof(float); int slice_size = (int)src0_end_row - (int)src0_start_row; if (slice_size > 0) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), + dma_queue_push(dma_q, dma_make_data(vtcm_src2_ptr, src2_addr), slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); - dma_queue_pop_nowait(dma_queue); + dma_queue_pop_nowait(dma_q); } } for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { @@ -919,7 +969,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { if (is0 >= n_prefetch) { break; } - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } @@ -932,7 +982,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { // Process src0 rows for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); mmctx->vec_dot_2x1(ne00, &tmp[ir0 - src0_start_row], ss0, ss0 + src0_stride, src1_col); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); @@ -941,7 +991,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { const uint32_t pr0 = (ir0 + n_prefetch); const uint32_t is0 = (pr0 - src0_start_row) & prefetch_mask; if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } @@ -950,9 +1000,9 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { if (src0_end_row != src0_end_row_x2) { const uint32_t ir0 = src0_end_row_x2; const uint32_t is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 1); - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); mmctx->vec_dot_1x1(ne00, &tmp[ir0 - src0_start_row], ss0, src1_col); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); @@ -973,6 +1023,154 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { } } +static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { + htp_matmul_preamble; + + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + const uint32_t prefetch_mask = n_prefetch - 1; + + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; + const uint32_t cur_m_rows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); + const uint32_t cur_m_start = mmctx->cur_m_start; + + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); + const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const size_t dst_row_size = nb1; + const size_t src0_row_size = nb01; + const size_t src0_stride = mmctx->vtcm_src0_stride; + const size_t src1_stride = mmctx->vtcm_src1_stride; + + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * restrict src1_data = mmctx->vtcm_src1; + + hvx_mm_run_quant_task(mmctx, ith); + + if (src0_start_row >= src0_end_row || cur_m_rows == 0) { + return; + } + + const uint32_t total_batches = ne12 * ne13; + const uint32_t b_start = fastdiv(cur_m_start, &kparams->div_ne1); + uint32_t b_end = fastdiv(cur_m_start + cur_m_rows + ne11 - 1, &kparams->div_ne1); + b_end = MIN(b_end, total_batches); + + uint32_t b_grp_start = b_start; + while (b_grp_start < b_end) { + const uint32_t b3 = fastdiv(b_grp_start, &kparams->div_ne12); + const uint32_t b2 = b_grp_start - b3 * ne12; + const uint32_t i02 = fastdiv(b2, &kparams->div_r2); + const uint32_t i03 = fastdiv(b3, &kparams->div_r3); + + uint32_t b_grp_end = b_grp_start + 1; + while (b_grp_end < b_end) { + const uint32_t cur_b3 = fastdiv(b_grp_end, &kparams->div_ne12); + const uint32_t cur_b2 = b_grp_end - cur_b3 * ne12; + const uint32_t cur_i02 = fastdiv(cur_b2, &kparams->div_r2); + const uint32_t cur_i03 = fastdiv(cur_b3, &kparams->div_r3); + if (cur_i02 != i02 || cur_i03 != i03) { + break; + } + b_grp_end++; + } + + const dma_addr_t src0_row = src0->data + ((size_t) i02 * nb02 + (size_t) i03 * nb03); + + for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { + const int is0 = (ir0 - src0_start_row); + if (is0 >= (int)n_prefetch) { + break; + } + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + (size_t) ir0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 2); + } + + for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + for (uint32_t b = b_grp_start; b < b_grp_end; b++) { + const uint32_t b_m_start = b * ne11; + const uint32_t m_first = MAX(cur_m_start, b_m_start); + const uint32_t m_last = MIN(cur_m_start + cur_m_rows, b_m_start + ne11); + if (m_first >= m_last) continue; + + const uint32_t cur_b3 = fastdiv(b, &kparams->div_ne12); + const uint32_t cur_b2 = b - cur_b3 * ne12; + uint8_t * dst_batch_base = (uint8_t *) dst->data + (size_t) cur_b2 * nb2 + (size_t) cur_b3 * nb3; + + const uint32_t chunk_m_offset = m_first - cur_m_start; + const uint32_t dst_m_offset = m_first - b_m_start; + const uint32_t batch_nrows = m_last - m_first; + + uint32_t ir1 = 0; + for (; ir1 + 1 < batch_nrows; ir1 += 2) { + const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (chunk_m_offset + ir1 + 0) * src1_stride); + const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (chunk_m_offset + ir1 + 1) * src1_stride); + float * restrict dst_row0 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 0) * dst_row_size); + float * restrict dst_row1 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 1) * dst_row_size); + mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); + } + for (; ir1 < batch_nrows; ++ir1) { + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (chunk_m_offset + ir1) * src1_stride); + float * restrict dst_row = (float *) (dst_batch_base + (dst_m_offset + ir1) * dst_row_size); + mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, src1_col); + } + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + + const int pr0 = (ir0 + n_prefetch); + const int is0 = (pr0 - src0_start_row) & prefetch_mask; + if (pr0 < src0_end_row_x2) { + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + (size_t) pr0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 2); + } + } + + if (src0_end_row != src0_end_row_x2) { + uint32_t ir0 = src0_end_row_x2; + const int is0 = (ir0 - src0_start_row) & prefetch_mask; + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + (size_t) ir0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 1); + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + for (uint32_t b = b_grp_start; b < b_grp_end; b++) { + const uint32_t b_m_start = b * ne11; + const uint32_t m_first = MAX(cur_m_start, b_m_start); + const uint32_t m_last = MIN(cur_m_start + cur_m_rows, b_m_start + ne11); + if (m_first >= m_last) continue; + + const uint32_t cur_b3 = fastdiv(b, &kparams->div_ne12); + const uint32_t cur_b2 = b - cur_b3 * ne12; + uint8_t * dst_batch_base = (uint8_t *) dst->data + (size_t) cur_b2 * nb2 + (size_t) cur_b3 * nb3; + + const uint32_t chunk_m_offset = m_first - cur_m_start; + const uint32_t dst_m_offset = m_first - b_m_start; + const uint32_t batch_nrows = m_last - m_first; + + for (uint32_t ir1 = 0; ir1 < batch_nrows; ++ir1) { + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (chunk_m_offset + ir1) * src1_stride); + float * restrict dst_row = (float *) (dst_batch_base + (dst_m_offset + ir1) * dst_row_size); + mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, src1_col); + } + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + } + + b_grp_start = b_grp_end; + } + + if (src2) { + hvx_tensor_add_f32_grid(dst, src2, cur_m_start, cur_m_start + cur_m_rows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); + } +} + #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)] static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { @@ -1018,7 +1216,7 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { continue; } - const uint8_t * src0_row = (const uint8_t *) src0->data + cur_a * nb02; + const dma_addr_t src0_row = src0->data + cur_a * nb02; const uint32_t tile_size = htp_mm_get_weight_tile_size(src0->type); const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src0->type); @@ -1032,12 +1230,12 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { uint32_t push_ct = ct_start; for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); } for (uint32_t ct = ct_start; ct < ct_end; ct++) { - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; int valid_rows = (int)ne01 - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); @@ -1057,7 +1255,7 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); if (push_ct < ct_end) { - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); push_ct++; } @@ -1105,7 +1303,7 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { } assert(eid < (int32_t) n_ids); - const uint8_t * restrict src0_row = (const uint8_t *) src0->data + eid * nb02; + const dma_addr_t src0_row = src0->data + eid * nb02; const uint8_t * restrict src1_col = (const uint8_t *) src1_data; float * restrict dst_row = (float *) (dst->data + ie1 * nb1); @@ -1121,12 +1319,12 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { uint32_t push_ct = ct_start; for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); } for (uint32_t ct = ct_start; ct < ct_end; ct++) { - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; int valid_rows = (int)ne01 - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); @@ -1136,7 +1334,7 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); if (push_ct < ct_end) { - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); push_ct++; } @@ -1147,7 +1345,6 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = (struct htp_mm_context *) data; struct htp_ops_context * octx = mmctx->octx; - dma_queue * dma_queue = octx->ctx->dma[ith]; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_weights = kparams->n_weights; const struct htp_tensor * restrict src0 = octx->src[0]; @@ -1176,6 +1373,7 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { const struct htp_tensor * restrict src_w = octx->src[p]; const struct htp_tensor * restrict dst = octx->dsts[p]; if (!src_w || !dst) continue; + dma_queue * dma_q = octx->ctx->dma[ith]; const uint32_t ne01 = src_w->ne[1]; uint32_t start_row = 0; @@ -1195,7 +1393,7 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, end_row); if (src0_start_row >= src0_end_row) continue; - const uint8_t * restrict src0_row = (const uint8_t *) src_w->data + eid * src_w->nb[2]; + const dma_addr_t src0_row = src_w->data + eid * src_w->nb[2]; const uint8_t * restrict src1_col = (const uint8_t *) src1_data; float * restrict dst_row = (float *) (dst->data + ie1 * dst->nb[1]); @@ -1211,12 +1409,12 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { uint32_t push_ct = ct_start; for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); } for (uint32_t ct = ct_start; ct < ct_end; ct++) { - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); @@ -1226,7 +1424,7 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); if (push_ct < ct_end) { - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); push_ct++; } @@ -1238,7 +1436,6 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = (struct htp_mm_context *) data; struct htp_ops_context * octx = mmctx->octx; - dma_queue * dma_queue = octx->ctx->dma[ith]; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_weights = kparams->n_weights; const struct htp_tensor * restrict src0 = octx->src[0]; @@ -1270,6 +1467,7 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { const struct htp_tensor * restrict src_w = octx->src[p]; const struct htp_tensor * restrict dst = octx->dsts[p]; if (!src_w || !dst) continue; + dma_queue * dma_q = octx->ctx->dma[ith]; const uint32_t ne01 = src_w->ne[1]; uint32_t start_row = 0; @@ -1289,7 +1487,7 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, end_row); if (src0_start_row >= src0_end_row) continue; - const uint8_t * src0_row = (const uint8_t *) src_w->data + cur_a * src_w->nb[2]; + const dma_addr_t src0_row = src_w->data + cur_a * src_w->nb[2]; const uint32_t tile_size = htp_mm_get_weight_tile_size(src_w->type); const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src_w->type); @@ -1303,12 +1501,12 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { uint32_t push_ct = ct_start; for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); } for (uint32_t ct = ct_start; ct < ct_end; ct++) { - const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); @@ -1328,7 +1526,7 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); if (push_ct < ct_end) { - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); push_ct++; } @@ -1384,6 +1582,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { const uint32_t src0_nrows = ne01; const uint32_t src1_nrows = ne11 * ne12 * ne13; + mmctx->src1_nrows = src1_nrows; uint32_t src0_row_start = 0; uint32_t src0_row_end = src0_nrows; @@ -1426,7 +1625,23 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { worker_callback_t quant_task_func; worker_callback_t matmul_job_func; uint32_t n_quant_tasks = 1; - if (src1_nrows > 1) { + const bool is_batched = (ne12 > 1 || ne13 > 1 || ne02 > 1 || ne03 > 1); + if (is_batched) { + if (is_repacked) { + switch (src0->type) { + case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_4d_repacked_q4_0; break; + case HTP_TYPE_Q4_1: + case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_4d_repacked_q4_1; break; + case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_4d_repacked_q8_0; break; + case HTP_TYPE_Q6_K: matmul_job_func = hvx_mm_4d_repacked_q6_k; break; + case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_4d_repacked_iq4nl; break; + case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_4d_repacked_mxfp4; break; + default: return HTP_STATUS_NO_SUPPORT; + } + } else { + matmul_job_func = hvx_mm_4d; + } + } else if (src1_nrows > 1) { if (is_repacked) { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_2d_repacked_q4_0; break; @@ -1462,7 +1677,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { switch (kparams->kernel_type) { case HTP_MM_KERNEL_HVX_F16_F16_VTCM: - quant_task_func = (src1->type == HTP_TYPE_F32) ? quantize_f32_f16_flat : quantize_f16_f16_flat; + quant_task_func = (src1->type == HTP_TYPE_F32) ? quantize_f32_f16 : quantize_f16_f16; mmctx->type = "f16-f16"; mmctx->vec_dot_1x1 = vec_dot_f16_f16_aa_1x1; mmctx->vec_dot_2x1 = vec_dot_f16_f16_aa_2x1; @@ -1470,34 +1685,8 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { src1_row_size = hex_round_up(ne10 * 2, 128); break; - case HTP_MM_KERNEL_HVX_F16_F32_DDR: - mmctx->type = "f16-f32"; - mmctx->vec_dot_1x1 = vec_dot_f16_f32_uu_1x1; - matmul_job_func = hvx_mm_4d; - mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1; - mmctx->mm_div_ne1 = kparams->div_ne1; - mmctx->mm_div_r2 = kparams->div_r2; - mmctx->mm_div_r3 = kparams->div_r3; - need_quant = false; - quant_task_func = NULL; - src1_row_size = nb11; - break; - - case HTP_MM_KERNEL_HVX_F16_F16_DDR: - mmctx->type = "f16-f16"; - mmctx->vec_dot_1x1 = vec_dot_f16_f16_uu_1x1; - matmul_job_func = hvx_mm_4d; - mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1; - mmctx->mm_div_ne1 = kparams->div_ne1; - mmctx->mm_div_r2 = kparams->div_r2; - mmctx->mm_div_r3 = kparams->div_r3; - src1_row_size = nb11; - need_quant = false; - quant_task_func = NULL; - break; - case HTP_MM_KERNEL_HVX_F32_F32_VTCM: - quant_task_func = quantize_f32_f32_flat; + quant_task_func = quantize_f32_f32; mmctx->type = "f32-f32"; mmctx->vec_dot_1x1 = vec_dot_f32_f32_aa_1x1; mmctx->vec_dot_2x1 = vec_dot_f32_f32_aa_2x1; @@ -1505,50 +1694,6 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { src1_row_size = hex_round_up(ne10 * 4, 128); break; - case HTP_MM_KERNEL_HVX_F32_F32_DDR: - quant_task_func = NULL; - mmctx->type = "f32-f32"; - mmctx->vec_dot_1x1 = vec_dot_f32_f32_uu_1x1; - mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1; - mmctx->mm_div_ne1 = kparams->div_ne1; - mmctx->mm_div_r2 = kparams->div_r2; - mmctx->mm_div_r3 = kparams->div_r3; - src1_row_size = nb11; - need_quant = false; - matmul_job_func = hvx_mm_4d; - break; - - case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat; - src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - - if (src1_nrows > 1) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_2d_repacked_q8_0_flat; break; - case HTP_TYPE_Q6_K: matmul_job_func = hvx_mm_2d_repacked_q6_k_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_2d_repacked_mxfp4_flat; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mv_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mv_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mv_2d_repacked_q8_0_flat; break; - case HTP_TYPE_Q6_K: matmul_job_func = hvx_mv_2d_repacked_q6_k_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mv_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mv_2d_repacked_mxfp4_flat; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } - break; - } - case HTP_MM_KERNEL_HVX_QUANT_BLOCK: case HTP_MM_KERNEL_HVX_QUANT_ROW: default: @@ -1560,7 +1705,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { const uint32_t nb = (ne10 + qk - 1) / qk; const uint32_t total_nb = src1_nrows * nb; - if (src1_nrows < octx->n_threads) { + if (src1_nrows < octx->n_threads && !is_batched) { n_quant_tasks = MIN(total_nb, octx->n_threads); quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { @@ -1579,8 +1724,12 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { break; } + const uint32_t m_chunk = (kparams->m_chunk > 0 && (uint32_t) kparams->m_chunk < src1_nrows) + ? (uint32_t) kparams->m_chunk : src1_nrows; + const uint32_t m_layout_rows = m_chunk; + struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, m_layout_rows, octx->n_threads, dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false); if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || @@ -1623,20 +1772,47 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { mmctx->vtcm_src0_stride = src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; - if (need_quant) { - mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; - mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); + if (kparams->m_chunk > 0 && (uint32_t) kparams->m_chunk < src1_nrows) { + atomic_init(&mmctx->quant_barrier, 0); + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); + + for (uint32_t m_start = 0; m_start < src1_nrows; m_start += m_chunk) { + const uint32_t cur_m_rows = MIN(src1_nrows - m_start, m_chunk); + mmctx->cur_m_start = m_start; + mmctx->cur_m_rows = cur_m_rows; + + if (need_quant) { + const uint32_t quant_tasks = MIN(cur_m_rows, octx->n_threads); + mmctx->n_quant_rows_per_thread = (cur_m_rows + quant_tasks - 1) / quant_tasks; + mmctx->n_quant_tasks = quant_tasks; + atomic_store(&mmctx->quant_barrier, quant_tasks); + mmctx->quant_task_func = quant_task_func; + } else { + mmctx->quant_task_func = NULL; + mmctx->n_quant_tasks = 0; + } + + worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads); + } } else { - mmctx->quant_task_func = NULL; - mmctx->n_quant_tasks = 0; + mmctx->cur_m_start = 0; + mmctx->cur_m_rows = src1_nrows; + + if (need_quant) { + mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; + mmctx->quant_task_func = quant_task_func; + mmctx->n_quant_tasks = n_quant_tasks; + atomic_init(&mmctx->quant_barrier, n_quant_tasks); + } else { + mmctx->quant_task_func = NULL; + mmctx->n_quant_tasks = 0; + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); + + worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads); } - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - - worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads); - return HTP_STATUS_OK; } @@ -1653,7 +1829,6 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; uint8_t * restrict src1_data = mmctx->vtcm_src1; - dma_queue * dma_queue = octx->ctx->dma[ith]; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t prefetch_mask = n_prefetch - 1; @@ -1666,6 +1841,7 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { const struct htp_tensor * restrict src_w = octx->src[widx]; const struct htp_tensor * restrict dst = octx->dsts[widx]; if (!src_w || !dst) continue; + dma_queue * dma_q = octx->ctx->dma[ith]; const uint32_t ne00 = src_w->ne[0]; const uint32_t ne01 = src_w->ne[1]; @@ -1691,17 +1867,17 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { const size_t src0_row_size = src_w->nb[1]; const size_t src0_stride = hex_round_up(src0_row_size, 128); - const uint8_t * restrict src0_row = (const uint8_t *) src_w->data; + const dma_addr_t src0_row = src_w->data; for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const int is0 = (ir0 - src0_start_row); if (is0 >= (int)n_prefetch) break; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); uint32_t ir1 = 0; for (; ir1 + 1 < src1_nrows; ir1 += 2) { @@ -1721,7 +1897,7 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { const int pr0 = (ir0 + n_prefetch); const int is0 = (pr0 - src0_start_row) & prefetch_mask; if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } @@ -1729,9 +1905,9 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { if (src0_end_row != src0_end_row_x2) { uint32_t ir0 = src0_end_row_x2; const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 1); - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); @@ -1868,7 +2044,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); - dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src + c_first), + dma_queue_push(dma_q, dma_make_data(thread_f32_act, src + c_first), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } // Push step 1 @@ -1878,7 +2054,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride + c_first; float *next_buf = thread_f32_act + 1 * R * c_len; - dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), + dma_queue_push(dma_q, dma_make_data(next_buf, next_src), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } } @@ -1909,7 +2085,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride + c_first; - dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src), + dma_queue_push(dma_q, dma_make_data(curr_buf, next_src), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } } @@ -2013,7 +2189,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); - dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src), + dma_queue_push(dma_q, dma_make_data(thread_f32_act, src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } // Push step 1 (if valid) @@ -2023,7 +2199,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride; float *next_buf = thread_f32_act + 1 * R * k_block; - dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), + dma_queue_push(dma_q, dma_make_data(next_buf, next_src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } @@ -2052,7 +2228,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride; - dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src), + dma_queue_push(dma_q, dma_make_data(curr_buf, next_src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } @@ -2459,10 +2635,12 @@ static inline void hmx_matmul_job_init(hmx_matmul_job_t * job, } static int hmx_mm_2d_f32(struct htp_context *ctx, + dma_queue *weight_dma, float *restrict dst, - const float *restrict src2, + dma_addr_t src2_addr, + size_t src2_bytes, const float *activation, - const uint8_t *weight, + dma_addr_t weight, int m, int k, int n, int act_stride, int weight_stride, @@ -2525,7 +2703,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size, src2_bytes); vtcm_used = L.total_bytes; if (vtcm_used > vtcm_budget) { @@ -2550,6 +2728,13 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 + const bool has_src2 = (src2_bytes > 0 && src2_addr != 0); + float *vtcm_src2 = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_src2, has_src2); + if (has_src2) { + dma_queue_push(weight_dma, dma_make_data(vtcm_src2, src2_addr), hex_align_up(src2_bytes, 128), 0, src2_bytes, 1); + dma_queue_pop(weight_dma); + } + FARF(HIGH, "hmx-mm-2d: m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu", m, k, n, weight_type, m_chunk_n_rows, n_chunk_n_cols, vtcm_used, vtcm_budget); @@ -2586,13 +2771,13 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // Prologue: push A0 and optionally A1 (if n_chunk_cnt > 1) const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), + dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0); if (1 < n_chunk_cnt) { const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols); const uint32_t height_A1 = is_quant ? (n_cols_A1 / 32) * n_k_tiles : n_cols_A1; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), + dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1); } @@ -2605,7 +2790,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t n_cols_p2 = hex_smin(n - nc_p2, n_chunk_n_cols); // 1. pop A_i - void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; // 2. dequantize A_i dequantize_tiled_weight_chunk_to_fp16_tiles( @@ -2616,7 +2801,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // 3. push A_{i+2} (if i+2 < n_chunk_cnt) if (i + 2 < n_chunk_cnt) { const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2; - dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_p2 * weight_stride), + dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_p2 * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2); } @@ -2633,7 +2818,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t nc_prev = (i - 1) * n_chunk_n_cols; const size_t n_cols_prev = hex_smin(n - nc_prev, n_chunk_n_cols); float *output_chunk = dst + (mr * dst_stride + nc_prev); - const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_prev) : NULL; + const float *src2_chunk = has_src2 ? (vtcm_src2 + mr * src2_stride + nc_prev) : NULL; int chunk_dst_cols = dst_cols - (int)nc_prev; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(i - 1) % 2], n_rows, n_cols_prev, dst_stride, src2_stride, chunk_dst_cols, n_threads); @@ -2646,7 +2831,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t nc_last = (n_chunk_cnt - 1) * n_chunk_n_cols; const size_t n_cols_last = hex_smin(n - nc_last, n_chunk_n_cols); float *output_chunk = dst + (mr * dst_stride + nc_last); - const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_last) : NULL; + const float *src2_chunk = has_src2 ? (vtcm_src2 + mr * src2_stride + nc_last) : NULL; int chunk_dst_cols = dst_cols - (int)nc_last; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(n_chunk_cnt - 1) % 2], n_rows, n_cols_last, dst_stride, src2_stride, chunk_dst_cols, n_threads); @@ -2678,7 +2863,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, if (n > 0) { const size_t n_cols = hex_smin(n, n_chunk_n_cols); const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); + dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { @@ -2687,7 +2872,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); // A: Wait for weight DMA - void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( @@ -2700,7 +2885,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, if (nc_next < n) { const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; - dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } // C: HMX Compute (Queue-based) @@ -2710,7 +2895,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // D: Output Store float *output_chunk = dst + (mr * dst_stride + nc); - const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc) : NULL; + const float *src2_chunk = has_src2 ? (vtcm_src2 + mr * src2_stride + nc) : NULL; int chunk_dst_cols = dst_cols - (int)nc; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output, n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads); @@ -2785,7 +2970,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size, 0); if (L.total_bytes > vtcm_budget) { FARF(ERROR, "hmx-mm-nx-2d: VTCM overflow: used %zu budget %zu, m %d k %d mc %d nc %d", @@ -2859,7 +3044,8 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const struct htp_tensor * restrict dst = octx->dsts[p]; if (!src_w || !dst) continue; - const uint8_t * weight = (const uint8_t *) src_w->data; + const dma_addr_t weight = src_w->data; + dma_queue * weight_dma = octx->ctx->dma[0]; float * dst_ptr = (float *) dst->data; const size_t n = src_w->ne[1]; if (n == 0) continue; @@ -2872,13 +3058,13 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), + dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0); if (1 < n_chunk_cnt) { const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols); const uint32_t height_A1 = is_quant ? (n_cols_A1 / 32) * n_k_tiles : n_cols_A1; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), + dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1); } @@ -2889,7 +3075,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); const size_t n_cols_p2 = hex_smin(n - nc_p2, n_chunk_n_cols); - void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; dequantize_tiled_weight_chunk_to_fp16_tiles( ctx, vtcm_weight_bufs[i % 2], curr_raw, @@ -2898,7 +3084,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k if (i + 2 < n_chunk_cnt) { const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2; - dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_p2 * weight_stride), + dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_p2 * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2); } @@ -2956,7 +3142,8 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const struct htp_tensor * restrict dst = octx->dsts[p]; if (!src_w || !dst) continue; - const uint8_t * weight = (const uint8_t *) src_w->data; + const dma_addr_t weight = src_w->data; + dma_queue * weight_dma = octx->ctx->dma[0]; float * dst_ptr = (float *) dst->data; const size_t n = src_w->ne[1]; if (n == 0) continue; @@ -2969,7 +3156,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k if (n > 0) { const size_t n_cols = hex_smin(n, n_chunk_n_cols); const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); + dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { @@ -2977,7 +3164,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); - void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; dequantize_tiled_weight_chunk_to_fp16_tiles( ctx, vtcm_scratch0, curr_raw, @@ -2988,7 +3175,7 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k if (nc_next < n) { const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; - dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); @@ -3008,13 +3195,11 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k return HTP_STATUS_OK; } -static inline const __fp16 *hmx_mm_weight_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, - int dst_b2, int dst_b3) { +static inline dma_addr_t hmx_mm_weight_batch_data(const hmx_mm_f16_f32_batched_params_t *params, + int dst_b2, int dst_b3) { const size_t b2_idx = (params->r2 <= 1) ? (size_t) dst_b2 : (size_t) fastdiv((uint32_t) dst_b2, ¶ms->div_r2); const size_t b3_idx = (params->r3 <= 1) ? (size_t) dst_b3 : (size_t) fastdiv((uint32_t) dst_b3, ¶ms->div_r3); - return (const __fp16 *) ((const uint8_t *) params->weight + - b2_idx * params->src0_nb2 + - b3_idx * params->src0_nb3); + return params->weight + b2_idx * params->src0_nb2 + b3_idx * params->src0_nb3; } static inline const float *hmx_mm_activation_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, @@ -3031,13 +3216,6 @@ static inline float *hmx_mm_dst_batch_ptr(const hmx_mm_f16_f32_batched_params_t (size_t) dst_b3 * params->dst_nb3); } -static inline const float *hmx_mm_src2_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, - int src2_b2, int src2_b3) { - return params->src2 ? (const float *) ((const uint8_t *) params->src2 + - (size_t) src2_b2 * params->src2_nb2 + - (size_t) src2_b3 * params->src2_nb3) : NULL; -} - static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, const hmx_mm_f16_f32_batched_params_t *params, int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size, @@ -3045,15 +3223,18 @@ static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, int ret = 0; for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) { for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) { - ret = hmx_mm_2d_f32(ctx, hmx_mm_dst_batch_ptr(params, b2, b3), - hmx_mm_src2_batch_ptr(params, b2, b3), - hmx_mm_activation_batch_ptr(params, b2, b3), - (const uint8_t *)hmx_mm_weight_batch_ptr(params, b2, b3), - params->m, params->k, params->n, - params->act_stride, params->weight_stride * (int)sizeof(__fp16), - HTP_TYPE_F16, params->k, params->dst_stride, params->src2_stride, params->n, - m_chunk, n_chunk, pipeline, n_threads, act_threads, - act_threads_div, k_div, 0, 0, vtcm_size); + dma_addr_t cur_src2_addr = params->src2_addr ? (params->src2_addr + + (dma_addr_t) b2 * params->src2_nb2 + + (dma_addr_t) b3 * params->src2_nb3) : 0; + ret = hmx_mm_2d_f32(ctx, params->weight_dma, hmx_mm_dst_batch_ptr(params, b2, b3), + cur_src2_addr, params->src2_bytes, + hmx_mm_activation_batch_ptr(params, b2, b3), + hmx_mm_weight_batch_data(params, b2, b3), + params->m, params->k, params->n, + params->act_stride, params->weight_stride * (int)sizeof(__fp16), + HTP_TYPE_F16, params->k, params->dst_stride, params->src2_stride, params->n, + m_chunk, n_chunk, pipeline, n_threads, act_threads, + act_threads_div, k_div, 0, 0, vtcm_size); } } return ret; @@ -3095,7 +3276,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ size_t vtcm_used = vtcm_size; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0, params->src2_bytes); if (L.total_bytes > vtcm_budget) { FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__); @@ -3112,6 +3293,13 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales); float *vtcm_f32_act = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_act_f32, use_dma_activation); + const bool has_src2 = (params->src2_bytes > 0 && params->src2_addr != 0); + float *vtcm_src2 = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_src2, has_src2); + if (has_src2) { + dma_queue_push(params->weight_dma, dma_make_data(vtcm_src2, params->src2_addr), hex_align_up(params->src2_bytes, 128), 0, params->src2_bytes, 1); + dma_queue_pop(params->weight_dma); + } + hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 FARF(HIGH, "%s: grouped path m=%d k=%d n=%d group=%d streams=%d mc=%zu nc=%zu vtcm=%zu/%zu", @@ -3128,7 +3316,8 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ for (int b3 = 0; b3 < params->ne13; ++b3) { for (int b2_base = 0; b2_base < params->ne12; b2_base += group_size) { - const __fp16 *weight_group = hmx_mm_weight_batch_ptr(params, b2_base, b3); + const dma_addr_t weight_group = hmx_mm_weight_batch_data(params, b2_base, b3); + dma_queue * weight_dma = params->weight_dma; for (size_t mr = 0; mr < (size_t) params->m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin((size_t) params->m - mr, m_chunk_n_rows); @@ -3162,12 +3351,12 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ // Prologue: Push A0 and A1 (if exists) { const size_t n_cols_first = hex_smin((size_t) params->n, n_chunk_n_cols); - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch0, weight_group), + dma_queue_push(weight_dma, dma_make_data(vtcm_scratch0, weight_group), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_first); } if (n_chunk_n_cols < (size_t) params->n) { const size_t n_cols_second = hex_smin((size_t) params->n - n_chunk_n_cols, n_chunk_n_cols); - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch1, weight_group + params->weight_stride), + dma_queue_push(weight_dma, dma_make_data(vtcm_scratch1, weight_group + params->weight_stride * sizeof(__fp16)), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_second); } @@ -3176,16 +3365,16 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ const size_t n_col_tiles = hmx_ceil_div((int) n_cols, HTP_MM_HMX_TILE_N_COLS); { - void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) curr_raw, n_cols, params->k, params->k, 0, n_cols); const size_t nc_next = nc + n_chunk_n_cols * 2; if (nc_next < (size_t) params->n) { const size_t n_cols_next = hex_smin((size_t) params->n - nc_next, n_chunk_n_cols); - const __fp16 *next_weight_chunk = weight_group + nc_next * params->weight_stride; + const dma_addr_t next_weight_chunk = weight_group + nc_next * params->weight_stride * sizeof(__fp16); - dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, next_weight_chunk), + dma_queue_push(weight_dma, dma_make_data(curr_raw, next_weight_chunk), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_next); } } @@ -3201,7 +3390,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ { float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc; - const float *src2_chunk = params->src2 ? (hmx_mm_src2_batch_ptr(params, b2_base + g, b3) + mr * params->src2_stride + nc) : NULL; + const float *src2_chunk = has_src2 ? (vtcm_src2 + mr * params->src2_stride + nc) : NULL; int chunk_dst_cols = params->n - (int)nc; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols, @@ -3314,9 +3503,10 @@ static void transfer_output_chunk_scattered_threaded( } static int hmx_mm_id_2d_f32(struct htp_context *ctx, + dma_queue *weight_dma, float *restrict dst, const float *activation, - const uint8_t *weight, + dma_addr_t weight, int m, int k, int n, int k_valid, int ne11, @@ -3429,7 +3619,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, if (n > 0) { const size_t n_cols = hex_smin((size_t) n, n_chunk_n_cols); const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols; - dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight), + dma_queue_push(weight_dma, dma_make_data(vtcm_weight, weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } @@ -3438,7 +3628,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); // A: Wait for weight DMA - void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( @@ -3452,7 +3642,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, if (nc_next < (size_t) n) { const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; - dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), + dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } @@ -3495,13 +3685,15 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k return HTP_STATUS_OK; } - const float * src2_ptr = NULL; + dma_addr_t src2_addr = 0; + size_t src2_bytes = 0; uint32_t src2_stride = 0; size_t src2_nb2 = 0; size_t src2_nb3 = 0; if (src2) { src2_stride = (src2->ne[1] == 1) ? 0 : (uint32_t) (src2->nb[1] / sizeof(float)); - src2_ptr = (const float *) src2->data + m_start * src2_stride; + src2_addr = src2->data + (dma_addr_t) m_start * src2_stride * sizeof(float); + src2_bytes = (size_t) kparams->vtcm_src2_size; src2_nb2 = (src2->ne[2] == 1) ? 0 : src2->nb[2]; src2_nb3 = (src2->ne[3] == 1) ? 0 : src2->nb[3]; } @@ -3515,9 +3707,11 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k if (kparams->kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) { hmx_mm_f16_f32_batched_params_t batch_params = { .dst = dst_ptr, - .src2 = src2_ptr, + .src2_addr = src2_addr, + .src2_bytes = src2_bytes, .activation = act_ptr, - .weight = (const __fp16 *) src0->data, + .weight = src0->data, + .weight_dma = octx->ctx->dma[0], .m = m_rows, .k = k, .n = n, @@ -3551,7 +3745,7 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k kparams->vtcm_size); } else { ret = hmx_mm_2d_f32( - octx->ctx, dst_ptr, src2_ptr, act_ptr, (const uint8_t *) src0->data, + octx->ctx, octx->ctx->dma[0], dst_ptr, src2_addr, src2_bytes, act_ptr, src0->data, m_rows, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0], dst_stride, src2_stride, (int)dst->ne[0], kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads, @@ -3609,8 +3803,8 @@ static int hmx_mm_op_matmul_id( } if (m_start >= m_end) continue; - int ret = hmx_mm_id_2d_f32(octx->ctx, (float*) dst->data, (float*) src1->data, - (const uint8_t *) src0->data + cur_a * nb02, + int ret = hmx_mm_id_2d_f32(octx->ctx, octx->ctx->dma[0], (float*) dst->data, (float*) src1->data, + src0->data + cur_a * nb02, cne1, ne00, ne01, ne10, ne11, @@ -3706,6 +3900,9 @@ static int hvx_mm_matmul_id( mmctx->vtcm_src2_size_per_thread = 0; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); + mmctx->cur_m_start = 0; + mmctx->cur_m_rows = src1_nrows; + mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; mmctx->n_quant_tasks = n_quant_tasks; @@ -3753,8 +3950,8 @@ static int hmx_mm_op_matmul_id_nx( const struct htp_tensor * restrict dst = octx->dsts[p]; if (!src_w || !dst) continue; - int ret = hmx_mm_id_2d_f32(octx->ctx, (float*) dst->data, (float*) act->data, - (const uint8_t *) src_w->data + cur_a * src_w->nb[2], + int ret = hmx_mm_id_2d_f32(octx->ctx, octx->ctx->dma[0], (float*) dst->data, (float*) act->data, + src_w->data + cur_a * src_w->nb[2], cne1, src_w->ne[0], src_w->ne[1], act->ne[0], act->ne[1], @@ -3844,6 +4041,9 @@ static int hvx_mm_matmul_id_nx( mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); + mmctx->cur_m_start = 0; + mmctx->cur_m_rows = src1_nrows; + mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; mmctx->n_quant_tasks = n_quant_tasks; @@ -3947,6 +4147,9 @@ int op_matmul_id(struct htp_ops_context * octx) { mmctx->act = src1; const struct htp_tensor * restrict ids = octx->src[2]; + if (htp_tensor_is_extended(ids) || htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } const size_t src0_row_size = nb01; const size_t dst_row_size = nb1; @@ -3997,9 +4200,11 @@ int op_matmul_id(struct htp_ops_context * octx) { mmctx->matrix_row_counts = matrix_row_counts; mmctx->matrix_rows = matrix_rows; mmctx->mapping_stride = mapping_stride; - mmctx->mm_div_ne11 = kparams->div_ne11; + mmctx->mm_div_ne11 = kparams->div_ne1; mmctx->src0_row_size_padded = src0_row_size_padded; mmctx->src1_nrows = src1_nrows; + mmctx->cur_m_start = 0; + mmctx->cur_m_rows = src1_nrows; htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); @@ -4062,6 +4267,14 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { const struct htp_tensor * restrict src0 = octx->src[0]; const struct htp_tensor * restrict act = octx->src[n_weights]; const struct htp_tensor * restrict ids = octx->src[n_weights + 1]; + if (htp_tensor_is_extended(ids) || htp_tensor_is_extended(act)) { + return HTP_STATUS_NO_SUPPORT; + } + for (uint32_t p = 0; p < n_weights; p++) { + if (octx->dsts[p] && htp_tensor_is_extended(octx->dsts[p])) { + return HTP_STATUS_NO_SUPPORT; + } + } mmctx->act = act; @@ -4110,9 +4323,11 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { mmctx->matrix_row_counts = matrix_row_counts; mmctx->matrix_rows = matrix_rows; mmctx->mapping_stride = mapping_stride; - mmctx->mm_div_ne11 = kparams->div_ne11; + mmctx->mm_div_ne11 = kparams->div_ne1; mmctx->src0_row_size_padded = src0_row_size_padded; mmctx->src1_nrows = src1_nrows; + mmctx->cur_m_start = 0; + mmctx->cur_m_rows = src1_nrows; htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); @@ -4163,6 +4378,9 @@ int op_matmul_nx(struct htp_ops_context * octx) { mmctx->act = act; const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; + mmctx->src1_nrows = src1_nrows; + mmctx->cur_m_start = 0; + mmctx->cur_m_rows = src1_nrows; const size_t src0_row_size = src0->nb[1]; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); @@ -4177,10 +4395,7 @@ int op_matmul_nx(struct htp_ops_context * octx) { worker_callback_t quant_task_func; uint32_t n_quant_tasks = 1; - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat; - } else if (src1_nrows < octx->n_threads) { + if (src1_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { @@ -4196,12 +4411,9 @@ int op_matmul_nx(struct htp_ops_context * octx) { quant_task_func = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; } - size_t src1_row_size; - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(act->ne[0]) : htp_mm_q8_0_flat_row_size(act->ne[0]); - } else { - src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); - } + const size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q4_K) + ? htp_mm_q8_1_tiled_row_size(act->ne[0]) + : htp_mm_q8_0_tiled_row_size(act->ne[0]); struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads, @@ -4242,26 +4454,14 @@ int op_matmul_nx(struct htp_ops_context * octx) { const uint32_t n_matmul_jobs = octx->n_threads; worker_callback_t matmul_job_func; if (is_repacked) { - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4_flat; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; - } + switch (src0->type) { + case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; + case HTP_TYPE_Q4_1: + case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; + case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; + case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; + case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; + default: return HTP_STATUS_NO_SUPPORT; } } else { matmul_job_func = hvx_mm_nx_2d; diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index 1df8c2933..fe9dbb61c 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -62,17 +62,11 @@ enum htp_mm_kernel_type { // HVX floating-point paths HTP_MM_KERNEL_HVX_F16_F16_VTCM, - HTP_MM_KERNEL_HVX_F16_F16_DDR, - HTP_MM_KERNEL_HVX_F16_F32_DDR, - HTP_MM_KERNEL_HVX_F32_F32_VTCM, - HTP_MM_KERNEL_HVX_F32_F32_DDR, - HTP_MM_KERNEL_HVX_F32_F16_DDR, // HVX quantized paths HTP_MM_KERNEL_HVX_QUANT_ROW, // standard row-wise parallel quantization HTP_MM_KERNEL_HVX_QUANT_BLOCK, // parallel block-wise quantization - HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, // row-wise fallback flat quantization }; // Op-specific struct for precomputed matmul params @@ -101,7 +95,7 @@ struct htp_mm_kernel_params { struct fastdiv_values div_ne1; struct fastdiv_values div_r2; struct fastdiv_values div_r3; - struct fastdiv_values div_ne11; + struct fastdiv_values div_ne12; struct fastdiv_values div_n_act_threads; struct fastdiv_values div_ne00_padded; }; @@ -246,20 +240,6 @@ static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) { return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1; } -static inline size_t htp_mm_q8_0_flat_row_size(uint32_t ne) { - const uint32_t quants_size = hex_align_up(ne, 128); - const uint32_t num_scales = (ne + 31) / 32; - const uint32_t scales_size = hex_align_up(num_scales * 2, 128); - return quants_size + scales_size; -} - -static inline size_t htp_mm_q8_1_flat_row_size(uint32_t ne) { - const uint32_t quants_size = hex_align_up(ne, 128); - const uint32_t num_scales = (ne + 31) / 32; - const uint32_t scales_size = hex_align_up(num_scales * 4, 128); - return quants_size + scales_size; -} - static inline size_t htp_mm_get_tiled_row_stride(int weight_type, uint32_t k) { uint32_t nb = (k + QK_Q4_0_TILED - 1) / QK_Q4_0_TILED; switch (weight_type) { @@ -331,6 +311,7 @@ struct htp_mm_hmx_vtcm_layout { size_t off_dst[2]; // [1] is only used when pipelined size_t off_scratch[2]; // dequantization scratch pads size_t off_scales; // HMX scales (256 bytes) + size_t off_src2; // src2 bias in VTCM // Cached sizes of regions for HMX kernel use size_t weight_area_bytes; @@ -339,6 +320,7 @@ struct htp_mm_hmx_vtcm_layout { size_t output_area_bytes; size_t scratch_bytes[2]; size_t act_head_stride; + size_t src2_bytes; size_t total_bytes; }; @@ -372,7 +354,8 @@ static inline void htp_mm_hmx_vtcm_layout_build( bool use_dma_activation, bool pipeline, uint32_t act_threads, - uint32_t aligned_tile_size + uint32_t aligned_tile_size, + size_t src2_size ) { size_t off = 0; @@ -390,6 +373,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( size_t off_group_a = 0; VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size); VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0); // Group B: Compute-only buffers (starts at off_group_a) size_t off_group_b = off_group_a; @@ -418,6 +402,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( L->scratch_bytes[0] = scratch_area_size; L->scratch_bytes[1] = scratch_area_size; L->act_head_stride = act_head_stride; + L->src2_bytes = src2_size; off = off_group_a + hex_smax(group_b_size, group_c_size); } else { @@ -441,6 +426,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( size_t off_group_a = 0; VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size); + VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0); // Group B: Compute-only buffers (starts at off_group_a) size_t off_group_b = off_group_a; @@ -468,6 +454,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( L->scratch_bytes[0] = scratch0_size; L->scratch_bytes[1] = scratch1_size; L->act_head_stride = 0; + L->src2_bytes = src2_size; off = off_group_a + hex_smax(group_b_size, group_c_size); } @@ -490,6 +477,7 @@ static inline void htp_mm_hvx_vtcm_layout_build( bool is_matmul_id, bool is_fused_nx ) { + (void)src1_row_size; size_t src0_sz = 0; size_t src1_sz = 0; size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0; @@ -517,12 +505,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); } - size_t flat_act_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - - size_t act_sz = (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) - ? hex_round_up(flat_act_row_size * src1_nrows, 128) - : hex_round_up(tiled_act_row_size * src1_nrows, 128); + size_t act_sz = hex_round_up(tiled_act_row_size * src1_nrows, 128); src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer src1_sz = act_sz; // quantized activation buffer @@ -547,6 +531,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( src0_sz = src0_sz_per_thread * n_threads; dst_sz = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; + src2_sz = 0; + src3_sz = 0; } else { const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1; @@ -559,15 +545,6 @@ static inline void htp_mm_hvx_vtcm_layout_build( dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; break; } - case HTP_MM_KERNEL_HVX_F16_F32_DDR: - case HTP_MM_KERNEL_HVX_F16_F16_DDR: - case HTP_MM_KERNEL_HVX_F32_F32_DDR: - case HTP_MM_KERNEL_HVX_F32_F16_DDR: { - src0_sz = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads; - src1_sz = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads; - dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } case HTP_MM_KERNEL_HVX_F32_F32_VTCM: { size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128); src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256); @@ -598,28 +575,6 @@ static inline void htp_mm_hvx_vtcm_layout_build( dst_sz = dst_size_per_thread * n_threads; break; } - case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: { - size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - - src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); - src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256); - - src0_sz = src0_sz * n_threads; - - if (is_repack) { - uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - uint32_t n_k_tiles = ne10 / 32; - uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - src0_sz = repacked_vtcm_size * n_threads; - } - - size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); - size_t dst_slice_per_thread = dst_nrows > 0 ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0; - size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread; - dst_sz = dst_size_per_thread * n_threads; - break; - } default: break; } @@ -640,19 +595,99 @@ static inline void htp_mm_hvx_vtcm_layout_build( L->total_bytes = off; } +static inline bool htp_mm_hvx_solve_vtcm_params( + int kernel_type, + int wtype, + uint32_t ne10, + uint32_t src1_nrows, + uint32_t n_threads, + size_t dst_row_size, + size_t src0_row_size, + size_t src1_row_size, + size_t src2_row_size, + uint32_t n_prefetch, + size_t vtcm_budget, + struct htp_mm_hvx_vtcm_layout * L_out, + uint32_t * m_chunk_out +) { + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, kernel_type, wtype, ne10, src1_nrows, n_threads, + dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false + ); + + if (L.total_bytes <= vtcm_budget) { + *L_out = L; + *m_chunk_out = src1_nrows; + return true; + } + + const size_t fixed_bytes = L.src0_bytes + L.src2_bytes + L.dst_bytes; + if (vtcm_budget <= fixed_bytes) { + return false; + } + + const size_t avail_act = vtcm_budget - fixed_bytes; + size_t row_size = 0; + if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { + row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) + ? htp_mm_q8_1_tiled_row_size(ne10) + : htp_mm_q8_0_tiled_row_size(ne10); + } else if (kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) { + row_size = hex_round_up(ne10 * 2, 128); + } else { + row_size = hex_round_up(ne10 * 4, 128); + } + if (row_size == 0) { + return false; + } + + uint32_t m_chunk = (uint32_t) (avail_act / row_size); + if (m_chunk > 1) { + m_chunk &= ~1U; + } + if (m_chunk > src1_nrows) { + m_chunk = src1_nrows; + } + if (m_chunk < 1) { + return false; + } + + htp_mm_hvx_vtcm_layout_build( + &L, kernel_type, wtype, ne10, m_chunk, n_threads, + dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false + ); + + while (m_chunk > 2 && L.total_bytes > vtcm_budget) { + m_chunk -= 2; + htp_mm_hvx_vtcm_layout_build( + &L, kernel_type, wtype, ne10, m_chunk, n_threads, + dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false + ); + } + + if (L.total_bytes <= vtcm_budget) { + *L_out = L; + *m_chunk_out = m_chunk; + return true; + } + + return false; +} + static inline size_t htp_mm_hmx_get_2d_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size + int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t src2_size ) { struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size, src2_size); return L.total_bytes; } static inline size_t htp_mm_hmx_get_batched_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) { + int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads, size_t src2_size) { (void)pipeline; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0, src2_size); return L.total_bytes; } @@ -665,6 +700,7 @@ static inline bool htp_mm_hmx_solve_batched_params( bool use_dma_activation, int n_threads, bool pipeline, + size_t src2_size, size_t vtcm_budget, size_t * m_chunk_out, size_t * n_chunk_out, @@ -679,7 +715,7 @@ static inline bool htp_mm_hmx_solve_batched_params( int act_threads = n_threads; while (act_threads >= 1) { - size_t group_overhead = htp_mm_hmx_get_batched_overhead(); + size_t group_overhead = htp_mm_hmx_get_batched_overhead() + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0); size_t group_size_per_n, group_size_per_m, group_size_per_mn; htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn); @@ -690,7 +726,7 @@ static inline bool htp_mm_hmx_solve_batched_params( if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded, (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT, &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) { - size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads); + size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads, src2_size); if (exact_size <= vtcm_budget) { size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate; if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) { @@ -730,6 +766,7 @@ static inline bool htp_mm_hmx_solve_2d_params( bool pipeline, bool is_matmul_id, uint32_t aligned_tile_size, + size_t src2_size, size_t vtcm_budget, size_t * m_chunk_out, size_t * n_chunk_out, @@ -746,7 +783,7 @@ static inline bool htp_mm_hmx_solve_2d_params( int act_threads = n_threads; while (act_threads >= 1) { - size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id); + size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id) + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0); size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn; htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn); @@ -757,7 +794,7 @@ static inline bool htp_mm_hmx_solve_2d_params( if (htp_mm_hmx_compute_chunks(vtcm_budget, simple_2d_overhead, simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn, m_for_chunks, ne01_padded, (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) m_for_cost * HTP_MM_HMX_COST_A_CONVERT, &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) { - size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size); + size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size, src2_size); if (exact_size <= vtcm_budget) { size_t mblocks = ((size_t) m_for_cost + m_chunk_candidate - 1) / m_chunk_candidate; if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) { diff --git a/ggml/src/ggml-hexagon/htp/pad-ops.c b/ggml/src/ggml-hexagon/htp/pad-ops.c index 0222f24dc..85f25a8eb 100644 --- a/ggml/src/ggml-hexagon/htp/pad-ops.c +++ b/ggml/src/ggml-hexagon/htp/pad-ops.c @@ -7,7 +7,7 @@ #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hvx-utils.h" #define GGML_COMMON_DECL_C @@ -51,6 +51,15 @@ static inline const uint8_t * pad_src_row_ptr(const struct htp_tensor * src, + (i3 - (uint32_t)lp3) * src->nb[3]; } +static inline dma_addr_t pad_src_row_data(const struct htp_tensor * src, + uint32_t i1, uint32_t i2, uint32_t i3, + int32_t lp1, int32_t lp2, int32_t lp3) { + return src->data + + (i1 - (uint32_t)lp1) * src->nb[1] + + (i2 - (uint32_t)lp2) * src->nb[2] + + (i3 - (uint32_t)lp3) * src->nb[3]; +} + /* Compute the DDR src row pointer for a circular row (wrap-around indexing) */ static inline const uint8_t * pad_circ_src_row_ptr(const struct htp_tensor * src, uint32_t i1, uint32_t i2, uint32_t i3, @@ -61,6 +70,15 @@ static inline const uint8_t * pad_circ_src_row_ptr(const struct htp_tensor * src + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3]; } +static inline dma_addr_t pad_circ_src_row_data(const struct htp_tensor * src, + uint32_t i1, uint32_t i2, uint32_t i3, + int32_t lp1, int32_t lp2, int32_t lp3) { + return src->data + + wrap_around((int32_t)i1 - lp1, src->ne[1]) * src->nb[1] + + wrap_around((int32_t)i2 - lp2, src->ne[2]) * src->nb[2] + + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3]; +} + struct htp_pad_context { struct htp_ops_context * octx; @@ -118,7 +136,7 @@ struct htp_pad_context { uint8_t * src_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; \ uint8_t * dst_spad_base = octx->dst_spad.data + ith * octx->dst_spad.size_per_thread; \ \ - dma_queue * dma = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; // --------------------------------------------------------------------------- // HVX vectorized PAD kernel @@ -196,9 +214,9 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void uint8_t * src_spad_cur = src_spad_base + spad_idx * src_row_size_aligned; uint8_t * dst_spad_cur = dst_spad_base + spad_idx * dst_row_size_aligned; - dma_queue_push_vtcm_to_ddr(dma, - dma_make_ptr((uint8_t *)dst->data, dst_spad_cur), - dst_row_size, dst_row_size_aligned, 0); + dma_queue_push(dma_q, + dma_make_data(dst->data, dst_spad_cur), + dst_row_size, dst_row_size_aligned, dst_row_size, 0); uint32_t i1, i2, i3; pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3); @@ -207,15 +225,14 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void lp2, rp2, ne2, lp3, rp3, ne3); - const uint8_t * src_ptr = interior - ? pad_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3) : NULL; + const dma_addr_t src_data = interior + ? pad_src_row_data(src, i1, i2, i3, lp1, lp2, lp3) : src->data; // Interior row: real DMA (1 row) from DDR to VTCM. // Border row: null DMA (nrows=0) - dma_queue_push_ddr_to_vtcm(dma, - dma_make_ptr(src_spad_cur, - src_ptr ? src_ptr : (const uint8_t *)src_spad_cur), - src_row_size_aligned, src_row_size, src_ptr ? 1 : 0); + dma_queue_push(dma_q, + dma_make_data(src_spad_cur, src_data), + src_row_size_aligned, src_row_size, src_row_size, interior ? 1 : 0); } // ----------------------------------------------------------------------- @@ -225,13 +242,13 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void struct htp_thread_trace * tr = &octx->ctx->trace[ith]; for (uint32_t ir = row_start; ir < row_end; ir++) { - uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src; - uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst; + uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma_q).dst; uint32_t i1, i2, i3; pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3); - uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3; + const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3; const int interior = pad_is_interior(i1, i2, i3, lp1, rp1, ne1, @@ -254,9 +271,9 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - dma_queue_push_vtcm_to_ddr(dma, - dma_make_ptr(dst_ptr, dst_spad_cur), - dst_row_size, dst_row_size_aligned, 1); + dma_queue_push(dma_q, + dma_make_data(dst_data, dst_spad_cur), + dst_row_size, dst_row_size_aligned, dst_row_size, 1); const uint32_t next_row = ir + 2; if (next_row < row_end) { @@ -266,17 +283,16 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void lp1, rp1, ne1, lp2, rp2, ne2, lp3, rp3, ne3); - const uint8_t * next_src_ptr = next_interior - ? pad_src_row_ptr(src, ni1, ni2, ni3, lp1, lp2, lp3) : NULL; + const dma_addr_t next_src_data = next_interior + ? pad_src_row_data(src, ni1, ni2, ni3, lp1, lp2, lp3) : src->data; - dma_queue_push_ddr_to_vtcm(dma, - dma_make_ptr(src_spad_cur, - next_src_ptr ? next_src_ptr : (const uint8_t *)src_spad_cur), - src_row_size_aligned, src_row_size, next_src_ptr ? 1 : 0); + dma_queue_push(dma_q, + dma_make_data(src_spad_cur, next_src_data), + src_row_size_aligned, src_row_size, src_row_size, next_interior ? 1 : 0); } } - dma_queue_flush(dma); + dma_queue_flush(dma_q); FARF(HIGH, "pad-hvx-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n", ith, nth, @@ -372,15 +388,16 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i uint8_t * src_spad_cur = src_spad_base + spad_idx * src_row_size_aligned; uint8_t * dst_spad_cur = dst_spad_base + spad_idx * dst_row_size_aligned; - dma_queue_push_vtcm_to_ddr(dma, - dma_make_ptr((uint8_t *)dst->data, dst_spad_cur), - dst_row_size, dst_row_size_aligned, 0); + dma_queue_push(dma_q, + dma_make_data(dst->data, dst_spad_cur), + dst_row_size, dst_row_size_aligned, dst_row_size, 0); uint32_t pi1, pi2, pi3; pad_decompose_row(ir, ne1, ne2, &pi1, &pi2, &pi3); - dma_queue_push_ddr_to_vtcm(dma, - dma_make_ptr(src_spad_cur, pad_circ_src_row_ptr(src, pi1, pi2, pi3, lp1, lp2, lp3)), - src_row_size_aligned, src_row_size, 1); + const dma_addr_t src_data = pad_circ_src_row_data(src, pi1, pi2, pi3, lp1, lp2, lp3); + dma_queue_push(dma_q, + dma_make_data(src_spad_cur, src_data), + src_row_size_aligned, src_row_size, src_row_size, 1); } // ----------------------------------------------------------------------- @@ -390,12 +407,12 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i struct htp_thread_trace * tr = &octx->ctx->trace[ith]; for (uint32_t ir = row_start; ir < row_end; ir++) { - uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src; - uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst; + uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma_q).src; + uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma_q).dst; uint32_t i1, i2, i3; pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3); - uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3; + const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); if (lp0 > 0) { @@ -431,22 +448,22 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); - dma_queue_push_vtcm_to_ddr(dma, - dma_make_ptr(dst_ptr, dst_spad_cur), - dst_row_size, dst_row_size_aligned, 1); + dma_queue_push(dma_q, + dma_make_data(dst_data, dst_spad_cur), + dst_row_size, dst_row_size_aligned, dst_row_size, 1); const uint32_t next_row = ir + 2; if (next_row < row_end) { uint32_t nri1, nri2, nri3; pad_decompose_row(next_row, ne1, ne2, &nri1, &nri2, &nri3); - dma_queue_push_ddr_to_vtcm(dma, - dma_make_ptr(src_spad_cur, - pad_circ_src_row_ptr(src, nri1, nri2, nri3, lp1, lp2, lp3)), - src_row_size_aligned, src_row_size, 1); + const dma_addr_t next_src_data = pad_circ_src_row_data(src, nri1, nri2, nri3, lp1, lp2, lp3); + dma_queue_push(dma_q, + dma_make_data(src_spad_cur, next_src_data), + src_row_size_aligned, src_row_size, src_row_size, 1); } } - dma_queue_flush(dma); + dma_queue_flush(dma_q); FARF(HIGH, "pad-hvx-circ-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n", ith, nth, @@ -468,10 +485,6 @@ int op_pad(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } - const int32_t lp0 = octx->op_params[0]; const int32_t rp0 = octx->op_params[1]; const int32_t lp1 = octx->op_params[2]; @@ -515,6 +528,10 @@ int op_pad(struct htp_ops_context * octx) { const int use_dma = (src0->nb[0] == (uint32_t)type_size) && (ne00 >= 512) && (octx->ctx->vtcm_size >= vtcm_needed); + if (!use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) { + return HTP_STATUS_NO_SUPPORT; + } + if (use_dma) { octx->src0_spad.size_per_thread = 2 * src_row_size_aligned; octx->dst_spad.size_per_thread = 2 * dst_row_size_aligned; diff --git a/ggml/src/ggml-hexagon/htp/repeat-ops.c b/ggml/src/ggml-hexagon/htp/repeat-ops.c index 530279d65..2551be225 100644 --- a/ggml/src/ggml-hexagon/htp/repeat-ops.c +++ b/ggml/src/ggml-hexagon/htp/repeat-ops.c @@ -122,8 +122,8 @@ int op_repeat(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; } const uint32_t total_dst_rows = dst->ne[1] * dst->ne[2] * dst->ne[3]; diff --git a/ggml/src/ggml-hexagon/htp/roll-ops.c b/ggml/src/ggml-hexagon/htp/roll-ops.c index 6faf2ac47..9c373f56d 100644 --- a/ggml/src/ggml-hexagon/htp/roll-ops.c +++ b/ggml/src/ggml-hexagon/htp/roll-ops.c @@ -67,8 +67,8 @@ static inline uint32_t htp_roll_wrap(int32_t i, uint32_t ne) { #define htp_roll_dma_preamble dma_queue * q = octx->ctx->dma[0]; static inline void roll_dma_push(dma_queue * q, - uintptr_t dst, - uintptr_t src, + dma_addr_t dst, + dma_addr_t src, uint32_t dst_stride, uint32_t src_stride, uint32_t bytes, @@ -77,10 +77,10 @@ static inline void roll_dma_push(dma_queue * q, return; } - if (!dma_queue_push(q, dma_make_ptr((void *) dst, (const void *) src), dst_stride, src_stride, bytes, nrows)) { + if (!dma_queue_push(q, dma_make_data(dst, src), dst_stride, src_stride, bytes, nrows)) { dma_queue_flush(q); - dma_queue_push(q, dma_make_ptr((void *) dst, (const void *) src), - dst_stride, src_stride, bytes, nrows); + dma_queue_push(q, dma_make_data(dst, src), + dst_stride, src_stride, bytes, nrows); } } @@ -92,29 +92,29 @@ static inline void roll_dma_push_rows(dma_queue * q, uint32_t nrows, uint32_t row_size, uint32_t i0_src0) { - const uintptr_t dst_base = dst->data + (uintptr_t) dst_row * row_size; - const uintptr_t src_base = src0->data + (uintptr_t) src_row * row_size; - const uint32_t n0 = src0->ne[0] - i0_src0; + const dma_addr_t dst_base = dst->data + (size_t) dst_row * row_size; + const dma_addr_t src_base = src0->data + (size_t) src_row * row_size; + const uint32_t n0 = src0->ne[0] - i0_src0; - roll_dma_push(q, dst_base, src_base + (uintptr_t) i0_src0 * sizeof(float), + roll_dma_push(q, dst_base, src_base + (size_t) i0_src0 * sizeof(float), row_size, row_size, n0 * sizeof(float), nrows); - roll_dma_push(q, dst_base + (uintptr_t) n0 * sizeof(float), src_base, + roll_dma_push(q, dst_base + (size_t) n0 * sizeof(float), src_base, row_size, row_size, i0_src0 * sizeof(float), nrows); } // Same row-wrap split as roll_dma_push_rows, but addressed with explicit byte strides so it // also works for a src0 that is row-contiguous only (e.g. a permuted view) rather than fully packed. static inline void roll_dma_push_range(dma_queue * q, - uintptr_t dst_row, - uintptr_t src_row, + dma_addr_t dst_row, + dma_addr_t src_row, uint32_t dst_stride, uint32_t src_stride, uint32_t nrows, uint32_t i0_src0, uint32_t n0) { - roll_dma_push(q, dst_row, src_row + (uintptr_t) i0_src0 * sizeof(float), + roll_dma_push(q, dst_row, src_row + (size_t) i0_src0 * sizeof(float), dst_stride, src_stride, n0 * sizeof(float), nrows); - roll_dma_push(q, dst_row + (uintptr_t) n0 * sizeof(float), src_row, + roll_dma_push(q, dst_row + (size_t) n0 * sizeof(float), src_row, dst_stride, src_stride, i0_src0 * sizeof(float), nrows); } @@ -184,12 +184,12 @@ static int roll_dma_f32_strided(struct htp_ops_context * octx) { for (uint32_t i2 = 0; i2 < ne2; i2++) { const uint32_t i02 = htp_roll_wrap((int32_t) i2 - s2, ne2); - const uintptr_t dst_row0 = dst->data + (uintptr_t) i2 * nb2 + (uintptr_t) i3 * nb3; - const uintptr_t src_row0 = src0->data + (uintptr_t) i02 * nb02 + (uintptr_t) i03 * nb03; + const dma_addr_t dst_row0 = dst->data + (size_t) i2 * nb2 + (size_t) i3 * nb3; + const dma_addr_t src_row0 = src0->data + (size_t) i02 * nb02 + (size_t) i03 * nb03; - roll_dma_push_range(q, dst_row0, src_row0 + (uintptr_t) i1_src0 * nb01, + roll_dma_push_range(q, dst_row0, src_row0 + (size_t) i1_src0 * nb01, nb1, nb01, n1_first, i0_src0, n0); - roll_dma_push_range(q, dst_row0 + (uintptr_t) n1_first * nb1, src_row0, + roll_dma_push_range(q, dst_row0 + (size_t) n1_first * nb1, src_row0, nb1, nb01, i1_src0, i0_src0, n0); } } @@ -223,8 +223,8 @@ static void roll_thread_f32(unsigned int nth, unsigned int ith, void * data) { const uint32_t i02 = htp_roll_wrap((int32_t) i2 - s2, ne2); const uint32_t i03 = htp_roll_wrap((int32_t) i3 - s3, ne3); - const uint8_t * src_row = (const uint8_t *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; - uint8_t * dst_row = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3; + const uint8_t * src_row = (const uint8_t *) (uintptr_t) src0->data + i01*nb01 + i02*nb02 + i03*nb03; + uint8_t * dst_row = (uint8_t *) (uintptr_t) dst->data + i1*nb1 + i2*nb2 + i3*nb3; hex_l2fetch(src_row + i0_src0 * sizeof(float), n0 * sizeof(float), ne0 * sizeof(float), 1); hvx_copy_uu(dst_row, src_row + i0_src0 * sizeof(float), n0, sizeof(float)); @@ -261,10 +261,6 @@ int execute_op_roll_f32(struct htp_ops_context * octx) { return HTP_STATUS_INVAL_PARAMS; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } - const uint32_t total_rows = ne1 * ne2 * ne3; const size_t dst_row_size = ne0 * sizeof(float); @@ -290,6 +286,10 @@ int execute_op_roll_f32(struct htp_ops_context * octx) { return roll_dma_f32_strided(octx); } + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; + } + const uint32_t n_threads = octx->n_threads; struct htp_roll_context rctx = { .octx = octx, diff --git a/ggml/src/ggml-hexagon/htp/rope-ops.c b/ggml/src/ggml-hexagon/htp/rope-ops.c index c36976ed0..f6b4d383c 100644 --- a/ggml/src/ggml-hexagon/htp/rope-ops.c +++ b/ggml/src/ggml-hexagon/htp/rope-ops.c @@ -9,7 +9,7 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hvx-utils.h" #include "hex-fastdiv.h" @@ -85,6 +85,8 @@ struct htp_rope_context { struct fastdiv_values div_ne2_ne1; struct fastdiv_values div_ne1; + + const float * freq_factors; }; static float rope_yarn_ramp(const float low, const float high, const int i0) { @@ -562,10 +564,10 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { float * theta_cache = (float *) (src0_spad_base); src0_spad_base = src0_spad_base + rctx->theta_cache_offset; - dma_queue * dma_queue = octx->ctx->dma[ith]; + dma_queue * dma_q = octx->ctx->dma[ith]; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - const int32_t * pos = (const int32_t *) src1->data; - const float * freq_factors = src2 ? (const float *) src2->data : NULL; + const int32_t * pos = (const int32_t *) (uintptr_t) src1->data; + const float * freq_factors = rctx->freq_factors; const uint32_t i3_start = fastdiv(src0_start_row, &rctx->div_ne2_ne1); const uint32_t rem = fastmodulo(src0_start_row, ne2 * ne1, &rctx->div_ne2_ne1); @@ -587,7 +589,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { const uint32_t nrows = MIN(src0_end_row - ir, ne1 - i1); // Depth before prefetch - const uint32_t dma_depth = dma_queue_depth(dma_queue); + const uint32_t dma_depth = dma_queue_depth(dma_q); // Prefetch up to 2 blocks const uint32_t p_nrows = MIN(nrows, 2 * HTP_ROPE_SPAD_BLOCK); @@ -595,12 +597,12 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { const uint32_t pnr = MIN(nrows - pr, HTP_ROPE_SPAD_BLOCK); const uint32_t slot = (cur_slot + pr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS; uint8_t * spad_slot = rope_spad_slot(src0_spad_base, slot, rctx->src0_row_size_aligned); - const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + (i1 + pr) * nb01; + const dma_addr_t src0_data = src0->data + i3 * nb03 + i2 * nb02 + (i1 + pr) * nb01; // Dummy DMA transaction for sequencing (interleaving wr, rd, wr, rd, ...) - dma_queue_push(dma_queue, dma_make_ptr((void *) dst->data, spad_slot), 0, 0, 0, 0); + dma_queue_push(dma_q, dma_make_data(dst->data, spad_slot), 0, 0, 0, 0); - dma_queue_push(dma_queue, dma_make_ptr(spad_slot, src_addr), + dma_queue_push(dma_q, dma_make_data(spad_slot, src0_data), rctx->src0_row_size_aligned, rctx->src0_row_stride, rctx->src0_row_size, pnr); } @@ -634,7 +636,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { } // Skip output DMA transactions from prev block (if any) - for (uint32_t d = 0; d < dma_depth; d++) { dma_queue_pop_nowait(dma_queue); } + for (uint32_t d = 0; d < dma_depth; d++) { dma_queue_pop_nowait(dma_q); } // Compute loop const uint32_t ne = is_vision ? ne0 : rctx->n_dims; @@ -647,8 +649,8 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { const uint32_t cur_ir = base_ir + cr; const uint32_t cur_i1 = base_i1 + cr; - dma_queue_pop(dma_queue); - uint8_t * cur_spad = (uint8_t *) dma_queue_pop(dma_queue).dst; + dma_queue_pop(dma_q); + uint8_t * cur_spad = (uint8_t *) dma_queue_pop(dma_q).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur_ir); if (is_neox || is_vision) { @@ -658,8 +660,8 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur_ir); - uint8_t * dst_addr = (uint8_t *) dst->data + i3 * nb3 + i2 * nb2 + cur_i1 * nb1; - dma_queue_push(dma_queue, dma_make_ptr(dst_addr, cur_spad), + const dma_addr_t dst_data = dst->data + i3 * nb3 + i2 * nb2 + cur_i1 * nb1; + dma_queue_push(dma_q, dma_make_data(dst_data, cur_spad), rctx->dst_row_stride, rctx->src0_row_size_aligned, rctx->dst_row_size, cnr); // Prefetch 2 blocks ahead into the slot just freed @@ -668,9 +670,9 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { const uint32_t pnr = MIN(nrows - p_cr, HTP_ROPE_SPAD_BLOCK); const uint32_t p_slot = (cur_slot + p_cr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS; uint8_t * p_spad = rope_spad_slot(src0_spad_base, p_slot, rctx->src0_row_size_aligned); - const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + (base_i1 + p_cr) * nb01; + const dma_addr_t p_src0_data = src0->data + i3 * nb03 + i2 * nb02 + (base_i1 + p_cr) * nb01; - dma_queue_push(dma_queue, dma_make_ptr(p_spad, src_addr), + dma_queue_push(dma_q, dma_make_data(p_spad, p_src0_data), rctx->src0_row_size_aligned, rctx->src0_row_stride, rctx->src0_row_size, pnr); } } @@ -685,7 +687,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) { } done: - dma_queue_flush(dma_queue); + dma_queue_flush(dma_q); FARF(HIGH, "rope-f32: %d/%d: (%u:%u)\n", ith, nth, src0_start_row, src0_end_row); } @@ -713,6 +715,10 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) { } assert(octx->ctx->vtcm_size >= kparams->vtcm_size); + if (htp_tensor_is_extended(src1)) { + return HTP_STATUS_NO_SUPPORT; + } + const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3]; const size_t dst_data_row_size = dst->ne[0] * sizeof(float); @@ -748,6 +754,16 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) { rctx.spad_per_thread = kparams->spad_per_thread; rctx.theta_cache_offset = kparams->theta_cache_offset; + if (src2) { + dma_queue * dma_q = octx->ctx->dma[0]; + const size_t ff_size = src2->ne[0] * sizeof(float); + float * vtcm_freq_factors = (float *) (rctx.vtcm_base + kparams->freq_factors_offset); + dma_queue_push(dma_q, dma_make_data(vtcm_freq_factors, src2->data), + kparams->freq_factors_size, 0, ff_size, 1); + dma_queue_pop(dma_q); + rctx.freq_factors = vtcm_freq_factors; + } + const int32_t * op_params = &octx->op_params[0]; rctx.n_dims = ((const int32_t *) op_params)[1]; rctx.mode = ((const int32_t *) op_params)[2]; diff --git a/ggml/src/ggml-hexagon/htp/rope-ops.h b/ggml/src/ggml-hexagon/htp/rope-ops.h index 476653d05..ee055ccbc 100644 --- a/ggml/src/ggml-hexagon/htp/rope-ops.h +++ b/ggml/src/ggml-hexagon/htp/rope-ops.h @@ -16,6 +16,8 @@ struct htp_rope_kernel_params { uint32_t spad_per_thread; uint32_t theta_cache_offset; uint32_t src0_row_size_aligned; + uint32_t freq_factors_offset; + uint32_t freq_factors_size; struct fastdiv_values div_ne2_ne1; struct fastdiv_values div_ne1; @@ -32,21 +34,25 @@ struct htp_rope_vtcm_layout { size_t bytes_per_thread; size_t theta_cache_size_aligned; size_t src0_row_size_aligned; + size_t freq_factors_size_aligned; }; static inline void htp_rope_vtcm_layout_build( struct htp_rope_vtcm_layout * layout, uint32_t ne00, - uint32_t n_threads + uint32_t n_threads, + uint32_t n_freq_factors ) { - const size_t src0_row_size = ne00 * sizeof(float); - const size_t src0_row_size_aligned = hex_round_up((uint32_t) src0_row_size, 128); - const size_t theta_cache_size_aligned = hex_round_up((uint32_t) src0_row_size, 256); + const size_t src0_row_size = ne00 * sizeof(float); + const size_t src0_row_size_aligned = hex_round_up((uint32_t) src0_row_size, 128); + const size_t theta_cache_size_aligned = hex_round_up((uint32_t) src0_row_size, 256); + const size_t freq_factors_size_aligned = hex_round_up(n_freq_factors * sizeof(float), 256); - layout->src0_row_size_aligned = src0_row_size_aligned; - layout->theta_cache_size_aligned = theta_cache_size_aligned; - layout->bytes_per_thread = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned; - layout->total_bytes = layout->bytes_per_thread * n_threads; + layout->src0_row_size_aligned = src0_row_size_aligned; + layout->theta_cache_size_aligned = theta_cache_size_aligned; + layout->freq_factors_size_aligned = freq_factors_size_aligned; + layout->bytes_per_thread = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned; + layout->total_bytes = layout->bytes_per_thread * n_threads + freq_factors_size_aligned; } static inline uint8_t * rope_spad_slot(uint8_t * base, uint32_t slot, size_t row_size_aligned) { diff --git a/ggml/src/ggml-hexagon/htp/set-rows-ops.c b/ggml/src/ggml-hexagon/htp/set-rows-ops.c index fbd5162a7..1d72538f1 100644 --- a/ggml/src/ggml-hexagon/htp/set-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/set-rows-ops.c @@ -77,7 +77,7 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig return; \ } \ const uint32_t ir1 = MIN(ir0 + dr, srctx->task_start + srctx->tasks); \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ + dma_queue * dma_q = octx->ctx->dma[ith]; \ const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout; \ uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ uint8_t * vtcm_dst = srctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ @@ -90,14 +90,14 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig uint32_t pi03 = 0; \ for (uint32_t step = 0, spad_idx = 0; step < total_steps && spad_idx < 2; ++step, spad_idx++) { \ uint32_t i = ir0 + pi_step; \ - const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03; \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)octx->dst->data, \ - vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ + const dma_addr_t src0_data = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03; \ + dma_queue_push(dma_q, \ + dma_make_data(octx->dst->data, \ + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0); \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size), \ - (const void *)src0_ptr), \ + dma_queue_push(dma_q, \ + dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, \ + src0_data), \ vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1); \ pi_step++; \ if (pi_step == nrows_per_thread) { \ @@ -115,8 +115,8 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig uint32_t ci11_base = 0; \ uint32_t ci12_base = 0; \ for (uint32_t step = 0; step < total_steps; ++step) { \ - void * dst_spad = (void *) dma_queue_pop(dma_queue).src; \ - void * src_spad = (void *) dma_queue_pop(dma_queue).dst; \ + void * dst_spad = (void *) dma_queue_pop(dma_q).src; \ + void * src_spad = (void *) dma_queue_pop(dma_q).dst; \ uint32_t i = ir0 + ci_step; \ const uintptr_t src1_addr = octx->src[1]->data + i*nb10 + ci11_base*nb11 + ci12_base*nb12; \ const IDX_TYPE i1 = *(const IDX_TYPE *)src1_addr; \ @@ -128,21 +128,21 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, step); \ if (valid_i1) { \ - const uintptr_t dst_ptr = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3; \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)dst_ptr, (const void *)dst_spad), \ + const dma_addr_t dst_data = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3; \ + dma_queue_push(dma_q, \ + dma_make_data(dst_data, dst_spad), \ dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 1); \ } else { \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)octx->dst->data, (const void *)dst_spad), \ + dma_queue_push(dma_q, \ + dma_make_data(octx->dst->data, dst_spad), \ dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0); \ } \ const uint32_t next_step = step + 2; \ if (next_step < total_steps) { \ uint32_t ni = ir0 + pi_step; \ - const uintptr_t psrc0_ptr = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03; \ - dma_queue_push(dma_queue, \ - dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr), \ + const dma_addr_t psrc0_data = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03; \ + dma_queue_push(dma_q, \ + dma_make_data(src_spad, psrc0_data), \ vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1); \ pi_step++; \ if (pi_step == nrows_per_thread) { \ @@ -172,7 +172,7 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig } \ } \ } \ - dma_queue_flush(dma_queue); \ + dma_queue_flush(dma_q); \ } SET_ROWS_THREAD_DMA_FN(f32, int32_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) @@ -196,8 +196,8 @@ int op_set_rows(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + if (htp_tensor_is_extended(octx->src[1])) { + return HTP_STATUS_NO_SUPPORT; } const struct htp_tensor * dst = octx->dst; diff --git a/ggml/src/ggml-hexagon/htp/softmax-ops.c b/ggml/src/ggml-hexagon/htp/softmax-ops.c index 2497ec763..48be5d725 100644 --- a/ggml/src/ggml-hexagon/htp/softmax-ops.c +++ b/ggml/src/ggml-hexagon/htp/softmax-ops.c @@ -8,54 +8,49 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" +#include "work-queue.h" #include "hvx-utils.h" #include "hex-fastdiv.h" +#include "hex-common.h" +#include "hex-profile.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" -#include "hex-common.h" -#include "hex-profile.h" #include "htp-ctx.h" #include "htp-ops.h" #include "htp-tensor.h" - -#define htp_softmax_preamble3 \ - const uint32_t ne00 = src0->ne[0]; \ - const uint32_t ne01 = src0->ne[1]; \ - const uint32_t ne02 = src0->ne[2]; \ - const uint32_t ne03 = src0->ne[3]; \ - \ - const uint32_t nb00 = src0->nb[0]; \ - const uint32_t nb01 = src0->nb[1]; \ - const uint32_t nb02 = src0->nb[2]; \ - const uint32_t nb03 = src0->nb[3]; \ - \ - const uint32_t ne10 = src1 ? src1->ne[0] : 1; \ - const uint32_t ne11 = src1 ? src1->ne[1] : 1; \ - const uint32_t ne12 = src1 ? src1->ne[2] : 1; \ - const uint32_t ne13 = src1 ? src1->ne[3] : 1; \ - \ - const uint32_t nb10 = src1 ? src1->nb[0] : 1; \ - const uint32_t nb11 = src1 ? src1->nb[1] : 1; \ - const uint32_t nb12 = src1 ? src1->nb[2] : 1; \ - const uint32_t nb13 = src1 ? src1->nb[3] : 1; \ - \ - 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 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]; +#include "htp-vtcm.h" +#include "htp/softmax-ops.h" +#include "hvx-flash-attn.h" struct htp_softmax_context { struct htp_ops_context * octx; + const struct htp_softmax_kernel_params * kparams; + + void * compute; + + dma_addr_t data_src0; + dma_addr_t data_src1; + dma_addr_t data_dst; + + uint8_t * vtcm_src0; + uint8_t * vtcm_src1; + uint8_t * vtcm_dst; + + uint32_t vtcm_src0_size_per_thread; + uint32_t vtcm_src1_size_per_thread; + uint32_t vtcm_dst_size_per_thread; + + uint32_t src0_spad_half_size; + uint32_t src1_spad_half_size; + uint32_t dst_spad_half_size; + + uint32_t src0_row_size_aligned; + uint32_t src1_row_size_aligned; + uint32_t dst_row_size_aligned; bool use_f16; - bool use_src1; uint32_t n_head; uint32_t n_head_log2; @@ -65,127 +60,148 @@ struct htp_softmax_context { float m0; float m1; - struct fastdiv_values fastdiv_ne01; - struct fastdiv_values fastdiv_ne02; - struct fastdiv_values fastdiv_ne12; // For mask broadcasting - struct fastdiv_values fastdiv_ne13; // For mask broadcasting + struct fastdiv_values div_ne01; + struct fastdiv_values div_ne02; + struct fastdiv_values div_ne12; + struct fastdiv_values div_ne13; uint32_t src0_nrows_per_thread; uint32_t row_start; uint32_t nrows; + + float slopes[512] __attribute__((aligned(128))); }; -static void apply_mask(float * restrict wp0, - const float * restrict mp_f32, - const __fp16 * restrict mp_f16, - uint32_t ne00, - float slope, - bool use_f16) { - if (!mp_f32) { - return; +typedef void (*softmax_compute_fn_t)( + void * restrict dst, + const void * restrict src0, + const void * restrict mask, + uint32_t ne00, + float scale, + float slope +); + +static void hvx_fast_softmax_prep_f16(const uint8_t * restrict src, + uint8_t * restrict dst, + const int num_elems, + float scale, + const uint8_t * restrict mask, + float slope) { + const HVX_Vector * restrict v_src = (const HVX_Vector *) src; + HVX_Vector * restrict v_dst = (HVX_Vector *) dst; + const HVX_Vector * restrict v_mask = (const HVX_Vector *) mask; + + HVX_Vector scale_vec = hvx_vec_splat_f32(scale); + HVX_Vector slope_vec = hvx_vec_splat_f32(slope); + + const int nvec_64 = num_elems / VLEN_FP16; + const int nloe_64 = num_elems % VLEN_FP16; + + #pragma unroll(2) + for (int i = 0; i < nvec_64; i++) { + HVX_VectorPair p = hvx_vec_f16_to_f32(v_mask[i]); + HVX_Vector m0 = Q6_V_lo_W(p); + HVX_Vector m1 = Q6_V_hi_W(p); + + HVX_Vector s0 = v_src[2 * i]; + HVX_Vector s1 = v_src[2 * i + 1]; + + HVX_Vector v0 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s0, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m0, slope_vec)); + HVX_Vector v1 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s1, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m1, slope_vec)); + + v_dst[2 * i] = Q6_Vsf_equals_Vqf32(v0); + v_dst[2 * i + 1] = Q6_Vsf_equals_Vqf32(v1); } - if (use_f16) { - for (uint32_t i = 0; i < ne00; ++i) { - wp0[i] += slope * (float) mp_f16[i]; - } - } else { - for (uint32_t i = 0; i < ne00; ++i) { - wp0[i] += slope * mp_f32[i]; + + if (nloe_64 > 0) { + HVX_VectorPair p = hvx_vec_f16_to_f32(v_mask[nvec_64]); + HVX_Vector m0 = Q6_V_lo_W(p); + + HVX_Vector s0 = v_src[2 * nvec_64]; + HVX_Vector v0 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s0, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m0, slope_vec)); + + if (nloe_64 <= VLEN_FP32) { + hvx_vec_store_a(&v_dst[2 * nvec_64], nloe_64 * sizeof(float), Q6_Vsf_equals_Vqf32(v0)); + } else { + v_dst[2 * nvec_64] = Q6_Vsf_equals_Vqf32(v0); + + HVX_Vector m1 = Q6_V_hi_W(p); + HVX_Vector s1 = v_src[2 * nvec_64 + 1]; + HVX_Vector v1 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(s1, scale_vec), Q6_Vqf32_vmpy_VsfVsf(m1, slope_vec)); + + hvx_vec_store_a(&v_dst[2 * nvec_64 + 1], (nloe_64 - VLEN_FP32) * sizeof(float), Q6_Vsf_equals_Vqf32(v1)); } } } -static void init_softmax_ctx(struct htp_softmax_context * smctx, struct htp_ops_context * octx) { - const struct htp_tensor * src0 = octx->src[0]; - const struct htp_tensor * src1 = octx->src[1]; - - memset(smctx, 0, sizeof(struct htp_softmax_context)); - - memcpy(&smctx->scale, (float *) octx->op_params, sizeof(float)); - memcpy(&smctx->max_bias, (float *) octx->op_params + 1, sizeof(float)); - - smctx->n_head = src0->ne[2]; - smctx->n_head_log2 = 1u << (uint32_t) floor(log2(smctx->n_head)); - - smctx->m0 = powf(2.0f, -(smctx->max_bias) / smctx->n_head_log2); - smctx->m1 = powf(2.0f, -(smctx->max_bias / 2.0f) / smctx->n_head_log2); - - smctx->use_src1 = (src1 != 0); - smctx->use_f16 = (src1 != 0) && (src1->type == HTP_TYPE_F16); - - smctx->octx = octx; - - // Initialize fastdiv values - const uint32_t ne01 = src0->ne[1]; - const uint32_t ne02 = src0->ne[2]; - - if (ne01 > 0) smctx->fastdiv_ne01 = init_fastdiv_values(ne01); - if (ne02 > 0) smctx->fastdiv_ne02 = init_fastdiv_values(ne02); - - const uint32_t ne12 = src1 ? src1->ne[2] : 1; - const uint32_t ne13 = src1 ? src1->ne[3] : 1; - - if (ne12 > 0) smctx->fastdiv_ne12 = init_fastdiv_values(ne12); - if (ne13 > 0) smctx->fastdiv_ne13 = init_fastdiv_values(ne13); -} - static void hvx_fast_softmax_prep_f32(const uint8_t * restrict src, uint8_t * restrict dst, const int num_elems, float scale, const uint8_t * restrict mask, float slope) { - const uint8_t * restrict src_curr = src; - uint8_t * restrict dst_curr = dst; - const uint8_t * restrict mask_curr = mask; + const HVX_Vector * restrict v_src = (const HVX_Vector *) src; + HVX_Vector * restrict v_dst = (HVX_Vector *) dst; + const HVX_Vector * restrict v_mask = (const HVX_Vector *) mask; HVX_Vector scale_vec = hvx_vec_splat_f32(scale); HVX_Vector slope_vec = hvx_vec_splat_f32(slope); - int step_of_1 = num_elems >> 5; + const int nvec = num_elems / VLEN_FP32; + const int nloe = num_elems % VLEN_FP32; #pragma unroll(4) - for (int i = 0; i < step_of_1; i++) { - HVX_Vector v1 = *(HVX_Vector *) src_curr; - - HVX_Vector v3 = *(HVX_Vector *) mask_curr; + for (int i = 0; i < nvec; i++) { + HVX_Vector v1 = v_src[i]; + HVX_Vector v3 = v_mask[i]; HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec); - HVX_Vector v4 = Q6_Vqf32_vmpy_VsfVsf(v3, slope_vec); - HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, v4); - *(HVX_Vector *) dst_curr = Q6_Vsf_equals_Vqf32(v5); + v_dst[i] = Q6_Vsf_equals_Vqf32(v5); + } - src_curr += VLEN; - dst_curr += VLEN; - mask_curr += VLEN; + if (nloe > 0) { + HVX_Vector v1 = v_src[nvec]; + HVX_Vector v3 = v_mask[nvec]; + + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec); + HVX_Vector v4 = Q6_Vqf32_vmpy_VsfVsf(v3, slope_vec); + HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, v4); + + hvx_vec_store_a(&v_dst[nvec], nloe * sizeof(float), Q6_Vsf_equals_Vqf32(v5)); } } -static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, uint8_t * restrict pad, const int num_elems) { - const HVX_Vector * restrict v_src = (HVX_Vector *) src; - HVX_Vector * restrict v_pad = (HVX_Vector *) pad; +static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, const int num_elems) { + const HVX_Vector * restrict v_src = (const HVX_Vector *) src; HVX_Vector * restrict v_dst = (HVX_Vector *) dst; - HVX_Vector sum_vec = Q6_V_vsplat_R(0x00000000); + const int nvec = num_elems / VLEN_FP32; + const int nloe = num_elems % VLEN_FP32; + HVX_Vector max_vec = hvx_vec_splat_f32(((const float *) src)[0]); - HVX_Vector zero_v = Q6_V_vzero(); - HVX_Vector one_v = hvx_vec_splat_f32(1.0); - int step_of_1 = num_elems >> 5; - - #pragma unroll(4) - for (int i = 0; i < step_of_1; i++) { + #pragma unroll(2) + for (int i = 0; i < nvec; i++) { HVX_Vector v1 = v_src[i]; max_vec = Q6_Vsf_vmax_VsfVsf(max_vec, v1); } - max_vec = hvx_vec_reduce_max_f32(max_vec); // replicated over all lanes + if (nloe > 0) { + HVX_VectorPred q_mask = Q6_Q_vsetq_R(nloe * sizeof(float)); + HVX_Vector neg_inf = hvx_vec_splat_f32(-INFINITY); + HVX_Vector v_tail = Q6_V_vmux_QVV(q_mask, v_src[nvec], neg_inf); + max_vec = Q6_Vsf_vmax_VsfVsf(max_vec, v_tail); + } - #pragma unroll(4) - for (int i = 0; i < step_of_1; i++) { + max_vec = hvx_vec_reduce_max_f32(max_vec); + + HVX_Vector sum_vec = Q6_V_vsplat_R(0x00000000); + + #pragma unroll(2) + for (int i = 0; i < nvec; i++) { HVX_Vector v1 = v_src[i]; HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, max_vec); @@ -193,39 +209,91 @@ static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restric sum_vec = Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(sum_vec), v3); - v_pad[i] = v3; + v_dst[i] = v3; } - sum_vec = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_vec)); // replicated over all lanes + if (nloe > 0) { + HVX_VectorPred q_mask = Q6_Q_vsetq_R(nloe * sizeof(float)); + HVX_Vector v1 = v_src[nvec]; + HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, max_vec); + HVX_Vector v3 = hvx_vec_exp_f32(Q6_Vsf_equals_Vqf32(v2)); + HVX_Vector v3_pad = Q6_V_vmux_QVV(q_mask, v3, Q6_V_vzero()); - HVX_VectorPred pos_sum = Q6_Q_vcmp_gt_VwVw(sum_vec, zero_v); + sum_vec = Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(sum_vec), v3_pad); + v_dst[nvec] = v3_pad; + } + + sum_vec = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_vec)); + + HVX_VectorPred pos_sum = Q6_Q_vcmp_gt_VwVw(sum_vec, Q6_V_vzero()); HVX_Vector v4 = hvx_vec_inverse_f32(sum_vec); - HVX_Vector scale_vec = Q6_V_vmux_QVV(pos_sum, v4, one_v); + HVX_Vector scale_vec = Q6_V_vmux_QVV(pos_sum, v4, hvx_vec_splat_f32(1.0f)); - #pragma unroll(4) - for (int i = 0; i < step_of_1; i++) { - HVX_Vector v1 = v_pad[i]; + #pragma unroll(2) + for (int i = 0; i < nvec; i++) { + HVX_Vector v1 = v_dst[i]; HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec); v_dst[i] = Q6_Vsf_equals_Vqf32(v2); } + + if (nloe > 0) { + HVX_Vector v1 = v_dst[nvec]; + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_vec); + hvx_vec_store_a(&v_dst[nvec], nloe * sizeof(float), Q6_Vsf_equals_Vqf32(v2)); + } } -static float hvx_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, uint8_t * restrict spad, const int num_elems, const float max) { - hvx_sub_scalar_f32(spad, src, max, num_elems); - - hvx_exp_f32(dst, spad, num_elems, false); - return hvx_reduce_sum_f32(dst, num_elems); +static void compute_fast_softmax_f32_nomask( + void * restrict dst, + const void * restrict src0, + const void * restrict mask, + uint32_t ne00, + float scale, + float slope +) { + (void) mask; + (void) slope; + hvx_scale_f32((uint8_t *) dst, (const uint8_t *) src0, ne00, scale); + hvx_fast_softmax_f32((const uint8_t *) dst, (uint8_t *) dst, ne00); } -static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) { - struct htp_softmax_context * smctx = (struct htp_softmax_context *) data; +static void compute_fast_softmax_f32_mask_f32( + void * restrict dst, + const void * restrict src0, + const void * restrict mask, + uint32_t ne00, + float scale, + float slope +) { + hvx_fast_softmax_prep_f32((const uint8_t *) src0, (uint8_t *) dst, ne00, scale, (const uint8_t *) mask, slope); + hvx_fast_softmax_f32((const uint8_t *) dst, (uint8_t *) dst, ne00); +} + +static void compute_fast_softmax_f32_mask_f16( + void * restrict dst, + const void * restrict src0, + const void * restrict mask, + uint32_t ne00, + float scale, + float slope +) { + hvx_fast_softmax_prep_f16((const uint8_t *) src0, (uint8_t *) dst, ne00, scale, (const uint8_t *) mask, slope); + hvx_fast_softmax_f32((const uint8_t *) dst, (uint8_t *) dst, ne00); +} + +static const softmax_compute_fn_t softmax_kernels[HTP_SOFTMAX_KERNEL_COUNT] = { + [HTP_SOFTMAX_KERNEL_NOMASK] = compute_fast_softmax_f32_nomask, + [HTP_SOFTMAX_KERNEL_MASK_F32] = compute_fast_softmax_f32_mask_f32, + [HTP_SOFTMAX_KERNEL_MASK_F16] = compute_fast_softmax_f32_mask_f16, +}; + +static void softmax_thread_dma(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + const struct htp_softmax_context * smctx = (const struct htp_softmax_context *) data; struct htp_ops_context * octx = smctx->octx; - const struct htp_tensor * src0 = octx->src[0]; - const struct htp_tensor * src1 = octx->src[1]; const struct htp_tensor * dst = octx->dst; - - htp_softmax_preamble3; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; const uint32_t src0_nrows = smctx->nrows; const uint32_t src0_nrows_per_thread = smctx->src0_nrows_per_thread; @@ -233,122 +301,213 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) { const uint32_t src0_start_row = smctx->row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, smctx->row_start + src0_nrows); - // no work for this thread if (src0_start_row >= src0_end_row) { return; } - int is_aligned = 1; - int opt_path = 0; + const dma_addr_t data_src0 = smctx->data_src0; + const dma_addr_t data_dst = smctx->data_dst; - if (!hex_is_aligned((void *) src0->data, VLEN) || !hex_is_aligned((void *) dst->data, VLEN)) { - is_aligned = 0; - FARF(HIGH, "softmax-f32: unaligned addresses in elementwise op, possibly slower execution\n"); + const size_t src0_row_size = src0->ne[0] * sizeof(float); + const size_t dst_row_size = src0->ne[0] * sizeof(float); + + uint8_t * src0_vtcm_base = smctx->vtcm_src0 + (ith * smctx->vtcm_src0_size_per_thread); + uint8_t * dst_vtcm_base = smctx->vtcm_dst + (ith * smctx->vtcm_dst_size_per_thread); + + const size_t src0_vtcm_half = smctx->src0_spad_half_size; + const size_t dst_vtcm_half = smctx->dst_spad_half_size; + + dma_queue * dma_q = octx->ctx->dma[ith]; + + for (uint32_t r = src0_start_row, idx = 0; r < src0_end_row && idx < 2; r++, idx++) { + dma_addr_t cur_dst = data_dst + r * dst_row_size; + dma_addr_t cur_src0 = data_src0 + r * src0_row_size; + void * d_spad = dst_vtcm_base + idx * dst_vtcm_half; + void * s_spad = src0_vtcm_base + idx * src0_vtcm_half; + + dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad), + dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 0); + dma_queue_push(dma_q, dma_make_data(s_spad, cur_src0), + smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1); } - // Only use the fast path when aligned AND row size is multiple of VLEN (128 bytes) - // The fast path (hvx_fast_softmax_f32) doesn't handle tail elements - // The non-opt path uses hvx_softmax_f32 which properly handles all sizes via its helper functions - if ((1 == is_aligned) && !(nb01 & (VLEN - 1))) { - opt_path = 1; - } - - uint8_t * src0_spad_data = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * src1_spad_data = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread); - uint8_t * dst_spad_data = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); - - float * wp0 = (float *) src0_spad_data; - float * wp1 = (float *) src1_spad_data; - float * wp2 = (float *) dst_spad_data; - - uint32_t prev_i2 = (uint32_t)-1; - float slope = 1.0f; - - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_start_row); + softmax_compute_fn_t compute = (softmax_compute_fn_t) smctx->compute; + const uint32_t ne00 = src0->ne[0]; for (uint32_t r = src0_start_row; r < src0_end_row; ++r) { - uint32_t i1 = fastmodulo(r, ne01, &smctx->fastdiv_ne01); - uint32_t r_div_ne01 = fastdiv(r, &smctx->fastdiv_ne01); - uint32_t i2 = fastmodulo(r_div_ne01, ne02, &smctx->fastdiv_ne02); - uint32_t i3 = fastdiv(r_div_ne01, &smctx->fastdiv_ne02); + void * d_spad = (void *) dma_queue_pop(dma_q).src; + void * s_spad = (void *) dma_queue_pop(dma_q).dst; - // Map to original logic indices - // i01 = i1 - // i02 = i2 - // i03 = i3 + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); + compute(d_spad, s_spad, NULL, ne00, smctx->scale, 1.0f); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); - const uint32_t i11 = i1; - // const uint32_t i12 = i2 % ne12; - // const uint32_t i13 = i3 % ne13; + dma_addr_t cur_dst = data_dst + r * dst_row_size; + dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad), + dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 1); - uint32_t i12, i13; - if (ne12 == ne02) { - i12 = i2; - } else { - i12 = fastmodulo(i2, ne12, &smctx->fastdiv_ne12); - } - - if (ne13 == ne03) { - i13 = i3; - } else { - i13 = fastmodulo(i3, ne13, &smctx->fastdiv_ne13); - } - - // ALiBi - if (i2 != prev_i2) { - const uint32_t h = i2; // head - slope = (smctx->max_bias > 0.0f) ? h < smctx->n_head_log2 ? powf(smctx->m0, h + 1) : powf(smctx->m1, 2 * (h - smctx->n_head_log2) + 1) : 1.0f; - prev_i2 = i2; - } - - float * sp = (float *) ((char *) src0->data + i1 * nb01 + i2 * nb02 + i3 * nb03); - float * dp = (float *) ((char *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3); - - // broadcast the mask across rows - __fp16 * mp_f16 = (smctx->use_src1) ? (__fp16 *) ((char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13) : NULL; - float * mp_f32 = (smctx->use_src1) ? (float *) ((char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13) : NULL; - - if ((1 == opt_path) && (mp_f32) && !(smctx->use_f16)) { - hvx_fast_softmax_prep_f32((const uint8_t *) sp, (uint8_t *) wp0, ne00, smctx->scale, (const uint8_t *) mp_f32, slope); - hvx_fast_softmax_f32((const uint8_t *) wp0, (uint8_t *) dp, (uint8_t *) wp1, ne00); - } else if (1 == opt_path) { - hvx_scale_f32((uint8_t *) wp0, (const uint8_t *) sp, ne00, smctx->scale); - apply_mask(wp0, mp_f32, mp_f16, ne00, slope, smctx->use_f16); - hvx_fast_softmax_f32((const uint8_t *) wp0, (uint8_t *) dp, (uint8_t *) wp1, ne00); - } else { - // Non-optimized path: uses HVX helper functions that properly handle all tensor sizes - // including non-multiples of 32 (the HVX vector lane count for f32) - hvx_scale_f32((uint8_t *) wp0, (const uint8_t *) sp, ne00, smctx->scale); - apply_mask(wp0, mp_f32, mp_f16, ne00, slope, smctx->use_f16); - float max = hvx_reduce_max_f32((const uint8_t *) wp0, ne00); - float sum = hvx_softmax_f32((const uint8_t *) wp0, (uint8_t *) wp2, (uint8_t *) wp1, ne00, max); - sum = sum > 0.0 ? (1.0 / sum) : 1; - hvx_scale_f32((uint8_t *) dp, (const uint8_t *) wp2, ne00, sum); + const uint32_t next_r = r + 2; + if (next_r < src0_end_row) { + dma_addr_t next_src0 = data_src0 + next_r * src0_row_size; + dma_queue_push(dma_q, dma_make_data(s_spad, next_src0), + smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1); } } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_start_row); - - FARF(HIGH, "softmax-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u : opt %u f16 %u\n", ith, nth, - ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13, - ne0, ne1, ne2, ne3, opt_path, smctx->use_f16); + dma_queue_flush(dma_q); } -static int execute_op_softmax_f32(struct htp_ops_context * octx) { - int err = HTP_STATUS_OK; - +static void softmax_thread_mask_dma(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + const struct htp_softmax_context * smctx = (const struct htp_softmax_context *) data; + struct htp_ops_context * octx = smctx->octx; const struct htp_tensor * src0 = octx->src[0]; const struct htp_tensor * src1 = octx->src[1]; const struct htp_tensor * dst = octx->dst; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const uint32_t src0_nrows = smctx->nrows; + const uint32_t src0_nrows_per_thread = smctx->src0_nrows_per_thread; + + const uint32_t src0_start_row = smctx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, smctx->row_start + src0_nrows); + + if (src0_start_row >= src0_end_row) { + return; + } + + const dma_addr_t data_src0 = smctx->data_src0; + const dma_addr_t data_src1 = smctx->data_src1; + const dma_addr_t data_dst = smctx->data_dst; + + const size_t src0_row_size = src0->ne[0] * sizeof(float); + const size_t dst_row_size = src0->ne[0] * sizeof(float); + const size_t mask_row_size = smctx->use_f16 ? (src1->ne[0] * sizeof(__fp16)) : (src1->ne[0] * sizeof(float)); + + uint8_t * src0_vtcm_base = smctx->vtcm_src0 + (ith * smctx->vtcm_src0_size_per_thread); + uint8_t * src1_vtcm_base = smctx->vtcm_src1 + (ith * smctx->vtcm_src1_size_per_thread); + uint8_t * dst_vtcm_base = smctx->vtcm_dst + (ith * smctx->vtcm_dst_size_per_thread); + + const size_t src0_vtcm_half = smctx->src0_spad_half_size; + const size_t src1_vtcm_half = smctx->src1_spad_half_size; + const size_t dst_vtcm_half = smctx->dst_spad_half_size; + + const uint32_t nb11 = src1->nb[1]; + const uint32_t nb12 = src1->nb[2]; + const uint32_t nb13 = src1->nb[3]; + + const uint32_t ne00 = src0->ne[0]; + const uint32_t ne01 = src0->ne[1]; + const uint32_t ne02 = src0->ne[2]; + const uint32_t ne03 = src0->ne[3]; + const uint32_t ne12 = src1->ne[2]; + const uint32_t ne13 = src1->ne[3]; + + const struct fastdiv_values * div_ne01 = &smctx->div_ne01; + const struct fastdiv_values * div_ne02 = &smctx->div_ne02; + const struct fastdiv_values * div_ne12 = &smctx->div_ne12; + const struct fastdiv_values * div_ne13 = &smctx->div_ne13; + + dma_queue * dma_q = octx->ctx->dma[ith]; + + for (uint32_t r = src0_start_row, idx = 0; r < src0_end_row && idx < 2; r++, idx++) { + dma_addr_t cur_dst = data_dst + r * dst_row_size; + dma_addr_t cur_src0 = data_src0 + r * src0_row_size; + + uint32_t i1 = fastmodulo(r, ne01, div_ne01); + uint32_t r_div_ne01 = fastdiv(r, div_ne01); + uint32_t i2 = fastmodulo(r_div_ne01, ne02, div_ne02); + uint32_t i3 = fastdiv(r_div_ne01, div_ne02); + uint32_t i12 = (ne12 == ne02) ? i2 : fastmodulo(i2, ne12, div_ne12); + uint32_t i13 = (ne13 == ne03) ? i3 : fastmodulo(i3, ne13, div_ne13); + dma_addr_t cur_src1 = data_src1 + i1 * nb11 + i12 * nb12 + i13 * nb13; + + void * d_spad = dst_vtcm_base + idx * dst_vtcm_half; + void * s_spad = src0_vtcm_base + idx * src0_vtcm_half; + void * m_spad = src1_vtcm_base + idx * src1_vtcm_half; + + dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad), + dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 0); + dma_queue_push(dma_q, dma_make_data(s_spad, cur_src0), + smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1); + dma_queue_push(dma_q, dma_make_data(m_spad, cur_src1), + smctx->src1_row_size_aligned, mask_row_size, mask_row_size, 1); + } + + softmax_compute_fn_t compute = (softmax_compute_fn_t) smctx->compute; + const bool has_bias = smctx->max_bias > 0.0f; + uint32_t prev_i2 = (uint32_t)-1; + float slope = 1.0f; + + for (uint32_t r = src0_start_row; r < src0_end_row; ++r) { + void * d_spad = (void *) (uintptr_t) dma_queue_pop(dma_q).src; + void * s_spad = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + void * m_spad = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + + if (has_bias) { + uint32_t r_div_ne01 = fastdiv(r, div_ne01); + uint32_t i2 = fastmodulo(r_div_ne01, ne02, div_ne02); + if (i2 != prev_i2) { + slope = smctx->slopes[i2]; + prev_i2 = i2; + } + } + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); + compute(d_spad, s_spad, m_spad, ne00, smctx->scale, slope); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); + + dma_addr_t cur_dst = data_dst + r * dst_row_size; + dma_queue_push(dma_q, dma_make_data(cur_dst, d_spad), + dst_row_size, smctx->dst_row_size_aligned, dst_row_size, 1); + + const uint32_t next_r = r + 2; + if (next_r < src0_end_row) { + dma_addr_t next_src0 = data_src0 + next_r * src0_row_size; + + uint32_t ni1 = fastmodulo(next_r, ne01, div_ne01); + uint32_t nr_div_ne01 = fastdiv(next_r, div_ne01); + uint32_t ni2 = fastmodulo(nr_div_ne01, ne02, div_ne02); + uint32_t ni3 = fastdiv(nr_div_ne01, div_ne02); + uint32_t ni12 = (ne12 == ne02) ? ni2 : fastmodulo(ni2, ne12, div_ne12); + uint32_t ni13 = (ne13 == ne03) ? ni3 : fastmodulo(ni3, ne13, div_ne13); + dma_addr_t next_src1 = data_src1 + ni1 * nb11 + ni12 * nb12 + ni13 * nb13; + + dma_queue_push(dma_q, dma_make_data(s_spad, next_src0), + smctx->src0_row_size_aligned, src0_row_size, src0_row_size, 1); + dma_queue_push(dma_q, dma_make_data(m_spad, next_src1), + smctx->src1_row_size_aligned, mask_row_size, mask_row_size, 1); + } + } + + dma_queue_flush(dma_q); +} + +static int execute_op_softmax_f32(struct htp_ops_context * octx) { + const struct htp_tensor * src0 = octx->src[0]; + const struct htp_tensor * dst = octx->dst; - struct htp_softmax_context smctx; const char * op_type = "softmax-f32"; - init_softmax_ctx(&smctx, octx); + const struct htp_softmax_kernel_params * kparams = + (const struct htp_softmax_kernel_params *) octx->kernel_params; + + if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) { + return HTP_STATUS_INVAL_PARAMS; + } + + if (kparams->kernel_id >= HTP_SOFTMAX_KERNEL_COUNT) { + return HTP_STATUS_INVAL_PARAMS; + } + + if (octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) { + FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %u\n", + op_type, octx->ctx->vtcm_size, kparams->vtcm_size); + return HTP_STATUS_VTCM_TOO_SMALL; + } const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; - const size_t elem_size = sizeof(float); + const size_t elem_size = sizeof(float); const size_t dst_row_size = dst->nb[1]; uint32_t row_start = 0; @@ -357,9 +516,13 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) { if (octx->ctx->mdev.count > 1) { uint32_t rows_per_chunk = 0; htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_row_size, &rows_per_chunk); - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition( + src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); row_start = range.start; nrows = range.count; + if (nrows < octx->n_threads) { + htp_ops_context_set_n_threads(octx, nrows ? nrows : 1); + } } if (nrows == 0) { @@ -367,50 +530,71 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) { } const uint32_t n_threads = octx->n_threads; + uint8_t * const vtcm_base = (uint8_t *) octx->ctx->vtcm_base; - smctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); - smctx.row_start = row_start; - smctx.nrows = nrows; + const uint32_t off_src0 = 0; + const uint32_t off_dst = off_src0 + kparams->vtcm_src0_size_per_thread * kparams->n_threads; + const uint32_t off_src1 = off_dst + kparams->vtcm_dst_size_per_thread * kparams->n_threads; - const size_t src0_row_size = src0->nb[1]; - const size_t src1_row_size = src0_row_size; + struct htp_softmax_context smctx = { + .octx = octx, + .kparams = kparams, + .compute = (void *) softmax_kernels[kparams->kernel_id], - // VTCM scratchpads for all tensors - // 4 rows per thread, padded to HVX vector size - octx->src0_spad.size_per_thread = hex_round_up(4 * src0_row_size, 128); - octx->src1_spad.size_per_thread = hex_round_up(4 * src1_row_size, 128); - octx->dst_spad.size_per_thread = hex_round_up(4 * dst_row_size, 128); + .data_src0 = src0->data, + .data_src1 = kparams->use_src1 ? octx->src[1]->data : 0, + .data_dst = dst->data, - octx->src0_spad.size = octx->src0_spad.size_per_thread * n_threads; - octx->src1_spad.size = octx->src1_spad.size_per_thread * n_threads; - octx->dst_spad.size = octx->dst_spad.size_per_thread * n_threads; + .vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, off_src0), + .vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, off_dst), + .vtcm_src1 = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, vtcm_base, off_src1, kparams->use_src1), - size_t spad_size = octx->src0_spad.size + octx->src1_spad.size + octx->dst_spad.size; + .vtcm_src0_size_per_thread = kparams->vtcm_src0_size_per_thread, + .vtcm_src1_size_per_thread = kparams->vtcm_src1_size_per_thread, + .vtcm_dst_size_per_thread = kparams->vtcm_dst_size_per_thread, - if (src1) { - FARF(HIGH, "%s: %ux%ux%ux%u x %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], src1->ne[0], src1->ne[1], src1->ne[2], - src1->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); - } else { - 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); + .src0_spad_half_size = kparams->src0_spad_half_size, + .src1_spad_half_size = kparams->src1_spad_half_size, + .dst_spad_half_size = kparams->dst_spad_half_size, + + .src0_row_size_aligned = kparams->src0_row_size_aligned, + .src1_row_size_aligned = kparams->src1_row_size_aligned, + .dst_row_size_aligned = kparams->dst_row_size_aligned, + + .use_f16 = kparams->use_f16 != 0, + + .n_head = kparams->n_head, + .n_head_log2 = kparams->n_head_log2, + + .scale = kparams->scale, + .max_bias = kparams->max_bias, + .m0 = kparams->m0, + .m1 = kparams->m1, + + .div_ne01 = kparams->div_ne01, + .div_ne02 = kparams->div_ne02, + .div_ne12 = kparams->div_ne12, + .div_ne13 = kparams->div_ne13, + + .src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div), + .row_start = row_start, + .nrows = nrows, + }; + + if (kparams->max_bias > 0.0f && kparams->use_src1) { + if (kparams->n_head > 512) { + return HTP_STATUS_INVAL_PARAMS; + } + for (uint32_t h = 0; h < kparams->n_head; h += 32) { + HVX_Vector v_slopes = hvx_alibi_slopes(h, 1, kparams->n_head_log2, kparams->m0, kparams->m1); + hvx_vmem(&smctx.slopes[h]) = v_slopes; + } } - // Make sure the reserved vtcm size is sufficient - if (octx->ctx->vtcm_size < spad_size) { - FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, spad_size); - return HTP_STATUS_VTCM_TOO_SMALL; - } + work_queue_func_t task_func = kparams->use_src1 ? softmax_thread_mask_dma : softmax_thread_dma; + work_queue_run(octx->ctx->work_queue, task_func, &smctx, n_threads); - octx->src0_spad.data = octx->ctx->vtcm_base; octx->src0_spad.src = NULL; - octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL; - octx->dst_spad.data = octx->src1_spad.data + octx->src1_spad.size; octx->dst_spad.src = NULL; - - work_queue_run(octx->ctx->work_queue, softmax_job_f32, &smctx, n_threads); - - return err; + return HTP_STATUS_OK; } int op_softmax(struct htp_ops_context * octx) { diff --git a/ggml/src/ggml-hexagon/htp/softmax-ops.h b/ggml/src/ggml-hexagon/htp/softmax-ops.h new file mode 100644 index 000000000..8d976adb8 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/softmax-ops.h @@ -0,0 +1,106 @@ +#ifndef HTP_SOFTMAX_OPS_H +#define HTP_SOFTMAX_OPS_H + +#include +#include +#include +#include +#include "hex-fastdiv.h" +#include "hex-common.h" + +enum htp_softmax_kernel_id { + HTP_SOFTMAX_KERNEL_NOMASK = 0, + HTP_SOFTMAX_KERNEL_MASK_F32, + HTP_SOFTMAX_KERNEL_MASK_F16, + HTP_SOFTMAX_KERNEL_COUNT, +}; + +struct htp_softmax_kernel_params { + uint32_t n_threads; + uint32_t src0_nrows; + uint32_t src0_nrows_per_thread; + uint32_t vtcm_size; + + uint32_t vtcm_src0_size_per_thread; + uint32_t vtcm_src1_size_per_thread; + uint32_t vtcm_dst_size_per_thread; + + uint32_t src0_row_size_aligned; + uint32_t src1_row_size_aligned; + uint32_t dst_row_size_aligned; + + uint32_t src0_spad_half_size; + uint32_t src1_spad_half_size; + uint32_t dst_spad_half_size; + + uint32_t n_head; + uint32_t n_head_log2; + uint32_t use_src1; + uint32_t use_f16; + uint32_t kernel_id; + + float scale; + float max_bias; + float m0; + float m1; + + struct fastdiv_values div_ne01; + struct fastdiv_values div_ne02; + struct fastdiv_values div_ne12; + struct fastdiv_values div_ne13; +}; + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_softmax_kernel_params) <= 128, "htp_softmax_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_softmax_kernel_params) <= 128, "htp_softmax_kernel_params is too large for kernel_params blob"); +#endif + +struct htp_softmax_vtcm_layout { + size_t total_bytes; + size_t off_src0; + size_t off_dst; + size_t off_src1; + + size_t src0_bytes_per_thread; + size_t dst_bytes_per_thread; + size_t src1_bytes_per_thread; + + size_t src0_spad_half_size; + size_t dst_spad_half_size; + size_t src1_spad_half_size; +}; + +static inline void htp_softmax_vtcm_layout_build( + struct htp_softmax_vtcm_layout * layout, + uint32_t ne00, + uint32_t ne10, + bool use_src1, + bool use_f16, + uint32_t n_threads +) { + size_t src0_row_size = ne00 * sizeof(float); + size_t dst_row_size = ne00 * sizeof(float); + size_t src1_row_size = use_src1 ? (ne10 * (use_f16 ? 2 : 4)) : 0; + + size_t src0_row_size_aligned = hex_round_up(src0_row_size, 128); + size_t dst_row_size_aligned = hex_round_up(dst_row_size, 128); + size_t src1_row_size_aligned = use_src1 ? hex_round_up(src1_row_size, 128) : 0; + + layout->src0_spad_half_size = src0_row_size_aligned; + layout->dst_spad_half_size = dst_row_size_aligned; + layout->src1_spad_half_size = src1_row_size_aligned; + + // Double buffering: 2 half-buffers per thread + layout->src0_bytes_per_thread = src0_row_size_aligned * 2; + layout->dst_bytes_per_thread = dst_row_size_aligned * 2; + layout->src1_bytes_per_thread = src1_row_size_aligned * 2; + + layout->off_src0 = 0; + layout->off_dst = layout->off_src0 + layout->src0_bytes_per_thread * n_threads; + layout->off_src1 = layout->off_dst + layout->dst_bytes_per_thread * n_threads; + + layout->total_bytes = layout->off_src1 + layout->src1_bytes_per_thread * n_threads; +} + +#endif // HTP_SOFTMAX_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/solve-tri-ops.c b/ggml/src/ggml-hexagon/htp/solve-tri-ops.c index 847a78712..182982fcd 100644 --- a/ggml/src/ggml-hexagon/htp/solve-tri-ops.c +++ b/ggml/src/ggml-hexagon/htp/solve-tri-ops.c @@ -218,8 +218,8 @@ int op_solve_tri(struct htp_ops_context * octx) { return HTP_STATUS_INVAL_PARAMS; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; } const uint32_t k = src1->ne[0]; diff --git a/ggml/src/ggml-hexagon/htp/ssm-conv.c b/ggml/src/ggml-hexagon/htp/ssm-conv.c index bef142536..931aa406e 100644 --- a/ggml/src/ggml-hexagon/htp/ssm-conv.c +++ b/ggml/src/ggml-hexagon/htp/ssm-conv.c @@ -14,124 +14,22 @@ #define GGML_COMMON_DECL_C #include "ggml-common.h" #include "htp-ctx.h" -#include "hex-dma.h" +#include "dma-queue.h" #include "hex-profile.h" #include "htp-ops.h" #include "htp-tensor.h" #include "hvx-utils.h" - -#define htp_ssm_conv_tensors_preamble \ - const struct htp_tensor * restrict src0 = octx->src[0]; \ - const struct htp_tensor * restrict src1 = octx->src[1]; \ - const struct htp_tensor * restrict dst = octx->dst; \ - struct htp_spad * restrict src0_spad = &octx->src0_spad; \ - struct htp_spad * restrict src1_spad = &octx->src1_spad; \ - struct htp_spad * restrict dst_spad = &octx->dst_spad; \ - \ - const uint32_t ne00 = src0->ne[0]; \ - const uint32_t ne01 = src0->ne[1]; \ - const uint32_t ne02 = src0->ne[2]; \ - const uint32_t ne03 = src0->ne[3]; \ - \ - const uint32_t ne10 = src1->ne[0]; \ - const uint32_t ne11 = src1->ne[1]; \ - const uint32_t ne12 = src1->ne[2]; \ - const uint32_t ne13 = src1->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 = src0->nb[0]; \ - const uint32_t nb01 = src0->nb[1]; \ - const uint32_t nb02 = src0->nb[2]; \ - const uint32_t nb03 = src0->nb[3]; \ - \ - const uint32_t nb10 = src1->nb[0]; \ - const uint32_t nb11 = src1->nb[1]; \ - const uint32_t nb12 = src1->nb[2]; \ - const uint32_t nb13 = src1->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]; +#include "ssm-conv.h" struct htp_ssm_conv_context { - struct htp_ops_context * octx; - uint32_t nrows_per_thread; - uint32_t d_inner_tile; - uint64_t t_start; - uint32_t row_start; - uint32_t nrows; + struct htp_ops_context * octx; + const struct htp_ssm_conv_kernel_params * kparams; + uint32_t nrows_per_thread; + uint32_t d_inner_tile; + uint32_t row_start; + uint32_t nrows; }; -#define htp_ssm_conv_preamble \ - struct htp_ssm_conv_context * scctx = (struct htp_ssm_conv_context *) data; \ - struct htp_ops_context * octx = scctx->octx; \ - htp_ssm_conv_tensors_preamble; \ - dma_queue * dma_queue = octx->ctx->dma[ith]; - -// Scalar FP32 SSM_CONV implementation -static void ssm_conv_thread_f32_f32(unsigned int nth, unsigned int ith, void *data) { - htp_ssm_conv_preamble; - - const uint32_t d_conv = src1->ne[0]; - const uint32_t d_inner = src0->ne[1]; - const uint32_t n_t = dst->ne[1]; - const uint32_t n_s = dst->ne[2]; - - const uint32_t src0_stride_inner = src0->nb[1] / sizeof(float); // stride for inner dimension - const uint32_t src0_stride_seq = src0->nb[2] / sizeof(float); // stride for sequence dimension - const uint32_t src1_stride_inner = src1->nb[1] / sizeof(float); // stride for inner dimension - const uint32_t dst_stride_token = dst->nb[1] / sizeof(float); // stride for token dimension - const uint32_t dst_stride_seq = dst->nb[2] / sizeof(float); // stride for sequence dimension - - const float * src0_data = (const float *) src0->data; - const float * src1_data = (const float *) src1->data; - float * dst_data = (float *) dst->data; - - // Calculate row range for this thread - const uint32_t d_inner_per_thread = scctx->nrows_per_thread; - const uint32_t d_inner_start = scctx->row_start + d_inner_per_thread * ith; - const uint32_t d_inner_end = MIN(d_inner_start + d_inner_per_thread, scctx->row_start + scctx->nrows); - - // No work for this thread - if (d_inner_start >= d_inner_end) { - return; - } - - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) d_inner_start); - - for (uint32_t i3 = 0; i3 < n_s; ++i3) { - for (uint32_t i2 = 0; i2 < n_t; ++i2) { - for (uint32_t i1 = d_inner_start; i1 < d_inner_end; ++i1) { - float sumf = 0.0f; - - for (uint32_t i0 = 0; i0 < d_conv; ++i0) { - const uint32_t src0_idx = (i2 + i0) + i1 * src0_stride_inner + i3 * src0_stride_seq; - const uint32_t src1_idx = i0 + i1 * src1_stride_inner; - - sumf += src0_data[src0_idx] * src1_data[src1_idx]; - } - - const uint32_t dst_idx = i1 + i2 * dst_stride_token + i3 * dst_stride_seq; - dst_data[dst_idx] = sumf; - } - } - } - - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) d_inner_end); - - FARF(HIGH, "ssm-conv-f32 %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n", - ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], d_inner_start, d_inner_end, - src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1], - dst->ne[2], dst->ne[3]); -} - - // In-register 32x32 fp32 transpose using std 5-stage HVX vshuff butterfly. static inline void hvx_transpose_32x32_f32(HVX_Vector m[32]) { HVX_Vector tmp[32]; @@ -181,40 +79,69 @@ static inline void hvx_transpose_32x32_f32(HVX_Vector m[32]) { } } -// HVX FP32 SSM_CONV implementation - channel-vectorized HVX kernel with src0/src1 -// transposed into VTCM. -// -// VTCM layouts (per thread): -// src1_T : {d_inner_stride, d_conv} - staged once per launch (small). -// src0_T : {d_inner_tile, ncs} - staged per d_inner-tile. -// -// d_inner_tile is chosen so that per-thread VTCM stays under the budget. -// Each thread iterates ceil(d_inner_per_thread d_inner_tile) tiles serially. -#define HTP_SSM_CONV_VTCM_BUDGET (1u << 20) // 1 MiB per thread +// HVX deinterleave for d_conv == 4: channel-major raw VTCM -> tap-major T VTCM +static inline void hvx_ssm_conv_unpack_to_T_4(const float * raw, float * T, uint32_t d_inner_per_thread, uint32_t d_inner_stride) { + for (uint32_t cb = 0; cb < d_inner_per_thread; cb += VLEN_FP32) { + HVX_Vector v0 = *(const HVX_Vector *)(raw + (cb + 0) * 4); + HVX_Vector v1 = *(const HVX_Vector *)(raw + (cb + 8) * 4); + HVX_Vector v2 = *(const HVX_Vector *)(raw + (cb + 16) * 4); + HVX_Vector v3 = *(const HVX_Vector *)(raw + (cb + 24) * 4); + + HVX_VectorPair p01 = Q6_W_vdeal_VVR(v1, v0, -4); + HVX_VectorPair p23 = Q6_W_vdeal_VVR(v3, v2, -4); + + HVX_VectorPair p_w02 = Q6_W_vdeal_VVR(Q6_V_lo_W(p23), Q6_V_lo_W(p01), -4); + HVX_VectorPair p_w13 = Q6_W_vdeal_VVR(Q6_V_hi_W(p23), Q6_V_hi_W(p01), -4); + + *(HVX_Vector *)(T + 0 * d_inner_stride + cb) = Q6_V_lo_W(p_w02); + *(HVX_Vector *)(T + 1 * d_inner_stride + cb) = Q6_V_lo_W(p_w13); + *(HVX_Vector *)(T + 2 * d_inner_stride + cb) = Q6_V_hi_W(p_w02); + *(HVX_Vector *)(T + 3 * d_inner_stride + cb) = Q6_V_hi_W(p_w13); + } +} + +// HVX transpose for general d_conv <= 32: channel-major raw VTCM -> tap-major T VTCM +static inline void hvx_ssm_conv_unpack_to_T_gen(const float * raw, float * T, uint32_t d_inner_per_thread, uint32_t d_inner_stride, uint32_t d_conv) { + uint32_t __attribute__((aligned(VLEN))) mask_buf[VLEN_FP32] = { 0 }; + for (uint32_t j = 0; j < d_conv; ++j) { + mask_buf[j] = 0xFFFFFFFF; + } + const HVX_Vector mask = *(const HVX_Vector *) mask_buf; + + for (uint32_t cb = 0; cb < d_inner_per_thread; cb += VLEN_FP32) { + const uint32_t cb_n = MIN(VLEN_FP32, d_inner_per_thread - cb); + HVX_Vector sub[32]; + for (uint32_t r = 0; r < cb_n; ++r) { + const float * ch_ptr = raw + (cb + r) * d_conv; + sub[r] = Q6_V_vand_VV(*(const HVX_UVector *) ch_ptr, mask); + } + for (uint32_t r = cb_n; r < 32; ++r) { + sub[r] = hvx_vec_splat_f32(0.0f); + } + + hvx_transpose_32x32_f32(sub); -// Scalar transpose: src1 {d_conv, d_inner} (DDR) -> {d_inner_stride, d_conv} (VTCM) -static inline void transpose_src1(const float * src1_data, - uint32_t src1_stride_inner, - uint32_t i1_off, - uint32_t d_inner_per_thread, - uint32_t d_inner_stride, - uint32_t d_conv, - float * src1_T) { - for (uint32_t i = 0; i < d_inner_per_thread; ++i) { - const float * src_row = src1_data + (i1_off + i) * src1_stride_inner; for (uint32_t j = 0; j < d_conv; ++j) { - src1_T[j * d_inner_stride + i] = src_row[j]; + *(HVX_Vector *)(T + j * d_inner_stride + cb) = sub[j]; } } } -// HVX 32x32 src0 transpose: src0 {ncs, d_inner} (DDR) -> src0_T {d_inner_tile, ncs} (VTCM) +static inline void hvx_ssm_conv_unpack_to_T(const float * raw, float * T, uint32_t d_inner_per_thread, uint32_t d_inner_stride, uint32_t d_conv) { + if (d_conv == 4 && (d_inner_per_thread % VLEN_FP32 == 0)) { + hvx_ssm_conv_unpack_to_T_4(raw, T, d_inner_per_thread, d_inner_stride); + } else { + hvx_ssm_conv_unpack_to_T_gen(raw, T, d_inner_per_thread, d_inner_stride, d_conv); + } +} + +// HVX 32x32 src0 transpose for prefill: src0 {tile_n, ncs} (VTCM) -> src0_T {ncs, d_inner_tile} (VTCM) static inline void transpose_src0_block(const float * src0_block, uint32_t ncs, uint32_t cb_n, uint32_t d_inner_tile, float * src0_T_block_dst, - uint32_t cb /* dst column offset */) { + uint32_t cb) { const uint32_t T_TILE = VLEN_FP32; HVX_Vector __attribute__((aligned(VLEN))) sub[32]; @@ -222,20 +149,15 @@ static inline void transpose_src0_block(const float * src0_block, for (uint32_t t0 = 0; t0 < ncs; t0 += T_TILE) { const uint32_t t_n = MIN(T_TILE, ncs - t0); - // Load 32 rows (channels) of T_TILE samples; pad missing channels with zeros. + uint32_t __attribute__((aligned(VLEN))) mask_buf[VLEN_FP32] = { 0 }; + for (uint32_t k = 0; k < t_n; ++k) { + mask_buf[k] = 0xFFFFFFFF; + } + const HVX_Vector mask = *(const HVX_Vector *) mask_buf; + for (uint32_t r = 0; r < cb_n; ++r) { const float * src_row = src0_block + r * ncs + t0; - if (t_n == T_TILE) { - sub[r] = *(const HVX_UVector *) src_row; - } else { - HVX_Vector v = hvx_vec_splat_f32(0.0f); - hvx_vec_store_u(&v, t_n * sizeof(float), hvx_vec_splat_f32(0.0f)); - - float __attribute__((aligned(VLEN))) tmp[VLEN_FP32] = { 0 }; - for (uint32_t k = 0; k < t_n; ++k) tmp[k] = src_row[k]; - v = *(const HVX_Vector *) tmp; - sub[r] = v; - } + sub[r] = (t_n == T_TILE) ? *(const HVX_UVector *) src_row : Q6_V_vand_VV(*(const HVX_UVector *) src_row, mask); } for (uint32_t r = cb_n; r < T_TILE; ++r) { sub[r] = hvx_vec_splat_f32(0.0f); @@ -243,8 +165,6 @@ static inline void transpose_src0_block(const float * src0_block, hvx_transpose_32x32_f32(sub); - // Store transposed sub-tile to src0_T at offsets (t0 + j) * d_inner_tile + cb. - // Only write the valid t_n rows of the transposed result. for (uint32_t r = 0; r < t_n; ++r) { float * dst = src0_T_block_dst + (t0 + r) * d_inner_tile + cb; if (cb_n == T_TILE) { @@ -256,20 +176,21 @@ static inline void transpose_src0_block(const float * src0_block, } } -static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void *data) { - htp_ssm_conv_preamble; +// Single-row decode worker (n_t == 1) +static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void * data) { + struct htp_ssm_conv_context * scctx = (struct htp_ssm_conv_context *) data; + struct htp_ops_context * octx = scctx->octx; + const struct htp_ssm_conv_kernel_params * kparams = scctx->kparams; - const uint32_t d_conv = src1->ne[0]; - const uint32_t d_inner = src0->ne[1]; - const uint32_t n_t = dst->ne[1]; - const uint32_t n_s = dst->ne[2]; - const uint32_t ncs = src0->ne[0]; + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict src1 = octx->src[1]; + const struct htp_tensor * restrict dst = octx->dst; - const uint32_t src0_stride_inner = src0->nb[1] / sizeof(float); - const uint32_t src0_stride_seq = src0->nb[2] / sizeof(float); - const uint32_t src1_stride_inner = src1->nb[1] / sizeof(float); - const uint32_t dst_stride_token = dst->nb[1] / sizeof(float); - const uint32_t dst_stride_seq = dst->nb[2] / sizeof(float); + dma_queue * dma_q = octx->ctx->dma[ith]; + + const uint32_t d_conv = kparams->d_conv; + const uint32_t d_inner = kparams->d_inner; + const uint32_t n_s = kparams->n_s; const uint32_t dr = scctx->nrows_per_thread; const uint32_t ir0 = scctx->row_start + dr * ith; @@ -279,23 +200,141 @@ static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void return; } + const uint32_t d_inner_per_thread = ir1 - ir0; + const uint32_t d_inner_stride = hex_round_up(d_inner_per_thread, VLEN_FP32); + + const size_t src0_stride_seq_bytes = src0->nb[2]; + const size_t dst_stride_seq_bytes = dst->nb[2]; + + uint8_t * src1_spad_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread; + uint8_t * src0_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; + uint8_t * dst_spad_base = octx->dst_spad.data + ith * octx->dst_spad.size_per_thread; + + const size_t weight_bytes = (size_t) d_inner_per_thread * d_conv * sizeof(float); + const size_t weight_raw_size = hex_round_up(weight_bytes, 128); + + float * src1_raw = (float *) src1_spad_base; + float * src1_T = (float *) (src1_spad_base + weight_raw_size); + + float * src0_raw = (float *) src0_spad_base; + float * src0_T = (float *) (src0_spad_base + weight_raw_size); + + float * dst_spad = (float *) dst_spad_base; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + // 1. Fetch weights src1 from DDR into VTCM via DMA (DMA64-safe) + const dma_addr_t src1_ddr = src1->data + ir0 * d_conv * sizeof(float); + dma_queue_push(dma_q, dma_make_data((uint8_t *) src1_raw, src1_ddr), weight_bytes, weight_bytes, weight_bytes, 1); + dma_queue_pop(dma_q); + + // 2. Unpack/transpose src1_raw into src1_T {d_conv, d_inner_stride} htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0); + hvx_ssm_conv_unpack_to_T(src1_raw, src1_T, d_inner_per_thread, d_inner_stride, d_conv); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0); + + const size_t input_bytes = (size_t) d_inner_per_thread * d_conv * sizeof(float); + const size_t output_bytes = (size_t) d_inner_per_thread * sizeof(float); + + // 3. Process each sequence + for (uint32_t s = 0; s < n_s; ++s) { + const dma_addr_t src0_ddr = src0->data + s * src0_stride_seq_bytes + ir0 * d_conv * sizeof(float); + dma_queue_push(dma_q, dma_make_data((uint8_t *) src0_raw, src0_ddr), input_bytes, input_bytes, input_bytes, 1); + dma_queue_pop(dma_q); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) s); + hvx_ssm_conv_unpack_to_T(src0_raw, src0_T, d_inner_per_thread, d_inner_stride, d_conv); + + for (uint32_t cb = 0; cb < d_inner_per_thread; cb += VLEN_FP32) { + const uint32_t cb_n = MIN(VLEN_FP32, d_inner_per_thread - cb); + HVX_Vector acc = hvx_vec_splat_f32(0.0f); + for (uint32_t j = 0; j < d_conv; ++j) { + HVX_Vector x = *(const HVX_Vector *)(src0_T + j * d_inner_stride + cb); + HVX_Vector w = *(const HVX_Vector *)(src1_T + j * d_inner_stride + cb); + acc = Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(x, w)); + } + HVX_Vector y = Q6_Vsf_equals_Vqf32(acc); + if (cb_n == VLEN_FP32) { + *(HVX_Vector *)(dst_spad + cb) = y; + } else { + hvx_vec_store_u(dst_spad + cb, cb_n * sizeof(float), y); + } + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) s); + + const dma_addr_t dst_ddr = dst->data + s * dst_stride_seq_bytes + ir0 * sizeof(float); + dma_queue_push(dma_q, dma_make_data(dst_ddr, (uint8_t *) dst_spad), output_bytes, output_bytes, output_bytes, 1); + dma_queue_pop(dma_q); + } + + FARF(HIGH, "ssm-conv-f32-decode %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n", + ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1, + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1], + dst->ne[2], dst->ne[3]); +} + +// Multi-token prefill worker (n_t > 1) +static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void * data) { + struct htp_ssm_conv_context * scctx = (struct htp_ssm_conv_context *) data; + struct htp_ops_context * octx = scctx->octx; + const struct htp_ssm_conv_kernel_params * kparams = scctx->kparams; + + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict src1 = octx->src[1]; + const struct htp_tensor * restrict dst = octx->dst; + + dma_queue * dma_q = octx->ctx->dma[ith]; + + const uint32_t d_conv = kparams->d_conv; + const uint32_t d_inner = kparams->d_inner; + const uint32_t n_t = kparams->n_t; + const uint32_t n_s = kparams->n_s; + const uint32_t ncs = src0->ne[0]; + + const uint32_t dr = scctx->nrows_per_thread; + const uint32_t ir0 = scctx->row_start + dr * ith; + const uint32_t ir1 = MIN(ir0 + dr, scctx->row_start + scctx->nrows); + + if (ir0 >= ir1) { + return; + } const uint32_t d_inner_per_thread = ir1 - ir0; - const uint32_t d_inner_stride = scctx->nrows_per_thread; + const uint32_t d_inner_stride = hex_round_up(d_inner_per_thread, VLEN_FP32); const uint32_t d_inner_tile = scctx->d_inner_tile; - const float * src0_data = (const float *) src0->data; - const float * src1_data = (const float *) src1->data; - float * dst_data = (float *) dst->data; + const size_t src0_stride_inner_bytes = src0->nb[1]; + const size_t src0_stride_seq_bytes = src0->nb[2]; + const size_t dst_stride_token_bytes = dst->nb[1]; + const size_t dst_stride_seq_bytes = dst->nb[2]; - // Per-thread VTCM regions. - float * src0_T = (float *)(octx->src0_spad.data + ith * octx->src0_spad.size_per_thread); - float * src1_T = (float *)(octx->src1_spad.data + ith * octx->src1_spad.size_per_thread); + uint8_t * src1_spad_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread; + uint8_t * src0_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; + uint8_t * dst_spad_base = octx->dst_spad.data + ith * octx->dst_spad.size_per_thread; - // Stage src1 weights once into VTCM in {d_inner_stride, d_conv} layout. - transpose_src1(src1_data, src1_stride_inner, ir0, d_inner_per_thread, d_inner_stride, d_conv, src1_T); + const size_t weight_bytes = (size_t) d_inner_per_thread * d_conv * sizeof(float); + const size_t weight_raw_size = hex_round_up(weight_bytes, 128); + + float * src1_raw = (float *) src1_spad_base; + float * src1_T = (float *) (src1_spad_base + weight_raw_size); + + const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * ncs * sizeof(float), 128); + float * src0_tile_raw = (float *) src0_spad_base; + float * src0_T = (float *) (src0_spad_base + src0_tile_raw_bytes); + + float * dst_tile = (float *) dst_spad_base; + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + // 1. Fetch weights src1 from DDR into VTCM via DMA (DMA64-safe) + const dma_addr_t src1_ddr = src1->data + ir0 * d_conv * sizeof(float); + dma_queue_push(dma_q, dma_make_data((uint8_t *) src1_raw, src1_ddr), weight_bytes, weight_bytes, weight_bytes, 1); + dma_queue_pop(dma_q); + + // 2. Unpack/transpose src1_raw into src1_T {d_conv, d_inner_stride} + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0); + hvx_ssm_conv_unpack_to_T(src1_raw, src1_T, d_inner_per_thread, d_inner_stride, d_conv); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0); const uint32_t C_TILE = VLEN_FP32; @@ -303,14 +342,24 @@ static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void for (uint32_t tile_off = 0; tile_off < d_inner_per_thread; tile_off += d_inner_tile) { const uint32_t tile_n = MIN(d_inner_tile, d_inner_per_thread - tile_off); - // Place src0 chunk into VTCM in {d_inner_tile, ncs} layout. - const float * src0_block = src0_data + i3 * src0_stride_seq + (ir0 + tile_off) * src0_stride_inner; + // Fetch src0 chunk from DDR to VTCM via 2D DMA + const dma_addr_t src0_tile_ddr = src0->data + + i3 * src0_stride_seq_bytes + + (ir0 + tile_off) * src0_stride_inner_bytes; + const size_t row_bytes = ncs * sizeof(float); + dma_queue_push(dma_q, dma_make_data((uint8_t *) src0_tile_raw, src0_tile_ddr), + row_bytes, src0_stride_inner_bytes, row_bytes, tile_n); + dma_queue_pop(dma_q); + + // Transpose src0 chunk in VTCM into {d_inner_tile, ncs} layout + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off); for (uint32_t cb = 0; cb < tile_n; cb += C_TILE) { const uint32_t cb_n = MIN(C_TILE, tile_n - cb); - transpose_src0_block(src0_block + cb * src0_stride_inner, ncs, cb_n, d_inner_tile, src0_T, cb); + transpose_src0_block(src0_tile_raw + cb * ncs, ncs, cb_n, d_inner_tile, src0_T, cb); } + // Compute convolution for (uint32_t t = 0; t < n_t; ++t) { for (uint32_t cb = 0; cb < tile_n; cb += C_TILE) { const uint32_t cb_n = MIN(C_TILE, tile_n - cb); @@ -323,21 +372,29 @@ static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void } HVX_Vector y = Q6_Vsf_equals_Vqf32(acc); - - float * dst_ptr = dst_data + (ir0 + tile_off + cb) + t * dst_stride_token + i3 * dst_stride_seq; + float * dst_tile_ptr = dst_tile + t * tile_n + cb; if (cb_n == C_TILE) { - *(HVX_UVector *) dst_ptr = y; + *(HVX_Vector *) dst_tile_ptr = y; } else { - hvx_vec_store_u(dst_ptr, cb_n * sizeof(float), y); + hvx_vec_store_u(dst_tile_ptr, cb_n * sizeof(float), y); } } } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off); + + // Writeback dst_tile from VTCM to DDR via 2D DMA + const dma_addr_t dst_tile_ddr = dst->data + + i3 * dst_stride_seq_bytes + + (ir0 + tile_off) * sizeof(float); + const size_t dst_row_bytes = tile_n * sizeof(float); + + dma_queue_push(dma_q, dma_make_data(dst_tile_ddr, (uint8_t *) dst_tile), + dst_stride_token_bytes, dst_row_bytes, dst_row_bytes, n_t); + dma_queue_pop(dma_q); } } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1); - - FARF(HIGH, "ssm-conv-f32-hvx %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n", + FARF(HIGH, "ssm-conv-f32-prefill %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n", ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]); @@ -352,21 +409,25 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - const uint32_t d_conv = src1->ne[0]; - const uint32_t d_inner = src0->ne[1]; - const uint32_t n_t = dst->ne[1]; // tokens per sequence - const uint32_t n_s = dst->ne[2]; // number of sequences in the batch + const struct htp_ssm_conv_kernel_params * kparams = (const struct htp_ssm_conv_kernel_params *) octx->kernel_params; - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + + if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) { + return HTP_STATUS_INVAL_PARAMS; } uint32_t row_start = 0; - uint32_t nrows = d_inner; + uint32_t nrows = kparams->d_inner; if (octx->ctx->mdev.count > 1) { const uint32_t elems_per_chunk = VLEN_FP32; - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(d_inner, htp_tensor_mdev_data_aligned(dst) ? elems_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition( + kparams->d_inner, + htp_tensor_mdev_data_aligned(dst) ? elems_per_chunk : 0, + octx->ctx->mdev.idx, + octx->ctx->mdev.count, + &octx->ctx->mdev.count_div + ); row_start = range.start; nrows = range.count; } @@ -375,64 +436,49 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) { return HTP_STATUS_OK; } + if (kparams->vtcm_size > octx->ctx->vtcm_size) { + return HTP_STATUS_VTCM_TOO_SMALL; + } + const uint32_t n_threads = octx->n_threads; - struct htp_ssm_conv_context scctx = { 0 }; - scctx.octx = octx; - scctx.row_start = row_start; - scctx.nrows = nrows; + octx->src0_spad.size_per_thread = kparams->vtcm_src0_size_per_thread; + octx->src1_spad.size_per_thread = kparams->vtcm_src1_size_per_thread; + octx->dst_spad.size_per_thread = kparams->vtcm_dst_size_per_thread; - uint32_t use_hvx = 0; - if (nrows >= VLEN_FP32 && n_t >= VLEN_FP32) { - use_hvx = 1; - } + octx->src0_spad.size = kparams->vtcm_src0_size; + octx->src1_spad.size = kparams->vtcm_src1_size; + octx->dst_spad.size = kparams->vtcm_dst_size; - const uint32_t raw_rpt = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); - scctx.nrows_per_thread = hex_round_up(raw_rpt, VLEN_FP32); + octx->src0_spad.data = octx->ctx->vtcm_base; + octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; + octx->dst_spad.data = octx->src1_spad.data + octx->src1_spad.size; + octx->src0_spad.src = NULL; + octx->src1_spad.src = NULL; + octx->dst_spad.src = NULL; - const uint32_t d_inner_per_thread = scctx.nrows_per_thread; - const uint32_t ncs = src0->ne[0]; + const uint32_t raw_rpt = fastdiv(nrows + n_threads - 1, &octx->n_threads_div); + const uint32_t d_inner_per_thread = hex_round_up(raw_rpt, VLEN_FP32); - const uint32_t src1_T_size = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 256); - const uint32_t src0_T_max = HTP_SSM_CONV_VTCM_BUDGET > src1_T_size ? HTP_SSM_CONV_VTCM_BUDGET - src1_T_size : 0; + struct htp_ssm_conv_context scctx = { + .octx = octx, + .kparams = kparams, + .nrows_per_thread = d_inner_per_thread, + .d_inner_tile = kparams->d_inner_tile, + .row_start = row_start, + .nrows = nrows, + }; - uint32_t d_inner_tile = (src0_T_max / sizeof(float)) / ncs; - d_inner_tile -= (d_inner_tile % VLEN_FP32); - if (d_inner_tile == 0) { - FARF(HIGH, "ssm_conv-f32: inner tile rounds to 0 (ncs=%u), falling back to scalar\n", ncs); - use_hvx = 0; + FARF(HIGH, "ssm-conv-f32: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : mode %s\n", + src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], + src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + kparams->n_t == 1 ? "decode" : "prefill"); + + if (kparams->n_t == 1) { + work_queue_run(octx->ctx->work_queue, ssm_conv_thread_f32_decode, &scctx, n_threads); } else { - scctx.d_inner_tile = d_inner_tile; - - octx->src0_spad.size_per_thread = hex_round_up(d_inner_tile * ncs * sizeof(float), 256); - octx->src1_spad.size_per_thread = src1_T_size; - octx->dst_spad.size_per_thread = 0; - - octx->src0_spad.size = octx->src0_spad.size_per_thread * n_threads; - octx->src1_spad.size = octx->src1_spad.size_per_thread * n_threads; - octx->dst_spad.size = 0; - - octx->src0_spad.data = octx->ctx->vtcm_base; - octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; - octx->src0_spad.src = NULL; - octx->src1_spad.src = NULL; - - const size_t total_spad = octx->src0_spad.size + octx->src1_spad.size; - if (total_spad > octx->ctx->vtcm_size) { - FARF(HIGH, "ssm_conv-f32: scratchpad %zu exceeds VTCM %zu, falling back to scalar\n", - total_spad, octx->ctx->vtcm_size); - use_hvx = 0; - } - } - - FARF(HIGH, "ssm-conv-f32: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : use_hvx %d\n", src0->ne[0], - src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], - dst->ne[1], dst->ne[2], dst->ne[3], use_hvx); - - if (use_hvx) { - work_queue_run(octx->ctx->work_queue, ssm_conv_thread_f32_f32_hvx, &scctx, n_threads); - } else { - work_queue_run(octx->ctx->work_queue, ssm_conv_thread_f32_f32, &scctx, n_threads); + work_queue_run(octx->ctx->work_queue, ssm_conv_thread_f32_prefill, &scctx, n_threads); } return HTP_STATUS_OK; @@ -441,16 +487,10 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) { int op_ssm_conv(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_ssm_conv_f32(octx); - break; + return op_ssm_conv_f32(octx); default: - err = HTP_STATUS_NO_SUPPORT; - break; + return HTP_STATUS_NO_SUPPORT; } - - return err; } diff --git a/ggml/src/ggml-hexagon/htp/ssm-conv.h b/ggml/src/ggml-hexagon/htp/ssm-conv.h new file mode 100644 index 000000000..be62d7bf5 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/ssm-conv.h @@ -0,0 +1,40 @@ +#ifndef HTP_SSM_CONV_H +#define HTP_SSM_CONV_H + +#include + +#include "hex-fastdiv.h" +#include "htp-ops.h" + +struct htp_ssm_conv_kernel_params { + uint32_t n_threads; + uint32_t d_conv; + uint32_t d_inner; + uint32_t n_t; + uint32_t n_s; + uint32_t d_inner_per_thread; + uint32_t d_inner_tile; + + uint32_t src0_row_size_aligned; + uint32_t src1_row_size_aligned; + uint32_t dst_row_size_aligned; + + uint32_t vtcm_src0_size_per_thread; + uint32_t vtcm_src1_size_per_thread; + uint32_t vtcm_dst_size_per_thread; + + uint32_t vtcm_src0_size; + uint32_t vtcm_src1_size; + uint32_t vtcm_dst_size; + uint32_t vtcm_size; + + struct fastdiv_values div_n_threads; +}; + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_ssm_conv_kernel_params) <= 128, "htp_ssm_conv_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_ssm_conv_kernel_params) <= 128, "htp_ssm_conv_kernel_params is too large for kernel_params blob"); +#endif + +#endif // HTP_SSM_CONV_H diff --git a/ggml/src/ggml-hexagon/htp/sum-rows-ops.c b/ggml/src/ggml-hexagon/htp/sum-rows-ops.c index faf716b4b..9b8e04a0f 100644 --- a/ggml/src/ggml-hexagon/htp/sum-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/sum-rows-ops.c @@ -8,7 +8,7 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hvx-utils.h" #define GGML_COMMON_DECL_C @@ -106,8 +106,8 @@ int op_sum_rows(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; + if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) { + return HTP_STATUS_NO_SUPPORT; } const uint32_t src0_nrows = ne01 * ne02 * ne03; diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c index cb82bfa3c..9a1479e29 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ b/ggml/src/ggml-hexagon/htp/unary-ops.c @@ -8,7 +8,7 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" #include "hex-fastdiv.h" #include "hvx-exp.h" #include "hvx-sigmoid.h" @@ -23,13 +23,47 @@ #include "htp-vtcm.h" #include "hex-profile.h" +struct htp_unary_context; + +typedef void (*unary_compute_fn_t)(const void * restrict src, + void * restrict dst, + uint32_t num_rows, + const struct htp_unary_context * uctx); + +typedef void (*unary_rms_norm_mul_compute_fn_t)(const void * restrict src, + const void * restrict weight, + void * restrict dst, + uint32_t num_rows, + const struct htp_unary_context * uctx); + +typedef void (*unary_tri_compute_fn_t)(const void * restrict src, + void * restrict dst, + uint32_t num_rows, + uint32_t ir, + const struct htp_unary_context * uctx); + +typedef void (*unary_tile_compute_fn_t)(void * restrict dst, + const void * restrict src, + uint32_t tw, + const struct htp_unary_context * uctx); + +typedef void (*unary_tiled_tri_compute_fn_t)(const void * restrict src, + void * restrict dst, + uint32_t tile_elems, + uint32_t col_start, + uint32_t i01, + uint32_t ne0, + int32_t ttype); + struct htp_unary_context { struct htp_ops_context * octx; const struct htp_unary_kernel_params * kparams; - const uint8_t * data_src0; - const uint8_t * data_src1; // weight/scale tensor for RMS_NORM_MUL - uint8_t * data_dst; + void * compute; + + dma_addr_t data_src0; + dma_addr_t data_src1; // weight/scale tensor for RMS_NORM_MUL + dma_addr_t data_dst; size_t src0_data_row_size; // actual data bytes per row size_t src1_data_row_size; @@ -121,8 +155,8 @@ static inline uint32_t unary_block_size(uint32_t ir, const size_t src0_row_size_aligned = uctx->src0_row_size_aligned; \ const size_t dst_row_size_aligned = uctx->dst_row_size_aligned; -static void scale_f32(const float * restrict src, - float * restrict dst, +static void scale_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -139,8 +173,8 @@ static void scale_f32(const float * restrict src, } } -static void clamp_f32(const float * restrict src, - float * restrict dst, +static void clamp_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -157,8 +191,8 @@ static void clamp_f32(const float * restrict src, } } -static void leaky_relu_f32(const float * restrict src, - float * restrict dst, +static void leaky_relu_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -173,8 +207,8 @@ static void leaky_relu_f32(const float * restrict src, } } -static void rms_norm_f32(const float * restrict src, - float * restrict dst, +static void rms_norm_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -189,9 +223,9 @@ static void rms_norm_f32(const float * restrict src, } } -static void rms_norm_mul_f32(const float * restrict src, - const float * restrict weight, - float * restrict dst, +static void rms_norm_mul_f32(const void * restrict src, + const void * restrict weight, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -207,8 +241,8 @@ static void rms_norm_mul_f32(const float * restrict src, } } -static void norm_f32(const float * restrict src, - float * restrict dst, +static void norm_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -223,8 +257,8 @@ static void norm_f32(const float * restrict src, } } -static void sqr_f32(const float * restrict src, - float * restrict dst, +static void sqr_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -237,8 +271,8 @@ static void sqr_f32(const float * restrict src, } } -static void sqrt_f32(const float * restrict src, - float * restrict dst, +static void sqrt_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -251,8 +285,8 @@ static void sqrt_f32(const float * restrict src, } } -static void scale_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void scale_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -269,8 +303,8 @@ static void scale_f16(const _Float16 * restrict src, } } -static void clamp_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void clamp_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -287,8 +321,8 @@ static void clamp_f16(const _Float16 * restrict src, } } -static void rms_norm_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void rms_norm_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -303,8 +337,8 @@ static void rms_norm_f16(const _Float16 * restrict src, } } -static void norm_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void norm_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -319,8 +353,8 @@ static void norm_f16(const _Float16 * restrict src, } } -static void sqr_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void sqr_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -333,8 +367,8 @@ static void sqr_f16(const _Float16 * restrict src, } } -static void sqrt_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void sqrt_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -347,8 +381,8 @@ static void sqrt_f16(const _Float16 * restrict src, } } -static void abs_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void abs_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -361,8 +395,8 @@ static void abs_f16(const _Float16 * restrict src, } } -static void log_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void log_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -375,8 +409,8 @@ static void log_f16(const _Float16 * restrict src, } } -static void l2_norm_f16(const _Float16 * restrict src, - _Float16 * restrict dst, +static void l2_norm_f16(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -391,8 +425,8 @@ static void l2_norm_f16(const _Float16 * restrict src, } } -static void neg_f32(const float * restrict src, - float * restrict dst, +static void neg_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -405,8 +439,8 @@ static void neg_f32(const float * restrict src, } } -static void exp_f32(const float * restrict src, - float * restrict dst, +static void exp_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -419,8 +453,8 @@ static void exp_f32(const float * restrict src, } } -static void sigmoid_f32(const float * restrict src, - float * restrict dst, +static void sigmoid_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -434,8 +468,8 @@ static void sigmoid_f32(const float * restrict src, } // silu(x) = x * sigmoid(x) -static void silu_f32(const float * restrict src, - float * restrict dst, +static void silu_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -450,8 +484,8 @@ static void silu_f32(const float * restrict src, } // gelu(x) = x * sigmoid(1.702 * x) (quick/sigmoid approximation, matches CPU GELU_QUICK reference) -static void gelu_f32(const float * restrict src, - float * restrict dst, +static void gelu_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -466,8 +500,8 @@ static void gelu_f32(const float * restrict src, } } -static void tri_f32(const float * restrict src, - float * restrict dst, +static void tri_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const uint32_t ir, const struct htp_unary_context * uctx) { @@ -551,8 +585,8 @@ static void tri_f32(const float * restrict src, } } -static void softplus_f32(const float * restrict src, - float * restrict dst, +static void softplus_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -570,8 +604,8 @@ static void softplus_f32(const float * restrict src, } } -static void l2_norm_f32(const float * restrict src, - float * restrict dst, +static void l2_norm_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -579,15 +613,15 @@ static void l2_norm_f32(const float * restrict src, memcpy(&epsilon, op_params, sizeof(float)); for (uint32_t ir = 0; ir < num_rows; ir++) { - const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * src0_row_size_aligned)); - float * restrict dst_f = (float *)((uint8_t *)dst + (ir * dst_row_size_aligned)); + const uint8_t * restrict src_f = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_f = (uint8_t *)dst + (ir * dst_row_size_aligned); hvx_fast_l2_norm_f32((const uint8_t *)src_f, (uint8_t *)dst_f, ne0, epsilon); } } -static void tanh_f32(const float * restrict src, - float * restrict dst, +static void tanh_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -600,8 +634,8 @@ static void tanh_f32(const float * restrict src, } } -static void abs_f32(const float * restrict src, - float * restrict dst, +static void abs_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -614,8 +648,8 @@ static void abs_f32(const float * restrict src, } } -static void relu_f32(const float * restrict src, - float * restrict dst, +static void relu_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -628,8 +662,8 @@ static void relu_f32(const float * restrict src, } } -static void log_f32(const float * restrict src, - float * restrict dst, +static void log_f32(const void * restrict src, + void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; @@ -642,369 +676,100 @@ static void log_f32(const float * restrict src, } } -#define DEFINE_UNARY_TASK_IMPL(NAME, TYPE, SUFFIX, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \ -static void unary_task_##SUFFIX##_##NAME(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; \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - htp_unary_preamble; \ - \ - 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 = uctx->row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - const uint8_t * restrict data_src = uctx->data_src0; \ - const uint8_t * restrict data_src1 = uctx->data_src1; \ - uint8_t * restrict data_dst = uctx->data_dst; \ - \ - const struct htp_tensor * src1 = (IS_RMS_NORM_MUL) ? octx->src[1] : NULL; \ - const uint32_t nb11 = src1 ? src1->nb[1] : 0; \ - const uint32_t nb12 = src1 ? src1->nb[2] : 0; \ - const uint32_t nb13 = src1 ? src1->nb[3] : 0; \ - const uint32_t nb11_bc = (src1 && src1->ne[1] > 1) ? nb11 : 0; \ - const uint32_t nb12_bc = (src1 && src1->ne[2] > 1) ? nb12 : 0; \ - const uint32_t nb13_bc = (src1 && src1->ne[3] > 1) ? nb13 : 0; \ - const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false; \ - \ - uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \ - uint8_t * src1_vtcm_data = uctx->vtcm_src1 ? (uctx->vtcm_src1 + (ith * uctx->vtcm_src1_size_per_thread)) : NULL;\ - uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); \ - \ - size_t src0_vtcm_half_size = uctx->src0_vtcm_half_size; \ - size_t src1_vtcm_half_size = uctx->src1_vtcm_half_size; \ - size_t dst_vtcm_half_size = uctx->dst_vtcm_half_size; \ - \ - 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 struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; \ - const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \ - const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \ - \ - const bool src1_needs_row_clip = (IS_RMS_NORM_MUL) && !uctx->broadcast_weight && !src1_contig; \ - const bool block_src0_contig = src0_contig && !src1_needs_row_clip; \ - const bool block_dst_contig = dst_contig && !src1_needs_row_clip; \ - \ - const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ - const uint32_t dst_max_block = 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-" #SUFFIX " : current VTCM reservation %zu is too small, needed at least %zu\n", \ - uctx->vtcm_src0_size_per_thread, src0_row_size_aligned); \ - return; \ - } \ - \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ - \ - if ((IS_RMS_NORM_MUL) && uctx->broadcast_weight) { \ - dma_queue_push(dma_queue, dma_make_ptr(src1_vtcm_data, data_src1), \ - uctx->src1_row_size_aligned, 0, uctx->src1_data_row_size, 1); \ - dma_queue_flush(dma_queue); \ - } \ - \ - for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ - ne01, div_ne01); \ - \ - dma_queue_push(dma_queue, \ - dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), \ - nb1, dst_row_size_aligned, dst_data_row_size, 0); \ - \ - const size_t src0_off = src0_contig ? (ir * nb01) : \ - unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \ - dma_queue_push(dma_queue, \ - dma_make_ptr(src0_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src + src0_off), \ - src0_row_size_aligned, nb01, src0_data_row_size, block_size); \ - \ - if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ - const size_t src1_off = src1_contig ? (ir * nb11) : \ - unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc); \ - dma_queue_push(dma_queue, \ - dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), \ - uctx->src1_row_size_aligned, nb11, uctx->src1_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, block_src0_contig, block_dst_contig, \ - ne01, div_ne01); \ - \ - TYPE * dst_vtcm = (TYPE *) dma_queue_pop(dma_queue).src; \ - TYPE * src0_vtcm = (TYPE *) dma_queue_pop(dma_queue).dst; \ - TYPE * src1_vtcm = NULL; \ - if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ - src1_vtcm = (TYPE *) dma_queue_pop(dma_queue).dst; \ - } \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); \ - CORE_EXPR; \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); \ - \ - const size_t dst_off = dst_contig ? (ir * nb1) : \ - unary_row_offset(ir, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3); \ - dma_queue_push(dma_queue, \ - dma_make_ptr(data_dst + dst_off, dst_vtcm), \ - nb1, dst_row_size_aligned, dst_data_row_size, block_size); \ - \ - 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, block_src0_contig, \ - block_dst_contig, ne01, div_ne01); \ - 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, block_src0_contig, \ - block_dst_contig, ne01, div_ne01); \ - const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : \ - unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \ - dma_queue_push(dma_queue, \ - dma_make_ptr(src0_vtcm, data_src + src0_pref_off), \ - src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); \ - \ - if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ - const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : \ - unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, \ - nb13_bc); \ - dma_queue_push(dma_queue, \ - dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off), \ - uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); \ - } \ - } \ - } \ - ir += block_size; \ - } \ - \ - dma_queue_flush(dma_queue); \ -} - -// F32 unary task: row-block DMA/VTCM plumbing, float-typed VTCM buffers. -#define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \ - DEFINE_UNARY_TASK_IMPL(NAME, float, f32, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) - -DEFINE_UNARY_TASK(norm, false, false, norm_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(rms_norm, false, false, rms_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(rms_norm_mul, true, false, rms_norm_mul_f32(src0_vtcm, uctx->broadcast_weight ? (const float *) src1_vtcm_data : src1_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(scale, false, false, scale_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(clamp, false, false, clamp_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(leaky_relu, false, false, leaky_relu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(sqr, false, false, sqr_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(sqrt, false, false, sqrt_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_neg, false, false, neg_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_exp, false, false, exp_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_sigmoid, false, false, sigmoid_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_silu, false, false, silu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_gelu, false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_tanh, false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_abs, false, false, abs_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_log, false, false, log_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(unary_relu, false, false, relu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(l2_norm, false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK(tri, false, true, tri_f32(src0_vtcm, dst_vtcm, block_size, ir, uctx)) - -// F16 unary tasks: same DMA/VTCM plumbing as DEFINE_UNARY_TASK, but VTCM buffers are -// _Float16-typed. None of the current F16 ops need RMS_NORM_MUL or TRI support. -DEFINE_UNARY_TASK_IMPL(norm, _Float16, f16, false, false, norm_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(rms_norm, _Float16, f16, false, false, rms_norm_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(scale, _Float16, f16, false, false, scale_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(clamp, _Float16, f16, false, false, clamp_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(sqr, _Float16, f16, false, false, sqr_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(sqrt, _Float16, f16, false, false, sqrt_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(l2_norm, _Float16, f16, false, false, l2_norm_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(unary_abs, _Float16, f16, false, false, abs_f16(src0_vtcm, dst_vtcm, block_size, uctx)) -DEFINE_UNARY_TASK_IMPL(unary_log, _Float16, f16, false, false, log_f16(src0_vtcm, dst_vtcm, block_size, uctx)) - -// Apply a pointwise unary op to one column tile that is already in VTCM. -#define DEFINE_UNARY_TILED_TASK(NAME, IS_TRI, CORE_TILE_EXPR) \ -static void unary_task_f32_tiled_##NAME(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; \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - htp_unary_preamble; \ - \ - uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread; \ - \ - int32_t * op_params = octx->op_params; \ - const uint32_t col_tile = uctx->col_tile; \ - \ - const uint32_t src0_nrows = uctx->src0_nrows; \ - const uint32_t src0_start_row = uctx->row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - const uint8_t * restrict data_src = uctx->data_src0; \ - uint8_t * restrict data_dst = uctx->data_dst; \ - \ - uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \ - uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); \ - \ - const size_t src0_half = uctx->src0_vtcm_half_size; \ - const size_t dst_half = uctx->dst_vtcm_half_size; \ - \ - dma_queue * dmaq = octx->ctx->dma[ith]; \ - \ - const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; \ - const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \ - const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \ - const struct fastdiv_values * div_tpr = &uctx->kparams->div_tpr; \ - \ - const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile; \ - const int32_t tri_ttype = (IS_TRI) ? op_params[0] : 0; \ - \ - 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 total_tiles = (src0_end_row - src0_start_row) * tiles_per_row; \ - \ - for (uint32_t t = 0, vtcm_idx = 0; t < total_tiles && vtcm_idx < 2; t++, vtcm_idx++) { \ - const uint32_t row = src0_start_row + t / tiles_per_row; \ - const uint32_t col = (t % tiles_per_row) * col_tile; \ - const uint32_t tw = MIN(col_tile, ne0 - col); \ - const size_t tb = (size_t) tw * sizeof(float); \ - const size_t soff = (src0_contig ? (row * nb01) : \ - unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) + \ - (size_t) col * sizeof(float); \ - \ - dma_queue_push(dmaq, dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_half)), 0, 0, 0, 0); \ - dma_queue_push(dmaq, dma_make_ptr(src0_vtcm_data + (vtcm_idx * src0_half), data_src + soff), tb, tb, tb, 1); \ - } \ - \ - uint32_t row = src0_start_row; \ - uint32_t col = 0; \ - uint32_t tile_in_row = 0; \ - uint32_t i01 = fastmodulo(row, ne01, div_ne01); \ - \ - uint32_t prow = src0_start_row + fastdiv(2, div_tpr); \ - uint32_t pcol = fastmodulo(2, tiles_per_row, div_tpr) * col_tile; \ - uint32_t ptile_in_row = fastmodulo(2, tiles_per_row, div_tpr); \ - \ - for (uint32_t t = 0; t < total_tiles; t++) { \ - uint8_t * dst_vtcm = (uint8_t *) dma_queue_pop(dmaq).src; \ - uint8_t * src_vtcm = (uint8_t *) dma_queue_pop(dmaq).dst; \ - \ - const uint32_t tw = MIN(col_tile, ne0 - col); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, t); \ - CORE_TILE_EXPR; \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, t); \ - \ - const size_t doff = (dst_contig ? (row * nb1) : \ - unary_row_offset(row, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3)) + \ - (size_t) col * sizeof(float); \ - const size_t tb = (size_t) tw * sizeof(float); \ - dma_queue_push(dmaq, dma_make_ptr(data_dst + doff, dst_vtcm), tb, tb, tb, 1); \ - \ - const uint32_t pt = t + 2; \ - if (pt < total_tiles) { \ - const uint32_t ptw = MIN(col_tile, ne0 - pcol); \ - const size_t ptb = (size_t) ptw * sizeof(float); \ - const size_t psoff = (src0_contig ? (prow * nb01) : \ - unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, \ - nb03)) + \ - (size_t) pcol * sizeof(float); \ - dma_queue_push(dmaq, dma_make_ptr(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1); \ - } \ - \ - tile_in_row++; \ - col += col_tile; \ - if (tile_in_row == tiles_per_row) { \ - tile_in_row = 0; \ - col = 0; \ - row++; \ - i01++; \ - if (i01 == ne01) { \ - i01 = 0; \ - } \ - } \ - \ - ptile_in_row++; \ - pcol += col_tile; \ - if (ptile_in_row == tiles_per_row) { \ - ptile_in_row = 0; \ - pcol = 0; \ - prow++; \ - } \ - } \ - \ - dma_queue_flush(dmaq); \ -} - -static inline void tile_scale_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) { +#// Pointwise unary ops on one column tile in VTCM. +static void tile_scale_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { float scale = 0.f; - float bias = 0.f; - memcpy(&scale, &op_params[0], sizeof(float)); - memcpy(&bias, &op_params[1], sizeof(float)); - hvx_scale_offset_f32_aa(dst_vtcm, src_vtcm, tw, scale, bias); + float bias = 0.f; + memcpy(&scale, &uctx->octx->op_params[0], sizeof(float)); + memcpy(&bias, &uctx->octx->op_params[1], sizeof(float)); + hvx_scale_offset_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw, scale, bias); } -static inline void tile_clamp_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) { +static void tile_clamp_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { float min = 0.f; float max = 0.f; - memcpy(&min, &op_params[0], sizeof(float)); - memcpy(&max, &op_params[1], sizeof(float)); - hvx_clamp_scalar_f32(dst_vtcm, src_vtcm, min, max, tw); + memcpy(&min, &uctx->octx->op_params[0], sizeof(float)); + memcpy(&max, &uctx->octx->op_params[1], sizeof(float)); + hvx_clamp_scalar_f32((uint8_t *) dst, (const uint8_t *) src, min, max, tw); } -static inline void tile_leaky_relu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) { +static void tile_leaky_relu_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { float negative_slope = 0.f; - memcpy(&negative_slope, &op_params[0], sizeof(float)); - hvx_leaky_relu_scalar_f32(dst_vtcm, src_vtcm, negative_slope, tw); + memcpy(&negative_slope, &uctx->octx->op_params[0], sizeof(float)); + hvx_leaky_relu_scalar_f32((uint8_t *) dst, (const uint8_t *) src, negative_slope, tw); } -static inline void tile_unary_softplus_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) { - const float * restrict sf = (const float *) src_vtcm; - float * restrict df = (float *) dst_vtcm; +static void tile_sqr_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_sqr_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); +} + +static void tile_sqrt_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_sqrt_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); +} + +static void tile_neg_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_scale_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw, -1.0f); +} + +static void tile_exp_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_exp_f32((uint8_t *) dst, (const uint8_t *) src, tw, false); +} + +static void tile_sigmoid_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_sigmoid_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); +} + +static void tile_silu_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_sigmoid_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); + hvx_mul_f32_aaa((uint8_t *) dst, (const uint8_t *) src, (uint8_t *) dst, tw); +} + +static void tile_gelu_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_mul_scalar_f32((uint8_t *) dst, (const uint8_t *) src, 1.702f, tw); + hvx_sigmoid_f32_aa((uint8_t *) dst, (uint8_t *) dst, tw); + hvx_mul_f32_aaa((uint8_t *) dst, (const uint8_t *) src, (uint8_t *) dst, tw); +} + +static void tile_softplus_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + const float * restrict sf = (const float *) src; + float * restrict df = (float *) dst; for (uint32_t i = 0; i < tw; i++) { float x = sf[i]; df[i] = (x > 20.0f) ? x : logf(1.0f + expf(x)); } } -// silu(x) = x * sigmoid(x) -static inline void tile_silu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) { - hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw); - hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw); +static void tile_tanh_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_tanh_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); } -// gelu(x) = x * sigmoid(1.702 * x) (quick/sigmoid approximation, matches CPU GELU_QUICK reference) -static inline void tile_gelu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) { - hvx_mul_scalar_f32(dst_vtcm, src_vtcm, 1.702f, tw); - hvx_sigmoid_f32_aa(dst_vtcm, dst_vtcm, tw); - hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw); +static void tile_abs_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_abs_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); } -// Triangular mask applied to one column tile. Boundary is an absolute column index, so -// each vector compares against its absolute column position (col_start + i*VLEN_FP32). -static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * restrict dst, - uint32_t tile_elems, uint32_t col_start, uint32_t i01, - uint32_t ne0, int32_t ttype) { +static void tile_log_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_log_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); +} + +static void tile_relu_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { + (void) uctx; + hvx_max_scalar_f32((uint8_t *) dst, (const uint8_t *) src, 0.0f, tw); +} + +static void tri_apply_tile_f32(const void * restrict src, void * restrict dst, + uint32_t tile_elems, uint32_t col_start, uint32_t i01, + uint32_t ne0, int32_t ttype) { const HVX_Vector * restrict v_src = (const HVX_Vector *) src; HVX_Vector * restrict v_dst = (HVX_Vector *) dst; const HVX_Vector zero = hvx_vec_splat_f32(0.0f); @@ -1074,22 +839,623 @@ static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * re } } -DEFINE_UNARY_TILED_TASK(scale, false, tile_scale_f32(dst_vtcm, src_vtcm, tw, op_params)) -DEFINE_UNARY_TILED_TASK(clamp, false, tile_clamp_f32(dst_vtcm, src_vtcm, tw, op_params)) -DEFINE_UNARY_TILED_TASK(leaky_relu, false, tile_leaky_relu_f32(dst_vtcm, src_vtcm, tw, op_params)) -DEFINE_UNARY_TILED_TASK(sqr, false, hvx_sqr_f32_aa(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(sqrt, false, hvx_sqrt_f32_aa(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_neg, false, hvx_scale_f32_aa(dst_vtcm, src_vtcm, tw, -1.0f)) -DEFINE_UNARY_TILED_TASK(unary_exp, false, hvx_exp_f32(dst_vtcm, src_vtcm, tw, false)) -DEFINE_UNARY_TILED_TASK(unary_sigmoid, false, hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_silu, false, tile_silu_f32(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_gelu, false, tile_gelu_f32(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_tanh, false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_abs, false, hvx_abs_f32_aa(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_log, false, hvx_log_f32_aa(dst_vtcm, src_vtcm, tw)) -DEFINE_UNARY_TILED_TASK(unary_relu, false, hvx_max_scalar_f32(dst_vtcm, src_vtcm, 0.0f, tw)) -DEFINE_UNARY_TILED_TASK(tri, true, tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype)) +// 1. Standard row-block unary task (F32 and F16). +static void unary_thread_row_block(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + 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; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + htp_unary_preamble; + + const 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 = uctx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); + + if (src0_start_row >= src0_end_row) { + return; + } + + const dma_addr_t data_src = uctx->data_src0; + const dma_addr_t data_dst = uctx->data_dst; + + uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); + uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); + + const size_t src0_vtcm_half_size = uctx->src0_vtcm_half_size; + const size_t dst_vtcm_half_size = uctx->dst_vtcm_half_size; + + 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 struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; + const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; + const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; + + 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-row-block : current VTCM reservation %zu is too small, needed at least %zu\n", + uctx->vtcm_src0_size_per_thread, src0_row_size_aligned); + return; + } + + dma_queue * dma_q = octx->ctx->dma[ith]; + + for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, + ne01, div_ne01); + + dma_queue_push(dma_q, + dma_make_data(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), + nb1, dst_row_size_aligned, dst_data_row_size, 0); + + const size_t src0_off = src0_contig ? (ir * nb01) : + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); + dma_queue_push(dma_q, + dma_make_data(src0_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src + src0_off), + src0_row_size_aligned, nb01, src0_data_row_size, block_size); + + ir += block_size; + } + + unary_compute_fn_t compute = (unary_compute_fn_t) uctx->compute; + + 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, div_ne01); + + void * dst_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).src; + void * src0_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); + compute(src0_vtcm, dst_vtcm, block_size, uctx); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); + + const size_t dst_off = dst_contig ? (ir * nb1) : + unary_row_offset(ir, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3); + dma_queue_push(dma_q, + dma_make_data(data_dst + dst_off, dst_vtcm), + nb1, dst_row_size_aligned, dst_data_row_size, block_size); + + 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, div_ne01); + 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, div_ne01); + const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); + dma_queue_push(dma_q, + dma_make_data(src0_vtcm, data_src + src0_pref_off), + src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); + } + } + ir += block_size; + } + + dma_queue_flush(dma_q); +} + +// 2. RMS_NORM_MUL row-block task with weight buffer. +static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + 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; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + htp_unary_preamble; + + const 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 = uctx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); + + if (src0_start_row >= src0_end_row) { + return; + } + + const dma_addr_t data_src = uctx->data_src0; + const dma_addr_t data_src1 = uctx->data_src1; + const dma_addr_t data_dst = uctx->data_dst; + + const struct htp_tensor * src1 = octx->src[1]; + const uint32_t nb11 = src1->nb[1]; + const uint32_t nb12 = src1->nb[2]; + const uint32_t nb13 = src1->nb[3]; + const uint32_t nb11_bc = (src1->ne[1] > 1) ? nb11 : 0; + const uint32_t nb12_bc = (src1->ne[2] > 1) ? nb12 : 0; + const uint32_t nb13_bc = (src1->ne[3] > 1) ? nb13 : 0; + const bool src1_contig = ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)); + + uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); + uint8_t * src1_vtcm_data = uctx->vtcm_src1 ? (uctx->vtcm_src1 + (ith * uctx->vtcm_src1_size_per_thread)) : NULL; + uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); + + const size_t src0_vtcm_half_size = uctx->src0_vtcm_half_size; + const size_t src1_vtcm_half_size = uctx->src1_vtcm_half_size; + const size_t dst_vtcm_half_size = uctx->dst_vtcm_half_size; + + 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 struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; + const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; + const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; + + const bool src1_needs_row_clip = !uctx->broadcast_weight && !src1_contig; + const bool block_src0_contig = src0_contig && !src1_needs_row_clip; + const bool block_dst_contig = dst_contig && !src1_needs_row_clip; + + const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); + const uint32_t dst_max_block = 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-rms-norm-mul : current VTCM reservation %zu is too small, needed at least %zu\n", + uctx->vtcm_src0_size_per_thread, src0_row_size_aligned); + return; + } + + dma_queue * dma_q = octx->ctx->dma[ith]; + + if (uctx->broadcast_weight) { + dma_queue_push(dma_q, dma_make_data(src1_vtcm_data, data_src1), + uctx->src1_row_size_aligned, 0, uctx->src1_data_row_size, 1); + dma_queue_flush(dma_q); + } + + for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, + ne01, div_ne01); + + dma_queue_push(dma_q, + dma_make_data(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), + nb1, dst_row_size_aligned, dst_data_row_size, 0); + + const size_t src0_off = src0_contig ? (ir * nb01) : + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); + dma_queue_push(dma_q, + dma_make_data(src0_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src + src0_off), + src0_row_size_aligned, nb01, src0_data_row_size, block_size); + + if (!uctx->broadcast_weight) { + const size_t src1_off = src1_contig ? (ir * nb11) : + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc); + dma_queue_push(dma_q, + dma_make_data(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), + uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size); + } + + ir += block_size; + } + + unary_rms_norm_mul_compute_fn_t compute = (unary_rms_norm_mul_compute_fn_t) uctx->compute; + + for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, + ne01, div_ne01); + + void * dst_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).src; + void * src0_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + void * src1_vtcm = NULL; + if (!uctx->broadcast_weight) { + src1_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + } + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); + const void * w = uctx->broadcast_weight ? (const void *) src1_vtcm_data : src1_vtcm; + compute(src0_vtcm, w, dst_vtcm, block_size, uctx); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); + + const size_t dst_off = dst_contig ? (ir * nb1) : + unary_row_offset(ir, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3); + dma_queue_push(dma_q, + dma_make_data(data_dst + dst_off, dst_vtcm), + nb1, dst_row_size_aligned, dst_data_row_size, block_size); + + 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, block_src0_contig, + block_dst_contig, ne01, div_ne01); + 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, block_src0_contig, + block_dst_contig, ne01, div_ne01); + const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); + dma_queue_push(dma_q, + dma_make_data(src0_vtcm, data_src + src0_pref_off), + src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); + + if (!uctx->broadcast_weight) { + const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, + nb13_bc); + dma_queue_push(dma_q, + dma_make_data(src1_vtcm, data_src1 + src1_pref_off), + uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); + } + } + } + ir += block_size; + } + + dma_queue_flush(dma_q); +} + +// 3. TRI row-block task with row index ir. +static void unary_thread_tri_f32(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + 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; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + htp_unary_preamble; + + const 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 = uctx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); + + if (src0_start_row >= src0_end_row) { + return; + } + + const dma_addr_t data_src = uctx->data_src0; + const dma_addr_t data_dst = uctx->data_dst; + + uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); + uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); + + const size_t src0_vtcm_half_size = uctx->src0_vtcm_half_size; + const size_t dst_vtcm_half_size = uctx->dst_vtcm_half_size; + + 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 struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; + const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; + const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; + + 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-tri : current VTCM reservation %zu is too small, needed at least %zu\n", + uctx->vtcm_src0_size_per_thread, src0_row_size_aligned); + return; + } + + dma_queue * dma_q = octx->ctx->dma[ith]; + + for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, + ne01, div_ne01); + + dma_queue_push(dma_q, + dma_make_data(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), + nb1, dst_row_size_aligned, dst_data_row_size, 0); + + const size_t src0_off = src0_contig ? (ir * nb01) : + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); + dma_queue_push(dma_q, + dma_make_data(src0_vtcm_data + (vtcm_idx * src0_vtcm_half_size), data_src + src0_off), + src0_row_size_aligned, nb01, src0_data_row_size, block_size); + + ir += block_size; + } + + unary_tri_compute_fn_t compute = (unary_tri_compute_fn_t) uctx->compute; + + 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, div_ne01); + + void * dst_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).src; + void * src0_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); + compute(src0_vtcm, dst_vtcm, block_size, ir, uctx); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); + + const size_t dst_off = dst_contig ? (ir * nb1) : + unary_row_offset(ir, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3); + dma_queue_push(dma_q, + dma_make_data(data_dst + dst_off, dst_vtcm), + nb1, dst_row_size_aligned, dst_data_row_size, block_size); + + 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, div_ne01); + 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, div_ne01); + const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); + dma_queue_push(dma_q, + dma_make_data(src0_vtcm, data_src + src0_pref_off), + src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); + } + } + ir += block_size; + } + + dma_queue_flush(dma_q); +} + +// 4. Pointwise tiled unary task. +static void unary_thread_tiled(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + 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; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + htp_unary_preamble; + + const uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread; + const uint32_t col_tile = uctx->col_tile; + + const uint32_t src0_nrows = uctx->src0_nrows; + const uint32_t src0_start_row = uctx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); + + if (src0_start_row >= src0_end_row) { + return; + } + + const dma_addr_t data_src = uctx->data_src0; + const dma_addr_t data_dst = uctx->data_dst; + + uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); + uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); + + const size_t src0_half = uctx->src0_vtcm_half_size; + const size_t dst_half = uctx->dst_vtcm_half_size; + + dma_queue * dma_q = octx->ctx->dma[ith]; + + const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; + const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; + const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; + const struct fastdiv_values * div_tpr = &uctx->kparams->div_tpr; + + const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile; + + 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 total_tiles = (src0_end_row - src0_start_row) * tiles_per_row; + + for (uint32_t t = 0, vtcm_idx = 0; t < total_tiles && vtcm_idx < 2; t++, vtcm_idx++) { + const uint32_t row = src0_start_row + t / tiles_per_row; + const uint32_t col = (t % tiles_per_row) * col_tile; + const uint32_t tw = MIN(col_tile, ne0 - col); + const size_t tb = (size_t) tw * sizeof(float); + const size_t soff = (src0_contig ? (row * nb01) : + unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) + + (size_t) col * sizeof(float); + + dma_queue_push(dma_q, dma_make_data(data_dst, dst_vtcm_data + (vtcm_idx * dst_half)), 0, 0, 0, 0); + dma_queue_push(dma_q, dma_make_data(src0_vtcm_data + (vtcm_idx * src0_half), data_src + soff), tb, tb, tb, 1); + } + + unary_tile_compute_fn_t compute = (unary_tile_compute_fn_t) uctx->compute; + + uint32_t row = src0_start_row; + uint32_t col = 0; + uint32_t tile_in_row = 0; + + uint32_t prow = src0_start_row + fastdiv(2, div_tpr); + uint32_t pcol = fastmodulo(2, tiles_per_row, div_tpr) * col_tile; + uint32_t ptile_in_row = fastmodulo(2, tiles_per_row, div_tpr); + + for (uint32_t t = 0; t < total_tiles; t++) { + void * dst_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).src; + void * src_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + + const uint32_t tw = MIN(col_tile, ne0 - col); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, t); + compute(dst_vtcm, src_vtcm, tw, uctx); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, t); + + const size_t doff = (dst_contig ? (row * nb1) : + unary_row_offset(row, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3)) + + (size_t) col * sizeof(float); + const size_t tb = (size_t) tw * sizeof(float); + dma_queue_push(dma_q, dma_make_data(data_dst + doff, dst_vtcm), tb, tb, tb, 1); + + const uint32_t pt = t + 2; + if (pt < total_tiles) { + const uint32_t ptw = MIN(col_tile, ne0 - pcol); + const size_t ptb = (size_t) ptw * sizeof(float); + const size_t psoff = (src0_contig ? (prow * nb01) : + unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, + nb03)) + + (size_t) pcol * sizeof(float); + dma_queue_push(dma_q, dma_make_data(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1); + } + + tile_in_row++; + col += col_tile; + if (tile_in_row == tiles_per_row) { + tile_in_row = 0; + col = 0; + row++; + } + + ptile_in_row++; + pcol += col_tile; + if (ptile_in_row == tiles_per_row) { + ptile_in_row = 0; + pcol = 0; + prow++; + } + } + + dma_queue_flush(dma_q); +} + +// 5. TRI tiled task. +static void unary_thread_tiled_tri_f32(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + 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; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + htp_unary_preamble; + + const uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread; + const int32_t * op_params = octx->op_params; + const uint32_t col_tile = uctx->col_tile; + + const uint32_t src0_nrows = uctx->src0_nrows; + const uint32_t src0_start_row = uctx->row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows); + + if (src0_start_row >= src0_end_row) { + return; + } + + const dma_addr_t data_src = uctx->data_src0; + const dma_addr_t data_dst = uctx->data_dst; + + uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); + uint8_t * dst_vtcm_data = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread); + + const size_t src0_half = uctx->src0_vtcm_half_size; + const size_t dst_half = uctx->dst_vtcm_half_size; + + dma_queue * dma_q = octx->ctx->dma[ith]; + + const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01; + const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; + const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; + const struct fastdiv_values * div_tpr = &uctx->kparams->div_tpr; + + const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile; + const int32_t tri_ttype = op_params[0]; + + 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 total_tiles = (src0_end_row - src0_start_row) * tiles_per_row; + + for (uint32_t t = 0, vtcm_idx = 0; t < total_tiles && vtcm_idx < 2; t++, vtcm_idx++) { + const uint32_t row = src0_start_row + t / tiles_per_row; + const uint32_t col = (t % tiles_per_row) * col_tile; + const uint32_t tw = MIN(col_tile, ne0 - col); + const size_t tb = (size_t) tw * sizeof(float); + const size_t soff = (src0_contig ? (row * nb01) : + unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) + + (size_t) col * sizeof(float); + + dma_queue_push(dma_q, dma_make_data(data_dst, dst_vtcm_data + (vtcm_idx * dst_half)), 0, 0, 0, 0); + dma_queue_push(dma_q, dma_make_data(src0_vtcm_data + (vtcm_idx * src0_half), data_src + soff), tb, tb, tb, 1); + } + + unary_tiled_tri_compute_fn_t compute = (unary_tiled_tri_compute_fn_t) uctx->compute; + + uint32_t row = src0_start_row; + uint32_t col = 0; + uint32_t tile_in_row = 0; + uint32_t i01 = fastmodulo(row, ne01, div_ne01); + + uint32_t prow = src0_start_row + fastdiv(2, div_tpr); + uint32_t pcol = fastmodulo(2, tiles_per_row, div_tpr) * col_tile; + uint32_t ptile_in_row = fastmodulo(2, tiles_per_row, div_tpr); + + for (uint32_t t = 0; t < total_tiles; t++) { + void * dst_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).src; + void * src_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst; + + const uint32_t tw = MIN(col_tile, ne0 - col); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, t); + compute(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, t); + + const size_t doff = (dst_contig ? (row * nb1) : + unary_row_offset(row, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3)) + + (size_t) col * sizeof(float); + const size_t tb = (size_t) tw * sizeof(float); + dma_queue_push(dma_q, dma_make_data(data_dst + doff, dst_vtcm), tb, tb, tb, 1); + + const uint32_t pt = t + 2; + if (pt < total_tiles) { + const uint32_t ptw = MIN(col_tile, ne0 - pcol); + const size_t ptb = (size_t) ptw * sizeof(float); + const size_t psoff = (src0_contig ? (prow * nb01) : + unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, + nb03)) + + (size_t) pcol * sizeof(float); + dma_queue_push(dma_q, dma_make_data(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1); + } + + tile_in_row++; + col += col_tile; + if (tile_in_row == tiles_per_row) { + tile_in_row = 0; + col = 0; + row++; + i01++; + if (i01 == ne01) { + i01 = 0; + } + } + + ptile_in_row++; + pcol += col_tile; + if (ptile_in_row == tiles_per_row) { + ptile_in_row = 0; + pcol = 0; + prow++; + } + } + + dma_queue_flush(dma_q); +} static int execute_op_unary(struct htp_ops_context * octx) { int err = HTP_STATUS_OK; @@ -1207,115 +1573,128 @@ static int execute_op_unary(struct htp_ops_context * octx) { src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], kparams->vtcm_src0_size, kparams->vtcm_src1_size, kparams->vtcm_dst_size); - if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) { - uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - struct htp_unary_context uctx = { - .octx = octx, - .kparams = kparams, - .src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div), - .src0_nrows = nrows, - .row_start = row_start, + uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; + struct htp_unary_context uctx = { + .octx = octx, + .kparams = kparams, + .src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div), + .src0_nrows = nrows, + .row_start = row_start, - .data_src0 = (const uint8_t *)src0->data, - .data_src1 = (octx->op == HTP_OP_RMS_NORM_MUL) ? (const uint8_t *)src1->data : NULL, - .data_dst = (uint8_t *)dst->data, + .data_src0 = src0->data, + .data_src1 = (octx->op == HTP_OP_RMS_NORM_MUL) ? src1->data : 0, + .data_dst = dst->data, - .src0_data_row_size = src0_data_row_size, - .src1_data_row_size = src1_data_row_size, - .dst_data_row_size = dst_data_row_size, + .src0_data_row_size = src0_data_row_size, + .src1_data_row_size = src1_data_row_size, + .dst_data_row_size = dst_data_row_size, - .src0_row_size_aligned = src0_row_size_aligned, - .src1_row_size_aligned = src1_row_size_aligned, - .dst_row_size_aligned = dst_row_size_aligned, + .src0_row_size_aligned = src0_row_size_aligned, + .src1_row_size_aligned = src1_row_size_aligned, + .dst_row_size_aligned = dst_row_size_aligned, - .src0_vtcm_half_size = kparams->vtcm_src0_size_per_thread / 2, - .src1_vtcm_half_size = (octx->op == HTP_OP_RMS_NORM_MUL) ? (kparams->vtcm_src1_size_per_thread / (broadcast_weight ? 1 : 2)) : 0, - .dst_vtcm_half_size = kparams->vtcm_dst_size_per_thread / 2, + .src0_vtcm_half_size = kparams->vtcm_src0_size_per_thread / 2, + .src1_vtcm_half_size = (octx->op == HTP_OP_RMS_NORM_MUL) ? (kparams->vtcm_src1_size_per_thread / (broadcast_weight ? 1 : 2)) : 0, + .dst_vtcm_half_size = kparams->vtcm_dst_size_per_thread / 2, - .block = kparams->block, - .nc = src0->ne[0], - .col_tile = col_tile, - .broadcast_weight = broadcast_weight, + .block = kparams->block, + .nc = src0->ne[0], + .col_tile = col_tile, + .broadcast_weight = broadcast_weight, - .vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, 0), - .vtcm_src1 = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, kparams->vtcm_src0_size, kparams->vtcm_src1_size > 0), - .vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, kparams->vtcm_src0_size + kparams->vtcm_src1_size), + .vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, 0), + .vtcm_src1 = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, kparams->vtcm_src0_size, kparams->vtcm_src1_size > 0), + .vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, kparams->vtcm_src0_size + kparams->vtcm_src1_size), - .vtcm_src0_size_per_thread = kparams->vtcm_src0_size_per_thread, - .vtcm_src1_size_per_thread = kparams->vtcm_src1_size_per_thread, - .vtcm_dst_size_per_thread = kparams->vtcm_dst_size_per_thread, - }; + .vtcm_src0_size_per_thread = kparams->vtcm_src0_size_per_thread, + .vtcm_src1_size_per_thread = kparams->vtcm_src1_size_per_thread, + .vtcm_dst_size_per_thread = kparams->vtcm_dst_size_per_thread, + }; - FARF(HIGH, "%s: %s mode (col_tile %u)\n", op_type, col_tile ? "tiled" : "row-block", col_tile); + FARF(HIGH, "%s: %s mode (col_tile %u)\n", op_type, col_tile ? "tiled" : "row-block", col_tile); - worker_callback_t task_func = NULL; - if (col_tile) { - switch (octx->op) { - case HTP_OP_SCALE: task_func = unary_task_f32_tiled_scale; break; - case HTP_OP_CLAMP: task_func = unary_task_f32_tiled_clamp; break; - case HTP_OP_LEAKY_RELU: task_func = unary_task_f32_tiled_leaky_relu; break; - case HTP_OP_SQR: task_func = unary_task_f32_tiled_sqr; break; - case HTP_OP_SQRT: task_func = unary_task_f32_tiled_sqrt; break; - case HTP_OP_UNARY_NEG: task_func = unary_task_f32_tiled_unary_neg; break; - case HTP_OP_UNARY_EXP: task_func = unary_task_f32_tiled_unary_exp; break; - case HTP_OP_UNARY_SIGMOID: task_func = unary_task_f32_tiled_unary_sigmoid; break; - case HTP_OP_UNARY_SILU: task_func = unary_task_f32_tiled_unary_silu; break; - case HTP_OP_UNARY_GELU: task_func = unary_task_f32_tiled_unary_gelu; break; - case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_tiled_unary_softplus; break; - case HTP_OP_UNARY_TANH: task_func = unary_task_f32_tiled_unary_tanh; break; - case HTP_OP_UNARY_ABS: task_func = unary_task_f32_tiled_unary_abs; break; - case HTP_OP_UNARY_LOG: task_func = unary_task_f32_tiled_unary_log; break; - case HTP_OP_UNARY_RELU: task_func = unary_task_f32_tiled_unary_relu; break; - case HTP_OP_TRI: task_func = unary_task_f32_tiled_tri; break; - default: break; - } - } else if (is_f16) { - switch (octx->op) { - case HTP_OP_NORM: task_func = unary_task_f16_norm; break; - case HTP_OP_RMS_NORM: task_func = unary_task_f16_rms_norm; break; - case HTP_OP_SCALE: task_func = unary_task_f16_scale; break; - case HTP_OP_CLAMP: task_func = unary_task_f16_clamp; break; - case HTP_OP_SQR: task_func = unary_task_f16_sqr; break; - case HTP_OP_SQRT: task_func = unary_task_f16_sqrt; break; - case HTP_OP_L2_NORM: task_func = unary_task_f16_l2_norm; break; - case HTP_OP_UNARY_ABS: task_func = unary_task_f16_unary_abs; break; - case HTP_OP_UNARY_LOG: task_func = unary_task_f16_unary_log; break; - default: break; - } - } else { - switch (octx->op) { - case HTP_OP_NORM: task_func = unary_task_f32_norm; break; - case HTP_OP_RMS_NORM: task_func = unary_task_f32_rms_norm; break; - case HTP_OP_RMS_NORM_MUL: task_func = unary_task_f32_rms_norm_mul; break; - case HTP_OP_SCALE: task_func = unary_task_f32_scale; break; - case HTP_OP_CLAMP: task_func = unary_task_f32_clamp; break; - case HTP_OP_LEAKY_RELU: task_func = unary_task_f32_leaky_relu; break; - case HTP_OP_SQR: task_func = unary_task_f32_sqr; break; - case HTP_OP_SQRT: task_func = unary_task_f32_sqrt; break; - case HTP_OP_UNARY_NEG: task_func = unary_task_f32_unary_neg; break; - case HTP_OP_UNARY_EXP: task_func = unary_task_f32_unary_exp; break; - case HTP_OP_UNARY_SIGMOID: task_func = unary_task_f32_unary_sigmoid; break; - case HTP_OP_UNARY_SILU: task_func = unary_task_f32_unary_silu; break; - case HTP_OP_UNARY_GELU: task_func = unary_task_f32_unary_gelu; break; - case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_unary_softplus; break; - case HTP_OP_UNARY_TANH: task_func = unary_task_f32_unary_tanh; break; - case HTP_OP_UNARY_ABS: task_func = unary_task_f32_unary_abs; break; - case HTP_OP_UNARY_LOG: task_func = unary_task_f32_unary_log; break; - case HTP_OP_UNARY_RELU: task_func = unary_task_f32_unary_relu; break; - case HTP_OP_L2_NORM: task_func = unary_task_f32_l2_norm; break; - case HTP_OP_TRI: task_func = unary_task_f32_tri; break; - default: break; - } + worker_callback_t task_func = NULL; + void * compute_func = NULL; + + if (col_tile) { + task_func = unary_thread_tiled; + switch (octx->op) { + case HTP_OP_SCALE: compute_func = (void *) tile_scale_f32; break; + case HTP_OP_CLAMP: compute_func = (void *) tile_clamp_f32; break; + case HTP_OP_LEAKY_RELU: compute_func = (void *) tile_leaky_relu_f32; break; + case HTP_OP_SQR: compute_func = (void *) tile_sqr_f32; break; + case HTP_OP_SQRT: compute_func = (void *) tile_sqrt_f32; break; + case HTP_OP_UNARY_NEG: compute_func = (void *) tile_neg_f32; break; + case HTP_OP_UNARY_EXP: compute_func = (void *) tile_exp_f32; break; + case HTP_OP_UNARY_SIGMOID: compute_func = (void *) tile_sigmoid_f32; break; + case HTP_OP_UNARY_SILU: compute_func = (void *) tile_silu_f32; break; + case HTP_OP_UNARY_GELU: compute_func = (void *) tile_gelu_f32; break; + case HTP_OP_UNARY_SOFTPLUS: compute_func = (void *) tile_softplus_f32; break; + case HTP_OP_UNARY_TANH: compute_func = (void *) tile_tanh_f32; break; + case HTP_OP_UNARY_ABS: compute_func = (void *) tile_abs_f32; break; + case HTP_OP_UNARY_LOG: compute_func = (void *) tile_log_f32; break; + case HTP_OP_UNARY_RELU: compute_func = (void *) tile_relu_f32; break; + case HTP_OP_TRI: + task_func = unary_thread_tiled_tri_f32; + compute_func = (void *) tri_apply_tile_f32; + break; + default: break; } - - if (task_func) { - work_queue_run(octx->ctx->work_queue, task_func, &uctx, n_threads); - } else { - FARF(ERROR, "execute_op_unary: task function is NULL for op %d\n", octx->op); - err = HTP_STATUS_NO_SUPPORT; + } else if (is_f16) { + task_func = unary_thread_row_block; + switch (octx->op) { + case HTP_OP_NORM: compute_func = (void *) norm_f16; break; + case HTP_OP_RMS_NORM: compute_func = (void *) rms_norm_f16; break; + case HTP_OP_SCALE: compute_func = (void *) scale_f16; break; + case HTP_OP_CLAMP: compute_func = (void *) clamp_f16; break; + case HTP_OP_SQR: compute_func = (void *) sqr_f16; break; + case HTP_OP_SQRT: compute_func = (void *) sqrt_f16; break; + case HTP_OP_L2_NORM: compute_func = (void *) l2_norm_f16; break; + case HTP_OP_UNARY_ABS: compute_func = (void *) abs_f16; break; + case HTP_OP_UNARY_LOG: compute_func = (void *) log_f16; break; + default: break; + } + } else { + task_func = unary_thread_row_block; + switch (octx->op) { + case HTP_OP_NORM: compute_func = (void *) norm_f32; break; + case HTP_OP_RMS_NORM: compute_func = (void *) rms_norm_f32; break; + case HTP_OP_RMS_NORM_MUL: + task_func = unary_thread_rms_norm_mul_f32; + compute_func = (void *) rms_norm_mul_f32; + break; + case HTP_OP_SCALE: compute_func = (void *) scale_f32; break; + case HTP_OP_CLAMP: compute_func = (void *) clamp_f32; break; + case HTP_OP_LEAKY_RELU: compute_func = (void *) leaky_relu_f32; break; + case HTP_OP_SQR: compute_func = (void *) sqr_f32; break; + case HTP_OP_SQRT: compute_func = (void *) sqrt_f32; break; + case HTP_OP_UNARY_NEG: compute_func = (void *) neg_f32; break; + case HTP_OP_UNARY_EXP: compute_func = (void *) exp_f32; break; + case HTP_OP_UNARY_SIGMOID: compute_func = (void *) sigmoid_f32; break; + case HTP_OP_UNARY_SILU: compute_func = (void *) silu_f32; break; + case HTP_OP_UNARY_GELU: compute_func = (void *) gelu_f32; break; + case HTP_OP_UNARY_SOFTPLUS: compute_func = (void *) softplus_f32; break; + case HTP_OP_UNARY_TANH: compute_func = (void *) tanh_f32; break; + case HTP_OP_UNARY_ABS: compute_func = (void *) abs_f32; break; + case HTP_OP_UNARY_LOG: compute_func = (void *) log_f32; break; + case HTP_OP_UNARY_RELU: compute_func = (void *) relu_f32; break; + case HTP_OP_L2_NORM: compute_func = (void *) l2_norm_f32; break; + case HTP_OP_TRI: + task_func = unary_thread_tri_f32; + compute_func = (void *) tri_f32; + break; + default: break; } } + if (!task_func || !compute_func) { + FARF(ERROR, "execute_op_unary: task function is NULL for op %d\n", octx->op); + return HTP_STATUS_NO_SUPPORT; + } + + uctx.compute = compute_func; + work_queue_run(octx->ctx->work_queue, task_func, &uctx, n_threads); + return err; } diff --git a/scripts/snapdragon/ggml-hexagon-inspect.py b/scripts/snapdragon/ggml-hexagon-inspect.py new file mode 100755 index 000000000..3afda8a09 --- /dev/null +++ b/scripts/snapdragon/ggml-hexagon-inspect.py @@ -0,0 +1,1106 @@ +#!/usr/bin/env python3 +""" +ggml-hexagon-inspect.py - Hexagon DSP binary inspection and diagnostic tool. + +Inspects Hexagon ELF binaries (libggml-htp-v*.so) for: + - Register spills (--spills): counts scalar and HVX vector stack spills, + separating in-loop spills from frame setup/teardown. + - Function disassembly (--disasm ): annotated disassembly showing + hardware loop bounds, packet boundaries, and spill instructions. + - Crash address resolution (--addr2line ): maps hex crash offsets + to function symbols, offsets, and source lines. + - CI verification (--strict): fails with non-zero exit if in-loop vector + spills or DMA worker vector instructions are detected. + +Usage: + # Check spills across all functions or specific operations + ./scripts/snapdragon/ggml-hexagon-inspect.py --spills + ./scripts/snapdragon/ggml-hexagon-inspect.py --spills --func "^compute_" + ./scripts/snapdragon/ggml-hexagon-inspect.py --spills --func "^compute_" --strict + + # Disassemble a function with annotated loop and spill markers + ./scripts/snapdragon/ggml-hexagon-inspect.py --disasm compute_same_shape_div_f32 + + # Resolve crash addresses (CLI arguments or piped logcat/FARF logs) + ./scripts/snapdragon/ggml-hexagon-inspect.py --addr2line 0x51a30 0x5ba54 + adb logcat | ./scripts/snapdragon/ggml-hexagon-inspect.py --addr2line +""" + +import argparse +import logging +import os +import platform +import re +import shutil +import signal +import subprocess +import sys +from pathlib import Path +from typing import Dict, List, NamedTuple, Optional, Tuple + +# Ignore SIGPIPE to handle pipes (e.g. head, grep) gracefully +if hasattr(signal, "SIGPIPE"): + signal.signal(signal.SIGPIPE, signal.SIG_DFL) + +logger = logging.getLogger("ggml-hexagon-inspect") + + +class InsnInfo(NamedTuple): + address: int + asm_text: str + is_vec: bool + is_vspill: bool + is_sspill: bool + is_store: bool + is_load: bool + in_loop: bool + + +class FuncStats: + def __init__(self, name: str, address: int, size: int): + self.name = name + self.address = address + self.size = size + self.packet_count = 0 + self.insn_count = 0 + self.vec_insn_count = 0 + self.loop_count = 0 + self.vspills_in_loop = 0 + self.vspills_total = 0 + self.sspills_in_loop = 0 + self.sspills_total = 0 + self.promotions_in_loop = 0 + self.promotions_total = 0 + self.promotion_targets: Dict[str, int] = {} + self.calls_in_loop = 0 + self.calls_total = 0 + self.insns: List[InsnInfo] = [] + + +class SymbolEntry(NamedTuple): + address: int + size: int + name: str + + +# Regular expression patterns for Hexagon disassembly parsing +RE_SYMBOL_HEADER = re.compile(r"^([0-9a-fA-F]+)\s+<([^>]+)>:", re.MULTILINE) +RE_INSN_LINE = re.compile( + r"^\s*([0-9a-fA-F]+):\s+([0-9a-fA-F]{2}(?:\s+[0-9a-fA-F]{2}){3})\s+([0-9a-fA-F]{8})\s*(.*)$" +) +RE_LOOP0_START = re.compile(r"\bloop0\((0x[0-9a-fA-F]+)") +RE_LOOP1_START = re.compile(r"\bloop1\((0x[0-9a-fA-F]+)") +RE_VSPILL = re.compile(r"\bvmemu?\s*\(\s*r(?:29|30)\b") +RE_SSPILL = re.compile(r"\bmem[bwhd]\s*\(\s*r(?:29|30)\b") +RE_VEC_OP = re.compile(r"\b(v[0-9]+|w[0-9]+|q[0-3]|vmemu?)\b") +RE_STORE = re.compile(r"=\s*(?:v[0-9]|r[0-9]|w[0-9]|#)") +RE_PROMOTION_CALL = re.compile( + r"\b(?:call|jump)\s+(?:0x[0-9a-fA-F]+\s+)?<(__(?:trunc|extend)[a-zA-Z0-9_]+)(?:@plt)?>" +) +RE_ANY_CALL = re.compile(r"\bcallr?\b") + + +def get_repo_root() -> Path: + # Resolve repository root from script location + return Path(__file__).resolve().parent.parent.parent + + +def extract_arch_num(p: Path) -> int: + # Extract integer architecture version (e.g. v81 -> 81) + m = re.search(r"-v([0-9]+)\.so$", p.name) + return int(m.group(1)) if m else 0 + + +def find_default_lib(repo_root: Path, arch_filter: Optional[str] = None) -> Optional[Path]: + # Search for built Hexagon shared libraries in build and pkg directories + candidates = [] + search_dirs = [ + repo_root / "build-adb" / "ggml" / "src" / "ggml-hexagon", + repo_root / "build-android" / "ggml" / "src" / "ggml-hexagon", + repo_root / "build-ubuntu" / "ggml" / "src" / "ggml-hexagon", + repo_root / "build-linux" / "ggml" / "src" / "ggml-hexagon", + repo_root / "pkg-adb" / "llama.cpp" / "lib", + repo_root / "pkg-android" / "llama.cpp" / "lib", + repo_root / "pkg-ubuntu" / "llama.cpp" / "lib", + ] + + arch_needle = None + if arch_filter: + arch_needle = arch_filter if arch_filter.startswith("v") else f"v{arch_filter}" + + for d in search_dirs: + if not d.is_dir(): + continue + for p in d.glob("libggml-htp-*.so"): + if arch_needle and arch_needle not in p.name: + continue + candidates.append(p) + + if not candidates: + for p in repo_root.glob("build-*/ggml/src/ggml-hexagon/libggml-htp-*.so"): + if arch_needle and arch_needle not in p.name: + continue + candidates.append(p) + + if not candidates: + return None + + # Group latest build candidates (within 60s of max mtime) and pick highest arch + max_mtime = max(p.stat().st_mtime for p in candidates) + recent = [p for p in candidates if max_mtime - p.stat().st_mtime <= 60] + recent.sort(key=lambda p: extract_arch_num(p), reverse=True) + return recent[0] + + +def translate_container_arg(arg: str, repo_root: Path) -> str: + # Do not translate non-path command flags + if arg.startswith("-") and "=" not in arg: + return arg + if arg.startswith("--") and "=" in arg: + flag, val = arg.split("=", 1) + return f"{flag}={translate_container_arg(val, repo_root)}" + try: + p = Path(arg) + if (p.is_absolute() and p.exists()) or (p.exists() and ("/" in arg or "\\" in arg)): + resolved = p.resolve() + if resolved.is_relative_to(repo_root): + rel = resolved.relative_to(repo_root) + return f"/workspace/{rel.as_posix()}" + except Exception: + pass + return arg + + +class HexagonToolchain: + def __init__( + self, + repo_root: Path, + use_docker: bool = False, + image_url: str = "ghcr.io/snapdragon-toolchain", + image_name: str = "arm64-android", + image_ver: str = "v0.7", + ): + self.repo_root = repo_root + self.image = f"{image_url}/{image_name}:{image_ver}" + self.docker_bin = shutil.which("docker") + self.use_docker = use_docker + + if not use_docker: + self.native_objdump, self.native_addr2line = self._discover_native_tools() + else: + self.native_objdump = None + self.native_addr2line = None + + if not self.native_objdump and not self.native_addr2line: + self.use_docker = True + + def _discover_native_tools(self) -> Tuple[Optional[str], Optional[str]]: + # Check system PATH + objdump = shutil.which("hexagon-llvm-objdump") + addr2line = shutil.which("hexagon-addr2line") or shutil.which("hexagon-llvm-addr2line") + + # Check HEXAGON_TOOLS_ROOT environment variable + tools_root = os.environ.get("HEXAGON_TOOLS_ROOT") + if tools_root: + bin_dir = Path(tools_root) / "Tools" / "bin" + objdump_path = bin_dir / "hexagon-llvm-objdump" + addr2line_path = bin_dir / "hexagon-addr2line" + if objdump_path.is_file() and not objdump: + objdump = str(objdump_path) + if addr2line_path.is_file() and not addr2line: + addr2line = str(addr2line_path) + + # Check HEXAGON_SDK_ROOT environment variable + sdk_root = os.environ.get("HEXAGON_SDK_ROOT") + if sdk_root: + tools_parent = Path(sdk_root) / "tools" / "HEXAGON_Tools" + if tools_parent.is_dir(): + for t_dir in tools_parent.iterdir(): + bin_dir = t_dir / "Tools" / "bin" + objdump_path = bin_dir / "hexagon-llvm-objdump" + addr2line_path = bin_dir / "hexagon-addr2line" + if objdump_path.is_file() and not objdump: + objdump = str(objdump_path) + if addr2line_path.is_file() and not addr2line: + addr2line = str(addr2line_path) + + return objdump, addr2line + + def run_tool(self, tool_name: str, args: List[str], stdin_data: Optional[str] = None) -> str: + # Execute tool either natively or inside Docker container + if not self.use_docker: + tool_path = self.native_objdump if "objdump" in tool_name else self.native_addr2line + if not tool_path: + tool_path = shutil.which(tool_name) + if not tool_path: + raise RuntimeError(f"Tool {tool_name} not found natively. Use Docker instead.") + + cmd = [tool_path] + args + res = subprocess.run(cmd, capture_output=True, text=True, input=stdin_data) + if res.returncode != 0: + raise RuntimeError(f"Tool {tool_name} failed: {res.stderr.strip()}") + return res.stdout + + # Running via Docker container + if not self.docker_bin: + raise RuntimeError("Docker is required but not installed or found on PATH.") + + container_tools_dir = "/opt/hexagon/6.6.0.0/tools/HEXAGON_Tools/19.0.07/Tools/bin" + if "objdump" in tool_name: + container_tool = f"{container_tools_dir}/hexagon-llvm-objdump" + elif "addr2line" in tool_name: + container_tool = f"{container_tools_dir}/hexagon-addr2line" + elif "nm" in tool_name: + container_tool = f"{container_tools_dir}/llvm-nm" + else: + container_tool = f"{container_tools_dir}/{tool_name}" + + # Translate file paths from host to /workspace + translated_args = [translate_container_arg(arg, self.repo_root) for arg in args] + + docker_cmd = [ + "docker", + "run", + "--rm", + "--platform", + "linux/amd64", + "-v", + f"{self.repo_root}:/workspace", + "-w", + "/workspace", + ] + + if platform.system() != "Windows": + docker_cmd += ["-u", f"{os.getuid()}:{os.getgid()}"] + + docker_cmd += [self.image, container_tool] + translated_args + + res = subprocess.run(docker_cmd, capture_output=True, text=True, input=stdin_data) + if res.returncode != 0: + raise RuntimeError(f"Docker command failed: {res.stderr.strip()}") + return res.stdout + + +def parse_symbols(toolchain: HexagonToolchain, lib_path: Path) -> List[SymbolEntry]: + # Parse function symbols from library using objdump -t + output = toolchain.run_tool("hexagon-llvm-objdump", ["-t", str(lib_path)]) + sym_re = re.compile(r"^([0-9a-fA-F]+)\s+[lgw! ]+\s+F\s+\.text\s+([0-9a-fA-F]+)\s+(.+)$") + + symbols = [] + for line in output.splitlines(): + m = sym_re.match(line.strip()) + if m: + addr = int(m.group(1), 16) + size = int(m.group(2), 16) + name = m.group(3).strip() + symbols.append(SymbolEntry(addr, size, name)) + + symbols.sort(key=lambda s: s.address) + return symbols + + +def find_enclosing_symbol(symbols: List[SymbolEntry], address: int) -> Optional[Tuple[str, int]]: + # Binary search enclosing function symbol and compute offset + low = 0 + high = len(symbols) - 1 + best = None + + while low <= high: + mid = (low + high) // 2 + s = symbols[mid] + if s.address <= address: + if address < s.address + s.size: + return (s.name, address - s.address) + best = s + low = mid + 1 + else: + high = mid - 1 + + if best and address < best.address + best.size: + return (best.name, address - best.address) + return None + + +def parse_disassembly( + disasm_text: str, func_filter: Optional[re.Pattern] = None +) -> List[FuncStats]: + # Parse disassembly text into structured function statistics + matches = list(RE_SYMBOL_HEADER.finditer(disasm_text)) + funcs: List[FuncStats] = [] + + for i, m in enumerate(matches): + name = m.group(2) + if func_filter and not func_filter.search(name): + continue + + addr = int(m.group(1), 16) + start_idx = m.end() + end_idx = matches[i + 1].start() if i + 1 < len(matches) else len(disasm_text) + chunk = disasm_text[start_idx:end_idx] + + # Calculate rough byte size from line addresses + stats = FuncStats(name=name, address=addr, size=0) + + loop0_target: Optional[int] = None + loop1_target: Optional[int] = None + loop0_active = False + loop1_active = False + + first_addr = None + last_addr = None + + for raw_line in chunk.splitlines(): + lm = RE_INSN_LINE.match(raw_line) + if not lm: + continue + + cur_addr = int(lm.group(1), 16) + asm_chunk = lm.group(4) + + if first_addr is None: + first_addr = cur_addr + last_addr = cur_addr + + # Track packet count + if "{" in asm_chunk: + stats.packet_count += 1 + + # Check loop starts + m0 = RE_LOOP0_START.search(asm_chunk) + if m0: + loop0_target = int(m0.group(1), 16) + stats.loop_count += 1 + + m1 = RE_LOOP1_START.search(asm_chunk) + if m1: + loop1_target = int(m1.group(1), 16) + stats.loop_count += 1 + + if loop0_target is not None and cur_addr >= loop0_target: + loop0_active = True + if loop1_target is not None and cur_addr >= loop1_target: + loop1_active = True + + in_loop = loop0_active or loop1_active + + # Parse instructions within packet line + cleaned = re.sub(r"[{}\s]|:endloop[01]", " ", asm_chunk) + sub_insns = [p.strip() for p in cleaned.split(";") if p.strip()] + + for insn in sub_insns: + stats.insn_count += 1 + is_vec = bool(RE_VEC_OP.search(insn)) + if is_vec: + stats.vec_insn_count += 1 + + is_vspill = bool(RE_VSPILL.search(insn)) + is_sspill = bool(RE_SSPILL.search(insn)) + + # Identify store vs load + is_store = False + is_load = False + if is_vspill or is_sspill: + if RE_STORE.search(insn): + is_store = True + else: + is_load = True + + if is_vspill: + stats.vspills_total += 1 + if in_loop: + stats.vspills_in_loop += 1 + elif is_sspill: + stats.sspills_total += 1 + if in_loop: + stats.sspills_in_loop += 1 + + is_call = bool(RE_ANY_CALL.search(insn)) + prom_m = RE_PROMOTION_CALL.search(insn) + if is_call: + stats.calls_total += 1 + if in_loop: + stats.calls_in_loop += 1 + if prom_m: + stats.promotions_total += 1 + ptarget = prom_m.group(1) + stats.promotion_targets[ptarget] = stats.promotion_targets.get(ptarget, 0) + 1 + if in_loop: + stats.promotions_in_loop += 1 + + stats.insns.append( + InsnInfo( + address=cur_addr, + asm_text=insn, + is_vec=is_vec, + is_vspill=is_vspill, + is_sspill=is_sspill, + is_store=is_store, + is_load=is_load, + in_loop=in_loop, + ) + ) + + # Check loop ends + if ":endloop0" in asm_chunk: + loop0_active = False + loop0_target = None + if ":endloop1" in asm_chunk: + loop1_active = False + loop1_target = None + + if first_addr is not None and last_addr is not None: + stats.size = (last_addr - first_addr) + 4 + + funcs.append(stats) + + return funcs + + +def annotate_disasm_line( + raw_line: str, + loop0_target: Optional[int], + loop1_target: Optional[int], + loop0_active: bool, + loop1_active: bool, + use_color: bool = True, +) -> Tuple[str, Optional[int], Optional[int], bool, bool]: + # Annotate disassembly line with spill and loop tags + lm = RE_INSN_LINE.match(raw_line) + if not lm: + return raw_line, loop0_target, loop1_target, loop0_active, loop1_active + + cur_addr = int(lm.group(1), 16) + asm_chunk = lm.group(4) + + # Check loop starts + m0 = RE_LOOP0_START.search(asm_chunk) + if m0: + loop0_target = int(m0.group(1), 16) + m1 = RE_LOOP1_START.search(asm_chunk) + if m1: + loop1_target = int(m1.group(1), 16) + + if loop0_target is not None and cur_addr >= loop0_target: + loop0_active = True + if loop1_target is not None and cur_addr >= loop1_target: + loop1_active = True + + in_loop = loop0_active or loop1_active + + tags = [] + if m0: + tags.append("[LOOP0-START]") + if m1: + tags.append("[LOOP1-START]") + + if RE_VSPILL.search(asm_chunk): + if in_loop: + tags.append("[V-SPILL:IN-LOOP]" if not use_color else "\033[1;31m[V-SPILL:IN-LOOP]\033[0m") + else: + tags.append("[V-SPILL]" if not use_color else "\033[1;33m[V-SPILL]\033[0m") + elif RE_SSPILL.search(asm_chunk): + if in_loop: + tags.append("[S-SPILL:IN-LOOP]" if not use_color else "\033[1;35m[S-SPILL:IN-LOOP]\033[0m") + + prom_m = RE_PROMOTION_CALL.search(asm_chunk) + if prom_m: + ptarget = prom_m.group(1) + if in_loop: + tags.append(f"[PROMOTION:{ptarget}:IN-LOOP]" if not use_color else f"\033[1;31m[PROMOTION:{ptarget}:IN-LOOP]\033[0m") + else: + tags.append(f"[PROMOTION:{ptarget}]" if not use_color else f"\033[1;35m[PROMOTION:{ptarget}]\033[0m") + elif RE_ANY_CALL.search(asm_chunk): + if in_loop: + tags.append("[CALL:IN-LOOP]" if not use_color else "\033[1;31m[CALL:IN-LOOP]\033[0m") + else: + tags.append("[CALL]" if not use_color else "\033[1;36m[CALL]\033[0m") + + if ":endloop0" in asm_chunk: + tags.append("[LOOP0-END]") + loop0_active = False + loop0_target = None + if ":endloop1" in asm_chunk: + tags.append("[LOOP1-END]") + loop1_active = False + loop1_target = None + + tag_str = " ".join(tags) + if tag_str: + annotated = f"{raw_line:<80} {tag_str}" + else: + annotated = raw_line + + return annotated, loop0_target, loop1_target, loop0_active, loop1_active + + +def run_spills( + toolchain: HexagonToolchain, + lib_path: Path, + args: argparse.Namespace, +) -> int: + # Scan and report register spills across binary functions + logger.info(f"Inspecting library: {lib_path}") + disasm_text = toolchain.run_tool("hexagon-llvm-objdump", ["-d", str(lib_path)]) + + func_re = re.compile(args.func) if args.func else None + funcs = parse_disassembly(disasm_text, func_re) + + # Filter functions + reported = [] + for f in funcs: + has_spills = f.vspills_total > 0 or f.sspills_in_loop > 0 or f.sspills_total > 0 + if args.all or args.func or has_spills: + reported.append(f) + + # Sort: in-loop vector spills desc, then total vector spills desc, then in-loop scalar spills desc + reported.sort( + key=lambda x: (x.vspills_in_loop, x.vspills_total, x.sspills_in_loop, x.sspills_total), + reverse=True, + ) + + use_color = not args.no_color and sys.stdout.isatty() + + # Print summary table + col_addr = "Address" + col_name = "Function" + col_pkts = "Packets" + col_insn = "Insns" + col_vec = "HVX Ops" + col_vloop = "V-Loop" + col_vtot = "V-Tot" + col_sloop = "S-Loop" + col_stot = "S-Tot" + + hdr = ( + f"{col_addr:<10} | {col_name:<44} | {col_pkts:>7} | {col_insn:>6} | " + f"{col_vec:>7} | {col_vloop:>6} | {col_vtot:>5} | {col_sloop:>6} | {col_stot:>5}" + ) + sep = "-" * len(hdr) + + logger.info("\n" + sep) + logger.info(hdr) + logger.info(sep) + + tot_vloop = 0 + tot_sloop = 0 + tot_funcs_with_vloop = 0 + strict_violations = [] + + dma_re: Optional[re.Pattern[str]] = re.compile(args.dma_pattern) if args.dma_pattern else None + + for f in reported: + tot_vloop += f.vspills_in_loop + tot_sloop += f.sspills_in_loop + if f.vspills_in_loop > 0: + tot_funcs_with_vloop += 1 + + # Check strict criteria + if args.strict: + if f.vspills_in_loop > args.max_inloop_vspills: + strict_violations.append( + f"{f.name}: {f.vspills_in_loop} in-loop vector spills (max allowed: {args.max_inloop_vspills})" + ) + if dma_re and dma_re.search(f.name): + if f.vec_insn_count > args.max_dma_vec_ops: + strict_violations.append( + f"{f.name}: DMA worker contains {f.vec_insn_count} HVX vector ops (max allowed: {args.max_dma_vec_ops})" + ) + + # Highlight in-loop vector spills + vloop_str = f"{f.vspills_in_loop:>6}" + if f.vspills_in_loop > 0 and use_color: + vloop_str = f"\033[1;31m{vloop_str}\033[0m" + + logger.info( + f"0x{f.address:08x} | {f.name:<44} | {f.packet_count:>7} | {f.insn_count:>6} | " + f"{f.vec_insn_count:>7} | {vloop_str} | {f.vspills_total:>5} | {f.sspills_in_loop:>6} | {f.sspills_total:>5}" + ) + + logger.info(sep) + logger.info( + f"Total functions analyzed: {len(funcs)} | Reported: {len(reported)} | " + f"Functions with in-loop vector spills: {tot_funcs_with_vloop} | " + f"Total in-loop vector spills: {tot_vloop} | Total in-loop scalar spills: {tot_sloop}" + ) + + if args.strict: + logger.info("\n" + "=" * 50) + if strict_violations: + if use_color: + logger.error("\033[1;31mSTRICT CHECK FAILED\033[0m") + else: + logger.error("STRICT CHECK FAILED") + for v in strict_violations: + logger.error(f" - {v}") + logger.info("=" * 50) + return 1 + else: + if use_color: + logger.info("\033[1;32mSTRICT CHECK PASSED: 0 violations\033[0m") + else: + logger.info("STRICT CHECK PASSED: 0 violations") + logger.info("=" * 50) + + return 0 + + +def run_promotions( + toolchain: HexagonToolchain, + lib_path: Path, + args: argparse.Namespace, +) -> int: + # Scan and report soft-float promotion calls across binary functions + logger.info(f"Inspecting library: {lib_path}") + disasm_text = toolchain.run_tool("hexagon-llvm-objdump", ["-d", str(lib_path)]) + + func_re = re.compile(args.func) if args.func else None + funcs = parse_disassembly(disasm_text, func_re) + + reported = [] + for f in funcs: + if args.all or f.promotions_total > 0: + reported.append(f) + + # Sort: in-loop promotions desc, then total promotions desc + reported.sort( + key=lambda x: (x.promotions_in_loop, x.promotions_total), + reverse=True, + ) + + use_color = not args.no_color and sys.stdout.isatty() + + col_addr = "Address" + col_name = "Function" + col_loop = "Loops" + col_inloop = "In-Loop" + col_tot = "Total" + col_targets = "Promotion Targets" + + hdr = f"{col_addr:<10} | {col_name:<44} | {col_loop:>5} | {col_inloop:>7} | {col_tot:>5} | {col_targets}" + sep = "-" * max(len(hdr), 110) + + logger.info("\n" + sep) + logger.info(hdr) + logger.info(sep) + + tot_inloop = 0 + tot_prom = 0 + tot_funcs_with_prom = 0 + strict_violations = [] + + for f in reported: + tot_inloop += f.promotions_in_loop + tot_prom += f.promotions_total + if f.promotions_total > 0: + tot_funcs_with_prom += 1 + + if args.strict: + max_p = args.max_promotions if args.max_promotions is not None else 0 + if f.promotions_total > max_p: + strict_violations.append( + f"{f.name}: {f.promotions_total} float promotion calls (max allowed: {max_p})" + ) + + inloop_str = f"{f.promotions_in_loop:>7}" + if f.promotions_in_loop > 0 and use_color: + inloop_str = f"\033[1;31m{inloop_str}\033[0m" + + targets_str = ", ".join(f"{t}: {c}" for t, c in sorted(f.promotion_targets.items())) + logger.info( + f"0x{f.address:08x} | {f.name:<44} | {f.loop_count:>5} | {inloop_str} | {f.promotions_total:>5} | {targets_str}" + ) + + logger.info(sep) + logger.info( + f"Total functions analyzed: {len(funcs)} | Reported: {len(reported)} | " + f"Functions with float promotions: {tot_funcs_with_prom} | " + f"Total promotion calls: {tot_prom} | In-loop: {tot_inloop}" + ) + + if args.strict: + logger.info("\n" + "=" * 50) + if strict_violations: + if use_color: + logger.error("\033[1;31mSTRICT CHECK FAILED\033[0m") + else: + logger.error("STRICT CHECK FAILED") + for v in strict_violations: + logger.error(f" - {v}") + logger.info("=" * 50) + return 1 + else: + if use_color: + logger.info("\033[1;32mSTRICT CHECK PASSED: 0 violations\033[0m") + else: + logger.info("STRICT CHECK PASSED: 0 violations") + logger.info("=" * 50) + + return 0 + + +def run_disasm( + toolchain: HexagonToolchain, + lib_path: Path, + args: argparse.Namespace, +) -> int: + # Disassemble matching function(s) with annotated loop and spill markers + func_pattern = args.disasm + logger.info(f"Inspecting library: {lib_path}") + logger.info(f"Disassembling functions matching: '{func_pattern}'\n") + + # Disassemble symbol + disasm_text = toolchain.run_tool( + "hexagon-llvm-objdump", + ["-d", f"--disassemble-symbols={func_pattern}", str(lib_path)], + ) + + # If --disassemble-symbols yielded nothing (e.g. pattern was a regex), dump whole binary and filter + matches = list(RE_SYMBOL_HEADER.finditer(disasm_text)) + if not matches: + all_disasm = toolchain.run_tool("hexagon-llvm-objdump", ["-d", str(lib_path)]) + pat = re.compile(func_pattern) + all_matches = list(RE_SYMBOL_HEADER.finditer(all_disasm)) + matched_symbols = [m.group(2) for m in all_matches if pat.search(m.group(2))] + if not matched_symbols: + logger.error(f"Error: No symbols found matching '{func_pattern}'.") + return 1 + # Re-run with symbol list bounded by limit + sym_limit = args.limit if hasattr(args, "limit") and args.limit and args.limit > 0 else len(matched_symbols) + sym_arg = ",".join(matched_symbols[:sym_limit]) + disasm_text = toolchain.run_tool( + "hexagon-llvm-objdump", + ["-d", f"--disassemble-symbols={sym_arg}", str(lib_path)], + ) + matches = list(RE_SYMBOL_HEADER.finditer(disasm_text)) + + use_color = not args.no_color and sys.stdout.isatty() + + # Parse and log annotated functions + for i, m in enumerate(matches): + name = m.group(2) + addr = int(m.group(1), 16) + start_idx = m.end() + end_idx = matches[i + 1].start() if i + 1 < len(matches) else len(disasm_text) + chunk = disasm_text[start_idx:end_idx] + + # Parse statistics for this function + func_stats = parse_disassembly(disasm_text[m.start():end_idx])[0] + + # Log header + hdr_border = "=" * 80 + logger.info(hdr_border) + logger.info(f"Function: {name}") + logger.info(f"Address: 0x{addr:08x} - 0x{addr + func_stats.size:08x} ({func_stats.size} bytes)") + logger.info(f"Packets: {func_stats.packet_count} | Instructions: {func_stats.insn_count} | Loops: {func_stats.loop_count}") + vec_pct = (func_stats.vec_insn_count / func_stats.insn_count * 100.0) if func_stats.insn_count else 0.0 + logger.info(f"HVX Ops: {func_stats.vec_insn_count} ({vec_pct:.1f}% of instructions)") + logger.info( + f"Spills: Vector in-loop: {func_stats.vspills_in_loop} | Vector total: {func_stats.vspills_total} | " + f"Scalar in-loop: {func_stats.sspills_in_loop} | Scalar total: {func_stats.sspills_total}" + ) + logger.info( + f"Calls: Total: {func_stats.calls_total} (in-loop: {func_stats.calls_in_loop}) | " + f"Float promotions: {func_stats.promotions_total} (in-loop: {func_stats.promotions_in_loop})" + ) + logger.info(hdr_border) + + # Log annotated disassembly + loop0_target: Optional[int] = None + loop1_target: Optional[int] = None + loop0_active = False + loop1_active = False + + for line in chunk.splitlines(): + ann_line, loop0_target, loop1_target, loop0_active, loop1_active = annotate_disasm_line( + line, loop0_target, loop1_target, loop0_active, loop1_active, use_color + ) + logger.info(ann_line) + logger.info("") + + return 0 + + +def extract_addresses_from_input(lines: List[str]) -> List[int]: + # Extract hex program counter addresses from input lines + re_pc = re.compile(r"\b(?:pc|PC|ip|IP)\s*(?:=|:|\s)\s*0*(?:0x)?([0-9a-fA-F]{3,8})\b") + re_plus_hex = re.compile(r"\+0x([0-9a-fA-F]{3,8})\b") + re_hex = re.compile(r"\b0x([0-9a-fA-F]{3,8})\b") + re_bare_hex = re.compile(r"^\s*0*([0-9a-fA-F]{3,8})\s*$") + + addrs = [] + seen = set() + + for line in lines: + matched = False + for m in re_pc.finditer(line): + val = int(m.group(1), 16) + if val not in seen: + seen.add(val) + addrs.append(val) + matched = True + + if not matched: + for m in re_plus_hex.finditer(line): + val = int(m.group(1), 16) + if val not in seen: + seen.add(val) + addrs.append(val) + matched = True + + if not matched: + for m in re_hex.finditer(line): + val = int(m.group(1), 16) + if val not in seen: + seen.add(val) + addrs.append(val) + matched = True + + if not matched: + m = re_bare_hex.match(line) + if m: + val = int(m.group(1), 16) + if val not in seen: + seen.add(val) + addrs.append(val) + + return addrs + + +def run_addr2line( + toolchain: HexagonToolchain, + lib_path: Path, + args: argparse.Namespace, +) -> int: + # Resolve addresses or crash logs to source locations and symbols + input_addrs: List[int] = [] + + if args.addr2line: + for arg in args.addr2line: + if arg == "-": + continue + try: + val = int(arg, 16) + input_addrs.append(val) + except ValueError: + # Treat as text line and search for hex addresses + input_addrs.extend(extract_addresses_from_input([arg])) + + # Read from stdin if piped or requested via '-' + if not sys.stdin.isatty() or "-" in (args.addr2line or []): + stdin_lines = sys.stdin.readlines() + input_addrs.extend(extract_addresses_from_input(stdin_lines)) + + if not input_addrs: + logger.error("Error: No addresses found to resolve. Provide hex addresses or pipe crash logs to stdin.") + logger.error("Example: ./scripts/snapdragon/ggml-hexagon-inspect.py --addr2line 0x51a30 0x5ba54") + return 1 + + logger.info(f"Resolving {len(input_addrs)} address(es) against: {lib_path}\n") + + # Load symbol table for symbol + offset fallback + symbols = parse_symbols(toolchain, lib_path) + + # Format addresses for addr2line tool (prefixed with 0x) + addr_strs = [f"0x{a:x}" for a in input_addrs] + tool_args = ["-e", str(lib_path), "-f", "-C", "-p", "-a"] + addr_strs + + raw_output = toolchain.run_tool("hexagon-addr2line", tool_args) + + # Parse output lines + # Format: 0x51a30: binary_thread_add_id_f32 at /path/file.c:123 + re_out = re.compile(r"^(0x[0-9a-fA-F]+):\s+(.*?)\s+at\s+(.*)$") + + for line in raw_output.splitlines(): + line = line.strip() + if not line: + continue + m = re_out.match(line) + if m: + addr_hex = m.group(1) + addr_val = int(addr_hex, 16) + func_name = m.group(2) + src_loc = m.group(3) + + # Check if function name is unknown or generic, look up symbol table + sym_info = find_enclosing_symbol(symbols, addr_val) + if sym_info: + sym_name, sym_offset = sym_info + sym_display = f"{sym_name}+0x{sym_offset:x}" + else: + sym_display = func_name + + logger.info(f"{addr_hex:<12} -> {sym_display:<40} ({src_loc})") + else: + logger.info(line) + + return 0 + + +def main(): + parser = argparse.ArgumentParser( + description="Inspect Hexagon DSP binaries for register spills, function disassembly, and crash analysis." + ) + + # Target library + parser.add_argument( + "--lib", + help="Path to Hexagon shared library (e.g. libggml-htp-v81.so). Auto-detected if omitted.", + ) + parser.add_argument( + "--arch", + help="Architecture version filter for auto-detection (e.g. v75, v79, v81).", + ) + + # Modes + parser.add_argument( + "--spills", + action="store_true", + help="Scan binary and report scalar/vector stack spills table.", + ) + parser.add_argument( + "--promotions", + action="store_true", + help="Scan binary and report functions with soft-float promotion calls (__trunc*, __extend*).", + ) + parser.add_argument( + "--disasm", + metavar="FUNC", + help="Disassemble function symbol or regex pattern with annotated loop and spill markers.", + ) + parser.add_argument( + "--limit", + type=int, + default=20, + help="Maximum symbols to disassemble when using pattern in --disasm (default: 20, 0 for unlimited).", + ) + parser.add_argument( + "--addr2line", + nargs="*", + metavar="ADDR", + help="Resolve hex addresses or piped crash traces to symbols and source lines.", + ) + + # Filtering & Display + parser.add_argument( + "--func", + "-f", + help="Regex filter for function names in --spills or --promotions.", + ) + parser.add_argument( + "--all", + "-a", + action="store_true", + help="Show all functions in table, even those with 0 spills/promotions.", + ) + parser.add_argument( + "--no-color", + action="store_true", + help="Disable ANSI color output.", + ) + + # Strict check options + parser.add_argument( + "--strict", + action="store_true", + help="CI mode: exit with non-zero status if violations (in-loop vector spills, DMA worker vector ops) occur.", + ) + parser.add_argument( + "--max-inloop-vspills", + type=int, + default=0, + help="Maximum allowed in-loop vector spills in --strict mode (default: 0).", + ) + parser.add_argument( + "--max-dma-vec-ops", + type=int, + default=0, + help="Maximum allowed vector instructions in DMA workers in --strict mode (default: 0).", + ) + parser.add_argument( + "--max-promotions", + type=int, + default=None, + help="Maximum allowed float promotion calls in --strict mode (default: 0).", + ) + parser.add_argument( + "--dma-pattern", + default=r"^.*_thread(?:_.*)?$", + help="Regex pattern identifying DMA worker functions (default: '^.*_thread(?:_.*)?$').", + ) + + # Toolchain options + parser.add_argument( + "--docker", + action="store_true", + help="Force execution inside Docker container.", + ) + parser.add_argument( + "--no-docker", + action="store_true", + help="Force native execution on host instead of Docker.", + ) + parser.add_argument( + "--toolchain-version", + default="v0.7", + help="Docker toolchain tag (default: v0.7).", + ) + parser.add_argument( + "--toolchain-url", + default="ghcr.io/snapdragon-toolchain", + help="Docker toolchain registry (default: ghcr.io/snapdragon-toolchain).", + ) + parser.add_argument( + "--image-name", + default="arm64-android", + help="Docker toolchain image name (default: arm64-android).", + ) + + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO, format="%(message)s") + + repo_root = get_repo_root() + + # Determine target library + lib_path = None + if args.lib: + lib_path = Path(args.lib).resolve() + if not lib_path.is_file(): + logger.error(f"Error: Specified library '{args.lib}' does not exist.") + sys.exit(1) + else: + lib_path = find_default_lib(repo_root, args.arch) + if not lib_path: + logger.error("Error: No Hexagon library found in build-* or pkg-* directories.") + logger.error("Build the project first via ./scripts/snapdragon/build.py --target adb or specify --lib.") + sys.exit(1) + + # Initialize toolchain wrapper + use_docker = args.docker or (not args.no_docker and platform.system() == "Darwin") + try: + toolchain = HexagonToolchain( + repo_root=repo_root, + use_docker=use_docker, + image_url=args.toolchain_url, + image_name=args.image_name, + image_ver=args.toolchain_version, + ) + except Exception as e: + logger.error(f"Error initializing toolchain: {e}") + sys.exit(1) + + # Dispatch commands + if args.addr2line is not None: + sys.exit(run_addr2line(toolchain, lib_path, args)) + elif args.disasm: + sys.exit(run_disasm(toolchain, lib_path, args)) + elif args.promotions: + sys.exit(run_promotions(toolchain, lib_path, args)) + else: + # Default action is --spills + sys.exit(run_spills(toolchain, lib_path, args)) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(message)s") + main() diff --git a/scripts/snapdragon/run.py b/scripts/snapdragon/run.py index dc71d4a32..8917febc1 100755 --- a/scripts/snapdragon/run.py +++ b/scripts/snapdragon/run.py @@ -21,6 +21,7 @@ MANAGED_ENV_NAMES = ( "GGML_HEXAGON_NHVX", "GGML_HEXAGON_NHMX", "GGML_HEXAGON_HOSTBUF", + "GGML_HEXAGON_DMA64", "GGML_HEXAGON_OPBATCH", "GGML_HEXAGON_OPQUEUE", "GGML_HEXAGON_OPPOLL", @@ -155,6 +156,7 @@ def main(): parser.add_argument("--hex-nhvx", help="Number of HVX units to use (GGML_HEXAGON_NHVX)") parser.add_argument("--hex-nhmx", help="Number of HMX units to use. 0 disables HMX power-up (GGML_HEXAGON_NHMX)") parser.add_argument("--hex-hostbuf", help="Enable host buffers (GGML_HEXAGON_HOSTBUF)") + parser.add_argument("--hex-dma64", nargs="?", const="1", help="Enable (1) or disable (0) 64-bit DMA for model weights (GGML_HEXAGON_DMA64)") parser.add_argument("--hex-opbatch", help="Maximum number of operations to batch into a single HTP execution (GGML_HEXAGON_OPBATCH)") parser.add_argument("--hex-opqueue", help="Size of the asynchronous NPU operation queue (GGML_HEXAGON_OPQUEUE)") parser.add_argument("--hex-oppoll", default="1", help="Enable (1) or Disable (0) polling for NPU opbatch completion (GGML_HEXAGON_OPPOLL) (default: 1)") @@ -162,7 +164,7 @@ def main(): parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)") parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)") parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)") - parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 3:HMX,2:HVX-tiled,1:HVX-flat,0:disable") + parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 2:HMX,1:HVX,0:disable") parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable") parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable") parser.add_argument("--hex-etm", help="Enable Embedded Trace Macrocell hardware tracing / trace logging (GGML_HEXAGON_ETM)") @@ -294,6 +296,7 @@ def main(): set_env("GGML_HEXAGON_NHVX", args.hex_nhvx) set_env("GGML_HEXAGON_NHMX", args.hex_nhmx) set_env("GGML_HEXAGON_HOSTBUF", args.hex_hostbuf) + set_env("GGML_HEXAGON_DMA64", args.hex_dma64) set_env("GGML_HEXAGON_OPBATCH", args.hex_opbatch) set_env("GGML_HEXAGON_OPQUEUE", args.hex_opqueue) set_env("GGML_HEXAGON_OPPOLL", args.hex_oppoll) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c75cb3c0f..5c7de196d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10348,6 +10348,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {200001, 2, 3, 1}, true, true, GGML_TYPE_F16, {1, 1}, 0.1f, 8.0f)); test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {200000, 1, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {200000, 4, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); + test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {4, 1, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); + test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {4, 1023, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {643251, 3, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); for (float max_bias : {0.0f, 8.0f}) { From 6ad1af56033cc9d4cecacf574ac8af8eb4b3249e Mon Sep 17 00:00:00 2001 From: Samriddha Sinha <74808231+sam-india-007@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:32:40 +0530 Subject: [PATCH 02/42] ci : Upgrade CUDA to 13.4 for Ubuntu CUDA Release Builds (#29202) --- .github/workflows/release.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 87f652427..363227f52 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -330,13 +330,13 @@ jobs: defines: '-DGGML_CUDA_CUB_3DOT2=ON' - build: 'x64' os: ubuntu-24.04 - cuda: '13.3.1' - label: '13.3' + cuda: '13.4.1' + label: '13.4' defines: '' - build: 'arm64' os: ubuntu-24.04-arm - cuda: '13.3.1' - label: '13.3' + cuda: '13.4.1' + label: '13.4' defines: '' runs-on: ${{ matrix.os }} @@ -1854,8 +1854,8 @@ jobs: - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz) - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz) - [Ubuntu x64 (CUDA 12)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz) - [CUDA 12.8 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz) - - [Ubuntu x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-x64.tar.gz) - [CUDA 13.3 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-x64.tar.gz) - - [Ubuntu arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-arm64.tar.gz) - [CUDA 13.3 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.3-arm64.tar.gz) + - [Ubuntu x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-x64.tar.gz) - [CUDA 13.4 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-x64.tar.gz) + - [Ubuntu arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-arm64.tar.gz) - [CUDA 13.4 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-arm64.tar.gz) - [Ubuntu x64 (ROCm 10.0)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-10.0-x64.tar.gz) - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz) - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz) From 8034c1d1f166ce365b4909f5e27c4b9ac4edc804 Mon Sep 17 00:00:00 2001 From: pl752 Date: Mon, 21 Sep 2026 13:04:51 +0500 Subject: [PATCH 03/42] ggml-cpu: ARM Repack kernels for Q1_0 (#23492) * Implemented ARM NEON DP q1 4x4 repack * Hoisted out scaling by b_d in gemm * Added 4x8 NEON I8MM repack kernels * Cleanup for q1 arm repack * Added missing aliases for arch fallback * Corrected unused var statements * Extended table guard condition to account for i8mm w/o dp build Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Moved new declarations and references to groups' top * Moved declarations for uniformity --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- ggml/src/ggml-cpu/arch-fallback.h | 28 ++ ggml/src/ggml-cpu/arch/arm/repack.cpp | 309 +++++++++++++++++++ ggml/src/ggml-cpu/repack.cpp | 412 ++++++++++++++++++++++++++ ggml/src/ggml-cpu/repack.h | 13 + 4 files changed, 762 insertions(+) diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index 2b9a42657..4dbd1982b 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/ggml/src/ggml-cpu/arch-fallback.h @@ -39,6 +39,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 #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 @@ -55,6 +57,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_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 @@ -87,6 +91,8 @@ // 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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 #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 @@ -98,6 +104,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_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 @@ -124,6 +132,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 #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 @@ -140,6 +150,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_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 @@ -171,6 +183,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 #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 @@ -187,6 +201,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_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 @@ -213,6 +229,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 #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 @@ -228,6 +246,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_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 @@ -262,6 +282,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_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 @@ -277,6 +299,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_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 @@ -314,6 +338,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0 +#define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0 #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 @@ -330,6 +356,8 @@ #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_q1_0_4x4_q8_0_generic ggml_gemm_q1_0_4x4_q8_0 +#define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_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 diff --git a/ggml/src/ggml-cpu/arch/arm/repack.cpp b/ggml/src/ggml-cpu/arch/arm/repack.cpp index a75344430..ad0e5ccaf 100644 --- a/ggml/src/ggml-cpu/arch/arm/repack.cpp +++ b/ggml/src/ggml-cpu/arch/arm/repack.cpp @@ -48,6 +48,24 @@ static inline void decode_q_Kx8_6bit_scales(const uint8_t * scales_in, int16x8_t } #endif +#if defined(__aarch64__) && defined(__ARM_NEON) && (defined(__ARM_FEATURE_DOTPROD) || defined(__ARM_FEATURE_MATMUL_INT8)) +#define B1(c,s,n) 0x ## n ## c , 0x ## n ## s +#define B2(c,s,n) B1(c,s,n ## c), B1(c,s,n ## s) +#define B3(c,s,n) B2(c,s,n ## c), B2(c,s,n ## s) +#define B4(c,s,n) B3(c,s,n ## c), B3(c,s,n ## s) +#define B5(c,s,n) B4(c,s,n ## c), B4(c,s,n ## s) +#define B6(c,s,n) B5(c,s,n ## c), B5(c,s,n ## s) +#define B7(c,s,n) B6(c,s,n ## c), B6(c,s,n ## s) +#define B8(c,s ) B7(c,s, c), B7(c,s, s) + +static const uint64_t table_q1_signs[256] = { B8(ff, 01) }; + +static inline int8x16_t ggml_q1_0_unpack_pair(uint8_t bits0, uint8_t bits1) { + return vreinterpretq_s8_u8(vcombine_u8(vcreate_u8(table_q1_signs[bits0]), + vcreate_u8(table_q1_signs[bits1]))); +} +#endif + void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { assert(QK8_0 == 32); assert(k % QK8_0 == 0); @@ -1823,6 +1841,132 @@ void ggml_gemv_q8_0_4x8_q8_0(int n, ggml_gemv_q8_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc); } +void ggml_gemv_q1_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) { + const int qk = QK1_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(nb); + UNUSED(ncols_interleaved); + +#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + for (int c = 0; c < nc; c += ncols_interleaved) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (c / ncols_interleaved) * nb; + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + float32x4_t acc = vdupq_n_f32(0); + + for (int l = 0; l < nb; l++) { + const float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d)); + float32x4_t accb = vdupq_n_f32(0); + + for (int k = 0; k < 4; k++) { + const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * 4 + k; + const float ad = GGML_CPU_FP16_TO_FP32(a_blk->d); + int32x4_t ret = vdupq_n_s32(0); + + for (int tile = 0; tile < 8; tile += 4) { + const int8x16_t signs0 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 0) + 0], + b_ptr[l].qs[k * 16 + 2 * (tile + 0) + 1]); + const int8x16_t signs1 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 1) + 0], + b_ptr[l].qs[k * 16 + 2 * (tile + 1) + 1]); + const int8x16_t signs2 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 2) + 0], + b_ptr[l].qs[k * 16 + 2 * (tile + 2) + 1]); + const int8x16_t signs3 = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * (tile + 3) + 0], + b_ptr[l].qs[k * 16 + 2 * (tile + 3) + 1]); + const int8x16_t q_tiles = vld1q_s8(a_blk->qs + tile * 4); + + ret = vdotq_laneq_s32(ret, signs0, q_tiles, 0); + ret = vdotq_laneq_s32(ret, signs1, q_tiles, 1); + ret = vdotq_laneq_s32(ret, signs2, q_tiles, 2); + ret = vdotq_laneq_s32(ret, signs3, q_tiles, 3); + } + + accb = vfmaq_n_f32(accb, vcvtq_f32_s32(ret), ad); + } + acc = vfmaq_f32(acc, accb, b_d); + } + vst1q_f32(s, acc); + s += ncols_interleaved; + } + return; +#endif + ggml_gemv_q1_0_4x4_q8_0_generic(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemv_q1_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) { + const int qk = QK1_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(nb); + UNUSED(ncols_interleaved); + +#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + for (int c = 0; c < nc; c += ncols_interleaved) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (c / ncols_interleaved) * nb; + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + float32x4_t acc = vdupq_n_f32(0); + + for (int l = 0; l < nb; l++) { + const float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d)); + float32x4_t accb = vdupq_n_f32(0); + + for (int k = 0; k < 4; ++k) { + const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * 4 + k; + const uint8_t * GGML_RESTRICT b_qs = (const uint8_t *) b_ptr[l].qs + k * 16; + const float ad = GGML_CPU_FP16_TO_FP32(a_blk->d); + + int8x8x4_t a_chunks = vld1_s8_x4(a_blk->qs); + int8x16_t a0 = vcombine_s8(a_chunks.val[0], a_chunks.val[0]); + int8x16_t a1 = vcombine_s8(a_chunks.val[1], a_chunks.val[1]); + int8x16_t a2 = vcombine_s8(a_chunks.val[2], a_chunks.val[2]); + int8x16_t a3 = vcombine_s8(a_chunks.val[3], a_chunks.val[3]); + + int32x4_t ret0 = vdupq_n_s32(0); + int32x4_t ret1 = vdupq_n_s32(0); + + ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[0], b_qs[1]), a0); + ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[2], b_qs[3]), a0); + ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[4], b_qs[5]), a1); + ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[6], b_qs[7]), a1); + ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[8], b_qs[9]), a2); + ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[10], b_qs[11]), a2); + ret0 = vdotq_s32(ret0, ggml_q1_0_unpack_pair(b_qs[12], b_qs[13]), a3); + ret1 = vdotq_s32(ret1, ggml_q1_0_unpack_pair(b_qs[14], b_qs[15]), a3); + + accb = vfmaq_n_f32(accb, vcvtq_f32_s32(vpaddq_s32(ret0, ret1)), ad); + } + + acc = vfmaq_f32(acc, accb, b_d); + } + + vst1q_f32(s, acc); + s += ncols_interleaved; + } + return; +#endif + + ggml_gemv_q1_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, 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) { const int qk = QK8_0; const int nb = n / qk; @@ -5154,3 +5298,168 @@ void ggml_gemm_q8_0_4x8_q8_0(int n, #endif // defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8) ggml_gemm_q8_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc); } + +void ggml_gemm_q1_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) { + const int qk = QK1_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(nb); + UNUSED(ncols_interleaved); + +#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb); + + float32x4_t sumf[4]; + for (int m = 0; m < 4; m++) { + sumf[m] = vdupq_n_f32(0); + } + + for (int l = 0; l < nb; l++) { + float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d)); + float32x4_t blockf_0 = vdupq_n_f32(0); + float32x4_t blockf_1 = vdupq_n_f32(0); + float32x4_t blockf_2 = vdupq_n_f32(0); + float32x4_t blockf_3 = vdupq_n_f32(0); + + for (int k = 0; k < 4; ++k) { + const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k; + float32x4_t a_d = vcvt_f32_f16(vld1_f16((const float16_t *) a_blk->d)); + + int32x4_t sumi_0 = vdupq_n_s32(0); + int32x4_t sumi_1 = vdupq_n_s32(0); + int32x4_t sumi_2 = vdupq_n_s32(0); + int32x4_t sumi_3 = vdupq_n_s32(0); + + for (int tile = 0; tile < 8; ++tile) { + const int8x16_t signs = ggml_q1_0_unpack_pair(b_ptr[l].qs[k * 16 + 2 * tile + 0], + b_ptr[l].qs[k * 16 + 2 * tile + 1]); + const int8x16_t a_tile = vld1q_s8(a_blk->qs + tile * 16); + + sumi_0 = vdotq_laneq_s32(sumi_0, signs, a_tile, 0); + sumi_1 = vdotq_laneq_s32(sumi_1, signs, a_tile, 1); + sumi_2 = vdotq_laneq_s32(sumi_2, signs, a_tile, 2); + sumi_3 = vdotq_laneq_s32(sumi_3, signs, a_tile, 3); + } + + blockf_0 = vfmaq_laneq_f32(blockf_0, vcvtq_f32_s32(sumi_0), a_d, 0); + blockf_1 = vfmaq_laneq_f32(blockf_1, vcvtq_f32_s32(sumi_1), a_d, 1); + blockf_2 = vfmaq_laneq_f32(blockf_2, vcvtq_f32_s32(sumi_2), a_d, 2); + blockf_3 = vfmaq_laneq_f32(blockf_3, vcvtq_f32_s32(sumi_3), a_d, 3); + } + + sumf[0] = vfmaq_f32(sumf[0], blockf_0, b_d); + sumf[1] = vfmaq_f32(sumf[1], blockf_1, b_d); + sumf[2] = vfmaq_f32(sumf[2], blockf_2, b_d); + sumf[3] = vfmaq_f32(sumf[3], blockf_3, b_d); + } + + for (int m = 0; m < 4; m++) { + vst1q_f32(s + (y * 4 + m) * bs + x * 4, sumf[m]); + } + } + } + return; +#endif + ggml_gemm_q1_0_4x4_q8_0_generic(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemm_q1_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) { + const int qk = QK1_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(nb); + UNUSED(ncols_interleaved); + +#if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8) + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb); + + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb); + + float32x4_t sumf[4]; + for (int m = 0; m < 4; ++m) { + sumf[m] = vdupq_n_f32(0); + } + + for (int l = 0; l < nb; l++) { + const float32x4_t b_d = vcvt_f32_f16(vld1_f16((const float16_t *) b_ptr[l].d)); + float32x4_t blockf[4]; + for (int m = 0; m < 4; ++m) { + blockf[m] = vdupq_n_f32(0); + } + + for (int k = 0; k < 4; ++k) { + const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k; + const uint8_t * GGML_RESTRICT b_qs = (const uint8_t *) b_ptr[l].qs + k * 16; + + int32x4_t acc[4]; + for (int i = 0; i < 4; ++i) { + acc[i] = vdupq_n_s32(0); + } + + for (int chunk = 0; chunk < 4; ++chunk) { + const int8x16_t a01 = vld1q_s8(a_blk->qs + chunk * 32); + const int8x16_t a23 = vld1q_s8(a_blk->qs + chunk * 32 + 16); + const int8x16_t b01 = ggml_q1_0_unpack_pair(b_qs[chunk * 4 + 0], b_qs[chunk * 4 + 1]); + const int8x16_t b23 = ggml_q1_0_unpack_pair(b_qs[chunk * 4 + 2], b_qs[chunk * 4 + 3]); + + acc[0] = vmmlaq_s32(acc[0], a01, b01); + acc[1] = vmmlaq_s32(acc[1], a01, b23); + acc[2] = vmmlaq_s32(acc[2], a23, b01); + acc[3] = vmmlaq_s32(acc[3], a23, b23); + } + + const int32x4_t row0 = vcombine_s32(vget_low_s32(acc[0]), vget_low_s32(acc[1])); + const int32x4_t row1 = vcombine_s32(vget_high_s32(acc[0]), vget_high_s32(acc[1])); + const int32x4_t row2 = vcombine_s32(vget_low_s32(acc[2]), vget_low_s32(acc[3])); + const int32x4_t row3 = vcombine_s32(vget_high_s32(acc[2]), vget_high_s32(acc[3])); + const float32x4_t a_d = vcvt_f32_f16(vld1_f16((const float16_t *) a_blk->d)); + + blockf[0] = vfmaq_laneq_f32(blockf[0], vcvtq_f32_s32(row0), a_d, 0); + blockf[1] = vfmaq_laneq_f32(blockf[1], vcvtq_f32_s32(row1), a_d, 1); + blockf[2] = vfmaq_laneq_f32(blockf[2], vcvtq_f32_s32(row2), a_d, 2); + blockf[3] = vfmaq_laneq_f32(blockf[3], vcvtq_f32_s32(row3), a_d, 3); + } + + sumf[0] = vfmaq_f32(sumf[0], blockf[0], b_d); + sumf[1] = vfmaq_f32(sumf[1], blockf[1], b_d); + sumf[2] = vfmaq_f32(sumf[2], blockf[2], b_d); + sumf[3] = vfmaq_f32(sumf[3], blockf[3], b_d); + } + + for (int m = 0; m < 4; ++m) { + vst1q_f32(s + (y * 4 + m) * bs + x * 4, sumf[m]); + } + } + } + return; +#endif + + ggml_gemm_q1_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc); +} diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp index f5e419c1e..d56db9802 100644 --- a/ggml/src/ggml-cpu/repack.cpp +++ b/ggml/src/ggml-cpu/repack.cpp @@ -1365,6 +1365,133 @@ void ggml_gemv_q8_0_4x8_q8_0_generic(int n, } } +void ggml_gemv_q1_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 = QK1_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[4]; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + } + + for (int l = 0; l < nb; l++) { + const float d0[4] = { + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]), + }; + + for (int k = 0; k < QK1_0 / QK8_0; ++k) { + const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * (QK1_0 / QK8_0) + k; + const float d1 = GGML_CPU_FP16_TO_FP32(a_blk->d); + const float scale[4] = { d0[0] * d1, d0[1] * d1, d0[2] * d1, d0[3] * d1 }; + + for (int tile = 0; tile < QK8_0 / 4; ++tile) { + const uint8_t bits_lo = b_ptr[l].qs[k * 16 + 2 * tile + 0]; + const uint8_t bits_hi = b_ptr[l].qs[k * 16 + 2 * tile + 1]; + + for (int p = 0; p < 4; ++p) { + const float q = (float) a_blk->qs[tile * 4 + p]; + + sumf[0] += ((bits_lo & (1u << p)) ? scale[0] : -scale[0]) * q; + sumf[1] += ((bits_lo & (1u << (4 + p))) ? scale[1] : -scale[1]) * q; + sumf[2] += ((bits_hi & (1u << p)) ? scale[2] : -scale[2]) * q; + sumf[3] += ((bits_hi & (1u << (4 + p))) ? scale[3] : -scale[3]) * q; + } + } + } + } + + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j]; + } + } +} + +void ggml_gemv_q1_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 = QK1_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]; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0f; + } + + for (int l = 0; l < nb; l++) { + const float d0[4] = { + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]), + }; + + for (int k = 0; k < qk / blocklen; ++k) { + const block_q8_0 * GGML_RESTRICT a_blk = a_ptr + l * (qk / QK8_0) + k / (QK8_0 / blocklen); + const float d1 = GGML_CPU_FP16_TO_FP32(a_blk->d); + const float scale[4] = { d0[0] * d1, d0[1] * d1, d0[2] * d1, d0[3] * d1 }; + const uint8_t bits0 = b_ptr[l].qs[k * ncols_interleaved + 0]; + const uint8_t bits1 = b_ptr[l].qs[k * ncols_interleaved + 1]; + const uint8_t bits2 = b_ptr[l].qs[k * ncols_interleaved + 2]; + const uint8_t bits3 = b_ptr[l].qs[k * ncols_interleaved + 3]; + const int q_offset = (k % (QK8_0 / blocklen)) * blocklen; + + for (int p = 0; p < blocklen; ++p) { + const float q = (float) a_blk->qs[q_offset + p]; + + sumf[0] += ((bits0 & (1u << p)) ? scale[0] : -scale[0]) * q; + sumf[1] += ((bits1 & (1u << p)) ? scale[1] : -scale[1]) * q; + sumf[2] += ((bits2 & (1u << p)) ? scale[2] : -scale[2]) * q; + sumf[3] += ((bits3 & (1u << p)) ? scale[3] : -scale[3]) * q; + } + } + } + + 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) { @@ -2383,6 +2510,176 @@ void ggml_gemm_q8_0_4x8_q8_0_generic(int n, } } +void ggml_gemm_q1_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 = QK1_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][4]; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q1_0x4 * b_ptr = (const block_q1_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++) { + const float d0[4] = { + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]), + }; + + for (int k = 0; k < QK1_0 / QK8_0; ++k) { + const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k; + const float a_d[4] = { + GGML_CPU_FP16_TO_FP32(a_blk->d[0]), + GGML_CPU_FP16_TO_FP32(a_blk->d[1]), + GGML_CPU_FP16_TO_FP32(a_blk->d[2]), + GGML_CPU_FP16_TO_FP32(a_blk->d[3]), + }; + + for (int tile = 0; tile < QK8_0 / 4; ++tile) { + const uint8_t bits_lo = b_ptr[l].qs[k * 16 + 2 * tile + 0]; + const uint8_t bits_hi = b_ptr[l].qs[k * 16 + 2 * tile + 1]; + const int tile_offset = tile * 16; + + for (int p = 0; p < 4; ++p) { + const int8_t q_row[4] = { + a_blk->qs[tile_offset + 0 * 4 + p], + a_blk->qs[tile_offset + 1 * 4 + p], + a_blk->qs[tile_offset + 2 * 4 + p], + a_blk->qs[tile_offset + 3 * 4 + p], + }; + const int sign[4] = { + (bits_lo & (1u << p)) ? 1 : -1, + (bits_lo & (1u << (4 + p))) ? 1 : -1, + (bits_hi & (1u << p)) ? 1 : -1, + (bits_hi & (1u << (4 + p))) ? 1 : -1, + }; + + for (int m = 0; m < 4; ++m) { + const float row_scale = a_d[m]; + sumf[m][0] += sign[0] * q_row[m] * d0[0] * row_scale; + sumf[m][1] += sign[1] * q_row[m] * d0[1] * row_scale; + sumf[m][2] += sign[2] * q_row[m] * d0[2] * row_scale; + sumf[m][3] += sign[3] * q_row[m] * d0[3] * row_scale; + } + } + } + } + } + + 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_q1_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 = QK1_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]; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (4 * y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q1_0x4 * b_ptr = (const block_q1_0x4 *) 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++) { + const float d0[4] = { + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[0]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[1]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[2]), + GGML_CPU_FP16_TO_FP32(b_ptr[l].d[3]), + }; + + for (int k = 0; k < qk / blocklen; ++k) { + const block_q8_0x4 * GGML_RESTRICT a_blk = a_ptr + 4 * l + k / (QK8_0 / blocklen); + const float a_d[4] = { + GGML_CPU_FP16_TO_FP32(a_blk->d[0]), + GGML_CPU_FP16_TO_FP32(a_blk->d[1]), + GGML_CPU_FP16_TO_FP32(a_blk->d[2]), + GGML_CPU_FP16_TO_FP32(a_blk->d[3]), + }; + const uint8_t bits0 = b_ptr[l].qs[k * ncols_interleaved + 0]; + const uint8_t bits1 = b_ptr[l].qs[k * ncols_interleaved + 1]; + const uint8_t bits2 = b_ptr[l].qs[k * ncols_interleaved + 2]; + const uint8_t bits3 = b_ptr[l].qs[k * ncols_interleaved + 3]; + const int q_offset = (k % (QK8_0 / blocklen)) * 4 * blocklen; + + for (int p = 0; p < blocklen; ++p) { + const int8_t q_row[4] = { + a_blk->qs[q_offset + 0 * blocklen + p], + a_blk->qs[q_offset + 1 * blocklen + p], + a_blk->qs[q_offset + 2 * blocklen + p], + a_blk->qs[q_offset + 3 * blocklen + p], + }; + const int sign[4] = { + (bits0 & (1u << p)) ? 1 : -1, + (bits1 & (1u << p)) ? 1 : -1, + (bits2 & (1u << p)) ? 1 : -1, + (bits3 & (1u << p)) ? 1 : -1, + }; + + for (int m = 0; m < 4; ++m) { + const float row_scale = a_d[m]; + sumf[m][0] += sign[0] * q_row[m] * d0[0] * row_scale; + sumf[m][1] += sign[1] * q_row[m] * d0[1] * row_scale; + sumf[m][2] += sign[2] * q_row[m] * d0[2] * row_scale; + sumf[m][3] += sign[3] * q_row[m] * d0[3] * row_scale; + } + } + } + } + + 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) { @@ -2739,6 +3036,50 @@ static block_q8_0x4 make_block_q8_0x4(block_q8_0 * in, unsigned int blck_size_in return out; } +static block_q1_0x4 make_block_q1_0x4(block_q1_0 * in, unsigned int blck_size_interleave) { + block_q1_0x4 out; + + for (int i = 0; i < 4; i++) { + out.d[i] = in[i].d; + } + + GGML_ASSERT(blck_size_interleave == 4 || blck_size_interleave == 8); + + if (blck_size_interleave == 4) { + for (int k = 0; k < QK1_0 / QK8_0; ++k) { + for (int tile = 0; tile < QK8_0 / 4; ++tile) { + uint8_t packed_lo = 0; + uint8_t packed_hi = 0; + + const int weight_base = k * QK8_0 + tile * 4; + for (int pos = 0; pos < 4; ++pos) { + const int weight_idx = weight_base + pos; + const int byte_idx = weight_idx / 8; + const int bit_idx = weight_idx % 8; + + packed_lo |= ((in[0].qs[byte_idx] >> bit_idx) & 1u) << pos; + packed_lo |= ((in[1].qs[byte_idx] >> bit_idx) & 1u) << (4 + pos); + packed_hi |= ((in[2].qs[byte_idx] >> bit_idx) & 1u) << pos; + packed_hi |= ((in[3].qs[byte_idx] >> bit_idx) & 1u) << (4 + pos); + } + + out.qs[k * 16 + 2 * tile + 0] = packed_lo; + out.qs[k * 16 + 2 * tile + 1] = packed_hi; + } + } + return out; + } + + for (int byte_idx = 0; byte_idx < QK1_0 / 8; ++byte_idx) { + out.qs[byte_idx * 4 + 0] = in[0].qs[byte_idx]; + out.qs[byte_idx * 4 + 1] = in[1].qs[byte_idx]; + out.qs[byte_idx * 4 + 2] = in[2].qs[byte_idx]; + out.qs[byte_idx * 4 + 3] = in[3].qs[byte_idx]; + } + + return out; +} + static block_q4_0x4 make_block_q4_0x4(block_q4_0 * in, int blck_size_interleave) { block_q4_0x4 out; @@ -3509,6 +3850,38 @@ static int repack_q8_0_to_q8_0_4_bl(struct ggml_tensor * t, return 0; } +static int repack_q1_0_to_q1_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_Q1_0); + GGML_ASSERT(interleave_block == 4 || interleave_block == 8); + constexpr int nrows_interleaved = 4; + + block_q1_0x4 * dst = (block_q1_0x4 *) t->data; + const block_q1_0 * src = (const block_q1_0 *) data; + block_q1_0 dst_tmp[4]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK1_0; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q1_0)); + + 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_q1_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; @@ -3865,6 +4238,14 @@ 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_q1_0_to_q1_0_4_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q1_0_to_q1_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_4_bl(t, 4, data, data_size); } @@ -3960,6 +4341,14 @@ template <> int repack(struct ggml_tensor * t, const void * d 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_q1_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_q1_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_4x4_q8_0(n, s, bs, vx, vy, nr, nc); } @@ -4057,6 +4446,14 @@ template <> void gemv(int n, float * s, size_ 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_q1_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_q1_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_4x4_q8_0(n, s, bs, vx, vy, nr, nc); } @@ -4526,6 +4923,10 @@ template q1_0_4x4_q8_0; + static const ggml::cpu::repack::tensor_traits q1_0_4x8_q8_0; + // 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; @@ -4723,6 +5124,17 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons } #endif } + } else if (cur->type == GGML_TYPE_Q1_0) { + if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { + if (cur->ne[1] % 4 == 0) { + return &q1_0_4x8_q8_0; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 4 == 0) { + return &q1_0_4x4_q8_0; + } + } } return nullptr; diff --git a/ggml/src/ggml-cpu/repack.h b/ggml/src/ggml-cpu/repack.h index cb21edf62..fc6715c39 100644 --- a/ggml/src/ggml-cpu/repack.h +++ b/ggml/src/ggml-cpu/repack.h @@ -11,6 +11,9 @@ ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void); template constexpr int QK_0() { + if constexpr (K == 1) { + return QK1_0; + } if constexpr (K == 4) { return QK4_0; } @@ -26,6 +29,7 @@ template struct block { }; // control size +static_assert(sizeof(block<1, 4>) == 4 * sizeof(ggml_half) + QK1_0 / 2, "wrong block<1,4> size/padding"); 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"); @@ -33,6 +37,7 @@ static_assert(sizeof(block<8, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 4, "wrong b 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_q1_0x4 = block<1, 4>; using block_q4_0x4 = block<4, 4>; using block_q4_0x8 = block<4, 8>; using block_q4_0x16 = block<4, 16>; @@ -141,6 +146,8 @@ void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTR 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_q1_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_q1_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_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); @@ -157,6 +164,8 @@ void ggml_gemv_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const v 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_q1_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_q1_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); @@ -193,6 +202,8 @@ void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GG 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_q1_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_q1_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_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); @@ -209,6 +220,8 @@ void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, 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_q1_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_q1_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); From 1aa2954bde90b1cb4d2dca96f90b07d7b155124b Mon Sep 17 00:00:00 2001 From: Anant Shrivastava Date: Mon, 21 Sep 2026 13:37:04 +0530 Subject: [PATCH 04/42] sycl : coalesce MKL-FA softmax loads instead of one work-item per row (#28918) * sycl : coalesce MKL-FA softmax loads instead of one work-item per row * better human readable variable name --- ggml/src/ggml-sycl/fattn-mkl.cpp | 77 +++++++++++++++++--------------- 1 file changed, 41 insertions(+), 36 deletions(-) diff --git a/ggml/src/ggml-sycl/fattn-mkl.cpp b/ggml/src/ggml-sycl/fattn-mkl.cpp index 2d164a084..30947b17b 100644 --- a/ggml/src/ggml-sycl/fattn-mkl.cpp +++ b/ggml/src/ggml-sycl/fattn-mkl.cpp @@ -110,8 +110,15 @@ static void mkl_fa_init_softmax_state( // The tile spans absolute rows [q0, q0 + q_rows). Score buffers // (KQ_f32/S_f16) are indexed RELATIVE to the tile; the persistent state // (VKQ_accum/KQ_max/KQ_sum) and mask are indexed by ABSOLUTE row. -// For each row: find local max → rescale previous VKQ_accum → -// compute exp(s - max) → write S_f16 → update running max/sum. +// One WORK-GROUP per query row (local size = wg_size): work-items stride +// over the chunk so adjacent items touch adjacent elements (coalesced), +// the row max/sum come from group reductions, and the DV-long VKQ +// rescale is spread across the items. Item 0 is the sole writer of +// KQ_max/KQ_sum; its writes are ordered after every other item's reads +// by the second group reduction (a collective). Per-element math is +// identical to the original one-item-per-row kernel: softcap before +// mask, native::exp, -1e30 sentinel, half-precision S. Only the float +// summation order differs (tree vs serial), i.e. last-ulp level. static void mkl_fa_online_softmax_chunk( dpct::queue_ptr stream, float * __restrict KQ_f32, @@ -126,25 +133,27 @@ static void mkl_fa_online_softmax_chunk( int64_t mask_row_stride, int mask_n_heads, float logit_softcap, int64_t wg_size) { - const int64_t wg = ((q_rows + wg_size - 1) / wg_size) * wg_size; - + // One work-group per query row: exactly q_rows groups of wg_size + // items. q_rows * wg_size is already a multiple of wg_size, so unlike + // the one-item-per-row kernels there is no round-up / tail guard. + const int64_t wg = q_rows * wg_size; + const int local_size = (int) wg_size; // stride in the loops below stream->submit([&](sycl::handler & cgh) { cgh.parallel_for(sycl::nd_range<1>(wg, wg_size), [=](sycl::nd_item<1> item) { - int jc_rel = item.get_global_id(0); - if (jc_rel >= q_rows) return; - int jc_abs = q0 + jc_rel; - + const int local_id = (int)item.get_local_id(0); + const int row = (int)item.get_group(0); // tile-relative + const int jc_abs = q0 + row; const int gqa_group = jc_abs / n_queries; const int q_row = jc_abs % n_queries; - // Score buffers are tile-local (relative index). const float * __restrict KQ_row = KQ_f32 - + jc_rel * (int64_t)chunk_size; + + row * (int64_t)chunk_size; + sycl::half * __restrict S_row = S_f16 + + row * (int64_t)chunk_size; // Persistent accumulator is full-sized (absolute index). float * __restrict vkq = VKQ_accum + jc_abs * (int64_t)DV; - const sycl::half * mask_h = nullptr; int64_t m_stride = 0; if (mask_data) { @@ -153,10 +162,8 @@ static void mkl_fa_online_softmax_chunk( mask_h = mask_data + (int64_t)m_head * mask_head_stride; m_stride = mask_row_stride; } - - // Row-wise local maximum (softcap before mask) - float local_max = -1e30f; - for (int i = 0; i < chunk_size; i++) { + // Score at chunk offset i — original per-element math. + auto score = [&](int i) { float s = KQ_row[i]; if (logit_softcap != 0.0f) { s = logit_softcap * sycl::tanh(s); @@ -165,40 +172,38 @@ static void mkl_fa_online_softmax_chunk( s += (float)mask_h[q_row * m_stride + (chunk_start + i)]; } + return s; + }; + // Pass 1: strided (coalesced) row-wise local maximum. + float local_max = -1e30f; + for (int i = local_id; i < chunk_size; i += local_size) { + float s = score(i); if (s > local_max) local_max = s; } - + const float final_local_max = sycl::reduce_over_group( + item.get_group(), local_max, sycl::maximum()); // Rescale previous accumulator by exp(old_max - new_max) float old_max = KQ_max[jc_abs]; - float new_max = (old_max > local_max) ? old_max : local_max; + float new_max = (old_max > final_local_max) ? old_max : final_local_max; float rescale = (old_max < -1e29f) ? 1.0f : sycl::native::exp(old_max - new_max); - - for (int v = 0; v < DV; v++) { + for (int v = local_id; v < DV; v += local_size) { vkq[v] *= rescale; } - - // Softmax and write S_f16 (tile-local index) + // Pass 2: softmax numerators, strided; S row written once. float local_sum = 0.0f; - sycl::half * __restrict S_row = S_f16 - + jc_rel * (int64_t)chunk_size; - - for (int i = 0; i < chunk_size; i++) { - float s = KQ_row[i]; - if (logit_softcap != 0.0f) { - s = logit_softcap * sycl::tanh(s); - } - if (mask_h) { - s += (float)mask_h[q_row * m_stride - + (chunk_start + i)]; - } + for (int i = local_id; i < chunk_size; i += local_size) { + float s = score(i); float val = sycl::native::exp(s - new_max); S_row[i] = sycl::half(val); local_sum += val; } - - KQ_sum[jc_abs] = KQ_sum[jc_abs] * rescale + local_sum; - KQ_max[jc_abs] = new_max; + const float total_sum = sycl::reduce_over_group( + item.get_group(), local_sum, sycl::plus()); + if (local_id == 0) { + KQ_sum[jc_abs] = KQ_sum[jc_abs] * rescale + total_sum; + KQ_max[jc_abs] = new_max; + } }); }); } From 26394b4e6749a41c3633db040e0987500a5f7013 Mon Sep 17 00:00:00 2001 From: Silverside Date: Mon, 21 Sep 2026 03:32:07 -0500 Subject: [PATCH 05/42] json: Fixed json enum handling (#28518) * Fixed json enum handling Added common_json_value handling for enum values. Added tests/test-json.cpp to cover testing of some aspects of common_json. * Removed tests as requested. * Applied recommended style and simplification Simplified by delegating enum constructor to the constructor of the underlying type Matched style of surrounding templating code --- common/json.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/common/json.h b/common/json.h index f3ad4edee..95b6d6f21 100644 --- a/common/json.h +++ b/common/json.h @@ -82,6 +82,9 @@ struct common_json_value { // note: a nested pair {"a", "b"} does not build, use common_json::array({"a", "b"}) for an array common_json_value(std::initializer_list items); + template ::value, int>::type = 0> + common_json_value(T val) : common_json_value((typename std::underlying_type::type) val) {} + template ::value && !std::is_same::value, int>::type = 0> common_json_value(T val) : type(std::is_signed::value ? VAL_INT : VAL_UINT) { if (std::is_signed::value) { @@ -111,6 +114,7 @@ struct common_json_item { // the types common_json_value holds on its own // anything else reaches its common_json ctor and recurses forever template struct common_json_is_value : std::integral_constant::value || std::is_arithmetic::value || std::is_same::value || std::is_same::value || From 335b21fcbda972777e4e9e69decad1f719cafeb3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 21 Sep 2026 12:37:24 +0300 Subject: [PATCH 06/42] ggml-metal : simplify fusion pattern op list declaration (#29206) * ggml-metal : derive non-empty fusion ops from ops_all Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * ggml-metal : drop _all suffix from fusion op pattern vectors Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp --- ggml/src/ggml-metal/ggml-metal-fusion.cpp | 131 ++++++++++++---------- 1 file changed, 70 insertions(+), 61 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-fusion.cpp b/ggml/src/ggml-metal/ggml-metal-fusion.cpp index 5a160a18d..eac3bd6fe 100644 --- a/ggml/src/ggml-metal/ggml-metal-fusion.cpp +++ b/ggml/src/ggml-metal/ggml-metal-fusion.cpp @@ -10,10 +10,23 @@ #include #include +// derive the non-empty op sequence from the raw `ops_all` sequence +static std::vector ggml_metal_fusion_filter_ops(const std::vector & ops_all) { + std::vector ops; + + for (ggml_op op : ops_all) { + if (!ggml_op_is_empty(op)) { + ops.push_back(op); + } + } + + return ops; +} + struct ggml_metal_fusion { ggml_metal_fusion_id id; - std::vector ops; // op sequence (fixed length, non-empty nodes) + std::vector ops; // non-empty op sequence, derived from ops_all std::vector ops_all; // full raw op sequence (may include empty RESHAPE/VIEW nodes) std::vector outs; // additional fused output nodes, relative to ops @@ -30,6 +43,25 @@ struct ggml_metal_fusion { const int * node_idxs, int idx, ggml_metal_fusion_mode mode); + + ggml_metal_fusion( + ggml_metal_fusion_id id, + const std::vector & ops_all, + const std::vector & outs, + bool unsafe, + bool (*check)(const struct ggml_metal_fusion * fusion, + const struct ggml_tensor * const * nodes, + const struct ggml_cgraph * gf, + const int * node_idxs, + int idx, + ggml_metal_fusion_mode mode)) + : id(id), + ops(ggml_metal_fusion_filter_ops(ops_all)), + ops_all(ops_all), + outs(outs), + unsafe(unsafe), + check(check) { + } }; ggml_metal_fusion_id ggml_metal_fusion_get_id(const ggml_metal_fusion * fusion) { @@ -297,17 +329,17 @@ static bool ggml_metal_fusion_check_snake( // SOFT_MAX + ARGSORT + GET_ROWS (plus optional norm/scale) for MoE routing. // This is a multi-output elision chain: the fused kernel writes both the selected // expert ids and the gathered/normalized routing weights. -static const std::vector ops_topk_moe_all = { +static const std::vector ops_topk_moe = { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }; -static const std::vector ops_topk_moe_scale_all = { +static const std::vector ops_topk_moe_scale = { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_SCALE }; -static const std::vector ops_topk_moe_norm_all = { +static const std::vector ops_topk_moe_norm = { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE }; -static const std::vector ops_topk_moe_norm_scale_all = { +static const std::vector ops_topk_moe_norm_scale = { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE }; @@ -607,86 +639,63 @@ static const std::vector ops_snake = { GGML_OP_MUL, GGML_OP_SIN, GGML_O static const std::vector ops_gdn_cache = { GGML_OP_GATED_DELTA_NET, GGML_OP_CPY }; -static const std::vector ops_topk_moe = { - GGML_OP_SOFT_MAX, GGML_OP_ARGSORT, GGML_OP_GET_ROWS -}; -static const std::vector ops_topk_moe_scale = { - GGML_OP_SOFT_MAX, GGML_OP_ARGSORT, GGML_OP_GET_ROWS, GGML_OP_SCALE -}; -static const std::vector ops_topk_moe_norm = { - GGML_OP_SOFT_MAX, GGML_OP_ARGSORT, GGML_OP_GET_ROWS, - GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV -}; -static const std::vector ops_topk_moe_norm_scale = { - GGML_OP_SOFT_MAX, GGML_OP_ARGSORT, GGML_OP_GET_ROWS, - GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_SCALE -}; - static const std::vector ops_ssm_conv_silu = { GGML_OP_SSM_CONV, GGML_OP_UNARY }; -static const std::vector ops_moe_reduce_2 = { GGML_OP_MUL, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_3 = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_4 = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_5 = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_6 = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_7 = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_8 = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; - -static const std::vector ops_moe_reduce_all_2 = { +static const std::vector ops_moe_reduce_2 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_all_3 = { +static const std::vector ops_moe_reduce_3 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_all_4 = { +static const std::vector ops_moe_reduce_4 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_all_5 = { +static const std::vector ops_moe_reduce_5 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_all_6 = { +static const std::vector ops_moe_reduce_6 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_all_7 = { +static const std::vector ops_moe_reduce_7 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; -static const std::vector ops_moe_reduce_all_8 = { +static const std::vector ops_moe_reduce_8 = { GGML_OP_MUL, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_VIEW, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; static const std::vector ggml_metal_fusions = { - { GGML_METAL_FUSION_NORM_MUL, ops_norm_mul, ops_norm_mul, {}, false, ggml_metal_fusion_check_norm }, - { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add, ops_norm_mul_add, {}, false, ggml_metal_fusion_check_norm }, - { GGML_METAL_FUSION_NORM_SCALE, ops_norm_scale, ops_norm_scale, {}, false, ggml_metal_fusion_check_norm }, - { GGML_METAL_FUSION_NORM_MUL, ops_rms_norm_mul, ops_rms_norm_mul, {}, false, ggml_metal_fusion_check_norm }, - { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, ops_rms_norm_mul_add, {}, false, ggml_metal_fusion_check_norm }, - { GGML_METAL_FUSION_NORM_SCALE, ops_rms_norm_scale, ops_rms_norm_scale, {}, false, ggml_metal_fusion_check_norm }, - { GGML_METAL_FUSION_ADD_CHAIN, ops_add_2, ops_add_2, {}, false, ggml_metal_fusion_check_add_chain }, - { GGML_METAL_FUSION_ADD_CHAIN, ops_add_3, ops_add_3, {}, false, ggml_metal_fusion_check_add_chain }, - { GGML_METAL_FUSION_ADD_CHAIN, ops_add_4, ops_add_4, {}, false, ggml_metal_fusion_check_add_chain }, - { GGML_METAL_FUSION_ADD_CHAIN, ops_add_5, ops_add_5, {}, false, ggml_metal_fusion_check_add_chain }, - { GGML_METAL_FUSION_ADD_CHAIN, ops_add_6, ops_add_6, {}, false, ggml_metal_fusion_check_add_chain }, - { GGML_METAL_FUSION_ADD_CHAIN, ops_add_7, ops_add_7, {}, false, ggml_metal_fusion_check_add_chain }, - { GGML_METAL_FUSION_SNAKE, ops_snake, ops_snake, {}, false, ggml_metal_fusion_check_snake }, - { GGML_METAL_FUSION_GDN_CACHE, ops_gdn_cache, ops_gdn_cache, {}, true, ggml_metal_fusion_check_gdn_cache }, - { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe, ops_topk_moe_all, {1}, true, ggml_metal_fusion_check_topk_moe }, - { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_scale, ops_topk_moe_scale_all, {1}, true, ggml_metal_fusion_check_topk_moe }, - { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm, ops_topk_moe_norm_all, {1}, true, ggml_metal_fusion_check_topk_moe }, - { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm_scale, ops_topk_moe_norm_scale_all, {1}, true, ggml_metal_fusion_check_topk_moe }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_2, ops_moe_reduce_all_2, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_3, ops_moe_reduce_all_3, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_4, ops_moe_reduce_all_4, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_5, ops_moe_reduce_all_5, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_6, ops_moe_reduce_all_6, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_7, ops_moe_reduce_all_7, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_8, ops_moe_reduce_all_8, {}, true, ggml_metal_fusion_check_moe_reduce }, - { GGML_METAL_FUSION_SSM_CONV_SILU, ops_ssm_conv_silu, ops_ssm_conv_silu, {}, false, ggml_metal_fusion_check_ssm_conv_silu }, + { GGML_METAL_FUSION_NORM_MUL, ops_norm_mul, {}, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add, {}, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_SCALE, ops_norm_scale, {}, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_MUL, ops_rms_norm_mul, {}, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, {}, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_SCALE, ops_rms_norm_scale, {}, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_2, {}, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_3, {}, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_4, {}, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_5, {}, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_6, {}, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_7, {}, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_SNAKE, ops_snake, {}, false, ggml_metal_fusion_check_snake }, + { GGML_METAL_FUSION_GDN_CACHE, ops_gdn_cache, {}, true, ggml_metal_fusion_check_gdn_cache }, + { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe, {1}, true, ggml_metal_fusion_check_topk_moe }, + { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_scale, {1}, true, ggml_metal_fusion_check_topk_moe }, + { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm, {1}, true, ggml_metal_fusion_check_topk_moe }, + { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm_scale, {1}, true, ggml_metal_fusion_check_topk_moe }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_2, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_3, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_4, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_5, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_6, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_7, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_8, {}, true, ggml_metal_fusion_check_moe_reduce }, + { GGML_METAL_FUSION_SSM_CONV_SILU, ops_ssm_conv_silu, {}, false, ggml_metal_fusion_check_ssm_conv_silu }, }; // ---- alloc deps ----------------------------------------------------------- From 711f60beeb9e983f6d8e01bda25302dee9226c8d Mon Sep 17 00:00:00 2001 From: Mostafa Date: Mon, 21 Sep 2026 13:38:48 +0300 Subject: [PATCH 07/42] tests : remove stale comment (#29140) --- tests/test-backend-ops.cpp | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 5c7de196d..887add245 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6219,13 +6219,6 @@ struct test_conv_2d : public test_case { // Whether the inputs are contiguous in the channel dim or the width dim const bool cwhn; - // If true, the direct CONV_2D will be used in the graph, otherwise it - // uses ggml_conv_2d: - // * if the program is called with -o CONV_2D_DIRECT_IMPL, the - // CONV_2D graph will be built, while - // * if the program is called with -o CONV_2D_INDIRECT_IMPL, the - // IM2COL -> MUL_MM graph will be built. - std::string vars() override { return VARS_TO_STR10(ne_input, ne_kernel, type_kernel, stride0, stride1, padding0, padding1, dilation0, dilation1, cwhn); } From 982a3329af8401772087d67c54d5948a5b270943 Mon Sep 17 00:00:00 2001 From: Nandan Vallamdasu Date: Mon, 21 Sep 2026 16:15:58 +0530 Subject: [PATCH 08/42] server : do not forward --api-key-file to router-spawned child instances (#28938) In router mode, authentication belongs to the router. unset_reserved_args() already unset LLAMA_API_KEY, but did not unset LLAMA_ARG_API_KEY_FILE. When --api-key-file was passed, children re-validated against file keys only, causing clients using --api-key to 401 on chat completions (#28820). In addition, router internal calls without auth headers (such as POST /v1/streams/lookup and DELETE /v1/stream) were silently rejected with 401. Unset LLAMA_ARG_API_KEY_FILE in unset_reserved_args() so no API keys reach child instances. This keeps keys out of child argv, ensures all keys the router accepts work end-to-end, and prevents router internal stream calls from 401ing. Fixes #28820 --- tools/server/server-models.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index b10d9bd8a..9c00036e0 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -470,6 +470,7 @@ static void unset_reserved_args(common_preset & preset, bool unset_model_args) { preset.unset_option("LLAMA_ARG_SSL_KEY_FILE"); preset.unset_option("LLAMA_ARG_SSL_CERT_FILE"); preset.unset_option("LLAMA_API_KEY"); + preset.unset_option("LLAMA_ARG_API_KEY_FILE"); preset.unset_option("LLAMA_ARG_MODELS_DIR"); preset.unset_option("LLAMA_ARG_MODELS_MAX"); preset.unset_option("LLAMA_ARG_MODELS_PRESET"); From e0dff58475bc9ed68eedcb265ee998f2fcabb3b1 Mon Sep 17 00:00:00 2001 From: Mikolaj Kucharski Date: Mon, 21 Sep 2026 10:47:38 +0000 Subject: [PATCH 09/42] args: add env vars for temperature, top-p, min-p and penalties (#27380) Allow configuring --temp, --top-p, --min-p, --repeat-penalty, --presence-penalty and --frequency-penalty via LLAMA_ARG_* so llama-server can be fully controlled from an EnvironmentFile (e.g. systemd on Debian). Use `llama-gen-docs` to regenerate the readme files. --- common/arg.cpp | 12 ++++++------ tools/cli/README.md | 12 ++++++------ tools/completion/README.md | 12 ++++++------ tools/server/README.md | 12 ++++++------ 4 files changed, 24 insertions(+), 24 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index c4c4e143c..996ea75fe 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2016,7 +2016,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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()); + ).set_sampling().set_env("LLAMA_ARG_TEMPERATURE")); add_opt(common_arg( {"--top-k"}, "N", string_format("top-k sampling (default: %d, 0 = disabled)", params.sampling.top_k), @@ -2032,7 +2032,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.sampling.top_p = std::stof(value); params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_P; } - ).set_sampling()); + ).set_sampling().set_env("LLAMA_ARG_TOP_P")); add_opt(common_arg( {"--min-p"}, "N", string_format("min-p sampling (default: %.2f, 0.0 = disabled)", (double)params.sampling.min_p), @@ -2040,7 +2040,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.sampling.min_p = std::stof(value); params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIN_P; } - ).set_sampling()); + ).set_sampling().set_env("LLAMA_ARG_MIN_P")); 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), @@ -2096,7 +2096,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.sampling.penalty_repeat = penalty_repeat; params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT; } - ).set_sampling()); + ).set_sampling().set_env("LLAMA_ARG_REPEAT_PENALTY")); add_opt(common_arg( {"--presence-penalty"}, "N", string_format("repeat alpha presence penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_present), @@ -2107,7 +2107,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } params.sampling.penalty_present = penalty_present; } - ).set_sampling()); + ).set_sampling().set_env("LLAMA_ARG_PRESENCE_PENALTY")); add_opt(common_arg( {"--frequency-penalty"}, "N", string_format("repeat alpha frequency penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_freq), @@ -2118,7 +2118,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } params.sampling.penalty_freq = penalty_freq; } - ).set_sampling()); + ).set_sampling().set_env("LLAMA_ARG_FREQUENCY_PENALTY")); add_opt(common_arg( {"--dry-multiplier"}, "N", string_format("set DRY sampling multiplier (default: %.2f, 0.0 = disabled)", (double)params.sampling.dry_multiplier), diff --git a/tools/cli/README.md b/tools/cli/README.md index ea1f7aaf8..acecc74a6 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -106,18 +106,18 @@ | `-s, --seed SEED` | RNG seed (default: -1, use random seed for -1) | | `--sampler-seq, --sampling-seq SEQUENCE` | simplified sequence for samplers that will be used (default: edskypmxt) | | `--ignore-eos` | ignore end of stream token and continue generating (implies --logit-bias EOS-inf) | -| `--temp, --temperature N` | temperature (default: 0.80) | +| `--temp, --temperature N` | temperature (default: 0.80)
(env: LLAMA_ARG_TEMPERATURE) | | `--top-k N` | top-k sampling (default: 40, 0 = disabled)
(env: LLAMA_ARG_TOP_K) | -| `--top-p N` | top-p sampling (default: 0.95, 1.0 = disabled) | -| `--min-p N` | min-p sampling (default: 0.05, 0.0 = disabled) | +| `--top-p N` | top-p sampling (default: 0.95, 1.0 = disabled)
(env: LLAMA_ARG_TOP_P) | +| `--min-p N` | min-p sampling (default: 0.05, 0.0 = disabled)
(env: LLAMA_ARG_MIN_P) | | `--top-nsigma, --top-n-sigma N` | top-n-sigma sampling (default: -1.00, -1.0 = disabled) | | `--xtc-probability N` | xtc probability (default: 0.00, 0.0 = disabled) | | `--xtc-threshold N` | xtc threshold (default: 0.10, 1.0 = disabled) | | `--typical, --typical-p N` | locally typical sampling, parameter p (default: 1.00, 1.0 = disabled) | | `--repeat-last-n N` | last n tokens to consider for penalize (default: 64, 0 = disabled) | -| `--repeat-penalty N` | penalize repeat sequence of tokens (default: 1.00, 1.0 = disabled) | -| `--presence-penalty N` | repeat alpha presence penalty (default: 0.00, 0.0 = disabled) | -| `--frequency-penalty N` | repeat alpha frequency penalty (default: 0.00, 0.0 = disabled) | +| `--repeat-penalty N` | penalize repeat sequence of tokens (default: 1.00, 1.0 = disabled)
(env: LLAMA_ARG_REPEAT_PENALTY) | +| `--presence-penalty N` | repeat alpha presence penalty (default: 0.00, 0.0 = disabled)
(env: LLAMA_ARG_PRESENCE_PENALTY) | +| `--frequency-penalty N` | repeat alpha frequency penalty (default: 0.00, 0.0 = disabled)
(env: LLAMA_ARG_FREQUENCY_PENALTY) | | `--dry-multiplier N` | set DRY sampling multiplier (default: 0.00, 0.0 = disabled) | | `--dry-base N` | set DRY sampling base value (default: 1.75) | | `--dry-allowed-length N` | set allowed length for DRY sampling (default: 2) | diff --git a/tools/completion/README.md b/tools/completion/README.md index c9a4cccfc..e2ac0668c 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -189,18 +189,18 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `-s, --seed SEED` | RNG seed (default: -1, use random seed for -1) | | `--sampler-seq, --sampling-seq SEQUENCE` | simplified sequence for samplers that will be used (default: edskypmxt) | | `--ignore-eos` | ignore end of stream token and continue generating (implies --logit-bias EOS-inf) | -| `--temp, --temperature N` | temperature (default: 0.80) | +| `--temp, --temperature N` | temperature (default: 0.80)
(env: LLAMA_ARG_TEMPERATURE) | | `--top-k N` | top-k sampling (default: 40, 0 = disabled)
(env: LLAMA_ARG_TOP_K) | -| `--top-p N` | top-p sampling (default: 0.95, 1.0 = disabled) | -| `--min-p N` | min-p sampling (default: 0.05, 0.0 = disabled) | +| `--top-p N` | top-p sampling (default: 0.95, 1.0 = disabled)
(env: LLAMA_ARG_TOP_P) | +| `--min-p N` | min-p sampling (default: 0.05, 0.0 = disabled)
(env: LLAMA_ARG_MIN_P) | | `--top-nsigma, --top-n-sigma N` | top-n-sigma sampling (default: -1.00, -1.0 = disabled) | | `--xtc-probability N` | xtc probability (default: 0.00, 0.0 = disabled) | | `--xtc-threshold N` | xtc threshold (default: 0.10, 1.0 = disabled) | | `--typical, --typical-p N` | locally typical sampling, parameter p (default: 1.00, 1.0 = disabled) | | `--repeat-last-n N` | last n tokens to consider for penalize (default: 64, 0 = disabled) | -| `--repeat-penalty N` | penalize repeat sequence of tokens (default: 1.00, 1.0 = disabled) | -| `--presence-penalty N` | repeat alpha presence penalty (default: 0.00, 0.0 = disabled) | -| `--frequency-penalty N` | repeat alpha frequency penalty (default: 0.00, 0.0 = disabled) | +| `--repeat-penalty N` | penalize repeat sequence of tokens (default: 1.00, 1.0 = disabled)
(env: LLAMA_ARG_REPEAT_PENALTY) | +| `--presence-penalty N` | repeat alpha presence penalty (default: 0.00, 0.0 = disabled)
(env: LLAMA_ARG_PRESENCE_PENALTY) | +| `--frequency-penalty N` | repeat alpha frequency penalty (default: 0.00, 0.0 = disabled)
(env: LLAMA_ARG_FREQUENCY_PENALTY) | | `--dry-multiplier N` | set DRY sampling multiplier (default: 0.00, 0.0 = disabled) | | `--dry-base N` | set DRY sampling base value (default: 1.75) | | `--dry-allowed-length N` | set allowed length for DRY sampling (default: 2) | diff --git a/tools/server/README.md b/tools/server/README.md index ef9033404..0ee8df291 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -123,18 +123,18 @@ For the full list of features, please refer to [server's changelog](https://gith | `-s, --seed SEED` | RNG seed (default: -1, use random seed for -1) | | `--sampler-seq, --sampling-seq SEQUENCE` | simplified sequence for samplers that will be used (default: edskypmxt) | | `--ignore-eos` | ignore end of stream token and continue generating (implies --logit-bias EOS-inf) | -| `--temp, --temperature N` | temperature (default: 0.80) | +| `--temp, --temperature N` | temperature (default: 0.80)
(env: LLAMA_ARG_TEMPERATURE) | | `--top-k N` | top-k sampling (default: 40, 0 = disabled)
(env: LLAMA_ARG_TOP_K) | -| `--top-p N` | top-p sampling (default: 0.95, 1.0 = disabled) | -| `--min-p N` | min-p sampling (default: 0.05, 0.0 = disabled) | +| `--top-p N` | top-p sampling (default: 0.95, 1.0 = disabled)
(env: LLAMA_ARG_TOP_P) | +| `--min-p N` | min-p sampling (default: 0.05, 0.0 = disabled)
(env: LLAMA_ARG_MIN_P) | | `--top-nsigma, --top-n-sigma N` | top-n-sigma sampling (default: -1.00, -1.0 = disabled) | | `--xtc-probability N` | xtc probability (default: 0.00, 0.0 = disabled) | | `--xtc-threshold N` | xtc threshold (default: 0.10, 1.0 = disabled) | | `--typical, --typical-p N` | locally typical sampling, parameter p (default: 1.00, 1.0 = disabled) | | `--repeat-last-n N` | last n tokens to consider for penalize (default: 64, 0 = disabled) | -| `--repeat-penalty N` | penalize repeat sequence of tokens (default: 1.00, 1.0 = disabled) | -| `--presence-penalty N` | repeat alpha presence penalty (default: 0.00, 0.0 = disabled) | -| `--frequency-penalty N` | repeat alpha frequency penalty (default: 0.00, 0.0 = disabled) | +| `--repeat-penalty N` | penalize repeat sequence of tokens (default: 1.00, 1.0 = disabled)
(env: LLAMA_ARG_REPEAT_PENALTY) | +| `--presence-penalty N` | repeat alpha presence penalty (default: 0.00, 0.0 = disabled)
(env: LLAMA_ARG_PRESENCE_PENALTY) | +| `--frequency-penalty N` | repeat alpha frequency penalty (default: 0.00, 0.0 = disabled)
(env: LLAMA_ARG_FREQUENCY_PENALTY) | | `--dry-multiplier N` | set DRY sampling multiplier (default: 0.00, 0.0 = disabled) | | `--dry-base N` | set DRY sampling base value (default: 1.75) | | `--dry-allowed-length N` | set allowed length for DRY sampling (default: 2) | From 542e9202d7d14863562e1e908be0bf32aeec2c2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Mon, 21 Sep 2026 12:51:53 +0200 Subject: [PATCH 10/42] ci : refactor build-self-hosted into backend-specific workflows (#28991) * refactor build-self-hosted into backends * update workflow names * build -> ci * bump openvino * trigger on cpu and generic ggml changes --- .github/workflows/build-self-hosted.yml | 539 ------------------ .github/workflows/ci-self-hosted-cpu.yml | 112 ++++ .github/workflows/ci-self-hosted-cuda.yml | 124 ++++ .github/workflows/ci-self-hosted-kleidiai.yml | 85 +++ .github/workflows/ci-self-hosted-metal.yml | 59 ++ .github/workflows/ci-self-hosted-openvino.yml | 75 +++ .github/workflows/ci-self-hosted-vulkan.yml | 201 +++++++ .github/workflows/ci-self-hosted-webgpu.yml | 130 +++++ 8 files changed, 786 insertions(+), 539 deletions(-) delete mode 100644 .github/workflows/build-self-hosted.yml create mode 100644 .github/workflows/ci-self-hosted-cpu.yml create mode 100644 .github/workflows/ci-self-hosted-cuda.yml create mode 100644 .github/workflows/ci-self-hosted-kleidiai.yml create mode 100644 .github/workflows/ci-self-hosted-metal.yml create mode 100644 .github/workflows/ci-self-hosted-openvino.yml create mode 100644 .github/workflows/ci-self-hosted-vulkan.yml create mode 100644 .github/workflows/ci-self-hosted-webgpu.yml diff --git a/.github/workflows/build-self-hosted.yml b/.github/workflows/build-self-hosted.yml deleted file mode 100644 index d54f71ac5..000000000 --- a/.github/workflows/build-self-hosted.yml +++ /dev/null @@ -1,539 +0,0 @@ -name: CI (self-hosted) - -on: - workflow_dispatch: # allows manual triggering - push: - branches: - - master - paths: [ - '.github/workflows/build-self-hosted.yml', - 'ci/run.sh', - '**/CMakeLists.txt', - '**/.cmake', - '**/*.h', - '**/*.hpp', - '**/*.c', - '**/*.cpp', - '**/*.cu', - '**/*.cuh', - '**/*.swift', - '**/*.m', - '**/*.metal', - '**/*.comp', - '**/*.glsl', - '**/*.wgsl' - ] - - pull_request: - types: [opened, synchronize, reopened] - paths: [ - '.github/workflows/build-self-hosted.yml', - 'ci/run.sh', - '**/CMakeLists.txt', - '**/.cmake', - '**/*.h', - '**/*.hpp', - '**/*.c', - '**/*.cpp', - '**/*.cu', - '**/*.cuh', - '**/*.swift', - '**/*.m', - '**/*.metal', - '**/*.comp', - '**/*.glsl', - '**/*.wgsl' - ] - -concurrency: - group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} - cancel-in-progress: true - -env: - # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) - HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} - GGML_NLOOP: 3 - GGML_N_THREADS: 1 - LLAMA_ARG_LOG_COLORS: 1 - LLAMA_ARG_LOG_PREFIX: 1 - LLAMA_ARG_LOG_TIMESTAMPS: 1 - -jobs: - gpu-cuda: - runs-on: "hf-jobs-t4-small:cuda13" - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Install dependencies - run: | - sudo apt update - sudo apt install -y cmake libssl-dev time unzip wget python3 python3-venv python3-pip - - - name: ccache - uses: ggml-org/ccache-action@v1.2.24 - with: - restore: false - save: false - - - name: ccache-buckets-restore - uses: ./.github/actions/ccache-buckets - with: - key: self-hosted-gpu-cuda - folder: llama.cpp - hf_bucket: ggml-org/cache - - - name: Test - id: ggml-ci - run: | - nvidia-smi - GG_BUILD_CUDA=1 CUDACXX=/usr/local/cuda/bin/nvcc bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - - name: ccache-buckets-save - if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} - uses: ./.github/actions/ccache-buckets - env: - HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} - with: - key: self-hosted-gpu-cuda - folder: llama.cpp - evict-old-files: 1d - hf_bucket: ggml-org/cache - save: true - - gpu-rocm: - runs-on: [self-hosted, Linux, AMD] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Test - id: ggml-ci - # HIP_LAUNCH_BLOCKING=1: workaround for an async-execution correctness - # issue on integrated RDNA3.5 (gfx1151) where batched inference returns - # incorrect output (perplexity ~88 vs ~9.4). Serializing kernel launches - # restores correctness. Remove once the underlying ROCm/HIP issue is fixed. - env: - HIP_LAUNCH_BLOCKING: "1" - run: | - rocminfo - GG_BUILD_ROCM=1 GG_BUILD_AMDGPU_TARGETS=gfx1151 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - gpu-vulkan-nvidia-cm: - # runs-on: "hf-jobs-t4-small:ubuntu26_04" - runs-on: [self-hosted, Linux, NVIDIA] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - # - name: Install dependencies - # run: | - # sudo apt update - # sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan-dev glslc spirv-headers vulkan-tools mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip - - # - name: ccache - # uses: ggml-org/ccache-action@v1.2.24 - # with: - # restore: false - # save: false - - # - name: ccache-buckets-restore - # uses: ./.github/actions/ccache-buckets - # with: - # key: self-hosted-vulkan-nvidia-cm - # folder: llama.cpp - # hf_bucket: ggml-org/cache - - - name: Test - id: ggml-ci - run: | - vulkaninfo --summary - GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - # - name: ccache-buckets-save - # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} - # uses: ./.github/actions/ccache-buckets - # env: - # HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} - # with: - # key: self-hosted-vulkan-nvidia-cm - # folder: llama.cpp - # evict-old-files: 1d - # hf_bucket: ggml-org/cache - # save: true - - gpu-vulkan-nvidia-cm2: - # runs-on: "hf-jobs-t4-small:ubuntu26_04" - runs-on: [self-hosted, Linux, NVIDIA, COOPMAT2] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - # - name: Install dependencies - # run: | - # sudo apt update - # sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan-dev glslc spirv-headers vulkan-tools mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip - - # - name: ccache - # uses: ggml-org/ccache-action@v1.2.24 - # with: - # restore: false - # save: false - - # - name: ccache-buckets-restore - # uses: ./.github/actions/ccache-buckets - # with: - # key: self-hosted-vulkan-nvidia-cm2 - # folder: llama.cpp - # hf_bucket: ggml-org/cache - - - name: Test - id: ggml-ci - run: | - vulkaninfo --summary - GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - # - name: ccache-buckets-save - # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} - # uses: ./.github/actions/ccache-buckets - # env: - # HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} - # with: - # key: self-hosted-vulkan-nvidia-cm2 - # folder: llama.cpp - # evict-old-files: 1d - # hf_bucket: ggml-org/cache - # save: true - - gpu-webgpu-nvidia: - runs-on: "hf-jobs-t4-small:ubuntu26_04" - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Install dependencies - run: | - sudo apt update - sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip - - - name: ccache - uses: ggml-org/ccache-action@v1.2.24 - with: - restore: false - save: false - - - name: ccache-buckets-restore - uses: ./.github/actions/ccache-buckets - with: - key: self-hosted-webgpu-nvidia - folder: llama.cpp - hf_bucket: ggml-org/cache - - - name: Dawn Dependency - id: dawn-depends - run: | - DAWN_VERSION="v20260908.214631" - DAWN_OWNER="google" - DAWN_REPO="dawn" - DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release" - echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" - curl -L -o artifact.tar.gz \ - "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" - mkdir dawn - tar -xvf artifact.tar.gz -C dawn --strip-components=1 - - - name: Test - id: ggml-ci - run: | - GG_BUILD_WEBGPU=1 \ - GG_BUILD_WEBGPU_DAWN_PREFIX="$GITHUB_WORKSPACE/dawn" \ - GG_BUILD_WEBGPU_DAWN_DIR="$GITHUB_WORKSPACE/dawn/lib64/cmake/Dawn" \ - bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - - name: ccache-buckets-save - if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} - uses: ./.github/actions/ccache-buckets - env: - HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} - with: - key: self-hosted-webgpu-nvidia - folder: llama.cpp - evict-old-files: 1d - hf_bucket: ggml-org/cache - save: true - - # TODO: provision AMX-compatible machine - #cpu-amx: - # runs-on: [self-hosted, Linux, CPU, AMX] - - # steps: - # - name: Clone - # id: checkout - # uses: actions/checkout@v6 - - # - name: Test - # id: ggml-ci - # run: | - # bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - # TODO: provision AMD GPU machine - # amd-vulkan: - # runs-on: [self-hosted, Linux, AMD] - - # steps: - # - name: Clone - # id: checkout - # uses: actions/checkout@v6 - - # - name: Test - # id: ggml-ci - # run: | - # vulkaninfo --summary - # GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - # TODO: provision AMD GPU machine - # amd-rocm: - # runs-on: [self-hosted, Linux, AMD] - - # steps: - # - name: Clone - # id: checkout - # uses: actions/checkout@v6 - - # - name: Test - # id: ggml-ci - # run: | - # amd-smi static - # GG_BUILD_ROCM=1 GG_BUILD_AMDGPU_TARGETS="gfx1101" bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - gpu-metal: - runs-on: [self-hosted, macOS, ARM64] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Test - id: ggml-ci - run: | - GG_BUILD_METAL=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - gpu-webgpu-apple: - runs-on: [self-hosted, macOS, ARM64] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Dawn Dependency - id: dawn-depends - run: | - DAWN_VERSION="v20260908.214631" - DAWN_OWNER="google" - DAWN_REPO="dawn" - DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release" - echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" - curl -L -o artifact.tar.gz \ - "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" - mkdir dawn - tar -xvf artifact.tar.gz -C dawn --strip-components=1 - - - name: Test - id: ggml-ci - run: | - GG_BUILD_WEBGPU=1 GG_BUILD_WEBGPU_DAWN_PREFIX="$GITHUB_WORKSPACE/dawn" \ - bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - gpu-vulkan-apple: - runs-on: [self-hosted, macOS, ARM64] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Test - id: ggml-ci - run: | - vulkaninfo --summary - GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - gpu-vulkan-intel-linux: - runs-on: [self-hosted, Linux, Intel] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - with: - persist-credentials: false - - - name: Test - id: ggml-ci - run: | - vulkaninfo --summary - GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - gpu-vulkan-intel-windows: - runs-on: [self-hosted, Windows, X64, Intel] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Test - id: ggml-ci - shell: C:\msys64\usr\bin\bash.exe --noprofile --norc -eo pipefail "{0}" - env: - MSYSTEM: UCRT64 - CHERE_INVOKING: 1 - PATH: C:\msys64\ucrt64\bin;C:\msys64\usr\bin;C:\Windows\System32;${{ env.PATH }} - run: | - vulkaninfo --summary - # Skip python related tests with GG_BUILD_LOW_PERF=1 since Windows MSYS2 UCRT64 currently fails to create - # a valid python environment for testing - LLAMA_FATAL_WARNINGS=OFF GG_BUILD_NINJA=1 GG_BUILD_VULKAN=1 GG_BUILD_LOW_PERF=1 ./ci/run.sh ./results/llama.cpp ./mnt/llama.cpp - - gpu-openvino-low-perf: - runs-on: [self-hosted, Linux, Intel, OpenVINO] - - env: - # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.4" - OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3" - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Setup OpenVINO Toolkit - uses: ./.github/actions/linux-setup-openvino - with: - path: ./openvino_toolkit - version_major: ${{ env.OPENVINO_VERSION_MAJOR }} - version_full: ${{ env.OPENVINO_VERSION_FULL }} - - - name: Install OpenVINO dependencies - run: | - cd ./openvino_toolkit - chmod +x ./install_dependencies/install_openvino_dependencies.sh - echo "Y" | sudo -E ./install_dependencies/install_openvino_dependencies.sh - - - name: Test - id: ggml-ci - run: | - source ./openvino_toolkit/setupvars.sh - GG_BUILD_OPENVINO=1 GGML_OPENVINO_DEVICE=GPU GG_BUILD_LOW_PERF=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - cpu-x64-high-perf: - runs-on: [self-hosted, Linux, X64] - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Test - id: ggml-ci - run: | - LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - cpu-arm64-high-perf-graviton4: - runs-on: ah-ubuntu_24_04-c8g_8x - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Dependencies - id: depends - run: | - set -euxo pipefail - sudo apt-get update - sudo DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a \ - apt-get install -y \ - build-essential \ - python3-venv \ - gpg \ - wget \ - time \ - git-lfs - - git lfs install - - # install the latest cmake - sudo install -d /usr/share/keyrings - wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc \ - | gpg --dearmor \ - | sudo tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null - echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ jammy main' \ - | sudo tee /etc/apt/sources.list.d/kitware.list - sudo apt-get update - sudo apt-get install -y cmake - - - name: Test - id: ggml-ci - run: | - LLAMA_ARG_THREADS=$(nproc) \ - GG_BUILD_HIGH_PERF=1 \ - GG_BUILD_NO_BF16=1 \ - GG_BUILD_EXTRA_TESTS_0=1 \ - bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp - - cpu-arm64-graviton4-kleidiai: - runs-on: ah-ubuntu_24_04-c8g_8x - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Dependencies - id: depends - run: | - set -euxo pipefail - sudo apt-get update - sudo DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a \ - apt-get install -y \ - build-essential \ - python3-venv \ - gpg \ - wget \ - time \ - git-lfs - - git lfs install - - # install the latest cmake - sudo install -d /usr/share/keyrings - wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc \ - | gpg --dearmor \ - | sudo tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null - echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ jammy main' \ - | sudo tee /etc/apt/sources.list.d/kitware.list - sudo apt-get update - sudo apt-get install -y cmake - - - name: Test - id: ggml-ci - run: | - LLAMA_ARG_THREADS=$(nproc) \ - GG_BUILD_KLEIDIAI=1 \ - GG_BUILD_EXTRA_TESTS_0=1 \ - GG_BUILD_HIGH_PERF=1 \ - bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-cpu.yml b/.github/workflows/ci-self-hosted-cpu.yml new file mode 100644 index 000000000..18bc9f144 --- /dev/null +++ b/.github/workflows/ci-self-hosted-cpu.yml @@ -0,0 +1,112 @@ +name: CI (self-hosted CPU backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-cpu.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-cpu.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + cpu-x64-high-perf: + runs-on: [self-hosted, Linux, X64] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Test + id: ggml-ci + run: | + LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + cpu-arm64-high-perf-graviton4: + runs-on: ah-ubuntu_24_04-c8g_8x + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Dependencies + id: depends + run: | + set -euxo pipefail + sudo apt-get update + sudo DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a \ + apt-get install -y \ + build-essential \ + python3-venv \ + gpg \ + wget \ + time \ + git-lfs + + git lfs install + + # install the latest cmake + sudo install -d /usr/share/keyrings + wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc \ + | gpg --dearmor \ + | sudo tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null + echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ jammy main' \ + | sudo tee /etc/apt/sources.list.d/kitware.list + sudo apt-get update + sudo apt-get install -y cmake + + - name: Test + id: ggml-ci + run: | + LLAMA_ARG_THREADS=$(nproc) \ + GG_BUILD_HIGH_PERF=1 \ + GG_BUILD_NO_BF16=1 \ + GG_BUILD_EXTRA_TESTS_0=1 \ + bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + # TODO: provision AMX-compatible machine + #cpu-amx: + # runs-on: [self-hosted, Linux, CPU, AMX] + + # steps: + # - name: Clone + # id: checkout + # uses: actions/checkout@v6 + + # - name: Test + # id: ggml-ci + # run: | + # bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-cuda.yml b/.github/workflows/ci-self-hosted-cuda.yml new file mode 100644 index 000000000..5951b52d3 --- /dev/null +++ b/.github/workflows/ci-self-hosted-cuda.yml @@ -0,0 +1,124 @@ +name: CI (self-hosted CUDA backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-cuda.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp', + '**/*.cu', + '**/*.cuh' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-cuda.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**', + 'ggml/src/ggml-cuda/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + gpu-cuda: + runs-on: "hf-jobs-t4-small:cuda13" + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Install dependencies + run: | + sudo apt update + sudo apt install -y cmake libssl-dev time unzip wget python3 python3-venv python3-pip + + - name: ccache + uses: ggml-org/ccache-action@v1.2.24 + with: + restore: false + save: false + + - name: ccache-buckets-restore + uses: ./.github/actions/ccache-buckets + with: + key: self-hosted-gpu-cuda + folder: llama.cpp + hf_bucket: ggml-org/cache + + - name: Test + id: ggml-ci + run: | + nvidia-smi + GG_BUILD_CUDA=1 CUDACXX=/usr/local/cuda/bin/nvcc bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + - name: ccache-buckets-save + if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} + uses: ./.github/actions/ccache-buckets + env: + HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} + with: + key: self-hosted-gpu-cuda + folder: llama.cpp + evict-old-files: 1d + hf_bucket: ggml-org/cache + save: true + + gpu-rocm: + runs-on: [self-hosted, Linux, AMD] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Test + id: ggml-ci + # HIP_LAUNCH_BLOCKING=1: workaround for an async-execution correctness + # issue on integrated RDNA3.5 (gfx1151) where batched inference returns + # incorrect output (perplexity ~88 vs ~9.4). Serializing kernel launches + # restores correctness. Remove once the underlying ROCm/HIP issue is fixed. + env: + HIP_LAUNCH_BLOCKING: "1" + run: | + rocminfo + GG_BUILD_ROCM=1 GG_BUILD_AMDGPU_TARGETS=gfx1151 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + # TODO: provision AMD GPU machine + # amd-rocm: + # runs-on: [self-hosted, Linux, AMD] + + # steps: + # - name: Clone + # id: checkout + # uses: actions/checkout@v6 + + # - name: Test + # id: ggml-ci + # run: | + # amd-smi static + # GG_BUILD_ROCM=1 GG_BUILD_AMDGPU_TARGETS="gfx1101" bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-kleidiai.yml b/.github/workflows/ci-self-hosted-kleidiai.yml new file mode 100644 index 000000000..c955aa002 --- /dev/null +++ b/.github/workflows/ci-self-hosted-kleidiai.yml @@ -0,0 +1,85 @@ +name: CI (self-hosted KleidiAI backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-kleidiai.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-kleidiai.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + cpu-arm64-graviton4-kleidiai: + runs-on: ah-ubuntu_24_04-c8g_8x + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Dependencies + id: depends + run: | + set -euxo pipefail + sudo apt-get update + sudo DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a \ + apt-get install -y \ + build-essential \ + python3-venv \ + gpg \ + wget \ + time \ + git-lfs + + git lfs install + + # install the latest cmake + sudo install -d /usr/share/keyrings + wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc \ + | gpg --dearmor \ + | sudo tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null + echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ jammy main' \ + | sudo tee /etc/apt/sources.list.d/kitware.list + sudo apt-get update + sudo apt-get install -y cmake + + - name: Test + id: ggml-ci + run: | + LLAMA_ARG_THREADS=$(nproc) \ + GG_BUILD_KLEIDIAI=1 \ + GG_BUILD_EXTRA_TESTS_0=1 \ + GG_BUILD_HIGH_PERF=1 \ + bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-metal.yml b/.github/workflows/ci-self-hosted-metal.yml new file mode 100644 index 000000000..7af30e294 --- /dev/null +++ b/.github/workflows/ci-self-hosted-metal.yml @@ -0,0 +1,59 @@ +name: CI (self-hosted Metal backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-metal.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp', + '**/*.swift', + '**/*.m', + '**/*.metal' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-metal.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**', + 'ggml/src/ggml-metal/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + gpu-metal: + runs-on: [self-hosted, macOS, ARM64] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Test + id: ggml-ci + run: | + GG_BUILD_METAL=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-openvino.yml b/.github/workflows/ci-self-hosted-openvino.yml new file mode 100644 index 000000000..e0947c46e --- /dev/null +++ b/.github/workflows/ci-self-hosted-openvino.yml @@ -0,0 +1,75 @@ +name: CI (self-hosted OpenVINO backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-openvino.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-openvino.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**', + 'ggml/src/ggml-openvino/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + gpu-openvino-low-perf: + runs-on: [self-hosted, Linux, Intel, OpenVINO] + + env: + # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile + OPENVINO_VERSION_MAJOR: "2026.4" + OPENVINO_VERSION_FULL: "2026.4.0.22959.99c81491cc3" + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Setup OpenVINO Toolkit + uses: ./.github/actions/linux-setup-openvino + with: + path: ./openvino_toolkit + version_major: ${{ env.OPENVINO_VERSION_MAJOR }} + version_full: ${{ env.OPENVINO_VERSION_FULL }} + + - name: Install OpenVINO dependencies + run: | + cd ./openvino_toolkit + chmod +x ./install_dependencies/install_openvino_dependencies.sh + echo "Y" | sudo -E ./install_dependencies/install_openvino_dependencies.sh + + - name: Test + id: ggml-ci + run: | + source ./openvino_toolkit/setupvars.sh + GG_BUILD_OPENVINO=1 GGML_OPENVINO_DEVICE=GPU GG_BUILD_LOW_PERF=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-vulkan.yml b/.github/workflows/ci-self-hosted-vulkan.yml new file mode 100644 index 000000000..ffed4b09b --- /dev/null +++ b/.github/workflows/ci-self-hosted-vulkan.yml @@ -0,0 +1,201 @@ +name: CI (self-hosted Vulkan backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-vulkan.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp', + '**/*.comp', + '**/*.glsl' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-vulkan.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**', + 'ggml/src/ggml-vulkan/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + gpu-vulkan-nvidia-cm: + # runs-on: "hf-jobs-t4-small:ubuntu26_04" + runs-on: [self-hosted, Linux, NVIDIA] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + # - name: Install dependencies + # run: | + # sudo apt update + # sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan-dev glslc spirv-headers vulkan-tools mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip + + # - name: ccache + # uses: ggml-org/ccache-action@v1.2.24 + # with: + # restore: false + # save: false + + # - name: ccache-buckets-restore + # uses: ./.github/actions/ccache-buckets + # with: + # key: self-hosted-vulkan-nvidia-cm + # folder: llama.cpp + # hf_bucket: ggml-org/cache + + - name: Test + id: ggml-ci + run: | + vulkaninfo --summary + GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + # - name: ccache-buckets-save + # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} + # uses: ./.github/actions/ccache-buckets + # env: + # HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} + # with: + # key: self-hosted-vulkan-nvidia-cm + # folder: llama.cpp + # evict-old-files: 1d + # hf_bucket: ggml-org/cache + # save: true + + gpu-vulkan-nvidia-cm2: + # runs-on: "hf-jobs-t4-small:ubuntu26_04" + runs-on: [self-hosted, Linux, NVIDIA, COOPMAT2] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + # - name: Install dependencies + # run: | + # sudo apt update + # sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan-dev glslc spirv-headers vulkan-tools mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip + + # - name: ccache + # uses: ggml-org/ccache-action@v1.2.24 + # with: + # restore: false + # save: false + + # - name: ccache-buckets-restore + # uses: ./.github/actions/ccache-buckets + # with: + # key: self-hosted-vulkan-nvidia-cm2 + # folder: llama.cpp + # hf_bucket: ggml-org/cache + + - name: Test + id: ggml-ci + run: | + vulkaninfo --summary + GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + # - name: ccache-buckets-save + # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} + # uses: ./.github/actions/ccache-buckets + # env: + # HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} + # with: + # key: self-hosted-vulkan-nvidia-cm2 + # folder: llama.cpp + # evict-old-files: 1d + # hf_bucket: ggml-org/cache + # save: true + + gpu-vulkan-apple: + runs-on: [self-hosted, macOS, ARM64] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Test + id: ggml-ci + run: | + vulkaninfo --summary + GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + gpu-vulkan-intel-linux: + runs-on: [self-hosted, Linux, Intel] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + with: + persist-credentials: false + + - name: Test + id: ggml-ci + run: | + vulkaninfo --summary + GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + gpu-vulkan-intel-windows: + runs-on: [self-hosted, Windows, X64, Intel] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Test + id: ggml-ci + shell: C:\msys64\usr\bin\bash.exe --noprofile --norc -eo pipefail "{0}" + env: + MSYSTEM: UCRT64 + CHERE_INVOKING: 1 + PATH: C:\msys64\ucrt64\bin;C:\msys64\usr\bin;C:\Windows\System32;${{ env.PATH }} + run: | + vulkaninfo --summary + # Skip python related tests with GG_BUILD_LOW_PERF=1 since Windows MSYS2 UCRT64 currently fails to create + # a valid python environment for testing + LLAMA_FATAL_WARNINGS=OFF GG_BUILD_NINJA=1 GG_BUILD_VULKAN=1 GG_BUILD_LOW_PERF=1 ./ci/run.sh ./results/llama.cpp ./mnt/llama.cpp + + # TODO: provision AMD GPU machine + # amd-vulkan: + # runs-on: [self-hosted, Linux, AMD] + + # steps: + # - name: Clone + # id: checkout + # uses: actions/checkout@v6 + + # - name: Test + # id: ggml-ci + # run: | + # vulkaninfo --summary + # GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp diff --git a/.github/workflows/ci-self-hosted-webgpu.yml b/.github/workflows/ci-self-hosted-webgpu.yml new file mode 100644 index 000000000..a6a36a8eb --- /dev/null +++ b/.github/workflows/ci-self-hosted-webgpu.yml @@ -0,0 +1,130 @@ +name: CI (self-hosted WebGPU backend) + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/ci-self-hosted-webgpu.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + '**/*.h', + '**/*.hpp', + '**/*.c', + '**/*.cpp', + '**/*.wgsl' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/ci-self-hosted-webgpu.yml', + 'ci/run.sh', + '**/CMakeLists.txt', + '**/.cmake', + 'ggml/src/*', + 'ggml/src/ggml-cpu/**', + 'ggml/src/ggml-webgpu/**' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302) + HF_TOKEN: ${{ secrets.HF_TOKEN_CI }} + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + gpu-webgpu-nvidia: + runs-on: "hf-jobs-t4-small:ubuntu26_04" + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Install dependencies + run: | + sudo apt update + sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip + + - name: ccache + uses: ggml-org/ccache-action@v1.2.24 + with: + restore: false + save: false + + - name: ccache-buckets-restore + uses: ./.github/actions/ccache-buckets + with: + key: self-hosted-webgpu-nvidia + folder: llama.cpp + hf_bucket: ggml-org/cache + + - name: Dawn Dependency + id: dawn-depends + run: | + DAWN_VERSION="v20260908.214631" + DAWN_OWNER="google" + DAWN_REPO="dawn" + DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release" + echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" + curl -L -o artifact.tar.gz \ + "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" + mkdir dawn + tar -xvf artifact.tar.gz -C dawn --strip-components=1 + + - name: Test + id: ggml-ci + run: | + GG_BUILD_WEBGPU=1 \ + GG_BUILD_WEBGPU_DAWN_PREFIX="$GITHUB_WORKSPACE/dawn" \ + GG_BUILD_WEBGPU_DAWN_DIR="$GITHUB_WORKSPACE/dawn/lib64/cmake/Dawn" \ + bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp + + - name: ccache-buckets-save + if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} + uses: ./.github/actions/ccache-buckets + env: + HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }} + with: + key: self-hosted-webgpu-nvidia + folder: llama.cpp + evict-old-files: 1d + hf_bucket: ggml-org/cache + save: true + + gpu-webgpu-apple: + runs-on: [self-hosted, macOS, ARM64] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Dawn Dependency + id: dawn-depends + run: | + DAWN_VERSION="v20260908.214631" + DAWN_OWNER="google" + DAWN_REPO="dawn" + DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release" + echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" + curl -L -o artifact.tar.gz \ + "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz" + mkdir dawn + tar -xvf artifact.tar.gz -C dawn --strip-components=1 + + - name: Test + id: ggml-ci + run: | + GG_BUILD_WEBGPU=1 GG_BUILD_WEBGPU_DAWN_PREFIX="$GITHUB_WORKSPACE/dawn" \ + bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp From 1d72b05d3899f2bd961cb30a29f54d7e15aacb2c Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 21 Sep 2026 13:57:19 +0300 Subject: [PATCH 11/42] tests/test-backend-ops : allow regex entries in the -o filter (#29204) * tests/test-backend-ops : allow regex entries in the -o filter so far -o only accepted a comma separated list of exact op names or full test case strings. entries that are not plain op names are now treated as regexes matched against the op name (e.g. "MUL_MAT.*"), while plain names keep their exact-matching behavior so that "-o ADD" does not match ADD_EX etc. Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : don't print the FA vec slice log when not needed * tests/test-backend-ops : reformat the help text use the same style as the other tools, with separate sections for modes, options, and examples Assisted-by: pi:llama.cpp/Qwen3.8-27B --- tests/test-backend-ops.cpp | 146 ++++++++++++++++++++++--------------- 1 file changed, 89 insertions(+), 57 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 887add245..cc68e9ca7 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1145,6 +1145,47 @@ static void print_test_result_locked(printer * output_printer, const test_result output_printer->print_test_result(result); } +// Splits the -o filter into comma separated entries. Commas inside parentheses +// (i.e. inside a full test case string) are not treated as separators. +static std::vector op_filter_entries(const char * op_names_filter) { + std::vector entries; + if (op_names_filter == nullptr) { + return entries; + } + std::string_view filter(op_names_filter); + while (!filter.empty()) { + auto comma_pos = filter.find_first_of(','); + const auto lparen_pos = filter.find_first_of('('); + if (lparen_pos < comma_pos) { + const auto rparen_pos = filter.find_first_of(')'); + comma_pos = filter.find_first_of(',', rparen_pos); + } + entries.push_back(filter.substr(0, comma_pos)); + filter = comma_pos != std::string_view::npos ? filter.substr(comma_pos + 1) : ""; + } + return entries; +} + +// An entry from the -o filter matches an op if it is either +// * an exact op name as given by ggml_op_desc() (e.g. "ADD"), or +// * a regex that matches the op name (e.g. "DSV4.*") +static bool op_filter_entry_matches(std::string_view entry, std::string_view op_name) { + if (entry == op_name) { + return true; + } + // plain op names are matched exactly, anything else is treated as a regex + if (std::regex_match(std::string(entry), std::regex("[A-Z0-9_]+"))) { + return false; + } + std::regex re; + try { + re = std::regex(std::string(entry)); + } catch (const std::regex_error &) { + return false; + } + return std::regex_search(op_name.data(), op_name.data() + op_name.size(), re); +} + struct test_case { virtual ~test_case() {} @@ -1308,34 +1349,24 @@ struct test_case { return t; } - // Checks an op against the test filter, which is a comma separated list of OP names or specific variations + // Checks an op against the test filter, which is a comma separated list of OP names, regexes, or specific variations bool matches_filter(ggml_tensor * op, const char * op_names_filter) { - if (op_names_filter) { - const auto op_name = op_desc(op); - const auto op_full_name = op_name + "(" + vars() + ")"; - std::string_view filter(op_names_filter); - while (!filter.empty()) { - auto comma_pos = filter.find_first_of(','); - const auto lparen_pos = filter.find_first_of('('); - if (lparen_pos < comma_pos) { - auto rparen_pos = filter.find_first_of(')'); - comma_pos = filter.find_first_of(',', rparen_pos); - const auto op_filter = filter.substr(0, comma_pos); - if (op_filter == op_full_name) { - return true; - } - } else { - const auto op_filter = filter.substr(0, comma_pos); - if (op_filter == op_name) { - return true; - } - } - filter = comma_pos != std::string_view::npos ? filter.substr(comma_pos + 1) : ""; - } - return false; - } else { + if (op_names_filter == nullptr) { return true; } + const auto op_name = op_desc(op); + const auto op_full_name = op_name + "(" + vars() + ")"; + for (const auto & entry : op_filter_entries(op_names_filter)) { + if (entry.find_first_of('(') != std::string_view::npos) { + // a full test case string, matched exactly + if (entry == op_full_name) { + return true; + } + } else if (op_filter_entry_matches(entry, op_name)) { + return true; + } + } + return false; } test_status_t eval(ggml_backend_t backend1, @@ -11597,25 +11628,16 @@ static std::vector fa_vec_legal_ne(int dk, int dv) { } static bool op_names_filter_selects(const char * op_names_filter, const char * op_name) { - if (!op_names_filter) { + if (op_names_filter == nullptr) { return true; } - std::string_view filter(op_names_filter); - while (!filter.empty()) { - auto comma_pos = filter.find_first_of(','); - const auto lparen_pos = filter.find_first_of('('); - std::string_view entry; - if (lparen_pos < comma_pos) { - const auto rparen_pos = filter.find_first_of(')'); - comma_pos = filter.find_first_of(',', rparen_pos); - entry = filter.substr(0, lparen_pos); - } else { - entry = filter.substr(0, comma_pos); - } - if (entry == op_name) { + for (const auto & entry : op_filter_entries(op_names_filter)) { + // a full test case string is matched by its op name prefix + const auto lparen_pos = entry.find_first_of('('); + const auto op_entry = lparen_pos != std::string_view::npos ? entry.substr(0, lparen_pos) : entry; + if (op_filter_entry_matches(op_entry, op_name)) { return true; } - filter = comma_pos != std::string_view::npos ? filter.substr(comma_pos + 1) : ""; } return false; } @@ -11632,8 +11654,6 @@ static bool run_fa_vec_slice(ggml_backend_t backend, ggml_backend_t backend_cpu, return true; } - printf("Running FA vec slice tests (env LLAMA_TEST_FA_VEC_DISABLE=1 to skip)\n"); - auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); auto set_ov = (set_fa_vec_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_set_fa_vec_override"); @@ -11642,13 +11662,16 @@ static bool run_fa_vec_slice(ggml_backend_t backend, ggml_backend_t backend_cpu, return true; // not the Metal backend: nothing to force } + printf("Running FA vec slice tests (env LLAMA_TEST_FA_VEC_DISABLE=1 to skip)\n"); + struct shape_t { int dk, dv; }; const shape_t shapes[] = { { 128, 128 }, { 576, 512 } }; // mainstream head size + MLA shared K/V view const int ne01_pts[] = { 1, 3 }; // decode, and padded rows for Q=2 and Q=4 const int ne11_pts[] = { 512, 4097 }; // nsg=1, and nsg>=2 together with kvpad const ggml_type types[] = { GGML_TYPE_F16, GGML_TYPE_Q4_0 }; - int n_run = 0, n_fail = 0; + int n_run = 0; + int n_fail = 0; for (auto s : shapes) { for (int ne : fa_vec_legal_ne(s.dk, s.dv)) { for (int Q : { 1, 2, 4 }) { @@ -11968,20 +11991,29 @@ static void show_test_coverage() { } static void usage(char ** argv) { - printf("Usage: %s [mode] [-o ] [-b ] [-p ] [--output ] [--list-ops]", argv[0]); - printf(" [--show-coverage] [--test-file ] [-j ]\n"); - printf(" valid modes:\n"); - printf(" - test (default, compare with CPU backend for correctness)\n"); - printf(" - grad (compare gradients from backpropagation with method of finite differences)\n"); - printf(" - perf (performance evaluation)\n"); - printf(" - support (probe backend operation support)\n"); - printf(" op names for -o are as given by ggml_op_desc() (e.g. ADD, MUL_MAT, etc),\n"); - printf(" optionally including the full test case string (e.g. \"ADD(type=f16,ne=[1,1,8,1],nr=[1,1,1,1],nf=1)\")\n"); - printf(" --output specifies output format (default: console, options: console, sql, csv)\n"); - printf(" --list-ops lists all available GGML operations\n"); - printf(" --show-coverage shows test coverage\n"); - printf(" --test-file reads test operators from a test file generated by test-export-graph-ops\n"); - printf(" -j runs tests using parallel worker threads (default: 1, test mode only)\n"); + printf("Usage: %s [mode] [options]\n\n", argv[0]); + printf("Valid modes:\n"); + printf(" test (default) compare with CPU backend for correctness\n"); + printf(" grad compare gradients from backpropagation with method of finite differences\n"); + printf(" perf performance evaluation\n"); + printf(" support probe backend operation support\n\n"); + printf("Options:\n"); + printf(" -o comma separated list of exact op names (as given by ggml_op_desc()),\n"); + printf(" full test case strings, and/or regexes matched against the op name\n"); + printf(" -b run tests on the given backend (e.g. CPU, MTL0, CUDA0)\n"); + printf(" -p filter test cases by a regex matched against their params\n"); + printf(" --output output format (default: console)\n"); + printf(" --list-ops list all available GGML operations\n"); + printf(" --show-coverage show test coverage\n"); + printf(" --test-file read test operators from a test file generated by test-export-graph-ops\n"); + printf(" -j run tests using parallel worker threads (default: 1, test mode only)\n\n"); + printf("Examples:\n"); + printf(" %s -j 8\n", argv[0]); + printf(" %s -o ADD,MUL_MAT\n", argv[0]); + printf(" %s -o ADD -p 'type=f16.*perm1=0'\n", argv[0]); + printf(" %s -b MTL0 -o 'DSV4.*'\n", argv[0]); + printf(" %s -b CUDA0 -o 'ADD(type=f16,ne=[1,1,1,1],nr=[32,1,1,1],nf=1,perm1=0,src_overlap=0)'\n", argv[0]); + printf(" %s perf -o 'MUL_MAT.*'\n", argv[0]); } int main(int argc, char ** argv) { From 161755f29e415e2c33efe906e91843c068efd664 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 21 Sep 2026 13:57:36 +0300 Subject: [PATCH 12/42] test-llama-archs : make tensor data stdev configurable and improve help (#29133) * test-llama-archs : make tensor data stdev configurable Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * test-llama-archs : expand usage and add examples Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * test-llama-archs : fail on unknown args and log usage Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * test-llama-archs : add test run summary Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * test-llama-archs : initialize Mamba ssm_a negative Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * test-recurrent-state-rollback : report NMSE for logits mismatches Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * test-recurrent-state-rollback : use NMSE for rollback logits checks Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * tests : disable invalid test * cont : adjust nmse_eps * tests : zero DSA indexer score projection in synthetic fixtures Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp * cont : add support for `--arch` regex * cont : alternative top-k stability * cont : indentation * cont : fix top-k value * cont : consistent logs --- src/models/minimax-m3.cpp | 7 +- tests/test-llama-archs.cpp | 198 +++++++++++++++++------- tests/test-recurrent-state-rollback.cpp | 123 ++++++++++----- 3 files changed, 232 insertions(+), 96 deletions(-) diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index f3b64b210..53caca1a5 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -23,7 +23,12 @@ void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size); ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); - msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks }; + + msa_p = { + /*.blk =*/ (int) hparams.indexer_block_size, + /*.topk_blocks =*/ (int) hparams.indexer_top_k, + /*.local =*/ (int) hparams.indexer_local_blocks, + }; GGML_ASSERT(hparams.indexer_block_size > 0); // avoid div by zero diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 80a045185..488770297 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -12,16 +12,25 @@ #include "../src/llama-model-saver.h" #include +#include #include #include #include #include #include +#include #include #include #include #include +static bool arch_matches(const std::string & filter, llm_arch arch) { + if (filter.empty()) { + return true; + } + return std::regex_search(llm_arch_name(arch), std::regex(filter)); +} + // normalized mean squared error = mse(a, b) / mse(a, 0) static double nmse(const std::vector & a, const std::vector & b) { GGML_ASSERT(a.size() == b.size()); @@ -39,24 +48,36 @@ static double nmse(const std::vector & a, const std::vector & b) { return mse_a_b / mse_a_0; } +struct tensor_data_params { + size_t seed; + float stdev; +}; + static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { - size_t seed = *(const size_t *) userdata; + const tensor_data_params & params = *(const tensor_data_params *) userdata; + size_t seed = params.seed; std::hash hasher; seed ^= hasher(tensor->name); std::mt19937 gen(seed); - std::normal_distribution dis(0.0f, 1.0e-2f); + std::normal_distribution dis(0.0f, params.stdev); + // TODO: refactor per-tensor initialization logic in a cleaner way + + // note: Mamba A must be negative (state decay) + const bool is_ssm_a = strstr(tensor->name, "ssm_a") != nullptr; const int64_t ne = ggml_nelements(tensor); if (tensor->type == GGML_TYPE_F32) { std::vector tmp(ne); for (int64_t i = 0; i < ne; i++) { - tmp[i] = dis(gen); + float val = dis(gen); + tmp[i] = is_ssm_a ? -fabsf(val) : val; } ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor)); } else if (tensor->type == GGML_TYPE_F16) { std::vector tmp(ne); for (int64_t i = 0; i < ne; i++) { - tmp[i] = ggml_fp32_to_fp16(dis(gen)); + float val = dis(gen); + tmp[i] = ggml_fp32_to_fp16(is_ssm_a ? -fabsf(val) : val); } ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor)); } else { @@ -65,7 +86,19 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help]\n", argv[0]); + LOG("Usage: %s [options]\n\n", argv[0]); + LOG("Options:\n"); + LOG(" -a, --arch Run only matching LLM architectures (default: all supported)\n"); + LOG(" -s, --seed Set the random seed for tensor initialization and token generation\n"); + LOG(" -d, --stdev Set the standard deviation of the tensor initialization distribution (default: 0.1f)\n"); + LOG(" -o, --out Save generated test models to instead of running backend tests\n"); + LOG(" -v Set log verbosity level\n"); + LOG(" -h, --help Show this help message\n\n"); + LOG("Examples:\n"); + LOG(" %s\n", argv[0]); + LOG(" %s -a qwen35moe\n", argv[0]); + LOG(" %s -a deepseek4 -o tests/test-models/\n", argv[0]); + LOG(" %s -a cohere2moe -v 5\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -292,7 +325,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); + // note: using a realistic top-k here makes the results unstable and hard to match between CPU and GPU + // a large value makes things deterministic since all data is selected by the indexer + //ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(131072)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); @@ -314,13 +351,13 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } if (arch == LLM_ARCH_DEEPSEEK4) { - ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8)); - ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32)); - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 0, 4, 128})); - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); - ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); - ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); - ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8)); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32)); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 0, 4, 128})); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0)); ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); @@ -381,19 +418,19 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK, n_head); ms.add_kv(LLM_KV_SSM_GROUP_COUNT, arch == LLM_ARCH_PLAMO2 ? 0 : uint32_t(2)); ms.add_kv(LLM_KV_KDA_HEAD_DIM, uint32_t(128)); - ms.add_kv(LLM_KV_KDA_SAFE_GATE, true); - ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); + ms.add_kv(LLM_KV_KDA_SAFE_GATE, true); + ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); if (arch == LLM_ARCH_BAILINGMOE3) { ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector({0.0f, 4.0f})); ms.add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector({0.0f, 5.0f})); } - ms.add_kv(LLM_KV_WKV_HEAD_SIZE, n_embd/n_head); - ms.add_kv(LLM_KV_SHORTCONV_L_CACHE, uint32_t(3)); - ms.add_kv(LLM_KV_RESIDUAL_SCALE, 3.5565588200778455f); - ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, uint32_t(12)); - ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA, 4.0f); + ms.add_kv(LLM_KV_WKV_HEAD_SIZE, n_embd/n_head); + ms.add_kv(LLM_KV_SHORTCONV_L_CACHE, uint32_t(3)); + ms.add_kv(LLM_KV_RESIDUAL_SCALE, 3.5565588200778455f); + ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, uint32_t(12)); + ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA, 4.0f); ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f); - ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); + ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); for (uint32_t il = 0; il < n_layer; il++) { ggml_tensor t; @@ -416,7 +453,8 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) } static std::pair get_model_and_ctx( - struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, + struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const float stdev, + const std::vector & devs, const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); @@ -434,9 +472,9 @@ static std::pair get_model_and_ctx( ctx_params.n_ubatch = 64; } - size_t tmp = seed; + tensor_data_params tensor_params = { seed, stdev }; llama_model_ptr model(gguf_ctx != nullptr ? - llama_model_init_from_user(gguf_ctx, set_tensor_data, &tmp, model_params) : + llama_model_init_from_user(gguf_ctx, set_tensor_data, &tensor_params, model_params) : llama_model_load_from_file_ptr(file, model_params)); if (!model) { throw std::runtime_error("failed to create llama model"); @@ -608,7 +646,7 @@ static bool arch_supported(const llm_arch arch) { return true; } -static int save_models(const llm_arch target_arch, const size_t seed, const int verbosity, const std::string & dir) { +static int save_models(const std::string & arch_filter, const size_t seed, const float stdev, const int verbosity, const std::string & dir) { struct user_data_t { struct { ggml_log_callback callback; @@ -635,7 +673,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const int if (arch == LLM_ARCH_UNKNOWN) { continue; } - if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) { + if (!arch_matches(arch_filter, arch)) { continue; } if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { @@ -656,7 +694,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const int continue; } gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe); - auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}); + auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, {}); const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf"); LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str()); llama_model_save_to_file(model_and_ctx.first.get(), path.c_str()); @@ -666,7 +704,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const int return 0; } -static int test_backends(const llm_arch target_arch, const size_t seed, const int verbosity) { +static int test_backends(const std::string & arch_filter, const size_t seed, const float stdev, const int verbosity) { struct user_data_t { struct { ggml_log_callback callback; @@ -731,22 +769,24 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in const std::string template_row_res = "%15s %10s|%20s|\n"; bool all_ok = true; + size_t n_tests = 0; + size_t n_failed = 0; common_log_flush(common_log_main()); - printf(template_header.c_str(), "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip"); - printf("|"); + LOG(template_header.c_str(), "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip"); + LOG("|"); for (size_t i = 0; i < max_arch_name_length; i++) { - printf("-"); + LOG("-"); } - printf("|"); + LOG("|"); for (size_t i = 0; i < max_device_label_length; i++) { - printf("-"); + LOG("-"); } - printf("|------|---------------|---------|\n"); + LOG("|------|---------------|---------|\n"); for (const llm_arch & arch : llm_arch_all()) { if (arch == LLM_ARCH_UNKNOWN) { continue; } - if (target_arch != LLM_ARCH_UNKNOWN && arch != target_arch) { + if (!arch_matches(arch_filter, arch)) { continue; } if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { @@ -773,8 +813,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in std::vector logits_cpu; for (device_config & dc : dev_configs) { // print test config first; should anything fail during model loading or inference, at least we know which test case caused it - printf(template_row_cfg.c_str(), - llm_arch_name(arch), dc.label.c_str(), config_name.c_str()); + LOG(template_row_cfg.c_str(), llm_arch_name(arch), dc.label.c_str(), config_name.c_str()); fflush(stdout); std::pair model_and_ctx_dev; @@ -784,19 +823,22 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in char nmse_str[12] = {0}; bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); + bool test_executed = false; + bool test_ok = true; if (!skip) { if (logits_cpu.empty()) { - model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, encode); + model_and_ctx_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, {}, LLAMA_SPLIT_MODE_LAYER, encode); logits_cpu = get_logits(model_and_ctx_cpu.first.get(), model_and_ctx_cpu.second.get(), tokens, encode); } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { - model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.devs, dc.split_mode, encode); + test_executed = true; + model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); status_nmse = "\033[1;32mOK\033[0m"; if (nmse_val > 1e-4) { - all_ok = false; + test_ok = false; status_nmse = "\033[1;31mFAIL\033[0m"; } } @@ -805,6 +847,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in // FIXME: when adding a tensor to a gguf_context a copy is made, this changes the pointer which the meta backend // in turn uses to map the tensors to their simple equivalents - this is fundamentally incompatible if (file != nullptr && llama_model_saver_supports_arch(arch) && dc.split_mode != LLAMA_SPLIT_MODE_TENSOR) { + test_executed = true; GGML_ASSERT(model_and_ctx_dev.first && model_and_ctx_dev.second); llama_model_saver ms = llama_model_saver(model_and_ctx_dev.first.get()); ms.add_kv_from_model(); @@ -812,14 +855,14 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in ms.save(file); rewind(file); - auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, dc.devs, dc.split_mode, encode); + auto model_and_ctx_roundtrip = get_model_and_ctx(nullptr, file, seed, stdev, dc.devs, dc.split_mode, encode); const std::vector logits_roundtrip = get_logits( model_and_ctx_roundtrip.first.get(), model_and_ctx_roundtrip.second.get(), tokens, encode); status_roundtrip = "\033[1;32mOK\033[0m"; GGML_ASSERT(logits_roundtrip.size() == logits_dev.size()); for (size_t i = 0; i < logits_roundtrip.size(); i++) { if (logits_roundtrip[i] != logits_dev[i]) { - all_ok = false; + test_ok = false; status_roundtrip = "\033[1;31mFAIL\033[0m"; break; } @@ -827,12 +870,28 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in } } + if (test_executed) { + n_tests++; + if (!test_ok) { + n_failed++; + all_ok = false; + } + } + // log the results for this test case - printf(template_row_res.c_str(), - status_nmse.c_str(), nmse_str, status_roundtrip.c_str()); + LOG(template_row_res.c_str(), status_nmse.c_str(), nmse_str, status_roundtrip.c_str()); } } } + + if (n_tests == 0) { + LOG("Summary: no tests executed\n"); + } else if (n_failed == 0) { + LOG("Summary: all %zu test(s) passed\n", n_tests); + } else { + LOG("Summary: %zu test(s) executed, %zu failed\n", n_tests, n_failed); + } + llama_log_set(ud.log_old.callback, ud.log_old.user_data); return all_ok ? 0 : 1; } @@ -844,8 +903,9 @@ int main(int argc, char ** argv) { std::random_device rd; - llm_arch arch = LLM_ARCH_UNKNOWN; + std::string arch_filter; size_t seed = rd(); + float stdev = 0.1f; std::string out; int verbosity = LOG_LEVEL_ERROR; @@ -854,52 +914,70 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { usage(argv); return 0; - } - if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) { + } else if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) { if (i + 1 < argc) { const std::string arch_name = argv[++i]; - arch = llm_arch_from_string(arch_name); - if (arch == LLM_ARCH_UNKNOWN) { - LOG_ERR("%s: unkown LLM architecture: %s\n", __func__, arch_name.c_str()); - return 1; + if (llm_arch_from_string(arch_name) != LLM_ARCH_UNKNOWN) { + // exact architecture name + arch_filter = "^" + arch_name + "$"; + } else { + try { + std::regex re(arch_name); + arch_filter = arch_name; + } catch (const std::regex_error & err) { + LOG_ERR("%s: invalid architecture regex: %s (%s)\n", __func__, arch_name.c_str(), err.what()); + return 1; + } } } else { usage(argv); return 1; } - } - if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { + } else if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { if (i + 1 < argc) { seed = std::stoull(argv[++i]); } else { usage(argv); return 1; } - } - if (strcmp(argv[i], "-v") == 0) { + } else if (strcmp(argv[i], "-d") == 0 || strcmp(argv[i], "--stdev") == 0) { + if (i + 1 < argc) { + stdev = std::stof(argv[++i]); + } else { + usage(argv); + return 1; + } + } else if (strcmp(argv[i], "-v") == 0) { if (i + 1 < argc) { verbosity = std::stoull(argv[++i]); } else { usage(argv); return 1; } - } - if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--out") == 0) { + } else if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--out") == 0) { if (i + 1 < argc) { out = argv[++i]; } else { usage(argv); return 1; } + } else { + LOG_ERR("%s: unknown argument: %s\n", __func__, argv[i]); + usage(argv); + return 1; } } - printf("%s: using seed %zu\n", __func__, seed); + if (stdev <= 0.0f) { + LOG_ERR("%s: stdev must be > 0\n", __func__); + return 1; + } + LOG_INF("%s: using seed %zu, stdev %f\n", __func__, seed, stdev); try { if (!out.empty()) { - return save_models(arch, seed, verbosity, out); + return save_models(arch_filter, seed, stdev, verbosity, out); } - return test_backends(arch, seed, verbosity); + return test_backends(arch_filter, seed, stdev, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what()); return -1; diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index ef05de67d..fdac344d8 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -90,6 +90,20 @@ static float logit_diff(float a, float b) { return std::isfinite(a) && std::isfinite(b) ? std::fabs(a - b) : std::numeric_limits::infinity(); } +static double nmse(const float * a, const float * b, int n) { + double mse_ab = 0.0; + double mse_a0 = 0.0; + for (int i = 0; i < n; i++) { + if (!std::isfinite(a[i]) || !std::isfinite(b[i])) { + return std::numeric_limits::infinity(); + } + const double diff = (double) a[i] - b[i]; + mse_ab += diff*diff; + mse_a0 += (double) a[i]*a[i]; + } + return mse_a0 == 0.0 ? (mse_ab == 0.0 ? 0.0 : std::numeric_limits::infinity()) : mse_ab/mse_a0; +} + // Roll back multiple sequences, then replay them in a single batch whose // per-seq token count exceeds n_ubatch: each seq's replay spans several // ubatches while its rollback restore is still pending. Compared against a @@ -182,13 +196,15 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode return false; } - // identical ubatch shapes from bit-exact states: a correct implementation - // matches bitwise, so eps only allows backend scheduling noise - constexpr float eps = 1e-7f; + // identical ubatch shapes should produce identical states, but the larger + // stdev makes the model sensitive to backend scheduling/rounding noise + constexpr float nmse_eps = 1e-5f; float diff_max = 0.0f; uint32_t seq_first = 0; int32_t pos_first = -1; + double nmse_ab = 0.0; + double nmse_a0 = 0.0; for (uint32_t i = 0; i < n_seqs*n_replay; ++i) { const float * l_roll = llama_get_logits_ith(ctx_roll, i); const float * l_ref = llama_get_logits_ith(ctx_ref, i); @@ -198,23 +214,34 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode return false; } for (int t = 0; t < n_vocab; ++t) { - const float diff = logit_diff(l_roll[t], l_ref[t]); - if (diff > eps && pos_first < 0) { + const float r = l_roll[t]; + const float f = l_ref[t]; + const float diff = logit_diff(r, f); + if (diff > 0.0f && pos_first < 0) { seq_first = i/n_replay; pos_first = p0 + (int32_t) (i%n_replay); } diff_max = std::max(diff_max, diff); + if (std::isfinite(r) && std::isfinite(f)) { + const double d = (double) r - f; + nmse_ab += d*d; + nmse_a0 += (double) r*r; + } else { + nmse_ab = std::numeric_limits::infinity(); + nmse_a0 = 1.0; + } } } + const double nmse_val = nmse_a0 == 0.0 ? (nmse_ab == 0.0 ? 0.0 : std::numeric_limits::infinity()) : nmse_ab/nmse_a0; - if (diff_max > eps) { - fprintf(stderr, "%s : multi-seq split replay logits mismatch (max diff %g, first at seq %u pos %d)\n", - __func__, (double) diff_max, seq_first, pos_first); + if (nmse_val > nmse_eps) { + fprintf(stderr, "%s : multi-seq split replay logits mismatch (max diff %g, nmse %g, first at seq %u pos %d)\n", + __func__, (double) diff_max, nmse_val, seq_first, pos_first); cleanup(); return false; } - fprintf(stderr, "%s : multi-seq split replay matched (max diff %g)\n", __func__, (double) diff_max); + fprintf(stderr, "%s : multi-seq split replay matched (max diff %g, nmse %g)\n", __func__, (double) diff_max, nmse_val); // seq-1-only decodes must be independent of seq 0's content: diverge seq 0 // in ctx_ref only, then compare identical seq-1-only continuations bitwise @@ -231,6 +258,8 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode } float diff_tail = 0.0f; + double nmse_tail_ab = 0.0; + double nmse_tail_a0 = 0.0; for (uint32_t i = 0; i < n_tail && ok; ++i) { const llama_pos pos = p0 + (llama_pos) (n_replay + i); llama_batch batch_one = llama_batch_init(1, 0, 1); @@ -246,18 +275,29 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode const float * l_ref = llama_get_logits_ith(ctx_ref, 0); ok = l_roll != nullptr && l_ref != nullptr; for (int t = 0; ok && t < n_vocab; ++t) { - diff_tail = std::max(diff_tail, logit_diff(l_roll[t], l_ref[t])); + const float r = l_roll[t]; + const float f = l_ref[t]; + diff_tail = std::max(diff_tail, logit_diff(r, f)); + if (std::isfinite(r) && std::isfinite(f)) { + const double d = (double) r - f; + nmse_tail_ab += d*d; + nmse_tail_a0 += (double) r*r; + } else { + nmse_tail_ab = std::numeric_limits::infinity(); + nmse_tail_a0 = 1.0; + } } } + const double nmse_tail = nmse_tail_a0 == 0.0 ? (nmse_tail_ab == 0.0 ? 0.0 : std::numeric_limits::infinity()) : nmse_tail_ab/nmse_tail_a0; - if (!ok || diff_tail > eps) { - fprintf(stderr, "%s : seq-1-only decode leaked seq 0 state (ok=%d, max diff %g)\n", - __func__, ok ? 1 : 0, (double) diff_tail); + if (!ok || nmse_tail > nmse_eps) { + fprintf(stderr, "%s : seq-1-only decode leaked seq 0 state (ok=%d, max diff %g, nmse %g)\n", + __func__, ok ? 1 : 0, (double) diff_tail, nmse_tail); cleanup(); return false; } - fprintf(stderr, "%s : seq-1-only decode independent of seq 0 (max diff %g)\n", __func__, (double) diff_tail); + fprintf(stderr, "%s : seq-1-only decode independent of seq 0 (max diff %g, nmse %g)\n", __func__, (double) diff_tail, nmse_tail); cleanup(); return true; } @@ -266,6 +306,7 @@ static int test_rollback(const common_params & params, llama_model * model, uint const llama_vocab * vocab = llama_model_get_vocab(model); const int n_vocab = llama_vocab_n_tokens(vocab); + // TODO: use smart pointers llama_context * ctx_src = make_ctx(params, model, fill); llama_context * ctx_dst = make_ctx(params, model, fill); if (ctx_src == nullptr || ctx_dst == nullptr) { @@ -320,7 +361,7 @@ static int test_rollback(const common_params & params, llama_model * model, uint ckpt.update_tgt(ctx_src, 0, 0); ckpt.load_tgt(ctx_dst, 0, 0); - constexpr float eps = 1e-5f; + constexpr float nmse_eps = 0.0; std::vector> logits_src_replay(n_rollback); const auto replay_and_compare = [&](const char * mode) { for (uint32_t i = 0; i < n_rollback; ++i) { @@ -339,13 +380,18 @@ static int test_rollback(const common_params & params, llama_model * model, uint } logits_src_replay[i].assign(logits_src, logits_src + n_vocab); + const double nmse_val = nmse(logits_src, logits_dst, n_vocab); + int token_first = -1; for (int token = 0; token < n_vocab; ++token) { - if (logit_diff(logits_src[token], logits_dst[token]) > eps) { - fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n", - __func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]); - return false; + if (logit_diff(logits_src[token], logits_dst[token]) > 0.0f && token_first < 0) { + token_first = token; } } + if (nmse_val > nmse_eps) { + fprintf(stderr, "%s : %s logits mismatch at position %d, first token %d, nmse %g\n", + __func__, mode, pos, token_first, nmse_val); + return false; + } } return true; }; @@ -353,20 +399,22 @@ static int test_rollback(const common_params & params, llama_model * model, uint return 1; } - if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) || - !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) { - fprintf(stderr, "%s : partial rollback failed\n", __func__); - return 1; - } + // TODO: this test is invalid because RS rollback is only correct once after a ubatch with more than n_rs_seq tokens + // this is not the case here. add asserts and guardrails to prevent such attempts + //if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) || + // !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) { + // fprintf(stderr, "%s : partial rollback failed\n", __func__); + // return 1; + //} - constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; - common_prompt_checkpoint ckpt_partial; - ckpt_partial.update_tgt(ctx_src, 0, partial_flags); - ckpt_partial.load_tgt(ctx_dst, 0, partial_flags); + //constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + //common_prompt_checkpoint ckpt_partial; + //ckpt_partial.update_tgt(ctx_src, 0, partial_flags); + //ckpt_partial.load_tgt(ctx_dst, 0, partial_flags); - if (!replay_and_compare("partial")) { - return 1; - } + //if (!replay_and_compare("partial")) { + // return 1; + //} // Repeat the load into a context that already has its own rollback state: // groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is @@ -408,13 +456,18 @@ static int test_rollback(const common_params & params, llama_model * model, uint return 1; } + const double nmse_dirty = nmse(logits_src_replay[i].data(), logits_dirty, n_vocab); + int token_first = -1; for (int token = 0; token < n_vocab; ++token) { - if (logit_diff(logits_src_replay[i][token], logits_dirty[token]) > eps) { - fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n", - __func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]); - return 1; + if (logit_diff(logits_src_replay[i][token], logits_dirty[token]) > 0.0f && token_first < 0) { + token_first = token; } } + if (nmse_dirty > nmse_eps) { + fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, first token %d, nmse %g\n", + __func__, pos, token_first, nmse_dirty); + return 1; + } } fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__); From 1884824fdaeb19af27cd8ba0ca50fcbed28eca44 Mon Sep 17 00:00:00 2001 From: ynankani Date: Mon, 21 Sep 2026 10:58:28 +0000 Subject: [PATCH 13/42] CUDA: Follow up of #25635, refactoring FA shared smem swizzle (#28536) * remove explicit swz value in config and rebase Signed-off-by: ynankani * address review comments Signed-off-by: ynankani --------- Signed-off-by: ynankani --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 86 ++++++++++-------- ggml/src/ggml-cuda/fattn-swizzle.cuh | 126 --------------------------- ggml/src/ggml-cuda/mma.cuh | 58 ++++++++++++ 3 files changed, 107 insertions(+), 163 deletions(-) delete mode 100644 ggml/src/ggml-cuda/fattn-swizzle.cuh diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index df7dd8615..84219ca92 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -2,7 +2,6 @@ #include "cp-async.cuh" #include "mma.cuh" #include "fattn-common.cuh" -#include "fattn-swizzle.cuh" using namespace ggml_cuda_mma; @@ -327,6 +326,32 @@ static constexpr __device__ bool ggml_cuda_fattn_mma_get_Q_in_reg(const int DKQ, return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).Q_in_reg; } +// Swizzling needs a tile stride that is a multiple of 32 half2 columns. +static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_bank_aligned(const int nbatch_2) { + return nbatch_2 >= 32 && nbatch_2 % 32 == 0; +} + +// Swizzling needs ldmatrix, on other hardware the tiles keep the row padding. +static __host__ bool ggml_cuda_fattn_mma_get_swizzled(const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) { + const fattn_mma_config cfg = ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols1*ncols2, cc); + return turing_mma_available(cc) && ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_K2) && ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_V2); +} + +static constexpr __device__ bool ggml_cuda_fattn_mma_get_swizzled(const int DKQ, const int DV, const int ncols1, const int ncols2) { +#if defined(TURING_MMA_AVAILABLE) + const fattn_mma_config cfg = ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols1*ncols2); + return ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_K2) && ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_V2); +#else + GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2); + return false; +#endif // defined(TURING_MMA_AVAILABLE) +} + +// Row padding is only needed if the tile is not swizzled. +static constexpr __host__ __device__ int ggml_cuda_fattn_mma_get_stride_tile(const int nbatch_2, const bool swizzled) { + return swizzled ? nbatch_2 : nbatch_2 + 4; +} + static constexpr __device__ int get_cols_per_thread() { #if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) return 1; // AMD has a single column per thread. @@ -411,12 +436,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) { const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k); - if constexpr (swz) { - const int smem_offs_b = ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk); - cp_async_cg_16(tile_KV_32 + smem_offs_b, KV + i_KV*stride_KV + k*h2_per_chunk); - } else { - cp_async_cg_16(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i_KV*stride_KV + k*h2_per_chunk); - } + cp_async_cg_16(tile_KV_32 + swizzle_bytes(i, k*h2_per_chunk, stride_tile), KV + i_KV*stride_KV + k*h2_per_chunk); } } }; @@ -458,11 +478,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile( } else { src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero; } - if constexpr (swz) { - ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc(i, k*h2_per_chunk), src); - } else { - ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, src); - } + ggml_cuda_memcpy_1<16>((char *) tile_KV + swizzle_bytes(i, k*h2_per_chunk, stride_tile), src); } } }; @@ -605,11 +621,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse); - // swizzle the tile stride for K and V based on the batch size. - constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2); - constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2); - constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2); - constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2); + constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2); + constexpr int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swz); + constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swz); const int k_VKQ_0 = kb0 * nbatch_fa; #if defined(TURING_MMA_AVAILABLE) @@ -627,7 +641,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool use_cp_async = true; cp_async_wait_all(); __syncthreads(); - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr); } else { // the sparse mask values are gathered per element, always load them synchronously @@ -647,7 +661,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices); if (use_cp_async) { cp_async_wait_all(); @@ -663,7 +677,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( #pragma unroll for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) { T_A_KQ K_A; - ggml_cuda_fattn_smem_swizzle::load_ldmatrix(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start); + load_ldmatrix(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start, stride_tile_K); if constexpr (cols_per_warp == 8) { mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]); } else { @@ -689,7 +703,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I; T_A_KQ K_A; - ggml_cuda_fattn_smem_swizzle::load_ldmatrix(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start); + load_ldmatrix(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start, stride_tile_K); if constexpr (cols_per_warp == 8) { mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[0]); @@ -984,7 +998,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( flash_attn_ext_f16_load_mask (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr); } } @@ -1000,7 +1014,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i0_diff = i0_stop - i0_start; if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_0, k_VKQ_sup, indices); if (use_cp_async) { cp_async_wait_all(); @@ -1019,7 +1033,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::J; T_A_VKQ A; // Transposed in SRAM but not in registers, gets transposed on load. - ggml_cuda_fattn_smem_swizzle::load_ldmatrix_trans(A, tile_V, (int)(tile_V_i - tile_V) + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2); + load_ldmatrix_trans(A, tile_V, 2*k0, (int)(tile_V_i - tile_V) + (i_VKQ_0 - i0_start)/2, stride_tile_V); if constexpr (T_B_KQ::I == 8) { mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]); } else { @@ -1045,7 +1059,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::I; T_A_VKQ A; // Transposed in both SRAM and registers, load normally. - ggml_cuda_fattn_smem_swizzle::load_ldmatrix(A, tile_V, (int)(tile_V_i - tile_V) + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2); + static_assert(!swz, "Volta has no ldmatrix"); + load_ldmatrix(A, tile_V_i + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V); mma(VKQ_C[i_VKQ_0/i0_stride], B[k00/(np*T_A_VKQ::I)], A); } } @@ -1236,12 +1251,10 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( static_assert(nwarps * (cols_per_warp/ncols2) % ncols1 == 0, "bad nwarps"); constexpr int stride_tile_Q = DKQ/2 + 4; - // swizzle the tile stride for K and V based on the batch size. - constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2); - constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2); + constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2); + constexpr int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swz); + constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swz); constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V; - constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2); - constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2); extern __shared__ half2 tile_Q[]; half2 * tile_K = Q_in_reg ? tile_Q : tile_Q + ncols * stride_tile_Q; @@ -1338,7 +1351,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_load_mask (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr); } - flash_attn_ext_f16_load_tile + flash_attn_ext_f16_load_tile (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr); } @@ -1503,14 +1516,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int tile_stride = nbatch_combine + 4; static_assert((DV/2) % nbatch_combine == 0, "bad nbatch_combine"); - constexpr bool combine_needs_sync = swz_K || swz_V; - if constexpr (cols_per_warp == 8) { const int jc_cwmo = (threadIdx.x % (2*T_C_VKQ::J)) / T_C_VKQ::J; // jc combine write meta offset const int jc_cwm = threadIdx.y*(2*T_C_VKQ::J) + 2*T_C_VKQ::get_j(-1) + jc_cwmo; // jc combine write meta const float2 KQ_cmr = make_float2(KQ_max[jc_cwmo], KQ_rowsum[jc_cwmo]); // KQ combine max rowsum - if constexpr (combine_needs_sync) { + if constexpr (swz) { __syncthreads(); } @@ -1550,7 +1561,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const bool thread_should_write = T_C_KQ::J == 8 || T_C_KQ::get_j(threadIdx.x & 2) < 8; #endif // defined(TURING_MMA_AVAILABLE) - if constexpr (combine_needs_sync) { + if constexpr (swz) { __syncthreads(); } @@ -2025,8 +2036,9 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu // KV tile strides must match flash_attn_ext_f16_iter / _process_tile. - const int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2, cc); - const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2, cc); + const bool swizzled = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2, cc); + const int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swizzled); + const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swizzled); const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(stride_tile_K, stride_tile_V) * sizeof(half2); const size_t nbytes_shared_KV_2stage = nbatch_fa * (stride_tile_K + stride_tile_V) * sizeof(half2); const size_t nbytes_shared_Q = ncols * (DKQ/2 + 4) * sizeof(half2); diff --git a/ggml/src/ggml-cuda/fattn-swizzle.cuh b/ggml/src/ggml-cuda/fattn-swizzle.cuh deleted file mode 100644 index 44338c8db..000000000 --- a/ggml/src/ggml-cuda/fattn-swizzle.cuh +++ /dev/null @@ -1,126 +0,0 @@ -#pragma once - -#include "common.cuh" -#include "mma.cuh" - -// XOR swizzle for K/V SMEM tiles to avoid bank conflicts without row padding (Turing+ only). -// Stride must be a multiple of 32 half2 columns, otherwise we keep +4 row padding. - -namespace ggml_cuda_fattn_smem_swizzle { - -static __host__ __device__ constexpr bool bank_aligned(const int nbatch_2) { - return nbatch_2 >= 32 && nbatch_2 % 32 == 0; -} - -static __device__ constexpr bool enabled(const int nbatch_2) { -#if defined(TURING_MMA_AVAILABLE) - return bank_aligned(nbatch_2); -#else - GGML_UNUSED(nbatch_2); - return false; -#endif // defined(TURING_MMA_AVAILABLE) -} - -static __host__ bool enabled(const int nbatch_2, const int cc) { -#ifdef GGML_USE_HIP - GGML_UNUSED(nbatch_2); - GGML_UNUSED(cc); - return false; -#else - return turing_mma_available(cc) && bank_aligned(nbatch_2); -#endif // GGML_USE_HIP -} - -static __device__ constexpr int tile_stride(const int nbatch_2) { - return enabled(nbatch_2) ? nbatch_2 : nbatch_2 + 4; -} - -static __host__ int tile_stride(const int nbatch_2, const int cc) { - return enabled(nbatch_2, cc) ? nbatch_2 : nbatch_2 + 4; -} - -// Swizzled byte offset for tile element (row, col_h2), same map used for writes and reads. -template -static __device__ __forceinline__ int bytes_rc(const int row, const int col_h2) { - static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32"); - return ((row * stride_h2 + col_h2) * (int) sizeof(half2)) ^ ((row & 7) << 4); -} - -// ldmatrix.x4 via 64-bit generic pointer. -static __device__ __forceinline__ void ldmatrix_x4(int * xi, const half2 * addr) { -#if defined(TURING_MMA_AVAILABLE) - asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" - : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) - : "l"(addr)); -#else - GGML_UNUSED_VARS(xi, addr); - NO_DEVICE_CODE; -#endif // defined(TURING_MMA_AVAILABLE) -} - -static __device__ __forceinline__ void ldmatrix_x4_trans(int * xi, const half2 * addr) { -#if defined(TURING_MMA_AVAILABLE) - asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" - : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) - : "l"(addr)); -#else - GGML_UNUSED_VARS(xi, addr); - NO_DEVICE_CODE; -#endif // defined(TURING_MMA_AVAILABLE) -} - -// Per-lane swizzled address for one tile<16, 8, half2> ldmatrix: 16 rows, 4 half2 columns per lane. -template -static __device__ __forceinline__ const half2 * lane_addr( - const half2 * tile_base, const int base_row, const int base_col_h2, const int I, const int J) { - static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32"); - const int lane_row = threadIdx.x % I; - const int lane_col = (threadIdx.x / I) * (J / 2); - uint32_t byte_off = (uint32_t) ((base_row + lane_row)*stride_h2 + base_col_h2 + lane_col) * (uint32_t) sizeof(half2); - byte_off ^= (uint32_t) (((base_row + lane_row) & 7) << 4); - return (const half2 *) ((const char *) tile_base + byte_off); -} - -template -static __device__ __forceinline__ void load_ldmatrix( - TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) { - if constexpr (swz) { - static_assert(std::is_same_v>, - "the swizzled layout is only supported for tile<16, 8, half2>"); - ldmatrix_x4((int *) t.x, lane_addr(tile_base, base_row, base_col_h2, TileT::I, TileT::J)); - } else { - ggml_cuda_mma::load_ldmatrix(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2); - } -} - -template -static __device__ __forceinline__ void load_ldmatrix(TileT & t, const half2 * tile_base, const int off_h2) { - if constexpr (swz) { - load_ldmatrix(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2); - } else { - ggml_cuda_mma::load_ldmatrix(t, tile_base + off_h2, stride_h2); - } -} - -template -static __device__ __forceinline__ void load_ldmatrix_trans( - TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) { - if constexpr (swz) { - static_assert(std::is_same_v>, - "the swizzled layout is only supported for tile<16, 8, half2>"); - ldmatrix_x4_trans((int *) t.x, lane_addr(tile_base, base_row, base_col_h2, TileT::I, TileT::J)); - } else { - ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2); - } -} - -template -static __device__ __forceinline__ void load_ldmatrix_trans(TileT & t, const half2 * tile_base, const int off_h2) { - if constexpr (swz) { - load_ldmatrix_trans(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2); - } else { - ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + off_h2, stride_h2); - } -} - -} // namespace ggml_cuda_fattn_smem_swizzle diff --git a/ggml/src/ggml-cuda/mma.cuh b/ggml/src/ggml-cuda/mma.cuh index 8d7c69dc3..6af2b6a14 100644 --- a/ggml/src/ggml-cuda/mma.cuh +++ b/ggml/src/ggml-cuda/mma.cuh @@ -782,6 +782,20 @@ namespace ggml_cuda_mma { } } + // Byte offset of tile element (i, j). If swz, XOR swizzle it to avoid bank conflicts without row padding. + template + static __device__ __forceinline__ int swizzle_bytes(const int i, const int j, const int stride) { + static_assert(!swz || sizeof(T) == 4, "swizzled tiles need 32 bit elements"); + const int off = (i*stride + j) * (int) sizeof(T); + return swz ? off ^ ((i & 7) << 4) : off; + } + + template + static __device__ __forceinline__ const T * swizzle( + const T * __restrict__ tile_base, const int i, const int j, const int stride) { + return (const T *) ((const char *) tile_base + swizzle_bytes(i, j, stride)); + } + template static __device__ __forceinline__ void load_ldmatrix( tile<8, 8, T> & t, const T * __restrict__ xs0, const int stride) { @@ -858,6 +872,27 @@ namespace ggml_cuda_mma { #endif // TURING_MMA_AVAILABLE } + // Load from tile element (i0, j0), swz tells if the tile is stored swizzled. + template + static __device__ __forceinline__ void load_ldmatrix( + tile<16, 8, T, dl> & t, const T * __restrict__ tile_base, const int i0, const int j0, const int stride) { + if constexpr (!swz) { + load_ldmatrix(t, tile_base + i0*stride + j0, stride); + return; + } +#if defined(TURING_MMA_AVAILABLE) + const int i = i0 + threadIdx.x % t.I; + const int j = j0 + (threadIdx.x / t.I) * (t.J / 2); + int * xi = (int *) t.x; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) + : "l"(swizzle(tile_base, i, j, stride))); +#else + GGML_UNUSED_VARS(t, tile_base, i0, j0, stride); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) + } + static __device__ __forceinline__ void load_ldmatrix( tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); @@ -917,6 +952,29 @@ namespace ggml_cuda_mma { #endif // TURING_MMA_AVAILABLE } + // Load from tile element (i0, j0), swz tells if the tile is stored swizzled. + template + static __device__ __forceinline__ void load_ldmatrix_trans( + tile & t, const T * __restrict__ tile_base, const int i0, const int j0, const int stride) { + if constexpr (!swz) { + load_ldmatrix_trans(t, tile_base + i0*stride + j0, stride); + return; + } +#if defined(TURING_MMA_AVAILABLE) + static_assert(I == 16, "bad tile width"); + static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); + const int i = i0 + threadIdx.x % t.I; + const int j = j0 + (threadIdx.x / t.I) * (t.J / 2); + int * xi = (int *) t.x; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) + : "l"(swizzle(tile_base, i, j, stride))); +#else + GGML_UNUSED_VARS(t, tile_base, i0, j0, stride); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) + } + static __device__ __forceinline__ void mma( tile<16, 8, int> & D, const tile<16, 4, int> & A, const tile<8, 4, int> & B) { #ifdef TURING_MMA_AVAILABLE From af911149c58a4a2c17470d7ccb2e91140bc73a9f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Mon, 21 Sep 2026 12:58:59 +0200 Subject: [PATCH 14/42] sycl : pinned memory use right device context instead of 0 (#28895) --- ggml/include/ggml-sycl.h | 2 + ggml/src/ggml-sycl/ggml-sycl.cpp | 71 +++++++++++++++++++++----------- 2 files changed, 48 insertions(+), 25 deletions(-) diff --git a/ggml/include/ggml-sycl.h b/ggml/include/ggml-sycl.h index 093fa4a7e..1c18f706a 100644 --- a/ggml/include/ggml-sycl.h +++ b/ggml/include/ggml-sycl.h @@ -36,6 +36,8 @@ GGML_BACKEND_API void ggml_backend_sycl_comm_free(void * comm_ctx); GGML_BACKEND_API bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx, struct ggml_tensor ** tensors); // pinned host buffer for use with the CPU backend for faster copies between CPU and GPU +// pins on device 0 - a copy between another device and this memory can fail, +// use ggml_backend_dev_host_buffer_type to pin on the device that does the copy GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_host_buffer_type(void); GGML_BACKEND_API void ggml_backend_sycl_print_sycl_devices(void); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index b24664a0b..a7fbd1644 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -1553,14 +1553,18 @@ static const char * ggml_backend_sycl_host_buffer_type_name(ggml_backend_buffer_ GGML_UNUSED(buft); } +static int ggml_backend_sycl_host_buffer_type_device(ggml_backend_buffer_type_t buft) { + return static_cast(buft->device->context)->device; +} + //host pinned memory -static void * ggml_backend_sycl_host_malloc(size_t size) { - GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_host_malloc\n"); +static void * ggml_backend_sycl_host_malloc(int device, size_t size) { + GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_host_malloc of size %.2f MiB on device %d\n", size / 1024.0 / 1024.0, device); void * ptr = nullptr; try { ggml_check_sycl(); // USM host memory is page-locked and device-accessible by construction - auto & q = dpct::dev_mgr::instance().get_device(0).default_queue(); + auto & q = dpct::dev_mgr::instance().get_device(device).default_queue(); ptr = sycl::malloc_host(size, q, sycl::property_list{}); } catch (...) { ptr = nullptr; @@ -1578,7 +1582,8 @@ static void ggml_backend_sycl_host_buffer_free_buffer(ggml_backend_buffer_t buff return; } if (g_ggml_sycl_enable_host_pinned_mem) { - auto & q = dpct::dev_mgr::instance().get_device(0).default_queue(); + const int device = ggml_backend_sycl_host_buffer_type_device(buffer->buft); + auto & q = dpct::dev_mgr::instance().get_device(device).default_queue(); SYCL_CHECK(CHECK_TRY_ERROR(sycl::free(buffer->context, q))); } else { free_aligned_mem_host((void *) buffer->context); @@ -1586,8 +1591,9 @@ static void ggml_backend_sycl_host_buffer_free_buffer(ggml_backend_buffer_t buff } static ggml_backend_buffer_t ggml_backend_sycl_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - void * ptr = g_ggml_sycl_enable_host_pinned_mem ? ggml_backend_sycl_host_malloc(size) : - aligned_malloc_host(TENSOR_ALIGNMENT, size); + void * ptr = g_ggml_sycl_enable_host_pinned_mem ? + ggml_backend_sycl_host_malloc(ggml_backend_sycl_host_buffer_type_device(buft), size) : + aligned_malloc_host(TENSOR_ALIGNMENT, size); if (ptr == nullptr) { // fallback to cpu buffer return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); @@ -1604,8 +1610,8 @@ static ggml_backend_buffer_t ggml_backend_sycl_host_buffer_type_alloc_buffer(ggm static size_t ggml_backend_sycl_host_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { if (g_ggml_sycl_enable_host_pinned_mem) { - ggml_backend_sycl_device_context * dev_ctx = (ggml_backend_sycl_device_context *) buft->device->context; - size_t max_alloc_size = dpct::dev_mgr::instance().get_device(dev_ctx->device).get_max_mem_alloc_size(); + const int device = ggml_backend_sycl_host_buffer_type_device(buft); + size_t max_alloc_size = dpct::dev_mgr::instance().get_device(device).get_max_mem_alloc_size(); if (g_ggml_sycl_host_pinned_mem_2g) { return std::min(max_alloc_size, (size_t) 2LL*1024*1024*1024); } else { @@ -1616,22 +1622,37 @@ static size_t ggml_backend_sycl_host_buffer_type_get_max_size(ggml_backend_buffe } } -ggml_backend_buffer_type_t ggml_backend_sycl_host_buffer_type() { - GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_host_buffer_type\n"); - static struct ggml_backend_buffer_type ggml_backend_sycl_buffer_type_host = { - /* .iface = */ { - /* .get_name = */ ggml_backend_sycl_host_buffer_type_name, - /* .alloc_buffer = */ ggml_backend_sycl_host_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, - /* .get_max_size = */ ggml_backend_sycl_host_buffer_type_get_max_size, - /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, - /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, - }, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_sycl_reg(), 0), - /* .context = */ nullptr, - }; +static ggml_backend_buffer_type_t ggml_backend_sycl_host_buffer_type_for_device(int device) { + GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_host_buffer_type_for_device on device %d\n", device); - return &ggml_backend_sycl_buffer_type_host; + // the vector is never resized after this, so the returned pointers stay valid + static std::vector buffer_types_host = [] { + std::vector bufts(ggml_backend_sycl_get_device_count()); + for (size_t i = 0; i < bufts.size(); i++) { + bufts[i] = { + /* .iface = */ { + /* .get_name = */ ggml_backend_sycl_host_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_sycl_host_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, + /* .get_max_size = */ ggml_backend_sycl_host_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, + /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_sycl_reg(), i), + /* .context = */ nullptr, + }; + } + return bufts; + }(); + + GGML_ASSERT(device >= 0 && device < (int) buffer_types_host.size()); + + return &buffer_types_host[device]; +} + +// TODO: this function is unused and is a temporary hack to avoid breaking changes +ggml_backend_buffer_type_t ggml_backend_sycl_host_buffer_type() { + return ggml_backend_sycl_host_buffer_type_for_device(0); } // buffer pool for sycl (legacy) @@ -6299,8 +6320,8 @@ static ggml_backend_buffer_type_t ggml_backend_sycl_device_get_buffer_type(ggml_ } static ggml_backend_buffer_type_t ggml_backend_sycl_device_get_host_buffer_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); - return ggml_backend_sycl_host_buffer_type(); + ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *) dev->context; + return ggml_backend_sycl_host_buffer_type_for_device(ctx->device); } static ggml_backend_buffer_t ggml_backend_sycl_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) { From bb3c853c300323913d8535eb371fe906cae9b932 Mon Sep 17 00:00:00 2001 From: cwriter Date: Mon, 21 Sep 2026 12:59:38 +0200 Subject: [PATCH 15/42] sycl : support gated DSV4_HC_PRE and optional HC_POST comb matrix (#29132) Co-authored-by: cwriter --- ggml/src/ggml-sycl/dsv4-hc.cpp | 102 ++++++++++++++++++++++++------- ggml/src/ggml-sycl/ggml-sycl.cpp | 5 +- 2 files changed, 83 insertions(+), 24 deletions(-) diff --git a/ggml/src/ggml-sycl/dsv4-hc.cpp b/ggml/src/ggml-sycl/dsv4-hc.cpp index bb66e8c1b..337f4af45 100644 --- a/ggml/src/ggml-sycl/dsv4-hc.cpp +++ b/ggml/src/ggml-sycl/dsv4-hc.cpp @@ -2,22 +2,30 @@ #include "dsv4-hc.hpp" #include +#include static constexpr int DSV4_HC = 4; +// tunable: one work-item per (embedding element, token) +static constexpr int dsv4_hc_pre_block_size = 256; + +// gated: the weight is a per-element gate [n_embd, hc, n_tokens] passed through a sigmoid. +// otherwise it is one weight per (stream, token). +template static void dsv4_hc_pre_f32_sycl( const float * x, const float * weights, float * dst, int64_t n_embd, int64_t hc, int64_t n_tokens, int64_t sx0, int64_t sx1, int64_t sx2, - int64_t sw0, int64_t sw1, + int64_t sw0, int64_t sw1, int64_t sw2, int64_t sd0, int64_t sd1, + float scale, queue_ptr stream) { const int64_t nr = n_embd * n_tokens; - const int64_t block_size = 256; - const int64_t num_blocks = (nr + block_size - 1) / block_size; + const int64_t num_blocks = (nr + dsv4_hc_pre_block_size - 1) / dsv4_hc_pre_block_size; stream->parallel_for( - sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)), + sycl::nd_range<1>(sycl::range<1>(num_blocks * dsv4_hc_pre_block_size), + sycl::range<1>(dsv4_hc_pre_block_size)), [=](sycl::nd_item<1> item) { const int64_t ir = item.get_global_id(0); if (ir >= nr) { @@ -27,14 +35,20 @@ static void dsv4_hc_pre_f32_sycl( const int64_t i0 = ir % n_embd; const int64_t it = ir / n_embd; - float sum = x[i0*sx0 + it*sx2] * weights[it*sw1]; - for (int64_t ih = 1; ih < hc; ++ih) { + float sum = 0.0f; + for (int64_t ih = 0; ih < hc; ++ih) { const float xv = x[i0*sx0 + ih*sx1 + it*sx2]; - const float wv = weights[ih*sw0 + it*sw1]; + float wv; + if constexpr (gated) { + const float gv = weights[i0*sw0 + ih*sw1 + it*sw2]; + wv = 1.0f / (1.0f + sycl::exp(-gv)); + } else { + wv = weights[ih*sw0 + it*sw1]; + } sum += xv * wv; } - dst[i0*sd0 + it*sd1] = sum; + dst[i0*sd0 + it*sd1] = scale * sum; }); } @@ -138,6 +152,12 @@ static void dsv4_hc_comb_f32_sycl( }); } +// tunable: one work-item per (embedding element, stream, token) +static constexpr int dsv4_hc_post_block_size = 256; + +// comb == nullptr is identity mixing: each destination stream keeps its own residual +// instead of summing across the streams. +template static void dsv4_hc_post_f32_sycl( const float * x, const float * residual, const float * post, const float * comb, float * dst, int64_t n_embd, int64_t hc, int64_t n_tokens, @@ -148,7 +168,7 @@ static void dsv4_hc_post_f32_sycl( int64_t sd0, int64_t sd1, int64_t sd2, queue_ptr stream) { const int64_t nr = n_embd * hc * n_tokens; - const int64_t block_size = 256; + const int64_t block_size = dsv4_hc_post_block_size; const int64_t num_blocks = (nr + block_size - 1) / block_size; stream->parallel_for( @@ -164,8 +184,12 @@ static void dsv4_hc_post_f32_sycl( const int64_t it = ir / (n_embd * hc); float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1]; - for (int64_t isrc = 0; isrc < hc; ++isrc) { - sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < hc; ++isrc) { + sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2]; + } + } else { + sum += residual[i0*sr0 + idst*sr1 + it*sr2]; } dst[i0*sd0 + idst*sd1 + it*sd2] = sum; @@ -189,15 +213,33 @@ void ggml_sycl_op_dsv4_hc_pre(ggml_backend_sycl_context & ctx, ggml_tensor * dst const int64_t hc = x->ne[1]; const int64_t n_tokens = x->ne[2]; + const float scale = ggml_get_op_params_f32(dst, 0); + const bool gated = ggml_get_op_params_i32(dst, 1) != 0; + queue_ptr stream = ctx.stream(); - dsv4_hc_pre_f32_sycl( - (const float *) x->data, (const float *) weights->data, (float *) dst->data, - n_embd, hc, n_tokens, - nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float), - nbw0 / sizeof(float), nbw1 / sizeof(float), - nbd0 / sizeof(float), nbd1 / sizeof(float), - stream); + if (gated) { + GGML_ASSERT(weights->ne[0] == n_embd); + GGML_ASSERT(weights->ne[1] == hc); + GGML_ASSERT(weights->ne[2] == n_tokens); + dsv4_hc_pre_f32_sycl( + (const float *) x->data, (const float *) weights->data, (float *) dst->data, + n_embd, hc, n_tokens, + nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float), + nbw0 / sizeof(float), nbw1 / sizeof(float), nbw2 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), + scale, stream); + } else { + GGML_ASSERT(weights->ne[0] == hc); + GGML_ASSERT(weights->ne[1] == n_tokens); + dsv4_hc_pre_f32_sycl( + (const float *) x->data, (const float *) weights->data, (float *) dst->data, + n_embd, hc, n_tokens, + nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float), + nbw0 / sizeof(float), nbw1 / sizeof(float), /*sw2=*/ 0, + nbd0 / sizeof(float), nbd1 / sizeof(float), + scale, stream); + } } void ggml_sycl_op_dsv4_hc_comb(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { @@ -252,24 +294,33 @@ void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * ds GGML_ASSERT(x->type == GGML_TYPE_F32); GGML_ASSERT(residual->type == GGML_TYPE_F32); GGML_ASSERT(post->type == GGML_TYPE_F32); - GGML_ASSERT(comb->type == GGML_TYPE_F32); GGML_ASSERT(dst->type == GGML_TYPE_F32); GGML_TENSOR_LOCALS(size_t, nbx, x, nb); GGML_TENSOR_LOCALS(size_t, nbr, residual, nb); GGML_TENSOR_LOCALS(size_t, nbp, post, nb); - GGML_TENSOR_LOCALS(size_t, nbc, comb, nb); GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + size_t nbc0 = 0; + size_t nbc1 = 0; + size_t nbc2 = 0; + if (comb) { + GGML_ASSERT(comb->type == GGML_TYPE_F32); + nbc0 = comb->nb[0]; + nbc1 = comb->nb[1]; + nbc2 = comb->nb[2]; + } + const int64_t n_embd = x->ne[0]; const int64_t n_tokens = x->ne[1]; const int64_t hc = residual->ne[1]; queue_ptr stream = ctx.stream(); - dsv4_hc_post_f32_sycl( + const auto launch = [&](auto has_comb) { + dsv4_hc_post_f32_sycl( (const float *) x->data, (const float *) residual->data, - (const float *) post->data, (const float *) comb->data, (float *) dst->data, + (const float *) post->data, comb ? (const float *) comb->data : nullptr, (float *) dst->data, n_embd, hc, n_tokens, nbx0 / sizeof(float), nbx1 / sizeof(float), nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float), @@ -277,4 +328,11 @@ void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * ds nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float), nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float), stream); + }; + + if (comb) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index a7fbd1644..e599d2d84 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6482,13 +6482,14 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons break; case GGML_OP_DSV4_HC_PRE: return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && - op->type == GGML_TYPE_F32 && ggml_get_op_params_i32(op, 1) == 0; + op->type == GGML_TYPE_F32; case GGML_OP_DSV4_HC_COMB: return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_DSV4_HC_POST: return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && - op->src[2]->type == GGML_TYPE_F32 && op->src[3] != nullptr && op->src[3]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + (op->src[3] == nullptr || op->src[3]->type == GGML_TYPE_F32) && op->type == GGML_TYPE_F32; case GGML_OP_LIGHTNING_INDEXER: return op->src[0]->type == GGML_TYPE_F32 && From ec91ab5add06555970f98d9c5361d884f3f530f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Mon, 21 Sep 2026 13:03:41 +0200 Subject: [PATCH 16/42] docker : bump cuda to 13.4.1 (#29207) --- .github/workflows/docker.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 1de25b522..b2aafc985 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -90,8 +90,8 @@ jobs: { "tag": "cpu", "dockerfile": ".devops/s390x.Dockerfile", "platforms": "linux/s390x", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-s390x", "prebuilt_ui": true }, { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" }, { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" }, - { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.3.0", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" }, - { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.3.0", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" }, + { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.4.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" }, + { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.4.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" }, { "tag": "musa", "dockerfile": ".devops/musa.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" }, { "tag": "intel", "dockerfile": ".devops/intel.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" }, { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" }, From 6f41ac59e0a49a00483a316a22ada6b04edd2950 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Mon, 21 Sep 2026 13:44:43 +0200 Subject: [PATCH 17/42] vendor : update cpp-httplib to 0.57.0 (#29214) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Adrien Gallouët --- scripts/sync_vendor.py | 2 +- vendor/cpp-httplib/httplib.cpp | 95 ++++++++++++++++++++++++++-------- vendor/cpp-httplib/httplib.h | 5 +- 3 files changed, 77 insertions(+), 25 deletions(-) diff --git a/scripts/sync_vendor.py b/scripts/sync_vendor.py index a73ae1193..7d5ab77dd 100755 --- a/scripts/sync_vendor.py +++ b/scripts/sync_vendor.py @@ -5,7 +5,7 @@ import os import sys import subprocess -HTTPLIB_VERSION = "refs/tags/v0.56.0" +HTTPLIB_VERSION = "refs/tags/v0.57.0" # used by examples/gguf-hash, these repos have no release tag, so we pin a commit XXHASH_COMMIT = "9f465f1ea932d6ad9a26cd77496311ffa544cd68" diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index c82ff1e71..e6d1db0f0 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -1248,6 +1248,13 @@ bool parse_trailers(stream_line_reader &line_reader, Headers &dest, // to look up. split(trailer_header.data(), trailer_header.data() + trailer_header.size(), ',', [&](const char *b, const char *e) { + // A legitimate message declares only a handful of trailers. Cap the + // set so a peer cannot grow it without bound: an oversized set only + // arises from an attempt to force many colliding names into + // quadratic lookups (case_ignore::hash is unkeyed). + if (declared_trailers.size() >= CPPHTTPLIB_HEADER_MAX_COUNT) { + return; + } std::string key(b, e); if (prohibited_trailers.find(key) == prohibited_trailers.end()) { declared_trailers.insert(key); @@ -1258,6 +1265,8 @@ bool parse_trailers(stream_line_reader &line_reader, Headers &dest, size_t trailer_header_count = 0; while (strcmp(line_reader.ptr(), "\r\n") != 0) { if (line_reader.size() > CPPHTTPLIB_HEADER_MAX_LENGTH) { return false; } + // Count every received trailer field, not only the declared ones stored in + // dest, so undeclared fields cannot keep this loop running past the limit. if (trailer_header_count >= CPPHTTPLIB_HEADER_MAX_COUNT) { return false; } constexpr auto line_terminator_len = 2; @@ -1270,12 +1279,13 @@ bool parse_trailers(stream_line_reader &line_reader, Headers &dest, if (declared_trailers.find(key) != declared_trailers.end()) { dest.emplace(key, val); - trailer_header_count++; } })) { return false; } + trailer_header_count++; + if (!line_reader.getline()) { return false; } } @@ -8626,7 +8636,10 @@ bool Server::write_response_core(Stream &strm, bool close_connection, // Prepare additional headers if (close_connection || detail::has_header_token(req.headers, "Connection", "close") || - 400 <= res.status) { // Don't leave connections open after errors + 400 <= res.status || // Don't leave connections open after errors + // The client withholds the body until `100 Continue`, which was never + // sent, so whether and when the body follows is unknown. + (req.expect_100_continue_pending_ && detail::has_framed_body(req))) { res.set_header("Connection", "close"); } else { std::string s = "timeout="; @@ -8877,6 +8890,13 @@ bool Server::read_content_core( } #endif + // The client is waiting for this before it sends the body. + if (req.expect_100_continue_pending_) { + req.expect_100_continue_pending_ = false; + detail::write_response_line(strm, StatusCode::Continue_100); + strm.write("\r\n"); + } + if (!detail::read_content(strm, req, payload_max_length_, res.status, nullptr, out, true)) { return false; @@ -9714,19 +9734,20 @@ Server::process_request(Stream &strm, const std::string &remote_addr, // case-insensitive, and a 100-continue expectation in an HTTP/1.0 request // must be ignored. An expectation we do not recognize is left alone; the // 417 the section allows for one is a MAY, not a requirement. + // + // `100 Continue` itself is deferred until the body is actually read (see + // read_content_core), so a request rejected by a later handler never + // invites the client to send a body nobody will read. if (req.version != "HTTP/1.0" && detail::has_header_token(req.headers, "Expect", "100-continue")) { int status = StatusCode::Continue_100; if (expect_100_continue_handler_) { status = expect_100_continue_handler_(req, res); } - switch (status) { - case StatusCode::Continue_100: - case StatusCode::ExpectationFailed_417: - detail::write_response_line(strm, status); - strm.write("\r\n"); - break; - default: + if (status == StatusCode::Continue_100) { + req.expect_100_continue_pending_ = true; + } else { + if (res.status == -1) { res.status = status; } connection_closed = true; return write_response(strm, true, req, res); } @@ -9739,18 +9760,25 @@ Server::process_request(Stream &strm, const std::string &remote_addr, }; // WebSocket upgrade - // Check pre_routing_handler_ before upgrading so that authentication - // and other middleware can reject the request with an HTTP response - // (e.g., 401) before the protocol switches. + // Run pre_routing_handler_ and pre_request_handler_ before upgrading so + // that authentication and other middleware can reject the request with an + // HTTP response (e.g., 401) before the protocol switches. if (detail::is_websocket_upgrade(req)) { if (pre_routing_handler_ && pre_routing_handler_(req, res) == HandlerResponse::Handled) { if (res.status == -1) { res.status = StatusCode::OK_200; } - return write_response(strm, close_connection, req, res); + return write_response_with_content(strm, close_connection, req, res); } // Find matching WebSocket handler for (const auto &entry : websocket_handlers_) { if (entry.matcher->match(req)) { + req.matched_route = entry.matcher->pattern(); + if (pre_request_handler_ && + pre_request_handler_(req, res) == HandlerResponse::Handled) { + if (res.status == -1) { res.status = StatusCode::OK_200; } + return write_response_with_content(strm, close_connection, req, res); + } + // Compute accept key auto client_key = req.get_header_value("Sec-WebSocket-Key"); auto accept_key = detail::websocket_accept_key(client_key); @@ -10610,22 +10638,45 @@ ssize_t ChunkedDecoder::read_payload(char *buf, size_t len, stream_line_reader lr(strm, line_buf, sizeof(line_buf)); if (!lr.getline()) { return -1; } + // Everything below is bounded by eol rather than by the buffer's NUL, so + // the line terminator is never mistaken for line content. + const char *eol = lr.ptr() + lr.size(); + if (lr.end_with_crlf()) { + eol -= 2; + } else if (eol != lr.ptr() && eol[-1] == '\n') { + // Only reachable under CPPHTTPLIB_ALLOW_LF_AS_LINE_TERMINATOR, where + // getline() ends the line on a bare LF. That LF is the terminator, so it + // has to come off here or the check below would reject the line. + eol -= 1; + } + // RFC 9112 §7.1: chunk-size = 1*HEXDIG const char *p = lr.ptr(); int v = 0; - if (!is_hex(*p, v)) { return -1; } + if (p == eol || !is_hex(*p, v)) { return -1; } size_t chunk_len = 0; constexpr size_t chunk_len_max = (std::numeric_limits::max)(); - for (; is_hex(*p, v); ++p) { + for (; p < eol && is_hex(*p, v); ++p) { if (chunk_len > (chunk_len_max >> 4)) { return -1; } chunk_len = (chunk_len << 4) | static_cast(v); } - while (is_space_or_tab(*p)) { + while (p < eol && is_space_or_tab(*p)) { ++p; } - if (*p != '\0' && *p != ';' && *p != '\r' && *p != '\n') { return -1; } + + // RFC 9112 §7.1.1: only a chunk-ext may sit between the size and the line + // terminator, and it is built from tokens and quoted-strings, so it never + // holds a CR, LF or any other control character. getline() reads up to the + // CRLF, so a bare LF left in here would be swallowed as extension text + // while an intermediary that ends the line on it delimits the chunks + // differently, and the two disagree on where the body ends (request + // smuggling). + if (p < eol && *p != ';') { return -1; } + for (; p < eol; ++p) { + if (!is_space_or_tab(*p) && !fields::is_field_vchar(*p)) { return -1; } + } if (chunk_len == 0) { chunk_remaining = 0; @@ -14650,11 +14701,11 @@ void shutdown(session_t session, bool graceful) { auto ssl = static_cast(session); if (graceful) { - // First call sends close_notify - if (SSL_shutdown(ssl) == 0) { - // Second call waits for peer's close_notify - SSL_shutdown(ssl); - } + // Send close_notify without waiting for the peer's. The connection is + // closed right after this, so a unidirectional shutdown is enough, and an + // idle peer that never answers would otherwise hold this thread until the + // read timeout. The other backends do not wait either. + SSL_shutdown(ssl); } } diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index a3a2ff45a..2c4382560 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.56.0" -#define CPPHTTPLIB_VERSION_NUM "0x003800" +#define CPPHTTPLIB_VERSION "0.57.0" +#define CPPHTTPLIB_VERSION_NUM "0x003900" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -1756,6 +1756,7 @@ struct Request { // private members... bool body_consumed_ = false; + bool expect_100_continue_pending_ = false; size_t redirect_count_ = CPPHTTPLIB_REDIRECT_MAX_COUNT; size_t content_length_ = 0; ContentProvider content_provider_; From c21284cdf5fd833b90ac1f824cdf8065e2700dc4 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 21 Sep 2026 22:25:44 +0800 Subject: [PATCH 18/42] ggml : fix dimension and stride truncation in ggml_permute (#29227) --- ggml/src/ggml.c | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 175281409..a286083b9 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -3897,8 +3897,8 @@ struct ggml_tensor * ggml_permute( 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]; + int64_t ne[GGML_MAX_DIMS]; + size_t nb[GGML_MAX_DIMS]; ne[axis0] = a->ne[0]; ne[axis1] = a->ne[1]; From e6cef8152f6e8351a870d8e1a98627139c9c379a Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 21 Sep 2026 23:11:43 +0800 Subject: [PATCH 19/42] cuda : accelerate conv2d with implicit GEMM (#29135) --- ggml/src/ggml-cuda/conv2d.cu | 301 +++++++++++++++++++++++++++++++++-- tests/test-backend-ops.cpp | 3 + 2 files changed, 295 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-cuda/conv2d.cu b/ggml/src/ggml-cuda/conv2d.cu index 14774d4a5..10109ad36 100644 --- a/ggml/src/ggml-cuda/conv2d.cu +++ b/ggml/src/ggml-cuda/conv2d.cu @@ -1,5 +1,6 @@ #include "conv2d.cuh" #include "convert.cuh" +#include "mma.cuh" struct conv_params { const int64_t IW, IH; @@ -111,6 +112,220 @@ static void conv2d_cuda(const float * X_D, const T * K_D, float * Y_D, const con conv2d_kernel<<>>(X_D, K_D, Y_D, P); } +static __global__ void +conv2d_pad_f16(const float * input, half * output, int iw, int ih, int pw, int ph, int px, int py, int total) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) { + return; + } + const int x = i % pw - px, y = i / pw % ph - py, nc = i / (pw * ph); + output[i] = __float2half( + (unsigned) x < (unsigned) iw && (unsigned) y < (unsigned) ih ? input[(nc * ih + y) * iw + x] : 0.0f); +} + +template +static __global__ void conv2d_implicit_gemm_f16(const half * __restrict__ input, + const half * __restrict__ weight, + float * __restrict__ output, + const conv_params P, + const int split_k) { + using namespace ggml_cuda_mma; + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nthreads = 4 * warp_size; + constexpr int BM = 64, BN = 64, BK = 64; + constexpr int AS = BK / 2 + 4; + constexpr int BS = BN / 2 + 4; + __shared__ __align__(16) half2 a_s[BM][AS]; + __shared__ __align__(16) half2 b_s[BK][BS]; + + const int tid = threadIdx.y * warp_size + threadIdx.x; + const int iw = int(P.IW), ih = int(P.IH), ow = int(P.OW), oh = int(P.OH); + const int kw = KW ? KW : int(P.KW), kh = KH ? KH : int(P.KH); + const int ic = int(P.IC), oc = int(P.OC); + const int sx = int(P.ST_X), sy = int(P.ST_Y); + const int dx = int(P.DL_X), dy = int(P.DL_Y); + const int n = blockIdx.z / split_k, split = blockIdx.z % split_k; + const int m0 = blockIdx.y * BM, n0 = blockIdx.x * BN; + + const int k_total = ic * kw * kh; + const int load_lane = warp_size == 32 ? threadIdx.x : threadIdx.x % (BN / 2); + const int load_row = threadIdx.y * (warp_size / (BN / 2)) + (warp_size == 32 ? 0 : threadIdx.x / (BN / 2)); + const int spatial = n0 + 2 * load_lane; + const int spatial0 = min(spatial, ow * oh - 1), spatial1 = min(spatial + 1, ow * oh - 1); + const int y0 = spatial0 / ow, x0 = spatial0 % ow; + const int y1 = spatial1 / ow, x1 = spatial1 % ow; + const int pos0 = y0 * sy * iw + x0 * sx, pos1 = y1 * sy * iw + x1 * sx; + + [[maybe_unused]] const int wm = threadIdx.y / 2 * 32, wn = threadIdx.y % 2 * 32; +#if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) + using tile_ab = tile<16, 8, half2, get_input_data_layout()>; +# if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) + // AMD accumulator fragments transpose the input fragment's row/column mapping. + using tile_c = tile<16, 16, float, DATA_LAYOUT_J_MAJOR>; +# else + using tile_c = tile<16, 16, float>; +# endif + [[maybe_unused]] tile_c c[2][2]; +#else + if constexpr (use_mma) { + NO_DEVICE_CODE; + return; + } +#endif + constexpr int RM = 4, RN = BM * BN / (nthreads * RM); + [[maybe_unused]] const int simt_m = tid / (BN / RN) * RM, simt_n = tid % (BN / RN) * RN; + [[maybe_unused]] float c_simt[RM][RN] = {}; + const int tiles = (k_total + BK - 1) / BK; + const int begin = int(int64_t(tiles) * split / split_k) * BK; + const int end = int(int64_t(tiles) * (split + 1) / split_k) * BK; + for (int k0 = begin; k0 < end; k0 += BK) { + if (k_total % 8 == 0 && uintptr_t(weight) % 16 == 0) { +#pragma unroll + for (int i = tid; i < BM * BK / 8; i += nthreads) { + const int row = i / (BK / 8), col = 8 * (i % (BK / 8)); + const int4 v = m0 + row < oc && k0 + col < k_total ? + ((const int4 *) weight)[((m0 + row) * k_total + k0 + col) / 8] : + make_int4(0, 0, 0, 0); + *(int4 *) &a_s[row][col / 2] = v; + } + } else { +#pragma unroll + for (int i = tid; i < BM * BK / 2; i += nthreads) { + const int row = i / (BK / 2), col = 2 * (i % (BK / 2)); + half lo = __float2half(0.0f), hi = lo; + if (m0 + row < oc && k0 + col < k_total) { + lo = weight[(m0 + row) * k_total + k0 + col]; + if (k0 + col + 1 < k_total) { + hi = weight[(m0 + row) * k_total + k0 + col + 1]; + } + } + a_s[row][col / 2] = __halves2half2(lo, hi); + } + } +#pragma unroll + for (int k = load_row; k < BK; k += nthreads / (BN / 2)) { + const int ki = k0 + k; + const int ci = ki / (kw * kh), ky = ki / kw % kh, kx = ki % kw; + const int offset = ki < k_total ? (n * ic + ci) * ih * iw + ky * dy * iw + kx * dx : 0; + half lo = __float2half(0.0f), hi = lo; + if (ki < k_total && spatial < ow * oh) { + lo = input[offset + pos0]; + } + if (ki < k_total && spatial + 1 < ow * oh) { + hi = input[offset + pos1]; + } + b_s[k][load_lane] = __halves2half2(lo, hi); + } + __syncthreads(); + if constexpr (use_mma) { +#if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) +# pragma unroll + for (int k = 0; k < BK; k += 16) { + tile_ab a[2], b[2]; +# pragma unroll + for (int i = 0; i < 2; ++i) { + load_ldmatrix(a[i], &a_s[wm + 16 * i][k / 2], AS); + load_ldmatrix_trans(b[i], &b_s[k][(wn + 16 * i) / 2], BS); + } +# pragma unroll + for (int i = 0; i < 2; ++i) { +# pragma unroll + for (int j = 0; j < 2; ++j) { + mma(c[i][j], a[i], b[j]); + } + } + } +#endif + } else { +#pragma unroll 4 + for (int k = 0; k < BK; ++k) { + float a[RM], b[RN]; +#pragma unroll + for (int i = 0; i < RM; ++i) { + a[i] = __half2float(((const half *) a_s[simt_m + i])[k]); + } +#pragma unroll + for (int j = 0; j < RN; ++j) { + b[j] = __half2float(((const half *) b_s[k])[simt_n + j]); + } +#pragma unroll + for (int i = 0; i < RM; ++i) { +#pragma unroll + for (int j = 0; j < RN; ++j) { + c_simt[i][j] += a[i] * b[j]; + } + } + } + } + __syncthreads(); + } + if constexpr (use_mma) { +#if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) +# pragma unroll + for (int i = 0; i < 2; ++i) { +# pragma unroll + for (int j = 0; j < 2; ++j) { +# pragma unroll + for (int l = 0; l < c[i][j].ne; ++l) { + const int co = m0 + wm + 16 * i + c[i][j].get_i(l); + const int pos = n0 + wn + 16 * j + c[i][j].get_j(l); + if (co < oc && pos < ow * oh) { + output[(int64_t(blockIdx.z) * oc + co) * ow * oh + pos] = c[i][j].x[l]; + } + } + } + } +#endif + } else { +#pragma unroll + for (int i = 0; i < RM; ++i) { +#pragma unroll + for (int j = 0; j < RN; ++j) { + const int co = m0 + simt_m + i, pos = n0 + simt_n + j; + if (co < oc && pos < ow * oh) { + output[(int64_t(blockIdx.z) * oc + co) * ow * oh + pos] = c_simt[i][j]; + } + } + } + } +} + +static __global__ void conv2d_reduce_split_k(const float * __restrict__ partial, + float * __restrict__ output, + const int total, + const int per_batch, + const int split_k) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) { + return; + } + const int n = i / per_batch; + const float * src = partial + int64_t(n) * (split_k - 1) * per_batch + i; + float sum = 0.0f; + for (int k = 0; k < split_k; ++k) { + sum += src[int64_t(k) * per_batch]; + } + output[i] = sum; +} + +template +static void conv2d_launch_implicit_gemm(const half * input, + const half * weight, + float * output, + const conv_params & params, + int split_k, + dim3 grid, + dim3 block, + cudaStream_t stream) { + if (params.KW == 3 && params.KH == 3) { + conv2d_implicit_gemm_f16<3, 3, use_mma><<>>(input, weight, output, params, split_k); + } else if (params.KW == 1 && params.KH == 1) { + conv2d_implicit_gemm_f16<1, 1, use_mma><<>>(input, weight, output, params, split_k); + } else { + conv2d_implicit_gemm_f16<0, 0, use_mma><<>>(input, weight, output, params, split_k); + } +} + static void conv2d_cuda_f16(const float * X_D, const half * K_D, float * Y_D, const conv_params P, cudaStream_t st) { conv2d_cuda(X_D, K_D, Y_D, P, st); } @@ -126,6 +341,7 @@ void ggml_cuda_op_conv2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const float * X_D = (const float *) input->data; float * Y_D = (float *) dst->data; + GGML_ASSERT(input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_contiguous(input)); GGML_ASSERT(ggml_is_contiguous(kernel)); GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32); @@ -146,19 +362,86 @@ void ggml_cuda_op_conv2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { // No cwhn GGML_ASSERT(p[6] == false); - const int IW = input->ne[0]; // input_w - const int IH = input->ne[1]; // input_h - const int OW = dst->ne[0]; // output_w - const int OH = dst->ne[1]; // output_h - const int KW = kernel->ne[0]; // kernel_w - const int KH = kernel->ne[1]; // kernel_h - const int IC = input->ne[2]; // input_channels - const int OC = kernel->ne[3]; // ouptut_chanles - const int B = input->ne[3]; // n_batches + const int64_t IW = input->ne[0]; // input_w + const int64_t IH = input->ne[1]; // input_h + const int64_t OW = dst->ne[0]; // output_w + const int64_t OH = dst->ne[1]; // output_h + const int64_t KW = kernel->ne[0]; // kernel_w + const int64_t KH = kernel->ne[1]; // kernel_h + const int64_t IC = input->ne[2]; // input_channels + const int64_t OC = kernel->ne[3]; // ouptut_chanles + const int64_t B = input->ne[3]; // n_batches const int64_t total = B * OC * OH * OW; conv_params params = { IW, IH, OW, OH, KW, KH, ST_X, ST_Y, PD_X, PD_Y, DL_X, DL_Y, IC, OC, B, total }; + const auto & device = ggml_cuda_info().devices[ctx.device]; + const bool use_mma = + turing_mma_available(device.cc) || amd_wmma_available(device.cc) || amd_mfma_available(device.cc); + // MUSA can share the tiling without a native fragment implementation in mma.cuh. + const bool use_simt = GGML_CUDA_CC_IS_MTHREADS(device.cc); + const bool pointwise = KW == 1 && KH == 1 && ST_X == 1 && ST_Y == 1 && PD_X == 0 && PD_Y == 0; + const bool use_blas = pointwise && fast_fp16_hardware_available(device.cc); + // Short reductions on small maps do not amortize conversion and launch costs. + const bool small_conv = IC * KW * KH < 64 && OW * OH < 512; + + const int64_t limit = INT_MAX - 256; + const int64_t padded_w = IW + 2 * int64_t(PD_X), padded_h = IH + 2 * int64_t(PD_Y); + const bool padded_fits = padded_w > 0 && padded_w <= limit && padded_h > 0 && padded_h <= limit && + padded_w * padded_h <= limit && IC * B <= limit / (padded_w * padded_h); + if (kernel->type == GGML_TYPE_F16 && (use_mma || use_blas || use_simt) && (use_blas || !small_conv) && + ggml_nelements(input) <= limit && ggml_nelements(kernel) <= limit && total <= limit && padded_fits && + PD_X >= 0 && PD_Y >= 0 && ST_X > 0 && ST_Y > 0 && DL_X > 0 && DL_Y > 0 && + (OW - 1) * ST_X + (KW - 1) * DL_X < padded_w && (OH - 1) * ST_Y + (KH - 1) * DL_Y < padded_h && + (OC + 63) / 64 <= 65535 && B <= 65535) { + const int pw = int(padded_w), ph = int(padded_h); + const int padded_total = int(padded_w * padded_h * IC * B); + + ggml_cuda_pool_alloc x_half(ctx.pool(), padded_total); + // Match im2col's F16 input precision, but expand patches only in shared memory and accumulate in F32. + if (PD_X == 0 && PD_Y == 0) { + ggml_get_to_fp16_cuda(input->type)(X_D, x_half.get(), padded_total, st); + } else { + conv2d_pad_f16<<<(padded_total + 255) / 256, 256, 0, st>>>(X_D, x_half.get(), int(IW), int(IH), pw, ph, + PD_X, PD_Y, padded_total); + } + const conv_params padded_params = { pw, ph, OW, OH, KW, KH, ST_X, ST_Y, 0, 0, DL_X, DL_Y, IC, OC, B, total }; + if (use_blas) { + const float alpha = 1.0f, beta = 0.0f; + const int positions = int(OW * OH); + cublasHandle_t cublas_h = ctx.cublas_handle(); + for (int n = 0; n < B; ++n) { + CUBLAS_CHECK(cublasGemmEx(cublas_h, CUBLAS_OP_N, CUBLAS_OP_N, positions, int(OC), int(IC), &alpha, + x_half.get() + int64_t(n) * IC * positions, CUDA_R_16F, positions, K_D, + CUDA_R_16F, int(IC), &beta, Y_D + int64_t(n) * OC * positions, CUDA_R_32F, + positions, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + } + return; + } + const int64_t blocks = ((OW * OH + 63) / 64) * ((OC + 63) / 64) * B; + const int target = 8 * ggml_cuda_info().devices[ctx.device].nsm; + // Split long reductions so small spatial maps still occupy the GPU. + const int split_k = int(std::min({ int64_t(32), int64_t(65535) / B, (IC * KW * KH + 63) / 64, + std::max(int64_t(1), (target + blocks - 1) / blocks) })); + + ggml_cuda_pool_alloc partial(ctx.pool()); + float * result = split_k == 1 ? Y_D : partial.alloc(total * split_k); + const dim3 block(device.warp_size, 4); + const dim3 grid(unsigned((OW * OH + 63) / 64), unsigned((OC + 63) / 64), unsigned(B * split_k)); + if (use_mma) { + conv2d_launch_implicit_gemm(x_half.get(), (const half *) K_D, result, padded_params, split_k, grid, + block, st); + } else { + conv2d_launch_implicit_gemm(x_half.get(), (const half *) K_D, result, padded_params, split_k, grid, + block, st); + } + if (split_k > 1) { + conv2d_reduce_split_k<<<(total + 255) / 256, 256, 0, st>>>(result, Y_D, int(total), int(OC * OW * OH), + split_k); + } + return; + } + if (kernel->type == GGML_TYPE_F16) { conv2d_cuda_f16(X_D, (half *) K_D, Y_D, params, st); } else { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index cc68e9ca7..e4d04f13b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9367,6 +9367,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_conv_2d({ 256, 256, 192, 1 }, { 3, 3, 192, 96 }, kernel_type, 1, 1, 1, 1, 1, 1, false)); // bool cwhn = false test_cases.emplace_back(new test_conv_2d({ 256, 256, 192, 1 }, { 3, 3, 192, 96 }, kernel_type, 1, 1, 1, 1, 1, 1, true)); // bool cwhn = true } + test_cases.emplace_back(new test_conv_2d({ 19, 17, 8, 2 }, { 3, 3, 8, 65 }, GGML_TYPE_F16, 1, 1, 1, 1, 1, 1)); + test_cases.emplace_back(new test_conv_2d({ 19, 17, 16, 3 }, { 3, 3, 16, 33 }, GGML_TYPE_F16, 2, 3, 4, 2, 2, 1)); + test_cases.emplace_back(new test_conv_2d({ 13, 11, 16, 3 }, { 1, 1, 16, 33 }, GGML_TYPE_F16, 1, 1, 0, 0, 1, 1)); // sycl backend will limit task global_range < MAX_INT // test cases for 2D im2col with large input W and H (occurs in stable-diffusion) From f4e276a2066a40cd200db17c1131826d6c0c7a94 Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Mon, 21 Sep 2026 18:00:51 +0200 Subject: [PATCH 20/42] ggml-cuda : convert contiguous tensors four elements at a time (#29155) convert_unary handles the contiguous case through the general strided kernel, one element per thread: each lane reads 4 bytes and writes 2. Converting the activations for a bf16 matrix multiplication that way moves 126 MB in 1021 us on gfx1151, about 65% of what the memory system can do. Give the contiguous path its own kernel that takes four elements per thread through a vector type, so a warp loads 512 bytes at a time instead of 128. It is used only when the element count is a multiple of four and both pointers carry the alignment the vector type needs, and falls back to the strided kernel otherwise. Model level, Qwen3.8-Next-Flash IQ3_XXS on gfx1151, llama-bench -ub 2048 -r 6, mean of the last 3 reps, ABBA counterbalanced: pp2048 688.0 680.0 -> 694.3 691.1 +1.26% tg128 24.8 24.8 -> 24.8 24.8 +0.14% Every conversion in a prefill takes the new kernel (kernel trace: 1146 convert_unary_cont_vec4, no convert_unary). Output is bit identical; MUL_MAT, MUL_MAT_ID, CPY, CONT, GET_ROWS and SET_ROWS pass. Assisted-by: Claude Opus 5 --- ggml/src/ggml-cuda/convert.cu | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index 360c614a4..0619f4760 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -439,6 +439,29 @@ static __global__ void convert_unary( } } +template struct alignas(sizeof(T)*4) cvt_vec4 { T v[4]; }; + +// four elements per thread, so a warp moves 512B (RDNA) / 1k (CDNA) per load +template +static __global__ void convert_unary_cont_vec4( + const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k4) { + const int64_t i = (int64_t)blockDim.x*blockIdx.x + threadIdx.x; + + if (i >= k4) { + return; + } + + const cvt_vec4 xv = ((const cvt_vec4 *) vx)[i]; + + cvt_vec4 yv; +#pragma unroll + for (int j = 0; j < 4; ++j) { + yv.v[j] = ggml_cuda_cast(xv.v[j]); + } + + ((cvt_vec4 *) y)[i] = yv; +} + template static void convert_unary_cuda(const void * vx, dst_t * y, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, @@ -452,6 +475,15 @@ static void convert_unary_cuda(const void * vx, dst_t * y, template static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + if (k % 4 == 0 && + (uintptr_t) vx % alignof(cvt_vec4) == 0 && + (uintptr_t) y % alignof(cvt_vec4) == 0) { + const int64_t k4 = k/4; + const int64_t num_blocks = (k4 + CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / CUDA_DEQUANTIZE_BLOCK_SIZE; + convert_unary_cont_vec4<<>>(vx, y, k4); + return; + } + convert_unary_cuda(vx, y, k, 1, 1, 1, k, k, k, stream); } From b1c2863e2c2c861ab3009aad8622826472100750 Mon Sep 17 00:00:00 2001 From: lingyezhixing <144504450+lingyezhixing@users.noreply.github.com> Date: Tue, 22 Sep 2026 00:11:29 +0800 Subject: [PATCH 21/42] cuda: fix sm_70 tile compilation error (#29224) The 5-argument load_ldmatrix added in 1884824fd only defines tile<16,8>, so the Volta tile<8,4> does not match. See https://github.com/ggml-org/llama.cpp/issues/29222 for details. Building on 1884824fd, generalize the tile shape of the 5-argument load_ldmatrix from <16,8> to , so the non-swizzle branch forwards to the 3-argument loader for any shape. Local compilation and testing passed. Assisted-by: DeepSeek V4.1 Flash (OpenCode) --- ggml/src/ggml-cuda/mma.cuh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/mma.cuh b/ggml/src/ggml-cuda/mma.cuh index 6af2b6a14..3583ba5e1 100644 --- a/ggml/src/ggml-cuda/mma.cuh +++ b/ggml/src/ggml-cuda/mma.cuh @@ -873,14 +873,16 @@ namespace ggml_cuda_mma { } // Load from tile element (i0, j0), swz tells if the tile is stored swizzled. - template + template static __device__ __forceinline__ void load_ldmatrix( - tile<16, 8, T, dl> & t, const T * __restrict__ tile_base, const int i0, const int j0, const int stride) { + tile & t, const T * __restrict__ tile_base, const int i0, const int j0, const int stride) { if constexpr (!swz) { load_ldmatrix(t, tile_base + i0*stride + j0, stride); return; } #if defined(TURING_MMA_AVAILABLE) + static_assert(I == 16, "bad tile width"); + static_assert(J == 8, "bad tile height"); const int i = i0 + threadIdx.x % t.I; const int j = j0 + (threadIdx.x / t.I) * (t.J / 2); int * xi = (int *) t.x; From 96550613656e7f024df65f91cf8b2d80a83cf09e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 21 Sep 2026 19:13:04 +0300 Subject: [PATCH 22/42] llama-context : report graph inputs and input tensors during sched reserve (#26625) * llama-context : report graph inputs and input tensors during sched reserve - fix the tg (token generation) graph bs label to use n_seqs instead of a hardcoded 1 - report the number of graph inputs from llm_graph_result::inputs for both the pp and tg graphs - report the number of input tensors (nodes and their src tensors flagged with GGML_TENSOR_FLAG_INPUT) - log a warning when an input tensor has an op other than GGML_OP_NONE - log a trace line for each input tensor and the nodes (name and op) that use it Assisted-by: llama.cpp:DeepSeek-v4-Flash-0731 * cont : count input tensors before reserving the sched * wip * llama-graph : name the unnamed graph input tensors - name the kv-cache idxs input tensors (attn_inp_k_idxs, attn_inp_v_idxs) - name the recurrent state copy idxs input tensor (rs_s_copy) - report the input tensor shape in the sched_reserve trace Assisted-by: pi:llama.cpp/Qwen3.8-27B * llama-context : rename "graph inputs" to "graph input objects" Assisted-by: pi:llama.cpp/Qwen3.8-27B * llama-context : report the sched reserve graph stats on a single line - print nodes, splits, input objects and input tensors in one line - when the pp and tg graphs differ, print each value as 'pp / tg' and annotate the line with the batch sizes used for each graph Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : pad logs --- src/llama-context.cpp | 83 +++++++++++++++++++++++++++++++++--------- src/llama-context.h | 1 + src/llama-graph.cpp | 11 +++++- src/llama-kv-cache.cpp | 2 + 4 files changed, 79 insertions(+), 18 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ef53728d1..fcd4dfb13 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -19,6 +19,7 @@ #include #include #include +#include // // llama_context @@ -579,6 +580,40 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 } } +static int llama_graph_n_input_tensors(ggml_cgraph * gf) { + std::unordered_map> users; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + if (node->flags & GGML_TENSOR_FLAG_INPUT) { + users[node].push_back(node); + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + ggml_tensor * src = node->src[j]; + if (!src) { + break; + } + if (src->flags & GGML_TENSOR_FLAG_INPUT) { + users[src].push_back(node); + } + } + } + + for (const auto & [tensor, nodes] : users) { + if (tensor->op != GGML_OP_NONE) { + LLAMA_LOG_WARN("%s: input tensor '%32s' has op %s, expected GGML_OP_NONE\n", + __func__, tensor->name, ggml_op_name(tensor->op)); + } + for (const ggml_tensor * node : nodes) { + LLAMA_LOG_DEBUG("%s: input tensor '%32s' [%s, ne = { %5" PRId64 ", %5" PRId64 ", %5" PRId64 ", %5" PRId64 " }] is used by node '%s' (%s)\n", + __func__, tensor->name, ggml_type_name(tensor->type), + tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3], + node->name, ggml_op_name(node->op)); + } + } + + return (int) users.size(); +} + void llama_context::sched_reserve() { if (!sched_need_reserve) { return; @@ -624,11 +659,15 @@ void llama_context::sched_reserve() { resolve_fused_ops(mctx.get(), n_seqs); // reserve worst-case graph - int n_splits_pp = -1; - int n_nodes_pp = -1; + int n_splits_pp = -1; + int n_nodes_pp = -1; + int n_inputs_pp = -1; + int n_input_tensors_pp = -1; - int n_splits_tg = -1; - int n_nodes_tg = -1; + int n_splits_tg = -1; + int n_nodes_tg = -1; + int n_inputs_tg = -1; + int n_input_tensors_tg = -1; const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max); @@ -648,8 +687,10 @@ void llama_context::sched_reserve() { } } - n_splits_pp = ggml_backend_sched_get_n_splits(sched.get()); - n_nodes_pp = ggml_graph_n_nodes(gf); + n_splits_pp = ggml_backend_sched_get_n_splits(sched.get()); + n_nodes_pp = ggml_graph_n_nodes(gf); + n_inputs_pp = get_gf_res_reserve()->inputs.size(); + n_input_tensors_pp = this->n_input_tensors; } // reserve with tg (token generation) graph to get the number of splits and nodes @@ -659,8 +700,10 @@ void llama_context::sched_reserve() { throw std::runtime_error("failed to allocate compute tg buffers"); } - n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); - n_nodes_tg = ggml_graph_n_nodes(gf); + n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); + n_nodes_tg = ggml_graph_n_nodes(gf); + n_inputs_tg = get_gf_res_reserve()->inputs.size(); + n_input_tensors_tg = this->n_input_tensors; } // reserve again with pp graph to avoid ggml-alloc reallocations during inference @@ -698,16 +741,21 @@ void llama_context::sched_reserve() { } } - if (n_nodes_pp == n_nodes_tg) { - LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); - } else { - LLAMA_LOG_INFO("%s: graph nodes = %d (with bs=%d), %d (with bs=1)\n", __func__, n_nodes_pp, n_tokens, n_nodes_tg); - } + { + const bool diff = n_nodes_pp != n_nodes_tg || n_splits_pp != n_splits_tg || + n_inputs_pp != n_inputs_tg || n_input_tensors_pp != n_input_tensors_tg; - if (n_splits_pp == n_splits_tg) { - LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits_pp); - } else { - LLAMA_LOG_INFO("%s: graph splits = %d (with bs=%d), %d (with bs=1)\n", __func__, n_splits_pp, n_tokens, n_splits_tg); + const auto val = [diff](int v_pp, int v_tg) -> std::string { + return diff ? format("%d / %d", v_pp, v_tg) : format("%d", v_pp); + }; + + LLAMA_LOG_INFO("%s: graph%s: nodes = %s, splits = %s, input objects = %s, input tensors = %s\n", + __func__, + diff ? format(" (pp bs=%d, tg bs=%d)", n_tokens, n_seqs).c_str() : "", + val(n_nodes_pp, n_nodes_tg).c_str(), + val(n_splits_pp, n_splits_tg).c_str(), + val(n_inputs_pp, n_inputs_tg).c_str(), + val(n_input_tensors_pp, n_input_tensors_tg).c_str()); } const int64_t t_end_us = ggml_time_us(); @@ -2475,6 +2523,7 @@ ggml_cgraph * llama_context::graph_reserve( auto * gf = model.build_graph(gparams); + this->n_input_tensors = llama_graph_n_input_tensors(gf); this->n_outputs = save_n_outputs; // initialize scheduler with the specified graph diff --git a/src/llama-context.h b/src/llama-context.h index b7a9db591..77ef92fc6 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -333,6 +333,7 @@ private: // reuse the batch_allocr to avoid unnecessary memory allocations std::unique_ptr balloc; + uint32_t n_input_tensors = 0; // number of tensors marked as input during the last graph reserve uint32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch std::vector output_ids; // map batch token positions to ids of the logits and embd buffers diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index fd4290cf0..02ae8bd92 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2451,6 +2451,7 @@ ggml_tensor * llm_graph_context::build_inp_pos() const { cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t)n_tokens*hparams.n_pos_per_embd()); ggml_set_input(cur); + cb(cur, "inp_pos", -1); res->add_input(std::move(inp)); @@ -2465,7 +2466,7 @@ ggml_tensor * llm_graph_context::build_inp_attn_scale() const { // this need to be 1x1xN for broadcasting cur = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); ggml_set_input(cur); - ggml_set_name(cur, "attn_scale"); + cb(cur, "inp_attn_scale", -1); res->add_input(std::move(inp)); @@ -2487,6 +2488,7 @@ ggml_tensor * llm_graph_context::build_inp_out_ids() const { cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_outputs); ggml_set_input(cur); + ggml_set_name(cur, "out_ids"); res->add_input(std::move(inp)); @@ -2500,6 +2502,7 @@ ggml_tensor * llm_graph_context::build_inp_mean() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, ubatch.n_seqs_unq); ggml_set_input(cur); + ggml_set_name(cur, "mean"); res->add_input(std::move(inp)); @@ -2513,6 +2516,7 @@ ggml_tensor * llm_graph_context::build_inp_cls() const { cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_seqs_unq); ggml_set_input(cur); + ggml_set_name(cur, "cls"); res->add_input(std::move(inp)); @@ -2537,6 +2541,7 @@ ggml_tensor * llm_graph_context::build_inp_cross_embd() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_enc); ggml_set_input(cur); + ggml_set_name(cur, "cross_embd"); res->add_input(std::move(inp)); @@ -2550,6 +2555,7 @@ ggml_tensor * llm_graph_context::build_inp_pos_bucket_enc() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_tokens, n_tokens); ggml_set_input(cur); + ggml_set_name(cur, "pos_bucket_enc"); res->add_input(std::move(inp)); @@ -2567,6 +2573,7 @@ ggml_tensor * llm_graph_context::build_inp_pos_bucket_dec() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_tokens); ggml_set_input(cur); + ggml_set_name(cur, "pos_bucket_dec"); res->add_input(std::move(inp)); @@ -2735,6 +2742,7 @@ llm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() con // note: there is no KV cache, so the number of KV values is equal to the number of tokens in the batch inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1); ggml_set_input(inp->self_kq_mask); + cb(inp->self_kq_mask, "self_kq_mask", -1); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3511,6 +3519,7 @@ static std::unique_ptr build_rs_inp_impl( inp->s_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs); ggml_set_input(inp->s_copy); + ggml_set_name(inp->s_copy, "rs_s_copy"); inp->s_copy_main = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0); inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a342ee119..332d1abe0 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1412,6 +1412,7 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); ggml_set_input(k_idxs); + ggml_set_name(k_idxs, "attn_inp_k_idxs"); return k_idxs; } @@ -1428,6 +1429,7 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama } ggml_set_input(v_idxs); + ggml_set_name(v_idxs, "attn_inp_v_idxs"); return v_idxs; } From c641dfa83338b717a63f2f3371cac0acee18b53b Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 21 Sep 2026 20:19:11 +0300 Subject: [PATCH 23/42] test-save-load-state : compare logits with NMSE and feed expected tokens (#29238) Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp --- tests/test-save-load-state.cpp | 137 +++++++++++++++++++++++++-------- 1 file changed, 106 insertions(+), 31 deletions(-) diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 74d1ba6c2..08c7c6772 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -11,6 +11,43 @@ #include #include +constexpr double NMSE_THRESHOLD = 1e-5; + +// normalized mean squared error = mse(a, b) / mse(a, 0) +static double nmse(const std::vector & a, const std::vector & b) { + GGML_ASSERT(a.size() == b.size()); + double mse_a_b = 0.0; + double mse_a_0 = 0.0; + + for (size_t i = 0; i < a.size(); i++) { + const float a_i = a[i]; + const float b_i = b[i]; + + mse_a_b += (double) (a_i - b_i) * (a_i - b_i); + mse_a_0 += (double) a_i * a_i; + } + + return mse_a_b / mse_a_0; +} + +struct generation_result { + llama_tokens tokens; + std::vector> logits; + + bool empty() const { return tokens.empty(); } +}; + +static bool get_current_logits(llama_context * ctx, std::vector & out) { + const auto * vocab = llama_model_get_vocab(llama_get_model(ctx)); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); + const float * logits = llama_get_logits_ith(ctx, -1); + if (logits == nullptr) { + return false; + } + out.assign(logits, logits + n_vocab); + return true; +} + struct llama_batch_ptr { llama_batch batch; @@ -28,15 +65,22 @@ struct llama_batch_ptr { const llama_batch & get() const { return batch; } }; -static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) { - llama_tokens result; +static generation_result generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) { + generation_result result; llama_batch_ptr batch(1, 0, 1); for (int i = 0; i < n_predict; i++) { + std::vector logits; + if (!get_current_logits(ctx, logits)) { + LOG_ERR("\n%s: failed to get logits\n", __func__); + return {}; + } + auto next_token = llama_sampler_sample(smpl, ctx, -1); LOG("%d ", next_token); - result.push_back(next_token); + result.tokens.push_back(next_token); + result.logits.push_back(std::move(logits)); common_batch_clear(batch.get()); common_batch_add(batch.get(), next_token, n_past, {seq_id}, true); @@ -51,12 +95,61 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i return result; } +static bool generate_tokens_compare( + llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id, + const generation_result & expected) { + if (expected.tokens.size() != expected.logits.size() || expected.tokens.size() < (size_t) n_predict) { + LOG_ERR("\n%s: invalid expected generation\n", __func__); + return false; + } + + llama_batch_ptr batch(1, 0, 1); + + for (int i = 0; i < n_predict; i++) { + std::vector logits; + if (!get_current_logits(ctx, logits)) { + LOG_ERR("\n%s: failed to get logits\n", __func__); + return false; + } + if (logits.size() != expected.logits[i].size()) { + LOG_ERR("\n%s: logits size mismatch at step %d: %zu != %zu\n", __func__, i, logits.size(), expected.logits[i].size()); + return false; + } + + const double nmse_val = nmse(expected.logits[i], logits); + LOG_TRC("%s: step %d nmse = %.6e\n", __func__, i, nmse_val); + if (nmse_val > NMSE_THRESHOLD) { + LOG_ERR("\n%s: error: NMSE at step %d is %.6e (threshold %.1e)\n", __func__, i, nmse_val, NMSE_THRESHOLD); + return false; + } + + const auto next_token = llama_sampler_sample(smpl, ctx, -1); + const auto expected_token = expected.tokens[i]; + + LOG("%d ", next_token); + if (next_token != expected_token) { + LOG_TRC("%s: sampled token %d differs from expected %d, using expected token\n", __func__, next_token, expected_token); + } + + common_batch_clear(batch.get()); + common_batch_add(batch.get(), expected_token, n_past, {seq_id}, true); + + if (llama_decode(ctx, batch.get())) { + LOG_ERR("\n%s: failed to evaluate\n", __func__); + return false; + } + n_past++; + } + + return true; +} + // Test 1: baseline // - decode all but the last token // - save state to disk // - decode the last token // - generate n_predict tokens -static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { +static generation_result test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { auto params_ctx = common_context_params_to_llama(params); params_ctx.n_seq_max = 2; auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; @@ -166,7 +259,7 @@ static bool test_seq_rm_isolated( // - load state from file // - replay the last prompt token // - generate n_predict tokens and compare against expected result -static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) { +static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const generation_result & expected_result) { auto params_ctx = common_context_params_to_llama(params); params_ctx.n_seq_max = 2; auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; @@ -195,14 +288,8 @@ static bool test_state_load(struct llama_model * model, const struct common_para } n_past++; - // Generate tokens - auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0); - if (result.empty()) { - return false; - } - - if (result != expected_result) { - LOG_ERR("\n%s: error: generation differs from expected\n", __func__); + // Generate tokens and compare logits against the baseline + if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, params.n_predict, 0, expected_result)) { return false; } @@ -217,7 +304,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para // - replay the last prompt token // - migrate KV cache from seq 0 to seq 1 via the CPU path // - generate n_predict tokens on seq 1 and compare against expected result -static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) { +static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const generation_result & expected_result) { auto params_ctx = common_context_params_to_llama(params); params_ctx.n_seq_max = 2; auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; @@ -267,14 +354,8 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset); } - // Generate tokens on seq 1 - auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1); - if (result.empty()) { - return false; - } - - if (result != expected_result) { - LOG_ERR("\n%s: error: generation differs from expected\n", __func__); + // Generate tokens and compare logits against the baseline + if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, params.n_predict, 1, expected_result)) { return false; } @@ -289,7 +370,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par // - replay the last prompt token // - migrate KV cache from seq 0 to seq 1 via the on-device path // - generate n_predict tokens on seq 1 and compare against expected result -static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) { +static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const generation_result & expected_result) { auto params_ctx = common_context_params_to_llama(params); params_ctx.n_seq_max = 2; auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; @@ -339,14 +420,8 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset); } - // Generate tokens on seq 1 - auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1); - if (result.empty()) { - return false; - } - - if (result != expected_result) { - LOG_ERR("\n%s: error: generation differs from expected\n", __func__); + // Generate tokens and compare logits against the baseline + if (!generate_tokens_compare(ctx.get(), smpl.get(), n_past, params.n_predict, 1, expected_result)) { return false; } From fb34fc262c1b43f1832c7472429fb2247d650493 Mon Sep 17 00:00:00 2001 From: Foad Abo Dahood <32059146+masterFoad@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:31:56 +0300 Subject: [PATCH 24/42] metal : fix mask bounds in flash attention block pre-pass (#29220) --- ggml/src/ggml-metal/kernels/fa.metal | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index 71e6e373e..f26d493d5 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -142,7 +142,7 @@ kernel void kernel_flash_attn_ext_blk( const int32_t i1 = tgpig[1]; const int32_t i0 = tgpig[0]; - char res = i0*C + C > args.ne30 ? 1 : 0; + char res = i0*C + C > args.ne30 || i1*Q + Q > args.ne31 ? 1 : 0; device const half * mask_src = (device const half *) (mask + (i1*Q)*args.nb31 + i2*args.nb32 + i3*args.nb33) + i0*C + tiisg; From ff0dbb975e93a9a2899efa34bdd32d1c5cfbc183 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Mon, 21 Sep 2026 22:50:20 +0200 Subject: [PATCH 25/42] vendor : update cpp-httplib to 0.57.1 (#29239) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Adrien Gallouët --- scripts/sync_vendor.py | 2 +- vendor/cpp-httplib/httplib.cpp | 19 +++++++++++-------- vendor/cpp-httplib/httplib.h | 4 ++-- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/scripts/sync_vendor.py b/scripts/sync_vendor.py index 7d5ab77dd..b320884bb 100755 --- a/scripts/sync_vendor.py +++ b/scripts/sync_vendor.py @@ -5,7 +5,7 @@ import os import sys import subprocess -HTTPLIB_VERSION = "refs/tags/v0.57.0" +HTTPLIB_VERSION = "refs/tags/v0.57.1" # used by examples/gguf-hash, these repos have no release tag, so we pin a commit XXHASH_COMMIT = "9f465f1ea932d6ad9a26cd77496311ffa544cd68" diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index e6d1db0f0..df79b3c7d 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -3854,11 +3854,13 @@ bool read_content(Stream &strm, T &x, size_t payload_max_length, int &status, ssize_t write_request_line(Stream &strm, const std::string &method, const std::string &path) { - // A request target must not carry CR/LF (or other control octets); otherwise - // a value smuggled into it splits the request line and injects headers or a - // whole request. The same field-value check already guards header values in - // check_and_write_headers and the request target in - // perform_websocket_handshake; apply it here too. + // Neither the method nor the request target may carry CR/LF (or other + // control octets); otherwise a value smuggled into either splits the request + // line and injects headers or a whole request. The method must be a token + // (RFC 9110 Section 9.1), which also rejects an empty method and embedded + // spaces. The target gets the same field-value check that already guards + // header values in check_and_write_headers. + if (!fields::is_token(method)) { return -1; } if (!fields::is_field_value(path)) { return -1; } std::string s = method; @@ -11105,9 +11107,10 @@ bool ClientImpl::write_request(Stream &strm, Request &req, // Write request line and headers if (detail::write_request_line(bstrm, req.method, path_with_query) < 0) { - // A rejected target (e.g. CR/LF smuggled in via a decoded redirect - // Location under set_path_encode(false)) must fail the request cleanly - // instead of emitting a request-line-less, header-injecting request. + // A rejected method (not a token, e.g. carrying CR/LF) or target (e.g. + // CR/LF smuggled in via a decoded redirect Location under + // set_path_encode(false)) must fail the request cleanly instead of + // emitting a request-line-less, header-injecting request. error = Error::Write; output_error_log(error, &req); return false; diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index 2c4382560..ba603ed25 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.57.0" -#define CPPHTTPLIB_VERSION_NUM "0x003900" +#define CPPHTTPLIB_VERSION "0.57.1" +#define CPPHTTPLIB_VERSION_NUM "0x003901" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 From 58367713a6935c0810103378144008df32e3d5db Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Mon, 21 Sep 2026 14:49:52 -0700 Subject: [PATCH 26/42] hexagon: new HMX-optimized GATED_DELTA_NET (#29199) * hex-gdn: start putting together HMX support for GDN * hex-gdn: working hmx but not-pipelined and slow for now * hex-gdn: re-write vtcm layout handling and prep for pipelining * hex-gdn: starting to pipeline hmx and dmas * hex-gdn: add hvx threading for most pipeline stages * hex-gdb: add detailed trace events * hex-gdn: vectorize expfs and use aligned hvx reads/writes * hex-gnd: vectorize the rest of expf * hex-gdn: optimize tail processing (pad partial chunks) * hex-gdb: avoid float up/down casts in hot loops * hex-fa: remove float up/down casts from inner loops * hex-gdn: do exp() in f16 to improve HVX utilization * hex-gdn: optimize tiler * hex-hmx: bump hmx-queue to 128 and dispatch all GDN gemms at once * hex-gdn: further pipeline improvements * hex-gdn: optimize gdn prep stage * hex-gdn: yet more tweaks to optimize GND_SOLVE task and pipeline * hex-gdn: improve accuracy and optmize gdn-prep further * hex-gdn: fix rebase conflict * hex-bufs: revert max_bufsize enforcement, it is enough to just enforce max_vmem * hex-scripts: improved inspect script to avoid false alarms in reg spill detector * hex-fa: improve inline softmax with in-reg VKQ32 accum * hex-fa: minor improvement for dma pipeline in hvx kernel * hex-fa: reduce ddr reads by 20-30% during token gen * hex-gdn: proper alignment for hvx vtcm spads --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 98 +- ggml/src/ggml-hexagon/htp-opnode.h | 4 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 571 ++++--- ggml/src/ggml-hexagon/htp/flash-attn-ops.h | 5 +- .../ggml-hexagon/htp/gated-delta-net-ops.c | 1324 ++++++++++++++++- .../ggml-hexagon/htp/gated-delta-net-ops.h | 10 +- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 4 +- ggml/src/ggml-hexagon/htp/htp-ops.h | 8 + ggml/src/ggml-hexagon/htp/main.c | 2 +- scripts/snapdragon/ggml-hexagon-inspect.py | 388 ++++- scripts/snapdragon/ggml-hexagon-profile.py | 3 +- scripts/snapdragon/ggml-hexagon-trace.py | 3 +- scripts/snapdragon/run.py | 3 + 13 files changed, 2011 insertions(+), 412 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index ec5a4aeb6..58806e37f 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -100,6 +100,7 @@ static bool opt_dma64 = false; static int opt_mm_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported) static int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported) +static int opt_gdn_select = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported) static int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 = unfused ALLREDUCE (DMA), 0 = fallback to CPY+FENCE // Default PMU events, if profiling with PMU (mode=2) is enabled @@ -182,6 +183,13 @@ static const char * htp_event_name(uint16_t id) { case HTP_TRACE_EVT_HVX_FA_Q_PREP: return "HVX_Q_PREP"; case HTP_TRACE_EVT_HVX_FA_K_PREP: return "HVX_K_PREP"; case HTP_TRACE_EVT_HVX_FA_V_PREP: return "HVX_V_PREP"; + case HTP_TRACE_EVT_HVX_GDN_PREP: return "HVX_GDN_PREP"; + case HTP_TRACE_EVT_HVX_GDN_SOLVE: return "HVX_GDN_SOLVE"; + case HTP_TRACE_EVT_HVX_GDN_V_PREP: return "HVX_GDN_V_PREP"; + case HTP_TRACE_EVT_HVX_GDN_D_PREP: return "HVX_GDN_D_PREP"; + case HTP_TRACE_EVT_HVX_GDN_OUT: return "HVX_GDN_OUT"; + case HTP_TRACE_EVT_HVX_GDN_STATE: return "HVX_GDN_STATE"; + case HTP_TRACE_EVT_HVX_GDN_REM: return "HVX_GDN_REM"; case HTP_TRACE_EVT_HMX_COMP: return "HMX_COMP"; case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH"; case HTP_TRACE_EVT_INIT: return "INIT"; @@ -472,7 +480,6 @@ struct ggml_hexagon_session { uint32_t n_hmx = 0; uint64_t vtcm_size = 0; size_t max_vmem = 0; - size_t max_bufsize = 0; uint32_t fence_seq = 0; std::atomic batch_req_seq{0}; @@ -538,7 +545,6 @@ struct ggml_backend_hexagon_device_context { int dev_id; ggml_hexagon_device_config config; ggml_backend_dev_t dev = nullptr; - size_t max_bufsize = 0; ggml_backend_buffer_type buffer_type = {}; ggml_backend_buffer_type host_buffer_type = {}; @@ -554,9 +560,6 @@ struct ggml_backend_hexagon_device_context { ggml_hexagon_session * session() { if (!sess) { sess = std::make_unique(config, dev); - if (max_bufsize > sess->max_vmem) { - max_bufsize = sess->max_vmem; - } } return sess.get(); } @@ -2076,11 +2079,6 @@ static const char * ggml_backend_hexagon_buffer_type_name(ggml_backend_buffer_ty static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto dev_ctx = static_cast(buffer_type->context)->dev_ctx; - if (size > dev_ctx->max_bufsize) { - GGML_LOG_ERROR("ggml-hex: %s buffer size %zu exceeds max_bufsize %zu\n", - dev_ctx->c_name(), size, dev_ctx->max_bufsize); - return nullptr; - } auto sess = dev_ctx->session(); if (sess && sess->max_vmem && size > sess->max_vmem) { GGML_LOG_ERROR("ggml-hex: %s buffer size %zu exceeds max_vmem %zu\n", @@ -2099,11 +2097,6 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto dev_ctx = static_cast(buffer_type->context)->dev_ctx; - if (size > dev_ctx->max_bufsize) { - GGML_LOG_ERROR("ggml-hex: %s host buffer size %zu exceeds max_bufsize %zu\n", - dev_ctx->c_name(), size, dev_ctx->max_bufsize); - return nullptr; - } auto sess = dev_ctx->session(); if (sess && sess->max_vmem && size > sess->max_vmem) { GGML_LOG_ERROR("ggml-hex: %s host buffer size %zu exceeds max_vmem %zu\n", @@ -2138,10 +2131,8 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe } static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { - auto * context = static_cast(buft->context); - auto dev_ctx = context->dev_ctx; - dev_ctx->session(); - return dev_ctx->max_bufsize; + return opt_mbuf; + GGML_UNUSED(buft); } static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) { @@ -2173,7 +2164,7 @@ static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interfac }; ggml_backend_hexagon_device_context::ggml_backend_hexagon_device_context(int dev_id, const ggml_hexagon_device_config & config, ggml_backend_dev_t dev) - : dev_id(dev_id), config(config), dev(dev), max_bufsize(opt_mbuf) { + : dev_id(dev_id), config(config), dev(dev) { buffer_type.device = dev; buffer_type.iface = ggml_backend_hexagon_buffer_type_interface; buffer_type.context = new ggml_backend_hexagon_buffer_type_context(config.name, this); @@ -3927,7 +3918,6 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n this->valid_handle = true; // Query HW info and resolve session options - this->max_bufsize = opt_mbuf; { unsigned int hw_n_threads = 0; unsigned int hw_n_hvx = 0; @@ -4340,6 +4330,10 @@ static bool ggml_hexagon_supported_flash_attn_ext(const struct ggml_hexagon_sess } static bool ggml_hexagon_supported_gated_delta_net(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { + if (opt_gdn_select < 1) { + return false; + } + const struct ggml_tensor * q = op->src[0]; const struct ggml_tensor * k = op->src[1]; const struct ggml_tensor * v = op->src[2]; @@ -4387,10 +4381,26 @@ static bool ggml_hexagon_supported_gated_delta_net(const struct ggml_hexagon_ses const uint32_t total_rows = (uint32_t) (H * n_seqs); const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, total_rows); - struct htp_gdn_vtcm_layout layout; - htp_gdn_vtcm_layout_build(&layout, (uint32_t) S_v, n_threads ? n_threads : 1); - if (layout.total_bytes > sess->vtcm_size) { - return false; + + const bool can_use_hmx = (opt_gdn_select >= 2) && + (sess->n_hmx > 0) && + (S_v % 64 == 0) && + (n_tokens >= HTP_GDN_MIN_TOKENS) && + (g->ne[0] == 1) && + (K == 1); + + if (can_use_hmx) { + struct htp_gdn_hmx_vtcm_layout layout; + uint32_t n_heads_batch = 0; + if (!htp_gdn_hmx_solve_layout(&layout, (uint32_t) S_v, HTP_GDN_CHUNK_SIZE, total_rows, sess->vtcm_size, n_threads, true, &n_heads_batch)) { + return false; + } + } else { + struct htp_gdn_vtcm_layout layout; + htp_gdn_vtcm_layout_build(&layout, (uint32_t) S_v, n_threads); + if (layout.total_bytes > sess->vtcm_size) { + return false; + } } return true; @@ -5206,10 +5216,37 @@ static void ggml_hexagon_precompute_gated_delta_net_params( const uint32_t total_rows = H * n_seqs; const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, total_rows); - struct htp_gdn_vtcm_layout layout; - htp_gdn_vtcm_layout_build(&layout, S_v, n_threads ? n_threads : 1); + const bool can_use_hmx = (opt_gdn_select >= 2) && + (sess->n_hmx > 0) && + (S_v % 64 == 0) && + (n_tokens >= HTP_GDN_MIN_TOKENS) && + (g->ne[0] == 1) && + (K == 1); - kparams->n_threads = n_threads ? n_threads : 1; + struct htp_gdn_hmx_vtcm_layout hmx_layout; + struct htp_gdn_vtcm_layout hvx_layout; + uint32_t n_heads_batch = 1; + + if (can_use_hmx && htp_gdn_hmx_solve_layout(&hmx_layout, S_v, HTP_GDN_CHUNK_SIZE, total_rows, sess->vtcm_size, n_threads, true, &n_heads_batch)) { + kparams->kernel_type = HTP_GDN_KERNEL_HMX_CHUNKED; + kparams->pipeline = hmx_layout.pipeline ? 1 : 0; + kparams->chunk_size = HTP_GDN_CHUNK_SIZE; + kparams->n_chunks = (n_tokens + HTP_GDN_CHUNK_SIZE - 1) / HTP_GDN_CHUNK_SIZE; + kparams->n_heads_batch = (uint16_t) n_heads_batch; + kparams->vtcm_size = (uint32_t) hmx_layout.total_bytes; + kparams->state_aligned = (uint32_t) hmx_layout.state_f32_bytes; + kparams->vtcm_per_thread = (uint32_t) (hmx_layout.total_bytes / (n_threads > 0 ? n_threads : 1)); + } else { + htp_gdn_vtcm_layout_build(&hvx_layout, S_v, n_threads); + kparams->kernel_type = HTP_GDN_KERNEL_HVX_RECURRENT; + kparams->pipeline = 0; + kparams->n_heads_batch = 1; + kparams->state_aligned = (uint32_t) hvx_layout.state_aligned; + kparams->vtcm_per_thread = (uint32_t) hvx_layout.bytes_per_thread; + kparams->vtcm_size = (uint32_t) hvx_layout.total_bytes; + } + + kparams->n_threads = n_threads; kparams->S_v = S_v; kparams->H = H; kparams->n_tokens = n_tokens; @@ -5218,9 +5255,6 @@ static void ggml_hexagon_precompute_gated_delta_net_params( kparams->total_rows = total_rows; kparams->rows_per_thread = (total_rows + kparams->n_threads - 1) / kparams->n_threads; kparams->kda = (g->ne[0] == S_v) ? 1 : 0; - kparams->state_aligned = (uint32_t) layout.state_aligned; - kparams->vtcm_per_thread = (uint32_t) layout.bytes_per_thread; - kparams->vtcm_size = (uint32_t) layout.total_bytes; kparams->state_seq_stride = (uint32_t) (state->nb[3] / sizeof(float)); kparams->state_size_per_snap = S_v * S_v * H * n_seqs; kparams->scale = 1.0f / sqrtf((float) S_v); @@ -7731,6 +7765,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { const char * str_nhmx = getenv("GGML_HEXAGON_NHMX"); const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT"); const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT"); + const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT"); const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT"); const char * str_ndev = getenv("GGML_HEXAGON_NDEV"); const char * str_arch = getenv("GGML_HEXAGON_ARCH"); @@ -7783,6 +7818,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { opt_nhmx = str_nhmx ? atoi(str_nhmx) : opt_nhmx; opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select; opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select; + opt_gdn_select = str_gdn_select ? atoi(str_gdn_select) : opt_gdn_select; opt_ar_select = str_ar_select ? atoi(str_ar_select) : opt_ar_select; opt_mbuf = str_mbuf ? strtoul(str_mbuf, NULL, 0) * MiB : opt_mbuf; opt_vmem = str_vmem ? strtoul(str_vmem, NULL, 0) * MiB : opt_vmem; diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index 0716a8d21..803aa3f5a 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -358,7 +358,9 @@ struct htp_opformat { snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_id, (int) kparams->n_threads, (int) kparams->vtcm_size); } else if (node.opcode == HTP_OP_GATED_DELTA_NET) { const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params; - snprintf(str, max_size, "%s vtcm %u", + const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent"; + snprintf(str, max_size, "%s-%s vtcm %u", + path, kparams->kda ? "kda" : "scalar", (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads)); } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID || diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 988886082..bfcf7cb0c 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -55,6 +55,7 @@ struct htp_fa_context { float scale; float max_bias; + bool has_softcap; __fp16 logit_softcap; uint32_t n_head_log2; @@ -103,6 +104,7 @@ struct hmx_fa_context { // Op parameters __fp16 scale; float max_bias; + bool has_softcap; __fp16 logit_softcap; uint32_t n_head_log2; float m0, m1; @@ -234,7 +236,10 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * dma_cache m_cache; dma_cache_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE); - for (uint32_t ir = ir0; ir < ir1; ++ir) { + const size_t size_vkq_acc_single = hex_round_up(DV * sizeof(float), 128); + + uint32_t ir = ir0; + while (ir < ir1) { const uint32_t iq3 = fastdiv(ir, &factx->src0_div21); const uint32_t iq2 = fastdiv(ir - iq3*neq2*neq1, &factx->src0_div1); const uint32_t iq1 = (ir - iq3*neq2*neq1 - iq2 * neq1); @@ -245,6 +250,59 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t iv3 = fastdiv(iq3, &factx->broadcast_rv3); const uint32_t iv2 = fastdiv(iq2, &factx->broadcast_rv2); + uint32_t G_local = 1; + if (neq1 == 1 && (mask == NULL || mask->ne[2] == 1)) { + while (ir + G_local < ir1 && G_local < FA_HVX_G_MAX) { + const uint32_t next_ir = ir + G_local; + const uint32_t next_iq3 = fastdiv(next_ir, &factx->src0_div21); + const uint32_t next_iq2 = fastdiv(next_ir - next_iq3*neq2*neq1, &factx->src0_div1); + const uint32_t next_iq1 = (next_ir - next_iq3*neq2*neq1 - next_iq2 * neq1); + + const uint32_t next_ik3 = fastdiv(next_iq3, &factx->broadcast_rk3); + const uint32_t next_ik2 = fastdiv(next_iq2, &factx->broadcast_rk2); + + const uint32_t next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3); + const uint32_t next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2); + + if (next_ik2 != ik2 || next_ik3 != ik3 || next_iv2 != iv2 || next_iv3 != iv3 || next_iq1 != iq1 || next_iq3 != iq3) { + break; + } + G_local++; + } + } + + uint32_t heads[FA_HVX_G_MAX]; + HVX_Vector slope_vecs[FA_HVX_G_MAX] __attribute__((aligned(128))); + HVX_Vector S_vec[FA_HVX_G_MAX] __attribute__((aligned(128))); + HVX_Vector M_vec[FA_HVX_G_MAX] __attribute__((aligned(128))); + uint8_t * q_ptrs[FA_HVX_G_MAX]; + float * vkq_ptrs[FA_HVX_G_MAX]; + + for (uint32_t g = 0; g < G_local; ++g) { + const uint32_t r = ir + g; + const uint32_t r_iq3 = fastdiv(r, &factx->src0_div21); + const uint32_t r_iq2 = fastdiv(r - r_iq3*neq2*neq1, &factx->src0_div1); + const uint32_t r_iq1 = (r - r_iq3*neq2*neq1 - r_iq2 * neq1); + + heads[g] = r_iq2; + const __fp16 slope = factx->slopes[r_iq2]; + slope_vecs[g] = hvx_vec_splat_f16(slope); + + S_vec[g] = hvx_vec_splat_f32(0.0f); + M_vec[g] = hvx_vec_splat_f32(HTP_FA_M_INITIAL_VAL); + + uint8_t * q_dst = spad_q + g * factx->size_q_row_padded; + q_ptrs[g] = q_dst; + + float * vkq_dst = (float *)(spad_a + g * size_vkq_acc_single); + vkq_ptrs[g] = vkq_dst; + hvx_splat_f32_a((uint8_t *) vkq_dst, 0, DV); + + // Fetch Q row g + const dma_addr_t q_row_ptr = q->data + r_iq1*nbq1 + r_iq2*nbq2 + r_iq3*nbq3; + dma_queue_push(dma_q, dma_make_data(q_dst, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); + } + dma_addr_t mp_base = 0; if (mask) { const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2); @@ -252,116 +310,44 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * mp_base = mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3]; } - // Precalculate next row variables if there is a next row - bool has_next_ir = (ir + 1 < ir1); - uint32_t next_ik2 = 0, next_ik3 = 0, next_iv2 = 0, next_iv3 = 0; - dma_addr_t next_q_row_ptr = 0; - dma_addr_t next_mp_base = 0; + // Prefetch first two blocks + for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) { + const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE; + const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); - dma_addr_t next_k_src0 = 0; - dma_addr_t next_v_src0 = 0; - dma_addr_t next_m_src0 = 0; - uint32_t next_block_size0 = 0; + // K + const dma_addr_t k_src = k->data + ic_start*nbk1 + ik2*nbk2 + ik3*nbk3; + uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block; + dma_queue_push(dma_q, dma_make_data(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size); - dma_addr_t next_k_src1 = 0; - dma_addr_t next_v_src1 = 0; - dma_addr_t next_m_src1 = 0; - uint32_t next_block_size1 = 0; - - if (has_next_ir) { - const uint32_t next_ir = ir + 1; - const uint32_t next_iq3 = fastdiv(next_ir, &factx->src0_div21); - const uint32_t next_iq2 = fastdiv(next_ir - next_iq3*neq2*neq1, &factx->src0_div1); - const uint32_t next_iq1 = (next_ir - next_iq3*neq2*neq1 - next_iq2 * neq1); - - next_ik3 = fastdiv(next_iq3, &factx->broadcast_rk3); - next_ik2 = fastdiv(next_iq2, &factx->broadcast_rk2); - - next_iv3 = fastdiv(next_iq3, &factx->broadcast_rv3); - next_iv2 = fastdiv(next_iq2, &factx->broadcast_rv2); - - next_q_row_ptr = q->data + next_iq1*nbq1 + next_iq2*nbq2 + next_iq3*nbq3; + // V + const dma_addr_t v_src = v->data + ic_start*nbv1 + iv2*nbv2 + iv3*nbv3; + uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block; + dma_queue_push(dma_q, dma_make_data(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size); + // Mask if (mask) { - const uint32_t next_im2 = fastmodulo(next_iq2, mask->ne[2], &factx->src3_div2); - const uint32_t next_im3 = fastmodulo(next_iq3, mask->ne[3], &factx->src3_div3); - next_mp_base = mask->data + next_iq1*mask->nb[1] + next_im2*mask->nb[2] + next_im3*mask->nb[3]; - } - - // Precalculate next K/V block 0 source pointers - { - const uint32_t ic_start = 0; - next_block_size0 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); - next_k_src0 = k->data + ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3; - next_v_src0 = v->data + ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3; - if (mask) { - next_m_src0 = next_mp_base + ic_start * sizeof(__fp16); - } - } - - // Precalculate next K/V block 1 source pointers (if n_blocks > 1) - if (factx->n_blocks > 1) { - const uint32_t ic_start = 1 * FLASH_ATTN_BLOCK_SIZE; - next_block_size1 = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); - next_k_src1 = k->data + ic_start*nbk1 + next_ik2*nbk2 + next_ik3*nbk3; - next_v_src1 = v->data + ic_start*nbv1 + next_iv2*nbv2 + next_iv3*nbv3; - if (mask) { - next_m_src1 = next_mp_base + ic_start * sizeof(__fp16); - } + const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16); + dma_cache_push(dma_q, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1); } } - if (ir == ir0) { - // Fetch Q row - const dma_addr_t q_row_ptr = q->data + iq1*nbq1 + iq2*nbq2 + iq3*nbq3; - dma_queue_push(dma_q, dma_make_data(spad_q, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); - - // Prefetch first two blocks - for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) { - const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE; - const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); - - // K - const dma_addr_t k_src = k->data + ic_start*nbk1 + ik2*nbk2 + ik3*nbk3; - uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block; - dma_queue_push(dma_q, dma_make_data(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size); - - // V - const dma_addr_t v_src = v->data + ic_start*nbv1 + iv2*nbv2 + iv3*nbv3; - uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block; - dma_queue_push(dma_q, dma_make_data(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size); - - // Mask - if (mask) { - const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16); - // Mask is 1D contiguous for this row - dma_cache_push(dma_q, &m_cache, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1); - } + // Pop all Q rows + for (uint32_t g = 0; g < G_local; ++g) { + uint8_t * q_ptr_vtcm = (void *) dma_queue_pop(dma_q).dst; + if (factx->is_q_fp32) { + hvx_copy_f16_f32_aa(q_ptr_vtcm, q_ptr_vtcm, DK); } } - const uint32_t h = iq2; // head index - const __fp16 slope = factx->slopes[h]; - - HVX_Vector S_vec = hvx_vec_splat_f32(0.0f); - HVX_Vector M_vec = hvx_vec_splat_f32(HTP_FA_M_INITIAL_VAL); - - // Clear accumulator - hvx_splat_f32_a(spad_a, 0, DV); - float * VKQ32 = (float *) (spad_a + 0); - - uint8_t * q_ptr_vtcm = (void *) dma_queue_pop(dma_q).dst; - if (factx->is_q_fp32) { - hvx_copy_f16_f32_aa(q_ptr_vtcm, q_ptr_vtcm, DK); // inplace convert f32 to f16 - } - - const HVX_Vector slope_vec = hvx_vec_splat_f16(slope); const HVX_Vector v_neg_inf = Q6_Vh_vsplat_R(0xfbff); - const HVX_Vector v_cap = (factx->logit_softcap != 0.0f) ? hvx_vec_splat_f16(factx->logit_softcap) : Q6_V_vzero(); + const bool has_softcap = factx->has_softcap; + const HVX_Vector v_cap = has_softcap ? hvx_vec_splat_f16(factx->logit_softcap) : Q6_V_vzero(); const HVX_Vector vinf = Q6_Vh_vsplat_R(0xFC00); const HVX_Vector vmin = Q6_Vh_vsplat_R(0xFBFF); const HVX_Vector v_log2e = hvx_vec_splat_f16(EXP_LOG2E_F); const uint32_t stride_v2 = factx->size_v_row_padded * 2; + for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) { const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE; const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start); @@ -388,235 +374,222 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); } - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir); + for (uint32_t g = 0; g < G_local; ++g) { + const uint32_t head_ir = ir + g; + uint8_t * q_ptr_vtcm = q_ptrs[g]; + float * VKQ32 = vkq_ptrs[g]; + const HVX_Vector slope_vec = slope_vecs[g]; - // Inner loop processing the block from VTCM - // 1. Compute scores (64 elements FP16) - HVX_Vector scores_f16 = Q6_V_vzero(); - if (current_block_size > 0) { - HVX_Vector scores0 = hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base, factx->size_k_row_padded, DK, factx->scale); - HVX_Vector scores1 = (current_block_size > 32) ? hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base + 32 * factx->size_k_row_padded, factx->size_k_row_padded, DK, factx->scale) : Q6_V_vzero(); - scores_f16 = hvx_vec_f32_to_f16(scores0, scores1); - } + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, head_ir); - // 2. Softcap (in FP16) - if (factx->logit_softcap != 0.0f) { - scores_f16 = hvx_vec_tanh_f16(scores_f16); - scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap); - } - - HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(current_block_size * sizeof(__fp16)); - - // 3. Mask (in FP16) - if (mask) { - HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base; - HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf); - m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16); - - HVX_Vector m_scaled = hvx_vec_mul_f16_f16(m_vals_f16, slope_vec); - scores_f16 = Q6_V_vmux_QVV(q_tail_keep, hvx_vec_add_f16_f16(scores_f16, m_scaled), v_neg_inf); - } else { - scores_f16 = Q6_V_vmux_QVV(q_tail_keep, scores_f16, v_neg_inf); - } - - // Compute block max in FP16 - HVX_Vector v_max_f16 = hvx_vec_reduce_max_f16(scores_f16); - HVX_Vector v_max = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); // splat block max in FP32 - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, ir); - - if (ib + 1 == factx->n_blocks && has_next_ir) { - // Queue next row's Q row! - dma_queue_push(dma_q, dma_make_data(spad_q, next_q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1); - - if (factx->n_blocks % 2 == 0) { - // Queue next row's block 0 (into buffer slot 0) - uint8_t * k_dst = spad_k + 0 * factx->size_k_block; - uint8_t * v_dst = spad_v + 0 * factx->size_v_block; - - // K (block 0 of next row) - dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0); - - // V (block 0 of next row) - dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0); - - // Mask (block 0 of next row) - if (mask) { - dma_cache_push(dma_q, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1); - } + HVX_Vector scores_f16 = Q6_V_vzero(); + if (current_block_size > 0) { + HVX_Vector scores0 = hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base, factx->size_k_row_padded, DK, factx->scale); + HVX_Vector scores1 = (current_block_size > 32) ? hvx_dot_f16_f16_aa_rx32(q_ptr_vtcm, k_base + 32 * factx->size_k_row_padded, factx->size_k_row_padded, DK, factx->scale) : Q6_V_vzero(); + scores_f16 = hvx_vec_f32_to_f16(scores0, scores1); } - } - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir); - { - // 4. Online Softmax Update - HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec); - HVX_Vector diff_vec = HVX_OP_SUB_F32(M_vec, M_new_vec); - - HVX_Vector diff_f16 = hvx_vec_f32_to_f16(diff_vec, diff_vec); - HVX_Vector diff_base2 = hvx_vec_mul_f16_f16(diff_f16, v_log2e); - HVX_Vector ms_f16 = hvx_vec_exp2_f16(diff_base2); - HVX_Vector ms_vec = Q6_V_lo_W(hvx_vec_f16_to_f32(ms_f16)); - - M_vec = M_new_vec; - - hvx_scale_vec_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms_vec); - - // Compute P = exp2((S - M) * log2(e)) in FP16 - HVX_Vector v_m_vec_f16 = hvx_vec_f32_to_f16(M_vec, M_vec); - HVX_Vector v_s_minus_m = Q6_Vqf16_vsub_VhfVhf(scores_f16, v_m_vec_f16); - - HVX_Vector v_s_minus_m_base2 = hvx_vec_mul_f16_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m), v_log2e); - - HVX_Vector P = hvx_vec_exp2_f16(v_s_minus_m_base2); - P = Q6_V_vmux_QVV(q_tail_keep, P, Q6_V_vzero()); - - // Convert P to FP32 to update the running sum S_vec - HVX_VectorPair P_pair = hvx_vec_f16_to_f32(P); - HVX_Vector P0 = Q6_V_lo_W(P_pair); - HVX_Vector P1 = Q6_V_hi_W(P_pair); - HVX_Vector p_sum_vec = hvx_vec_reduce_sum_f32(HVX_OP_ADD_F32(P0, P1)); - - S_vec = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec, ms_vec), p_sum_vec); - - // 5. Accumulate V (F16 * F16 -> F32 accumulator) - const uint8_t * v_ptr = v_base; - - for (uint32_t j = 0; j < current_block_size; j += 2) { - if (j + 1 == current_block_size) { - HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2)); - hvx_mad_f32_f16_aa_vec(VKQ32, v_ptr, S0, DV); - break; - } - - HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2)); - HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2)); - - hvx_mad_f32_f16_aa_rx2_vec(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, S0, S1, DV); - v_ptr += stride_v2; + if (has_softcap) { + scores_f16 = hvx_vec_tanh_f16(scores_f16); + scores_f16 = hvx_vec_mul_f16_f16(scores_f16, v_cap); } - } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, ir); - // Issue DMA for next+1 block (if exists) - if (ib + 2 < factx->n_blocks) { - const uint32_t next_ib = ib + 2; - const uint32_t next_ic_start = next_ib * FLASH_ATTN_BLOCK_SIZE; - const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start); + HVX_VectorPred q_tail_keep = Q6_Q_vsetq2_R(current_block_size * sizeof(__fp16)); - // K - const dma_addr_t k_src = k->data + next_ic_start*nbk1 + ik2*nbk2 + ik3*nbk3; - dma_queue_push(dma_q, dma_make_data(k_base, k_src), factx->size_k_row_padded, nbk1, size_k_row, next_block_size); - - // V - const dma_addr_t v_src = v->data + next_ic_start*nbv1 + iv2*nbv2 + iv3*nbv3; - dma_queue_push(dma_q, dma_make_data(v_base, v_src), factx->size_v_row_padded, nbv1, size_v_row, next_block_size); - - // Mask if (mask) { - const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16); - dma_cache_push(dma_q, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1); + HVX_Vector m_vals_f16 = *(const HVX_UVector *) m_base; + HVX_VectorPred is_inf = Q6_Q_vcmp_eq_VhVh(m_vals_f16, vinf); + m_vals_f16 = Q6_V_vmux_QVV(is_inf, vmin, m_vals_f16); + + HVX_Vector m_scaled = hvx_vec_mul_f16_f16(m_vals_f16, slope_vec); + scores_f16 = Q6_V_vmux_QVV(q_tail_keep, hvx_vec_add_f16_f16(scores_f16, m_scaled), v_neg_inf); + } else { + scores_f16 = Q6_V_vmux_QVV(q_tail_keep, scores_f16, v_neg_inf); } - } - } - if (has_next_ir) { - if (factx->n_blocks % 2 == 0) { - // Queue next row's block 1 (into buffer slot 1, if n_blocks > 1) - if (factx->n_blocks > 1) { - uint8_t * k_dst = spad_k + 1 * factx->size_k_block; - uint8_t * v_dst = spad_v + 1 * factx->size_v_block; + HVX_Vector v_max_f16 = hvx_vec_reduce_max_f16(scores_f16); + HVX_Vector v_max = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16)); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, head_ir); - // K (block 1 of next row) - dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1); + // prefetch K for block ib + 2 after last head finished QK + if (g + 1 == G_local && ib + 2 < factx->n_blocks) { + const uint32_t next_ib = ib + 2; + const uint32_t next_ic_start = next_ib * FLASH_ATTN_BLOCK_SIZE; + const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start); - // V (block 1 of next row) - dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1); - - // Mask (block 1 of next row) - if (mask) { - dma_cache_push(dma_q, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1); - } + const dma_addr_t k_src = k->data + next_ic_start*nbk1 + ik2*nbk2 + ik3*nbk3; + dma_queue_push(dma_q, dma_make_data(k_base, k_src), factx->size_k_row_padded, nbk1, size_k_row, next_block_size); } - } else { - // Queue next row's block 0 (into buffer slot 0) + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, head_ir); { - uint8_t * k_dst = spad_k + 0 * factx->size_k_block; - uint8_t * v_dst = spad_v + 0 * factx->size_v_block; + HVX_Vector M_new_vec = Q6_Vsf_vmax_VsfVsf(v_max, M_vec[g]); + HVX_Vector diff_vec = HVX_OP_SUB_F32(M_vec[g], M_new_vec); - // K (block 0 of next row) - dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src0), factx->size_k_row_padded, nbk1, size_k_row, next_block_size0); + HVX_Vector diff_f16 = hvx_vec_f32_to_f16(diff_vec, diff_vec); + HVX_Vector diff_base2 = hvx_vec_mul_f16_f16(diff_f16, v_log2e); + HVX_Vector ms_f16 = hvx_vec_exp2_f16(diff_base2); + HVX_Vector ms_vec = Q6_V_lo_W(hvx_vec_f16_to_f32(ms_f16)); - // V (block 0 of next row) - dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src0), factx->size_v_row_padded, nbv1, size_v_row, next_block_size0); + M_vec[g] = M_new_vec; - // Mask (block 0 of next row) - if (mask) { - dma_cache_push(dma_q, &m_cache, next_m_src0, next_block_size0 * 2, next_block_size0 * 2, next_block_size0 * 2, 1); + HVX_Vector v_m_vec_f16 = hvx_vec_f32_to_f16(M_vec[g], M_vec[g]); + HVX_Vector v_s_minus_m = Q6_Vqf16_vsub_VhfVhf(scores_f16, v_m_vec_f16); + HVX_Vector v_s_minus_m_base2 = hvx_vec_mul_f16_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m), v_log2e); + + HVX_Vector P = hvx_vec_exp2_f16(v_s_minus_m_base2); + P = Q6_V_vmux_QVV(q_tail_keep, P, Q6_V_vzero()); + + HVX_VectorPair P_pair = hvx_vec_f16_to_f32(P); + HVX_Vector P0 = Q6_V_lo_W(P_pair); + HVX_Vector P1 = Q6_V_hi_W(P_pair); + HVX_Vector p_sum_vec = hvx_vec_reduce_sum_f32(HVX_OP_ADD_F32(P0, P1)); + + S_vec[g] = HVX_OP_ADD_F32(HVX_OP_MUL_F32(S_vec[g], ms_vec), p_sum_vec); + + const uint8_t * v_ptr = v_base; + + if (DV == 64) { + HVX_VectorPair vkq0 = *((const HVX_VectorPair *) VKQ32); + vkq0 = Q6_W_vcombine_VV( + HVX_OP_MUL_F32(Q6_V_hi_W(vkq0), ms_vec), + HVX_OP_MUL_F32(Q6_V_lo_W(vkq0), ms_vec) + ); + + for (uint32_t j = 0; j < current_block_size; j += 2) { + HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2)); + const HVX_Vector * vx0 = (const HVX_Vector *) v_ptr; + if (j + 1 == current_block_size) { + vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx0[0]), S0); + break; + } + + HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2)); + const HVX_Vector * vx1 = (const HVX_Vector *) (v_ptr + factx->size_v_row_padded); + vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx0[0]), S0); + vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx1[0]), S1); + v_ptr += stride_v2; + } + + *((HVX_VectorPair *) VKQ32) = vkq0; + } else if (DV == 128) { + HVX_VectorPair vkq0 = ((const HVX_VectorPair *) VKQ32)[0]; + HVX_VectorPair vkq1 = ((const HVX_VectorPair *) VKQ32)[1]; + vkq0 = Q6_W_vcombine_VV( + HVX_OP_MUL_F32(Q6_V_hi_W(vkq0), ms_vec), + HVX_OP_MUL_F32(Q6_V_lo_W(vkq0), ms_vec) + ); + vkq1 = Q6_W_vcombine_VV( + HVX_OP_MUL_F32(Q6_V_hi_W(vkq1), ms_vec), + HVX_OP_MUL_F32(Q6_V_lo_W(vkq1), ms_vec) + ); + + for (uint32_t j = 0; j < current_block_size; j += 2) { + HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2)); + const HVX_Vector * vx0 = (const HVX_Vector *) v_ptr; + if (j + 1 == current_block_size) { + vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx0[0]), S0); + vkq1 = hvx_vec_mpyacc_f32_f16(vkq1, Q6_Vh_vshuff_Vh(vx0[1]), S0); + break; + } + + HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2)); + const HVX_Vector * vx1 = (const HVX_Vector *) (v_ptr + factx->size_v_row_padded); + vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx0[0]), S0); + vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx1[0]), S1); + vkq1 = hvx_vec_mpyacc_f32_f16(vkq1, Q6_Vh_vshuff_Vh(vx0[1]), S0); + vkq1 = hvx_vec_mpyacc_f32_f16(vkq1, Q6_Vh_vshuff_Vh(vx1[1]), S1); + v_ptr += stride_v2; + } + + ((HVX_VectorPair *) VKQ32)[0] = vkq0; + ((HVX_VectorPair *) VKQ32)[1] = vkq1; + } else { + hvx_scale_vec_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms_vec); + + for (uint32_t j = 0; j < current_block_size; j += 2) { + if (j + 1 == current_block_size) { + HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2)); + hvx_mad_f32_f16_aa_vec(VKQ32, v_ptr, S0, DV); + break; + } + + HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2)); + HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2)); + + hvx_mad_f32_f16_aa_rx2_vec(VKQ32, v_ptr, v_ptr + factx->size_v_row_padded, S0, S1, DV); + v_ptr += stride_v2; + } } } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, head_ir); - // Queue next row's block 1 (into buffer slot 1, if n_blocks > 1) - if (factx->n_blocks > 1) { - uint8_t * k_dst = spad_k + 1 * factx->size_k_block; - uint8_t * v_dst = spad_v + 1 * factx->size_v_block; + // prefetch V and mask for block ib + 2 after last head finished V accumulation + if (g + 1 == G_local && ib + 2 < factx->n_blocks) { + const uint32_t next_ib = ib + 2; + const uint32_t next_ic_start = next_ib * FLASH_ATTN_BLOCK_SIZE; + const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start); - // K (block 1 of next row) - dma_queue_push(dma_q, dma_make_data(k_dst, next_k_src1), factx->size_k_row_padded, nbk1, size_k_row, next_block_size1); + // V + const dma_addr_t v_src = v->data + next_ic_start*nbv1 + iv2*nbv2 + iv3*nbv3; + dma_queue_push(dma_q, dma_make_data(v_base, v_src), factx->size_v_row_padded, nbv1, size_v_row, next_block_size); - // V (block 1 of next row) - dma_queue_push(dma_q, dma_make_data(v_dst, next_v_src1), factx->size_v_row_padded, nbv1, size_v_row, next_block_size1); - - // Mask (block 1 of next row) + // Mask if (mask) { - dma_cache_push(dma_q, &m_cache, next_m_src1, next_block_size1 * 2, next_block_size1 * 2, next_block_size1 * 2, 1); + const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16); + dma_cache_push(dma_q, &m_cache, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1); } } + } // end for g + } // end for ib + + for (uint32_t g = 0; g < G_local; ++g) { + const uint32_t head_ir = ir + g; + const uint32_t h = heads[g]; + float * VKQ32 = vkq_ptrs[g]; + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, head_ir); + + float M = hvx_vec_get_f32(M_vec[g]); + float S = hvx_vec_get_f32(S_vec[g]); + + if (sinks) { + const float s = factx->spad_sinks[h]; + + float vs = 1.0f; + + if (s > M) { + HVX_Vector diff_vec = hvx_vec_splat_f32(M - s); + HVX_Vector ms_vec = hvx_vec_exp_f32(diff_vec); + hvx_scale_vec_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms_vec); + + float ms = hvx_vec_get_f32(ms_vec); + S = S * ms + vs; + } else { + HVX_Vector diff_vec = hvx_vec_splat_f32(s - M); + vs = hvx_vec_get_f32(hvx_vec_exp_f32(diff_vec)); + S += vs; + } } - } - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, ir); - // sinks - float M = hvx_vec_get_f32(M_vec); - float S = hvx_vec_get_f32(S_vec); + const float S_inv = S == 0.0f ? 0.0f : 1.0f/S; + hvx_scale_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, S_inv); - if (sinks) { - const float s = factx->spad_sinks[h]; + const uint32_t r_iq3 = fastdiv(head_ir, &factx->src0_div21); + const uint32_t r_iq2 = fastdiv(head_ir - r_iq3*neq2*neq1, &factx->src0_div1); + const uint32_t r_iq1 = (head_ir - r_iq3*neq2*neq1 - r_iq2 * neq1); - float vs = 1.0f; + uint8_t * dst_ptr = (uint8_t *) dst->data + r_iq2 * dst->nb[1] + r_iq1 * dst->nb[2] + r_iq3 * dst->nb[3]; - if (s > M) { - HVX_Vector diff_vec = hvx_vec_splat_f32(M - s); - HVX_Vector ms_vec = hvx_vec_exp_f32(diff_vec); - hvx_scale_vec_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms_vec); - - float ms = hvx_vec_get_f32(ms_vec); - S = S * ms + vs; - } else { - HVX_Vector diff_vec = hvx_vec_splat_f32(s - M); - vs = hvx_vec_get_f32(hvx_vec_exp_f32(diff_vec)); - S += vs; + if (dst->type == HTP_TYPE_F32) { + hvx_copy_f32_ua(dst_ptr, (uint8_t *) VKQ32, DV); + } else if (dst->type == HTP_TYPE_F16) { + hvx_copy_f16_f32_ua(dst_ptr, (uint8_t *) VKQ32, DV); } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, head_ir); } - const float S_inv = S == 0.0f ? 0.0f : 1.0f/S; - hvx_scale_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, S_inv); - - // Store result - // dst indices - const uint32_t i1 = iq1; - const uint32_t i2 = iq2; - const uint32_t i3 = iq3; - - // dst is permuted: [DV, n_heads, n_tokens, n_seq] - // head stride is nb[1], token stride is nb[2], batch stride is nb[3] - uint8_t * dst_ptr = (uint8_t *) dst->data + i2 * dst->nb[1] + i1 * dst->nb[2] + i3 * dst->nb[3]; - - if (dst->type == HTP_TYPE_F32) { - hvx_copy_f32_ua(dst_ptr, (uint8_t *) VKQ32, DV); - } else if (dst->type == HTP_TYPE_F16) { - hvx_copy_f16_f32_ua(dst_ptr, (uint8_t *) VKQ32, DV); - } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, ir); + ir += G_local; } } @@ -1554,7 +1527,7 @@ static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) { const bool mask_broadcast = factx->mask_broadcast; const bool is_g1 = (args->G == 1); const bool has_alibi = args->has_alibi; - const bool has_softcap = (factx->logit_softcap != 0.0f); + const bool has_softcap = factx->has_softcap; fa_softmax_impl(n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap); } @@ -1589,9 +1562,9 @@ static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx, const size_t n_row_vec_cnt = hmx_ceil_div(sargs->n_rows_g, 64); worker_callback_t softmax_fn = fa_softmax_thread; - if (sargs->mask == NULL && factx->logit_softcap == 0.0f && !sargs->has_alibi) { + if (sargs->mask == NULL && !factx->has_softcap && !sargs->has_alibi) { softmax_fn = fa_softmax_thread_nomask; - } else if (sargs->mask != NULL && factx->mask_broadcast && factx->logit_softcap == 0.0f && !sargs->has_alibi) { + } else if (sargs->mask != NULL && factx->mask_broadcast && !factx->has_softcap && !sargs->has_alibi) { if (sargs->G == 1) { softmax_fn = fa_softmax_thread_mask_broadcast_g1; } else { @@ -1905,13 +1878,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { factx.src3_div3 = kparams->src3_div3; } - if (kparams->logit_softcap == 0.0f) { + factx.has_softcap = (kparams->logit_softcap != 0.0f); + if (!factx.has_softcap) { factx.scale = (__fp16) (kparams->scale * EXP_LOG2E_F); // log2(e) } else { factx.scale = (__fp16) kparams->scale; } factx.max_bias = kparams->max_bias; - factx.logit_softcap = (__fp16) (kparams->logit_softcap * EXP_LOG2E_F); + factx.logit_softcap = factx.has_softcap ? (__fp16) (kparams->logit_softcap * EXP_LOG2E_F) : 0; factx.n_head_log2 = kparams->n_head_log2; factx.m0 = kparams->m0; @@ -2513,7 +2487,8 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { factx.scale = kparams->scale; factx.max_bias = kparams->max_bias; - factx.logit_softcap = (__fp16) kparams->logit_softcap; + factx.has_softcap = (kparams->logit_softcap != 0.0f); + factx.logit_softcap = factx.has_softcap ? (__fp16) kparams->logit_softcap : 0; factx.n_head_log2 = kparams->n_head_log2; factx.m0 = kparams->m0; diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h index 2bd232190..22bb8c53d 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.h +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.h @@ -247,6 +247,7 @@ static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, siz } #define FA_HVX_BLOCK_SIZE 64 +#define FA_HVX_G_MAX 8 struct hvx_fa_vtcm_layout { size_t off_q; @@ -275,11 +276,11 @@ static inline void hvx_fa_vtcm_layout_build(struct hvx_fa_vtcm_layout * L, const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128); const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128); - const size_t size_q_block = size_q_row_padded * 1; + const size_t size_q_block = size_q_row_padded * FA_HVX_G_MAX; const size_t size_k_block = size_k_row_padded * FA_HVX_BLOCK_SIZE; const size_t size_v_block = size_v_row_padded * FA_HVX_BLOCK_SIZE; const size_t size_m_block = hex_round_up(FA_HVX_BLOCK_SIZE * sizeof(__fp16), 128); - const size_t size_vkq_acc = hex_round_up(DV * sizeof(float), 128); + const size_t size_vkq_acc = hex_round_up(DV * sizeof(float), 128) * FA_HVX_G_MAX; const size_t size_sinks = hex_round_up(n_heads * sizeof(float), 128); size_t off = 0; diff --git a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c index b37313370..1dd828db7 100644 --- a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c +++ b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c @@ -2,6 +2,7 @@ #include #include #include +#include #include "hvx-base.h" #include "hvx-copy.h" @@ -11,6 +12,10 @@ #include "ggml-common.h" #include "htp-ctx.h" #include "htp-tensor.h" +#include "htp-vtcm.h" +#include "hmx-utils.h" +#include "hmx-fa-kernels.h" +#include "hmx-queue.h" #include "gated-delta-net-ops.h" #ifndef MIN @@ -55,9 +60,8 @@ static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const HVX_Vector return hvx_vec_reduce_sum_f32(acc); } -static inline HVX_Vector gdn_mul_scalar_dot_f32(float * restrict dst, float mul, const HVX_Vector * restrict dot, uint32_t n) { +static inline HVX_Vector gdn_mul_scalar_dot_f32(float * restrict dst, HVX_Vector vmul, const HVX_Vector * restrict dot, uint32_t n) { HVX_Vector acc = Q6_V_vzero(); - const HVX_Vector vmul = hvx_vec_splat_f32(mul); const uint32_t epv = 128 / sizeof(float); const uint32_t nvec = n / epv; const uint32_t nloe = n % epv; @@ -589,20 +593,15 @@ static inline void gdn_step_kda_f32( HVX_Vector vk[4]; HVX_Vector vg[4]; - 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); - for (uint32_t i = 0; i < nvec; ++i) { vq[i] = hvx_vmemu(q_t + i * epv); vk[i] = hvx_vmemu(k_t + i * epv); - vg[i] = hvx_vec_exp_f32_guard(hvx_vmemu(g_t + i * epv), max_exp, inf); + vg[i] = hvx_vec_exp_f32(hvx_vmemu(g_t + i * epv)); } if (nloe) { vq[nvec] = hvx_vmemu(q_t + nvec * epv); vk[nvec] = hvx_vmemu(k_t + nvec * epv); - vg[nvec] = hvx_vec_exp_f32_guard(hvx_vmemu(g_t + nvec * epv), max_exp, inf); + vg[nvec] = hvx_vec_exp_f32(hvx_vmemu(g_t + nvec * epv)); } const HVX_Vector vbeta = hvx_vec_splat_f32(beta_val); @@ -690,9 +689,8 @@ static inline void gdn_step_scalar_f32( vk[nvec] = hvx_vmemu(k_t + nvec * epv); } - const float gate = expf(g_t[0]); - const HVX_Vector vgate = hvx_vec_splat_f32(gate); - const HVX_Vector vbeta = hvx_vec_splat_f32(beta_val); + const HVX_Vector vgate = hvx_vec_exp_f32(hvx_vec_splat_f32(g_t[0])); + const HVX_Vector vbeta = hvx_vec_splat_f32(beta_val); const HVX_Vector vscale = hvx_vec_splat_f32(scale); float delta[8] __attribute__((aligned(128))); @@ -742,7 +740,7 @@ static inline void gdn_step_scalar_f32( } for (; j < S_v; ++j) { float * row = s_work + (uint64_t) j * S_v; - HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, gate, vk, S_v); + HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, vgate, vk, S_v); HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]); HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), vbeta); HVX_Vector vres = gdn_add_scaled_dot_f32(row, vk, vdj, vq, S_v); @@ -1022,6 +1020,1255 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo dma_queue_flush(dma_q); } +struct htp_gdn_hmx_gemm_task { + const __fp16 * row_tiles; + const __fp16 * col_tiles; + __fp16 * out_tiles; + uint32_t n_row_tiles; + uint32_t n_col_tiles; + uint32_t n_dot_tiles; + uint32_t dot_stride; + uint8_t * hmx_scales; +}; + +static void htp_gdn_hmx_gemm_worker(void * data) { + struct htp_gdn_hmx_gemm_task * task = (struct htp_gdn_hmx_gemm_task *) data; + asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)task->hmx_scales)); + + const size_t dot_stride = task->dot_stride; + for (uint32_t r = 0; r < task->n_row_tiles; ++r) { + const __fp16 * r_tiles = task->row_tiles + r * dot_stride; + const __fp16 * c_tiles = task->col_tiles; + __fp16 * o_tile = task->out_tiles + r * task->n_col_tiles * HMX_FP16_TILE_N_ELMS; + + for (uint32_t c = 0; c < task->n_col_tiles; ++c) { + hmx_fa_qk_dot_tile(r_tiles, c_tiles, o_tile, task->n_dot_tiles); + c_tiles += dot_stride; + o_tile += HMX_FP16_TILE_N_ELMS; + } + } +} + +static inline void htp_gdn_push_hmx_gemm_task( + hmx_queue_t q, + struct htp_gdn_hmx_gemm_task * task, + const __fp16 * row_tiles, + const __fp16 * col_tiles, + __fp16 * out_tiles, + uint32_t n_row_tiles, + uint32_t n_col_tiles, + uint32_t n_dot_tiles, + uint8_t * scales +) { + task->row_tiles = row_tiles; + task->col_tiles = col_tiles; + task->out_tiles = out_tiles; + task->n_row_tiles = n_row_tiles; + task->n_col_tiles = n_col_tiles; + task->n_dot_tiles = n_dot_tiles; + task->dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS; + task->hmx_scales = scales; + + hmx_queue_push(q, hmx_queue_make_desc(htp_gdn_hmx_gemm_worker, task)); +} + +static inline void gdn_unpack_64x64_tiles_to_vectors( + HVX_Vector * restrict rows, + const __fp16 * restrict tiles +) { + const HVX_Vector * t00 = (const HVX_Vector *) (tiles + 0 * HMX_FP16_TILE_N_ELMS); + const HVX_Vector * t01 = (const HVX_Vector *) (tiles + 1 * HMX_FP16_TILE_N_ELMS); + const HVX_Vector * t10 = (const HVX_Vector *) (tiles + 2 * HMX_FP16_TILE_N_ELMS); + const HVX_Vector * t11 = (const HVX_Vector *) (tiles + 3 * HMX_FP16_TILE_N_ELMS); + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp0 = Q6_W_vdeal_VVR(t01[r], t00[r], -2); + rows[2 * r + 0] = Q6_V_lo_W(vp0); + rows[2 * r + 1] = Q6_V_hi_W(vp0); + + HVX_VectorPair vp1 = Q6_W_vdeal_VVR(t11[r], t10[r], -2); + rows[32 + 2 * r + 0] = Q6_V_lo_W(vp1); + rows[32 + 2 * r + 1] = Q6_V_hi_W(vp1); + } +} + +static inline void gdn_pack_64x64_vectors_to_tiles( + __fp16 * restrict tiles, + const HVX_Vector * restrict rows +) { + HVX_Vector * t00 = (HVX_Vector *) (tiles + 0 * HMX_FP16_TILE_N_ELMS); + HVX_Vector * t01 = (HVX_Vector *) (tiles + 1 * HMX_FP16_TILE_N_ELMS); + HVX_Vector * t10 = (HVX_Vector *) (tiles + 2 * HMX_FP16_TILE_N_ELMS); + HVX_Vector * t11 = (HVX_Vector *) (tiles + 3 * HMX_FP16_TILE_N_ELMS); + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp0 = Q6_W_vshuff_VVR(rows[2 * r + 1], rows[2 * r + 0], -2); + t00[r] = Q6_V_lo_W(vp0); + t01[r] = Q6_V_hi_W(vp0); + + HVX_VectorPair vp1 = Q6_W_vshuff_VVR(rows[32 + 2 * r + 1], rows[32 + 2 * r + 0], -2); + t10[r] = Q6_V_lo_W(vp1); + t11[r] = Q6_V_hi_W(vp1); + } +} + +static inline void gdn_unpack_64xS_tiles_to_f32( + float * restrict dst_f32, + const __fp16 * restrict tiles, + uint32_t S_v +) { + const uint32_t n_col_tiles = S_v / 32; + for (uint32_t r0 = 0; r0 < 2; ++r0) { + for (uint32_t d = 0; d < S_v / 64; ++d) { + const HVX_Vector * t0 = (const HVX_Vector *) (tiles + (r0 * n_col_tiles + 2 * d + 0) * HMX_FP16_TILE_N_ELMS); + const HVX_Vector * t1 = (const HVX_Vector *) (tiles + (r0 * n_col_tiles + 2 * d + 1) * HMX_FP16_TILE_N_ELMS); + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp01 = Q6_W_vdeal_VVR(t1[r], t0[r], -2); + HVX_VectorPair p0 = hvx_vec_f16_to_f32(Q6_V_lo_W(vp01)); + HVX_VectorPair p1 = hvx_vec_f16_to_f32(Q6_V_hi_W(vp01)); + + float * out0 = dst_f32 + (r0 * 32 + 2 * r + 0) * S_v + d * 64; + float * out1 = dst_f32 + (r0 * 32 + 2 * r + 1) * S_v + d * 64; + + hvx_vmem(out0 + 0) = Q6_V_lo_W(p0); + hvx_vmem(out0 + 32) = Q6_V_hi_W(p0); + hvx_vmem(out1 + 0) = Q6_V_lo_W(p1); + hvx_vmem(out1 + 32) = Q6_V_hi_W(p1); + } + } + } +} + +static inline void gdn_unpack_64xS_tiles_to_f16( + __fp16 * restrict dst_f16, + const __fp16 * restrict tiles, + uint32_t S_v +) { + const uint32_t n_col_tiles = S_v / 32; + for (uint32_t r0 = 0; r0 < 2; ++r0) { + for (uint32_t d = 0; d < S_v / 64; ++d) { + const HVX_Vector * t0 = (const HVX_Vector *) (tiles + (r0 * n_col_tiles + 2 * d + 0) * HMX_FP16_TILE_N_ELMS); + const HVX_Vector * t1 = (const HVX_Vector *) (tiles + (r0 * n_col_tiles + 2 * d + 1) * HMX_FP16_TILE_N_ELMS); + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp01 = Q6_W_vdeal_VVR(t1[r], t0[r], -2); + __fp16 * out0 = dst_f16 + (r0 * 32 + 2 * r + 0) * S_v + d * 64; + __fp16 * out1 = dst_f16 + (r0 * 32 + 2 * r + 1) * S_v + d * 64; + + hvx_vmem(out0) = Q6_V_lo_W(vp01); + hvx_vmem(out1) = Q6_V_hi_W(vp01); + } + } + } +} + +static inline void gdn_unpack_SxS_tiles_to_f32( + float * restrict dst_f32, + const __fp16 * restrict tiles, + uint32_t S_v +) { + const uint32_t n_tiles = S_v / 32; + for (uint32_t r0 = 0; r0 < n_tiles; ++r0) { + for (uint32_t d = 0; d < S_v / 64; ++d) { + const HVX_Vector * t0 = (const HVX_Vector *) (tiles + (r0 * n_tiles + 2 * d + 0) * HMX_FP16_TILE_N_ELMS); + const HVX_Vector * t1 = (const HVX_Vector *) (tiles + (r0 * n_tiles + 2 * d + 1) * HMX_FP16_TILE_N_ELMS); + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp01 = Q6_W_vdeal_VVR(t1[r], t0[r], -2); + HVX_VectorPair p0 = hvx_vec_f16_to_f32(Q6_V_lo_W(vp01)); + HVX_VectorPair p1 = hvx_vec_f16_to_f32(Q6_V_hi_W(vp01)); + + float * out0 = dst_f32 + (r0 * 32 + 2 * r + 0) * S_v + d * 64; + float * out1 = dst_f32 + (r0 * 32 + 2 * r + 1) * S_v + d * 64; + + hvx_vmem(out0 + 0) = Q6_V_lo_W(p0); + hvx_vmem(out0 + 32) = Q6_V_hi_W(p0); + hvx_vmem(out1 + 0) = Q6_V_lo_W(p1); + hvx_vmem(out1 + 32) = Q6_V_hi_W(p1); + } + } + } +} + +static inline void gdn_f32_to_hmx_row_tiles_and_f16( + __fp16 * restrict dst_tiles, + __fp16 * restrict dst_prime_tiles, + __fp16 * restrict dst_f16, + const float * restrict src, + const __fp16 * restrict scale_per_row, + uint32_t n_rows, + uint32_t n_cols +) { + const uint32_t n_col_tiles = n_cols / 32; + const uint32_t * scale_pairs = (const uint32_t *) scale_per_row; + + for (uint32_t r = 0; r < n_rows; r += 2) { + uint32_t r0 = r / 32; + uint32_t r1 = (r % 32) / 2; + const float * p0 = src + (r + 0) * n_cols; + const float * p1 = src + (r + 1) * n_cols; + + HVX_Vector v_scale; + if (dst_prime_tiles) { + uint32_t scale_pair = scale_pairs ? scale_pairs[r / 2] : 0x3c003c00; + v_scale = Q6_V_vsplat_R(scale_pair); + } + + for (uint32_t c = 0; c < n_col_tiles; c += 2) { + HVX_Vector v0_0 = hvx_vmem(p0 + (c + 0) * 32); + HVX_Vector v1_0 = hvx_vmem(p1 + (c + 0) * 32); + HVX_Vector v0_1 = hvx_vmem(p0 + (c + 1) * 32); + HVX_Vector v1_1 = hvx_vmem(p1 + (c + 1) * 32); + + HVX_Vector vh0 = hvx_vec_f32_to_f16_shuff(v0_0, v1_0); + HVX_Vector vh1 = hvx_vec_f32_to_f16_shuff(v0_1, v1_1); + __fp16 * tile0 = dst_tiles + (r0 * n_col_tiles + c + 0) * HMX_FP16_TILE_N_ELMS; + __fp16 * tile1 = dst_tiles + (r0 * n_col_tiles + c + 1) * HMX_FP16_TILE_N_ELMS; + ((HVX_Vector *) tile0)[r1] = vh0; + ((HVX_Vector *) tile1)[r1] = vh1; + + if (dst_prime_tiles) { + HVX_Vector vh0_s = hvx_vec_mul_f16_f16(vh0, v_scale); + HVX_Vector vh1_s = hvx_vec_mul_f16_f16(vh1, v_scale); + __fp16 * tile0_s = dst_prime_tiles + (r0 * n_col_tiles + c + 0) * HMX_FP16_TILE_N_ELMS; + __fp16 * tile1_s = dst_prime_tiles + (r0 * n_col_tiles + c + 1) * HMX_FP16_TILE_N_ELMS; + ((HVX_Vector *) tile0_s)[r1] = vh0_s; + ((HVX_Vector *) tile1_s)[r1] = vh1_s; + } + + if (dst_f16) { + HVX_VectorPair vp01 = Q6_W_vdeal_VVR(vh1, vh0, -2); + hvx_vmem(dst_f16 + (r + 0) * n_cols + c * 32) = Q6_V_lo_W(vp01); + hvx_vmem(dst_f16 + (r + 1) * n_cols + c * 32) = Q6_V_hi_W(vp01); + } + } + } +} + +static inline void hvx_transpose_32x32_words(HVX_Vector * restrict m, HVX_Vector * restrict tmp) { + for (int i = 0; i < 16; ++i) { + HVX_VectorPair p = Q6_W_vshuff_VVR(m[2*i + 1], m[2*i], -4); + tmp[2*i + 0] = Q6_V_lo_W(p); + tmp[2*i + 1] = Q6_V_hi_W(p); + } + + for (int b = 0; b < 32; b += 4) { + HVX_VectorPair p0 = Q6_W_vshuff_VVR(tmp[b + 2], tmp[b + 0], -8); + HVX_VectorPair p1 = Q6_W_vshuff_VVR(tmp[b + 3], tmp[b + 1], -8); + m[b + 0] = Q6_V_lo_W(p0); m[b + 1] = Q6_V_hi_W(p0); + m[b + 2] = Q6_V_lo_W(p1); m[b + 3] = Q6_V_hi_W(p1); + } + + for (int b = 0; b < 32; b += 8) { + for (int i = 0; i < 4; ++i) { + HVX_VectorPair p = Q6_W_vshuff_VVR(m[b + i + 4], m[b + i], -16); + tmp[b + 2*i + 0] = Q6_V_lo_W(p); + tmp[b + 2*i + 1] = Q6_V_hi_W(p); + } + } + + for (int b = 0; b < 32; b += 16) { + for (int i = 0; i < 8; ++i) { + HVX_VectorPair p = Q6_W_vshuff_VVR(tmp[b + i + 8], tmp[b + i], -32); + m[b + 2*i + 0] = Q6_V_lo_W(p); + m[b + 2*i + 1] = Q6_V_hi_W(p); + } + } + + for (int i = 0; i < 16; ++i) { + HVX_VectorPair p = Q6_W_vshuff_VVR(m[i + 16], m[i], -64); + tmp[2 * i + 0] = Q6_V_lo_W(p); + tmp[2 * i + 1] = Q6_V_hi_W(p); + } + + for (int i = 0; i < 32; ++i) { + m[i] = tmp[i]; + } +} + +static inline void gdn_pack_d_t_row_tiles( + __fp16 * restrict dst_tiles, + const __fp16 * restrict src_d, + uint32_t S_v, + HVX_Vector * restrict m, + HVX_Vector * restrict tmp +) { + for (uint32_t col_half = 0; col_half < S_v / 64; ++col_half) { + uint32_t r0_base = col_half * 2; + for (uint32_t c0 = 0; c0 < 2; ++c0) { + for (uint32_t s_local = 0; s_local < 32; ++s_local) { + uint32_t s = c0 * 32 + s_local; + m[s_local] = hvx_vmem(src_d + s * S_v + col_half * 64); + } + + hvx_transpose_32x32_words(m, tmp); + + uint32_t tile0_idx = (r0_base + 0) * 2 + c0; + uint32_t tile1_idx = (r0_base + 1) * 2 + c0; + HVX_Vector * t0 = (HVX_Vector *)(dst_tiles + tile0_idx * HMX_FP16_TILE_N_ELMS); + HVX_Vector * t1 = (HVX_Vector *)(dst_tiles + tile1_idx * HMX_FP16_TILE_N_ELMS); + + for (uint32_t r = 0; r < 16; ++r) { + t0[r] = m[r]; + t1[r] = m[16 + r]; + } + } + } +} + +static __attribute__((noinline)) void gdn_build_inv_l_blocks( + __fp16 * restrict inv_row_tiles, + const HVX_Vector * restrict rows_kk, + const __fp16 * restrict decay_m, + const float * restrict beta, + __fp16 * restrict l10_tile, + __fp16 * restrict neg_a11_tile +) { + const HVX_Vector v_one_f16 = hvx_vec_splat_f16(1.0f); + const HVX_VectorPred q_mask64 = Q6_Q_vsetq2_R(64); + + uint16_t beta_u16[64] __attribute__((aligned(128))); + uint16_t l00[32][32] __attribute__((aligned(128))); + uint16_t l11[32][32] __attribute__((aligned(128))); + + HVX_Vector * restrict p_l00 = (HVX_Vector *) l00; + HVX_Vector * restrict p_l11 = (HVX_Vector *) l11; + HVX_Vector * restrict p_l10_tile = (HVX_Vector *) l10_tile; + + HVX_Vector * restrict tile00 = (HVX_Vector *) (inv_row_tiles + 0 * HMX_FP16_TILE_N_ELMS); + HVX_Vector * restrict tile01 = (HVX_Vector *) (inv_row_tiles + 1 * HMX_FP16_TILE_N_ELMS); + HVX_Vector * restrict tile11 = (HVX_Vector *) (inv_row_tiles + 3 * HMX_FP16_TILE_N_ELMS); + HVX_Vector * restrict p_neg_a11 = (HVX_Vector *) neg_a11_tile; + + hvx_vmem(beta_u16) = hvx_vec_f32_to_f16(hvx_vmem(beta + 0), hvx_vmem(beta + 32)); + + for (uint32_t r = 0; r < 16; ++r) { + tile01[r] = Q6_V_vzero(); + } + + for (uint32_t r = 0; r < 16; ++r) { + uint32_t t0 = 2 * r; + uint32_t t1 = t0 + 1; + + HVX_Vector v_d0 = hvx_vmem(decay_m + t0 * 64); + HVX_Vector v_d1 = hvx_vmem(decay_m + t1 * 64); + HVX_Vector v_b0 = Q6_Vh_vsplat_R(beta_u16[t0]); + HVX_Vector v_b1 = Q6_Vh_vsplat_R(beta_u16[t1]); + + HVX_Vector r0 = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(rows_kk[t0], v_d0), v_b0); + HVX_Vector r1 = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(rows_kk[t1], v_d1), v_b1); + + p_l00[r] = Q6_V_vmux_QVV(q_mask64, r0, Q6_V_vror_VR(r1, 64)); + } + + for (uint32_t r = 0; r < 16; ++r) { + uint32_t t0 = 32 + 2 * r; + uint32_t t1 = t0 + 1; + + HVX_Vector v_d0 = hvx_vmem(decay_m + t0 * 64); + HVX_Vector v_d1 = hvx_vmem(decay_m + t1 * 64); + HVX_Vector v_b0 = Q6_Vh_vsplat_R(beta_u16[t0]); + HVX_Vector v_b1 = Q6_Vh_vsplat_R(beta_u16[t1]); + + HVX_Vector r0 = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(rows_kk[t0], v_d0), v_b0); + HVX_Vector r1 = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(rows_kk[t1], v_d1), v_b1); + + HVX_VectorPair vp_l10 = Q6_W_vshuff_VVR(r1, r0, -2); + p_l10_tile[r] = Q6_V_lo_W(vp_l10); + p_l11[r] = Q6_V_vmux_QVV(q_mask64, Q6_V_vror_VR(r0, 64), r1); + } + + HVX_Vector a_rows[32]; + for (uint32_t t = 0; t < 32; ++t) { + HVX_Vector v_inv = Q6_V_vzero(); + for (uint32_t k = 0; k < t; ++k) { + HVX_Vector v_lk = Q6_Vh_vsplat_R(l00[t][k]); + v_inv = hvx_vec_sub_f16_f16(v_inv, hvx_vec_mul_f16_f16(v_lk, a_rows[k])); + } + HVX_VectorPred q_diag = (t == 0) ? Q6_Q_vsetq2_R(2) : Q6_Q_and_QQn(Q6_Q_vsetq2_R(2 * (t + 1)), Q6_Q_vsetq2_R(2 * t)); + a_rows[t] = Q6_V_vand_QV(q_mask64, Q6_V_vmux_QVV(q_diag, v_one_f16, v_inv)); + } + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp = Q6_W_vshuff_VVR(a_rows[2 * r + 1], a_rows[2 * r + 0], -2); + tile00[r] = Q6_V_lo_W(vp); + } + + for (uint32_t t = 0; t < 32; ++t) { + HVX_Vector v_inv = Q6_V_vzero(); + for (uint32_t k = 0; k < t; ++k) { + HVX_Vector v_lk = Q6_Vh_vsplat_R(l11[t][k]); + v_inv = hvx_vec_sub_f16_f16(v_inv, hvx_vec_mul_f16_f16(v_lk, a_rows[k])); + } + HVX_VectorPred q_diag = (t == 0) ? Q6_Q_vsetq2_R(2) : Q6_Q_and_QQn(Q6_Q_vsetq2_R(2 * (t + 1)), Q6_Q_vsetq2_R(2 * t)); + a_rows[t] = Q6_V_vand_QV(q_mask64, Q6_V_vmux_QVV(q_diag, v_one_f16, v_inv)); + } + + for (uint32_t r = 0; r < 16; ++r) { + HVX_VectorPair vp = Q6_W_vshuff_VVR(a_rows[2 * r + 1], a_rows[2 * r + 0], -2); + tile11[r] = Q6_V_lo_W(vp); + + HVX_Vector n0 = hvx_vec_sub_f16_f16(Q6_V_vzero(), a_rows[2 * r + 0]); + HVX_Vector n1 = hvx_vec_sub_f16_f16(Q6_V_vzero(), a_rows[2 * r + 1]); + HVX_VectorPair vp_neg = Q6_W_vshuff_VVR(n1, n0, -2); + p_neg_a11[r] = Q6_V_lo_W(vp_neg); + } +} + + +static inline void gdn_dma_push_chunk_inputs( + dma_queue * dma_q, + float * vtcm_q, + float * vtcm_k, + float * vtcm_v, + const struct htp_tensor * q, + const struct htp_tensor * k, + const struct htp_tensor * v, + uint32_t iq3, uint32_t iq1, + uint32_t ik3, uint32_t ik1, + uint32_t iv3, uint32_t iv1, + uint32_t t_chunk, + uint32_t chunk_size, + uint32_t S_v +) { + const dma_addr_t q_dma = q->data + (uint64_t) iq3 * q->nb[3] + (uint64_t) t_chunk * q->nb[2] + (uint64_t) iq1 * q->nb[1]; + const dma_addr_t k_dma = k->data + (uint64_t) ik3 * k->nb[3] + (uint64_t) t_chunk * k->nb[2] + (uint64_t) ik1 * k->nb[1]; + const dma_addr_t v_dma = v->data + (uint64_t) iv3 * v->nb[3] + (uint64_t) t_chunk * v->nb[2] + (uint64_t) iv1 * v->nb[1]; + + dma_queue_push(dma_q, dma_make_data(vtcm_q, q_dma), S_v * sizeof(float), q->nb[2], S_v * sizeof(float), chunk_size); + dma_queue_push(dma_q, dma_make_data(vtcm_k, k_dma), S_v * sizeof(float), k->nb[2], S_v * sizeof(float), chunk_size); + dma_queue_push(dma_q, dma_make_data(vtcm_v, v_dma), S_v * sizeof(float), v->nb[2], S_v * sizeof(float), chunk_size); +} + +static inline void gdn_dma_push_chunk_gb( + dma_queue * dma_q, + float * vtcm_g_raw, + float * vtcm_b_raw, + const struct htp_tensor * g, + const struct htp_tensor * beta, + uint32_t iv3, + uint32_t iv1, + uint32_t t_chunk, + uint32_t chunk_size, + uint32_t n_batch +) { + const dma_addr_t g_dma = g->data + (uint64_t) iv3 * g->nb[3] + (uint64_t) t_chunk * g->nb[2] + (uint64_t) iv1 * g->nb[1]; + const dma_addr_t beta_dma = beta->data + (uint64_t) iv3 * beta->nb[3] + (uint64_t) t_chunk * beta->nb[2] + (uint64_t) iv1 * beta->nb[1]; + const uint32_t row_bytes = n_batch * sizeof(float); + + dma_queue_push(dma_q, dma_make_data(vtcm_g_raw, g_dma), row_bytes, g->nb[2], row_bytes, chunk_size); + dma_queue_push(dma_q, dma_make_data(vtcm_b_raw, beta_dma), row_bytes, beta->nb[2], row_bytes, chunk_size); +} + +static inline void gdn_pack_s_col_tiles( + __fp16 * restrict vtcm_s_col_tiles, + __fp16 * restrict vtcm_s_f16, + const float * restrict vtcm_s_state, + uint32_t S_v +) { + for (uint32_t j = 0; j < S_v; ++j) { + for (uint32_t i = 0; i < S_v; i += 64) { + HVX_Vector v0 = hvx_vmem(vtcm_s_state + j * S_v + i + 0); + HVX_Vector v1 = (i + 32 < S_v) ? hvx_vmem(vtcm_s_state + j * S_v + i + 32) : Q6_V_vzero(); + hvx_vmem(vtcm_s_f16 + j * S_v + i) = hvx_vec_f32_to_f16(v0, v1); + } + } + hmx_interleave_rows_to_tiles(vtcm_s_col_tiles, vtcm_s_f16, S_v, S_v, S_v, 0, S_v); +} + +struct htp_gdn_head_ptrs { + float * s_state; + __fp16 * s_f16; + __fp16 * s_col_tiles; + float * s_update_f32; + __fp16 * s_update_tiles; + + float * q_f32[2]; + float * k_f32[2]; + float * v_f32[2]; + float * g_f32[2]; + float * b_f32[2]; + float * o_f32[2]; + + float * v_inter_f32; + float * o_inter_f32; + float * o_intra_f32; + + __fp16 * k_f16; + __fp16 * v_prime_f16; + __fp16 * delta_f16; + __fp16 * d_f16; + + __fp16 * q_row_tiles; + __fp16 * q_prime_row_tiles; + __fp16 * k_row_tiles; + __fp16 * k_col_tiles; + __fp16 * k_prime_row_tiles; + __fp16 * k_col_tiles_64x128; + __fp16 * kk_tiles; + __fp16 * qk_tiles; + __fp16 * v_inter_tiles; + __fp16 * o_inter_tiles; + __fp16 * inv_row_tiles; + __fp16 * a_row_tiles; + __fp16 * v_prime_col_tiles; + __fp16 * delta_tiles; + __fp16 * delta_col_tiles; + __fp16 * o_intra_tiles; + __fp16 * d_row_tiles; + + __fp16 * gamma; + float * lambda_init; + __fp16 * lambda_init_f16; + __fp16 * decay_m; + __fp16 * decay_a; + + HVX_Vector * rows_kk; + HVX_Vector * rows_qk; + HVX_Vector * rows_inv; + HVX_Vector * rows_a; + + HVX_Vector * vtcm_m; + HVX_Vector * vtcm_tmp; + + uint32_t iv1; + uint32_t iv3; + uint32_t iq1; + uint32_t ik1; + uint32_t iq3; + uint32_t ik3; + dma_addr_t state_in_dma; + dma_addr_t state_out_dma; +}; + +static inline void gdn_init_head_ptrs( + struct htp_gdn_head_ptrs * head, + const struct htp_gdn_hmx_vtcm_layout * L, + uint8_t * vtcm_base, + uint32_t h, + uint32_t base_iv1, + uint32_t iv3, + const struct htp_tensor * q, + const struct htp_tensor * k, + const struct htp_tensor * v, + const struct htp_tensor * state, + const struct htp_tensor * dst, + const struct htp_tensor * dst_cache, + const struct htp_gdn_kernel_params * kparams, + uint32_t S_v, + uint32_t H, + uint32_t n_tokens, + uint32_t chunk_size +) { + const size_t dma_scalar_sz = hex_round_up(chunk_size * sizeof(float), 128); + const size_t decay_sz = 64 * 64 * sizeof(__fp16); + const size_t row_vecs_sz = 64 * 128; + + head->s_state = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_s_state + h * L->state_f32_bytes); + head->s_f16 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_s_f16 + h * L->state_f16_bytes); + head->s_col_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_s_col_tiles + h * L->state_tiles_bytes); + head->s_update_f32 = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_s_update_f32 + h * L->state_f32_bytes); + head->s_update_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_s_update_tiles + h * L->state_tiles_bytes); + + head->q_f32[0] = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_q_f32[0] + h * L->dma_chunk_bytes); + head->q_f32[1] = L->pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L->off_q_f32[1] + h * L->dma_chunk_bytes) : head->q_f32[0]; + head->k_f32[0] = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_k_f32[0] + h * L->dma_chunk_bytes); + head->k_f32[1] = L->pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L->off_k_f32[1] + h * L->dma_chunk_bytes) : head->k_f32[0]; + head->v_f32[0] = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_v_f32[0] + h * L->dma_chunk_bytes); + head->v_f32[1] = L->pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L->off_v_f32[1] + h * L->dma_chunk_bytes) : head->v_f32[0]; + head->g_f32[0] = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_g_f32[0] + h * dma_scalar_sz); + head->g_f32[1] = L->pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L->off_g_f32[1] + h * dma_scalar_sz) : head->g_f32[0]; + head->b_f32[0] = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_b_f32[0] + h * dma_scalar_sz); + head->b_f32[1] = L->pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L->off_b_f32[1] + h * dma_scalar_sz) : head->b_f32[0]; + head->o_f32[0] = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_o_f32[0] + h * L->dma_chunk_bytes); + head->o_f32[1] = L->pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L->off_o_f32[1] + h * L->dma_chunk_bytes) : head->o_f32[0]; + + head->v_inter_f32 = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_v_inter_f32 + h * L->dma_chunk_bytes); + head->o_inter_f32 = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_o_inter_f32 + h * L->dma_chunk_bytes); + head->o_intra_f32 = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_o_intra_f32 + h * L->dma_chunk_bytes); + + head->k_f16 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_k_f16 + h * L->act_f16_bytes); + head->v_prime_f16 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_v_prime_f16 + h * L->act_f16_bytes); + head->delta_f16 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_delta_f16 + h * L->act_f16_bytes); + head->d_f16 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_d_f16 + h * L->act_f16_bytes); + + head->q_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_q_row_tiles + h * L->tile_64xSv_bytes); + head->q_prime_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_q_prime_row_tiles + h * L->tile_64xSv_bytes); + head->k_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_k_row_tiles + h * L->tile_64xSv_bytes); + head->k_col_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_k_col_tiles + h * L->tile_64xSv_bytes); + head->k_prime_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_k_prime_row_tiles + h * L->tile_64xSv_bytes); + head->k_col_tiles_64x128 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_k_col_tiles_64x128 + h * L->tile_64xSv_bytes); + head->kk_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_kk_tiles + h * L->tile_64x64_bytes); + head->qk_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_qk_tiles + h * L->tile_64x64_bytes); + head->v_inter_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_v_inter_tiles + h * L->tile_64xSv_bytes); + head->o_inter_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_o_inter_tiles + h * L->tile_64xSv_bytes); + head->inv_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_inv_row_tiles + h * L->tile_64x64_bytes); + head->a_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_a_row_tiles + h * L->tile_64x64_bytes); + head->v_prime_col_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_v_prime_col_tiles + h * L->tile_64xSv_bytes); + head->delta_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_delta_tiles + h * L->tile_64xSv_bytes); + head->delta_col_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_delta_col_tiles + h * L->tile_64xSv_bytes); + head->o_intra_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_o_intra_tiles + h * L->tile_64xSv_bytes); + head->d_row_tiles = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_d_row_tiles + h * L->tile_64xSv_bytes); + + head->gamma = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_gamma + h * dma_scalar_sz); + head->lambda_init_f16 = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_gamma + h * dma_scalar_sz + 128); + head->lambda_init = VTCM_LAYOUT_PTR(float, vtcm_base, L->off_lambda_init + h * dma_scalar_sz); + head->decay_m = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_decay_m + h * decay_sz); + head->decay_a = VTCM_LAYOUT_PTR(__fp16, vtcm_base, L->off_decay_a + h * decay_sz); + + head->rows_kk = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_rows_kk + h * row_vecs_sz); + head->rows_qk = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_rows_qk + h * row_vecs_sz); + head->rows_inv = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_rows_inv + h * row_vecs_sz); + head->rows_a = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_rows_a + h * row_vecs_sz); + + head->vtcm_m = VTCM_LAYOUT_PTR(HVX_Vector, vtcm_base, L->off_thread_scratch + h * (64 * 128)); + head->vtcm_tmp = head->vtcm_m + 32; + + head->iv1 = base_iv1 + h; + head->iv3 = iv3; + head->iq1 = fastmodulo(head->iv1, q->ne[1], &kparams->div_q1); + head->ik1 = fastmodulo(head->iv1, k->ne[1], &kparams->div_k1); + head->iq3 = fastdiv(head->iv3, &kparams->div_rq3); + head->ik3 = fastdiv(head->iv3, &kparams->div_rk3); + + head->state_in_dma = state->data + + ((uint64_t) head->iv3 * kparams->state_seq_stride + (uint64_t) head->iv1 * S_v * S_v) * sizeof(float); + + head->state_out_dma = dst_cache ? + (dst_cache->data + ((uint64_t) head->iv3 * H + head->iv1) * S_v * S_v * sizeof(float)) : + (dst->data + ((uint64_t) S_v * H * n_tokens * kparams->n_seqs + (uint64_t) (head->iv3 * H + head->iv1) * S_v * S_v) * sizeof(float)); +} + +struct htp_gdn_batch_context { + struct htp_gdn_head_ptrs * heads; + const float * vtcm_g_raw; + const float * vtcm_b_raw; + uint32_t curr_buf; + uint32_t c; + uint32_t n_batch; + uint32_t S_v; + float scale; + struct htp_ops_context * octx; + const struct htp_gdn_kernel_params * kparams; +}; + +static void gdn_hvx_init_state_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + gdn_pack_s_col_tiles(head->s_col_tiles, head->s_f16, head->s_state, bctx->S_v); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0); +} + +static inline __attribute__((unused)) HVX_Vector hvx_clamp_neg20_0(HVX_Vector v, HVX_Vector v_zero, HVX_Vector v_neg20) { + HVX_VectorPred p_gt = Q6_Q_vcmp_gt_VsfVsf(v, v_zero); + v = Q6_V_vmux_QVV(p_gt, v_zero, v); + HVX_VectorPred p_lt = Q6_Q_vcmp_gt_VsfVsf(v_neg20, v); + return Q6_V_vmux_QVV(p_lt, v_neg20, v); +} + +static inline HVX_Vector hvx_prefix_scan_f32(HVX_Vector v, HVX_Vector carry_in) { + const HVX_Vector zero = Q6_V_vzero(); + + v = hvx_vec_add_f32_f32(v, Q6_V_vlalign_VVR(v, zero, 4)); + v = hvx_vec_add_f32_f32(v, Q6_V_vlalign_VVR(v, zero, 8)); + v = hvx_vec_add_f32_f32(v, Q6_V_vlalign_VVR(v, zero, 16)); + v = hvx_vec_add_f32_f32(v, Q6_V_vlalign_VVR(v, zero, 32)); + v = hvx_vec_add_f32_f32(v, Q6_V_vlalign_VVR(v, zero, 64)); + v = hvx_vec_add_f32_f32(v, carry_in); + + return v; +} + +static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) { + return hvx_vec_repl4(Q6_V_vror_VR(v, 124)); +} + +static void gdn_hvx_phase1a_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t curr_buf = bctx->curr_buf; + const uint32_t S_v = bctx->S_v; + const uint32_t n_batch = bctx->n_batch; + + if (n_batch == 1) { + hvx_vmem(head->g_f32[curr_buf] + 0) = hvx_vmem(bctx->vtcm_g_raw + 0); + hvx_vmem(head->g_f32[curr_buf] + 32) = hvx_vmem(bctx->vtcm_g_raw + 32); + hvx_vmem(head->b_f32[curr_buf] + 0) = hvx_vmem(bctx->vtcm_b_raw + 0); + hvx_vmem(head->b_f32[curr_buf] + 32) = hvx_vmem(bctx->vtcm_b_raw + 32); + } else { + int32_t offsets[32] __attribute__((aligned(128))); + for (int k = 0; k < 32; ++k) { + offsets[k] = k * n_batch * sizeof(float); + } + HVX_Vector vv = *(const HVX_Vector *) offsets; + const size_t rt_g = (size_t) ((const uint8_t *) bctx->vtcm_g_raw + i * sizeof(float)); + const size_t rt_b = (size_t) ((const uint8_t *) bctx->vtcm_b_raw + i * sizeof(float)); + const size_t mu = 64 * n_batch * sizeof(float); + + Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 0), rt_g, mu, vv); + Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 32), rt_g + 32 * n_batch * sizeof(float), mu, vv); + Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 0), rt_b, mu, vv); + Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 32), rt_b + 32 * n_batch * sizeof(float), mu, vv); + } + + const uint32_t t_chunk = bctx->c * 64; + const uint32_t valid_tokens = hex_smin(64, bctx->kparams->n_tokens - t_chunk); + if (valid_tokens < 64) { + for (uint32_t t = valid_tokens; t < 64; ++t) { + head->g_f32[curr_buf][t] = 0.0f; + head->b_f32[curr_buf][t] = 0.0f; + } + const HVX_Vector vzero = Q6_V_vzero(); + for (uint32_t t = valid_tokens; t < 64; ++t) { + for (uint32_t j = 0; j < S_v; j += 32) { + hvx_vmem(head->q_f32[curr_buf] + t * S_v + j) = vzero; + hvx_vmem(head->k_f32[curr_buf] + t * S_v + j) = vzero; + hvx_vmem(head->v_f32[curr_buf] + t * S_v + j) = vzero; + } + } + } + + const HVX_Vector v_g0 = hvx_vmem(head->g_f32[curr_buf] + 0); + const HVX_Vector v_g1 = hvx_vmem(head->g_f32[curr_buf] + 32); + + HVX_Vector v_gamma0 = hvx_prefix_scan_f32(v_g0, Q6_V_vzero()); + HVX_Vector v_carry = hvx_splat_last_f32(v_gamma0); + HVX_Vector v_gamma1 = hvx_prefix_scan_f32(v_g1, v_carry); + + const HVX_Vector v_zero = Q6_V_vzero(); + const HVX_Vector v_neg20 = hvx_vec_splat_f32(-20.0f); + + hvx_vmem(head->gamma) = hvx_vec_f32_to_f16(v_gamma0, v_gamma1); + + HVX_Vector v_l0 = hvx_vec_exp_f32(hvx_clamp_neg20_0(v_gamma0, v_zero, v_neg20)); + HVX_Vector v_l1 = hvx_vec_exp_f32(hvx_clamp_neg20_0(v_gamma1, v_zero, v_neg20)); + + hvx_vmem(head->lambda_init + 0) = v_l0; + hvx_vmem(head->lambda_init + 32) = v_l1; + hvx_vmem(head->lambda_init_f16) = hvx_vec_f32_to_f16(v_l0, v_l1); + + gdn_f32_to_hmx_row_tiles_and_f16(head->k_row_tiles, head->k_prime_row_tiles, head->k_f16, + head->k_f32[curr_buf], head->lambda_init_f16, 64, S_v); + hmx_interleave_rows_to_tiles(head->k_col_tiles, head->k_f16, 64, S_v, S_v, 0, 64); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info); +} + +static void gdn_hvx_phase1b_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t curr_buf = bctx->curr_buf; + const uint32_t S_v = bctx->S_v; + + gdn_f32_to_hmx_row_tiles_and_f16(head->q_row_tiles, head->q_prime_row_tiles, NULL, + head->q_f32[curr_buf], head->lambda_init_f16, 64, S_v); + + hmx_interleave_cols_to_tiles(head->k_col_tiles_64x128, head->k_f16, 64, S_v, S_v, 2, 0, 64); + + const uint16_t * gamma_u16 = (const uint16_t *) head->gamma; + const HVX_Vector v_gamma = hvx_vmem(head->gamma); + + const HVX_Vector v_zero_f16 = Q6_V_vzero(); + const HVX_Vector v_neg20_f16 = hvx_vec_splat_f16(-20.0f); + const HVX_Vector v_log2e_f16 = hvx_vec_splat_f16(1.4426950408889634f); + const HVX_Vector v_one_f16 = hvx_vec_splat_f16(1.0f); + + hvx_vmem(head->decay_m + 0) = Q6_V_vzero(); + hvx_vmem(head->decay_a + 0) = Q6_V_vand_QV(Q6_Q_vsetq2_R(2), v_one_f16); + + for (uint32_t t = 1; t < 63; t += 2) { + uint32_t t0 = t; + uint32_t t1 = t + 1; + + HVX_Vector v_gamma_t0 = Q6_Vh_vsplat_R(gamma_u16[t0]); + HVX_Vector v_gamma_t1 = Q6_Vh_vsplat_R(gamma_u16[t1]); + + HVX_Vector diff0 = hvx_vec_sub_f16_f16(v_gamma_t0, v_gamma); + HVX_Vector diff1 = hvx_vec_sub_f16_f16(v_gamma_t1, v_gamma); + + HVX_VectorPred p_gt0 = Q6_Q_vcmp_gt_VhfVhf(diff0, v_zero_f16); + HVX_VectorPred p_gt1 = Q6_Q_vcmp_gt_VhfVhf(diff1, v_zero_f16); + + diff0 = Q6_V_vmux_QVV(p_gt0, v_zero_f16, diff0); + diff1 = Q6_V_vmux_QVV(p_gt1, v_zero_f16, diff1); + + diff0 = Q6_Vhf_vmax_VhfVhf(v_neg20_f16, diff0); + diff1 = Q6_Vhf_vmax_VhfVhf(v_neg20_f16, diff1); + + HVX_Vector diff_log2e0 = hvx_vec_mul_f16_f16(diff0, v_log2e_f16); + HVX_Vector diff_log2e1 = hvx_vec_mul_f16_f16(diff1, v_log2e_f16); + + HVX_Vector v_exp0 = hvx_vec_exp2_f16(diff_log2e0); + HVX_Vector v_exp1 = hvx_vec_exp2_f16(diff_log2e1); + + HVX_VectorPred mask_lt0 = Q6_Q_vsetq2_R(2 * t0); + HVX_VectorPred mask_lt1 = Q6_Q_vsetq2_R(2 * t1); + + HVX_Vector v_m0 = Q6_V_vand_QV(mask_lt0, v_exp0); + HVX_Vector v_m1 = Q6_V_vand_QV(mask_lt1, v_exp1); + + HVX_VectorPred mask_le0 = Q6_Q_vsetq2_R(2 * (t0 + 1)); + HVX_VectorPred mask_le1 = Q6_Q_vsetq2_R(2 * (t1 + 1)); + + HVX_VectorPred mask_diag0 = Q6_Q_and_QQn(mask_le0, mask_lt0); + HVX_VectorPred mask_diag1 = Q6_Q_and_QQn(mask_le1, mask_lt1); + + HVX_Vector v_a0 = Q6_V_vmux_QVV(mask_diag0, v_one_f16, v_m0); + HVX_Vector v_a1 = Q6_V_vmux_QVV(mask_diag1, v_one_f16, v_m1); + + hvx_vmem(head->decay_m + t0 * 64) = v_m0; + hvx_vmem(head->decay_a + t0 * 64) = v_a0; + hvx_vmem(head->decay_m + t1 * 64) = v_m1; + hvx_vmem(head->decay_a + t1 * 64) = v_a1; + } + + { + HVX_Vector v_gamma_t = Q6_Vh_vsplat_R(gamma_u16[63]); + HVX_Vector diff = hvx_vec_sub_f16_f16(v_gamma_t, v_gamma); + HVX_VectorPred p_gt = Q6_Q_vcmp_gt_VhfVhf(diff, v_zero_f16); + diff = Q6_V_vmux_QVV(p_gt, v_zero_f16, diff); + diff = Q6_Vhf_vmax_VhfVhf(v_neg20_f16, diff); + + HVX_Vector diff_log2e = hvx_vec_mul_f16_f16(diff, v_log2e_f16); + HVX_Vector v_exp = hvx_vec_exp2_f16(diff_log2e); + + HVX_VectorPred mask_lt_t = Q6_Q_vsetq2_R(2 * 63); + HVX_Vector v_m = Q6_V_vand_QV(mask_lt_t, v_exp); + + HVX_VectorPred mask_le_t = Q6_Q_vcmp_eq_VhVh(v_zero_f16, v_zero_f16); + HVX_VectorPred mask_diag = Q6_Q_and_QQn(mask_le_t, mask_lt_t); + HVX_Vector v_a = Q6_V_vmux_QVV(mask_diag, v_one_f16, v_m); + + hvx_vmem(head->decay_m + 63 * 64) = v_m; + hvx_vmem(head->decay_a + 63 * 64) = v_a; + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_PREP, info); +} + +static void gdn_hvx_phase2_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_SOLVE, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t curr_buf = bctx->curr_buf; + + gdn_unpack_64x64_tiles_to_vectors(head->rows_kk, head->kk_tiles); + + gdn_build_inv_l_blocks( + head->inv_row_tiles, + head->rows_kk, + head->decay_m, + head->b_f32[curr_buf], + (__fp16 *) head->vtcm_m, + (__fp16 *) head->vtcm_m + HMX_FP16_TILE_N_ELMS + ); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_SOLVE, info); +} + +static void gdn_hvx_phase3_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_V_PREP, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t curr_buf = bctx->curr_buf; + const uint32_t S_v = bctx->S_v; + + gdn_unpack_64xS_tiles_to_f32(head->v_inter_f32, head->v_inter_tiles, S_v); + + HVX_VectorAlias local_b[2]; + local_b[0].v = hvx_vmem(head->b_f32[curr_buf] + 0); + local_b[1].v = hvx_vmem(head->b_f32[curr_buf] + 32); + + for (uint32_t t = 0; t < 64; ++t) { + HVX_Vector vb = hvx_vec_splat_f32(local_b[t / 32].fp32[t % 32]); + for (uint32_t j = 0; j < S_v; j += 64) { + HVX_Vector vv0 = hvx_vmem(head->v_f32[curr_buf] + t * S_v + j + 0); + HVX_Vector vv1 = (j + 32 < S_v) ? hvx_vmem(head->v_f32[curr_buf] + t * S_v + j + 32) : Q6_V_vzero(); + HVX_Vector vi0 = hvx_vmem(head->v_inter_f32 + t * S_v + j + 0); + HVX_Vector vi1 = (j + 32 < S_v) ? hvx_vmem(head->v_inter_f32 + t * S_v + j + 32) : Q6_V_vzero(); + + HVX_Vector vp0 = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv0, vi0), vb); + HVX_Vector vp1 = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv1, vi1), vb); + + hvx_vmem(head->v_prime_f16 + t * S_v + j) = hvx_vec_f32_to_f16(vp0, vp1); + } + } + + hmx_interleave_cols_to_tiles(head->v_prime_col_tiles, head->v_prime_f16, 64, S_v, S_v, 2, 0, 64); + + gdn_unpack_64x64_tiles_to_vectors(head->rows_qk, head->qk_tiles); + for (uint32_t t = 0; t < 64; ++t) { + HVX_Vector v_decay_a = hvx_vmem(head->decay_a + t * 64); + head->rows_a[t] = hvx_vec_mul_f16_f16(head->rows_qk[t], v_decay_a); + } + gdn_pack_64x64_vectors_to_tiles(head->a_row_tiles, head->rows_a); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_V_PREP, info); +} + +static void gdn_hvx_phase4_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_D_PREP, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t S_v = bctx->S_v; + + gdn_unpack_64xS_tiles_to_f16(head->delta_f16, head->delta_tiles, S_v); + hmx_interleave_cols_to_tiles(head->delta_col_tiles, head->delta_f16, 64, S_v, S_v, 2, 0, 64); + + const uint16_t * decay_last = (const uint16_t *) (head->decay_a + 63 * 64); + const HVX_Vector vzero = Q6_V_vzero(); + + for (uint32_t s = 0; s < 64; ++s) { + HVX_Vector vs = Q6_Vh_vsplat_R(decay_last[s]); + HVX_VectorPred p_zero = Q6_Q_vcmp_eq_VhVh(vs, vzero); + for (uint32_t j = 0; j < S_v; j += 64) { + HVX_Vector vd = hvx_vmem(head->delta_f16 + s * S_v + j); + HVX_Vector prod = hvx_vec_mul_f16_f16(vd, vs); + hvx_vmem(head->d_f16 + s * S_v + j) = Q6_V_vmux_QVV(p_zero, vzero, prod); + } + } + + gdn_pack_d_t_row_tiles(head->d_row_tiles, head->d_f16, S_v, head->vtcm_m, head->vtcm_tmp); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_D_PREP, info); +} + +static void gdn_hvx_phase5_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_OUT, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t curr_buf = bctx->curr_buf; + const uint32_t S_v = bctx->S_v; + const float scale = bctx->scale; + + gdn_unpack_64xS_tiles_to_f32(head->o_inter_f32, head->o_inter_tiles, S_v); + gdn_unpack_64xS_tiles_to_f32(head->o_intra_f32, head->o_intra_tiles, S_v); + + HVX_Vector vscale = hvx_vec_splat_f32(scale); + for (uint32_t j = 0; j < 64 * S_v / 32; ++j) { + HVX_Vector vi = hvx_vmem(head->o_inter_f32 + j * 32); + HVX_Vector va = hvx_vmem(head->o_intra_f32 + j * 32); + hvx_vmem(head->o_f32[curr_buf] + j * 32) = hvx_vec_mul_f32_f32(hvx_vec_add_f32_f32(vi, va), vscale); + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_OUT, info); +} + +static void gdn_hvx_phase6_worker(unsigned int n, unsigned int i, void * data) { + (void) n; + struct htp_gdn_batch_context * bctx = (struct htp_gdn_batch_context *) data; + struct htp_thread_trace * tr = &bctx->octx->ctx->trace[i]; + const uint16_t info = (uint16_t) bctx->c; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_GDN_STATE, info); + + struct htp_gdn_head_ptrs * head = &bctx->heads[i]; + const uint32_t S_v = bctx->S_v; + const uint32_t c = bctx->c; + const uint32_t n_chunks = bctx->kparams->n_chunks; + + gdn_unpack_SxS_tiles_to_f32(head->s_update_f32, head->s_update_tiles, S_v); + + HVX_VectorAlias last_lambda; + last_lambda.v = hvx_vmem(head->lambda_init + 32); + HVX_Vector v_l_final = hvx_vec_splat_f32(last_lambda.fp32[31]); + + for (uint32_t j = 0; j < S_v * S_v / 32; ++j) { + HVX_Vector vs_old = hvx_vmem(head->s_state + j * 32); + HVX_Vector vsu = hvx_vmem(head->s_update_f32 + j * 32); + hvx_vmem(head->s_state + j * 32) = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vs_old, v_l_final), vsu); + } + + if (c + 1 < n_chunks) { + gdn_pack_s_col_tiles(head->s_col_tiles, head->s_f16, head->s_state, S_v); + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_GDN_STATE, info); +} + + +static int gated_delta_net_f32_hmx_chunked( + struct htp_ops_context * octx, + const struct htp_gdn_kernel_params * kparams, + uint32_t row_start, + uint32_t nrows +) { + const struct htp_tensor * q = octx->src[0]; + const struct htp_tensor * k = octx->src[1]; + const struct htp_tensor * v = octx->src[2]; + const struct htp_tensor * g = octx->src[3]; + const struct htp_tensor * beta = octx->src[4]; + const struct htp_tensor * state = octx->src[5]; + const struct htp_tensor * dst = octx->dst; + const struct htp_tensor * dst_cache = octx->dsts[1]; + + const uint32_t S_v = kparams->S_v; + const uint32_t H = kparams->H; + const uint32_t n_tokens = kparams->n_tokens; + const float scale = kparams->scale; + const uint32_t chunk_size = kparams->chunk_size; + const uint32_t n_chunks = kparams->n_chunks; + const uint32_t n_sv_tiles = S_v / 32; + + struct htp_gdn_hmx_vtcm_layout L; + htp_gdn_hmx_vtcm_layout_build(&L, S_v, chunk_size, kparams->n_heads_batch, kparams->n_threads, kparams->pipeline != 0); + + if (L.total_bytes > octx->ctx->vtcm_size) { + return HTP_STATUS_VTCM_TOO_SMALL; + } + + uint8_t * const vtcm_base = (uint8_t *) octx->ctx->vtcm_base; + + float * vtcm_g_raw[2] = { + VTCM_LAYOUT_PTR(float, vtcm_base, L.off_g_raw[0]), + L.pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L.off_g_raw[1]) : VTCM_LAYOUT_PTR(float, vtcm_base, L.off_g_raw[0]) + }; + float * vtcm_b_raw[2] = { + VTCM_LAYOUT_PTR(float, vtcm_base, L.off_b_raw[0]), + L.pipeline ? VTCM_LAYOUT_PTR(float, vtcm_base, L.off_b_raw[1]) : VTCM_LAYOUT_PTR(float, vtcm_base, L.off_b_raw[0]) + }; + + uint8_t * vtcm_scales_1 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L.off_scales_1); + hmx_init_column_scales(vtcm_scales_1, Q6_V_vsplat_R(0x3c00)); + + hmx_queue_t hmx_q = octx->ctx->hmx_queue; + dma_queue * dma_q = octx->ctx->dma[0]; + work_queue_t wp = octx->ctx->work_queue; + + struct htp_gdn_head_ptrs heads[8]; + struct htp_gdn_hmx_gemm_task gemm_tasks[8][9]; + + uint32_t n_batch = 1; + for (uint32_t r = row_start; r < row_start + nrows; r += n_batch) { + const uint32_t head_in_seq = fastmodulo(r, H, &kparams->div_H); + const uint32_t iv3 = fastdiv(r, &kparams->div_H); + const uint32_t heads_left_in_seq = H - head_in_seq; + const uint32_t heads_left_in_range = (row_start + nrows) - r; + n_batch = hex_smin((uint32_t) kparams->n_heads_batch, hex_smin(heads_left_in_seq, heads_left_in_range)); + + for (uint32_t h = 0; h < n_batch; ++h) { + gdn_init_head_ptrs(&heads[h], &L, vtcm_base, h, head_in_seq, iv3, + q, k, v, state, dst, dst_cache, kparams, S_v, H, n_tokens, chunk_size); + } + + struct htp_gdn_batch_context bctx; + bctx.heads = heads; + bctx.vtcm_g_raw = NULL; + bctx.vtcm_b_raw = NULL; + bctx.curr_buf = 0; + bctx.c = 0; + bctx.n_batch = n_batch; + bctx.S_v = S_v; + bctx.scale = scale; + bctx.octx = octx; + bctx.kparams = kparams; + + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_push(dma_q, dma_make_data(heads[h].s_state, heads[h].state_in_dma), + S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); + } + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_pop(dma_q); + } + + if (n_chunks > 0) { + work_queue_run(wp, gdn_hvx_init_state_worker, &bctx, n_batch); + + const uint32_t chunk0_tokens = hex_smin(chunk_size, n_tokens); + for (uint32_t h = 0; h < n_batch; ++h) { + gdn_dma_push_chunk_inputs(dma_q, heads[h].q_f32[0], heads[h].k_f32[0], heads[h].v_f32[0], + q, k, v, heads[h].iq3, heads[h].iq1, heads[h].ik3, heads[h].ik1, + heads[h].iv3, heads[h].iv1, 0, chunk0_tokens, S_v); + } + gdn_dma_push_chunk_gb(dma_q, vtcm_g_raw[0], vtcm_b_raw[0], g, beta, iv3, head_in_seq, 0, chunk0_tokens, n_batch); + } + + for (uint32_t c = 0; c < n_chunks; ++c) { + const uint32_t curr_buf = c & 1; + const uint32_t next_buf = (c + 1) & 1; + const uint32_t t_chunk = c * chunk_size; + + bctx.curr_buf = curr_buf; + bctx.c = c; + bctx.vtcm_g_raw = vtcm_g_raw[curr_buf]; + bctx.vtcm_b_raw = vtcm_b_raw[curr_buf]; + + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_pop(dma_q); + dma_queue_pop(dma_q); + dma_queue_pop(dma_q); + } + dma_queue_pop(dma_q); + dma_queue_pop(dma_q); + + if (c + 1 < n_chunks) { + const uint32_t next_t_chunk = (c + 1) * chunk_size; + const uint32_t next_tokens = hex_smin(chunk_size, n_tokens - next_t_chunk); + for (uint32_t h = 0; h < n_batch; ++h) { + gdn_dma_push_chunk_inputs(dma_q, heads[h].q_f32[next_buf], heads[h].k_f32[next_buf], heads[h].v_f32[next_buf], + q, k, v, heads[h].iq3, heads[h].iq1, heads[h].ik3, heads[h].ik1, + heads[h].iv3, heads[h].iv1, next_t_chunk, next_tokens, S_v); + } + gdn_dma_push_chunk_gb(dma_q, vtcm_g_raw[next_buf], vtcm_b_raw[next_buf], + g, beta, iv3, head_in_seq, next_t_chunk, next_tokens, n_batch); + } + + if (c > 0) { + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_pop(dma_q); + } + } + + work_queue_run(wp, gdn_hvx_phase1a_worker, &bctx, n_batch); + + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][0], heads[h].k_row_tiles, heads[h].k_col_tiles, heads[h].kk_tiles, 2, 2, n_sv_tiles, vtcm_scales_1); + } + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][2], heads[h].k_prime_row_tiles, heads[h].s_col_tiles, heads[h].v_inter_tiles, 2, n_sv_tiles, n_sv_tiles, vtcm_scales_1); + } + + work_queue_run(wp, gdn_hvx_phase1b_worker, &bctx, n_batch); + + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][1], heads[h].q_row_tiles, heads[h].k_col_tiles, heads[h].qk_tiles, 2, 2, n_sv_tiles, vtcm_scales_1); + } + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][3], heads[h].q_prime_row_tiles, heads[h].s_col_tiles, heads[h].o_inter_tiles, 2, n_sv_tiles, n_sv_tiles, vtcm_scales_1); + } + + work_queue_run(wp, gdn_hvx_phase2_worker, &bctx, n_batch); + + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task( + hmx_q, &gemm_tasks[h][7], + (__fp16 *) heads[h].vtcm_m, + heads[h].inv_row_tiles + 0 * HMX_FP16_TILE_N_ELMS, + (__fp16 *) heads[h].vtcm_tmp, + 1, 1, 1, vtcm_scales_1 + ); + } + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task( + hmx_q, &gemm_tasks[h][8], + (__fp16 *) heads[h].vtcm_m + HMX_FP16_TILE_N_ELMS, + (__fp16 *) heads[h].vtcm_tmp, + heads[h].inv_row_tiles + 2 * HMX_FP16_TILE_N_ELMS, + 1, 1, 1, vtcm_scales_1 + ); + } + + work_queue_run(wp, gdn_hvx_phase3_worker, &bctx, n_batch); + + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][4], heads[h].inv_row_tiles, heads[h].v_prime_col_tiles, heads[h].delta_tiles, 2, n_sv_tiles, 2, vtcm_scales_1); + } + + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + + work_queue_run(wp, gdn_hvx_phase4_worker, &bctx, n_batch); + + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][5], heads[h].a_row_tiles, heads[h].delta_col_tiles, heads[h].o_intra_tiles, 2, n_sv_tiles, 2, vtcm_scales_1); + } + for (uint32_t h = 0; h < n_batch; ++h) { + htp_gdn_push_hmx_gemm_task(hmx_q, &gemm_tasks[h][6], heads[h].d_row_tiles, heads[h].k_col_tiles_64x128, heads[h].s_update_tiles, n_sv_tiles, n_sv_tiles, 2, vtcm_scales_1); + } + + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + + work_queue_run(wp, gdn_hvx_phase5_worker, &bctx, n_batch); + + const uint32_t valid_tokens = hex_smin(chunk_size, n_tokens - t_chunk); + for (uint32_t h = 0; h < n_batch; ++h) { + const dma_addr_t attn_chunk_dma = dst->data + + ((uint64_t) heads[h].iv3 * n_tokens * H + (uint64_t) t_chunk * H + heads[h].iv1) * S_v * sizeof(float); + dma_queue_push(dma_q, dma_make_data(attn_chunk_dma, heads[h].o_f32[curr_buf]), + dst->nb[1], S_v * sizeof(float), S_v * sizeof(float), valid_tokens); + } + + for (uint32_t h = 0; h < n_batch; ++h) { + hmx_queue_pop(hmx_q); + } + + work_queue_run(wp, gdn_hvx_phase6_worker, &bctx, n_batch); + } + + if (n_chunks > 0) { + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_pop(dma_q); + } + } + + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_push(dma_q, dma_make_data(heads[h].state_out_dma, heads[h].s_state), + S_v * sizeof(float), S_v * sizeof(float), S_v * sizeof(float), S_v); + } + for (uint32_t h = 0; h < n_batch; ++h) { + dma_queue_pop(dma_q); + } + } + + dma_queue_flush(dma_q); + return HTP_STATUS_OK; +} + int op_gated_delta_net(struct htp_ops_context * octx) { const struct htp_tensor * q = octx->src[0]; const struct htp_tensor * k = octx->src[1]; @@ -1097,11 +2344,34 @@ int op_gated_delta_net(struct htp_ops_context * octx) { kparams_local.K = K; kparams_local.total_rows = total_rows; kparams_local.rows_per_thread = (total_rows + n_threads - 1) / n_threads; - struct htp_gdn_vtcm_layout layout_local; - htp_gdn_vtcm_layout_build(&layout_local, S_v, n_threads); - kparams_local.state_aligned = (uint32_t) layout_local.state_aligned; - kparams_local.vtcm_per_thread = (uint32_t) layout_local.bytes_per_thread; - kparams_local.vtcm_size = (uint32_t) layout_local.total_bytes; + const bool can_use_hmx = (octx->ctx->hmx_enabled) && + (S_v % 64 == 0) && + (n_tokens >= HTP_GDN_MIN_TOKENS) && + (g->ne[0] == 1) && + (K == 1); + + struct htp_gdn_hmx_vtcm_layout hmx_layout_local; + struct htp_gdn_vtcm_layout hvx_layout_local; + uint32_t n_heads_batch = 1; + + if (can_use_hmx && htp_gdn_hmx_solve_layout(&hmx_layout_local, S_v, HTP_GDN_CHUNK_SIZE, total_rows, octx->ctx->vtcm_size, n_threads, true, &n_heads_batch)) { + kparams_local.kernel_type = HTP_GDN_KERNEL_HMX_CHUNKED; + kparams_local.pipeline = hmx_layout_local.pipeline ? 1 : 0; + kparams_local.chunk_size = HTP_GDN_CHUNK_SIZE; + kparams_local.n_chunks = (n_tokens + HTP_GDN_CHUNK_SIZE - 1) / HTP_GDN_CHUNK_SIZE; + kparams_local.n_heads_batch = (uint16_t) n_heads_batch; + kparams_local.vtcm_size = (uint32_t) hmx_layout_local.total_bytes; + kparams_local.state_aligned = (uint32_t) hmx_layout_local.state_f32_bytes; + kparams_local.vtcm_per_thread = (uint32_t) (hmx_layout_local.total_bytes / (n_threads > 0 ? n_threads : 1)); + } else { + htp_gdn_vtcm_layout_build(&hvx_layout_local, S_v, n_threads); + kparams_local.kernel_type = HTP_GDN_KERNEL_HVX_RECURRENT; + kparams_local.pipeline = 0; + kparams_local.n_heads_batch = 1; + kparams_local.state_aligned = (uint32_t) hvx_layout_local.state_aligned; + kparams_local.vtcm_per_thread = (uint32_t) hvx_layout_local.bytes_per_thread; + kparams_local.vtcm_size = (uint32_t) hvx_layout_local.total_bytes; + } kparams_local.kda = (g->ne[0] == S_v) ? 1 : 0; kparams_local.scale = 1.0f / sqrtf((float) S_v); kparams_local.state_seq_stride = (uint32_t) (state->nb[3] / sizeof(float)); @@ -1121,7 +2391,19 @@ int op_gated_delta_net(struct htp_ops_context * octx) { uint32_t row_start = 0; uint32_t nrows = total_rows; - if (octx->op_params[1] != 0) { + if (octx->ctx->mdev.count > 1) { + const bool can_split = htp_tensor_mdev_data_aligned(dst) && + ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition( + total_rows, + can_split ? 1 : 0, + octx->ctx->mdev.idx, + octx->ctx->mdev.count, + &octx->ctx->mdev.count_div + ); + row_start = range.start; + nrows = range.count; + } else if (octx->op_params[1] != 0) { row_start = octx->op_params[1]; nrows = octx->op_params[2]; } @@ -1130,6 +2412,10 @@ int op_gated_delta_net(struct htp_ops_context * octx) { return HTP_STATUS_OK; } + if (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) { + return gated_delta_net_f32_hmx_chunked(octx, kparams, row_start, nrows); + } + const uint32_t n_threads = (nrows < kparams->n_threads) ? nrows : kparams->n_threads; struct htp_gdn_context gctx; diff --git a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h index fd703142e..32fb7d24b 100644 --- a/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h +++ b/ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h @@ -11,6 +11,7 @@ #define HTP_GDN_MAX_SV 128 #define HTP_GDN_CHUNK_SIZE 64 +#define HTP_GDN_MIN_TOKENS 8 #ifndef HMX_FP16_TILE_SIZE #define HMX_FP16_TILE_SIZE 2048 @@ -132,7 +133,6 @@ struct htp_gdn_hmx_vtcm_layout { size_t off_rows_a; size_t off_thread_scratch; - size_t off_attn_rem; size_t off_scales_1; size_t state_f32_bytes; @@ -192,8 +192,10 @@ static inline void htp_gdn_hmx_vtcm_layout_build( VTCM_LAYOUT_ALLOC(off, off_s_state, bh * state_f32_sz); VTCM_LAYOUT_ALLOC(off, off_s_f16, bh * state_f16_sz); + off = hex_align_up(off, HMX_FP16_TILE_SIZE); VTCM_LAYOUT_ALLOC(off, off_s_col_tiles, bh * state_tiles_sz); VTCM_LAYOUT_ALLOC(off, off_s_update_f32, bh * state_f32_sz); + off = hex_align_up(off, HMX_FP16_TILE_SIZE); VTCM_LAYOUT_ALLOC(off, off_s_update_tiles, bh * state_tiles_sz); VTCM_LAYOUT_ALLOC(off, off_q_f32[0], bh * dma_chunk_sz); @@ -222,6 +224,7 @@ static inline void htp_gdn_hmx_vtcm_layout_build( VTCM_LAYOUT_ALLOC(off, off_delta_f16, bh * act_f16_sz); VTCM_LAYOUT_ALLOC(off, off_d_f16, bh * act_f16_sz); + off = hex_align_up(off, HMX_FP16_TILE_SIZE); VTCM_LAYOUT_ALLOC(off, off_q_row_tiles, bh * tile_64xSv_sz); VTCM_LAYOUT_ALLOC(off, off_q_prime_row_tiles, bh * tile_64xSv_sz); VTCM_LAYOUT_ALLOC(off, off_k_row_tiles, bh * tile_64xSv_sz); @@ -250,9 +253,10 @@ static inline void htp_gdn_hmx_vtcm_layout_build( VTCM_LAYOUT_ALLOC(off, off_rows_a, bh * row_vecs_sz); const size_t thread_scratch_sz = 64 * 128; + off = hex_align_up(off, HMX_FP16_TILE_SIZE); VTCM_LAYOUT_ALLOC(off, off_thread_scratch, nth * thread_scratch_sz); - VTCM_LAYOUT_ALLOC(off, off_attn_rem, nth * (128 * sizeof(float))); - VTCM_LAYOUT_ALLOC(off, off_scales_1, 256); + off = hex_align_up(off, HMX_FP16_TILE_SIZE); + VTCM_LAYOUT_ALLOC(off, off_scales_1, HMX_FP16_TILE_SIZE); L->total_bytes = off; } diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h index 8fd299795..d5fb48ad9 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h +++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h @@ -48,7 +48,7 @@ static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) = }; // Inner HMX tile computation kernels -static void hmx_fa_qk_dot_tile( +static inline void hmx_fa_qk_dot_tile( const __fp16 * row_tiles, const __fp16 * col_tiles, __fp16 * out_tile, @@ -116,7 +116,7 @@ static void hmx_fa_qk_dot_tile( ); } -static void hmx_fa_o_update_tile( +static inline void hmx_fa_o_update_tile( const __fp16 * d_diag, const __fp16 * o_rc, const __fp16 * p_tile_in, diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index 0e63febdd..ee5b92441 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -204,6 +204,14 @@ enum htp_trace_event_id { HTP_TRACE_EVT_HVX_FA_K_PREP = 29, HTP_TRACE_EVT_HVX_FA_V_PREP = 30, + HTP_TRACE_EVT_HVX_GDN_PREP = 31, + HTP_TRACE_EVT_HVX_GDN_SOLVE = 32, + HTP_TRACE_EVT_HVX_GDN_V_PREP = 33, + HTP_TRACE_EVT_HVX_GDN_D_PREP = 34, + HTP_TRACE_EVT_HVX_GDN_OUT = 35, + HTP_TRACE_EVT_HVX_GDN_STATE = 36, + HTP_TRACE_EVT_HVX_GDN_REM = 37, + HTP_TRACE_EVT_HMX_COMP = 40, }; diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 653c9a250..b4b352b20 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -36,7 +36,7 @@ #include "allreduce-ops.h" #include "htp-fence.h" -#define HMX_QUEUE_CAPACITY 16 +#define HMX_QUEUE_CAPACITY 128 #define HMX_QUEUE_STACK_SIZE 16384 #define WORK_QUEUE_CAPACITY 16 #define WORK_QUEUE_STACK_SIZE 16384 diff --git a/scripts/snapdragon/ggml-hexagon-inspect.py b/scripts/snapdragon/ggml-hexagon-inspect.py index 3afda8a09..c977f6a17 100755 --- a/scripts/snapdragon/ggml-hexagon-inspect.py +++ b/scripts/snapdragon/ggml-hexagon-inspect.py @@ -36,12 +36,13 @@ import signal import subprocess import sys from pathlib import Path -from typing import Dict, List, NamedTuple, Optional, Tuple +from typing import Dict, List, NamedTuple, Optional, Set, Tuple # Ignore SIGPIPE to handle pipes (e.g. head, grep) gracefully if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL) +logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) logger = logging.getLogger("ggml-hexagon-inspect") @@ -56,6 +57,37 @@ class InsnInfo(NamedTuple): in_loop: bool +class LoopStats: + def __init__(self, loop_type: str, start_addr: int, end_addr: Optional[int] = None, loop_id: int = 0): + self.loop_id = loop_id + self.loop_type = loop_type # "loop0" or "loop1" + self.start_addr = start_addr + self.end_addr = end_addr + self.packet_count = 0 + self.insn_count = 0 + self.vec_insn_count = 0 + self.vspills_st = 0 + self.vspills_ld = 0 + self.sspills_st = 0 + self.sspills_ld = 0 + + @property + def vspills_total(self) -> int: + return self.vspills_st + self.vspills_ld + + @property + def sspills_total(self) -> int: + return self.sspills_st + self.sspills_ld + + @property + def has_v_roundtrip(self) -> bool: + return self.vspills_st > 0 and self.vspills_ld > 0 + + @property + def vec_density(self) -> float: + return (self.vec_insn_count / self.packet_count) if self.packet_count > 0 else 0.0 + + class FuncStats: def __init__(self, name: str, address: int, size: int): self.name = name @@ -66,14 +98,19 @@ class FuncStats: self.vec_insn_count = 0 self.loop_count = 0 self.vspills_in_loop = 0 + self.vspills_in_loop_st = 0 + self.vspills_in_loop_ld = 0 self.vspills_total = 0 self.sspills_in_loop = 0 + self.sspills_in_loop_st = 0 + self.sspills_in_loop_ld = 0 self.sspills_total = 0 self.promotions_in_loop = 0 self.promotions_total = 0 self.promotion_targets: Dict[str, int] = {} self.calls_in_loop = 0 self.calls_total = 0 + self.loops: List[LoopStats] = [] self.insns: List[InsnInfo] = [] @@ -90,16 +127,43 @@ RE_INSN_LINE = re.compile( ) RE_LOOP0_START = re.compile(r"\bloop0\((0x[0-9a-fA-F]+)") RE_LOOP1_START = re.compile(r"\bloop1\((0x[0-9a-fA-F]+)") -RE_VSPILL = re.compile(r"\bvmemu?\s*\(\s*r(?:29|30)\b") -RE_SSPILL = re.compile(r"\bmem[bwhd]\s*\(\s*r(?:29|30)\b") +RE_VMEM_BASE = re.compile(r"\bvmemu?\s*\(\s*([a-z0-9]+)\b") +RE_SMEM_BASE = re.compile(r"\bmem[bwhd](?:_locked|_fifo)?\s*\(\s*([a-z0-9]+)\b") +RE_MEM_STORE = re.compile(r"\bv?mem[bwhdu]?(?:_[a-z]+)?\s*\([^)]*\)\s*(\+|-)?=") +RE_ADD_OP = re.compile(r"\b(r[0-9]+)\s*=\s*add\s*\(\s*([^,()]+)\s*,\s*([^,()]+)\s*\)") +RE_ASSIGN_LHS = re.compile(r"^\s*(?:if\s*\([^)]+\)\s*)?(r[0-9]+)(?::(r[0-9]+))?\s*(?:[+\-*/&|^]?=)") RE_VEC_OP = re.compile(r"\b(v[0-9]+|w[0-9]+|q[0-3]|vmemu?)\b") -RE_STORE = re.compile(r"=\s*(?:v[0-9]|r[0-9]|w[0-9]|#)") RE_PROMOTION_CALL = re.compile( r"\b(?:call|jump)\s+(?:0x[0-9a-fA-F]+\s+)?<(__(?:trunc|extend)[a-zA-Z0-9_]+)(?:@plt)?>" ) RE_ANY_CALL = re.compile(r"\bcallr?\b") +def is_mem_store(insn: str) -> bool: + return bool(RE_MEM_STORE.search(insn)) + + +def update_sp_regs(insn: str, sp_regs: Set[str]) -> None: + # Track registers derived from stack frame (r29/r30) + m_add = RE_ADD_OP.search(insn) + if m_add: + dest = m_add.group(1) + op1 = m_add.group(2).strip() + op2 = m_add.group(3).strip() + if op1 in sp_regs or op2 in sp_regs: + sp_regs.add(dest) + return + + m_assign = RE_ASSIGN_LHS.match(insn.strip()) + if m_assign: + r1 = m_assign.group(1) + r2 = m_assign.group(2) + if r1 and r1 not in ("r29", "r30"): + sp_regs.discard(r1) + if r2 and r2 not in ("r29", "r30"): + sp_regs.discard(r2) + + def get_repo_root() -> Path: # Resolve repository root from script location return Path(__file__).resolve().parent.parent.parent @@ -338,13 +402,15 @@ def parse_disassembly( end_idx = matches[i + 1].start() if i + 1 < len(matches) else len(disasm_text) chunk = disasm_text[start_idx:end_idx] - # Calculate rough byte size from line addresses stats = FuncStats(name=name, address=addr, size=0) loop0_target: Optional[int] = None loop1_target: Optional[int] = None loop0_active = False loop1_active = False + current_loop0: Optional[LoopStats] = None + current_loop1: Optional[LoopStats] = None + sp_regs: Set[str] = {"r29", "r30"} first_addr = None last_addr = None @@ -364,6 +430,10 @@ def parse_disassembly( # Track packet count if "{" in asm_chunk: stats.packet_count += 1 + if current_loop0: + current_loop0.packet_count += 1 + if current_loop1: + current_loop1.packet_count += 1 # Check loop starts m0 = RE_LOOP0_START.search(asm_chunk) @@ -378,8 +448,23 @@ def parse_disassembly( if loop0_target is not None and cur_addr >= loop0_target: loop0_active = True + if current_loop0 is None: + current_loop0 = LoopStats( + loop_id=len(stats.loops) + 1, + loop_type="loop0", + start_addr=loop0_target, + end_addr=0, + ) + if loop1_target is not None and cur_addr >= loop1_target: loop1_active = True + if current_loop1 is None: + current_loop1 = LoopStats( + loop_id=len(stats.loops) + 1, + loop_type="loop1", + start_addr=loop1_target, + end_addr=0, + ) in_loop = loop0_active or loop1_active @@ -388,31 +473,70 @@ def parse_disassembly( sub_insns = [p.strip() for p in cleaned.split(";") if p.strip()] for insn in sub_insns: + update_sp_regs(insn, sp_regs) + stats.insn_count += 1 + if current_loop0: + current_loop0.insn_count += 1 + if current_loop1: + current_loop1.insn_count += 1 + is_vec = bool(RE_VEC_OP.search(insn)) if is_vec: stats.vec_insn_count += 1 + if current_loop0: + current_loop0.vec_insn_count += 1 + if current_loop1: + current_loop1.vec_insn_count += 1 - is_vspill = bool(RE_VSPILL.search(insn)) - is_sspill = bool(RE_SSPILL.search(insn)) + vm = RE_VMEM_BASE.search(insn) + is_vspill = bool(vm and vm.group(1) in sp_regs) + + sm = RE_SMEM_BASE.search(insn) + is_sspill = bool(sm and sm.group(1) in sp_regs) - # Identify store vs load is_store = False is_load = False if is_vspill or is_sspill: - if RE_STORE.search(insn): - is_store = True - else: - is_load = True + is_store = is_mem_store(insn) + is_load = not is_store if is_vspill: stats.vspills_total += 1 if in_loop: stats.vspills_in_loop += 1 + if is_store: + stats.vspills_in_loop_st += 1 + else: + stats.vspills_in_loop_ld += 1 + if current_loop0: + if is_store: + current_loop0.vspills_st += 1 + else: + current_loop0.vspills_ld += 1 + if current_loop1: + if is_store: + current_loop1.vspills_st += 1 + else: + current_loop1.vspills_ld += 1 elif is_sspill: stats.sspills_total += 1 if in_loop: stats.sspills_in_loop += 1 + if is_store: + stats.sspills_in_loop_st += 1 + else: + stats.sspills_in_loop_ld += 1 + if current_loop0: + if is_store: + current_loop0.sspills_st += 1 + else: + current_loop0.sspills_ld += 1 + if current_loop1: + if is_store: + current_loop1.sspills_st += 1 + else: + current_loop1.sspills_ld += 1 is_call = bool(RE_ANY_CALL.search(insn)) prom_m = RE_PROMOTION_CALL.search(insn) @@ -444,9 +568,29 @@ def parse_disassembly( if ":endloop0" in asm_chunk: loop0_active = False loop0_target = None + if current_loop0: + current_loop0.end_addr = cur_addr + stats.loops.append(current_loop0) + current_loop0 = None + if ":endloop1" in asm_chunk: loop1_active = False loop1_target = None + if current_loop1: + current_loop1.end_addr = cur_addr + stats.loops.append(current_loop1) + current_loop1 = None + + if current_loop0: + current_loop0.end_addr = last_addr or 0 + stats.loops.append(current_loop0) + if current_loop1: + current_loop1.end_addr = last_addr or 0 + stats.loops.append(current_loop1) + + stats.loops.sort(key=lambda lp: lp.start_addr) + for idx, loop in enumerate(stats.loops, 1): + loop.loop_id = idx if first_addr is not None and last_addr is not None: stats.size = (last_addr - first_addr) + 4 @@ -463,14 +607,19 @@ def annotate_disasm_line( loop0_active: bool, loop1_active: bool, use_color: bool = True, -) -> Tuple[str, Optional[int], Optional[int], bool, bool]: + sp_regs: Optional[Set[str]] = None, +) -> Tuple[str, Optional[int], Optional[int], bool, bool, bool]: # Annotate disassembly line with spill and loop tags lm = RE_INSN_LINE.match(raw_line) if not lm: - return raw_line, loop0_target, loop1_target, loop0_active, loop1_active + return raw_line, loop0_target, loop1_target, loop0_active, loop1_active, False cur_addr = int(lm.group(1), 16) asm_chunk = lm.group(4) + is_event = False + + if sp_regs is None: + sp_regs = {"r29", "r30"} # Check loop starts m0 = RE_LOOP0_START.search(asm_chunk) @@ -490,39 +639,75 @@ def annotate_disasm_line( tags = [] if m0: tags.append("[LOOP0-START]") + is_event = True if m1: tags.append("[LOOP1-START]") + is_event = True - if RE_VSPILL.search(asm_chunk): - if in_loop: - tags.append("[V-SPILL:IN-LOOP]" if not use_color else "\033[1;31m[V-SPILL:IN-LOOP]\033[0m") - else: - tags.append("[V-SPILL]" if not use_color else "\033[1;33m[V-SPILL]\033[0m") - elif RE_SSPILL.search(asm_chunk): - if in_loop: - tags.append("[S-SPILL:IN-LOOP]" if not use_color else "\033[1;35m[S-SPILL:IN-LOOP]\033[0m") + cleaned = re.sub(r"[{}\s]|:endloop[01]", " ", asm_chunk) + sub_insns = [p.strip() for p in cleaned.split(";") if p.strip()] + + for insn in sub_insns: + update_sp_regs(insn, sp_regs) + + for insn in sub_insns: + vm = RE_VMEM_BASE.search(insn) + if vm and vm.group(1) in sp_regs: + base = vm.group(1) + is_st = is_mem_store(insn) + op = "STORE" if is_st else "LOAD" + tgt = f"({base})" if base not in ("r29", "r30") else "" + if in_loop: + tag = f"[V-SPILL:{op}{tgt}:IN-LOOP]" + tags.append(f"\033[1;31m{tag}\033[0m" if use_color else tag) + else: + tag = f"[V-SPILL:{op}{tgt}]" + tags.append(f"\033[1;33m{tag}\033[0m" if use_color else tag) + is_event = True + + sm = RE_SMEM_BASE.search(insn) + if sm and sm.group(1) in sp_regs: + base = sm.group(1) + is_st = is_mem_store(insn) + op = "STORE" if is_st else "LOAD" + tgt = f"({base})" if base not in ("r29", "r30") else "" + if in_loop: + tag = f"[S-SPILL:{op}{tgt}:IN-LOOP]" + tags.append(f"\033[1;35m{tag}\033[0m" if use_color else tag) + else: + tag = f"[S-SPILL:{op}{tgt}]" + tags.append(f"\033[0;35m{tag}\033[0m" if use_color else tag) + is_event = True prom_m = RE_PROMOTION_CALL.search(asm_chunk) if prom_m: ptarget = prom_m.group(1) if in_loop: - tags.append(f"[PROMOTION:{ptarget}:IN-LOOP]" if not use_color else f"\033[1;31m[PROMOTION:{ptarget}:IN-LOOP]\033[0m") + tag = f"[PROMOTION:{ptarget}:IN-LOOP]" + tags.append(f"\033[1;31m{tag}\033[0m" if use_color else tag) else: - tags.append(f"[PROMOTION:{ptarget}]" if not use_color else f"\033[1;35m[PROMOTION:{ptarget}]\033[0m") + tag = f"[PROMOTION:{ptarget}]" + tags.append(f"\033[1;35m{tag}\033[0m" if use_color else tag) + is_event = True elif RE_ANY_CALL.search(asm_chunk): if in_loop: - tags.append("[CALL:IN-LOOP]" if not use_color else "\033[1;31m[CALL:IN-LOOP]\033[0m") + tag = "[CALL:IN-LOOP]" + tags.append(f"\033[1;31m{tag}\033[0m" if use_color else tag) + is_event = True else: - tags.append("[CALL]" if not use_color else "\033[1;36m[CALL]\033[0m") + tag = "[CALL]" + tags.append(f"\033[1;36m{tag}\033[0m" if use_color else tag) if ":endloop0" in asm_chunk: tags.append("[LOOP0-END]") loop0_active = False loop0_target = None + is_event = True if ":endloop1" in asm_chunk: tags.append("[LOOP1-END]") loop1_active = False loop1_target = None + is_event = True tag_str = " ".join(tags) if tag_str: @@ -530,7 +715,7 @@ def annotate_disasm_line( else: annotated = raw_line - return annotated, loop0_target, loop1_target, loop0_active, loop1_active + return annotated, loop0_target, loop1_target, loop0_active, loop1_active, is_event def run_spills( @@ -566,14 +751,15 @@ def run_spills( col_pkts = "Packets" col_insn = "Insns" col_vec = "HVX Ops" - col_vloop = "V-Loop" + col_vloop = "V-Loop (st/ld)" col_vtot = "V-Tot" - col_sloop = "S-Loop" + col_sloop = "S-Loop (st/ld)" col_stot = "S-Tot" + col_notes = "Notes" hdr = ( - f"{col_addr:<10} | {col_name:<44} | {col_pkts:>7} | {col_insn:>6} | " - f"{col_vec:>7} | {col_vloop:>6} | {col_vtot:>5} | {col_sloop:>6} | {col_stot:>5}" + f"{col_addr:<10} | {col_name:<40} | {col_pkts:>7} | {col_insn:>6} | " + f"{col_vec:>7} | {col_vloop:>14} | {col_vtot:>5} | {col_sloop:>14} | {col_stot:>5} | {col_notes}" ) sep = "-" * len(hdr) @@ -596,9 +782,11 @@ def run_spills( # Check strict criteria if args.strict: - if f.vspills_in_loop > args.max_inloop_vspills: + inloop_v = f.vspills_in_loop_st if getattr(args, "strict_stores_only", False) else f.vspills_in_loop + if inloop_v > args.max_inloop_vspills: + lbl = "in-loop vector store spills" if getattr(args, "strict_stores_only", False) else "in-loop vector spills" strict_violations.append( - f"{f.name}: {f.vspills_in_loop} in-loop vector spills (max allowed: {args.max_inloop_vspills})" + f"{f.name}: {inloop_v} {lbl} (max allowed: {args.max_inloop_vspills})" ) if dma_re and dma_re.search(f.name): if f.vec_insn_count > args.max_dma_vec_ops: @@ -606,14 +794,27 @@ def run_spills( f"{f.name}: DMA worker contains {f.vec_insn_count} HVX vector ops (max allowed: {args.max_dma_vec_ops})" ) - # Highlight in-loop vector spills - vloop_str = f"{f.vspills_in_loop:>6}" + vloop_detail = f"{f.vspills_in_loop} ({f.vspills_in_loop_st}s,{f.vspills_in_loop_ld}l)" if f.vspills_in_loop > 0 else "0" + sloop_detail = f"{f.sspills_in_loop} ({f.sspills_in_loop_st}s,{f.sspills_in_loop_ld}l)" if f.sspills_in_loop > 0 else "0" + + notes = "" + if f.vspills_in_loop_st > 0 and f.vspills_in_loop_ld > 0: + notes = "\033[1;31m[V-ROUNDTRIP!]\033[0m" if use_color else "[V-ROUNDTRIP!]" + elif f.vspills_in_loop_st == 0 and f.vspills_in_loop_ld > 0: + notes = "v-readonly" + + vloop_str = f"{vloop_detail:>14}" if f.vspills_in_loop > 0 and use_color: - vloop_str = f"\033[1;31m{vloop_str}\033[0m" + if f.vspills_in_loop_st > 0 and f.vspills_in_loop_ld > 0: + vloop_str = f"\033[1;31m{vloop_str}\033[0m" + else: + vloop_str = f"\033[1;33m{vloop_str}\033[0m" + + sloop_str = f"{sloop_detail:>14}" logger.info( - f"0x{f.address:08x} | {f.name:<44} | {f.packet_count:>7} | {f.insn_count:>6} | " - f"{f.vec_insn_count:>7} | {vloop_str} | {f.vspills_total:>5} | {f.sspills_in_loop:>6} | {f.sspills_total:>5}" + f"0x{f.address:08x} | {f.name:<40} | {f.packet_count:>7} | {f.insn_count:>6} | " + f"{f.vec_insn_count:>7} | {vloop_str} | {f.vspills_total:>5} | {sloop_str} | {f.sspills_total:>5} | {notes}" ) logger.info(sep) @@ -744,7 +945,7 @@ def run_disasm( args: argparse.Namespace, ) -> int: # Disassemble matching function(s) with annotated loop and spill markers - func_pattern = args.disasm + func_pattern = args.disasm if args.disasm else (args.func or ".*") logger.info(f"Inspecting library: {lib_path}") logger.info(f"Disassembling functions matching: '{func_pattern}'\n") @@ -794,9 +995,11 @@ def run_disasm( logger.info(f"Packets: {func_stats.packet_count} | Instructions: {func_stats.insn_count} | Loops: {func_stats.loop_count}") vec_pct = (func_stats.vec_insn_count / func_stats.insn_count * 100.0) if func_stats.insn_count else 0.0 logger.info(f"HVX Ops: {func_stats.vec_insn_count} ({vec_pct:.1f}% of instructions)") + vloop_info = f"{func_stats.vspills_in_loop} ({func_stats.vspills_in_loop_st} st, {func_stats.vspills_in_loop_ld} ld)" + sloop_info = f"{func_stats.sspills_in_loop} ({func_stats.sspills_in_loop_st} st, {func_stats.sspills_in_loop_ld} ld)" logger.info( - f"Spills: Vector in-loop: {func_stats.vspills_in_loop} | Vector total: {func_stats.vspills_total} | " - f"Scalar in-loop: {func_stats.sspills_in_loop} | Scalar total: {func_stats.sspills_total}" + f"Spills: Vector in-loop: {vloop_info} | Vector total: {func_stats.vspills_total} | " + f"Scalar in-loop: {sloop_info} | Scalar total: {func_stats.sspills_total}" ) logger.info( f"Calls: Total: {func_stats.calls_total} (in-loop: {func_stats.calls_in_loop}) | " @@ -804,18 +1007,76 @@ def run_disasm( ) logger.info(hdr_border) - # Log annotated disassembly - loop0_target: Optional[int] = None - loop1_target: Optional[int] = None + # Print Loop Breakdown Table if function has loops + if func_stats.loops: + logger.info(f"\n--- Loops ({len(func_stats.loops)}) " + "-" * 67) + loop_hdr = ( + f"{'#':<3} | {'Type':<5} | {'Address Range':<25} | {'Packets':>7} | " + f"{'HVX Ops':>7} | {'Vec/Pkt':>7} | {'V-Spills (st, ld)':>17} | {'S-Spills (st, ld)':>17} | Notes" + ) + logger.info(loop_hdr) + logger.info("-" * len(loop_hdr)) + for loop in func_stats.loops: + vspill_str = f"{loop.vspills_total} ({loop.vspills_st}s,{loop.vspills_ld}l)" + sspill_str = f"{loop.sspills_total} ({loop.sspills_st}s,{loop.sspills_ld}l)" + notes = [] + if loop.has_v_roundtrip: + notes.append("\033[1;31m[V-ROUNDTRIP!]\033[0m" if use_color else "[V-ROUNDTRIP!]") + elif loop.vspills_st == 0 and loop.vspills_ld > 0: + notes.append("v-readonly") + if loop.vec_density >= 1.5: + notes.append("\033[1;32mdual-hvx\033[0m" if use_color else "dual-hvx") + notes_str = ", ".join(notes) + logger.info( + f"{loop.loop_id:<3} | {loop.loop_type:<5} | 0x{loop.start_addr:08x} - 0x{loop.end_addr:08x} | " + f"{loop.packet_count:>7} | {loop.vec_insn_count:>7} | {loop.vec_density:>7.2f} | " + f"{vspill_str:>17} | {sspill_str:>17} | {notes_str}" + ) + logger.info("-" * len(loop_hdr) + "\n") + + # Parse lines and annotations + lines = chunk.splitlines() + annotated_lines = [] + is_event_list = [] + loop0_target = None + loop1_target = None loop0_active = False loop1_active = False + sp_regs = {"r29", "r30"} - for line in chunk.splitlines(): - ann_line, loop0_target, loop1_target, loop0_active, loop1_active = annotate_disasm_line( - line, loop0_target, loop1_target, loop0_active, loop1_active, use_color + for line in lines: + ann_line, loop0_target, loop1_target, loop0_active, loop1_active, is_ev = annotate_disasm_line( + line, loop0_target, loop1_target, loop0_active, loop1_active, use_color, sp_regs ) - logger.info(ann_line) - logger.info("") + annotated_lines.append(ann_line) + is_event_list.append(is_ev) + + # Filter output if --spills-only + if getattr(args, "spills_only", False): + ctx = args.context if args.context is not None else 2 + to_show = [False] * len(annotated_lines) + for idx, ev in enumerate(is_event_list): + if ev: + for j in range(max(0, idx - ctx), min(len(annotated_lines), idx + ctx + 1)): + to_show[j] = True + + if not any(to_show): + logger.info(" (No spills, promotions, or in-loop calls detected in this function)\n") + else: + in_gap = False + for idx, show in enumerate(to_show): + if show: + in_gap = False + logger.info(annotated_lines[idx]) + else: + if not in_gap: + logger.info(" ...") + in_gap = True + logger.info("") + else: + for ann_line in annotated_lines: + logger.info(ann_line) + logger.info("") return 0 @@ -964,9 +1225,24 @@ def main(): ) parser.add_argument( "--disasm", + nargs="?", + const="", metavar="FUNC", help="Disassemble function symbol or regex pattern with annotated loop and spill markers.", ) + parser.add_argument( + "--spills-only", + action="store_true", + help="In --disasm, only display packets containing spills, promotions, or in-loop calls, with surrounding context.", + ) + parser.add_argument( + "-C", + "--context", + type=int, + default=None, + metavar="N", + help="Number of context packets before and after spills in --disasm --spills-only (default: 2).", + ) parser.add_argument( "--limit", type=int, @@ -983,8 +1259,9 @@ def main(): # Filtering & Display parser.add_argument( "--func", + "--fn", "-f", - help="Regex filter for function names in --spills or --promotions.", + help="Regex filter for function names in --spills, --promotions, or --disasm.", ) parser.add_argument( "--all", @@ -1010,6 +1287,11 @@ def main(): default=0, help="Maximum allowed in-loop vector spills in --strict mode (default: 0).", ) + parser.add_argument( + "--strict-stores-only", + action="store_true", + help="In --strict mode, only count vector store spills (st > 0) towards violations, ignoring readonly stack loads.", + ) parser.add_argument( "--max-dma-vec-ops", type=int, @@ -1057,7 +1339,7 @@ def main(): args = parser.parse_args() - logging.basicConfig(level=logging.INFO, format="%(message)s") + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) repo_root = get_repo_root() @@ -1092,7 +1374,7 @@ def main(): # Dispatch commands if args.addr2line is not None: sys.exit(run_addr2line(toolchain, lib_path, args)) - elif args.disasm: + elif args.disasm is not None: sys.exit(run_disasm(toolchain, lib_path, args)) elif args.promotions: sys.exit(run_promotions(toolchain, lib_path, args)) @@ -1102,5 +1384,5 @@ def main(): if __name__ == "__main__": - logging.basicConfig(level=logging.INFO, format="%(message)s") + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) main() diff --git a/scripts/snapdragon/ggml-hexagon-profile.py b/scripts/snapdragon/ggml-hexagon-profile.py index 48b3fe479..4ac227678 100755 --- a/scripts/snapdragon/ggml-hexagon-profile.py +++ b/scripts/snapdragon/ggml-hexagon-profile.py @@ -54,6 +54,7 @@ def device_matches(record_device, target_device): return False +logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) logger = logging.getLogger("ggml-hexagon-profile") @@ -648,7 +649,7 @@ def main(): args = parser.parse_args() - logging.basicConfig(level=logging.INFO, format='%(message)s') + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) if "pmu" in args.sort and args.pmu_index is None: logger.error(f"Cannot sort by '{args.sort}' without --pmu-index.") diff --git a/scripts/snapdragon/ggml-hexagon-trace.py b/scripts/snapdragon/ggml-hexagon-trace.py index 99bf771b8..760eb57d9 100755 --- a/scripts/snapdragon/ggml-hexagon-trace.py +++ b/scripts/snapdragon/ggml-hexagon-trace.py @@ -10,6 +10,7 @@ import bisect from typing import Any, Dict, List, Optional from collections import defaultdict +logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) logger = logging.getLogger("ggml-hexagon-trace") op_pattern = re.compile( @@ -732,7 +733,7 @@ def main(): group.add_argument("--tail", type=int, help="Limit to last N ops") args = parser.parse_args() - logging.basicConfig(level=logging.INFO, format='%(message)s') + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) op_filter_re = None if args.filter: diff --git a/scripts/snapdragon/run.py b/scripts/snapdragon/run.py index 8917febc1..6d845c341 100755 --- a/scripts/snapdragon/run.py +++ b/scripts/snapdragon/run.py @@ -31,6 +31,7 @@ MANAGED_ENV_NAMES = ( "GGML_HEXAGON_MBUF", "GGML_HEXAGON_MM_SELECT", "GGML_HEXAGON_FA_SELECT", + "GGML_HEXAGON_GDN_SELECT", "GGML_HEXAGON_AR_SELECT", "GGML_HEXAGON_ETM", "GGML_HEXAGON_ARCH", @@ -166,6 +167,7 @@ def main(): parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)") parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 2:HMX,1:HVX,0:disable") parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable") + parser.add_argument("--hex-gdn-select", help="Select Gated Delta Net kernel (GGML_HEXAGON_GDN_SELECT) 2:HMX,1:HVX,0:disable") parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable") parser.add_argument("--hex-etm", help="Enable Embedded Trace Macrocell hardware tracing / trace logging (GGML_HEXAGON_ETM)") parser.add_argument("--hex-arch", help="Target Hexagon NPU architecture version override (v73, v75, v79, v81, etc.) (GGML_HEXAGON_ARCH)") @@ -306,6 +308,7 @@ def main(): set_env("GGML_HEXAGON_MBUF", args.hex_mbuf) set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select) set_env("GGML_HEXAGON_FA_SELECT", args.hex_fa_select) + set_env("GGML_HEXAGON_GDN_SELECT", args.hex_gdn_select) set_env("GGML_HEXAGON_AR_SELECT", args.hex_ar_select) set_env("GGML_HEXAGON_ETM", args.hex_etm) set_env("GGML_HEXAGON_ARCH", args.hex_arch) From c550d2f60bde72df19fcef1fef627895095b8ba8 Mon Sep 17 00:00:00 2001 From: Asahi-Prv Date: Tue, 22 Sep 2026 13:03:14 +0900 Subject: [PATCH 27/42] ci : update Level Zero SDK to v1.33.1 and enable the L0/oneDNN CMake flags in the SYCL job (#29230) --- .github/workflows/build-sycl.yml | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build-sycl.yml b/.github/workflows/build-sycl.yml index 9ddb894f2..5424a8607 100644 --- a/.github/workflows/build-sycl.yml +++ b/.github/workflows/build-sycl.yml @@ -49,7 +49,7 @@ jobs: env: ONEAPI_ROOT: /opt/intel/oneapi/ ONEAPI_INSTALLER_VERSION: "2025.3.3" - LEVEL_ZERO_VERSION: "1.28.2" + LEVEL_ZERO_VERSION: "1.33.1" LEVEL_ZERO_UBUNTU_VERSION: "u24.04" continue-on-error: true @@ -70,9 +70,10 @@ jobs: shell: bash run: | cd /tmp - wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero.deb - wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero-devel_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero-devel.deb - sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb + # v1.33.x renamed the Debian packages to libze1 / libze-dev + wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze1_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze1.deb + wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze-dev_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze-dev.deb + sudo apt-get install -y ./libze1.deb ./libze-dev.deb - name: ccache uses: ggml-org/ccache-action@v1.2.24 @@ -101,7 +102,11 @@ jobs: -DCMAKE_CXX_COMPILER=icpx \ -DLLAMA_OPENSSL=OFF \ -DGGML_NATIVE=OFF \ - -DGGML_SYCL_F16=${{ matrix.fp16 }} + -DGGML_SYCL_F16=${{ matrix.fp16 }} \ + -DGGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON \ + -DGGML_SYCL_DNN=ON \ + -DCMAKE_CXX_FLAGS="-fsycl-unnamed-lambda" \ + -DCMAKE_EXE_LINKER_FLAGS="-fsycl-unnamed-lambda" time cmake --build build --config Release -j $(nproc) - name: ccache-buckets-save @@ -126,7 +131,7 @@ jobs: env: WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b60765d1-2b85-4e85-86b6-cb0e9563a699/intel-deep-learning-essentials-2025.3.3.18_offline.exe WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel - LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero-win-sdk-1.28.2.zip + LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI" ONEAPI_INSTALLER_VERSION: "2025.3.3" steps: From ec5a12b85ae32fbccfa4276051382330a8e6458b Mon Sep 17 00:00:00 2001 From: shaofeiqi Date: Mon, 21 Sep 2026 23:14:00 -0700 Subject: [PATCH 28/42] opencl: add A8 Q4_0 non-MoE dp4a binary kernel (#29055) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 79 ++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index fe7377b2d..db5d510a6 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -1212,6 +1212,7 @@ struct ggml_backend_opencl_context { cl_kernel kernel_gemv_noshuffle_q4_0_f32; cl_kernel kernel_gemv_noshuffle_q4_0_f32_mc3; // multi-column (N=3) verify GEMV (spec/MTP) cl_kernel kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin; + cl_kernel kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8_bin; cl_kernel kernel_gemv_noshuffle_q4_0_f32_32b_trans; cl_kernel kernel_gemv_noshuffle_q4_0_f32_4096_1_11008; cl_kernel kernel_gemv_noshuffle_q4_0_f32_4096_1_4096; @@ -3873,6 +3874,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans = nullptr; backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin = nullptr; + backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8_bin = nullptr; if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) { { std::string opts = std::string("-cl-std=") + opencl_c_std + @@ -3905,6 +3907,17 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { CL_CHECK(clReleaseProgram(bin_prog)); GGML_LOG_CONT("."); } + + kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8", &bin_size); + if (kernel_bin && bin_size > 0) { + cl_program bin_prog = + build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, "", bin_size); + + CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8_bin = + clCreateKernel(bin_prog, "kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8", &err), err)); + CL_CHECK(clReleaseProgram(bin_prog)); + GGML_LOG_CONT("."); + } } } @@ -19327,6 +19340,72 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno_ila(ggml_backend_t backend, const gg cl_mem s_img = extra0_q4_0->d_img; GGML_ASSERT(a_img && s_img && "ILA Q4_0 weight images missing; set_tensor should have built them"); + static const char * q4_0_ila_dp4a_env = getenv("GGML_OPENCL_Q4_0_ILA_DP4A"); + bool q4_0_ila_dp4a_on = q4_0_ila_dp4a_env + ? (atoi(q4_0_ila_dp4a_env) != 0) + : true; + // dot prod has to be available + q4_0_ila_dp4a_on = backend_ctx->has_integer_dot && q4_0_ila_dp4a_on; + + if (q4_0_ila_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8_bin) { + const int dp4a_N_pad = CEIL_DIV(N, 32) * 32; + const size_t n_blocks = (size_t)dp4a_N_pad * (K / 32); + + backend_ctx->prealloc_moe_qa.allocate(context, (size_t)dp4a_N_pad * K * sizeof(cl_char)); + backend_ctx->prealloc_moe_da.allocate(context, n_blocks * sizeof(cl_half)); + backend_ctx->prealloc_moe_sa.allocate(context, n_blocks * sizeof(cl_half)); + + cl_mem b_sub = nullptr; + region.origin = offset1; + region.size = (size_t)K * N * sizeof(float); + CL_CHECK((b_sub = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err)); + + cl_int tb = (cl_int)((size_t)N * (K / 32)); + cl_kernel qk = backend_ctx->kernel_quant_a_q8_1; + CL_CHECK(clSetKernelArg(qk, 0, sizeof(cl_mem), &b_sub)); + CL_CHECK(clSetKernelArg(qk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer)); + CL_CHECK(clSetKernelArg(qk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer)); + CL_CHECK(clSetKernelArg(qk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_sa.buffer)); + CL_CHECK(clSetKernelArg(qk, 4, sizeof(cl_int), &tb)); + size_t q_local[1] = { 64 }; + size_t q_global[1] = { (size_t)CEIL_DIV(tb, 64) * 64 }; + backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst); + + cl_mem d_sub = nullptr; + cl_mem d_img = nullptr; + region.origin = offsetd; + region.size = (size_t)M * N * sizeof(float); + CL_CHECK((d_sub = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err), err)); + + img_fmt = { CL_R, CL_FLOAT }; + memset(&img_desc, 0, sizeof(img_desc)); + img_desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER; + img_desc.image_width = (size_t)M * N; + img_desc.buffer = d_sub; + CL_CHECK((d_img = clCreateImage(context, CL_MEM_WRITE_ONLY, &img_fmt, &img_desc, NULL, &err), err)); + + kernel = backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_ila_a8_bin; + + cl_uint k_arg = 0; + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &a_img)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extra0_q4_0->d)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &d_img)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &K)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &M)); + CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &N)); + + size_t local_work_size[3] = { 64, 1, 1 }; + size_t global_work_size[3] = { 64, (size_t)(M / 64), (size_t)(dp4a_N_pad / 32) }; + backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst); + + CL_CHECK(clReleaseMemObject(b_sub)); + CL_CHECK(clReleaseMemObject(d_img)); + CL_CHECK(clReleaseMemObject(d_sub)); + return; + } + // Pad B through a zero-filled scratch buffer when N needs // padding, since the GEMM kernel always reads a full N-tile. const bool need_pad = N_pad > N; From 8cfc315a8abd654938f21d64ec683f8035d403fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Agust=C3=ADn=20Mista?= Date: Tue, 22 Sep 2026 09:23:41 +0200 Subject: [PATCH 29/42] Add close button to UI toasts (#28246) This commit tweaks the Toaster element to include a close button. These toasts often cover other UI elements like the model selector, and this change avoids having to wait for them to disappear on their own (e.g. after a load failure). --- tools/ui/src/routes/+layout.svelte | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/ui/src/routes/+layout.svelte b/tools/ui/src/routes/+layout.svelte index 625035d71..5de4f82c1 100644 --- a/tools/ui/src/routes/+layout.svelte +++ b/tools/ui/src/routes/+layout.svelte @@ -335,7 +335,7 @@ - + From 0ee9435b8f9143b002dcb485f169aa5070ff6809 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Tue, 22 Sep 2026 03:35:03 -0400 Subject: [PATCH 30/42] ci : publish snapdragon builds in release workflow (#29007) The snapdragon CI builds packages only to feed the QDC device tests, so Hexagon NPU binaries never reached the releases page. Build both targets in release.yml and attach them as release assets. --- .github/workflows/release.yml | 118 ++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 363227f52..be13d9118 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -544,6 +544,120 @@ jobs: path: llama-${{ steps.tag.outputs.name }}-bin-android-arm64.tar.gz name: llama-bin-android-arm64.tar.gz + android-arm64-snapdragon: + needs: [check-release, ui-build] + if: ${{ needs.check-release.outputs.should_release == 'true' }} + + runs-on: ubuntu-latest + container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.7' + + defaults: + run: + shell: bash + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + # checkout runs as the host user; in-container steps run as root, so git + # refuses to touch a repo it does not own. Mark the workspace as safe. + - name: Git safe directory + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" + + - name: Download UI build + uses: actions/download-artifact@v7 + with: + name: llama-ui.zip + path: tools/ui/dist + + - name: Build + id: cmake_build + run: | + cp docs/backend/snapdragon/CMakeUserPresets.json . + cmake --preset arm64-android-snapdragon-release -B build \ + -DCMAKE_INSTALL_RPATH='$ORIGIN;$ORIGIN/../lib' \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + -DLLAMA_BUILD_BORINGSSL=ON \ + ${{ env.CMAKE_ARGS }} + cmake --build build -j $(nproc) + cmake --install build --prefix pkg-snapdragon/llama.cpp + + - name: Determine tag name + id: tag + uses: ./.github/actions/get-tag-name + + - name: Pack artifacts + id: pack_artifacts + run: | + cp LICENSE pkg-snapdragon/llama.cpp/ + tar -czvf llama-${{ steps.tag.outputs.name }}-bin-android-arm64-snapdragon.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C pkg-snapdragon/llama.cpp . + + - name: Upload artifacts + uses: actions/upload-artifact@v6 + with: + path: llama-${{ steps.tag.outputs.name }}-bin-android-arm64-snapdragon.tar.gz + name: llama-bin-android-arm64-snapdragon.tar.gz + + linux-arm64-snapdragon: + needs: [check-release, ui-build] + if: ${{ needs.check-release.outputs.should_release == 'true' }} + + runs-on: ubuntu-latest + container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.7' + + defaults: + run: + shell: bash + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + # checkout runs as the host user; in-container steps run as root, so git + # refuses to touch a repo it does not own. Mark the workspace as safe. + - name: Git safe directory + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" + + - name: Download UI build + uses: actions/download-artifact@v7 + with: + name: llama-ui.zip + path: tools/ui/dist + + - name: Build + id: cmake_build + run: | + cp docs/backend/snapdragon/CMakeUserPresets.json . + cmake --preset arm64-linux-snapdragon-release -B build -DGGML_OPENCL=ON \ + -DCMAKE_INSTALL_RPATH='$ORIGIN;$ORIGIN/../lib' \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + -DLLAMA_BUILD_BORINGSSL=ON \ + ${{ env.CMAKE_ARGS }} + cmake --build build -j $(nproc) + cmake --install build --prefix pkg-snapdragon/llama.cpp + + - name: Determine tag name + id: tag + uses: ./.github/actions/get-tag-name + + - name: Pack artifacts + id: pack_artifacts + run: | + cp LICENSE pkg-snapdragon/llama.cpp/ + tar -czvf llama-${{ steps.tag.outputs.name }}-bin-linux-arm64-snapdragon.tar.gz --transform "s,^\.,llama-${{ steps.tag.outputs.name }}," -C pkg-snapdragon/llama.cpp . + + - name: Upload artifacts + uses: actions/upload-artifact@v6 + with: + path: llama-${{ steps.tag.outputs.name }}-bin-linux-arm64-snapdragon.tar.gz + name: llama-bin-linux-arm64-snapdragon.tar.gz + ubuntu-24-openvino: needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1725,6 +1839,8 @@ jobs: - ubuntu-24-openvino - ubuntu-24-sycl - android-arm64 + - android-arm64-snapdragon + - linux-arm64-snapdragon - macos-cpu - ios-xcode #- openEuler-cann @@ -1860,9 +1976,11 @@ jobs: - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz) - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz) - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz) + - [Linux arm64 (Snapdragon: CPU, Adreno GPU, Hexagon NPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-linux-arm64-snapdragon.tar.gz) - [setup guide](https://github.com/ggml-org/llama.cpp/blob/master/docs/backend/snapdragon/linux.md) **Android:** - [Android arm64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-android-arm64.tar.gz) + - [Android arm64 (Snapdragon: CPU, Adreno GPU, Hexagon NPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-android-arm64-snapdragon.tar.gz) - [setup guide](https://github.com/ggml-org/llama.cpp/blob/master/docs/backend/snapdragon/README.md) **Windows:** - [Windows x64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cpu-x64.zip) From 7ab4ee7baad2d920464cbacfad4f4b07cf111fd2 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Tue, 22 Sep 2026 01:38:52 -0600 Subject: [PATCH 31/42] chat : Fix Muse Glimmer tool-call first parser error (#29242) * Fix Muse Glimmer tool-call first parser error * Add test to verify * Condense patterns * remove test for trigger patterns --- common/parsers/muse-glimmer.cpp | 2 +- tests/test-chat.cpp | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/common/parsers/muse-glimmer.cpp b/common/parsers/muse-glimmer.cpp index add95697e..d03bf2d58 100644 --- a/common/parsers/muse-glimmer.cpp +++ b/common/parsers/muse-glimmer.cpp @@ -130,7 +130,7 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa }); data.grammar_triggers = { { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, - "<\\|start\\|>assistant( to=(?!self<\\|message\\|>)(?!user<\\|message\\|>)[^<]*?<\\|message\\|>)" }, + "(?:^|<\\|start\\|>assistant)( to=(?!self<\\|message\\|>)(?!user<\\|message\\|>)[^<]*?<\\|message\\|>)" }, }; } diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 13733c1d9..f2728f7cc 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -6357,6 +6357,14 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect(message_assist) .run(); + // A tool call as the first message of the turn: "<|start|>assistant" is the + // generation prompt, so the output starts at " to=". + tst.test(" to=special_function<|message|>" + call_markup) + .tools({ special_function_tool }) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist_call) + .run(); + // "Inform then act": the model answers the user and calls a tool in ONE generation, // closing the answer with <|eom|>. The answer must stop there rather than swallow it. tst.test(" to=user<|message|>Hello, world!\nWhat's up?<|eom|>" From a60f9aead0047b104807930b75fbc0a649f57779 Mon Sep 17 00:00:00 2001 From: miyan <1138989048@qq.com> Date: Tue, 22 Sep 2026 19:58:42 +0800 Subject: [PATCH 32/42] cmake : allow repeated find_package calls for llama (#29228) --- cmake/llama-config.cmake.in | 20 +++++++++++--------- examples/simple-cmake-pkg/CMakeLists.txt | 3 +++ 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/cmake/llama-config.cmake.in b/cmake/llama-config.cmake.in index 6db73577a..6500203b1 100644 --- a/cmake/llama-config.cmake.in +++ b/cmake/llama-config.cmake.in @@ -17,14 +17,16 @@ find_library(llama_LIBRARY llama NO_CMAKE_FIND_ROOT_PATH ) -add_library(llama UNKNOWN IMPORTED) -set_target_properties(llama - PROPERTIES - INTERFACE_INCLUDE_DIRECTORIES "${LLAMA_INCLUDE_DIR}" - INTERFACE_LINK_LIBRARIES "ggml::ggml;ggml::ggml-base;" - IMPORTED_LINK_INTERFACE_LANGUAGES "CXX" - IMPORTED_LOCATION "${llama_LIBRARY}" - INTERFACE_COMPILE_FEATURES c_std_90 - POSITION_INDEPENDENT_CODE ON) +if(NOT TARGET llama) + add_library(llama UNKNOWN IMPORTED) + set_target_properties(llama + PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${LLAMA_INCLUDE_DIR}" + INTERFACE_LINK_LIBRARIES "ggml::ggml;ggml::ggml-base;" + IMPORTED_LINK_INTERFACE_LANGUAGES "CXX" + IMPORTED_LOCATION "${llama_LIBRARY}" + INTERFACE_COMPILE_FEATURES c_std_90 + POSITION_INDEPENDENT_CODE ON) +endif() check_required_components(Llama) diff --git a/examples/simple-cmake-pkg/CMakeLists.txt b/examples/simple-cmake-pkg/CMakeLists.txt index 128e38c8f..04fe1291c 100644 --- a/examples/simple-cmake-pkg/CMakeLists.txt +++ b/examples/simple-cmake-pkg/CMakeLists.txt @@ -5,6 +5,9 @@ set(TARGET llama-simple-cmake-pkg) find_package(Llama REQUIRED) +# Check that repeated package discovery does not redefine imported targets. +find_package(Llama 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}) From bfd73a876ed25c5a760d6b60b8fe12a65115fedf Mon Sep 17 00:00:00 2001 From: AesSedai <7980540+AesSedai@users.noreply.github.com> Date: Tue, 22 Sep 2026 05:38:09 -0700 Subject: [PATCH 33/42] convert: add MiMo-V2.6 support (#29257) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * convert: add MiMo-V2.6 support Hoist the K3 mxfp4 conversion repack into base.py so it can be reused Remove decoder from mmproj convert * Update conversion/mimo.py * fix: use autoparser --------- Co-authored-by: Sigbjørn Skjæret Co-authored-by: Piotr Wilkin --- common/chat.cpp | 4 +- conversion/base.py | 30 +++++++++++++++ conversion/kimi_k3.py | 35 +---------------- conversion/mimo.py | 88 ++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 121 insertions(+), 36 deletions(-) diff --git a/common/chat.cpp b/common/chat.cpp index 6c8099cf2..ed1942e15 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1212,7 +1212,9 @@ std::optional common_chat_try_specialized_template( // Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash if (src.find("") != std::string::npos && src.find("'") == std::string::npos) { LOG_DBG("Using specialized template: Qwen3-Coder\n"); return common_chat_params_init_qwen3_coder(tmpl, params); } diff --git a/conversion/base.py b/conversion/base.py index 6aca7f1d3..9fba5a86b 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -776,6 +776,36 @@ class ModelBase: raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) return raw.reshape(rows, n_blocks * 17).cpu().numpy() + def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]): + """ + One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. + + gguf_writer holds every added tensor until the final write, so building + this eagerly (like the DeepSeek-V4 path does) keeps every expert in + memory at once. lazy means only the tensor being written is resident. + """ + # meta shapes, so this does not read any weights + rows, packed_cols = loaders[0][0]().shape + n_blocks = (packed_cols * 2) // 32 + byte_shape = (len(loaders), rows, n_blocks * 17) + + def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: + out = np.empty(byte_shape, dtype=np.uint8) + for eid, (packed_fn, scale_fn) in enumerate(fns): + out[eid] = self.repack_mxfp4_blocks( + LazyTorchTensor.to_eager(packed_fn()), + LazyTorchTensor.to_eager(scale_fn()), + ) + return out + + # loaders goes through args, not the closure, so that `func` matches + # LazyBase's single-argument shape + return gguf.LazyNumpyTensor( + meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), + args=(loaders,), + func=load, + ) + @staticmethod def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]: """Repack NVFP4 ModelOpt tensors into ggml super-block layout. diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index d15d1d64b..70aabb705 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -2,15 +2,14 @@ from __future__ import annotations import re from pathlib import Path -from typing import Callable, Iterable, Iterator, TYPE_CHECKING +from typing import Iterable, Iterator, TYPE_CHECKING -import numpy as np import torch if TYPE_CHECKING: from torch import Tensor -from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger +from .base import ModelBase, TextModel, gguf, logger from .kimi_linear import KimiLinearModel @@ -104,36 +103,6 @@ class KimiK3Model(TextModel): "only the routed experts have a repack path" ) - def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]): - """ - One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. - - gguf_writer holds every added tensor until the final write, so building - this eagerly (like the DeepSeek-V4 path does) keeps all ~1.38 TB of - experts in memory. lazy means only the tensor being written is resident. - """ - # meta shapes, so this does not read any weights - rows, packed_cols = loaders[0][0]().shape - n_blocks = (packed_cols * 2) // 32 - byte_shape = (len(loaders), rows, n_blocks * 17) - - def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: - out = np.empty(byte_shape, dtype=np.uint8) - for eid, (packed_fn, scale_fn) in enumerate(fns): - out[eid] = self.repack_mxfp4_blocks( - LazyTorchTensor.to_eager(packed_fn()), - LazyTorchTensor.to_eager(scale_fn()), - ) - return out - - # loaders goes through args, not the closure, so that `func` matches - # LazyBase's single-argument shape - return gguf.LazyNumpyTensor( - meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), - args=(loaders,), - func=load, - ) - def _write_mxfp4_experts(self) -> None: n_experts = self.hparams["num_experts"] diff --git a/conversion/mimo.py b/conversion/mimo.py index 15dbeb7e7..8a2689b96 100644 --- a/conversion/mimo.py +++ b/conversion/mimo.py @@ -10,7 +10,7 @@ import torch if TYPE_CHECKING: from torch import Tensor -from .base import MmprojModel, ModelBase, TextModel, gguf +from .base import MmprojModel, ModelBase, TextModel, gguf, logger @ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM") @@ -167,6 +167,84 @@ class MimoV2Model(TextModel): self.gguf_writer.add_nextn_predict_layers(self._n_nextn) + _MXFP4_EXPERT_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" + ) + _MXFP4_PROJ = { + "gate": gguf.MODEL_TENSOR.FFN_GATE_EXP, + "up": gguf.MODEL_TENSOR.FFN_UP_EXP, + "down": gguf.MODEL_TENSOR.FFN_DOWN_EXP, + } + + def _is_mxfp4_packed(self) -> bool: + quant_config = self.hparams.get("quantization_config") or {} + if quant_config.get("store_dtype") != "mxfp4": + return False + # repack_mxfp4_blocks assumes ggml's 32-element group + block_size = quant_config.get("mxfp4_block_size", 32) + if block_size != 32: + raise NotImplementedError( + f"MXFP4 block size {block_size} is not ggml's QK_MXFP4 (32)") + return True + + def _write_mxfp4_experts(self) -> None: + n_experts = self.hparams["n_routed_experts"] + + # the FP8 half uses `weight_scale_inv` and is left to dequant_model + stray = [n for n in self.model_tensors + if n.endswith(".weight_scale") and not self._MXFP4_EXPERT_RE.match(n.removesuffix("_scale"))] + if stray: + raise NotImplementedError( + f"{len(stray)} MXFP4 tensor(s) outside the routed experts, e.g. {stray[0]!r}; " + "only the routed experts have a repack path" + ) + + # (bid, proj) -> {expert id: (weight name, scale name)} + groups: dict[tuple[int, str], dict[int, tuple[str, str]]] = {} + for name in self.model_tensors: + m = self._MXFP4_EXPERT_RE.match(name) + if m is None: + continue + bid, eid, proj = int(m.group(1)), int(m.group(2)), m.group(3) + scale_name = name + "_scale" + if scale_name not in self.model_tensors: + raise KeyError(f"missing {scale_name} for {name}") + groups.setdefault((bid, proj), {})[eid] = (name, scale_name) + + consumed: list[str] = [] + for (bid, proj), experts in sorted(groups.items()): + missing = [e for e in range(n_experts) if e not in experts] + if missing or len(experts) != n_experts: + raise KeyError( + f"layer {bid} {proj}_proj: {len(experts)} of {n_experts} experts present" + + (f", first missing is {missing[0]}" if missing else "") + ) + + loaders = [] + for eid in range(n_experts): + weight_name, scale_name = experts[eid] + loaders.append((self.model_tensors[weight_name], self.model_tensors[scale_name])) + consumed += [weight_name, scale_name] + + data = self._mxfp4_expert_tensor(loaders) + new_name = self.format_tensor_name(self._MXFP4_PROJ[proj], bid) + shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4) + logger.info( + f"{new_name}: repacked {n_experts} experts to MXFP4, " + f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}" + ) + self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4) + + for name in consumed: + del self.model_tensors[name] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # not a generator on purpose: base.py chains this with get_tensors(), so the + # tensors used here must be removed from model_tensors before that starts + if self._is_mxfp4_packed(): + self._write_mxfp4_experts() + return () + _experts: list[dict[str, Tensor]] | None = None @classmethod @@ -192,7 +270,7 @@ class MimoV2Model(TextModel): bid = new_bid # process the experts separately - if name.find("mlp.experts") != -1: + if ".mlp.experts." in name and name.endswith(".weight"): n_experts = self.hparams["n_routed_experts"] assert bid is not None @@ -229,6 +307,10 @@ class MimoV2Model(TextModel): if len(experts) > 0: raise ValueError(f"Unprocessed experts: {experts}") + if self._is_mxfp4_packed(): + self._is_mxfp4 = True + self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE + @ModelBase.register("MiMoV2ForCausalLM") @ModelBase.example("XiaomiMiMo/MiMo-V2.5") @@ -382,6 +464,8 @@ class MiMoV2VisionAudioModel(MmprojModel): "_codebook.inited", ) for name, tensor in state_dict.items(): + if name.startswith("decoder."): + continue if name.endswith(skip_suffixes): continue if m := codebook_re.match(name): From 828fdf282e195300c2965bd9511807e24ed53bdb Mon Sep 17 00:00:00 2001 From: wendadawen <130649302+wendadawen@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:04:52 +0800 Subject: [PATCH 34/42] spec : support DFlash for HunyuanOCR (#28890) * model : add DFlash layer-input taps for HunyuanVL DFlash speculative decoding needs the target graph to expose the residual stream entering each layer (res->t_layer_inp[il]) - the draft model reads those tensors to build its cross-context. Qwen3 and the other DFlash-capable targets register them, but the Hunyuan graphs do not, so serving a DFlash draft against a HunyuanOCR target aborts during the first graph build: GGML_ASSERT(t_layer_inp[il] != nullptr && "layer input tensor is null") Register the tensor at the top of the layer loop, mirroring qwen3. The layer input is the residual stream entering layer il, i.e. the output of layer il-1, which is what the draft's target_layers metadata refers to (the converter writes target_layer_ids+1). hunyuan-dense.cpp reuses this graph, so it is covered as well; hunyuan-moe has a separate graph and is untouched. The vector is only read when a speculative implementation enables those layer ids, so there is no behaviour change without a draft model. Tested with tencent/HunyuanOCR 1.5 and its DFlash draft: image requests now run, draft acceptance is ~0.5 and the OCR output is byte-identical to the non-speculative run. Co-authored-by: wendadawen * convert : fix DFlash draft conversion against HunYuan targets Converting a DFlash draft with a HunYuan target failed in two ways. 1. DFlashModel.set_vocab() reuses the target class' vocab handling by calling it unbound with the draft instance, but HunYuanModel.set_vocab() called self._fix_special_tokens(), a method that only exists on HunYuanModel, so the conversion always aborted with AttributeError: 'DFlashModel' object has no attribute '_fix_special_tokens' Make the vocab helpers module-level functions taking the model explicitly, so they do not depend on the instance being a HunYuanModel. They have no other callers, so the two id lookups are folded into _fix_special_tokens(). 2. The delegated call runs with self.dir_model pointed at the target but keeps the draft's self.hparams, so config lookups inside the target's vocab code (the pad_token_id < 0 guard, eod_token_id) read the draft's config instead of the target's. That aborts on targets with pad_token_id = -1 (e.g. the HunyuanOCR v1.0 checkpoint) and otherwise writes special token ids that disagree with the target. Add _vocab_hparams(): it returns the target's config (with text_config merged to the root, as TextModel does) when the model is a draft converted with --target-model-dir, and the model's own hparams otherwise, so a normal conversion is unaffected. Tested: converting tencent/HunyuanOCR/dflash succeeds with both the 1.5 and the v1.0 target; converting the base model without --target-model-dir produces a byte-identical GGUF to before. Co-authored-by: wendadawen * convert : fix DFlash draft vocab against HunYuan targets Switch hparams to the target config for the duration of the borrowed set_vocab(), matching the existing dir_model swap, instead of teaching HunYuanModel::set_vocab about draft models. * convert : fix HunYuan special token ids for DFlash drafts * convert : use load_hparams for HunYuan special token ids --- conversion/hunyuan.py | 46 ++++++++++++++++----------------------- src/models/hunyuan-vl.cpp | 2 ++ 2 files changed, 21 insertions(+), 27 deletions(-) diff --git a/conversion/hunyuan.py b/conversion/hunyuan.py index ee1a10654..58dee2ec8 100644 --- a/conversion/hunyuan.py +++ b/conversion/hunyuan.py @@ -159,32 +159,14 @@ class HunYuanMoEModel(TextModel): class HunYuanModel(TextModel): model_arch = gguf.MODEL_ARCH.HUNYUAN_DENSE - def _get_eod_token_id(self) -> int | None: - """Get the actual end-of-generation token from config (eod_token_id).""" - return self.hparams.get("eod_token_id") - - def _get_eot_token_id(self) -> int | None: - """Get the end-of-turn token from generation_config.json. - This is the first entry in eos_token_id when it's a list.""" - gen_cfg_path = self.dir_model / "generation_config.json" - if gen_cfg_path.is_file(): - with open(gen_cfg_path, encoding="utf-8") as f: - gen_cfg = json.load(f) - eos = gen_cfg.get("eos_token_id") - if isinstance(eos, list) and len(eos) >= 2: - return eos[0] - return None - - def _fix_special_tokens(self): - """Fix EOS/EOT tokens that are incorrect in upstream configs.""" - eod_id = self._get_eod_token_id() - if eod_id is not None: - self.gguf_writer.add_eos_token_id(eod_id) - eot_id = self._get_eot_token_id() - if eot_id is not None: - self.gguf_writer.add_eot_token_id(eot_id) - def set_vocab(self): + # Also called by draft models (e.g. DFlash), with dir_model pointing at + # the target model. + config = ModelBase.load_hparams(self.dir_model, self.is_mistral_format) + config = {**config, **config.get("text_config", {})} + self.hparams["pad_token_id"] = config.get("pad_token_id") + self.hparams["eod_token_id"] = config.get("eod_token_id") + if (self.dir_model / "tokenizer.json").is_file(): tokens, toktypes, tokpre = self.get_vocab_base() self.gguf_writer.add_tokenizer_model("gpt2") @@ -199,7 +181,6 @@ class HunYuanModel(TextModel): token_types = ('bos', 'eos', 'unk', 'sep', 'cls', 'mask') special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True, special_token_types=token_types) special_vocab.add_to_gguf(self.gguf_writer) - self._fix_special_tokens() else: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) @@ -251,7 +232,18 @@ class HunYuanModel(TextModel): # FIX for BOS token: Overwrite incorrect id read from config.json if self.hparams['hidden_size'] == 4096: self.gguf_writer.add_bos_token_id(127958) # only for 7b dense, fix <|bos|> token - self._fix_special_tokens() + + # Fix EOS/EOT tokens that are incorrect in upstream configs. + eod_id = self.hparams.get("eod_token_id") + if eod_id is not None: + self.gguf_writer.add_eos_token_id(eod_id) + + gen_cfg = self.dir_model / "generation_config.json" + if gen_cfg.is_file(): + with open(gen_cfg, encoding="utf-8") as f: + eos = json.load(f).get("eos_token_id") + if isinstance(eos, list) and len(eos) >= 2: + self.gguf_writer.add_eot_token_id(eos[0]) def set_gguf_parameters(self): # Some HunYuanVL variants set num_experts=1 (not real MoE); diff --git a/src/models/hunyuan-vl.cpp b/src/models/hunyuan-vl.cpp index da9bb74de..18b6eaf8c 100644 --- a/src/models/hunyuan-vl.cpp +++ b/src/models/hunyuan-vl.cpp @@ -83,6 +83,8 @@ llama_model_hunyuan_vl::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * inp_out_ids = build_inp_out_ids(); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + ggml_tensor * inpSA = inpL; // norm From 217f81c266a7b7c986ee3d2c58e1cccee05a0744 Mon Sep 17 00:00:00 2001 From: Emanuil Rusev Date: Tue, 22 Sep 2026 16:16:40 +0300 Subject: [PATCH 35/42] server: Add support for binding to multiple addresses (#28690) * Add support for binding llama-server to multiple addresses Assisted-by: Codex * remove redundant thread handler * make it clear about overlapping addr * reject --port 0 with multiple tcp addr * improve arg handler * nits * fix test * nits 2 * nits * nits 2 --------- Co-authored-by: Xuan Son Nguyen --- common/arg.cpp | 13 +- common/common.h | 2 +- tools/server/README.md | 2 +- tools/server/server-http.cpp | 191 +++++++++++++++++++------- tools/server/server-http.h | 8 +- tools/server/server.cpp | 41 +++--- tools/server/tests/unit/test_basic.py | 32 +++++ tools/server/tests/utils.py | 14 +- 8 files changed, 218 insertions(+), 85 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 996ea75fe..63e342776 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3308,9 +3308,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).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()), + string_format("IP addresses to listen on, comma-separated, or UNIX socket paths ending in .sock; with multiple TCP addresses, :: binds IPv6 only; overlapping addresses result in undefined behavior (default: %s)", params.hostnames[0].c_str()), [](common_params & params, const std::string & value) { - params.hostname = value; + params.hostnames.clear(); + for (auto & host : parse_csv_row(value)) { + host = string_strip(host); + if (!host.empty()) { + params.hostnames.push_back(host); + } + } + if (params.hostnames.empty()) { + throw std::invalid_argument("--host requires at least one address"); + } } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_HOST")); add_opt(common_arg( diff --git a/common/common.h b/common/common.h index 63d0badd0..7afc266ac 100644 --- a/common/common.h +++ b/common/common.h @@ -631,10 +631,10 @@ struct common_params { int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints 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 + std::vector hostnames = {"127.0.0.1"}; bool use_jinja = true; // NOLINT // server CORS params diff --git a/tools/server/README.md b/tools/server/README.md index 0ee8df291..e665904ba 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -189,7 +189,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-a, --alias STRING` | set model name aliases, comma-separated (to be used by API)
(env: LLAMA_ARG_ALIAS) | | `--tags STRING` | set model tags, comma-separated (informational, not used for routing)
(env: LLAMA_ARG_TAGS) | | `--embd-normalize N` | normalisation for embeddings (default: 2) (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm) | -| `--host HOST` | ip address to listen, or bind to an UNIX socket if the address ends with .sock (default: 127.0.0.1)
(env: LLAMA_ARG_HOST) | +| `--host HOST` | IP addresses to listen on, comma-separated, or UNIX socket paths ending in .sock; with multiple TCP addresses, :: binds IPv6 only; overlapping addresses result in undefined behavior (default: 127.0.0.1)
(env: LLAMA_ARG_HOST) | | `--port PORT` | port to listen (default: 8080)
(env: LLAMA_ARG_PORT) | | `--reuse-port` | allow multiple sockets to bind to the same port (default: disabled)
(env: LLAMA_ARG_REUSE_PORT) | | `--path PATH` | path to serve static files from (default: )
(env: LLAMA_ARG_STATIC_PATH) | diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index 2ec137aa0..e46a1f6f1 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -18,14 +18,37 @@ class server_http_context::Impl { public: - std::unique_ptr srv; + std::vector> servers; + std::vector hosts; + std::vector threads; // one thread per listener + std::unique_ptr pool; // single pool shared among all listeners + int n_threads_http = 0; +}; + +class server_http_task_queue : public httplib::TaskQueue { + httplib::ThreadPool & pool; +public: + explicit server_http_task_queue(httplib::ThreadPool & pool) : pool(pool) {} + bool enqueue(std::function fn) override { return pool.enqueue(std::move(fn)); } + // note: must call join() to drain the pool + void shutdown() override { /* no-op */ } }; server_http_context::server_http_context() : pimpl(std::make_unique()) {} -server_http_context::~server_http_context() = default; +server_http_context::~server_http_context() { + // just in case any exit paths that forget to call join() + try { + stop(); + join(); + } catch (const std::exception & e) { + SRV_ERR("failed to stop HTTP server: %s\n", e.what()); + } catch (...) { + SRV_ERR("%s", "failed to stop HTTP server\n"); + } +} static void log_server_request(const httplib::Request & req, const httplib::Response & res) { // skip logging requests that are regularly sent, to avoid log spam @@ -90,7 +113,6 @@ bool server_http_context::init(const common_params & params) { path_prefix = params.api_prefix; port = params.port; - hostname = params.hostname; if (gcp.enabled) { SRV_TRC("Google Cloud Platform compat: health route = %s, predict route = %s, port = %d\n", gcp.path_health.c_str(), gcp.path_predict.c_str(), gcp.port); @@ -102,7 +124,39 @@ bool server_http_context::init(const common_params & params) { port = gcp.port; } - auto & srv = pimpl->srv; + pimpl->hosts = params.hostnames; + size_t n_tcp_hosts = 0; + for (const auto & host : pimpl->hosts) { + if (!string_ends_with(host, ".sock")) { + n_tcp_hosts++; + } + } + if (port == 0 && n_tcp_hosts > 1) { + SRV_ERR("%s", "--port 0 is not supported with multiple TCP addresses\n"); + return false; + } + for (size_t i = 0; i < pimpl->hosts.size(); ++i) { + pimpl->servers.emplace_back(); + if (!init_listener(params)) { + return false; + } + // with multiple TCP addresses, [::] must not also claim 0.0.0.0 + if (n_tcp_hosts > 1) { + pimpl->servers.back()->set_ipv6_v6only(true); + } + } + + pimpl->n_threads_http = params.n_threads_http; + if (pimpl->n_threads_http < 1) { + // +4 threads for monitoring, health and MCP. + pimpl->n_threads_http = std::max(params.n_parallel + 4, static_cast(std::thread::hardware_concurrency() - 1)); + } + SRV_TRC("using %d threads for HTTP server\n", pimpl->n_threads_http); + return true; +} + +bool server_http_context::init_listener(const common_params & params) { + auto & srv = pimpl->servers.back(); #ifdef CPPHTTPLIB_OPENSSL_SUPPORT if (!params.ssl_file_key.empty() && !params.ssl_file_cert.empty()) { @@ -306,18 +360,8 @@ bool server_http_context::init(const common_params & params) { return httplib::Server::HandlerResponse::Unhandled; }); - auto n_threads_http = params.n_threads_http; - if (n_threads_http < 1) { - // +4 threads for monitoring, health and some threads reserved for MCP and other tasks in the future - n_threads_http = std::max(params.n_parallel + 4, static_cast(std::thread::hardware_concurrency() - 1)); - } - SRV_TRC("using %d threads for HTTP server\n", n_threads_http); - srv->new_task_queue = [n_threads_http] { - // spawn n_threads_http fixed thread (always alive), while allow up to 1024 max possible additional threads - // when n_threads_http is used, server will create new "dynamic" threads that will be destroyed after processing each request - // ref: https://github.com/yhirose/cpp-httplib/pull/2368 - const auto max_threads = static_cast(n_threads_http + 1024); - return new httplib::ThreadPool(n_threads_http, max_threads); + srv->new_task_queue = [this] { + return new server_http_task_queue(*pimpl->pool); }; // @@ -432,47 +476,76 @@ bool server_http_context::init(const common_params & params) { bool server_http_context::start() { // Bind and listen - const auto & srv = pimpl->srv; - auto was_bound = false; - auto is_sock = false; - if (string_ends_with(std::string(hostname), ".sock")) { - is_sock = true; - SRV_TRC("%s", "setting address family to AF_UNIX\n"); - srv->set_address_family(AF_UNIX); - // bind_to_port requires a second arg, any value other than 0 should - // simply get ignored - was_bound = srv->bind_to_port(hostname, 8080); - } else { - SRV_TRC("%s", "binding port with default address family\n"); - // bind HTTP listen port - if (port == 0) { - const auto bound_port = srv->bind_to_any_port(hostname); - was_bound = (bound_port >= 0); + listening_addresses.clear(); + for (size_t i = 0; i < pimpl->servers.size(); ++i) { + const auto & srv = pimpl->servers[i]; + const auto & host = pimpl->hosts[i]; + const bool is_sock = string_ends_with(host, ".sock"); + bool was_bound; + if (is_sock) { + SRV_TRC("%s", "setting address family to AF_UNIX\n"); + srv->set_address_family(AF_UNIX); + // AF_UNIX ignores the port, but bind_to_port requires a nonzero value. + was_bound = srv->bind_to_port(host, 8080); + } else if (port == 0) { + const auto bound_port = srv->bind_to_any_port(host); + was_bound = bound_port >= 0; if (was_bound) { port = bound_port; } } else { - was_bound = srv->bind_to_port(hostname, port); + was_bound = srv->bind_to_port(host, port); + } + if (!was_bound) { + SRV_ERR("couldn't bind HTTP server socket, hostname: %s, port: %d\n", host.c_str(), port); + stop(); + listening_addresses.clear(); + return false; + } + listening_addresses.push_back(is_sock ? string_format("unix://%s", host.c_str()) + : string_format("%s://%s:%d", is_ssl ? "https" : "http", common_http_format_host(host).c_str(), port)); + } + + // n_threads_http fixed threads (always alive), plus up to 1024 dynamic threads destroyed after each request + // ref: https://github.com/yhirose/cpp-httplib/pull/2368 + pimpl->pool = std::make_unique(pimpl->n_threads_http, pimpl->n_threads_http + 1024); + for (size_t i = 0; i < pimpl->servers.size(); ++i) { + const auto & srv = pimpl->servers[i]; + pimpl->threads.emplace_back([srv = srv.get(), addr = listening_addresses[i]] { + if (!srv->listen_after_bind()) { + SRV_ERR("listener on %s stopped unexpectedly\n", addr.c_str()); + } + }); + srv->wait_until_ready(); + if (!srv->is_running()) { + SRV_ERR("couldn't start HTTP listener on %s\n", listening_addresses[i].c_str()); + stop(); + join(); + listening_addresses.clear(); + return false; } } - - if (!was_bound) { - SRV_ERR("couldn't bind HTTP server socket, hostname: %s, port: %d\n", hostname.c_str(), port); - return false; - } - - // run the HTTP server in a thread - thread = std::thread([this] { pimpl->srv->listen_after_bind(); }); - srv->wait_until_ready(); - - listening_address = is_sock ? string_format("unix://%s", hostname.c_str()) - : string_format("%s://%s:%d", is_ssl ? "https" : "http", common_http_format_host(hostname).c_str(), port); return true; } void server_http_context::stop() const { - if (pimpl->srv) { - pimpl->srv->stop(); + for (const auto & srv : pimpl->servers) { + if (srv) { + srv->stop(); + } + } +} + +void server_http_context::join() { + for (auto & thread : pimpl->threads) { + if (thread.joinable()) { + thread.join(); + } + } + // Queued requests still refer to their servers until the workers finish. + if (pimpl->pool) { + pimpl->pool->shutdown(); + pimpl->pool.reset(); } } @@ -584,7 +657,7 @@ static void process_handler_response(server_http_req_ptr && request, server_http void server_http_context::get(const std::string & path, const server_http_context::handler_t & handler) const { handlers.emplace(path, handler); - pimpl->srv->Get(path_prefix + path, [handler](const httplib::Request & req, httplib::Response & res) { + auto callback = [handler](const httplib::Request & req, httplib::Response & res) { server_http_req_ptr request = std::make_unique(server_http_req{ get_params(req), get_headers(req), @@ -596,12 +669,16 @@ void server_http_context::get(const std::string & path, const server_http_contex }); server_http_res_ptr response = handler(*request); process_handler_response(std::move(request), response, res); - }); + }; + const std::string full_path = path_prefix + path; + for (const auto & srv : pimpl->servers) { + srv->Get(full_path, callback); + } } void server_http_context::post(const std::string & path, const server_http_context::handler_t & handler) const { handlers.emplace(path, handler); - pimpl->srv->Post(path_prefix + path, [handler](const httplib::Request & req, httplib::Response & res) { + auto callback = [handler](const httplib::Request & req, httplib::Response & res) { std::string body = req.body; std::map files; @@ -643,12 +720,16 @@ void server_http_context::post(const std::string & path, const server_http_conte }); server_http_res_ptr response = handler(*request); process_handler_response(std::move(request), response, res); - }); + }; + const std::string full_path = path_prefix + path; + for (const auto & srv : pimpl->servers) { + srv->Post(full_path, callback); + } } void server_http_context::del(const std::string & path, const server_http_context::handler_t & handler) const { handlers.emplace(path, handler); - pimpl->srv->Delete(path_prefix + path, [handler](const httplib::Request & req, httplib::Response & res) { + auto callback = [handler](const httplib::Request & req, httplib::Response & res) { server_http_req_ptr request = std::make_unique(server_http_req{ get_params(req), get_headers(req), @@ -660,7 +741,11 @@ void server_http_context::del(const std::string & path, const server_http_contex }); server_http_res_ptr response = handler(*request); process_handler_response(std::move(request), response, res); - }); + }; + const std::string full_path = path_prefix + path; + for (const auto & srv : pimpl->servers) { + srv->Delete(full_path, callback); + } } // diff --git a/tools/server/server-http.h b/tools/server/server-http.h index 032b08d0d..4554b20f4 100644 --- a/tools/server/server-http.h +++ b/tools/server/server-http.h @@ -68,7 +68,6 @@ struct server_http_context { class Impl; std::unique_ptr pimpl; - std::thread thread; // server thread std::atomic is_ready = false; // note: the handler should never throw exceptions @@ -76,7 +75,6 @@ struct server_http_context { mutable std::unordered_map handlers; std::string path_prefix; - std::string hostname; int port = 8080; bool is_ssl = false; @@ -86,6 +84,7 @@ struct server_http_context { bool init(const common_params & params); bool start(); void stop() const; + void join(); void get(const std::string & path, const handler_t & handler) const; void post(const std::string & path, const handler_t & handler) const; @@ -96,5 +95,8 @@ struct server_http_context { void register_gcp_compat() const; // for debugging - std::string listening_address; + std::vector listening_addresses; + +private: + bool init_listener(const common_params & params); }; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 1167c0aea..049bdcebb 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -111,7 +111,9 @@ int llama_server(int argc, char ** argv) { llama_backend_init(); llama_numa_init(params.numa); - return llama_server(params, argc, argv); + const int result = llama_server(params, argc, argv); + common_log_flush(common_log_main()); + return result; } int llama_server(common_params & params, int argc, char ** argv) { @@ -183,12 +185,6 @@ int llama_server(common_params & params, int argc, char ** argv) { // struct that contains llama context and inference server_context ctx_server; - server_http_context ctx_http; - if (!ctx_http.init(params)) { - SRV_ERR("%s", "failed to initialize HTTP server\n"); - return 1; - } - // // Router // @@ -199,6 +195,13 @@ int llama_server(common_params & params, int argc, char ** argv) { server_tools tools; std::optional models_routes{}; + + server_http_context ctx_http; + if (!ctx_http.init(params)) { + SRV_ERR("%s", "failed to initialize HTTP server\n"); + return 1; + } + if (is_router_server) { // setup server instances manager try { @@ -438,9 +441,7 @@ int llama_server(common_params & params, int argc, char ** argv) { } catch (const std::exception & e) { SRV_ERR("failed to load models on startup: %s\n", e.what()); ctx_http.stop(); - if (ctx_http.thread.joinable()) { - ctx_http.thread.join(); - } + ctx_http.join(); clean_up(); return 1; } @@ -473,9 +474,7 @@ int llama_server(common_params & params, int argc, char ** argv) { if (!ctx_server.load_model(params)) { clean_up(); - if (ctx_http.thread.joinable()) { - ctx_http.thread.join(); - } + ctx_http.join(); SRV_ERR("%s", "exiting due to model loading error\n"); return 1; } @@ -509,11 +508,15 @@ int llama_server(common_params & params, int argc, char ** argv) { #endif } - SRV_INF("listening on %s\n", ctx_http.listening_address.c_str()); + bool uses_default_port = false; + for (const auto & address : ctx_http.listening_addresses) { + SRV_INF("listening on %s\n", address.c_str()); + uses_default_port |= string_ends_with(address, ":8080"); + } // TODO: remove this in the future // check the string to also handle the .sock case - if (string_ends_with(ctx_http.listening_address, ":8080")) { + if (uses_default_port) { SRV_WRN("%s", "notice: server default port will be changed to :9931 in a future release (ref: https://github.com/ggml-org/llama.cpp/pull/26508)\n"); } @@ -523,9 +526,7 @@ int llama_server(common_params & params, int argc, char ** argv) { SRV_WRN("%s", " please only use presets that you can trust! Unknown presets may be unsafe\n"); } - if (ctx_http.thread.joinable()) { - ctx_http.thread.join(); // keep the main thread alive - } + ctx_http.join(); // keep the main thread alive // when the HTTP server stops, clean up and exit clean_up(); @@ -541,9 +542,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_server.start_loop(); clean_up(); - if (ctx_http.thread.joinable()) { - ctx_http.thread.join(); - } + ctx_http.join(); if (monitor_thread.joinable()) { monitor_thread.join(); } diff --git a/tools/server/tests/unit/test_basic.py b/tools/server/tests/unit/test_basic.py index 285726abf..b9e9f84f6 100644 --- a/tools/server/tests/unit/test_basic.py +++ b/tools/server/tests/unit/test_basic.py @@ -1,5 +1,6 @@ import pytest import requests +import socket from utils import * server = ServerPreset.tinyllama2() @@ -18,6 +19,37 @@ def test_server_start_simple(): assert res.status_code == 200 +def test_server_multiple_addresses(monkeypatch): + # The CLI value replaces the environment value, including an unavailable address. + monkeypatch.setenv("LLAMA_ARG_HOST", "192.0.2.1") + try: + with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as probe: + probe.bind(("::1", 0)) + except OSError: + pytest.skip("IPv6 loopback is unavailable") # ty: ignore[too-many-positional-arguments] + + server.server_host = "127.0.0.1,::1" + server.api_key = "test-multiple-addresses" + server.start() + + def check_address(host): + res = server.make_request("GET", "/health", host=host) + assert res.status_code == 200 + res = server.make_request("POST", "/v1/completions", data={}, host=host) + assert res.status_code == 401 + events = list(server.make_stream_request("POST", "/v1/completions", data={ + "prompt": "Once upon a time", + "max_tokens": 8, + "stream": True, + }, headers={"Authorization": f"Bearer {server.api_key}"}, host=host)) + assert len(events) > 1 + return True + + # parallel_function_calls swallows exceptions, a failed check leaves None in the results + results = parallel_function_calls([(check_address, (host,)) for host in ["127.0.0.1", "[::1]"]]) + assert all(results) + + def test_server_props(): global server server.start() diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 826aef2d5..3a50ae5c3 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -155,8 +155,6 @@ class ServerProcess: else: server_path = "../../../build/bin/llama-server" server_args = [ - "--host", - self.server_host, "--port", self.server_port, "--temp", @@ -164,6 +162,7 @@ class ServerProcess: "--seed", self.seed, ] + server_args.extend(["--host", self.server_host]) if self.offline: server_args.append("--offline") if self.model_file: @@ -365,6 +364,11 @@ class ServerProcess: if hasattr(self, '_log') and self._log != sys.stdout: self._log.close() + def make_url(self, path: str, host: str | None = None) -> str: + if host is None: + host = self.server_host.split(",")[0].strip() + return f"http://{host}:{self.server_port}{path}" + def make_request( self, method: str, @@ -372,8 +376,9 @@ class ServerProcess: data: dict | Any | None = None, headers: dict | None = None, timeout: float | None = DEFAULT_REQUEST_TIMEOUT, + host: str | None = None, ) -> ServerResponse: - url = f"http://{self.server_host}:{self.server_port}{path}" + url = self.make_url(path, host) parse_body = False if method == "GET": response = requests.get(url, headers=headers, timeout=timeout) @@ -407,8 +412,9 @@ class ServerProcess: path: str, data: dict | None = None, headers: dict | None = None, + host: str | None = None, ) -> Iterator[dict]: - url = f"http://{self.server_host}:{self.server_port}{path}" + url = self.make_url(path, host) if method == "POST": response = requests.post(url, headers=headers, json=data, stream=True) else: From 348f853b7adc7374a4dec989750eaa6ea563535e Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Tue, 22 Sep 2026 15:27:59 +0200 Subject: [PATCH 36/42] jinja: use const for statement::execute and ::visit (#29271) --- common/jinja/runtime.cpp | 53 +++++++++--------- common/jinja/runtime.h | 117 ++++++++++++++++++--------------------- 2 files changed, 80 insertions(+), 90 deletions(-) diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp index 252ab55de..227f6c094 100644 --- a/common/jinja/runtime.cpp +++ b/common/jinja/runtime.cpp @@ -51,7 +51,7 @@ static void ensure_key_type_allowed(const value & val) { } // execute with error handling -value statement::execute(context & ctx) { +value statement::execute(context & ctx) const { try { return execute_impl(ctx); } catch (const continue_statement::signal & /* ex */) { @@ -80,7 +80,7 @@ value statement::execute(context & ctx) { } } -value identifier::execute_impl(context & ctx) { +value identifier::execute_impl(context & ctx) const { auto it = ctx.get_val(val); auto builtins = global_builtins(); if (!it->is_undefined()) { @@ -98,7 +98,7 @@ value identifier::execute_impl(context & ctx) { } } -value object_literal::execute_impl(context & ctx) { +value object_literal::execute_impl(context & ctx) const { auto obj = mk_val(); for (const auto & pair : val) { value key = pair.first->execute(ctx); @@ -109,7 +109,7 @@ value object_literal::execute_impl(context & ctx) { return obj; } -value binary_expression::execute_impl(context & ctx) { +value binary_expression::execute_impl(context & ctx) const { value left_val = left->execute(ctx); // Logical operators @@ -317,9 +317,7 @@ static value try_builtin_func(context & ctx, const std::string & name, value & i 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; - +static value apply_filter(context & ctx, const statement_ptr & filter, value input) { JJ_DEBUG("Applying filter to %s", input->type().c_str()); auto set_filter_alias = [](auto & filter_id) { @@ -375,22 +373,21 @@ value filter_expression::execute_impl(context & ctx) { } } -value filter_statement::execute_impl(context & ctx) { +value filter_expression::execute_impl(context & ctx) const { + return apply_filter(ctx, filter, operand->execute(ctx)); +} + +value filter_statement::execute_impl(context & ctx) const { // 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; + return apply_filter(ctx, filter, parts); } -value test_expression::execute_impl(context & ctx) { +value test_expression::execute_impl(context & ctx) const { // NOTE: "value is something" translates to function call "test_is_something(value)" const auto & builtins = global_builtins(); @@ -439,7 +436,7 @@ value test_expression::execute_impl(context & ctx) { } } -value unary_expression::execute_impl(context & ctx) { +value unary_expression::execute_impl(context & ctx) const { value operand_val = argument->execute(ctx); JJ_DEBUG("Executing unary expression with operator '%s'", op.value.c_str()); @@ -458,7 +455,7 @@ value unary_expression::execute_impl(context & ctx) { throw std::runtime_error("Unknown unary operator '" + op.value + "'"); } -value if_statement::execute_impl(context & ctx) { +value if_statement::execute_impl(context & ctx) const { value test_val = test->execute(ctx); auto out = mk_val(); @@ -479,17 +476,17 @@ value if_statement::execute_impl(context & ctx) { return str; } -value for_statement::execute_impl(context & ctx) { +value for_statement::execute_impl(context & ctx) const { context scope(ctx); // new scope for loop variables - jinja::select_expression * select_expr = cast_stmt(iterable); + const jinja::select_expression * select_expr = cast_stmt(iterable); statement_ptr test_expr_nullptr; - statement_ptr & iter_expr = [&]() -> statement_ptr & { + const statement_ptr & iter_expr = [&]() -> const statement_ptr & { auto tmp = cast_stmt(iterable); return tmp ? tmp->lhs : iterable; }(); - statement_ptr & test_expr = [&]() -> statement_ptr & { + const statement_ptr & test_expr = [&]() -> const statement_ptr & { auto tmp = cast_stmt(iterable); return tmp ? tmp->test : test_expr_nullptr; }(); @@ -645,7 +642,7 @@ value for_statement::execute_impl(context & ctx) { return str; } -value set_statement::execute_impl(context & ctx) { +value set_statement::execute_impl(context & ctx) const { auto rhs = val ? val->execute(ctx) : exec_statements(body, ctx); if (is_stmt(assignee)) { @@ -744,7 +741,7 @@ static inline void bind_parameters(const std::string & name, const statements & } } -value macro_statement::execute_impl(context & ctx) { +value macro_statement::execute_impl(context & ctx) const { if (!is_stmt(this->name)) { throw std::runtime_error("Macro name must be an identifier"); } @@ -767,7 +764,7 @@ value macro_statement::execute_impl(context & ctx) { return mk_val(); } -value call_statement::execute_impl(context & ctx) { +value call_statement::execute_impl(context & ctx) const { auto call_expr = cast_stmt(this->call); if (!call_expr) { throw std::runtime_error("Call statement requires a valid call expression"); @@ -807,7 +804,7 @@ value call_statement::execute_impl(context & ctx) { return callee_func->invoke(args); } -value member_expression::execute_impl(context & ctx) { +value member_expression::execute_impl(context & ctx) const { value object = this->object->execute(ctx); value property; @@ -940,7 +937,7 @@ value member_expression::execute_impl(context & ctx) { return val; } -value call_expression::execute_impl(context & ctx) { +value call_expression::execute_impl(context & ctx) const { // gather arguments func_args args(ctx); for (auto & arg_stmt : this->args) { @@ -958,7 +955,7 @@ value call_expression::execute_impl(context & ctx) { return callee_func->invoke(args); } -value keyword_argument_expression::execute_impl(context & ctx) { +value keyword_argument_expression::execute_impl(context & ctx) const { if (!is_stmt(key)) { throw std::runtime_error("Keyword argument key must be identifiers"); } @@ -982,7 +979,7 @@ std::string runtime::debug_dump_program(const program & prog, const std::string return std::string(lvl * 2, ' '); }; - ctx.visitor = [&](bool is_leaf, statement * node, std::vector children) { + ctx.visitor = [&](bool is_leaf, const statement * node, std::vector children) { oss << indent(lvl) << node->type() << ":\n"; lvl++; if (is_leaf) { diff --git a/common/jinja/runtime.h b/common/jinja/runtime.h index 69bd683c6..bbd0c5cae 100644 --- a/common/jinja/runtime.h +++ b/common/jinja/runtime.h @@ -48,9 +48,9 @@ const T * cast_stmt(const statement_ptr & ptr) { void enable_debug(bool enable); // for visiting AST nodes -// function signature: void(bool is_leaf, statement * node, pair of ) -using visitor_pair = std::pair>; -using visitor_fn = std::function)>; +// function signature: void(bool is_leaf, const statement * node, pair of ) +using visitor_pair = std::pair>; +using visitor_fn = std::function)>; struct context { std::shared_ptr src; // for debugging; use shared_ptr to avoid copying on scope creation @@ -107,8 +107,8 @@ private: }; // utils for visiting AST nodes -static std::vector stmts_to_ptr(const statements & stmts) { - std::vector children; +static std::vector stmts_to_ptr(const statements & stmts) { + std::vector children; for (const auto & stmt : stmts) { children.push_back(stmt.get()); } @@ -117,17 +117,18 @@ static std::vector stmts_to_ptr(const statements & stmts) { /** * Base class for all nodes in the AST. + * The AST is shared between threads, so visit and execute must be const. */ struct statement { size_t pos; // position in source, for debugging virtual ~statement() = default; virtual std::string type() const { return "Statement"; } - virtual void visit(context & ctx) { ctx.visitor(true, this, {}); } + virtual void visit(context & ctx) const { ctx.visitor(true, this, {}); } // execute_impl must be overridden by derived classes - virtual value execute_impl(context &) { throw_exec_error(); } + virtual value execute_impl(context &) const { throw_exec_error(); } // execute is the public method to execute a statement with error handling - value execute(context &); + value execute(context &) const; private: [[noreturn]] void throw_exec_error() const { @@ -166,7 +167,7 @@ struct program : public statement { program() = default; explicit program(statements && body) : body(std::move(body)) {} std::string type() const override { return "Program"; } - [[noreturn]] value execute_impl(context &) override { + [[noreturn]] value execute_impl(context &) const override { throw std::runtime_error("Cannot execute program directly, use jinja::runtime instead"); } }; @@ -182,8 +183,8 @@ struct if_statement : public statement { } std::string type() const override { return "If"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"test", {test.get()}}, {"body", stmts_to_ptr(body)}, @@ -213,8 +214,8 @@ struct for_statement : public statement { } std::string type() const override { return "For"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"loopvar", {loopvar.get()}}, {"iterable", {iterable.get()}}, @@ -233,7 +234,7 @@ struct break_statement : public statement { } }; - [[noreturn]] value execute_impl(context &) override { + [[noreturn]] value execute_impl(context &) const override { throw break_statement::signal(); } }; @@ -247,7 +248,7 @@ struct continue_statement : public statement { } }; - [[noreturn]] value execute_impl(context &) override { + [[noreturn]] value execute_impl(context &) const override { throw continue_statement::signal(); } }; @@ -255,7 +256,7 @@ struct continue_statement : public statement { // do nothing struct noop_statement : public statement { std::string type() const override { return "Noop"; } - value execute_impl(context &) override { + value execute_impl(context &) const override { return mk_val(); } }; @@ -272,8 +273,8 @@ struct set_statement : public statement { } std::string type() const override { return "Set"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"assignee", {assignee.get()}}, {"value", {val.get()}}, @@ -294,8 +295,8 @@ struct macro_statement : public statement { } std::string type() const override { return "Macro"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"name", {name.get()}}, {"args", stmts_to_ptr(args)}, @@ -308,7 +309,7 @@ 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 { + value execute_impl(context &) const override { return mk_val(); } }; @@ -318,7 +319,7 @@ struct comment_statement : public statement { // 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 { + value execute_impl(context &) const override { return mk_val(); } }; @@ -334,8 +335,8 @@ struct member_expression : public expression { chk_type(this->property); } std::string type() const override { return "MemberExpression"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"object", {object.get()}}, {"property", {property.get()}} @@ -353,8 +354,8 @@ struct call_expression : public expression { for (const auto& arg : this->args) chk_type(arg); } std::string type() const override { return "CallExpression"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"callee", {callee.get()}}, {"args", stmts_to_ptr(args)} @@ -369,7 +370,7 @@ 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; + value execute_impl(context & ctx) const override; }; // Literals @@ -378,7 +379,7 @@ 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 { + value execute_impl(context &) const override { return mk_val(val); } }; @@ -387,7 +388,7 @@ 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 { + value execute_impl(context &) const override { return mk_val(val); } }; @@ -396,7 +397,7 @@ 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 { + value execute_impl(context &) const override { return mk_val(val); } }; @@ -407,7 +408,7 @@ struct array_literal : public expression { for (const auto& item : this->val) chk_type(item); } std::string type() const override { return "ArrayLiteral"; } - value execute_impl(context & ctx) override { + value execute_impl(context & ctx) const override { auto arr = mk_val(); for (const auto & item_stmt : val) { arr->push_back(item_stmt->execute(ctx)); @@ -422,7 +423,7 @@ struct tuple_literal : public expression { for (const auto& item : this->val) chk_type(item); } std::string type() const override { return "TupleLiteral"; } - value execute_impl(context & ctx) override { + value execute_impl(context & ctx) const override { auto arr = mk_val(); for (const auto & item_stmt : val) { arr->push_back(item_stmt->execute(ctx)); @@ -441,7 +442,7 @@ struct object_literal : public expression { } } std::string type() const override { return "ObjectLiteral"; } - value execute_impl(context & ctx) override; + value execute_impl(context & ctx) const override; }; // Complex Expressions @@ -462,8 +463,8 @@ struct binary_expression : public expression { chk_type(this->right); } std::string type() const override { return "BinaryExpression"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"left", {left.get()}}, {"right", {right.get()}} @@ -476,10 +477,7 @@ struct binary_expression : public expression { * 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) @@ -488,14 +486,9 @@ struct filter_expression : public expression { 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; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"operand", {operand.get()}}, {"filter", {filter.get()}} @@ -512,8 +505,8 @@ struct filter_statement : public statement { chk_type(this->filter); } std::string type() const override { return "FilterStatement"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"filter", {filter.get()}}, {"body", stmts_to_ptr(body)} @@ -537,14 +530,14 @@ struct select_expression : public expression { chk_type(this->test); } std::string type() const override { return "SelectExpression"; } - value execute_impl(context & ctx) override { + value execute_impl(context & ctx) const override { auto predicate = test->execute_impl(ctx); if (!predicate->as_bool()) { return mk_val(); } return lhs->execute_impl(ctx); } - void visit(context & ctx) override { + void visit(context & ctx) const override { ctx.visitor(false, this, { {"lhs", {lhs.get()}}, {"test", {test.get()}} @@ -567,8 +560,8 @@ struct test_expression : public expression { chk_type(this->test); } std::string type() const override { return "TestExpression"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"operand", {operand.get()}}, {"test", {test.get()}} @@ -588,8 +581,8 @@ struct unary_expression : public expression { chk_type(this->argument); } std::string type() const override { return "UnaryExpression"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"argument", {argument.get()}} }); @@ -608,10 +601,10 @@ struct slice_expression : public expression { chk_type(this->step_expr); } std::string type() const override { return "SliceExpression"; } - [[noreturn]] value execute_impl(context &) override { + [[noreturn]] value execute_impl(context &) const override { throw std::runtime_error("must be handled by MemberExpression"); } - void visit(context & ctx) override { + void visit(context & ctx) const override { ctx.visitor(false, this, { {"start_expr", {start_expr.get()}}, {"stop_expr", {stop_expr.get()}}, @@ -630,8 +623,8 @@ struct keyword_argument_expression : public expression { chk_type(this->val); } std::string type() const override { return "KeywordArgumentExpression"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"key", {key.get()}}, {"val", {val.get()}} @@ -645,7 +638,7 @@ struct spread_expression : public expression { chk_type(this->argument); } std::string type() const override { return "SpreadExpression"; } - void visit(context & ctx) override { + void visit(context & ctx) const override { ctx.visitor(false, this, { {"argument", {argument.get()}} }); @@ -663,8 +656,8 @@ struct call_statement : public statement { for (const auto & arg : this->caller_args) chk_type(arg); } std::string type() const override { return "CallStatement"; } - value execute_impl(context & ctx) override; - void visit(context & ctx) override { + value execute_impl(context & ctx) const override; + void visit(context & ctx) const override { ctx.visitor(false, this, { {"call", {call.get()}}, {"caller_args", stmts_to_ptr(caller_args)}, @@ -685,7 +678,7 @@ struct ternary_expression : public expression { chk_type(this->false_expr); } std::string type() const override { return "Ternary"; } - value execute_impl(context & ctx) override { + value execute_impl(context & ctx) const override { value cond_val = condition->execute(ctx); if (cond_val->as_bool()) { return true_expr->execute(ctx); @@ -693,7 +686,7 @@ struct ternary_expression : public expression { return false_expr->execute(ctx); } } - void visit(context & ctx) override { + void visit(context & ctx) const override { ctx.visitor(false, this, { {"condition", {condition.get()}}, {"true_expr", {true_expr.get()}}, From 9b421fa946def08fb0bd726db2f181081b3eb9d0 Mon Sep 17 00:00:00 2001 From: Eric Rodrigues Pires Date: Tue, 22 Sep 2026 10:40:24 -0300 Subject: [PATCH 37/42] ui : Accept WEBM video files (#28622) --- .../src/lib/constants/supported-file-types.constants.ts | 4 ++++ tools/ui/src/lib/enums/files.enums.ts | 9 ++++++--- tools/ui/src/lib/utils/file-type.ts | 1 + 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tools/ui/src/lib/constants/supported-file-types.constants.ts b/tools/ui/src/lib/constants/supported-file-types.constants.ts index a6bcefaa1..5675d682f 100644 --- a/tools/ui/src/lib/constants/supported-file-types.constants.ts +++ b/tools/ui/src/lib/constants/supported-file-types.constants.ts @@ -40,6 +40,10 @@ export const VIDEO_FILE_TYPES = { [FileTypeVideo.OGG]: { extensions: [FileExtensionVideo.OGG], mimeTypes: [MimeTypeVideo.OGG] + }, + [FileTypeVideo.WEBM]: { + extensions: [FileExtensionVideo.WEBM], + mimeTypes: [MimeTypeVideo.WEBM] } } as const; diff --git a/tools/ui/src/lib/enums/files.enums.ts b/tools/ui/src/lib/enums/files.enums.ts index 0185da478..b8cde6061 100644 --- a/tools/ui/src/lib/enums/files.enums.ts +++ b/tools/ui/src/lib/enums/files.enums.ts @@ -38,7 +38,8 @@ export enum FileTypeAudio { export enum FileTypeVideo { MP4 = 'mp4', - OGG = 'ogg' + OGG = 'ogg', + WEBM = 'webm' } export enum FileTypePdf { @@ -104,7 +105,8 @@ export enum FileExtensionAudio { export enum FileExtensionVideo { MP4 = '.mp4', - OGG = '.ogg' + OGG = '.ogg', + WEBM = '.webm' } export enum FileExtensionPdf { @@ -203,7 +205,8 @@ export enum MimeTypeAudio { export enum MimeTypeVideo { MP4 = 'video/mp4', - OGG = 'video/ogg' + OGG = 'video/ogg', + WEBM = 'video/webm' } export enum MimeTypeImage { diff --git a/tools/ui/src/lib/utils/file-type.ts b/tools/ui/src/lib/utils/file-type.ts index fd8828fc1..60f7c67d2 100644 --- a/tools/ui/src/lib/utils/file-type.ts +++ b/tools/ui/src/lib/utils/file-type.ts @@ -51,6 +51,7 @@ export function getFileTypeCategory(mimeType: string): FileTypeCategory | null { // Video case MimeTypeVideo.MP4: case MimeTypeVideo.OGG: + case MimeTypeVideo.WEBM: return FileTypeCategory.VIDEO; // PDF From c350a40bbd0ba0658793f0fc74a8b3b3ab65135e Mon Sep 17 00:00:00 2001 From: "David M. Rogers" Date: Tue, 22 Sep 2026 09:43:29 -0400 Subject: [PATCH 38/42] Performance tune for gemma4-26b-a4b flash attention shape. (#28450) --- ggml/src/ggml-sycl/fattn-tile.hpp | 4 ++++ tests/test-backend-ops.cpp | 3 +++ 2 files changed, 7 insertions(+) diff --git a/ggml/src/ggml-sycl/fattn-tile.hpp b/ggml/src/ggml-sycl/fattn-tile.hpp index 9ba529696..dcdcad88a 100644 --- a/ggml/src/ggml-sycl/fattn-tile.hpp +++ b/ggml/src/ggml-sycl/fattn-tile.hpp @@ -1173,6 +1173,10 @@ static void launch_fattn_tile_switch_ncols2(ggml_backend_sycl_context & ctx, ggm launch_fattn_tile_switch_ncols1(ctx, dst); return; } + if (use_gqa_opt && gqa_ratio % 8 == 0) { + launch_fattn_tile_switch_ncols1(ctx, dst); + return; + } if (use_gqa_opt && gqa_ratio % 4 == 0) { launch_fattn_tile_switch_ncols1(ctx, dst); return; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e4d04f13b..5f030c406 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -11334,6 +11334,9 @@ static std::vector> make_test_cases_perf() { // sparse decode at long context test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + // gemma-4-26b-a4b global-attn layers: head_count_kv=2, 16 query heads (gqa_ratio=8) + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 2, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 2, { 8, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 0)); test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 49152, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 2048)); From f95b0d95394d5e311ba8228689972843178c5e28 Mon Sep 17 00:00:00 2001 From: Bartowski <3266127+bartowski1182@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:54:45 +0200 Subject: [PATCH 39/42] ggml : IQ1_M build prefix sums once per block (#28706) --- ggml/src/ggml-quants.c | 117 ++++++++++++++++------------------------- 1 file changed, 44 insertions(+), 73 deletions(-) diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 1ebc50a76..55db802c0 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -4771,80 +4771,51 @@ static void quantize_row_iq1_m_impl(const float * GGML_RESTRICT x, void * GGML_R // 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]; + // prefix sums are kept per half of the block because each half can use a different sign (x_p or x_m) + // since v[0]-v[1] = v[1]-v[2] = -1 for both x_p and x_m, the 3-group sum for a split collapses to T*v[2] - px[i1] - px[i2] + { + float px[2][IQ1M_BLOCK_SIZE+1]; + float pw[2][IQ1M_BLOCK_SIZE+1]; + px[0][0] = px[1][0] = 0; + pw[0][0] = pw[1][0] = 0; + for (int j = 0; j < block_size; ++j) { + const int i = idx[2*j]; + const int h = i < block_size/2 ? 0 : 1; + px[h][j+1] = px[h][j] + weight[i]*xb[i]; + px[1-h][j+1] = px[1-h][j]; + pw[h][j+1] = pw[h][j] + weight[i]; + pw[1-h][j+1] = pw[1-h][j]; + } + const float txs[2] = {px[0][block_size], px[1][block_size]}; // total weight*x per half + const float tws[2] = {pw[0][block_size], pw[1][block_size]}; // total weight per half + const float p2 = x_p[2], m2 = x_m[2]; + const float cp1 = x_p[0]*x_p[0] - x_p[1]*x_p[1]; + const float cp2 = x_p[1]*x_p[1] - x_p[2]*x_p[2]; + const float cm1 = x_m[0]*x_m[0] - x_m[1]*x_m[1]; + const float cm2 = x_m[1]*x_m[1] - x_m[2]*x_m[2]; + for (int i1 = 0; i1 <= block_size; ++i1) { + for (int i2 = i1; i2 <= block_size; ++i2) { + float qx_p[2], qx_m[2], q2_p[2], q2_m[2]; + for (int h = 0; h < 2; ++h) { + const float sx = px[h][i1] + px[h][i2]; + qx_p[h] = txs[h]*p2 - sx; + qx_m[h] = txs[h]*m2 - sx; + q2_p[h] = tws[h]*p2*p2 + pw[h][i1]*cp1 + pw[h][i2]*cp2; + q2_m[h] = tws[h]*m2*m2 + pw[h][i1]*cm1 + pw[h][i2]*cm2; } - } - 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; + sumqx[0] = qx_p[0] + qx_p[1]; + sumqx[1] = qx_p[0] + qx_m[1]; + sumqx[2] = qx_m[0] + qx_p[1]; + sumqx[3] = qx_m[0] + qx_m[1]; + sumq2[0] = q2_p[0] + q2_p[1]; + sumq2[1] = q2_p[0] + q2_m[1]; + sumq2[2] = q2_m[0] + q2_p[1]; + sumq2[3] = q2_m[0] + q2_m[1]; + 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; + } } } } From 0f8a414b7587bc412e44611d4c9e2fea876449a6 Mon Sep 17 00:00:00 2001 From: Michael de Gans Date: Tue, 22 Sep 2026 17:32:28 +0200 Subject: [PATCH 40/42] metal : gate mul_mm_id src1 rescale behind ggml_prec (#29029) * metal : gate mul_mm_id src1 rescale behind ggml_prec Assisted-by: Claude Fable 5.1 * ggml-webgpu: reject MUL_MAT_ID when src1 precision is F32 * cuda/vulkan: reject MUL_MAT_ID in supports_op when src1 prec is F32 fix `supports_op` to return false for failing backends when the specified src1 precision is f32 Assisted-by: Claude Fable 5.1 --------- Co-authored-by: yomaytk --- ggml/src/ggml-cuda/ggml-cuda.cu | 3 +++ ggml/src/ggml-metal/ggml-metal-device.cpp | 6 +++++- ggml/src/ggml-metal/ggml-metal-ops.cpp | 13 ++++++++----- ggml/src/ggml-metal/kernels/mul_mm.metal | 5 +++-- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 +++ ggml/src/ggml-webgpu/ggml-webgpu.cpp | 3 +++ src/llama-graph.cpp | 4 ++++ tests/test-backend-ops.cpp | 14 +++++++++----- 8 files changed, 38 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a9038f1f4..c8b23b2d5 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5131,6 +5131,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { return false; } + if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) { + return false; + } #ifdef GGML_USE_MUSA const int cc = ggml_cuda_info().devices[dev_ctx->device].cc; if (b->ne[2]*b->ne[3] > 1 && !ggml_is_transposed(a) && !ggml_is_transposed(b)) { diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 2d3887588..dc6b695eb 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1156,14 +1156,18 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_m const bool bc_inp = op->src[0]->ne[0] % 32 != 0; + // src1 prec [TAG_GGML_PREC] + const bool amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32; + snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); - snprintf(name, 256, "%s_bci=%d", base, bc_inp); + snprintf(name, 256, "%s_bci=%d_amax=%d", base, bc_inp, amax); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); if (!res.pipeline) { ggml_metal_cv_t cv = ggml_metal_cv_init(); ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0); + ggml_metal_cv_set_bool(cv, amax, FC_MUL_MM + 6); res = ggml_metal_library_compile_pipeline(lib, base, name, cv); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 29db37f87..527892e18 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2719,9 +2719,12 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) { ggml_metal_buffer_id bid_amax = bid_ids; bid_amax.offs += ggml_metal_op_mul_mat_id_extra_ids(op); + // src1 prec [TAG_GGML_PREC] + const bool use_amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32; + // src1 rescale factors, computed before the matmul // ref: https://github.com/ggml-org/llama.cpp/pull/26223 - { + if (use_amax) { ggml_metal_kargs_mul_mm_id_amax args = { /*.ne00 =*/ ne10, /*.ne01 =*/ ne11, @@ -2779,17 +2782,17 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) { ggml_metal_op_concurrency_reset(ctx); - { + if (use_amax) { auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax(lib); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_buffer (enc, bid_amax, 0); ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1); - } - // the next kernel has to wait for the amax data - ggml_metal_op_concurrency_reset(ctx); + // the next kernel has to wait for the amax data + ggml_metal_op_concurrency_reset(ctx); + } { auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id(lib, op); diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index 71d991149..a25838f92 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -7,6 +7,7 @@ constant short FC_mul_mm_ne12 [[function_constant(FC_MUL_MM + 2)]]; constant short FC_mul_mm_ne13 [[function_constant(FC_MUL_MM + 3)]]; constant short FC_mul_mm_r2 [[function_constant(FC_MUL_MM + 4)]]; constant short FC_mul_mm_r3 [[function_constant(FC_MUL_MM + 5)]]; +constant bool FC_mul_mm_id_amax [[function_constant(FC_MUL_MM + 6)]]; // each block_q contains 16*nl weights #ifdef GGML_METAL_HAS_TENSOR @@ -584,8 +585,8 @@ kernel void kernel_mul_mm_id( const short lb1 = (short) tiitg/NL1; // 0 .. NR1-1, this thread's row of the B tile // power-of-two rescaling - const float s1_inv = ((device const float *) amax)[0]; - const float s1_scale = ((device const float *) amax)[1]; + const float s1_inv = FC_mul_mm_id_amax ? ((device const float *) amax)[0] : 1.0f; + const float s1_scale = FC_mul_mm_id_amax ? ((device const float *) amax)[1] : 1.0f; #ifndef GGML_METAL_HAS_TENSOR S0_8x8 ma[4]; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index da0e24fcd..f7e27703f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -14953,6 +14953,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm // If there's not enough shared memory for row_ids and the result tile, fallback to CPU return false; } + if (ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) { + return false; + } } switch (src0_type) { case GGML_TYPE_F32: diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 86f0e958a..9c5dc768e 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4506,6 +4506,9 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const default: break; } + if (ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) { + supports_op = false; + } break; case GGML_OP_FLASH_ATTN_EXT: { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 02ae8bd92..07ca49ad0 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2302,6 +2302,10 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] + if (arch == LLM_ARCH_MISTRAL4) { + // src1 can exceed F16 range + ggml_prec_set_src(experts, GGML_PREC_F32, 1); + } cb(experts, "ffn_moe_down", il); if (down_exps_s) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 5f030c406..505112f6b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5186,6 +5186,11 @@ struct test_mul_mat_id : public test_case { ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); ggml_set_name(out, "out"); + if (amax > 65504.0f) { + // src1 exceeds F16 range + ggml_prec_set_src(out, GGML_PREC_F32, 1); + } + return out; } @@ -10185,11 +10190,10 @@ static std::vector> make_test_cases_eval() { } // test src1 f16 overflow - // TODO: https://github.com/ggml-org/llama.cpp/pull/26223#issuecomment-5585815365 - //for (int n : {16, 32, 64}) { - // test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_K, GGML_TYPE_F32, 128, 4, false, 4096, n, 2048, 1e5f)); - // test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8, 2, false, 512, n, 256, 1e5f)); - //} + for (int n : {16, 32, 64}) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_K, GGML_TYPE_F32, 128, 4, false, 4096, n, 2048, 1e5f)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8, 2, false, 512, n, 256, 1e5f)); + } for (ggml_type type_a : base_types) { for (ggml_type type_b : {GGML_TYPE_F32 /*, GGML_TYPE_F16 */}) { From 73c941b11165cc0f7a36ba17380e79e8fe9797dc Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Tue, 22 Sep 2026 17:53:05 +0200 Subject: [PATCH 41/42] mtmd: add various sanity checks (#29276) --- tools/mtmd/clip.cpp | 27 ++++++++++++++++++++------- tools/mtmd/mtmd-image.cpp | 20 +++++++++++++------- 2 files changed, 33 insertions(+), 14 deletions(-) diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 1783c7ddf..feceb7ff7 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1408,9 +1408,14 @@ struct clip_model_loader { } // Load the vision/audio feature layer indices if they are explicitly provided - // NOTE: gguf conversions should standardize the values of the vision feature layer to - // be non-negative, since we use -1 to mark values as unset here. + // NOTE: gguf conversions should standardize the values of the vision feature layer to be non-negative, since we use -1 to mark values as unset here. get_arr_int(string_format(KEY_FEATURE_LAYERS, prefix), hparams.feature_layers, false); + for (const auto & v : hparams.feature_layers) { + if (v > (int) hparams.n_layer) { + throw std::runtime_error(string_format("%s: feature layer index %d is out of range (n_layer: %d)", + __func__, v, hparams.n_layer)); + } + } // model-specific params switch (model.proj_type) { @@ -1456,7 +1461,12 @@ struct clip_model_loader { std::vector wa_layer_indexes_vec; get_arr_int(KEY_WIN_ATTN_LAYER_INDEXES, wa_layer_indexes_vec, false); if (!wa_layer_indexes_vec.empty()) { - hparams.insert_layer_id = wa_layer_indexes_vec[0]; + const int insert_lid = wa_layer_indexes_vec[0]; + if (insert_lid < 0 || insert_lid >= (int) hparams.n_layer) { + throw std::runtime_error(string_format("%s: layer index %d is out of range (n_layer: %d)", + __func__, insert_lid, hparams.n_layer)); + } + hparams.insert_layer_id = insert_lid; } } break; case PROJECTOR_TYPE_INTERNVL: @@ -3226,6 +3236,7 @@ struct clip_model_loader { model.pos_embed = get_tensor(string_format(TN_SAM_POS_EMBD, "weight")); model.patch_embed_proj_w = get_tensor(string_format(TN_SAM_PATCH_EMBD, "weight")); model.patch_embed_proj_b = get_tensor(string_format(TN_SAM_PATCH_EMBD, "bias")); + model.n_sam_layers = hparams.sam_n_layer; model.sam_layers.resize(model.n_sam_layers); for (int il = 0; il < model.n_sam_layers; ++il) { auto & layer = model.sam_layers[il]; @@ -4652,8 +4663,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // -> https://huggingface.co/HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit // -> https://huggingface.co/HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit/blob/d66538faeba44480d0bfaa42145eef26f9423199/modeling_siglip.py#L316 std::vector positions(pos_h * pos_w); - int bucket_coords_h[1024]; - int bucket_coords_w[1024]; + // note: sized by the actual patch counts; a tall/wide image produces more + // than 1024 patches per side and a fixed [1024] array would be overrun + std::vector bucket_coords_h(pos_h); + std::vector bucket_coords_w(pos_w); for (int i = 0; i < pos_h; i++){ bucket_coords_h[i] = std::floor(70.0*i/pos_h); } @@ -4696,8 +4709,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // SigLIP position buckets (same as resampler path) std::vector positions(pos_h * pos_w); - int bucket_coords_h[1024]; - int bucket_coords_w[1024]; + std::vector bucket_coords_h(pos_h); + std::vector bucket_coords_w(pos_w); for (int i = 0; i < pos_h; i++){ bucket_coords_h[i] = std::floor(70.0*i/pos_h); } diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index c11d35c87..9b2e98862 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -294,8 +294,8 @@ private: support = filter_support * filterscale; // Widen filter when downsampling ksize = static_cast(std::ceil(support)) * 2 + 1; // Total pixels in kernel - std::vector pre_weights(outSize * ksize); // Temporary weights - bounds.resize(outSize * 2); + std::vector pre_weights((size_t) outSize * ksize); // Temporary weights + bounds.resize((size_t) outSize * 2); // For each output pixel, compute its filter coefficients @@ -322,20 +322,20 @@ private: for (x = 0; x < xmax; x++) { // Distance from input pixel center to output pixel center in input space double w = resample_filter((x + xmin - center + 0.5) * ss); - pre_weights[xx * ksize + x] = w; + pre_weights[(size_t) xx * ksize + x] = w; ww += w; // Accumulate for normalization } // Normalize weights to sum to 1.0 (preserves brightness) for (x = 0; x < xmax; x++) { if (ww != 0.0) { - pre_weights[xx * ksize + x] /= ww; + pre_weights[(size_t) xx * ksize + x] /= ww; } } // Zero-pad remaining kernel positions for (; x < ksize; x++) { - pre_weights[xx * ksize + x] = 0; + pre_weights[(size_t) xx * ksize + x] = 0; } // Store input pixel range for this output pixel @@ -345,11 +345,11 @@ private: // Convert floating-point coefficients to fixed-point integers // Formula: int32 = round(float * 2^PRECISION_BITS) - weights.resize(outSize * ksize); + weights.resize((size_t) outSize * ksize); const double fxp_scale = std::ldexp(1.0, PRECISION_BITS); // 1.0 * 2^PRECISION_BITS - for (int i = 0; i < outSize * ksize; i++) { + for (size_t i = 0; i < (size_t) outSize * ksize; i++) { // Pillow adds +/- 0.5 then truncates toward zero; std::round would round twice const double rounded = pre_weights[i] * fxp_scale + (pre_weights[i] < 0 ? -0.5 : 0.5); weights[i] = static_cast(rounded); @@ -442,6 +442,12 @@ private: const int src_width = img.get_size().width; const int src_height = img.get_size().height; + // sanity check on the target size + if (target_width <= 0 || target_width > 65536 || target_height <= 0 || target_height > 65536) { + throw std::runtime_error("resize target " + std::to_string(target_width) + "x" + + std::to_string(target_height) + " is out of range (max 65536)"); + } + bool need_horizontal = (target_width != src_width); bool need_vertical = (target_height != src_height); From 4ceb1719101f32637b841206c172f3f058ffc182 Mon Sep 17 00:00:00 2001 From: "Jiang, Fish" Date: Wed, 23 Sep 2026 00:05:37 +0800 Subject: [PATCH 42/42] vulkan: add Intel Xe flash attention optimization kernels (2/3, Xe-LPG Plus/Xe2/Xe3) (#24406) * vulkan : Intel FA kernel optimization for split k path * vulkan : Host code update for Intel split k FA kernel path selection, fix A770 Linux op test failures * vulkan : use symmetric coopMatMulAdd() in flash_attn_decode_phase_1 shader to resolve test op failre on A770 Linux with 26.2.3 mesa driver * vulkan : fix editorconfig issue in flash_attn_decode_phase_2.comp --------- Co-authored-by: Liu, Russell --- .../ggml-vulkan/ggml-vulkan-push-constants.h | 18 + ggml/src/ggml-vulkan/ggml-vulkan-types.h | 1 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 146 ++++++- .../flash_attn_decode_phase_1.comp | 263 +++++++++++ .../flash_attn_decode_phase_2.comp | 408 ++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 4 + 6 files changed, 839 insertions(+), 1 deletion(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 68b3200b3..8446e313c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h @@ -124,6 +124,24 @@ struct vk_flash_attn_push_constants { static_assert(sizeof(vk_flash_attn_push_constants) <= 128, "sizeof(vk_flash_attn_push_constants) must be <= 128"); +struct vk_fa_xe_opt_push_constants { + uint32_t kv_seq_len; + uint32_t activation_length; + uint32_t q_head; + uint32_t kv_head; + uint32_t qk_ratio; + uint32_t qk_sub_groups; + uint32_t flag; + uint32_t nbkv_tok; + uint32_t nbkv_head; + uint32_t batch_stride_q; + uint32_t batch_stride_k; + uint32_t batch_stride_v; + uint32_t batch_stride_m; + uint32_t batch_stride_o; + float softmax_scale; +}; + struct vk_op_push_constants { uint32_t KX; uint32_t KY; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index 67e3361ed..5df1c3900 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -996,6 +996,7 @@ struct vk_device_struct { bool fa_sparse_compact_use_subgroups; vk_pipeline pipeline_flash_attn_split_k_reduce; + std::map, std::pair> pipeline_xe_fa_decode_dual_phases; vk_pipeline pipeline_count_experts; // [2] is for whether to take n_experts from spec constant (0) or push constant (1) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f7e27703f..1a83ac320 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2973,6 +2973,46 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_flash_attn_split_k_reduce, "fa_split_k_reduce", fa_split_k_reduce_len, fa_split_k_reduce_data, "main", 3, sizeof(vk_op_flash_attn_split_k_reduce_push_constants), {1, device->subgroup_size, 1}, {device->subgroup_size}, 1, true); + if (device->vendor_id == VK_VENDOR_ID_INTEL && (device->architecture == INTEL_XE2 || (device->architecture == INTEL_XE1 && device->coopmat_support && device->uma))) { + auto upper_power_of_2 = [&](uint32_t in) { + GGML_ASSERT(in != 0); + if (in <= 1) return 1u; + uint32_t ret = in - 1; + ret |= ret >> 1; + ret |= ret >> 2; + ret |= ret >> 4; + ret |= ret >> 8; + ret |= ret >> 16; + return ret + 1; + }; + + uint32_t xe_native_sub_group_size = 16; + if (device->architecture == INTEL_XE1) { + xe_native_sub_group_size = 8; + } + + for (auto& it : device->pipeline_xe_fa_decode_dual_phases) { + const uint32_t split_p_chunk = 32; + auto HdQk = it.first; + auto& pipelines = it.second; + uint32_t head_dim_qk = std::get<0>(HdQk); + uint32_t head_dim_pv = std::get<1>(HdQk); + uint32_t gqa_ratio = std::get<2>(HdQk); + uint32_t q_len = std::get<3>(HdQk); + const uint32_t out_dim_per_wg = gqa_ratio > 16 ? 8 : 16; + uint32_t aligned_q_len = upper_power_of_2(q_len); + uint32_t group_sz_ph1 = std::min(std::max(aligned_q_len * xe_native_sub_group_size, 64u), 256u); + uint32_t out_per_wg_ph1 = std::min(q_len, 256u / xe_native_sub_group_size); + uint32_t aligned_gqa_ratio = upper_power_of_2(gqa_ratio); + uint32_t split_p_per_iter_ph2 = 256; + uint32_t split_p_per_warp = 16; + uint32_t group_sz_ph2 = (split_p_per_iter_ph2 / split_p_per_warp) * xe_native_sub_group_size; + uint32_t out_per_wg_ph2 = std::min(std::max(16u / aligned_gqa_ratio, 1u), q_len); + ggml_vk_create_pipeline(device, pipelines.first, "xe_fa_decode_ph1", fa_decode_ph1_cm1_len, fa_decode_ph1_cm1_data, "main", 5, sizeof(vk_fa_xe_opt_push_constants), { 1, 32, 1 }, { group_sz_ph1, gqa_ratio, head_dim_qk, xe_native_sub_group_size, split_p_chunk, out_per_wg_ph1 }, 1, false, true, xe_native_sub_group_size); + ggml_vk_create_pipeline(device, pipelines.second, "xe_fa_decode_ph2", fa_decode_ph2_cm1_len, fa_decode_ph2_cm1_data, "main", 5, sizeof(vk_fa_xe_opt_push_constants), { 1, 1, 1 }, { group_sz_ph2, gqa_ratio, head_dim_pv, out_per_wg_ph2, xe_native_sub_group_size, split_p_per_iter_ph2, split_p_chunk, out_dim_per_wg }, 1, false, true, xe_native_sub_group_size); + } + } + for (auto &it : device->pipeline_fa_mask_opt) { auto BrBc = it.first; ggml_vk_create_pipeline(device, it.second, "fa_mask_opt", fa_mask_opt_len, fa_mask_opt_data, "main", 2, sizeof(vk_op_flash_attn_mask_opt_push_constants), {1, 1, 1}, {128, 128 / device->subgroup_size, BrBc.first, BrBc.second}, 1, true, true, device->subgroup_size); @@ -7899,6 +7939,18 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const vk_pipeline pipeline = nullptr; + bool xe_fa_opt = false; + bool fa_copy_qstate = false; + bool xe_fa_supported_platform = + (ctx->device.get()->architecture == INTEL_XE2 && ctx->device.get()->properties.deviceID != 0xFD80 && ctx->device.get()->properties.deviceID != 0xFD81) || + (ctx->device.get()->architecture == INTEL_XE1 && ctx->device.get()->coopmat_support && ctx->device.get()->uma); + bool xe_fa_supported_usage = neq0 % 32 == 0 && nev0 % 16 == 0 && q->nb[1] > q->nb[2] && k->nb[1] > k->nb[2] && v->nb[1] > v->nb[2] && mask != nullptr; + bool xe_fa_supported_dtype = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && (mask != nullptr && mask->type == GGML_TYPE_F16); + std::pair xe_fa_pipeline_dual_phases = { nullptr , nullptr }; + vk_pipeline xe_fa_pipeline = nullptr; + size_t size_p = 0; + size_t size_group_max = 0; + { std::lock_guard guard(ctx->device->compile_mutex); auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16; @@ -7956,6 +8008,37 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const // of "align", so recompute split_k based on that. split_kv = ROUNDUP_POW2(std::max(1u, KV / split_k), alignment); split_k = CEIL_DIV(KV, split_kv); + xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype; + if (xe_fa_opt) { + std::lock_guard guard(ctx->device->compile_mutex); + const uint32_t split_p_size = 32; + const size_t max_dim = (nek1 + split_p_size - 1) / split_p_size; + const size_t p_dim = max_dim * split_p_size; + auto& pipelines = ctx->device->pipeline_xe_fa_decode_dual_phases; + auto it = pipelines.find({ (uint32_t)neq0, (uint32_t)nev0, qk_ratio, (uint32_t)neq1 }); + if (it != pipelines.end()) { + xe_fa_pipeline_dual_phases = it->second; + } else { + pipelines[{(uint32_t)neq0, (uint32_t)nev0, qk_ratio, (uint32_t)neq1}] = xe_fa_pipeline_dual_phases = std::make_pair(std::make_shared(), std::make_shared()); + } + + size_p = neq1 * neq2 * p_dim * neq3 * sizeof(ggml_fp16_t); + size_group_max = neq1 * neq2 * max_dim * neq3 * sizeof(float); + size_t temp_size = ggml_nelements(q) * sizeof(ggml_fp16_t) + size_p + size_group_max; + fa_copy_qstate = true; + if (ctx->prealloc_size_x < temp_size) { + ctx->prealloc_size_x = temp_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + } + } + + if (xe_fa_opt == true) { + use_mask_opt = false; } // Reserve space for split_k temporaries. For each split x batch, we need to store the O matrix (D x ne1) @@ -8111,7 +8194,68 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const mask_n_head_log2, m0, m1, gqa_ratio, split_kv, split_k }; - if (split_k > 1) { + if (xe_fa_opt && split_k > 1) { + auto upper_power_of_2 = [&](uint32_t in) { + GGML_ASSERT(in != 0); + if (in <= 1) return 1u; + uint32_t ret = in - 1; + ret |= ret >> 1; + ret |= ret >> 2; + ret |= ret >> 4; + ret |= ret >> 8; + ret |= ret >> 16; + return ret + 1; + }; + auto to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, q->type); + const uint32_t out_dim_per_wg = qk_ratio > 16 ? 8 : 16; + size_t x_ne = ggml_nelements(q); + size_t temp_buf_offset = 0; + uint32_t head_stride_k = uint32_t(nbk2 / ggml_type_size(k->type)); + uint32_t head_stride_v = uint32_t(nbv2 / ggml_type_size(v->type)); + uint32_t batch_stride_q = uint32_t(nbq3 / ggml_type_size(q->type)); + uint32_t batch_stride_k = uint32_t(nbk3 / ggml_type_size(k->type)); + uint32_t batch_stride_v = uint32_t(nbv3 / ggml_type_size(v->type)); + uint32_t batch_stride_m = mask ? uint32_t(mask->nb[3] / ggml_type_size(mask->type)) : 0u; + uint32_t batch_stride_o = uint32_t(nb3 / ggml_type_size(dst->type)); + vk_fa_xe_opt_push_constants pc_ph1 = { (uint32_t)nek1, (uint32_t)neq1, (uint32_t)neq2, (uint32_t)nek2, qk_ratio, 1, (sinks != nullptr) ? 1u : 0u, (uint32_t)k_stride, head_stride_k, + batch_stride_q, batch_stride_k, batch_stride_v, batch_stride_m, batch_stride_o, scale }; + vk_fa_xe_opt_push_constants pc_ph2 = pc_ph1; + pc_ph2.nbkv_tok = v_stride; + pc_ph2.nbkv_head = head_stride_v; + vk_subbuffer q_temp_buf = fa_copy_qstate ? ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset) : q_buf; + temp_buf_offset += fa_copy_qstate ? x_ne * sizeof(ggml_fp16_t) : 0; + vk_subbuffer p_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset); + temp_buf_offset += size_p; + vk_subbuffer max_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset); + temp_buf_offset += size_group_max; + uint32_t xe_native_sub_group_size = ctx->device.get()->architecture == INTEL_XE1 ? 8 : 16; + uint32_t aligned_gqa_ratio = upper_power_of_2(qk_ratio); + uint32_t out_per_wg_ph1 = std::min(256u / xe_native_sub_group_size, (uint32_t)neq1); + uint32_t out_per_wg_ph2 = std::min(std::max(16u / aligned_gqa_ratio, 1u), (uint32_t)neq1); + uint32_t ph1_wg = ((neq1 + out_per_wg_ph1 - 1) / out_per_wg_ph1) * nek2; + uint32_t ph2_wg = ((neq1 + out_per_wg_ph2 - 1) / out_per_wg_ph2) * ne0 / out_dim_per_wg; + if (fa_copy_qstate) { + const std::vector pc_cpy_fp16 = + { (uint32_t)q->ne[0], (uint32_t)q->ne[1], (uint32_t)q->ne[2], (uint32_t)q->ne[3], (uint32_t)(x_ne) }; + ggml_vk_sync_buffers(ctx, subctx); + ggml_pipeline_request_descriptor_sets(ctx, to_fp16_vk_0, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { q_buf, q_temp_buf }, pc_cpy_fp16, { (uint32_t)(x_ne), 1, 1 }); + } + + ggml_vk_sync_buffers(ctx, subctx); + ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.first, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.first, + { q_temp_buf, k_buf, mask_buf, p_temp_buf, max_temp_buf }, + pc_ph1, { (uint32_t)ph1_wg, (uint32_t)nek1, (uint32_t)neq3 }); + + ggml_vk_sync_buffers(ctx, subctx); + ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.second, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.second, + { p_temp_buf, v_buf, max_temp_buf, sinks_buf, dst_buf }, + pc_ph2, { (uint32_t)ph2_wg, (uint32_t)nev2, (uint32_t)neq3 }); + + ctx->prealloc_x_need_sync = true; + } else if (split_k > 1) { ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); if (ctx->prealloc_split_k_need_sync) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp new file mode 100644 index 000000000..b5f95aaa0 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp @@ -0,0 +1,263 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable +#extension GL_KHR_cooperative_matrix : enable +#extension GL_EXT_shared_memory_block : enable + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer Q {float16_t qState[];}; +layout (binding = 1) readonly buffer K_VEC4 {f16vec4 kStateVec4[];}; +layout (binding = 2) buffer MASK_F16 {float16_t mState_f16[];}; +layout (binding = 3) buffer P_FP16 {float16_t matP_f16[];}; +layout (binding = 4) buffer OUT_MAX {float out_max_f32[];}; + +layout (push_constant) uniform parameter +{ + uint kvSeqLen; + uint activationLength; + uint qHead; + uint kvHead; + uint qkRatio; + uint qkSubGroups; + uint flag; + uint kvStride1; + uint kvStride2; + uint batchStrideQ; + uint batchStrideK; + uint batchStrideV; + uint batchStrideM; + uint batchStrideO; + float softMaxScale; +} p; + +layout (constant_id = 0) const uint GROUPSIZE = 128; +layout (constant_id = 1) const uint GQA_RATIO = 8; +layout (constant_id = 2) const uint HEAD_DIM = 128; +layout (constant_id = 3) const uint WARPSIZE = 16; +layout (constant_id = 4) const uint MATP_REDUCE = 32; +layout (constant_id = 5) const uint N_TOK = 1; +layout (constant_id = 6) const uint COOP_MAT_P_PER_LOOP = 4; + +#define MAX_HEADS 8 + +#define TN WARPSIZE +#define TM 8 +#define TK 16 +#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE) +#define MATP_PER_LOOP (COOP_MAT_P_PER_LOOP * TM) +#define P_LOOP_COUNT (MATP_REDUCE / MATP_PER_LOOP) + +#define COOP_MAT_Q_PER_TOKEN ((GQA_RATIO + TN - 1) / TN) +#define COOP_MAT_P_M COOP_MAT_Q_PER_TOKEN +#define COOP_MAT_P_N (MATP_REDUCE / TM) +#define SLM_PV_SIZE (MATP_REDUCE * COOP_MAT_P_M * TN) +#define SLM_MASK_SIZE (N_TOK * MATP_REDUCE) +#define SLM_POOL_SIZE_K (MATP_PER_LOOP * HEAD_DIM) +#define K_LOAD_PER_LOOP (GROUPSIZE * 4) +#define HEAD_DIM_VEC4 (HEAD_DIM / 4) +#define SLM_CHUNK_SIZE (TK / 4) +#define K_LOAD_LOOPS ((SLM_POOL_SIZE_K + K_LOAD_PER_LOOP - 1) / K_LOAD_PER_LOOP) +#define O_COUNT ((GQA_RATIO + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT) + +shared slm_pool_block { + float slm_pool_pv[SLM_PV_SIZE + SLM_MASK_SIZE]; +} slm_pool_f32; + +shared slm_pool_alias_block { + float16_t slm_pool_k[SLM_POOL_SIZE_K]; +} slm_pool_f16; + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint kHeadIdx = gl_WorkGroupID.x % p.kvHead; + const uint outGroupIdx = gl_WorkGroupID.x / p.kvHead; + const uint v = gl_WorkGroupID.y; + const uint d = gl_WorkGroupID.z; + const uint localLinearId = gl_SubgroupID; + const uint wgLane = localLinearId * WARPSIZE + lane; + const uint qDim = p.qHead * HEAD_DIM; + const uint kvDim = p.kvStride1; + const uint maskDim = p.kvSeqLen; + const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE; + const uint pDim = maxDim * MATP_REDUCE; + const uint tokFlatIdx = localLinearId + outGroupIdx * N_TOK; + uint offsetBaseQ = min(tokFlatIdx, p.activationLength - 1) * qDim; + offsetBaseQ = offsetBaseQ + d * p.batchStrideQ + kHeadIdx * HEAD_DIM * GQA_RATIO; + const uint offsetBaseK = (d * p.batchStrideK + (v * MATP_REDUCE) * kvDim + kHeadIdx * p.kvStride2) / 4; + uint offsetOut = d * p.qHead * p.activationLength * pDim + v * MATP_REDUCE + kHeadIdx * GQA_RATIO * pDim + (localLinearId * O_COUNT + outGroupIdx * N_TOK * p.qHead) * pDim + lane; + uint offsetMax = d * p.qHead * p.activationLength * maxDim + v + kHeadIdx * GQA_RATIO * maxDim + (localLinearId * O_COUNT + outGroupIdx * N_TOK * p.qHead) * maxDim; + const uint offsetSlmLoadPv = (localLinearId * O_COUNT * MATP_REDUCE + lane); + const uint offsetBaseM = v * MATP_REDUCE + lane; + const float fp32Min = uintBitsToFloat(0xFEFFFFFF); + + const uint loopCount = HEAD_DIM / TK; + float maskFp32[MATP_REDUCE / WARPSIZE]; + + if (tokFlatIdx < p.activationLength) { + [[unroll]] for (uint mk = 0; mk < MATP_REDUCE / WARPSIZE; mk++) { + const uint maskOffset = mk * WARPSIZE + offsetBaseM; + if (maskOffset < maskDim) { + maskFp32[mk] = float(mState_f16[d * p.batchStrideM + tokFlatIdx * maskDim + maskOffset]); + } else { + maskFp32[mk] = fp32Min; + } + } + } + + coopmat matP[COOP_MAT_P_M][COOP_MAT_P_N]; + + [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) { + [[unroll]] for (uint np = 0; np < COOP_MAT_P_N; np++) { + matP[mp][np] = coopmat(0.0f); + } + } + + [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) { + const uint flatOffset = kLoad * GROUPSIZE + wgLane; + const uint kRowIdx = flatOffset / HEAD_DIM_VEC4; + const uint kColIdx = flatOffset % HEAD_DIM_VEC4; + const uint slmChunkCol = kColIdx % SLM_CHUNK_SIZE; + const uint slmChunkRow = kColIdx / SLM_CHUNK_SIZE; + const uint offsetK = offsetBaseK + kRowIdx * kvDim / 4 + kColIdx; + const uint offsetSlmK = kRowIdx * TK + slmChunkRow * TK * MATP_PER_LOOP + slmChunkCol * 4; + slm_pool_f16.slm_pool_k[offsetSlmK + 0] = kStateVec4[offsetK].x; + slm_pool_f16.slm_pool_k[offsetSlmK + 1] = kStateVec4[offsetK].y; + slm_pool_f16.slm_pool_k[offsetSlmK + 2] = kStateVec4[offsetK].z; + slm_pool_f16.slm_pool_k[offsetSlmK + 3] = kStateVec4[offsetK].w; + } + + [[unroll]] for (uint pLoop = 0; pLoop < P_LOOP_COUNT; pLoop++) { + f16vec4 kTemp[K_LOAD_LOOPS]; + + if (pLoop + 1 < P_LOOP_COUNT) { + [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) { + const uint flatOffset = kLoad * GROUPSIZE + wgLane; + const uint kRowIdx = flatOffset / HEAD_DIM_VEC4 + (pLoop + 1) * MATP_PER_LOOP; + const uint kColIdx = flatOffset % HEAD_DIM_VEC4; + const uint offsetK = offsetBaseK + kRowIdx * kvDim / 4 + kColIdx; + kTemp[kLoad] = kStateVec4[offsetK]; + } + } + + barrier(); + if (localLinearId < N_TOK) { + [[unroll]] for (uint loop = 0; loop < loopCount; loop++) { + coopmat matQ[COOP_MAT_P_M]; + coopmat matK[COOP_MAT_P_PER_LOOP]; + + [[unroll]] for (uint mq = 0; mq < COOP_MAT_P_M; mq++) { + coopMatLoad( + matQ[mq], + qState, + offsetBaseQ + mq * TN * HEAD_DIM + loop * TK, + HEAD_DIM, + gl_CooperativeMatrixLayoutColumnMajor); + } + + [[unroll]] for (uint np = 0; np < COOP_MAT_P_PER_LOOP; np++) { + coopMatLoad( + matK[np], + slm_pool_f16.slm_pool_k, + loop * TK * MATP_PER_LOOP + np * TM * TK, + TK, + gl_CooperativeMatrixLayoutRowMajor); + } + + [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) { + [[unroll]] for (uint np = 0; np < COOP_MAT_P_PER_LOOP; np++) { + matP[mp][pLoop * COOP_MAT_P_PER_LOOP + np] = coopMatMulAdd(matK[np], matQ[mp], matP[mp][pLoop * COOP_MAT_P_PER_LOOP + np]); + } + } + } + } + + barrier(); + + if (pLoop + 1 < P_LOOP_COUNT) { + [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) { + const uint flatOffset = kLoad * GROUPSIZE + wgLane; + const uint kRowIdx = flatOffset / HEAD_DIM_VEC4; + const uint kColIdx = flatOffset % HEAD_DIM_VEC4; + const uint slmChunkCol = kColIdx % SLM_CHUNK_SIZE; + const uint slmChunkRow = kColIdx / SLM_CHUNK_SIZE; + const uint offsetSlmK = kRowIdx * TK + slmChunkRow * TK * MATP_PER_LOOP + slmChunkCol * 4; + slm_pool_f16.slm_pool_k[offsetSlmK + 0] = kTemp[kLoad].x; + slm_pool_f16.slm_pool_k[offsetSlmK + 1] = kTemp[kLoad].y; + slm_pool_f16.slm_pool_k[offsetSlmK + 2] = kTemp[kLoad].z; + slm_pool_f16.slm_pool_k[offsetSlmK + 3] = kTemp[kLoad].w; + } + } + } + + barrier(); + + if (tokFlatIdx < p.activationLength) { + [[unroll]] for (uint mk = 0; mk < MATP_REDUCE / WARPSIZE; mk++) { + slm_pool_f32.slm_pool_pv[SLM_PV_SIZE + localLinearId * MATP_REDUCE + mk * WARPSIZE + lane] = maskFp32[mk]; + } + } + + [[unroll]] for (uint oLoop = 0; oLoop < N_TOK; oLoop++) { + if (oLoop + outGroupIdx * N_TOK < p.activationLength) { + if (localLinearId == oLoop) { + [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) { + [[unroll]] for (uint np = 0; np < COOP_MAT_P_N; np++) { + coopMatStore(matP[mp][np], slm_pool_f32.slm_pool_pv, mp * MATP_REDUCE * TN + np * TM, MATP_REDUCE, gl_CooperativeMatrixLayoutColumnMajor); + } + } + } + + barrier(); + + [[unroll]] for (uint maskIdx = 0; maskIdx < MATP_REDUCE / WARPSIZE; maskIdx++) { + maskFp32[maskIdx] = slm_pool_f32.slm_pool_pv[SLM_PV_SIZE + oLoop * MATP_REDUCE + maskIdx * WARPSIZE + lane]; + } + + float fp32O[O_COUNT][MATP_REDUCE / WARPSIZE]; + float maxOut[O_COUNT]; + + [[unroll]] for (uint oc = 0; oc < O_COUNT; oc++) { + [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) { + fp32O[oc][os] = slm_pool_f32.slm_pool_pv[offsetSlmLoadPv + os * WARPSIZE + oc * MATP_REDUCE] * p.softMaxScale; + } + + [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) { + fp32O[oc][os] = fp32O[oc][os] + maskFp32[os]; + } + + float maxTemp = fp32Min; + [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) { + maxTemp = max(maxTemp, fp32O[oc][os]); + } + maxOut[oc] = subgroupMax(maxTemp); + [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) { + fp32O[oc][os] = exp(fp32O[oc][os] - maxOut[oc]); + } + } + + [[unroll]] for (uint oc = 0; oc < O_COUNT; oc++) { + if (localLinearId * O_COUNT + oc < GQA_RATIO) { + [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) { + matP_f16[offsetOut + oc * pDim + os * WARPSIZE] = float16_t(fp32O[oc][os]); + } + + if (lane == 0) { + out_max_f32[offsetMax + oc * maxDim] = maxOut[oc]; + } + } + } + + offsetOut = offsetOut + p.qHead * pDim; + offsetMax = offsetMax + p.qHead * maxDim; + barrier(); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp new file mode 100644 index 000000000..60a1c2ce7 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp @@ -0,0 +1,408 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_KHR_memory_scope_semantics : enable +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_ballot : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable +#extension GL_KHR_cooperative_matrix : enable +#extension GL_EXT_shared_memory_block : enable + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer P {f16vec4 pStateVec4[];}; +layout (binding = 1) readonly buffer V {float16_t vState[];}; +layout (binding = 1) readonly buffer V_VEC4 {f16vec4 vStateVec4[];}; +layout (binding = 2) buffer MAX_FP32 {float max_f32[];}; +layout (binding = 3) buffer SINK_FP32 {float sink_f32[];}; +layout (binding = 4) buffer OUT_FP32 {float out_f32[];}; +layout (binding = 4) buffer OUT_VEC4 {vec4 out_f32_vec4[];}; +layout (binding = 4) buffer OUT_F16 {float16_t out_f16[];}; + +layout (push_constant) uniform parameter +{ + uint kvSeqLen; + uint activationLength; + uint qHead; + uint kvHead; + uint qkRatio; + uint qkSubGroups; + uint flag; + uint kvStride1; + uint kvStride2; + uint batchStrideQ; + uint batchStrideK; + uint batchStrideV; + uint batchStrideM; + uint batchStrideO; + float softMaxScale; +} p; + +layout (constant_id = 0) const uint GROUPSIZE = 256; +layout (constant_id = 1) const uint GQA_RATIO = 8; +layout (constant_id = 2) const uint HEAD_DIM = 128; +layout (constant_id = 3) const uint N_TOKS_PER_GROUP = 1; +layout (constant_id = 4) const uint WARPSIZE = 16; +layout (constant_id = 5) const uint MATP_PER_LOOP = 64; +layout (constant_id = 6) const uint MATP_REDUCE = 32; +layout (constant_id = 7) const uint WARP_V_DIM = 16; + +#define TN WARPSIZE +#define TM 8 +#define TK 16 +#define MAT_O_N (WARP_V_DIM / TM) +#define MAT_P_M (GQA_RATIO * N_TOKS_PER_GROUP) +#define ALIGNED_P_M ((MAT_P_M + WARPSIZE - 1) / WARPSIZE) +#define V_HEAD_GROUPS (HEAD_DIM / WARP_V_DIM) + +#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE) +#define SPLIT_P_GROUPS (MATP_PER_LOOP / TK) + +#define SLM_POOL_SIZE_O (SUBGROUP_COUNT * ALIGNED_P_M * TN * MAT_O_N * TM) + +#define P_LOAD_PER_LOOP (GROUPSIZE * 4) +#define P_LOAD_LOOPS ((MAT_P_M * MATP_PER_LOOP + P_LOAD_PER_LOOP - 1) / P_LOAD_PER_LOOP) +#define SLM_POOL_SIZE_P (P_LOAD_LOOPS * P_LOAD_PER_LOOP) +#define SIZE_LOCAL_MAX (MAT_P_M * MATP_PER_LOOP / MATP_REDUCE) +#define MAX_LOAD_LOOPS ((SIZE_LOCAL_MAX + GROUPSIZE - 1) / GROUPSIZE) +#define SLM_POOL_SIZE_LOCAL_MAX (MAX_LOAD_LOOPS * GROUPSIZE) +#define MAX_REDUCE_COUNT ((MAT_P_M + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT) +#define GLOBAL_MAX_SIZE (MAX_REDUCE_COUNT * SUBGROUP_COUNT) + +#define SLM_POOL_SIZE_SOFTMAX_SUM (SUBGROUP_COUNT * P_LOAD_LOOPS) + +#define SLM_OFFSET_P (GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2) + +#define SLM_OFFSET_GLOBAL_MAX 0 +#define SLM_OFFSET_SOFTMAX_SUM (SLM_OFFSET_GLOBAL_MAX + GLOBAL_MAX_SIZE) +#define SLM_OFFSET_O (SLM_OFFSET_SOFTMAX_SUM + SLM_POOL_SIZE_SOFTMAX_SUM) +#define SLM_OFFSET_LOCAL_MAX (GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM) + +#define P_REDUCE_VEC4 (MATP_PER_LOOP / 4) +#define MAX_PER_LOOP (MATP_PER_LOOP / MATP_REDUCE) +#define SLM_MAX_STRIDE (MATP_REDUCE / 4) +#define SUB_GROUPS_PER_LINE (MATP_PER_LOOP / WARPSIZE / 4) + +shared slm_pool_block { + float slm_pool_o[GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_O]; +} slm_pool_f32; + +shared slm_pool_alias_block { + float16_t slm_pool_pv[GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2 + SLM_POOL_SIZE_P * 2]; +} slm_pool_alias_f16; + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint v = gl_WorkGroupID.y; + const uint d = gl_WorkGroupID.z; + const uint vWarpIdx = gl_WorkGroupID.x % V_HEAD_GROUPS; + const uint outTokIdx = gl_WorkGroupID.x / V_HEAD_GROUPS; + const uint localLinearId = gl_SubgroupID; + const uint wgLane = localLinearId * WARPSIZE + lane; + const uint splitIdx = localLinearId; + const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE; + const uint pDim = maxDim * MATP_REDUCE; + const uint kvDim = p.kvStride1; + const uint oDim = p.qHead * HEAD_DIM; + const uint offsetBaseP = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * pDim / 4; + const uint offsetBaseMax = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * maxDim; + const uint offsetBaseV = (d * p.batchStrideV + v * p.kvStride2 + vWarpIdx * WARP_V_DIM + splitIdx * TK * kvDim); + const uint offsetSlmP = (SLM_OFFSET_P + wgLane * 4); + const float fp32Min = uintBitsToFloat(0xFEFFFFFF); + const float fp32Max = uintBitsToFloat(0x7EFFFFFF); + uint offsetV = offsetBaseV; + + coopmat sums[ALIGNED_P_M][MAT_O_N]; + f16vec4 pStateTemp[P_LOAD_LOOPS]; + + float fp32CompensationP[P_LOAD_LOOPS]; + + uint loadRowBase[P_LOAD_LOOPS]; + uint loadColBase[P_LOAD_LOOPS]; + float fp32SoftMaxSum[P_LOAD_LOOPS]; + float fp32GlobalMaxP[P_LOAD_LOOPS]; + uint maxRowBase[MAX_LOAD_LOOPS]; + uint maxColBase[MAX_LOAD_LOOPS]; + uint outOffsets[ALIGNED_P_M]; + bool outputMask[ALIGNED_P_M]; + float fp32SinkCoeff[ALIGNED_P_M]; + + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + const uint flatOffset = pm * WARPSIZE + lane; + const uint inGroupTokIdx = flatOffset / GQA_RATIO; + const uint inGroupHeadIdx = flatOffset % GQA_RATIO; + outputMask[pm] = (N_TOKS_PER_GROUP * outTokIdx + inGroupTokIdx < p.activationLength) && (inGroupHeadIdx < GQA_RATIO) && (inGroupTokIdx < N_TOKS_PER_GROUP); + outOffsets[pm] = (inGroupTokIdx * oDim + inGroupHeadIdx * HEAD_DIM) / 4; + if ((0x1 & p.flag) != 0) { + fp32SinkCoeff[pm] = sink_f32[inGroupHeadIdx + v * GQA_RATIO]; + } + } + + [[unroll]] for (uint maxCount = 0; maxCount < MAX_REDUCE_COUNT; maxCount++) { + const uint flatIdx = maxCount * SUBGROUP_COUNT + localLinearId; + const uint rowIdx = flatIdx % GQA_RATIO; + const uint tokIdx = flatIdx / GQA_RATIO; + + if (tokIdx < N_TOKS_PER_GROUP) { + float fp32MaxReduce = fp32Min; + const uint maxOffset = offsetBaseMax + (tokIdx * p.qHead + rowIdx) * maxDim; + [[unroll]] for (uint maxReduce = 0; maxReduce < (maxDim + WARPSIZE - 1) / WARPSIZE; maxReduce++) { + if (maxReduce * WARPSIZE + lane < maxDim) { + fp32MaxReduce = max(fp32MaxReduce, max_f32[maxOffset + maxReduce * WARPSIZE + lane]); + } + } + fp32MaxReduce = subgroupMax(fp32MaxReduce); + if (lane == 0) { + slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32MaxReduce; + } + } else { + if (lane == 0) { + slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32Max; + } + } + } + + barrier(); + + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + const uint flatOffset = (pLoad * GROUPSIZE + wgLane) / P_REDUCE_VEC4; + const uint rowIdxFlat = flatOffset % GQA_RATIO; + const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1); + loadColBase[pLoad] = (pLoad * GROUPSIZE + wgLane) % P_REDUCE_VEC4; + loadRowBase[pLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat); + fp32SoftMaxSum[pLoad] = 0.0f; + fp32GlobalMaxP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + flatOffset]; + } + + [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) { + const uint flatOffset = (maxLoad * GROUPSIZE + wgLane) / MAX_PER_LOOP; + const uint rowIdxFlat = flatOffset % GQA_RATIO; + const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1); + maxColBase[maxLoad] = (maxLoad * GROUPSIZE + wgLane) % MAX_PER_LOOP; + maxRowBase[maxLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat); + } + + [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) { + const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad]; + slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset]; + maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE; + } + + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad]; + pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset]; + } + + [[unroll]] for (uint n = 0; n < ALIGNED_P_M; n++) { + [[unroll]] for (uint i = 0; i < MAT_O_N; i++) { + sums[n][i] = coopmat(0.0f); + } + } + + barrier(); + + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE; + if (loadColBase[pLoad] < pDim / 4) { + fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxOffset]; + float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w); + float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]); + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + pTemp[kk] = pTemp[kk] * compTemp; + fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk]; + slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]); + } + } else { + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f); + } + } + + loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4; + } + + const uint loopCount = (p.kvSeqLen + MATP_PER_LOOP - 1) / MATP_PER_LOOP; + + for (uint loop = 0; loop < loopCount; loop++) { + const uint slmPingPongLoad = (loop & 0x1); + const uint slmPingPongStore = ((loop + 1) & 0x1); + + if (loop + 1 < loopCount) { + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad]; + pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset]; + } + + [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) { + const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad]; + slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset]; + maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE; + } + } + + barrier(); + + { + const uint coopMatOffsetP = SLM_OFFSET_P + slmPingPongLoad * SLM_POOL_SIZE_P + splitIdx * TK; + coopmat matV[MAT_O_N]; + [[unroll]] for (uint cc = 0; cc < MAT_O_N; cc++) { + coopMatLoad( + matV[cc], + vState, + offsetV + TM * cc, + kvDim, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) { + coopmat matP; + coopMatLoad( + matP, + slm_pool_alias_f16.slm_pool_pv, + coopMatOffsetP + mo * TN * MATP_PER_LOOP, + MATP_PER_LOOP, + gl_CooperativeMatrixLayoutColumnMajor); + + [[unroll]] for (uint no = 0; no < MAT_O_N; no++) { + sums[mo][no] = coopMatMulAdd(matV[no], matP, sums[mo][no]); + } + } + } + + offsetV += MATP_PER_LOOP * kvDim; + if (loop * MATP_PER_LOOP + splitIdx * TK >= p.kvSeqLen) { + offsetV = 0; + } + if (loop + 1 < loopCount) { + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE; + if (loadColBase[pLoad] < pDim / 4) { + fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxOffset]; + float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w); + float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]); + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + pTemp[kk] = pTemp[kk] * compTemp; + fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk]; + slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]); + } + } else { + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f); + } + } + loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4; + } + } + } + + barrier(); + + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + fp32SoftMaxSum[pLoad] = subgroupAdd(fp32SoftMaxSum[pLoad]); + } + + [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) { + [[unroll]] for (uint no = 0; no < MAT_O_N; no++) { + coopMatStore( + sums[mo][no], + slm_pool_f32.slm_pool_o, + SLM_OFFSET_O + mo * TN * WARP_V_DIM + TM * no + localLinearId * ALIGNED_P_M * TN * WARP_V_DIM, + WARP_V_DIM, + gl_CooperativeMatrixLayoutColumnMajor); + } + } + + [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pLoad * SUBGROUP_COUNT + localLinearId] = fp32SoftMaxSum[pLoad]; + } + + barrier(); + + if (localLinearId == 1) { + const uint sumBase = SLM_OFFSET_SOFTMAX_SUM + lane * SUB_GROUPS_PER_LINE; + float sumTemp[ALIGNED_P_M][SUB_GROUPS_PER_LINE]; + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + [[unroll]] for (uint reduce = 0; reduce < SUB_GROUPS_PER_LINE; reduce++) { + sumTemp[pm][reduce] = slm_pool_f32.slm_pool_o[sumBase + pm * WARPSIZE * SUB_GROUPS_PER_LINE + reduce]; + } + } + + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + [[unroll]] for (uint reduce = 1; reduce < SUB_GROUPS_PER_LINE; reduce++) { + sumTemp[pm][0] = sumTemp[pm][0] + sumTemp[pm][reduce]; + } + } + + if ((0x1 & p.flag) != 0) { + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + float fp32GlobalMax = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + pm * WARPSIZE + lane]; + float sinkCompensation = fp32GlobalMax - fp32SinkCoeff[pm]; + sinkCompensation = exp(sinkCompensation); + float softmaxSumTemp = sumTemp[pm][0] * sinkCompensation; + sumTemp[pm][0] = sumTemp[pm][0] + 1.0f / sinkCompensation; + sumTemp[pm][0] = 1.0f / sumTemp[pm][0]; + sinkCompensation = sinkCompensation / (1.0f + softmaxSumTemp); + sumTemp[pm][0] = fp32GlobalMax < fp32SinkCoeff[pm] ? sinkCompensation : sumTemp[pm][0]; + slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = sumTemp[pm][0]; + } + } else { + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = 1.0f / sumTemp[pm][0]; + } + } + } + + [[unroll]] for (uint reduce = 2; reduce < SPLIT_P_GROUPS; reduce = reduce << 1 ) { + const uint stride = (reduce >> 1) * ALIGNED_P_M * TN * MAT_O_N * TM; + if ((localLinearId % reduce) == 0) { + const uint reduceBase = localLinearId * ALIGNED_P_M * TN * MAT_O_N * TM + SLM_OFFSET_O; + float sumTemp0[4]; + float sumTemp1[4]; + const uint reduceVec4Count = ALIGNED_P_M * TN * MAT_O_N * TM / 4 / WARPSIZE; + [[unroll]] for (uint totalLoads = 0; totalLoads < reduceVec4Count; totalLoads++) { + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + sumTemp0[kk] = slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk]; + sumTemp1[kk] = slm_pool_f32.slm_pool_o[reduceBase + stride + totalLoads * 4 * WARPSIZE + 4 * lane + kk]; + } + + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + sumTemp0[kk] = sumTemp0[kk] + sumTemp1[kk]; + } + + [[unroll]] for (uint kk = 0; kk < 4; kk++) { + slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk] = sumTemp0[kk]; + } + } + } + barrier(); + } + + if (localLinearId == 0) { + const uint slmBase0 = SLM_OFFSET_O + lane * WARP_V_DIM; + const uint slmBase1 = slmBase0 + SPLIT_P_GROUPS / 2 * ALIGNED_P_M * TN * MAT_O_N * TM; + + const uint offsetOutBase = (d * p.batchStrideO + vWarpIdx * WARP_V_DIM + v * GQA_RATIO * HEAD_DIM + outTokIdx * oDim * N_TOKS_PER_GROUP) / 4; + float fp32SoftMaxMul[ALIGNED_P_M]; + float fp32Output[ALIGNED_P_M][4]; + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + fp32SoftMaxMul[pm] = slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane]; + } + + [[unroll]] for (uint vg = 0; vg < WARP_V_DIM / 4; vg++) { + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + [[unroll]] for (uint vc = 0; vc < 4; vc++) { + fp32Output[pm][vc] = slm_pool_f32.slm_pool_o[slmBase0 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm]; + fp32Output[pm][vc] = fp32Output[pm][vc] + slm_pool_f32.slm_pool_o[slmBase1 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm]; + } + } + + [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) { + if (outputMask[pm] == true) { + out_f32_vec4[offsetOutBase + outOffsets[pm] + vg] = vec4(fp32Output[pm][0], fp32Output[pm][1], fp32Output[pm][2], fp32Output[pm][3]); + } + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 73cef00b0..5b2479da2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -922,6 +922,10 @@ void process_shaders() { string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {}); string_to_spv("fa_mask_opt", "flash_attn_mask_opt.comp", {}); + + string_to_spv("fa_decode_ph1", "flash_attn_decode_phase_1.comp", {}, true, true, false, false); + string_to_spv("fa_decode_ph2", "flash_attn_decode_phase_2.comp", {}, true, true, false, false); + string_to_spv("fa_sparse_compact", "flash_attn_sparse_compact.comp", {}); string_to_spv("fa_sparse_compact_subgroup", "flash_attn_sparse_compact.comp", {{"USE_SUBGROUPS", "1"}});