hexagon: tiling, tracing and optimizations for unary ops (#25474)
* hexagon: tile wide rows in pointwise unary ops to avoid VTCM overflow * unary: reject permuted tensors for now (not used by models) * hex-unary: replace divs with fastdiv * hex-unary: add vtcm layout and host computed kernel params * hex-unary: move fastdiv init into kernel params * hex-unary: add specialized thread functions to improve generated code * hex-unary: tracing instrumentation for unary ops * hex-unary: factor out hvx kernels, streamline and remove more duplication * ggml-hexagon: fix std::min collision with Windows min macro * hex-cmake: make lto build happy --------- Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>
This commit is contained in:
Aparna M P
co-authored by
Max Krasnyansky
Max Krasnyansky
parent
82fce65d8b
commit
fb30ba9a6c
@@ -39,8 +39,8 @@ add_library(${HTP_LIB} SHARED
|
||||
diag-ops.c
|
||||
solve-tri-ops.c
|
||||
pad-ops.c
|
||||
flash-attn-ops.c
|
||||
matmul-ops.c
|
||||
flash-attn-ops.c
|
||||
)
|
||||
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE
|
||||
|
||||
@@ -120,7 +120,6 @@ int op_concat(struct htp_ops_context * octx);
|
||||
int op_diag(struct htp_ops_context * octx);
|
||||
int op_solve_tri(struct htp_ops_context * octx);
|
||||
int op_gated_delta_net(struct htp_ops_context * octx);
|
||||
int op_tri(struct htp_ops_context * octx);
|
||||
int op_pad(struct htp_ops_context * octx);
|
||||
|
||||
#endif /* HTP_CTX_H */
|
||||
|
||||
@@ -0,0 +1,257 @@
|
||||
#ifndef HVX_NORM_H
|
||||
#define HVX_NORM_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include "hvx-base.h"
|
||||
#include "hvx-reduce.h"
|
||||
#include "hvx-inverse.h"
|
||||
#include "hvx-sqrt.h"
|
||||
#include "hvx-repl.h"
|
||||
|
||||
static inline void hvx_fast_rms_norm_f32(const uint8_t * restrict src,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
|
||||
const HVX_Vector * restrict v_src = (HVX_Vector *) src;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32; // number of full vectors
|
||||
const int nloe = num_elems % VLEN_FP32; // leftover elements
|
||||
|
||||
// Compute sum of squares for full vectors
|
||||
HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Reduce HVX sum
|
||||
sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
|
||||
|
||||
HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
|
||||
HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
|
||||
HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
|
||||
HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
|
||||
|
||||
// Scale full vectors
|
||||
HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(v2);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
HVX_Vector result = Q6_Vsf_equals_Vqf32(v2);
|
||||
|
||||
// Store with masking to avoid overwriting memory beyond the tensor
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
|
||||
const uint8_t * restrict weight,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
const HVX_Vector * restrict v_src = (const HVX_Vector *) src;
|
||||
const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32; // number of full vectors
|
||||
const int nloe = num_elems % VLEN_FP32; // leftover elements
|
||||
|
||||
// Compute sum of squares for full vectors
|
||||
HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
|
||||
}
|
||||
|
||||
// Reduce HVX sum
|
||||
sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
|
||||
|
||||
HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
|
||||
HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
|
||||
HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
|
||||
HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
|
||||
|
||||
// Scale and multiply
|
||||
HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
|
||||
HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[i]);
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(result);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
|
||||
HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
|
||||
HVX_Vector result = Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[nvec]);
|
||||
HVX_Vector res_v = Q6_Vsf_equals_Vqf32(result);
|
||||
|
||||
// Store with masking to avoid overwriting memory beyond the tensor
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, res_v);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_fast_norm_f32(const uint8_t * restrict src,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
|
||||
const HVX_Vector * restrict v_src = (HVX_Vector *) src;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32; // number of full vectors
|
||||
const int nloe = num_elems % VLEN_FP32; // leftover elements
|
||||
|
||||
// Compute sum of squares and sum of values for full vectors
|
||||
HVX_Vector sum_sq_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector sum_x_v = Q6_V_vsplat_R(0x00000000);
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
|
||||
sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, v2);
|
||||
sum_x_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v, Q6_Vqf32_vadd_VsfVsf(v1, Q6_V_vzero()));
|
||||
}
|
||||
|
||||
// Reduce HVX sums
|
||||
sum_sq_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_sq_v));
|
||||
sum_x_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_x_v));
|
||||
|
||||
HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems);
|
||||
HVX_Vector denom_v = hvx_vec_inverse_f32(t_v);
|
||||
HVX_Vector mean_sq_v = Q6_Vqf32_vmpy_VsfVsf(sum_sq_v, denom_v);
|
||||
HVX_Vector mean_x_v = Q6_Vqf32_vmpy_VsfVsf(sum_x_v, denom_v);
|
||||
HVX_Vector mean_x_sq_v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(mean_x_v), Q6_Vsf_equals_Vqf32(mean_x_v));
|
||||
HVX_Vector var_v = Q6_Vqf32_vsub_Vqf32Vqf32(mean_sq_v, mean_x_sq_v);
|
||||
HVX_Vector var_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(var_v, epsilon_v);
|
||||
|
||||
// scale = rsqrt(variance + epsilon), mean_x broadcast for subtraction
|
||||
HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(var_epsilon_v));
|
||||
HVX_Vector mean_x_b = hvx_vec_repl_f32(Q6_Vsf_equals_Vqf32(mean_x_v));
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
|
||||
HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(v3);
|
||||
}
|
||||
|
||||
// Handle tail elements using vectorized ops with masking
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector v2 = Q6_Vqf32_vsub_VsfVsf(v1, mean_x_b);
|
||||
HVX_Vector v3 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v2), scale_v);
|
||||
HVX_Vector result = Q6_Vsf_equals_Vqf32(v3);
|
||||
|
||||
// Store with masking to avoid overwriting memory beyond the tensor
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_fast_l2_norm_f32(const uint8_t * restrict src,
|
||||
uint8_t * restrict dst,
|
||||
const int num_elems,
|
||||
float epsilon) {
|
||||
|
||||
const HVX_Vector * restrict v_src = (HVX_Vector *) src;
|
||||
HVX_Vector * restrict v_dst = (HVX_Vector *) dst;
|
||||
|
||||
HVX_Vector sum_v = hvx_vec_splat_f32(0.0f);
|
||||
|
||||
const int nvec = num_elems / VLEN_FP32;
|
||||
const int nloe = num_elems % VLEN_FP32;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
|
||||
}
|
||||
|
||||
// Include tail elements in the sum-of-squares using a predicate mask
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector sq = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
|
||||
sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, sq);
|
||||
}
|
||||
|
||||
// Compute scale = 1/fmax(sqrt(sum), epsilon) entirely in HVX registers.
|
||||
// hvx_vec_rsqrt_f32 + hvx_vec_inverse_f32 avoids scalar extraction.
|
||||
HVX_Vector sum_sf = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
|
||||
HVX_Vector rsqrt_v = hvx_vec_rsqrt_f32(sum_sf); // 1/sqrt(sum)
|
||||
HVX_Vector sqrt_v = hvx_vec_inverse_f32(rsqrt_v); // sqrt(sum)
|
||||
HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
|
||||
HVX_Vector denom_v = Q6_Vsf_vmax_VsfVsf(sqrt_v, epsilon_v); // fmax(sqrt(sum), epsilon)
|
||||
HVX_Vector scale_v = hvx_vec_inverse_f32(denom_v); // 1/fmax(sqrt(sum), epsilon)
|
||||
|
||||
#pragma unroll(4)
|
||||
for (int i = 0; i < nvec; i++) {
|
||||
HVX_Vector v1 = v_src[i];
|
||||
v_dst[i] = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
|
||||
}
|
||||
|
||||
if (nloe > 0) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
|
||||
HVX_Vector result = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v1, scale_v));
|
||||
hvx_vec_store_a(&v_dst[nvec], nloe * 4, result);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // HVX_NORM_H
|
||||
@@ -19,5 +19,6 @@
|
||||
#include "hvx-base.h"
|
||||
#include "hvx-pow.h"
|
||||
#include "hvx-log.h"
|
||||
#include "hvx-norm.h"
|
||||
|
||||
#endif /* HVX_UTILS_H */
|
||||
|
||||
@@ -667,7 +667,7 @@ static int execute_op(struct htp_ops_context * octx) {
|
||||
return op_gated_delta_net(octx);
|
||||
|
||||
case HTP_OP_TRI:
|
||||
return op_tri(octx);
|
||||
return op_unary(octx);
|
||||
|
||||
case HTP_OP_INVALID:
|
||||
break;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,162 @@
|
||||
#ifndef HTP_UNARY_OPS_H
|
||||
#define HTP_UNARY_OPS_H
|
||||
|
||||
#include "hex-common.h"
|
||||
#include "htp-ops.h"
|
||||
|
||||
// Op-specific struct for precomputed unary params
|
||||
struct htp_unary_kernel_params {
|
||||
uint32_t n_threads;
|
||||
uint32_t col_tile;
|
||||
uint32_t vtcm_row_per_thread;
|
||||
uint32_t block;
|
||||
uint32_t broadcast_weight;
|
||||
|
||||
uint32_t vtcm_src0_size_per_thread;
|
||||
uint32_t vtcm_src1_size_per_thread;
|
||||
uint32_t vtcm_dst_size_per_thread;
|
||||
|
||||
uint32_t vtcm_src0_size;
|
||||
uint32_t vtcm_src1_size;
|
||||
uint32_t vtcm_dst_size;
|
||||
|
||||
uint32_t src0_row_size_aligned;
|
||||
uint32_t src1_row_size_aligned;
|
||||
uint32_t dst_row_size_aligned;
|
||||
|
||||
uint32_t vtcm_size;
|
||||
|
||||
// Fastdiv helpers
|
||||
struct fastdiv_values div_ne01;
|
||||
struct fastdiv_values div_ne02;
|
||||
struct fastdiv_values div_ne012;
|
||||
struct fastdiv_values div_tpr;
|
||||
};
|
||||
|
||||
#if defined(__cplusplus)
|
||||
static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
|
||||
#else
|
||||
_Static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_params is too large for kernel_params blob");
|
||||
#endif
|
||||
|
||||
static inline bool htp_op_is_unary(uint32_t opcode) {
|
||||
switch (opcode) {
|
||||
case HTP_OP_NORM:
|
||||
case HTP_OP_RMS_NORM:
|
||||
case HTP_OP_RMS_NORM_MUL:
|
||||
case HTP_OP_SCALE:
|
||||
case HTP_OP_SQR:
|
||||
case HTP_OP_SQRT:
|
||||
case HTP_OP_UNARY_NEG:
|
||||
case HTP_OP_UNARY_EXP:
|
||||
case HTP_OP_UNARY_SIGMOID:
|
||||
case HTP_OP_UNARY_SOFTPLUS:
|
||||
case HTP_OP_UNARY_TANH:
|
||||
case HTP_OP_L2_NORM:
|
||||
case HTP_OP_TRI:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
struct htp_unary_vtcm_layout {
|
||||
size_t total_bytes;
|
||||
size_t off_src0;
|
||||
size_t off_src1;
|
||||
size_t off_dst;
|
||||
|
||||
size_t src0_bytes;
|
||||
size_t src1_bytes;
|
||||
size_t dst_bytes;
|
||||
};
|
||||
|
||||
static inline void htp_unary_vtcm_layout_build(
|
||||
struct htp_unary_vtcm_layout * L,
|
||||
uint32_t op,
|
||||
uint32_t ne00,
|
||||
uint32_t ne10,
|
||||
uint32_t ne11,
|
||||
bool broadcast_weight,
|
||||
uint32_t n_threads,
|
||||
size_t vtcm_size,
|
||||
uint32_t * out_col_tile,
|
||||
uint32_t * out_vtcm_row_per_thread
|
||||
) {
|
||||
const size_t src0_data_row_size = ne00 * sizeof(float);
|
||||
const size_t dst_data_row_size = ne10 * sizeof(float);
|
||||
|
||||
const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
|
||||
const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, 128);
|
||||
|
||||
size_t src1_row_size_aligned = 0;
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
const size_t src1_data_row_size = ne11 * sizeof(float);
|
||||
src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
|
||||
}
|
||||
|
||||
size_t vtcm_size_per_row = 0;
|
||||
size_t vtcm_row_per_thread = 0;
|
||||
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
if (broadcast_weight) {
|
||||
size_t available_vtcm = vtcm_size;
|
||||
size_t src1_vtcm_total = n_threads * src1_row_size_aligned;
|
||||
if (available_vtcm > src1_vtcm_total) {
|
||||
available_vtcm -= src1_vtcm_total;
|
||||
} else {
|
||||
available_vtcm = 0;
|
||||
}
|
||||
vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
|
||||
vtcm_row_per_thread = available_vtcm / (n_threads * vtcm_size_per_row);
|
||||
} else {
|
||||
vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned + src1_row_size_aligned);
|
||||
vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
|
||||
}
|
||||
} else {
|
||||
vtcm_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned);
|
||||
vtcm_row_per_thread = vtcm_size / (n_threads * vtcm_size_per_row);
|
||||
}
|
||||
|
||||
const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
|
||||
op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM);
|
||||
uint32_t col_tile = 0;
|
||||
|
||||
if (vtcm_row_per_thread == 0 && !is_reduction) {
|
||||
const size_t per_thread_budget = vtcm_size / n_threads;
|
||||
const size_t col_tile_bytes = hex_align_down(per_thread_budget / 4, 128);
|
||||
col_tile = (uint32_t) (col_tile_bytes / sizeof(float));
|
||||
|
||||
L->src0_bytes = col_tile_bytes * 2;
|
||||
L->dst_bytes = col_tile_bytes * 2;
|
||||
L->src1_bytes = 0;
|
||||
} else {
|
||||
L->src0_bytes = src0_row_size_aligned * vtcm_row_per_thread * 2;
|
||||
L->dst_bytes = dst_row_size_aligned * vtcm_row_per_thread * 2;
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
if (broadcast_weight) {
|
||||
L->src1_bytes = src1_row_size_aligned;
|
||||
} else {
|
||||
L->src1_bytes = src1_row_size_aligned * vtcm_row_per_thread * 2;
|
||||
}
|
||||
} else {
|
||||
L->src1_bytes = 0;
|
||||
}
|
||||
}
|
||||
|
||||
L->off_src0 = 0;
|
||||
if (op == HTP_OP_RMS_NORM_MUL) {
|
||||
L->off_src1 = L->off_src0 + L->src0_bytes * n_threads;
|
||||
L->off_dst = L->off_src1 + L->src1_bytes * n_threads;
|
||||
} else {
|
||||
L->off_src1 = 0;
|
||||
L->off_dst = L->off_src0 + L->src0_bytes * n_threads;
|
||||
}
|
||||
|
||||
L->total_bytes = L->off_dst + L->dst_bytes * n_threads;
|
||||
|
||||
*out_col_tile = col_tile;
|
||||
*out_vtcm_row_per_thread = vtcm_row_per_thread;
|
||||
}
|
||||
|
||||
#endif /* HTP_UNARY_OPS_H */
|
||||
Reference in New Issue
Block a user