Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

hexagon: MUL_MAT and MUL_MAT_ID fusion and fixes (#28202)

Browse Source 
* 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
This commit is contained in:
Max Krasnyansky
2026-09-02 09:15:21 +03:00
committed by GitHub
parent ba8818cbf3
commit 960dffab05
7 changed files with 1241 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-hexagon/ggml-hexagon.cpp
+383 -27
View File
@@ -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 // 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 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_opbatch = 1280; // max number of ops in a batch
static int opt_opqueue = 64; // max number of pending batches 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_optrace = 0; // trace buffer size per thread (0 means default)
static int opt_oppoll = 0; // polling for batch completions static int opt_oppoll = 0; // polling for batch completions
static int opt_opfusion = 1; // enable/disable op fusion 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 static std::regex* opt_opfilter = NULL; // regex of ops to not claim
#define HEX_VERBOSE(...) \ #define HEX_VERBOSE(...) \
@@ -293,6 +307,15 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
struct htp_mm_kernel_params * kparams 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( static bool ggml_hexagon_precompute_allreduce_params(
const struct ggml_hexagon_session * sess, const struct ggml_hexagon_session * sess,
const struct ggml_tensor * dst, 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 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(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_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 // ** backend sessions
@@ -1832,6 +1859,42 @@ struct ggml_hexagon_opbatch {
} }
} }
void sort_buffers() {
if (n_bufs <= 1) return;
std::vector<int> 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<uint16_t> remap(n_bufs);
std::vector<htp_buf_desc> 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) { bool try_fuse_allreduce_add(const htp_opnode & node) {
if (n_ops == 0 || opt_ar_select != 2) return false; if (n_ops == 0 || opt_ar_select != 2) return false;
if (node.opcode != HTP_OP_ADD) 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]) { if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0]) {
return false; 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; struct htp_mm_kernel_params kparams;
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &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) { if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n",
sess->c_name(), kparams.vtcm_size, sess->vtcm_size); sess->c_name(), kparams.vtcm_size, sess->vtcm_size);
@@ -2210,6 +2279,9 @@ struct ggml_hexagon_opbatch {
struct htp_mm_kernel_params kparams; struct htp_mm_kernel_params kparams;
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &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) { if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n",
sess->c_name(), kparams.vtcm_size, sess->vtcm_size); sess->c_name(), kparams.vtcm_size, sess->vtcm_size);
@@ -2272,17 +2344,171 @@ struct ggml_hexagon_opbatch {
return false; return false;
} }
enum ggml_hexagon_fusion_flags { bool try_fuse_mul_mat_id_nx(const htp_opnode & node) {
GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2 if (n_ops == 0 || node.opcode != HTP_OP_MUL_MAT_ID) return false;
GGML_HEXAGON_FUSE_RMS_NORM_MUL = (1 << 2), // 4 if (!is_mergeable_mul_mat_id(node.node)) return false;
GGML_HEXAGON_FUSE_MUL_MAT_ADD = (1 << 3), // 8
GGML_HEXAGON_FUSE_MUL_MAT_NX = (1 << 4), // 16
};
static inline bool ggml_hexagon_is_fusion_enabled(int flag) { const ggml_tensor * w_in = node.src0();
if (opt_opfusion <= 0) return false; const ggml_tensor * x_in = node.src1();
if (opt_opfusion == 1) return true; // 1 enables all const ggml_tensor * ids_in = node.node->src[2];
return (opt_opfusion & flag) != 0; 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<ggml_hexagon_shared_buffer *>(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<ggml_hexagon_shared_buffer *>(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) { bool try_fuse(const htp_opnode & node) {
@@ -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_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_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_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; return false;
} }
}; };
@@ -2350,6 +2577,8 @@ struct ggml_hexagon_opqueue {
delete shm_buf; delete shm_buf;
} }
size_t shm_size() const { return shm_buf ? shm_buf->size() : 0; }
// push new batch // push new batch
bool push(htp_opbatch_req& req, dspqueue_buffer& dbuf, ggml_hexagon_opbatch* op_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"); 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 * t_ptr = m_ptr; m_ptr += t_size;
uint8_t * o_ptr = m_ptr; uint8_t * o_ptr = m_ptr;
op_batch->sort_buffers();
memcpy(b_ptr, (void *) op_batch->h_bufs.data(), b_size); memcpy(b_ptr, (void *) op_batch->h_bufs.data(), b_size);
memcpy(t_ptr, (void *) op_batch->h_tens.data(), t_size); memcpy(t_ptr, (void *) op_batch->h_tens.data(), t_size);
memcpy(o_ptr, (void *) op_batch->h_ops.data(), o_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); opt_vmem = ggml_hexagon_measure_max_vmem(this);
GGML_LOG_INFO("ggml-hex: %s measured max vmem %zu\n", this->c_name(), opt_vmem); 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); 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_matmul_id,
bool is_batched bool is_batched
) { ) {
if (src1->type != GGML_TYPE_F32) {
return false;
}
const int ne00 = src0->ne[0]; const int ne00 = src0->ne[0];
const int ne11 = src1->ne[1]; const int ne11 = src1->ne[1];
const int ne12 = src1->ne[2]; const int ne12 = src1->ne[2];
@@ -3408,7 +3644,8 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
return false; 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; const int m = is_matmul_id ? ne12 : ne11;
if (m <= HTP_MM_HMX_MIN_NROWS) { if (m <= HTP_MM_HMX_MIN_NROWS) {
return false; return false;
@@ -3460,7 +3697,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
if (!use_grouped) { if (!use_grouped) {
// Fallback to simple 2D path (group_size = 1) // 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)) { 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; return false;
} }
@@ -3918,11 +4155,40 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
) { ) {
memset(kparams, 0, sizeof(*kparams)); memset(kparams, 0, sizeof(*kparams));
const int wtype = src0->type; const int ne00 = src0->ne[0];
const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); 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 ne10 = src1->ne[0];
const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; const int ne11 = src1->ne[1];
const int ne12 = src1->ne[2];
const int ne13 = src1->ne[3];
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);
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;
}
}
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 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 size_t src0_row_size = src0->nb[1];
@@ -3978,6 +4244,26 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
} }
} }
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) { static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) {
return t && t->buffer && ggml_backend_buft_is_host(t->buffer->buft); return t && t->buffer && ggml_backend_buft_is_host(t->buffer->buft);
GGML_UNUSED(sess); GGML_UNUSED(sess);
@@ -4010,11 +4296,6 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s
return false; 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) { if (src1->ne[2] != 1 || src1->ne[3] != 1) {
return false; // no broadcasting (for now) 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 * src1 = op->src[1]; // indices
const struct ggml_tensor * dst = op; 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) { if (src0->type != GGML_TYPE_F32 && src0->ne[0] < 32) {
return false; 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); 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) { static bool is_mergeable_mul_mat(const ggml_tensor * t) {
if (!t || t->op != GGML_OP_MUL_MAT) return false; 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); 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) { 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) { if (n1->src[0]->type != n2->src[0]->type) {
return false; 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; 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 })) { 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; extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE;
} else if (graph->nodes[i]->op == GGML_OP_MUL_MAT) { } 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, { GGML_OP_MUL_MAT, GGML_OP_ADD })) || 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)) { ggml_node_has_n_uses(graph, i, 1)) {
extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE;
} }
ggml/src/ggml-hexagon/htp-opnode.h
+2 -1
View File
@@ -315,7 +315,8 @@ struct htp_opformat {
} }
void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) { 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 || 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 auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
const char * path = "unknown"; const char * path = "unknown";
int32_t type = kparams->kernel_type; int32_t type = kparams->kernel_type;
ggml/src/ggml-hexagon/htp/htp-ctx.h
+1
View File
@@ -118,6 +118,7 @@ struct htp_context {
int op_matmul(struct htp_ops_context * octx); int op_matmul(struct htp_ops_context * octx);
int op_matmul_id(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_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_binary(struct htp_ops_context * octx);
int op_unary(struct htp_ops_context * octx); int op_unary(struct htp_ops_context * octx);
int op_sum_rows(struct htp_ops_context * octx); int op_sum_rows(struct htp_ops_context * octx);
ggml/src/ggml-hexagon/htp/htp-ops.h
+1
View File
@@ -52,6 +52,7 @@ enum htp_op_code {
HTP_OP_MUL_MAT, HTP_OP_MUL_MAT,
HTP_OP_MUL_MAT_ID, HTP_OP_MUL_MAT_ID,
HTP_OP_MUL_MAT_NX, HTP_OP_MUL_MAT_NX,
HTP_OP_MUL_MAT_ID_NX,
HTP_OP_MUL_MAT_ADD, HTP_OP_MUL_MAT_ADD,
HTP_OP_RMS_NORM, HTP_OP_RMS_NORM,
HTP_OP_RMS_NORM_MUL, HTP_OP_RMS_NORM_MUL,
ggml/src/ggml-hexagon/htp/main.c
+31 -8
View File
@@ -753,6 +753,9 @@ static int execute_op(struct htp_ops_context * octx) {
case HTP_OP_MUL_MAT_ID: case HTP_OP_MUL_MAT_ID:
return op_matmul_id(octx); return op_matmul_id(octx);
case HTP_OP_MUL_MAT_ID_NX:
return op_matmul_id_nx(octx);
case HTP_OP_MUL_MAT_NX: case HTP_OP_MUL_MAT_NX:
return op_matmul_nx(octx); 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) { static inline bool mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
if (b->base) return; // already mapped if (b->base) return true; // already mapped
// find unused mapping // find unused mapping
for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { 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) { if (!m->size) {
void *va = htp_mmap(b->fd, b->size); void *va = htp_mmap(b->fd, b->size);
if (va == NULL) { if (va == NULL) {
FARF(ERROR, "mmap failed : fd %u size %u", b->fd, (uint32_t) b->size); FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %u", b->fd, (uint32_t) b->size);
abort(); // can't do much else at this point return false;
} }
m->base = b->base = (uint64_t) va; 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; m->size = b->size;
FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->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); 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) { 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++) { for (uint32_t i=0; i < n_bufs; i++) {
struct htp_buf_desc *b = 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); 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) { static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) {
ggml/src/ggml-hexagon/htp/matmul-ops.c
+720 -52
View File
@@ -55,10 +55,14 @@ typedef struct {
size_t src0_nb3; size_t src0_nb3;
size_t src1_nb2; size_t src1_nb2;
size_t src1_nb3; size_t src1_nb3;
size_t dst_nb2;
size_t dst_nb3;
size_t src2_nb2; size_t src2_nb2;
size_t src2_nb3; 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; } hmx_mm_f16_f32_batched_params_t;
struct htp_mm_context { 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; const uint32_t nr1 = ne1 * ne2 * ne3;
// distribute the thread work across the inner or outer loop based on which one is larger // 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 dr0, dr1, ith0, ith1;
uint32_t nchunk1 = nr0 > nr1 ? 1 : nth; // parallelize by src1 rows if (nr0 > nr1) {
dr0 = fastdiv(nr0 + nth - 1, &octx->ctx->n_threads_div);
// The number of elements in each chunk dr1 = nr1;
const uint32_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; ith0 = ith;
const uint32_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; ith1 = 0;
} else {
uint32_t current_chunk = ith; dr0 = nr0;
dr1 = fastdiv(nr1 + nth - 1, &octx->ctx->n_threads_div);
const uint32_t ith0 = current_chunk % nchunk0; ith0 = 0;
const uint32_t ith1 = current_chunk / nchunk0; ith1 = ith;
}
const uint32_t ir0_start = dr0 * ith0; const uint32_t ir0_start = dr0 * ith0;
const uint32_t ir0_end = MIN(ir0_start + dr0, nr0); 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; \ 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]; \ 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); \ src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); \
\ \
const uint32_t start_row = src0_nrows_per_thread * ith; \ 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) { static int hvx_mm_init_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type type) {
switch (type) { switch (type) {
case HTP_TYPE_Q4_0: case HTP_TYPE_Q4_0:
@@ -1153,7 +1331,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
src0->type == HTP_TYPE_MXFP4); src0->type == HTP_TYPE_MXFP4);
// Compute src0_nrows_per_thread // 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) { if (is_repacked) {
mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
} else { } else {
@@ -1325,11 +1503,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
} else { } 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_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
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);
size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; 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 ne01 = src_w->ne[1];
const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; 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); src0_nrows_per_thread += (src0_nrows_per_thread & 1);
const uint32_t src0_start_row = src0_nrows_per_thread * ith; 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 { typedef struct {
const struct mmid_row_mapping *matrix_rows; struct htp_context * ctx;
struct htp_thread_trace * traces;
__fp16 * dst; __fp16 * dst;
const float * src; const float * src;
const struct mmid_row_mapping * matrix_rows;
float * vtcm_f32_act;
uint32_t n_tasks; uint32_t n_tasks;
uint32_t n_tot_chunks; uint32_t n_tot_chunks;
uint32_t n_chunks_per_task; uint32_t n_chunks_per_task;
uint32_t k_block; uint32_t k_block;
uint32_t k_stride; uint32_t k_stride;
uint32_t k_valid; 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; size_t vtcm_f32_act_bytes_per_thread;
uint32_t dma_step_rows; uint32_t dma_step_rows;
uint32_t dma_step_rows_shift; uint32_t dma_step_rows_shift;
} activation_transfer_task_state_t; } activation_transfer_task_state_t;
typedef struct { typedef struct {
struct htp_context * ctx;
struct htp_thread_trace * traces;
__fp16 * dst; __fp16 * dst;
const float * src; const float * src;
float * vtcm_f32_act;
uint32_t n_rows; uint32_t n_rows;
uint32_t k_block; uint32_t k_block;
uint32_t k_stride; uint32_t k_stride;
uint32_t k_valid; uint32_t k_valid;
uint32_t n_col_chunks; uint32_t n_col_chunks;
struct fastdiv_values n_threads_div; struct fastdiv_values n_threads_div;
float *vtcm_f32_act;
size_t vtcm_f32_act_bytes; 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;
uint32_t dma_step_rows_shift; uint32_t dma_step_rows_shift;
} activation_transfer_col_chunk_state_t; } activation_transfer_col_chunk_state_t;
@@ -1811,6 +1989,7 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i,
} }
typedef struct { typedef struct {
struct htp_thread_trace * traces;
const struct mmid_row_mapping * matrix_rows; const struct mmid_row_mapping * matrix_rows;
__fp16 * dst; __fp16 * dst;
const float * src; const float * src;
@@ -1827,10 +2006,10 @@ typedef struct {
uint32_t start_row; uint32_t start_row;
uint32_t cne1; uint32_t cne1;
uint32_t k_valid; uint32_t k_valid;
struct htp_thread_trace *traces;
} activation_transfer_gathered_task_state_t; } activation_transfer_gathered_task_state_t;
typedef struct { typedef struct {
struct htp_thread_trace * traces;
const struct mmid_row_mapping * matrix_rows; const struct mmid_row_mapping * matrix_rows;
const __fp16 * vtcm_src; const __fp16 * vtcm_src;
float * dst; float * dst;
@@ -1844,7 +2023,6 @@ typedef struct {
size_t dst_nb2; size_t dst_nb2;
uint32_t start_row; uint32_t start_row;
uint32_t cne1; uint32_t cne1;
struct htp_thread_trace *traces;
} output_transfer_scattered_task_state_t; } output_transfer_scattered_task_state_t;
static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) { static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) {
@@ -1946,17 +2124,17 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles(
} }
typedef struct { typedef struct {
struct htp_context * ctx;
struct htp_thread_trace * traces;
float * dst; float * dst;
const float *src2;
const __fp16 * vtcm_src; const __fp16 * vtcm_src;
const float * src2;
uint32_t n_rows; uint32_t n_rows;
uint32_t n_cols; uint32_t n_cols;
uint32_t dst_stride; uint32_t dst_stride;
uint32_t src2_stride; uint32_t src2_stride;
uint32_t dst_cols; uint32_t dst_cols;
struct fastdiv_values n_threads_div; struct fastdiv_values n_threads_div;
struct htp_thread_trace *traces;
struct htp_context *ctx;
} output_transfer_col_chunk_state_t; } output_transfer_col_chunk_state_t;
static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
@@ -1975,9 +2153,9 @@ static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned i
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first); 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 __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; int chunk_dst_cols = (int)st->dst_cols - (int)c_first;
if (chunk_dst_cols > 0) { 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; uint32_t n_blocks = (uint32_t)n_cols / 32;
if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) { 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; output_transfer_col_chunk_state_t col_state;
col_state.dst = dst; col_state.dst = dst;
col_state.src2 = src2; col_state.src2 = src2;
@@ -2128,8 +2306,7 @@ static void transfer_activation_chunk_threaded(const struct activation_transfer_
state.ctx = ctx; state.ctx = ctx;
state.vtcm_f32_act = vtcm_f32_act; 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(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128);
state.vtcm_f32_act_bytes_per_thread = hex_align_down(vtcm_f32_act_bytes / active_threads, 128);
uint32_t dma_step_rows = 2; uint32_t dma_step_rows = 2;
uint32_t dma_step_rows_shift = 1; 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 = dma_step_rows;
state.dma_step_rows_shift = dma_step_rows_shift; 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) { if (state.n_tasks == 1 || n_threads == 1) {
transfer_activation_chunk_worker_fn(1, 0, &state); transfer_activation_chunk_worker_fn(1, 0, &state);
} else { } else {
@@ -2447,21 +2625,286 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
return 0; return 0;
} }
static inline int hmx_mm_batch_r2(const hmx_mm_f16_f32_batched_params_t *params) { static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) {
return params->ne02 > 0 ? params->ne12 / params->ne02 : 1; struct htp_context * ctx = octx->ctx;
struct htp_thread_trace * tr = &ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
const uint32_t n_weights = kparams->n_weights;
if (n_weights == 0 || n_weights > HTP_OP_MAX_OUTPUTS) {
return HTP_STATUS_INVAL_PARAMS;
} }
static inline int hmx_mm_batch_r3(const hmx_mm_f16_f32_batched_params_t *params) { const struct htp_tensor * restrict src0 = octx->src[0];
return params->ne03 > 0 ? params->ne13 / params->ne03 : 1; 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, static inline const __fp16 *hmx_mm_weight_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params,
int dst_b2, int dst_b3) { int dst_b2, int dst_b3) {
const int r2 = hmx_mm_batch_r2(params); const size_t b2_idx = (params->r2 <= 1) ? (size_t) dst_b2 : (size_t) fastdiv((uint32_t) dst_b2, &params->div_r2);
const int r3 = hmx_mm_batch_r3(params); const size_t b3_idx = (params->r3 <= 1) ? (size_t) dst_b3 : (size_t) fastdiv((uint32_t) dst_b3, &params->div_r3);
return (const __fp16 *) ((const uint8_t *) params->weight + return (const __fp16 *) ((const uint8_t *) params->weight +
(size_t) (dst_b2 / r2) * params->src0_nb2 + b2_idx * params->src0_nb2 +
(size_t) (dst_b3 / r3) * params->src0_nb3); b3_idx * params->src0_nb3);
} }
static inline const float *hmx_mm_activation_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, 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 (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; } 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; const size_t vtcm_budget = ctx->vtcm_size;
// Check if the precomputed parameters are grouped or simple. // 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, htp_mm_hmx_get_2d_chunk_costs(weight_type, k, /*pipeline=*/false, aligned_tile_size,
&size_per_n, &size_per_m, &size_per_mn); &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; 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_padded, n,
/*m_block_cost=*/(size_t) n * HTP_MM_HMX_COST_W_DEQUANT, /*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)) { /*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], .dst_nb3 = dst->nb[3],
.src2_nb2 = src2_nb2, .src2_nb2 = src2_nb2,
.src2_nb3 = src2_nb3, .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, ret = hmx_mm_f16_f32_batched(octx->ctx, &batch_params,
kparams->m_chunk, kparams->n_chunk, 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_src0_stride = src0_row_size_padded;
mmctx->vtcm_src1_stride = src1_row_size; 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_src1_size_per_thread = L.src1_bytes;
mmctx->vtcm_src2_size_per_thread = 0; 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->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
mmctx->quant_task_func = quant_task_func; mmctx->quant_task_func = quant_task_func;
@@ -3123,6 +3571,134 @@ static int hvx_mm_matmul_id(
return HTP_STATUS_OK; 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( static inline void scan_expert_ids_n(
const struct htp_tensor * ids, const struct htp_tensor * ids,
const uint32_t n_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 src0_nrows = ne01; // per expert
const uint32_t src1_nrows = ne11 * ne12 * ne13; 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); mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
// row groups // row groups
@@ -3280,11 +3856,103 @@ int op_matmul_id(struct htp_ops_context * octx) {
return s; 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]; struct htp_thread_trace * tr = &octx->ctx->trace[0];
htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 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 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 uint32_t n_weights = kparams->n_weights;
const struct htp_tensor * restrict src0 = octx->src[0]; // first weight 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_src0_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src1_stride = src1_row_size; 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_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->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
mmctx->quant_task_func = quant_task_func; mmctx->quant_task_func = quant_task_func;
ggml/src/ggml-hexagon/htp/matmul-ops.h
+17 -11
View File
@@ -134,7 +134,8 @@ static inline int htp_mm_hmx_compute_chunks(size_t vtcm_total,
size_t best_mn = 0; size_t best_mn = 0;
size_t best_m = 0, best_n = 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) { 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; size_t n_fixed = 0, ncmn = 0, mc_denom = 0;
if (hex_mul_overflow(nc, per_n_cost, &n_fixed)) continue; 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); *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 { struct htp_mm_hmx_vtcm_layout {
// Byte offsets from vtcm_base for each region // Byte offsets from vtcm_base for each region
size_t off_weight[2]; // [1] is only used when pipelined 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 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; 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;
if (dst_size_per_thread < quant_scratch_size_per_thread) { 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_size_per_thread = quant_scratch_size_per_thread;
}
dst_sz = dst_size_per_thread * n_threads; dst_sz = dst_size_per_thread * n_threads;
break; 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 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; size_t dst_slice_per_thread = dst_nrows > 0 ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
if (dst_size_per_thread < quant_scratch_size_per_thread) { 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_size_per_thread = quant_scratch_size_per_thread;
}
dst_sz = dst_size_per_thread * n_threads; dst_sz = dst_size_per_thread * n_threads;
break; break;
} }
@@ -658,7 +664,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
int act_threads = n_threads; int act_threads = n_threads;
while (act_threads >= 1) { 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; 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); 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; int act_threads = n_threads;
while (act_threads >= 1) { 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; 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); 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);