CANN: missing const qualifiers added for consistency (#30933)

* spec : add probabilistic sampling for DFlash2 drafts

* spec : apply review suggestions

---------

Co-authored-by: Pranesh Gonegandla <pgonegandla@nvidia.com>

common/speculative.cpp

@@ -19,6 +19,7 @@
 #include <fstream>
 #include <iomanip>
 #include <map>
+#include <random>
 #include <cinttypes>
 
 #define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -1010,6 +1011,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     bool    is_mrope       = false;
     int32_t selector_top_k = 0;
 
+    // DFlash2 probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
+    std::vector<std::mt19937> dflash2_rng;
+    std::vector<int64_t>      dflash2_seed;
+
     // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
     const bool is_dspark;
 
@@ -1060,6 +1065,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 
         selector_top_k = llama_model_dflash_selector_top_k(model_dft);
         is_dflash2     = selector_top_k > 0;
+        dflash2_rng.resize(n_seq);
+        dflash2_seed.assign(n_seq, -1);
         char lilicorr_meta[32] = {};
         lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
                 ? std::atoi(lilicorr_meta) : 0;
@@ -1158,6 +1165,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             return;
         }
 
+        // reseed per request rather than per round, or two identical requests differ
+        dflash2_seed[seq_id] = -1;
+
         const int32_t N = (int32_t) prompt.size();
         if (N <= 0) {
             return;
@@ -1338,6 +1348,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                 const float * lattice = llama_get_embeddings_nextn(ctx_dft);
                 GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
 
+                // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
+                const bool sample = is_dflash2 && params.probabilistic && dp.result_q;
+
+                // greedy keeps temperature 1, sampling uses the target's
+                const float temp = sample ? dp.temp : 1.0f;
+
+                if (sample && dflash2_seed[seq_id] != dp.seed) {
+                    dflash2_seed[seq_id] = dp.seed;
+                    // the default seed means random; mixed so it differs from the target's, as in spec_retune
+                    dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+                }
+
+                std::vector<float>   probs;
+                std::vector<int32_t> sorted_idx;
+
                 int32_t predecessor = 0;
                 for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) {
                     const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
@@ -1346,21 +1371,53 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 
                     predecessor = (int32_t) std::distance(scores,
                             std::max_element(scores, scores + top_k));
-                    if (params.p_min > 0.0f) {
-                        // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+                    if (params.p_min > 0.0f || sample) {
+                        // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
+                        probs.resize(top_k);
                         float sum = 0.0f;
                         for (int32_t k = 0; k < top_k; ++k) {
-                            sum += std::exp(scores[k] - scores[predecessor]);
+                            probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
+                            sum += probs[k];
                         }
                         if (1.0f / sum < params.p_min) {
                             break;
                         }
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            probs[k] /= sum;
+                        }
+                    }
+                    if (sample) {
+                        // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
+                        sorted_idx.resize(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            sorted_idx[k] = k;
+                        }
+                        std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
+
+                        // a round-off overshoot keeps the argmax, which has q > 0
+                        double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
+                        for (const int32_t k : sorted_idx) {
+                            u -= probs[k];
+                            if (u < 0.0) {
+                                predecessor = k;
+                                break;
+                            }
+                        }
+
+                        auto & q = dp.result_q->emplace_back();
+                        q.reserve(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            q.push_back({ (llama_token) row[k], scores[k], probs[k] });
+                        }
                     }
                     result.push_back((llama_token) row[predecessor]);
                 }
 
                 if (result.size() < (size_t) params.n_min) {
                     result.clear();
+                    if (sample) {
+                        dp.result_q->clear();
+                    }
                 }
                 continue;
             }

common/speculative.h

@@ -73,7 +73,7 @@ struct common_speculative_draft_params {
     // the generated draft from the last _draft() call
     llama_tokens * result;
 
-    // candidate distribution per drafted token; set it to make draft-simple and draft-mtp sample
+    // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash (DFlash2 only) sample
     std::vector<std::vector<llama_token_data>> * result_q = nullptr;
 
     // the target's temp and seed, read only when the drafter samples probabilistically

conversion/mistral3.py

@@ -39,6 +39,11 @@ class Mistral3Model(TextModel):
             self.gguf_writer.add_add_bos_token(False)
             self.gguf_writer.add_add_eos_token(False)
 
+        def get_vocab_base_pre(self, tokenizer) -> str:
+            if self.origin_hf_arch == "Mistral3Model":
+                return "nemotron-3.5-embed"
+            return super().get_vocab_base_pre(tokenizer)
+
         def set_gguf_parameters(self):
             super().set_gguf_parameters()
             self._try_set_pooling_type()

docs/backend/SYCL.md

@@ -783,7 +783,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL          | ON (mandatory)                        | Enable build with SYCL code path.           |
 | GGML_SYCL_TARGET   | INTEL *(default)*                     | Set the SYCL target device type.            |
 | GGML_SYCL_DEVICE_ARCH | Optional                           | Set the SYCL device architecture. Setting the device architecture can improve the performance. See the table [--offload-arch](https://github.com/intel/llvm/blob/sycl/sycl/doc/design/OffloadDesign.md#--offload-arch) for a list of valid architectures. |
-| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. (1.) |
+| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. [1.](#sycl-f16-footnote) |
 | GGML_SYCL_GRAPH    | ON *(default)* \|OFF *(Optional)*     | Enable build with [SYCL Graph extension](https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/experimental/sycl_ext_oneapi_graph.asciidoc). |
 | GGML_SYCL_DNN      | ON *(default)* \|OFF *(Optional)*     | Enable build with oneDNN.                   |
 | GGML_SYCL_HOST_MEM_FALLBACK | ON *(default)* \|OFF *(Optional)* | Allow host memory fallback when device memory is full during quantized weight reorder. Enables inference to continue at reduced speed (reading over PCIe) instead of failing. Requires Linux kernel 6.8+. |
@@ -791,7 +791,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | CMAKE_C_COMPILER   | `icx` *(Linux)*, `icx/cl` *(Windows)* | Set `icx` compiler for SYCL code path.      |
 | CMAKE_CXX_COMPILER | `icpx` *(Linux)*, `icx` *(Windows)*   | Set `icpx/icx` compiler for SYCL code path. |
 
-1. FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.
+1. <a id="sycl-f16-footnote"></a>FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.
 
 ### Runtime
 
@@ -802,7 +802,8 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_DEV2DEV_MEMCPY | 0 (default), 1, 2 | Choose the method of dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default), only support dGPUs.<br>* 1: L0 API -- Better performance, only support dGPUs, found to lead to abnormal crash in some case. <br>* 2: Host Forward -- Most stable method for all cases (including iGPU + dGPU*N), but with lower performance (-2% to -5%).<br>SYCL & L0 API are easy to be impacted by Intel GPU driver issue. When you meet the garbled output or crash issues in multiple GPUs case, try with this debug flag to work around or check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
-| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
+| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still in development, performance depends on model and device. [1.](#sycl-graph-footnote) |
+| GGML_SYCL_GRAPH_EVICTION_TIMEOUT | 10 (default) or positive integer | Seconds a cached SYCL graph can go unused before it is evicted. Checked on a sweep that runs every half of this timeout. |
 | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
 | GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
 | GGML_SYCL_UPLOAD_STAGING_SLOTS | 4 (default) or non-negative integer | Number of 8 MiB pinned host slots used to stage tensor uploads (model loading), so the host copy of one slot overlaps the transfer of the previous one. Set to 0 to use the old path: a malloc'd bounce buffer and a blocking copy per tensor. |
@@ -834,6 +835,17 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
 | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. |
 
+1. <a id="sycl-graph-footnote"></a>SYCL Graph support (`GGML_SYCL_ENABLE_GRAPH=1`) is limited to single-GPU workloads and may also depend on the model architecture and kernels in use.
+   Performance depends on how often a recorded graph can be reused, which will vary by workload.
+   For this reason, it is recommended to measure performance both with and without graphs enabled on your workloads.
+
+   To check whether graphs are supported for your workload, enable INFO logging (`-lv 4`) and look for `enabling SYCL graphs` or `disabling SYCL graphs`.
+   Support may differ between prefill and decode, so check the status for the phases that matter to you.
+   DEBUG logging (`-lv 5`) reports why graphs are disabled.
+
+   Then benchmark with your preferred tool, such as `llama-bench`.
+   For graph activity details (including reuse), set `GGML_SYCL_DEBUG=1` and look for `[SYCL-GRAPH]` messages.
+
 ## Compile-time Flags
 
 Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spot.

ggml/src/ggml-cann/ggml-cann.cpp

@@ -920,7 +920,7 @@ static void * ggml_backend_cann_buffer_get_base(ggml_backend_buffer_t buffer) {
  * @param dst Pointer to the destination buffer where transformed data will be
  * stored.
  */
-static void ggml_backend_cann_transform_q4_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q4_0(const ggml_tensor * tensor, const void * src, void * dst) {
     int64_t n_elems     = ggml_nelements(tensor);
     int64_t groups      = n_elems / QK4_0;
     size_t  quant_bytes = n_elems * sizeof(uint8_t) / 2;
@@ -1014,7 +1014,7 @@ static void ggml_backend_cann_transform_back_q4_0(const ggml_tensor * tensor, vo
  * @param dst Pointer to the destination buffer where transformed data will be
  * stored.
  */
-static void ggml_backend_cann_transform_q8_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q8_0(const ggml_tensor * tensor, const void * src, void * dst) {
     int64_t n_elems     = ggml_nelements(tensor);
     int64_t groups      = n_elems / QK8_0;
     size_t  quant_bytes = n_elems * sizeof(uint8_t);
@@ -1075,7 +1075,7 @@ static void ggml_backend_cann_transform_back_q8_0(const ggml_tensor * tensor, co
  * @param dst Pointer to the destination buffer where transformed data will be
  * stored.
  */
-static void ggml_backend_cann_transform(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform(const ggml_tensor * tensor, const void * src, void * dst) {
     switch (tensor->type) {
         case GGML_TYPE_Q4_0:
             ggml_backend_cann_transform_q4_0(tensor, src, dst);

ggml/src/ggml-hexagon/CMakeLists.txt

@@ -51,6 +51,13 @@ endif()
 set(TARGET_NAME ggml-hexagon)
 ggml_add_backend_library(${TARGET_NAME}
     ggml-hexagon.cpp
+    htp-common.h
+    htp-kparams.cpp
+    htp-kparams.h
+    htp-opnode.cpp
+    htp-opnode.h
+    htp-fusion.cpp
+    htp-fusion.h
     htp-repack.cpp
     htp-repack.h
     htp-drv.cpp

ggml/src/ggml-hexagon/htp-common.h

@@ -0,0 +1,106 @@
+#ifndef HTP_COMMON_H
+#define HTP_COMMON_H
+
+#include "htp-opnode.h"
+#include "htp-repack.h"
+
+struct htp_tensor_extra {
+    std::vector<uint8_t> shadow_buf;
+    size_t               shadow_size { 0 };
+    uint32_t             flags { 0 };
+};
+
+extern int    opt_arch;
+extern size_t opt_ndev;
+extern size_t opt_nhvx;
+extern int    opt_nhmx;
+extern size_t opt_vmem;
+extern int    opt_etm;
+extern int    opt_verbose;
+extern int    opt_profile;
+extern bool   opt_hostbuf;
+extern bool   opt_dma64;
+extern int    opt_power_mode;
+
+extern size_t opt_mbuf_dyn;
+extern size_t opt_mbuf_static;
+extern size_t opt_mbuf_total;
+
+extern int    opt_mm_select;
+extern int    opt_fa_select;
+extern int    opt_fa_head_split;
+extern int    opt_gdn_select;
+extern int    opt_ar_select;
+extern int    opt_ar_scatter;
+
+extern int    opt_opbatch;
+extern int    opt_opqueue;
+extern int    opt_optrace;
+extern int    opt_oppoll;
+extern int    opt_opfusion;
+extern int    opt_rt_repack;
+
+#ifndef HEX_VERBOSE
+#define HEX_VERBOSE(...) \
+    if (opt_verbose) GGML_LOG_DEBUG(__VA_ARGS__)
+#endif
+
+static inline bool htp_tensor_is_fuseable(const struct ggml_tensor * t) {
+    if (!t || !t->extra) return false;
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    return (extra->flags & HTP_TENSOR_FUSEABLE) != 0;
+}
+
+static inline bool htp_can_rt_repack(const struct ggml_tensor * t) {
+    if (opt_rt_repack == 0 || !htp_is_tiled_weight_tensor(t)) {
+        return false;
+    }
+    if (t->buffer && ggml_backend_buft_is_host(t->buffer->buft)) {
+        return false;
+    }
+    return opt_rt_repack > 1 || ggml_is_quantized(t->type);
+}
+
+static inline bool htp_tensor_is_tiled(const struct ggml_tensor * t) {
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+        return true;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_needs_repack(const struct ggml_tensor * t) {
+    const auto * extra = (const struct htp_tensor_extra *) t->extra;
+    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+        return false;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_tensors_overlap(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    const uintptr_t a0 = (uintptr_t) a->data;
+    const uintptr_t b0 = (uintptr_t) b->data;
+    const uintptr_t a1 = a0 + ggml_nbytes(a);
+    const uintptr_t b1 = b0 + ggml_nbytes(b);
+
+    return a0 < b1 && b0 < a1;
+}
+
+static inline bool htp_can_row_partition(const struct ggml_tensor * t) {
+    if (t->ne[1] > 1 && (t->nb[1] & 127) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & 127) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & 127) != 0) return false;
+    return true;
+}
+
+struct ggml_hexagon_session;
+
+const char * htp_session_name(const struct ggml_hexagon_session * sess);
+size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess);
+
+htp_op_code  htp_op_remap(const struct ggml_tensor * t);
+
+#endif // HTP_COMMON_H

ggml/src/ggml-hexagon/htp-fusion.cpp

@@ -0,0 +1,650 @@
+#include "htp-fusion.h"
+#include "htp-kparams.h"
+#include "ggml-impl.h"
+
+#include <algorithm>
+#include <cstring>
+
+static inline bool htp_same_tensor_data(const ggml_tensor * a, const ggml_tensor * b) {
+    return a && b && (a == b || a->data == b->data);
+}
+
+static inline bool htp_same_shape_and_contiguity(const ggml_tensor * a, const ggml_tensor * b) {
+    return ggml_are_same_shape(a, b) && (ggml_is_contiguous(a) == ggml_is_contiguous(b));
+}
+
+static inline bool htp_inplace_overlap_unsafe(const ggml_tensor * dst, const ggml_tensor * in) {
+    return htp_tensors_overlap(dst, in) && (dst->data != in->data || !ggml_are_same_stride(dst, in));
+}
+
+static inline bool htp_inputs_overlap_unsafe(const htp_opnode & prev, const ggml_tensor * dst, const ggml_tensor * extra = nullptr) {
+    if (extra && htp_inplace_overlap_unsafe(dst, extra)) return true;
+    for (const auto * in : prev.inputs) {
+        if (htp_inplace_overlap_unsafe(dst, in)) return true;
+    }
+    return false;
+}
+
+static inline bool htp_match_consumer(
+    const htp_opnode & prev,
+    const htp_opnode & cur,
+    bool commutative,
+    const ggml_tensor ** other_src = nullptr
+) {
+    if (prev.outputs.size() != 1) {
+        return false;
+    }
+    const ggml_tensor * prev_out = prev.dst();
+    if (!htp_tensor_is_fuseable(prev_out)) {
+        return false;
+    }
+    if (htp_same_tensor_data(cur.src0(), prev_out)) {
+        if (other_src) *other_src = cur.src1();
+        return true;
+    }
+    if (commutative && htp_same_tensor_data(cur.src1(), prev_out)) {
+        if (other_src) *other_src = cur.src0();
+        return true;
+    }
+    return false;
+}
+
+static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_tensor_is_tiled(src0) && (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) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_tensor_is_tiled(src0) &&
+           (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK);
+}
+
+// --- Fusion rules ---
+
+static bool fuse_allreduce_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (opt_ar_select != 2 || prev.opcode != HTP_OP_ALLREDUCE) return false;
+
+    const auto * ar_kparams = prev.kparams<htp_allreduce_kernel_params>();
+    const uint32_t rank    = (uint32_t) ar_kparams->rank;
+    const uint32_t n_ranks = (uint32_t) ar_kparams->n_ranks;
+
+    const ggml_tensor * ar_local = prev.inputs[rank];
+    if (!htp_tensor_is_fuseable(ar_local)) return false;
+
+    const ggml_tensor * add_dst    = cur.dst();
+    const ggml_tensor * res_tensor = nullptr;
+    if (htp_same_tensor_data(cur.src0(), ar_local)) {
+        res_tensor = cur.src1();
+    } else if (htp_same_tensor_data(cur.src1(), ar_local)) {
+        res_tensor = cur.src0();
+    } else {
+        return false;
+    }
+
+    if (ar_local->type != res_tensor->type) return false;
+
+    const bool is_same_shape = ggml_are_same_shape(ar_local, res_tensor);
+    const bool is_row_bcast  = !is_same_shape && (ar_local->ne[0] == res_tensor->ne[0] &&
+                                                  res_tensor->ne[1] == 1 && res_tensor->ne[2] == 1 && res_tensor->ne[3] == 1);
+    if (!is_same_shape && !is_row_bcast) return false;
+
+    if (is_same_shape && !ggml_are_same_stride(ar_local, res_tensor)) {
+        return false;
+    }
+    if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(add_dst)) return false;
+
+    const bool is_shard_ok = n_ranks <= HTP_OP_MAX_OUTPUTS && add_dst->data == ar_local->data;
+
+    struct htp_allreduce_kernel_params kparams;
+    if (!htp_kparams_solve_allreduce(sess, add_dst, rank, n_ranks, true, is_row_bcast, is_shard_ok, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", htp_session_name(sess));
+        return false;
+    }
+
+    const bool scatter_ok = kparams.n_dsts > 1;
+    kparams.mode = scatter_ok ? HTP_ALLREDUCE_SHARDED_FANOUT : HTP_ALLREDUCE_FULL;
+
+    for (uint32_t r = 0; r < n_ranks; r++) {
+        const ggml_tensor * ar_src = prev.inputs[r];
+        if (!htp_tensors_overlap(add_dst, ar_src)) continue;
+        if (scatter_ok && r == rank) continue;
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: dst overlaps allreduce src %u\n", htp_session_name(sess), r);
+        return false;
+    }
+
+    prev.opcode = HTP_OP_ALLREDUCE_ADD;
+    prev.inputs.push_back(res_tensor);
+    if (kparams.mode != HTP_ALLREDUCE_SHARDED_FANOUT) {
+        prev.outputs = { add_dst };
+    }
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_mul(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_RMS_NORM) return false;
+
+    const ggml_tensor * weight = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &weight)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || src0->ne[0] != weight->ne[0] || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, weight)) return false;
+
+    const bool is_row_bcast  = (weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1);
+    const bool is_same_shape = ggml_are_same_shape(src0, weight);
+    if (!is_row_bcast && !is_same_shape) return false;
+
+    struct htp_unary_kernel_params kparams;
+    if (!htp_kparams_solve_rms_norm_mul(sess, src0, weight, dst, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip RMS_NORM_MUL fusion: VTCM needed (%d) > budget (%zu)\n",
+                    htp_session_name(sess), kparams.vtcm_size, htp_session_vtcm_size(sess));
+        return false;
+    }
+
+    prev.opcode  = HTP_OP_RMS_NORM_MUL;
+    prev.inputs  = { src0, weight };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_scale(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_RMS_NORM && prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst)) return false;
+
+    float scale = 0.0f;
+    float bias  = 0.0f;
+    memcpy(&scale, (const float *) cur.params + 0, sizeof(float));
+    memcpy(&bias,  (const float *) cur.params + 1, sizeof(float));
+    if (bias != 0.0f || scale == 0.0f) return false;
+
+    float prev_scale = 0.0f;
+    memcpy(&prev_scale, &prev.params[1], sizeof(float));
+    if (prev_scale != 0.0f) {
+        scale *= prev_scale;
+    }
+    memcpy(&prev.params[1], &scale, sizeof(float));
+
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &src2)) return false;
+    if (src2->type != GGML_TYPE_F32) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(prev.dst(), dst)) return false;
+
+    const auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    struct htp_mm_kernel_params kparams;
+    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams) ||
+        orig_kparams->collapse != kparams.collapse) {
+        return false;
+    }
+
+    const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
+    if (kparams.n_hmx == 0 && src1_nrows != 1) return false;
+
+    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",
+                        htp_session_name(sess), 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;
+        }
+    }
+
+    prev.opcode  = HTP_OP_MUL_MAT_ADD;
+    prev.inputs  = { src0, src1, src2 };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT && prev.opcode != HTP_OP_MUL_MAT_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0   = prev.inputs[0];
+    const ggml_tensor * x    = prev.inputs[curr_n];
+    const ggml_tensor * w_in = cur.src0();
+    const ggml_tensor * x_in = cur.src1();
+    const ggml_tensor * d_in = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
+        w0->ne[2] != 1 || w0->ne[3] != 1 || w_in->ne[2] != 1 || w_in->ne[3] != 1 ||
+        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    const bool collapse = orig_kparams->collapse && htp_matmul_can_collapse(w_in, x, d_in);
+    if ((x->ne[2] > 1 || x->ne[3] > 1) && !collapse) return false;
+
+    if (prev.opcode == HTP_OP_MUL_MAT_NX) {
+        orig_kparams->n_weights = curr_n + 1;
+    } else {
+        struct htp_mm_kernel_params kparams;
+        if (collapse) {
+            const ggml_tensor x_collapsed = htp_tensor_collapse_rows(x);
+            if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
+            kparams.collapse = 1;
+        } else if (!htp_kparams_solve_mmnx(sess, w0, x, curr_n + 1, &kparams)) {
+            return false;
+        }
+        if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) return false;
+
+        prev.set_kparams(kparams);
+        prev.opcode = HTP_OP_MUL_MAT_NX;
+    }
+
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_id_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT_ID && prev.opcode != HTP_OP_MUL_MAT_ID_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 2 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0     = prev.inputs[0];
+    const ggml_tensor * x      = prev.inputs[curr_n];
+    const ggml_tensor * ids    = prev.inputs[curr_n + 1];
+    const ggml_tensor * w_in   = cur.src0();
+    const ggml_tensor * x_in   = cur.src1();
+    const ggml_tensor * ids_in = cur.node->src[2];
+    const ggml_tensor * d_in   = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    if (x_in != x || ids_in != ids || w_in->type != w0->type ||
+        w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2] ||
+        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    struct htp_mm_kernel_params kparams;
+    if (d_in->nb[1] == prev.node->nb[1]) {
+        kparams = *orig_kparams;
+        kparams.n_weights = curr_n + 1;
+    } else if (!htp_kparams_solve_mmidnx(sess, w0, x, d_in, curr_n + 1, &kparams)) {
+        return false;
+    }
+    if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) return false;
+
+    prev.opcode = HTP_OP_MUL_MAT_ID_NX;
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_gdn_cpy(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_GATED_DELTA_NET || prev.outputs.size() != 1) return false;
+
+    const ggml_tensor * gdn_out  = prev.dst();
+    const ggml_tensor * cpy_node = cur.node;
+    const ggml_tensor * cpy_src  = cur.src0();
+    const ggml_tensor * cpy_dst  = cur.dst();
+
+    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 = prev.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(prev.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<int64_t>(n_tokens, K);
+
+    if (cpy_src->op != GGML_OP_VIEW || !htp_same_tensor_data(cpy_src->view_src, gdn_out) ||
+        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->ne[2] != n_written || !ggml_is_contiguous(cpy_dst)) {
+        return false;
+    }
+
+    prev.outputs.push_back(cpy_dst);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_ssm_conv_silu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_SSM_CONV) return false;
+
+    auto * kparams = prev.kparams<htp_ssm_conv_kernel_params>();
+    if (kparams->silu) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * conv_dst = prev.dst();
+    const ggml_tensor * dst      = cur.dst();
+    if (!ggml_are_same_layout(conv_dst, dst)) return false;
+    if (htp_tensors_overlap(dst, prev.src0()) || htp_tensors_overlap(dst, prev.src1())) return false;
+
+    kparams->silu   = 1;
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_repeat_get_rows(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_REPEAT) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * repeat_out  = prev.dst();
+    const ggml_tensor * repeat_src0 = prev.src0();
+    const ggml_tensor * gr_src1     = cur.src1();
+    const ggml_tensor * gr_dst      = cur.dst();
+
+    if (repeat_src0->ne[0] != repeat_out->ne[0] || repeat_src0->ne[1] != repeat_out->ne[1]) return false;
+    if (repeat_src0->type != repeat_out->type) return false;
+
+    struct htp_get_rows_kernel_params kparams;
+    if (!htp_kparams_solve_repeat_get_rows(sess, repeat_src0, gr_src1, gr_dst, &kparams) || kparams.vtcm_size == 0) return false;
+
+    prev.opcode  = HTP_OP_GET_ROWS;
+    prev.inputs  = { repeat_src0, gr_src1 };
+    prev.outputs = { gr_dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_unary_mul_glu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    htp_op_code glu_op = HTP_OP_INVALID;
+    switch (prev.opcode) {
+        case HTP_OP_UNARY_SILU:       glu_op = HTP_OP_GLU_SWIGLU;      break;
+        case HTP_OP_UNARY_GELU:       glu_op = HTP_OP_GLU_GEGLU;       break;
+        case HTP_OP_UNARY_GELU_QUICK: glu_op = HTP_OP_GLU_GEGLU_QUICK; break;
+        case HTP_OP_UNARY_GELU_ERF:   glu_op = HTP_OP_GLU_GEGLU_ERF;   break;
+        case HTP_OP_UNARY_RELU:       glu_op = HTP_OP_GLU_REGLU;       break;
+        case HTP_OP_UNARY_SIGMOID:    glu_op = HTP_OP_GLU_SIGMOID;     break;
+        default: return false;
+    }
+
+    const ggml_tensor * up = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &up)) return false;
+
+    const ggml_tensor * gate = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (htp_inputs_overlap_unsafe(prev, dst, up)) return false;
+
+    struct htp_glu_kernel_params kparams;
+    if (!htp_kparams_solve_unary_mul_glu(sess, glu_op, gate, up, dst, &kparams)) return false;
+
+    prev.opcode  = glu_op;
+    prev.inputs  = { gate, up };
+    prev.outputs = { dst };
+    prev.fused.push_back(cur.node);
+    memset(prev.params, 0, sizeof(prev.params));
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mul_scalar(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL || prev.inputs.size() != 2) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    if (k0->kernel_type != HTP_BINARY_KERNEL_SCALAR || k1->kernel_type != HTP_BINARY_KERNEL_SCALAR) return false;
+    if (k0->src1_strided || k1->src1_strided) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false, &src2)) return false;
+
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src1, src2) || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, src2)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    kparams.src2_size       = kparams.src1_size;
+    kparams.rows_per_buffer = 0;
+    if (kparams.src1_size > 65536 || kparams.src1_size + kparams.src2_size > vtcm_size) return false;
+
+    struct htp_binary_vtcm_layout L;
+    htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    if (L.rows_per_buffer == 0 || L.total_bytes > vtcm_size) return false;
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = L.total_bytes;
+
+    prev.inputs.push_back(src2);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_binary_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != cur.opcode) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.inputs.size();
+    if (curr_n < 2 || curr_n >= HTP_OP_MAX_INPUTS) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    const bool k0_ok = (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k0->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k0->is_scalar);
+    const bool k1_ok = (k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k1->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k1->is_scalar);
+    if (!k0_ok || !k1_ok) return false;
+
+    const bool is_comm = (cur.opcode == HTP_OP_ADD || cur.opcode == HTP_OP_MUL);
+    const ggml_tensor * new_src = nullptr;
+    if (!htp_match_consumer(prev, cur, is_comm, &new_src)) return false;
+
+    const ggml_tensor * src0 = prev.inputs[0];
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src0, prev.inputs[1]) ||
+        !ggml_are_same_shape(src0, new_src) ||
+        !htp_same_shape_and_contiguity(prev.dst(), dst)) {
+        return false;
+    }
+    if (htp_inputs_overlap_unsafe(prev, dst, new_src)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    struct htp_binary_vtcm_layout L;
+    memset(&L, 0, sizeof(L));
+
+    const uint32_t n_srcs = curr_n + 1;
+    if (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE || k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) {
+        if (k0->kernel_type != HTP_BINARY_KERNEL_SAME_SHAPE) {
+            kparams = *k1;
+        }
+        kparams.n_srcs          = (uint8_t) n_srcs;
+        kparams.rows_per_buffer = 0;
+
+        auto check_split = [&](const ggml_tensor * t) {
+            if (src0->ne[2] > 1 && t->nb[2] != t->ne[1] * t->nb[1]) kparams.split_at_ne01 = 1;
+            if (src0->ne[3] > 1 && t->nb[3] != t->ne[2] * t->nb[2]) kparams.split_at_ne02 = 1;
+        };
+        for (const auto * in : prev.inputs) check_split(in);
+        check_split(new_src);
+        check_split(dst);
+
+        htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    }
+
+    if (L.rows_per_buffer == 0 || L.total_bytes == 0 || L.total_bytes > vtcm_size) {
+        bool all_contig = ggml_is_contiguous(new_src) && ggml_is_contiguous(dst);
+        for (const auto * in : prev.inputs) {
+            all_contig = all_contig && ggml_is_contiguous(in);
+        }
+        if (!all_contig) return false;
+
+        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+        const uint32_t elem_size   = (uint32_t) ggml_type_size(src0->type);
+        const uint32_t max_threads = htp_session_n_threads(sess);
+        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_threads, n_srcs, false, vtcm_size)) {
+            return false;
+        }
+    }
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = (uint32_t) L.total_bytes;
+
+    prev.inputs.push_back(new_src);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a) {
+    if (!a || !a->src[0] || !a->src[1]) return false;
+    switch (a->op) {
+        case GGML_OP_MUL_MAT:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX)) return false;
+            if (htp_is_hadamard(a)) return false;
+            if (a->src[0]->ne[2] != 1 || a->src[0]->ne[3] != 1) return false;
+            return ggml_is_quantized(a->src[0]->type);
+        case GGML_OP_MUL_MAT_ID:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX)) return false;
+            return a->src[2] != nullptr && ggml_is_quantized(a->src[0]->type);
+        default:
+            return false;
+    }
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    if (!htp_fusion_can_stack(b) || a->op != b->op || a->src[1] != b->src[1]) {
+        return false;
+    }
+    if (a->op == GGML_OP_MUL_MAT_ID && a->src[2] != b->src[2]) {
+        return false;
+    }
+    const ggml_tensor * w0 = a->src[0];
+    const ggml_tensor * w1 = b->src[0];
+    return w0->type == w1->type && w0->ne[0] == w1->ne[0] && w0->ne[2] == w1->ne[2];
+}
+
+int htp_fusion_max(const struct ggml_cgraph * graph, int idx) {
+    const int n = graph->n_nodes;
+    const ggml_op op0 = graph->nodes[idx]->op;
+    if (op0 != GGML_OP_ADD && op0 != GGML_OP_NORM && op0 != GGML_OP_RMS_NORM) {
+        return 1;
+    }
+
+    enum ggml_op ops[HTP_OP_MAX_FUSED];
+    ops[0] = op0;
+
+    int f = idx + 1;
+    while (f < n && f < idx + HTP_OP_MAX_FUSED) {
+        const ggml_op op = graph->nodes[f]->op;
+        if (op != GGML_OP_ADD && op != GGML_OP_MUL && op != GGML_OP_SCALE &&
+            op != GGML_OP_NORM && op != GGML_OP_RMS_NORM) {
+            break;
+        }
+        ops[f - idx] = op;
+        f++;
+    }
+
+    f -= idx;
+    for (; f > 1; f--) {
+        if (ggml_can_fuse(graph, idx, ops, f)) {
+            break;
+        }
+    }
+    return f;
+}
+
+void htp_fusion_tag_graph(struct ggml_cgraph * graph) {
+    for (int i = 0; i < graph->n_nodes; i++) {
+        auto * extra = (htp_tensor_extra *) graph->nodes[i]->extra;
+        if (!extra) continue;
+
+        extra->flags &= ~HTP_TENSOR_FUSEABLE;
+
+        const ggml_op op = graph->nodes[i]->op;
+        if ((op == GGML_OP_RMS_NORM || op == GGML_OP_SCALE || op == GGML_OP_REPEAT ||
+             op == GGML_OP_UNARY    || op == GGML_OP_ADD   || op == GGML_OP_SUB    ||
+             op == GGML_OP_MUL      || op == GGML_OP_DIV   || op == GGML_OP_MUL_MAT_ID ||
+             op == GGML_OP_SSM_CONV ||
+             (op == GGML_OP_MUL_MAT && !htp_is_hadamard(graph->nodes[i]))) &&
+            ggml_node_has_n_uses(graph, i, 1)) {
+            extra->flags |= HTP_TENSOR_FUSEABLE;
+        }
+    }
+}
+
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur) {
+    if (!opt_opfusion) return false;
+
+    switch (cur.opcode) {
+        case HTP_OP_ADD:
+            return (htp_fusion_enabled(HTP_FUSE_ALLREDUCE_ADD) && fuse_allreduce_add(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MAT_ADD)   && fuse_mul_mat_add(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)     && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_MUL:
+            return (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL)   && fuse_rms_norm_mul(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_UNARY_MUL_GLU)  && fuse_unary_mul_glu(sess, prev, cur))  ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MUL_SCALAR) && fuse_mul_mul_scalar(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)      && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_SUB:
+        case HTP_OP_DIV:
+            return htp_fusion_enabled(HTP_FUSE_BINARY_NX) && fuse_binary_nx(sess, prev, cur);
+        case HTP_OP_SCALE:
+            return htp_fusion_enabled(HTP_FUSE_RMS_NORM_SCALE) && fuse_rms_norm_scale(sess, prev, cur);
+        case HTP_OP_UNARY_SILU:
+            return htp_fusion_enabled(HTP_FUSE_SSM_CONV_SILU) && fuse_ssm_conv_silu(sess, prev, cur);
+        case HTP_OP_GET_ROWS:
+            return htp_fusion_enabled(HTP_FUSE_REPEAT_GET_ROWS) && fuse_repeat_get_rows(sess, prev, cur);
+        case HTP_OP_CPY:
+            return htp_fusion_enabled(HTP_FUSE_GDN_CPY) && fuse_gdn_cpy(sess, prev, cur);
+        case HTP_OP_MUL_MAT:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX) && fuse_mul_mat_nx(sess, prev, cur);
+        case HTP_OP_MUL_MAT_ID:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX) && fuse_mul_mat_id_nx(sess, prev, cur);
+        default:
+            return false;
+    }
+}

ggml/src/ggml-hexagon/htp-fusion.h

@@ -0,0 +1,33 @@
+#ifndef HTP_FUSION_H
+#define HTP_FUSION_H
+
+#include "htp-common.h"
+
+enum htp_fusion_flags {
+    HTP_FUSE_ALLREDUCE_ADD   = (1 << 1),  // 2
+    HTP_FUSE_RMS_NORM_MUL    = (1 << 2),  // 4
+    HTP_FUSE_MUL_MAT_ADD     = (1 << 3),  // 8
+    HTP_FUSE_MUL_MAT_NX      = (1 << 4),  // 16
+    HTP_FUSE_MUL_MAT_ID_NX   = (1 << 5),  // 32
+    HTP_FUSE_GDN_CPY         = (1 << 6),  // 64
+    HTP_FUSE_REPEAT_GET_ROWS = (1 << 7),  // 128
+    HTP_FUSE_UNARY_MUL_GLU   = (1 << 8),  // 256
+    HTP_FUSE_MUL_MUL_SCALAR  = (1 << 9),  // 512
+    HTP_FUSE_BINARY_NX       = (1 << 10), // 1024
+    HTP_FUSE_RMS_NORM_SCALE  = (1 << 11), // 2048
+    HTP_FUSE_SSM_CONV_SILU   = (1 << 12), // 4096
+};
+
+static inline bool htp_fusion_enabled(int flag) {
+    if (opt_opfusion <= 0) return false;
+    if (opt_opfusion == 1) return true; // 1 enables all
+    return (opt_opfusion & flag) != 0;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a);
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b);
+int  htp_fusion_max(const struct ggml_cgraph * graph, int idx);
+void htp_fusion_tag_graph(struct ggml_cgraph * graph);
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur);
+
+#endif // HTP_FUSION_H

ggml/src/ggml-hexagon/htp-kparams.h

@@ -0,0 +1,123 @@
+#ifndef HTP_KPARAMS_H
+#define HTP_KPARAMS_H
+
+#include "htp-common.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/concat-ops.h"
+#include "htp/cpy-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-ops.h"
+#include "htp/fill-ops.h"
+#include "htp/flash-attn-ops.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/get-rows-ops.h"
+#include "htp/glu-ops.h"
+#include "htp/im2col-ops.h"
+#include "htp/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/pool-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/rope-ops.h"
+#include "htp/set-rows-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <type_traits>
+#include <utility>
+
+bool        htp_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
+bool        htp_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
+ggml_tensor htp_tensor_collapse_rows(const struct ggml_tensor * t);
+
+template <typename... Args>
+auto htp_kparams_type(bool (*)(Args...)) -> std::remove_pointer_t<std::remove_reference_t<decltype((std::declval<Args>(), ...))>>;
+
+template <auto SolveFn, typename... Args>
+static inline bool htp_kparams_check(Args &&... args) {
+    decltype(htp_kparams_type(SolveFn)) kparams;
+    return SolveFn(std::forward<Args>(args)..., &kparams);
+}
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    void * kparams = nullptr
+);
+
+static inline bool htp_kparams_check(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op
+) {
+    return htp_kparams_solve(sess, op, nullptr);
+}
+
+bool htp_kparams_solve_allreduce(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    uint32_t rank,
+    uint32_t n_ranks,
+    bool has_add,
+    bool is_row_bcast,
+    bool is_shard_ok,
+    struct htp_allreduce_kernel_params * kparams
+);
+
+bool htp_kparams_solve_rms_norm_mul(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_unary_kernel_params * kparams
+);
+
+bool htp_kparams_solve_matmul_add(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * src2,
+    const struct ggml_tensor * dst,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmidnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    const struct ggml_tensor * dst,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_repeat_get_rows(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_get_rows_kernel_params * kparams
+);
+
+bool htp_kparams_solve_unary_mul_glu(
+    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_glu_kernel_params * kparams
+);
+
+#endif // HTP_KPARAMS_H

ggml/src/ggml-hexagon/htp-opnode.cpp

@@ -0,0 +1,260 @@
+#include "htp-opnode.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-ops.h"
+#include "htp/fill-ops.h"
+#include "htp/flash-attn-ops.h"
+#include "htp/gated-delta-net-ops.h"
+#include "htp/get-rows-ops.h"
+#include "htp/glu-ops.h"
+#include "htp/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <stdio.h>
+
+std::string htp_opnode::name() const {
+    switch (opcode) {
+        case HTP_OP_MDEV_GROUP:    return "MDEV_GROUP";
+        case HTP_OP_CPY_FENCE:     return "CPY+FENCE";
+        case HTP_OP_FENCE:         return (node && node->op_params[1] == 0) ? "FENCE_WAIT" : "FENCE_SIGNAL";
+        case HTP_OP_ALLREDUCE:     return "ALLREDUCE";
+        case HTP_OP_ALLREDUCE_ADD: return "ALLREDUCE+ADD";
+        case HTP_OP_HADAMARD:      return "MUL_MAT_HADAMARD";
+        case HTP_OP_MUL_MAT_NX:    return "MUL_MAT_NX";
+        case HTP_OP_MUL_MAT_ID_NX: return "MUL_MAT_ID_NX";
+        case HTP_OP_ADD:
+        case HTP_OP_SUB:
+        case HTP_OP_MUL:
+        case HTP_OP_DIV:
+            if (!fused.empty()) {
+                const auto * kp = kparams<htp_binary_kernel_params>();
+                if (kp->kernel_type != HTP_BINARY_KERNEL_SCALAR) {
+                    switch (opcode) {
+                        case HTP_OP_ADD: return "ADD_NX";
+                        case HTP_OP_SUB: return "SUB_NX";
+                        case HTP_OP_MUL: return "MUL_NX";
+                        case HTP_OP_DIV: return "DIV_NX";
+                        default: break;
+                    }
+                }
+            }
+            break;
+        default:
+            break;
+    }
+
+    if (!node) {
+        return "";
+    }
+    if (htp_is_hadamard(node)) {
+        return "MUL_MAT_HADAMARD";
+    }
+
+    std::string res = ggml_op_desc(node);
+    for (const auto * t : fused) {
+        res += "+";
+        res += ggml_op_desc(t);
+    }
+    return res;
+}
+
+static const char * format_tensor_name(const struct ggml_tensor * t) {
+    return t ? t->name : "NONE";
+}
+
+static std::string format_tensor_dims(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%d:%d", (int) t->ne[0], (int) t->ne[1]);
+    } else {
+        snprintf(buf, sizeof(buf), "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
+    }
+    return buf;
+}
+
+static const char * format_tensor_type(const struct ggml_tensor * t) {
+    return t ? ggml_type_name(t->type) : "NONE";
+}
+
+static std::string format_tensor_strides(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    const char * c = ggml_is_contiguous(t) ? "" : "!";
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
+    } else {
+        snprintf(buf, sizeof(buf), "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
+    }
+    return buf;
+}
+
+static const char * format_tensor_buff(const struct ggml_tensor * t) {
+    if (t && t->buffer) {
+        return ggml_backend_buffer_name(t->buffer);
+    }
+    return "NONE";
+}
+
+template <typename Fn>
+static std::string format_op_tensors(const htp_opnode & node, Fn fn) {
+    std::string res;
+    const auto & inputs  = node.get_inputs();
+    const auto & outputs = node.get_outputs();
+
+    if (!inputs.empty()) {
+        for (size_t i = 0; i < inputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(inputs[i]);
+        }
+        res += " -> ";
+    }
+
+    if (outputs.empty()) {
+        res += fn(node.dst());
+    } else {
+        for (size_t i = 0; i < outputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(outputs[i]);
+        }
+    }
+    return res;
+}
+
+static std::string format_kernel_params(const htp_opnode & node) {
+    char str[128];
+    const size_t max_size = sizeof(str);
+
+    if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
+        node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
+        node.opcode == HTP_OP_MUL_MAT_ADD) {
+        const auto * kparams = node.kparams<htp_mm_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
+            path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
+                   type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
+            path = "hvx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
+            path = "hvx-tiled";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+        const auto * kparams = node.kparams<htp_fa_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_FA_KERNEL_HMX) {
+            path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
+        } else if (type == HTP_FA_KERNEL_HVX) {
+            path = "hvx";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (htp_op_is_unary(node.opcode)) {
+        const auto * kparams = node.kparams<htp_unary_kernel_params>();
+        snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
+        snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
+    } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
+        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 = node.kparams<htp_ssm_conv_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 = node.kparams<htp_softmax_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_type, (int) kparams->n_threads, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
+        const auto * kparams = node.kparams<htp_gdn_kernel_params>();
+        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);
+    } 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 = node.kparams<htp_binary_kernel_params>();
+        snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+        const auto * kparams = node.kparams<htp_sort_kernel_params>();
+        snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+                 node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
+                 (int) kparams->n_threads, (int) kparams->n_chunks,
+                 (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GET_ROWS) {
+        const auto * kparams = node.kparams<htp_get_rows_kernel_params>();
+        const char * ktype_str = "unknown";
+        switch (kparams->kernel_type) {
+            case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
+            case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
+            case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
+        }
+        snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
+    } else if (htp_op_is_glu(node.opcode)) {
+        const auto * kparams = node.kparams<htp_glu_kernel_params>();
+        snprintf(str, max_size, "nth %d vtcm %u", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_PAD) {
+        const auto * kparams = node.kparams<htp_pad_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %u", (int) kparams->kernel_type, (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_CUMSUM) {
+        const auto * kparams = node.kparams<htp_cumsum_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FILL) {
+        const auto * kparams = node.kparams<htp_fill_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_ROLL) {
+        const auto * kparams = node.kparams<htp_roll_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d", (int) kparams->kernel_type, (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_DIAG) {
+        const auto * kparams = node.kparams<htp_diag_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_SOLVE_TRI) {
+        const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+    } else {
+        snprintf(str, max_size, "----");
+    }
+    return str;
+}
+
+static void append_section(std::string & res, const std::string & sec) {
+    if (!res.empty()) {
+        res += "|";
+    }
+    res += sec;
+}
+
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_tensor_name(t));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_tensor_dims(t));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_tensor_type(t));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_tensor_strides(t));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_tensor_buff(t));
+    return res;
+}
+
+std::string htp_format_op(const htp_opnode & node, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_OP)      append_section(res, node.name());
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_op_tensors(node, format_tensor_name));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_op_tensors(node, format_tensor_dims));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_op_tensors(node, format_tensor_type));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_op_tensors(node, format_tensor_strides));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_op_tensors(node, format_tensor_buff));
+    if (flags & HTP_OPFMT_KPARAMS) append_section(res, format_kernel_params(node));
+    return res;
+}

ggml/src/ggml-hexagon/htp-opnode.h

@@ -1,38 +1,74 @@
 #ifndef HTP_OPNODE_H
 #define HTP_OPNODE_H
 
-#define GGML_COMMON_IMPL_CPP
 #include "ggml-backend-impl.h"
-#include "ggml-common.h"
+#include "ggml-impl.h"
+#include "htp/htp-ops.h"
 
-#include <algorithm>
+#include <cstring>
+#include <initializer_list>
+#include <memory>
 #include <string>
 #include <vector>
-#include <memory>
-#include <stdio.h>
-#include "htp-ops.h"
-#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"
-#include "htp/argsort-ops.h"
-#include "htp/get-rows-ops.h"
+
+#define HTP_OP_MAX_FUSED 16
+
+// local std::inplace_vector replacement
+template <typename T, size_t MAX_N>
+struct htp_opvec {
+    T      data[MAX_N] {};
+    size_t n = 0;
+
+    htp_opvec() = default;
+    htp_opvec(std::initializer_list<T> init) {
+        GGML_ASSERT(init.size() <= MAX_N);
+        for (const auto & v : init) {
+            data[n++] = v;
+        }
+    }
+
+    size_t    size()  const { return n; }
+    bool      empty() const { return n == 0; }
+    void      clear()       { n = 0; }
+    T *       begin()       { return data; }
+    T *       end()         { return data + n; }
+    const T * begin() const { return data; }
+    const T * end()   const { return data + n; }
+    T &       back()        { return data[n - 1]; }
+    const T & back()  const { return data[n - 1]; }
+    T &       operator[](size_t i)       { return data[i]; }
+    const T & operator[](size_t i) const { return data[i]; }
+
+    void push_back(T val) {
+        GGML_ASSERT(n < MAX_N);
+        data[n++] = val;
+    }
+
+    void insert(T * pos, T val) {
+        GGML_ASSERT(n < MAX_N);
+        for (T * p = data + n; p > pos; --p) {
+            *p = *(p - 1);
+        }
+        *pos = val;
+        n++;
+    }
+};
+
+static inline bool htp_is_hadamard(const struct ggml_tensor * t) {
+    return t && t->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(t, 1) == GGML_HINT_SRC0_IS_HADAMARD;
+}
 
 struct htp_opnode {
-    ggml_tensor * node   { nullptr };
-    htp_op_code   opcode { HTP_OP_INVALID };
-    int32_t       kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
+    const ggml_tensor * node   { nullptr };
+    htp_op_code         opcode { HTP_OP_INVALID };
+    int32_t             params[HTP_OP_MAX_PARAMS] {0};
+    int32_t             kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
 
-    std::vector<ggml_tensor *>                fused;
-    std::vector<std::shared_ptr<ggml_tensor>> dummy;
+    std::vector<std::shared_ptr<ggml_tensor>>          dummy;
 
-    std::vector<const ggml_tensor *> inputs;
-    std::vector<const ggml_tensor *> outputs;
-    std::string                      name;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_FUSED>   fused;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS>  inputs;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> outputs;
 
     int n_active_src(const ggml_tensor * t) const {
         if (!t) return 0;
@@ -44,34 +80,36 @@ struct htp_opnode {
         return 0;
     }
 
-    void init(ggml_tensor * node) {
+    void init(const ggml_tensor * node) {
+        static_assert(sizeof(this->params) == sizeof(this->node->op_params), "params size mismatch");
+
         this->node = node;
-        if (this->node) {
-            if (this->node->op == GGML_OP_MUL_MAT &&
-                ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                this->name = "MUL_MAT_HADAMARD";
-                this->inputs.push_back(this->node->src[1]);
-            } else {
-                this->name = ggml_op_desc(this->node);
+        if (!node) return;
 
-                // Build inputs (preserving optional nullptrs)
-                int n_inputs = n_active_src(this->node);
-                this->inputs.resize(n_inputs, nullptr);
-                for (int i = 0; i < n_inputs; i++) {
-                    this->inputs[i] = this->node->src[i];
-                }
+        memcpy(this->params, this->node->op_params, sizeof(this->params));
+
+        // Build inputs (preserving optional nullptrs)
+        if (htp_is_hadamard(this->node)) {
+            this->inputs.push_back(this->node->src[1]);
+        } else {
+            int n_inputs = n_active_src(this->node);
+            for (int i = 0; i < n_inputs; i++) {
+                this->inputs.push_back(this->node->src[i]);
             }
-
-            // Build outputs
-            this->outputs.push_back(this->dst());
         }
+
+        // Build outputs
+        this->outputs.push_back(this->dst());
     }
 
-    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) {
+    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, const ggml_tensor * node = nullptr) : opcode(opcode) {
+        init(node);
+    }
+
+    htp_opnode(const ggml_tensor * node) : opcode(HTP_OP_INVALID) {
         init(node);
     }
 
-    ggml_op             op()   const { return node->op; }
     const ggml_tensor * src0() const { return node->src[0]; }
     const ggml_tensor * src1() const { return node->src[1]; }
     const ggml_tensor * dst()  const { return outputs.empty() ? node : outputs.back(); }
@@ -81,343 +119,54 @@ struct htp_opnode {
         return dummy.back().get();
     }
 
-    void add_fused(ggml_tensor * t, bool extra_dst = false) {
-        fused.push_back(t);
-
-        name += "+";
-        name += ggml_op_desc(t);
-
-        if (extra_dst) {
-            outputs.push_back(t);
-        } else {
-            outputs.clear();
-            outputs.push_back(t);
-        }
-
-        // Remove the newly fused intermediate output tensor t from inputs (if it was there)
-        inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end());
-
-        // Append new inputs from t, preserving middle nullptrs
-        int n_inputs = n_active_src(t);
-        for (int i = 0; i < n_inputs; i++) {
-            const auto * src = t->src[i];
-            if (!src) {
-                inputs.push_back(nullptr);
-            } else if (src != node &&
-                       std::find(fused.begin(), fused.end(), src) == fused.end() &&
-                       std::find(inputs.begin(), inputs.end(), src) == inputs.end()) {
-                inputs.push_back(src);
-            }
-        }
-    }
-
-    const std::vector<const ggml_tensor *> & get_inputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS> & get_inputs() const {
         return inputs;
     }
 
-    const std::vector<const ggml_tensor *> & get_outputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> & get_outputs() const {
         return outputs;
     }
 
-    std::string op_name() const {
-        return name;
+    template <typename T>
+    T * kparams() {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<T *>(kernel_params);
     }
 
-    bool is_empty() const {
-        return ggml_op_is_empty(node->op);
+    template <typename T>
+    const T * kparams() const {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<const T *>(kernel_params);
     }
 
-    bool stackable() const {
-        switch (this->op()) {
-            case GGML_OP_MUL_MAT:
-                if (ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                    return false;
-                }
-                return ggml_is_quantized(this->src0()->type);
-            case GGML_OP_MUL_MAT_ID:
-                return ggml_is_quantized(this->src0()->type);
-            default:
-                return false;
-        }
+    template <typename T>
+    void set_kparams(const T & kp) {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        memcpy(kernel_params, &kp, sizeof(T));
     }
 
-    bool same_input(const htp_opnode& n) const {
-        return n.src1() == this->src1();
-    }
+    std::string name() const;
 };
 
-struct htp_opformat {
-    char strides[64 * GGML_MAX_SRC];
-    char dims[64 * GGML_MAX_SRC];
-    char types[16 * GGML_MAX_SRC];
-    char buffs[64 * GGML_MAX_SRC];
-    char names[64 * GGML_MAX_SRC];
-    char kparams[128];
+enum htp_opformat_flags {
+    HTP_OPFMT_OP      = (1 << 0),
+    HTP_OPFMT_NAMES   = (1 << 1),
+    HTP_OPFMT_DIMS    = (1 << 2),
+    HTP_OPFMT_TYPES   = (1 << 3),
+    HTP_OPFMT_STRIDES = (1 << 4),
+    HTP_OPFMT_BUFFS   = (1 << 5),
+    HTP_OPFMT_KPARAMS = (1 << 6),
 
-    int format_tensor_dims(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%d:%d", (int) t->ne[0], (int) t->ne[1]);
-        } else {
-            return snprintf(str, max_size, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
-        }
-    }
-
-    void format_op_dims(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_dims(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    int format_tensor_strides(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        const char * c = ggml_is_contiguous(t) ? "" : "!";
-
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
-        } else {
-            return snprintf(str, max_size, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
-        }
-    }
-
-    void format_op_strides(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_strides(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_types(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? ggml_type_name(inputs[0]->type) : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? ggml_type_name(inputs[i]->type) : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", ggml_type_name(node.dst()->type)), (size_t)(p_end - p));
-        }
-    }
-
-    const char * tensor_buff_name(const struct ggml_tensor * t) {
-        if (t && t->buffer) {
-            return ggml_backend_buffer_name(t->buffer);
-        }
-        return "NONE";
-    }
-
-    void format_op_buffs(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[0])), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[i])), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(node.dst())), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_names(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? inputs[0]->name : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? inputs[i]->name : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", node.dst()->name), (size_t)(p_end - p));
-        }
-    }
-    void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
-        if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
-            node.opcode == HTP_OP_MUL_MAT_ADD) {
-            const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
-                path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
-                       type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
-                path = "hvx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-                path = "hvx-tiled";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
-            const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_FA_KERNEL_HMX) {
-                path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
-            } else if (type == HTP_FA_KERNEL_HVX) {
-                path = "hvx";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (htp_op_is_unary(node.opcode)) {
-            const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
-            snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
-        } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
-            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;
-            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 ||
-                   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 if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
-            const auto * kparams = (const struct htp_sort_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
-                     node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
-                     (int) kparams->n_threads, (int) kparams->n_chunks,
-                     (int) kparams->chunk_elems, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_GET_ROWS) {
-            const auto * kparams = (const struct htp_get_rows_kernel_params *) node.kernel_params;
-            const char * ktype_str = "unknown";
-            switch (kparams->kernel_type) {
-                case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
-                case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
-                case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
-            }
-            snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
-        } else {
-            snprintf(str, max_size, "----");
-        }
-
-    }
-
-    void format(const htp_opnode & node) {
-        format_op_dims(dims, sizeof(dims), node);
-        format_op_strides(strides, sizeof(strides), node);
-        format_op_types(types, sizeof(types), node);
-        format_op_buffs(buffs, sizeof(buffs), node);
-        format_op_names(names, sizeof(names), node);
-        format_kernel_params(kparams, sizeof(kparams), node);
-    }
-
-    htp_opformat() {
-        strides[0] = '\0';
-        dims[0]    = '\0';
-        types[0]   = '\0';
-        buffs[0]   = '\0';
-        names[0]   = '\0';
-        kparams[0] = '\0';
-    }
-    htp_opformat(const htp_opnode & node) { format(node); }
+    HTP_OPFMT_EXEC    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_SUPP    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS,
+    HTP_OPFMT_PROF    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_TENSOR  = HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES,
 };
 
+std::string htp_format_op(const htp_opnode & node, uint32_t flags = HTP_OPFMT_EXEC);
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags = HTP_OPFMT_TENSOR);
+
 #endif // HTP_OPNODE_H

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -28,7 +28,7 @@ add_library(${HTP_LIB} SHARED
     unary-ops.c
     sum-rows-ops.c
     softmax-ops.c
-    act-ops.c
+    glu-ops.c
     rope-ops.c
     set-rows-ops.c
     get-rows-ops.c

ggml/src/ggml-hexagon/htp/allreduce-ops.h

@@ -32,6 +32,7 @@ static inline size_t htp_allreduce_vtcm_buffer_count(
 }
 
 struct htp_allreduce_kernel_params {
+    uint8_t kernel_type;
     int32_t rank;
     int32_t n_ranks;
     int32_t n_threads;
@@ -39,7 +40,6 @@ struct htp_allreduce_kernel_params {
     int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
     int32_t vtcm_size_per_thread;
     int32_t vtcm_size;
-    int32_t kernel_type;
     int32_t ne0;
     int32_t ne1;
     int32_t row_size_aligned;
@@ -50,6 +50,12 @@ struct htp_allreduce_kernel_params {
     int32_t mode;                 // enum htp_allreduce_mode: FULL or SHARDED_FANOUT
 };
 
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#endif
+
 #ifdef __cplusplus
 }
 #endif

ggml/src/ggml-hexagon/htp/argsort-ops.c

@@ -453,7 +453,10 @@ static void merge_runs(
 struct htp_sort_chunk_ctx {
     struct htp_ops_context *              octx;
     const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    float *                               merge_values;
+    int32_t *                             merge_indices;
     uint32_t                              real_count[HTP_MAX_NTHREADS];
 };
 
@@ -462,26 +465,27 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
     struct htp_ops_context * octx = cctx->octx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
     const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
 
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-    uint32_t ne00        = (uint32_t) kparams->ne00;
+    uint32_t ne00        = src0->ne[0];
     uint32_t chunk_base  = i * chunk_elems;
 
     uint32_t real_count = (chunk_base < ne00) ? MIN(chunk_elems, ne00 - chunk_base) : 0;
     cctx->real_count[i] = real_count;
 
-    uint8_t * spad = cctx->vtcm_base + (size_t) kparams->phase1_slot_size * i;
+    uint8_t * spad = cctx->phase1_base + cctx->phase1_slot_size * i;
     float *   values_buf  = (float *) spad;
     int32_t * indices_buf = (int32_t *) (spad + chunk_elems * sizeof(float));
 
-    float *   merge_values  = (float *)   (cctx->vtcm_base + kparams->merge_values_off);
-    int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off);
+    float *   merge_values  = cctx->merge_values;
+    int32_t * merge_indices = cctx->merge_indices;
 
     struct htp_thread_trace * tr = &octx->ctx->trace[i];
     const dma_addr_t src_addr = src0->data + (size_t) chunk_base * sizeof(float);
 
     if (kparams->is_top_k) {
-        uint32_t k = (uint32_t) kparams->k;
+        uint32_t k = dst->ne[0];
         uint32_t local_k = MIN(k, chunk_elems);
 
         if (real_count == 0) {
@@ -559,17 +563,17 @@ static void htp_argsort_merge_l1_job(unsigned int n, unsigned int i, void * data
     uint32_t count0 = cctx->real_count[c0];
     uint32_t count1 = cctx->real_count[c1];
 
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
-    const float *   val0 = (const float *)   (cctx->vtcm_base + slot_size * c0);
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + slot_size * c0 + chunk_elems * sizeof(float));
-    const float *   val1 = (const float *)   (cctx->vtcm_base + slot_size * c1);
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + slot_size * c1 + chunk_elems * sizeof(float));
+    const float *   val0 = (const float *)   (cctx->phase1_base + slot_size * c0);
+    const int32_t * idx0 = (const int32_t *) (cctx->phase1_base + slot_size * c0 + chunk_elems * sizeof(float));
+    const float *   val1 = (const float *)   (cctx->phase1_base + slot_size * c1);
+    const int32_t * idx1 = (const int32_t *) (cctx->phase1_base + slot_size * c1 + chunk_elems * sizeof(float));
 
     uint32_t out_offset = (2 * i) * chunk_elems;
-    float *   out_val = (float *)   (cctx->vtcm_base + kparams->merge_values_off)  + out_offset;
-    int32_t * out_idx = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + out_offset;
+    float *   out_val = cctx->merge_values  + out_offset;
+    int32_t * out_idx = cctx->merge_indices + out_offset;
 
     struct htp_thread_trace * tr = &cctx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
@@ -587,21 +591,21 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
     struct htp_sort_chunk_ctx * cctx = mctx->cctx;
     const struct htp_sort_kernel_params * kparams = cctx->kparams;
 
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
     uint32_t off0 = (4 * i) * chunk_elems;
     uint32_t off1 = (4 * i + 2) * chunk_elems;
 
-    const float *   val0 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off0;
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off0;
-    const float *   val1 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off1;
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off1;
+    const float *   val0 = cctx->merge_values  + off0;
+    const int32_t * idx0 = cctx->merge_indices + off0;
+    const float *   val1 = cctx->merge_values  + off1;
+    const int32_t * idx1 = cctx->merge_indices + off1;
 
     uint32_t count0 = cctx->real_count[4*i + 0] + cctx->real_count[4*i + 1];
     uint32_t count1 = cctx->real_count[4*i + 2] + cctx->real_count[4*i + 3];
 
-    uint8_t * slot_base = cctx->vtcm_base + (4 * i) * slot_size;
+    uint8_t * slot_base = cctx->phase1_base + (4 * i) * slot_size;
     float *   out_val   = (float *) slot_base;
     int32_t * out_idx   = (int32_t *) (slot_base + 4 * chunk_elems * sizeof(float));
 
@@ -614,7 +618,10 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
 struct htp_sort_multi_row_ctx {
     struct htp_ops_context *              octx;
     const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    uint32_t                              row_start;
+    uint32_t                              row_end;
     uint32_t                              nrows_per_thread;
 };
 
@@ -624,12 +631,18 @@ struct htp_tensor_row_coord {
     uint32_t i03;
 };
 
-static inline struct htp_tensor_row_coord init_row_coord(uint32_t r, uint32_t ne01, uint32_t ne02) {
+static inline struct htp_tensor_row_coord init_row_coord(
+    uint32_t r,
+    uint32_t ne01,
+    uint32_t ne02,
+    const struct fastdiv_values * div_ne01,
+    const struct fastdiv_values * div_ne01_ne02
+) {
     uint32_t ne01_ne02 = ne01 * ne02;
-    uint32_t i03 = (ne01_ne02 > 0) ? (r / ne01_ne02) : 0;
-    uint32_t rem = (ne01_ne02 > 0) ? (r % ne01_ne02) : 0;
-    uint32_t i02 = (ne01 > 0) ? (rem / ne01) : 0;
-    uint32_t i01 = (ne01 > 0) ? (rem % ne01) : 0;
+    uint32_t i03 = (ne01_ne02 > 0) ? fastdiv(r, div_ne01_ne02) : 0;
+    uint32_t rem = (ne01_ne02 > 0) ? (r - i03 * ne01_ne02) : 0;
+    uint32_t i02 = (ne01 > 0) ? fastdiv(rem, div_ne01) : 0;
+    uint32_t i01 = (ne01 > 0) ? (rem - i02 * ne01) : 0;
     return (struct htp_tensor_row_coord){ i01, i02, i03 };
 }
 
@@ -701,7 +714,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    uint32_t ne00 = (uint32_t) kparams->ne00;
+    uint32_t ne00 = src0->ne[0];
     uint32_t ne01 = src0->ne[1];
     uint32_t ne02 = src0->ne[2];
 
@@ -713,16 +726,16 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     uint32_t nb2 = dst->nb[2];
     uint32_t nb3 = dst->nb[3];
 
-    uint32_t row_start = (uint32_t) kparams->row_start;
-    uint32_t row_end   = (uint32_t) kparams->row_end;
+    uint32_t row_start = mctx->row_start;
+    uint32_t row_end   = mctx->row_end;
     uint32_t r_start   = row_start + ith * mctx->nrows_per_thread;
     uint32_t r_end     = MIN(r_start + mctx->nrows_per_thread, row_end);
 
     if (r_start >= r_end) return;
 
-    size_t spad_offset = ith * (size_t) (kparams->phase1_slot_size * kparams->n_slots);
-    uint8_t * thread_spad = mctx->vtcm_base + spad_offset;
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = mctx->phase1_slot_size;
+    size_t spad_offset = ith * slot_size * (size_t) kparams->n_slots;
+    uint8_t * thread_spad = mctx->phase1_base + spad_offset;
 
     uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
     float *   values_buf  = (float *) thread_spad;
@@ -732,14 +745,14 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     bool is_top_k = kparams->is_top_k != 0;
-    uint32_t k = (uint32_t) kparams->k;
+    uint32_t k = dst->ne[0];
     enum ggml_sort_order order = (enum ggml_sort_order) kparams->order;
 
     const uint32_t src_bytes = ne00 * sizeof(float);
     const uint32_t dst_bytes = (is_top_k ? k : ne00) * sizeof(int32_t);
 
     if (kparams->n_slots == 2) {
-        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
         for (uint32_t r = r_start, spad_idx = 0; r < r_end && spad_idx < 2; r++, spad_idx++) {
             uint8_t * cur_spad = thread_spad + spad_idx * slot_size;
             float *   cur_values  = (float *) cur_spad;
@@ -759,8 +772,8 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
             step_row_coord(&prime_coord, ne01, ne02);
         }
 
-        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02);
-        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02);
+        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
+        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
 
         for (uint32_t r = r_start; r < r_end; r++) {
             uint32_t cur_slot = (r - r_start) & 1;
@@ -796,7 +809,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
 
         dma_queue_flush(dma_q);
     } else {
-        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
         for (uint32_t r = r_start; r < r_end; r++) {
             dma_addr_t src_addr = src0->data + coord_offset(coord, nb01, nb02, nb03);
             dma_addr_t dst_addr = dst->data + coord_offset(coord, nb1, nb2, nb3);
@@ -847,52 +860,39 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     uint32_t ne00 = src0->ne[0];
     uint32_t k    = dst->ne[0];
 
-    struct htp_sort_kernel_params kparams_local;
     const struct htp_sort_kernel_params * kparams = (const struct htp_sort_kernel_params *) octx->kernel_params;
 
-    if (kparams->n_threads == 0) {
-        struct htp_sort_vtcm_layout layout;
-        if (!htp_sort_solve_layout(&layout, ne00, total_rows, k, octx->n_threads, octx->ctx->vtcm_size, is_top_k)) {
-            return HTP_STATUS_VTCM_TOO_SMALL;
-        }
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads         = (int32_t) layout.n_threads;
-        kparams_local.total_rows        = (int32_t) total_rows;
-        kparams_local.row_start         = (int32_t) row_start;
-        kparams_local.row_end           = (int32_t) row_end;
-        kparams_local.ne00              = (int32_t) ne00;
-        kparams_local.k                 = (int32_t) k;
-        kparams_local.order             = is_top_k ? GGML_SORT_ORDER_DESC : (int32_t) octx->op_params[0];
-        kparams_local.is_top_k          = is_top_k ? 1 : 0;
-        kparams_local.use_dma           = 1;
-        kparams_local.chunk_elems       = (int32_t) layout.chunk_elems;
-        kparams_local.n_chunks          = (int32_t) layout.n_chunks;
-        kparams_local.vtcm_size         = (int32_t) layout.total_bytes;
-        kparams_local.phase1_slot_size  = (int32_t) layout.phase1_slot_size;
-        kparams_local.merge_values_off  = (int32_t) layout.merge_values_off;
-        kparams_local.merge_indices_off = (int32_t) layout.merge_indices_off;
-        kparams_local.merge_elems       = (int32_t) layout.merge_elems;
-        kparams_local.n_slots           = (int32_t) layout.n_slots;
-        kparams = &kparams_local;
-    } else if (octx->ctx->mdev.count > 1) {
-        memcpy(&kparams_local, kparams, sizeof(kparams_local));
-        kparams_local.row_start = (int32_t) row_start;
-        kparams_local.row_end   = (int32_t) row_end;
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
     }
 
-    if (octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) {
-        FARF(ERROR, "sort: VTCM size too small. Needed %d, have %zu", kparams->vtcm_size, octx->ctx->vtcm_size);
+    struct htp_sort_vtcm_layout layout;
+    htp_sort_vtcm_layout_build(&layout,
+                               total_rows,
+                               (uint32_t) kparams->chunk_elems,
+                               (uint32_t) kparams->n_chunks,
+                               (uint32_t) kparams->merge_elems,
+                               (uint32_t) kparams->n_threads,
+                               (uint32_t) kparams->n_slots);
+
+    if (octx->ctx->vtcm_size < layout.total_bytes) {
+        FARF(ERROR, "sort: VTCM size too small. Needed %zu, have %zu", layout.total_bytes, octx->ctx->vtcm_size);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * vtcm_base     = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * phase1_base   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_phase1);
+    float *   merge_values  = VTCM_LAYOUT_PTR(float,   vtcm_base, layout.off_merge_values);
+    int32_t * merge_indices = VTCM_LAYOUT_PTR(int32_t, vtcm_base, layout.off_merge_indices);
 
     if (total_rows == 1 && kparams->n_chunks > 1) {
         struct htp_sort_chunk_ctx cctx;
-        cctx.octx      = octx;
-        cctx.kparams   = kparams;
-        cctx.vtcm_base = vtcm_base;
+        cctx.octx             = octx;
+        cctx.kparams          = kparams;
+        cctx.phase1_base      = phase1_base;
+        cctx.phase1_slot_size = layout.phase1_slot_size;
+        cctx.merge_values     = merge_values;
+        cctx.merge_indices    = merge_indices;
         memset(cctx.real_count, 0, sizeof(cctx.real_count));
 
         work_queue_run(octx->ctx->work_queue, htp_sort_chunk_job, &cctx, (uint32_t) kparams->n_chunks);
@@ -902,8 +902,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
             uint32_t local_k     = MIN(k, chunk_elems);
 
             if (k == 1 && ne00 >= 128*1024) {
-                float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-                int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                float *   cand_vals = merge_values;
+                int32_t * cand_idxs = merge_indices;
                 float best_val = cand_vals[0];
                 int32_t best_idx = cand_idxs[0];
                 for (int32_t c = 1; c < kparams->n_chunks; c++) {
@@ -918,8 +918,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
             }
 
             uint32_t total_candidates = (uint32_t) kparams->n_chunks * local_k;
-            float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-            int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+            float *   cand_vals = merge_values;
+            int32_t * cand_idxs = merge_indices;
             uint32_t merge_elems = (uint32_t) kparams->merge_elems;
 
             if (merge_elems > total_candidates) {
@@ -943,18 +943,18 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 struct htp_argsort_merge_l2_ctx l2ctx = { &cctx, asc };
                 work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l2_job, &l2ctx, 2);
 
-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
-                const float *   val0 = (const float *) vtcm_base;
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + 4 * chunk_elems * sizeof(float));
-                const float *   val1 = (const float *) (vtcm_base + 4 * slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));
+                const float *   val0 = (const float *) phase1_base;
+                const int32_t * idx0 = (const int32_t *) (phase1_base + 4 * chunk_elems * sizeof(float));
+                const float *   val1 = (const float *) (phase1_base + 4 * slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));
 
                 uint32_t n01 = cctx.real_count[0] + cctx.real_count[1] + cctx.real_count[2] + cctx.real_count[3];
                 uint32_t n23 = cctx.real_count[4] + cctx.real_count[5] + cctx.real_count[6] + cctx.real_count[7];
 
-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                 struct htp_thread_trace * tr = &octx->ctx->trace[0];
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 merge_runs(val0, idx0, n01, val1, idx1, n23, NULL, final_idx, asc);
@@ -968,27 +968,27 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 uint32_t n01 = cctx.real_count[0] + cctx.real_count[1];
                 uint32_t n23 = cctx.real_count[2] + cctx.real_count[3];
 
-                const float *   val01 = (const float *)   (vtcm_base + kparams->merge_values_off);
-                const int32_t * idx01 = (const int32_t *) (vtcm_base + kparams->merge_indices_off);
-                const float *   val23 = (const float *)   (vtcm_base + kparams->merge_values_off)  + 2 * chunk_elems;
-                const int32_t * idx23 = (const int32_t *) (vtcm_base + kparams->merge_indices_off) + 2 * chunk_elems;
+                const float *   val01 = merge_values;
+                const int32_t * idx01 = merge_indices;
+                const float *   val23 = merge_values  + 2 * chunk_elems;
+                const int32_t * idx23 = merge_indices + 2 * chunk_elems;
 
-                int32_t * final_idx = (int32_t *) vtcm_base;
+                int32_t * final_idx = (int32_t *) phase1_base;
                 struct htp_thread_trace * tr = &octx->ctx->trace[0];
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 merge_runs(val01, idx01, n01, val23, idx23, n23, NULL, final_idx, asc);
                 htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             } else if (kparams->n_chunks == 2) {
-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
 
-                const float *   val0 = (const float *)   (vtcm_base);
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + chunk_elems * sizeof(float));
-                const float *   val1 = (const float *)   (vtcm_base + slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + slot_size + chunk_elems * sizeof(float));
+                const float *   val0 = (const float *)   (phase1_base);
+                const int32_t * idx0 = (const int32_t *) (phase1_base + chunk_elems * sizeof(float));
+                const float *   val1 = (const float *)   (phase1_base + slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + slot_size + chunk_elems * sizeof(float));
 
-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                 struct htp_thread_trace * tr = &octx->ctx->trace[0];
                 htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                 merge_runs(val0, idx0, cctx.real_count[0], val1, idx1, cctx.real_count[1], NULL, final_idx, asc);
@@ -996,7 +996,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             } else {
                 uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-                int32_t * final_idx = (int32_t *) (vtcm_base + chunk_elems * sizeof(float));
+                int32_t * final_idx = (int32_t *) (phase1_base + chunk_elems * sizeof(float));
                 dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
             }
         }
@@ -1006,7 +1006,10 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     struct htp_sort_multi_row_ctx mctx;
     mctx.octx             = octx;
     mctx.kparams          = kparams;
-    mctx.vtcm_base        = vtcm_base;
+    mctx.phase1_base      = phase1_base;
+    mctx.phase1_slot_size = layout.phase1_slot_size;
+    mctx.row_start        = row_start;
+    mctx.row_end          = row_end;
 
     uint32_t dst_row_bytes = (uint32_t) dst_row_size;
     uint32_t rows_per_chunk = 0;
@@ -1014,13 +1017,13 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
     if (rows_per_chunk == 0) {
         rows_per_chunk = 1;
     }
-    uint32_t dr = (nrows + (uint32_t) kparams->n_threads - 1) / (uint32_t) kparams->n_threads;
+    uint32_t dr = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
     if (rows_per_chunk > 1) {
         dr = hex_round_up(dr, rows_per_chunk);
     }
     mctx.nrows_per_thread = dr;
 
-    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, (uint32_t) kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, octx->n_threads);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/argsort-ops.h

@@ -7,33 +7,27 @@
 #include <string.h>
 
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 
 struct htp_sort_kernel_params {
-    int32_t  n_threads;
-    int32_t  total_rows;
-    int32_t  row_start;
-    int32_t  row_end;
-    int32_t  ne00;
-    int32_t  k;
-    int32_t  order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
-    int32_t  is_top_k;          // 1 if TOP_K, 0 if ARGSORT
-    int32_t  use_dma;           // 1 if DMA enabled
-    int32_t  chunk_elems;
-    int32_t  n_chunks;
-    int32_t  vtcm_size;
-    int32_t  phase1_slot_size;
-    int32_t  merge_values_off;
-    int32_t  merge_indices_off;
-    int32_t  merge_elems;
-    int32_t  n_slots;
-    int32_t  pad[15];
+    uint32_t n_threads;
+    uint32_t order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
+    uint32_t is_top_k;          // 1 if TOP_K, 0 if ARGSORT
+    uint32_t chunk_elems;
+    uint32_t n_chunks;
+    uint32_t vtcm_size;
+    uint32_t merge_elems;
+    uint32_t n_slots;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne01_ne02;
 };
 
 struct htp_sort_vtcm_layout {
     size_t   total_bytes;
     size_t   phase1_slot_size;
-    size_t   merge_values_off;
-    size_t   merge_indices_off;
+    size_t   off_phase1;
+    size_t   off_merge_values;
+    size_t   off_merge_indices;
     uint32_t chunk_elems;
     uint32_t n_chunks;
     uint32_t merge_elems;
@@ -41,6 +35,39 @@ struct htp_sort_vtcm_layout {
     uint32_t n_slots;
 };
 
+static inline void htp_sort_vtcm_layout_build(
+    struct htp_sort_vtcm_layout * L,
+    uint32_t total_rows,
+    uint32_t chunk_elems,
+    uint32_t n_chunks,
+    uint32_t merge_elems,
+    uint32_t n_threads,
+    uint32_t n_slots) {
+
+    memset(L, 0, sizeof(*L));
+    L->chunk_elems = chunk_elems;
+    L->n_chunks    = n_chunks;
+    L->merge_elems = merge_elems;
+    L->n_threads   = n_threads;
+    L->n_slots     = n_slots;
+
+    const size_t values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+    const size_t indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+    L->phase1_slot_size       = ((values_size + indices_size) + 255) & ~255;
+
+    size_t off = 0;
+    if (total_rows > 1) {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_slots * n_threads);
+    } else {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_chunks * n_slots);
+        const size_t merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
+        const size_t merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_values,  merge_values_size,  n_chunks > 1);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_indices, merge_indices_size, n_chunks > 1);
+    }
+    L->total_bytes = off;
+}
+
 static inline bool htp_sort_solve_layout(
     struct htp_sort_vtcm_layout * layout,
     uint32_t ne00,
@@ -75,18 +102,8 @@ static inline bool htp_sort_solve_layout(
             n_threads--;
         }
 
-        size_t total_bytes = spad_per_thread * n_threads;
-        if (total_bytes > vtcm_budget) {
-            return false;
-        }
-
-        layout->total_bytes      = total_bytes;
-        layout->phase1_slot_size = spad_per_slot;
-        layout->chunk_elems      = ne00_padded;
-        layout->n_chunks         = 1;
-        layout->n_threads        = n_threads;
-        layout->n_slots          = n_slots;
-        return true;
+        htp_sort_vtcm_layout_build(layout, total_rows, ne00_padded, 1, 0, n_threads, n_slots);
+        return layout->total_bytes <= vtcm_budget;
     }
 
     uint32_t n_vec = (ne00 + 31) / 32;
@@ -102,17 +119,7 @@ static inline bool htp_sort_solve_layout(
 
     uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
     uint32_t chunk_elems = chunk_n_vec * 32;
-
-    size_t phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-    size_t phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-    size_t phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-    size_t phase1_total_size   = phase1_slot_size * n_chunks;
-
-    size_t merge_values_size  = 0;
-    size_t merge_indices_size = 0;
-    size_t merge_values_off   = phase1_total_size;
-    size_t merge_indices_off  = 0;
-    uint32_t merge_elems      = 0;
+    uint32_t merge_elems = 0;
 
     if (n_chunks > 1) {
         if (is_top_k) {
@@ -125,47 +132,25 @@ static inline bool htp_sort_solve_layout(
         } else {
             merge_elems = n_vec_pow2 * 32;
         }
-        merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
-        merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
-        merge_indices_off  = merge_values_off + merge_values_size;
     }
 
-    size_t total_bytes = phase1_total_size + merge_values_size + merge_indices_size;
+    htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, n_chunks, 1);
 
-    if (total_bytes > vtcm_budget && n_chunks > 1) {
-        n_chunks = 1;
+    if (layout->total_bytes > vtcm_budget && n_chunks > 1) {
+        n_chunks    = 1;
         chunk_elems = n_vec_pow2 * 32;
-        phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-        phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-        phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-        phase1_total_size   = phase1_slot_size;
-        merge_values_size   = 0;
-        merge_indices_size  = 0;
-        merge_values_off    = phase1_total_size;
-        merge_indices_off   = 0;
-        merge_elems         = 0;
-        total_bytes         = phase1_total_size;
+        merge_elems = 0;
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, 1, 1);
     }
 
-    if (total_bytes > vtcm_budget) {
+    if (layout->total_bytes > vtcm_budget) {
         return false;
     }
 
-    uint32_t n_slots = 1;
-    if (n_chunks == 1 && phase1_slot_size * 2 <= vtcm_budget) {
-        n_slots = 2;
-        total_bytes = phase1_slot_size * 2;
+    if (n_chunks == 1 && layout->phase1_slot_size * 2 <= vtcm_budget) {
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, 1, 0, 1, 2);
     }
 
-    layout->total_bytes       = total_bytes;
-    layout->phase1_slot_size  = phase1_slot_size;
-    layout->merge_values_off  = merge_values_off;
-    layout->merge_indices_off = merge_indices_off;
-    layout->chunk_elems       = chunk_elems;
-    layout->n_chunks          = n_chunks;
-    layout->merge_elems       = merge_elems;
-    layout->n_threads         = n_chunks;
-    layout->n_slots           = n_slots;
     return true;
 }
 

ggml/src/ggml-hexagon/htp/binary-ops.c

@@ -42,6 +42,7 @@ struct htp_binary_context {
     uint32_t nrows_per_thread;
     uint32_t total_rows;
     uint32_t row_start;
+    uint32_t n_srcs;
     size_t   src0_row_size_aligned;
     size_t   src1_row_size_aligned;
     size_t   dst_row_size_aligned;
@@ -380,6 +381,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
     size_t dst_spad_half  = layout->dst_spad_half_size;
 
     dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = bctx->n_srcs;
     uint32_t ir_prefetch = start_row;
     int spad_idx = 0;
 
@@ -391,12 +393,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
         i02 = fastdiv(rem, &bctx->src0_dim1_div);
         i01 = rem - i02 * ne01;
 
-        uint32_t i13 = (ne13 == 1) ? 0 : i03;
-        uint32_t i12 = (ne12 == 1) ? 0 : i02;
-        uint32_t i11 = (ne11 == 1) ? 0 : i01;
-
         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;
@@ -405,7 +402,15 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
 
         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);
+        for (uint32_t s = 1; s < n_srcs; s++) {
+            const struct htp_tensor * src_s = octx->src[s];
+            uint32_t is3 = (src_s->ne[3] == 1) ? 0 : i03;
+            uint32_t is2 = (src_s->ne[2] == 1) ? 0 : i02;
+            uint32_t is1 = (src_s->ne[1] == 1) ? 0 : i01;
+            dma_addr_t s_curr = src_s->data + is3 * src_s->nb[3] + is2 * src_s->nb[2] + is1 * src_s->nb[1];
+            uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, current_block_size);
+        }
         ir_prefetch += current_block_size;
         spad_idx ^= 1;
     }
@@ -422,6 +427,11 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         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);
+        for (uint32_t s = 2; s < n_srcs; s++) {
+            uint8_t * ss_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+            compute(d_spad, d_spad, ss_spad, current_block_size,
+                    bctx->dst_row_size_aligned, bctx->dst_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;
@@ -440,15 +450,18 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
             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;
-
             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(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);
+
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                const struct htp_tensor * src_s = octx->src[s];
+                uint32_t ps3 = (src_s->ne[3] == 1) ? 0 : p03;
+                uint32_t ps2 = (src_s->ne[2] == 1) ? 0 : p02;
+                uint32_t ps1 = (src_s->ne[1] == 1) ? 0 : p01;
+                dma_addr_t s_next = src_s->data + ps3 * src_s->nb[3] + ps2 * src_s->nb[2] + ps1 * src_s->nb[1];
+                uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_next), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, next_block_size);
+            }
 
             ir_prefetch += next_block_size;
         }
@@ -859,6 +872,7 @@ struct binary_chunked_context {
     uint32_t                 chunk_size;
     uint32_t                 chunks_per_thread;
     uint32_t                 total_chunks;
+    uint32_t                 n_srcs;
     bool                     is_scalar;
     float                    scalar_f32;
     _Float16                 scalar_f16;
@@ -872,7 +886,6 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
     struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
     struct htp_ops_context * octx = ctx->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 uint32_t start_chunk = ctx->chunks_per_thread * ith;
@@ -895,6 +908,7 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
     uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
 
     dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = ctx->n_srcs;
     uint32_t prefetch_chunk = start_chunk;
     int spad_idx = 0;
 
@@ -913,9 +927,12 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
         dma_queue_push(dma_q, dma_make_data(d_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 0);
         dma_queue_push(dma_q, dma_make_data(s0_spad, s0_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
         if (!ctx->is_scalar) {
-            dma_addr_t s1_curr = src1->data + (size_t) c_start * elem_size;
             uint8_t * s1_spad = src1_spad_base + spad_idx * chunk_bytes;
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                dma_addr_t s_curr = octx->src[s]->data + (size_t) c_start * elem_size;
+                uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            }
         }
 
         prefetch_chunk++;
@@ -943,6 +960,10 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
             }
         } else {
             ctx->compute(d_spad, s0_spad, s1_spad, cur_elems);
+            for (uint32_t s = 2; s < n_srcs; s++) {
+                uint8_t * s_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+                ctx->compute(d_spad, d_spad, s_spad, cur_elems);
+            }
         }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
 
@@ -958,8 +979,11 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
             dma_addr_t s0_next = src0->data + (size_t) pc_start * elem_size;
             dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), chunk_bytes, chunk_bytes, p_bytes, 1);
             if (!ctx->is_scalar) {
-                dma_addr_t s1_next = src1->data + (size_t) pc_start * elem_size;
-                dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                for (uint32_t s = 1; s < n_srcs; s++) {
+                    dma_addr_t s_next = octx->src[s]->data + (size_t) pc_start * elem_size;
+                    uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                    dma_queue_push(dma_q, dma_make_data(s_spad, s_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                }
             }
 
             prefetch_chunk++;
@@ -974,6 +998,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     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 n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;
 
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
 
@@ -985,7 +1010,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const size_t dst_row_size  = dst->ne[0]  * elem_size;
 
     if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
-        const bool is_scalar = kparams->is_scalar || (src1->ne[0] == 1 && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
+        const bool is_scalar = kparams->is_scalar;
 
         compute_binary_chunked_t compute = NULL;
         compute_binary_scalar_chunked_f32_t compute_scalar_f32 = NULL;
@@ -1045,12 +1070,16 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         const uint32_t elems_per_line = (src0_type == HTP_TYPE_F32) ? 32 : 64;
 
         if (octx->ctx->mdev.count > 1) {
-            const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
-                                   htp_tensor_mdev_data_aligned(src0) &&
-                                   (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
-                                   htp_tensor_is_contiguous(dst, elem_size) &&
-                                   htp_tensor_is_contiguous(src0, elem_size) &&
-                                   (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                             htp_tensor_mdev_data_aligned(src0) &&
+                             (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
+                             htp_tensor_is_contiguous(dst, elem_size) &&
+                             htp_tensor_is_contiguous(src0, elem_size) &&
+                             (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            for (uint32_t s = 2; can_split && s < n_srcs; s++) {
+                can_split = htp_tensor_mdev_data_aligned(octx->src[s]) &&
+                            htp_tensor_is_contiguous(octx->src[s], elem_size);
+            }
             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);
@@ -1072,7 +1101,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             return HTP_STATUS_VTCM_TOO_SMALL;
         }
 
-        const uint32_t chunk_size = kparams->chunk_size > 0 ? kparams->chunk_size : (32768 / elem_size);
+        const uint32_t chunk_size = kparams->chunk_size;
         const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
         const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
         const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
@@ -1101,6 +1130,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
             .chunk_size          = chunk_size,
             .chunks_per_thread   = chunks_per_thread,
             .total_chunks        = total_chunks,
+            .n_srcs              = n_srcs,
             .is_scalar           = is_scalar,
             .scalar_f32          = scalar_f32,
             .scalar_f16          = scalar_f16,
@@ -1146,65 +1176,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         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 ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
-        if (!kparams->src1_strided) {
-            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
-            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
-        } else {
-            const uint32_t ne11 = src1->ne[1];
-            const uint32_t ne12 = src1->ne[2];
-            const uint32_t ne13 = src1->ne[3];
-            const size_t row_bytes = src1->ne[0] * elem_size;
-            for (uint32_t i13 = 0; i13 < ne13; i13++) {
-                for (uint32_t i12 = 0; i12 < ne12; i12++) {
-                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
-                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
-                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
-                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
-                        dma_queue_pop(dma_q);
-                    }
-                }
-            }
-        }
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
-        const struct htp_tensor * src2 = octx->src[2];
-        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
-        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]);
-        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
-                       src2_row_bytes, src2->nb[1],
-                       src2_row_bytes, src2->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             = 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]);
-    bctx.src0_dim12_div = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
-
-    bctx.src1_dim1_div  = init_fastdiv_values(src1->ne[1]);
-    bctx.src1_dim2_div  = init_fastdiv_values(src1->ne[2]);
-    bctx.src1_dim3_div  = init_fastdiv_values(src1->ne[3]);
-
-    bctx.split_at_ne01  = kparams->split_at_ne01;
-    bctx.split_at_ne02  = kparams->split_at_ne02;
-
     worker_callback_t worker_func = NULL;
     void * compute_func = NULL;
 
@@ -1322,7 +1293,70 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    bctx.compute = compute_func;
+    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 ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
+        if (!kparams->src1_strided) {
+            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
+            if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL) {
+                uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+                dma_queue_push(dma_q, dma_make_data(vtcm_src2, octx->src[2]->data), bctx.vtcm_layout.src2_size, 0, src1_bytes, 1);
+            }
+        } else {
+            const uint32_t ne11 = src1->ne[1];
+            const uint32_t ne12 = src1->ne[2];
+            const uint32_t ne13 = src1->ne[3];
+            const size_t row_bytes = src1->ne[0] * elem_size;
+            for (uint32_t i13 = 0; i13 < ne13; i13++) {
+                for (uint32_t i12 = 0; i12 < ne12; i12++) {
+                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
+                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
+                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
+                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
+                        dma_queue_pop(dma_q);
+                    }
+                }
+            }
+        }
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
+        const struct htp_tensor * src2 = octx->src[2];
+        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+        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]);
+        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
+                       src2_row_bytes, src2->nb[1],
+                       src2_row_bytes, src2->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.n_srcs                = n_srcs;
+    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  = kparams->div_ne01;
+    bctx.src0_dim2_div  = kparams->div_ne02;
+    bctx.src0_dim12_div = kparams->div_ne012;
+
+    bctx.src1_dim1_div  = kparams->div_ne11;
+    bctx.src1_dim2_div  = kparams->div_ne12;
+    bctx.src1_dim3_div  = kparams->div_ne13;
+
+    bctx.split_at_ne01  = kparams->split_at_ne01;
+    bctx.split_at_ne02  = kparams->split_at_ne02;
+    bctx.compute        = compute_func;
 
     if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST ||
         kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
@@ -1330,6 +1364,15 @@ static int execute_op_binary(struct htp_ops_context * octx) {
         kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX ||
         kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
         dma_queue_flush(dma_q);
+        if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL && octx->op == HTP_OP_MUL) {
+            uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+            const uint32_t n_scalars = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+            if (src0_type == HTP_TYPE_F32) {
+                hvx_mul_f32_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            } else if (src0_type == HTP_TYPE_F16) {
+                hvx_mul_f16_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            }
+        }
     }
 
     work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);

ggml/src/ggml-hexagon/htp/binary-ops.h

@@ -6,6 +6,7 @@
 #include <string.h>
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
 #include "htp-ops.h"
 #include "htp-vtcm.h"
 
@@ -20,7 +21,15 @@ enum htp_binary_kernel_type {
 };
 
 struct htp_binary_kernel_params {
-    uint32_t kernel_type;
+    uint8_t  kernel_type;
+    uint8_t  split_at_ne01;
+    uint8_t  split_at_ne02;
+    uint8_t  src1_strided;
+
+    uint8_t  n_srcs;
+    uint8_t  is_scalar;
+    uint8_t  pad[2];
+
     uint32_t n_threads;
     uint32_t rows_per_buffer;
 
@@ -34,12 +43,14 @@ struct htp_binary_kernel_params {
 
     uint32_t chunk_size;
     uint32_t chunk_bytes;
-    uint32_t is_scalar;
 
-    uint8_t  split_at_ne01;
-    uint8_t  split_at_ne02;
-    uint8_t  src1_strided;
-    uint8_t  pad;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne012;
+
+    struct fastdiv_values div_ne11;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne13;
 };
 
 #if defined(__cplusplus)
@@ -100,6 +111,8 @@ static inline void htp_binary_vtcm_layout_build(
         return;
     }
 
+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;
+
     if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
         const size_t chunk_bytes = kparams->chunk_bytes;
         if (chunk_bytes == 0) {
@@ -107,7 +120,7 @@ static inline void htp_binary_vtcm_layout_build(
         }
 
         L->src0_bytes_per_thread = 2 * chunk_bytes;
-        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * chunk_bytes);
+        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * (n_srcs - 1) * chunk_bytes);
         L->dst_bytes_per_thread  = 2 * chunk_bytes;
 
         L->src0_spad_half_size = chunk_bytes;
@@ -135,7 +148,7 @@ static inline void htp_binary_vtcm_layout_build(
     }
 
     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 + (n_srcs - 1) * kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
         : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);
 
     const size_t static_vtcm_size = kparams->src1_size + kparams->src2_size;
@@ -143,7 +156,9 @@ static inline void htp_binary_vtcm_layout_build(
         return;
     }
 
-    const size_t rows_per_buffer = (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
+    const size_t rows_per_buffer = (kparams->rows_per_buffer > 0)
+        ? kparams->rows_per_buffer
+        : (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
     if (rows_per_buffer == 0) {
         return;
     }
@@ -155,11 +170,13 @@ static inline void htp_binary_vtcm_layout_build(
     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
+        ? rows_per_buffer * 2 * (n_srcs - 1) * 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->src1_spad_half_size = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+        ? rows_per_buffer * kparams->src1_row_size_aligned
+        : 0;
     L->dst_spad_half_size  = L->dst_bytes_per_thread / 2;
 
     const size_t src0_total = n_threads * L->src0_bytes_per_thread;
@@ -178,4 +195,53 @@ static inline void htp_binary_vtcm_layout_build(
     L->total_bytes = off;
 }
 
+static inline bool htp_binary_solve_chunked(
+    struct htp_binary_vtcm_layout * L,
+    struct htp_binary_kernel_params * kparams,
+    uint32_t total_elems,
+    uint32_t elem_size,
+    uint32_t max_threads,
+    uint32_t n_srcs,
+    bool is_scalar,
+    size_t vtcm_size
+) {
+    const size_t n_bufs = is_scalar ? 2 : (size_t) n_srcs + 1;
+    const size_t scalar_bytes = is_scalar ? 128 : 0;
+    const size_t budget_per_buf = (max_threads > 0 && vtcm_size > scalar_bytes)
+        ? hex_align_down((vtcm_size - scalar_bytes) / (2 * max_threads * n_bufs), 128)
+        : 0;
+    const uint32_t max_chunk_bytes = (budget_per_buf > 0 && budget_per_buf < 32768) ? (uint32_t) budget_per_buf : 32768u;
+    const uint32_t max_chunk_elems = max_chunk_bytes / elem_size;
+    if (max_chunk_elems == 0) {
+        return false;
+    }
+
+    const uint32_t min_chunk_elems = 256 < max_chunk_elems ? 256 : max_chunk_elems;
+    const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * max_threads) - 1) / (2 * max_threads), 32);
+    const uint32_t chunk_size  = (target_chunk_elems < min_chunk_elems) ? min_chunk_elems
+                               : (target_chunk_elems > max_chunk_elems) ? max_chunk_elems
+                               : target_chunk_elems;
+    const uint32_t chunk_bytes  = hex_round_up(chunk_size * elem_size, 128);
+    const uint32_t total_chunks = (total_elems + chunk_size - 1) / chunk_size;
+    const uint32_t n_threads    = (total_chunks >= 2) ? (max_threads < total_chunks ? max_threads : total_chunks) : 1;
+
+    kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+    kparams->n_threads             = n_threads;
+    kparams->n_srcs                = (uint8_t) n_srcs;
+    kparams->rows_per_buffer       = 1;
+    kparams->src1_row_size_aligned = is_scalar ? 0 : kparams->src0_row_size_aligned;
+    kparams->src1_size             = 0;
+    kparams->chunk_size            = chunk_size;
+    kparams->chunk_bytes           = chunk_bytes;
+    kparams->is_scalar             = is_scalar ? 1 : 0;
+
+    htp_binary_vtcm_layout_build(L, kparams, vtcm_size);
+    if (L->total_bytes == 0 || L->total_bytes > vtcm_size) {
+        return false;
+    }
+
+    kparams->vtcm_size = (uint32_t) L->total_bytes;
+    return true;
+}
+
 #endif

ggml/src/ggml-hexagon/htp/concat-ops.c

@@ -294,22 +294,21 @@ static int concat_transposed(struct htp_ops_context * octx, const struct htp_con
 
     const uint32_t n_threads = octx->n_threads;
 
-    // layout precomputed on host; kept for reference:
-    // struct htp_concat_transposed_vtcm_layout layout;
-    // htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
+    struct htp_concat_transposed_vtcm_layout layout;
+    htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
 
     struct htp_concat_context cctx;
     cctx.octx                  = octx;
-    cctx.spad0_base            = vtcm_base;
-    cctx.spad1_base            = vtcm_base + n_threads * kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad0_size_per_thread = kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad1_size_per_thread = kparams->u.transposed.spad1_size_per_thread;
+    cctx.spad0_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad0);
+    cctx.spad1_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad1);
+    cctx.spad0_size_per_thread = layout.src0_spad_size_per_thread;
+    cctx.spad1_size_per_thread = layout.src1_spad_size_per_thread;
     cctx.row_start             = row_start;
     cctx.nrows                 = nrows;
     cctx.nplanes               = dst->ne[2] * dst->ne[3];
-    cctx.div_ne2               = init_fastdiv_values(dst->ne[2]);
+    cctx.div_ne2               = kparams->u.transposed.div_ne2;
     cctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     work_queue_func_t worker_func = (type_size == 4) ? concat_2d_f32_transposed : concat_2d_f16_transposed;

ggml/src/ggml-hexagon/htp/concat-ops.h

@@ -2,6 +2,8 @@
 #define HTP_CONCAT_OPS_H
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>
 
@@ -13,8 +15,7 @@ enum htp_concat_kernel_type {
 };
 
 struct htp_concat_transposed_params {
-    uint32_t spad0_size_per_thread;
-    uint32_t spad1_size_per_thread;
+    struct fastdiv_values div_ne2;
 };
 
 struct htp_concat_kernel_params {
@@ -40,11 +41,13 @@ _Static_assert(sizeof(struct htp_concat_kernel_params) <= 128, "htp_concat_kerne
 struct htp_concat_transposed_vtcm_layout {
     uint32_t src0_spad_size_per_thread;
     uint32_t src1_spad_size_per_thread;
+    size_t   off_spad0;
+    size_t   off_spad1;
     uint32_t total_bytes;
 };
 
 static inline void htp_concat_transposed_vtcm_layout_build(
-    struct htp_concat_transposed_vtcm_layout * layout,
+    struct htp_concat_transposed_vtcm_layout * L,
     uint32_t src0_ne0,
     uint32_t src1_ne0,
     uint32_t type_size,
@@ -55,9 +58,13 @@ static inline void htp_concat_transposed_vtcm_layout_build(
     uint32_t src1_ne0_padded = hex_round_up(src1_ne0, block_i);
     uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + src1_ne0_padded * type_size;
 
-    layout->src0_spad_size_per_thread = block_i * spad0_row_bytes;
-    layout->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
-    layout->total_bytes = n_threads * (layout->src0_spad_size_per_thread + layout->src1_spad_size_per_thread);
+    L->src0_spad_size_per_thread = block_i * spad0_row_bytes;
+    L->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_spad1, (size_t) n_threads * L->src1_spad_size_per_thread);
+    L->total_bytes = (uint32_t) off;
 }
 
 #endif // HTP_CONCAT_OPS_H

ggml/src/ggml-hexagon/htp/cpy-ops.c

@@ -32,6 +32,7 @@ struct htp_copy_context {
     uint32_t nelem;
     uint32_t elem_per_thread;
 
+    struct htp_copy_convert_vtcm_layout layout;
     uint8_t * vtcm_src0;
     uint8_t * vtcm_dst;
 };
@@ -176,13 +177,14 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                   \
                                                                                              \
     const struct htp_copy_convert_params * cvt = &ct->kparams->u.convert;                    \
-    const uint32_t src0_buf_size = cvt->src0_buf_size;                                       \
-    const uint32_t dst_buf_size  = cvt->dst_buf_size;                                        \
-    const uint32_t src0_row_stride = cvt->src0_row_stride;                                   \
-    const uint32_t dst_row_stride  = cvt->dst_row_stride;                                    \
-    const uint32_t blk_rows        = cvt->blk_rows;                                          \
-    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * cvt->spad0_size_per_thread;             \
-    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * cvt->spad1_size_per_thread;             \
+    const struct htp_copy_convert_vtcm_layout * L = &ct->layout;                             \
+    const uint32_t src0_buf_size   = L->src0_buf_size;                                       \
+    const uint32_t dst_buf_size    = L->dst_buf_size;                                        \
+    const uint32_t src0_row_stride = L->src0_row_stride;                                     \
+    const uint32_t dst_row_stride  = L->dst_row_stride;                                      \
+    const uint32_t blk_rows        = L->blk_rows;                                            \
+    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * L->spad0_size_per_thread;               \
+    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * L->spad1_size_per_thread;               \
     const uint32_t src0_row_size = ne00 * ct->kparams->src0_type_size;                       \
     const uint32_t dst_row_size  = ne00 * ct->kparams->dst_type_size;                        \
                                                                                              \
@@ -254,17 +256,16 @@ static int cpy_scalar(struct htp_ops_context * octx, const struct htp_copy_kerne
     }
 
     dma_queue * dma_q = octx->ctx->dma[0];
-    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) octx->ctx->vtcm_base;
-    dma_addr_t d_vtcm = s_vtcm + VLEN;
+    uint8_t * s_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, 0);
+    uint8_t * d_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, VLEN);
+    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) s_ptr;
+    dma_addr_t d_vtcm = (dma_addr_t)(uintptr_t) d_ptr;
     const uint32_t s_size = kparams->src0_type_size;
     const uint32_t d_size = kparams->dst_type_size;
 
     dma_queue_push(dma_q, dma_make_data(s_vtcm, src0->data), s_size, s_size, s_size, 1);
     dma_queue_pop(dma_q);
 
-    uint8_t * s_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    uint8_t * d_ptr = s_ptr + VLEN;
-
     const HVX_Vector v_src = hvx_vmem(s_ptr);
     HVX_Vector v_dst;
 
@@ -366,6 +367,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
     const struct htp_copy_convert_params * cvt = &kparams->u.convert;
 
     struct htp_copy_context ct;
+    htp_copy_convert_vtcm_layout_build(&ct.layout, src0->ne[0], kparams->src0_type_size, kparams->dst_type_size, cvt->blk_rows, n_threads);
+
     ct.octx                  = octx;
     ct.kparams               = kparams;
     ct.row_start             = row_start;
@@ -373,8 +376,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
     ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
-    ct.vtcm_src0 = vtcm_base;
-    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * cvt->spad0_size_per_thread;
+    ct.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_src0);
+    ct.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_dst);
 
     work_queue_func_t copy_fun = NULL;
     if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {

ggml/src/ggml-hexagon/htp/cpy-ops.h

@@ -3,6 +3,7 @@
 
 #include "hex-common.h"
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>
 
@@ -17,12 +18,6 @@ enum htp_copy_kernel_type {
 };
 
 struct htp_copy_convert_params {
-    uint32_t              src0_buf_size;
-    uint32_t              dst_buf_size;
-    uint32_t              spad0_size_per_thread;
-    uint32_t              spad1_size_per_thread;
-    uint32_t              src0_row_stride;  // row pitch in the VTCM buffers
-    uint32_t              dst_row_stride;
     uint32_t              blk_rows;         // rows moved per DMA descriptor
     struct fastdiv_values div_ne01;
     struct fastdiv_values div_ne02_ne01;
@@ -60,6 +55,9 @@ struct htp_copy_kernel_params {
 #define HTP_COPY_VTCM_BLK_BYTES (32 * 1024)
 
 struct htp_copy_convert_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src0;
+    size_t   off_dst;
     uint32_t src0_row_stride;
     uint32_t dst_row_stride;
     uint32_t blk_rows;
@@ -67,11 +65,35 @@ struct htp_copy_convert_vtcm_layout {
     uint32_t dst_buf_size;
     uint32_t spad0_size_per_thread;
     uint32_t spad1_size_per_thread;
-    uint32_t total_bytes;
 };
 
-// ne01 bounds a block: rows of one block share the dim-1 stride.
 static inline void htp_copy_convert_vtcm_layout_build(
+    struct htp_copy_convert_vtcm_layout * L,
+    uint32_t ne00,
+    uint32_t src_type_size,
+    uint32_t dst_type_size,
+    uint32_t blk_rows,
+    uint32_t n_threads) {
+
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+
+    L->src0_row_stride       = src0_row_stride;
+    L->dst_row_stride        = dst_row_stride;
+    L->blk_rows              = blk_rows;
+    L->src0_buf_size         = blk_rows * src0_row_stride;
+    L->dst_buf_size          = blk_rows * dst_row_stride;
+    L->spad0_size_per_thread = 2 * L->src0_buf_size;
+    L->spad1_size_per_thread = 2 * L->dst_buf_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, (size_t) n_threads * L->spad0_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  (size_t) n_threads * L->spad1_size_per_thread);
+    L->total_bytes = off;
+}
+
+// ne01 bounds a block: rows of one block share the dim-1 stride.
+static inline void htp_copy_convert_solve_layout(
     struct htp_copy_convert_vtcm_layout * layout,
     uint32_t ne00,
     uint32_t ne01,
@@ -79,10 +101,10 @@ static inline void htp_copy_convert_vtcm_layout_build(
     uint32_t dst_type_size,
     uint32_t n_threads) {
 
-    layout->src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
-    layout->dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
 
-    const uint32_t max_row = layout->src0_row_stride > layout->dst_row_stride ? layout->src0_row_stride : layout->dst_row_stride;
+    const uint32_t max_row = src0_row_stride > dst_row_stride ? src0_row_stride : dst_row_stride;
     uint32_t blk_rows = HTP_COPY_VTCM_BLK_BYTES / max_row;
     if (blk_rows > ne01) {
         blk_rows = ne01;
@@ -90,13 +112,8 @@ static inline void htp_copy_convert_vtcm_layout_build(
     if (blk_rows == 0) {
         blk_rows = 1;
     }
-    layout->blk_rows = blk_rows;
 
-    layout->src0_buf_size = blk_rows * layout->src0_row_stride;
-    layout->dst_buf_size  = blk_rows * layout->dst_row_stride;
-    layout->spad0_size_per_thread = 2 * layout->src0_buf_size;
-    layout->spad1_size_per_thread = 2 * layout->dst_buf_size;
-    layout->total_bytes = n_threads * (layout->spad0_size_per_thread + layout->spad1_size_per_thread);
+    htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
 }
 
 #if defined(__cplusplus)

ggml/src/ggml-hexagon/htp/cumsum-ops.c

@@ -12,9 +12,11 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
 #include "dma-queue.h"
+#include "cumsum-ops.h"
 
 #define htp_cumsum_tensors_preamble                         \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -89,8 +91,8 @@ static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) {
 }
 
 static inline void hvx_cumsum_row_f32(const float * restrict src, float * restrict dst, uint32_t n) {
-    const uint32_t nvec = n / VLEN_FP32;
-    const uint32_t nloe = n % VLEN_FP32;
+    const uint32_t nvec = n >> 5;
+    const uint32_t nloe = n & (VLEN_FP32 - 1);
 
     HVX_Vector carry = Q6_V_vsplat_R(0);
 
@@ -210,16 +212,19 @@ static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_cumsum_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
+    const struct htp_cumsum_kernel_params * kparams =
+        (const struct htp_cumsum_kernel_params *) octx->kernel_params;
 
-    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);
+    const uint32_t total_rows   = kparams->total_rows;
+    const size_t   src_row_size = src0->nb[1];
+    const size_t   dst_row_size = dst->nb[1];
 
     uint32_t row_start = 0;
     uint32_t nrows     = total_rows;
 
     if (octx->ctx->mdev.count > 1) {
         uint32_t rows_per_chunk = 0;
-        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_start = range.start;
         nrows     = range.count;
@@ -229,43 +234,45 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
 
-    const size_t src_row_size         = src0->nb[1];
-    const size_t dst_row_size         = dst->nb[1];
-    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
-    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
 
-    // 2 ping-pong buffers per thread for src and dst
-    const size_t spad_per_thread = 2 * (src_row_size_aligned + dst_row_size_aligned);
+    struct htp_cumsum_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_cumsum_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);
 
-    octx->src0_spad.size_per_thread = src_row_size_aligned * 2;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned * 2;
-
-    octx->src0_spad.size  = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size   = n_threads * octx->dst_spad.size_per_thread;
-
-    octx->src0_spad.data  = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data   = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }
 
     struct htp_cumsum_context cctx = {
         .octx                 = octx,
         .src_row_size         = src_row_size,
         .dst_row_size         = dst_row_size,
-        .src_row_size_aligned = src_row_size_aligned,
-        .dst_row_size_aligned = dst_row_size_aligned,
+        .src_row_size_aligned = layout.src_row_size_aligned,
+        .dst_row_size_aligned = layout.dst_row_size_aligned,
         .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
         .total_rows           = nrows,
         .row_start            = row_start,
     };
 
-    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 {
+    if (kparams->use_dma) {
         work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/cumsum-ops.h

@@ -0,0 +1,49 @@
+#ifndef HTP_CUMSUM_OPS_H
+#define HTP_CUMSUM_OPS_H
+
+#include "hex-common.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_cumsum_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_row_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_cumsum_vtcm_layout_build(
+    struct htp_cumsum_vtcm_layout * L,
+    size_t                          src_row_size,
+    size_t                          dst_row_size,
+    uint32_t                        n_threads
+) {
+    L->src_row_size_aligned  = hex_round_up(src_row_size, 128);
+    L->dst_row_size_aligned  = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread = 2 * L->src_row_size_aligned;
+    L->dst_bytes_per_thread  = 2 * L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_cumsum_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t use_dma;
+    uint32_t vtcm_size;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_cumsum_kernel_params) <= 128, "htp_cumsum_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_cumsum_kernel_params) <= 128, "htp_cumsum_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_CUMSUM_OPS_H

ggml/src/ggml-hexagon/htp/diag-ops.c

@@ -10,10 +10,12 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
 #include "dma-queue.h"
+#include "diag-ops.h"
 
 #define htp_diag_tensors_preamble                           \
     const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -33,6 +35,7 @@
 
 struct htp_diag_context {
     struct htp_ops_context * octx;
+    const struct htp_diag_kernel_params * kparams;
     size_t          src_batch_size;
     size_t          dst_row_size;
     size_t          src_batch_size_aligned;
@@ -83,8 +86,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;
 
         const dma_addr_t src_batch = src_data + i3 * nb03 + i2 * nb02;
 
@@ -134,8 +137,8 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
 
     for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;
 
         const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02);
 
@@ -155,29 +158,17 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_diag_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
+    const struct htp_diag_kernel_params * kparams =
+        (const struct htp_diag_kernel_params *) octx->kernel_params;
 
-    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
+    const uint32_t total_batches = kparams->total_batches;
 
     uint32_t batch_start = 0;
     uint32_t nbatches    = total_batches;
 
     if (octx->ctx->mdev.count > 1) {
-        bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->ne[0] == 1 || dst->nb[0] == sizeof(float)) && !htp_tensor_is_permuted(dst);
-        uint32_t batches_per_chunk = 1;
-        if (can_split) {
-            if (dst->ne[2] > 1 && (dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
-                (dst->ne[3] <= 1 || (dst->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
-                batches_per_chunk = 1;
-            } else if (dst->nb[2] == dst_batch_size &&
-                       (dst->ne[3] <= 1 || dst->nb[3] == dst->nb[2] * dst->ne[2])) {
-                batches_per_chunk = (dst_batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(dst_batch_size, HEX_L2_LINE_SIZE)) : 1;
-            } else {
-                can_split = false;
-            }
-        }
-
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const bool can_split = (kparams->can_split != 0) && htp_tensor_mdev_data_aligned(dst);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? kparams->batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         batch_start = range.start;
         nbatches    = range.count;
     }
@@ -186,43 +177,49 @@ int op_diag_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads     = octx->n_threads;
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nbatches))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
 
-    const size_t src_batch_size         = src0->ne[0] * sizeof(float);
-    const size_t dst_row_size           = dst->ne[0] * sizeof(float);
-    const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN);
-    const size_t dst_row_size_aligned   = hex_round_up(dst_row_size, VLEN);
+    const uint32_t n_threads = octx->n_threads;
 
-    // 1 src buffer + 1 dst row buffer per thread
-    const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned;
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
 
-    octx->src0_spad.size_per_thread = src_batch_size_aligned;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned;
+    const size_t src_batch_size = src0->ne[0] * sizeof(float);
+    const size_t dst_row_size   = dst->ne[0] * sizeof(float);
 
-    octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
+    struct htp_diag_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_diag_vtcm_layout_build(&layout, src_batch_size, dst_row_size, n_threads);
 
-    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }
 
     struct htp_diag_context dctx = {
         .octx                   = octx,
+        .kparams                = kparams,
         .src_batch_size         = src_batch_size,
         .dst_row_size           = dst_row_size,
-        .src_batch_size_aligned = src_batch_size_aligned,
-        .dst_row_size_aligned   = dst_row_size_aligned,
+        .src_batch_size_aligned = layout.src_batch_size_aligned,
+        .dst_row_size_aligned   = layout.dst_row_size_aligned,
         .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
         .total_batches          = nbatches,
         .batch_start            = batch_start,
     };
 
-    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 {
+    if (kparams->use_dma) {
         work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/diag-ops.h

@@ -0,0 +1,54 @@
+#ifndef HTP_DIAG_OPS_H
+#define HTP_DIAG_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_diag_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_batch_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_diag_vtcm_layout_build(
+    struct htp_diag_vtcm_layout * L,
+    size_t                        src_batch_size,
+    size_t                        dst_row_size,
+    uint32_t                      n_threads
+) {
+    L->src_batch_size_aligned = hex_round_up(src_batch_size, 128);
+    L->dst_row_size_aligned   = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread  = L->src_batch_size_aligned;
+    L->dst_bytes_per_thread   = L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_diag_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_batches;
+    uint32_t use_dma;
+    uint32_t can_split;
+    uint32_t batches_per_chunk;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_ne02;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_diag_kernel_params) <= 128, "htp_diag_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_diag_kernel_params) <= 128, "htp_diag_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_DIAG_OPS_H

ggml/src/ggml-hexagon/htp/dma-queue.h

@@ -436,10 +436,10 @@ static inline void dma_sync_write(dma_queue * dma_q, dma_addr_t dst, const void
     }
 }
 
-#define DMA_CACHE_MAX_SIZE 256U
+#define DMA_CACHE_MAX_SIZE 128U
 
 // Fully assoc LRU cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
     uint8_t *base;
     uint32_t line_size;
     uint32_t capacity;
@@ -485,7 +485,7 @@ static inline bool dma_cache_fa_push(dma_queue *q, dma_cache_fa *c, dma_addr_t s
 }
 
 // Direct mapped cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
     uint8_t *base;
     uint32_t line_size;
     uint32_t capacity;

ggml/src/ggml-hexagon/htp/fill-ops.c

@@ -16,6 +16,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "fill-ops.h"
 
 // ggml op_params layout for FILL:
 //   op_params[0] (as float) - the scalar fill value
@@ -30,21 +31,19 @@
                                                \
     const uint32_t nb1 = dst->nb[1];           \
     const uint32_t nb2 = dst->nb[2];           \
-    const uint32_t nb3 = dst->nb[3];           \
-                                               \
-    const uint32_t nr = ne1 * ne2 * ne3;
+    const uint32_t nb3 = dst->nb[3];
 
 struct htp_fill_context {
     struct htp_ops_context * octx;
+    const struct htp_fill_kernel_params * kparams;
     uint32_t nrows_per_thread;
     uint32_t total_rows;  // ne1 * ne2 * ne3
     uint32_t row_start;
-    bool     opt_path;
     HVX_Vector splat_vec;
     uint32_t   elem_size;
 };
 
-static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
+static void fill_thread_contiguous(unsigned int nth, unsigned int ith, void * data) {
     const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
     struct htp_ops_context        * octx = fctx->octx;
     fill_preamble;
@@ -60,45 +59,59 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
 
-    if (fctx->opt_path) {
-        // Opt path: tensor is fully contiguous, treat as flat array
-        const uint32_t elem_start = ir0 * ne0;
-        const uint32_t elem_end = ir1 * ne0;
-        uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
-        hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
-    } else {
-        // Non-contiguous path: must respect strides
-        for (uint32_t ir = ir0; ir < ir1; ++ir) {
-            const uint32_t i1 = ir % ne1;
-            const uint32_t i2 = (ir / ne1) % ne2;
-            const uint32_t i3 = ir / (ne1 * ne2);
-            uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
-            hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
-        }
+    const uint32_t elem_start = ir0 * ne0;
+    const uint32_t elem_end   = ir1 * ne0;
+    uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
+    hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
+    FARF(HIGH, "fill-contig %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+}
+
+static void fill_thread_strided(unsigned int nth, unsigned int ith, void * data) {
+    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
+    struct htp_ops_context        * octx = fctx->octx;
+    const struct htp_fill_kernel_params * kparams = fctx->kparams;
+    fill_preamble;
+
+    // Parallelise over the flat row index spanning ne1*ne2*ne3
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
+
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+
+    for (uint32_t ir = ir0; ir < ir1; ++ir) {
+        const uint32_t i3  = fastdiv(ir, &kparams->div_ne1_ne2);
+        const uint32_t rem = ir - i3 * (ne1 * ne2);
+        const uint32_t i2  = fastdiv(rem, &kparams->div_ne1);
+        const uint32_t i1  = rem - i2 * ne1;
+        uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
+        hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
     }
 
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
-    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+    FARF(HIGH, "fill-strided %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }
 
 int op_fill(struct htp_ops_context * octx) {
-    fill_preamble;
-
-    if (dst->type != HTP_TYPE_F32 && dst->type != HTP_TYPE_F16) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
+    const struct htp_tensor * dst = octx->dst;
+    const struct htp_fill_kernel_params * kparams =
+        (const struct htp_fill_kernel_params *) octx->kernel_params;
 
     if (htp_tensor_is_extended(dst)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
     uint32_t row_start = 0;
-    uint32_t nrows     = nr;
+    uint32_t nrows     = kparams->total_rows;
 
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = nb1;
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? 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(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_start = range.start;
         nrows     = range.count;
     }
@@ -107,40 +120,43 @@ int op_fill(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
 
-    // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
-    const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
-
-    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u opt=%d\n",
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, (int) opt_path);
+    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u ktype=%u\n",
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, kparams->kernel_type);
 
     float val_f32 = 0.f;
     memcpy(&val_f32, &octx->op_params[0], sizeof(float));
 
     struct htp_fill_context fctx = {
         .octx             = octx,
+        .kparams          = kparams,
         .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
         .total_rows       = nrows,
         .row_start        = row_start,
-        .opt_path         = opt_path,
+        .elem_size        = kparams->elem_size,
     };
 
     switch (dst->type) {
     case HTP_TYPE_F32:
         fctx.splat_vec = hvx_vec_splat_f32(val_f32);
-        fctx.elem_size = sizeof(float);
         break;
     case HTP_TYPE_F16:
         fctx.splat_vec = hvx_vec_splat_f16((_Float16) val_f32);
-        fctx.elem_size = sizeof(_Float16);
         break;
     default:
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);
+    work_queue_func_t worker = (kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS)
+        ? fill_thread_contiguous
+        : fill_thread_strided;
+
+    work_queue_run(octx->ctx->work_queue, worker, &fctx, n_threads);
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/fill-ops.h

@@ -0,0 +1,29 @@
+#ifndef HTP_FILL_OPS_H
+#define HTP_FILL_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+enum htp_fill_kernel_type {
+    HTP_FILL_KERNEL_CONTIGUOUS = 0,
+    HTP_FILL_KERNEL_STRIDED    = 1,
+};
+
+struct htp_fill_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t elem_size;
+    uint32_t rows_per_chunk;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne1_ne2;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_fill_kernel_params) <= 128, "htp_fill_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fill_kernel_params) <= 128, "htp_fill_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_FILL_OPS_H

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -39,6 +39,26 @@
 // Must be multiple of 32
 #define FLASH_ATTN_BLOCK_SIZE (32 * 2)
 
+// Large scratch lives in ctx->ddr_spad to keep stack usage low.
+// Ops can run on the dspqueue callback thread, which has a small stack.
+struct hvx_fa_ddr_spad {
+    __fp16   slopes[512];
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    uint16_t thread_active_blocks[HTP_MAX_NTHREADS][FA_MAX_KV_BLOCKS] __attribute__((aligned(128)));
+};
+
+struct hmx_fa_ddr_spad {
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    __fp16   slopes[] __attribute__((aligned(128))); // [g_br]
+};
+
+static inline void * fa_get_ddr_spad(const struct htp_ops_context * octx, size_t size) {
+    if (!octx->ctx->ddr_spad_base || octx->ctx->ddr_spad_size < size) {
+        return NULL;
+    }
+    return octx->ctx->ddr_spad_base;
+}
+
 struct htp_fa_context {
     const struct htp_ops_context * octx;
 
@@ -61,7 +81,6 @@ struct htp_fa_context {
     uint32_t n_head_log2;
     float m0;
     float m1;
-    __fp16 slopes[512];
 
     uint32_t n_blocks;
     uint32_t n_active_blocks;
@@ -90,7 +109,9 @@ struct htp_fa_context {
     uint8_t *  spad_m;
     uint8_t *  spad_a;
     float *    spad_sinks;
-    uint16_t   active_blocks[FA_MAX_KV_BLOCKS];
+
+    struct hvx_fa_ddr_spad * ddr;
+    dma_cache_dm *           m_cache; // [n_threads]
 
     const struct htp_tensor * k;
     const struct htp_tensor * v;
@@ -156,6 +177,7 @@ struct hmx_fa_context {
     __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
+    __fp16 *     ddr_slopes;           // ALiBi slopes staging buffer in DDR [g_br]
     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;
@@ -164,7 +186,7 @@ struct hmx_fa_context {
     size_t       col_vec_bytes;
     size_t       d_tile_bytes;
     bool         mask_broadcast;       // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
-    dma_cache_fa m_cache;
+    dma_cache_fa * m_cache;
 };
 
 // Scan mask rows [q_start, q_start + n_rows_q) and build compact list of KV blocks with any entry != -inf.
@@ -308,11 +330,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
     uint8_t * spad_m = factx->spad_m + (mask ? factx->size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * ith;
     uint8_t * spad_a = factx->spad_a + factx->size_vkq_acc * ith;
 
-    dma_cache_dm m_cache;
-    dma_cache_dm_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
+    dma_cache_dm * m_cache = &factx->m_cache[ith];
+    dma_cache_dm_init(m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
 
-    uint16_t thread_active_blocks[FA_MAX_KV_BLOCKS];
-    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->active_blocks : thread_active_blocks;
+    uint16_t * thread_active_blocks = factx->ddr->thread_active_blocks[ith];
+    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->ddr->active_blocks : thread_active_blocks;
     uint32_t n_active_blocks       = factx->n_active_blocks;
     dma_addr_t prev_mp_base        = 0;
 
@@ -385,7 +407,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             const uint32_t r_iq1 = (r - r_iq3*neq2*neq1 - r_iq2 * neq1);
 
             heads[g] = r_iq2;
-            const __fp16 slope = factx->slopes[r_iq2];
+            const __fp16 slope = factx->ddr->slopes[r_iq2];
             slope_vecs[g] = hvx_vec_splat_f16(slope);
 
             S_vec[g] = hvx_vec_splat_f32(0.0f);
@@ -422,7 +444,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             // Mask
             if (mask) {
                 const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16);
-                dma_cache_dm_push(dma_q, &m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                dma_cache_dm_push(dma_q, m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
             }
         }
 
@@ -682,7 +704,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     // Mask
                     if (mask) {
                         const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16);
-                        dma_cache_dm_push(dma_q, &m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
+                        dma_cache_dm_push(dma_q, m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
                     }
                 }
             } // end for g
@@ -1887,7 +1909,7 @@ static __attribute__((noinline)) void fa_compute_slopes(
     const float    m0          = factx->m0;
     const float    m1          = factx->m1;
 
-    __fp16 temp_slopes[512] __attribute__((aligned(128)));
+    __fp16 * temp_slopes = factx->ddr_slopes;
     if (G <= 32) {
         // Fast path: Compute G unique slope values in vector registers
         HVX_Vector v_val = hvx_alibi_slopes(kv_head, G, n_head_log2, m0, m1);
@@ -1905,14 +1927,13 @@ static __attribute__((noinline)) void fa_compute_slopes(
         }
     }
 
-    // Allocate stack buffer to avoid scalar writes to VTCM (which generates L2 misses)
-    __fp16 local_slopes[n_rows_g] __attribute__((aligned(128)));
-    for (size_t r = 0; r < n_rows_g; ++r) {
-        local_slopes[r] = temp_slopes[fastmodulo(r, G, &factx->div_G)];
+    // Replicate the G values in DDR to avoid scalar writes to VTCM (which generates L2 misses)
+    for (size_t r = G; r < n_rows_g; ++r) {
+        temp_slopes[r] = temp_slopes[r - G];
     }
 
     // Copy to VTCM slopes using HVX block copy (both are aligned to 128 bytes)
-    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)local_slopes, n_rows_g);
+    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)temp_slopes, n_rows_g);
 }
 
 static void fa_push_mask_dma_gqa(
@@ -1958,7 +1979,7 @@ static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor
     if (mask) {
         if (__builtin_expect(factx->mask_broadcast, true)) {
             const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
-            dma_cache_fa_push(dma_q, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+            dma_cache_fa_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_q, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
         }
@@ -2151,7 +2172,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
     const size_t m_line_bytes = L.m_line_bytes;  // used by the mask DMAs in the KV loop
 
-    dma_cache_fa_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+    struct hmx_fa_ddr_spad * ddr = (struct hmx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hmx_fa_ddr_spad) + factx.g_br * sizeof(__fp16));
+    if (!ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+
+    factx.ddr_slopes = ddr->slopes;
+    factx.m_cache    = &octx->dma_cache_fa[0];
+    dma_cache_fa_init(factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
 
     // Head-dim padding: the K/V DMA staging buffers and the flat-Q buffer are laid out
     // with padded row strides (size_{k,v,q}_row_padded, covering D_pad columns) but the
@@ -2210,7 +2238,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
             const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
             const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
 
-            uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+            uint16_t * active_blocks = ddr->active_blocks;
             uint32_t n_active_blocks = factx.n_kv_blocks;
             if (mask && factx.mask_broadcast && factx.n_kv_blocks > 1) {
                 const size_t scratch_bytes = factx.k_tile_bytes * (factx.pipeline ? 2 : 1);
@@ -2252,7 +2280,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (factx.pipeline && mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                             } else {
                                 fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                             }
@@ -2453,7 +2481,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_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_q, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                             }
@@ -2577,7 +2605,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (factx.pipeline && mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                             } else {
                                 fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, next_kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                             }
@@ -2680,8 +2708,15 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     if (n_head > 512) {
         return HTP_STATUS_NO_SUPPORT;
     }
+
+    factx.ddr = (struct hvx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hvx_fa_ddr_spad));
+    if (!factx.ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+    factx.m_cache = octx->dma_cache_dm;
+
     for (uint32_t h = 0; h < n_head; ++h) {
-        factx.slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
+        factx.ddr->slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
     }
 
     // total rows in q
@@ -2754,7 +2789,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     if (has_mask && neq1 == 1 && neq3 == 1 && mask->ne[2] == 1 && factx.n_blocks > 1 && factx.n_blocks <= FA_MAX_KV_BLOCKS) {
         factx.n_active_blocks = fa_scan_active_blocks(
             octx->ctx->dma[0], mask, 0, 1, 0, 0, k->ne[1], FLASH_ATTN_BLOCK_SIZE,
-            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.active_blocks);
+            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.ddr->active_blocks);
         factx.shared_active_blocks = true;
     }
 

ggml/src/ggml-hexagon/htp/flash-attn-ops.h

@@ -35,7 +35,6 @@ enum htp_fa_kernel_type {
 struct htp_fa_kernel_params {
     uint8_t  kernel_type;        // enum htp_fa_kernel_type
     uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
-    uint8_t  flags;              // reserved
     uint8_t  n_threads;          // Number of threads to run
 
     // Common parameters
@@ -49,9 +48,6 @@ struct htp_fa_kernel_params {
     float    logit_softcap;
     uint32_t vtcm_size;
 
-    uint32_t qrows;
-    uint32_t qrows_per_thread;
-    uint32_t qrow_start;
     float    m0;
     float    m1;
     uint32_t n_head_log2;
@@ -67,8 +63,6 @@ struct htp_fa_kernel_params {
     union {
         struct {
             uint32_t g_br;
-            uint32_t row_buf_stride;
-            uint32_t mask_buf_row_stride;
             int32_t  mask_broadcast;
             int32_t  pipeline;
             struct fastdiv_values div_G;
@@ -85,6 +79,8 @@ struct htp_fa_kernel_params {
 
 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
 #endif
 
 // VTCM region layout for the HMX flash-attention kernel.

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c

@@ -2270,46 +2270,9 @@ static int gated_delta_net_f32_hmx_chunked(
 }
 
 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];
-    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 = octx->dst;
 
-    if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
-        g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
-        dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    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];
-
-    if (S_v == 0 || S_v > HTP_GDN_MAX_SV || H == 0 || n_tokens == 0 || n_seqs == 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if ((g->ne[0] != 1 && g->ne[0] != S_v) || beta->ne[0] != 1) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (q->ne[0] != S_v || k->ne[0] != S_v || q->ne[1] == 0 || k->ne[1] == 0 ||
-        q->ne[2] != n_tokens || k->ne[2] != n_tokens || q->ne[3] == 0 || k->ne[3] == 0 ||
-        (n_seqs % q->ne[3]) != 0 || (n_seqs % k->ne[3]) != 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    // state holds s0 only: [S_v, S_v, H, n_seqs]
-    if (state->ne[0] != S_v || state->ne[1] != S_v || state->ne[2] != H || state->ne[3] != n_seqs) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (dst->ne[0] != S_v * H || dst->ne[1] != n_tokens * n_seqs + S_v * n_seqs * K) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    for (int i = 0; i < 5; i++) {
+    for (int i = 0; i < 6; i++) {
         if (htp_tensor_is_extended(octx->src[i])) {
             return HTP_STATUS_NO_SUPPORT;
         }
@@ -2325,68 +2288,12 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
     }
 
     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;
-        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));
-        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;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
     }
 
+    const uint32_t S_v        = kparams->S_v;
+    const uint32_t n_tokens   = kparams->n_tokens;
     const uint32_t total_rows = kparams->total_rows;
     uint32_t row_start = 0;
     uint32_t nrows     = total_rows;
@@ -2416,16 +2323,16 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
         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;
+    const uint32_t n_threads = (nrows < octx->n_threads) ? nrows : octx->n_threads;
 
     struct htp_gdn_context gctx;
     gctx.octx      = octx;
     gctx.kparams   = kparams;
     gctx.row_start = row_start;
     gctx.nrows     = nrows;
-    gctx.vtcm_base = octx->ctx->vtcm_base;
 
     htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_threads);
+    gctx.vtcm_base = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, gctx.layout.off_state);
 
     if (gctx.layout.total_bytes > octx->ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2433,10 +2340,10 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
 
     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",
-         q->ne[0], q->ne[1], q->ne[2], q->ne[3],
-         k->ne[0], k->ne[1], k->ne[2], k->ne[3],
-         v->ne[0], v->ne[1], v->ne[2], v->ne[3],
-         state->ne[0], state->ne[1], state->ne[2], state->ne[3],
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
+         octx->src[2]->ne[0], octx->src[2]->ne[1], octx->src[2]->ne[2], octx->src[2]->ne[3],
+         octx->src[5]->ne[0], octx->src[5]->ne[1], octx->src[5]->ne[2], octx->src[5]->ne[3],
          dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
          gctx.layout.total_bytes, n_threads);
 

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h

@@ -37,13 +37,7 @@ struct htp_gdn_kernel_params {
     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;
@@ -55,7 +49,6 @@ struct htp_gdn_kernel_params {
     struct fastdiv_values div_k1;
     struct fastdiv_values div_rq3;
     struct fastdiv_values div_rk3;
-    struct fastdiv_values div_n_threads;
 };
 
 #if defined(__cplusplus)
@@ -67,18 +60,21 @@ _Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_para
 struct htp_gdn_vtcm_layout {
     size_t state_aligned;
     size_t bytes_per_thread;
+    size_t off_state;
     size_t total_bytes;
 };
 
 static inline void htp_gdn_vtcm_layout_build(
-    struct htp_gdn_vtcm_layout * layout,
+    struct htp_gdn_vtcm_layout * L,
     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;
+    L->state_aligned = hex_round_up(state_bytes, 128);
+    L->bytes_per_thread = 2 * L->state_aligned;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_state, L->bytes_per_thread * n_threads);
+    L->total_bytes = off;
 }
 
 struct htp_gdn_hmx_vtcm_layout {

ggml/src/ggml-hexagon/htp/gather-rows.c

@@ -206,14 +206,22 @@ static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * da
 
     struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];
 
-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 4);
 
-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };
 
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
     gather_rows_init_tables_f32(tab, inc, kparams, a, b);
@@ -279,14 +287,22 @@ static void gather_rows_thread_f16(unsigned int nth, unsigned int ith, void * da
 
     struct htp_thread_trace * tr = &task->octx->ctx->trace[ith];
 
-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 2);
 
-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };
 
     htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
     gather_rows_init_tables_f16(tab, inc, kparams, a, b);

ggml/src/ggml-hexagon/htp/gather-rows.h

@@ -2,6 +2,7 @@
 #define HTP_GATHER_ROWS_H
 
 #include "hex-common.h"
+#include "htp-vtcm.h"
 #include <stdint.h>
 #include <stdbool.h>
 
@@ -19,8 +20,6 @@ struct htp_gather_rows_params {
     uint32_t v1;
     uint32_t span1;
     uint32_t region_size;
-    uint32_t out_off;
-    uint32_t buf_stride;
     uint32_t slice;
     uint32_t total_rows;
     uint32_t rows_per_thread;
@@ -28,6 +27,40 @@ struct htp_gather_rows_params {
     uint8_t  pad[3];
 };
 
+struct htp_gather_rows_vtcm_layout {
+    size_t off_tab;
+    size_t off_inc;
+    size_t off_cur;
+    size_t off_region0;
+    size_t off_out0;
+    size_t off_region1;
+    size_t off_out1;
+    size_t total_per_thread;
+};
+
+static inline void htp_gather_rows_vtcm_layout_build(
+    struct htp_gather_rows_vtcm_layout * L,
+    uint32_t vg,
+    uint32_t region_size,
+    uint32_t nr_max,
+    uint32_t ne,
+    uint32_t elem_size) {
+
+    const size_t vec_bytes      = (size_t) vg * 128;
+    const size_t region_aligned = hex_round_up(region_size, 128);
+    const size_t out_aligned    = hex_round_up(region_aligned + (size_t) nr_max * ne * elem_size, 128) - region_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_tab,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_inc,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_cur,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_region0, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out0,    out_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_region1, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out1,    out_aligned);
+    L->total_per_thread = off;
+}
+
 static inline bool htp_gather_rows_solve_layout(
     struct htp_gather_rows_params * kparams,
     uint32_t na, uint32_t a_nb1,
@@ -80,8 +113,6 @@ static inline bool htp_gather_rows_solve_layout(
     uint32_t v1          = 0;
     uint32_t span1       = 0;
     uint32_t region_size = 0;
-    uint32_t out_off     = 0;
-    uint32_t buf_stride  = 0;
 
     while (nr_max > 0) {
         v1 = hex_round_up((nr_max - 1) * a_nb1 + na * elem_size, 128);
@@ -98,10 +129,10 @@ static inline bool htp_gather_rows_solve_layout(
             region_size = v1;
         }
 
-        out_off    = hex_round_up(region_size, 128);
-        buf_stride = hex_round_up(out_off + nr_max * ne * elem_size, 128);
+        struct htp_gather_rows_vtcm_layout layout;
+        htp_gather_rows_vtcm_layout_build(&layout, vg, region_size, nr_max, ne, elem_size);
 
-        if ((elem_size != 2 || region_size <= 32768) && fixed + 2 * buf_stride <= slice) {
+        if ((elem_size != 2 || region_size <= 32768) && layout.total_per_thread <= slice) {
             break;
         }
 
@@ -121,8 +152,6 @@ static inline bool htp_gather_rows_solve_layout(
     kparams->v1              = v1;
     kparams->span1           = span1;
     kparams->region_size     = region_size;
-    kparams->out_off         = out_off;
-    kparams->buf_stride      = buf_stride;
     kparams->slice           = slice;
     kparams->total_rows      = total_rows;
     kparams->rows_per_thread = rows_per_thread;

ggml/src/ggml-hexagon/htp/get-rows-ops.c

@@ -404,7 +404,7 @@ int op_get_rows(struct htp_ops_context * octx) {
 
     if (octx->ctx->mdev.count > 1) {
         uint32_t tasks_per_chunk = 1;
-        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], dst_row_size, &tasks_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) dst->nb[0], dst_row_size, &tasks_per_chunk);
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         task_start = range.start;
         tasks      = range.count;
@@ -414,7 +414,7 @@ int op_get_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -428,7 +428,7 @@ int op_get_rows(struct htp_ops_context * octx) {
     grctx.compute_tiled = NULL;
     grctx.task_start = task_start;
     grctx.tasks = tasks;
-    grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
     grctx.tile_size   = 0;
     grctx.tile_stride = 0;
     grctx.src0_row_size = src0_row_size;

ggml/src/ggml-hexagon/htp/get-rows-ops.h

@@ -15,11 +15,10 @@ enum htp_get_rows_kernel_type {
 };
 
 struct htp_get_rows_kernel_params {
-    int32_t  n_threads;
-    int32_t  kernel_type;
-    int32_t  total_tasks;
-    int32_t  tasks_per_thread;
-    int32_t  vtcm_size;
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_tasks;
+    uint32_t vtcm_size;
 
     // Fastdiv helpers
     struct fastdiv_values div_ne10;

ggml/src/ggml-hexagon/htp/glu-ops.c

@@ -17,11 +17,12 @@
 #include "hex-common.h"
 #include "htp-tensor.h"
 #include "htp-vtcm.h"
+#include "glu-ops.h"
 
-#define htp_act_preamble                                 \
-    const struct htp_tensor * src0 = actx->octx->src[0]; \
-    const struct htp_tensor * src1 = actx->octx->src[1]; \
-    const struct htp_tensor * dst  = actx->octx->dst;    \
+#define htp_glu_preamble                                 \
+    const struct htp_tensor * src0 = gctx->octx->src[0]; \
+    const struct htp_tensor * src1 = gctx->octx->src[1]; \
+    const struct htp_tensor * dst  = gctx->octx->dst;    \
                                                          \
     const uint32_t ne00 = src0->ne[0];                   \
     const uint32_t ne01 = src0->ne[1];                   \
@@ -53,21 +54,21 @@
     const uint32_t nb2 = dst->nb[2];                     \
     const uint32_t nb3 = dst->nb[3];
 
-struct htp_act_context;
+struct htp_glu_context;
 
 typedef void (*glu_compute_f32_fn_t)(const float * restrict src0,
                                      const float * restrict src1,
                                      float * restrict dst,
                                      const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+                                     const struct htp_glu_context * gctx);
 
 typedef void (*glu_compute_f16_fn_t)(const _Float16 * restrict src0,
                                      const _Float16 * restrict src1,
                                      _Float16 * restrict dst,
                                      const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+                                     const struct htp_glu_context * gctx);
 
-struct htp_act_context {
+struct htp_glu_context {
     struct htp_ops_context * octx;
 
     void *                   compute;
@@ -108,53 +109,18 @@ struct htp_act_context {
     size_t                   vtcm_dst_size_per_thread;
 };
 
-struct htp_act_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;
-
-    uint32_t vtcm_row_per_thread;
-};
-
-static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L,
-                                             size_t                       src0_row_size_aligned,
-                                             size_t                       src1_row_size_aligned,
-                                             size_t                       dst_row_size_aligned,
-                                             uint32_t                     n_threads,
-                                             size_t                       vtcm_size) {
-    const size_t   spad_size_per_row   = src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned;
-    const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * spad_size_per_row));
-
-    L->vtcm_row_per_thread = vtcm_row_per_thread;
-
-    L->src0_bytes_per_thread = src0_row_size_aligned * vtcm_row_per_thread;
-    L->src1_bytes_per_thread = src1_row_size_aligned * vtcm_row_per_thread;
-    L->dst_bytes_per_thread  = dst_row_size_aligned * vtcm_row_per_thread;
-
-    L->off_src0 = 0;
-    L->off_src1 = L->off_src0 + L->src0_bytes_per_thread * n_threads;
-    L->off_dst  = L->off_src1 + L->src1_bytes_per_thread * n_threads;
-
-    L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
-}
-
 #define htp_glu_op_preamble                                            \
-    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;   \
-    const int    nc                    = actx->nc;
+    const size_t src0_row_size_aligned = gctx->src0_row_size_aligned;  \
+    const size_t src1_row_size_aligned = gctx->src1_row_size_aligned;  \
+    const size_t dst_row_size_aligned  = gctx->dst_row_size_aligned;   \
+    const int    nc                    = gctx->nc;
 
 // 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 struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -172,10 +138,10 @@ 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 struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const float alpha = ((const float *) (actx->octx->op_params))[2];
-    const float limit = ((const float *) (actx->octx->op_params))[3];
+    const float alpha = ((const float *) (gctx->octx->op_params))[2];
+    const float limit = ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -201,9 +167,9 @@ 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 struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const float limit = ((const float *) (actx->octx->op_params))[3];
+    const float limit = ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -455,7 +421,7 @@ 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 struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -472,7 +438,7 @@ static void geglu_quick_f32(const float * restrict src0,
                             const float * restrict src1,
                             float * restrict dst,
                             const uint32_t num_rows,
-                            const struct htp_act_context * actx) {
+                            const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -489,7 +455,7 @@ static void geglu_erf_f32(const float * restrict src0,
                           const float * restrict src1,
                           float * restrict dst,
                           const uint32_t num_rows,
-                          const struct htp_act_context * actx) {
+                          const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -501,12 +467,124 @@ static void geglu_erf_f32(const float * restrict src0,
     }
 }
 
+static inline void hvx_reglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(float);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    const HVX_Vector v_zero = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
+        vdst[i]      = hvx_vec_mul_f32_f32(x, vsrc1[i]);
+    }
+
+    if (nloe) {
+        HVX_Vector x   = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
+        HVX_Vector res = hvx_vec_mul_f32_f32(x, vsrc1[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
+    }
+}
+
+// reglu(x, g) = relu(x) * g
+static void reglu_f32(const float * restrict src0,
+                      const float * restrict src1,
+                      float * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_reglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f32_vec(HVX_Vector * restrict vdst,
+                                                         const HVX_Vector * restrict vsrc0,
+                                                         const HVX_Vector * restrict vsrc1,
+                                                         uint32_t nvec) {
+    const HVX_Vector v_one     = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
+
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[i], v_one, v_max_exp, v_min_exp);
+        vdst[i]        = hvx_vec_mul_f32_f32(sig, vsrc1[i]);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f32_tail(HVX_Vector * restrict vdst,
+                                                          const HVX_Vector * restrict vsrc0,
+                                                          const HVX_Vector * restrict vsrc1,
+                                                          uint32_t nloe) {
+    const HVX_Vector v_one     = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
+
+    HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[0], v_one, v_max_exp, v_min_exp);
+    HVX_Vector res = hvx_vec_mul_f32_f32(sig, vsrc1[0]);
+    hvx_vec_store_a((void *) vdst, nloe * sizeof(float), res);
+}
+
+static inline void hvx_sigglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(float);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    if (nvec) {
+        hvx_sigglu_f32_vec(vdst, vsrc0, vsrc1, nvec);
+    }
+    if (nloe) {
+        hvx_sigglu_f32_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
+    }
+}
+
+// sigglu(x, g) = sigmoid(x) * g
+static void sigglu_f32(const float * restrict src0,
+                       const float * restrict src1,
+                       float * restrict dst,
+                       const uint32_t num_rows,
+                       const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_sigglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
 // swiglu(x) = x1 * sigmoid(x0)
 static void swiglu_f16(const _Float16 * restrict src0,
                        const _Float16 * restrict src1,
                        _Float16 * restrict dst,
                        const uint32_t num_rows,
-                       const struct htp_act_context * actx) {
+                       const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -524,10 +602,10 @@ static void swiglu_oai_f16(const _Float16 * restrict src0,
                            const _Float16 * restrict src1,
                            _Float16 * restrict dst,
                            const uint32_t num_rows,
-                           const struct htp_act_context * actx) {
+                           const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const _Float16 alpha = (_Float16) ((const float *) (actx->octx->op_params))[2];
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
+    const _Float16 alpha = (_Float16) ((const float *) (gctx->octx->op_params))[2];
+    const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -553,9 +631,9 @@ static void swiglu_clamp_f16(const _Float16 * restrict src0,
                              const _Float16 * restrict src1,
                              _Float16 * restrict dst,
                              const uint32_t                 num_rows,
-                             const struct htp_act_context * actx) {
+                             const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
+    const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
         const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
@@ -574,7 +652,7 @@ static void geglu_f16(const _Float16 * restrict src0,
                       const _Float16 * restrict src1,
                       _Float16 * restrict dst,
                       const uint32_t num_rows,
-                      const struct htp_act_context * actx) {
+                      const struct htp_glu_context * gctx) {
     htp_glu_op_preamble;
 
     for (uint32_t ib = 0; ib < num_rows; ib++) {
@@ -586,56 +664,155 @@ static void geglu_f16(const _Float16 * restrict src0,
     }
 }
 
+static inline void hvx_reglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(_Float16);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    const HVX_Vector v_zero = Q6_V_vzero();
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_Vector x = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
+        vdst[i]      = hvx_vec_mul_f16_f16(x, vsrc1[i]);
+    }
+
+    if (nloe) {
+        HVX_Vector x   = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
+        HVX_Vector res = hvx_vec_mul_f16_f16(x, vsrc1[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(_Float16), res);
+    }
+}
+
+// reglu(x, g) = relu(x) * g
+static void reglu_f16(const _Float16 * restrict src0,
+                      const _Float16 * restrict src1,
+                      _Float16 * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_reglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f16_vec(HVX_Vector * restrict vdst,
+                                                         const HVX_Vector * restrict vsrc0,
+                                                         const HVX_Vector * restrict vsrc1,
+                                                         uint32_t nvec) {
+    for (uint32_t i = 0; i < nvec; i++) {
+        HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[i]);
+        vdst[i]        = hvx_vec_mul_f16_f16(sig, vsrc1[i]);
+    }
+}
+
+static __attribute__((noinline)) void hvx_sigglu_f16_tail(HVX_Vector * restrict vdst,
+                                                          const HVX_Vector * restrict vsrc0,
+                                                          const HVX_Vector * restrict vsrc1,
+                                                          uint32_t nloe) {
+    HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[0]);
+    HVX_Vector res = hvx_vec_mul_f16_f16(sig, vsrc1[0]);
+    hvx_vec_store_a((void *) vdst, nloe * sizeof(_Float16), res);
+}
+
+static inline void hvx_sigglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
+    assert((unsigned long) dst  % 128 == 0);
+    assert((unsigned long) src0 % 128 == 0);
+    assert((unsigned long) src1 % 128 == 0);
+
+    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
+    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
+    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
+
+    const uint32_t epv  = 128 / sizeof(_Float16);
+    const uint32_t nvec = n / epv;
+    const uint32_t nloe = n % epv;
+
+    if (nvec) {
+        hvx_sigglu_f16_vec(vdst, vsrc0, vsrc1, nvec);
+    }
+    if (nloe) {
+        hvx_sigglu_f16_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
+    }
+}
+
+// sigglu(x, g) = sigmoid(x) * g
+static void sigglu_f16(const _Float16 * restrict src0,
+                       const _Float16 * restrict src1,
+                       _Float16 * restrict dst,
+                       const uint32_t num_rows,
+                       const struct htp_glu_context * gctx) {
+    htp_glu_op_preamble;
+
+    for (uint32_t ib = 0; ib < num_rows; ib++) {
+        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
+        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
+        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+
+        hvx_sigglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    }
+}
+
 #define DEFINE_GLU_PER_THREAD(TYPE, SUFFIX)                                                                            \
     static void glu_##SUFFIX##_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_glu_context * gctx = (struct htp_glu_context *) data;                                               \
+        htp_glu_preamble;                                                                                              \
                                                                                                                        \
-        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                          \
+        struct htp_thread_trace * tr = gctx->octx->ctx ? &gctx->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_size = gctx->src0_row_size;                                                                    \
+        size_t src1_row_size = gctx->src1_row_size;                                                                    \
+        size_t dst_row_size  = gctx->dst_row_size;                                                                     \
                                                                                                                        \
-        size_t src0_row_stride = actx->src0_row_stride;                                                                \
-        size_t src1_row_stride = actx->src1_row_stride;                                                                \
+        size_t src0_row_stride = gctx->src0_row_stride;                                                                \
+        size_t src1_row_stride = gctx->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_nrows            = gctx->src0_nrows;                                                       \
+        const uint32_t src0_nrows_per_thread = gctx->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);     \
+        const uint32_t src0_start_row = gctx->row_start + src0_nrows_per_thread * ith;                                 \
+        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, gctx->row_start + src0_nrows);     \
                                                                                                                        \
         /* no work for this thread */                                                                                  \
         if (src0_start_row >= src0_end_row) {                                                                          \
             return;                                                                                                    \
         }                                                                                                              \
                                                                                                                        \
-        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 dma_addr_t data_src0 = gctx->data_src0;                                                                  \
+        const dma_addr_t data_src1 = gctx->data_src1;                                                                  \
+        const dma_addr_t data_dst  = gctx->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;                                               \
+        const size_t src0_row_size_aligned = gctx->src0_row_size_aligned;                                              \
+        const size_t src1_row_size_aligned = gctx->src1_row_size_aligned;                                              \
+        const size_t dst_row_size_aligned  = gctx->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);                  \
+        uint8_t * restrict src0_spad_data = gctx->vtcm_src0 + (ith * gctx->vtcm_src0_size_per_thread);                 \
+        uint8_t * restrict src1_spad_data = gctx->vtcm_src1 + (ith * gctx->vtcm_src1_size_per_thread);                 \
+        uint8_t * restrict dst_spad_data  = gctx->vtcm_dst  + (ith * gctx->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;                                                         \
+        size_t src0_spad_half_size = gctx->src0_spad_half_size;                                                        \
+        size_t src1_spad_half_size = gctx->src1_spad_half_size;                                                        \
+        size_t dst_spad_half_size  = gctx->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;                                                                                                    \
-        }                                                                                                              \
+        const int BLOCK = gctx->block;                                                                                 \
                                                                                                                        \
-        dma_queue * dma_q = actx->octx->ctx->dma[ith];                                                                 \
-        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) actx->compute;                             \
+        dma_queue * dma_q = gctx->octx->ctx->dma[ith];                                                                 \
+        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) gctx->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);                                                 \
@@ -662,7 +839,7 @@ static void geglu_f16(const _Float16 * restrict src0,
             TYPE * src1_spad = (TYPE *) 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);                                                 \
+            compute(src0_spad, src1_spad, dst_spad, block_size, gctx);                                                 \
             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),                             \
@@ -685,16 +862,11 @@ static void geglu_f16(const _Float16 * restrict src0,
 DEFINE_GLU_PER_THREAD(float, f32)
 DEFINE_GLU_PER_THREAD(_Float16, f16)
 
-static int execute_op_activations_f32(struct htp_ops_context * octx) {
+static int execute_op_glu_f32(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;
 
-    if ((dst->ne[0] * SIZEOF_FP32) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     glu_compute_f32_fn_t compute_fn = NULL;
     const char *         op_type    = NULL;
 
@@ -729,11 +901,27 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
             op_type    = "geglu-erf-f32";
             break;
 
+        case HTP_OP_GLU_REGLU:
+            compute_fn = reglu_f32;
+            op_type    = "reglu-f32";
+            break;
+
+        case HTP_OP_GLU_SIGMOID:
+            compute_fn = sigglu_f32;
+            op_type    = "sigglu-f32";
+            break;
+
         default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
+            FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
             return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
     const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;
 
@@ -762,17 +950,16 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     const size_t src0_row_stride = src0->nb[1];
     const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
 
-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
+    struct htp_glu_vtcm_layout L;
+    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);
 
-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
+    const size_t src0_row_size_aligned = L.row_size_aligned;
+    const size_t src1_row_size_aligned = L.row_size_aligned;
+    const size_t dst_row_size_aligned  = L.row_size_aligned;
 
     // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
+    if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
+        FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
              (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
@@ -789,58 +976,53 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
     }
 
     // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
+    struct htp_glu_context gctx;
+    gctx.octx    = octx;
+    gctx.compute = (void *) compute_fn;
+    gctx.op_str  = op_type;
 
-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
+    gctx.src0_row_size = src0_row_size;
+    gctx.src1_row_size = src1_row_size;
+    gctx.dst_row_size  = dst_row_size;
 
-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
+    gctx.src0_row_size_aligned = src0_row_size_aligned;
+    gctx.src1_row_size_aligned = src1_row_size_aligned;
+    gctx.dst_row_size_aligned  = dst_row_size_aligned;
 
-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
+    gctx.src0_row_stride = src0_row_stride;
+    gctx.src1_row_stride = src1_row_stride;
 
     uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+    gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    gctx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
+    gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
+    gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
+    gctx.vtcm_dst_size_per_thread  = (uint32_t) L.dst_bytes_per_thread;
 
-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
+    gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
+    gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
+    gctx.dst_spad_half_size  = gctx.vtcm_dst_size_per_thread  >> 1;
 
-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
+    gctx.block      = L.block;
+    gctx.src0_nrows = nrows;
+    gctx.row_start  = row_start;
 
-    actx.nc = dst->ne[0];
+    gctx.nc = dst->ne[0];
 
     // 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 ||
-                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
-                  octx->op == HTP_OP_GLU_GEGLU ||
-                  octx->op == HTP_OP_GLU_GEGLU_QUICK ||
-                  octx->op == HTP_OP_GLU_GEGLU_ERF)) {
+    if (!src1) {
          const int32_t swapped = octx->op_params[1];
          data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
+         gctx.src1_row_size = gctx.src0_row_size;
 
-         size_t nc_in_bytes = actx.nc * SIZEOF_FP32;
+         size_t nc_in_bytes = gctx.nc * SIZEOF_FP32;
          if (swapped) {
              data_src0 += nc_in_bytes;
          } else {
@@ -848,24 +1030,19 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
          }
     }
 
-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
+    gctx.data_src0 = data_src0;
+    gctx.data_src1 = data_src1;
+    gctx.data_dst  = dst->data;
 
-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &gctx, n_threads);
     return HTP_STATUS_OK;
 }
 
-static int execute_op_activations_f16(struct htp_ops_context * octx) {
+static int execute_op_glu_f16(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;
 
-    if ((dst->ne[0] * SIZEOF_FP16) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     glu_compute_f16_fn_t compute_fn = NULL;
     const char *         op_type    = NULL;
 
@@ -890,11 +1067,27 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
             op_type    = "geglu-f16";
             break;
 
+        case HTP_OP_GLU_REGLU:
+            compute_fn = reglu_f16;
+            op_type    = "reglu-f16";
+            break;
+
+        case HTP_OP_GLU_SIGMOID:
+            compute_fn = sigglu_f16;
+            op_type    = "sigglu-f16";
+            break;
+
         default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
+            FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
             return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
     const size_t dst_row_size = dst->ne[0] * SIZEOF_FP16;
 
@@ -923,17 +1116,16 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
     const size_t src0_row_stride = src0->nb[1];
     const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
 
-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
+    struct htp_glu_vtcm_layout L;
+    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);
 
-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
+    const size_t src0_row_size_aligned = L.row_size_aligned;
+    const size_t src1_row_size_aligned = L.row_size_aligned;
+    const size_t dst_row_size_aligned  = L.row_size_aligned;
 
     // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
+    if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
+        FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
              (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
@@ -950,56 +1142,53 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
     }
 
     // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
+    struct htp_glu_context gctx;
+    gctx.octx    = octx;
+    gctx.compute = (void *) compute_fn;
+    gctx.op_str  = op_type;
 
-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
+    gctx.src0_row_size = src0_row_size;
+    gctx.src1_row_size = src1_row_size;
+    gctx.dst_row_size  = dst_row_size;
 
-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
+    gctx.src0_row_size_aligned = src0_row_size_aligned;
+    gctx.src1_row_size_aligned = src1_row_size_aligned;
+    gctx.dst_row_size_aligned  = dst_row_size_aligned;
 
-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
+    gctx.src0_row_stride = src0_row_stride;
+    gctx.src1_row_stride = src1_row_stride;
 
     uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
+    gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    gctx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
+    gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
+    gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
+    gctx.vtcm_dst_size_per_thread  = (uint32_t) L.dst_bytes_per_thread;
 
-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
+    gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
+    gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
+    gctx.dst_spad_half_size  = gctx.vtcm_dst_size_per_thread  >> 1;
 
-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
+    gctx.block      = L.block;
+    gctx.src0_nrows = nrows;
+    gctx.row_start  = row_start;
 
-    actx.nc = dst->ne[0];
+    gctx.nc = dst->ne[0];
 
     // 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 ||
-                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
-                  octx->op == HTP_OP_GLU_GEGLU)) {
+    if (!src1) {
          const int32_t swapped = octx->op_params[1];
          data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
+         gctx.src1_row_size = gctx.src0_row_size;
 
-         size_t nc_in_bytes = actx.nc * SIZEOF_FP16;
+         size_t nc_in_bytes = gctx.nc * SIZEOF_FP16;
          if (swapped) {
              data_src0 += nc_in_bytes;
          } else {
@@ -1007,22 +1196,22 @@ static int execute_op_activations_f16(struct htp_ops_context * octx) {
          }
     }
 
-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
+    gctx.data_src0 = data_src0;
+    gctx.data_src1 = data_src1;
+    gctx.data_dst  = dst->data;
 
-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &gctx, n_threads);
     return HTP_STATUS_OK;
 }
 
 
-int op_activations(struct htp_ops_context * octx) {
+int op_glu(struct htp_ops_context * octx) {
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-            return execute_op_activations_f32(octx);
+            return execute_op_glu_f32(octx);
 
         case HTP_TYPE_F16:
-            return execute_op_activations_f16(octx);
+            return execute_op_glu_f16(octx);
 
         default:
             return HTP_STATUS_NO_SUPPORT;

ggml/src/ggml-hexagon/htp/glu-ops.h

@@ -0,0 +1,76 @@
+#ifndef HTP_GLU_OPS_H
+#define HTP_GLU_OPS_H
+
+#include "hex-common.h"
+#include "htp-ops.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_glu_kernel_params {
+    uint32_t n_threads;
+    uint32_t vtcm_row_per_thread;
+    uint32_t vtcm_size;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_glu_kernel_params) <= 128, "htp_glu_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_glu_kernel_params) <= 128, "htp_glu_kernel_params is too large for kernel_params blob");
+#endif
+
+static inline bool htp_op_is_glu(uint32_t opcode) {
+    switch (opcode) {
+        case HTP_OP_GLU_SWIGLU:
+        case HTP_OP_GLU_SWIGLU_OAI:
+        case HTP_OP_GLU_SWIGLU_CLAMP:
+        case HTP_OP_GLU_GEGLU:
+        case HTP_OP_GLU_GEGLU_QUICK:
+        case HTP_OP_GLU_GEGLU_ERF:
+        case HTP_OP_GLU_REGLU:
+        case HTP_OP_GLU_SIGMOID:
+            return true;
+        default:
+            return false;
+    }
+}
+
+struct htp_glu_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_src1;
+    size_t off_dst;
+
+    size_t row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t src1_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+
+    uint32_t vtcm_row_per_thread;
+    uint32_t block;
+};
+
+static inline void htp_glu_vtcm_layout_build(
+    struct htp_glu_vtcm_layout * L,
+    size_t                       dst_row_size,
+    uint32_t                     n_threads,
+    uint32_t                     vtcm_row_per_thread
+) {
+    const size_t row_size_aligned = hex_round_up(dst_row_size, 128);
+
+    L->row_size_aligned      = row_size_aligned;
+    L->vtcm_row_per_thread   = vtcm_row_per_thread;
+    L->block                 = vtcm_row_per_thread >> 1;
+
+    L->src0_bytes_per_thread = row_size_aligned * vtcm_row_per_thread;
+    L->src1_bytes_per_thread = row_size_aligned * vtcm_row_per_thread;
+    L->dst_bytes_per_thread  = row_size_aligned * vtcm_row_per_thread;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_src1, L->src1_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+#endif // HTP_GLU_OPS_H

ggml/src/ggml-hexagon/htp/hex-common.h

@@ -17,6 +17,10 @@
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #endif
 
+#ifndef HEX_L2_LINE_SIZE
+#define HEX_L2_LINE_SIZE 128
+#endif
+
 static inline uint32_t hex_ceil_pow2(uint32_t x) {
     if (x <= 1) { return 1; }
     int p = 2;

ggml/src/ggml-hexagon/htp/hex-profile.h

@@ -44,12 +44,14 @@ struct htp_thread_trace {
 };
 
 static inline void htp_trace_event(struct htp_thread_trace * tr, uint16_t id, uint16_t info, uint32_t type) {
-    if (tr->count < tr->max_events) {
-        uint32_t i = tr->count;
+    if (!tr->max_events) {
+        return;
+    }
+    uint32_t i = tr->count++;
+    if (i < tr->max_events) {
         tr->events[i].id     = id;
         tr->events[i].info   = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0);
         tr->events[i].cycles = (uint32_t) hex_get_cycles();
-        tr->count++;
     }
 }
 

ggml/src/ggml-hexagon/htp/htp-ctx.h

@@ -12,8 +12,10 @@
 #include <assert.h>
 #include <dspqueue.h>
 #include <stdatomic.h>
+#include <stddef.h>
 #include <stdint.h>
 #include <stdbool.h>
+#include <string.h>
 
 #ifndef HTP_MAX_NTHREADS
 #define HTP_MAX_NTHREADS 10
@@ -66,14 +68,18 @@ struct htp_ops_context {
     // TODO convert these to an array
     struct htp_spad src0_spad;
     struct htp_spad src1_spad;
-    struct htp_spad src2_spad;
-    struct htp_spad src3_spad;
     struct htp_spad dst_spad;
 
     uint32_t              flags;
     uint32_t              n_threads;
     struct fastdiv_values n_threads_div;
     int                   status;
+
+    // Must stay last: not cleared by htp_ops_context_reset(), ops init the caches they use
+    union {
+        dma_cache_dm dma_cache_dm[HTP_MAX_NTHREADS];
+        dma_cache_fa dma_cache_fa[HTP_MAX_NTHREADS];
+    } __attribute__((aligned(128)));
 };
 
 // Main context for htp DSP backend
@@ -118,6 +124,13 @@ struct htp_context {
     size_t                 footprint;
 };
 
+static inline void htp_ops_context_reset(struct htp_ops_context * octx, struct htp_context * ctx) {
+    memset(octx, 0, offsetof(struct htp_ops_context, dma_cache_dm));
+    octx->ctx           = ctx;
+    octx->n_threads     = ctx->n_threads;
+    octx->n_threads_div = ctx->n_threads_div;
+}
+
 static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
     if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
         return false;
@@ -148,7 +161,7 @@ int op_unary(struct htp_ops_context * octx);
 int op_sum(struct htp_ops_context * octx);
 int op_sum_rows(struct htp_ops_context * octx);
 int op_argmax(struct htp_ops_context * octx);
-int op_activations(struct htp_ops_context * octx);
+int op_glu(struct htp_ops_context * octx);
 int op_softmax(struct htp_ops_context * octx);
 int op_add_id(struct htp_ops_context * octx);
 int op_rope(struct htp_ops_context * octx);

ggml/src/ggml-hexagon/htp/htp-ops.h

@@ -38,11 +38,6 @@ enum htp_data_type {
     HTP_TYPE_INVALID
 };
 
-// Constants for internal types
-#define QK_Q8_0_TILED  128  // 32x32 Q8_0 tiled layout
-
-
-
 // Do not reorder first 4 (used as an index)
 enum htp_op_code {
     HTP_OP_MUL = 0,
@@ -115,6 +110,8 @@ enum htp_op_code {
     HTP_OP_UNARY_GELU_QUICK,
     HTP_OP_IM2COL_3D,
     HTP_OP_HADAMARD,
+    HTP_OP_GLU_REGLU,
+    HTP_OP_GLU_SIGMOID,
 
     HTP_OP_INVALID
 };
@@ -125,8 +122,8 @@ enum htp_op_code {
 #define HTP_OP_MAX_PARAMS  16   // aka GGML_MAX_OP_PARAMS
 #define HTP_OP_MAX_KERN_PARAMS 32
 
-#define HTP_OP_MAX_BUFS    16
-#define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)
+#define HTP_OP_MAX_BUFS    64
+#define HTP_OP_MAX_TENSORS 65534 // must stay under 64K (uint16)
 
 #define HTP_FENCE_TIMEOUT  (1000000000ULL)
 #define HTP_FENCE_SLOT_SIZE 128
@@ -136,9 +133,10 @@ enum htp_op_code {
 #define HTP_MMAP_MAX_VMEM  (2147483648u)
 
 enum htp_tensor_flags {
-    HTP_TENSOR_WEIGHT  = (1U << 0), // Tensor buffer model weight data (not compute)
-    HTP_TENSOR_TILED   = (1U << 1), // Tensor is in tiled format
-    HTP_TENSOR_FENCE   = (1U << 2)  // Tensor is synchronization fence (explicitly managed)
+    HTP_TENSOR_WEIGHT   = (1U << 0), // Tensor buffer model weight data (not compute)
+    HTP_TENSOR_TILED    = (1U << 1), // Tensor is in tiled format
+    HTP_TENSOR_FENCE    = (1U << 2), // Tensor is synchronization fence (explicitly managed)
+    HTP_TENSOR_FUSEABLE = (1U << 3), // Tensor can be fused (host-only)
 };
 
 enum htp_buf_flags {

ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h

@@ -233,7 +233,7 @@ static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * r
 
 static void quantize_row_f32_q8_0_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -244,7 +244,7 @@ static void quantize_row_f32_q8_0_tiled(float * restrict x, uint8_t * restrict y
 
 static void quantize_row_f32_q8_1_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -255,7 +255,7 @@ static void quantize_row_f32_q8_1_tiled(float * restrict x, uint8_t * restrict y
 
 static void quantize_row_f32_q8_1_s16_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -310,7 +310,7 @@ static inline void quantize_block_f16_f16_tiled(const __fp16 * restrict x, uint8
 
 static void quantize_row_f32_f16_tiled(const float * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -321,7 +321,7 @@ static void quantize_row_f32_f16_tiled(const float * restrict x, uint8_t * restr
 
 static void quantize_row_f16_f16_tiled(const __fp16 * restrict x, uint8_t * restrict y, uint32_t k) {
     assert(k % 32 == 0);
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (k + qk - 1) / qk;
 
     for (uint32_t i = 0; i < nb; i++) {
@@ -1963,7 +1963,7 @@ static inline void quantize_f32_q8_0_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -1993,7 +1993,7 @@ static inline void quantize_f32_q8_1_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -2023,7 +2023,7 @@ static inline void quantize_f32_q8_1_s16_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -2053,7 +2053,7 @@ static inline void quantize_f32_f16_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
@@ -2083,7 +2083,7 @@ static inline void quantize_f16_f16_tiled_block_kernel(
     uint32_t c
 ) {
     (void) tmp_data;
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne0 + qk - 1) / qk;
 
     for (uint32_t ib = ib_first; ib < ib_last; ++ib) {

ggml/src/ggml-hexagon/htp/hvx-norm.h

@@ -11,7 +11,8 @@
 static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
                                          uint8_t * restrict dst,
                                          const int num_elems,
-                                         float     epsilon) {
+                                         float     epsilon,
+                                         float     scale) {
 
     const HVX_Vector * restrict v_src = (HVX_Vector *) src;
     HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
@@ -48,6 +49,9 @@ static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
 
     // Scale full vectors
     HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+    if (scale != 0.0f) {
+        scale_v = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(scale_v, hvx_vec_splat_f32(scale)));
+    }
 
     #pragma unroll(4)
     for (int i = 0; i < nvec; i++) {
@@ -72,7 +76,8 @@ static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
                                              const uint8_t * restrict weight,
                                              uint8_t * restrict dst,
                                              const int num_elems,
-                                             float     epsilon) {
+                                             float     epsilon,
+                                             float     scale) {
     const HVX_Vector * restrict v_src    = (const HVX_Vector *) src;
     const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
     HVX_Vector * restrict v_dst          = (HVX_Vector *) dst;
@@ -109,6 +114,9 @@ static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
 
     // Scale and multiply
     HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+    if (scale != 0.0f) {
+        scale_v = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(scale_v, hvx_vec_splat_f32(scale)));
+    }
 
     #pragma unroll(4)
     for (int i = 0; i < nvec; i++) {

ggml/src/ggml-hexagon/htp/hvx-sigmoid.h

@@ -68,6 +68,14 @@ static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
     return Q6_Vsf_equals_Vqf32(res);
 }
 
+static inline HVX_Vector hvx_vec_silu_f32(HVX_Vector x) {
+    HVX_Vector one     = hvx_vec_splat_f32(1.0f);
+    HVX_Vector max_exp = hvx_vec_splat_f32(87.f);
+    HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);
+    HVX_Vector sig     = hvx_vec_fast_sigmoid_f32_guard(x, one, max_exp, min_exp);
+    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(x, sig));
+}
+
 #define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)                                  \
     do {                                                                                      \
         dst_type * restrict vdst = (dst_type *) dst;                                          \

ggml/src/ggml-hexagon/htp/im2col-ops.c

@@ -768,19 +768,13 @@ int op_im2col(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     const struct htp_tensor * dst = octx->dst;
 
     uint32_t row_base = 0;
     uint32_t nrows    = kparams->total_rows;
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = dst->nb[2];
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
         const struct htp_tensor_mdev_range range =
-            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
+            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0,
                                       octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_base = range.start;
         nrows    = range.count;
@@ -790,20 +784,27 @@ int op_im2col(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN((uint32_t) kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    struct htp_im2col_vtcm_layout layout;
+    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads);
+
     struct htp_im2col_context ictx = { 0 };
     uint8_t * const base = octx->ctx->vtcm_base;
 
-    const uint32_t pth = MIN(octx->n_threads, nrows);
+    const uint32_t pth = octx->n_threads;
     ictx.octx                   = octx;
     ictx.pe_row_base            = row_base;
     ictx.pe_nrows               = nrows;
-    ictx.pe_rows_per_thread     = (nrows + pth - 1) / pth;
+    ictx.pe_rows_per_thread     = fastdiv(nrows + pth - 1, &octx->n_threads_div);
     ictx.pe_src_row_bytes       = kparams->src_row_bytes;
     ictx.pe_dst_row_bytes       = kparams->dst_row_bytes;
-    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_src);
-    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_dst);
-    ictx.pe_src_size_per_thread = kparams->src_bytes_per_thread;
-    ictx.pe_dst_size_per_thread = kparams->dst_bytes_per_thread;
+    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_src);
+    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_dst);
+    ictx.pe_src_size_per_thread = layout.src_bytes_per_thread;
+    ictx.pe_dst_size_per_thread = layout.dst_bytes_per_thread;
     ictx.pe_owb                 = kparams->owb;
     ictx.pe_wb                  = kparams->wb;
     ictx.pe_nbpr                = kparams->nbpr;
@@ -821,9 +822,10 @@ int op_im2col(struct htp_ops_context * octx) {
     ictx.patch_stride           = kparams->patch_stride;
 
     for (uint32_t ith = 0; ith < pth; ith++) {
-        uint32_t r = row_base + ictx.pe_rows_per_thread * ith;
-        ictx.thread_ioh[ith] = r % kparams->OH;
-        ictx.thread_in[ith]  = r / kparams->OH;
+        const uint32_t r  = row_base + ictx.pe_rows_per_thread * ith;
+        const uint32_t in = fastdiv(r, &kparams->div_OH);
+        ictx.thread_ioh[ith] = r - in * kparams->OH;
+        ictx.thread_in[ith]  = in;
     }
 
     ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);
@@ -850,19 +852,13 @@ int op_im2col_3d(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     const struct htp_tensor * dst = octx->dst;
 
     uint32_t row_base = 0;
     uint32_t nrows    = kparams->total_rows;
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = dst->nb[2];
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
         const struct htp_tensor_mdev_range range =
-            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
+            htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0,
                                       octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         row_base = range.start;
         nrows    = range.count;
@@ -872,20 +868,27 @@ int op_im2col_3d(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN((uint32_t) kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    struct htp_im2col_vtcm_layout layout;
+    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads);
+
     struct htp_im2col_context ictx = { 0 };
     uint8_t * const base = octx->ctx->vtcm_base;
 
-    const uint32_t pth = MIN(octx->n_threads, nrows);
+    const uint32_t pth = octx->n_threads;
     ictx.octx                   = octx;
     ictx.pe_row_base            = row_base;
     ictx.pe_nrows               = nrows;
-    ictx.pe_rows_per_thread     = (nrows + pth - 1) / pth;
+    ictx.pe_rows_per_thread     = fastdiv(nrows + pth - 1, &octx->n_threads_div);
     ictx.pe_src_row_bytes       = kparams->src_row_bytes;
     ictx.pe_dst_row_bytes       = kparams->dst_row_bytes;
-    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_src);
-    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_dst);
-    ictx.pe_src_size_per_thread = kparams->src_bytes_per_thread;
-    ictx.pe_dst_size_per_thread = kparams->dst_bytes_per_thread;
+    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_src);
+    ictx.pe_vtcm_dst            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_dst);
+    ictx.pe_src_size_per_thread = layout.src_bytes_per_thread;
+    ictx.pe_dst_size_per_thread = layout.dst_bytes_per_thread;
     ictx.pe_owb                 = kparams->owb;
     ictx.pe_wb                  = kparams->wb;
     ictx.pe_nbpr                = kparams->nbpr;
@@ -903,11 +906,12 @@ int op_im2col_3d(struct htp_ops_context * octx) {
     ictx.patch_stride           = kparams->patch_stride;
 
     for (uint32_t ith = 0; ith < pth; ith++) {
-        uint32_t rem = row_base + ictx.pe_rows_per_thread * ith;
-        ictx.thread_ioh[ith] = rem % kparams->OH;
-        rem /= kparams->OH;
-        ictx.thread_iod[ith] = rem % kparams->OD;
-        ictx.thread_in[ith]  = rem / kparams->OD;
+        const uint32_t rem  = row_base + ictx.pe_rows_per_thread * ith;
+        const uint32_t q_oh = fastdiv(rem, &kparams->div_OH);
+        ictx.thread_ioh[ith] = rem - q_oh * kparams->OH;
+        const uint32_t in   = fastdiv(q_oh, &kparams->div_OD);
+        ictx.thread_iod[ith] = q_oh - in * kparams->OD;
+        ictx.thread_in[ith]  = in;
     }
 
     ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);

ggml/src/ggml-hexagon/htp/im2col-ops.h

@@ -2,6 +2,8 @@
 #define HTP_IM2COL_OPS_H
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include <stdint.h>
 
 enum htp_im2col_kernel_type {
@@ -11,38 +13,33 @@ enum htp_im2col_kernel_type {
 };
 
 struct htp_im2col_vtcm_layout {
-    uint32_t off_src;
-    uint32_t off_dst;
+    size_t   total_bytes;
+    size_t   off_src;
+    size_t   off_dst;
     uint32_t src_bytes_per_thread;
     uint32_t dst_bytes_per_thread;
-    uint32_t total_bytes;
 };
 
 static inline void htp_im2col_vtcm_layout_build(
-    struct htp_im2col_vtcm_layout * layout,
+    struct htp_im2col_vtcm_layout * L,
     uint32_t src_row_bytes,
     uint32_t dst_row_bytes,
     uint32_t n_threads) {
 
-    layout->src_bytes_per_thread = 2 * src_row_bytes + 256;
-    layout->dst_bytes_per_thread = 2 * dst_row_bytes;
+    L->src_bytes_per_thread = 2 * src_row_bytes + 256;
+    L->dst_bytes_per_thread = 2 * dst_row_bytes;
 
-    layout->off_src     = 0;
-    layout->off_dst     = layout->off_src + layout->src_bytes_per_thread * n_threads;
-    layout->total_bytes = layout->off_dst + layout->dst_bytes_per_thread * n_threads;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src, (size_t) L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst, (size_t) L->dst_bytes_per_thread * n_threads);
+    L->total_bytes = off;
 }
 
 struct htp_im2col_kernel_params {
     uint8_t  kernel_type;
-    uint8_t  is_2d;
-    uint8_t  is_3d;
     uint8_t  n_threads;
 
     uint32_t vtcm_size;
-    uint32_t off_src;
-    uint32_t off_dst;
-    uint32_t src_bytes_per_thread;
-    uint32_t dst_bytes_per_thread;
     uint32_t src_row_bytes;
     uint32_t dst_row_bytes;
 
@@ -57,6 +54,10 @@ struct htp_im2col_kernel_params {
     uint32_t OD, OH, OW;
     uint32_t patch_stride;
     uint32_t total_rows;
+    uint32_t rows_per_chunk;
+
+    struct fastdiv_values div_OH;
+    struct fastdiv_values div_OD;
 };
 
 #if defined(__cplusplus)

ggml/src/ggml-hexagon/htp/main.c

@@ -955,7 +955,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_GLU_GEGLU:
         case HTP_OP_GLU_GEGLU_QUICK:
         case HTP_OP_GLU_GEGLU_ERF:
-            return op_activations(octx);
+        case HTP_OP_GLU_REGLU:
+        case HTP_OP_GLU_SIGMOID:
+            return op_glu(octx);
 
         case HTP_OP_SOFTMAX:
             return op_softmax(octx);
@@ -1244,8 +1246,6 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs
 
     octx->src0_spad.src = NULL;
     octx->src1_spad.src = NULL;
-    octx->src2_spad.src = NULL;
-    octx->src3_spad.src = NULL;
     octx->dst_spad.src  = NULL;
 
     return octx->status;
@@ -1306,10 +1306,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
     prep_tensors(ctx, bufs, tens, n_tens);
 
     struct htp_ops_context *octx = &ctx->octx;
-    memset(octx, 0, sizeof(*octx));
-    octx->n_threads     = ctx->n_threads;
-    octx->n_threads_div = ctx->n_threads_div;
-    octx->ctx           = ctx;
+    htp_ops_context_reset(octx, ctx);
 
     mdev_group_init(ctx, req);
 

ggml/src/ggml-hexagon/htp/matmul-ops.c

@@ -1816,7 +1816,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
                 return HTP_STATUS_NO_SUPPORT;
             }
 
-            const uint32_t qk = QK_Q8_0_TILED;
+            const uint32_t qk = HTP_MM_ACT_BLOCK_K;
             const uint32_t nb = (ne10 + qk - 1) / qk;
             const uint32_t total_nb = act_nrows * nb;
 
@@ -1886,7 +1886,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
     mmctx->vtcm_src0_stride = src0_row_size_padded;
     mmctx->vtcm_act_stride  = act_row_size;
     if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW) {
-        mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(ne10);
     } else if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) {
         mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), 128);
     } else {
@@ -1904,7 +1904,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
             if (need_quant) {
                 hvx_mm_transfer_act_dma(octx, kparams, act, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, m_start, cur_m_rows);
 
-                const uint32_t qk = QK_Q8_0_TILED;
+                const uint32_t qk = HTP_MM_ACT_BLOCK_K;
                 const uint32_t nb = (ne10 + qk - 1) / qk;
                 const uint32_t total_nb = cur_m_rows * nb;
                 uint32_t quant_tasks;
@@ -3914,7 +3914,7 @@ static int hvx_mm_matmul_id(
     const struct htp_tensor * restrict ids = octx->src[2];
     const size_t src0_row_size = nb01;
 
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (ne10 + qk - 1) / qk;
     const uint32_t total_nb = act_nrows * nb;
 
@@ -3966,7 +3966,7 @@ static int hvx_mm_matmul_id(
 
     mmctx->vtcm_src0_stride    = src0_row_size_padded;
     mmctx->vtcm_act_stride     = act_row_size;
-    mmctx->vtcm_act_raw_stride = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+    mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(ne10);
 
     mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_act_size_per_thread  = L.act_bytes;
@@ -4081,7 +4081,7 @@ static int hvx_mm_matmul_id_nx(
     const struct htp_tensor * restrict ids  = octx->src[n_weights + 1];
     const size_t src0_row_size = src0->nb[1];
 
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (act->ne[0] + qk - 1) / qk;
     const uint32_t total_nb = act_nrows * nb;
 
@@ -4124,7 +4124,7 @@ static int hvx_mm_matmul_id_nx(
 
     mmctx->vtcm_src0_stride    = 0;
     mmctx->vtcm_act_stride     = act_row_size;
-    mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+    mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(act->ne[0]);
 
     mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_act_size_per_thread  = L.act_bytes;
@@ -4474,7 +4474,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    const uint32_t qk = QK_Q8_0_TILED;
+    const uint32_t qk = HTP_MM_ACT_BLOCK_K;
     const uint32_t nb = (act->ne[0] + qk - 1) / qk;
     const uint32_t total_nb = act_nrows * nb;
 
@@ -4518,7 +4518,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) {
 
     mmctx->vtcm_src0_stride    = 0;
     mmctx->vtcm_act_stride     = act_row_size;
-    mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+    mmctx->vtcm_act_raw_stride = htp_mm_act_raw_row_size(act->ne[0]);
 
     mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
     mmctx->vtcm_act_size_per_thread  = L.act_bytes;

ggml/src/ggml-hexagon/htp/matmul-ops.h

@@ -25,6 +25,7 @@ extern "C" {
 #define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q2_0  640
 
 // --- Activation Tiled Block Sizes (including padding) ---
+#define HTP_MM_ACT_BLOCK_K             128
 #define HTP_MM_ACT_TILE_SIZE_Q8_0      1152
 #define HTP_MM_ACT_TILE_SIZE_Q8_1      1280
 #define HTP_MM_ACT_TILE_SIZE_F16       2048
@@ -74,10 +75,7 @@ struct htp_mm_kernel_params {
     int32_t  aligned_tile_size;  // Aligned weight tile size (padded to 128)
     int32_t  act_row_size;       // Row size for activation scratchpad
     int32_t  vtcm_size;          // Total required scratchpad size in VTCM
-    int32_t  vtcm_src0_size;     // src0 scratchpad size in VTCM
-    int32_t  vtcm_act_size;      // activation scratchpad size in VTCM
     int32_t  vtcm_bias_size;     // bias scratchpad size in VTCM (fused only)
-    int32_t  vtcm_dst_size;      // dst scratchpad size in VTCM
 
     // Precomputed division values
     struct fastdiv_values div_ne12_ne1;
@@ -192,19 +190,19 @@ static inline bool htp_mm_weight_has_offset(int weight_type) {
 
 // --- Activation/Row Size Helpers ---
 static inline size_t htp_mm_q8_0_tiled_row_size(uint32_t ne) {
-    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t ne_padded = hex_round_up(ne, HTP_MM_ACT_BLOCK_K);
     const uint32_t nb_32 = ne_padded / 32;
     return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_0;
 }
 
 static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) {
-    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t ne_padded = hex_round_up(ne, HTP_MM_ACT_BLOCK_K);
     const uint32_t nb_32 = ne_padded / 32;
     return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1;
 }
 
 static inline size_t htp_mm_f16_tiled_row_size(uint32_t ne) {
-    const uint32_t ne_padded = ((ne + 127) / 128) * 128;
+    const uint32_t ne_padded = hex_round_up(ne, HTP_MM_ACT_BLOCK_K);
     const uint32_t nb_32 = ne_padded / 32;
     return nb_32 * HTP_MM_ACT_TILE_SIZE_F16;
 }
@@ -216,6 +214,10 @@ static inline size_t htp_mm_tiled_act_row_size(int weight_type, uint32_t ne) {
     return htp_mm_weight_has_offset(weight_type) ? htp_mm_q8_1_tiled_row_size(ne) : htp_mm_q8_0_tiled_row_size(ne);
 }
 
+static inline size_t htp_mm_act_raw_row_size(uint32_t ne) {
+    return hex_round_up(ne, HTP_MM_ACT_BLOCK_K) * sizeof(float);
+}
+
 static inline size_t htp_mm_round_up(size_t n, size_t m) {
     return ((n + m - 1) / m) * m;
 }
@@ -445,7 +447,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
         size_t tiled_act_row_size = htp_mm_tiled_act_row_size(wtype, ne10);
         size_t q_act_sz = hex_round_up(tiled_act_row_size * act_nrows, 128);
-        size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        size_t raw_row_size = htp_mm_act_raw_row_size(ne10);
 
         src0_sz    = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
         act_sz     = q_act_sz;                         // quantized activation buffer
@@ -461,7 +463,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
         act_sz = htp_mm_round_up(act_row_size_tiled * act_nrows, 256);
 
-        size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        size_t raw_row_size = htp_mm_act_raw_row_size(ne10);
 
         src0_sz    = src0_sz_per_thread * n_threads;
         dst_sz     = 0;
@@ -502,7 +504,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
                 size_t dst_slice_per_thread = (dst_nrows > 0 && act_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
                 dst_sz = dst_slice_per_thread * n_threads;
-                size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+                size_t raw_row_size = htp_mm_act_raw_row_size(ne10);
                 act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
                 break;
             }
@@ -582,7 +584,7 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
 
     size_t eff_row_size = row_size;
     if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-        eff_row_size += hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+        eff_row_size += htp_mm_act_raw_row_size(ne10);
     }
 
     uint32_t m_chunk = (uint32_t) (avail_act / eff_row_size);

ggml/src/ggml-hexagon/htp/pad-ops.c

@@ -17,18 +17,28 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
+#include "pad-ops.h"
 
 /* Circular wrap: maps any integer x into [0, n) */
-static inline uint32_t wrap_around(int32_t x, uint32_t n) {
-    return (uint32_t)(((x % (int32_t)n) + (int32_t)n) % (int32_t)n);
+static inline uint32_t wrap_around(int32_t x, uint32_t n, const struct fastdiv_values * div_n) {
+    if (x >= 0) {
+        const uint32_t ux = (uint32_t) x;
+        return ux < n ? ux : fastmodulo(ux, n, div_n);
+    }
+    const uint32_t ux  = (uint32_t) (-x);
+    const uint32_t rem = ux < n ? ux : fastmodulo(ux, n, div_n);
+    return rem ? (n - rem) : 0;
 }
 
 /* Decompose a flat dst row index into (i1, i2, i3) */
 static inline void pad_decompose_row(uint32_t ir, uint32_t ne1, uint32_t ne2,
+                                     const struct htp_pad_kernel_params * kparams,
                                      uint32_t *i1, uint32_t *i2, uint32_t *i3) {
-    *i1 = ir % ne1;
-    *i2 = (ir / ne1) % ne2;
-    *i3 = ir / (ne1 * ne2);
+    *i3 = fastdiv(ir, &kparams->div_ne1_ne2);
+    const uint32_t rem = ir - (*i3) * (ne1 * ne2);
+    *i2 = fastdiv(rem, &kparams->div_ne1);
+    *i1 = rem - (*i2) * ne1;
 }
 
 /* Return non-zero if row (i1,i2,i3) falls in the non-padded interior */
@@ -62,25 +72,28 @@ static inline dma_addr_t pad_src_row_data(const struct htp_tensor * src,
 
 /* 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,
+                                                    const struct htp_pad_kernel_params * kparams,
                                                     uint32_t i1, uint32_t i2, uint32_t i3,
                                                     int32_t lp1, int32_t lp2, int32_t lp3) {
     return (const uint8_t *) 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];
+        + wrap_around((int32_t)i1 - lp1, src->ne[1], &kparams->div_ne01) * src->nb[1]
+        + wrap_around((int32_t)i2 - lp2, src->ne[2], &kparams->div_ne02) * src->nb[2]
+        + wrap_around((int32_t)i3 - lp3, src->ne[3], &kparams->div_ne03) * src->nb[3];
 }
 
 static inline dma_addr_t pad_circ_src_row_data(const struct htp_tensor * src,
+                                               const struct htp_pad_kernel_params * kparams,
                                                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];
+        + wrap_around((int32_t)i1 - lp1, src->ne[1], &kparams->div_ne01) * src->nb[1]
+        + wrap_around((int32_t)i2 - lp2, src->ne[2], &kparams->div_ne02) * src->nb[2]
+        + wrap_around((int32_t)i3 - lp3, src->ne[3], &kparams->div_ne03) * src->nb[3];
 }
 
 struct htp_pad_context {
     struct htp_ops_context * octx;
+    const struct htp_pad_kernel_params * kparams;
 
     int32_t  lp0, rp0;
     int32_t  lp1, rp1;
@@ -152,7 +165,7 @@ static void pad_job_per_thread_hvx(unsigned int nth, unsigned int ith, void * da
 
     for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
         uint32_t i1, i2, i3;
-        pad_decompose_row(dst_row, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(dst_row, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
 
         uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
@@ -219,7 +232,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
             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);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
         const int interior = pad_is_interior(i1, i2, i3,
                                              lp1, rp1, ne1,
                                              lp2, rp2, ne2,
@@ -246,7 +259,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
         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);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
 
         const dma_addr_t dst_data = dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
 
@@ -278,7 +291,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
         const uint32_t next_row = ir + 2;
         if (next_row < row_end) {
             uint32_t ni1, ni2, ni3;
-            pad_decompose_row(next_row, ne1, ne2, &ni1, &ni2, &ni3);
+            pad_decompose_row(next_row, ne1, ne2, pctx->kparams, &ni1, &ni2, &ni3);
             const int next_interior = pad_is_interior(ni1, ni2, ni3,
                                                       lp1, rp1, ne1,
                                                       lp2, rp2, ne2,
@@ -315,10 +328,10 @@ static void pad_job_per_thread_hvx_circular(unsigned int nth, unsigned int ith,
 
     for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
         uint32_t i1, i2, i3;
-        pad_decompose_row(dst_row, ne1, ne2, &i1, &i2, &i3);
+        pad_decompose_row(dst_row, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
 
         uint8_t       * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
-        const uint8_t * src_row = pad_circ_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3);
+        const uint8_t * src_row = pad_circ_src_row_ptr(src, pctx->kparams, i1, i2, i3, lp1, lp2, lp3);
 
         if (nb00 == type_size) {
 
@@ -393,8 +406,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
             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);
-        const dma_addr_t src_data = pad_circ_src_row_data(src, pi1, pi2, pi3, lp1, lp2, lp3);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &pi1, &pi2, &pi3);
+        const dma_addr_t src_data = pad_circ_src_row_data(src, pctx->kparams, 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);
@@ -411,7 +424,7 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
         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);
+        pad_decompose_row(ir, ne1, ne2, pctx->kparams, &i1, &i2, &i3);
         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);
@@ -455,8 +468,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
         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);
-            const dma_addr_t next_src_data = pad_circ_src_row_data(src, nri1, nri2, nri3, lp1, lp2, lp3);
+            pad_decompose_row(next_row, ne1, ne2, pctx->kparams, &nri1, &nri2, &nri3);
+            const dma_addr_t next_src_data = pad_circ_src_row_data(src, pctx->kparams, 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);
@@ -475,15 +488,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
 int op_pad(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
-
-    // Only F32 supported
-    size_t type_size;
-    switch (src0->type) {
-        case HTP_TYPE_F32: type_size = 4; break;
-        default:
-            FARF(ERROR, "pad-hvx: unsupported type %u\n", src0->type);
-            return HTP_STATUS_NO_SUPPORT;
-    }
+    const struct htp_pad_kernel_params * kparams =
+        (const struct htp_pad_kernel_params *) octx->kernel_params;
 
     const int32_t lp0 = octx->op_params[0];
     const int32_t rp0 = octx->op_params[1];
@@ -495,11 +501,10 @@ int op_pad(struct htp_ops_context * octx) {
     const int32_t rp3 = octx->op_params[7];
     const int32_t circular = octx->op_params[8];
 
-    const uint32_t ne0  = dst->ne[0];
-    const uint32_t ne00 = src0->ne[0];
-
-    const uint32_t total_dst_rows = dst->ne[1] * dst->ne[2] * dst->ne[3];
-    const size_t dst_row_size     = (size_t)ne0 * type_size;
+    const uint32_t total_dst_rows = kparams->total_rows;
+    const size_t   type_size      = kparams->elem_size;
+    const size_t   src_row_size   = src0->ne[0] * type_size;
+    const size_t   dst_row_size   = dst->ne[0] * type_size;
 
     uint32_t row_start = 0;
     uint32_t nrows     = total_dst_rows;
@@ -516,60 +521,68 @@ int op_pad(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
 
-    const size_t src_row_size         = (size_t)ne00 * type_size;
-    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
-    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
-
-    // Total VTCM needed: 2 buffers (ping+pong) for src and dst, per thread
-    const size_t vtcm_needed = (size_t)n_threads * 2 * (src_row_size_aligned + dst_row_size_aligned);
-
-    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))) {
+    if (!kparams->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;
-        octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
-        octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
-        octx->src0_spad.data = octx->ctx->vtcm_base;
-        octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size;
+    struct htp_pad_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_pad_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);
+
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
         octx->src0_spad.src  = NULL;
         octx->dst_spad.src   = NULL;
     }
 
     struct htp_pad_context pctx = {
-        .octx             = octx,
+        .octx                 = octx,
+        .kparams              = kparams,
         .lp0 = lp0, .rp0 = rp0,
         .lp1 = lp1, .rp1 = rp1,
         .lp2 = lp2, .rp2 = rp2,
         .lp3 = lp3, .rp3 = rp3,
-        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
-        .total_dst_rows   = nrows,
-        .row_start        = row_start,
-        .type_size        = type_size,
+        .nrows_per_thread     = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_dst_rows       = nrows,
+        .row_start            = row_start,
+        .type_size            = type_size,
         .src_row_size         = src_row_size,
-        .src_row_size_aligned = src_row_size_aligned,
+        .src_row_size_aligned = layout.src_row_size_aligned,
         .dst_row_size         = dst_row_size,
-        .dst_row_size_aligned = dst_row_size_aligned,
+        .dst_row_size_aligned = layout.dst_row_size_aligned,
     };
 
     FARF(HIGH, "pad-hvx%s%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) pads=(%d,%d,%d,%d,%d,%d,%d,%d)\n",
-         circular ? "-circ" : "",
-         use_dma   ? "-dma"  : "",
+         circular         ? "-circ" : "",
+         kparams->use_dma ? "-dma"  : "",
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
          dst->ne[0],  dst->ne[1],  dst->ne[2],  dst->ne[3],
          lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
 
-    if      (circular && use_dma) { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads); }
-    else if (circular)            { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular,     &pctx, n_threads); }
-    else if (use_dma)             { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_dma,          &pctx, n_threads); }
-    else                          { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx,              &pctx, n_threads); }
+    switch (kparams->kernel_type) {
+        case HTP_PAD_KERNEL_HVX_CIRCULAR_DMA:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads);
+            break;
+        case HTP_PAD_KERNEL_HVX_CIRCULAR:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular,     &pctx, n_threads);
+            break;
+        case HTP_PAD_KERNEL_HVX_DMA:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_dma,          &pctx, n_threads);
+            break;
+        default:
+            work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx,              &pctx, n_threads);
+            break;
+    }
 
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/pad-ops.h

@@ -0,0 +1,65 @@
+#ifndef HTP_PAD_OPS_H
+#define HTP_PAD_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+enum htp_pad_kernel_type {
+    HTP_PAD_KERNEL_HVX              = 0,
+    HTP_PAD_KERNEL_HVX_DMA          = 1,
+    HTP_PAD_KERNEL_HVX_CIRCULAR     = 2,
+    HTP_PAD_KERNEL_HVX_CIRCULAR_DMA = 3,
+};
+
+struct htp_pad_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_row_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_pad_vtcm_layout_build(
+    struct htp_pad_vtcm_layout * L,
+    size_t                       src_row_size,
+    size_t                       dst_row_size,
+    uint32_t                     n_threads
+) {
+    L->src_row_size_aligned  = hex_round_up(src_row_size, 128);
+    L->dst_row_size_aligned  = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread = 2 * L->src_row_size_aligned;
+    L->dst_bytes_per_thread  = 2 * L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_pad_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t use_dma;
+    uint32_t elem_size;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne1_ne2;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne03;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_pad_kernel_params) <= 128, "htp_pad_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_pad_kernel_params) <= 128, "htp_pad_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_PAD_OPS_H

ggml/src/ggml-hexagon/htp/pool-ops.c

@@ -247,6 +247,7 @@ typedef void (*pool_plane_fn_t)(const float * src, float * dst, const struct htp
 struct pool_2d_context {
     struct htp_ops_context * octx;
     const struct htp_pool_2d_kernel_params * kparams;
+    struct htp_pool_vtcm_layout layout;
     pool_plane_fn_t pool_plane;
     uint32_t n_threads;
     uint32_t plane_start;
@@ -257,6 +258,7 @@ struct pool_2d_context {
 static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
     struct pool_2d_context * ctx = (struct pool_2d_context *) data;
     const struct htp_pool_2d_kernel_params * p = ctx->kparams;
+    const struct htp_pool_vtcm_layout * L = &ctx->layout;
     const struct htp_tensor * src0 = ctx->octx->src[0];
     const struct htp_tensor * dst = ctx->octx->dst;
     pool_plane_fn_t pool_plane = ctx->pool_plane;
@@ -271,15 +273,14 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
     struct htp_thread_trace * tr = &ctx->octx->ctx->trace[ith];
     dma_queue * dma_queue = ctx->octx->ctx->dma[ith];
 
-    const uint32_t src_spad_half = p->src_plane_bytes_aligned;
-    const uint32_t dst_spad_half = p->dst_plane_bytes_aligned;
-    const uint32_t src_bytes_per_thread = 2 * src_spad_half;
-    const uint32_t dst_bytes_per_thread = 2 * dst_spad_half;
-    const size_t off_dst = (size_t) ctx->n_threads * src_bytes_per_thread;
+    const uint32_t src_plane_bytes = L->src_plane_bytes;
+    const uint32_t dst_plane_bytes = L->dst_plane_bytes;
+    const uint32_t src_spad_half   = (uint32_t) L->src_spad_half_size;
+    const uint32_t dst_spad_half   = (uint32_t) L->dst_spad_half_size;
 
     uint8_t * vtcm_base = (uint8_t *) ctx->octx->ctx->vtcm_base;
-    uint8_t * src_spad  = vtcm_base + ith * src_bytes_per_thread;
-    uint8_t * dst_spad  = vtcm_base + off_dst + ith * dst_bytes_per_thread;
+    uint8_t * src_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_src) + ith * L->src_bytes_per_thread;
+    uint8_t * dst_spad  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_dst) + ith * L->dst_bytes_per_thread;
 
     float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) };
     float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) };
@@ -290,14 +291,14 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
     for (uint32_t i = 0; i < total && i < 2; ++i) {
         dma_queue_push(dma_queue,
                        dma_make_data(dst->data, dstb2[i]),
-                       p->dst_plane_bytes, dst_spad_half,
-                       p->dst_plane_bytes, 0);
+                       dst_plane_bytes, dst_spad_half,
+                       dst_plane_bytes, 0);
 
-        const dma_addr_t src_addr = src0->data + (first + i) * p->src_plane_bytes;
+        const dma_addr_t src_addr = src0->data + (first + i) * src_plane_bytes;
         dma_queue_push(dma_queue,
                        dma_make_data(srcb2[i], src_addr),
-                       src_spad_half, p->src_plane_bytes,
-                       p->src_plane_bytes, 1);
+                       src_spad_half, src_plane_bytes,
+                       src_plane_bytes, 1);
     }
 
     for (uint32_t i = 0; i < total; ++i) {
@@ -313,18 +314,18 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
         pool_plane(srcb, dstb, p);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) plane);
 
-        const dma_addr_t dst_addr = dst->data + plane * p->dst_plane_bytes;
+        const dma_addr_t dst_addr = dst->data + plane * dst_plane_bytes;
         dma_queue_push(dma_queue,
                        dma_make_data(dst_addr, dstb),
-                       p->dst_plane_bytes, dst_spad_half,
-                       p->dst_plane_bytes, 1);
+                       dst_plane_bytes, dst_spad_half,
+                       dst_plane_bytes, 1);
 
         if (i + 2 < total) {
-            const dma_addr_t next_src_addr = src0->data + (plane + 2) * p->src_plane_bytes;
+            const dma_addr_t next_src_addr = src0->data + (plane + 2) * src_plane_bytes;
             dma_queue_push(dma_queue,
                            dma_make_data(srcb, next_src_addr),
-                           src_spad_half, p->src_plane_bytes,
-                           p->src_plane_bytes, 1);
+                           src_spad_half, src_plane_bytes,
+                           src_plane_bytes, 1);
         }
     }
 
@@ -337,22 +338,15 @@ static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
 }
 
 int op_pool_2d(struct htp_ops_context * octx) {
-    const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst = octx->dst;
     const struct htp_pool_2d_kernel_params * p =
         (const struct htp_pool_2d_kernel_params *) octx->kernel_params;
 
-    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32 ||
-        (p->pool_op != HTP_POOL_AVG && p->pool_op != HTP_POOL_MAX)) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     uint32_t plane_start = 0;
     uint32_t plane_count = p->planes;
     if (octx->ctx->mdev.count > 1) {
-        const uint32_t planes_per_chunk = (p->dst_plane_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(p->dst_plane_bytes, HEX_L2_LINE_SIZE)) : 1;
         const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
-            plane_count, htp_tensor_mdev_data_aligned(dst) ? planes_per_chunk : 0,
+            plane_count, htp_tensor_mdev_data_aligned(dst) ? p->planes_per_chunk : 0,
             octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
         plane_start = range.start;
         plane_count = range.count;
@@ -369,16 +363,12 @@ int op_pool_2d(struct htp_ops_context * octx) {
     const uint32_t planes_per_thread = fastdiv(plane_count + n_threads - 1, &octx->n_threads_div);
 
     pool_plane_fn_t pool_plane;
-    if (p->global_path) {
-        pool_plane = pool_plane_global;
-    } else if (p->block_path) {
-        pool_plane = pool_plane_block;
-    } else if (p->narrow_path) {
-        pool_plane = pool_plane_hvx_narrow;
-    } else if (p->fast_path) {
-        pool_plane = pool_plane_hvx;
-    } else {
-        pool_plane = pool_plane_general;
+    switch (p->kernel_type) {
+        case HTP_POOL_KERNEL_GLOBAL:     pool_plane = pool_plane_global;     break;
+        case HTP_POOL_KERNEL_BLOCK:      pool_plane = pool_plane_block;      break;
+        case HTP_POOL_KERNEL_HVX_NARROW: pool_plane = pool_plane_hvx_narrow; break;
+        case HTP_POOL_KERNEL_HVX:        pool_plane = pool_plane_hvx;        break;
+        default:                         pool_plane = pool_plane_general;    break;
     }
 
     struct pool_2d_context ctx = {
@@ -390,6 +380,7 @@ int op_pool_2d(struct htp_ops_context * octx) {
         .plane_count = plane_count,
         .planes_per_thread = planes_per_thread,
     };
+    htp_pool_vtcm_layout_build(&ctx.layout, p->src_x, p->src_y, p->dst_x, p->dst_y, n_threads);
     work_queue_run(octx->ctx->work_queue, pool_2d_thread, &ctx, n_threads);
     return HTP_STATUS_OK;
 }

ggml/src/ggml-hexagon/htp/pool-ops.h

@@ -7,8 +7,18 @@
 #include <string.h>
 
 #include "hex-common.h"
+#include "htp-vtcm.h"
+
+enum htp_pool_kernel_type {
+    HTP_POOL_KERNEL_GENERAL = 0,
+    HTP_POOL_KERNEL_HVX,
+    HTP_POOL_KERNEL_HVX_NARROW,
+    HTP_POOL_KERNEL_GLOBAL,
+    HTP_POOL_KERNEL_BLOCK,
+};
 
 struct htp_pool_2d_kernel_params {
+    uint8_t  kernel_type;
     uint32_t src_x;
     uint32_t src_y;
     uint32_t dst_x;
@@ -19,17 +29,10 @@ struct htp_pool_2d_kernel_params {
     uint32_t stride_y;
     int32_t  pad_x;
     int32_t  pad_y;
-    uint32_t src_plane_bytes;
-    uint32_t dst_plane_bytes;
-    uint32_t src_plane_bytes_aligned;
-    uint32_t dst_plane_bytes_aligned;
     uint32_t n_threads;
     uint32_t planes;
+    uint32_t planes_per_chunk;
     uint32_t pool_op;
-    uint32_t fast_path;
-    uint32_t narrow_path;
-    uint32_t global_path;
-    uint32_t block_path;
     uint32_t avg_divide_count;
     uint32_t ox_lo;
     uint32_t ox_hi;
@@ -46,12 +49,36 @@ struct htp_pool_vtcm_layout {
     size_t   total_bytes;
     size_t   off_src;
     size_t   off_dst;
+    uint32_t src_plane_bytes;
+    uint32_t dst_plane_bytes;
     size_t   src_bytes_per_thread;
     size_t   dst_bytes_per_thread;
     size_t   src_spad_half_size;
     size_t   dst_spad_half_size;
 };
 
+static inline void htp_pool_vtcm_layout_build(
+    struct htp_pool_vtcm_layout * L,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    uint32_t n_threads
+) {
+    // Full-plane double buffering (256 bytes guard space for vector loads)
+    L->src_plane_bytes      = src_x * src_y * sizeof(float);
+    L->dst_plane_bytes      = dst_x * dst_y * sizeof(float);
+    L->src_spad_half_size   = hex_round_up(L->src_plane_bytes + 256, 128);
+    L->dst_spad_half_size   = hex_round_up(L->dst_plane_bytes, 128);
+    L->src_bytes_per_thread = 2 * L->src_spad_half_size;
+    L->dst_bytes_per_thread = 2 * L->dst_spad_half_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src, L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes_per_thread * n_threads);
+    L->total_bytes = off;
+}
+
 static inline bool htp_pool_solve_layout(
     struct htp_pool_vtcm_layout * layout,
     uint32_t src_x,
@@ -61,25 +88,8 @@ static inline bool htp_pool_solve_layout(
     uint32_t n_threads,
     size_t   vtcm_budget
 ) {
-    // Full-plane double buffering (256 bytes guard space for vector loads)
-    const size_t src_plane_bytes   = (size_t) src_x * src_y * sizeof(float);
-    const size_t dst_plane_bytes   = (size_t) dst_x * dst_y * sizeof(float);
-    const size_t src_plane_aligned = hex_round_up((uint32_t) src_plane_bytes + 256, 128);
-    const size_t dst_plane_aligned = hex_round_up((uint32_t) dst_plane_bytes, 128);
-
-    const size_t spad_per_thread = 2 * (src_plane_aligned + dst_plane_aligned);
-    if (spad_per_thread * n_threads > vtcm_budget) {
-        return false;
-    }
-
-    layout->src_spad_half_size   = src_plane_aligned;
-    layout->dst_spad_half_size   = dst_plane_aligned;
-    layout->src_bytes_per_thread = 2 * src_plane_aligned;
-    layout->dst_bytes_per_thread = 2 * dst_plane_aligned;
-    layout->off_src              = 0;
-    layout->off_dst              = layout->src_bytes_per_thread * n_threads;
-    layout->total_bytes          = layout->off_dst + layout->dst_bytes_per_thread * n_threads;
-    return true;
+    htp_pool_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, n_threads);
+    return layout->total_bytes <= vtcm_budget;
 }
 
 #endif // HTP_POOL_OPS_H

ggml/src/ggml-hexagon/htp/repeat-ops.c

@@ -16,14 +16,11 @@
 #include "hex-profile.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "repeat-ops.h"
 
 struct htp_repeat_context {
-    struct htp_ops_context * octx;
-
-    uint32_t nr0;
-    uint32_t nr1;
-    uint32_t nr2;
-    uint32_t nr3;
+    struct htp_ops_context *               octx;
+    const struct htp_repeat_kernel_params * kparams;
 
     uint32_t nrows_per_thread;
     uint32_t total_dst_rows;  // ne1 * ne2 * ne3
@@ -35,6 +32,7 @@ struct htp_repeat_context {
 static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * data) {
     const struct htp_repeat_context * rctx = (const struct htp_repeat_context *) data;
     struct htp_ops_context * octx = rctx->octx;
+    const struct htp_repeat_kernel_params * kparams = rctx->kparams;
     const struct htp_tensor * src = octx->src[0];
     const struct htp_tensor * dst = octx->dst;
 
@@ -58,29 +56,30 @@ static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * dat
     const uint32_t nb2 = dst->nb[2];
     const uint32_t nb3 = dst->nb[3];
 
-    const uint32_t nr0 = rctx->nr0;
-    const uint32_t nr1 = rctx->nr1;
-    const uint32_t nr2 = rctx->nr2;
-    const uint32_t nr3 = rctx->nr3;
+    const uint32_t nr0 = kparams->nr0;
 
     const size_t row_bytes = ne00 * rctx->type_size;
 
     const uint32_t row_start = rctx->row_start + rctx->nrows_per_thread * ith;
     const uint32_t row_end   = MIN(row_start + rctx->nrows_per_thread, rctx->row_start + rctx->total_dst_rows);
+    if (row_start >= row_end) {
+        return;
+    }
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
 
     for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
         // Decompose flat dst row index into (i1, i2, i3)
-        const uint32_t i1 = dst_row % ne1;
-        const uint32_t i2 = (dst_row / ne1) % ne2;
-        const uint32_t i3 = dst_row / (ne1 * ne2);
+        const uint32_t dst_row_div_ne1 = fastdiv(dst_row, &kparams->div_ne1);
+        const uint32_t i1              = dst_row - dst_row_div_ne1 * ne1;
+        const uint32_t i3              = fastdiv(dst_row, &kparams->div_ne12);
+        const uint32_t i2              = dst_row_div_ne1 - i3 * ne2;
 
         // Map to source indices (tiling)
-        const uint32_t k1 = i1 % ne01;
-        const uint32_t k2 = i2 % ne02;
-        const uint32_t k3 = i3 % ne03;
+        const uint32_t k1 = fastmodulo(i1, ne01, &kparams->div_ne01);
+        const uint32_t k2 = fastmodulo(i2, ne02, &kparams->div_ne02);
+        const uint32_t k3 = fastmodulo(i3, ne03, &kparams->div_ne03);
 
         const uint8_t * src_row = (const uint8_t *) src->data + k1 * nb01 + k2 * nb02 + k3 * nb03;
         uint8_t * dst_base      = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
@@ -104,15 +103,6 @@ int op_repeat(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * dst  = octx->dst;
 
-    // Validate that dst dims are multiples of src dims
-    if (dst->ne[0] % src0->ne[0] != 0 ||
-        dst->ne[1] % src0->ne[1] != 0 ||
-        dst->ne[2] % src0->ne[2] != 0 ||
-        dst->ne[3] % src0->ne[3] != 0) {
-        FARF(ERROR, "repeat: dst dims must be multiples of src dims\n");
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
     size_t type_size;
     switch (src0->type) {
         case HTP_TYPE_F32: type_size = 4; break;
@@ -126,6 +116,13 @@ int op_repeat(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_repeat_kernel_params * kparams =
+        (const struct htp_repeat_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t total_dst_rows  = dst->ne[1] * dst->ne[2] * dst->ne[3];
     const size_t dst_row_size = dst->ne[0] * type_size;
 
@@ -148,10 +145,7 @@ int op_repeat(struct htp_ops_context * octx) {
 
     struct htp_repeat_context rctx = {
         .octx             = octx,
-        .nr0              = dst->ne[0] / src0->ne[0],
-        .nr1              = dst->ne[1] / src0->ne[1],
-        .nr2              = dst->ne[2] / src0->ne[2],
-        .nr3              = dst->ne[3] / src0->ne[3],
+        .kparams          = kparams,
         .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
         .total_dst_rows   = nrows,
         .row_start        = row_start,
@@ -161,7 +155,7 @@ int op_repeat(struct htp_ops_context * octx) {
     FARF(HIGH, "repeat: (%ux%ux%ux%u) -> (%ux%ux%ux%u) nr=(%u,%u,%u,%u)\n",
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
          dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         rctx.nr0, rctx.nr1, rctx.nr2, rctx.nr3);
+         kparams->nr0, kparams->nr1, kparams->nr2, kparams->nr3);
 
     work_queue_run(octx->ctx->work_queue, repeat_job_per_thread, &rctx, n_threads);
 

ggml/src/ggml-hexagon/htp/repeat-ops.h

@@ -0,0 +1,27 @@
+#ifndef HTP_REPEAT_OPS_H
+#define HTP_REPEAT_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+struct htp_repeat_kernel_params {
+    uint32_t n_threads;
+    uint32_t nr0;
+    uint32_t nr1;
+    uint32_t nr2;
+    uint32_t nr3;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne03;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_repeat_kernel_params) <= 128, "htp_repeat_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_repeat_kernel_params) <= 128, "htp_repeat_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_REPEAT_OPS_H

ggml/src/ggml-hexagon/htp/roll-ops.c

@@ -17,16 +17,15 @@
 #include "hex-profile.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "roll-ops.h"
 
 struct htp_roll_context {
     struct htp_ops_context * octx;
+    const struct htp_roll_kernel_params * kparams;
 
     uint32_t row_start;
     uint32_t nrows;
     uint32_t nrows_per_thread;
-
-    struct fastdiv_values div_ne1;
-    struct fastdiv_values div_ne2_ne1;
 };
 
 static inline uint32_t htp_roll_wrap(int32_t i, uint32_t ne) {
@@ -39,30 +38,32 @@ static inline uint32_t htp_roll_wrap(int32_t i, uint32_t ne) {
     return (uint32_t) i;
 }
 
-#define htp_roll_preamble                             \
-    const struct htp_tensor * src0 = octx->src[0];    \
-    const struct htp_tensor * dst  = octx->dst;       \
-                                                      \
-    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 nb01 = src0->nb[1];                \
-    const uint32_t nb02 = src0->nb[2];                \
-    const uint32_t nb03 = src0->nb[3];                \
-                                                      \
-    const uint32_t nb1 = dst->nb[1];                  \
-    const uint32_t nb2 = dst->nb[2];                  \
-    const uint32_t nb3 = dst->nb[3];                  \
-                                                      \
-    const int32_t s0 = octx->op_params[0];            \
-    const int32_t s1 = octx->op_params[1];            \
-    const int32_t s2 = octx->op_params[2];            \
-    const int32_t s3 = octx->op_params[3];            \
-                                                      \
-    const uint32_t i0_src0 = htp_roll_wrap(-s0, ne0); \
-    const uint32_t n0      = ne0 - i0_src0;
+#define htp_roll_preamble                                            \
+    const struct htp_tensor * src0 = octx->src[0];                   \
+    const struct htp_tensor * dst  = octx->dst;                      \
+    const struct htp_roll_kernel_params * kparams =                  \
+        (const struct htp_roll_kernel_params *) octx->kernel_params; \
+                                                                     \
+    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 nb01 = src0->nb[1];                               \
+    const uint32_t nb02 = src0->nb[2];                               \
+    const uint32_t nb03 = src0->nb[3];                               \
+                                                                     \
+    const uint32_t nb1 = dst->nb[1];                                 \
+    const uint32_t nb2 = dst->nb[2];                                 \
+    const uint32_t nb3 = dst->nb[3];                                 \
+                                                                     \
+    const int32_t s0 = octx->op_params[0];                           \
+    const int32_t s1 = octx->op_params[1];                           \
+    const int32_t s2 = octx->op_params[2];                           \
+    const int32_t s3 = octx->op_params[3];                           \
+                                                                     \
+    const uint32_t i0_src0 = kparams->i0_src0;                       \
+    const uint32_t n0      = kparams->n0;
 
 #define htp_roll_dma_preamble dma_queue * q = octx->ctx->dma[0];
 
@@ -91,10 +92,10 @@ static inline void roll_dma_push_rows(dma_queue *               q,
                                       uint32_t                  src_row,
                                       uint32_t                  nrows,
                                       uint32_t                  row_size,
-                                      uint32_t                  i0_src0) {
+                                      uint32_t                  i0_src0,
+                                      uint32_t                  n0) {
     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 + (size_t) i0_src0 * sizeof(float),
                   row_size, row_size, n0 * sizeof(float), nrows);
@@ -125,7 +126,7 @@ static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
     const uint32_t row_size = ne0 * sizeof(float);
 
     if (s1 == 0 && s2 == 0 && s3 == 0) {
-        roll_dma_push_rows(q, dst, src0, 0, 0, ne1 * ne2 * ne3, row_size, i0_src0);
+        roll_dma_push_rows(q, dst, src0, 0, 0, ne1 * ne2 * ne3, row_size, i0_src0, n0);
         dma_queue_flush(q);
         return HTP_STATUS_OK;
     }
@@ -139,9 +140,9 @@ static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
             const uint32_t n2_first = ne2 - i2_src0;
 
             roll_dma_push_rows(q, dst, src0, dst_row0, src_row0, n2_first * ne1,
-                               row_size, i0_src0);
+                               row_size, i0_src0, n0);
             roll_dma_push_rows(q, dst, src0, dst_row0 + n2_first * ne1, i03 * ne2 * ne1,
-                               i2_src0 * ne1, row_size, i0_src0);
+                               i2_src0 * ne1, row_size, i0_src0, n0);
         }
 
         dma_queue_flush(q);
@@ -159,9 +160,9 @@ static int roll_dma_f32_contiguous(struct htp_ops_context * octx) {
             const uint32_t src_row0 = (i03 * ne2 + i02) * ne1;
 
             roll_dma_push_rows(q, dst, src0, dst_row0, src_row0 + i1_src0,
-                               n1_first, row_size, i0_src0);
+                               n1_first, row_size, i0_src0, n0);
             roll_dma_push_rows(q, dst, src0, dst_row0 + n1_first, src_row0,
-                               i1_src0, row_size, i0_src0);
+                               i1_src0, row_size, i0_src0, n0);
         }
     }
 
@@ -214,9 +215,9 @@ static void roll_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
 
     for (uint32_t row = row_start; row < row_end; row++) {
-        const uint32_t i3  = fastdiv(row, &rctx->div_ne2_ne1);
+        const uint32_t i3  = fastdiv(row, &kparams->div_ne2_ne1);
         const uint32_t rem = row - i3 * ne2 * ne1;
-        const uint32_t i2  = fastdiv(rem, &rctx->div_ne1);
+        const uint32_t i2  = fastdiv(rem, &kparams->div_ne1);
         const uint32_t i1  = rem - i2 * ne1;
 
         const uint32_t i01 = htp_roll_wrap((int32_t) i1 - s1, ne1);
@@ -245,23 +246,14 @@ static void roll_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int execute_op_roll_f32(struct htp_ops_context * octx) {
     htp_roll_preamble;
 
-    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        FARF(ERROR, "roll: unsupported type %u -> %u\n", src0->type, dst->type);
-        return HTP_STATUS_NO_SUPPORT;
+    if (kparams->kernel_type == HTP_ROLL_KERNEL_DMA_CONTIGUOUS) {
+        return roll_dma_f32_contiguous(octx);
+    }
+    if (kparams->kernel_type == HTP_ROLL_KERNEL_DMA_STRIDED) {
+        return roll_dma_f32_strided(octx);
     }
 
-    if (src0->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
-        FARF(ERROR, "roll: unsupported nb0 %u -> %u\n", src0->nb[0], dst->nb[0]);
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    if (src0->ne[0] != ne0 || src0->ne[1] != ne1 ||
-        src0->ne[2] != ne2 || src0->ne[3] != ne3) {
-        FARF(ERROR, "roll: shape mismatch\n");
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
-    const uint32_t total_rows = ne1 * ne2 * ne3;
+    const uint32_t total_rows = kparams->total_rows;
     const size_t dst_row_size = ne0 * sizeof(float);
 
     uint32_t row_start = 0;
@@ -279,25 +271,21 @@ int execute_op_roll_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    if (octx->ctx->mdev.count <= 1) {
-        if (htp_tensor_is_contiguous(src0, sizeof(float)) && htp_tensor_is_contiguous(dst, sizeof(float))) {
-            return roll_dma_f32_contiguous(octx);
-        }
-        return roll_dma_f32_strided(octx);
-    }
-
     if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t n_threads = octx->n_threads;
     struct htp_roll_context rctx = {
         .octx             = octx,
+        .kparams          = kparams,
         .row_start        = row_start,
         .nrows            = nrows,
         .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
-        .div_ne1          = init_fastdiv_values(dst->ne[1]),
-        .div_ne2_ne1      = init_fastdiv_values(dst->ne[2] * dst->ne[1]),
     };
 
     work_queue_run(octx->ctx->work_queue, roll_thread_f32, &rctx, n_threads);

ggml/src/ggml-hexagon/htp/roll-ops.h

@@ -0,0 +1,30 @@
+#ifndef HTP_ROLL_OPS_H
+#define HTP_ROLL_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+enum htp_roll_kernel_type {
+    HTP_ROLL_KERNEL_DMA_CONTIGUOUS = 1,
+    HTP_ROLL_KERNEL_DMA_STRIDED    = 2,
+    HTP_ROLL_KERNEL_HVX            = 3,
+};
+
+struct htp_roll_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t i0_src0;
+    uint32_t n0;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne2_ne1;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_roll_kernel_params) <= 128, "htp_roll_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_roll_kernel_params) <= 128, "htp_roll_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_ROLL_OPS_H

ggml/src/ggml-hexagon/htp/rope-ops.c

@@ -700,15 +700,6 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src2 = octx->src[2];
     const struct htp_tensor * dst  = octx->dst;
 
-    switch (octx->op) {
-        case HTP_OP_ROPE:
-            break;
-
-        default:
-            FARF(ERROR, "Unsupported Op %u\n", octx->op);
-            return HTP_STATUS_NO_SUPPORT;
-    }
-
     const struct htp_rope_kernel_params * kparams = (const struct htp_rope_kernel_params *) octx->kernel_params;
     if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
         return HTP_STATUS_INVAL_PARAMS;
@@ -746,20 +737,23 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     const size_t dst_row_size    = dst->ne[0] * sizeof(float);
     const size_t dst_row_stride  = dst->nb[1];
 
+    struct htp_rope_vtcm_layout layout;
+    htp_rope_vtcm_layout_build(&layout, src0->ne[0], kparams->n_threads, src2 ? src2->ne[0] : 0);
+
     struct htp_rope_context rctx;
     memset(&rctx, 0, sizeof(struct htp_rope_context));
 
     rctx.octx                  = octx;
-    rctx.vtcm_base             = (uint8_t *) octx->ctx->vtcm_base;
-    rctx.spad_per_thread       = kparams->spad_per_thread;
-    rctx.theta_cache_offset    = kparams->theta_cache_offset;
+    rctx.vtcm_base             = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_spad);
+    rctx.spad_per_thread       = layout.bytes_per_thread;
+    rctx.theta_cache_offset    = layout.theta_cache_size_aligned;
 
     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);
+        float * vtcm_freq_factors = VTCM_LAYOUT_PTR(float, octx->ctx->vtcm_base, layout.off_freq_factors);
         dma_queue_push(dma_q, dma_make_data(vtcm_freq_factors, src2->data),
-                       kparams->freq_factors_size, 0, ff_size, 1);
+                       layout.freq_factors_size_aligned, 0, ff_size, 1);
         dma_queue_pop(dma_q);
         rctx.freq_factors = vtcm_freq_factors;
     }
@@ -791,7 +785,7 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
     rctx.src0_row_stride       = src0_row_stride;
     rctx.dst_row_size          = dst_row_size;
     rctx.dst_row_stride        = dst_row_stride;
-    rctx.src0_row_size_aligned = kparams->src0_row_size_aligned;
+    rctx.src0_row_size_aligned = layout.src0_row_size_aligned;
 
     rctx.src0_nrows            = nrows;
     rctx.nrows                 = nrows;

ggml/src/ggml-hexagon/htp/rope-ops.h

@@ -3,6 +3,7 @@
 
 #include "hex-common.h"
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 
 #define HTP_ROPE_SPAD_BLOCK  8
 #define HTP_ROPE_SPAD_NSLOTS 4
@@ -10,14 +11,7 @@
 
 struct htp_rope_kernel_params {
     uint32_t n_threads;
-    uint32_t src0_nrows;
-    uint32_t src0_nrows_per_thread;
     uint32_t vtcm_size;
-    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;
@@ -31,6 +25,8 @@ _Static_assert(sizeof(struct htp_rope_kernel_params) <= 128, "htp_rope_kernel_pa
 
 struct htp_rope_vtcm_layout {
     size_t total_bytes;
+    size_t off_spad;
+    size_t off_freq_factors;
     size_t bytes_per_thread;
     size_t theta_cache_size_aligned;
     size_t src0_row_size_aligned;
@@ -38,7 +34,7 @@ struct htp_rope_vtcm_layout {
 };
 
 static inline void htp_rope_vtcm_layout_build(
-    struct htp_rope_vtcm_layout * layout,
+    struct htp_rope_vtcm_layout * L,
     uint32_t ne00,
     uint32_t n_threads,
     uint32_t n_freq_factors
@@ -48,11 +44,15 @@ static inline void htp_rope_vtcm_layout_build(
     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->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;
+    L->src0_row_size_aligned     = src0_row_size_aligned;
+    L->theta_cache_size_aligned  = theta_cache_size_aligned;
+    L->freq_factors_size_aligned = freq_factors_size_aligned;
+    L->bytes_per_thread          = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_spad, L->bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_freq_factors, freq_factors_size_aligned, n_freq_factors > 0);
+    L->total_bytes = off;
 }
 
 static inline uint8_t * rope_spad_slot(uint8_t * base, uint32_t slot, size_t row_size_aligned) {

ggml/src/ggml-hexagon/htp/set-rows-ops.c

@@ -79,8 +79,10 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
     const uint32_t ir1 = MIN(ir0 + dr, srctx->task_start + srctx->tasks);                                        \
     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;    \
+    uint8_t * vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, srctx->vtcm_base, vtcm_layout->off_src0) +                    \
+                          ith * vtcm_layout->src0_bytes_per_thread;                                              \
+    uint8_t * vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, srctx->vtcm_base, vtcm_layout->off_dst) +                     \
+                          ith * vtcm_layout->dst_bytes_per_thread;                                               \
     const uint32_t src0_row_size = ne00 * sizeof(float);                                                         \
     const uint32_t dst_row_size  = htp_tensor_get_row_size(octx->dst->type, ne00);                               \
     const uint32_t nrows_per_thread = ir1 - ir0;                                                                 \
@@ -188,14 +190,6 @@ int op_set_rows(struct htp_ops_context * octx) {
     const struct htp_set_rows_kernel_params * kparams = (const struct htp_set_rows_kernel_params *)octx->kernel_params;
     set_rows_preamble;
 
-    if (octx->src[0]->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_Q8_0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     if (htp_tensor_is_extended(octx->src[1])) {
         return HTP_STATUS_NO_SUPPORT;
     }

ggml/src/ggml-hexagon/htp/set-rows-ops.h

@@ -2,18 +2,13 @@
 #define HTP_SET_ROWS_OPS_H
 
 #include "hex-fastdiv.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
 
 struct htp_set_rows_kernel_params {
-    int32_t  n_threads;
-    int32_t  total_tasks;
-    int32_t  tasks_per_thread;
-    int32_t  vtcm_size;
-
-    // Fastdiv helpers
-    struct fastdiv_values div_ne11;
-    struct fastdiv_values div_ne12;
-    struct fastdiv_values div_tasks_per_thread;
-    struct fastdiv_values div_ne02;
+    uint32_t n_threads;
+    uint32_t total_tasks;
+    uint32_t vtcm_size;
 };
 
 struct htp_set_rows_vtcm_layout {
@@ -29,40 +24,27 @@ struct htp_set_rows_vtcm_layout {
 };
 
 static inline void htp_set_rows_vtcm_layout_build(
-    struct htp_set_rows_vtcm_layout * vtcm_layout,
+    struct htp_set_rows_vtcm_layout * L,
     int dst_type,
     uint32_t ne00,
     uint32_t n_threads) {
 
-    size_t src0_row_size = ne00 * 4;
-    size_t dst_row_size = 0;
-    switch (dst_type) {
-        case 0: // HTP_TYPE_F32
-            dst_row_size = ne00 * 4;
-            break;
-        case 1: // HTP_TYPE_F16
-            dst_row_size = ne00 * 2;
-            break;
-        case 8: // HTP_TYPE_Q8_0
-            dst_row_size = (ne00 / 32) * 34;
-            break;
-        default:
-            dst_row_size = 0;
-            break;
-    }
+    size_t src0_row_size = ne00 * sizeof(float);
+    size_t dst_row_size  = htp_tensor_get_row_size(dst_type, ne00);
 
     size_t src0_row_size_aligned = (src0_row_size + 255) & ~255;
     size_t dst_row_size_aligned  = (dst_row_size + 255) & ~255;
 
-    vtcm_layout->src0_spad_half_size = src0_row_size_aligned;
-    vtcm_layout->dst_spad_half_size  = dst_row_size_aligned;
+    L->src0_spad_half_size = src0_row_size_aligned;
+    L->dst_spad_half_size  = dst_row_size_aligned;
 
-    vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
-    vtcm_layout->dst_bytes_per_thread  = dst_row_size_aligned * 2;
+    L->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    L->dst_bytes_per_thread  = dst_row_size_aligned * 2;
 
-    vtcm_layout->off_src0 = 0;
-    vtcm_layout->off_dst  = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads;
-    vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread  * n_threads);
+    L->total_bytes = off;
 }
 
 #if defined(__cplusplus)

ggml/src/ggml-hexagon/htp/softmax-ops.c

@@ -496,7 +496,7 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
-    if (kparams->kernel_id >= HTP_SOFTMAX_KERNEL_COUNT) {
+    if (kparams->kernel_type >= HTP_SOFTMAX_KERNEL_COUNT) {
         return HTP_STATUS_INVAL_PARAMS;
     }
 
@@ -532,34 +532,38 @@ 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;
 
-    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;
+    struct htp_softmax_vtcm_layout layout;
+    htp_softmax_vtcm_layout_build(&layout,
+                                  src0->ne[0],
+                                  kparams->use_src1 ? octx->src[1]->ne[0] : 0,
+                                  kparams->use_src1 != 0,
+                                  kparams->use_f16 != 0,
+                                  kparams->n_threads);
 
     struct htp_softmax_context smctx = {
         .octx                  = octx,
         .kparams               = kparams,
-        .compute               = (void *) softmax_kernels[kparams->kernel_id],
+        .compute               = (void *) softmax_kernels[kparams->kernel_type],
 
         .data_src0             = src0->data,
         .data_src1             = kparams->use_src1 ? octx->src[1]->data : 0,
         .data_dst              = dst->data,
 
-        .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),
+        .vtcm_src0             = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_src0),
+        .vtcm_dst              = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_dst),
+        .vtcm_src1             = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, vtcm_base, layout.off_src1, kparams->use_src1),
 
-        .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 = (uint32_t) layout.src0_bytes_per_thread,
+        .vtcm_src1_size_per_thread = (uint32_t) layout.src1_bytes_per_thread,
+        .vtcm_dst_size_per_thread  = (uint32_t) layout.dst_bytes_per_thread,
 
-        .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_spad_half_size   = (uint32_t) layout.src0_spad_half_size,
+        .src1_spad_half_size   = (uint32_t) layout.src1_spad_half_size,
+        .dst_spad_half_size    = (uint32_t) layout.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,
+        .src0_row_size_aligned = (uint32_t) layout.src0_spad_half_size,
+        .src1_row_size_aligned = (uint32_t) layout.src1_spad_half_size,
+        .dst_row_size_aligned  = (uint32_t) layout.dst_spad_half_size,
 
         .use_f16               = kparams->use_f16 != 0,
 

ggml/src/ggml-hexagon/htp/softmax-ops.h

@@ -7,8 +7,9 @@
 #include <math.h>
 #include "hex-fastdiv.h"
 #include "hex-common.h"
+#include "htp-vtcm.h"
 
-enum htp_softmax_kernel_id {
+enum htp_softmax_kernel_type {
     HTP_SOFTMAX_KERNEL_NOMASK = 0,
     HTP_SOFTMAX_KERNEL_MASK_F32,
     HTP_SOFTMAX_KERNEL_MASK_F16,
@@ -16,28 +17,14 @@ enum htp_softmax_kernel_id {
 };
 
 struct htp_softmax_kernel_params {
+    uint8_t  kernel_type;
     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;
@@ -72,7 +59,7 @@ struct htp_softmax_vtcm_layout {
 };
 
 static inline void htp_softmax_vtcm_layout_build(
-    struct htp_softmax_vtcm_layout * layout,
+    struct htp_softmax_vtcm_layout * L,
     uint32_t ne00,
     uint32_t ne10,
     bool use_src1,
@@ -87,20 +74,21 @@ static inline void htp_softmax_vtcm_layout_build(
     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;
+    L->src0_spad_half_size = src0_row_size_aligned;
+    L->dst_spad_half_size  = dst_row_size_aligned;
+    L->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;
+    L->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    L->dst_bytes_per_thread  = dst_row_size_aligned  * 2;
+    L->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;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread  * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_src1, L->src1_bytes_per_thread * n_threads, use_src1);
 
-    layout->total_bytes = layout->off_src1 + layout->src1_bytes_per_thread * n_threads;
+    L->total_bytes = off;
 }
 
 #endif // HTP_SOFTMAX_OPS_H

ggml/src/ggml-hexagon/htp/solve-tri-ops.c

@@ -13,14 +13,16 @@
 #include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
+#include "solve-tri-ops.h"
 
 struct htp_solve_tri_context {
-    struct htp_ops_context * octx;
-    uint32_t                 jobs_per_thread;
-    uint32_t                 total_jobs;
-    uint32_t                 job_start;
-    uint32_t                 k_chunks;
-    uint32_t                 col_block;
+    struct htp_ops_context *               octx;
+    const struct htp_solve_tri_kernel_params * kparams;
+    uint32_t                               jobs_per_thread;
+    uint32_t                               total_jobs;
+    uint32_t                               job_start;
+    uint32_t                               k_chunks;
+    uint32_t                               col_block;
 };
 
 static inline void solve_tri_row_scalar(const float * A_row,
@@ -80,6 +82,7 @@ static inline void solve_tri_row_hvx(const float * A_row,
 static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     struct htp_solve_tri_context * sctx = (struct htp_solve_tri_context *) data;
     struct htp_ops_context *       octx = sctx->octx;
+    const struct htp_solve_tri_kernel_params * kparams = sctx->kparams;
 
     const struct htp_tensor * src0 = octx->src[0];  // A
     const struct htp_tensor * src1 = octx->src[1];  // B
@@ -91,7 +94,7 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
     const uint32_t ne02 = src0->ne[2];
 
     const uint32_t col_block = VLEN_FP32;
-    const uint32_t k_full    = (k / col_block) * col_block;
+    const uint32_t k_full    = k & ~(col_block - 1);
 
     const uint32_t start_batch = sctx->job_start + sctx->jobs_per_thread * ith;
     const uint32_t end_batch   = MIN(start_batch + sctx->jobs_per_thread, sctx->job_start + sctx->total_jobs);
@@ -100,7 +103,7 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_batch);
 
     for (uint32_t batch = start_batch; batch < end_batch; ++batch) {
-        const uint32_t i03 = batch / ne02;
+        const uint32_t i03 = fastdiv(batch, &kparams->div_ne02);
         const uint32_t i02 = batch - i03 * ne02;
 
         const float * A_batch =
@@ -141,6 +144,7 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
 static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void * data) {
     struct htp_solve_tri_context * sctx = (struct htp_solve_tri_context *) data;
     struct htp_ops_context *       octx = sctx->octx;
+    const struct htp_solve_tri_kernel_params * kparams = sctx->kparams;
 
     const struct htp_tensor * src0 = octx->src[0];  // A
     const struct htp_tensor * src1 = octx->src[1];  // B
@@ -158,10 +162,10 @@ static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_job);
 
     for (uint32_t job = start_job; job < end_job; ++job) {
-        const uint32_t batch = job / sctx->k_chunks;
+        const uint32_t batch = fastdiv(job, &kparams->div_k_chunks);
         const uint32_t chunk = job - batch * sctx->k_chunks;
 
-        const uint32_t i03 = batch / ne02;
+        const uint32_t i03 = fastdiv(batch, &kparams->div_ne02);
         const uint32_t i02 = batch - i03 * ne02;
 
         const float * A_batch =
@@ -198,36 +202,17 @@ int op_solve_tri(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];  // A
     const struct htp_tensor * src1 = octx->src[1];  // B
     const struct htp_tensor * dst  = octx->dst;     // X
-
-    if (src0->type != HTP_TYPE_F32 || src1->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    // left=true, lower=true, uni=false only
-    if (src0->ne[0] != src0->ne[1]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-    if (src0->ne[1] != src1->ne[1]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-    if (src0->ne[2] != src1->ne[2] || src0->ne[3] != src1->ne[3]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-    if (dst->ne[0] != src1->ne[0] || dst->ne[1] != src1->ne[1] || dst->ne[2] != src1->ne[2] ||
-        dst->ne[3] != src1->ne[3]) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
+    const struct htp_solve_tri_kernel_params * kparams =
+        (const struct htp_solve_tri_kernel_params *) octx->kernel_params;
 
     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];
-
-    const uint32_t col_block     = VLEN_FP32;
-    const uint32_t k_chunks      = (k + col_block - 1) / col_block;
-    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const bool     batched       = total_batches >= (uint32_t) octx->n_threads;
+    const uint32_t col_block     = kparams->col_block;
+    const uint32_t k_chunks      = kparams->k_chunks;
+    const uint32_t total_batches = kparams->total_batches;
+    const bool     batched       = (kparams->batched != 0);
 
     FARF(HIGH, "solve-tri: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : batched %d\n",
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
@@ -239,9 +224,7 @@ int op_solve_tri(struct htp_ops_context * octx) {
         uint32_t njobs     = total_batches;
 
         if (octx->ctx->mdev.count > 1) {
-            const uint32_t batch_size = dst->nb[2];
-            const uint32_t batches_per_chunk = (batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(batch_size, HEX_L2_LINE_SIZE)) : 1;
-            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, htp_tensor_mdev_data_aligned(dst) ? batches_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_batches, htp_tensor_mdev_data_aligned(dst) ? kparams->batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
             job_start = range.start;
             njobs     = range.count;
         }
@@ -250,11 +233,15 @@ int op_solve_tri(struct htp_ops_context * octx) {
             return HTP_STATUS_OK;
         }
 
-        // Batch-level parallelism
+        if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, njobs))) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+
         const uint32_t n_threads = octx->n_threads;
 
         struct htp_solve_tri_context sctx = {
             .octx            = octx,
+            .kparams         = kparams,
             .jobs_per_thread = fastdiv(njobs + n_threads - 1, &octx->n_threads_div),
             .total_jobs      = njobs,
             .job_start       = job_start,
@@ -264,8 +251,7 @@ int op_solve_tri(struct htp_ops_context * octx) {
 
         work_queue_run(octx->ctx->work_queue, solve_tri_batch_thread_f32, &sctx, n_threads);
     } else {
-        // Chunk-level parallelism
-        const uint32_t total_jobs = total_batches * k_chunks;
+        const uint32_t total_jobs = kparams->total_jobs;
 
         uint32_t job_start = 0;
         uint32_t njobs     = total_jobs;
@@ -281,10 +267,15 @@ int op_solve_tri(struct htp_ops_context * octx) {
             return HTP_STATUS_OK;
         }
 
+        if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, njobs))) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+
         const uint32_t n_threads = octx->n_threads;
 
         struct htp_solve_tri_context sctx = {
             .octx            = octx,
+            .kparams         = kparams,
             .jobs_per_thread = fastdiv(njobs + n_threads - 1, &octx->n_threads_div),
             .total_jobs      = njobs,
             .job_start       = job_start,

ggml/src/ggml-hexagon/htp/solve-tri-ops.h

@@ -0,0 +1,26 @@
+#ifndef HTP_SOLVE_TRI_OPS_H
+#define HTP_SOLVE_TRI_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+struct htp_solve_tri_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_batches;
+    uint32_t total_jobs;
+    uint32_t k_chunks;
+    uint32_t col_block;
+    uint32_t batched;
+    uint32_t batches_per_chunk;
+
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_k_chunks;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_solve_tri_kernel_params) <= 128, "htp_solve_tri_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_solve_tri_kernel_params) <= 128, "htp_solve_tri_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_SOLVE_TRI_OPS_H

ggml/src/ggml-hexagon/htp/ssm-conv.c

@@ -135,6 +135,13 @@ static inline void hvx_ssm_conv_unpack_to_T(const float * raw, float * T, uint32
     }
 }
 
+static __attribute__((noinline)) void hvx_ssm_conv_silu_vtcm(float * restrict buf, uint32_t n_vecs) {
+    HVX_Vector * restrict vbuf = (HVX_Vector *) buf;
+    for (uint32_t i = 0; i < n_vecs; ++i) {
+        vbuf[i] = hvx_vec_silu_f32(vbuf[i]);
+    }
+}
+
 // Decode dot product specialization for d_conv == 4: multiply in the raw channel-major layout,
 // then deinterleave the products so each vector holds one tap of 32 channels, and sum.
 // Keeps both operands in DMA layout - no transpose, no scratch.
@@ -189,9 +196,10 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
 
     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 d_conv     = kparams->d_conv;
+    const uint32_t d_inner    = kparams->d_inner;
+    const uint32_t n_s        = kparams->n_s;
+    const bool     apply_silu = kparams->silu != 0;
 
     const uint32_t dr  = scctx->nrows_per_thread;
     const uint32_t ir0 = scctx->row_start + dr * ith;
@@ -232,7 +240,7 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
     // raw_dot keeps both operands in the DMA layout, so src1 needs no prep pass
     const bool raw_dot = (d_conv == 4) && (d_inner_per_thread % VLEN_FP32 == 0);
 
-    const uint32_t n_tiles   = (d_inner_per_thread + d_inner_tile - 1) / d_inner_tile;
+    const uint32_t n_tiles   = fastdiv(d_inner_per_thread + d_inner_tile - 1, &kparams->div_d_inner_tile);
     const uint32_t n_chunks  = n_s * n_tiles;
     const size_t   row_bytes = d_conv * sizeof(float);
 
@@ -295,21 +303,18 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
             hvx_ssm_conv_unpack_to_T(src0_tile_raw[c & 1], src0_T, tile_n, tile_stride, d_conv);
 
             for (uint32_t cb = 0; cb < tile_n; cb += VLEN_FP32) {
-                const uint32_t cb_n = MIN(VLEN_FP32, tile_n - 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 * tile_stride + cb);
                     HVX_Vector w = *(const HVX_Vector *)(src1_T + j * d_inner_stride + tile_off + 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 *)(out + cb) = y;
-                } else {
-                    hvx_vec_store_u(out + cb, cb_n * sizeof(float), y);
-                }
+                *(HVX_Vector *)(out + cb) = Q6_Vsf_equals_Vqf32(acc);
             }
         }
+        if (apply_silu) {
+            hvx_ssm_conv_silu_vtcm(out, hex_round_up(tile_n, VLEN_FP32) / VLEN_FP32);
+        }
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i3);
 
         const dma_addr_t dst_ddr        = dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float);
@@ -352,11 +357,12 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
 
     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 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 bool     apply_silu = kparams->silu != 0;
 
     const uint32_t dr  = scctx->nrows_per_thread;
     const uint32_t ir0 = scctx->row_start + dr * ith;
@@ -397,7 +403,7 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
-    const uint32_t n_tiles   = (d_inner_per_thread + d_inner_tile - 1) / d_inner_tile;
+    const uint32_t n_tiles   = fastdiv(d_inner_per_thread + d_inner_tile - 1, &kparams->div_d_inner_tile);
     const uint32_t n_chunks  = n_s * n_tiles;
     const size_t   row_bytes = ncs * sizeof(float);
 
@@ -513,10 +519,14 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
         }
 
+        if (apply_silu) {
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
+            hvx_ssm_conv_silu_vtcm(out, n_t * (dst_tile_stride / C_TILE));
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
+        }
+
         // Writeback dst_tile VTCM -> DDR via 2D DMA
-        dma_queue_push(dma_q,
-                       dma_make_data(dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float),
-                                     (uint8_t *) out),
+        dma_queue_push(dma_q, dma_make_data(dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float), (uint8_t *) out),
                        dst_stride_token_bytes, dst_tile_stride * sizeof(float), tile_n * sizeof(float), n_t);
 
         if (c + 2 < n_chunks) {
@@ -547,13 +557,8 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) {
     const struct htp_tensor * src1 = octx->src[1];
     const struct htp_tensor * dst  = octx->dst;
 
-    if (src0->type != HTP_TYPE_F32 || src1->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     const struct htp_ssm_conv_kernel_params * kparams = (const struct htp_ssm_conv_kernel_params *) octx->kernel_params;
 
-
     if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
         return HTP_STATUS_INVAL_PARAMS;
     }
@@ -582,26 +587,30 @@ int op_ssm_conv_f32(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_threads          = octx->n_threads;
+    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);
 
-    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;
+    struct htp_ssm_conv_vtcm_layout L;
+    htp_ssm_conv_vtcm_layout_build(&L, d_inner_per_thread, kparams->d_conv, kparams->n_t, kparams->d_inner_tile, n_threads);
 
-    octx->src0_spad.size = kparams->vtcm_src0_size;
-    octx->src1_spad.size = kparams->vtcm_src1_size;
-    octx->dst_spad.size  = kparams->vtcm_dst_size;
+    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
 
-    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.size_per_thread = L.src0_bytes_per_thread;
+    octx->src1_spad.size_per_thread = L.src1_bytes_per_thread;
+    octx->dst_spad.size_per_thread  = L.dst_bytes_per_thread;
+
+    octx->src0_spad.size = L.src0_bytes;
+    octx->src1_spad.size = L.src1_bytes;
+    octx->dst_spad.size  = L.dst_bytes;
+
+    octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    octx->src1_spad.data = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
     octx->src0_spad.src  = NULL;
     octx->src1_spad.src  = NULL;
     octx->dst_spad.src   = NULL;
 
-    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);
-
     struct htp_ssm_conv_context scctx = {
         .octx             = octx,
         .kparams          = kparams,

ggml/src/ggml-hexagon/htp/ssm-conv.h

@@ -3,8 +3,10 @@
 
 #include <stdint.h>
 
+#include "hex-common.h"
 #include "hex-fastdiv.h"
 #include "htp-ops.h"
+#include "htp-vtcm.h"
 
 struct htp_ssm_conv_kernel_params {
     uint32_t n_threads;
@@ -14,16 +16,62 @@ struct htp_ssm_conv_kernel_params {
     uint32_t n_s;
     uint32_t d_inner_tile;
 
-    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;
+    uint32_t silu;
+
+    struct fastdiv_values div_d_inner_tile;
 };
 
+struct htp_ssm_conv_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_bytes;
+    size_t src1_bytes;
+    size_t dst_bytes;
+};
+
+static inline void htp_ssm_conv_vtcm_layout_build(
+    struct htp_ssm_conv_vtcm_layout * L,
+    uint32_t d_inner_per_thread,
+    uint32_t d_conv,
+    uint32_t n_t,
+    uint32_t d_inner_tile,
+    uint32_t n_threads
+) {
+    const size_t src1_raw_bytes = hex_round_up((size_t) d_inner_per_thread * d_conv * sizeof(float), 128) + 128;
+    const size_t src1_T_bytes   = hex_round_up((size_t) d_conv * d_inner_per_thread * sizeof(float), 128);
+
+    L->src1_bytes_per_thread = src1_raw_bytes + src1_T_bytes;
+
+    if (n_t == 1) {
+        const size_t src0_tile_raw_bytes = hex_round_up((size_t) d_inner_tile * d_conv * sizeof(float), 128);
+        const size_t src0_T_bytes        = hex_round_up((size_t) d_conv * d_inner_tile * sizeof(float), 128);
+        L->src0_bytes_per_thread         = 2 * src0_tile_raw_bytes + src0_T_bytes;
+        L->dst_bytes_per_thread          = 2 * hex_round_up((size_t) d_inner_tile * sizeof(float), 128);
+    } else {
+        const uint32_t ncs          = d_conv - 1 + n_t;
+        const size_t   src0_block_T = hex_round_up((size_t) ncs * 32 * sizeof(float), 128);
+        const size_t   src0_tile_raw = hex_round_up((size_t) d_inner_tile * ncs * sizeof(float), 128);
+        L->src0_bytes_per_thread    = 2 * src0_tile_raw + src0_block_T;
+        L->dst_bytes_per_thread     = 2 * hex_round_up((size_t) d_inner_tile * n_t * sizeof(float), 128);
+    }
+
+    L->src0_bytes = L->src0_bytes_per_thread * n_threads;
+    L->src1_bytes = L->src1_bytes_per_thread * n_threads;
+    L->dst_bytes  = L->dst_bytes_per_thread * n_threads;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_src1, L->src1_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes);
+    L->total_bytes = off;
+}
+
 #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

ggml/src/ggml-hexagon/htp/sum-rows-ops.c

@@ -18,6 +18,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "sum-rows-ops.h"
 
 #define sum_rows_preamble                         \
     const struct htp_tensor *src0 = octx->src[0]; \
@@ -52,7 +53,6 @@ struct sum_rows_context {
     size_t          dst_stride;
     uint32_t        rows_per_thread;
     uint32_t        total_rows;
-    bool            opt_path;
 };
 
 static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data) {
@@ -71,7 +71,6 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
     const size_t   src_stride = smctx->src_stride;
     const size_t   dst_stride = smctx->dst_stride;
     const uint32_t ne00       = smctx->ne00;
-    const bool     opt_path   = smctx->opt_path;
 
     const float * restrict src_th = (const float *) (smctx->src_data + (start_row * src_stride));
     float       * restrict dst_th = (float *)       (smctx->dst_data + (start_row * dst_stride));
@@ -90,11 +89,7 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
             hex_l2fetch((const uint8_t *) src_local + src_stride, src_stride, src_stride, 1);
         }
 
-        if (opt_path) {
-            *dst_local = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
-        } else {
-            *dst_local = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
-        }
+        *dst_local = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
     }
 
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
@@ -111,6 +106,13 @@ int op_sum_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_sum_rows_kernel_params * kparams =
+        (const struct htp_sum_rows_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t src0_nrows      = ne01 * ne02 * ne03;
     const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
 
@@ -129,14 +131,9 @@ int op_sum_rows(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads = octx->n_threads;
+    const uint32_t n_threads       = octx->n_threads;
     const uint32_t rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
-    bool opt_path = false;
-    if ((0 == hex_is_aligned((void *) src0->data, VLEN)) && !(nb01 & (VLEN - 1))) {
-        opt_path = true;
-    }
-
     struct sum_rows_context smctx = {
         .octx            = octx,
         .src_data        = (const uint8_t *) src0->data + row_start * nb01,
@@ -146,7 +143,6 @@ int op_sum_rows(struct htp_ops_context * octx) {
         .dst_stride      = nb1,
         .rows_per_thread = rows_per_thread,
         .total_rows      = nrows,
-        .opt_path        = opt_path,
     };
 
     work_queue_run(octx->ctx->work_queue, sum_rows_thread_f32, &smctx, n_threads);
@@ -197,6 +193,13 @@ int op_sum(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_sum_rows_kernel_params * kparams =
+        (const struct htp_sum_rows_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
         return HTP_STATUS_OK;
     }
@@ -207,9 +210,8 @@ int op_sum(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads = (total_elems >= 1024) ? MIN(octx->n_threads, HTP_MAX_NTHREADS) : 1;
-    const uint32_t raw_chunk = (total_elems + n_threads - 1) / n_threads;
-    const uint32_t elems_per_thread = hex_round_up(raw_chunk, 32);
+    const uint32_t n_threads        = octx->n_threads;
+    const uint32_t elems_per_thread = kparams->elems_per_thread;
 
     struct sum_context sctx = {
         .octx             = octx,
@@ -324,6 +326,13 @@ int op_argmax(struct htp_ops_context * octx) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
+    const struct htp_sum_rows_kernel_params * kparams =
+        (const struct htp_sum_rows_kernel_params *) octx->kernel_params;
+
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
     const uint32_t ne00 = src0->ne[0];
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
 
@@ -341,9 +350,8 @@ int op_argmax(struct htp_ops_context * octx) {
             return HTP_STATUS_OK;
         }
 
-        const uint32_t n_threads = (ne00 >= 1024) ? MIN(octx->n_threads, HTP_MAX_NTHREADS) : 1;
-        const uint32_t raw_chunk = (ne00 + n_threads - 1) / n_threads;
-        const uint32_t elems_per_thread = hex_round_up(raw_chunk, 32);
+        const uint32_t n_threads        = octx->n_threads;
+        const uint32_t elems_per_thread = kparams->elems_per_thread;
 
         struct argmax_context actx = {
             .octx             = octx,
@@ -389,8 +397,8 @@ int op_argmax(struct htp_ops_context * octx) {
         return HTP_STATUS_OK;
     }
 
-    const uint32_t n_threads = MIN(octx->n_threads, nrows);
-    const uint32_t rows_per_thread = (nrows + n_threads - 1) / n_threads;
+    const uint32_t n_threads       = octx->n_threads;
+    const uint32_t rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
 
     struct argmax_context actx = {
         .octx             = octx,

ggml/src/ggml-hexagon/htp/sum-rows-ops.h

@@ -0,0 +1,17 @@
+#ifndef HTP_SUM_ROWS_OPS_H
+#define HTP_SUM_ROWS_OPS_H
+
+#include <stdint.h>
+
+struct htp_sum_rows_kernel_params {
+    uint32_t n_threads;
+    uint32_t elems_per_thread;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_sum_rows_kernel_params) <= 128, "htp_sum_rows_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_sum_rows_kernel_params) <= 128, "htp_sum_rows_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_SUM_ROWS_OPS_H

ggml/src/ggml-hexagon/htp/unary-ops.c

@@ -109,18 +109,17 @@ static inline size_t unary_row_offset(uint32_t ir,
     return i1 * nb1 + i2 * nb2 + i3 * nb3;
 }
 
-// Safe DMA block size from row `ir`: clamp to the tighter dim-1 slice
-// boundary of src and dst so the nb1 stride stays valid for all rows.
+// Safe DMA block size from row `ir`: clamp to the dim-1 slice boundary
+// when non-contiguous so the nb1 stride stays valid for all rows.
 static inline uint32_t unary_block_size(uint32_t ir,
                                         uint32_t end_row,
                                         uint32_t block,
-                                        bool src_contig,
-                                        bool dst_contig,
+                                        bool contig,
                                         uint32_t ne1,
                                         const struct fastdiv_values * div_ne1) {
     uint32_t limit = MIN(block, end_row - ir);
 
-    if (!src_contig || !dst_contig) {
+    if (!contig) {
         const uint32_t slice_end = (fastdiv(ir, div_ne1) + 1) * ne1;
         limit = MIN(limit, slice_end - ir);
     }
@@ -213,13 +212,15 @@ static void rms_norm_f32(const void * restrict src,
                          const struct htp_unary_context * uctx) {
     htp_unary_op_preamble;
     float epsilon = 0.f;
-    memcpy(&epsilon, op_params, sizeof(float));
+    float scale   = 0.f;
+    memcpy(&epsilon, &op_params[0], sizeof(float));
+    memcpy(&scale,   &op_params[1], sizeof(float));
 
     for (uint32_t ir = 0; ir < num_rows; ir++) {
         const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
         uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
 
-        hvx_fast_rms_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, ne0, epsilon);
+        hvx_fast_rms_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, ne0, epsilon, scale);
     }
 }
 
@@ -230,14 +231,16 @@ static void rms_norm_mul_f32(const void * restrict src,
                              const struct htp_unary_context * uctx) {
     htp_unary_op_preamble;
     float epsilon = 0.f;
-    memcpy(&epsilon, op_params, sizeof(float));
+    float scale   = 0.f;
+    memcpy(&epsilon, &op_params[0], sizeof(float));
+    memcpy(&scale,   &op_params[1], sizeof(float));
 
     for (uint32_t ir = 0; ir < num_rows; ir++) {
         const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
         const uint8_t * restrict w_local   = (const uint8_t *)weight + (uctx->broadcast_weight ? 0 : ir * uctx->src1_row_size_aligned);
         uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
 
-        hvx_fast_rms_norm_mul_f32(src_local, w_local, dst_local, ne0, epsilon);
+        hvx_fast_rms_norm_mul_f32(src_local, w_local, dst_local, ne0, epsilon, scale);
     }
 }
 
@@ -617,10 +620,11 @@ static void tri_f32(const void * restrict src,
     const uint32_t nloe  = ne0 % VLEN_FP32;
 
     const uint32_t ne01 = uctx->octx->src[0]->ne[1];
+    const struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01;
 
     for (uint32_t b = 0; b < num_rows; b++) {
         const uint32_t abs_row = ir + b;
-        const uint32_t i01     = abs_row % ne01;
+        const uint32_t i01     = fastmodulo(abs_row, ne01, div_ne01);
 
         const HVX_Vector * restrict v_src = (const HVX_Vector *) ((const uint8_t *) src + b * src0_row_size_aligned);
         HVX_Vector * restrict v_dst       = (HVX_Vector *) ((uint8_t *) dst + b * dst_row_size_aligned);
@@ -1163,35 +1167,23 @@ static void unary_thread_row_block(unsigned int nth, unsigned int ith, void * da
     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 bool src0_contig = uctx->kparams->src0_contig;
+    const uint32_t BLOCK   = uctx->block;
 
     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);
+        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_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) :
+        const size_t src0_off = src0_contig ? ((size_t) 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),
@@ -1203,8 +1195,7 @@ static void unary_thread_row_block(unsigned int nth, unsigned int ith, void * da
     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);
+        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_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;
@@ -1213,8 +1204,7 @@ static void unary_thread_row_block(unsigned int nth, unsigned int ith, void * da
         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);
+        const size_t dst_off = (size_t) ir * nb1;
         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);
@@ -1222,12 +1212,12 @@ static void unary_thread_row_block(unsigned int nth, unsigned int ith, void * da
         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);
+                                                              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) :
+                                                                  ne01, div_ne01);
+                const size_t src0_pref_off = src0_contig ? ((size_t) 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),
@@ -1273,10 +1263,6 @@ static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, vo
     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;
@@ -1286,28 +1272,15 @@ static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, vo
     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 bool src0_contig  = uctx->kparams->src0_contig;
+    const bool src1_contig  = uctx->kparams->src1_contig;
+    const bool block_contig = src0_contig && src1_contig;
+    const uint32_t BLOCK    = uctx->block;
 
     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) {
@@ -1317,22 +1290,21 @@ static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, vo
     }
 
     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);
+        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_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) :
+        const size_t src0_off = src0_contig ? ((size_t) 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);
+            const size_t src1_off = src1_contig ? ((size_t) ir * nb11) :
+                unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13);
             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);
@@ -1344,8 +1316,7 @@ static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, vo
     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);
+        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_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;
@@ -1359,30 +1330,28 @@ static void unary_thread_rms_norm_mul_f32(unsigned int nth, unsigned int ith, vo
         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);
+        const size_t dst_off = (size_t) ir * nb1;
         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 next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, block_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) :
+                const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, block_contig,
+                                                                  ne01, div_ne01);
+                const size_t src0_pref_off = src0_contig ? ((size_t) 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);
+                    const size_t src1_pref_off = src1_contig ? ((size_t) pref_ir * nb11) :
+                        unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13);
                     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);
@@ -1428,37 +1397,18 @@ static void unary_thread_tri_f32(unsigned int nth, unsigned int ith, void * data
 
     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;
-    }
+    const uint32_t BLOCK             = uctx->block;
 
     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);
+        const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
 
         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);
+        const size_t src0_off = (size_t) ir * nb01;
         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);
@@ -1469,8 +1419,7 @@ static void unary_thread_tri_f32(unsigned int nth, unsigned int ith, void * data
     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);
+        const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
 
         void * dst_vtcm  = (void *) (uintptr_t) dma_queue_pop(dma_q).src;
         void * src0_vtcm = (void *) (uintptr_t) dma_queue_pop(dma_q).dst;
@@ -1479,22 +1428,18 @@ static void unary_thread_tri_f32(unsigned int nth, unsigned int ith, void * data
         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);
+        const size_t dst_off = (size_t) ir * nb1;
         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 next_block_size = MIN(BLOCK, src0_end_row - next_ir);
             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);
+                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_ir);
+                const size_t src0_pref_off = (size_t) pref_ir * nb01;
                 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);
@@ -1544,21 +1489,17 @@ static void unary_thread_tiled(unsigned int nth, unsigned int ith, void * data)
     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 tiles_per_row = uctx->kparams->tiles_per_row;
+    const bool     src0_contig   = uctx->kparams->src0_contig;
 
     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 row  = src0_start_row + fastdiv(t, div_tpr);
+        const uint32_t col  = fastmodulo(t, tiles_per_row, div_tpr) * 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) :
+        const size_t   soff = (src0_contig ? ((size_t) row * nb01) :
                                unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) +
                                (size_t) col * sizeof(float);
 
@@ -1586,9 +1527,7 @@ static void unary_thread_tiled(unsigned int nth, unsigned int ith, void * data)
         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 doff = (size_t) row * nb1 + (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);
 
@@ -1596,7 +1535,7 @@ static void unary_thread_tiled(unsigned int nth, unsigned int ith, void * data)
         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) :
+            const size_t   psoff = (src0_contig ? ((size_t) prow * nb01) :
                                     unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02,
                                                      nb03)) +
                                    (size_t) pcol * sizeof(float);
@@ -1657,29 +1596,20 @@ static void unary_thread_tiled_tri_f32(unsigned int nth, unsigned int ith, void
 
     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 struct fastdiv_values * div_ne01 = &uctx->kparams->div_ne01;
+    const struct fastdiv_values * div_tpr  = &uctx->kparams->div_tpr;
 
-    const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile;
+    const uint32_t tiles_per_row = uctx->kparams->tiles_per_row;
     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 row  = src0_start_row + fastdiv(t, div_tpr);
+        const uint32_t col  = fastmodulo(t, tiles_per_row, div_tpr) * 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);
+        const size_t   soff = (size_t) row * nb01 + (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);
@@ -1706,20 +1636,15 @@ static void unary_thread_tiled_tri_f32(unsigned int nth, unsigned int ith, void
         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 doff = (size_t) row * nb1 + (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);
+            const uint32_t ptw   = MIN(col_tile, ne0 - pcol);
+            const size_t   ptb   = (size_t) ptw * sizeof(float);
+            const size_t   psoff = (size_t) prow * nb01 + (size_t) pcol * sizeof(float);
             dma_queue_push(dma_q, dma_make_data(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1);
         }
 
@@ -1787,28 +1712,6 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             return HTP_STATUS_NO_SUPPORT;
     }
 
-    // F16 only has row-block kernels for this subset of ops (see the dispatch switch
-    // below) - reject everything else up front, before touching kparams/VTCM.
-    if (is_f16) {
-        switch (octx->op) {
-            case HTP_OP_NORM:
-            case HTP_OP_RMS_NORM:
-            case HTP_OP_SCALE:
-            case HTP_OP_CLAMP:
-            case HTP_OP_SQR:
-            case HTP_OP_SQRT:
-            case HTP_OP_L2_NORM:
-            case HTP_OP_UNARY_ABS:
-            case HTP_OP_UNARY_LOG:
-            case HTP_OP_UNARY_STEP:
-            case HTP_OP_HADAMARD:
-                break;
-            default:
-                FARF(ERROR, "unary-%s: not supported for F16\n", op_type);
-                return HTP_STATUS_NO_SUPPORT;
-        }
-    }
-
     const struct htp_unary_kernel_params * kparams = (const struct htp_unary_kernel_params *) octx->kernel_params;
 
     if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
@@ -1838,15 +1741,11 @@ static int execute_op_unary(struct htp_ops_context * octx) {
 
     const uint32_t n_threads = octx->n_threads;
 
-    const size_t src0_row_size_aligned = kparams->src0_row_size_aligned;
-    const size_t dst_row_size_aligned  = kparams->dst_row_size_aligned;
-
-    // Always 0 for F16 - htp_unary_vtcm_layout_build() keeps F16 on the row-block path,
+    // Always 0 for F16 - htp_unary_solve_layout() keeps F16 on the row-block path,
     // since only F32 has unary_task_f32_tiled_* kernels.
     const uint32_t col_tile = kparams->col_tile;
 
     size_t src1_data_row_size = 0;
-    size_t src1_row_size_aligned = kparams->src1_row_size_aligned;
     bool broadcast_weight = kparams->broadcast_weight;
     const struct htp_tensor * src1 = NULL;
 
@@ -1857,8 +1756,8 @@ static int execute_op_unary(struct htp_ops_context * octx) {
         src1_data_row_size = src1->ne[0] * sizeof(float);
     }
 
-    if (octx->ctx->vtcm_size < (size_t)kparams->vtcm_size) {
-        FARF(ERROR, "unary-%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, (size_t)kparams->vtcm_size);
+    if ((!col_tile && kparams->block == 0) || octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) {
+        FARF(ERROR, "unary-%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, (size_t) kparams->vtcm_size);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
@@ -1866,9 +1765,13 @@ static int execute_op_unary(struct htp_ops_context * octx) {
     octx->src1_spad.src = NULL;
     octx->dst_spad.src  = NULL;
 
-    FARF(HIGH, "%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %u src1-vtcm-size %u dst-vtcm-size %u\n", op_type,
+    struct htp_unary_vtcm_layout L;
+    htp_unary_vtcm_layout_build(&L, octx->op, src0_data_row_size, dst_data_row_size, src1_data_row_size,
+                                src0_elem_size, col_tile, kparams->block, broadcast_weight, kparams->n_threads);
+
+    FARF(HIGH, "%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\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],
-         kparams->vtcm_src0_size, kparams->vtcm_src1_size, kparams->vtcm_dst_size);
+         L.src0_bytes * n_threads, L.src1_bytes * n_threads, L.dst_bytes * n_threads);
 
     uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
     struct htp_unary_context uctx = {
@@ -1886,26 +1789,26 @@ static int execute_op_unary(struct htp_ops_context * octx) {
         .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 = L.src0_row_size_aligned,
+        .src1_row_size_aligned = L.src1_row_size_aligned,
+        .dst_row_size_aligned  = L.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   = L.src0_bytes >> 1,
+        .src1_vtcm_half_size   = (octx->op == HTP_OP_RMS_NORM_MUL) ? (L.src1_bytes >> (broadcast_weight ? 0 : 1)) : 0,
+        .dst_vtcm_half_size    = L.dst_bytes >> 1,
 
         .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, L.off_src0),
+        .vtcm_src1             = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_src1, L.src1_bytes > 0),
+        .vtcm_dst              = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst),
 
-        .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 = L.src0_bytes,
+        .vtcm_src1_size_per_thread = L.src1_bytes,
+        .vtcm_dst_size_per_thread  = L.dst_bytes,
     };
 
     FARF(HIGH, "%s: %s mode (col_tile %u)\n", op_type, col_tile ? "tiled" : "row-block", col_tile);
@@ -2007,8 +1910,25 @@ static int execute_op_unary(struct htp_ops_context * octx) {
 int op_unary(struct htp_ops_context * octx) {
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-        case HTP_TYPE_F16:
-            return execute_op_unary(octx);
+        case HTP_TYPE_F16: {
+            if (!htp_unary_is_elementwise(octx->op)) {
+                return execute_op_unary(octx);
+            }
+
+            struct htp_tensor src0_c = *octx->src[0];
+            struct htp_tensor dst_c  = *octx->dst;
+            const size_t elem_size = (src0_c.type == HTP_TYPE_F16) ? sizeof(_Float16) : sizeof(float);
+            HTP_UNARY_COLLAPSE_DIMS(src0_c.ne, src0_c.nb, dst_c.ne, dst_c.nb, elem_size);
+
+            const struct htp_tensor * orig_src0 = octx->src[0];
+            const struct htp_tensor * orig_dst  = octx->dst;
+            octx->src[0] = &src0_c;
+            octx->dst    = &dst_c;
+            const int status = execute_op_unary(octx);
+            octx->src[0] = orig_src0;
+            octx->dst    = orig_dst;
+            return status;
+        }
 
         default:
             return HTP_STATUS_NO_SUPPORT;

ggml/src/ggml-hexagon/htp/unary-ops.h

@@ -2,29 +2,22 @@
 #define HTP_UNARY_OPS_H
 
 #include "hex-common.h"
+#include "hex-fastdiv.h"
 #include "htp-ops.h"
+#include "htp-vtcm.h"
 
 // Op-specific struct for precomputed unary params
 struct htp_unary_kernel_params {
     uint32_t  n_threads;
     uint32_t  col_tile;
-    uint32_t  vtcm_row_per_thread;
     uint32_t  block;
     uint32_t  broadcast_weight;
 
-    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  src0_row_size_aligned;
-    uint32_t  src1_row_size_aligned;
-    uint32_t  dst_row_size_aligned;
+    uint32_t  src0_contig;
+    uint32_t  src1_contig;
 
     uint32_t  vtcm_size;
+    uint32_t  tiles_per_row;
 
     // Fastdiv helpers
     struct fastdiv_values div_ne01;
@@ -71,23 +64,105 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
     }
 }
 
+static inline bool htp_unary_is_elementwise(uint32_t opcode) {
+    switch (opcode) {
+        case HTP_OP_NORM:
+        case HTP_OP_RMS_NORM:
+        case HTP_OP_RMS_NORM_MUL:
+        case HTP_OP_L2_NORM:
+        case HTP_OP_HADAMARD:
+        case HTP_OP_TRI:
+            return false;
+        default:
+            return true;
+    }
+}
+
+#define HTP_UNARY_COLLAPSE_DIMS(src_ne, src_nb, dst_ne, dst_nb, elem_size)                   \
+    do {                                                                                     \
+        while (((src_ne)[1] > 1 || (src_ne)[2] > 1 || (src_ne)[3] > 1) &&                    \
+               (size_t) (src_nb)[0] == (elem_size) &&                                        \
+               (size_t) (dst_nb)[0] == (elem_size) &&                                        \
+               (((src_ne)[1] == 1 && (dst_ne)[1] == 1) ||                                    \
+                ((size_t) (src_nb)[1] == (size_t) (src_ne)[0] * (size_t) (src_nb)[0] &&      \
+                 (size_t) (dst_nb)[1] == (size_t) (dst_ne)[0] * (size_t) (dst_nb)[0] &&      \
+                 ((src_ne)[0] == 1 ||                                                        \
+                  (uint64_t) (src_ne)[0] * (uint64_t) (src_ne)[1] * (elem_size) <= 128)))) { \
+            (src_ne)[0] *= (src_ne)[1];                                                      \
+            (src_ne)[1]  = (src_ne)[2];                                                      \
+            (src_ne)[2]  = (src_ne)[3];                                                      \
+            (src_ne)[3]  = 1;                                                                \
+            (dst_ne)[0] *= (dst_ne)[1];                                                      \
+            (dst_ne)[1]  = (dst_ne)[2];                                                      \
+            (dst_ne)[2]  = (dst_ne)[3];                                                      \
+            (dst_ne)[3]  = 1;                                                                \
+            (src_nb)[1]  = (src_nb)[2];                                                      \
+            (src_nb)[2]  = (src_nb)[3];                                                      \
+            (dst_nb)[1]  = (dst_nb)[2];                                                      \
+            (dst_nb)[2]  = (dst_nb)[3];                                                      \
+        }                                                                                    \
+    } while (0)
+
 struct htp_unary_vtcm_layout {
     size_t total_bytes;
     size_t off_src0;
     size_t off_src1;
     size_t off_dst;
 
+    size_t src0_row_size_aligned;
+    size_t src1_row_size_aligned;
+    size_t dst_row_size_aligned;
+
     size_t src0_bytes;
     size_t src1_bytes;
     size_t dst_bytes;
 };
 
 static inline void htp_unary_vtcm_layout_build(
+    struct htp_unary_vtcm_layout * L,
+    uint32_t op,
+    size_t src0_data_row_size,
+    size_t dst_data_row_size,
+    size_t src1_data_row_size,
+    size_t elem_size,
+    uint32_t col_tile,
+    uint32_t block,
+    bool broadcast_weight,
+    uint32_t n_threads
+) {
+    L->src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
+    L->dst_row_size_aligned  = hex_round_up(dst_data_row_size,  128);
+    L->src1_row_size_aligned = (op == HTP_OP_RMS_NORM_MUL) ? hex_round_up(src1_data_row_size, 128) : 0;
+
+    if (col_tile > 0) {
+        const size_t col_tile_bytes = col_tile * elem_size;
+        L->src0_bytes = col_tile_bytes * 2;
+        L->dst_bytes  = col_tile_bytes * 2;
+        L->src1_bytes = 0;
+    } else {
+        L->src0_bytes = L->src0_row_size_aligned * block * 2;
+        L->dst_bytes  = L->dst_row_size_aligned  * block * 2;
+        if (op == HTP_OP_RMS_NORM_MUL) {
+            L->src1_bytes = broadcast_weight ? L->src1_row_size_aligned : (L->src1_row_size_aligned * block * 2);
+        } else {
+            L->src1_bytes = 0;
+        }
+    }
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes * n_threads);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_src1, L->src1_bytes * n_threads, L->src1_bytes > 0);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes * n_threads);
+    L->total_bytes = off;
+}
+
+static inline void htp_unary_solve_layout(
     struct htp_unary_vtcm_layout * L,
     uint32_t op,
     uint32_t ne00,
     uint32_t ne10,
     uint32_t ne11,
+    uint32_t max_block,
     bool broadcast_weight,
     uint32_t n_threads,
     size_t vtcm_size,
@@ -98,16 +173,11 @@ static inline void htp_unary_vtcm_layout_build(
     const size_t dst_elem_size      = (op == HTP_OP_HADAMARD) ? sizeof(float) : elem_size;
     const size_t src0_data_row_size = ne00 * elem_size;
     const size_t dst_data_row_size  = ne10 * dst_elem_size;
+    const size_t src1_data_row_size = (op == HTP_OP_RMS_NORM_MUL) ? (ne11 * sizeof(float)) : 0;
 
     const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
     const size_t dst_row_size_aligned  = hex_round_up(dst_data_row_size,  128);
-
-    size_t src1_row_size_aligned = 0;
-    if (op == HTP_OP_RMS_NORM_MUL) {
-        // RMS_NORM_MUL fusion is F32-only; its weight tensor is always F32.
-        const size_t src1_data_row_size = ne11 * sizeof(float);
-        src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
-    }
+    const size_t src1_row_size_aligned = (op == HTP_OP_RMS_NORM_MUL) ? hex_round_up(src1_data_row_size, 128) : 0;
 
     size_t vtcm_size_per_row = 0;
     size_t vtcm_row_per_thread = 0;
@@ -132,6 +202,10 @@ static inline void htp_unary_vtcm_layout_build(
         vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
     }
 
+    if (max_block > 0 && vtcm_row_per_thread > max_block) {
+        vtcm_row_per_thread = max_block;
+    }
+
     const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
                                op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM ||
                                op == HTP_OP_HADAMARD);
@@ -148,37 +222,14 @@ static inline void htp_unary_vtcm_layout_build(
         const size_t per_thread_budget = vtcm_size / n_threads;
         const size_t col_tile_bytes = hex_align_down(per_thread_budget / 4, 128);
         col_tile = (uint32_t) (col_tile_bytes / elem_size);
-
-        L->src0_bytes = col_tile_bytes * 2;
-        L->dst_bytes  = col_tile_bytes * 2;
-        L->src1_bytes = 0;
-    } else {
-        L->src0_bytes = src0_row_size_aligned * vtcm_row_per_thread * 2;
-        L->dst_bytes  = dst_row_size_aligned * vtcm_row_per_thread * 2;
-        if (op == HTP_OP_RMS_NORM_MUL) {
-            if (broadcast_weight) {
-                L->src1_bytes = src1_row_size_aligned;
-            } else {
-                L->src1_bytes = src1_row_size_aligned * vtcm_row_per_thread * 2;
-            }
-        } else {
-            L->src1_bytes = 0;
-        }
     }
 
-    L->off_src0 = 0;
-    if (op == HTP_OP_RMS_NORM_MUL) {
-        L->off_src1 = L->off_src0 + L->src0_bytes * n_threads;
-        L->off_dst  = L->off_src1 + L->src1_bytes * n_threads;
-    } else {
-        L->off_src1 = 0;
-        L->off_dst  = L->off_src0 + L->src0_bytes * n_threads;
-    }
-
-    L->total_bytes = L->off_dst + L->dst_bytes * n_threads;
+    const uint32_t block = col_tile ? 0 : (uint32_t) vtcm_row_per_thread;
+    htp_unary_vtcm_layout_build(L, op, src0_data_row_size, dst_data_row_size, src1_data_row_size,
+                                elem_size, col_tile, block, broadcast_weight, n_threads);
 
     *out_col_tile = col_tile;
-    *out_vtcm_row_per_thread = vtcm_row_per_thread;
+    *out_vtcm_row_per_thread = (uint32_t) vtcm_row_per_thread;
 }
 
 #endif /* HTP_UNARY_OPS_H */

ggml/src/ggml-sycl/common.hpp

@@ -16,7 +16,10 @@
 #include <cstddef>
 #include <fstream>
 #include <iostream>
+#include <optional>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 
 #include "base.hpp"
 #include "dpct/helper.hpp"
@@ -134,6 +137,7 @@ extern int g_ggml_sycl_fa_onednn_max_kv;
 extern int g_ggml_sycl_enable_mkl_fa;
 extern int g_ggml_sycl_memtrace;
 extern int g_ggml_sycl_memtrace_step;
+extern int g_ggml_sycl_graph_eviction_timeout;
 
 
 #define CHECK_TRY_ERROR(expr)                                            \
@@ -315,6 +319,7 @@ struct optimize_feature {
 };
 
 struct sycl_device_info {
+    bool active; //the device is activated/used in application.
     int cc;  // compute capability
     int nsm; // number of streaming multiprocessors (CUDA) maps to the maximum
              // number of compute units on a SYCL device.
@@ -331,10 +336,15 @@ struct sycl_device_info {
     sycl_hw_info hw_info;
     optimize_feature opt_feature;
     bool    usm_system_support; // support for USM system allocations
+#ifdef GGML_SYCL_GRAPH
+    bool    graph_support;        // command graphs can be recorded and replayed
+    bool    graph_update_support; // a finalized command graph can be updated
+#endif
 };
 
 
 struct ggml_sycl_device_info {
+    int active_device_count; //the number of active devices in the application
     int device_count;
 
     sycl_device_info devices[GGML_SYCL_MAX_DEVICES] = {};
@@ -346,7 +356,7 @@ struct ggml_sycl_device_info {
     bool ext_oneapi_level_zero = true; // sycl::backend::ext_oneapi_level_zero used by all enumerated GPU devices
 };
 
-const ggml_sycl_device_info & ggml_sycl_info();
+ggml_sycl_device_info & ggml_sycl_info();
 
 static inline bool ggml_sycl_xmx_enabled(int device) {
     return g_ggml_sycl_enable_xmx_esimd && ggml_sycl_xmx_supports_arch(ggml_sycl_info().devices[device].hw_info.arch);
@@ -441,6 +451,32 @@ struct ggml_sycl_gg_tile {
 };
 
 namespace sycl_ex = sycl::ext::oneapi::experimental;
+
+#ifdef GGML_SYCL_GRAPH
+struct ggml_sycl_graph {
+    // src data/ne/nb are kept next to the node copy: the scheduler can hand back the same src
+    // pointer with different contents or shape, see https://github.com/ggml-org/llama.cpp/pull/21736
+    struct node_properties {
+        ggml_tensor node;
+        void *      node_src_data_ptrs[GGML_MAX_SRC];
+        int64_t     node_src_ne[GGML_MAX_SRC][GGML_MAX_DIMS];
+        size_t      node_src_nb[GGML_MAX_SRC][GGML_MAX_DIMS];
+    };
+
+    std::unique_ptr<sycl_ex::command_graph<sycl_ex::graph_state::executable>> exec_graph = nullptr;
+    std::vector<node_properties> node_props;
+    bool     warmup_complete = false;
+    uint64_t uid             = 0;
+    int64_t  last_used_time  = 0;
+
+    // result of check_graph_compatibility() and the uid it was made for
+    bool     compatible     = false;
+    uint64_t compatible_uid = 0;
+};
+
+static_assert(std::is_trivial<ggml_sycl_graph::node_properties>::value, "node_properties must be trivial");
+#endif
+
 struct ggml_backend_sycl_context {
     int device;
     std::string name;
@@ -535,7 +571,36 @@ struct ggml_backend_sycl_context {
     }
 
 #ifdef GGML_SYCL_GRAPH
-    std::unique_ptr<sycl_ex::command_graph<sycl_ex::graph_state::executable>> exec_graph = nullptr;
+    // Map from first node pointer to graph - allows multiple graphs per context when the
+    // computation is split across CPU/GPU (e.g. with --n-cpu-moe)
+    std::unordered_map<const void *, std::unique_ptr<ggml_sycl_graph>> sycl_graphs;
+
+    int64_t last_graph_eviction_sweep = 0;
+    std::optional<bool> graph_compatible;
+
+    ggml_sycl_graph * sycl_graph(const void * first_node_ptr) {
+        const int64_t time_now = ggml_time_us();
+        const int64_t eviction_timeout_us = g_ggml_sycl_graph_eviction_timeout * 1'000'000LL;
+
+        // sweep every half the eviction timeout, evicting graphs unused for >= the eviction timeout
+        if (time_now - last_graph_eviction_sweep >= eviction_timeout_us / 2) {
+            last_graph_eviction_sweep = time_now;
+            for (auto it = sycl_graphs.begin(); it != sycl_graphs.end(); ) {
+                if (time_now - it->second->last_used_time >= eviction_timeout_us) {
+                    it = sycl_graphs.erase(it);
+                } else {
+                    ++it;
+                }
+            }
+        }
+
+        auto it = sycl_graphs.find(first_node_ptr);
+        if (it == sycl_graphs.end()) {
+            it = sycl_graphs.emplace(first_node_ptr, std::make_unique<ggml_sycl_graph>()).first;
+        }
+        it->second->last_used_time = time_now;
+        return it->second.get();
+    }
 #endif
 
     ggml_sycl_pool & host_pool(int device) {

ggml/src/ggml-sycl/fattn-mkl.cpp

@@ -14,6 +14,7 @@
 #include <oneapi/mkl.hpp>
 #include <cstdio>
 #include <chrono>
+#include <vector>
 
 #define MKL_FA_CHUNK_SIZE_KV 8192
 
@@ -130,7 +131,8 @@ static void mkl_fa_online_softmax_chunk(
     int kvh_head,
     const sycl::half * mask_data, int64_t mask_head_stride,
     int64_t mask_row_stride, int mask_n_heads,
-    float logit_softcap, int64_t wg_size) {
+    float logit_softcap, int64_t wg_size,
+    const sycl::event & dep_event) {
 
     // 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
@@ -138,6 +140,7 @@ static void mkl_fa_online_softmax_chunk(
     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.depends_on(dep_event);
         cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
             [=](sycl::nd_item<1> item) {
                 const int local_id = (int)item.get_local_id(0);
@@ -563,9 +566,6 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor *
                 KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr,
                 n_query_rows, DV, wg_size);
 
-            // Sync before MKL GEMM (MKL may use an internal queue)
-            stream->wait();
-
             // 3. KV chunk loop (OUTER): dequant each chunk once, then tile queries.
             for (int chunk_start = 0; chunk_start < n_kv; chunk_start += chunk_size) {
                 int this_chunk = std::min(chunk_size, n_kv - chunk_start);
@@ -579,29 +579,26 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor *
                         mkl_fa_dequant_chunk(stream, V_desc, KQV,
                             V_chunk_f16_ptr, ikvh, chunk_start, this_chunk);
                     }
-                    stream->wait();  // dequant must be ready before MKL GEMM
                     MKL_ACCUM(dequant_time_us, t0);
                 }
 
                 // 3b. Query tile loop (INNER) — bounds KQ_f32/S_f16 footprint.
                 for (int q0 = 0; q0 < n_query_rows; q0 += q_tile_rows) {
                     int q_rows = std::min(q_tile_rows, n_query_rows - q0);
+                    sycl::event kq_done, vkq_done;
 
                     // GEMM: KQ = Q_tile × K_chunk^T
                     {
                         MKL_TAKE_TIME(t0);
-                        sycl::event ev = gemm(*stream,
+                        kq_done = gemm(*stream,
                             transpose::trans, transpose::nontrans,
                             this_chunk, q_rows, DKQ,
                             alpha,
                             K_chunk_f16_ptr, DKQ,
                             Q_head_f16_ptr + (int64_t)q0 * DKQ, DKQ,
                             beta,
-                            KQ_f32_ptr, this_chunk);
-                        try { ev.wait_and_throw(); } catch (sycl::exception & e) {
-                            GGML_LOG_INFO("[MKL-FA] GEMM KQ: %s\n", e.what());
-                            GGML_ABORT("MKL GEMM KQ failed");
-                        }
+                            KQ_f32_ptr, this_chunk,
+                            std::vector<sycl::event>{*stream->ext_oneapi_get_last_event()});
                         MKL_ACCUM(gemm_kq_time_us, t0);
                     }
                     // Online softmax over this chunk for this query tile
@@ -615,26 +612,22 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor *
                             kvh_base_head,
                             mask_batch, mask_head_stride,
                             mask_row_stride, mask_n_heads,
-                            logit_softcap, wg_size);
-                        stream->wait();  // S_f16 must be ready for GEMM
+                            logit_softcap, wg_size, kq_done);
                         MKL_ACCUM(softmax_time_us, t0);
                     }
 
                     // GEMM: VKQ_chunk = S × V_chunk
                     {
                         MKL_TAKE_TIME(t0);
-                        sycl::event ev = gemm(*stream,
+                        vkq_done = gemm(*stream,
                             transpose::nontrans, transpose::nontrans,
                             DV, q_rows, this_chunk,
                             alpha,
                             V_chunk_f16_ptr, DV,
                             S_f16_ptr, this_chunk,
                             beta,
-                            VKQ_chunk_ptr, DV);
-                        try { ev.wait_and_throw(); } catch (sycl::exception & e) {
-                            GGML_LOG_INFO("[MKL-FA] GEMM VKQ: %s\n", e.what());
-                            GGML_ABORT("MKL GEMM VKQ failed");
-                        }
+                            VKQ_chunk_ptr, DV,
+                            std::vector<sycl::event>{*stream->ext_oneapi_get_last_event()});
                         MKL_ACCUM(gemm_vkq_time_us, t0);
                     }
                     // VKQ_accum[q0..] += VKQ_chunk
@@ -644,6 +637,7 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor *
                             * wg_size;
                         float * accum = VKQ_accum_ptr + (int64_t)q0 * DV;
                         stream->submit([&](sycl::handler & cgh) {
+                            cgh.depends_on(vkq_done);
                             cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
                                 [=](sycl::nd_item<1> item) {
                                     int64_t i = item.get_global_id(0);

ggml/src/ggml-sycl/fattn-sparse.cpp

@@ -218,7 +218,7 @@ bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tenso
                           d_idx, d_cnt, n_kv_g, n_rows_m,
                           mask->nb[1] / sizeof(sycl::half));
 
-    if (sparse_fa_debug()) {
+    if (sparse_fa_debug() && stream->ext_oneapi_get_state() != sycl::ext::oneapi::experimental::queue_state::recording) {
         int32_t h_cnt = 0;
         SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(&h_cnt, d_cnt, sizeof(int32_t))));
         SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));

ggml/src/ggml-sycl/fattn.cpp

@@ -365,10 +365,10 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
         const ggml_tensor * K_diag = dst->src[1];
         const ggml_tensor * V_diag = dst->src[2];
         const ggml_tensor * Q_diag = dst->src[0];
-        if (K_diag->ne[1] >= 1024) {
+        dpct::queue_ptr q = ctx.stream();
+        if (K_diag->ne[1] >= 1024 && q->ext_oneapi_get_state() != sycl::ext::oneapi::experimental::queue_state::recording) {
             fa_diag_count++;
             float diag_buf[64];
-            dpct::queue_ptr q = ctx.stream();
             q->memcpy(diag_buf, dst->data, 64 * sizeof(float));
             q->wait();
             const char * kname = "???";
@@ -404,6 +404,14 @@ bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) {
     return ggml_sycl_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE;
 }
 
+bool ggml_sycl_flash_attn_ext_needs_sync(int device, const ggml_tensor * dst) {
+    return ggml_sycl_get_best_fattn_kernel(device, dst) == BEST_FATTN_KERNEL_MKL;
+}
+
+bool ggml_sycl_flash_attn_ext_uses_onednn(int device, const ggml_tensor * dst) {
+    return ggml_sycl_get_best_fattn_kernel(device, dst) == BEST_FATTN_KERNEL_ONEDNN;
+}
+
 static uintptr_t ggml_sycl_fattn_reserve_halves(ggml_sycl_fattn_extra & extra, size_t n_halves) {
     if (n_halves == 0) {
         return 0;

ggml/src/ggml-sycl/fattn.hpp

@@ -19,6 +19,12 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
 
 bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst);
 
+// Reports if flash attention runs this node with the oneMKL kernel, which waits on the host
+bool ggml_sycl_flash_attn_ext_needs_sync(int device, const ggml_tensor * dst);
+
+// Reports if flash attention runs this node with the oneDNN kernel
+bool ggml_sycl_flash_attn_ext_uses_onednn(int device, const ggml_tensor * dst);
+
 // Scratch that flash attention needs beyond the output tensor
 struct ggml_sycl_fattn_extra {
     uintptr_t K_buffer_ptr     = 0;   // F16 copy of the K cache

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -26,6 +26,7 @@
 #include <stdio.h>
 #include <vector>
 #include <cmath>
+#include <cstring>
 #include <iostream>
 #include <fstream>
 #include <stdio.h>
@@ -46,6 +47,13 @@
 #    include <sycl/ext/oneapi/virtual_mem/virtual_mem.hpp>
 #    define GGML_SYCL_SUPPORT_VMM
 #endif
+// beginning in oneDNN 3.12.3, "Enabled matmul and inner product primitives compatibility with SYCL Graph native recording mode on Intel GPUs."
+#if GGML_SYCL_DNNL
+#    include "dnnl_version.h"
+#    if (DNNL_VERSION_MAJOR > 3 || (DNNL_VERSION_MAJOR == 3 && (DNNL_VERSION_MINOR > 12 || (DNNL_VERSION_MINOR == 12 && DNNL_VERSION_PATCH >= 3))))
+#        define GGML_SYCL_GRAPH_DNNL_SUPPORTED
+#    endif
+#endif
 #include <sycl/half_type.hpp>
 
 #include "ggml.h"
@@ -144,6 +152,7 @@ int g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_LEVEL_ZERO;
 int g_ggml_sycl_enable_sparse_fa = 0;
 int g_ggml_sycl_debug_sparse_fa = 0;
 int g_ggml_sycl_sparse_fa_margin = 256;
+int g_ggml_sycl_graph_eviction_timeout = 10;
 
 #if GGML_SYCL_DNNL
 // ask oneDNN which matmul it picks for a small f16 problem
@@ -245,6 +254,11 @@ static ggml_sycl_device_info ggml_sycl_init() {
         info.devices[i].warp_size = WARP_SIZE;
         info.devices[i].usm_system_support = device.has(sycl::aspect::usm_system_allocations);
 
+#ifdef GGML_SYCL_GRAPH
+        info.devices[i].graph_support        = device.has(sycl::aspect::ext_oneapi_limited_graph);
+        info.devices[i].graph_update_support = device.has(sycl::aspect::ext_oneapi_graph);
+#endif
+
         info.max_work_group_sizes[i] = prop.get_max_work_group_size();
         info.devices[i].max_wg_per_cu = info.max_work_group_sizes[i] / prop.get_max_compute_units();
         info.devices[i].hw_info = get_device_hw_info(&device);
@@ -283,7 +297,7 @@ static ggml_sycl_device_info ggml_sycl_init() {
     return info;
 }
 
-const ggml_sycl_device_info & ggml_sycl_info() {
+ggml_sycl_device_info & ggml_sycl_info() {
     static ggml_sycl_device_info info = ggml_sycl_init();
     return info;
 }
@@ -432,6 +446,7 @@ static void ggml_check_sycl() try {
         g_ggml_sycl_xmx_glu_max_cols =
             std::clamp(ggml_sycl_get_env("GGML_SYCL_XMX_GLU_MAX_COLS", GGML_SYCL_XMX_GLU_MAX_COLS), 0, GGML_SYCL_XMX_GLU_MAX_COLS);
         g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
+        g_ggml_sycl_graph_eviction_timeout = ggml_sycl_get_env("GGML_SYCL_GRAPH_EVICTION_TIMEOUT", 10);
         g_ggml_sycl_xmx_gather_types = ggml_sycl_get_env("GGML_SYCL_XMX_GATHER_TYPES", GGML_SYCL_XMX_GATHER_TYPES_DEFAULT);
         g_ggml_sycl_xmx_gather_shapes = ggml_sycl_get_env("GGML_SYCL_XMX_GATHER_SHAPES", GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT);
         g_ggml_sycl_dynamic_precision =
@@ -542,6 +557,7 @@ static void ggml_check_sycl() try {
 
 #ifdef GGML_SYCL_GRAPH
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_GRAPH: %d\n", g_ggml_sycl_enable_graph);
+        GGML_LOG_INFO("  GGML_SYCL_GRAPH_EVICTION_TIMEOUT: %d\n", g_ggml_sycl_graph_eviction_timeout);
 #else
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_GRAPH: graph disabled by compile flag\n");
 #endif
@@ -601,10 +617,19 @@ static void ggml_check_sycl() try {
         fprintf(stderr, "%s: SYCL_USE_XMX: no\n", __func__);
 #endif
 */
+#ifdef GGML_SYCL_GRAPH
+        if (g_ggml_sycl_enable_graph) {
+            for (unsigned int i = 0; i < dpct::dev_mgr::instance().device_count(); ++i) {
+                if (!dpct::dev_mgr::instance().get_device(i).has(sycl::aspect::ext_oneapi_limited_graph)) {
+                    GGML_LOG_WARN("SYCL device %u does not support SYCL graphs\n", i);
+                }
+            }
+        }
+
         // Async USM allocation/free is also useful outside the graph path: it avoids the host waits in the reorder
         // staging path while preserving queue ordering semantics. Graph support still depends on the extension being
         // available, but it no longer needs to control the non-graph fast path.
-#if defined(GGML_SYCL_GRAPH) && SYCL_EXT_ONEAPI_ASYNC_MEMORY_ALLOC
+#if SYCL_EXT_ONEAPI_ASYNC_MEMORY_ALLOC
         g_ggml_sycl_use_async_mem_op = g_ggml_sycl_use_async_mem_op_requested || g_ggml_sycl_enable_graph;
         if (g_ggml_sycl_use_async_mem_op) {
             for (unsigned int i = 0; i < dpct::dev_mgr::instance().device_count(); ++i) {
@@ -614,6 +639,7 @@ static void ggml_check_sycl() try {
                 }
             }
         }
+#endif
 #endif
         if (CHECK_TRY_ERROR(g_all_sycl_device_count =
                             dpct::dev_mgr::instance().device_count()) != 0) {
@@ -3108,6 +3134,11 @@ catch (sycl::exception const &exc) {
   std::exit(1);
 }
 
+#if GGML_SYCL_DNNL
+// F32 GEMMs below this size (row_diff * src1_ncols * ne10) call oneMKL directly instead of oneDNN
+static constexpr int64_t SYCL_F32_DNNL_MIN_GEMM_FLOPS = 256 * 256 * 256;
+#endif
+
 inline void ggml_sycl_op_mul_mat_sycl(
     ggml_backend_sycl_context & ctx,
     const ggml_tensor *src0, const ggml_tensor *src1, ggml_tensor *dst,
@@ -3257,7 +3288,7 @@ inline void ggml_sycl_op_mul_mat_sycl(
         {
 #if GGML_SYCL_DNNL
             const int64_t gemm_flops = (int64_t)row_diff * src1_ncols * ne10;
-            const bool use_mkl_direct = gemm_flops < 256 * 256 * 256;
+            const bool use_mkl_direct = gemm_flops < SYCL_F32_DNNL_MIN_GEMM_FLOPS;
             if (g_ggml_sycl_enable_dnn && !use_mkl_direct && ggml_sycl_dnnl_has_optimized_gemm(ggml_sycl_get_device())) {
                 DnnlGemmWrapper::row_gemm(ctx, row_diff, src1_ncols, ne10, src0_ddf_i,
                                           DnnlGemmWrapper::to_dt<float>(), src1_ddf1_i, DnnlGemmWrapper::to_dt<float>(),
@@ -4348,6 +4379,36 @@ static bool ggml_sycl_supports_dmmv(enum ggml_type type) {
     }
 }
 
+static bool ggml_sycl_supports_mmvq_id(enum ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q2_0:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+            return true;
+        default:
+            return false;
+    }
+}
+
 // Helper functions to unify device memory allocation for both async and sync paths
 static inline void * sycl_ext_malloc_device(dpct::queue_ptr stream, size_t size) {
     bool use_async = g_ggml_sycl_use_async_mem_op;
@@ -5198,6 +5259,20 @@ static bool can_use_mul_mat_vec_q(int device, const ggml_tensor * src0, const gg
            (src1->ne[1] <= MMVQ_MAX_BATCH_SIZE || can_use_xmx_batch(device, src0, src1));
 }
 
+static bool can_use_mul_mat_q(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    // mmvq and mmq need the __dp4a instruction which is available for gen12+
+    // Workaround in https://github.com/ggml-org/llama.cpp/commit/95f84d5ce8b449a9b16009434aca800df504a02e
+    if (!ggml_sycl_supports_mmq(src0->type) || src0->type == GGML_TYPE_IQ2_XXS) {
+        return false;
+    }
+#ifdef SYCL_USE_XMX
+    if (src1->ne[1] > MMQ_MAX_BATCH_SIZE) {
+        return false;
+    }
+#endif // SYCL_USE_XMX
+    return src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32;
+}
+
 static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
 
@@ -5245,16 +5320,7 @@ static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor
 
     bool use_mul_mat_vec_q = can_use_mul_mat_vec_q(ctx.device, src0, src1, dst);
 
-    bool use_mul_mat_q =  ggml_sycl_supports_mmq(src0->type)
-        && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32;
-
-
-    // mmvq and mmq need the __dp4a instruction which is available for gen12+
-    // Workaround in https://github.com/ggml-org/llama.cpp/commit/95f84d5ce8b449a9b16009434aca800df504a02e
-    use_mul_mat_q = use_mul_mat_q && (src0->type != GGML_TYPE_IQ2_XXS);
-#ifdef SYCL_USE_XMX
-    use_mul_mat_q = use_mul_mat_q && (src1->ne[1] <= MMQ_MAX_BATCH_SIZE);
-#endif // SYCL_USE_XMX
+    bool use_mul_mat_q = can_use_mul_mat_q(src0, src1, dst);
 
     // When reorder is enabled, both ESIMD, MMVQ and DMMV kernels may be used. For
     // best performance use ESIMD when supported, followed by MMVQ, and finally DMMV.
@@ -5634,22 +5700,39 @@ __dpct_inline__ static void k_copy_dst_from_contiguous(
     }
 }
 
+// returns true when ggml_sycl_mul_mat_id takes the fallback path that requires stream synchronization
+static bool ggml_sycl_mul_mat_id_needs_sync(const ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * ids  = dst->src[2];
+
+    const int64_t ne10 = src1->ne[0];
+    const int64_t ne11 = src1->ne[1];
+    const int64_t ne12 = src1->ne[2];
+
+    if (ne12 != 1) return true;
+    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) return true;
+    if (ne10 != src0->ne[0] || ne10 % QK8_1 != 0) return true;
+    if (!ggml_is_contiguous(src1)) return true;
+
+    const int64_t n_ids_per_group = ids->ne[0];
+    if (ids->ne[1] != 1) return true;
+    if (ne11 != 1 && ne11 != n_ids_per_group) return true;
+
+    return !ggml_sycl_supports_mmvq_id(src0->type);
+}
+
 // Fused MoE TG fast path. Returns false to fall back to the per-expert loop below.
 static bool ggml_sycl_mul_mat_id_mmvq_fused(
     ggml_backend_sycl_context & ctx, const ggml_tensor * src0,
     const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst)
 {
+    if (ggml_sycl_mul_mat_id_needs_sync(dst)) return false;
+
     const int64_t ne10 = src1->ne[0];
     const int64_t ne11 = src1->ne[1];
-    const int64_t ne12 = src1->ne[2];
-    if (ne12 != 1) return false;
-    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) return false;
-    if (ne10 != src0->ne[0] || ne10 % QK8_1 != 0) return false;
-    if (!ggml_is_contiguous(src1)) return false;
 
     const int64_t n_ids_per_group = ids->ne[0];
-    if (ids->ne[1] != 1) return false;
-    if (ne11 != 1 && ne11 != n_ids_per_group) return false;
 
     const queue_ptr stream           = ctx.stream();
     const int       src1_padded_cols = GGML_PAD((int) ne10, MATRIX_ROW_PADDING);
@@ -5756,6 +5839,7 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
             return;
         }
     }
+    GGML_ASSERT(ggml_sycl_mul_mat_id_needs_sync(dst));
 
     // The per-expert loop below reads the experts in whatever layout they have: reorder MXFP4 here as well, so prompt processing does not depend on a single-token decode having run first.
     if (src0->type == GGML_TYPE_MXFP4) {
@@ -6627,8 +6711,6 @@ static int ggml_sycl_try_gdn_cache_fusion(const ggml_cgraph * cgraph, int node_i
 }
 
 static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * sycl_ctx, ggml_cgraph * cgraph) {
-    ggml_sycl_set_main_device(sycl_ctx->device);
-
     for (int i = 0; i < cgraph->n_nodes; i++) {
         ggml_tensor * node = cgraph->nodes[i];
         if (ggml_sycl_is_view_or_noop(node)) {
@@ -6743,42 +6825,167 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
 }
 
 #ifdef GGML_SYCL_GRAPH
-static bool check_graph_compatibility(ggml_cgraph * cgraph) {
-    if (ggml_sycl_info().device_count > 1) {
+static const void * ggml_sycl_graph_get_key(ggml_cgraph * cgraph) {
+    return cgraph->nodes[0];
+}
+
+static bool ggml_sycl_graph_update_required(ggml_sycl_graph * graph, ggml_cgraph * cgraph) {
+    bool res = false;
+
+    if (cgraph->uid != 0 && cgraph->uid == graph->uid) {
+        GGML_SYCL_DEBUG("[SYCL-GRAPH] graph id %" PRIu64 " reused\n", cgraph->uid);
+        GGML_ASSERT((int) graph->node_props.size() == cgraph->n_nodes);
+        return false;
+    }
+
+    graph->uid = cgraph->uid;
+
+    if ((int) graph->node_props.size() != cgraph->n_nodes) {
+        res = true;
+        graph->node_props.resize(cgraph->n_nodes);
+    }
+
+    for (int i = 0; i < cgraph->n_nodes; i++) {
+        ggml_sycl_graph::node_properties prop = {};
+        memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor));
+
+        for (int j = 0; j < GGML_MAX_SRC; ++j) {
+            if (cgraph->nodes[i]->src[j]) {
+                prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data;
+                memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j]));
+                memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j]));
+            }
+        }
+
+        if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) {
+            graph->node_props[i] = prop;
+            res = true;
+        }
+    }
+
+    return res;
+}
+
+// Reports if ggml_sycl_mul_mat() gives this node to oneDNN. oneDNN chains the submission on events
+// made before recording started, which SYCL graphs do not allow.
+static bool mul_mat_uses_onednn(int device, ggml_tensor * dst) {
+#if GGML_SYCL_DNNL
+    if (!g_ggml_sycl_enable_dnn) {
+        return false;
+    }
+
+    ggml_tensor * src0 = dst->src[0];
+    ggml_tensor * src1 = dst->src[1];
+
+    // the branch order below follows the dispatch in ggml_sycl_mul_mat(); the F16 branches end in
+    // a vector kernel or the batched oneMKL path, neither of which uses oneDNN
+    if (!ggml_backend_buffer_is_sycl_split(src0->buffer) && src0->type == GGML_TYPE_F16) {
+        if (ggml_is_permuted(src0) && ggml_is_permuted(src1) && src1->ne[1] == 1) {
+            return false;
+        }
+        if (!ggml_is_contiguous(src0) && !ggml_is_transposed(src1) && src1->ne[1] == 1 && src1->ne[3] == 1) {
+            return false;
+        }
+        if (!ggml_is_transposed(src0) && !ggml_is_transposed(src1) && src1->ne[2] * src1->ne[3] > 1) {
+            return false;
+        }
+    }
+
+    // ggml_sycl_op_mul_mat_sycl() is the last resort of the dispatch and the only user of oneDNN.
+    // The reorder decision that ggml_sycl_mul_mat() makes between DMMV and MMVQ is left out here:
+    // neither uses a library.
+    if (can_use_dequantize_mul_mat_vec(src0, src1, dst) || can_use_mul_mat_vec_q(device, src0, src1, dst) ||
+        can_use_mul_mat_q(src0, src1, dst)) {
+        return false;
+    }
+
+    // the branch order below follows ggml_sycl_op_mul_mat_sycl(); graphs are single device only, so
+    // row_diff == src0->ne[1] and src1_ncols == src1->ne[1] there
+#if defined(GGML_SYCL_HAS_BF16)
+    if (src0->type == GGML_TYPE_BF16 && ggml_is_contiguous(src0)) {
+        return true;
+    }
+#endif
+#ifdef GGML_SYCL_F16
+    if ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && ggml_is_contiguous(src0) &&
+        dst->op_params[0] == GGML_PREC_DEFAULT) {
+        return true;
+    }
+#endif
+    // the F32 branch only uses oneDNN for large GEMMs and calls oneMKL directly otherwise
+    const int64_t gemm_flops = src0->ne[1] * src1->ne[1] * src0->ne[0];
+    return gemm_flops >= SYCL_F32_DNNL_MIN_GEMM_FLOPS;
+#else
+    GGML_UNUSED(device);
+    GGML_UNUSED(dst);
+    return false;
+#endif
+}
+
+static bool check_graph_compatibility(ggml_backend_sycl_context * ctx, ggml_cgraph * cgraph) {
+    if (ggml_sycl_info().active_device_count > 1) {
         // A sycl_ex::command_graph object can only be created for a single device
-        GGML_LOG_INFO("%s: disabling SYCL graphs due to multiple devices\n", __func__);
+        GGML_LOG_DEBUG("%s: disabling SYCL graphs due to multiple devices\n", __func__);
         return false;
     }
 
     for (int i = 0; i < cgraph->n_nodes; i++) {
-        const ggml_op node_op = cgraph->nodes[i]->op;
+        ggml_tensor * node = cgraph->nodes[i];
+
+        // skip the nodes that ggml_backend_sycl_graph_compute_impl() does not run
+        if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || ggml_sycl_is_view_or_noop(node)) {
+            continue;
+        }
+
+        const ggml_op node_op = node->op;
+
         switch (node_op) {
             default:
                 break;
-            case GGML_OP_CONCAT:
-                // ggml_sycl_op_concat() does a blocking host wait after memcpy operations,
-                // but wait() can't be called on the events returned by a queue recording
-                // to a graph.
-                [[fallthrough]];
             case GGML_OP_MUL_MAT_ID:
-                // ggml_sycl_mul_mat_id() does a blocking host wait on the sycl queue after
-                // submitting a memcpy operation, but wait() can't be called on a queue that
-                // is recording to a graph.
-                GGML_LOG_INFO("%s: disabling SYCL graphs due to unsupported node type %s\n", __func__,
-                              ggml_op_name(node_op));
-                return false;
+                // The ggml_sycl_mul_mat_id() fallback path does a blocking host wait on the sycl
+                // queue after submitting a memcpy operation, but wait() can't be called on a queue
+                // that is recording to a graph.
+                if (ggml_sycl_mul_mat_id_needs_sync(node)) {
+                    GGML_LOG_DEBUG("%s: disabling SYCL graphs due to host wait in %s\n", __func__,
+                                   ggml_op_name(node_op));
+                    return false;
+                }
+                break;
+            case GGML_OP_FLASH_ATTN_EXT:
+                if (ggml_sycl_flash_attn_ext_needs_sync(ctx->device, node)) {
+                    GGML_LOG_DEBUG("%s: disabling SYCL graphs due to host wait in %s\n", __func__,
+                                   ggml_op_name(node_op));
+                    return false;
+                }
+#ifndef GGML_SYCL_GRAPH_DNNL_SUPPORTED
+                if (ggml_sycl_flash_attn_ext_uses_onednn(ctx->device, node)) {
+                    GGML_LOG_DEBUG("%s: disabling SYCL graphs due to %s using oneDNN\n", __func__,
+                                   ggml_op_name(node_op));
+                    return false;
+                }
+#endif
+                break;
             case GGML_OP_MUL_MAT:
                 // We cannot use graphs with ggml_sycl_mul_mat() when SYCL async memory allocation extensions are not available,
                 // as SYCL malloc / free and host wait calls are not supported when recording to a graph which are all present
                 // in reordering.
                 if (!g_ggml_sycl_use_async_mem_op) {
-                    GGML_LOG_INFO(
+                    GGML_LOG_DEBUG(
                         "%s: disabling SYCL graphs due to unsupported node type when using a compiler without the "
                         "oneAPI async memory allocation extension "
                         "%s\n",
                         __func__, ggml_op_name(node_op));
                     return false;
                 }
+#ifndef GGML_SYCL_GRAPH_DNNL_SUPPORTED
+                if (mul_mat_uses_onednn(ctx->device, node)) {
+                    GGML_LOG_DEBUG("%s: disabling SYCL graphs due to %s using oneDNN\n", __func__,
+                                   ggml_op_name(node_op));
+                    return false;
+                }
+#endif
+                break;
         }
     }
     return true;
@@ -6788,44 +6995,102 @@ static bool check_graph_compatibility(ggml_cgraph * cgraph) {
 static ggml_status ggml_backend_sycl_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
     auto * sycl_ctx = static_cast<ggml_backend_sycl_context *>(backend->context);
 
+    // set the device on every call: a graph replay does not go through
+    // ggml_backend_sycl_graph_compute_impl()
+    ggml_sycl_set_main_device(sycl_ctx->device);
+
 #ifdef GGML_SYCL_GRAPH
-    bool use_sycl_graph = false;
+    bool              use_sycl_graph = false;
+    ggml_sycl_graph * graph          = nullptr;
+
     if (g_ggml_sycl_enable_graph) {
-        use_sycl_graph = check_graph_compatibility(cgraph);
+        if (!ggml_sycl_info().devices[sycl_ctx->device].graph_support) {
+            // the device does not change, so report this once
+            if (sycl_ctx->graph_compatible != false) {
+                GGML_LOG_DEBUG("%s: disabling SYCL graphs due to unsupported graph feature on device\n", __func__);
+                GGML_LOG_INFO("%s: disabling SYCL graphs\n", __func__);
+                sycl_ctx->graph_compatible = false;
+            }
+        } else {
+            graph = sycl_ctx->sycl_graph(ggml_sycl_graph_get_key(cgraph));
+
+            // The nodes of a cgraph do not change while its uid does not, so the compatibility scan
+            // runs again only when the uid changes. A uid of 0 means it is not set.
+            if (cgraph->uid == 0 || cgraph->uid != graph->compatible_uid) {
+                graph->compatible_uid = cgraph->uid;
+                graph->compatible     = check_graph_compatibility(sycl_ctx, cgraph);
+
+                if (sycl_ctx->graph_compatible != graph->compatible) {
+                    if (graph->compatible) {
+                        GGML_LOG_INFO("%s: enabling SYCL graphs\n", __func__);
+                    } else {
+                        GGML_LOG_INFO("%s: disabling SYCL graphs\n", __func__);
+                    }
+                    sycl_ctx->graph_compatible = graph->compatible;
+                }
+            }
+
+            use_sycl_graph = graph->compatible;
+        }
     }
     if (use_sycl_graph) {
-        const bool graph_support = dpct::get_device(sycl_ctx->device).has(sycl::aspect::ext_oneapi_limited_graph);
-        if (!graph_support) {
-            GGML_SYCL_DEBUG("[SYCL-GRAPH] can not use graphs on device:%d\n", sycl_ctx->device);
-            ggml_backend_sycl_graph_compute_impl(sycl_ctx, cgraph);
-            return GGML_STATUS_SUCCESS;
-        }
+        const bool properties_changed = ggml_sycl_graph_update_required(graph, cgraph);
+        bool replay_graph = false;
+        bool sycl_graph_update_required = false;
 
-        sycl_ex::command_graph model_sycl_graph(*(sycl_ctx->stream()), {sycl_ex::property::graph::assume_buffer_outlives_graph{}});
-
-        model_sycl_graph.begin_recording(*(sycl_ctx->stream()));
-        ggml_backend_sycl_graph_compute_impl(sycl_ctx, cgraph);
-        model_sycl_graph.end_recording();
-
-        const bool graph_update_support = dpct::get_device(sycl_ctx->device).has(sycl::aspect::ext_oneapi_graph);
-        if (!sycl_ctx->exec_graph || !graph_update_support) {
-            auto exec_graph = graph_update_support ? model_sycl_graph.finalize(sycl_ex::property::graph::updatable{}) :
-                                                     model_sycl_graph.finalize();
-            sycl_ctx->exec_graph = std::make_unique<
-                sycl_ex::command_graph<sycl_ex::graph_state::executable>>(exec_graph);
+        if (!graph->warmup_complete) {
+            // Warmup: need at least 2 calls with no property change on the 2nd call.
+            if (!properties_changed) {
+                graph->warmup_complete = true;
+                replay_graph = true;
+                sycl_graph_update_required = true;
+                GGML_SYCL_DEBUG("[SYCL-GRAPH] warmup complete\n");
+            }
         } else {
-            try {
-                sycl_ctx->exec_graph->update(model_sycl_graph);
-                GGML_SYCL_DEBUG("[SYCL-GRAPH] update success\n");
-            } catch (sycl::exception const & e) {
-                GGML_SYCL_DEBUG("[SYCL-GRAPH] Exception when updating graph, %s\n", e.what());
-                auto exec_graph = model_sycl_graph.finalize({sycl_ex::property::graph::updatable{}});
-                sycl_ctx->exec_graph = std::make_unique<
-                    sycl_ex::command_graph<sycl_ex::graph_state::executable>>(exec_graph);
+            // Post-warmup: replay graph while properties remain stable.
+            if (properties_changed) {
+                graph->warmup_complete = false;
+                GGML_SYCL_DEBUG("[SYCL-GRAPH] warmup reset\n");
+            } else {
+                replay_graph = true;
+                sycl_graph_update_required = graph->exec_graph == nullptr;
             }
         }
 
-        sycl_ctx->stream()->ext_oneapi_graph(*(sycl_ctx->exec_graph));
+        if (replay_graph) {
+            if (sycl_graph_update_required) {
+                const bool graph_update_support = ggml_sycl_info().devices[sycl_ctx->device].graph_update_support;
+
+                sycl_ex::command_graph model_sycl_graph(*(sycl_ctx->stream()), {sycl_ex::property::graph::assume_buffer_outlives_graph{}});
+
+                model_sycl_graph.begin_recording(*(sycl_ctx->stream()));
+                ggml_backend_sycl_graph_compute_impl(sycl_ctx, cgraph);
+                model_sycl_graph.end_recording();
+
+                if (!graph->exec_graph || !graph_update_support) {
+                    auto exec_graph = graph_update_support ? model_sycl_graph.finalize(sycl_ex::property::graph::updatable{}) :
+                                                             model_sycl_graph.finalize();
+                    graph->exec_graph = std::make_unique<
+                        sycl_ex::command_graph<sycl_ex::graph_state::executable>>(exec_graph);
+                } else {
+                    try {
+                        graph->exec_graph->update(model_sycl_graph);
+                        GGML_SYCL_DEBUG("[SYCL-GRAPH] update success\n");
+                    } catch (sycl::exception const & e) {
+                        // the executable graph cannot be updated due to violated constraints,
+                        // so finalize the recorded graph again
+                        GGML_SYCL_DEBUG("[SYCL-GRAPH] Exception when updating graph, %s\n", e.what());
+                        auto exec_graph = model_sycl_graph.finalize({sycl_ex::property::graph::updatable{}});
+                        graph->exec_graph = std::make_unique<
+                            sycl_ex::command_graph<sycl_ex::graph_state::executable>>(exec_graph);
+                    }
+                }
+            }
+
+            sycl_ctx->stream()->ext_oneapi_graph(*(graph->exec_graph));
+        } else {
+            ggml_backend_sycl_graph_compute_impl(sycl_ctx, cgraph);
+        }
     } else
 #endif
     {
@@ -7893,8 +8158,19 @@ ggml_backend_reg_t ggml_backend_sycl_reg() {
     return &reg;
 }
 
+static void set_device_active(int device) {
+    ggml_sycl_device_info & info = ggml_sycl_info();
+    info.devices[device].active = true;
+    int cnt = 0;
+    for (int i = 0; i < info.device_count; i++) {
+        if (info.devices[i].active)
+            cnt++;
+    }
+    info.active_device_count = cnt;
+}
+
 ggml_backend_t ggml_backend_sycl_init(int device) {
-    GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_init\n");
+    GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_init. device=%d\n", device);
     ggml_check_sycl();
 
     check_allow_gpu_index(device);
@@ -7905,6 +8181,7 @@ ggml_backend_t ggml_backend_sycl_init(int device) {
         return nullptr;
     };
 
+    set_device_active(device);
     ggml_backend_t sycl_backend = new ggml_backend {
         /* .guid    = */ ggml_backend_sycl_guid(),
         /* .iface   = */ ggml_backend_sycl_interface,

ggml/src/ggml-sycl/mmvq.cpp

@@ -3221,6 +3221,7 @@ bool ggml_sycl_mul_mat_vec_q_id(
     size_t             dst_row_stride,
     size_t             src1_row_stride,
     dpct::queue_ptr    stream) {
+    // if adding a new type, ensure it is also added to ggml_sycl_supports_mmvq_id
     switch (src0_type) {
         case GGML_TYPE_Q4_0:
             launch_mul_mat_vec_q_moe<QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>(

src/llama-vocab.cpp

@@ -422,6 +422,7 @@ struct llm_tokenizer_bpe : llm_tokenizer {
                 };
                 break;
             case LLAMA_VOCAB_PRE_TYPE_TEKKEN:
+            case LLAMA_VOCAB_PRE_TYPE_NEMOTRON_3_5_EMBD:
                 // original regex from tokenizer.json
                 // "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
                 regex_exprs = {
@@ -637,11 +638,22 @@ struct llm_tokenizer_bpe_session {
     }
 
     virtual void tokenize(const std::string & text, std::vector<llama_token> & output) {
+        auto tok_pre = vocab.get_pre_type();
+
+        std::vector<std::string> word_collection;
+        if (tok_pre == LLAMA_VOCAB_PRE_TYPE_NEMOTRON_3_5_EMBD) {
+            static const std::string zwnj = "\xe2\x80\x8c";
+            std::string normalized_text = text;
+            replace_all(normalized_text, "[" + zwnj, "[");
+            replace_all(normalized_text, "<" + zwnj, "<");
+            word_collection = unicode_regex_split(normalized_text, tokenizer.regex_exprs, tokenizer.byte_encode);
+        } else {
+            word_collection = unicode_regex_split(text, tokenizer.regex_exprs, tokenizer.byte_encode);
+        }
+
         int final_prev_index = -1;
-        auto word_collection = unicode_regex_split(text, tokenizer.regex_exprs, tokenizer.byte_encode);
 
         symbols_final.clear();
-        auto tok_pre = vocab.get_pre_type();
 
         if (tok_pre == LLAMA_VOCAB_PRE_TYPE_MMBERT) {
             // Metaspace pre-tokenizer: a text starts with an escaped space, and each escaped space starts a new word
@@ -2713,6 +2725,11 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 tokenizer_pre == "k2-horizon") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_K2_HORIZON;
                 clean_spaces = false;
+            } else if (
+                tokenizer_pre == "nemotron-3.5-embed") {
+                pre_type = LLAMA_VOCAB_PRE_TYPE_NEMOTRON_3_5_EMBD;
+                ignore_merges = true;
+                add_bos = true;
             } else {
                 throw std::runtime_error(format("unknown pre-tokenizer type: '%s'", tokenizer_pre.c_str()));
             }

src/llama-vocab.h

@@ -72,6 +72,7 @@ enum llama_vocab_pre_type {
     LLAMA_VOCAB_PRE_TYPE_K2_HORIZON        = 61,
     LLAMA_VOCAB_PRE_TYPE_FLINT             = 62,
     LLAMA_VOCAB_PRE_TYPE_DOMYN_SMALL       = 63,
+    LLAMA_VOCAB_PRE_TYPE_NEMOTRON_3_5_EMBD = 64,
 };
 
 // HF Metaspace pre-tokenizer prepend_scheme

tests/test-backend-ops.cpp

@@ -9971,7 +9971,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     // fused unary + mul (gated activations that are not expressed as GGML_OP_GLU)
-    for (ggml_unary_op op : { GGML_UNARY_OP_GELU, GGML_UNARY_OP_SILU, GGML_UNARY_OP_SIGMOID, GGML_UNARY_OP_SOFTPLUS }) {
+    for (ggml_unary_op op : { GGML_UNARY_OP_GELU, GGML_UNARY_OP_SILU, GGML_UNARY_OP_RELU, GGML_UNARY_OP_SIGMOID, GGML_UNARY_OP_SOFTPLUS }) {
         for (ggml_type type : { GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16 }) {
             for (bool swap : { false, true }) {
                 test_cases.emplace_back(new test_unary_mul(op, type, { 128, 2, 2, 2 }, swap));