#pragma clang diagnostic ignored "-Wgnu-zero-variadic-macro-arguments" #pragma clang diagnostic ignored "-Wunused-function" #pragma clang diagnostic ignored "-Wunused-variable" #pragma clang diagnostic ignored "-Wunused-but-set-variable" #include #include #include #include #include #include #include "hex-dma.h" #include "hvx-utils.h" #include "hvx-dump.h" #include "hvx-arith.h" #include "hvx-reduce.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" #include "htp-ctx.h" #include "htp-ops.h" #include "htp-tensor.h" #include "matmul-ops.h" #include "htp-vtcm.h" static void hvx_tensor_add_f32_grid( const struct htp_tensor * restrict dst, const struct htp_tensor * restrict src2, uint32_t start_row, uint32_t end_row, uint32_t start_col, uint32_t end_col, const struct fastdiv_values * div_ne11_12, const struct fastdiv_values * div_ne11 ); typedef struct { float *dst; const float *src2; const float *activation; const __fp16 *weight; int m; int k; int n; int act_stride; int weight_stride; int dst_stride; uint32_t src2_stride; int ne02; int ne03; int ne12; int ne13; size_t src0_nb2; size_t src0_nb3; size_t src1_nb2; size_t src1_nb3; size_t src2_nb2; size_t src2_nb3; size_t dst_nb2; size_t dst_nb3; int r2; int r3; struct fastdiv_values div_r2; struct fastdiv_values div_r3; } hmx_mm_f16_f32_batched_params_t; struct htp_mm_context { const char * type; struct htp_ops_context * octx; const struct htp_tensor * act; void (*vec_dot_1x1)(const uint32_t n, float * restrict s0, const void * restrict vx0, const void * restrict vy0); void (*vec_dot_2x1)(const uint32_t n, float * restrict s0, const void * restrict vx0, const void * restrict vx1, const void * restrict vy0); void (*vec_dot_2x2)(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx0, const void * restrict vx1, const void * restrict vy0, const void * restrict vy1); void (*vec_dot_32x1)(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz); // Precomputed values uint32_t src0_nrows_per_thread; uint32_t src0_row_start; uint32_t src0_row_end; uint32_t src0_row_size_padded; uint32_t src1_nrows; struct fastdiv_values mm_div_ne12_ne1; struct fastdiv_values mm_div_ne1; struct fastdiv_values mm_div_r2; struct fastdiv_values mm_div_r3; struct fastdiv_values mm_div_ne11; // Per thread quant tasks // Precomputed block-parallel quantization values worker_callback_t quant_task_func; uint32_t quant_ib_first[WORK_QUEUE_MAX_N_THREADS]; uint32_t quant_ib_last[WORK_QUEUE_MAX_N_THREADS]; uint32_t quant_r[WORK_QUEUE_MAX_N_THREADS]; uint32_t quant_c[WORK_QUEUE_MAX_N_THREADS]; uint32_t n_quant_tasks; uint32_t n_quant_rows_per_thread; atomic_uint quant_barrier; // Fields for scattered mapping & HMX support in MUL_MAT_ID const uint32_t * matrix_row_counts; const struct mmid_row_mapping * matrix_rows; uint32_t mapping_stride; // Dynamic VTCM pointers allocated sequentially uint8_t * vtcm_src0; uint8_t * vtcm_src1; uint8_t * vtcm_src2; uint8_t * vtcm_src3; uint8_t * vtcm_dst; // Cached strides uint32_t vtcm_src0_stride; uint32_t vtcm_src1_stride; uint32_t vtcm_src2_stride; uint32_t vtcm_src3_stride; // Cached thread offsets/sizes uint32_t vtcm_src0_size_per_thread; uint32_t vtcm_src1_size_per_thread; uint32_t vtcm_src2_size_per_thread; uint32_t vtcm_src3_size_per_thread; uint32_t vtcm_dst_size_per_thread; }; static int htp_mm_init_context( struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams ) { if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) { return HTP_STATUS_INVAL_PARAMS; } if (kparams->n_hmx) { if (kparams->n_act_threads <= 0 || kparams->n_act_threads > (int32_t) octx->n_threads) { return HTP_STATUS_INVAL_PARAMS; } } return HTP_STATUS_OK; } // vdelta control to expand first 32 e8m0 values into 32 uint32 elements static const uint8_t __attribute__((aligned(128))) expand_x32_e8m0[128] = { 0x00, 0x00, 0x00, 0x00, 0x01, 0x04, 0x00, 0x00, 0x02, 0x00, 0x08, 0x08, 0x01, 0x02, 0x00, 0x04, 0x04, 0x00, 0x00, 0x00, 0x11, 0x10, 0x10, 0x10, 0x02, 0x00, 0x04, 0x00, 0x01, 0x02, 0x08, 0x08, 0x08, 0x08, 0x00, 0x00, 0x01, 0x04, 0x00, 0x00, 0x22, 0x20, 0x20, 0x20, 0x21, 0x22, 0x20, 0x24, 0x04, 0x00, 0x00, 0x00, 0x09, 0x08, 0x00, 0x00, 0x02, 0x00, 0x04, 0x00, 0x11, 0x12, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x01, 0x04, 0x00, 0x00, 0x02, 0x00, 0x08, 0x08, 0x01, 0x02, 0x00, 0x04, 0x44, 0x40, 0x40, 0x40, 0x41, 0x40, 0x40, 0x40, 0x42, 0x40, 0x44, 0x40, 0x41, 0x42, 0x48, 0x48, 0x08, 0x08, 0x00, 0x00, 0x01, 0x04, 0x00, 0x00, 0x12, 0x10, 0x10, 0x10, 0x01, 0x02, 0x00, 0x04, 0x04, 0x00, 0x00, 0x00, 0x09, 0x08, 0x00, 0x00, 0x22, 0x20, 0x24, 0x20, 0x21, 0x22, 0x20, 0x20, }; // IQ4_NL dequantization LUT: maps 4-bit index (0-15) to int8 kvalue // kvalues: -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113 static const uint8_t __attribute__((aligned(VLEN))) kvalues_iq4nl_lut[] = { 0x81, 0, 0x98, 0, 0xAD, 0, 0xBF, 0, 0xCF, 0, 0xDD, 0, 0xEA, 0, 0xF6, 0, 0x01, 0, 0x0D, 0, 0x19, 0, 0x26, 0, 0x35, 0, 0x45, 0, 0x59, 0, 0x71, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, }; static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = { 0, 0, 1, 0, 2, 0, 3, 0, 4, 0, 6, 0, 8, 0, 12, 0, 0, 0, 0xff, 0, 0xfe, 0, 0xfd, 0, 0xfc, 0, 0xfa, 0, 0xf8, 0, 0xf4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, }; #define htp_matmul_tensors_preamble \ const struct htp_tensor * restrict src0 = octx->src[0]; \ const struct htp_tensor * restrict src1 = octx->src[1]; \ const struct htp_tensor * restrict src2 = octx->src[2]; \ const struct htp_tensor * restrict dst = octx->dst; \ \ const uint32_t ne00 = src0->ne[0]; \ const uint32_t ne01 = src0->ne[1]; \ const uint32_t ne02 = src0->ne[2]; \ const uint32_t ne03 = src0->ne[3]; \ \ const uint32_t ne10 = src1->ne[0]; \ const uint32_t ne11 = src1->ne[1]; \ const uint32_t ne12 = src1->ne[2]; \ const uint32_t ne13 = src1->ne[3]; \ \ const uint32_t ne20 = src2 ? src2->ne[0] : 0; \ const uint32_t ne21 = src2 ? src2->ne[1] : 0; \ const uint32_t ne22 = src2 ? src2->ne[2] : 0; \ const uint32_t ne23 = src2 ? src2->ne[3] : 0; \ \ const uint32_t ne0 = dst->ne[0]; \ const uint32_t ne1 = dst->ne[1]; \ const uint32_t ne2 = dst->ne[2]; \ const uint32_t ne3 = dst->ne[3]; \ \ const uint32_t nb00 = src0->nb[0]; \ const uint32_t nb01 = src0->nb[1]; \ const uint32_t nb02 = src0->nb[2]; \ const uint32_t nb03 = src0->nb[3]; \ \ const uint32_t nb10 = src1->nb[0]; \ const uint32_t nb11 = src1->nb[1]; \ const uint32_t nb12 = src1->nb[2]; \ const uint32_t nb13 = src1->nb[3]; \ \ const uint32_t nb0 = dst->nb[0]; \ const uint32_t nb1 = dst->nb[1]; \ const uint32_t nb2 = dst->nb[2]; \ const uint32_t nb3 = dst->nb[3]; #define htp_matmul_preamble \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ dma_queue *dma_queue = octx->ctx->dma[ith]; \ uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; \ htp_matmul_tensors_preamble; static inline void hvx_mm_run_quant_task(struct htp_mm_context * mmctx, unsigned int ith) { if (mmctx->quant_task_func) { if (ith < mmctx->n_quant_tasks) { mmctx->quant_task_func(mmctx->n_quant_tasks, ith, mmctx); atomic_fetch_sub(&mmctx->quant_barrier, 1); } while (atomic_load(&mmctx->quant_barrier) > 0) { // spin } } } // *** matmul with support for 4d tensors and full broadcasting static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; assert(ne12 % ne02 == 0); assert(ne13 % ne03 == 0); // This is the size of the first dimension of the result, so we can iterate that way. (see the ASSERT above, these are the same numbers) const uint32_t nr0 = ne0; // This is the size of the rest of the dimensions of the result const uint32_t nr1 = ne1 * ne2 * ne3; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; // distribute the thread work across the inner or outer loop based on which one is larger uint32_t dr0, dr1, ith0, ith1; if (nr0 > nr1) { dr0 = fastdiv(src0_nrows + nth - 1, &octx->n_threads_div); dr1 = nr1; ith0 = ith; ith1 = 0; } else { dr0 = src0_nrows; dr1 = fastdiv(nr1 + nth - 1, &octx->n_threads_div); ith0 = 0; ith1 = ith; } const uint32_t ir0_start = mmctx->src0_row_start + dr0 * ith0; const uint32_t ir0_end = MIN(ir0_start + dr0, mmctx->src0_row_end); const uint32_t ir1_start = dr1 * ith1; const uint32_t ir1_end = MIN(ir1_start + dr1, nr1); // no work for this thread if (ir0_start >= ir0_end || ir1_start >= ir1_end) { return; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start); const uint32_t blck_0 = 64; const uint32_t blck_1 = 64; for (uint32_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) { for (uint32_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) { for (uint32_t ir1 = iir1; ir1 < MIN(iir1 + blck_1, ir1_end); ir1++) { const uint32_t i13 = fastdiv(ir1, &mmctx->mm_div_ne12_ne1); const uint32_t i12 = fastdiv(ir1 - i13 * ne12 * ne1, &mmctx->mm_div_ne1); const uint32_t i11 = (ir1 - i13 * ne12 * ne1 - i12 * ne1); // broadcast src0 into src1 const uint32_t i03 = fastdiv(i13, &mmctx->mm_div_r3); const uint32_t i02 = fastdiv(i12, &mmctx->mm_div_r2); const uint32_t i1 = i11; const uint32_t i2 = i12; const uint32_t i3 = i13; const uint8_t * restrict src0_base = (const uint8_t *) src0->data + (0 + i02 * nb02 + i03 * nb03); const uint8_t * restrict src1_col = (const uint8_t *) src1->data + (i11 * nb11 + i12 * nb12 + i13 * nb13); float * dst_col = (float *) ((uint8_t * restrict) dst->data + (i1 * nb1 + i2 * nb2 + i3 * nb3)); const uint32_t ir0_block_end = MIN(iir0 + blck_0, ir0_end); for (uint32_t ir0 = iir0; ir0 < ir0_block_end; ir0++) { const uint8_t * restrict src0_row = src0_base + ir0 * nb01; mmctx->vec_dot_1x1(ne00, &dst_col[ir0], src0_row, src1_col); } } } } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start); if (src2) { hvx_tensor_add_f32_grid(dst, src2, ir1_start, ir1_end, ir0_start, ir0_end, &mmctx->mm_div_ne12_ne1, &mmctx->mm_div_ne1); } } #include "hmx-mm-kernels-tiled.h" #include "hvx-mm-kernels-tiled.h" #include "hvx-mm-kernels-flat.h" // Specialized repacked matmul macros #define MATMUL_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ htp_matmul_preamble; \ \ const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ const uint32_t src1_nrows = ne11 * ne12 * ne13; \ \ const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ \ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ \ const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ const uint32_t n_prefetch = kparams->n_prefetch; \ assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ \ const size_t dst_row_size = nb1; \ const size_t src1_row_size = nb11; \ const size_t src1_stride = mmctx->vtcm_src1_stride; \ const size_t src2_stride = src2 ? ((src2->ne[1] == 1) ? 0 : src2->nb[1]) : 0; \ \ uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ uint8_t * restrict src1_data = mmctx->vtcm_src1; \ \ const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ \ const uint32_t tile_size = TILE_SIZE; \ const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ \ uint32_t n_k_tiles_w = ne00 / 32; \ uint32_t n_k_tiles_a = ne10 / 32; \ uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ uint32_t ct_start = src0_start_row / 32; \ uint32_t ct_end = (src0_end_row + 31) / 32; \ \ uint32_t push_ct = ct_start; \ if (src0_start_row < src0_end_row) { \ for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ } \ } \ \ hvx_mm_run_quant_task(mmctx, ith); \ \ if (src0_start_row >= src0_end_row) { \ return; \ } \ \ for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ \ int valid_rows = (int)ne0 - (int)(ct * 32); \ valid_rows = MIN(32, MAX(0, valid_rows)); \ \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ uint32_t ir1 = 0; \ for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); \ float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); \ \ float * dst_ptr0 = &dst_row0[ct * 32]; \ float * dst_ptr1 = &dst_row1[ct * 32]; \ \ const float * src2_ptr0 = NULL; \ const float * src2_ptr1 = NULL; \ if (src2) { \ const float * restrict src2_row0 = (const float *) ((const uint8_t *) src2->data + ((ir1+0) * src2_stride)); \ const float * restrict src2_row1 = (const float *) ((const uint8_t *) src2->data + ((ir1+1) * src2_stride)); \ src2_ptr0 = &src2_row0[ct * 32]; \ src2_ptr1 = &src2_row1[ct * 32]; \ } \ DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, src2_ptr0, src2_ptr1); \ } \ \ for (; ir1 < src1_nrows; ++ir1) { \ const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \ float * dst_ptr = &dst_row[ct * 32]; \ \ const float * src2_ptr = NULL; \ if (src2) { \ const float * restrict src2_row = (const float *) ((const uint8_t *) src2->data + (ir1 * src2_stride)); \ src2_ptr = &src2_row[ct * 32]; \ } \ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, src2_ptr); \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ if (push_ct < ct_end) { \ dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ push_ct++; \ } \ } \ } #define MATVEC_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X1) \ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ htp_matmul_preamble; \ \ const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ \ const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ \ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ \ const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ const uint32_t n_prefetch = kparams->n_prefetch; \ assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ \ const size_t dst_row_size = nb1; \ const size_t src1_row_size = nb11; \ const size_t src1_stride = mmctx->vtcm_src1_stride; \ \ uint8_t * vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ uint8_t * vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ uint8_t * src1_data = mmctx->vtcm_src1; \ \ float * tmp = (float *) vtcm_dst_ptr; \ \ const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ \ const uint8_t * restrict src1_col = (const uint8_t *) src1_data; \ float * restrict dst_col = (float *) dst->data; \ \ const uint32_t tile_size = TILE_SIZE; \ const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ \ uint32_t n_k_tiles_w = ne00 / 32; \ uint32_t n_k_tiles_a = ne10 / 32; \ uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ uint32_t ct_start = src0_start_row / 32; \ uint32_t ct_end = (src0_end_row + 31) / 32; \ \ uint32_t push_ct = ct_start; \ if (src0_start_row < src0_end_row) { \ if (src2) { \ float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; \ const float * src2_ptr = (const float *) src2->data + src0_start_row; \ int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ if (slice_size > 0) { \ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), \ slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \ dma_queue_pop_nowait(dma_queue); \ } \ } \ for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ } \ } \ \ hvx_mm_run_quant_task(mmctx, ith); \ \ if (src0_start_row >= src0_end_row) { \ return; \ } \ \ for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ \ float * dst_ptr = &tmp[ct * 32 - src0_start_row]; \ int valid_rows = (int)ne0 - (int)(ct * 32); \ valid_rows = MIN(32, MAX(0, valid_rows)); \ \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ if (push_ct < ct_end) { \ dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ push_ct++; \ } \ } \ \ int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ if (copy_cnt > 0) { \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ if (src2) { \ hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], \ (const uint8_t *) tmp, \ (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), \ copy_cnt); \ } else { \ hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ } \ } #define MATMUL_NX_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ const uint32_t n_weights = kparams->n_weights; \ \ const struct htp_tensor * restrict act = octx->src[n_weights]; /* x */ \ const uint32_t ne10 = act->ne[0]; \ const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; \ const size_t src1_stride = mmctx->vtcm_src1_stride; \ \ uint8_t * restrict vtcm_weight_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ uint8_t * restrict src1_data = mmctx->vtcm_src1; \ \ 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 tile_size = TILE_SIZE; \ const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ uint32_t n_k_tiles_a = ne10 / 32; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ dma_queue * dma_queue = octx->ctx->dma[ith]; \ \ hvx_mm_run_quant_task(mmctx, ith); \ \ for (uint32_t widx = 0; widx < n_weights; widx++) { \ const struct htp_tensor * restrict src_w = octx->src[widx]; \ const struct htp_tensor * restrict dst = octx->dsts[widx]; \ if (!src_w || !dst) continue; \ \ const uint32_t ne00 = src_w->ne[0]; \ const uint32_t ne01 = src_w->ne[1]; \ const size_t dst_row_size = dst->nb[1]; \ const uint8_t * restrict src_w_row = (const uint8_t *) src_w->data; \ \ uint32_t n_k_tiles_w = ne00 / 32; \ uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ \ uint32_t src0_start_row = 0; \ uint32_t src0_end_row = ne01; \ if (octx->ctx->mdev.count > 1) { \ const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); \ const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, \ octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); \ src0_start_row = range.start; \ src0_end_row = range.start + range.count; \ } \ \ const uint32_t nrows = src0_end_row - src0_start_row; \ uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); \ src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); \ \ const uint32_t start_row = src0_start_row + src0_nrows_per_thread * ith; \ const uint32_t end_row = MIN(start_row + src0_nrows_per_thread, src0_end_row); \ if (start_row >= end_row) continue; \ \ uint32_t ct_start = start_row / 32; \ uint32_t ct_end = (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_weight_ptr + d * tile_row_transfer_size_aligned, \ src_w_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ } \ \ for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ int valid_rows = (int)ne01 - (int)(ct * 32); \ valid_rows = MIN(32, MAX(0, valid_rows)); \ \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ uint32_t ir1 = 0; \ for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ \ float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); \ float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); \ float * dst_ptr0 = &dst_row0[ct * 32]; \ float * dst_ptr1 = &dst_row1[ct * 32]; \ \ DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, NULL, NULL); \ } \ \ for (; ir1 < src1_nrows; ++ir1) { \ const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \ float * dst_ptr = &dst_row[ct * 32]; \ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ if (push_ct < ct_end) { \ dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src_w_row + push_ct * tile_row_stride), \ aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ push_ct++; \ } \ } \ } \ } MATMUL_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) MATMUL_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) MATMUL_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) MATMUL_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) MATMUL_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) MATMUL_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) MATMUL_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) MATMUL_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) MATMUL_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) #define QUANTIZE_IMPL(name, log_name, kernel_fn, dst_row_size_expr) \ static void name(unsigned int nth, unsigned int ith, void * data) { \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ const struct htp_tensor * src = mmctx->act; \ const uint32_t ne0 = src->ne[0]; \ const uint32_t ne1 = src->ne[1]; \ const uint32_t ne2 = src->ne[2]; \ const uint32_t ne3 = src->ne[3]; \ const uint32_t nrows = ne1 * ne2 * ne3; \ const uint32_t nrows_per_thread = mmctx->n_quant_rows_per_thread; \ \ const uint32_t ir_first = nrows_per_thread * ith; \ if (ir_first >= nrows) { \ return; \ } \ \ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ \ uint8_t * restrict dst = mmctx->vtcm_src1; \ const uint32_t ir_last = MIN(ir_first + nrows_per_thread, nrows); \ const size_t src_row_size = src->nb[1]; \ const size_t dst_row_size = (dst_row_size_expr); \ const uint8_t * restrict src_data = (const uint8_t *) src->data + (src_row_size * ir_first); \ uint8_t * restrict dst_data = (uint8_t *) dst + (dst_row_size * ir_first); \ uint8_t * restrict tmp_data = (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith); \ kernel_fn(src_data, dst_data, tmp_data, ne0, ir_last - ir_first, src_row_size, dst_row_size); \ \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ } QUANTIZE_IMPL(quantize_f32_q8_0_tiled, "quantize-f32-q8_0_tiled", quantize_f32_q8_0_tiled_kernel, htp_mm_q8_0_tiled_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_q8_1_tiled, "quantize-f32-q8_1_tiled", quantize_f32_q8_1_tiled_kernel, htp_mm_q8_1_tiled_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_q8_0_flat, "quantize-f32-q8_0_flat", quantize_f32_q8_0_flat_kernel, htp_mm_q8_0_flat_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_q8_1_flat, "quantize-f32-q8_1_flat", quantize_f32_q8_1_flat_kernel, htp_mm_q8_1_flat_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_f32_flat, "quantize-f32-f32", quantize_f32_f32_flat_kernel, mmctx->vtcm_src1_stride) QUANTIZE_IMPL(quantize_f32_f16_flat, "quantize-f32-f16", quantize_f32_f16_flat_kernel, mmctx->vtcm_src1_stride) QUANTIZE_IMPL(quantize_f16_f16_flat, "quantize-f16-f16", quantize_f16_f16_flat_kernel, mmctx->vtcm_src1_stride) static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = data; struct htp_ops_context * octx = mmctx->octx; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); const struct htp_tensor * src = mmctx->act; quantize_f32_q8_0_tiled_block_kernel( (const float *) src->data, mmctx->vtcm_src1, (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith), src->ne[0], mmctx->quant_ib_first[ith], mmctx->quant_ib_last[ith], src->nb[1], htp_mm_q8_0_tiled_row_size(src->ne[0]), mmctx->quant_r[ith], mmctx->quant_c[ith] ); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); } static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = data; struct htp_ops_context * octx = mmctx->octx; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); const struct htp_tensor * src = mmctx->act; quantize_f32_q8_1_tiled_block_kernel( (const float *) src->data, mmctx->vtcm_src1, (uint8_t *) mmctx->vtcm_dst + (mmctx->vtcm_dst_size_per_thread * ith), src->ne[0], mmctx->quant_ib_first[ith], mmctx->quant_ib_last[ith], src->nb[1], htp_mm_q8_1_tiled_row_size(src->ne[0]), mmctx->quant_r[ith], mmctx->quant_c[ith] ); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); } MATVEC_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x1) MATVEC_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x1) MATVEC_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x1) MATVEC_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x1) MATVEC_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x1) MATVEC_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x1) MATVEC_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x1) MATVEC_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x1) MATVEC_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x1) MATVEC_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x1) MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) MATMUL_NX_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) MATMUL_NX_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) MATMUL_NX_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) MATMUL_NX_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) MATMUL_NX_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t prefetch_mask = n_prefetch - 1; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; // src0 rows const uint32_t src1_nrows = ne11 * ne12 * ne13; // src1 rows const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); struct htp_thread_trace * tr = &octx->ctx->trace[ith]; const size_t dst_row_size = nb1; const size_t src0_row_size = nb01; const size_t src1_row_size = nb11; const size_t src0_stride = mmctx->vtcm_src0_stride; const size_t src1_stride = mmctx->vtcm_src1_stride; // Per-thread VTCMs for all tensors uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; uint8_t * restrict src1_data = mmctx->vtcm_src1; const uint8_t * restrict src0_row = (const uint8_t *) src0->data; // Prefill vtcm with src0 rows if (src0_start_row < src0_end_row) { for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const int is0 = (ir0 - src0_start_row); if (is0 >= (int)n_prefetch) { break; } dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; } // Process src0 rows for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); // Process src1 columns in pairs (2x2 tiling) uint32_t ir1 = 0; for (; ir1 + 1 < src1_nrows; ir1 += 2) { const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); } // Handle remaining src1 rows (fallback to 2x1) for (; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, src1_col); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); // Prefetch next (n + vtcm_nrows) row const int pr0 = (ir0 + n_prefetch); const int is0 = (pr0 - src0_start_row) & prefetch_mask; if (pr0 < src0_end_row_x2) { dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } // Process the last row (if any) if (src0_end_row != src0_end_row_x2) { uint32_t ir0 = src0_end_row_x2; const int is0 = (ir0 - src0_start_row) & prefetch_mask; dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 1); const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); #pragma unroll(2) for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, src1_col); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); } if (src2) { hvx_tensor_add_f32_grid(dst, src2, 0, src1_nrows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); } } static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); struct htp_thread_trace * tr = &octx->ctx->trace[ith]; const size_t dst_row_size = nb1; const size_t src0_row_size = nb01; const size_t src1_row_size = nb11; const size_t src0_stride = mmctx->vtcm_src0_stride; const size_t src1_stride = mmctx->vtcm_src1_stride; // Per-thread VTCMs for all tensors uint8_t * vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; uint8_t * vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; uint8_t * src1_data = mmctx->vtcm_src1; float * tmp = (float *) vtcm_dst_ptr; const uint8_t * restrict src0_row = (const uint8_t *) src0->data; const uint8_t * restrict src1_col = (const uint8_t *) src1_data; float * restrict dst_col = (float *) dst->data; const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t prefetch_mask = n_prefetch - 1; // Prefill vtcm with 2x src0 rows if (src0_start_row < src0_end_row) { if (src2) { float * vtcm_src2_ptr = (float *) mmctx->vtcm_src2 + src0_start_row; const float * src2_ptr = (const float *) src2->data + src0_start_row; int slice_size = (int)src0_end_row - (int)src0_start_row; if (slice_size > 0) { dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr, src2_ptr), slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); dma_queue_pop_nowait(dma_queue); } } for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const uint32_t is0 = (ir0 - src0_start_row); if (is0 >= n_prefetch) { break; } dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; } // Process src0 rows for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); mmctx->vec_dot_2x1(ne00, &tmp[ir0 - src0_start_row], ss0, ss0 + src0_stride, src1_col); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); // Prefetch next (n + vtcm_nrows) row const uint32_t pr0 = (ir0 + n_prefetch); const uint32_t is0 = (pr0 - src0_start_row) & prefetch_mask; if (pr0 < src0_end_row_x2) { dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } // Process the last row (if any) if (src0_end_row != src0_end_row_x2) { const uint32_t ir0 = src0_end_row_x2; const uint32_t is0 = (ir0 - src0_start_row) & prefetch_mask; dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 1); const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); mmctx->vec_dot_1x1(ne00, &tmp[ir0 - src0_start_row], ss0, src1_col); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); } int copy_cnt = src0_end_row - src0_start_row; if (copy_cnt > 0) { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row); if (src2) { hvx_add_f32_uaa((uint8_t *) &dst_col[src0_start_row], (const uint8_t *) tmp, (const uint8_t *) ((const float *) mmctx->vtcm_src2 + src0_start_row), copy_cnt); } else { hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_end_row); } } #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)] static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_tensor * restrict ids = octx->src[2]; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; // src0 rows per expert const uint32_t src1_nrows = ne11; const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t n_ids = ids->ne[0]; // n_expert_used const uint32_t n_as = ne02; // n_expert const uint32_t * matrix_row_counts = mmctx->matrix_row_counts; const struct mmid_row_mapping * matrix_rows = mmctx->matrix_rows; const size_t dst_row_size = nb1; const size_t src1_row_size = htp_mm_q8_0_tiled_row_size(ne10); const size_t src1_stride = mmctx->vtcm_src1_stride; // Per-thread VTCMs for all tensors 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; } const uint8_t * src0_row = (const uint8_t *) src0->data + cur_a * nb02; const uint32_t tile_size = htp_mm_get_weight_tile_size(src0->type); const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src0->type); const uint32_t n_k_tiles_w = ne00 / 32; const uint32_t n_k_tiles_a = ne10 / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; const uint32_t 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)ne01 - (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 < cne1; ++cid) { struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, cid); const int rm1 = row_mapping.i1; // expert idx const int rm2 = row_mapping.i2; // token idx const uint32_t ir1 = fastmodulo(rm1, ne11, &mmctx->mm_div_ne11); // src1 row idx const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (ir1 + rm2 * ne11 + 0) * src1_stride); float * restrict dst_row = (float *) (dst->data + (rm1 * nb1 + rm2 * nb2 + 0)); mmctx->vec_dot_32x1(ne10, &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_mv_id(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_tensor * restrict ids = octx->src[2]; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; // src0 rows per expert const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); hvx_mm_run_quant_task(mmctx, ith); if (src0_start_row >= src0_end_row) { return; } struct htp_thread_trace * tr = &octx->ctx->trace[ith]; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); assert(ne13 % ne03 == 0); const size_t dst_row_size = nb1; const size_t src1_row_size = htp_mm_q8_0_tiled_row_size(ne10); const uint32_t n_aids = src2->ne[0]; // num activated experts const uint32_t n_ids = ne02; // num experts // Per-thread VTCMs for all tensors 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) { // for each expert const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]); if (eid < 0) { continue; } assert(eid < (int32_t) n_ids); const uint8_t * restrict src0_row = (const uint8_t *) src0->data + eid * nb02; const uint8_t * restrict src1_col = (const uint8_t *) src1_data; float * restrict dst_row = (float *) (dst->data + ie1 * nb1); const uint32_t tile_size = htp_mm_get_weight_tile_size(src0->type); const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src0->type); const uint32_t n_k_tiles_w = ne00 / 32; const uint32_t n_k_tiles_a = ne10 / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; const uint32_t 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)ne01 - (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(ne10, &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_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 ne01 = src_w->ne[1]; uint32_t start_row = 0; uint32_t end_row = ne01; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); start_row = range.start; end_row = range.start + range.count; } const uint32_t nrows = end_row - start_row; uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, end_row); if (src0_start_row >= src0_end_row) continue; const uint8_t * restrict src0_row = (const uint8_t *) src_w->data + eid * src_w->nb[2]; const 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 ne01 = src_w->ne[1]; uint32_t start_row = 0; uint32_t end_row = ne01; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); start_row = range.start; end_row = range.start + range.count; } const uint32_t nrows = end_row - start_row; uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, end_row); if (src0_start_row >= src0_end_row) continue; const uint8_t * src0_row = (const uint8_t *) src_w->data + cur_a * src_w->nb[2]; const uint32_t tile_size = htp_mm_get_weight_tile_size(src_w->type); const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src_w->type); const uint32_t n_k_tiles_w = src_w->ne[0] / 32; const uint32_t n_k_tiles_a = act->ne[0] / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; const uint32_t ct_start = src0_start_row / 32; const uint32_t ct_end = (src0_end_row + 31) / 32; uint32_t push_ct = ct_start; for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); } for (uint32_t ct = ct_start; ct < ct_end; ct++) { const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); for (uint32_t cid = 0; cid < (uint32_t) cne1; ++cid) { struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, cid); const int rm1 = row_mapping.i1; const int rm2 = row_mapping.i2; const uint32_t ir1 = fastmodulo(rm1, act->ne[1], &mmctx->mm_div_ne11); const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + (ir1 + rm2 * act->ne[1]) * src1_stride); float * restrict dst_row = (float *) (dst->data + (rm1 * dst->nb[1] + rm2 * dst->nb[2])); mmctx->vec_dot_32x1(act->ne[0], &dst_row[ct * 32], w_tile, src1_col, valid_rows, NULL); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); if (push_ct < ct_end) { dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); push_ct++; } } } } } static int hvx_mm_init_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type type) { switch (type) { case HTP_TYPE_Q4_0: mmctx->type = "q4_0_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q4_0_32x1; return 0; case HTP_TYPE_Q4_1: mmctx->type = "q4_1_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q4_1_32x1; return 0; case HTP_TYPE_Q8_0: mmctx->type = "q8_0_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q8_0_32x1; return 0; case HTP_TYPE_IQ4_NL: mmctx->type = "iq4nl_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_iq4nl_32x1; return 0; case HTP_TYPE_MXFP4: mmctx->type = "mxfp4_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_mxfp4_32x1; return 0; default: return -1; } } static int hvx_mm_matmul(struct htp_ops_context * octx) { htp_matmul_tensors_preamble; struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; mmctx->act = src1; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t src0_nrows = ne01; const uint32_t src1_nrows = ne11 * ne12 * ne13; uint32_t src0_row_start = 0; uint32_t src0_row_end = src0_nrows; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); src0_row_start = range.start; src0_row_end = range.start + range.count; } if (src0_row_start >= src0_row_end) { return HTP_STATUS_OK; } const uint32_t nrows = src0_row_end - src0_row_start; mmctx->src0_row_start = src0_row_start; mmctx->src0_row_end = src0_row_end; bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || src0->type == HTP_TYPE_MXFP4); // Compute src0_nrows_per_thread mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div); if (is_repacked) { mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); } else { mmctx->src0_nrows_per_thread += (mmctx->src0_nrows_per_thread & 1); // round up to even } const size_t src0_row_size = nb01; const size_t dst_row_size = nb1; size_t src1_row_size = nb11; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); size_t src1_row_size_padded; worker_callback_t quant_task_func; worker_callback_t matmul_job_func; uint32_t n_quant_tasks = 1; if (src1_nrows > 1) { if (is_repacked) { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_2d_repacked_q4_0; break; case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_2d_repacked_q4_1; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_2d_repacked_q8_0; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_2d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_2d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; } } else { matmul_job_func = hvx_mm_2d; } } else { if (is_repacked) { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mv_2d_repacked_q4_0; break; case HTP_TYPE_Q4_1: matmul_job_func = hvx_mv_2d_repacked_q4_1; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mv_2d_repacked_q8_0; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mv_2d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mv_2d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; } } else { matmul_job_func = hvx_mv_2d; } } bool need_quant = true; switch (kparams->kernel_type) { case HTP_MM_KERNEL_HVX_F16_F16_VTCM: quant_task_func = (src1->type == HTP_TYPE_F32) ? quantize_f32_f16_flat : quantize_f16_f16_flat; mmctx->type = "f16-f16"; mmctx->vec_dot_1x1 = vec_dot_f16_f16_aa_1x1; mmctx->vec_dot_2x1 = vec_dot_f16_f16_aa_2x1; mmctx->vec_dot_2x2 = vec_dot_f16_f16_aa_2x2; src1_row_size = hex_round_up(ne10 * 2, 128); break; case HTP_MM_KERNEL_HVX_F16_F32_DDR: mmctx->type = "f16-f32"; mmctx->vec_dot_1x1 = vec_dot_f16_f32_uu_1x1; matmul_job_func = hvx_mm_4d; mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1; mmctx->mm_div_ne1 = kparams->div_ne1; mmctx->mm_div_r2 = kparams->div_r2; mmctx->mm_div_r3 = kparams->div_r3; need_quant = false; quant_task_func = NULL; src1_row_size = nb11; break; case HTP_MM_KERNEL_HVX_F16_F16_DDR: mmctx->type = "f16-f16"; mmctx->vec_dot_1x1 = vec_dot_f16_f16_uu_1x1; matmul_job_func = hvx_mm_4d; mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1; mmctx->mm_div_ne1 = kparams->div_ne1; mmctx->mm_div_r2 = kparams->div_r2; mmctx->mm_div_r3 = kparams->div_r3; src1_row_size = nb11; need_quant = false; quant_task_func = NULL; break; case HTP_MM_KERNEL_HVX_F32_F32_VTCM: quant_task_func = quantize_f32_f32_flat; mmctx->type = "f32-f32"; mmctx->vec_dot_1x1 = vec_dot_f32_f32_aa_1x1; mmctx->vec_dot_2x1 = vec_dot_f32_f32_aa_2x1; mmctx->vec_dot_2x2 = vec_dot_f32_f32_aa_2x2; src1_row_size = hex_round_up(ne10 * 4, 128); break; case HTP_MM_KERNEL_HVX_F32_F32_DDR: quant_task_func = NULL; mmctx->type = "f32-f32"; mmctx->vec_dot_1x1 = vec_dot_f32_f32_uu_1x1; mmctx->mm_div_ne12_ne1 = kparams->div_ne12_ne1; mmctx->mm_div_ne1 = kparams->div_ne1; mmctx->mm_div_r2 = kparams->div_r2; mmctx->mm_div_r3 = kparams->div_r3; src1_row_size = nb11; need_quant = false; matmul_job_func = hvx_mm_4d; break; case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: { n_quant_tasks = MIN(src1_nrows, octx->n_threads); quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat; src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); if (src1_nrows > 1) { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_2d_repacked_q4_0_flat; break; case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_2d_repacked_q4_1_flat; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_2d_repacked_q8_0_flat; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_2d_repacked_iq4nl_flat; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_2d_repacked_mxfp4_flat; break; default: return HTP_STATUS_NO_SUPPORT; } } else { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mv_2d_repacked_q4_0_flat; break; case HTP_TYPE_Q4_1: matmul_job_func = hvx_mv_2d_repacked_q4_1_flat; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mv_2d_repacked_q8_0_flat; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mv_2d_repacked_iq4nl_flat; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mv_2d_repacked_mxfp4_flat; break; default: return HTP_STATUS_NO_SUPPORT; } } break; } case HTP_MM_KERNEL_HVX_QUANT_BLOCK: case HTP_MM_KERNEL_HVX_QUANT_ROW: default: if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) { return HTP_STATUS_NO_SUPPORT; } const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (ne10 + qk - 1) / qk; const uint32_t total_nb = src1_nrows * nb; 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; } src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); break; } struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false); if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { mmctx->vtcm_src1_size_per_thread = L.src1_bytes; } else { mmctx->vtcm_src1_size_per_thread = fastdiv(L.src1_bytes, &octx->n_threads_div); } mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div); mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); const size_t vtcm_size = L.total_bytes; FARF(HIGH, "matmul-%s : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu (%zu)\n", mmctx->type, L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size); FARF(HIGH, "matmul-%s : %ux%ux%ux%u * %ux%ux%ux%u-> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], src0->data, src1->data, dst->data); if (octx->ctx->vtcm_size < vtcm_size) { FARF(ERROR, "matmul-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size); return HTP_STATUS_VTCM_TOO_SMALL; } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; octx->dst_spad.src = NULL; mmctx->vtcm_src0_stride = src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; if (need_quant) { mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; mmctx->n_quant_tasks = n_quant_tasks; atomic_init(&mmctx->quant_barrier, n_quant_tasks); } else { mmctx->quant_task_func = NULL; mmctx->n_quant_tasks = 0; } htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads); return HTP_STATUS_OK; } static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = data; struct htp_ops_context * octx = mmctx->octx; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_weights = kparams->n_weights; const struct htp_tensor * restrict act = octx->src[n_weights]; const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; 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; dma_queue * dma_queue = octx->ctx->dma[ith]; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t prefetch_mask = n_prefetch - 1; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; hvx_mm_run_quant_task(mmctx, ith); for (uint32_t widx = 0; widx < n_weights; widx++) { const struct htp_tensor * restrict src_w = octx->src[widx]; const struct htp_tensor * restrict dst = octx->dsts[widx]; if (!src_w || !dst) continue; const uint32_t ne00 = src_w->ne[0]; const uint32_t ne01 = src_w->ne[1]; uint32_t start_row = 0; uint32_t end_row = ne01; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); start_row = range.start; end_row = range.start + range.count; } const uint32_t nrows = end_row - start_row; uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); src0_nrows_per_thread += (src0_nrows_per_thread & 1); const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith; const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, end_row); const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); if (src0_start_row >= src0_end_row) continue; const size_t dst_row_size = dst->nb[1]; const size_t src0_row_size = src_w->nb[1]; const size_t src0_stride = hex_round_up(src0_row_size, 128); const uint8_t * restrict src0_row = (const uint8_t *) src_w->data; for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const int is0 = (ir0 - src0_start_row); if (is0 >= (int)n_prefetch) break; dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); uint32_t ir1 = 0; for (; ir1 + 1 < src1_nrows; ir1 += 2) { const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); } for (; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, src1_col); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); const int pr0 = (ir0 + n_prefetch); const int is0 = (pr0 - src0_start_row) & prefetch_mask; if (pr0 < src0_end_row_x2) { dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); } } if (src0_end_row != src0_end_row_x2) { uint32_t ir0 = src0_end_row_x2; const int is0 = (ir0 - src0_start_row) & prefetch_mask; dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 1); const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, src1_col); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); } } } #define DEQUANTIZE_WORKER_LOOP_IMPL(SUFFIX) \ static void dequantize_tiled_worker_loop_##SUFFIX(unsigned int n, unsigned int i, void *data) { \ tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; \ struct htp_thread_trace * tr = &state->traces[i]; \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); \ for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { \ int start = task_id * state->n_tiles_per_task; \ int end = hex_smin(start + state->n_tiles_per_task, state->n_tot_tiles); \ dequantize_tiled_weight_to_fp16_task_##SUFFIX(state, start, end); \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); \ } DEQUANTIZE_WORKER_LOOP_IMPL(q4_0) DEQUANTIZE_WORKER_LOOP_IMPL(q4_1) DEQUANTIZE_WORKER_LOOP_IMPL(iq4_nl) DEQUANTIZE_WORKER_LOOP_IMPL(mxfp4) DEQUANTIZE_WORKER_LOOP_IMPL(q8_0) static void convert_f16_worker_loop(unsigned int n, unsigned int i, void *data) { tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; struct htp_thread_trace * tr = &state->traces[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { int start = task_id * state->n_tiles_per_task; int end = hex_smin(start + state->n_tiles_per_task, state->n_tot_tiles); convert_f16_weight_to_fp16_tiles_task(state, start, end); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); } static void quantize_f32_worker_loop(unsigned int n, unsigned int i, void *data) { tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; struct htp_thread_trace * tr = &state->traces[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, i); for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { int start = task_id * state->n_tiles_per_task; int end = hex_smin(start + state->n_tiles_per_task, state->n_tot_tiles); quantize_f32_weight_to_fp16_tiles_task(state, start, end); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, i); } static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { output_transfer_task_state_t *st = (output_transfer_task_state_t *) data; struct htp_thread_trace * tr = &st->traces[i]; int start_chunk_idx = i * st->n_chunks_per_task; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx); for (unsigned int task_id = i; task_id < (unsigned int)st->n_tasks; task_id += n) { int chunk_idx = task_id * st->n_chunks_per_task; size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task); float *dst = st->dst + chunk_idx * st->dst_stride; const float *src2 = st->src2 ? (st->src2 + chunk_idx * st->src2_stride) : NULL; transfer_output_chunk_fp16_to_fp32(dst, src2, st->vtcm_src, chunk_idx, chunk_size, st->n_cols, st->dst_stride, st->src2_stride, st->dst_cols); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx); } typedef struct { struct htp_context * ctx; struct htp_thread_trace * traces; __fp16 * dst; const float * src; const struct mmid_row_mapping * matrix_rows; float * vtcm_f32_act; uint32_t n_tasks; uint32_t n_tot_chunks; uint32_t n_chunks_per_task; uint32_t k_block; uint32_t k_stride; uint32_t k_valid; size_t vtcm_f32_act_bytes_per_thread; uint32_t dma_step_rows; uint32_t dma_step_rows_shift; } activation_transfer_task_state_t; typedef struct { struct htp_context * ctx; struct htp_thread_trace * traces; __fp16 * dst; const float * src; float * vtcm_f32_act; uint32_t n_rows; uint32_t k_block; uint32_t k_stride; uint32_t k_valid; uint32_t n_col_chunks; struct fastdiv_values n_threads_div; size_t vtcm_f32_act_bytes; uint32_t dma_step_rows; uint32_t dma_step_rows_shift; } activation_transfer_col_chunk_state_t; static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( dma_queue *dma_q, __fp16 *restrict vtcm_dst, const float *restrict src, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, uint32_t k_chunk_valid, uint32_t c_first, uint32_t c_len, float *thread_f32_act, struct htp_thread_trace *tr, uint32_t dma_step_rows, uint32_t dma_step_rows_shift) { const uint32_t R = dma_step_rows; const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); const uint32_t n_steps = n_rows_padded >> dma_step_rows_shift; // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src + c_first), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } // Push step 1 if (n_steps > 1) { uint32_t next_r = R * 1; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride + c_first; float *next_buf = thread_f32_act + 1 * R * c_len; dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } } for (uint32_t s = 0; s < n_steps; ++s) { uint32_t r = s << dma_step_rows_shift; float *curr_buf = thread_f32_act; if (r < n_rows) { curr_buf = (float *) dma_queue_pop(dma_q).dst; } htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r); for (uint32_t p = 0; p < (R >> 1); ++p) { uint32_t row_idx = r + (p << 1); float *pair_buf = curr_buf + (p << 1) * c_len; bool r0_valid = ((row_idx + 0) < n_rows); bool r1_valid = ((row_idx + 1) < n_rows); transfer_activation_row_pair_fp32_to_fp16_col_chunk( vtcm_dst, pair_buf, pair_buf + c_len, row_idx, k_block, c_first, c_len, k_chunk_valid, r0_valid, r1_valid ); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r); // Push step s + 2 uint32_t next_s = s + 2; uint32_t next_r = next_s << dma_step_rows_shift; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride + c_first; dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src), c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch); } } } static void transfer_activation_chunk_fp32_to_fp16_col_chunk( __fp16 *restrict vtcm_dst, const float *restrict src, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, uint32_t c_first, uint32_t c_len, uint32_t k_chunk_valid) { const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); const uint32_t n_rows_tiled = (n_rows / HTP_MM_HMX_TILE_N_ROWS) * HTP_MM_HMX_TILE_N_ROWS; uint32_t r = 0; #pragma unroll(2) for (r = 0; r < n_rows_tiled; r += 2) { const float *ptr_in0 = src + (r + 0) * k_stride + c_first; const float *ptr_in1 = src + (r + 1) * k_stride + c_first; transfer_activation_row_pair_fp32_to_fp16_col_chunk( vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, true, true ); } for (; r < n_rows_padded; r += 2) { const bool row0_valid = r < n_rows; const bool row1_valid = (r + 1) < n_rows; const float *ptr_in0 = row0_valid ? (src + (r + 0) * k_stride + c_first) : NULL; const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride + c_first) : NULL; transfer_activation_row_pair_fp32_to_fp16_col_chunk( vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, row0_valid, row1_valid ); } } static void transfer_activation_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_col_chunk_state_t *st = (activation_transfer_col_chunk_state_t *) data; struct htp_thread_trace * tr = &st->traces[i]; uint32_t n_blocks = st->k_block / 32; uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div); uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div); uint32_t c_first = b_first * 32; uint32_t c_last = b_last * 32; uint32_t c_len = c_last - c_first; if (c_len == 0) { return; } uint32_t k_chunk_valid = 0; if (st->k_valid > c_first) { k_chunk_valid = hex_smin(st->k_valid, c_last) - c_first; } __fp16 *dst = st->dst; const float *src = st->src; if (st->vtcm_f32_act) { size_t thread_scratch_bytes = hex_align_down(fastdiv(st->vtcm_f32_act_bytes, &st->n_threads_div), 128); float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * thread_scratch_bytes); transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk( st->ctx->dma[i], dst, src, st->n_rows, st->k_block, st->k_stride, k_chunk_valid, c_first, c_len, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift ); } else { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first); transfer_activation_chunk_fp32_to_fp16_col_chunk( dst, src, st->n_rows, st->k_block, st->k_stride, c_first, c_len, k_chunk_valid ); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first); } } static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined( dma_queue *dma_q, __fp16 *restrict vtcm_dst, const float *restrict src, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, uint32_t k_valid, float *thread_f32_act, struct htp_thread_trace *tr, uint32_t dma_step_rows, uint32_t dma_step_rows_shift) { const uint32_t R = dma_step_rows; const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS); const uint32_t n_steps = n_rows_padded >> dma_step_rows_shift; // Push step 0 if (n_steps > 0 && n_rows > 0) { uint32_t nrows_to_fetch = hex_smin(n_rows, R); dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } // Push step 1 (if valid) if (n_steps > 1) { uint32_t next_r = R * 1; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride; float *next_buf = thread_f32_act + 1 * R * k_block; dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } for (uint32_t s = 0; s < n_steps; ++s) { uint32_t r = s << dma_step_rows_shift; float *curr_buf = thread_f32_act; if (r < n_rows) { curr_buf = (float *) dma_queue_pop(dma_q).dst; } htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r); for (uint32_t p = 0; p < (R >> 1); ++p) { uint32_t row_idx = r + (p << 1); float *pair_buf = curr_buf + (p << 1) * k_block; bool r0_valid = ((row_idx + 0) < n_rows); bool r1_valid = ((row_idx + 1) < n_rows); transfer_activation_row_pair_fp32_to_fp16(vtcm_dst, pair_buf, pair_buf + k_block, row_idx, k_block, k_valid, r0_valid, r1_valid); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r); // Push step s + 2 uint32_t next_s = s + 2; uint32_t next_r = next_s << dma_step_rows_shift; if (next_r < n_rows) { uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R); const float *next_src = src + next_r * k_stride; dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src), k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch); } } } static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_task_state_t *st = (activation_transfer_task_state_t *) data; struct htp_thread_trace * tr = &st->traces[i]; for (unsigned int task_id = i; task_id < (unsigned int)st->n_tasks; task_id += n) { int chunk_idx = task_id * st->n_chunks_per_task; size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task); __fp16 *dst = st->dst + chunk_idx * st->k_block; const float *src = st->src + chunk_idx * st->k_stride; if (st->vtcm_f32_act) { float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * st->vtcm_f32_act_bytes_per_thread); transfer_activation_chunk_fp32_to_fp16_dma_pipelined( st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift ); } else { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); transfer_activation_chunk_fp32_to_fp16(dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); } } } typedef struct { struct htp_thread_trace * traces; const struct mmid_row_mapping * matrix_rows; __fp16 * dst; const float * src; uint32_t n_tasks; uint32_t n_tot_chunks; uint32_t n_chunks_per_task; uint32_t k_block; uint32_t cur_a; uint32_t mapping_stride; uint32_t ne11; struct fastdiv_values ne11_div; size_t nb11; size_t nb12; uint32_t start_row; uint32_t cne1; uint32_t k_valid; } activation_transfer_gathered_task_state_t; typedef struct { struct htp_thread_trace * traces; const struct mmid_row_mapping * matrix_rows; const __fp16 * vtcm_src; float * dst; uint32_t n_tasks; uint32_t n_tot_chunks; uint32_t n_chunks_per_task; uint32_t n_cols; uint32_t cur_a; uint32_t mapping_stride; size_t dst_nb1; size_t dst_nb2; uint32_t start_row; uint32_t cne1; } output_transfer_scattered_task_state_t; static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_gathered_task_state_t *st = data; struct htp_thread_trace * tr = &st->traces[i]; int chunk_idx = i; int chunk_size = st->n_chunks_per_task; int vtcm_start_row = chunk_idx * chunk_size; int start_row = st->start_row + vtcm_start_row; int n_rows = hex_smin(st->cne1 - start_row, chunk_size); if (n_rows > 0) { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); transfer_activation_chunk_fp32_to_fp16_gathered( st->dst, st->src, start_row, vtcm_start_row, n_rows, st->k_block, st->matrix_rows, st->cur_a, st->mapping_stride, st->ne11, &st->ne11_div, st->nb11, st->nb12, st->cne1, st->k_valid); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); } } static void transfer_activation_chunk_gathered_worker_flat_fn(unsigned int n, unsigned int i, void *data) { activation_transfer_gathered_task_state_t *st = data; struct htp_thread_trace * tr = &st->traces[i]; int chunk_idx = i; int chunk_size = st->n_chunks_per_task; int vtcm_start_row = chunk_idx * chunk_size; int start_row = st->start_row + vtcm_start_row; int n_rows = hex_smin(st->cne1 - start_row, chunk_size); if (n_rows > 0) { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); transfer_activation_chunk_fp32_to_fp16_gathered_flat( st->dst, st->src, start_row, vtcm_start_row, n_rows, st->k_block, st->matrix_rows, st->cur_a, st->mapping_stride, st->nb12, st->cne1, st->k_valid); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx); } } static void transfer_output_chunk_scattered_worker_fn(unsigned int n, unsigned int i, void *data) { output_transfer_scattered_task_state_t *st = data; struct htp_thread_trace * tr = &st->traces[i]; int chunk_idx = i; int chunk_size = st->n_chunks_per_task; int vtcm_start_row = chunk_idx * chunk_size; int start_row = st->start_row + vtcm_start_row; int n_rows = hex_smin(st->cne1 - start_row, chunk_size); if (n_rows > 0) { htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, chunk_idx); transfer_output_chunk_fp16_to_fp32_scattered( st->dst, st->vtcm_src, start_row, vtcm_start_row, n_rows, st->n_cols, st->matrix_rows, st->cur_a, st->mapping_stride, st->dst_nb1, st->dst_nb2, st->cne1); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, chunk_idx); } } // --- HMX Dispatchers & Entry Points --- static void dequantize_tiled_weight_chunk_to_fp16_tiles( struct htp_context *ctx, __fp16 *vtcm_dst, const void *weight_src_ddr, int n_cols, int k_block, size_t row_stride, int weight_type, int n_k_tiles, struct fastdiv_values n_k_tiles_div, worker_callback_t dequant_worker_fn, int n_threads) { assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0); assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0); size_t n_col_tiles = n_cols / HTP_MM_HMX_TILE_N_COLS; size_t n_tot_tiles = n_col_tiles * n_k_tiles; size_t n_tiles_per_task = (n_threads == 1) ? n_tot_tiles : hmx_ceil_div(n_tot_tiles, n_threads); tiled_dequantize_state_t state; state.n_tasks = (n_tot_tiles + n_tiles_per_task - 1) / n_tiles_per_task; state.n_tot_tiles = n_tot_tiles; state.n_tiles_per_task = n_tiles_per_task; state.dst = vtcm_dst; state.src = (const uint8_t *)weight_src_ddr; state.n_cols = n_cols; state.k_block = k_block; state.row_stride = row_stride; state.weight_type = weight_type; state.n_k_tiles = n_k_tiles; state.n_k_tiles_div = n_k_tiles_div; state.traces = ctx->trace; state.ctx = ctx; state.tile_size = htp_mm_get_weight_tile_size(weight_type); state.aligned_tile_size = htp_mm_get_weight_aligned_tile_size(weight_type); if (state.n_tasks == 1 || n_threads == 1) { dequant_worker_fn(1, 0, &state); } else { int n_tasks = hex_smin((int) state.n_tasks, n_threads); worker_pool_run_func(ctx->worker_pool, dequant_worker_fn, &state, n_tasks); } } typedef struct { struct htp_context * ctx; struct htp_thread_trace * traces; float * dst; const __fp16 * vtcm_src; const float * src2; uint32_t n_rows; uint32_t n_cols; uint32_t dst_stride; uint32_t src2_stride; uint32_t dst_cols; struct fastdiv_values n_threads_div; } output_transfer_col_chunk_state_t; static void transfer_output_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { (void) n; output_transfer_col_chunk_state_t *st = (output_transfer_col_chunk_state_t *) data; struct htp_thread_trace * tr = &st->traces[i]; uint32_t n_blocks = st->n_cols / 32; uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div); uint32_t b_last = fastdiv(n_blocks * (i + 1), &st->n_threads_div); uint32_t c_first = b_first * 32; uint32_t c_last = b_last * 32; uint32_t c_len = c_last - c_first; if (c_len == 0) return; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first); const __fp16 *vtcm_src = st->vtcm_src + b_first * HTP_MM_HMX_TILE_N_ELMS; const float *src2 = st->src2 ? (st->src2 + c_first) : NULL; float *dst = st->dst + c_first; int chunk_dst_cols = (int)st->dst_cols - (int)c_first; if (chunk_dst_cols > 0) { transfer_output_chunk_fp16_to_fp32_col_chunk( dst, src2, vtcm_src, 0, st->n_rows, c_len, st->n_cols, st->dst_stride, st->src2_stride, (uint32_t)chunk_dst_cols ); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, c_first); } static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst, const float *src2, const __fp16 *vtcm_src, int n_rows, int n_cols, int dst_stride, uint32_t src2_stride, int dst_cols, int n_threads) { assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0); if (n_rows <= 0) return; uint32_t n_blocks = (uint32_t)n_cols / 32; if (n_threads > 1 && n_blocks >= (uint32_t)n_threads) { struct fastdiv_values n_threads_div = (n_threads == (int)ctx->n_threads) ? ctx->n_threads_div : init_fastdiv_values(n_threads); output_transfer_col_chunk_state_t col_state; col_state.dst = dst; col_state.src2 = src2; col_state.vtcm_src = vtcm_src; col_state.n_rows = (uint32_t)n_rows; col_state.n_cols = (uint32_t)n_cols; col_state.dst_stride = (uint32_t)dst_stride; col_state.src2_stride = src2_stride; col_state.dst_cols = (uint32_t)dst_cols; col_state.n_threads_div = n_threads_div; col_state.traces = ctx->trace; col_state.ctx = ctx; worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_col_chunk_worker_fn, &col_state, n_threads); return; } size_t n_tot_chunks = n_rows; size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : hmx_ceil_div(n_rows, n_threads); n_chunks_per_task = hex_align_up(n_chunks_per_task, 2); int actual_threads = hmx_ceil_div(n_rows, n_chunks_per_task); output_transfer_task_state_t state; state.n_tasks = actual_threads; state.n_tot_chunks = n_tot_chunks; state.n_chunks_per_task = n_chunks_per_task; state.dst = dst; state.src2 = src2; state.vtcm_src = vtcm_src; state.n_cols = n_cols; state.dst_stride = dst_stride; state.src2_stride = src2_stride; state.dst_cols = dst_cols; state.traces = ctx->trace; if (actual_threads <= 1) { transfer_output_chunk_worker_fn(1, 0, &state); } else { worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_worker_fn, &state, actual_threads); } } struct activation_transfer_params { struct htp_context * ctx; __fp16 * dst; const float * src; int n_rows; int k_block; int k_stride; int n_threads; const struct fastdiv_values * act_threads_div; const struct fastdiv_values * k_div; int k_valid; float * vtcm_f32_act; size_t vtcm_f32_act_bytes; }; static void transfer_activation_chunk_threaded(const struct activation_transfer_params * params) { struct htp_context * ctx = params->ctx; __fp16 * dst = params->dst; const float * src = params->src; int n_rows = params->n_rows; int k_block = params->k_block; int k_stride = params->k_stride; int n_threads = params->n_threads; const struct fastdiv_values * act_threads_div = params->act_threads_div; const struct fastdiv_values * k_div = params->k_div; int k_valid = params->k_valid; float * vtcm_f32_act = params->vtcm_f32_act; size_t vtcm_f32_act_bytes = params->vtcm_f32_act_bytes; if (n_rows <= 0) { return; } const size_t n_tasks = (n_rows + 31) >> 5; if (n_threads > 1 && k_block > 32 && n_tasks < (size_t) n_threads) { // Calculate step rows parameters for column-chunked dma pipelining uint32_t dma_step_rows = 2; uint32_t dma_step_rows_shift = 1; if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) { size_t thread_scratch_bytes = hex_align_down(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128); size_t thread_scratch_elements = thread_scratch_bytes / sizeof(float); size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div); if (dma_step_rows_max >= 4) { dma_step_rows = 4; dma_step_rows_shift = 2; } } activation_transfer_col_chunk_state_t col_state; col_state.dst = dst; col_state.src = src; col_state.n_rows = n_rows; col_state.k_block = k_block; col_state.k_stride = k_stride; col_state.k_valid = k_valid; col_state.n_col_chunks = n_threads; col_state.n_threads_div = *act_threads_div; col_state.vtcm_f32_act = vtcm_f32_act; col_state.vtcm_f32_act_bytes = vtcm_f32_act_bytes; col_state.traces = ctx->trace; col_state.ctx = ctx; col_state.dma_step_rows = dma_step_rows; col_state.dma_step_rows_shift = dma_step_rows_shift; worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_col_chunk_worker_fn, &col_state, n_threads); return; } assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0 && k_stride % HTP_MM_HMX_TILE_N_COLS == 0); size_t n_tot_chunks = n_rows; size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : 32; // must be multiple of 32 to ensure correct destination address activation_transfer_task_state_t state; state.n_tasks = (n_threads == 1) ? 1 : hmx_ceil_div(n_tot_chunks, 32); state.n_tot_chunks = n_tot_chunks; state.n_chunks_per_task = n_chunks_per_task; state.dst = dst; state.src = src; state.k_block = k_block; state.k_stride = k_stride; state.k_valid = k_valid; state.traces = ctx->trace; state.ctx = ctx; state.vtcm_f32_act = vtcm_f32_act; state.vtcm_f32_act_bytes_per_thread = hex_align_down(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128); uint32_t dma_step_rows = 2; uint32_t dma_step_rows_shift = 1; if (vtcm_f32_act && state.vtcm_f32_act_bytes_per_thread > 0 && k_block > 0) { size_t thread_scratch_elements = state.vtcm_f32_act_bytes_per_thread / sizeof(float); size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div); if (dma_step_rows_max >= 4) { dma_step_rows = 4; dma_step_rows_shift = 2; } } state.dma_step_rows = dma_step_rows; state.dma_step_rows_shift = dma_step_rows_shift; int active_threads = hex_smin(n_threads, (int)state.n_tasks); if (state.n_tasks == 1 || n_threads == 1) { transfer_activation_chunk_worker_fn(1, 0, &state); } else { worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, active_threads); } } // --- Async HMX matmul job (for pipeline overlap) --- typedef struct { __fp16 * output; const __fp16 * activation; const __fp16 * weight; const __fp16 * scales; uint32_t n_row_tiles; uint32_t n_col_tiles; uint32_t n_dot_tiles; } hmx_matmul_job_t; static void hmx_matmul_worker_fn(void * data) { hmx_matmul_job_t * job = (hmx_matmul_job_t *) data; FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles); core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles); } static inline void hmx_matmul_job_init(hmx_matmul_job_t * job, __fp16 * output, const __fp16 * activation, const __fp16 * weight, const __fp16 * scales, uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { job->output = output; job->activation = activation; job->weight = weight; job->scales = scales; job->n_row_tiles = n_row_tiles; job->n_col_tiles = n_col_tiles; job->n_dot_tiles = n_dot_tiles; } static int hmx_mm_2d_f32(struct htp_context *ctx, float *restrict dst, const float *restrict src2, const float *activation, const uint8_t *weight, int m, int k, int n, int act_stride, int weight_stride, int weight_type, int k_valid, int dst_stride, uint32_t src2_stride, int dst_cols, int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div, int tile_size, int aligned_tile_size, int vtcm_size) { struct htp_thread_trace * tr = &ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); if (k % 32 != 0 || n % 32 != 0) { return -1; } if (!hex_is_aligned(dst, VLEN) || !hex_is_aligned(activation, VLEN)) { return -1; } size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); if (row_stride == 0) { return -1; } 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 -1; } 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 vec_dot_size = k * sizeof(__fp16); const size_t vtcm_budget = ctx->vtcm_size; const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride; const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; size_t m_chunk_n_rows = m_chunk; size_t n_chunk_n_cols = n_chunk; size_t vtcm_used = vtcm_size; const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; struct htp_mm_hmx_vtcm_layout L; htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size); vtcm_used = L.total_bytes; if (vtcm_used > vtcm_budget) { FARF(ERROR, "hmx-mm-2d-precomputed: VTCM overflow: used %zu budget %zu, m %d k %d n %d mc %zu nc %zu", vtcm_used, vtcm_budget, m, k, n, m_chunk_n_rows, n_chunk_n_cols); return -1; } 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-2d: m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu", m, k, n, weight_type, m_chunk_n_rows, n_chunk_n_cols, vtcm_used, vtcm_budget); int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); if (pipeline) { // --- Asynchronous Pipelined Loop --- hmx_matmul_job_t job_slots[2]; // persistent double-buffered job descriptors for (size_t mr = 0; mr < 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); // Prologue: push A0 and optionally A1 (if n_chunk_cnt > 1) const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0; dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_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); } // Main loop: pop A_i -> dequantize A_i -> push A_{i+2} -> submit C_i -> wait C_{i-1} and store D_{i-1} 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); // 1. pop A_i void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; // 2. dequantize A_i 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); // 3. push A_{i+2} (if i+2 < n_chunk_cnt) if (i + 2 < n_chunk_cnt) { const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2; dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_p2 * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2); } // 4. submit C_i 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])); // 5. wait C_{i-1} and store D_{i-1} (multi-thread HVX, parallel with C_i) 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 + (mr * dst_stride + nc_prev); const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_prev) : NULL; int chunk_dst_cols = dst_cols - (int)nc_prev; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(i - 1) % 2], n_rows, n_cols_prev, dst_stride, src2_stride, chunk_dst_cols, n_threads); } } } // Epilogue: wait C_{last} and store D_{last} 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 + (mr * dst_stride + nc_last); const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_last) : NULL; int chunk_dst_cols = dst_cols - (int)nc_last; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(n_chunk_cnt - 1) % 2], n_rows, n_cols_last, dst_stride, src2_stride, chunk_dst_cols, n_threads); } } } else { // --- Synchronous loop (m <= 32 or fallback) --- hmx_matmul_job_t job; for (size_t mr = 0; mr < 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); // A0: Pre-fetch the first weight chunk (nc = 0) 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); // A: Wait for weight DMA void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; // B: Weight Dequantize (Threaded) 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); // Start weight DMA for the next chunk early 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); } // C: HMX Compute (Queue-based) 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); // D: Output Store float *output_chunk = dst + (mr * dst_stride + nc); const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc) : NULL; int chunk_dst_cols = dst_cols - (int)nc; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output, n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads); } } } } return 0; } static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) { struct htp_context * ctx = octx->ctx; struct htp_thread_trace * tr = &ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); const uint32_t n_weights = kparams->n_weights; if (n_weights == 0 || n_weights > HTP_OP_MAX_OUTPUTS) { return HTP_STATUS_INVAL_PARAMS; } const struct htp_tensor * restrict src0 = octx->src[0]; const struct htp_tensor * restrict act = octx->src[n_weights]; 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 int m_start = 0; int m_rows = m; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(octx->dsts[0], sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); m_start = (int) range.start; m_rows = (int) range.count; } if (m_rows == 0) { return HTP_STATUS_OK; } FARF(HIGH, "hmx-mm-nx-2d: n_weights %u m %d (%d..%d) k %d wtype %d mc %d nc %d vtcm %zu/%zu", n_weights, m, m_start, m_start + m_rows, 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); const size_t mr_end = (size_t)(m_start + m_rows); if (pipeline) { hmx_matmul_job_t job_slots[2]; for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(mr_end - 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 = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows); struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, .src = activation + mr * act_stride, .n_rows = (int) n_rows, .k_block = k, .k_stride = act_stride, .n_threads = act_threads, .act_threads_div = act_threads_div, .k_div = k_div, .k_valid = k_valid, .vtcm_f32_act = vtcm_f32_act, .vtcm_f32_act_bytes = L.act_f32_bytes, }; transfer_activation_chunk_threaded(&act_params); for (uint32_t p = 0; p < n_weights; p++) { const struct htp_tensor * restrict src_w = octx->src[p]; const struct htp_tensor * restrict dst = octx->dsts[p]; if (!src_w || !dst) continue; const uint8_t * weight = (const uint8_t *) src_w->data; float * dst_ptr = (float *) dst->data; const size_t n = src_w->ne[1]; if (n == 0) continue; const size_t weight_stride = src_w->nb[1]; const size_t dst_stride = dst->nb[1] / sizeof(float); const int dst_cols = (int) dst->ne[0]; const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; if (n > 0) { const size_t n_cols = hex_smin(n, n_chunk_n_cols); const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols; dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; dequantize_tiled_weight_chunk_to_fp16_tiles( ctx, vtcm_scratch0, curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); const size_t nc_next = nc + n_chunk_n_cols; if (nc_next < n) { const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); hmx_queue_pop(ctx->hmx_queue); float *output_chunk = dst_ptr + (mr * dst_stride + nc); int chunk_dst_cols = dst_cols - (int)nc; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output_chunk, NULL, vtcm_output, n_rows, n_cols, dst_stride, 0, chunk_dst_cols, n_threads); } } } } } return HTP_STATUS_OK; } static inline const __fp16 *hmx_mm_weight_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, int dst_b2, int dst_b3) { const size_t b2_idx = (params->r2 <= 1) ? (size_t) dst_b2 : (size_t) fastdiv((uint32_t) dst_b2, ¶ms->div_r2); const size_t b3_idx = (params->r3 <= 1) ? (size_t) dst_b3 : (size_t) fastdiv((uint32_t) dst_b3, ¶ms->div_r3); return (const __fp16 *) ((const uint8_t *) params->weight + b2_idx * params->src0_nb2 + b3_idx * params->src0_nb3); } static inline const float *hmx_mm_activation_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, int dst_b2, int dst_b3) { return (const float *) ((const uint8_t *) params->activation + (size_t) dst_b2 * params->src1_nb2 + (size_t) dst_b3 * params->src1_nb3); } static inline float *hmx_mm_dst_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, int dst_b2, int dst_b3) { return (float *) ((uint8_t *) params->dst + (size_t) dst_b2 * params->dst_nb2 + (size_t) dst_b3 * params->dst_nb3); } static inline const float *hmx_mm_src2_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, int src2_b2, int src2_b3) { return params->src2 ? (const float *) ((const uint8_t *) params->src2 + (size_t) src2_b2 * params->src2_nb2 + (size_t) src2_b3 * params->src2_nb3) : NULL; } static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, const hmx_mm_f16_f32_batched_params_t *params, int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size, const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div) { int ret = 0; for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) { for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) { ret = hmx_mm_2d_f32(ctx, hmx_mm_dst_batch_ptr(params, b2, b3), hmx_mm_src2_batch_ptr(params, b2, b3), hmx_mm_activation_batch_ptr(params, b2, b3), (const uint8_t *)hmx_mm_weight_batch_ptr(params, b2, b3), params->m, params->k, params->n, params->act_stride, params->weight_stride * (int)sizeof(__fp16), HTP_TYPE_F16, params->k, params->dst_stride, params->src2_stride, params->n, m_chunk, n_chunk, pipeline, n_threads, act_threads, act_threads_div, k_div, 0, 0, vtcm_size); } } return ret; } static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_batched_params_t *params, int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div, int vtcm_size) { if (params->act_stride < params->k || params->weight_stride < params->k || params->dst_stride < params->n) { return -1; } if (params->ne02 <= 0 || params->ne03 <= 0 || params->ne12 <= 0 || params->ne13 <= 0) { return -1; } if (params->ne12 % params->ne02 != 0 || params->ne13 % params->ne03 != 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; } const int group_size = params->r2; const size_t vtcm_budget = ctx->vtcm_size; // Check if the precomputed parameters are grouped or simple. // If simple, or if group_size <= 1, we use simple fallback loop. // Grouped path is only valid if group_size > 1 and it fits within VTCM budget. bool run_grouped = (group_size > 1 && (size_t)vtcm_size <= vtcm_budget); if (!run_grouped) { return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div); } struct htp_thread_trace * tr = &ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); const size_t vec_dot_size = params->k * sizeof(__fp16); const bool use_dma_activation = (params->act_stride > params->k); const size_t f32_scratch_size = use_dma_activation ? hex_align_up((size_t)act_threads * HTP_MM_DMA_ACT_MULTIPLIER * (size_t) params->k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0; size_t m_chunk_n_rows = m_chunk; size_t n_chunk_n_cols = n_chunk; size_t vtcm_used = vtcm_size; struct htp_mm_hmx_vtcm_layout L; htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0); if (L.total_bytes > vtcm_budget) { FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div); } uint8_t * const base = (uint8_t *) ctx->vtcm_base; __fp16 *vtcm_weight = VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]); __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act); __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(void, base, L.off_scratch[1]); __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales); float *vtcm_f32_act = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_act_f32, use_dma_activation); hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16 FARF(HIGH, "%s: grouped path m=%d k=%d n=%d group=%d streams=%d mc=%zu nc=%zu vtcm=%zu/%zu", __func__, params->m, params->k, params->n, group_size, params->ne13, m_chunk_n_rows, n_chunk_n_cols, L.total_bytes, vtcm_budget); const size_t fp16_row_bytes = (size_t) params->k * sizeof(__fp16); const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); hmx_matmul_job_t job; for (int b3 = 0; b3 < params->ne13; ++b3) { for (int b2_base = 0; b2_base < params->ne12; b2_base += group_size) { const __fp16 *weight_group = hmx_mm_weight_batch_ptr(params, b2_base, b3); for (size_t mr = 0; mr < (size_t) params->m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin((size_t) params->m - mr, m_chunk_n_rows); const size_t n_row_tiles = hmx_ceil_div((int) n_rows, HTP_MM_HMX_TILE_N_ROWS); // Pre-load activations for all heads in the group (once per m_chunk). // When the source is strided (permuted Q), use 2D DMA to gather // contiguous rows into a VTCM scratch buffer first, then HVX // converts from the contiguous VTCM buffer. This avoids L2 cache // thrashing from HVX loads at large strides. for (int g = 0; g < group_size; ++g) { const float *activation_chunk = hmx_mm_activation_batch_ptr(params, b2_base + g, b3) + mr * params->act_stride; __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_act_g, .src = activation_chunk, .n_rows = (int) n_rows, .k_block = params->k, .k_stride = params->act_stride, .n_threads = act_threads, .act_threads_div = act_threads_div, .k_div = k_div, .k_valid = params->k, .vtcm_f32_act = vtcm_f32_act, .vtcm_f32_act_bytes = L.act_f32_bytes, }; transfer_activation_chunk_threaded(&act_params); } // Prologue: Push A0 and A1 (if exists) { const size_t n_cols_first = hex_smin((size_t) params->n, n_chunk_n_cols); dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch0, weight_group), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_first); } if (n_chunk_n_cols < (size_t) params->n) { const size_t n_cols_second = hex_smin((size_t) params->n - n_chunk_n_cols, n_chunk_n_cols); dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch1, weight_group + params->weight_stride), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_second); } for (size_t nc = 0; nc < (size_t) params->n; nc += n_chunk_n_cols) { const size_t n_cols = hex_smin((size_t) params->n - nc, n_chunk_n_cols); const size_t n_col_tiles = hmx_ceil_div((int) n_cols, HTP_MM_HMX_TILE_N_COLS); { void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) curr_raw, n_cols, params->k, params->k, 0, n_cols); const size_t nc_next = nc + n_chunk_n_cols * 2; if (nc_next < (size_t) params->n) { const size_t n_cols_next = hex_smin((size_t) params->n - nc_next, n_chunk_n_cols); const __fp16 *next_weight_chunk = weight_group + nc_next * params->weight_stride; dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, next_weight_chunk), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_next); } } // Reuse the interleaved weight for every q_head in this GQA group for (int g = 0; g < group_size; ++g) { { const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride; hmx_matmul_job_init(&job, vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); hmx_queue_pop(ctx->hmx_queue); } { float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc; const float *src2_chunk = params->src2 ? (hmx_mm_src2_batch_ptr(params, b2_base + g, b3) + mr * params->src2_stride + nc) : NULL; int chunk_dst_cols = params->n - (int)nc; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, params->src2_stride, chunk_dst_cols, n_threads); } } } } } } } return 0; } static void transfer_activation_chunk_gathered_threaded( struct htp_context *ctx, __fp16 *dst, const float *src, int start_row, int n_rows, int k_block, const struct mmid_row_mapping *matrix_rows, int cur_a, int mapping_stride, int ne11, size_t nb11, size_t nb12, int cne1, int n_threads, int k_valid) { if (n_rows <= 0) return; int chunks_per_thread = hmx_ceil_div(n_rows, n_threads); chunks_per_thread = hex_align_up(chunks_per_thread, 2); int actual_threads = hmx_ceil_div(n_rows, chunks_per_thread); activation_transfer_gathered_task_state_t state = { .dst = dst, .src = src, .n_tasks = actual_threads, .n_tot_chunks = n_rows, .n_chunks_per_task = chunks_per_thread, .k_block = k_block, .matrix_rows = matrix_rows, .cur_a = cur_a, .mapping_stride = mapping_stride, .ne11 = ne11, .ne11_div = ne11 > 1 ? init_fastdiv_values(ne11) : (struct fastdiv_values){0, 0}, .nb11 = nb11, .nb12 = nb12, .start_row = start_row, .cne1 = cne1, .k_valid = k_valid, .traces = ctx->trace, }; worker_callback_t worker_fn = ne11 == 1 ? transfer_activation_chunk_gathered_worker_flat_fn : transfer_activation_chunk_gathered_worker_fn; if (actual_threads <= 1) { worker_fn(1, 0, &state); } else { worker_pool_run_func(ctx->worker_pool, worker_fn, &state, actual_threads); } } static void transfer_output_chunk_scattered_threaded( struct htp_context *ctx, float *dst, const __fp16 *vtcm_src, int start_row, int n_rows, int n_cols, const struct mmid_row_mapping *matrix_rows, int cur_a, int mapping_stride, size_t dst_nb1, size_t dst_nb2, int cne1, int n_threads) { if (n_rows <= 0) return; int chunks_per_thread = hmx_ceil_div(n_rows, n_threads); chunks_per_thread = hex_align_up(chunks_per_thread, 2); int actual_threads = hmx_ceil_div(n_rows, chunks_per_thread); output_transfer_scattered_task_state_t state = { .vtcm_src = vtcm_src, .dst = dst, .n_tasks = actual_threads, .n_tot_chunks = n_rows, .n_chunks_per_task = chunks_per_thread, .n_cols = n_cols, .matrix_rows = matrix_rows, .cur_a = cur_a, .mapping_stride = mapping_stride, .dst_nb1 = dst_nb1, .dst_nb2 = dst_nb2, .start_row = start_row, .cne1 = cne1, .traces = ctx->trace, }; if (actual_threads <= 1) { transfer_output_chunk_scattered_worker_fn(1, 0, &state); } else { worker_pool_run_func(ctx->worker_pool, transfer_output_chunk_scattered_worker_fn, &state, actual_threads); } } static int hmx_mm_id_2d_f32(struct htp_context *ctx, float *restrict dst, const float *activation, const uint8_t *weight, int m, int k, int n, int k_valid, int ne11, size_t act_nb1, size_t act_nb2, size_t dst_nb1, size_t dst_nb2, int weight_stride, int weight_type, const struct mmid_row_mapping *matrix_rows, int cur_a, int mapping_stride, int m_start, int m_end, int n_threads) { struct htp_thread_trace * tr = &ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); const int cne1 = m; const int m_padded = hex_align_up(m, 32); if (k % 32 != 0 || n % 32 != 0) { return -1; } if (!hex_is_aligned(dst, VLEN) || !hex_is_aligned(activation, VLEN)) { return -1; } size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); if (row_stride == 0) { return -1; } 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 -1; } 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 vec_dot_size = k * sizeof(__fp16); const size_t vtcm_budget = ctx->vtcm_size; size_t vtcm_used = 0; int tile_size = htp_mm_get_weight_tile_size(weight_type); int aligned_tile_size = htp_mm_get_weight_aligned_tile_size(weight_type); const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride; const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; const size_t weight_row_stride = is_quant ? qweight_row_stride : row_stride; size_t size_per_n = 0, size_per_m = 0, size_per_mn = 0; htp_mm_hmx_get_2d_chunk_costs(weight_type, k, /*pipeline=*/false, aligned_tile_size, &size_per_n, &size_per_m, &size_per_mn); const size_t overhead = htp_mm_hmx_get_2d_overhead(/*pipeline=*/false, /*is_matmul_id=*/true); size_t m_chunk_n_rows = 0, n_chunk_n_cols = 0; if (htp_mm_hmx_compute_chunks(vtcm_budget, overhead, size_per_n, size_per_m, size_per_mn, m_padded, n, /*m_block_cost=*/(size_t) n * HTP_MM_HMX_COST_W_DEQUANT, /*n_block_cost=*/(size_t) m_padded * HTP_MM_HMX_COST_A_CONVERT, &m_chunk_n_rows, &n_chunk_n_cols, &vtcm_used)) { FARF(ERROR, "hmx-mm-id-2d: VTCM too small : m %d k %d n %d budget %zu", m_padded, k, n, vtcm_budget); return -1; } const size_t weight_area_size = hex_align_up(n_chunk_n_cols * weight_row_stride, HTP_MM_HMX_TILE_SIZE); const size_t act_area_size = hex_align_up(m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE); const size_t output_area_size = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE); size_t scratch0_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE); uint8_t *vtcm_ptr = (uint8_t *) ctx->vtcm_base; __fp16 *vtcm_weight = weight_area_size ? (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size) : NULL; __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size); __fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size); void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch0_size); __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256); vtcm_used = vtcm_ptr - (uint8_t *) ctx->vtcm_base; if (vtcm_used > vtcm_budget) { FARF(ERROR, "hmx-mm-id-2d: VTCM overflow: used %zu budget %zu", vtcm_used, vtcm_budget); return -1; } hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); hmx_matmul_job_t job; for (size_t mr = (size_t) m_start; mr < (size_t) m_end; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin((size_t) m_end - mr, m_chunk_n_rows); const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); transfer_activation_chunk_gathered_threaded( ctx, vtcm_f16_act, activation, (int) mr, (int) n_rows, k, matrix_rows, cur_a, mapping_stride, ne11, act_nb1, act_nb2, cne1, n_threads, k_valid); // A0: Pre-fetch the first weight chunk (nc = 0) if (n > 0) { const size_t n_cols = hex_smin((size_t) n, n_chunk_n_cols); const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols; dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } for (size_t nc = 0; nc < (size_t) n; nc += n_chunk_n_cols) { const size_t n_cols = hex_smin((size_t) n - nc, n_chunk_n_cols); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); // A: Wait for weight DMA void * curr_raw = dma_queue_pop(ctx->dma[0]).dst; // B: Weight Dequantize (Threaded) 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 ); // Start weight DMA for the next chunk early const size_t nc_next = nc + n_chunk_n_cols; if (nc_next < (size_t) n) { const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols); const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next; dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } // C: HMX Compute (Queue-based) 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); // D: Output Store transfer_output_chunk_scattered_threaded( ctx, dst + nc, vtcm_output, (int) mr, (int) n_rows, (int) n_cols, matrix_rows, cur_a, mapping_stride, dst_nb1, dst_nb2, cne1, n_threads); } } return 0; } // --- Dispatchers and Public Entry Points --- static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) { htp_matmul_tensors_preamble; int k = (int) src0->ne[0]; int n = (int) src0->ne[1]; const int m_total = (int) src1->ne[1]; const int act_stride = (int)(src1->nb[1] / sizeof(float)); const int wgt_stride = (int)(src0->nb[1] / sizeof(__fp16)); int m_start = 0; int m_rows = m_total; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_total, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); m_start = (int) range.start; m_rows = (int) range.count; } if (m_rows == 0) { return HTP_STATUS_OK; } const float * src2_ptr = NULL; uint32_t src2_stride = 0; size_t src2_nb2 = 0; size_t src2_nb3 = 0; if (src2) { src2_stride = (src2->ne[1] == 1) ? 0 : (uint32_t) (src2->nb[1] / sizeof(float)); src2_ptr = (const float *) src2->data + m_start * src2_stride; src2_nb2 = (src2->ne[2] == 1) ? 0 : src2->nb[2]; src2_nb3 = (src2->ne[3] == 1) ? 0 : src2->nb[3]; } const int dst_stride = (int)(dst->nb[1] / sizeof(float)); float * dst_ptr = (float *) dst->data + m_start * dst_stride; const float * act_ptr = (const float *) src1->data + m_start * act_stride; int ret = -1; const int n_threads = kparams->n_threads; if (kparams->kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) { hmx_mm_f16_f32_batched_params_t batch_params = { .dst = dst_ptr, .src2 = src2_ptr, .activation = act_ptr, .weight = (const __fp16 *) src0->data, .m = m_rows, .k = k, .n = n, .act_stride = act_stride, .weight_stride = wgt_stride, .dst_stride = dst_stride, .src2_stride = src2_stride, .ne02 = ne02, .ne03 = ne03, .ne12 = ne12, .ne13 = ne13, .src0_nb2 = src0->nb[2], .src0_nb3 = src0->nb[3], .src1_nb2 = src1->nb[2], .src1_nb3 = src1->nb[3], .dst_nb2 = dst->nb[2], .dst_nb3 = dst->nb[3], .src2_nb2 = src2_nb2, .src2_nb3 = src2_nb3, .r2 = (ne02 > 0) ? (ne12 / ne02) : 1, .r3 = (ne03 > 0) ? (ne13 / ne03) : 1, .div_r2 = kparams->div_r2, .div_r3 = kparams->div_r3, }; ret = hmx_mm_f16_f32_batched(octx->ctx, &batch_params, kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads, kparams->n_act_threads, &kparams->div_n_act_threads, &kparams->div_ne00_padded, kparams->vtcm_size); } else { ret = hmx_mm_2d_f32( octx->ctx, dst_ptr, src2_ptr, act_ptr, (const uint8_t *) src0->data, m_rows, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0], dst_stride, src2_stride, (int)dst->ne[0], kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads, kparams->n_act_threads, &kparams->div_n_act_threads, &kparams->div_ne00_padded, kparams->tile_size, kparams->aligned_tile_size, kparams->vtcm_size ); } if (ret != 0) { FARF(ERROR, "HMX matmul failed (ret=%d)\n", ret); return HTP_STATUS_INTERNAL_ERR; } return HTP_STATUS_OK; } int op_matmul(struct htp_ops_context * octx) { const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const int status = htp_mm_init_context(octx, kparams); if (status != HTP_STATUS_OK) { return status; } if (kparams->n_hmx) { return hmx_mm_op_matmul(octx, kparams); } return hvx_mm_matmul(octx); } static int hmx_mm_op_matmul_id( 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; htp_matmul_tensors_preamble; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const int n_ids = octx->src[2]->ne[0]; const int n_as = ne02; 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; const int m_padded = hex_align_up(cne1, 32); int m_start = 0, m_end = m_padded; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_mdev_data_aligned(dst) && (uint32_t) cne1 >= octx->ctx->mdev.count; const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_padded, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); m_start = (int) range.start; m_end = (int) (range.start + range.count); } if (m_start >= m_end) continue; int ret = hmx_mm_id_2d_f32(octx->ctx, (float*) dst->data, (float*) src1->data, (const uint8_t *) src0->data + cur_a * nb02, cne1, ne00, ne01, ne10, ne11, nb11, nb12, nb1, nb2, (int) src0->nb[1], (int) src0->type, matrix_rows, cur_a, mmctx->mapping_stride, m_start, m_end, (int) octx->n_threads); if (ret != 0) { FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret); return HTP_STATUS_NO_SUPPORT; } } return HTP_STATUS_OK; } static int hvx_mm_matmul_id( struct htp_ops_context * octx, struct htp_mm_context * mmctx, work_queue_func_t hvx_mmid_task_func ) { htp_matmul_tensors_preamble; 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 struct htp_tensor * restrict ids = octx->src[2]; const size_t src0_row_size = nb01; const uint32_t qk = QK_Q8_0_TILED; const uint32_t nb = (ne10 + 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(ne10) : htp_mm_q8_0_tiled_row_size(ne10); struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false); const size_t vtcm_size = L.total_bytes; FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size 0 dst-spad-size %zu (%zu)\n", mmctx->type, L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size); FARF(HIGH, "matmul-id-%s : %ux%ux%ux%u * %ux%ux%ux%u (%ux%ux%ux%u) -> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], ids->ne[0], ids->ne[1], ids->ne[2], ids->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], src0->data, src1->data, dst->data); // Make sure the reserved vtcm size is sufficient if (octx->ctx->vtcm_size < vtcm_size) { FARF(ERROR, "matmul-id-%s : current VTCM reservation %zu is too small, needed %zu\n", mmctx->type, octx->ctx->vtcm_size, vtcm_size); return HTP_STATUS_VTCM_TOO_SMALL; } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); mmctx->vtcm_src2 = NULL; mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; octx->src2_spad.src = NULL; octx->dst_spad.src = NULL; mmctx->vtcm_src0_stride = src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_src2_size_per_thread = 0; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->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); 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 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; const int m_padded = hex_align_up(cne1, 32); int m_start = 0, m_end = m_padded; if (octx->ctx->mdev.count > 1) { bool can_split = (uint32_t) cne1 >= octx->ctx->mdev.count; for (uint32_t p = 0; p < n_weights && can_split; ++p) { const struct htp_tensor * restrict dst = octx->dsts[p]; can_split = !dst || htp_tensor_mdev_data_aligned(dst); } const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_padded, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); m_start = (int) range.start; m_end = (int) (range.start + range.count); } if (m_start >= m_end) 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, m_start, m_end, (int) octx->n_threads); 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); const size_t 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->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; mmctx->quant_task_func = quant_task_func; mmctx->n_quant_tasks = n_quant_tasks; atomic_init(&mmctx->quant_barrier, n_quant_tasks); FARF(HIGH, "matmul-id-nx: src0 %d:%d:%d type %s nrows %u, src1 %d:%d:%d nrows %u, vtcm %zu/%zu, threads %d\n", src0->ne[0], src0->ne[1], src0->ne[2], mmctx->type, src0->ne[1], act->ne[0], act->ne[1], act->ne[2], src1_nrows, L.total_bytes, octx->ctx->vtcm_size, octx->n_threads); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); worker_pool_run_func(octx->ctx->worker_pool, hvx_mmid_task_func, mmctx, octx->n_threads); return HTP_STATUS_OK; } static inline void scan_expert_ids_n( const struct htp_tensor * ids, const uint32_t n_ids, uint32_t n_as, uint32_t * counts, struct mmid_row_mapping * matrix_rows, uint32_t mapping_stride ) { const size_t ids_nb1 = ids->nb[1]; const uint8_t * ids_data = (const uint8_t *) ids->data; 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 < 0) { continue; } assert(i02 < n_as); if (matrix_rows) { matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 }; } counts[i02] += 1; } } } static inline void scan_expert_ids( const struct htp_tensor * ids, uint32_t n_ids, uint32_t n_as, uint32_t * counts, struct mmid_row_mapping * matrix_rows, uint32_t mapping_stride ) { const size_t ids_nb0 = ids->nb[0]; 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; } } else { // Strided fallback const size_t ids_nb1 = ids->nb[1]; const uint8_t * ids_data = (const uint8_t *) ids->data; 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 < 0) { continue; } assert(i02 < n_as); if (matrix_rows) { matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 }; } counts[i02] += 1; } } } } int op_matmul_id(struct htp_ops_context * octx) { htp_matmul_tensors_preamble; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; const int status = htp_mm_init_context(octx, kparams); if (status != HTP_STATUS_OK) { return status; } struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); mmctx->octx = octx; mmctx->act = src1; const struct htp_tensor * restrict ids = octx->src[2]; const size_t src0_row_size = nb01; const size_t dst_row_size = nb1; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); const uint32_t src0_nrows = ne01; // per expert const uint32_t src1_nrows = ne11 * ne12 * ne13; // row groups 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; 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(octx, mmctx); } else { uint32_t src0_row_start = 0; uint32_t src0_row_end = src0_nrows; if (octx->ctx->mdev.count > 1) { const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); src0_row_start = range.start; src0_row_end = range.start + range.count; } if (src0_row_start >= src0_row_end) { if (mapping_buf != octx->ctx->ddr_spad_base) { free(mapping_buf); } return HTP_STATUS_OK; } const uint32_t nrows = src0_row_end - src0_row_start; mmctx->src0_row_start = src0_row_start; mmctx->src0_row_end = src0_row_end; mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div); mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) { s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id); } else { s = HTP_STATUS_NO_SUPPORT; } } if (mapping_buf != octx->ctx->ddr_spad_base) { free(mapping_buf); } return s; } int op_matmul_id_nx(struct htp_ops_context * octx) { const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; const int status = htp_mm_init_context(octx, kparams); if (status != HTP_STATUS_OK) { return status; } struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); mmctx->octx = octx; 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]; 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 src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; 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; const int status = htp_mm_init_context(octx, kparams); if (status != HTP_STATUS_OK) { return status; } if (kparams->n_hmx) { return hmx_mm_nx_2d_f32(octx, kparams); } struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); const uint32_t n_weights = kparams->n_weights; const struct htp_tensor * restrict src0 = octx->src[0]; // first weight const struct htp_tensor * restrict act = octx->src[n_weights]; // activation x bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || src0->type == HTP_TYPE_MXFP4); struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; mmctx->act = act; const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; const size_t src0_row_size = src0->nb[1]; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) { return HTP_STATUS_NO_SUPPORT; } 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; worker_callback_t quant_task_func; uint32_t n_quant_tasks = 1; if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { n_quant_tasks = MIN(src1_nrows, octx->n_threads); quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat; } else 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; if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(act->ne[0]) : htp_mm_q8_0_flat_row_size(act->ne[0]); } else { src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? 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, false, true); const size_t vtcm_size = L.total_bytes; if (octx->ctx->vtcm_size < vtcm_size) { FARF(ERROR, "matmul-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 = is_repacked ? 0 : src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div); mmctx->vtcm_src1_size_per_thread = L.src1_bytes; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->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); // Run fused matmul const uint32_t n_matmul_jobs = octx->n_threads; worker_callback_t matmul_job_func; if (is_repacked) { if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0_flat; break; case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1_flat; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0_flat; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl_flat; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4_flat; break; default: return HTP_STATUS_NO_SUPPORT; } } else { switch (src0->type) { case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; } } } else { matmul_job_func = hvx_mm_nx_2d; } htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs); return HTP_STATUS_OK; }