hexagon: add support for CONCAT op (#23648)
* hexagon: add support for CONCAT with optimized concat_2d_transposed qwen3.5 models are quite heavy on the CONCAT with large and transposed src1. * hex-concat: use fastdiv in generic version * hex-concat: make checks for transposed a bit more readable * hex-concat: reoder dma ops for better pipelining * hex-cont/cpy: optimize CPY and CONT ops The primary change is to avoid scalar divs in the inner loops. We were calling hvx_copy_uu(... type_size) where type_size is non a constexpr. This causes runtime divs by that value which is normally just 4 or 2 (f32/f16). * hex-get-rows: optimize GET_ROWS for large rows We now use DMA for larger rows and also split them into chunks to improve perf for Qwen3.5 and other models that do lots of GET_ROWS with huge (2MB+ rows). Also bump the DMA queue depth now that we can take advantage of it. * hex-concat: unroll the inner loops of concat_2d * hex-concat: more updates to concat_2d to improve perf a bit further * hex-cpy: fixed n_rows per thread checks in the copy ops * hmx-fa: fix alignment issues while computing dma sizes * hex-set-rows: add early returns for idle threads * hvx-rope: minor optimization to replace loops with fastdiv logic * hex-rope: replace scalar tail processing with HVX * hex-rope: optimize rope cache init with HVX Add hvx-utils sin/cos helpers that use an aprox method (similar to rsqrt, inverse, etc) Use the helpers to optimize ROPE.
This commit is contained in:
Max Krasnyansky
+164
-144
@@ -28,158 +28,170 @@ struct htp_copy_context {
|
||||
uint32_t dst_blocks_per_row;
|
||||
|
||||
uint32_t src0_nrows_per_thread;
|
||||
|
||||
void (*copy)(struct htp_copy_context * ct, struct htp_ops_context * octx, int nth, int ith);
|
||||
};
|
||||
|
||||
#define cpy_preamble \
|
||||
const struct htp_tensor *src0 = octx->src[0]; \
|
||||
const struct htp_tensor *dst = octx->dst; \
|
||||
\
|
||||
const uint32_t ne00 = src0->ne[0]; \
|
||||
const uint32_t ne01 = src0->ne[1]; \
|
||||
const uint32_t ne02 = src0->ne[2]; \
|
||||
const uint32_t ne03 = src0->ne[3]; \
|
||||
\
|
||||
const uint32_t nb00 = src0->nb[0]; \
|
||||
const uint32_t nb01 = src0->nb[1]; \
|
||||
const uint32_t nb02 = src0->nb[2]; \
|
||||
const uint32_t nb03 = src0->nb[3]; \
|
||||
\
|
||||
const uint32_t ne0 = dst->ne[0]; \
|
||||
const uint32_t ne1 = dst->ne[1]; \
|
||||
const uint32_t ne2 = dst->ne[2]; \
|
||||
const uint32_t ne3 = dst->ne[3]; \
|
||||
\
|
||||
const uint32_t nb0 = dst->nb[0]; \
|
||||
const uint32_t nb1 = dst->nb[1]; \
|
||||
const uint32_t nb2 = dst->nb[2]; \
|
||||
const uint32_t nb3 = dst->nb[3]; \
|
||||
\
|
||||
const uint32_t ne00 = src0->ne[0]; \
|
||||
const uint32_t ne01 = src0->ne[1]; \
|
||||
const uint32_t ne02 = src0->ne[2]; \
|
||||
const uint32_t ne03 = src0->ne[3]; \
|
||||
\
|
||||
const uint32_t nb00 = src0->nb[0]; \
|
||||
const uint32_t nb01 = src0->nb[1]; \
|
||||
const uint32_t nb02 = src0->nb[2]; \
|
||||
const uint32_t nb03 = src0->nb[3]; \
|
||||
\
|
||||
const uint32_t ne0 = dst->ne[0]; \
|
||||
const uint32_t ne1 = dst->ne[1]; \
|
||||
const uint32_t ne2 = dst->ne[2]; \
|
||||
const uint32_t ne3 = dst->ne[3]; \
|
||||
\
|
||||
const uint32_t nb0 = dst->nb[0]; \
|
||||
const uint32_t nb1 = dst->nb[1]; \
|
||||
const uint32_t nb2 = dst->nb[2]; \
|
||||
const uint32_t nb3 = dst->nb[3]; \
|
||||
\
|
||||
const uint32_t nr = ne01;
|
||||
|
||||
static void cpy_thread_sametype_sameshape(struct htp_copy_context * ct, struct htp_ops_context * octx, const int nth, const int ith) {
|
||||
cpy_preamble;
|
||||
|
||||
// parallelize by src0 rows
|
||||
const uint32_t dr = ct->src0_nrows_per_thread;
|
||||
const uint32_t ir0 = dr * ith;
|
||||
const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
|
||||
|
||||
// copy by rows
|
||||
for (uint32_t i03 = 0; i03 < ne03; i03++) {
|
||||
for (uint32_t i02 = 0; i02 < ne02; i02++) {
|
||||
#pragma unroll(2)
|
||||
for (uint32_t i01 = ir0; i01 < ir1; i01++) {
|
||||
uint8_t* dst_ptr = (uint8_t*) dst->data + i01*nb1 + i02*nb2 + i03*nb3;
|
||||
uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
|
||||
hex_l2fetch(src0_ptr, ne00 * ct->src0_type_size, nb01, 2);
|
||||
hvx_copy_uu(dst_ptr, src0_ptr, ne00, ct->src0_type_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \
|
||||
static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) { \
|
||||
struct htp_copy_context * ct = (struct htp_copy_context *) data; \
|
||||
struct htp_ops_context * octx = ct->octx; \
|
||||
cpy_preamble; \
|
||||
const uint32_t dr = ct->src0_nrows_per_thread; \
|
||||
const uint32_t ir0 = dr * ith; \
|
||||
const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr; \
|
||||
if (ir0 >= nr) return; \
|
||||
for (uint32_t i03 = 0; i03 < ne03; i03++) { \
|
||||
for (uint32_t i02 = 0; i02 < ne02; i02++) { \
|
||||
_Pragma("unroll(4)") \
|
||||
for (uint32_t i01 = ir0; i01 < ir1; i01++) { \
|
||||
uint8_t* dst_ptr = (uint8_t*) dst->data + i01*nb1 + i02*nb2 + i03*nb3; \
|
||||
uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03; \
|
||||
hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2); \
|
||||
hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE); \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
}
|
||||
|
||||
static void cpy_thread_sametype_reshape(struct htp_copy_context * ct, struct htp_ops_context * octx, int nth, int ith) {
|
||||
cpy_preamble;
|
||||
DEFINE_CPY_SAMESHAPE(f32, float, 4)
|
||||
DEFINE_CPY_SAMESHAPE(f16, __fp16, 2)
|
||||
|
||||
// parallelize by src0 rows
|
||||
const uint32_t dr = ct->src0_nrows_per_thread;
|
||||
const uint32_t ir0 = dr * ith;
|
||||
const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
|
||||
|
||||
// Fast path: when both src0 and dst are contiguous in memory
|
||||
// Replace the element-by-element loop with a single bulk HVX copy per (i03, i02) slice.
|
||||
const bool src0_contig = (nb00 == ct->src0_type_size) &&
|
||||
(nb01 == ne00 * nb00) &&
|
||||
(nb02 == ne01 * nb01) &&
|
||||
(nb03 == ne02 * nb02);
|
||||
const bool dst_contig = (nb0 == ct->dst_type_size) &&
|
||||
(nb1 == ne0 * nb0) &&
|
||||
(nb2 == ne1 * nb1) &&
|
||||
(nb3 == ne2 * nb2);
|
||||
|
||||
if (src0_contig && dst_contig) {
|
||||
for (int64_t i03 = 0; i03 < ne03; i03++) {
|
||||
for (int64_t i02 = 0; i02 < ne02; i02++) {
|
||||
uint8_t * src_ptr = (uint8_t *) src0->data + i03*nb03 + i02*nb02 + ir0*nb01;
|
||||
uint32_t flat = ((i03*ne02 + i02)*ne01 + ir0) * ne00;
|
||||
uint8_t * dst_ptr = (uint8_t *) dst->data + flat * ct->src0_type_size;
|
||||
hvx_copy_uu(dst_ptr, src_ptr, (ir1 - ir0) * ne00, ct->src0_type_size);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// dst counters
|
||||
int64_t k10 = 0;
|
||||
int64_t i11 = 0;
|
||||
int64_t i12 = 0;
|
||||
int64_t i13 = 0;
|
||||
|
||||
// number of blocks in a row
|
||||
const int64_t nk00 = ct->src0_blocks_per_row;
|
||||
const int64_t nk0 = ct->dst_blocks_per_row;
|
||||
|
||||
for (int64_t i03 = 0; i03 < ne03; i03++) {
|
||||
for (int64_t i02 = 0; i02 < ne02; i02++) {
|
||||
k10 += nk00 * ir0;
|
||||
while (k10 >= nk0) {
|
||||
k10 -= nk0;
|
||||
if (++i11 == ne1) {
|
||||
i11 = 0;
|
||||
if (++i12 == ne2) {
|
||||
i12 = 0;
|
||||
if (++i13 == ne3) {
|
||||
i13 = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int64_t i01 = ir0; i01 < ir1; i01++) {
|
||||
for (int64_t k00 = 0; k00 < nk00; k00++) {
|
||||
const char * src0_ptr = ((char *) src0->data + k00*nb00 + i01*nb01 + i02*nb02 + i03*nb03);
|
||||
char * dst_ptr = ((char *) dst->data + k10*nb0 + i11*nb1 + i12*nb2 + i13*nb3);
|
||||
memcpy(dst_ptr, src0_ptr, ct->dst_type_size);
|
||||
|
||||
if (++k10 == nk0) {
|
||||
k10 = 0;
|
||||
if (++i11 == ne1) {
|
||||
i11 = 0;
|
||||
if (++i12 == ne2) {
|
||||
i12 = 0;
|
||||
if (++i13 == ne3) {
|
||||
i13 = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
k10 += nk00 * (ne01 - ir1);
|
||||
while (k10 >= nk0) {
|
||||
k10 -= nk0;
|
||||
if (++i11 == ne1) {
|
||||
i11 = 0;
|
||||
if (++i12 == ne2) {
|
||||
i12 = 0;
|
||||
if (++i13 == ne3) {
|
||||
i13 = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \
|
||||
static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) { \
|
||||
struct htp_copy_context * ct = (struct htp_copy_context *) data; \
|
||||
struct htp_ops_context * octx = ct->octx; \
|
||||
cpy_preamble; \
|
||||
const uint32_t dr = ct->src0_nrows_per_thread; \
|
||||
const uint32_t ir0 = dr * ith; \
|
||||
const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr; \
|
||||
if (ir0 >= nr) return; \
|
||||
const bool src0_contig = (nb00 == ELEM_SIZE) && \
|
||||
(nb01 == ne00 * nb00) && \
|
||||
(nb02 == ne01 * nb01) && \
|
||||
(nb03 == ne02 * nb02); \
|
||||
const bool dst_contig = (nb0 == ELEM_SIZE) && \
|
||||
(nb1 == ne0 * nb0) && \
|
||||
(nb2 == ne1 * nb1) && \
|
||||
(nb3 == ne2 * nb2); \
|
||||
if (src0_contig && dst_contig) { \
|
||||
for (int64_t i03 = 0; i03 < ne03; i03++) { \
|
||||
for (int64_t i02 = 0; i02 < ne02; i02++) { \
|
||||
uint8_t * src_ptr = (uint8_t *) src0->data + i03*nb03 + i02*nb02 + ir0*nb01; \
|
||||
uint32_t flat = ((i03*ne02 + i02)*ne01 + ir0) * ne00; \
|
||||
uint8_t * dst_ptr = (uint8_t *) dst->data + flat * ELEM_SIZE; \
|
||||
hvx_copy_uu(dst_ptr, src_ptr, (ir1 - ir0) * ne00, ELEM_SIZE); \
|
||||
} \
|
||||
} \
|
||||
return; \
|
||||
} \
|
||||
const bool reshape_flat_fast = (ne03 == 1 && ne2 == 1 && ne3 == 1) && \
|
||||
(ne0 == ne00 * ne01) && (ne1 == ne02) && \
|
||||
(nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE); \
|
||||
if (reshape_flat_fast) { \
|
||||
for (uint32_t i02 = 0; i02 < ne02; i02++) { \
|
||||
for (uint32_t i01 = ir0; i01 < ir1; i01++) { \
|
||||
uint8_t * src0_ptr = (uint8_t *) src0->data + i01 * nb01 + i02 * nb02; \
|
||||
uint8_t * dst_ptr = (uint8_t *) dst->data + i01 * ne00 * ELEM_SIZE + i02 * nb1; \
|
||||
hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE); \
|
||||
} \
|
||||
} \
|
||||
return; \
|
||||
} \
|
||||
int64_t k10 = 0; \
|
||||
int64_t i11 = 0; \
|
||||
int64_t i12 = 0; \
|
||||
int64_t i13 = 0; \
|
||||
const int64_t nk00 = ct->src0_blocks_per_row; \
|
||||
const int64_t nk0 = ct->dst_blocks_per_row; \
|
||||
for (int64_t i03 = 0; i03 < ne03; i03++) { \
|
||||
for (int64_t i02 = 0; i02 < ne02; i02++) { \
|
||||
k10 += nk00 * ir0; \
|
||||
while (k10 >= nk0) { \
|
||||
k10 -= nk0; \
|
||||
if (++i11 == ne1) { \
|
||||
i11 = 0; \
|
||||
if (++i12 == ne2) { \
|
||||
i12 = 0; \
|
||||
if (++i13 == ne3) { \
|
||||
i13 = 0; \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
for (int64_t i01 = ir0; i01 < ir1; i01++) { \
|
||||
for (int64_t k00 = 0; k00 < nk00; k00++) { \
|
||||
const char * src0_ptr = ((char *) src0->data + k00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); \
|
||||
char * dst_ptr = ((char *) dst->data + k10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); \
|
||||
memcpy(dst_ptr, src0_ptr, ELEM_SIZE); \
|
||||
if (++k10 == nk0) { \
|
||||
k10 = 0; \
|
||||
if (++i11 == ne1) { \
|
||||
i11 = 0; \
|
||||
if (++i12 == ne2) { \
|
||||
i12 = 0; \
|
||||
if (++i13 == ne3) { \
|
||||
i13 = 0; \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
k10 += nk00 * (ne01 - ir1); \
|
||||
while (k10 >= nk0) { \
|
||||
k10 -= nk0; \
|
||||
if (++i11 == ne1) { \
|
||||
i11 = 0; \
|
||||
if (++i12 == ne2) { \
|
||||
i12 = 0; \
|
||||
if (++i13 == ne3) { \
|
||||
i13 = 0; \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
} \
|
||||
}
|
||||
|
||||
static void cpy_thread_f16_f32_sameshape(struct htp_copy_context * ct, struct htp_ops_context * octx, const int nth, const int ith) {
|
||||
DEFINE_CPY_RESHAPE(f32, float, 4)
|
||||
DEFINE_CPY_RESHAPE(f16, __fp16, 2)
|
||||
|
||||
static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) {
|
||||
struct htp_copy_context * ct = (struct htp_copy_context *) data;
|
||||
struct htp_ops_context * octx = ct->octx;
|
||||
cpy_preamble;
|
||||
|
||||
// parallelize by src0 rows
|
||||
const uint32_t dr = ct->src0_nrows_per_thread;
|
||||
const uint32_t ir0 = dr * ith;
|
||||
const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
|
||||
if (ir0 >= nr) return;
|
||||
|
||||
// copy by rows
|
||||
for (uint32_t i03 = 0; i03 < ne03; i03++) {
|
||||
@@ -195,13 +207,16 @@ static void cpy_thread_f16_f32_sameshape(struct htp_copy_context * ct, struct ht
|
||||
}
|
||||
}
|
||||
|
||||
static void cpy_thread_f32_f16_sameshape(struct htp_copy_context * ct, struct htp_ops_context * octx, const int nth, const int ith) {
|
||||
static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, void * data) {
|
||||
struct htp_copy_context * ct = (struct htp_copy_context *) data;
|
||||
struct htp_ops_context * octx = ct->octx;
|
||||
cpy_preamble;
|
||||
|
||||
// parallelize by src0 rows
|
||||
const uint32_t dr = ct->src0_nrows_per_thread;
|
||||
const uint32_t ir0 = dr * ith;
|
||||
const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
|
||||
if (ir0 >= nr) return;
|
||||
|
||||
// copy by rows
|
||||
for (uint32_t i03 = 0; i03 < ne03; i03++) {
|
||||
@@ -217,11 +232,6 @@ static void cpy_thread_f32_f16_sameshape(struct htp_copy_context * ct, struct ht
|
||||
}
|
||||
}
|
||||
|
||||
static void cpy_work_func(unsigned int n, unsigned int i, void *data) {
|
||||
struct htp_copy_context *ct = (struct htp_copy_context *) data;
|
||||
ct->copy(ct, ct->octx, n, i);
|
||||
}
|
||||
|
||||
int op_cpy(struct htp_ops_context * octx) {
|
||||
cpy_preamble;
|
||||
|
||||
@@ -254,22 +264,32 @@ int op_cpy(struct htp_ops_context * octx) {
|
||||
|
||||
ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;
|
||||
|
||||
worker_callback_t copy_fun;
|
||||
|
||||
if (sametype && sameshape) {
|
||||
ct.copy = cpy_thread_sametype_sameshape;
|
||||
if (src0->type == HTP_TYPE_F32) {
|
||||
copy_fun = cpy_thread_f32_sameshape;
|
||||
} else {
|
||||
copy_fun = cpy_thread_f16_sameshape;
|
||||
}
|
||||
} else if (sameshape) {
|
||||
/**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32)
|
||||
ct.copy = cpy_thread_f16_f32_sameshape;
|
||||
copy_fun = cpy_thread_f16_f32_sameshape;
|
||||
else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16)
|
||||
ct.copy = cpy_thread_f32_f16_sameshape;
|
||||
copy_fun = cpy_thread_f32_f16_sameshape;
|
||||
else
|
||||
return HTP_STATUS_NO_SUPPORT;
|
||||
} else if (sametype) {
|
||||
ct.copy = cpy_thread_sametype_reshape;
|
||||
if (src0->type == HTP_TYPE_F32) {
|
||||
copy_fun = cpy_thread_f32_reshape;
|
||||
} else {
|
||||
copy_fun = cpy_thread_f16_reshape;
|
||||
}
|
||||
} else {
|
||||
return HTP_STATUS_NO_SUPPORT;
|
||||
}
|
||||
|
||||
worker_pool_run_func(octx->ctx->worker_pool, cpy_work_func, &ct, n_threads);
|
||||
worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
|
||||
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user