Skip to content

Commit 7223ee4

Browse files
aparmp-quicaparmp-quic
committed
hex-unary: fix RMS_NORM_MUL weight-offset bugs for grouped/broadcast norms
1 parent 0ef6e55 commit 7223ee4

1 file changed

Lines changed: 20 additions & 12 deletions

File tree

‎ggml/src/ggml-hexagon/htp/unary-ops.c‎

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -478,6 +478,9 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
478478
const uint32_t nb11 = src1 ? src1->nb[1] : 0; \
479479
const uint32_t nb12 = src1 ? src1->nb[2] : 0; \
480480
const uint32_t nb13 = src1 ? src1->nb[3] : 0; \
481+
const uint32_t nb11_bc = (src1 && src1->ne[1] > 1) ? nb11 : 0; \
482+
const uint32_t nb12_bc = (src1 && src1->ne[2] > 1) ? nb12 : 0; \
483+
const uint32_t nb13_bc = (src1 && src1->ne[3] > 1) ? nb13 : 0; \
481484
const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false; \
482485
\
483486
uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \
@@ -497,8 +500,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
497500
const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \
498501
const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \
499502
\
500-
const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \
501-
const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \
503+
const bool src1_needs_row_clip = (IS_RMS_NORM_MUL) && !uctx->broadcast_weight && !src1_contig; \
504+
const bool block_src0_contig = src0_contig && !src1_needs_row_clip; \
505+
const bool block_dst_contig = dst_contig && !src1_needs_row_clip; \
506+
\
507+
const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \
508+
const uint32_t dst_max_block = block_dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \
502509
const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); \
503510
if (BLOCK == 0) { \
504511
FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small, needed at least %zu\n", \
@@ -515,8 +522,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
515522
} \
516523
\
517524
for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { \
518-
const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \
519-
div_ne01); \
525+
const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \
526+
ne01, div_ne01); \
520527
\
521528
dma_queue_push(dma_queue, \
522529
dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), \
@@ -530,7 +537,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
530537
\
531538
if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \
532539
const size_t src1_off = src1_contig ? (ir * nb11) : \
533-
unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \
540+
unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc); \
534541
dma_queue_push(dma_queue, \
535542
dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), \
536543
uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size); \
@@ -540,8 +547,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
540547
} \
541548
\
542549
for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { \
543-
const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \
544-
div_ne01); \
550+
const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \
551+
ne01, div_ne01); \
545552
\
546553
float * dst_vtcm = (float *) dma_queue_pop(dma_queue).src; \
547554
float * src0_vtcm = (float *) dma_queue_pop(dma_queue).dst; \
@@ -562,12 +569,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
562569
\
563570
const uint32_t next_ir = ir + block_size; \
564571
if (next_ir < src0_end_row) { \
565-
const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig,\
566-
ne01, div_ne01); \
572+
const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, block_src0_contig, \
573+
block_dst_contig, ne01, div_ne01); \
567574
const uint32_t pref_ir = next_ir + next_block_size; \
568575
if (pref_ir < src0_end_row) { \
569-
const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, \
570-
dst_contig, ne01, div_ne01); \
576+
const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, block_src0_contig, \
577+
block_dst_contig, ne01, div_ne01); \
571578
const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : \
572579
unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \
573580
dma_queue_push(dma_queue, \
@@ -576,7 +583,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
576583
\
577584
if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \
578585
const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : \
579-
unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \
586+
unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, \
587+
nb13_bc); \
580588
dma_queue_push(dma_queue, \
581589
dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off), \
582590
uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); \

0 commit comments

Comments
 (0)