whispercpp 1.3.6 → 1.3.8

data/ext/sources/ggml/src/ggml-hexagon/htp/set-rows-ops.c

@@ -14,33 +14,37 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
-#include "htp-msg.h"
+#include "htp-ops.h"
 #include "htp-ops.h"
 
-#define set_rows_preamble \
-    const uint32_t ne00 = octx->src0.ne[0]; \
-    const uint32_t ne01 = octx->src0.ne[1]; \
-    const uint32_t ne02 = octx->src0.ne[2]; \
-    const uint32_t ne03 = octx->src0.ne[3]; \
-                                            \
-    const uint32_t ne10 = octx->src1.ne[0]; \
-    const uint32_t ne11 = octx->src1.ne[1]; \
-    const uint32_t ne12 = octx->src1.ne[2]; \
-                                            \
-    const uint32_t nb01 = octx->src0.nb[1]; \
-    const uint32_t nb02 = octx->src0.nb[2]; \
-    const uint32_t nb03 = octx->src0.nb[3]; \
-                                            \
-    const uint32_t nb10 = octx->src1.nb[0]; \
-    const uint32_t nb11 = octx->src1.nb[1]; \
-    const uint32_t nb12 = octx->src1.nb[2]; \
-                                            \
-    const uint32_t nb1 = octx->dst.nb[1];   \
-    const uint32_t nb2 = octx->dst.nb[2];   \
-    const uint32_t nb3 = octx->dst.nb[3];   \
-                                            \
-    const uint32_t ne1 = octx->dst.ne[1];   \
-                                            \
+#define set_rows_preamble                      \
+    const uint32_t ne00 = octx->src[0]->ne[0]; \
+    const uint32_t ne01 = octx->src[0]->ne[1]; \
+    const uint32_t ne02 = octx->src[0]->ne[2]; \
+    const uint32_t ne03 = octx->src[0]->ne[3]; \
+                                               \
+    const uint32_t ne10 = octx->src[1]->ne[0]; \
+    const uint32_t ne11 = octx->src[1]->ne[1]; \
+    const uint32_t ne12 = octx->src[1]->ne[2]; \
+    const uint32_t ne13 = octx->src[1]->ne[3]; \
+                                               \
+    const uint32_t nb01 = octx->src[0]->nb[1]; \
+    const uint32_t nb02 = octx->src[0]->nb[2]; \
+    const uint32_t nb03 = octx->src[0]->nb[3]; \
+                                               \
+    const uint32_t nb10 = octx->src[1]->nb[0]; \
+    const uint32_t nb11 = octx->src[1]->nb[1]; \
+    const uint32_t nb12 = octx->src[1]->nb[2]; \
+                                               \
+    const uint32_t nb1 = octx->dst->nb[1];     \
+    const uint32_t nb2 = octx->dst->nb[2];     \
+    const uint32_t nb3 = octx->dst->nb[3];     \
+                                               \
+    const uint32_t ne0 = octx->dst->ne[0];     \
+    const uint32_t ne1 = octx->dst->ne[1];     \
+    const uint32_t ne2 = octx->dst->ne[2];     \
+    const uint32_t ne3 = octx->dst->ne[3];     \
+                                               \
     const uint32_t nr  = ne01;
 
 struct htp_set_rows_context {
@@ -56,12 +60,17 @@ static void set_rows_thread_f32_f32(unsigned int nth, unsigned int ith, void *da
 
     set_rows_preamble;
 
+    uint64_t qt = HAP_perf_get_qtimer_count();
+
     // parallelize by rows of src0
     const uint32_t dr  = srctx->src0_nrows_per_thread;
     const uint32_t ir0 = dr * ith;
+    if (ir0 >= nr) {
+        return;
+    }
     const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr;
 
-    const bool is_i32 = (octx->src1.type == HTP_TYPE_I32);
+    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
 
     for (uint32_t i03 = 0; i03 < ne03; ++i03) {
         for (uint32_t i02 = 0; i02 < ne02; ++i02) {
@@ -70,7 +79,7 @@ static void set_rows_thread_f32_f32(unsigned int nth, unsigned int ith, void *da
                 const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11);
                 const uint32_t i10 = i;
 
-                const uintptr_t src1_addr = octx->src1.data + i10*nb10 + i11*nb11 + i12*nb12;
+                const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
 
                 uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr;
                 if (i1 >= ne1) {
@@ -78,14 +87,18 @@ static void set_rows_thread_f32_f32(unsigned int nth, unsigned int ith, void *da
                     continue;
                 }
 
-                const uintptr_t src0_ptr = octx->src0.data + i*nb01 + i02*nb02 + i03*nb03;
-                const uintptr_t dst_ptr  = octx->dst.data  + i1*nb1 + i02*nb2  + i03*nb3;
+                const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03;
+                const uintptr_t dst_ptr  = octx->dst->data  + i1*nb1 + i02*nb2  + i03*nb3;
 
                 // copy row
                 hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, ne00);
             }
         }
     }
+
+    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
+    FARF(HIGH, "set-rows-f32-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
+         ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt);
 }
 
 static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *data) {
@@ -94,12 +107,17 @@ static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *da
 
     set_rows_preamble;
 
+    uint64_t qt = HAP_perf_get_qtimer_count();
+
     // parallelize by rows of src0
     const uint32_t dr  = srctx->src0_nrows_per_thread;
     const uint32_t ir0 = dr * ith;
+    if (ir0 >= nr) {
+        return;
+    }
     const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr;
 
-    const bool is_i32 = (octx->src1.type == HTP_TYPE_I32);
+    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
 
     for (uint32_t i03 = 0; i03 < ne03; ++i03) {
         for (uint32_t i02 = 0; i02 < ne02; ++i02) {
@@ -108,7 +126,7 @@ static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *da
                 const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11);
                 const uint32_t i10 = i;
 
-                const uintptr_t src1_addr = octx->src1.data + i10*nb10 + i11*nb11 + i12*nb12;
+                const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
 
                 uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr;
                 if (i1 >= ne1) {
@@ -116,13 +134,17 @@ static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *da
                     continue;
                 }
 
-                const uint8_t* src0_ptr = (const uint8_t *) octx->src0.data + i*nb01 + i02*nb02 + i03*nb03;
-                uint8_t*       dst_ptr  = (uint8_t *)       octx->dst.data  + i1*nb1 + i02*nb2  + i03*nb3;
+                const uint8_t* src0_ptr = (const uint8_t *) octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03;
+                uint8_t*       dst_ptr  = (uint8_t *)       octx->dst->data  + i1*nb1 + i02*nb2  + i03*nb3;
 
                 hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
             }
         }
     }
+
+    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
+    FARF(HIGH, "set-rows-f16-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
+         ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt);
 }
 
 int op_set_rows(struct htp_ops_context * octx) {
@@ -130,15 +152,15 @@ int op_set_rows(struct htp_ops_context * octx) {
 
     const uint32_t n_threads = MIN(nr, octx->n_threads);
 
-    if (octx->src0.type != HTP_TYPE_F32) {
+    if (octx->src[0]->type != HTP_TYPE_F32) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->dst.type != HTP_TYPE_F32 && octx->dst.type != HTP_TYPE_F16) {
+    if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if (octx->src1.type != HTP_TYPE_I32 && octx->src1.type != HTP_TYPE_I64) {
+    if (octx->src[1]->type != HTP_TYPE_I32 && octx->src[1]->type != HTP_TYPE_I64) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
@@ -153,7 +175,7 @@ int op_set_rows(struct htp_ops_context * octx) {
 
     srctx.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;
 
-    switch(octx->dst.type) {
+    switch(octx->dst->type) {
     case HTP_TYPE_F32:
         worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f32_f32, &srctx, n_threads);
         break;

data/ext/sources/ggml/src/ggml-hexagon/htp/softmax-ops.c

@@ -15,68 +15,89 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
-#include "htp-msg.h"
+#include "htp-ops.h"
 #include "htp-ops.h"
 
-#define htp_softmax_preamble3                              \
-    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 ne10 = (src1->ne[0]) ? src1->ne[0] : 1; \
-    const uint32_t ne11 = (src1->ne[0]) ? src1->ne[1] : 1; \
-    const uint32_t ne12 = (src1->ne[0]) ? src1->ne[2] : 1; \
-    const uint32_t ne13 = (src1->ne[0]) ? src1->ne[3] : 1; \
-                                                           \
-    const uint32_t nb10 = (src1->ne[0]) ? src1->nb[0] : 1; \
-    const uint32_t nb11 = (src1->ne[0]) ? src1->nb[1] : 1; \
-    const uint32_t nb12 = (src1->ne[0]) ? src1->nb[2] : 1; \
-    const uint32_t nb13 = (src1->ne[0]) ? src1->nb[3] : 1; \
-                                                           \
-    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];                       \
+#define htp_softmax_preamble3                     \
+    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 ne10 = src1 ? src1->ne[0] : 1; \
+    const uint32_t ne11 = src1 ? src1->ne[1] : 1; \
+    const uint32_t ne12 = src1 ? src1->ne[2] : 1; \
+    const uint32_t ne13 = src1 ? src1->ne[3] : 1; \
+                                                  \
+    const uint32_t nb10 = src1 ? src1->nb[0] : 1; \
+    const uint32_t nb11 = src1 ? src1->nb[1] : 1; \
+    const uint32_t nb12 = src1 ? src1->nb[2] : 1; \
+    const uint32_t nb13 = src1 ? src1->nb[3] : 1; \
+                                                  \
+    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];
 
 struct htp_softmax_context {
+    struct htp_ops_context * octx;
+
     bool     use_f16;
     bool     use_src1;
+
     uint32_t n_head;
     uint32_t n_head_log2;
 
-    float scale;
-    float max_bias;
-    float m0;
-    float m1;
+    float    scale;
+    float    max_bias;
+    float    m0;
+    float    m1;
 
-    uint32_t src0_nrows_per_thread;
     struct fastdiv_values fastdiv_ne01;
     struct fastdiv_values fastdiv_ne02;
     struct fastdiv_values fastdiv_ne12; // For mask broadcasting
     struct fastdiv_values fastdiv_ne13; // For mask broadcasting
-    size_t spad_stride;
 
-    struct htp_ops_context * octx;
+    uint32_t src0_nrows_per_thread;
 };
 
+static void apply_mask(float * restrict wp0,
+                       const float * restrict mp_f32,
+                       const __fp16 * restrict mp_f16,
+                       uint32_t ne00,
+                       float slope,
+                       bool use_f16) {
+    if (!mp_f32) {
+        return;
+    }
+    if (use_f16) {
+        for (uint32_t i = 0; i < ne00; ++i) {
+            wp0[i] += slope * (float) mp_f16[i];
+        }
+    } else {
+        for (uint32_t i = 0; i < ne00; ++i) {
+            wp0[i] += slope * mp_f32[i];
+        }
+    }
+}
+
 static void init_softmax_ctx(struct htp_softmax_context * smctx, struct htp_ops_context * octx) {
-    const struct htp_tensor * src0 = &octx->src0;
-    const struct htp_tensor * src1 = &octx->src1;
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * src1 = octx->src[1];
 
     memset(smctx, 0, sizeof(struct htp_softmax_context));
 
-    memcpy(&smctx->scale, (float *) octx->op_params, sizeof(float));
+    memcpy(&smctx->scale,    (float *) octx->op_params,     sizeof(float));
     memcpy(&smctx->max_bias, (float *) octx->op_params + 1, sizeof(float));
 
     smctx->n_head      = src0->ne[2];
@@ -85,8 +106,8 @@ static void init_softmax_ctx(struct htp_softmax_context * smctx, struct htp_ops_
     smctx->m0 = powf(2.0f, -(smctx->max_bias) / smctx->n_head_log2);
     smctx->m1 = powf(2.0f, -(smctx->max_bias / 2.0f) / smctx->n_head_log2);
 
-    smctx->use_src1 = (src1->ne[0] != 0);
-    smctx->use_f16  = (src1->ne[0] != 0) && (src1->type == HTP_TYPE_F16);
+    smctx->use_src1 = (src1 != 0);
+    smctx->use_f16  = (src1 != 0) && (src1->type == HTP_TYPE_F16);
 
     smctx->octx = octx;
 
@@ -97,8 +118,8 @@ static void init_softmax_ctx(struct htp_softmax_context * smctx, struct htp_ops_
     if (ne01 > 0) smctx->fastdiv_ne01 = init_fastdiv_values(ne01);
     if (ne02 > 0) smctx->fastdiv_ne02 = init_fastdiv_values(ne02);
 
-    const uint32_t ne12 = (src1->ne[0]) ? src1->ne[2] : 1;
-    const uint32_t ne13 = (src1->ne[0]) ? src1->ne[3] : 1;
+    const uint32_t ne12 = src1 ? src1->ne[2] : 1;
+    const uint32_t ne13 = src1 ? src1->ne[3] : 1;
 
     if (ne12 > 0) smctx->fastdiv_ne12 = init_fastdiv_values(ne12);
     if (ne13 > 0) smctx->fastdiv_ne13 = init_fastdiv_values(ne13);
@@ -139,10 +160,7 @@ static void hvx_fast_softmax_prep_f32(const uint8_t * restrict src,
     }
 }
 
-static void hvx_fast_softmax_f32(const uint8_t * restrict src,
-                                 uint8_t * restrict dst,
-                                 uint8_t * restrict pad,
-                                 const int num_elems) {
+static void hvx_fast_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, uint8_t * restrict pad, const int num_elems) {
     const HVX_Vector * restrict v_src = (HVX_Vector *) src;
     HVX_Vector * restrict v_pad       = (HVX_Vector *) pad;
     HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
@@ -188,27 +206,20 @@ static void hvx_fast_softmax_f32(const uint8_t * restrict src,
     }
 }
 
-static float hvx_softmax_f32(const uint8_t * restrict src,
-                             uint8_t * restrict dst,
-                             uint8_t * restrict spad,
-                             const int   num_elems,
-                             const float max) {
+static float hvx_softmax_f32(const uint8_t * restrict src, uint8_t * restrict dst, uint8_t * restrict spad, const int  num_elems, const float max) {
     hvx_sub_scalar_f32(spad, src, max, num_elems);
 
-    hvx_exp_f32(spad, dst, num_elems, false);
-
-    float sum = hvx_reduce_sum_f32(dst, num_elems);
-
-    return sum;
+    hvx_exp_f32(dst, spad, num_elems, false);
+    return hvx_reduce_sum_f32(dst, num_elems);
 }
 
 static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
     struct htp_softmax_context * smctx = (struct htp_softmax_context *) data;
     struct htp_ops_context * octx = smctx->octx;
 
-    const struct htp_tensor * src0 = &octx->src0;
-    const struct htp_tensor * src1 = &octx->src1;
-    struct htp_tensor *       dst  = &octx->dst;
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * src1 = octx->src[1];
+    const struct htp_tensor * dst  = octx->dst;
 
     htp_softmax_preamble3;
 
@@ -223,22 +234,26 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
         return;
     }
 
-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
+    uint64_t qt = HAP_perf_get_qtimer_count();
 
     int is_aligned = 1;
     int opt_path   = 0;
+
     if (!hex_is_aligned((void *) src0->data, VLEN) || !hex_is_aligned((void *) dst->data, VLEN)) {
         is_aligned = 0;
         FARF(HIGH, "softmax-f32: unaligned addresses in elementwise op, possibly slower execution\n");
     }
+
+    // Only use the fast path when aligned AND row size is multiple of VLEN (128 bytes)
+    // The fast path (hvx_fast_softmax_f32) doesn't handle tail elements
+    // The non-opt path uses hvx_softmax_f32 which properly handles all sizes via its helper functions
     if ((1 == is_aligned) && !(nb01 & (VLEN - 1))) {
         opt_path = 1;
     }
 
-    uint8_t * src0_spad_data = octx->src0_spad.data + (ith * smctx->spad_stride);
-    uint8_t * src1_spad_data = octx->src1_spad.data + (ith * smctx->spad_stride);
-    uint8_t * dst_spad_data  = octx->dst_spad.data + (ith * smctx->spad_stride);
+    uint8_t * src0_spad_data = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
+    uint8_t * src1_spad_data = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread);
+    uint8_t * dst_spad_data  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
 
     float * wp0 = (float *) src0_spad_data;
     float * wp1 = (float *) src1_spad_data;
@@ -278,47 +293,29 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
         // ALiBi
         if (i2 != prev_i2) {
             const uint32_t h = i2;  // head
-
-            slope = (smctx->max_bias > 0.0f) ?
-                        h < smctx->n_head_log2 ?
-                        powf(smctx->m0, h + 1) :
-                        powf(smctx->m1, 2 * (h - smctx->n_head_log2) + 1) :
-                        1.0f;
+            slope = (smctx->max_bias > 0.0f) ? h < smctx->n_head_log2 ? powf(smctx->m0, h + 1) : powf(smctx->m1, 2 * (h - smctx->n_head_log2) + 1) : 1.0f;
             prev_i2 = i2;
         }
 
-        float * sp = (float *) ((char *) octx->src0.data + i1 * nb01 + i2 * nb02 + i3 * nb03);
-        float * dp = (float *) ((char *) octx->dst.data + i1 * nb1 + i2 * nb2 + i3 * nb3);
+        float * sp = (float *) ((char *) src0->data + i1 * nb01 + i2 * nb02 + i3 * nb03);
+        float * dp = (float *) ((char *) dst->data  + i1 * nb1  + i2 * nb2  + i3 * nb3);
 
         // broadcast the mask across rows
-        __fp16 * mp_f16 = (smctx->use_src1) ?
-                              (__fp16 *) ((char *) octx->src1.data + i11 * nb11 + i12 * nb12 + i13 * nb13) :
-                              NULL;
-        float *  mp_f32 = (smctx->use_src1) ?
-                              (float *) ((char *) octx->src1.data + i11 * nb11 + i12 * nb12 + i13 * nb13) :
-                              NULL;
+        __fp16 * mp_f16 = (smctx->use_src1) ? (__fp16 *) ((char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13) : NULL;
+        float *  mp_f32 = (smctx->use_src1) ? (float *)  ((char *) src1->data + i11 * nb11 + i12 * nb12 + i13 * nb13) : NULL;
 
         if ((1 == opt_path) && (mp_f32) && !(smctx->use_f16)) {
-            hvx_fast_softmax_prep_f32((const uint8_t *) sp, (uint8_t *) wp0, ne00, smctx->scale,
-                                      (const uint8_t *) mp_f32, slope);
-        } else {
+            hvx_fast_softmax_prep_f32((const uint8_t *) sp, (uint8_t *) wp0, ne00, smctx->scale, (const uint8_t *) mp_f32, slope);
+            hvx_fast_softmax_f32((const uint8_t *) wp0, (uint8_t *) dp, (uint8_t *) wp1, ne00);
+        } else if (1 == opt_path) {
             hvx_scale_f32((uint8_t *) wp0, (const uint8_t *) sp, ne00, smctx->scale);
-            if (mp_f32) {
-                if (smctx->use_f16) {
-                    for (int i = 0; i < ne00; ++i) {
-                        wp0[i] += slope * (float) mp_f16[i];
-                    }
-                } else {
-                    for (int i = 0; i < ne00; ++i) {
-                        wp0[i] += slope * mp_f32[i];
-                    }
-                }
-            }
-        }
-
-        if (1 == opt_path) {
+            apply_mask(wp0, mp_f32, mp_f16, ne00, slope, smctx->use_f16);
             hvx_fast_softmax_f32((const uint8_t *) wp0, (uint8_t *) dp, (uint8_t *) wp1, ne00);
         } else {
+            // Non-optimized path: uses HVX helper functions that properly handle all tensor sizes
+            // including non-multiples of 32 (the HVX vector lane count for f32)
+            hvx_scale_f32((uint8_t *) wp0, (const uint8_t *) sp, ne00, smctx->scale);
+            apply_mask(wp0, mp_f32, mp_f16, ne00, slope, smctx->use_f16);
             float max = hvx_reduce_max_f32((const uint8_t *) wp0, ne00);
             float sum = hvx_softmax_f32((const uint8_t *) wp0, (uint8_t *) wp2, (uint8_t *) wp1, ne00, max);
             sum       = sum > 0.0 ? (1.0 / sum) : 1;
@@ -326,54 +323,47 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
         }
     }
 
-    t2 = HAP_perf_get_qtimer_count();
-
-    FARF(HIGH, "softmax-f32 %d/%d/%d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
-         smctx->use_f16, opt_path, ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13,
-         ne0, ne1, ne2, ne3, (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
+    FARF(HIGH, "softmax-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u : opt %u f16 %u usec %u\n", ith, nth,
+         ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13,
+         ne0, ne1, ne2, ne3, opt_path, smctx->use_f16, (unsigned) qt);
 }
 
 static int execute_op_softmax_f32(struct htp_ops_context * octx) {
     int err = HTP_STATUS_OK;
 
-    const struct htp_tensor * src0 = &octx->src0;
-    const struct htp_tensor * src1 = &octx->src1;
-    struct htp_tensor *       dst  = &octx->dst;
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * src1 = octx->src[1];
+    const struct htp_tensor * dst  = octx->dst;
 
     struct htp_softmax_context smctx;
     const char * op_type = "softmax-f32";
 
-    switch (octx->op) {
-        case HTP_OP_SOFTMAX:
-            init_softmax_ctx(&smctx, octx);
-            break;
-
-        default:
-            FARF(ERROR, "Unsupported Op %u\n", octx->op);
-            return HTP_STATUS_NO_SUPPORT;
-    }
+    init_softmax_ctx(&smctx, octx);
 
     const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
     const uint32_t n_threads  = MIN(octx->n_threads, src0_nrows);
 
+    smctx.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+
     const size_t src0_row_size = src0->nb[1];
     const size_t src1_row_size = src0_row_size;
     const size_t dst_row_size  = dst->nb[1];
 
     // VTCM scratchpads for all tensors
-    // N rows per thread, padded to HVX vector size
-    octx->dst_spad.size  = hex_round_up(dst_row_size, 128) * n_threads;
-    octx->src0_spad.size = hex_round_up(src0_row_size, 128) * n_threads;
-    octx->src1_spad.size = hex_round_up(src1_row_size, 128) * n_threads;
+    // 4 rows per thread, padded to HVX vector size
+    octx->src0_spad.size_per_thread = hex_round_up(4 * src0_row_size, 128);
+    octx->src1_spad.size_per_thread = hex_round_up(4 * src1_row_size, 128);
+    octx->dst_spad.size_per_thread  = hex_round_up(4 * dst_row_size, 128);
 
-    // Use stride for calculating offset
-    smctx.spad_stride = hex_round_up(src0_row_size, 128);
+    octx->src0_spad.size = octx->src0_spad.size_per_thread * n_threads;
+    octx->src1_spad.size = octx->src1_spad.size_per_thread * n_threads;
+    octx->dst_spad.size  = octx->dst_spad.size_per_thread  * n_threads;
 
     size_t spad_size = octx->src0_spad.size + octx->src1_spad.size + octx->dst_spad.size;
 
-    if (src1->ne[0]) {
-        FARF(HIGH,
-             "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n",
+    if (src1) {
+        FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n",
              op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
              src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], octx->src0_spad.size, octx->src1_spad.size,
              octx->dst_spad.size);
@@ -385,19 +375,17 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
 
     // Make sure the reserved vtcm size is sufficient
     if (octx->ctx->vtcm_size < spad_size) {
-        FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size,
-             spad_size);
+        FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, spad_size);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    octx->src0_spad.data = octx->ctx->vtcm_base;
-    octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
-    octx->dst_spad.data  = octx->src1_spad.data + octx->src1_spad.size;
+    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
+    octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL;
+    octx->dst_spad.data  = octx->src1_spad.data + octx->src1_spad.size; octx->dst_spad.src  = NULL;
 
-    if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        smctx.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
-        worker_pool_run_func(octx->ctx->worker_pool, softmax_job_f32, &smctx, n_threads);
-    }
+    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) return err;
+
+    worker_pool_run_func(octx->ctx->worker_pool, softmax_job_f32, &smctx, n_threads);
 
     return err;
 }
@@ -405,7 +393,7 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
 int op_softmax(struct htp_ops_context * octx) {
     int err = HTP_STATUS_OK;
 
-    switch (octx->src0.type) {
+    switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
             err = execute_op_softmax_f32(octx);
             break;