From 960dffab0583556351d829308a0d8f533624dc0a Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Tue, 1 Sep 2026 23:15:21 -0700 Subject: [PATCH] hexagon: MUL_MAT and MUL_MAT_ID fusion and fixes (#28202) * hex-mm: fuse QKV and FFN matmuls that land on HMX * hex-mm: remove hardcoded ne[1] < 32K restriction * hex-get-rows: explicitly reject repacked Q8_0 just in case somebody decided to add an override * hex-mm: correct overhead sizing to make sure we dont exceed vtcm budget for large dims * hex-mm: fuse MUL_MAT_ID into MUL_MAT_ID_NX (2x,3x,...) where possible * hex-fusion: update opbatch and opqueue sizing to acount for new fusion and reduce overhead for trace buffer alloc * hex-bufs: sort buffers while finalizing opbatch, helps avoid va space fragmentation * hex-bufs: add simple va defrag to make sure we dont abort just because the va space is fragmented * hex-mm: replaced more scalar divs with fastdiv and minor cleanup * hex-mm: tighten up supported fusion checks to exactly match supported kernels --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 508 ++++++++++++--- ggml/src/ggml-hexagon/htp-opnode.h | 3 +- ggml/src/ggml-hexagon/htp/htp-ctx.h | 1 + ggml/src/ggml-hexagon/htp/htp-ops.h | 1 + ggml/src/ggml-hexagon/htp/main.c | 39 +- ggml/src/ggml-hexagon/htp/matmul-ops.c | 846 ++++++++++++++++++++++--- ggml/src/ggml-hexagon/htp/matmul-ops.h | 28 +- 7 files changed, 1241 insertions(+), 185 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 3eb84fd2a..04fb9a223 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -98,12 +98,26 @@ static int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 = // https://docs.qualcomm.com/doc/80-N2040-61/topic/hvx-pmu-events.html static u32vec opt_pmu_evt { 0x3, 0x111, 0x100, 0x105, 0x240, 0x256, 0x7D, 0x8C }; -static int opt_opbatch = 1024; // max number of ops in a batch -static int opt_opqueue = 64; // max number of pending batches +static int opt_opbatch = 1280; // max number of ops in a batch +static int opt_opqueue = 32; // max number of pending batches static int opt_optrace = 0; // trace buffer size per thread (0 means default) static int opt_oppoll = 0; // polling for batch completions static int opt_opfusion = 1; // enable/disable op fusion +enum ggml_hexagon_fusion_flags { + GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2 + GGML_HEXAGON_FUSE_RMS_NORM_MUL = (1 << 2), // 4 + GGML_HEXAGON_FUSE_MUL_MAT_ADD = (1 << 3), // 8 + GGML_HEXAGON_FUSE_MUL_MAT_NX = (1 << 4), // 16 + GGML_HEXAGON_FUSE_MUL_MAT_ID_NX = (1 << 5), // 32 +}; + +static inline bool ggml_hexagon_is_fusion_enabled(int flag) { + if (opt_opfusion <= 0) return false; + if (opt_opfusion == 1) return true; // 1 enables all + return (opt_opfusion & flag) != 0; +} + static std::regex* opt_opfilter = NULL; // regex of ops to not claim #define HEX_VERBOSE(...) \ @@ -293,6 +307,15 @@ static void ggml_hexagon_precompute_fused_mmnx_params( struct htp_mm_kernel_params * kparams ); +static void ggml_hexagon_precompute_fused_mmidnx_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + int32_t n_weights, + struct htp_mm_kernel_params * kparams +); + static bool ggml_hexagon_precompute_allreduce_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * dst, @@ -304,8 +327,12 @@ static bool ggml_hexagon_precompute_allreduce_params( ); static bool mm_is_hmx_eligible(const ggml_tensor * t); +static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams); +static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams); static bool is_mergeable_mul_mat(const ggml_tensor * t); static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2); +static bool is_mergeable_mul_mat_id(const ggml_tensor * t); +static bool is_mergeable_mul_mat_id_pair(const ggml_tensor * n1, const ggml_tensor * n2); // ** backend sessions @@ -1832,6 +1859,42 @@ struct ggml_hexagon_opbatch { } } + void sort_buffers() { + if (n_bufs <= 1) return; + + std::vector order(n_bufs); + for (unsigned int i = 0; i < n_bufs; i++) { order[i] = (int) i; } + + std::stable_sort(order.begin(), order.end(), [&](int a, int b) { + return h_bufs[a].size > h_bufs[b].size; + }); + + bool already_sorted = true; + for (unsigned int i = 0; i < n_bufs; i++) { + if (order[i] != (int) i) { + already_sorted = false; + break; + } + } + if (already_sorted) return; + + std::vector remap(n_bufs); + std::vector sorted_bufs(n_bufs); + for (unsigned int new_bi = 0; new_bi < n_bufs; new_bi++) { + int old_bi = order[new_bi]; + remap[old_bi] = (uint16_t) new_bi; + sorted_bufs[new_bi] = h_bufs[old_bi]; + } + + for (unsigned int i = 0; i < n_bufs; i++) { + h_bufs[i] = sorted_bufs[i]; + } + + for (unsigned int i = 0; i < n_tens; i++) { + h_tens[i].bi = remap[h_tens[i].bi]; + } + } + bool try_fuse_allreduce_add(const htp_opnode & node) { if (n_ops == 0 || opt_ar_select != 2) return false; if (node.opcode != HTP_OP_ADD) return false; @@ -2144,9 +2207,15 @@ struct ggml_hexagon_opbatch { if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0]) { return false; } + if (!last_node.fused.empty() && (mm_is_hmx_eligible(last_node.fused[0]) != mm_is_hmx_eligible(node.node))) { + return false; + } struct htp_mm_kernel_params kparams; ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams); + if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) { + return false; + } if ((size_t) kparams.vtcm_size > sess->vtcm_size) { HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", sess->c_name(), kparams.vtcm_size, sess->vtcm_size); @@ -2210,6 +2279,9 @@ struct ggml_hexagon_opbatch { struct htp_mm_kernel_params kparams; ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams); + if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) { + return false; + } if ((size_t) kparams.vtcm_size > sess->vtcm_size) { HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", sess->c_name(), kparams.vtcm_size, sess->vtcm_size); @@ -2272,18 +2344,172 @@ struct ggml_hexagon_opbatch { return false; } -enum ggml_hexagon_fusion_flags { - GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2 - GGML_HEXAGON_FUSE_RMS_NORM_MUL = (1 << 2), // 4 - GGML_HEXAGON_FUSE_MUL_MAT_ADD = (1 << 3), // 8 - GGML_HEXAGON_FUSE_MUL_MAT_NX = (1 << 4), // 16 -}; + bool try_fuse_mul_mat_id_nx(const htp_opnode & node) { + if (n_ops == 0 || node.opcode != HTP_OP_MUL_MAT_ID) return false; + if (!is_mergeable_mul_mat_id(node.node)) return false; -static inline bool ggml_hexagon_is_fusion_enabled(int flag) { - if (opt_opfusion <= 0) return false; - if (opt_opfusion == 1) return true; // 1 enables all - return (opt_opfusion & flag) != 0; -} + const ggml_tensor * w_in = node.src0(); + const ggml_tensor * x_in = node.src1(); + const ggml_tensor * ids_in = node.node->src[2]; + const ggml_tensor * d_in = node.dst(); + if (!w_in || !x_in || !ids_in || !d_in) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + + // Case 1: last_node is already MUL_MAT_ID_NX + if (last_node.opcode == HTP_OP_MUL_MAT_ID_NX) { + const uint32_t curr_n = (uint32_t) last_node.outputs.size(); + if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 2 >= HTP_OP_MAX_INPUTS) { + return false; + } + + const ggml_tensor * w0 = last_node.inputs[0]; + const ggml_tensor * x = last_node.inputs[curr_n]; + const ggml_tensor * ids = last_node.inputs[curr_n + 1]; + + 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]) { + return false; + } + if (!last_node.fused.empty() && (mm_is_hmx_eligible(last_node.fused[0]) != mm_is_hmx_eligible(node.node))) { + return false; + } + + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_mmidnx_params(sess, w0, x, d_in, curr_n + 1, &kparams); + if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) { + return false; + } + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip ID NX fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(w_in); + fit_t(d_in); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.inputs[curr_n] = w_in; + last_node.inputs[curr_n + 1] = x; + last_node.inputs.push_back(ids); + last_node.outputs.push_back(d_in); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + for (uint32_t s = 0; s <= curr_n + 2; s++) { + o.src[s] = add_tensor(last_node.inputs[s]); + } + for (uint32_t s = curr_n + 3; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + for (uint32_t d = 0; d <= curr_n; d++) { + o.dst[d] = add_tensor(last_node.outputs[d]); + } + for (uint32_t d = curr_n + 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_ID_NX (N=%u, #%u)\n", sess->c_name(), curr_n + 1, n_ops - 1); + return true; + } + + // Case 2: last_node is single MUL_MAT_ID + if (last_node.opcode == HTP_OP_MUL_MAT_ID) { + if (!is_mergeable_mul_mat_id_pair(last_node.node, node.node)) { + return false; + } + + const ggml_tensor * w0 = last_node.src0(); + const ggml_tensor * x = last_node.src1(); + const ggml_tensor * ids = last_node.node->src[2]; + const ggml_tensor * w1 = node.src0(); + if (!w0 || !x || !ids || !w1) return false; + + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_mmidnx_params(sess, w0, x, node.dst(), 2, &kparams); + if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) { + return false; + } + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip ID NX fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(w1); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + const ggml_tensor * dst_0 = last_node.dst(); + const ggml_tensor * dst_1 = node.dst(); + + last_node.opcode = HTP_OP_MUL_MAT_ID_NX; + last_node.name = "MUL_MAT_ID_NX"; + last_node.inputs.clear(); + last_node.inputs.push_back(w0); + last_node.inputs.push_back(w1); + last_node.inputs.push_back(x); + last_node.inputs.push_back(ids); + last_node.outputs.clear(); + last_node.outputs.push_back(dst_0); + last_node.outputs.push_back(dst_1); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_MUL_MAT_ID_NX; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + o.src[0] = add_tensor(w0); + o.src[1] = add_tensor(w1); + o.src[2] = add_tensor(x); + o.src[3] = add_tensor(ids); + for (uint32_t s = 4; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(dst_0); + o.dst[1] = add_tensor(dst_1); + for (uint32_t d = 2; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_ID_NX (N=2, #%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + return false; + } bool try_fuse(const htp_opnode & node) { if (!opt_opfusion) return false; @@ -2291,6 +2517,7 @@ static inline bool ggml_hexagon_is_fusion_enabled(int flag) { if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_RMS_NORM_MUL) && try_fuse_rms_norm_mul(node)) return true; if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ADD) && try_fuse_mul_mat_add(node)) return true; if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_NX) && try_fuse_mul_mat_nx(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ID_NX) && try_fuse_mul_mat_id_nx(node)) return true; return false; } }; @@ -2350,6 +2577,8 @@ struct ggml_hexagon_opqueue { delete shm_buf; } + size_t shm_size() const { return shm_buf ? shm_buf->size() : 0; } + // push new batch bool push(htp_opbatch_req& req, dspqueue_buffer& dbuf, ggml_hexagon_opbatch* op_batch) { static_assert(sizeof(htp_opbatch_req) % 8 == 0, "sizeof(htp_opbatch_req) must be multiple of 8"); @@ -2396,6 +2625,8 @@ struct ggml_hexagon_opqueue { uint8_t * t_ptr = m_ptr; m_ptr += t_size; uint8_t * o_ptr = m_ptr; + op_batch->sort_buffers(); + memcpy(b_ptr, (void *) op_batch->h_bufs.data(), b_size); memcpy(t_ptr, (void *) op_batch->h_tens.data(), t_size); memcpy(o_ptr, (void *) op_batch->h_ops.data(), o_size); @@ -3018,7 +3249,8 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n opt_vmem = ggml_hexagon_measure_max_vmem(this); GGML_LOG_INFO("ggml-hex: %s measured max vmem %zu\n", this->c_name(), opt_vmem); } - this->max_vmem = opt_vmem; + const size_t shm_size = this->op_queue->shm_size(); + this->max_vmem = (opt_vmem > shm_size) ? (opt_vmem - shm_size) : opt_vmem; this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem); @@ -3378,6 +3610,10 @@ static bool ggml_hexagon_matmul_is_hmx_eligible( bool is_matmul_id, bool is_batched ) { + if (src1->type != GGML_TYPE_F32) { + return false; + } + const int ne00 = src0->ne[0]; const int ne11 = src1->ne[1]; const int ne12 = src1->ne[2]; @@ -3408,7 +3644,8 @@ static bool ggml_hexagon_matmul_is_hmx_eligible( return false; } - // M alignment: Use HMX when M > HTP_MM_HMX_MIN_NROWS + // M alignment: Use HMX when M > HTP_MM_HMX_MIN_NROWS. + // For MUL_MAT_ID, src1 shape is [K, n_expert_used, n_tokens, 1], so n_tokens is ne12. const int m = is_matmul_id ? ne12 : ne11; if (m <= HTP_MM_HMX_MIN_NROWS) { return false; @@ -3460,7 +3697,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( if (!use_grouped) { // Fallback to simple 2D path (group_size = 1) - const int m_id_rows = (int) ((size_t) dst->ne[1] * dst->ne[2]); + const int m_id_rows = (dst && is_matmul_id) ? (int) ((size_t) dst->ne[1] * dst->ne[2]) : 0; if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { return false; } @@ -3918,64 +4155,113 @@ static void ggml_hexagon_precompute_fused_mmnx_params( ) { memset(kparams, 0, sizeof(*kparams)); - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); + const int ne00 = src0->ne[0]; + const int ne01 = src0->ne[1]; + const int ne02 = src0->ne[2]; + const int ne03 = src0->ne[3]; const int ne10 = src1->ne[0]; - const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - const size_t src0_row_size = src0->nb[1]; + const int ne11 = src1->ne[1]; + const int ne12 = src1->ne[2]; + const int ne13 = src1->ne[3]; - uint32_t best_n_prefetch = 16; + const int wtype = src0->type; + const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); + const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00; + const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01; + const int ne11_padded = hex_round_up(ne11, 32); - if (is_repack) { - const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; - best_n_prefetch = 2; - for (uint32_t d = max_prefetch; d >= 2; d /= 2) { - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, d, false, true - ); - if (L.total_bytes <= sess->vtcm_size) { - best_n_prefetch = d; - break; - } + const size_t vtcm_budget = sess->vtcm_size; + const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1); + + bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 3); + if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, nullptr, ne01_padded, false, is_batched)) { + if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, nullptr, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, false, is_batched, vtcm_budget, kparams)) { + kparams->n_weights = n_weights; + goto finalize; } } - struct htp_mm_hvx_vtcm_layout L; - bool try_tiled = (opt_mm_select >= 2); - - // Test tiled first - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true - ); - - if (try_tiled && L.total_bytes <= sess->vtcm_size) { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - kparams->n_weights = n_weights; - } else { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; - size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true - ); - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - kparams->n_weights = n_weights; + if (!is_repack) { + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } + + { + const int src1_nrows = ne11 * ne12 * ne13; + const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + const size_t src0_row_size = src0->nb[1]; + + uint32_t best_n_prefetch = 16; + + if (is_repack) { + const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; + best_n_prefetch = 2; + for (uint32_t d = max_prefetch; d >= 2; d /= 2) { + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, src1_row_size, 0, d, false, true + ); + if (L.total_bytes <= sess->vtcm_size) { + best_n_prefetch = d; + break; + } + } + } + + struct htp_mm_hvx_vtcm_layout L; + bool try_tiled = (opt_mm_select >= 2); + + // Test tiled first + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true + ); + + if (try_tiled && L.total_bytes <= sess->vtcm_size) { + kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; + kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; + } else { + kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; + size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); + + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, + 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true + ); + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_src1_size = L.src1_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; + kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; + } + } + +finalize: + kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11); + kparams->div_ne1 = init_fastdiv_values(ne11); + kparams->div_r2 = init_fastdiv_values(ne02 > 0 ? ne12 / ne02 : 1); + kparams->div_r3 = init_fastdiv_values(ne03 > 0 ? ne13 / ne03 : 1); + kparams->div_ne11 = init_fastdiv_values(ne11); +} + +static void ggml_hexagon_precompute_fused_mmidnx_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, // W0 + const struct ggml_tensor * src1, // x + const struct ggml_tensor * dst, // dst0 + int32_t n_weights, + struct htp_mm_kernel_params * kparams +) { + ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams); + kparams->n_weights = n_weights; } static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { @@ -4010,11 +4296,6 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s return false; } - // hardcoded limit to refuse the lm-head for now - if (src0->ne[1] > 32768) { - return false; - } - if (src1->ne[2] != 1 || src1->ne[3] != 1) { return false; // no broadcasting (for now) } @@ -4348,6 +4629,13 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * const struct ggml_tensor * src1 = op->src[1]; // indices const struct ggml_tensor * dst = op; + if (src0->extra) { + const auto * extra = (const ggml_hexagon_tensor_extra *) src0->extra; + if (extra->flags & GGML_HEXAGON_TENSOR_REPACK) { + return false; + } + } + if (src0->type != GGML_TYPE_F32 && src0->ne[0] < 32) { return false; } @@ -4734,10 +5022,43 @@ static bool mm_is_hmx_eligible(const ggml_tensor * t) { return ggml_hexagon_matmul_is_hmx_eligible(src0, src1, t, ne01_padded, is_matmul_id, is_batched); } +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; + } + + if (!ggml_hexagon_is_repack_type(src0->type)) { + return false; + } + + return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; +} + +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; + } + + if (!ggml_hexagon_is_repack_type(src0->type)) { + return false; + } + + return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK; +} + static bool is_mergeable_mul_mat(const ggml_tensor * t) { - if (!t || t->op != GGML_OP_MUL_MAT) return false; - if (t->src[1]->type != GGML_TYPE_F32) return false; - return ggml_is_quantized(t->src[0]->type) && !mm_is_hmx_eligible(t); + if (!t || t->op != GGML_OP_MUL_MAT) return false; + + const ggml_tensor * src0 = t->src[0]; + const ggml_tensor * src1 = t->src[1]; + if (src1->type != GGML_TYPE_F32) return false; + if (src0->ne[2] != 1 || src0->ne[3] != 1) return false; + + if (mm_is_hmx_eligible(t)) { + return ggml_hexagon_is_hmx_weight_type(src0->type); + } + + return ggml_hexagon_is_repack_type(src0->type); } static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) { @@ -4753,6 +5074,41 @@ static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor if (n1->src[0]->type != n2->src[0]->type) { return false; } + if (mm_is_hmx_eligible(n1) != mm_is_hmx_eligible(n2)) { + return false; + } + return true; +} + +static bool is_mergeable_mul_mat_id(const ggml_tensor * t) { + if (!t || t->op != GGML_OP_MUL_MAT_ID) return false; + + const ggml_tensor * src0 = t->src[0]; + return ggml_hexagon_is_repack_type(src0->type); +} + +static bool is_mergeable_mul_mat_id_pair(const ggml_tensor * n1, const ggml_tensor * n2) { + if (!is_mergeable_mul_mat_id(n1) || !is_mergeable_mul_mat_id(n2)) { + return false; + } + if (n1->src[1] != n2->src[1]) { + return false; + } + if (n1->src[2] != n2->src[2]) { + return false; + } + if (n1->src[0]->ne[0] != n2->src[0]->ne[0]) { + return false; + } + if (n1->src[0]->ne[2] != n2->src[0]->ne[2]) { + return false; + } + if (n1->src[0]->type != n2->src[0]->type) { + return false; + } + if (mm_is_hmx_eligible(n1) != mm_is_hmx_eligible(n2)) { + return false; + } return true; } @@ -4776,8 +5132,8 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg if (graph->nodes[i]->op == GGML_OP_RMS_NORM && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; - } else if (graph->nodes[i]->op == GGML_OP_MUL_MAT) { - if ((i + 1 < graph->n_nodes && graph->nodes[i + 1]->op == GGML_OP_ADD && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) || + } else if (graph->nodes[i]->op == GGML_OP_MUL_MAT || graph->nodes[i]->op == GGML_OP_MUL_MAT_ID) { + if ((i + 1 < graph->n_nodes && graph->nodes[i + 1]->op == GGML_OP_ADD && ggml_can_fuse(graph, i, { graph->nodes[i]->op, GGML_OP_ADD })) || ggml_node_has_n_uses(graph, i, 1)) { extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; } diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index 741b5e04e..b083e2671 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -315,7 +315,8 @@ struct htp_opformat { } 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_ADD) { + 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; diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index 88ecf144b..c8a909d61 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -118,6 +118,7 @@ struct htp_context { int op_matmul(struct htp_ops_context * octx); int op_matmul_id(struct htp_ops_context * octx); int op_matmul_nx(struct htp_ops_context * octx); +int op_matmul_id_nx(struct htp_ops_context * octx); int op_binary(struct htp_ops_context * octx); int op_unary(struct htp_ops_context * octx); int op_sum_rows(struct htp_ops_context * octx); diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index 53c95f28d..cf938f7ee 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -52,6 +52,7 @@ enum htp_op_code { HTP_OP_MUL_MAT, HTP_OP_MUL_MAT_ID, HTP_OP_MUL_MAT_NX, + HTP_OP_MUL_MAT_ID_NX, HTP_OP_MUL_MAT_ADD, HTP_OP_RMS_NORM, HTP_OP_RMS_NORM_MUL, diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 72cf02a32..3ab4613cf 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -753,6 +753,9 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_MUL_MAT_ID: return op_matmul_id(octx); + case HTP_OP_MUL_MAT_ID_NX: + return op_matmul_id_nx(octx); + case HTP_OP_MUL_MAT_NX: return op_matmul_nx(octx); @@ -878,8 +881,8 @@ static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) { } } -static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { - if (b->base) return; // already mapped +static inline bool mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { + if (b->base) return true; // already mapped // find unused mapping for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { @@ -887,8 +890,8 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { if (!m->size) { void *va = htp_mmap(b->fd, b->size); if (va == NULL) { - FARF(ERROR, "mmap failed : fd %u size %u", b->fd, (uint32_t) b->size); - abort(); // can't do much else at this point + FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %u", b->fd, (uint32_t) b->size); + return false; } m->base = b->base = (uint64_t) va; @@ -896,12 +899,12 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { m->size = b->size; FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); - return; + return true; } } FARF(ERROR, "mmap failed : exceeded mapping capacity limit of %u", HTP_MAX_MMAPS); - abort(); + return false; } static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t n_bufs) { @@ -934,12 +937,32 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin } } - // Create missing mappings + // Create missing mappings (pass 1) + bool mmap_ok = true; for (uint32_t i=0; i < n_bufs; i++) { struct htp_buf_desc *b = bufs + i; - mmap_buf(ctx, b); + if (!mmap_buf(ctx, b)) { + mmap_ok = false; + break; + } FARF(HIGH, "prep-buf #%u : pass1 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags); } + + if (!mmap_ok) { + // Attempt clean defragmentation: drop all mappings and remap (pass 2) + FARF(HIGH, "prep-bufs : dropping all mappings to defragment address space"); + for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { drop_mmap(ctx, ctx->mmap + i); } + + for (uint32_t i=0; i < n_bufs; i++) { + struct htp_buf_desc *b = bufs + i; + b->base = 0; + if (!mmap_buf(ctx, b)) { + FARF(ERROR, "prep-bufs : mmap failed after defragmentation (fd %u size %u)", b->fd, (uint32_t) b->size); + abort(); + } + FARF(HIGH, "prep-buf #%u : pass2 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags); + } + } } static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) { diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index a6adc0e61..2a87dd19e 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -55,10 +55,14 @@ typedef struct { size_t src0_nb3; size_t src1_nb2; size_t src1_nb3; - size_t dst_nb2; - size_t dst_nb3; size_t src2_nb2; size_t src2_nb3; + size_t dst_nb2; + size_t dst_nb3; + int r2; + int r3; + struct fastdiv_values div_r2; + struct fastdiv_values div_r3; } hmx_mm_f16_f32_batched_params_t; struct htp_mm_context { @@ -235,17 +239,18 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { const uint32_t nr1 = ne1 * ne2 * ne3; // distribute the thread work across the inner or outer loop based on which one is larger - uint32_t nchunk0 = nr0 > nr1 ? nth : 1; // parallelize by src0 rows - uint32_t nchunk1 = nr0 > nr1 ? 1 : nth; // parallelize by src1 rows - - // The number of elements in each chunk - const uint32_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; - const uint32_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; - - uint32_t current_chunk = ith; - - const uint32_t ith0 = current_chunk % nchunk0; - const uint32_t ith1 = current_chunk / nchunk0; + uint32_t dr0, dr1, ith0, ith1; + if (nr0 > nr1) { + dr0 = fastdiv(nr0 + nth - 1, &octx->ctx->n_threads_div); + dr1 = nr1; + ith0 = ith; + ith1 = 0; + } else { + dr0 = nr0; + dr1 = fastdiv(nr1 + nth - 1, &octx->ctx->n_threads_div); + ith0 = 0; + ith1 = ith; + } const uint32_t ir0_start = dr0 * ith0; const uint32_t ir0_end = MIN(ir0_start + dr0, nr0); @@ -545,7 +550,7 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ \ const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; \ - uint32_t src0_nrows_per_thread = (src0_nrows + nth - 1) / nth; \ + uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div); \ src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); \ \ const uint32_t start_row = src0_nrows_per_thread * ith; \ @@ -1105,6 +1110,179 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { } } +static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { + struct htp_mm_context * mmctx = (struct htp_mm_context *) data; + struct htp_ops_context * octx = mmctx->octx; + dma_queue * dma_queue = octx->ctx->dma[ith]; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict act = octx->src[n_weights]; + const struct htp_tensor * restrict ids = octx->src[n_weights + 1]; + + hvx_mm_run_quant_task(mmctx, ith); + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + + const uint32_t n_aids = ids->ne[0]; + const uint32_t n_ids = src0->ne[2]; + + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * restrict src1_data = mmctx->vtcm_src1; + + for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) { + const int32_t eid = *(const int32_t *) ((const uint8_t *) ids->data + ie1 * ids->nb[0]); + if (eid < 0) continue; + assert(eid < (int32_t) n_ids); + + for (uint32_t p = 0; p < n_weights; ++p) { + const struct htp_tensor * restrict src_w = octx->src[p]; + const struct htp_tensor * restrict dst = octx->dsts[p]; + if (!src_w || !dst) continue; + + const uint32_t src0_nrows = src_w->ne[1]; + uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div); + src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); + + const uint32_t src0_start_row = src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); + if (src0_start_row >= src0_end_row) continue; + + const uint8_t * restrict src0_row = (const uint8_t *) src_w->data + eid * src_w->nb[2]; + const uint8_t * restrict src1_col = (const uint8_t *) src1_data; + float * restrict dst_row = (float *) (dst->data + ie1 * dst->nb[1]); + + const uint32_t tile_size = htp_mm_get_weight_tile_size(src_w->type); + const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src_w->type); + const uint32_t n_k_tiles_w = src_w->ne[0] / 32; + const uint32_t n_k_tiles_a = act->ne[0] / 32; + const uint32_t tile_row_stride = n_k_tiles_w * tile_size; + const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; + + const uint32_t ct_start = src0_start_row / 32; + const uint32_t ct_end = (src0_end_row + 31) / 32; + + uint32_t push_ct = ct_start; + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { + dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + } + + for (uint32_t ct = ct_start; ct < ct_end; ct++) { + const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; + + int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); + valid_rows = MIN(32, MAX(0, valid_rows)); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); + mmctx->vec_dot_32x1(act->ne[0], &dst_row[ct * 32], w_tile, src1_col, valid_rows, NULL); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); + + if (push_ct < ct_end) { + dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + push_ct++; + } + } + } + } +} + +static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { + struct htp_mm_context * mmctx = (struct htp_mm_context *) data; + struct htp_ops_context * octx = mmctx->octx; + dma_queue * dma_queue = octx->ctx->dma[ith]; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict act = octx->src[n_weights]; + const struct htp_tensor * restrict ids = octx->src[n_weights + 1]; + + hvx_mm_run_quant_task(mmctx, ith); + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + + const uint32_t n_as = src0->ne[2]; + + const uint32_t * matrix_row_counts = mmctx->matrix_row_counts; + const struct mmid_row_mapping * matrix_rows = mmctx->matrix_rows; + + const size_t src1_stride = mmctx->vtcm_src1_stride; + + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * restrict src1_data = mmctx->vtcm_src1; + + for (uint32_t cur_a = 0; cur_a < n_as; ++cur_a) { + const int32_t cne1 = matrix_row_counts[cur_a]; + if (cne1 == 0) continue; + + for (uint32_t p = 0; p < n_weights; ++p) { + const struct htp_tensor * restrict src_w = octx->src[p]; + const struct htp_tensor * restrict dst = octx->dsts[p]; + if (!src_w || !dst) continue; + + const uint32_t src0_nrows = src_w->ne[1]; + uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div); + src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); + + const uint32_t src0_start_row = src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); + if (src0_start_row >= src0_end_row) continue; + + const uint8_t * src0_row = (const uint8_t *) src_w->data + cur_a * src_w->nb[2]; + + const uint32_t tile_size = htp_mm_get_weight_tile_size(src_w->type); + const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src_w->type); + const uint32_t n_k_tiles_w = src_w->ne[0] / 32; + const uint32_t n_k_tiles_a = act->ne[0] / 32; + const uint32_t tile_row_stride = n_k_tiles_w * tile_size; + const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; + + const uint32_t ct_start = src0_start_row / 32; + const uint32_t ct_end = (src0_end_row + 31) / 32; + + uint32_t push_ct = ct_start; + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { + dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + } + + for (uint32_t ct = ct_start; ct < ct_end; ct++) { + const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; + + int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); + valid_rows = MIN(32, MAX(0, valid_rows)); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); + for (uint32_t cid = 0; cid < (uint32_t) cne1; ++cid) { + struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, cid); + const int rm1 = row_mapping.i1; + const int rm2 = row_mapping.i2; + + const uint32_t ir1 = fastmodulo(rm1, act->ne[1], &mmctx->mm_div_ne11); + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (ir1 + rm2 * act->ne[1]) * src1_stride); + float * restrict dst_row = (float *) (dst->data + (rm1 * dst->nb[1] + rm2 * dst->nb[2])); + + mmctx->vec_dot_32x1(act->ne[0], &dst_row[ct * 32], w_tile, src1_col, valid_rows, NULL); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); + + if (push_ct < ct_end) { + dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + push_ct++; + } + } + } + } +} + static int hvx_mm_init_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type type) { switch (type) { case HTP_TYPE_Q4_0: @@ -1153,7 +1331,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { src0->type == HTP_TYPE_MXFP4); // Compute src0_nrows_per_thread - mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads; + mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div); if (is_repacked) { mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); } else { @@ -1325,11 +1503,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { mmctx->vtcm_src1_size_per_thread = L.src1_bytes; } else { - mmctx->vtcm_src1_size_per_thread = L.src1_bytes / octx->n_threads; + mmctx->vtcm_src1_size_per_thread = fastdiv(L.src1_bytes, &octx->ctx->n_threads_div); } - mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; - mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; + mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div); + mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div); size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; @@ -1407,7 +1585,7 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { const uint32_t ne01 = src_w->ne[1]; const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; - uint32_t src0_nrows_per_thread = (src0_nrows + nth - 1) / nth; + uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div); src0_nrows_per_thread += (src0_nrows_per_thread & 1); const uint32_t src0_start_row = src0_nrows_per_thread * ith; @@ -1538,36 +1716,36 @@ static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void } typedef struct { - const struct mmid_row_mapping *matrix_rows; - __fp16 *dst; - const float *src; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t k_block; - uint32_t k_stride; - uint32_t k_valid; - struct htp_thread_trace * traces; - struct htp_context * ctx; - float * vtcm_f32_act; - size_t vtcm_f32_act_bytes_per_thread; - uint32_t dma_step_rows; - uint32_t dma_step_rows_shift; + struct htp_context * ctx; + struct htp_thread_trace * traces; + __fp16 * dst; + const float * src; + const struct mmid_row_mapping * matrix_rows; + float * vtcm_f32_act; + uint32_t n_tasks; + uint32_t n_tot_chunks; + uint32_t n_chunks_per_task; + uint32_t k_block; + uint32_t k_stride; + uint32_t k_valid; + size_t vtcm_f32_act_bytes_per_thread; + uint32_t dma_step_rows; + uint32_t dma_step_rows_shift; } activation_transfer_task_state_t; typedef struct { - __fp16 *dst; - const float *src; + struct htp_context * ctx; + struct htp_thread_trace * traces; + __fp16 * dst; + const float * src; + float * vtcm_f32_act; uint32_t n_rows; uint32_t k_block; uint32_t k_stride; uint32_t k_valid; uint32_t n_col_chunks; struct fastdiv_values n_threads_div; - float *vtcm_f32_act; size_t vtcm_f32_act_bytes; - struct htp_thread_trace *traces; - struct htp_context *ctx; uint32_t dma_step_rows; uint32_t dma_step_rows_shift; } activation_transfer_col_chunk_state_t; @@ -1811,9 +1989,10 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, } typedef struct { - const struct mmid_row_mapping *matrix_rows; - __fp16 *dst; - const float *src; + struct htp_thread_trace * traces; + const struct mmid_row_mapping * matrix_rows; + __fp16 * dst; + const float * src; uint32_t n_tasks; uint32_t n_tot_chunks; uint32_t n_chunks_per_task; @@ -1827,13 +2006,13 @@ typedef struct { uint32_t start_row; uint32_t cne1; uint32_t k_valid; - struct htp_thread_trace *traces; } activation_transfer_gathered_task_state_t; typedef struct { - const struct mmid_row_mapping *matrix_rows; - const __fp16 *vtcm_src; - float *dst; + struct htp_thread_trace * traces; + const struct mmid_row_mapping * matrix_rows; + const __fp16 * vtcm_src; + float * dst; uint32_t n_tasks; uint32_t n_tot_chunks; uint32_t n_chunks_per_task; @@ -1844,17 +2023,16 @@ typedef struct { size_t dst_nb2; uint32_t start_row; uint32_t cne1; - struct htp_thread_trace *traces; } output_transfer_scattered_task_state_t; static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_gathered_task_state_t *st = data; struct htp_thread_trace * tr = &st->traces[i]; - int chunk_idx = i; - int chunk_size = st->n_chunks_per_task; + int chunk_idx = i; + int chunk_size = st->n_chunks_per_task; int vtcm_start_row = chunk_idx * chunk_size; - int start_row = st->start_row + vtcm_start_row; - int n_rows = hex_smin(st->cne1 - start_row, chunk_size); + int start_row = st->start_row + vtcm_start_row; + int n_rows = hex_smin(st->cne1 - start_row, chunk_size); if (n_rows > 0) { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); transfer_activation_chunk_fp32_to_fp16_gathered( @@ -1946,17 +2124,17 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles( } typedef struct { - float *dst; - const float *src2; - const __fp16 *vtcm_src; - uint32_t n_rows; - uint32_t n_cols; - uint32_t dst_stride; - uint32_t src2_stride; - uint32_t dst_cols; - struct fastdiv_values n_threads_div; - struct htp_thread_trace *traces; - struct htp_context *ctx; + struct htp_context * ctx; + struct htp_thread_trace * traces; + float * dst; + const __fp16 * vtcm_src; + const float * src2; + uint32_t n_rows; + uint32_t n_cols; + uint32_t dst_stride; + uint32_t src2_stride; + uint32_t dst_cols; + struct fastdiv_values n_threads_div; } output_transfer_col_chunk_state_t; static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { @@ -1965,19 +2143,19 @@ static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned i struct htp_thread_trace * tr = &st->traces[i]; uint32_t n_blocks = st->n_cols / 32; - uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div); - uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div); - uint32_t c_first = b_first * 32; - uint32_t c_last = b_last * 32; - uint32_t c_len = c_last - c_first; + uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div); + uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div); + uint32_t c_first = b_first * 32; + uint32_t c_last = b_last * 32; + uint32_t c_len = c_last - c_first; if (c_len == 0) return; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first); - float *dst = st->dst + c_first; - const float *src2 = st->src2 ? (st->src2 + c_first) : NULL; const __fp16 *vtcm_src = st->vtcm_src + b_first * HTP_MM_HMX_TILE_N_ELMS; + const float *src2 = st->src2 ? (st->src2 + c_first) : NULL; + float *dst = st->dst + c_first; int chunk_dst_cols = (int)st->dst_cols - (int)c_first; if (chunk_dst_cols > 0) { @@ -1998,7 +2176,7 @@ static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, uint32_t n_blocks = (uint32_t)n_cols / 32; if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) { - struct fastdiv_values n_threads_div = init_fastdiv_values(n_threads); + struct fastdiv_values n_threads_div = (n_threads == (int)ctx->n_threads) ? ctx->n_threads_div : init_fastdiv_values(n_threads); output_transfer_col_chunk_state_t col_state; col_state.dst = dst; col_state.src2 = src2; @@ -2128,8 +2306,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_ state.ctx = ctx; state.vtcm_f32_act = vtcm_f32_act; - int active_threads = hex_smin(n_threads, (int)state.n_tasks); - state.vtcm_f32_act_bytes_per_thread = hex_align_down(vtcm_f32_act_bytes / active_threads, 128); + state.vtcm_f32_act_bytes_per_thread = hex_align_down(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128); uint32_t dma_step_rows = 2; uint32_t dma_step_rows_shift = 1; @@ -2144,6 +2321,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_ state.dma_step_rows = dma_step_rows; state.dma_step_rows_shift = dma_step_rows_shift; + int active_threads = hex_smin(n_threads, (int)state.n_tasks); if (state.n_tasks == 1 || n_threads == 1) { transfer_activation_chunk_worker_fn(1, 0, &state); } else { @@ -2447,21 +2625,286 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, return 0; } -static inline int hmx_mm_batch_r2(const hmx_mm_f16_f32_batched_params_t *params) { - return params->ne02 > 0 ? params->ne12 / params->ne02 : 1; -} +static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) { + struct htp_context * ctx = octx->ctx; + struct htp_thread_trace * tr = &ctx->trace[0]; + htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); -static inline int hmx_mm_batch_r3(const hmx_mm_f16_f32_batched_params_t *params) { - return params->ne03 > 0 ? params->ne13 / params->ne03 : 1; + const uint32_t n_weights = kparams->n_weights; + if (n_weights == 0 || n_weights > HTP_OP_MAX_OUTPUTS) { + return HTP_STATUS_INVAL_PARAMS; + } + + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict act = octx->src[n_weights]; + + if (!src0 || !act) { + return HTP_STATUS_INVAL_PARAMS; + } + + const int weight_type = (int) src0->type; + const int k = (int) act->ne[0]; + const int k_valid = (int) act->ne[0]; + const int m = (int) (act->ne[1] * act->ne[2] * act->ne[3]); + const int act_stride = (int) (act->nb[1] / sizeof(float)); + const float * activation = (const float *) act->data; + + if (k % 32 != 0) { return HTP_STATUS_NO_SUPPORT; } + if (!hex_is_aligned(activation, VLEN)) { return HTP_STATUS_NO_SUPPORT; } + + size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); + if (row_stride == 0) { + return HTP_STATUS_NO_SUPPORT; + } + + worker_callback_t dequant_worker_fn = NULL; + switch (weight_type) { + case HTP_TYPE_Q4_0: dequant_worker_fn = dequantize_tiled_worker_loop_q4_0; break; + case HTP_TYPE_IQ4_NL: dequant_worker_fn = dequantize_tiled_worker_loop_iq4_nl; break; + case HTP_TYPE_Q4_1: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; + case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; + case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; + case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; + case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; + default: + return HTP_STATUS_NO_SUPPORT; + } + + const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; + const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles); + + const bool is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32); + const size_t vtcm_budget = ctx->vtcm_size; + + const int m_chunk_n_rows = kparams->m_chunk; + const int n_chunk_n_cols = kparams->n_chunk; + const int pipeline = kparams->pipeline; + const int n_threads = octx->n_threads; + const int act_threads = kparams->n_act_threads; + const struct fastdiv_values * act_threads_div = &kparams->div_n_act_threads; + const struct fastdiv_values * k_div = &kparams->div_ne00_padded; + const int tile_size = kparams->tile_size; + const int aligned_tile_size = kparams->aligned_tile_size; + + const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride; + const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; + + struct htp_mm_hmx_vtcm_layout L; + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size); + + if (L.total_bytes > vtcm_budget) { + FARF(ERROR, "hmx-mm-nx-2d: VTCM overflow: used %zu budget %zu, m %d k %d mc %d nc %d", + L.total_bytes, vtcm_budget, m, k, m_chunk_n_rows, n_chunk_n_cols); + return HTP_STATUS_VTCM_TOO_SMALL; + } + + uint8_t * const base = (uint8_t *) ctx->vtcm_base; + __fp16 *vtcm_weight_raw[2] = { + VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]), + VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_weight[1], pipeline) + }; + + __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act); + float *vtcm_f32_act = VTCM_LAYOUT_PTR(float, base, L.off_act_f32); + __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]); + void *vtcm_scratch0 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]); + void *vtcm_scratch1 = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_scratch[1], pipeline); + void *vtcm_scratch2 = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_dst[1], pipeline); + __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales); + + hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 + + FARF(HIGH, "hmx-mm-nx-2d: n_weights %u m %d k %d wtype %d mc %d nc %d vtcm %zu/%zu", + n_weights, m, k, weight_type, m_chunk_n_rows, n_chunk_n_cols, L.total_bytes, vtcm_budget); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); + + if (pipeline) { + hmx_matmul_job_t job_slots[2]; + + for (size_t mr = 0; mr < (size_t) m; mr += m_chunk_n_rows) { + const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); + + void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 }; + void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 }; + + struct activation_transfer_params act_params = { + .ctx = ctx, + .dst = vtcm_f16_act, + .src = activation + mr * act_stride, + .n_rows = (int) n_rows, + .k_block = k, + .k_stride = act_stride, + .n_threads = act_threads, + .act_threads_div = act_threads_div, + .k_div = k_div, + .k_valid = k_valid, + .vtcm_f32_act = vtcm_f32_act, + .vtcm_f32_act_bytes = L.act_f32_bytes, + }; + transfer_activation_chunk_threaded(&act_params); + + for (uint32_t p = 0; p < n_weights; p++) { + const struct htp_tensor * restrict src_w = octx->src[p]; + const struct htp_tensor * restrict dst = octx->dsts[p]; + if (!src_w || !dst) continue; + + const uint8_t * weight = (const uint8_t *) src_w->data; + float * dst_ptr = (float *) dst->data; + const size_t n = src_w->ne[1]; + if (n == 0) continue; + const size_t weight_stride = src_w->nb[1]; + const size_t dst_stride = dst->nb[1] / sizeof(float); + const int dst_cols = (int) dst->ne[0]; + const int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols); + + const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; + + const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); + const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0; + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), + dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0); + + if (1 < n_chunk_cnt) { + const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols); + const uint32_t height_A1 = is_quant ? (n_cols_A1 / 32) * n_k_tiles : n_cols_A1; + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), + dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1); + } + + for (int i = 0; i < n_chunk_cnt; ++i) { + const size_t nc = i * n_chunk_n_cols; + const size_t nc_p2 = nc + 2 * n_chunk_n_cols; + + const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); + const size_t n_cols_p2 = hex_smin(n - nc_p2, n_chunk_n_cols); + + void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + + dequantize_tiled_weight_chunk_to_fp16_tiles( + ctx, vtcm_weight_bufs[i % 2], curr_raw, + n_cols, k, row_stride, weight_type, + n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); + + if (i + 2 < n_chunk_cnt) { + const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2; + dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_p2 * weight_stride), + dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2); + } + + hmx_matmul_job_init(&job_slots[i % 2], (__fp16 *) vtcm_output_bufs[i % 2], + (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], + vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS), + hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS); + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[i % 2])); + + if (i > 0) { + hmx_queue_pop(ctx->hmx_queue); + const size_t nc_prev = (i - 1) * n_chunk_n_cols; + const size_t n_cols_prev = hex_smin(n - nc_prev, n_chunk_n_cols); + float *output_chunk = dst_ptr + (mr * dst_stride + nc_prev); + int chunk_dst_cols = dst_cols - (int)nc_prev; + if (chunk_dst_cols > 0) { + transfer_output_chunk_threaded(ctx, output_chunk, NULL, vtcm_output_bufs[(i - 1) % 2], n_rows, n_cols_prev, dst_stride, 0, chunk_dst_cols, n_threads); + } + } + } + + hmx_queue_pop(ctx->hmx_queue); + const size_t nc_last = (n_chunk_cnt - 1) * n_chunk_n_cols; + const size_t n_cols_last = hex_smin(n - nc_last, n_chunk_n_cols); + float *output_chunk = dst_ptr + (mr * dst_stride + nc_last); + int chunk_dst_cols = dst_cols - (int)nc_last; + if (chunk_dst_cols > 0) { + transfer_output_chunk_threaded(ctx, output_chunk, NULL, vtcm_output_bufs[(n_chunk_cnt - 1) % 2], n_rows, n_cols_last, dst_stride, 0, chunk_dst_cols, n_threads); + } + } + } + } else { + hmx_matmul_job_t job; + for (size_t mr = 0; mr < (size_t) m; mr += m_chunk_n_rows) { + const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); + + struct activation_transfer_params act_params = { + .ctx = ctx, + .dst = vtcm_f16_act, + .src = activation + mr * act_stride, + .n_rows = (int) n_rows, + .k_block = k, + .k_stride = act_stride, + .n_threads = act_threads, + .act_threads_div = act_threads_div, + .k_div = k_div, + .k_valid = k_valid, + .vtcm_f32_act = vtcm_f32_act, + .vtcm_f32_act_bytes = L.act_f32_bytes, + }; + transfer_activation_chunk_threaded(&act_params); + + for (uint32_t p = 0; p < n_weights; p++) { + const struct htp_tensor * restrict src_w = octx->src[p]; + const struct htp_tensor * restrict dst = octx->dsts[p]; + if (!src_w || !dst) continue; + + const uint8_t * weight = (const uint8_t *) src_w->data; + float * dst_ptr = (float *) dst->data; + const size_t n = src_w->ne[1]; + if (n == 0) continue; + const size_t weight_stride = src_w->nb[1]; + const size_t dst_stride = dst->nb[1] / sizeof(float); + const int dst_cols = (int) dst->ne[0]; + + const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; + + if (n > 0) { + const size_t n_cols = hex_smin(n, n_chunk_n_cols); + const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols; + dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); + } + + for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { + const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); + const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); + const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); + + void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; + + dequantize_tiled_weight_chunk_to_fp16_tiles( + ctx, vtcm_scratch0, curr_raw, + n_cols, k, row_stride, weight_type, + n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); + + const size_t nc_next = nc + n_chunk_n_cols; + if (nc_next < n) { + const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); + const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; + dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + } + + hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); + hmx_queue_pop(ctx->hmx_queue); + + float *output_chunk = dst_ptr + (mr * dst_stride + nc); + int chunk_dst_cols = dst_cols - (int)nc; + if (chunk_dst_cols > 0) { + transfer_output_chunk_threaded(ctx, output_chunk, NULL, vtcm_output, n_rows, n_cols, dst_stride, 0, chunk_dst_cols, n_threads); + } + } + } + } + } + + return HTP_STATUS_OK; } static inline const __fp16 *hmx_mm_weight_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, int dst_b2, int dst_b3) { - const int r2 = hmx_mm_batch_r2(params); - const int r3 = hmx_mm_batch_r3(params); + const size_t b2_idx = (params->r2 <= 1) ? (size_t) dst_b2 : (size_t) fastdiv((uint32_t) dst_b2, ¶ms->div_r2); + const size_t b3_idx = (params->r3 <= 1) ? (size_t) dst_b3 : (size_t) fastdiv((uint32_t) dst_b3, ¶ms->div_r3); return (const __fp16 *) ((const uint8_t *) params->weight + - (size_t) (dst_b2 / r2) * params->src0_nb2 + - (size_t) (dst_b3 / r3) * params->src0_nb3); + b2_idx * params->src0_nb2 + + b3_idx * params->src0_nb3); } static inline const float *hmx_mm_activation_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, @@ -2517,7 +2960,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ if (params->k % 32 != 0 || params->n % 32 != 0) { return -1; } if (!hex_is_aligned(params->dst, VLEN) || !hex_is_aligned(params->activation, VLEN)) { return -1; } - const int group_size = hmx_mm_batch_r2(params); + const int group_size = params->r2; const size_t vtcm_budget = ctx->vtcm_size; // Check if the precomputed parameters are grouped or simple. @@ -2825,8 +3268,9 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, htp_mm_hmx_get_2d_chunk_costs(weight_type, k, /*pipeline=*/false, aligned_tile_size, &size_per_n, &size_per_m, &size_per_mn); + const size_t overhead = htp_mm_hmx_get_2d_overhead(/*pipeline=*/false, /*is_matmul_id=*/true); size_t m_chunk_n_rows = 0, n_chunk_n_cols = 0; - if (htp_mm_hmx_compute_chunks(vtcm_budget, /*overhead=*/256, size_per_n, size_per_m, size_per_mn, + if (htp_mm_hmx_compute_chunks(vtcm_budget, overhead, size_per_n, size_per_m, size_per_mn, m_padded, n, /*m_block_cost=*/(size_t) n * HTP_MM_HMX_COST_W_DEQUANT, /*n_block_cost=*/(size_t) m_padded * HTP_MM_HMX_COST_A_CONVERT, &m_chunk_n_rows, &n_chunk_n_cols, &vtcm_used)) { @@ -2962,6 +3406,10 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k .dst_nb3 = dst->nb[3], .src2_nb2 = src2_nb2, .src2_nb3 = src2_nb3, + .r2 = (ne02 > 0) ? (ne12 / ne02) : 1, + .r3 = (ne03 > 0) ? (ne13 / ne03) : 1, + .div_r2 = kparams->div_r2, + .div_r3 = kparams->div_r3, }; ret = hmx_mm_f16_f32_batched(octx->ctx, &batch_params, kparams->m_chunk, kparams->n_chunk, @@ -3106,10 +3554,10 @@ static int hvx_mm_matmul_id( mmctx->vtcm_src0_stride = src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; - mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; + mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_src2_size_per_thread = 0; - mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; + mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div); mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; @@ -3123,6 +3571,134 @@ static int hvx_mm_matmul_id( return HTP_STATUS_OK; } +static int hmx_mm_op_matmul_id_nx( + struct htp_ops_context * octx, + struct htp_mm_context * mmctx +) { + const uint32_t * matrix_row_counts = mmctx->matrix_row_counts; + const struct mmid_row_mapping * matrix_rows = mmctx->matrix_rows; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict act = octx->src[n_weights]; + const int n_as = src0->ne[2]; + + for (uint32_t cur_a = 0; cur_a < (uint32_t) n_as; ++cur_a) { + const int32_t cne1 = matrix_row_counts[cur_a]; + if (cne1 == 0) continue; + + for (uint32_t p = 0; p < n_weights; ++p) { + const struct htp_tensor * restrict src_w = octx->src[p]; + const struct htp_tensor * restrict dst = octx->dsts[p]; + if (!src_w || !dst) continue; + + int ret = hmx_mm_id_2d_f32(octx->ctx, (float*) dst->data, (float*) act->data, + (const uint8_t *) src_w->data + cur_a * src_w->nb[2], + cne1, src_w->ne[0], src_w->ne[1], + act->ne[0], + act->ne[1], + act->nb[1], act->nb[2], + dst->nb[1], dst->nb[2], + (int) src_w->nb[1], (int) src_w->type, + matrix_rows, cur_a, mmctx->mapping_stride); + if (ret != 0) { + FARF(ERROR, "HMX matmul ID NX failed for expert %u weight %u, error %d\n", cur_a, p, ret); + return HTP_STATUS_NO_SUPPORT; + } + } + } + + return HTP_STATUS_OK; +} + +static int hvx_mm_matmul_id_nx( + struct htp_ops_context * octx, + struct htp_mm_context * mmctx, + work_queue_func_t hvx_mmid_task_func +) { + const uint32_t src0_row_size_padded = mmctx->src0_row_size_padded; + const uint32_t src1_nrows = mmctx->src1_nrows; + + struct htp_thread_trace * tr = &octx->ctx->trace[0]; + htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); + + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict act = octx->src[n_weights]; + 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 nb = (act->ne[0] + qk - 1) / qk; + const uint32_t total_nb = src1_nrows * nb; + + work_queue_func_t quant_task_func; + uint32_t n_quant_tasks = 1; + if (src1_nrows < octx->n_threads) { + n_quant_tasks = MIN(total_nb, octx->n_threads); + quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; + for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { + uint32_t ib_first = (total_nb * ith) / n_quant_tasks; + uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; + mmctx->quant_ib_first[ith] = ib_first; + mmctx->quant_ib_last[ith] = ib_last; + mmctx->quant_r[ith] = ib_first / nb; + mmctx->quant_c[ith] = ib_first % nb; + } + } else { + n_quant_tasks = MIN(src1_nrows, octx->n_threads); + quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; + } + size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); + + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads, + 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false); + + size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; + + if (octx->ctx->vtcm_size < vtcm_size) { + FARF(ERROR, "matmul-id-nx: current VTCM reservation %zu is too small, needed %zu\n", + octx->ctx->vtcm_size, vtcm_size); + return HTP_STATUS_VTCM_TOO_SMALL; + } + + uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; + mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); + mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); + mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + + 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; + + mmctx->vtcm_src0_stride = 0; + mmctx->vtcm_src1_stride = src1_row_size; + + mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div); + mmctx->vtcm_src1_size_per_thread = L.src1_bytes; + mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div); + + mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; + mmctx->quant_task_func = quant_task_func; + mmctx->n_quant_tasks = n_quant_tasks; + atomic_init(&mmctx->quant_barrier, n_quant_tasks); + + FARF(HIGH, "matmul-id-nx: src0 %d:%d:%d type %s nrows %u, src1 %d:%d:%d nrows %u, vtcm %zu/%zu, threads %d\n", + src0->ne[0], src0->ne[1], src0->ne[2], mmctx->type, src0->ne[1], + act->ne[0], act->ne[1], act->ne[2], src1_nrows, + L.total_bytes, octx->ctx->vtcm_size, octx->n_threads); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); + + worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads); + + return HTP_STATUS_OK; +} + static inline void scan_expert_ids_n( const struct htp_tensor * ids, const uint32_t n_ids, @@ -3213,7 +3789,7 @@ int op_matmul_id(struct htp_ops_context * octx) { const uint32_t src0_nrows = ne01; // per expert const uint32_t src1_nrows = ne11 * ne12 * ne13; - mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads; + mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div); mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); // row groups @@ -3280,11 +3856,103 @@ int op_matmul_id(struct htp_ops_context * octx) { return s; } -int op_matmul_nx(struct htp_ops_context * octx) { + +int op_matmul_id_nx(struct htp_ops_context * octx) { struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + const struct htp_tensor * restrict src0 = octx->src[0]; + const struct htp_tensor * restrict act = octx->src[n_weights]; + const struct htp_tensor * restrict ids = octx->src[n_weights + 1]; + + struct htp_mm_context mmctx_struct = {0}; + struct htp_mm_context * mmctx = &mmctx_struct; + mmctx->octx = octx; + mmctx->act = act; + + const size_t src0_row_size = src0->nb[1]; + const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); + + const uint32_t src0_nrows = src0->ne[1]; + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; + + mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div); + mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); + + const int n_ids = ids->ne[0]; + const int n_as = src0->ne[2]; + + uint8_t * mapping_buf = octx->ctx->ddr_spad_base; + uint32_t mapping_stride = 1; + uint32_t * matrix_row_counts = (uint32_t *) mapping_buf; + struct mmid_row_mapping * matrix_rows = NULL; + + if (src1_nrows > 1) { + const size_t matrix_row_counts_size = n_as * sizeof(uint32_t); + assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size); + + hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]); + + memset(matrix_row_counts, 0, matrix_row_counts_size); + scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0); + + uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as); + mapping_stride = max_count > 0 ? max_count : 1; + + size_t matrix_row_map_size = n_as * mapping_stride * sizeof(struct mmid_row_mapping); + const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size; + + if (total_map_size > octx->ctx->ddr_spad_size) { + mapping_buf = memalign(128, total_map_size); + if (!mapping_buf) { + return HTP_STATUS_INTERNAL_ERR; + } + } + + matrix_row_counts = (uint32_t *) mapping_buf; + matrix_rows = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size); + + memset(matrix_row_counts, 0, n_as * sizeof(uint32_t)); + scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride); + } + + mmctx->matrix_row_counts = matrix_row_counts; + mmctx->matrix_rows = matrix_rows; + mmctx->mapping_stride = mapping_stride; + mmctx->mm_div_ne11 = kparams->div_ne11; + mmctx->src0_row_size_padded = src0_row_size_padded; + mmctx->src1_nrows = src1_nrows; + + htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); + + int s; + if (kparams->n_hmx) { + s = hmx_mm_op_matmul_id_nx(octx, mmctx); + } else { + if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) { + s = hvx_mm_matmul_id_nx(octx, mmctx, src1_nrows > 1 ? hvx_mm_id_nx : hvx_mv_id_nx); + } else { + s = HTP_STATUS_NO_SUPPORT; + } + } + + if (mapping_buf != octx->ctx->ddr_spad_base) { + free(mapping_buf); + } + + return s; +} +int op_matmul_nx(struct htp_ops_context * octx) { + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + if (kparams->n_hmx) { + return hmx_mm_nx_2d_f32(octx, kparams); + } + + struct htp_thread_trace * tr = &octx->ctx->trace[0]; + htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); + const uint32_t n_weights = kparams->n_weights; const struct htp_tensor * restrict src0 = octx->src[0]; // first weight @@ -3366,9 +4034,9 @@ int op_matmul_nx(struct htp_ops_context * octx) { mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; - mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; + mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; - mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; + mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div); mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index dbc8e3590..2dbcb0c2e 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -134,7 +134,8 @@ static inline int htp_mm_hmx_compute_chunks(size_t vtcm_total, size_t best_mn = 0; size_t best_m = 0, best_n = 0; - const size_t n_max = hex_align_down((size_t)n, HTP_MM_HMX_TILE_N_COLS); + const size_t max_nc_budget = (usable / per_n_cost); + const size_t n_max = hex_align_down(hex_smin((size_t)n, max_nc_budget), HTP_MM_HMX_TILE_N_COLS); for (size_t nc = n_max; nc >= HTP_MM_HMX_TILE_N_COLS; nc -= HTP_MM_HMX_TILE_N_COLS) { size_t n_fixed = 0, ncmn = 0, mc_denom = 0; if (hex_mul_overflow(nc, per_n_cost, &n_fixed)) continue; @@ -299,6 +300,15 @@ static inline void htp_mm_hmx_get_batched_chunk_costs( *size_per_mn_out = sizeof(uint16_t); } +static inline size_t htp_mm_hmx_get_2d_overhead(bool pipeline, bool is_matmul_id) { + size_t num_regions = pipeline ? 7 : (is_matmul_id ? 4 : 5); + return num_regions * HTP_MM_HMX_TILE_SIZE + 256; +} + +static inline size_t htp_mm_hmx_get_batched_overhead(void) { + return 5 * HTP_MM_HMX_TILE_SIZE + 256; +} + struct htp_mm_hmx_vtcm_layout { // Byte offsets from vtcm_base for each region size_t off_weight[2]; // [1] is only used when pipelined @@ -568,10 +578,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( } size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); - size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0; - if (dst_size_per_thread < quant_scratch_size_per_thread) { - dst_size_per_thread = quant_scratch_size_per_thread; - } + size_t dst_slice_per_thread = (dst_nrows > 0 && src1_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0; + size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread; dst_sz = dst_size_per_thread * n_threads; break; } @@ -592,10 +600,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( } size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); - size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0; - if (dst_size_per_thread < quant_scratch_size_per_thread) { - dst_size_per_thread = quant_scratch_size_per_thread; - } + size_t dst_slice_per_thread = dst_nrows > 0 ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0; + size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread; dst_sz = dst_size_per_thread * n_threads; break; } @@ -658,7 +664,7 @@ static inline bool htp_mm_hmx_solve_batched_params( int act_threads = n_threads; while (act_threads >= 1) { - size_t group_overhead = 256; + size_t group_overhead = htp_mm_hmx_get_batched_overhead(); size_t group_size_per_n, group_size_per_m, group_size_per_mn; htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn); @@ -725,7 +731,7 @@ static inline bool htp_mm_hmx_solve_2d_params( int act_threads = n_threads; while (act_threads >= 1) { - size_t simple_2d_overhead = 256; + size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id); size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn; htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);