diff --git a/common/arg.cpp b/common/arg.cpp index c3c312736..ac0e9b71c 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -720,9 +720,8 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context // model is required (except for server) // TODO @ngxson : maybe show a list of available models in CLI in this case - if (params.model.path.empty() - && !params.usage - && !params.completion) { + bool can_skip_model = params.usage || params.completion || !params.server_base.empty(); + if (!can_skip_model && params.model.path.empty()) { throw std::invalid_argument("error: --model is required\n"); } } @@ -1242,6 +1241,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.completion = true; } )); + add_opt(common_arg( + {"--server-base"}, "URL", + string_format("connect to this server instead of starting a new one, example: 'http://localhost:8080' (default: none)"), + [](common_params & params, const std::string & value) { + params.server_base = value; + } + ).set_examples({LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"--verbose-prompt"}, string_format("print a verbose prompt before generation (default: %s)", params.verbose_prompt ? "true" : "false"), @@ -2844,7 +2850,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.out_file = value; } ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_CVECTOR_GENERATOR, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_FINETUNE, - LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS})); + LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"-ofreq", "--output-frequency"}, "N", string_format("output the imatrix every N iterations (default: %d)", params.n_out_freq), diff --git a/common/common.h b/common/common.h index 996ec511d..295451970 100644 --- a/common/common.h +++ b/common/common.h @@ -645,6 +645,9 @@ struct common_params { std::map default_template_kwargs; + // CLI params + std::string server_base; // if set, connect to this server instead of starting a new one + // UI configs bool ui = true; bool ui_mcp_proxy = false; diff --git a/common/http.h b/common/http.h index e88bc6a5e..878ad1ce2 100644 --- a/common/http.h +++ b/common/http.h @@ -2,6 +2,16 @@ #include +#ifdef _WIN32 +#include +#include +#else +#include +#include +#include +#include +#endif + struct common_http_url { std::string scheme; std::string user; @@ -119,3 +129,63 @@ static std::pair common_http_client(const std: static std::string common_http_show_masked_url(const common_http_url & parts) { return parts.scheme + "://" + (parts.user.empty() ? "" : "****:****@") + common_http_format_host(parts.host) + parts.path; } + +static int common_http_get_free_port() { +#ifdef _WIN32 + WSADATA wsaData; + if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) { + return -1; + } + typedef SOCKET native_socket_t; +#define INVALID_SOCKET_VAL INVALID_SOCKET +#define CLOSE_SOCKET(s) closesocket(s) +#else + typedef int native_socket_t; +#define INVALID_SOCKET_VAL -1 +#define CLOSE_SOCKET(s) close(s) +#endif + + native_socket_t sock = socket(AF_INET, SOCK_STREAM, 0); + if (sock == INVALID_SOCKET_VAL) { +#ifdef _WIN32 + WSACleanup(); +#endif + return -1; + } + + struct sockaddr_in serv_addr; + std::memset(&serv_addr, 0, sizeof(serv_addr)); + serv_addr.sin_family = AF_INET; + serv_addr.sin_addr.s_addr = htonl(INADDR_ANY); + serv_addr.sin_port = htons(0); + + if (bind(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0) { + CLOSE_SOCKET(sock); +#ifdef _WIN32 + WSACleanup(); +#endif + return -1; + } + +#ifdef _WIN32 + int namelen = sizeof(serv_addr); +#else + socklen_t namelen = sizeof(serv_addr); +#endif + if (getsockname(sock, (struct sockaddr*)&serv_addr, &namelen) != 0) { + CLOSE_SOCKET(sock); +#ifdef _WIN32 + WSACleanup(); +#endif + return -1; + } + + int port = ntohs(serv_addr.sin_port); + + CLOSE_SOCKET(sock); +#ifdef _WIN32 + WSACleanup(); +#endif + + return port; +} diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 0dfdacf66..0376eb476 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -7299,6 +7299,13 @@ struct ggml_conv_2d_dw_params { int dilation_y; }; +static inline float ggml_conv_2d_dw_knl_f32(const char * data, int64_t i, ggml_type type) { + if (type == GGML_TYPE_F16) { + return GGML_FP16_TO_FP32(((const ggml_fp16_t *)data)[i]); + } + return ((const float *)data)[i]; +} + static void ggml_compute_forward_conv_2d_dw_cwhn( const ggml_compute_params * params, const ggml_tensor * src, @@ -7307,7 +7314,8 @@ static void ggml_compute_forward_conv_2d_dw_cwhn( const ggml_conv_2d_dw_params & p) { const int64_t c = p.channels; - const float * knl_data = (const float *)kernel->data; + const char * knl_data = (const char *)kernel->data; + const ggml_type knl_type = kernel->type; const int64_t rows_total = p.dst_h * p.batch; const int64_t rows_per_thread = (rows_total + params->nth - 1) / params->nth; @@ -7315,13 +7323,16 @@ static void ggml_compute_forward_conv_2d_dw_cwhn( const int64_t row_end = MIN(row_start + rows_per_thread, rows_total); #ifdef GGML_SIMD + int64_t c_pkg_end = 0; + int64_t pkg_size = GGML_F32_EPR; + if (knl_type == GGML_TYPE_F32) { #if defined(__ARM_FEATURE_SVE) - const int64_t pkg_size = svcntw(); + pkg_size = svcntw(); #else - const int64_t pkg_size = GGML_F32_EPR; + pkg_size = GGML_F32_EPR; #endif - const int64_t pkg_count = c / pkg_size; - const int64_t c_pkg_end = pkg_count * pkg_size; + c_pkg_end = (c / pkg_size) * pkg_size; + } #else const int64_t c_pkg_end = 0; #endif @@ -7335,7 +7346,6 @@ static void ggml_compute_forward_conv_2d_dw_cwhn( const int64_t src_x_base = dst_x * p.stride_x - p.pad_x; #ifdef GGML_SIMD - // Vectorized loop for (int64_t c_i = 0; c_i < c_pkg_end; c_i += pkg_size) { GGML_F32_VEC sum = GGML_F32_VEC_ZERO; for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { @@ -7348,7 +7358,8 @@ static void ggml_compute_forward_conv_2d_dw_cwhn( if (src_x < 0 || src_x >= p.src_w) { continue; } - GGML_F32_VEC k = GGML_F32_VEC_LOAD(knl_data + (knl_y * p.knl_w + knl_x) * c + c_i); + const float * kp = (const float *)knl_data + (knl_y * p.knl_w + knl_x) * c + c_i; + GGML_F32_VEC k = GGML_F32_VEC_LOAD(kp); GGML_F32_VEC s = GGML_F32_VEC_LOAD(src_data + (src_y * p.src_w + src_x) * c + c_i); sum = GGML_F32_VEC_FMA(sum, k, s); } @@ -7356,7 +7367,6 @@ static void ggml_compute_forward_conv_2d_dw_cwhn( GGML_F32_VEC_STORE(dst_data + c_i, sum); } #endif - // Scalar loop for (int64_t c_i = c_pkg_end; c_i < c; ++c_i) { float sum = 0.0f; for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { @@ -7369,7 +7379,7 @@ static void ggml_compute_forward_conv_2d_dw_cwhn( if (src_x < 0 || src_x >= p.src_w) { continue; } - sum += knl_data[(knl_y * p.knl_w + knl_x) * c + c_i] + sum += ggml_conv_2d_dw_knl_f32(knl_data, (knl_y * p.knl_w + knl_x) * c + c_i, knl_type) * src_data[(src_y * p.src_w + src_x) * c + c_i]; } } @@ -7390,9 +7400,11 @@ static void ggml_compute_forward_conv_2d_dw_whcn( const int64_t per_thread = (n + params->nth - 1) / params->nth; const int64_t start = params->ith * per_thread; const int64_t end = MIN(start + per_thread, n); + const char * knl_base = (const char *)kernel->data; + const ggml_type knl_type = kernel->type; for (int64_t i = start; i < end; ++i) { - const float * knl_data = (const float *)kernel->data + (i % p.channels) * p.knl_w * p.knl_h; + const int64_t knl_offset = (i % p.channels) * p.knl_w * p.knl_h; const float * src_data = (const float *)src->data + i * p.src_w * p.src_h; float * dst_data = (float *)dst->data + i * p.dst_w * p.dst_h; @@ -7410,7 +7422,7 @@ static void ggml_compute_forward_conv_2d_dw_whcn( if (src_x < 0 || src_x >= p.src_w) { continue; } - sum += knl_data[knl_y * p.knl_w + knl_x] + sum += ggml_conv_2d_dw_knl_f32(knl_base, knl_offset + knl_y * p.knl_w + knl_x, knl_type) * src_data[src_y * p.src_w + src_x]; } } @@ -7442,13 +7454,13 @@ void ggml_compute_forward_conv_2d_dw( p.dilation_x = dst->op_params[4]; p.dilation_y = dst->op_params[5]; + GGML_ASSERT(kernel->type == GGML_TYPE_F32 || kernel->type == GGML_TYPE_F16); GGML_ASSERT(kernel->ne[3] == p.channels); GGML_ASSERT(dst->ne[3] == p.batch); if (ggml_is_contiguous(src)) { ggml_compute_forward_conv_2d_dw_whcn(params, src, kernel, dst, p); } else if (ggml_is_contiguous_channels(src)) { - // kernel should also have channels most contiguous in memory GGML_ASSERT(kernel->nb[0] >= kernel->nb[2] && kernel->nb[1] >= kernel->nb[0]); ggml_compute_forward_conv_2d_dw_cwhn(params, src, kernel, dst, p); } else { diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index c4f08091e..26af90025 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -28,6 +28,20 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr #endif // STRIDED_ITERATOR_AVAILABLE #ifdef GGML_CUDA_USE_CUB + +// returns the suggested maximum number of rows to process during one argsort_f32_i32_cuda_cub() call +int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows) { + // perform argsort in chunks up to approximately this size (currently 64MB) + // to avoid excessive temporary buffers memory usage + const int chunk_bytes = 1 << 26; + + // calculate how many rows will fit in one chunk (must be at least one) + const int chunk_nrows = std::max((int) (chunk_bytes / nb01), 1); + + // limit the resulting amount to total nrows + return std::min((int64_t) chunk_nrows, nrows); +} + void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -254,11 +268,23 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; - if (shared_mem > max_shared_mem || ncols > 1024) { - ggml_cuda_pool & pool = ctx.pool(); - argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream); - } else { + // early return if we can use bitonic argsort + if (shared_mem <= max_shared_mem && ncols <= 1024) { argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); + return; + } + + const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows); + + ggml_cuda_pool & pool = ctx.pool(); + + for (int64_t i = 0; i < nrows; i += chunk_nrows) { + int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i); + + argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, iter_nrows, order, stream); + + src0_d += ncols * iter_nrows; + dst_d += ncols * iter_nrows; } #else argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream); diff --git a/ggml/src/ggml-cuda/argsort.cuh b/ggml/src/ggml-cuda/argsort.cuh index 22b7306f2..3abb6448a 100644 --- a/ggml/src/ggml-cuda/argsort.cuh +++ b/ggml/src/ggml-cuda/argsort.cuh @@ -3,6 +3,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst); #ifdef GGML_CUDA_USE_CUB +int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows); void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 928c1965d..490bcf885 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3173,18 +3173,21 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph (a->ne[2] == 1 && a->ne[3] == 1); const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; - // x must be in the supported whitelist and every operand / intermediate - // result must share x's type, since launch_snake casts a / inv_b as - // float and templates the kernel on a single T. Mixed precision chains - // fall back to the naive path. + // x is in the supported whitelist and every chain intermediate shares + // x's type. launch_snake reads a and inv_b as const float *, so they + // stay F32. const ggml_tensor * sin1 = cgraph->nodes[i + 1]; const bool types_ok = (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && - (a->type == x->type) && (inv_b->type == x->type) && + (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && (mul0->type == x->type) && (sin1->type == x->type) && (sqr->type == x->type) && (mul1->type == x->type) && (add->type == x->type); - if (types_ok && shape_ok && dim_ok && x_in_add == x) { + // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous + const bool contig_ok = ggml_is_contiguous(x) && ggml_is_contiguous(add) && + ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); + + if (types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x) { ggml_cuda_op_snake_fused(*cuda_ctx, x, a, inv_b, add); return 4; } @@ -4722,10 +4725,16 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g } break; case GGML_OP_SET_ROWS: { - return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || - op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || - op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) && - op->src[0]->type == GGML_TYPE_F32 && + return ( + ( + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || + op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || + op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) && + op->src[0]->type == GGML_TYPE_F32 + ) || ( + op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F16 + ) + ) && (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); } break; case GGML_OP_SET: @@ -4921,7 +4930,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_IM2COL: case GGML_OP_IM2COL_3D: case GGML_OP_CONV_2D: + return true; case GGML_OP_CONV_2D_DW: + return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_CONV_TRANSPOSE_2D: case GGML_OP_POOL_2D: return true; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index a48cc48b2..e18ada537 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -549,8 +549,8 @@ static __global__ void mul_mat_vec_q( [[maybe_unused]] float x_biases[ncols_dst] = { 0.0f }; [[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f }; - [[maybe_unused]] float x_scales; - [[maybe_unused]] float gate_scales; + [[maybe_unused]] float x_scales = 1.0f; + [[maybe_unused]] float gate_scales = 1.0f; if constexpr (has_fusion) { // 1. Hide latency by prefetching bias, gates and scales here // 2. load only on threads that won't die after partial sum calculation @@ -655,47 +655,38 @@ static __global__ void mul_mat_vec_q( tmp_gate[j][i] = warp_reduce_sum(tmp_gate[j][i]); } } - } - if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { - float result = tmp[j][threadIdx.x]; - if constexpr (has_fusion) { - if constexpr (type == GGML_TYPE_NVFP4) { - if (use_scale) { + if (threadIdx.x == i && (rows_per_cuda_block == 1 || uint32_t(row0 + i) < stride_col_dst)) { + float result = tmp[j][i]; + if constexpr (has_fusion) { + if constexpr (type == GGML_TYPE_NVFP4) { result *= x_scales; } - } - if (use_bias) { result += x_biases[j]; - } - if (use_gate) { - float gate_value = tmp_gate[j][threadIdx.x]; - if constexpr (type == GGML_TYPE_NVFP4) { - if (use_gate_scale) { + if (use_gate) { + float gate_value = tmp_gate[j][i]; + if constexpr (type == GGML_TYPE_NVFP4) { gate_value *= gate_scales; } - } - if (use_gate_bias) { gate_value += gate_biases[j]; - } - switch (active_glu) { - case GGML_GLU_OP_SWIGLU: - result *= ggml_cuda_op_silu_single(gate_value); - break; - case GGML_GLU_OP_GEGLU: - result *= ggml_cuda_op_gelu_single(gate_value); - break; - case GGML_GLU_OP_SWIGLU_OAI: { - result = ggml_cuda_op_swiglu_oai_single(gate_value, result); - break; + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + default: + result = result * gate_value; + break; } - default: - result = result * gate_value; - break; } } + dst[j*stride_col_dst + i] = result; } - dst[j*stride_col_dst + threadIdx.x] = result; } } diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu index 3b4f004c9..465997065 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/ggml/src/ggml-cuda/set-rows.cu @@ -322,17 +322,77 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s } } +template<> +void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const half * src0_d = (const half *)src0->data; + const int32_t * src1_d = (const int32_t *)src1->data; + + GGML_TENSOR_BINARY_OP_LOCALS + + cudaStream_t stream = ctx.stream(); + + + if (dst->type == GGML_TYPE_F16) { + set_rows_cuda( + src0_d, src1_d, (half*)dst->data, + ne00, ne01, ne02, ne03, + ne10, ne11, ne12, ne13, + nb01, nb02, nb03, + nb10, nb11, nb12, + nb1, nb2, nb3, + stream + ); + } else { + GGML_ABORT("unsupported type %s", ggml_type_name(dst->type)); + } +} + +template<> +void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const half * src0_d = (const half *)src0->data; + const int64_t * src1_d = (const int64_t *)src1->data; + + GGML_TENSOR_BINARY_OP_LOCALS + + cudaStream_t stream = ctx.stream(); + + + if (dst->type == GGML_TYPE_F16) { + set_rows_cuda( + src0_d, src1_d, (half*)dst->data, + ne00, ne01, ne02, ne03, + ne10, ne11, ne12, ne13, + nb01, nb02, nb03, + nb10, nb11, nb12, + nb1, nb2, nb3, + stream + ); + } else { + GGML_ABORT("unsupported type %s", ggml_type_name(dst->type)); + } +} + void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16)); GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32); - if (src1->type == GGML_TYPE_I64) { - set_rows_cuda(ctx, src0, src1, dst); + if (src0->type == GGML_TYPE_F32) { + if (src1->type == GGML_TYPE_I64) { + set_rows_cuda(ctx, src0, src1, dst); + } else { + set_rows_cuda(ctx, src0, src1, dst); + } + } else if (src0->type == GGML_TYPE_F16) { + if (src1->type == GGML_TYPE_I64) { + set_rows_cuda(ctx, src0, src1, dst); + } else { + set_rows_cuda(ctx, src0, src1, dst); + } } else { - set_rows_cuda(ctx, src0, src1, dst); + GGML_ABORT("unsupported type %s", ggml_type_name(src0->type)); } } diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index db1d39e2d..9681cd293 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -75,17 +75,26 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; + const bool use_bitonic = shared_mem <= max_shared_mem && ncols <= 1024; + const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows); - ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); + ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * chunk_nrows); int * tmp_dst = temp_dst_alloc.get(); - if (shared_mem > max_shared_mem || ncols > 1024) { - argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); - } else { - argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); + for (int64_t i = 0; i < nrows; i += chunk_nrows) { + int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i); + + if (use_bitonic) { + argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream); + } else { + argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream); + } + CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), iter_nrows, + cudaMemcpyDeviceToDevice, stream)); + + src0_d += ncols * iter_nrows; + dst_d += k * iter_nrows; } - CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, - cudaMemcpyDeviceToDevice, stream)); #else // GGML_CUDA_USE_CUB ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); int * tmp_dst = temp_dst_alloc.get(); diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h deleted file mode 100644 index 19a2504c7..000000000 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ /dev/null @@ -1,390 +0,0 @@ -#ifndef HTP_OPNODE_H -#define HTP_OPNODE_H - -#define GGML_COMMON_IMPL_CPP -#include "ggml-backend-impl.h" -#include "ggml-common.h" - -#include -#include -#include -#include -#include "htp-ops.h" -#include "htp/matmul-ops.h" -#include "htp/flash-attn-ops.h" - -struct htp_opnode { - ggml_tensor * node = nullptr; - - std::vector fused; - - htp_op_code opcode = HTP_OP_INVALID; - - std::vector extra_dsts; - - int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] = {0}; - - htp_opnode(ggml_tensor * node = nullptr, std::vector fused = {}, htp_op_code opcode = HTP_OP_INVALID, std::vector extra_dsts = {}) - : node(node), fused(std::move(fused)), opcode(opcode), extra_dsts(std::move(extra_dsts)) {} - - ggml_op op() const { - return node->op; - } - - const ggml_tensor * dst() const { - return fused.empty() ? node : fused.back(); - } - - void add_fused(ggml_tensor * t, bool extra_dst = false) { - fused.push_back(t); - if (extra_dst) { - extra_dsts.push_back(t); - } - } - - std::vector get_outputs() const { - std::vector res; - if (extra_dsts.empty()) { - res.push_back(dst()); - } else { - res.push_back(node); - for (const auto * x : extra_dsts) { - res.push_back(x); - } - } - return res; - } - - const ggml_tensor * src0() const { - return node->src[0]; - } - - const ggml_tensor * src1() const { - return node->src[1]; - } - - bool is_empty() const { - return ggml_op_is_empty(node->op); - } - - bool stackable() const { - switch (this->op()) { - case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_ID: - return ggml_is_quantized(this->src0()->type); - default: - return false; - } - } - - bool same_input(const htp_opnode& n) const { - return n.src1() == this->src1(); - } - - std::vector get_inputs() const { - if (fused.empty()) { - int last_non_null = -1; - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - last_non_null = i; - } - } - std::vector inputs(last_non_null + 1, nullptr); - for (int i = 0; i <= last_non_null; i++) { - inputs[i] = node->src[i]; - } - return inputs; - } - - std::vector inputs(GGML_MAX_SRC, nullptr); - std::vector outputs; - outputs.push_back(node); - for (const auto * f : fused) { - outputs.push_back(f); - } - - auto contains = [&](const std::vector & vec, const ggml_tensor * t) { - for (const auto * x : vec) { - if (x == t) return true; - } - return false; - }; - - int count = 0; - auto add_input = [&](const ggml_tensor * t) { - if (t && !contains(outputs, t) && !contains(inputs, t)) { - if (count < (int)inputs.size()) { - inputs[count++] = t; - } else { - inputs.push_back(t); - } - } - }; - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - add_input(node->src[i]); - } - } - for (const auto * f : fused) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (f->src[i]) { - add_input(f->src[i]); - } - } - } - - inputs.resize(count); - return inputs; - } - - std::string op_name() const { - if (fused.empty()) { - return ggml_op_desc(node); - } - std::string name = ggml_op_desc(node); - for (const auto * f : fused) { - name += "+"; - name += ggml_op_desc(f); - } - return name; - } -}; - -struct htp_opformat { - char strides[64 * GGML_MAX_SRC]; - char dims[64 * GGML_MAX_SRC]; - char types[16 * GGML_MAX_SRC]; - char buffs[64 * GGML_MAX_SRC]; - char names[64 * GGML_MAX_SRC]; - char kparams[128]; - - int format_tensor_dims(char * str, size_t max_size, const struct ggml_tensor * t) { - if (!t) { - return snprintf(str, max_size, "NONE"); - } - if (t->ne[2] == 1 && t->ne[3] == 1) { - return snprintf(str, max_size, "%d:%d", (int) t->ne[0], (int) t->ne[1]); - } else { - return snprintf(str, max_size, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]); - } - } - - void format_op_dims(char * str, size_t max_size, const htp_opnode & node) { - char * p = str; - char * p_end = str + max_size; - auto inputs = node.get_inputs(); - - if (!inputs.empty()) { - p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[0]), (size_t)(p_end - p)); - - for (size_t i = 1; i < inputs.size(); i++) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p)); - } - if (p < p_end) { - p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[i]), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p)); - } - } - - char self[64]; - format_tensor_dims(self, sizeof(self), node.dst()); - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p)); - } - } - - int format_tensor_strides(char * str, size_t max_size, const struct ggml_tensor * t) { - if (!t) { - return snprintf(str, max_size, "NONE"); - } - const char * c = ggml_is_contiguous(t) ? "" : "!"; - - if (t->ne[2] == 1 && t->ne[3] == 1) { - return snprintf(str, max_size, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c); - } else { - return snprintf(str, max_size, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c); - } - } - - void format_op_strides(char * str, size_t max_size, const htp_opnode & node) { - char * p = str; - char * p_end = str + max_size; - auto inputs = node.get_inputs(); - - if (!inputs.empty()) { - p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[0]), (size_t)(p_end - p)); - - for (size_t i = 1; i < inputs.size(); i++) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p)); - } - if (p < p_end) { - p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[i]), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p)); - } - } - - char self[64]; - format_tensor_strides(self, sizeof(self), node.dst()); - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p)); - } - } - - void format_op_types(char * str, size_t max_size, const htp_opnode & node) { - char * p = str; - char * p_end = str + max_size; - auto inputs = node.get_inputs(); - - if (!inputs.empty()) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? ggml_type_name(inputs[0]->type) : "NONE"), (size_t)(p_end - p)); - } - - for (size_t i = 1; i < inputs.size(); i++) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p)); - } - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? ggml_type_name(inputs[i]->type) : "NONE"), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", ggml_type_name(node.dst()->type)), (size_t)(p_end - p)); - } - } - - const char * tensor_buff_name(const struct ggml_tensor * t) { - if (t && t->buffer) { - return ggml_backend_buffer_name(t->buffer); - } - return "NONE"; - } - - void format_op_buffs(char * str, size_t max_size, const htp_opnode & node) { - char * p = str; - char * p_end = str + max_size; - auto inputs = node.get_inputs(); - - if (!inputs.empty()) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[0])), (size_t)(p_end - p)); - } - - for (size_t i = 1; i < inputs.size(); i++) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p)); - } - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[i])), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(node.dst())), (size_t)(p_end - p)); - } - } - - void format_op_names(char * str, size_t max_size, const htp_opnode & node) { - char * p = str; - char * p_end = str + max_size; - auto inputs = node.get_inputs(); - - if (!inputs.empty()) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? inputs[0]->name : "NONE"), (size_t)(p_end - p)); - } - - for (size_t i = 1; i < inputs.size(); i++) { - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p)); - } - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? inputs[i]->name : "NONE"), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p)); - } - } - - if (p < p_end) { - p += std::min((size_t)snprintf(p, p_end - p, "%s", node.dst()->name), (size_t)(p_end - p)); - } - } - void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) { - if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID || - node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN || - node.opcode == HTP_OP_MUL_MAT_ADD) { - const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params; - const char * path = "unknown"; - int32_t type = kparams->kernel_type; - if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_F16_BATCHED) { - path = "hmx-tiled"; - } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || - type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { - path = "hvx-tiled"; - } else if (type == HTP_MM_KERNEL_HVX_F16_F16_DDR || type == HTP_MM_KERNEL_HVX_F16_F32_DDR || - type == HTP_MM_KERNEL_HVX_F32_F32_DDR || type == HTP_MM_KERNEL_HVX_F32_F16_DDR || - type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - path = "hvx-flat"; - } - snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size); - } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) { - const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params; - const char * path = "unknown"; - int32_t type = kparams->kernel_type; - if (type == HTP_FA_KERNEL_HMX) { - path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq"; - } else if (type == HTP_FA_KERNEL_HVX) { - path = "hvx"; - } - snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size); - } else { - snprintf(str, max_size, "----"); - } - } - - void format(const htp_opnode & node) { - format_op_dims(dims, sizeof(dims), node); - format_op_strides(strides, sizeof(strides), node); - format_op_types(types, sizeof(types), node); - format_op_buffs(buffs, sizeof(buffs), node); - format_op_names(names, sizeof(names), node); - format_kernel_params(kparams, sizeof(kparams), node); - } - - htp_opformat() { - strides[0] = '\0'; - dims[0] = '\0'; - types[0] = '\0'; - buffs[0] = '\0'; - names[0] = '\0'; - kparams[0] = '\0'; - } - htp_opformat(const htp_opnode & node) { format(node); } -}; - -#endif // HTP_OPNODE_H diff --git a/ggml/src/ggml-hexagon/htp/hex-common.h b/ggml/src/ggml-hexagon/htp/hex-common.h deleted file mode 100644 index 4714486a0..000000000 --- a/ggml/src/ggml-hexagon/htp/hex-common.h +++ /dev/null @@ -1,80 +0,0 @@ -#ifndef HEX_COMMON_H -#define HEX_COMMON_H - -#include -#include -#include - -#ifndef SIZE_MAX -#define SIZE_MAX ((size_t)-1) -#endif - -#ifndef MAX -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#endif - -#ifndef MIN -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#endif - -static inline uint32_t hex_ceil_pow2(uint32_t x) { - if (x <= 1) { return 1; } - int p = 2; - x--; - while (x >>= 1) { p <<= 1; } - return p; -} - -static inline size_t hmx_ceil_div(size_t num, size_t den) { - return (num + den - 1) / den; -} - -static inline int32_t hex_is_aligned(const void * addr, uint32_t align) { - return ((size_t) addr & (align - 1)) == 0; -} - -static inline size_t hex_align_up(size_t v, size_t align) { - return hmx_ceil_div(v, align) * align; -} - -static inline size_t hex_align_down(size_t v, size_t align) { - return (v / align) * align; -} - -static inline int32_t hex_is_one_chunk(void * addr, uint32_t n, uint32_t chunk_size) { - uint32_t left_off = (size_t) addr & (chunk_size - 1); - uint32_t right_off = left_off + n; - return right_off <= chunk_size; -} - -static inline uint32_t hex_round_up(uint32_t n, uint32_t m) { - return m * ((n + m - 1) / m); -} - -static inline size_t hex_smin(size_t a, size_t b) { - return a < b ? a : b; -} - -static inline size_t hex_smax(size_t a, size_t b) { - return a > b ? a : b; -} - -static inline void hex_swap_ptr(void ** p1, void ** p2) { - void * t = *p1; - *p1 = *p2; - *p2 = t; -} - -static inline bool hex_mul_overflow(size_t a, size_t b, size_t *out) { - if (a != 0 && b > SIZE_MAX / a) return true; - *out = a * b; - return false; -} - -static inline bool hex_add_overflow(size_t a, size_t b, size_t *out) { - if (a > SIZE_MAX - b) return true; - *out = a + b; - return false; -} - -#endif // HEX_COMMON_H diff --git a/ggml/src/ggml-hexagon/htp/hex-profile.h b/ggml/src/ggml-hexagon/htp/hex-profile.h deleted file mode 100644 index 8a37a4a06..000000000 --- a/ggml/src/ggml-hexagon/htp/hex-profile.h +++ /dev/null @@ -1,64 +0,0 @@ -#ifndef HEX_PROFILE_H -#define HEX_PROFILE_H - -#include -#include -#include - -#include "hex-utils.h" -#include "htp-ops.h" - -#define HTP_TRACE_EVT_START 0 -#define HTP_TRACE_EVT_STOP 1 - -#ifndef HEX_NUM_PMU_COUNTERS -#define HEX_NUM_PMU_COUNTERS 8 -#endif - -static inline void hex_get_pmu(uint32_t counters[]) { -#if __HVX_ARCH__ >= 79 - asm volatile("%0 = upmucnt0" : "=r"(counters[0])); - asm volatile("%0 = upmucnt1" : "=r"(counters[1])); - asm volatile("%0 = upmucnt2" : "=r"(counters[2])); - asm volatile("%0 = upmucnt3" : "=r"(counters[3])); - asm volatile("%0 = upmucnt4" : "=r"(counters[4])); - asm volatile("%0 = upmucnt5" : "=r"(counters[5])); - asm volatile("%0 = upmucnt6" : "=r"(counters[6])); - asm volatile("%0 = upmucnt7" : "=r"(counters[7])); -#else - counters[0] = qurt_pmu_get(QURT_PMUCNT0); - counters[1] = qurt_pmu_get(QURT_PMUCNT1); - counters[2] = qurt_pmu_get(QURT_PMUCNT2); - counters[3] = qurt_pmu_get(QURT_PMUCNT3); - counters[4] = qurt_pmu_get(QURT_PMUCNT4); - counters[5] = qurt_pmu_get(QURT_PMUCNT5); - counters[6] = qurt_pmu_get(QURT_PMUCNT6); - counters[7] = qurt_pmu_get(QURT_PMUCNT7); -#endif -} - -struct htp_thread_trace { - uint32_t count; - uint32_t max_events; - struct htp_trace_desc * events; -}; - -static inline void htp_trace_event(struct htp_thread_trace * tr, uint16_t id, uint16_t info, uint32_t type) { - if (tr && tr->events && tr->count < tr->max_events) { - uint32_t idx = tr->count; - tr->events[idx].id = id; - tr->events[idx].info = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0); - tr->events[idx].cycles = (uint32_t) hex_get_cycles(); - tr->count++; - } -} - -static inline void htp_trace_event_start(struct htp_thread_trace * tr, uint16_t id, uint16_t info) { - htp_trace_event(tr, id, info, HTP_TRACE_EVT_START); -} - -static inline void htp_trace_event_stop(struct htp_thread_trace * tr, uint16_t id, uint16_t info) { - htp_trace_event(tr, id, info, HTP_TRACE_EVT_STOP); -} - -#endif /* HEX_PROFILE_H */ diff --git a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h deleted file mode 100644 index 4a0ca7885..000000000 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ /dev/null @@ -1,1341 +0,0 @@ -#include "hmx-utils.h" -#include "hmx-queue.h" - -// MXFP4 dequantization LUT: maps 4-bit index to fp16 mantissa value -// kvalues: 0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6 -static const __fp16 mxfp4_to_fp16_lut[64] __attribute__((aligned(VLEN))) = { - 0, 0, 0.5, 0, 1, 0, 1.5, 0, 2, 0, 3, 0, 4, 0, 6, 0, 0, 0, -0.5, 0, -1, 0, -1.5, 0, -2, 0, -3, 0, -4, 0, -6, 0, -}; - -static const __fp16 iq4_nl_to_fp16_lut[64] __attribute__((aligned(VLEN))) = { - -127, 0, -104, 0, -83, 0, -65, 0, -49, 0, -35, 0, -22, 0, -10, 0, - 1, 0, 13, 0, 25, 0, 38, 0, 53, 0, 69, 0, 89, 0, 113, 0, -}; - -// --- tiled format dequantizers --- - -typedef struct { - struct htp_context * ctx; - struct htp_thread_trace * traces; - __fp16 * dst; - const uint8_t * src; - - struct fastdiv_values n_k_tiles_div; - uint32_t n_k_tiles; - uint32_t n_tot_tiles; - uint32_t n_tiles_per_task; - uint32_t tile_size; - uint32_t aligned_tile_size; - uint32_t n_tasks; - uint32_t n_cols; - uint32_t k_block; - size_t row_stride; - uint32_t weight_type; -} tiled_dequantize_state_t; - -// Dequantize a single tile from tiled weight data (already in VTCM) to tile-major FP16. -static void dequantize_tiled_weight_to_fp16_task_q4_0( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - const HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - for (uint32_t t = start_tile; t < end_tile; t++) { - const uint8_t * tile_src = state->src + t * state->aligned_tile_size; - __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - - HVX_Vector v_sc = hvx_vmem(tile_src + 512); - HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(v_sc, v_sc, -2)); - - // Load all 4 groups in parallel - HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128); - HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128); - HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128); - HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128); - - // Nibble extraction - HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4); - HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4); - HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4); - HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4); - HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4); - HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4); - HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4); - HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4); - - // Offsetting (-8) - v_lo0 = Q6_Vb_vsub_VbVb(v_lo0, i8); - v_hi0 = Q6_Vb_vsub_VbVb(v_hi0, i8); - v_lo1 = Q6_Vb_vsub_VbVb(v_lo1, i8); - v_hi1 = Q6_Vb_vsub_VbVb(v_hi1, i8); - v_lo2 = Q6_Vb_vsub_VbVb(v_lo2, i8); - v_hi2 = Q6_Vb_vsub_VbVb(v_hi2, i8); - v_lo3 = Q6_Vb_vsub_VbVb(v_lo3, i8); - v_hi3 = Q6_Vb_vsub_VbVb(v_hi3, i8); - - // Shuffling - HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1); - HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1); - HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1); - HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1); - - // Unpack to 16-bit - HVX_VectorPair vp_int16_lo0 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf0)); - HVX_VectorPair vp_int16_hi0 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf0)); - HVX_VectorPair vp_int16_lo1 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf1)); - HVX_VectorPair vp_int16_hi1 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf1)); - HVX_VectorPair vp_int16_lo2 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf2)); - HVX_VectorPair vp_int16_hi2 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf2)); - HVX_VectorPair vp_int16_lo3 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf3)); - HVX_VectorPair vp_int16_hi3 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf3)); - - // Convert and scale multiplication - HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo0)), v_scale_duplicated)); - HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo0)), v_scale_duplicated)); - HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi0)), v_scale_duplicated)); - HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi0)), v_scale_duplicated)); - - HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo1)), v_scale_duplicated)); - HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo1)), v_scale_duplicated)); - HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi1)), v_scale_duplicated)); - HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi1)), v_scale_duplicated)); - - HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo2)), v_scale_duplicated)); - HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo2)), v_scale_duplicated)); - HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi2)), v_scale_duplicated)); - HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi2)), v_scale_duplicated)); - - HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo3)), v_scale_duplicated)); - HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo3)), v_scale_duplicated)); - HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi3)), v_scale_duplicated)); - HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi3)), v_scale_duplicated)); - - hvx_vmem(dst_ptr + 0 * 64) = v_grp0_0; - hvx_vmem(dst_ptr + 1 * 64) = v_grp0_1; - hvx_vmem(dst_ptr + 2 * 64) = v_grp0_2; - hvx_vmem(dst_ptr + 3 * 64) = v_grp0_3; - - hvx_vmem(dst_ptr + 4 * 64) = v_grp1_0; - hvx_vmem(dst_ptr + 5 * 64) = v_grp1_1; - hvx_vmem(dst_ptr + 6 * 64) = v_grp1_2; - hvx_vmem(dst_ptr + 7 * 64) = v_grp1_3; - - hvx_vmem(dst_ptr + 8 * 64) = v_grp2_0; - hvx_vmem(dst_ptr + 9 * 64) = v_grp2_1; - hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2; - hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3; - - hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0; - hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1; - hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2; - hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3; - } -} - -static void dequantize_tiled_weight_to_fp16_task_q4_1( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - - for (uint32_t t = start_tile; t < end_tile; t++) { - const uint8_t * tile_src = state->src + t * state->aligned_tile_size; - __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - - HVX_Vector vscale_offset = hvx_vmem(tile_src + 512); - HVX_VectorPair dm_deal = Q6_W_vdeal_VVR(vscale_offset, vscale_offset, -2); - HVX_Vector vd = Q6_V_lo_W(dm_deal); - HVX_Vector vm = Q6_V_hi_W(dm_deal); - - HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(vd, vd, -2)); - HVX_Vector v_offset_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(vm, vm, -2)); - - // Load all 4 groups in parallel - HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128); - HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128); - HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128); - HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128); - - // Nibble extraction - HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4); - HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4); - HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4); - HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4); - HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4); - HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4); - HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4); - HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4); - - // Shuffling - HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1); - HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1); - HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1); - HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1); - - // Unpack to 16-bit - HVX_VectorPair vp_int16_lo0 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf0)); - HVX_VectorPair vp_int16_hi0 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf0)); - HVX_VectorPair vp_int16_lo1 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf1)); - HVX_VectorPair vp_int16_hi1 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf1)); - HVX_VectorPair vp_int16_lo2 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf2)); - HVX_VectorPair vp_int16_hi2 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf2)); - HVX_VectorPair vp_int16_lo3 = Q6_Wh_vunpack_Vb(Q6_V_lo_W(vp_shuf3)); - HVX_VectorPair vp_int16_hi3 = Q6_Wh_vunpack_Vb(Q6_V_hi_W(vp_shuf3)); - - // Convert, multiply, add offset - HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo0)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo0)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi0)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi0)), v_scale_duplicated), v_offset_duplicated)); - - HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo1)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo1)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi1)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi1)), v_scale_duplicated), v_offset_duplicated)); - - HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo2)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo2)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi2)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi2)), v_scale_duplicated), v_offset_duplicated)); - - HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_lo3)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_lo3)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_hi3)), v_scale_duplicated), v_offset_duplicated)); - HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vadd_Vqf16Vhf(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_hi3)), v_scale_duplicated), v_offset_duplicated)); - - // Parallel Stores - hvx_vmem(dst_ptr + 0 * 64) = v_grp0_0; - hvx_vmem(dst_ptr + 1 * 64) = v_grp0_1; - hvx_vmem(dst_ptr + 2 * 64) = v_grp0_2; - hvx_vmem(dst_ptr + 3 * 64) = v_grp0_3; - - hvx_vmem(dst_ptr + 4 * 64) = v_grp1_0; - hvx_vmem(dst_ptr + 5 * 64) = v_grp1_1; - hvx_vmem(dst_ptr + 6 * 64) = v_grp1_2; - hvx_vmem(dst_ptr + 7 * 64) = v_grp1_3; - - hvx_vmem(dst_ptr + 8 * 64) = v_grp2_0; - hvx_vmem(dst_ptr + 9 * 64) = v_grp2_1; - hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2; - hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3; - - hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0; - hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1; - hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2; - hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3; - } -} - -static void dequantize_tiled_weight_to_fp16_task_iq4_nl( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - const HVX_Vector vlut_cvt = hvx_vmem(iq4_nl_to_fp16_lut); - - for (uint32_t t = start_tile; t < end_tile; t++) { - const uint8_t * tile_src = state->src + t * state->aligned_tile_size; - __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - - HVX_Vector v_sc = hvx_vmem(tile_src + 512); - HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(v_sc, v_sc, -2)); - - // Load all 4 groups in parallel - HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128); - HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128); - HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128); - HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128); - - // Nibble extraction - HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4); - HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4); - HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4); - HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4); - HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4); - HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4); - HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4); - HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4); - - // Shuffling - HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1); - HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1); - HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1); - HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1); - - // Shuffle for LUT lookup - HVX_Vector v_q_lo0 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf0)); - HVX_Vector v_q_hi0 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf0)); - HVX_Vector v_q_lo1 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf1)); - HVX_Vector v_q_hi1 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf1)); - HVX_Vector v_q_lo2 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf2)); - HVX_Vector v_q_hi2 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf2)); - HVX_Vector v_q_lo3 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf3)); - HVX_Vector v_q_hi3 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf3)); - - // LUT lookup - HVX_VectorPair vp_lo0 = Q6_Wh_vlut16_VbVhR(v_q_lo0, vlut_cvt, 0); - HVX_VectorPair vp_hi0 = Q6_Wh_vlut16_VbVhR(v_q_hi0, vlut_cvt, 0); - HVX_VectorPair vp_lo1 = Q6_Wh_vlut16_VbVhR(v_q_lo1, vlut_cvt, 0); - HVX_VectorPair vp_hi1 = Q6_Wh_vlut16_VbVhR(v_q_hi1, vlut_cvt, 0); - HVX_VectorPair vp_lo2 = Q6_Wh_vlut16_VbVhR(v_q_lo2, vlut_cvt, 0); - HVX_VectorPair vp_hi2 = Q6_Wh_vlut16_VbVhR(v_q_hi2, vlut_cvt, 0); - HVX_VectorPair vp_lo3 = Q6_Wh_vlut16_VbVhR(v_q_lo3, vlut_cvt, 0); - HVX_VectorPair vp_hi3 = Q6_Wh_vlut16_VbVhR(v_q_hi3, vlut_cvt, 0); - - // Convert and scale multiplication - HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo0), v_scale_duplicated)); - HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo0), v_scale_duplicated)); - HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi0), v_scale_duplicated)); - HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi0), v_scale_duplicated)); - - HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo1), v_scale_duplicated)); - HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo1), v_scale_duplicated)); - HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi1), v_scale_duplicated)); - HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi1), v_scale_duplicated)); - - HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo2), v_scale_duplicated)); - HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo2), v_scale_duplicated)); - HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi2), v_scale_duplicated)); - HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi2), v_scale_duplicated)); - - HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo3), v_scale_duplicated)); - HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo3), v_scale_duplicated)); - HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi3), v_scale_duplicated)); - HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi3), v_scale_duplicated)); - - hvx_vmem(dst_ptr + 0 * 64) = v_grp0_0; - hvx_vmem(dst_ptr + 1 * 64) = v_grp0_1; - hvx_vmem(dst_ptr + 2 * 64) = v_grp0_2; - hvx_vmem(dst_ptr + 3 * 64) = v_grp0_3; - - hvx_vmem(dst_ptr + 4 * 64) = v_grp1_0; - hvx_vmem(dst_ptr + 5 * 64) = v_grp1_1; - hvx_vmem(dst_ptr + 6 * 64) = v_grp1_2; - hvx_vmem(dst_ptr + 7 * 64) = v_grp1_3; - - hvx_vmem(dst_ptr + 8 * 64) = v_grp2_0; - hvx_vmem(dst_ptr + 9 * 64) = v_grp2_1; - hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2; - hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3; - - hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0; - hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1; - hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2; - hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3; - } -} - -static void dequantize_tiled_weight_to_fp16_task_mxfp4( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - const HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - const HVX_Vector vlut_cvt = hvx_vmem(mxfp4_to_fp16_lut); - - for (uint32_t t = start_tile; t < end_tile; t++) { - const uint8_t * tile_src = state->src + t * state->aligned_tile_size; - __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - - HVX_Vector v = hvx_vmem(tile_src + 512); - HVX_Vector vh = Q6_V_lo_W(Q6_Wuh_vunpack_Vub(v)); - vh = Q6_Vh_vsub_VhVh(vh, Q6_Vh_vsplat_R(112)); - vh = Q6_Vh_vmax_VhVh(vh, Q6_V_vzero()); - vh = Q6_Vh_vmin_VhVh(vh, Q6_Vh_vsplat_R(30)); - vh = Q6_Vh_vasl_VhR(vh, 10); - - HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(vh, vh, -2)); - - // Load all 4 groups in parallel - HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128); - HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128); - HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128); - HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128); - - // Nibble extraction - HVX_Vector v_lo0 = Q6_V_vand_VV(vq0, mask_h4); - HVX_Vector v_hi0 = Q6_Vub_vlsr_VubR(vq0, 4); - HVX_Vector v_lo1 = Q6_V_vand_VV(vq1, mask_h4); - HVX_Vector v_hi1 = Q6_Vub_vlsr_VubR(vq1, 4); - HVX_Vector v_lo2 = Q6_V_vand_VV(vq2, mask_h4); - HVX_Vector v_hi2 = Q6_Vub_vlsr_VubR(vq2, 4); - HVX_Vector v_lo3 = Q6_V_vand_VV(vq3, mask_h4); - HVX_Vector v_hi3 = Q6_Vub_vlsr_VubR(vq3, 4); - - // Shuffling - HVX_VectorPair vp_shuf0 = Q6_W_vshuff_VVR(v_hi0, v_lo0, -1); - HVX_VectorPair vp_shuf1 = Q6_W_vshuff_VVR(v_hi1, v_lo1, -1); - HVX_VectorPair vp_shuf2 = Q6_W_vshuff_VVR(v_hi2, v_lo2, -1); - HVX_VectorPair vp_shuf3 = Q6_W_vshuff_VVR(v_hi3, v_lo3, -1); - - // Shuffle for LUT lookup - HVX_Vector v_q_lo0 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf0)); - HVX_Vector v_q_hi0 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf0)); - HVX_Vector v_q_lo1 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf1)); - HVX_Vector v_q_hi1 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf1)); - HVX_Vector v_q_lo2 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf2)); - HVX_Vector v_q_hi2 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf2)); - HVX_Vector v_q_lo3 = Q6_Vb_vshuff_Vb(Q6_V_lo_W(vp_shuf3)); - HVX_Vector v_q_hi3 = Q6_Vb_vshuff_Vb(Q6_V_hi_W(vp_shuf3)); - - // LUT lookup - HVX_VectorPair vp_lo0 = Q6_Wh_vlut16_VbVhR(v_q_lo0, vlut_cvt, 0); - HVX_VectorPair vp_hi0 = Q6_Wh_vlut16_VbVhR(v_q_hi0, vlut_cvt, 0); - HVX_VectorPair vp_lo1 = Q6_Wh_vlut16_VbVhR(v_q_lo1, vlut_cvt, 0); - HVX_VectorPair vp_hi1 = Q6_Wh_vlut16_VbVhR(v_q_hi1, vlut_cvt, 0); - HVX_VectorPair vp_lo2 = Q6_Wh_vlut16_VbVhR(v_q_lo2, vlut_cvt, 0); - HVX_VectorPair vp_hi2 = Q6_Wh_vlut16_VbVhR(v_q_hi2, vlut_cvt, 0); - HVX_VectorPair vp_lo3 = Q6_Wh_vlut16_VbVhR(v_q_lo3, vlut_cvt, 0); - HVX_VectorPair vp_hi3 = Q6_Wh_vlut16_VbVhR(v_q_hi3, vlut_cvt, 0); - - // Convert and scale multiplication - HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo0), v_scale_duplicated)); - HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo0), v_scale_duplicated)); - HVX_Vector v_grp0_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi0), v_scale_duplicated)); - HVX_Vector v_grp0_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi0), v_scale_duplicated)); - - HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo1), v_scale_duplicated)); - HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo1), v_scale_duplicated)); - HVX_Vector v_grp1_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi1), v_scale_duplicated)); - HVX_Vector v_grp1_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi1), v_scale_duplicated)); - - HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo2), v_scale_duplicated)); - HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo2), v_scale_duplicated)); - HVX_Vector v_grp2_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi2), v_scale_duplicated)); - HVX_Vector v_grp2_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi2), v_scale_duplicated)); - - HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_lo3), v_scale_duplicated)); - HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_lo3), v_scale_duplicated)); - HVX_Vector v_grp3_2 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_lo_W(vp_hi3), v_scale_duplicated)); - HVX_Vector v_grp3_3 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_V_hi_W(vp_hi3), v_scale_duplicated)); - - hvx_vmem(dst_ptr + 0 * 64) = v_grp0_0; - hvx_vmem(dst_ptr + 1 * 64) = v_grp0_1; - hvx_vmem(dst_ptr + 2 * 64) = v_grp0_2; - hvx_vmem(dst_ptr + 3 * 64) = v_grp0_3; - - hvx_vmem(dst_ptr + 4 * 64) = v_grp1_0; - hvx_vmem(dst_ptr + 5 * 64) = v_grp1_1; - hvx_vmem(dst_ptr + 6 * 64) = v_grp1_2; - hvx_vmem(dst_ptr + 7 * 64) = v_grp1_3; - - hvx_vmem(dst_ptr + 8 * 64) = v_grp2_0; - hvx_vmem(dst_ptr + 9 * 64) = v_grp2_1; - hvx_vmem(dst_ptr + 10 * 64) = v_grp2_2; - hvx_vmem(dst_ptr + 11 * 64) = v_grp2_3; - - hvx_vmem(dst_ptr + 12 * 64) = v_grp3_0; - hvx_vmem(dst_ptr + 13 * 64) = v_grp3_1; - hvx_vmem(dst_ptr + 14 * 64) = v_grp3_2; - hvx_vmem(dst_ptr + 15 * 64) = v_grp3_3; - } -} - -static void dequantize_tiled_weight_to_fp16_task_q8_0( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - for (uint32_t t = start_tile; t < end_tile; t++) { - const uint8_t * tile_src = state->src + t * state->aligned_tile_size; - __fp16 * dst_ptr = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - - HVX_Vector v_sc = hvx_vmem(tile_src + 1024); - HVX_Vector v_scale_duplicated = Q6_V_lo_W(Q6_W_vshuff_VVR(v_sc, v_sc, -2)); - - // Load groups 0-3 in parallel - HVX_Vector vq0 = hvx_vmem(tile_src + 0 * 128); - HVX_Vector vq1 = hvx_vmem(tile_src + 1 * 128); - HVX_Vector vq2 = hvx_vmem(tile_src + 2 * 128); - HVX_Vector vq3 = hvx_vmem(tile_src + 3 * 128); - - HVX_VectorPair vp_int16_0 = Q6_Wh_vunpack_Vb(vq0); - HVX_VectorPair vp_int16_1 = Q6_Wh_vunpack_Vb(vq1); - HVX_VectorPair vp_int16_2 = Q6_Wh_vunpack_Vb(vq2); - HVX_VectorPair vp_int16_3 = Q6_Wh_vunpack_Vb(vq3); - - // Load groups 4-7 in parallel - HVX_Vector vq4 = hvx_vmem(tile_src + 4 * 128); - HVX_Vector vq5 = hvx_vmem(tile_src + 5 * 128); - HVX_Vector vq6 = hvx_vmem(tile_src + 6 * 128); - HVX_Vector vq7 = hvx_vmem(tile_src + 7 * 128); - - HVX_VectorPair vp_int16_4 = Q6_Wh_vunpack_Vb(vq4); - HVX_VectorPair vp_int16_5 = Q6_Wh_vunpack_Vb(vq5); - HVX_VectorPair vp_int16_6 = Q6_Wh_vunpack_Vb(vq6); - HVX_VectorPair vp_int16_7 = Q6_Wh_vunpack_Vb(vq7); - - // Convert and scale multiply for groups 0-3 - HVX_Vector v_grp0_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_0)), v_scale_duplicated)); - HVX_Vector v_grp0_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_0)), v_scale_duplicated)); - HVX_Vector v_grp1_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_1)), v_scale_duplicated)); - HVX_Vector v_grp1_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_1)), v_scale_duplicated)); - HVX_Vector v_grp2_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_2)), v_scale_duplicated)); - HVX_Vector v_grp2_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_2)), v_scale_duplicated)); - HVX_Vector v_grp3_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_3)), v_scale_duplicated)); - HVX_Vector v_grp3_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_3)), v_scale_duplicated)); - - // Store groups 0-3 - hvx_vmem(dst_ptr + 0 * 64) = v_grp0_0; - hvx_vmem(dst_ptr + 1 * 64) = v_grp0_1; - hvx_vmem(dst_ptr + 2 * 64) = v_grp1_0; - hvx_vmem(dst_ptr + 3 * 64) = v_grp1_1; - hvx_vmem(dst_ptr + 4 * 64) = v_grp2_0; - hvx_vmem(dst_ptr + 5 * 64) = v_grp2_1; - hvx_vmem(dst_ptr + 6 * 64) = v_grp3_0; - hvx_vmem(dst_ptr + 7 * 64) = v_grp3_1; - - // Convert and scale multiply for groups 4-7 - HVX_Vector v_grp4_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_4)), v_scale_duplicated)); - HVX_Vector v_grp4_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_4)), v_scale_duplicated)); - HVX_Vector v_grp5_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_5)), v_scale_duplicated)); - HVX_Vector v_grp5_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_5)), v_scale_duplicated)); - HVX_Vector v_grp6_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_6)), v_scale_duplicated)); - HVX_Vector v_grp6_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_6)), v_scale_duplicated)); - HVX_Vector v_grp7_0 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_lo_W(vp_int16_7)), v_scale_duplicated)); - HVX_Vector v_grp7_1 = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(Q6_Vhf_equals_Vh(Q6_V_hi_W(vp_int16_7)), v_scale_duplicated)); - - // Store groups 4-7 - hvx_vmem(dst_ptr + 8 * 64) = v_grp4_0; - hvx_vmem(dst_ptr + 9 * 64) = v_grp4_1; - hvx_vmem(dst_ptr + 10 * 64) = v_grp5_0; - hvx_vmem(dst_ptr + 11 * 64) = v_grp5_1; - hvx_vmem(dst_ptr + 12 * 64) = v_grp6_0; - hvx_vmem(dst_ptr + 13 * 64) = v_grp6_1; - hvx_vmem(dst_ptr + 14 * 64) = v_grp7_0; - hvx_vmem(dst_ptr + 15 * 64) = v_grp7_1; - } -} - -static void convert_f16_weight_to_fp16_tiles_task( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - const uint32_t n_k_tiles = state->n_k_tiles; - const struct fastdiv_values n_k_tiles_div = state->n_k_tiles_div; - - const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets); - const HVX_Vector v_scat_step = Q6_V_vsplat_R(4); - const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64); - - unsigned ct = fastdiv((unsigned)start_tile, &n_k_tiles_div); - unsigned kt = fastmodulo((unsigned)start_tile, n_k_tiles, &n_k_tiles_div); - - for (unsigned t = start_tile; t < (unsigned)end_tile; ) { - if (kt >= (unsigned)n_k_tiles) { kt = 0; ct++; } - - __fp16 *tile_base = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - { - uint32_t byte_off = kt * 32 * sizeof(__fp16); - - HVX_Vector v_off = v_scat_base; - for (uint32_t r = 0; r < HTP_MM_HMX_TILE_N_ROWS; r += 2) { - uint32_t row0 = ct * HTP_MM_HMX_TILE_N_COLS + r; - uint32_t row1 = row0 + 1; - - const uint8_t *r0 = state->src + row0 * state->row_stride; - const uint8_t *r1 = state->src + row1 * state->row_stride; - - HVX_Vector v0 = hvx_vmemu((const __fp16 *)(r0 + byte_off)); - HVX_Vector v1 = (row1 < state->n_cols) ? hvx_vmemu((const __fp16 *)(r1 + byte_off)) : Q6_V_vzero(); - - Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v0); - v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); - Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v1); - v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); - } - (void) *(volatile HVX_Vector *)(tile_base); - } - ++t; ++kt; - } - - if (start_tile < end_tile) { - (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS); - } -} - -static void quantize_f32_weight_to_fp16_tiles_task( - const tiled_dequantize_state_t *state, - uint32_t start_tile, uint32_t end_tile) { - - const uint32_t n_k_tiles = state->n_k_tiles; - const struct fastdiv_values n_k_tiles_div = state->n_k_tiles_div; - - const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets); - const HVX_Vector v_scat_step = Q6_V_vsplat_R(4); - const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64); - - unsigned ct = fastdiv((unsigned)start_tile, &n_k_tiles_div); - unsigned kt = fastmodulo((unsigned)start_tile, n_k_tiles, &n_k_tiles_div); - - for (unsigned t = start_tile; t < (unsigned)end_tile; ) { - if (kt >= (unsigned)n_k_tiles) { kt = 0; ct++; } - - __fp16 *tile_base = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; - { - uint32_t byte_off = kt * 32 * sizeof(float); - - HVX_Vector v_off = v_scat_base; - for (uint32_t r = 0; r < HTP_MM_HMX_TILE_N_ROWS; r += 2) { - uint32_t row0 = ct * HTP_MM_HMX_TILE_N_COLS + r; - uint32_t row1 = row0 + 1; - - const uint8_t *r0 = state->src + row0 * state->row_stride; - const uint8_t *r1 = state->src + row1 * state->row_stride; - - HVX_Vector v0_f32 = hvx_vmem((const float *)(r0 + byte_off)); - HVX_Vector v1_f32 = (row1 < state->n_cols) ? hvx_vmem((const float *)(r1 + byte_off)) : Q6_V_vzero(); - - HVX_Vector v_out = hvx_vec_f32_to_f16(v0_f32, v1_f32); - - Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out); - v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); - - HVX_Vector v_out_hi = Q6_V_vror_VR(v_out, 64); - Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out_hi); - v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); - } - (void) *(volatile HVX_Vector *)(tile_base); - } - ++t; ++kt; - } - - if (start_tile < end_tile) { - (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS); - } -} - -// --- End tiled dequantizers --- - -// requires external HMX lock -static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) { - __builtin_assume(n_row_tiles > 0); - __builtin_assume(n_col_tiles > 0); - __builtin_assume(n_dot_tiles > 0); - - Q6_bias_mxmem2_A((void *)scales); - for (uint32_t r = 0; r < n_row_tiles; ++r) { - for (size_t c = 0; c < n_col_tiles; ++c) { - Q6_mxclracc_hf(); - - const __fp16 *row_tiles = activation + r * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; - const __fp16 *col_tiles = weight + c * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; - - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * (uint32_t)k_block - 1; - Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); - Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - } - - __fp16 *out_tile = output + (r * n_col_tiles + c) * HTP_MM_HMX_TILE_N_ELMS; - Q6_mxmem_AR_after_hf(out_tile, 0); - } - } -} - -// C += AB -static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b, - const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile, - uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) { - __builtin_assume(n_row_tiles > 0); - __builtin_assume(n_col_tiles > 0); - __builtin_assume(n_dot_tiles > 0); - - Q6_bias_mxmem2_A((void *)col_scales); - - const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS; - for (size_t i = 0; i < n_row_tiles; ++i) { - const __fp16 *row_base = a + i * dot_tile_stride; - __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS; - for (size_t j = 0; j < n_col_tiles; ++j) { - Q6_mxclracc_hf(); - - const __fp16 *col_tiles = b + j * dot_tile_stride; - const __fp16 *row_tiles = row_base; - __fp16 *accum_tile = res_base + j * HTP_MM_HMX_TILE_N_ELMS; - if (!zero_init) { - Q6_activation_hf_mxmem_RR((unsigned int)accum_tile, 2047); - Q6_weight_hf_mxmem_RR((unsigned int)eye_tile, 2047); - } - - for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) { - k_block = hex_smin(n_dot_tiles - k, 32); - const uint32_t range = 2048u * k_block - 1; - Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range); - Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range); - row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS; - } - - Q6_mxmem_AR_after_hf(accum_tile, 0); - } - } -} - -// --- 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; -} - -// output : fp16 -> f32p - -static void transfer_output_chunk_fp16_to_fp32( - float *restrict dst, - const float *restrict src2, - const __fp16 *restrict vtcm_src, - uint32_t start_row, - uint32_t n_rows, - uint32_t n_cols, - uint32_t dst_stride, - uint32_t src2_stride, - uint32_t dst_cols -) { - assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0); - const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS; - - const HVX_Vector one = hvx_vec_splat_f16(1.0); - - const size_t limit_c = hex_smin(n_cols, dst_cols); - const size_t limit_c_aligned = (limit_c & ~31); - - for (size_t r = 0; r < n_rows; r += 2) { - const size_t r_idx0 = start_row + r + 0; - const size_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; - const size_t r1 = (r_idx0 % HTP_MM_HMX_TILE_N_ROWS) / 2; // index of the row pair within the tile - const __fp16 *row_base = vtcm_src + r0 * tile_row_stride; - float *output_row_base = dst + r * dst_stride; // global memory row base for row r (and r+1) - const float *src2_row_base = src2 ? (src2 + r * src2_stride) : NULL; - - #pragma unroll(4) - for (size_t c = 0; c < limit_c_aligned; c += HTP_MM_HMX_TILE_N_COLS) { - const size_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - const __fp16 *tile = row_base + c0 * HTP_MM_HMX_TILE_N_ELMS; - HVX_Vector v = ((const HVX_Vector *) tile)[r1]; - HVX_VectorPair vp = Q6_Wqf32_vmpy_VhfVhf(v, one); - - HVX_Vector *pv_out0 = (HVX_Vector *) (output_row_base + c + 0); - HVX_Vector *pv_out1 = (HVX_Vector *) (output_row_base + c + dst_stride); - - HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp)); - if (src2_row_base) { - HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0); - v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0); - } - *pv_out0 = v_out0; - - if (r + 1 < n_rows) { - HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp)); - if (src2_row_base) { - HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride); - v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1); - } - *pv_out1 = v_out1; - } - } - - if (limit_c_aligned < limit_c) { - size_t c = limit_c_aligned; - size_t valid_c = limit_c - c; - const size_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - const __fp16 *tile = row_base + c0 * HTP_MM_HMX_TILE_N_ELMS; - HVX_Vector v = ((const HVX_Vector *) tile)[r1]; - HVX_VectorPair vp = Q6_Wqf32_vmpy_VhfVhf(v, one); - - HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp)); - if (src2_row_base) { - HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0); - v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0); - } - hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), v_out0); - - if (r + 1 < n_rows) { - HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp)); - if (src2_row_base) { - HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride); - v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1); - } - hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), v_out1); - } - } - } -} - -typedef struct { - const __fp16 *vtcm_src; - float *dst; - const float *src2; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t n_cols; - uint32_t dst_stride; // DDR row stride - uint32_t src2_stride; // DDR row stride for residual - uint32_t dst_cols; // Actual output columns - struct htp_thread_trace * traces; -} output_transfer_task_state_t; - -// activations : fp32 -> fp16 - -static void transfer_activation_chunk_fp32_to_fp16(__fp16 *restrict vtcm_dst, const float *restrict src, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, uint32_t k_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) { - uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - const float *ptr_in0 = src + (r + 0) * k_stride; - const float *ptr_in1 = src + (r + 1) * k_stride; - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = *(const HVX_Vector *)(ptr_in0 + c); - HVX_Vector v1 = *(const HVX_Vector *)(ptr_in1 + c); - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = *(const HVX_Vector *)(ptr_in0 + c); - HVX_Vector v1 = *(const HVX_Vector *)(ptr_in1 + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } - - for (; r < n_rows_padded; r += 2) { - uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - 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) : NULL; - const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride) : NULL; - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } -} - -typedef struct { - __fp16 *dst; - const float *src; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t k_block; - uint32_t k_stride; - uint32_t k_valid; - struct htp_thread_trace * traces; - struct htp_context * ctx; - float * vtcm_f32_act; -} activation_transfer_task_state_t; - -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) { - - const uint32_t R = HTP_MM_DMA_ACT_ROWS_PER_STEP; - 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 / R; - - // pre-fetch 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); - } - - for (uint32_t s = 0; s < n_steps; ++s) { - uint32_t r = R * s; - float *curr_buf = thread_f32_act + (s % 2) * R * k_block; - - if (r < n_rows) { - dma_queue_pop(dma_q); - } - - uint32_t next_s = s + 1; - uint32_t next_r = R * next_s; - 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 + (next_s % 2) * 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); - } - - #pragma unroll - for (uint32_t i = 0; i < HTP_MM_DMA_ACT_ROWS_PER_STEP; i += 2) { - uint32_t curr_r = r + i; - const bool row0_valid = (curr_r < n_rows); - const bool row1_valid = (curr_r + 1) < n_rows; - - const float *ptr_in0 = curr_buf + i * k_block; - const float *ptr_in1 = curr_buf + (i + 1) * k_block; - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t r0 = curr_r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = curr_r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c); - if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t r0 = curr_r / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = curr_r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } - } -} - -typedef struct { - const struct mmid_row_mapping *matrix_rows; - __fp16 *dst; - const float *src; - uint32_t n_tasks; - uint32_t n_tot_chunks; - uint32_t n_chunks_per_task; - uint32_t k_block; - uint32_t 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; - struct htp_thread_trace *traces; -} activation_transfer_gathered_task_state_t; - -typedef struct { - const struct mmid_row_mapping *matrix_rows; - const __fp16 *vtcm_src; - float *dst; - 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; - struct htp_thread_trace *traces; -} output_transfer_scattered_task_state_t; - -static void transfer_activation_chunk_fp32_to_fp16_gathered( - __fp16 *restrict vtcm_dst, - const float *restrict src, - uint32_t start_row, - uint32_t n_rows, - uint32_t k_block, - const struct mmid_row_mapping *matrix_rows, - uint32_t cur_a, - uint32_t mapping_stride, - uint32_t ne11, - const struct fastdiv_values * ne11_div, - size_t nb11, - size_t nb12, - uint32_t cne1, - uint32_t k_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) { - uint32_t r_idx0 = start_row + r + 0; - uint32_t r_idx1 = start_row + r + 1; - uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0]; - struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1]; - - uint32_t i11_0 = fastmodulo(mapping0.i1, ne11, ne11_div); - uint32_t i11_1 = fastmodulo(mapping1.i1, ne11, ne11_div); - - const float *row0_ptr = (const float *) ((const uint8_t *) src + i11_0 * nb11 + mapping0.i2 * nb12); - const float *row1_ptr = (const float *) ((const uint8_t *) src + i11_1 * nb11 + mapping1.i2 * nb12); - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = *(const HVX_Vector *)(row0_ptr + c); - HVX_Vector v1 = *(const HVX_Vector *)(row1_ptr + c); - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = *(const HVX_Vector *)(row0_ptr + c); - HVX_Vector v1 = *(const HVX_Vector *)(row1_ptr + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } - - for (; r < n_rows_padded; r += 2) { - uint32_t r_idx0 = start_row + r; - uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - const bool row0_valid = (start_row + r + 0) < cne1; - const bool row1_valid = (start_row + r + 1) < cne1; - - const float *row0_ptr = NULL; - const float *row1_ptr = NULL; - - if (row0_valid) { - struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + (start_row + r + 0)]; - uint32_t i11_0 = fastmodulo(mapping0.i1, ne11, ne11_div); - row0_ptr = (const float *) ((const uint8_t *) src + i11_0 * nb11 + mapping0.i2 * nb12); - } - if (row1_valid) { - struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + (start_row + r + 1)]; - uint32_t i11_1 = fastmodulo(mapping1.i1, ne11, ne11_div); - row1_ptr = (const float *) ((const uint8_t *) src + i11_1 * nb11 + mapping1.i2 * nb12); - } - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(row0_ptr + c); - if (row1_valid) v1 = *(const HVX_Vector *)(row1_ptr + c); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(row0_ptr + c); - if (row1_valid) v1 = *(const HVX_Vector *)(row1_ptr + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } -} - -static void transfer_activation_chunk_fp32_to_fp16_gathered_flat( - __fp16 *restrict vtcm_dst, - const float *restrict src, - uint32_t start_row, - uint32_t n_rows, - uint32_t k_block, - const struct mmid_row_mapping *matrix_rows, - uint32_t cur_a, - uint32_t mapping_stride, - size_t nb12, - uint32_t cne1, - uint32_t k_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) { - uint32_t r_idx0 = start_row + r + 0; - uint32_t r_idx1 = start_row + r + 1; - uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0]; - struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1]; - - const float *row0_ptr = (const float *) ((const uint8_t *) src + mapping0.i2 * nb12); - const float *row1_ptr = (const float *) ((const uint8_t *) src + mapping1.i2 * nb12); - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = *(const HVX_Vector *)(row0_ptr + c); - HVX_Vector v1 = *(const HVX_Vector *)(row1_ptr + c); - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = *(const HVX_Vector *)(row0_ptr + c); - HVX_Vector v1 = *(const HVX_Vector *)(row1_ptr + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } - - for (; r < n_rows_padded; r += 2) { - uint32_t r_idx0 = start_row + r; - uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; // tile row index - uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx - - const bool row0_valid = (start_row + r + 0) < cne1; - const bool row1_valid = (start_row + r + 1) < cne1; - - const float *row0_ptr = NULL; - const float *row1_ptr = NULL; - - if (row0_valid) { - struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + (start_row + r + 0)]; - row0_ptr = (const float *) ((const uint8_t *) src + mapping0.i2 * nb12); - } - if (row1_valid) { - struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + (start_row + r + 1)]; - row1_ptr = (const float *) ((const uint8_t *) src + mapping1.i2 * nb12); - } - - uint32_t c = 0; - for (; c + 32 <= k_valid; c += 32) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(row0_ptr + c); - if (row1_valid) v1 = *(const HVX_Vector *)(row1_ptr + c); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - if (c < k_block) { - HVX_Vector v0 = Q6_V_vzero(); - HVX_Vector v1 = Q6_V_vzero(); - if (row0_valid) v0 = *(const HVX_Vector *)(row0_ptr + c); - if (row1_valid) v1 = *(const HVX_Vector *)(row1_ptr + c); - - uint32_t rem = k_valid - c; - HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0); - v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero()); - v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero()); - - HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1); - - uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0; - - HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS); - tile[r1 / 2] = v_out; - } - } -} - -static void transfer_output_chunk_fp16_to_fp32_scattered( - float *restrict dst, - const __fp16 *restrict vtcm_src, - uint32_t start_row, - uint32_t n_rows, - uint32_t n_cols, - const struct mmid_row_mapping *matrix_rows, - uint32_t cur_a, - uint32_t mapping_stride, - size_t dst_nb1, - size_t dst_nb2, - uint32_t cne1) { - assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0); - const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS; - - const HVX_Vector one = hvx_vec_splat_f16(1.0); - - for (size_t r = 0; r < n_rows; r += 2) { - uint32_t r_idx0 = start_row + r + 0; - uint32_t r_idx1 = start_row + r + 1; - const size_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS; - const size_t r1 = (r_idx0 % HTP_MM_HMX_TILE_N_ROWS) / 2; // index of the row pair within the tile - const __fp16 *row_base = vtcm_src + r0 * tile_row_stride; - - if (r_idx0 >= cne1) break; - - struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0]; - float *output_row0 = (float *) ((uint8_t *) dst + mapping0.i1 * dst_nb1 + mapping0.i2 * dst_nb2); - - float *output_row1 = NULL; - if (r_idx1 < cne1) { - struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1]; - output_row1 = (float *) ((uint8_t *) dst + mapping1.i1 * dst_nb1 + mapping1.i2 * dst_nb2); - } - - #pragma unroll(4) - for (size_t c = 0; c < (size_t)n_cols; c += HTP_MM_HMX_TILE_N_COLS) { - const size_t c0 = c / HTP_MM_HMX_TILE_N_COLS; - const __fp16 *tile = row_base + c0 * HTP_MM_HMX_TILE_N_ELMS; - HVX_Vector v = ((const HVX_Vector *) tile)[r1]; - HVX_VectorPair vp = Q6_Wqf32_vmpy_VhfVhf(v, one); - - HVX_Vector *pv_out0 = (HVX_Vector *) (output_row0 + c); - HVX_Vector *pv_out1 = output_row1 ? (HVX_Vector *) (output_row1 + c) : NULL; - - *pv_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp)); - if (pv_out1) { - *pv_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp)); - } - } - } -} diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h deleted file mode 100644 index 328a83118..000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-flat.h +++ /dev/null @@ -1,1511 +0,0 @@ -// Dynamic quantizers that produce flat (non-tiled) activations - -static inline void quantize_block_f32_q8_0_flat( - float * restrict x, - uint8_t * restrict y_quants, - __fp16 * restrict y_scales, - uint32_t block_idx -) { - HVX_Vector * vx = (HVX_Vector *) x; - HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); - HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); - HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); - HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); - - HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); - HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); - HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); - HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); - - HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); - HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); - HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); - HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); - - HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); - HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); - - HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); - HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - - HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); - HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); - - HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); - HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); - - HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); - HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); - HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - - * (HVX_Vector *) (y_quants + block_idx * 128) = vx_i8; - - HVX_VectorPair vp1 = Q6_W_vshuff_VVR(vd23_hf, vd01_hf, -2); - HVX_VectorPair vp2 = Q6_W_vshuff_VVR(Q6_V_hi_W(vp1), Q6_V_lo_W(vp1), -2); - HVX_Vector v_scales = Q6_V_lo_W(vp2); - hvx_vec_store_u(y_scales + block_idx * 4, 8, v_scales); -} - -static inline void quantize_block_f32_q8_1_flat( - float * restrict x, - uint8_t * restrict y_quants, - __fp16 * restrict y_scales, - uint32_t block_idx -) { - HVX_Vector * vx = (HVX_Vector *) x; - HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); - HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); - HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); - HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); - - HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); - HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); - HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); - HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); - - HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); - HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); - HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); - HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); - - HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); - HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); - - HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); - HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - - HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); - HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); - - HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); - HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); - - HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); - HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); - HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - - const HVX_Vector ones = Q6_Vb_vsplat_R(1); - HVX_Vector v_sums = Q6_Vw_vrmpy_VbVb(vx_i8, ones); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 4)); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 8)); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 16)); - - * (HVX_Vector *) (y_quants + block_idx * 128) = vx_i8; - - HVX_VectorPair vp1 = Q6_W_vshuff_VVR(vd23_hf, vd01_hf, -2); - HVX_VectorPair vp2 = Q6_W_vshuff_VVR(Q6_V_hi_W(vp1), Q6_V_lo_W(vp1), -2); - HVX_Vector v_scales = Q6_V_lo_W(vp2); - - HVX_VectorPair v_deal1 = Q6_W_vdeal_VVR(v_sums, v_sums, -4); - HVX_Vector v_even1 = Q6_V_lo_W(v_deal1); - HVX_VectorPair v_deal2 = Q6_W_vdeal_VVR(v_even1, v_even1, -4); - HVX_Vector v_even2 = Q6_V_lo_W(v_deal2); - HVX_VectorPair v_deal3 = Q6_W_vdeal_VVR(v_even2, v_even2, -4); - HVX_Vector v_sums_shuffled = Q6_V_lo_W(v_deal3); - - HVX_Vector v_sums_sf = Q6_Vsf_equals_Vw(v_sums_shuffled); - HVX_Vector v_sums_hf = hvx_vec_f32_to_f16(v_sums_sf, Q6_V_vzero()); - - HVX_Vector v_prod = hvx_vec_mul_f16_f16(v_scales, v_sums_hf); - - HVX_VectorPair vp_scales = Q6_W_vshuff_VVR(v_prod, v_scales, -2); - HVX_Vector v_final = Q6_V_lo_W(vp_scales); - - hvx_vec_store_u(y_scales + block_idx * 8, 16, v_final); -} - -static inline void quantize_row_f32_q8_0_flat(float * restrict x, uint8_t * restrict y, uint32_t k) { - assert(k % 32 == 0); - const uint32_t quants_size = hex_round_up(k, 128); - uint8_t * restrict y_quants = y; - __fp16 * restrict y_scales = (__fp16 *) (y + quants_size); - - const uint32_t nb = (k + 127) / 128; - for (uint32_t i = 0; i < nb; i++) { - quantize_block_f32_q8_0_flat(x + i * 128, y_quants, y_scales, i); - } -} - -static inline void quantize_row_f32_q8_1_flat(float * restrict x, uint8_t * restrict y, uint32_t k) { - assert(k % 32 == 0); - const uint32_t quants_size = hex_round_up(k, 128); - uint8_t * restrict y_quants = y; - __fp16 * restrict y_scales = (__fp16 *) (y + quants_size); - - const uint32_t nb = (k + 127) / 128; - for (uint32_t i = 0; i < nb; i++) { - quantize_block_f32_q8_1_flat(x + i * 128, y_quants, y_scales, i); - } -} - -static inline void quantize_f32_q8_0_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_row_size, - size_t dst_row_size -) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_0_flat((float *) tmp_data, dst_data, ne0); - dst_data += dst_row_size; - src_data += src_row_size; - } -} - -static inline void quantize_f32_q8_1_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_row_size, - size_t dst_row_size -) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_1_flat((float *) tmp_data, dst_data, ne0); - dst_data += dst_row_size; - src_data += src_row_size; - } -} - -static inline void quantize_f32_f32_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_stride, - size_t dst_stride -) { - (void) tmp_data; - const size_t src_row_size = ne0 * sizeof(float); - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_stride, 2); - hvx_copy_f32_au(dst_data, src_data, ne0); - - dst_data += dst_stride; - src_data += src_stride; - } -} - -static inline void quantize_f32_f16_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_stride, - size_t dst_stride -) { - (void) tmp_data; - const size_t src_row_size = ne0 * sizeof(float); - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_stride, 2); - hvx_copy_f16_f32_au(dst_data, src_data, ne0); - - dst_data += dst_stride; - src_data += src_stride; - } -} - -static inline void quantize_f16_f16_flat_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_stride, - size_t dst_stride -) { - (void) tmp_data; - const size_t src_row_size = ne0 * sizeof(float); - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_stride, 2); - hvx_copy_f16_au(dst_data, src_data, ne0); - - dst_data += dst_stride; - src_data += src_stride; - } -} - -// Dot kernels that consume flat (non-tiled) activations - -static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act_rep, i8); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0_rep, v_act1_rep, i8); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act_rep, Q6_V_vzero()); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_offset = vptr[4]; - HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); - HVX_Vector v_scale = Q6_V_lo_W(p_deal); - HVX_Vector v_offset = Q6_V_hi_W(p_deal); - - __fp16 scale_a_val = y_scales[kt * 2 + 0]; - __fp16 sum_a_val = y_scales[kt * 2 + 1]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - HVX_Vector v_sum_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a); - HVX_Vector v_offset_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a); - - HVX_Vector v_scaled_dot = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - HVX_Vector v_sum_scaled = hvx_vec_add_f32_f32(v_scaled_dot, v_offset_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0_rep, v_act1_rep, Q6_V_vzero()); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_offset = vptr[4]; - HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); - HVX_Vector v_scale = Q6_V_lo_W(p_deal); - HVX_Vector v_offset = Q6_V_hi_W(p_deal); - - __fp16 scale_a0_val = y0_scales[kt * 2 + 0]; - __fp16 sum_a0_val = y0_scales[kt * 2 + 1]; - __fp16 scale_a1_val = y1_scales[kt * 2 + 0]; - __fp16 sum_a1_val = y1_scales[kt * 2 + 1]; - - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_sum_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - HVX_Vector v_sum_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&sum_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a0); - HVX_Vector v_offset_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a1); - HVX_Vector v_offset_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a1); - - HVX_Vector v_scaled_dot_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c0 = hvx_vec_add_f32_f32(v_scaled_dot_c0, v_offset_comb_c0); - - HVX_Vector v_scaled_dot_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - HVX_Vector v_sum_scaled_c1 = hvx_vec_add_f32_f32(v_scaled_dot_c1, v_offset_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx_i8 = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx_i8, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_q8_0_32x1(vptr, v_act_rep); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[8]; - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0_i8 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1_i8 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0_i8, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1_i8, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_q8_0_32x2(vptr, v_act0_rep, v_act1_rep); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[8]; - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act_rep, mask_h4, lut); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0_rep, v_act1_rep, mask_h4, lut); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[4]; - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut; - HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0; - HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y_scales = (const __fp16 *) (y_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx = * (const HVX_Vector *) (y_q + block_idx * 128); - HVX_Vector v_act_raw = Q6_V_vror_VR(vx, sub_idx * 32); - - HVX_Vector v_act_rep[8]; - v_act_rep[0] = Q6_V_vdelta_VV(v_act_raw, v_repl_ctrl); - v_act_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 4), v_repl_ctrl); - v_act_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 8), v_repl_ctrl); - v_act_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 12), v_repl_ctrl); - v_act_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 16), v_repl_ctrl); - v_act_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 20), v_repl_ctrl); - v_act_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 24), v_repl_ctrl); - v_act_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act_raw, 28), v_repl_ctrl); - - HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act_rep, mask_h4, lut); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512); - HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand); - r0_d = Q6_V_vand_VV(r0_d, e8m0_mask); - HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23); - - __fp16 scale_a_val = y_scales[kt]; - HVX_Vector v_scale_a_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a_val)); - HVX_VectorPair p_scale_a_f32 = hvx_vec_f16_to_f32(v_scale_a_f16); - HVX_Vector v_scale_a = Q6_V_lo_W(p_scale_a_f32); - - HVX_Vector v_scale_comb = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f)); - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut; - HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0; - HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - const uint32_t quants_size = hex_round_up(n, 128); - const __fp16 * restrict y0_scales = (const __fp16 *) (y0_q + quants_size); - const __fp16 * restrict y1_scales = (const __fp16 *) (y1_q + quants_size); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - - uint32_t block_idx = kt / 4; - uint32_t sub_idx = kt % 4; - - HVX_Vector vx0 = * (const HVX_Vector *) (y0_q + block_idx * 128); - HVX_Vector vx1 = * (const HVX_Vector *) (y1_q + block_idx * 128); - - HVX_Vector v_act0_raw = Q6_V_vror_VR(vx0, sub_idx * 32); - HVX_Vector v_act1_raw = Q6_V_vror_VR(vx1, sub_idx * 32); - - HVX_Vector v_act0_rep[8]; - v_act0_rep[0] = Q6_V_vdelta_VV(v_act0_raw, v_repl_ctrl); - v_act0_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 4), v_repl_ctrl); - v_act0_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 8), v_repl_ctrl); - v_act0_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 12), v_repl_ctrl); - v_act0_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 16), v_repl_ctrl); - v_act0_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 20), v_repl_ctrl); - v_act0_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 24), v_repl_ctrl); - v_act0_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act0_raw, 28), v_repl_ctrl); - - HVX_Vector v_act1_rep[8]; - v_act1_rep[0] = Q6_V_vdelta_VV(v_act1_raw, v_repl_ctrl); - v_act1_rep[1] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 4), v_repl_ctrl); - v_act1_rep[2] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 8), v_repl_ctrl); - v_act1_rep[3] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 12), v_repl_ctrl); - v_act1_rep[4] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 16), v_repl_ctrl); - v_act1_rep[5] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 20), v_repl_ctrl); - v_act1_rep[6] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 24), v_repl_ctrl); - v_act1_rep[7] = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act1_raw, 28), v_repl_ctrl); - - HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0_rep, v_act1_rep, mask_h4, lut); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512); - HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand); - r0_d = Q6_V_vand_VV(r0_d, e8m0_mask); - HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23); - - __fp16 scale_a0_val = y0_scales[kt]; - __fp16 scale_a1_val = y1_scales[kt]; - HVX_Vector v_scale_a0_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a0_val)); - HVX_Vector v_scale_a1_f16 = hvx_vec_repl_f16(Q6_Vh_vsplat_R(*(const int16_t *)&scale_a1_val)); - HVX_VectorPair p_scale_a0_f32 = hvx_vec_f16_to_f32(v_scale_a0_f16); - HVX_VectorPair p_scale_a1_f32 = hvx_vec_f16_to_f32(v_scale_a1_f16); - HVX_Vector v_scale_a0 = Q6_V_lo_W(p_scale_a0_f32); - HVX_Vector v_scale_a1 = Q6_V_lo_W(p_scale_a1_f32); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f)); - v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f)); - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -#if __HVX_ARCH__ < 79 -#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b)) -#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) -#else -#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b) -#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) -#endif - -static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { - const HVX_Vector * restrict x = (const HVX_Vector *) vx; - const HVX_Vector * restrict y = (const HVX_Vector *) vy; - - uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors - uint32_t nloe = n % VLEN_FP32; // leftover elements - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(4) - for (i = 0; i < nvec; i++) { - HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]); - rsum = HVX_OP_ADD_F32(rsum, prod); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]); - HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]); - HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf); - rsum = HVX_OP_ADD_F32(rsum, prod); - } - - *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum)); -} - -static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0, - const void * restrict vx0, const void * restrict vx1, - const void * restrict vy0) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y = (const HVX_Vector *) vy0; - - uint32_t nvec = n / VLEN_FP32; - uint32_t nloe = n % VLEN_FP32; - - HVX_Vector rsum0 = Q6_V_vzero(); - HVX_Vector rsum1 = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector y_sf = y[i]; - HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf); - HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf); - rsum0 = HVX_OP_ADD_F32(rsum0, prod0); - rsum1 = HVX_OP_ADD_F32(rsum1, prod1); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]); - HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]); - HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf); - HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf); - rsum0 = HVX_OP_ADD_F32(rsum0, prod0); - rsum1 = HVX_OP_ADD_F32(rsum1, prod1); - } - - HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); - hvx_vec_store_u(s0, 8, rsum); -} - -static inline void vec_dot_f32_f32_aa_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) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; - const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; - - uint32_t nvec = n / VLEN_FP32; - uint32_t nloe = n % VLEN_FP32; - - HVX_Vector r0_c0_sum = Q6_V_vzero(); - HVX_Vector r0_c1_sum = Q6_V_vzero(); - HVX_Vector r1_c0_sum = Q6_V_vzero(); - HVX_Vector r1_c1_sum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector r0_sf = x0[i]; - HVX_Vector r1_sf = x1[i]; - HVX_Vector c0_sf = y0[i]; - HVX_Vector c1_sf = y1[i]; - - r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf)); - r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf)); - r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf)); - r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf)); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - - HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]); - HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]); - HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]); - - r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf)); - r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf)); - r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf)); - r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf)); - } - - // Reduce and store results - HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); - HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); - - hvx_vec_store_u(s0, 8, r0_r1_c0_sum); - hvx_vec_store_u(s1, 8, r0_r1_c1_sum); -} - -static inline void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) { - const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x; - const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y; - - uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors - uint32_t nloe = n % VLEN_FP32; // leftover elements - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector x_sf = vx[i]; - HVX_Vector y_sf = vy[i]; - - rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf)); - } - - if (nloe) { - HVX_Vector x_sf = vx[i]; - HVX_Vector y_sf = vy[i]; - - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - x_sf = Q6_V_vand_QV(bmask, x_sf); - y_sf = Q6_V_vand_QV(bmask, y_sf); - - rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf)); - } - - rsum = hvx_vec_reduce_sum_f32(rsum); - hvx_vec_store_u(&s[0], 4, rsum); -} - -#undef HVX_OP_ADD_F32 -#undef HVX_OP_MUL_F32 - -static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { - const HVX_Vector * restrict x = (const HVX_Vector *) vx; - const HVX_Vector * restrict y = (const HVX_Vector *) vy; - - uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors - uint32_t nloe = n % VLEN_FP16; // leftover elements - - HVX_VectorPair rsum_p = Q6_W_vzero(); - - uint32_t i = 0; - - #pragma unroll(4) - for (i = 0; i < nvec; i++) { - rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]); - HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); - rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf); - } - - HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p))); - hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum)); -} - -static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0, - const void * restrict vx0, const void * restrict vx1, - const void * restrict vy0) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y = (const HVX_Vector *) vy0; - - uint32_t nvec = n / VLEN_FP16; - uint32_t nloe = n % VLEN_FP16; - - HVX_VectorPair rsum0_p = Q6_W_vzero(); - HVX_VectorPair rsum1_p = Q6_W_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector y_hf = y[i]; - rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf); - rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); - HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]); - rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf); - rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf); - } - - HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p))); - HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p))); - HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); - hvx_vec_store_u(s0, 8, rsum); -} - -static inline void vec_dot_f16_f16_aa_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) { - const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; - const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; - const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; - const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; - - uint32_t nvec = n / VLEN_FP16; - uint32_t nloe = n % VLEN_FP16; - - // Row sums (sf) - 4 accumulators for 2x2 tile - HVX_VectorPair r0_c0_sum_p = Q6_W_vzero(); - HVX_VectorPair r0_c1_sum_p = Q6_W_vzero(); - HVX_VectorPair r1_c0_sum_p = Q6_W_vzero(); - HVX_VectorPair r1_c1_sum_p = Q6_W_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - HVX_Vector r0_hf = x0[i]; - HVX_Vector r1_hf = x1[i]; - HVX_Vector c0_hf = y0[i]; - HVX_Vector c1_hf = y1[i]; - - // Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1 - r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf); - r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf); - r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf); - r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - - HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]); - HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]); - HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]); - HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]); - - r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf); - r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf); - r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf); - r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf); - } - - HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p))); - HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p))); - HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p))); - HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p))); - - // Reduce and store results - HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); - HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); - - hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum); // row0,col0 row1,col0 - hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum); // row0,col1 row1,col1 -} - -static inline void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { - const HVX_UVector * restrict x = (const HVX_UVector *) vx; - const HVX_UVector * restrict y = (const HVX_UVector *) vy; - - uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors - uint32_t nloe = n % VLEN_FP16; // leftover elements - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(4) - for (i = 0; i < nvec; i++) { - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]); - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - if (nloe) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]); - HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]); - - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf); - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum)); - hvx_vec_store_u(&s[0], 4, rsum); -} - -static inline void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) { - const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x; - const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y; - - uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors - uint32_t nloe = n % VLEN_FP16; // leftover elements - - const HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector rsum = Q6_V_vzero(); - - uint32_t i = 0; - - #pragma unroll(2) - for (i = 0; i < nvec; i++) { - // Load y (fp32) and convert into fp16 - HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements - HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements - HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf))); - - // Load x (fp16) - HVX_Vector x_hf = vx[i]; - - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf); - - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - if (nloe) { - // Load y (fp32) and convert into fp16 - HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements - HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements - HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf))); - - // Load x (fp16) - HVX_Vector x_hf = vx[i]; - - // Zero-out unused elements - // Note that we need to clear both x and y because they may contain NANs - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2); - x_hf = Q6_V_vand_QV(bmask, x_hf); - y_hf = Q6_V_vand_QV(bmask, y_hf); - - HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf); - - rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf))); - } - - // Convert into fp32 and reduce - rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum)); - hvx_vec_store_u(&s[0], 4, rsum); -} - -static inline 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 -) { - if (start_row >= end_row || start_col >= end_col) return; - const uint32_t nb1 = dst->nb[1]; // row stride in bytes - - const uint32_t ne11 = dst->ne[1]; - const uint32_t ne12 = dst->ne[2]; - const uint32_t ne11_12 = ne11 * ne12; - - const bool is_broadcast1 = (src2->ne[1] == 1); - const bool is_broadcast2 = (src2->ne[2] == 1); - const bool is_broadcast3 = (src2->ne[3] == 1); - - for (uint32_t r = start_row; r < end_row; r++) { - float * dst_row = (float *) ((uint8_t *) dst->data + r * nb1); - - uint32_t i13 = fastdiv(r, div_ne11_12); - uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11); - uint32_t i11 = r - i13 * ne11_12 - i12 * ne11; - - uint32_t i23 = is_broadcast3 ? 0 : i13; - uint32_t i22 = is_broadcast2 ? 0 : i12; - uint32_t i21 = is_broadcast1 ? 0 : i11; - - const float * src2_row = (const float *) ((const uint8_t *) src2->data + - i21 * src2->nb[1] + i22 * src2->nb[2] + i23 * src2->nb[3]); - - float * dst_ptr = &dst_row[start_col]; - const float * src2_ptr = &src2_row[start_col]; - int remaining = end_col - start_col; - while (remaining >= 32) { - HVX_Vector v_out = hvx_vmemu(dst_ptr); - HVX_Vector v_z = hvx_vmemu(src2_ptr); - hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z); - dst_ptr += 32; - src2_ptr += 32; - remaining -= 32; - } - if (remaining > 0) { - HVX_Vector v_out = hvx_vmemu(dst_ptr); - HVX_Vector v_z = hvx_vmemu(src2_ptr); - hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z)); - } - } -} - diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h deleted file mode 100644 index 40b65aa3b..000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h +++ /dev/null @@ -1,1200 +0,0 @@ -// Dynamic quantizers that produce tiled activations - -static inline void quantize_block_f32_q8_1_tiled(float * restrict x, uint8_t * restrict y_block) { - assert((unsigned long) x % 128 == 0); - assert((unsigned long) y_block % 128 == 0); - - HVX_Vector * vx = (HVX_Vector *) x; - HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); - HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); - HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); - HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); - - HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); - HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); - HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); - HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); - - HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); - HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); - HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); - HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); - - HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); - HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); - - HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); - HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - - HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 - HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); - HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); - - HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); - HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); - - HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); - HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); - HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - - const HVX_Vector ones = Q6_Vb_vsplat_R(1); - HVX_Vector v_sums = Q6_Vw_vrmpy_VbVb(vx_i8, ones); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 4)); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 8)); - v_sums = Q6_Vw_vadd_VwVw(v_sums, Q6_V_vror_VR(v_sums, 16)); - - float vmax0[32] __attribute__((aligned(128))); - float vmax1[32] __attribute__((aligned(128))); - float vmax2[32] __attribute__((aligned(128))); - float vmax3[32] __attribute__((aligned(128))); - int32_t sums[32] __attribute__((aligned(128))); - - hvx_vec_store_u(vmax0, 128, vmax0_sf); - hvx_vec_store_u(vmax1, 128, vmax1_sf); - hvx_vec_store_u(vmax2, 128, vmax2_sf); - hvx_vec_store_u(vmax3, 128, vmax3_sf); - hvx_vec_store_u(sums, 128, v_sums); - - float d0 = vmax0[0] / 127.0f; - float d1 = vmax1[0] / 127.0f; - float d2 = vmax2[0] / 127.0f; - float d3 = vmax3[0] / 127.0f; - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - for (int b = 0; b < 4; b++) { - HVX_Vector v_act = Q6_V_vror_VR(vx_i8, b * 32); - - HVX_Vector r0 = Q6_V_vdelta_VV(v_act, v_repl_ctrl); - HVX_Vector r1 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 4), v_repl_ctrl); - HVX_Vector r2 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 8), v_repl_ctrl); - HVX_Vector r3 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 12), v_repl_ctrl); - HVX_Vector r4 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 16), v_repl_ctrl); - HVX_Vector r5 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 20), v_repl_ctrl); - HVX_Vector r6 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 24), v_repl_ctrl); - HVX_Vector r7 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 28), v_repl_ctrl); - - __fp16 scale_h, offset_h; - if (b == 0) { - scale_h = (__fp16) d0; - offset_h = (__fp16) (sums[0] * d0); - } else if (b == 1) { - scale_h = (__fp16) d1; - offset_h = (__fp16) (sums[8] * d1); - } else if (b == 2) { - scale_h = (__fp16) d2; - offset_h = (__fp16) (sums[16] * d2); - } else { - scale_h = (__fp16) d3; - offset_h = (__fp16) (sums[24] * d3); - } - - HVX_Vector r_scale = Q6_Vh_vsplat_R(*(int16_t *)&scale_h); - HVX_Vector r_offset = Q6_Vh_vsplat_R(*(int16_t *)&offset_h); - - HVX_Vector * restrict dst = (HVX_Vector *) (y_block + b * 1280); - dst[0] = r0; - dst[1] = r1; - dst[2] = r2; - dst[3] = r3; - dst[4] = r4; - dst[5] = r5; - dst[6] = r6; - dst[7] = r7; - dst[8] = r_scale; - dst[9] = r_offset; - } -} - -static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * restrict y_block) { - assert((unsigned long) x % 128 == 0); - assert((unsigned long) y_block % 128 == 0); - - HVX_Vector * vx = (HVX_Vector *) x; - HVX_Vector zero = Q6_V_vzero(); - - HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); - HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); - HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); - HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); - - HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); - HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); - - HVX_Vector vmax_hf = hvx_vec_reduce_max_f16(hvx_vec_abs_f16(vx01_hf)); - vmax_hf = hvx_vec_reduce_max2_f16(hvx_vec_abs_f16(vx23_hf), vmax_hf); - - HVX_Vector vd_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax_hf, Q6_Vh_vsplat_R(0x2008)); - HVX_Vector vd_hf = Q6_Vhf_equals_Vqf16(vd_qf16); - - HVX_Vector vd_inv_hf = hvx_vec_inverse_f16(vd_hf); - vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd_inv_hf)); - vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd_inv_hf)); - - HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); - HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); - HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); - - HVX_Vector r_scale = hvx_vec_repl_f16(vd_hf); - - static const uint8_t __attribute__((aligned(128))) repl[128] = { - 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, - }; - HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl; - - for (int b = 0; b < 4; b++) { - HVX_Vector v_act = Q6_V_vror_VR(vx_i8, b * 32); - - HVX_Vector r0 = Q6_V_vdelta_VV(v_act, v_repl_ctrl); - HVX_Vector r1 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 4), v_repl_ctrl); - HVX_Vector r2 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 8), v_repl_ctrl); - HVX_Vector r3 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 12), v_repl_ctrl); - HVX_Vector r4 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 16), v_repl_ctrl); - HVX_Vector r5 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 20), v_repl_ctrl); - HVX_Vector r6 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 24), v_repl_ctrl); - HVX_Vector r7 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 28), v_repl_ctrl); - - HVX_Vector * restrict dst = (HVX_Vector *) (y_block + b * 1152); - dst[0] = r0; - dst[1] = r1; - dst[2] = r2; - dst[3] = r3; - dst[4] = r4; - dst[5] = r5; - dst[6] = r6; - dst[7] = r7; - dst[8] = r_scale; - } -} - -static void quantize_row_f32_q8_0_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) { - assert(k % 32 == 0); - const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (k + qk - 1) / qk; - - for (uint32_t i = 0; i < nb; i++) { - uint8_t * restrict y_block = y + i * 4 * 1152; - quantize_block_f32_q8_0_tiled(x + i * qk, y_block); - } -} - -static void quantize_row_f32_q8_1_tiled(float * restrict x, uint8_t * restrict y, uint32_t k) { - assert(k % 32 == 0); - const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (k + qk - 1) / qk; - - for (uint32_t i = 0; i < nb; i++) { - uint8_t * restrict y_block = y + i * 4 * 1280; - quantize_block_f32_q8_1_tiled(x + i * qk, y_block); - } -} - -// Dot kernels & helpers that consume tiled activations - -static inline HVX_Vector hvx_vec_mul_f16_f16_to_f32_lower32(HVX_Vector v1, HVX_Vector v2) { -#if __HVX_ARCH__ >= 79 - HVX_VectorPair p = Q6_Wsf_vmpy_VhfVhf(v1, v2); - return Q6_V_lo_W(Q6_W_vshuff_VVR(Q6_V_hi_W(p), Q6_V_lo_W(p), -4)); -#else - HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(v1, v2); - HVX_Vector hi = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(p)); - HVX_Vector lo = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(p)); - return Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -4)); -#endif -} - -static inline HVX_Vector unpack_and_interleave_4bit(HVX_Vector v_a, HVX_Vector v_b, HVX_Vector mask_h4) { - HVX_Vector v_W0 = Q6_V_vand_VV(v_a, mask_h4); - HVX_Vector v_W1 = Q6_Vub_vlsr_VubR(v_a, 4); - HVX_Vector v_W2 = Q6_V_vand_VV(v_b, mask_h4); - HVX_Vector v_W3 = Q6_Vub_vlsr_VubR(v_b, 4); - - HVX_VectorPair v01_pair = Q6_W_vshuff_VVR(v_W1, v_W0, -1); - HVX_VectorPair v23_pair = Q6_W_vshuff_VVR(v_W3, v_W2, -1); - HVX_VectorPair v0123_pair = Q6_W_vshuff_VVR(Q6_V_lo_W(v23_pair), Q6_V_lo_W(v01_pair), -2); - return Q6_V_lo_W(v0123_pair); -} - -static inline HVX_VectorPair unpack_and_interleave_4bit_x2(HVX_Vector v_src, HVX_Vector mask_h4) { - HVX_Vector v_lo = Q6_V_vand_VV(v_src, mask_h4); - HVX_Vector v_hi = Q6_Vub_vlsr_VubR(v_src, 4); - HVX_VectorPair v01_pair = Q6_W_vshuff_VVR(v_hi, v_lo, -1); - HVX_Vector v01_lo = Q6_V_lo_W(v01_pair); - HVX_Vector v01_hi = Q6_V_hi_W(v01_pair); - - HVX_Vector v23_lo = Q6_V_valign_VVR(v01_hi, v01_lo, 64); - HVX_Vector v_W0 = Q6_V_lo_W(Q6_W_vshuff_VVR(v23_lo, v01_lo, -2)); - - HVX_Vector v67_lo = Q6_V_valign_VVR(v01_lo, v01_hi, 64); - HVX_Vector v_W1 = Q6_V_lo_W(Q6_W_vshuff_VVR(v67_lo, v01_hi, -2)); - - return Q6_W_vcombine_VV(v_W1, v_W0); -} - -static inline HVX_Vector accum_4bit_32x1( - const HVX_Vector * restrict vptr, - const HVX_Vector * restrict v_act, - HVX_Vector i8 -) { - HVX_Vector v_sum0 = Q6_V_vzero(); - HVX_Vector v_sum1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - - #pragma unroll - for (int i = 0; i < 4; i++) { - HVX_VectorPair v_W_pair = unpack_and_interleave_4bit_x2(vptr[i], mask_h4); - HVX_Vector v_W0 = Q6_Vb_vsub_VbVb(Q6_V_lo_W(v_W_pair), i8); - HVX_Vector v_W1 = Q6_Vb_vsub_VbVb(Q6_V_hi_W(v_W_pair), i8); - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act[i * 2 + 0]); - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act[i * 2 + 1]); - } - - return Q6_Vw_vadd_VwVw(v_sum0, v_sum1); -} - -static inline HVX_Vector accum_4bit_32x1_lut( - const HVX_Vector * restrict vptr, - const HVX_Vector * restrict v_act, - HVX_Vector mask_h4, - HVX_Vector lut -) { - HVX_Vector v_sum0 = Q6_V_vzero(); - HVX_Vector v_sum1 = Q6_V_vzero(); - - #pragma unroll - for (int i = 0; i < 4; i++) { - HVX_VectorPair v_W_pair = unpack_and_interleave_4bit_x2(vptr[i], mask_h4); - HVX_Vector v_W0 = Q6_Vb_vlut32_VbVbI(Q6_V_lo_W(v_W_pair), lut, 0); - HVX_Vector v_W1 = Q6_Vb_vlut32_VbVbI(Q6_V_hi_W(v_W_pair), lut, 0); - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act[i * 2 + 0]); - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act[i * 2 + 1]); - } - - return Q6_Vw_vadd_VwVw(v_sum0, v_sum1); -} - -static inline HVX_VectorPair accum_4bit_32x2( - const HVX_Vector * restrict vptr, - const HVX_Vector * restrict v_act0, - const HVX_Vector * restrict v_act1, - HVX_Vector i8 -) { - HVX_Vector v_sum0 = Q6_V_vzero(); - HVX_Vector v_sum1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - - #pragma unroll - for (int i = 0; i < 4; i++) { - HVX_VectorPair v_W_pair = unpack_and_interleave_4bit_x2(vptr[i], mask_h4); - HVX_Vector v_W0 = Q6_Vb_vsub_VbVb(Q6_V_lo_W(v_W_pair), i8); - HVX_Vector v_W1 = Q6_Vb_vsub_VbVb(Q6_V_hi_W(v_W_pair), i8); - - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act0[i * 2 + 0]); - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W1, v_act0[i * 2 + 1]); - - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W0, v_act1[i * 2 + 0]); - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act1[i * 2 + 1]); - } - - return Q6_W_vcombine_VV(v_sum1, v_sum0); -} - -static inline HVX_VectorPair accum_4bit_32x2_lut( - const HVX_Vector * restrict vptr, - const HVX_Vector * restrict v_act0, - const HVX_Vector * restrict v_act1, - HVX_Vector mask_h4, - HVX_Vector lut -) { - HVX_Vector v_sum0 = Q6_V_vzero(); - HVX_Vector v_sum1 = Q6_V_vzero(); - - #pragma unroll - for (int i = 0; i < 4; i++) { - HVX_VectorPair v_W_pair = unpack_and_interleave_4bit_x2(vptr[i], mask_h4); - HVX_Vector v_W0 = Q6_Vb_vlut32_VbVbI(Q6_V_lo_W(v_W_pair), lut, 0); - HVX_Vector v_W1 = Q6_Vb_vlut32_VbVbI(Q6_V_hi_W(v_W_pair), lut, 0); - - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W0, v_act0[i * 2 + 0]); - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W1, v_act0[i * 2 + 1]); - - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W0, v_act1[i * 2 + 0]); - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W1, v_act1[i * 2 + 1]); - } - - return Q6_W_vcombine_VV(v_sum1, v_sum0); -} - -static inline HVX_Vector accum_q8_0_32x1( - const HVX_Vector * restrict vptr, - const HVX_Vector * restrict v_act -) { - HVX_Vector v_sum = Q6_V_vzero(); - #pragma unroll - for (int g = 0; g < 8; g++) { - HVX_Vector v_rot = Q6_V_vror_VR(vptr[g], 64); - HVX_Vector v_W = Q6_V_lo_W(Q6_W_vshuff_VVR(v_rot, vptr[g], -2)); - v_sum = Q6_Vw_vrmpyacc_VwVbVb(v_sum, v_W, v_act[g]); - } - return v_sum; -} - -static inline HVX_VectorPair accum_q8_0_32x2( - const HVX_Vector * restrict vptr, - const HVX_Vector * restrict v_act0, - const HVX_Vector * restrict v_act1 -) { - HVX_Vector v_sum0 = Q6_V_vzero(); - HVX_Vector v_sum1 = Q6_V_vzero(); - #pragma unroll - for (int g = 0; g < 8; g++) { - HVX_Vector v_rot = Q6_V_vror_VR(vptr[g], 64); - HVX_Vector v_W = Q6_V_lo_W(Q6_W_vshuff_VVR(v_rot, vptr[g], -2)); - v_sum0 = Q6_Vw_vrmpyacc_VwVbVb(v_sum0, v_W, v_act0[g]); - v_sum1 = Q6_Vw_vrmpyacc_VwVbVb(v_sum1, v_W, v_act1[g]); - } - return Q6_W_vcombine_VV(v_sum1, v_sum0); -} - -static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1152); - - HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act, i8); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[4]; - HVX_Vector v_scale_a = v_act[8]; - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector i8 = Q6_Vb_vsplat_R(8); - - uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_4bit_32x2(vptr0, v_act0_0, v_act1_0, i8); - HVX_VectorPair v_sums1 = accum_4bit_32x2(vptr1, v_act0_1, v_act1_1, i8); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = vptr0[4]; - HVX_Vector v_scale_w1 = vptr1[4]; - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); - const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); - - HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0, v_act1, i8); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[4]; - HVX_Vector v_scale_a_c0 = v_act0[8]; - HVX_Vector v_scale_a_c1 = v_act1[8]; - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1280); - - HVX_Vector v_sum = accum_4bit_32x1(vptr, v_act, Q6_V_vzero()); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_offset = vptr[4]; - HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); - HVX_Vector v_scale = Q6_V_lo_W(p_deal); - HVX_Vector v_offset = Q6_V_hi_W(p_deal); - - HVX_Vector v_scale_a = v_act[8]; - HVX_Vector v_sum_a = v_act[9]; - - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a); - HVX_Vector v_offset_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a); - - HVX_Vector v_scaled_dot = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - HVX_Vector v_sum_scaled = hvx_vec_add_f32_f32(v_scaled_dot, v_offset_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - - uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1280); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1280); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1280); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1280); - - HVX_VectorPair v_sums0 = accum_4bit_32x2(vptr0, v_act0_0, v_act1_0, Q6_V_vzero()); - HVX_VectorPair v_sums1 = accum_4bit_32x2(vptr1, v_act0_1, v_act1_1, Q6_V_vzero()); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_offset0 = vptr0[4]; - HVX_VectorPair p_deal0 = Q6_W_vdeal_VVR(v_scale_offset0, v_scale_offset0, -2); - HVX_Vector v_scale0 = Q6_V_lo_W(p_deal0); - HVX_Vector v_offset0 = Q6_V_hi_W(p_deal0); - - HVX_Vector v_scale_offset1 = vptr1[4]; - HVX_VectorPair p_deal1 = Q6_W_vdeal_VVR(v_scale_offset1, v_scale_offset1, -2); - HVX_Vector v_scale1 = Q6_V_lo_W(p_deal1); - HVX_Vector v_offset1 = Q6_V_hi_W(p_deal1); - - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_sum_a_c0_0 = v_act0_0[9]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_sum_a_c1_0 = v_act1_0[9]; - - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_sum_a_c0_1 = v_act0_1[9]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - HVX_Vector v_sum_a_c1_1 = v_act1_1[9]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale0, v_scale_a_c0_0); - HVX_Vector v_offset_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset0, v_sum_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale0, v_scale_a_c1_0); - HVX_Vector v_offset_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset0, v_sum_a_c1_0); - - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale1, v_scale_a_c0_1); - HVX_Vector v_offset_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset1, v_sum_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale1, v_scale_a_c1_1); - HVX_Vector v_offset_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset1, v_sum_a_c1_1); - - HVX_Vector v_scaled_dot_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_add_f32_f32(v_scaled_dot_c0_0, v_offset_comb_c0_0); - - HVX_Vector v_scaled_dot_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_add_f32_f32(v_scaled_dot_c1_0, v_offset_comb_c1_0); - - HVX_Vector v_scaled_dot_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_add_f32_f32(v_scaled_dot_c0_1, v_offset_comb_c0_1); - - HVX_Vector v_scaled_dot_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_add_f32_f32(v_scaled_dot_c1_1, v_offset_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1280); - const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1280); - - HVX_VectorPair v_sums = accum_4bit_32x2(vptr, v_act0, v_act1, Q6_V_vzero()); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_offset = vptr[4]; - HVX_VectorPair p_deal = Q6_W_vdeal_VVR(v_scale_offset, v_scale_offset, -2); - HVX_Vector v_scale = Q6_V_lo_W(p_deal); - HVX_Vector v_offset = Q6_V_hi_W(p_deal); - - HVX_Vector v_scale_a_c0 = v_act0[8]; - HVX_Vector v_sum_a_c0 = v_act0[9]; - HVX_Vector v_scale_a_c1 = v_act1[8]; - HVX_Vector v_sum_a_c1 = v_act1[9]; - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a_c0); - HVX_Vector v_offset_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a_c0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale, v_scale_a_c1); - HVX_Vector v_offset_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_offset, v_sum_a_c1); - - HVX_Vector v_scaled_dot_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c0 = hvx_vec_add_f32_f32(v_scaled_dot_c0, v_offset_comb_c0); - - HVX_Vector v_scaled_dot_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - HVX_Vector v_sum_scaled_c1 = hvx_vec_add_f32_f32(v_scaled_dot_c1, v_offset_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); - const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1152); - - HVX_Vector v_sum = accum_q8_0_32x1(vptr, v_act); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[8]; - HVX_Vector v_scale_a = v_act[8]; - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - - uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 1152); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 1152); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_q8_0_32x2(vptr0, v_act0_0, v_act1_0); - HVX_VectorPair v_sums1 = accum_q8_0_32x2(vptr1, v_act0_1, v_act1_1); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = vptr0[8]; - HVX_Vector v_scale_w1 = vptr1[8]; - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 1152); - const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); - const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); - - HVX_VectorPair v_sums = accum_q8_0_32x2(vptr, v_act0, v_act1); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[8]; - HVX_Vector v_scale_a_c0 = v_act0[8]; - HVX_Vector v_scale_a_c1 = v_act1[8]; - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1152); - - HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act, mask_h4, lut); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = vptr[4]; - HVX_Vector v_scale_a = v_act[8]; - HVX_Vector v_scale_comb = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_iq4nl_lut; - - uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_4bit_32x2_lut(vptr0, v_act0_0, v_act1_0, mask_h4, lut); - HVX_VectorPair v_sums1 = accum_4bit_32x2_lut(vptr1, v_act0_1, v_act1_1, mask_h4, lut); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = vptr0[4]; - HVX_Vector v_scale_w1 = vptr1[4]; - HVX_Vector v_scale_a_c0_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_1 = v_act1_1[8]; - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); - const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); - - HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0, v_act1, mask_h4, lut); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = vptr[4]; - HVX_Vector v_scale_a_c0 = v_act0[8]; - HVX_Vector v_scale_a_c1 = v_act1[8]; - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f16_f16_to_f32_lower32(v_scale_w, v_scale_a_c1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y_q = vy; - - HVX_Vector v_sum_float = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut; - HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0; - HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); - - uint32_t n_k_tiles = n / 32; - for (uint32_t kt = 0; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act = (const HVX_Vector *) (y_q + kt * 1152); - - HVX_Vector v_sum = accum_4bit_32x1_lut(vptr, v_act, mask_h4, lut); - HVX_Vector v_sum_sf = Q6_Vsf_equals_Vw(v_sum); - - HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512); - HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand); - r0_d = Q6_V_vand_VV(r0_d, e8m0_mask); - HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23); - - HVX_Vector v_scale_a_f16 = v_act[8]; - HVX_VectorPair p_scale_a_f32 = hvx_vec_f16_to_f32_shuff(v_scale_a_f16); - HVX_Vector v_scale_a = Q6_V_lo_W(p_scale_a_f32); - - HVX_Vector v_scale_comb = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a); - HVX_Vector v_sum_scaled = hvx_vec_mul_f32_f32(v_sum_sf, v_scale_comb); - - v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled); - } - - v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f)); - - if (sz) { - hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz))); - } else { - hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float); - } -} - -static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { - const uint8_t * restrict tile_ptr = vx; - const uint8_t * restrict y0_q = vy0; - const uint8_t * restrict y1_q = vy1; - - HVX_Vector v_sum_float_c0 = Q6_V_vzero(); - HVX_Vector v_sum_float_c1 = Q6_V_vzero(); - HVX_Vector mask_h4 = Q6_Vb_vsplat_R(0x0F); - HVX_Vector lut = *(const HVX_Vector *) kvalues_mxfp4_lut; - HVX_Vector expand = *(const HVX_Vector *) expand_x32_e8m0; - HVX_Vector e8m0_mask = Q6_V_vsplat_R(0x000000ff); - - uint32_t n_k_tiles = n / 32; - uint32_t kt = 0; - for (; kt + 1 < n_k_tiles; kt += 2) { - const HVX_Vector * restrict vptr0 = (const HVX_Vector *) (tile_ptr + (kt + 0) * 640); - const HVX_Vector * restrict v_act0_0 = (const HVX_Vector *) (y0_q + (kt + 0) * 1152); - const HVX_Vector * restrict v_act1_0 = (const HVX_Vector *) (y1_q + (kt + 0) * 1152); - - const HVX_Vector * restrict vptr1 = (const HVX_Vector *) (tile_ptr + (kt + 1) * 640); - const HVX_Vector * restrict v_act0_1 = (const HVX_Vector *) (y0_q + (kt + 1) * 1152); - const HVX_Vector * restrict v_act1_1 = (const HVX_Vector *) (y1_q + (kt + 1) * 1152); - - HVX_VectorPair v_sums0 = accum_4bit_32x2_lut(vptr0, v_act0_0, v_act1_0, mask_h4, lut); - HVX_VectorPair v_sums1 = accum_4bit_32x2_lut(vptr1, v_act0_1, v_act1_1, mask_h4, lut); - - HVX_Vector v_sum_c0_0 = Q6_V_lo_W(v_sums0); - HVX_Vector v_sum_c1_0 = Q6_V_hi_W(v_sums0); - HVX_Vector v_sum_c0_1 = Q6_V_lo_W(v_sums1); - HVX_Vector v_sum_c1_1 = Q6_V_hi_W(v_sums1); - - HVX_Vector v_sum_sf_c0_0 = Q6_Vsf_equals_Vw(v_sum_c0_0); - HVX_Vector v_sum_sf_c1_0 = Q6_Vsf_equals_Vw(v_sum_c1_0); - HVX_Vector v_sum_sf_c0_1 = Q6_Vsf_equals_Vw(v_sum_c0_1); - HVX_Vector v_sum_sf_c1_1 = Q6_Vsf_equals_Vw(v_sum_c1_1); - - HVX_Vector v_scale_w0 = hvx_vmem(tile_ptr + (kt + 0) * 640 + 512); - HVX_Vector r0_d0 = Q6_V_vdelta_VV(v_scale_w0, expand); - r0_d0 = Q6_V_vand_VV(r0_d0, e8m0_mask); - HVX_Vector v_scale_w_f32_0 = Q6_Vw_vasl_VwR(r0_d0, 23); - - HVX_Vector v_scale_w1 = hvx_vmem(tile_ptr + (kt + 1) * 640 + 512); - HVX_Vector r0_d1 = Q6_V_vdelta_VV(v_scale_w1, expand); - r0_d1 = Q6_V_vand_VV(r0_d1, e8m0_mask); - HVX_Vector v_scale_w_f32_1 = Q6_Vw_vasl_VwR(r0_d1, 23); - - HVX_Vector v_scale_a_c0_f16_0 = v_act0_0[8]; - HVX_Vector v_scale_a_c1_f16_0 = v_act1_0[8]; - HVX_Vector v_scale_a_c0_f16_1 = v_act0_1[8]; - HVX_Vector v_scale_a_c1_f16_1 = v_act1_1[8]; - - HVX_VectorPair p_scale_a_c0_f32_0 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16_0); - HVX_VectorPair p_scale_a_c1_f32_0 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16_0); - HVX_VectorPair p_scale_a_c0_f32_1 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16_1); - HVX_VectorPair p_scale_a_c1_f32_1 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16_1); - - HVX_Vector v_scale_a_c0_0 = Q6_V_lo_W(p_scale_a_c0_f32_0); - HVX_Vector v_scale_a_c1_0 = Q6_V_lo_W(p_scale_a_c1_f32_0); - HVX_Vector v_scale_a_c0_1 = Q6_V_lo_W(p_scale_a_c0_f32_1); - HVX_Vector v_scale_a_c1_1 = Q6_V_lo_W(p_scale_a_c1_f32_1); - - HVX_Vector v_scale_comb_c0_0 = hvx_vec_mul_f32_f32(v_scale_w_f32_0, v_scale_a_c0_0); - HVX_Vector v_scale_comb_c1_0 = hvx_vec_mul_f32_f32(v_scale_w_f32_0, v_scale_a_c1_0); - HVX_Vector v_scale_comb_c0_1 = hvx_vec_mul_f32_f32(v_scale_w_f32_1, v_scale_a_c0_1); - HVX_Vector v_scale_comb_c1_1 = hvx_vec_mul_f32_f32(v_scale_w_f32_1, v_scale_a_c1_1); - - HVX_Vector v_sum_scaled_c0_0 = hvx_vec_mul_f32_f32(v_sum_sf_c0_0, v_scale_comb_c0_0); - HVX_Vector v_sum_scaled_c1_0 = hvx_vec_mul_f32_f32(v_sum_sf_c1_0, v_scale_comb_c1_0); - HVX_Vector v_sum_scaled_c0_1 = hvx_vec_mul_f32_f32(v_sum_sf_c0_1, v_scale_comb_c0_1); - HVX_Vector v_sum_scaled_c1_1 = hvx_vec_mul_f32_f32(v_sum_sf_c1_1, v_scale_comb_c1_1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vec_add_f32_f32(v_sum_scaled_c0_0, v_sum_scaled_c0_1)); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vec_add_f32_f32(v_sum_scaled_c1_0, v_sum_scaled_c1_1)); - } - - for (; kt < n_k_tiles; kt++) { - const HVX_Vector * restrict vptr = (const HVX_Vector *) (tile_ptr + kt * 640); - const HVX_Vector * restrict v_act0 = (const HVX_Vector *) (y0_q + kt * 1152); - const HVX_Vector * restrict v_act1 = (const HVX_Vector *) (y1_q + kt * 1152); - - HVX_VectorPair v_sums = accum_4bit_32x2_lut(vptr, v_act0, v_act1, mask_h4, lut); - HVX_Vector v_sum_c0 = Q6_V_lo_W(v_sums); - HVX_Vector v_sum_c1 = Q6_V_hi_W(v_sums); - - HVX_Vector v_sum_sf_c0 = Q6_Vsf_equals_Vw(v_sum_c0); - HVX_Vector v_sum_sf_c1 = Q6_Vsf_equals_Vw(v_sum_c1); - - HVX_Vector v_scale_w = hvx_vmem(tile_ptr + kt * 640 + 512); - HVX_Vector r0_d = Q6_V_vdelta_VV(v_scale_w, expand); - r0_d = Q6_V_vand_VV(r0_d, e8m0_mask); - HVX_Vector v_scale_w_f32 = Q6_Vw_vasl_VwR(r0_d, 23); - - HVX_Vector v_scale_a_c0_f16 = v_act0[8]; - HVX_Vector v_scale_a_c1_f16 = v_act1[8]; - - HVX_VectorPair p_scale_a_c0_f32 = hvx_vec_f16_to_f32_shuff(v_scale_a_c0_f16); - HVX_VectorPair p_scale_a_c1_f32 = hvx_vec_f16_to_f32_shuff(v_scale_a_c1_f16); - - HVX_Vector v_scale_a_c0 = Q6_V_lo_W(p_scale_a_c0_f32); - HVX_Vector v_scale_a_c1 = Q6_V_lo_W(p_scale_a_c1_f32); - - HVX_Vector v_scale_comb_c0 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a_c0); - HVX_Vector v_scale_comb_c1 = hvx_vec_mul_f32_f32(v_scale_w_f32, v_scale_a_c1); - - HVX_Vector v_sum_scaled_c0 = hvx_vec_mul_f32_f32(v_sum_sf_c0, v_scale_comb_c0); - HVX_Vector v_sum_scaled_c1 = hvx_vec_mul_f32_f32(v_sum_sf_c1, v_scale_comb_c1); - - v_sum_float_c0 = hvx_vec_add_f32_f32(v_sum_float_c0, v_sum_scaled_c0); - v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1); - } - - v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f)); - v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f)); - - if (sz0) { - hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0))); - } else { - hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0); - } - if (sz1) { - hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1))); - } else { - hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1); - } -} - -static inline void quantize_f32_q8_0_tiled_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_row_size, - size_t dst_row_size -) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_0_tiled((float *) tmp_data, dst_data, ne0); - dst_data += dst_row_size; - src_data += src_row_size; - } -} - -static inline void quantize_f32_q8_1_tiled_kernel( - const uint8_t * restrict src_data, - uint8_t * restrict dst_data, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t nrows, - size_t src_row_size, - size_t dst_row_size -) { - const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float)); - hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float)); - - for (uint32_t i = 0; i < nrows; ++i) { - hex_l2fetch(src_data, src_row_size, src_row_size, 2); - hvx_copy_f32_aa(tmp_data, src_data, ne0); - - quantize_row_f32_q8_1_tiled((float *) tmp_data, dst_data, ne0); - dst_data += dst_row_size; - src_data += src_row_size; - } -} - -static inline void quantize_f32_q8_0_tiled_block_kernel( - const float * restrict src, - uint8_t * restrict dst, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t ib_first, - uint32_t ib_last, - size_t src_row_size, - size_t dst_row_size, - uint32_t r, - uint32_t c -) { - const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (ne0 + qk - 1) / qk; - - for (uint32_t ib = ib_first; ib < ib_last; ++ib) { - const uint8_t * restrict src_ptr = (const uint8_t *) src + r * src_row_size + c * qk * sizeof(float); - uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 1152; - - hex_l2fetch(src_ptr, qk * sizeof(float), qk * sizeof(float), 1); - - if (c == nb - 1) { - uint32_t active_elements = ne0 - c * qk; - hvx_splat_f32_a(tmp_data, 0.0f, qk); - hvx_copy_f32_aa(tmp_data, src_ptr, active_elements); - } else { - hvx_copy_f32_aa(tmp_data, src_ptr, qk); - } - - quantize_block_f32_q8_0_tiled((float *) tmp_data, dst_ptr); - - c++; - if (c == nb) { - c = 0; - r++; - } - } -} - -static inline void quantize_f32_q8_1_tiled_block_kernel( - const float * restrict src, - uint8_t * restrict dst, - uint8_t * restrict tmp_data, - uint32_t ne0, - uint32_t ib_first, - uint32_t ib_last, - size_t src_row_size, - size_t dst_row_size, - uint32_t r, - uint32_t c -) { - const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (ne0 + qk - 1) / qk; - - for (uint32_t ib = ib_first; ib < ib_last; ++ib) { - const uint8_t * restrict src_ptr = (const uint8_t *) src + r * src_row_size + c * qk * sizeof(float); - uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 1280; - - hex_l2fetch(src_ptr, qk * sizeof(float), qk * sizeof(float), 1); - - if (c == nb - 1) { - uint32_t active_elements = ne0 - c * qk; - hvx_splat_f32_a(tmp_data, 0.0f, qk); - hvx_copy_f32_aa(tmp_data, src_ptr, active_elements); - } else { - hvx_copy_f32_aa(tmp_data, src_ptr, qk); - } - - quantize_block_f32_q8_1_tiled((float *) tmp_data, dst_ptr); - - c++; - if (c == nb) { - c = 0; - r++; - } - } -} diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h deleted file mode 100644 index 96369825a..000000000 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ /dev/null @@ -1,495 +0,0 @@ -#ifndef HTP_MATMUL_OPS_H -#define HTP_MATMUL_OPS_H - -#include -#include -#include "htp-ops.h" -#include "hex-fastdiv.h" -#include "hex-common.h" - -#ifdef __cplusplus -extern "C" { -#endif - -// --- HMX Tile Constraints --- -#define HTP_MM_HMX_TILE_N_COLS 32 -#define HTP_MM_HMX_TILE_N_ROWS 32 -#define HTP_MM_HMX_TILE_SIZE (32 * 32 * sizeof(__fp16)) // 2048 bytes -#define HTP_MM_HMX_TILE_N_ELMS 1024 -#define HTP_MM_HMX_MIN_NROWS 4 - -// --- Weight Repacked Tile Sizes --- -#define HTP_MM_WEIGHT_TILE_SIZE_Q4_0 576 -#define HTP_MM_WEIGHT_TILE_SIZE_Q4_1 640 -#define HTP_MM_WEIGHT_TILE_SIZE_Q8_0 1088 -#define HTP_MM_WEIGHT_TILE_SIZE_IQ4_NL 576 -#define HTP_MM_WEIGHT_TILE_SIZE_MXFP4 544 - -// --- Weight Repacked Aligned Tile Sizes --- -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_0 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0 1152 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_IQ4_NL 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4 640 - -// --- Activation Tiled Block Sizes (including padding) --- -#define HTP_MM_ACT_TILE_SIZE_Q8_0 1152 -#define HTP_MM_ACT_TILE_SIZE_Q8_1 1280 - -#define HTP_MM_MAX_PREFETCH 16 - -// --- Solver Cost Model Penalty Weights (HMX-specific) --- -#define HTP_MM_HMX_COST_W_DEQUANT 3 // cost penalty for quantized weight loading/dequantization -#define HTP_MM_HMX_COST_A_CONVERT 2 // cost penalty for activation loading/conversion - -// --- DMA Activation Transfer Configuration --- -#define HTP_MM_DMA_ACT_ROWS_PER_STEP 2 -#define HTP_MM_DMA_ACT_MULTIPLIER 4 - -enum htp_mm_kernel_type { - HTP_MM_KERNEL_UNSUPPORTED = 0, - - // HMX paths - HTP_MM_KERNEL_HMX_2D, - HTP_MM_KERNEL_HMX_F16_BATCHED, - - // HVX floating-point paths - HTP_MM_KERNEL_HVX_F16_F16_VTCM, - HTP_MM_KERNEL_HVX_F16_F16_DDR, - HTP_MM_KERNEL_HVX_F16_F32_DDR, - - HTP_MM_KERNEL_HVX_F32_F32_VTCM, - HTP_MM_KERNEL_HVX_F32_F32_DDR, - HTP_MM_KERNEL_HVX_F32_F16_DDR, - - // HVX quantized paths - HTP_MM_KERNEL_HVX_QUANT_ROW, // standard row-wise parallel quantization - HTP_MM_KERNEL_HVX_QUANT_BLOCK, // parallel block-wise quantization - HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, // row-wise fallback flat quantization -}; - -// Op-specific struct for precomputed matmul params -struct htp_mm_kernel_params { - int32_t kernel_type; // enum htp_mm_kernel_type - int32_t pipeline; // 1 = pipelined execution, 0 = standard - int32_t m_chunk; // Row chunk size (M chunk) - int32_t n_chunk; // Col chunk size (N chunk) - int32_t n_threads; // Number of threads to spawn - int32_t n_act_threads; // Number of threads for activation preparation - int32_t n_hmx; // 1 = use HMX, 0 = use HVX - int32_t n_prefetch; // Prefetch lookahead buffers/rows in VTCM - int32_t tile_size; // Weight tile size - int32_t aligned_tile_size; // Aligned weight tile size (padded to 128) - int32_t src1_row_size; // Row size for quantized activation - int32_t vtcm_size; // Total required scratchpad size in VTCM - int32_t vtcm_src0_size; // src0 scratchpad size in VTCM - int32_t vtcm_src1_size; // src1 scratchpad size in VTCM - int32_t vtcm_src2_size; // src2 scratchpad size in VTCM (fused only) - int32_t vtcm_src3_size; // src3 scratchpad size in VTCM (fused only) - int32_t vtcm_dst_size; // dst scratchpad size in VTCM - - // Precomputed division values - struct fastdiv_values div_ne12_ne1; - struct fastdiv_values div_ne1; - struct fastdiv_values div_r2; - struct fastdiv_values div_r3; - struct fastdiv_values div_ne11; -}; - -#if defined(__cplusplus) -static_assert(sizeof(struct htp_mm_kernel_params) <= 128, "htp_matmul_kernel_params is too large for kernel_params blob"); -#else -_Static_assert(sizeof(struct htp_mm_kernel_params) <= 128, "htp_matmul_kernel_params is too large for kernel_params blob"); -#endif - -struct mmid_row_mapping { - uint32_t i1; - uint32_t i2; -}; - -// Search for optimal (mc, nc) chunk sizes within VTCM budget. -static inline int htp_mm_hmx_compute_chunks(size_t vtcm_total, - size_t overhead, - size_t per_n_cost, - size_t per_m_cost, - size_t per_mn_cost, - size_t m, - size_t n, - size_t m_block_cost, - size_t n_block_cost, - size_t * m_chunk_out, - size_t * n_chunk_out, - size_t * total_out) { - if (m == 0 || n == 0) return -1; - if (vtcm_total <= overhead) return -1; - if (per_n_cost == 0 || per_m_cost == 0 || per_mn_cost == 0) return -1; - - const size_t usable = vtcm_total - overhead; - - size_t best_cost = SIZE_MAX; - size_t best_mn = 0; - size_t best_m = 0, best_n = 0; - - const size_t n_max = hex_align_down((size_t)n, HTP_MM_HMX_TILE_N_COLS); - for (size_t nc = n_max; nc >= HTP_MM_HMX_TILE_N_COLS; nc -= HTP_MM_HMX_TILE_N_COLS) { - size_t n_fixed = 0, ncmn = 0, mc_denom = 0; - if (hex_mul_overflow(nc, per_n_cost, &n_fixed)) continue; - if (n_fixed >= usable) goto next_nc; - - if (hex_mul_overflow(nc, per_mn_cost, &ncmn)) goto next_nc; - if (hex_add_overflow(per_m_cost, ncmn, &mc_denom) || mc_denom == 0) goto next_nc; - - { - size_t remain = usable - n_fixed; - size_t mc = remain / mc_denom; - mc = hex_align_down(mc, HTP_MM_HMX_TILE_N_ROWS); - mc = hex_smin(mc, m); - - if (mc == 0) { - goto next_nc; - } - - size_t mblocks = ((size_t) m + mc - 1) / mc; - size_t nblocks = ((size_t) n + nc - 1) / nc; - size_t cost = mblocks * m_block_cost + nblocks * n_block_cost; - size_t mn = mc * nc; - if (cost < best_cost || (cost == best_cost && mn > best_mn)) { - best_cost = cost; - best_mn = mn; - best_m = mc; - best_n = nc; - } - } - -next_nc: - if (nc == HTP_MM_HMX_TILE_N_COLS) break; // avoid size_t underflow - } - - if (best_m == 0 || best_n == 0) return -1; - - // Compute exact total (with overflow checks) - size_t t0 = 0, t1 = 0, t2 = 0, mn = 0, total = 0; - if (hex_mul_overflow(best_n, per_n_cost, &t0)) return -1; - if (hex_mul_overflow(best_m, per_m_cost, &t1)) return -1; - if (hex_mul_overflow(best_m, best_n, &mn)) return -1; - if (hex_mul_overflow(mn, per_mn_cost, &t2)) return -1; - if (hex_add_overflow(t0, t1, &total)) return -1; - if (hex_add_overflow(total, t2, &total)) return -1; - if (hex_add_overflow(total, overhead, &total)) return -1; - - *m_chunk_out = best_m; - *n_chunk_out = best_n; - *total_out = total; - return 0; -} - -// --- Tile Size Helpers --- -static inline uint32_t htp_mm_get_weight_tile_size(int weight_type) { - switch (weight_type) { - case HTP_TYPE_Q4_0: - case HTP_TYPE_IQ4_NL: - return HTP_MM_WEIGHT_TILE_SIZE_Q4_0; - case HTP_TYPE_Q4_1: - return HTP_MM_WEIGHT_TILE_SIZE_Q4_1; - case HTP_TYPE_Q8_0: - return HTP_MM_WEIGHT_TILE_SIZE_Q8_0; - case HTP_TYPE_MXFP4: - return HTP_MM_WEIGHT_TILE_SIZE_MXFP4; - default: - return 0; - } -} - -static inline uint32_t htp_mm_get_weight_aligned_tile_size(int weight_type) { - switch (weight_type) { - case HTP_TYPE_Q4_0: - case HTP_TYPE_IQ4_NL: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_0; - case HTP_TYPE_Q4_1: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1; - case HTP_TYPE_Q8_0: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0; - case HTP_TYPE_MXFP4: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4; - default: - return 0; - } -} - -// --- Activation/Row Size Helpers --- -static inline size_t htp_mm_q8_0_tiled_row_size(uint32_t ne) { - const uint32_t ne_padded = ((ne + 127) / 128) * 128; - const uint32_t nb_32 = ne_padded / 32; - return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_0; -} - -static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) { - const uint32_t ne_padded = ((ne + 127) / 128) * 128; - const uint32_t nb_32 = ne_padded / 32; - return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1; -} - -static inline size_t htp_mm_q8_0_flat_row_size(uint32_t ne) { - const uint32_t quants_size = hex_align_up(ne, 128); - const uint32_t num_scales = (ne + 31) / 32; - const uint32_t scales_size = hex_align_up(num_scales * 2, 128); - return quants_size + scales_size; -} - -static inline size_t htp_mm_q8_1_flat_row_size(uint32_t ne) { - const uint32_t quants_size = hex_align_up(ne, 128); - const uint32_t num_scales = (ne + 31) / 32; - const uint32_t scales_size = hex_align_up(num_scales * 4, 128); - return quants_size + scales_size; -} - -static inline size_t htp_mm_get_tiled_row_stride(int weight_type, uint32_t k) { - uint32_t nb = (k + QK_Q4_0_TILED - 1) / QK_Q4_0_TILED; - switch (weight_type) { - case HTP_TYPE_Q4_0: - case HTP_TYPE_IQ4_NL: - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q8_0: - case HTP_TYPE_MXFP4: - return (size_t) nb * htp_mm_get_weight_tile_size(weight_type); - case HTP_TYPE_F16: - return (size_t) k * sizeof(__fp16); - case HTP_TYPE_F32: - return (size_t) k * sizeof(float); - default: - return 0; - } -} - -static inline size_t htp_mm_round_up(size_t n, size_t m) { - return ((n + m - 1) / m) * m; -} - -static inline bool htp_mm_hmx_pipeline(uint32_t m) { - return m > 32; -} - -static inline void htp_mm_hmx_get_2d_chunk_costs( - int wtype, uint32_t k, bool pipeline, uint32_t aligned_tile_size, - size_t * size_per_n_out, size_t * size_per_m_out, size_t * size_per_mn_out -) { - const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); - const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k); - const size_t vec_dot_size = k * sizeof(uint16_t); - const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; - const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; - - *size_per_n_out = (pipeline ? 2 : 1) * (is_quant ? qweight_row_stride : row_stride) + - (pipeline ? 2 * vec_dot_size : vec_dot_size); - *size_per_m_out = vec_dot_size; - *size_per_mn_out = (pipeline ? 2 : 1) * sizeof(uint16_t); -} - -static inline void htp_mm_hmx_get_batched_chunk_costs( - uint32_t k, uint32_t group_size, - size_t * size_per_n_out, size_t * size_per_m_out, size_t * size_per_mn_out -) { - const size_t vec_dot_size = k * sizeof(uint16_t); - *size_per_n_out = 3 * vec_dot_size; - *size_per_m_out = group_size * vec_dot_size; - *size_per_mn_out = sizeof(uint16_t); -} - -static inline size_t htp_mm_hmx_get_2d_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size -) { - const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; - const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); - const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k); - const size_t vec_dot_size = k * sizeof(uint16_t); - - const size_t act_f32_size = htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE); - size_t weight_area_size = is_quant - ? htp_mm_round_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE) - : htp_mm_round_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE); - if (pipeline) { - weight_area_size *= 2; - } - const size_t act_area_size = htp_mm_round_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t output_area_size = htp_mm_round_up(mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); - - size_t scratch0_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - size_t scratch1_size = pipeline ? scratch0_size : 0; - size_t scratch2_size = pipeline ? output_area_size : 0; - - return weight_area_size + act_area_size + act_f32_size + output_area_size + - scratch0_size + scratch1_size + scratch2_size + 256; -} - -static inline size_t htp_mm_hmx_get_batched_vtcm_size( - int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) { - (void)wtype; - (void)pipeline; - const size_t vec_dot_size = k * sizeof(uint16_t); - const size_t f32_scratch_size = use_dma_activation - ? htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0; - - const size_t act_head_stride = mc * k; - const size_t weight_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - const size_t act_area_size = htp_mm_round_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); - const size_t output_area_size = htp_mm_round_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE); - const size_t scratch_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); - - return weight_area_size + act_area_size + output_area_size + - 2 * scratch_area_size + 256 + f32_scratch_size; -} - -static inline size_t htp_mm_hvx_get_vtcm_sizes( - int kernel_type, - int wtype, - uint32_t ne10, // k - uint32_t src1_nrows, // m_total (or act_nrows) - uint32_t n_threads, - size_t dst_row_size, - size_t src0_row_size, - size_t src1_row_size, - uint32_t n_prefetch, - size_t * vtcm_src0_size_out, - size_t * vtcm_src1_size_out, - size_t * vtcm_dst_size_out -) { - size_t vtcm_src0_size = 0; - size_t vtcm_src1_size = 0; - size_t vtcm_dst_size = 0; - - const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 || - wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || - wtype == HTP_TYPE_MXFP4); - - const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); - const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1; - - switch (kernel_type) { - case HTP_MM_KERNEL_HVX_F16_F16_VTCM: { - size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128); - vtcm_src1_size = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256); - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; - vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } - case HTP_MM_KERNEL_HVX_F16_F32_DDR: - case HTP_MM_KERNEL_HVX_F16_F16_DDR: - case HTP_MM_KERNEL_HVX_F32_F32_DDR: - case HTP_MM_KERNEL_HVX_F32_F16_DDR: { - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads; - vtcm_src1_size = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads; - vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } - case HTP_MM_KERNEL_HVX_F32_F32_VTCM: { - size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128); - vtcm_src1_size = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256); - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; - vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0; - break; - } - case HTP_MM_KERNEL_HVX_QUANT_BLOCK: - case HTP_MM_KERNEL_HVX_QUANT_ROW: { - size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); - vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256); - - vtcm_src0_size = vtcm_src0_size * n_threads; - - if (is_repack) { - uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - uint32_t n_k_tiles = ne10 / 32; - uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - vtcm_src0_size = repacked_vtcm_size * n_threads; - } - - size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); - size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0; - if (dst_size_per_thread < quant_scratch_size_per_thread) { - dst_size_per_thread = quant_scratch_size_per_thread; - } - vtcm_dst_size = dst_size_per_thread * n_threads; - break; - } - case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: { - size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - - vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); - vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256); - - vtcm_src0_size = vtcm_src0_size * n_threads; - - if (is_repack) { - uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - uint32_t n_k_tiles = ne10 / 32; - uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - vtcm_src0_size = repacked_vtcm_size * n_threads; - } - - size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); - size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0; - if (dst_size_per_thread < quant_scratch_size_per_thread) { - dst_size_per_thread = quant_scratch_size_per_thread; - } - vtcm_dst_size = dst_size_per_thread * n_threads; - break; - } - default: - break; - } - - *vtcm_src0_size_out = vtcm_src0_size; - *vtcm_src1_size_out = vtcm_src1_size; - *vtcm_dst_size_out = vtcm_dst_size; - - return vtcm_src0_size + vtcm_src1_size + vtcm_dst_size; -} - -static inline size_t htp_mm_hvx_id_get_vtcm_sizes( - int wtype, - uint32_t ne10, // k - uint32_t src1_nrows, - uint32_t n_threads, - size_t src0_row_size, // nb01 - uint32_t n_prefetch, - size_t * vtcm_src0_size_out, - size_t * vtcm_src1_size_out, - size_t * vtcm_dst_size_out -) { - const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 || - wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || - wtype == HTP_TYPE_MXFP4); - - const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); - const size_t src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) - : htp_mm_q8_0_tiled_row_size(ne10); - - size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); - size_t src1_sz = htp_mm_round_up(src1_row_size * src1_nrows, 256); - - if (is_repack) { - const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - const uint32_t n_k_tiles = ne10 / 32; - const uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - src0_sz_per_thread = repacked_vtcm_size; - } - - const size_t vtcm_src0_size = src0_sz_per_thread * n_threads; - const size_t vtcm_dst_size = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; - - *vtcm_src0_size_out = vtcm_src0_size; - *vtcm_src1_size_out = src1_sz; - *vtcm_dst_size_out = vtcm_dst_size; - - return vtcm_src0_size + src1_sz + vtcm_dst_size; -} - -#ifdef __cplusplus -} -#endif - -#endif // HTP_MATMUL_OPS_H diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index e7ac21a2b..15290c3d1 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1869,6 +1869,29 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d(ggml_met return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw(ggml_metal_library_t lib, const ggml_tensor * op, bool tiled) { + assert(op->op == GGML_OP_CONV_2D_DW); + + GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_conv_2d_dw%s_%s_%s", + tiled ? "_tiled" : "", + ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_CONV_3D); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index dc75a34b2..9d4aca121 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -152,6 +152,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_tran struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tiled); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_upscale (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index d51f30b13..732eb9a85 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1204,6 +1204,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); + case GGML_OP_CONV_2D_DW: + return op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); case GGML_OP_UPSCALE: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_POOL_1D: diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 6bf61423c..d6761023b 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -656,6 +656,34 @@ typedef struct { int32_t d1; } ggml_metal_kargs_conv_2d; +typedef struct { + uint64_t nb00; // kernel strides + uint64_t nb01; + uint64_t nb02; + uint64_t nb10; // input strides + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb0; // output strides + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t IW; // input width + int32_t IH; // input height + int32_t KW; // kernel width + int32_t KH; // kernel height + int32_t C; // channels (IC == OC for depthwise) + int32_t OW; // output width + int32_t OH; // output height + int32_t N; // batch size + int32_t s0; // stride x + int32_t s1; // stride y + int32_t p0; // padding x + int32_t p1; // padding y + int32_t d0; // dilation x + int32_t d1; // dilation y +} ggml_metal_kargs_conv_2d_dw; + typedef struct { uint64_t ofs0; uint64_t ofs1; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index d196bae4f..45909c477 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -387,6 +387,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_conv_2d(ctx, idx); } break; + case GGML_OP_CONV_2D_DW: + { + n_fuse = ggml_metal_op_conv_2d_dw(ctx, idx); + } break; case GGML_OP_CONV_TRANSPOSE_1D: { n_fuse = ggml_metal_op_conv_transpose_1d(ctx, idx); @@ -3742,6 +3746,86 @@ int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_conv_2d_dw(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + + GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); + + const int32_t s0 = ((const int32_t *) op->op_params)[0]; + const int32_t s1 = ((const int32_t *) op->op_params)[1]; + const int32_t p0 = ((const int32_t *) op->op_params)[2]; + const int32_t p1 = ((const int32_t *) op->op_params)[3]; + const int32_t d0 = ((const int32_t *) op->op_params)[4]; + const int32_t d1 = ((const int32_t *) op->op_params)[5]; + + ggml_metal_kargs_conv_2d_dw args = { + /*.nb00 =*/ nb00, + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb03, + /*.nb10 =*/ nb10, + /*.nb11 =*/ nb11, + /*.nb12 =*/ nb12, + /*.nb13 =*/ nb13, + /*.nb0 =*/ nb0, + /*.nb1 =*/ nb1, + /*.nb2 =*/ nb2, + /*.nb3 =*/ nb3, + /*.IW =*/ ne10, + /*.IH =*/ ne11, + /*.KW =*/ ne00, + /*.KH =*/ ne01, + /*.C =*/ ne12, + /*.OW =*/ ne0, + /*.OH =*/ ne1, + /*.N =*/ ne13, + /*.s0 =*/ s0, + /*.s1 =*/ s1, + /*.p0 =*/ p0, + /*.p1 =*/ p1, + /*.d0 =*/ d0, + /*.d1 =*/ d1, + }; + + const bool use_tiled = (nb12 < nb10); + + auto pipeline = ggml_metal_library_get_pipeline_conv_2d_dw(lib, op, use_tiled); + + int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); + nth = std::min(nth, 256); + nth = std::max(nth, 1); + + const int32_t OW = ne0; + const int32_t OH = ne1; + const int32_t C = ne12; + const int32_t N = ne13; + + const int tg_x = use_tiled ? (C + nth - 1) / nth : (OW + nth - 1) / nth; + const int tg_y = OH; + const int tg_z = use_tiled ? OW * N : C * N; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + + ggml_metal_encoder_dispatch_threadgroups(enc, tg_x, tg_y, tg_z, nth, 1, 1); + + return 1; +} + int ggml_metal_op_conv_3d(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 13f5274c7..0bebd836a 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -75,6 +75,7 @@ int ggml_metal_op_norm (ggml_metal_op_t ctx, int idx); int ggml_metal_op_rope (ggml_metal_op_t ctx, int idx); int ggml_metal_op_im2col (ggml_metal_op_t ctx, int idx); int ggml_metal_op_conv_2d (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_conv_2d_dw (ggml_metal_op_t ctx, int idx); int ggml_metal_op_conv_3d (ggml_metal_op_t ctx, int idx); int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx); int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index dcb6803f5..6b6f9fd87 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -4908,6 +4908,202 @@ kernel void kernel_conv_2d( uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]); +// grid: x = C tile, y = OH, z = OW * N (for channel-contiguous layouts) +template +kernel void kernel_conv_2d_dw_tiled( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]) { + + const int32_t c = (int32_t)(tgpig.x * ntg.x + tpitg.x); + if (c >= args.C) { + return; + } + + const int32_t oh = tgpig.y; + const int32_t own = tgpig.z; + const int32_t ow = own % args.OW; + const int32_t n = own / args.OW; + + const int32_t base_y = oh*args.s1 - args.p1; + + int32_t ky_start = 0; + if (base_y < 0) { + ky_start = (-base_y + args.d1 - 1)/args.d1; + } + int32_t ky_end = args.KH; + const int32_t y_max = args.IH - 1 - base_y; + if (y_max < 0) { + ky_end = ky_start; + } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) { + ky_end = min(ky_end, y_max/args.d1 + 1); + } + + const int32_t base_x = ow*args.s0 - args.p0; + + int32_t kx_start = 0; + if (base_x < 0) { + kx_start = (-base_x + args.d0 - 1)/args.d0; + } + int32_t kx_end = args.KW; + const int32_t x_max = args.IW - 1 - base_x; + if (x_max < 0) { + kx_end = kx_start; + } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) { + kx_end = min(kx_end, x_max/args.d0 + 1); + } + + float acc = 0.0f; + + if (ky_start < ky_end && kx_start < kx_end) { + const uint64_t w_base = (uint64_t) c * args.nb02; + const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12; + + for (int32_t ky = ky_start; ky < ky_end; ++ky) { + const int32_t iy = base_y + ky*args.d1; + const uint64_t src_row = src_base + (uint64_t) iy * args.nb11; + const uint64_t w_row = w_base + (uint64_t) ky * args.nb01; + + for (int32_t kx = kx_start; kx < kx_end; ++kx) { + const int32_t ix = base_x + kx*args.d0; + const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10); + const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00)); + acc += x * w; + } + } + } + + const uint64_t dst_offs = + (uint64_t) n * args.nb3 + + (uint64_t) c * args.nb2 + + (uint64_t) oh * args.nb1 + + (uint64_t) ow * args.nb0; + + *(device float *)(dst + dst_offs) = acc; +} + +// grid: x = OW tile, y = OH, z = C * N (for spatially-contiguous layouts) +template +kernel void kernel_conv_2d_dw( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]) { + + const int32_t oh = tgpig.y; + const int32_t cn = tgpig.z; + const int32_t c = cn % args.C; + const int32_t n = cn / args.C; + + const int32_t base_y = oh*args.s1 - args.p1; + + int32_t ky_start = 0; + if (base_y < 0) { + ky_start = (-base_y + args.d1 - 1)/args.d1; + } + int32_t ky_end = args.KH; + const int32_t y_max = args.IH - 1 - base_y; + if (y_max < 0) { + ky_end = ky_start; + } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) { + ky_end = min(ky_end, y_max/args.d1 + 1); + } + + const uint64_t w_base = (uint64_t) c * args.nb02; + const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12; + + const int32_t ow = (int32_t)(tgpig.x * ntg.x + tpitg.x); + if (ow >= args.OW) { + return; + } + + float acc = 0.0f; + + const int32_t base_x = ow*args.s0 - args.p0; + + int32_t kx_start = 0; + if (base_x < 0) { + kx_start = (-base_x + args.d0 - 1)/args.d0; + } + int32_t kx_end = args.KW; + const int32_t x_max = args.IW - 1 - base_x; + if (x_max < 0) { + kx_end = kx_start; + } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) { + kx_end = min(kx_end, x_max/args.d0 + 1); + } + + if (ky_start < ky_end && kx_start < kx_end) { + for (int32_t ky = ky_start; ky < ky_end; ++ky) { + const int32_t iy = base_y + ky*args.d1; + const uint64_t src_row = src_base + (uint64_t) iy * args.nb11; + const uint64_t w_row = w_base + (uint64_t) ky * args.nb01; + + for (int32_t kx = kx_start; kx < kx_end; ++kx) { + const int32_t ix = base_x + kx*args.d0; + const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10); + const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00)); + acc += x * w; + } + } + } + + const uint64_t dst_offs = + (uint64_t) n * args.nb3 + + (uint64_t) c * args.nb2 + + (uint64_t) oh * args.nb1 + + (uint64_t) ow * args.nb0; + + *(device float *)(dst + dst_offs) = acc; +} + +template [[host_name("kernel_conv_2d_dw_f32_f32")]] +kernel void kernel_conv_2d_dw( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + +template [[host_name("kernel_conv_2d_dw_f16_f32")]] +kernel void kernel_conv_2d_dw( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + +template [[host_name("kernel_conv_2d_dw_tiled_f32_f32")]] +kernel void kernel_conv_2d_dw_tiled( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + +template [[host_name("kernel_conv_2d_dw_tiled_f16_f32")]] +kernel void kernel_conv_2d_dw_tiled( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + typedef void (conv_transpose_1d_t)( constant ggml_metal_kargs_conv_transpose_1d & args, device const float * src0, diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8d8eb4916..d955ba9ab 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -10343,7 +10343,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } // Only use mask opt when the mask is fairly large. This hasn't been tuned extensively. - bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16; + bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 + && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256); vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc, mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type); @@ -16341,7 +16342,18 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg uint32_t submit_count = 0; uint64_t batch_flops = 0; uint64_t total_flops = 0; - uint64_t flops_per_submit = std::min(uint64_t(200'000'000'000), ctx->last_total_flops / 40u); + uint64_t flops_cap = 200'000'000'000ULL; + + // On weaker AMD GPUs larger submissions can hit a driver timeout, submit more often to avoid this + if (ctx->device->vendor_id == VK_VENDOR_ID_AMD && ctx->device->shader_core_count > 0) { + if (ctx->device->architecture == AMD_GCN && ctx->device->shader_core_count < 32) { + flops_cap = 500'000'000ULL * ctx->device->shader_core_count; + } else if (ctx->device->architecture != AMD_GCN && ctx->device->shader_core_count < 24) { + flops_cap = 2'000'000'000ULL * ctx->device->shader_core_count; + } + } + uint64_t flops_per_submit = std::min(flops_cap, ctx->last_total_flops / 40u); + for (int i = 0; i < cgraph->n_nodes; i++) { if (first_node_in_batch) { submit_node_idx = i; diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/set_rows_quant.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/set_rows_quant.wgsl deleted file mode 100644 index 876e65b6a..000000000 --- a/ggml/src/ggml-webgpu/wgsl-shaders/set_rows_quant.wgsl +++ /dev/null @@ -1,224 +0,0 @@ -#ifdef DST_Q8_0 -#define BLOCK_SIZE 32u -#define BLOCK_BYTES 34u -#define QS_WORDS 8u -#elif defined(DST_Q4_0) -#define BLOCK_SIZE 32u -#define BLOCK_BYTES 18u -#define QS_WORDS 4u -#endif - -@group(0) @binding(0) -var src: array; - -@group(0) @binding(1) -var idx: array; - -@group(0) @binding(2) -#ifdef PAIR_BLOCKS -var dst: array; -#else -var dst: array>; -#endif - -#ifdef I64_IDX -@group(0) @binding(3) -var error: atomic; -#define PARAMS_BINDING 4 -#else -#define PARAMS_BINDING 3 -#endif - -struct Params { - offset_src: u32, // in elements - offset_idx: u32, // in elements - offset_dst: u32, // in blocks - - // Strides (in elements / blocks) - stride_src1: u32, - stride_src2: u32, - stride_src3: u32, - - stride_idx0: u32, - stride_idx1: u32, - stride_idx2: u32, - - stride_dst1: u32, - stride_dst2: u32, - stride_dst3: u32, - - // Shape of src - ne0: u32, - n_rows: u32, - ne2: u32, - ne3: u32, - - // Shape of idx - idx1: u32, - idx2: u32, -}; - -@group(0) @binding(PARAMS_BINDING) -var params: Params; - -// if the quantization type is unaligned and there are an odd number of blocks per row, we need to store atomically -#ifndef PAIR_BLOCKS -fn merge_store_dst_word(word_idx: u32, mask: u32, bits: u32) { - loop { - let old = atomicLoad(&dst[word_idx]); - let merged = (old & ~mask) | (bits & mask); - let result = atomicCompareExchangeWeak(&dst[word_idx], old, merged); - if (result.exchanged) { - return; - } - } -} -#else -fn merge_store_dst_word(word_idx: u32, mask: u32, bits: u32) { - let old = dst[word_idx]; - dst[word_idx] = (old & ~mask) | (bits & mask); -} -#endif - -fn store_u16(dst_word_idx: u32, block_byte_offset: u32, byte_offset: u32, value: u32) { - let total_byte_offset = block_byte_offset + byte_offset; - let word_idx = dst_word_idx + total_byte_offset / 4u; - let shift = (total_byte_offset & 2u) * 8u; - let mask = 0xFFFFu << shift; - merge_store_dst_word(word_idx, mask, (value & 0xFFFFu) << shift); -} - -fn store_u32(dst_word_idx: u32, block_byte_offset: u32, byte_offset: u32, value: u32) { - let total_byte_offset = block_byte_offset + byte_offset; - let word_idx = dst_word_idx + total_byte_offset / 4u; - let shift = (total_byte_offset & 3u) * 8u; - - if (shift == 0u) { -#ifdef PAIR_BLOCKS - dst[word_idx] = value; -#else - atomicStore(&dst[word_idx], value); -#endif - return; - } - - let lo_mask = 0xFFFFFFFFu << shift; - let hi_mask = (1u << shift) - 1u; - merge_store_dst_word(word_idx, lo_mask, value << shift); - merge_store_dst_word(word_idx + 1u, hi_mask, value >> (32u - shift)); -} - -fn quantize_block_params(src_block: u32) -> vec2 { -#ifdef DST_Q8_0 - var amax = 0.0; - for (var j: u32 = 0u; j < BLOCK_SIZE; j++) { - amax = max(amax, abs(src[src_block + j])); - } - - let d = amax / 127.0; - let id = select(0.0, 1.0 / d, d > 0.0); - return vec2(d, id); -#elif defined(DST_Q4_0) - var amax = 0.0; - var max_val = 0.0; - for (var j: u32 = 0u; j < BLOCK_SIZE; j++) { - let v = src[src_block + j]; - let av = abs(v); - if (amax < av) { - amax = av; - max_val = v; - } - } - - let d = max_val / -8.0; - let id = select(0.0, 1.0 / d, d != 0.0); - return vec2(d, id); -#endif -} - -fn quantize_block_word(src_block: u32, j: u32, id: f32) -> u32 { -#ifdef DST_Q8_0 - let base = src_block + j * 4u; - return (u32(i32(round(src[base + 0u] * id)) & 0xFF) << 0u) | - (u32(i32(round(src[base + 1u] * id)) & 0xFF) << 8u) | - (u32(i32(round(src[base + 2u] * id)) & 0xFF) << 16u) | - (u32(i32(round(src[base + 3u] * id)) & 0xFF) << 24u); -#elif defined(DST_Q4_0) - var packed_q = 0u; - for (var k: u32 = 0u; k < 4u; k++) { - let x0 = src[src_block + j * 4u + k] * id; - let x1 = src[src_block + 16u + j * 4u + k] * id; - let q0 = u32(clamp(i32(x0 + 8.5), 0, 15)); - let q1 = u32(clamp(i32(x1 + 8.5), 0, 15)); - packed_q |= (q0 & 0xFu) << (8u * k); - packed_q |= (q1 & 0xFu) << (8u * k + 4u); - } - return packed_q; -#endif -} - -fn quantize_block(src_block: u32, dst_word_idx: u32, block_byte_offset: u32) { - let params = quantize_block_params(src_block); - let d = params.x; - let id = params.y; - let packed_d = pack2x16float(vec2(d, 0.0)) & 0xFFFFu; - store_u16(dst_word_idx, block_byte_offset, 0u, packed_d); - - for (var j: u32 = 0u; j < QS_WORDS; j++) { - store_u32(dst_word_idx, block_byte_offset, 2u + j * 4u, quantize_block_word(src_block, j, id)); - } -} - -@compute @workgroup_size(WG_SIZE) -fn main(@builtin(global_invocation_id) gid: vec3) { - let blocks_per_row = params.ne0 / BLOCK_SIZE; -#ifdef PAIR_BLOCKS - let blocks_per_invocation = 2u; -#else - let blocks_per_invocation = 1u; -#endif - let invocations_per_row = blocks_per_row / blocks_per_invocation; - let total_invocations = params.ne3 * params.ne2 * params.n_rows * invocations_per_row; - if (gid.x >= total_invocations) { - return; - } - - var i = gid.x / invocations_per_row; - let block_in_row = (gid.x % invocations_per_row) * blocks_per_invocation; - - let i_src3 = i / (params.ne2 * params.n_rows); - i = i % (params.ne2 * params.n_rows); - let i_src2 = i / params.n_rows; - let i_src1 = i % params.n_rows; - - let i_idx2 = i_src3 % params.idx2; - let i_idx1 = i_src2 % params.idx1; - let i_idx0 = i_src1; - -#ifdef I64_IDX - let idx_high = (params.offset_idx + i_idx0 * params.stride_idx0 + i_idx1 * params.stride_idx1 + i_idx2 * params.stride_idx2) * 2u; - let idx_val = idx[idx_high]; - let idx_low_val = idx[idx_high + 1u]; - - if (idx_low_val != 0u) { - atomicStore(&error, 1u); - return; - } -#else - let idx_i = params.offset_idx + i_idx0 * params.stride_idx0 + i_idx1 * params.stride_idx1 + i_idx2 * params.stride_idx2; - let idx_val = idx[idx_i]; -#endif - - let dst_row_blocks = params.offset_dst + idx_val * params.stride_dst1 + i_src2 * params.stride_dst2 + i_src3 * params.stride_dst3; - let src_row = params.offset_src + i_src1 * params.stride_src1 + i_src2 * params.stride_src2 + i_src3 * params.stride_src3; - let src_block = src_row + block_in_row * BLOCK_SIZE; - let dst_block_byte = (dst_row_blocks + block_in_row) * BLOCK_BYTES; - - let dst_word_idx = dst_block_byte / 4u; -#ifdef PAIR_BLOCKS - quantize_block(src_block, dst_word_idx, 0u); - quantize_block(src_block + BLOCK_SIZE, dst_word_idx, BLOCK_BYTES); -#else - quantize_block(src_block, dst_word_idx, dst_block_byte & 3u); -#endif -} diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 5436717c4..2b98a552f 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -379,6 +379,8 @@ bool llama_batch_allocr::init( LLAMA_LOG_ERROR("%s: sequence %d positions are decreasing (not allowed)\n", __func__, seq_id); return false; } + + cur_seq_pos[seq_id] = pos; } } } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 49c9a9911..dc120ea15 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -17,6 +17,7 @@ #include #include #include +#include // // llama_context @@ -33,6 +34,30 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) { throw std::runtime_error("Unsupported ctx type"); } +struct llm_fused_op_probe { + llm_fused_op op; + const char * name; + uint32_t n_tokens_per_seq; +}; + +static const llm_fused_op_probe llm_fused_op_flash_attn_probe = { + /*.op =*/ LLM_FUSED_OP_FLASH_ATTN, + /*.name =*/ "Flash Attention", + /*.n_tokens_per_seq =*/ 1, +}; + +static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = { + /*.op =*/ LLM_FUSED_OP_GDN_AR, + /*.name =*/ "fused Gated Delta Net (autoregressive)", + /*.n_tokens_per_seq =*/ 1, +}; + +static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = { + /*.op =*/ LLM_FUSED_OP_GDN_CH, + /*.name =*/ "fused Gated Delta Net (chunked)", + /*.n_tokens_per_seq =*/ 16, +}; + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -444,6 +469,69 @@ llama_context::~llama_context() { ggml_opt_free(opt_ctx); } +void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) { + const char * func = __func__; + auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) { + if (!enabled) { + return; + } + + const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs; + + auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true); + if (!gf) { + throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check"); + } + + bool device_mismatch = false; + for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) { + if (node.op != probe.op) { + continue; + } + + GGML_ASSERT(node.il >= 0); + + ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor); + ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr; + + // TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior, + // but is still wrong for cases like --no-kv-offload. + ggml_backend_dev_t device_layer = model.dev_layer(node.il); + + if (device_fused != device_layer) { + LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s " + "is assigned to device %s (usually due to missing support)\n", + func, node.il, + device_layer ? ggml_backend_dev_name(device_layer) : "none", + probe.name, + device_fused ? ggml_backend_dev_name(device_fused) : "none"); + device_mismatch = true; + break; + } + } + + if (device_mismatch) { + enabled = false; + LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name); + } else { + enabled = true; + LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name); + } + }; + + if (cparams.auto_fa) { + resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn); + cparams.auto_fa = false; + } + + if (cparams.auto_fgdn) { + LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func); + resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar); + resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch); + cparams.auto_fgdn = false; + } +} + void llama_context::sched_reserve() { if (!sched_need_reserve) { return; @@ -484,128 +572,7 @@ void llama_context::sched_reserve() { LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs); - // resolve automatic Flash Attention use - if (cparams.auto_fa) { - auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true); - if (!gf) { - throw std::runtime_error("failed to reserve graph for Flash Attention check"); - } - - const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FATTN) + 1; - bool fa_device_mismatch = false; - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - ggml_tensor * n = ggml_graph_node(gf, i); - if (n->op != GGML_OP_FLASH_ATTN_EXT) { - continue; - } - ggml_backend_dev_t device_fa = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n)); - - // TODO: instead of the tensor names, use a map to keep track of which (FA) tensors belong to which layer - GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FATTN "-", prefix_len) == 0); - const int il = std::stoi(n->name + prefix_len); - ggml_backend_dev_t device_kv = model.dev_layer(il); - if (device_fa != device_kv) { - LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the Flash Attention tensor " - "is assigned to device %s (usually due to missing support)\n", - __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_fa)); - // FIXME: fa_device_mismatch logic is wrong for --no-kv-offload, but this is broken anyways - fa_device_mismatch = true; - break; - } - } - - if (fa_device_mismatch) { - cparams.flash_attn = false; - LLAMA_LOG_WARN("%s: Flash Attention was auto, set to disabled\n", __func__); - } else { - cparams.flash_attn = true; - LLAMA_LOG_INFO("%s: Flash Attention was auto, set to enabled\n", __func__); - } - - cparams.auto_fa = false; - } - - if (cparams.auto_fgdn) { - LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__); - - if (cparams.fused_gdn_ar) { - auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true); - if (!gf) { - throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (autoregressive)"); - } - - const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_AR) + 1; - bool gdn_device_mismatch = false; - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - ggml_tensor * n = ggml_graph_node(gf, i); - if (n->op != GGML_OP_GATED_DELTA_NET) { - continue; - } - ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n)); - - GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_AR "-", prefix_len) == 0); - const int il = std::stoi(n->name + prefix_len); - ggml_backend_dev_t device_kv = model.dev_layer(il); - if (device_gdn != device_kv) { - LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor " - "is assigned to device %s (usually due to missing support)\n", - __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn)); - gdn_device_mismatch = true; - break; - } - } - - if (gdn_device_mismatch) { - cparams.fused_gdn_ar = false; - LLAMA_LOG_WARN("%s: fused Gated Delta Net (autoregressive) not supported, set to disabled\n", __func__); - } else { - LLAMA_LOG_INFO("%s: fused Gated Delta Net (autoregressive) enabled\n", __func__); - } - } - - if (cparams.fused_gdn_ch) { - // more than one token in the batch per sequence in order to take the chunked path - // note: n_outputs must match n_tokens for embedding models with mean/rank pooling, - // because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies - // it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens, - // the ggml_mul_mat assertion fails. - const uint32_t n_tokens_ch = 16*n_seqs; - auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true); - if (!gf) { - throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (chunked)"); - } - - const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_CH) + 1; - bool gdn_device_mismatch = false; - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - ggml_tensor * n = ggml_graph_node(gf, i); - if (n->op != GGML_OP_GATED_DELTA_NET) { - continue; - } - ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n)); - - GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_CH "-", prefix_len) == 0); - const int il = std::stoi(n->name + prefix_len); - ggml_backend_dev_t device_kv = model.dev_layer(il); - if (device_gdn != device_kv) { - LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor " - "is assigned to device %s (usually due to missing support)\n", - __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn)); - gdn_device_mismatch = true; - break; - } - } - - if (gdn_device_mismatch) { - cparams.fused_gdn_ch = false; - LLAMA_LOG_WARN("%s: fused Gated Delta Net (chunked) not supported, set to disabled\n", __func__); - } else { - LLAMA_LOG_INFO("%s: fused Gated Delta Net (chunked) enabled\n", __func__); - } - } - - cparams.auto_fgdn = false; - } + resolve_fused_ops(mctx.get(), n_seqs); // reserve worst-case graph int n_splits_pp = -1; diff --git a/src/llama-context.h b/src/llama-context.h index f8b780587..bf91daa8b 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -262,6 +262,10 @@ private: llm_graph_cb graph_get_cb() const; + // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device + // that differs from the layer it belongs to (usually due to missing backend support) + void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs); + // TODO: read/write lora adapters and cvec size_t state_write_data(llama_io_write_i & io); size_t state_read_data (llama_io_read_i & io); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 39b828c1b..4a1fe0f46 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1193,6 +1193,7 @@ void llm_graph_result::reset() { params = {}; inputs.clear(); + fused_nodes.clear(); buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false)); @@ -1294,6 +1295,10 @@ llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) { return inputs.back().get(); } +void llm_graph_result::add_fused_node(llm_graph_fused_node result) { + fused_nodes.push_back(result); +} + void llm_graph_result::set_params(const llm_graph_params & params) { this->params = params; } @@ -1353,6 +1358,8 @@ void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const { } } + + ggml_tensor * llm_graph_context::build_cvec( ggml_tensor * cur, int il) const { @@ -2403,7 +2410,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); - cb(cur, LLAMA_TENSOR_NAME_FATTN, il); + res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); diff --git a/src/llama-graph.h b/src/llama-graph.h index 4b5b75c63..97141ef93 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -38,6 +38,12 @@ enum llm_graph_type { LLM_GRAPH_TYPE_DECODER_MTP, }; +enum llm_fused_op { + LLM_FUSED_OP_FLASH_ATTN, + LLM_FUSED_OP_GDN_AR, + LLM_FUSED_OP_GDN_CH, +}; + enum llm_ffn_op_type : int { LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use LLM_FFN_SILU, @@ -775,6 +781,12 @@ struct llm_graph_params { } }; +struct llm_graph_fused_node { + llm_fused_op op; + ggml_tensor * tensor; + int il; +}; + class llm_graph_result { public: llm_graph_result(int64_t max_nodes); @@ -808,6 +820,10 @@ public: llm_graph_input_i * add_input(llm_graph_input_ptr input); + void add_fused_node(llm_graph_fused_node result); + + const std::vector & get_fused_nodes() const { return fused_nodes; } + void set_params(const llm_graph_params & params); // important graph nodes @@ -826,6 +842,7 @@ public: std::map t_sampled_probs; std::vector inputs; + std::vector fused_nodes; ggml_context_ptr ctx_compute; diff --git a/src/llama-impl.h b/src/llama-impl.h index 2d06752c9..4988b06d2 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -103,7 +103,3 @@ std::string llama_format_tensor_shape(const std::vector & ne); std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); - -#define LLAMA_TENSOR_NAME_FATTN "__fattn__" -#define LLAMA_TENSOR_NAME_FGDN_AR "__fgdn_ar__" -#define LLAMA_TENSOR_NAME_FGDN_CH "__fgdn_ch__" diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index cf5e38095..ad6612647 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -401,9 +401,9 @@ std::pair llm_build_delta_net_base::build_delta_ne // K=1: output carries the final state only. state s is 4D [S_v, S_v, H_v, n_seqs]. ggml_tensor * result = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, /*K=*/1); if (n_tokens == 1) { - cb(result, LLAMA_TENSOR_NAME_FGDN_AR, il); + res->add_fused_node({LLM_FUSED_OP_GDN_AR, result, il}); } else { - cb(result, LLAMA_TENSOR_NAME_FGDN_CH, il); + res->add_fused_node({LLM_FUSED_OP_GDN_CH, result, il}); } ggml_tensor * output = ggml_view_4d(ctx0, result, @@ -566,9 +566,9 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); if (n_seq_tokens > 1) { - cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_CH, il); + res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il}); } else { - cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_AR, il); + res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il}); } const int64_t attn_score_elems = S_v * H_v * n_seq_tokens * n_seqs; diff --git a/tools/cli/cli-client.cpp b/tools/cli/cli-client.cpp new file mode 100644 index 000000000..1c563335b --- /dev/null +++ b/tools/cli/cli-client.cpp @@ -0,0 +1,130 @@ +#include "cli-client.h" + +#include "http.h" + +#include +#include +#include + +// generation can stall for a long time during prompt processing, so the +// read timeout must be generous +static constexpr time_t CLI_HTTP_READ_TIMEOUT_SEC = 3600; + +// upper bound for the accumulated response body kept for error reporting +static constexpr size_t CLI_HTTP_MAX_ERROR_BODY = 1024 * 1024; + +// returns the path with the base url's path prefix prepended (if any) +static std::string join_path(const common_http_url & parts, const std::string & path) { + if (parts.path.empty() || parts.path == "/") { + return path; + } + std::string prefix = parts.path; + if (prefix.back() == '/') { + prefix.pop_back(); + } + return prefix + path; +} + +std::string cli_client::get(const std::string & path) { + auto [cli, parts] = common_http_client(server_base); + cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0); + auto path_with_model = path + (model.empty() ? "" : ("?model=" + model)); + auto res = cli.Get(join_path(parts, path_with_model)); + if (!res) { + throw std::runtime_error("failed to connect to " + server_base + ": " + httplib::to_string(res.error())); + } + if (res->status < 200 || res->status >= 300) { + throw std::runtime_error("GET " + path + " failed with status " + std::to_string(res->status) + ": " + res->body); + } + return res->body; +} + +std::string cli_client::post(const std::string & path, const std::string & body) { + auto [cli, parts] = common_http_client(server_base); + cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0); + auto res = cli.Post(join_path(parts, path), body, "application/json"); + if (!res) { + throw std::runtime_error("failed to connect to " + server_base + ": " + httplib::to_string(res.error())); + } + if (res->status < 200 || res->status >= 300) { + throw std::runtime_error("POST " + path + " failed with status " + std::to_string(res->status) + ": " + res->body); + } + return res->body; +} + +std::string cli_client::post_sse(const std::string & path, + const std::string & body, + const std::function & should_stop, + const std::function & on_data) { + auto [cli, parts] = common_http_client(server_base); + cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0); + + std::string pending; // buffer for incomplete SSE lines + std::string raw_body; // accumulated body, used only for error reporting + + auto receiver = [&](const char * data, size_t len) -> bool { + if (should_stop()) { + return false; // aborts the request + } + if (raw_body.size() < CLI_HTTP_MAX_ERROR_BODY) { + raw_body.append(data, std::min(len, CLI_HTTP_MAX_ERROR_BODY - raw_body.size())); + } + pending.append(data, len); + size_t pos; + while ((pos = pending.find('\n')) != std::string::npos) { + std::string line = pending.substr(0, pos); + pending.erase(0, pos + 1); + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.rfind("data: ", 0) != 0) { + continue; + } + std::string payload = line.substr(6); + if (payload == "[DONE]") { + continue; + } + on_data(payload); + } + return true; + }; + + httplib::Headers headers = {{"Accept", "text/event-stream"}}; + auto res = cli.Post(join_path(parts, path), headers, body, "application/json", receiver); + + if (!res) { + if (res.error() == httplib::Error::Canceled && should_stop()) { + return ""; // cancelled by the user + } + return "failed to connect to " + server_base + ": " + httplib::to_string(res.error()); + } + if (res->status < 200 || res->status >= 300) { + if (!raw_body.empty()) { + return raw_body; + } + return "request failed with status " + std::to_string(res->status); + } + return ""; +} + +bool cli_client::wait_health(const std::function & is_aborted) { + int connect_attempts = 0; + while (!is_aborted()) { + auto [cli, parts] = common_http_client(server_base); + cli.set_connection_timeout(1, 0); + auto res = cli.Get(join_path(parts, "/health")); + if (res) { + if (res->status == 200) { + return true; + } + // any other status means the server is up but not ready yet + // (e.g. 503 while the model is still loading) + } else if (++connect_attempts >= 10) { + last_error = "failed to connect to " + server_base + ": " + httplib::to_string(res.error()); + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + } + last_error = "aborted while waiting for the server to become ready"; + return false; +} diff --git a/tools/cli/cli-client.h b/tools/cli/cli-client.h new file mode 100644 index 000000000..9493b4fe6 --- /dev/null +++ b/tools/cli/cli-client.h @@ -0,0 +1,33 @@ +#pragma once + +#include +#include + +// openai-like client for CLI +struct cli_client { + std::string server_base; // base url, for example "http://127.0.0.1:8080" + std::string last_error; // set when wait_health() fails + + std::string model; // optional, set when the server has multiple models (router mode) + + // simple GET request, returns the raw response body + // throws std::runtime_error on transport error or non-2xx status + std::string get(const std::string & path); + + // simple POST request, returns the raw response body + // throws std::runtime_error on transport error or non-2xx status + std::string post(const std::string & path, const std::string & body); + + // POST request with an SSE streaming response + // on_data is invoked per "data:" event with the raw event payload + // returns after the stream is finished (empty string on graceful exit) + // otherwise, the raw error response body + std::string post_sse(const std::string & path, + const std::string & body, + const std::function & should_stop, + const std::function & on_data); + + // poll /health until the server is ready to accept requests + // returns false if is_aborted returned true or the server is unreachable + bool wait_health(const std::function & is_aborted); +}; diff --git a/tools/cli/cli-context.cpp b/tools/cli/cli-context.cpp new file mode 100644 index 000000000..74d60eb19 --- /dev/null +++ b/tools/cli/cli-context.cpp @@ -0,0 +1,664 @@ +#include "cli-context.h" +#include "cli-ui.h" + +#include "arg.h" +#include "base64.hpp" +#include "log.h" +#include "console.h" + +#define JSON_ASSERT GGML_ASSERT +#include + +#include +#include +#include +#include +#include +#include + +using json = nlohmann::ordered_json; + +struct cli_context_impl { + json messages = json::array(); + json pending_media = json::array(); // staged multimodal content parts +}; + +cli_context::cli_context(const common_params & params) : params(params), impl(new cli_context_impl()) {} + +cli_context::~cli_context() { + shutdown(); +} + +std::atomic & cli_context::interrupted() { + static std::atomic flag = false; + return flag; +} + +static bool should_stop() { + return cli_context::interrupted().load(); +} + +static constexpr size_t FILE_GLOB_MAX_RESULTS = 100; + +const char * LLAMA_ASCII_LOGO = R"( +▄▄ ▄▄ +██ ██ +██ ██ ▀▀█▄ ███▄███▄ ▀▀█▄ ▄████ ████▄ ████▄ +██ ██ ▄█▀██ ██ ██ ██ ▄█▀██ ██ ██ ██ ██ ██ +██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀ + ██ ██ + ▀▀ ▀▀ +)"; + +// number of values an arg consumes on the command line +static int arg_num_values(const common_arg & opt) { + if (opt.value_hint_2 != nullptr) { + return 2; + } + if (opt.value_hint != nullptr) { + return 1; + } + return 0; +} + +static std::string format_error_message(const json & err) { + if (err.contains("error") && err.at("error").is_object()) { + const auto & e = err.at("error"); + if (e.contains("message") && e.at("message").is_string()) { + return e.at("message").get(); + } + } + return err.dump(); +} + +// err is the raw response body of a failed request; it may or may not be JSON +static std::string format_error_message(const std::string & err) { + json parsed = json::parse(err, nullptr, false); + if (!parsed.is_discarded()) { + return format_error_message(parsed); + } + return err; +} + +static std::string media_type_from_ext(const std::string & fname) { + std::string ext = std::filesystem::path(fname).extension().string(); + std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); }); + if (ext == ".wav" || ext == ".mp3") { + return "audio"; + } + if (ext == ".mp4" || ext == ".avi" || ext == ".mkv" || ext == ".mov" || ext == ".webm") { + return "video"; + } + return "image"; +} + +bool cli_context::init() { + ui::init(params); + + std::optional spinner; + + bool use_external_server = !params.server_base.empty(); + if (use_external_server) { + std::string base = params.server_base; + while (!base.empty() && base.back() == '/') { + base.pop_back(); + } + client.server_base = base; + + spinner.emplace("Connecting to server at " + base); + } else { + if (params.model.path.empty() && params.model.url.empty() && + params.model.hf_repo.empty() && params.model.docker_repo.empty()) { + ui::show_error( + "no model specified", + "use -m or -hf to run a local model,\n" + "or --server-base to connect to a running llama-server" + ); + return false; + } + + spinner.emplace("\n\nLoading model..."); + + server.emplace(); + if (!server->start(params)) { + ui::show_error("server start failed"); + return false; + } + if (!server->wait_ready(should_stop)) { + if (!should_stop()) { + ui::show_error("the server exited before becoming ready"); + } + return false; + } + client.server_base = server->address(); + } + + // for --server-base this is the main availability check; for a spawned + // server it is a cheap sanity check on top of the ready signal + auto is_aborted = [this]() { + return should_stop() || (server && !server->alive()); + }; + bool healthy = false; + try { + healthy = client.wait_health(is_aborted); + } catch (const std::exception & e) { + client.last_error = e.what(); + } + if (!healthy) { + if (!should_stop()) { + ui::show_error(client.last_error); + } + return false; + } + + if (use_external_server) { + spinner.reset(); + if (!list_and_ask_models()) { + return false; + } + // restore the spinner for the next step + spinner.emplace("Waiting for server..."); + } + + fetch_server_props(); + + if (!params.out_file.empty()) { + output_file.emplace(params.out_file); + if (!output_file->is_open()) { + ui::show_error(string_format("failed to open output file '%s'", params.out_file.c_str())); + return false; + } + } + + return true; +} + +void cli_context::fetch_server_props() { + try { + json props = json::parse(client.get("/props")); + model_name = props.value("model_alias", ""); + if (model_name.empty()) { + const std::string path = props.value("model_path", ""); + if (!path.empty()) { + model_name = std::filesystem::path(path).filename().string(); + } + } + model_ftype = props.value("model_ftype", ""); + build_info = props.value("build_info", ""); + if (props.contains("modalities") && props.at("modalities").is_object()) { + const auto & modalities = props.at("modalities"); + has_vision = modalities.value("vision", false); + has_audio = modalities.value("audio", false); + has_video = modalities.value("video", false); + } + } catch (const std::exception & e) { + // /props can be disabled on remote servers; not fatal + LOG_DBG("failed to fetch /props: %s\n", e.what()); + } +} + +bool cli_context::list_and_ask_models() { + json resp = json::parse(client.get("/v1/models")); + if (!resp.contains("data") || !resp.at("data").is_array()) { + throw std::runtime_error("invalid response from /v1/models"); + } + std::vector models; + std::vector models_display; + for (const auto & m : resp.at("data")) { + if (!m.contains("id") || !m.at("id").is_string()) { + continue; + } + std::string name = m.at("id").get(); + std::string display = name; + if (m.contains("aliases") && m.at("aliases").is_array()) { + std::vector aliases; + for (const auto & a : m.at("aliases")) { + if (a.is_string()) { + aliases.push_back(a.get()); + } + } + if (!aliases.empty()) { + display += " (" + string_join(aliases, ", ") + ")"; + } + } + models.push_back(name); + models_display.push_back(display); + } + + // only one model: use it without asking + if (models.size() == 1) { + model_name = models[0]; + client.model = model_name; + return true; + } + + std::string message = "\nAvailable models:"; + for (size_t i = 0; i < models_display.size(); ++i) { + message += "\n " + std::to_string(i + 1) + ". " + models_display[i]; + } + message += "\n"; + ui::show_message(message); + std::string selection; + while (selection.empty()) { + if (should_stop()) { + return false; + } + ui::user_turn user_turn; + selection = user_turn.read_input(false, "Select model by number: "); + if (selection.empty()) { + continue; + } + try { + size_t idx = std::stoul(selection); + if (idx > 0 && idx <= models.size()) { + model_name = models[idx - 1]; + client.model = model_name; + ui::show_message("Selected model: " + model_name); + break; + } + } catch (...) { + // ignore + } + ui::show_error("Invalid selection. Please enter a valid number."); + selection.clear(); + continue; + } + return true; +} + +void cli_context::add_system_prompt() { + if (!params.system_prompt.empty()) { + impl->messages.push_back({ + {"role", "system"}, + {"content", params.system_prompt} + }); + } +} + +void cli_context::push_user_message(const std::string & text) { + json content; + if (impl->pending_media.empty()) { + content = text; + } else { + // multimodal message: media parts first, then the text + content = impl->pending_media; + content.push_back({ + {"type", "text"}, + {"text", text} + }); + impl->pending_media = json::array(); + } + impl->messages.push_back({ + {"role", "user"}, + {"content", content} + }); +} + +bool cli_context::stage_media_file(const std::string & fname, const std::string & type) { + std::ifstream file(fname, std::ios::binary); + if (!file) { + return false; + } + std::string data((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + std::string encoded = base64::encode(data); + + if (type == "audio") { + std::string ext = std::filesystem::path(fname).extension().string(); + std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); }); + impl->pending_media.push_back({ + {"type", "input_audio"}, + {"input_audio", { + {"data", encoded}, + {"format", ext == ".mp3" ? "mp3" : "wav"} + }} + }); + } else if (type == "video") { + impl->pending_media.push_back({ + {"type", "input_video"}, + {"input_video", { + {"data", encoded} + }} + }); + } else { + // the server detects the actual image type from the data + impl->pending_media.push_back({ + {"type", "image_url"}, + {"image_url", { + {"url", "data:image/unknown;base64," + encoded} + }} + }); + } + return true; +} + +void cli_context::write_output_file(const std::string & content) { + if (output_file) { + (*output_file) << content; + output_file->flush(); + } +} + +bool cli_context::generate_completion(generated_content & content_out, cli_timings & timings) { + json body = { + {"messages", impl->messages}, + {"stream", true}, + // in order to get timings even when we cancel mid-way + {"timings_per_token", true}, + }; + if (!client.model.empty()) { + body["model"] = client.model; + } + + bool stream_error = false; + + ui::assistant_turn a; + + std::string err = client.post_sse("/v1/chat/completions", body.dump(), should_stop, [&](const std::string & payload) { + json chunk = json::parse(payload, nullptr, false); + if (chunk.is_discarded()) { + return; + } + if (chunk.contains("error")) { + stream_error = true; + ui::show_error(format_error_message(chunk)); + return; + } + if (chunk.contains("timings")) { + const auto & t = chunk.at("timings"); + timings.prompt_per_second = t.value("prompt_per_second", 0.0); + timings.predicted_per_second = t.value("predicted_per_second", 0.0); + } + if (!chunk.contains("choices") || !chunk.at("choices").is_array() || chunk.at("choices").empty()) { + return; + } + const auto & choice = chunk.at("choices").at(0); + if (!choice.contains("delta")) { + return; + } + const auto & delta = choice.at("delta"); + if (delta.contains("reasoning_content") && delta.at("reasoning_content").is_string()) { + const std::string text = delta.at("reasoning_content").get(); + if (!text.empty()) { + content_out.reasoning += text; + a.push(ui::ASSISTANT_DISPLAY_MODE_REASONING, text); + } + } + if (delta.contains("content") && delta.at("content").is_string()) { + const std::string text = delta.at("content").get(); + if (!text.empty()) { + content_out.content += text; + a.push(ui::ASSISTANT_DISPLAY_MODE_CONTENT, text); + } + } + }); + + cli_context::interrupted().store(false); + + if (!err.empty()) { + ui::show_error(format_error_message(err)); + return false; + } + return !stream_error; +} + +int cli_context::run() { + add_system_prompt(); + + std::string modalities = "text"; + if (has_vision) { + modalities += ", vision"; + } + if (has_audio) { + modalities += ", audio"; + } + if (has_video) { + modalities += ", video"; + } + + std::string banner; + banner += "\n"; + banner += LLAMA_ASCII_LOGO; + banner += "\n"; + banner += "build : " + build_info + "\n"; + banner += "model : " + model_name + "\n"; + if (!model_ftype.empty()) { + banner += "ftype : " + model_ftype + "\n"; + } + banner += "modalities : " + modalities + "\n"; + if (!params.system_prompt.empty()) { + banner += "using custom system prompt\n"; + } + banner += "\n"; + banner += "available commands:\n"; + banner += " /exit or Ctrl+C stop or exit\n"; + banner += " /regen regenerate the last response\n"; + banner += " /clear clear the chat history\n"; + banner += " /read add a text file\n"; + banner += " /glob add text files using globbing pattern\n"; + if (has_vision) { + banner += " /image add an image file\n"; + } + if (has_audio) { + banner += " /audio add an audio file\n"; + } + if (has_video) { + banner += " /video add a video file\n"; + } + banner += "\n"; + + ui::show_message(banner); + + // interactive loop + std::string cur_msg; + + auto add_text_file = [&](const std::string & fname) -> bool { + std::ifstream file(fname, std::ios::binary); + if (!file) { + ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str())); + return false; + } + std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + cur_msg += "--- File: "; + cur_msg += fname; + cur_msg += " ---\n"; + cur_msg += content; + ui::show_message(string_format("Loaded text from '%s'", fname.c_str())); + return true; + }; + + while (true) { + std::string buffer; + { + ui::user_turn user_turn; + + if (params.prompt.empty()) { + buffer = user_turn.read_input(params.multiline_input); + } else { + // process input prompt from args + for (auto & fname : params.image) { + if (!stage_media_file(fname, media_type_from_ext(fname))) { + ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str())); + break; + } + ui::show_message(string_format("Loaded media from '%s'", fname.c_str())); + } + buffer = params.prompt; + user_turn.echo(buffer); + params.prompt.clear(); // only use it once + } + } + + if (should_stop()) { + cli_context::interrupted().store(false); + break; + } + + // remove trailing newline + if (!buffer.empty() && buffer.back() == '\n') { + buffer.pop_back(); + } + + // skip empty messages + if (buffer.empty()) { + continue; + } + + bool add_user_msg = true; + + // process commands + if (string_starts_with(buffer, "/exit")) { + break; + } else if (string_starts_with(buffer, "/regen")) { + if (impl->messages.size() >= 2) { + size_t last_idx = impl->messages.size() - 1; + impl->messages.erase(last_idx); + add_user_msg = false; + } else { + ui::show_error("No message to regenerate."); + continue; + } + } else if (string_starts_with(buffer, "/clear")) { + impl->messages.clear(); + add_system_prompt(); + + impl->pending_media = json::array(); + ui::show_message("Chat history cleared."); + continue; + } else if ( + (string_starts_with(buffer, "/image ") && has_vision) || + (string_starts_with(buffer, "/audio ") && has_audio) || + (string_starts_with(buffer, "/video ") && has_video)) { + std::string type = buffer.substr(1, 5); + // just in case (bad copy-paste for example), we strip all trailing/leading spaces + std::string fname = string_strip(buffer.substr(7)); + if (!stage_media_file(fname, type)) { + ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str())); + continue; + } + ui::show_message(string_format("Loaded media from '%s'", fname.c_str())); + write_output_file(string_format("User: Added media: %s\n", fname.c_str())); + continue; + } else if (string_starts_with(buffer, "/read ")) { + std::string fname = string_strip(buffer.substr(6)); + add_text_file(fname); + write_output_file(string_format("User: Added text file: %s\n", fname.c_str())); + continue; + } else if (string_starts_with(buffer, "/glob ")) { + std::error_code ec; + size_t count = 0; + auto curdir = std::filesystem::current_path(); + std::string pattern = string_strip(buffer.substr(6)); + std::filesystem::path rel_path; + + auto startglob = pattern.find_first_of("![*?"); + if (startglob != std::string::npos && startglob != 0) { + auto endpath = pattern.substr(0, startglob).find_last_of('/'); + if (endpath != std::string::npos) { + std::string rel_pattern = pattern.substr(0, endpath); +#if !defined(_WIN32) + if (string_starts_with(rel_pattern, '~')) { + const char * home = std::getenv("HOME"); + if (home && home[0]) { + rel_pattern = home + rel_pattern.substr(1); + } + } +#endif + rel_path = rel_pattern; + pattern.erase(0, endpath + 1); + curdir /= rel_path; + } + } + + for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir, + std::filesystem::directory_options::skip_permission_denied, ec)) { + if (!entry.is_regular_file()) { + continue; + } + + std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string(); + if (ec) { + ec.clear(); + continue; + } + std::replace(rel.begin(), rel.end(), '\\', '/'); + + if (!glob_match(pattern, rel)) { + continue; + } + + const std::string full_path = (curdir / rel).string(); + if (!add_text_file(full_path)) { + continue; + } + write_output_file(string_format("User: Added text file: %s\n", full_path.c_str())); + + if (++count >= FILE_GLOB_MAX_RESULTS) { + ui::show_error(string_format("Maximum number of globbed files allowed (%zu) reached.", FILE_GLOB_MAX_RESULTS)); + break; + } + } + continue; + } else { + // not a command + cur_msg += buffer; + } + + // generate response + if (add_user_msg) { + push_user_message(cur_msg); + write_output_file(string_format("User:\n%s\n\n", cur_msg.c_str())); + cur_msg.clear(); + } + + cli_timings timings; + generated_content content; + generate_completion(content, timings); + + impl->messages.push_back({ + {"role", "assistant"}, + {"content", content.content} + }); + + if (output_file) { + std::string out_content = "Assistant:\n"; + if (!content.reasoning.empty()) { + out_content += "[Start thinking]\n\n"; + out_content += content.reasoning; + out_content += "[End thinking]\n\n"; + } + out_content += content.content; + if (!out_content.empty() && out_content.back() != '\n') { + out_content += "\n"; + } + out_content += "\n"; + write_output_file(out_content); + } + + if (params.show_timings) { + ui::show_info(string_format( + "\n[ Prompt: %.1f t/s | Generation: %.1f t/s ]", + timings.prompt_per_second, + timings.predicted_per_second + )); + } + + if (params.single_turn) { + break; + } + } + + ui::show_message("\n\nExiting..."); + + return 0; +} + +void cli_context::shutdown() { + if (server) { + server->stop(); + server.reset(); + } + if (output_file) { + output_file->close(); + output_file.reset(); + } +} diff --git a/tools/cli/cli-context.h b/tools/cli/cli-context.h new file mode 100644 index 000000000..15ce4efee --- /dev/null +++ b/tools/cli/cli-context.h @@ -0,0 +1,76 @@ +#pragma once + +#include "common.h" + +#include "cli-client.h" +#include "cli-server.h" + +#include +#include +#include +#include +#include + +struct cli_timings { + double prompt_per_second = 0.0; + double predicted_per_second = 0.0; +}; + +struct cli_context_impl; + +struct cli_context { + common_params params; + + cli_client client; // always initialized + std::optional server; // only set when no --server-base is given + + // properties of the connected server + // will be populated by fetch_server_props() + std::string model_name; + std::string model_ftype; + std::string build_info; + bool has_vision = false; + bool has_audio = false; + bool has_video = false; + + std::optional output_file; + + cli_context(const common_params & params); + ~cli_context(); + + // connect to --server-base or spawn a local llama-server child; + // argc/argv are needed to forward the server-relevant args to the child + bool init(); + + // run the interactive chat loop, returns the process exit code + int run(); + + // stop the local server child (if any) + void shutdown(); + + // set by the SIGINT handler; cleared once the interrupt has been handled + static std::atomic & interrupted(); + +private: + struct generated_content { + std::string reasoning; + std::string content; + }; + bool generate_completion(generated_content & content_out, cli_timings & timings); + void fetch_server_props(); + void add_system_prompt(); + void push_user_message(const std::string & text); + + // check if server have multiple models (router mode) + // if yes, list them then ask; do nothing otherwise + bool list_and_ask_models(); + + // read a file and stage it as a multimodal content part; type is one of + // "image", "audio", "video"; returns false if the file cannot be read + bool stage_media_file(const std::string & fname, const std::string & type); + + // no-op if output file is not set + void write_output_file(const std::string & content); + + std::unique_ptr impl; +}; diff --git a/tools/cli/cli-server.h b/tools/cli/cli-server.h new file mode 100644 index 000000000..7596efb01 --- /dev/null +++ b/tools/cli/cli-server.h @@ -0,0 +1,89 @@ +#pragma once + +#include + +#include "http.h" + +// llama_server will be available as a dynamic library symbol +int llama_server(common_params & params, int argc, char ** argv); +void llama_server_terminate(); + +struct cli_server { + std::thread th; + int port = -1; + std::atomic is_alive = false; + std::atomic is_stopping = false; + + ~cli_server() { + stop(); + } + + void stop() { + if (is_stopping.exchange(true)) { + return; + } + if (alive()) { + llama_server_terminate(); + } + if (th.joinable()) { + th.join(); + } + } + + // spawn llama-server in a thread and interact with it via a random port + bool start(common_params & params) { + port = common_http_get_free_port(); + if (port <= 0) { + fprintf(stderr, "failed to get a free port\n"); + exit(1); + } + + is_alive.store(true, std::memory_order_release); + + common_params server_params = params; // copy + server_params.port = port; + + th = std::thread([this, server_params]() mutable { + // argc / argv are only used in router mode, we can skip them for now + int res = llama_server(server_params, 0, nullptr); + if (res != 0) { + fprintf(stderr, "llama_server exited with code %d\n", res); + } + is_alive.store(false, std::memory_order_release); + }); + + return true; + } + + std::string address() const { + return "http://127.0.0.1:" + std::to_string(port); + } + + bool wait_ready(std::function should_stop) { + if (!alive()) { + return false; + } + while (!should_stop()) { + auto [cli, parts] = common_http_client(address()); + cli.set_connection_timeout(1, 0); + auto res = cli.Get("/health"); + if (res) { + if (res->status == 200) { + return true; + } + // any other status means the server is up but not ready yet + // (e.g. 503 while the model is still loading) + } + if (!alive()) { + // in case server die permanently + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + return true; + } + + bool alive() const { + return is_alive.load(std::memory_order_acquire); + } +}; diff --git a/tools/cli/cli-ui.h b/tools/cli/cli-ui.h new file mode 100644 index 000000000..43aaba6f8 --- /dev/null +++ b/tools/cli/cli-ui.h @@ -0,0 +1,251 @@ +#pragma once + +#include "common.h" +#include "console.h" + +#include +#include +#include +#include +#include + +// TODO?: Make this reusable, enums, docs +static const std::array cmds = { + "/audio ", + "/clear", + "/exit", + "/glob ", + "/image ", + "/read ", + "/regen", + "/video ", +}; + +static std::vector> auto_completion_callback(std::string_view line, size_t cursor_byte_pos) { + std::vector> matches; + std::string cmd; + + if (line.length() > 1 && line.front() == '/' && !std::any_of(cmds.begin(), cmds.end(), [line](std::string_view prefix) { + return string_starts_with(line, prefix); + })) { + auto it = cmds.begin(); + + while ((it = std::find_if(it, cmds.end(), [line](std::string_view cmd_line) { + return string_starts_with(cmd_line, line); + })) != cmds.end()) { + matches.emplace_back(*it, it->length()); + ++it; + } + } else { + auto it = std::find_if(cmds.begin(), cmds.end(), [line](std::string_view prefix) { + return prefix.back() == ' ' && string_starts_with(line, prefix); + }); + + if (it != cmds.end()) { + cmd = *it; + } + } + + if (!cmd.empty() && cmd != "/glob " && line.length() >= cmd.length() && cursor_byte_pos >= cmd.length()) { + const std::string path_prefix = std::string(line.substr(cmd.length(), cursor_byte_pos - cmd.length())); + const std::string path_postfix = std::string(line.substr(cursor_byte_pos)); + auto cur_dir = std::filesystem::current_path(); + std::string cur_dir_str = cur_dir.string(); + std::string expanded_prefix = path_prefix; + +#if !defined(_WIN32) + if (string_starts_with(path_prefix, '~')) { + const char * home = std::getenv("HOME"); + if (home && home[0]) { + expanded_prefix = home + path_prefix.substr(1); + } + } + if (string_starts_with(expanded_prefix, '/')) { +#else + if (std::isalpha(static_cast(expanded_prefix[0])) && expanded_prefix.find(':') == 1) { +#endif + cur_dir = std::filesystem::path(expanded_prefix).parent_path(); + cur_dir_str.clear(); + } else if (!path_prefix.empty()) { + cur_dir /= std::filesystem::path(path_prefix).parent_path(); + } + + std::error_code ec; + for (const auto & entry : std::filesystem::directory_iterator(cur_dir, ec)) { + if (ec) { + break; + } + if (!entry.exists(ec)) { + ec.clear(); + continue; + } + + const std::string path_full = entry.path().string(); + std::string path_entry = !cur_dir_str.empty() && string_starts_with(path_full, cur_dir_str) ? path_full.substr(cur_dir_str.length() + 1) : path_full; + + if (entry.is_directory(ec)) { + path_entry.push_back(std::filesystem::path::preferred_separator); + } + + if (expanded_prefix.empty() || string_starts_with(path_entry, expanded_prefix)) { + const std::string updated_line = cmd + path_entry; + matches.emplace_back(updated_line + path_postfix, updated_line.length()); + } + + if (ec) { + ec.clear(); + } + } + + if (matches.empty()) { + const std::string updated_line = cmd + path_prefix; + matches.emplace_back(updated_line + path_postfix, updated_line.length()); + } + + // Add the longest common prefix + if (!expanded_prefix.empty() && matches.size() > 1) { + const std::string_view match0(matches[0].first); + const std::string_view match1(matches[1].first); + auto it = std::mismatch(match0.begin(), match0.end(), match1.begin(), match1.end()); + size_t len = it.first - match0.begin(); + + for (size_t i = 2; i < matches.size(); ++i) { + const std::string_view matchi(matches[i].first); + auto cmp = std::mismatch(match0.begin(), match0.end(), matchi.begin(), matchi.end()); + len = std::min(len, static_cast(cmp.first - match0.begin())); + } + + const std::string updated_line = std::string(match0.substr(0, len)); + matches.emplace_back(updated_line + path_postfix, updated_line.length()); + } + + std::sort(matches.begin(), matches.end(), [](const auto & a, const auto & b) { + return a.first.compare(0, a.second, b.first, 0, b.second) < 0; + }); + } + + return matches; +} + +// note: make this view implementation generic, so that we can move to TUI in the future if we want to +namespace ui { + static void init(const common_params & params) { + // TODO: avoid using atexit() here by making `console` a singleton + console::init(params.simple_io, params.use_color); + atexit([]() { console::cleanup(); }); + + console::set_completion_callback(auto_completion_callback); + } + + struct spinner { + spinner(const std::string & message) { + if (!message.empty()) { + console::log("%s ", message.c_str()); + } + console::spinner::start(); + } + ~spinner() { + console::spinner::stop(); + } + }; + + struct user_turn { + user_turn() { + console::set_display(DISPLAY_TYPE_USER_INPUT); + } + ~user_turn() { + console::set_display(DISPLAY_TYPE_RESET); + } + void echo(const std::string & buffer) { + if (buffer.size() > 500) { + console::log("\n> %s ... (truncated)\n", buffer.substr(0, 500).c_str()); + } else { + console::log("\n> %s\n", buffer.c_str()); + } + } + std::string read_input(bool multiline_input, const char * prompt = nullptr) { + if (prompt) { + console::log("%s", prompt); + } else { + console::log("\n> "); + } + std::string buffer; + std::string line; + bool another_line = true; + do { + another_line = console::readline(line, multiline_input); + buffer += line; + } while (another_line); + return buffer; + } + }; + + enum assistant_display_mode { + ASSISTANT_DISPLAY_MODE_REASONING, + ASSISTANT_DISPLAY_MODE_CONTENT, + }; + struct assistant_turn { + assistant_display_mode mode = ASSISTANT_DISPLAY_MODE_CONTENT; + bool trailing_newline = true; + bool is_inside_reasoning = false; + assistant_turn() { + console::set_display(DISPLAY_TYPE_RESET); + } + ~assistant_turn() { + console::set_display(DISPLAY_TYPE_RESET); + add_newline_if_needed(); + } + void push(assistant_display_mode m, const std::string & buffer) { + if (m != mode) { + add_newline_if_needed(); + switch (m) { + case ASSISTANT_DISPLAY_MODE_CONTENT: + { + if (is_inside_reasoning) { + console::log("[End thinking]\n\n"); + is_inside_reasoning = false; + } + console::set_display(DISPLAY_TYPE_RESET); + } break; + case ASSISTANT_DISPLAY_MODE_REASONING: + { + console::set_display(DISPLAY_TYPE_REASONING); + is_inside_reasoning = true; + console::log("\n[Start thinking]\n\n"); + } break; + } + } + mode = m; + if (buffer.empty()) { + return; + } + trailing_newline = buffer.back() == '\n'; + console::log("%s", buffer.c_str()); + console::flush(); + } + void add_newline_if_needed() { + if (!trailing_newline) { + console::log("\n"); + console::flush(); + } + } + }; + + static void show_error(const std::string & title, const std::string & message = "") { + console::spinner::stop(); + console::error("Error: %s\n", title.c_str()); + if (!message.empty()) { + console::log("%s\n", message.c_str()); + } + } + + static void show_message(const std::string & message) { + console::log("%s\n", message.c_str()); + } + + static void show_info(const std::string & message) { + console::set_display(DISPLAY_TYPE_INFO); + console::log("%s\n", message.c_str()); + console::set_display(DISPLAY_TYPE_RESET); + } +} diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index aa5d0a2ab..98d0cca1c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1444,6 +1444,7 @@ private: // populate chat template params { common_chat_templates_ptr chat_templates; + bool enable_thinking = false; try { chat_templates = common_chat_templates_init(model_tgt, params_base.chat_template); @@ -1451,6 +1452,12 @@ private: SRV_TRC("%s: chat template, example_format: '%s'\n", __func__, common_chat_format_example(chat_templates.get(), params_base.use_jinja, params_base.default_template_kwargs).c_str()); + // thinking is enabled if: + // 1. It's not explicitly disabled via --reasoning off + // 2. The chat template supports it + const bool template_supports_thinking = params_base.use_jinja && common_chat_templates_support_enable_thinking(chat_templates.get()); + enable_thinking = params_base.enable_reasoning != 0 && template_supports_thinking; + SRV_TRC("%s: chat template, thinking = %d\n", __func__, enable_thinking); } catch (const std::exception & e) { SRV_ERR("%s: chat template parsing error: %s\n", __func__, e.what()); SRV_ERR("%s: please consider disabling jinja via --no-jinja, or use a custom chat template via --chat-template\n", __func__); @@ -1458,13 +1465,6 @@ private: return false; } - // thinking is enabled if: - // 1. It's not explicitly disabled via --reasoning off - // 2. The chat template supports it - const bool template_supports_thinking = params_base.use_jinja && common_chat_templates_support_enable_thinking(chat_templates.get()); - const bool enable_thinking = params_base.enable_reasoning != 0 && template_supports_thinking; - SRV_TRC("%s: chat template, thinking = %d\n", __func__, enable_thinking); - // IMPORTANT: chat_params is reused across sleeping / resuming states, // never store llama_context/llama_model pointers in chat_params, // as they may be invalidated after sleeping @@ -3435,9 +3435,14 @@ private: slot.n_prompt_tokens_processed++; - // stop the prompt batch exactly before a user message - if (spans.is_user_start(slot.prompt.n_tokens())) { - break; + // break at the last user message, or at user messages at least min step past the last checkpoint + if (do_checkpoint && spans.is_user_start(slot.prompt.n_tokens())) { + const auto pos = slot.prompt.n_tokens(); + const auto & checkpoints = slot.prompt.checkpoints; + + if (pos == last_user_pos || checkpoints.empty() || pos > checkpoints.back().n_tokens + params_base.checkpoint_min_step) { + break; + } } // process the last few tokens of the prompt separately in order to allow for a checkpoint to be created. @@ -4521,6 +4526,7 @@ void server_routes::init_routes() { { "default_generation_settings", default_generation_settings_for_props }, { "total_slots", params.n_parallel }, { "model_alias", meta->model_name }, + { "model_ftype", meta->model_ftype }, { "model_path", meta->model_path }, { "modalities", json { {"vision", meta->has_inp_image}, diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 0cbf520af..d1fdc0607 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -7,6 +7,7 @@ #include "build-info.h" #include "preset.h" #include "download.h" +#include "http.h" #include // TODO: remove this once we use HTTP client from download.h #include @@ -28,14 +29,7 @@ #include #include -#ifdef _WIN32 -#include -#include -#else -#include -#include -#include -#include +#ifndef _WIN32 extern char **environ; #endif @@ -716,66 +710,6 @@ std::optional server_models::get_meta(const std::string & nam return std::nullopt; } -static int get_free_port() { -#ifdef _WIN32 - WSADATA wsaData; - if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) { - return -1; - } - typedef SOCKET native_socket_t; -#define INVALID_SOCKET_VAL INVALID_SOCKET -#define CLOSE_SOCKET(s) closesocket(s) -#else - typedef int native_socket_t; -#define INVALID_SOCKET_VAL -1 -#define CLOSE_SOCKET(s) close(s) -#endif - - native_socket_t sock = socket(AF_INET, SOCK_STREAM, 0); - if (sock == INVALID_SOCKET_VAL) { -#ifdef _WIN32 - WSACleanup(); -#endif - return -1; - } - - struct sockaddr_in serv_addr; - std::memset(&serv_addr, 0, sizeof(serv_addr)); - serv_addr.sin_family = AF_INET; - serv_addr.sin_addr.s_addr = htonl(INADDR_ANY); - serv_addr.sin_port = htons(0); - - if (bind(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0) { - CLOSE_SOCKET(sock); -#ifdef _WIN32 - WSACleanup(); -#endif - return -1; - } - -#ifdef _WIN32 - int namelen = sizeof(serv_addr); -#else - socklen_t namelen = sizeof(serv_addr); -#endif - if (getsockname(sock, (struct sockaddr*)&serv_addr, &namelen) != 0) { - CLOSE_SOCKET(sock); -#ifdef _WIN32 - WSACleanup(); -#endif - return -1; - } - - int port = ntohs(serv_addr.sin_port); - - CLOSE_SOCKET(sock); -#ifdef _WIN32 - WSACleanup(); -#endif - - return port; -} - // helper to convert vector to char ** // pointers are only valid as long as the original vector is valid static std::vector to_char_ptr_array(const std::vector & vec) { @@ -879,7 +813,7 @@ void server_models::load(const std::string & name, const load_options & opts) { // prepare new instance info instance_t inst; inst.meta = meta; - inst.meta.port = get_free_port(); + inst.meta.port = common_http_get_free_port(); inst.meta.status = SERVER_MODEL_STATUS_LOADING; inst.meta.loaded_info = json{}; inst.meta.last_used = ggml_time_ms(); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 9e8603be6..c2b21120a 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -36,6 +36,19 @@ static inline void signal_handler(int signal) { shutdown_handler(signal); } +// satisfies -Wmissing-declarations (used by llama command) +int llama_server(int argc, char ** argv); + +// to be used via CLI (argc / argv are used by router mode only) +int llama_server(common_params & params, int argc, char ** argv); +void llama_server_terminate(); +void llama_server_terminate() { + if (shutdown_handler) { + shutdown_handler(0); + } +} + + // wrapper function that handles exceptions and logs errors // this is to make sure handler_t never throws exceptions; instead, it returns an error response static server_http_context::handler_t ex_wrapper(server_http_context::handler_t func) { @@ -72,9 +85,6 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t }; } -// satisfies -Wmissing-declarations -int llama_server(int argc, char ** argv); - int llama_server(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -94,16 +104,26 @@ int llama_server(int argc, char ** argv) { llama_backend_init(); llama_numa_init(params.numa); + return llama_server(params, argc, argv); +} + +int llama_server(common_params & params, int argc, char ** argv) { + bool is_run_by_cli = (argv == nullptr); + common_models_handler models_handler; - try { - models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER); - if (common_models_handler_is_preset_repo(models_handler)) { - // apply the preset and start the server in router mode - common_models_handler_apply(models_handler, params); + + // note: router mode also accepts -hf remote-preset, so we need to check that first + if (!is_run_by_cli && !params.model.hf_repo.empty()) { + try { + models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER); + if (common_models_handler_is_preset_repo(models_handler)) { + // apply the preset and start the server in router mode + common_models_handler_apply(models_handler, params); + } + } catch (const std::exception & e) { + SRV_ERR("failed to fetch model metadata: %s\n", e.what()); + return 1; } - } catch (const std::exception & e) { - SRV_ERR("failed to fetch model metadata: %s\n", e.what()); - return 1; } // router server never loads a model and must not touch the GPU @@ -321,8 +341,9 @@ int llama_server(int argc, char ** argv) { if (child.is_child() && child.get_mode() == SERVER_CHILD_MODE_DOWNLOAD) { return child.run_download(params); - } else if (!is_router_server) { + } else if (!is_router_server && !is_run_by_cli) { // single-model mode (NOT spawned by router) + // if this is invoked by CLI, model downloading should be already handled try { common_models_handler_apply(models_handler, params); } catch (const std::exception & e) { @@ -411,20 +432,22 @@ int llama_server(int argc, char ** argv) { }; } - // TODO: refactor in common/console + // register signal handler if not running by CLI + if (!is_run_by_cli) { #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) - struct sigaction sigint_action; - sigint_action.sa_handler = signal_handler; - sigemptyset (&sigint_action.sa_mask); - sigint_action.sa_flags = 0; - sigaction(SIGINT, &sigint_action, NULL); - sigaction(SIGTERM, &sigint_action, NULL); + struct sigaction sigint_action; + sigint_action.sa_handler = signal_handler; + sigemptyset (&sigint_action.sa_mask); + sigint_action.sa_flags = 0; + sigaction(SIGINT, &sigint_action, NULL); + sigaction(SIGTERM, &sigint_action, NULL); #elif defined (_WIN32) - auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL { - return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false; - }; - SetConsoleCtrlHandler(reinterpret_cast(console_ctrl_handler), true); + auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL { + return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false; + }; + SetConsoleCtrlHandler(reinterpret_cast(console_ctrl_handler), true); #endif + } SRV_INF("listening on %s\n", ctx_http.listening_address.c_str());