diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 784bb37b9..b4b943058 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -16,6 +16,7 @@ #include "hvx-dump.h" #include "hvx-arith.h" #include "hvx-reduce.h" +#include "hvx-copy.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" @@ -1413,7 +1414,7 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) { for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) { // for each expert const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]); - if (eid == -1) { + if (eid < 0) { continue; } assert(eid >= 0 && eid < (int32_t) n_ids); @@ -4111,7 +4112,7 @@ static int hvx_mm_matmul_id_nx( return HTP_STATUS_OK; } -static inline void scan_expert_ids_n( +static __attribute__((noinline)) uint32_t scan_expert_ids_n( const struct htp_tensor * ids, const uint32_t n_ids, uint32_t n_as, @@ -4121,15 +4122,17 @@ static inline void scan_expert_ids_n( ) { const size_t ids_nb1 = ids->nb[1]; const uint8_t * ids_data = (const uint8_t *) ids->data; + uint32_t n_skipped = 0; for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1); for (uint32_t id = 0; id < n_ids; ++id) { const int32_t i02 = row_ptr[id]; - if (i02 == -1) { + if (i02 < 0) { + n_skipped++; continue; } - assert(i02 >= 0 && i02 < n_as); + assert(i02 >= 0 && i02 < (int32_t) n_as); if (matrix_rows) { matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 }; @@ -4137,31 +4140,10 @@ static inline void scan_expert_ids_n( counts[i02] += 1; } } + return n_skipped; } -// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here -static inline void zero_skipped_dst_rows( - const struct htp_tensor * ids, - uint32_t n_ids, - const struct htp_tensor * dst, - uint32_t ne0, - size_t nb1, - size_t nb2 -) { - const uint8_t * ids_data = (const uint8_t *) ids->data; - - for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { - for (uint32_t id = 0; id < n_ids; ++id) { - const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]); - if (i02 != -1) { - continue; - } - memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float)); - } - } -} - -static inline void scan_expert_ids( +static __attribute__((noinline)) uint32_t scan_expert_ids( const struct htp_tensor * ids, uint32_t n_ids, uint32_t n_as, @@ -4173,23 +4155,25 @@ static inline void scan_expert_ids( if (ids_nb0 == 4) { switch (n_ids) { - case 8: scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride); break; - case 4: scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride); break; - case 2: scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride); break; - default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break; + case 8: return scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride); + case 4: return scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride); + case 2: return scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride); + default: return scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); } } else { // Strided fallback const size_t ids_nb1 = ids->nb[1]; const uint8_t * ids_data = (const uint8_t *) ids->data; + uint32_t n_skipped = 0; for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1); for (uint32_t id = 0; id < n_ids; ++id) { const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0); - if (i02 == -1) { + if (i02 < 0) { + n_skipped++; continue; } - assert(i02 >= 0 && i02 < n_as); + assert(i02 >= 0 && i02 < (int32_t) n_as); if (matrix_rows) { matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 }; @@ -4197,9 +4181,124 @@ static inline void scan_expert_ids( counts[i02] += 1; } } + return n_skipped; } } +static __attribute__((noinline)) uint32_t count_skipped_ids_single(const struct htp_tensor * ids, uint32_t n_ids) { + const uint8_t * ids_data = (const uint8_t *) ids->data; + const size_t ids_nb0 = ids->nb[0]; + uint32_t n_skipped = 0; + for (uint32_t id = 0; id < n_ids; ++id) { + const int32_t i02 = *(const int32_t *) (ids_data + id * ids_nb0); + if (i02 < 0) { + n_skipped++; + } + } + return n_skipped; +} + +// Zero rows for skipped slots (ids < 0). Matmuls only write rows that an expert owns. +static void zero_skipped_dst_rows( + struct htp_ops_context * octx, + const struct htp_tensor * ids, + uint32_t n_ids, + const struct htp_tensor * dst +) { + const uint32_t elem_size = (dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); + uint32_t col_start = 0; + uint32_t col_count = dst->ne[0]; + + if (octx->ctx->mdev.count > 1) { + const bool can_split = htp_tensor_can_row_partition(dst, elem_size); + const uint32_t chunk_elems = (dst->type == HTP_TYPE_F16) ? 64 : 32; + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition( + dst->ne[0], can_split ? chunk_elems : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + col_start = range.start; + col_count = range.count; + } + + if (col_count == 0) { + return; + } + + const uint8_t * ids_data = (const uint8_t *) ids->data; + const size_t ids_nb0 = ids->nb[0]; + const size_t ids_nb1 = ids->nb[1]; + const size_t dst_nb1 = dst->nb[1]; + const size_t dst_nb2 = dst->nb[2]; + const uint32_t col_offset_bytes = col_start * elem_size; + + for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { + const uint8_t * ids_row = ids_data + iid1 * ids_nb1; + uint8_t * dst_row_group = (uint8_t *) dst->data + iid1 * dst_nb2 + col_offset_bytes; + for (uint32_t id = 0; id < n_ids; ++id) { + const int32_t i02 = *(const int32_t *) (ids_row + id * ids_nb0); + if (i02 < 0) { + hvx_splat_u(dst_row_group + id * dst_nb1, Q6_V_vzero(), col_count, elem_size); + } + } + } +} + +static int htp_mmid_prepare_mapping( + struct htp_ops_context * octx, + const struct htp_tensor * ids, + uint32_t n_ids, + uint32_t n_as, + uint32_t act_nrows, + uint8_t ** out_mapping_buf, + uint32_t ** out_matrix_row_counts, + struct mmid_row_mapping ** out_matrix_rows, + uint32_t * out_mapping_stride, + uint32_t * out_n_skipped +) { + 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; + uint32_t n_skipped = 0; + + if (act_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]); + + hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t)); + n_skipped = 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); + + hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t)); + scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride); + } else { + n_skipped = count_skipped_ids_single(ids, n_ids); + } + + *out_mapping_buf = mapping_buf; + *out_matrix_row_counts = matrix_row_counts; + *out_matrix_rows = matrix_rows; + *out_mapping_stride = mapping_stride; + *out_n_skipped = n_skipped; + + return HTP_STATUS_OK; +} + int op_matmul_id(struct htp_ops_context * octx) { htp_matmul_tensors_preamble; @@ -4235,41 +4334,22 @@ int op_matmul_id(struct htp_ops_context * octx) { const int n_ids = ids->ne[0]; // n_expert_used const int n_as = ne02; // n_expert - uint8_t * mapping_buf = octx->ctx->ddr_spad_base; - uint32_t mapping_stride = 1; - uint32_t * matrix_row_counts = (uint32_t *) mapping_buf; + uint8_t * mapping_buf = NULL; + uint32_t * matrix_row_counts = NULL; struct mmid_row_mapping * matrix_rows = NULL; + uint32_t mapping_stride = 1; + uint32_t n_skipped = 0; - if (act_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); + int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows, + &mapping_buf, &matrix_row_counts, &matrix_rows, + &mapping_stride, &n_skipped); + if (map_status != HTP_STATUS_OK) { + return map_status; } - zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2); + if (n_skipped > 0) { + zero_skipped_dst_rows(octx, ids, n_ids, dst); + } mmctx->matrix_row_counts = matrix_row_counts; mmctx->matrix_rows = matrix_rows; @@ -4360,38 +4440,26 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { 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; + uint8_t * mapping_buf = NULL; + uint32_t * matrix_row_counts = NULL; struct mmid_row_mapping * matrix_rows = NULL; + uint32_t mapping_stride = 1; + uint32_t n_skipped = 0; - if (act_nrows > 1) { - const size_t matrix_row_counts_size = n_as * sizeof(uint32_t); - assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size); + int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows, + &mapping_buf, &matrix_row_counts, &matrix_rows, + &mapping_stride, &n_skipped); + if (map_status != HTP_STATUS_OK) { + return map_status; + } - 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; + if (n_skipped > 0) { + for (uint32_t p = 0; p < n_weights; ++p) { + const struct htp_tensor * restrict d = octx->dsts[p]; + if (d) { + zero_skipped_dst_rows(octx, ids, n_ids, d); } } - - 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;