[softmax_acc]: Improve Accuarcy and Fix Issue in LUT.

mfarag13 · dzakhar · commit 8a60337fa8fa · 2021-08-19T13:56:48.000+03:00
diff --git a/lib/src/bricks/impl/mli_prv_lut_dsp.h b/lib/src/bricks/impl/mli_prv_lut_dsp.h
@@ -47,7 +47,9 @@ static MLI_FORCE_INLINE v2q15_t activation_lut_two_elem_interpolate(
     int shift_in = in_frac_bits - lut->in_frac_bits;
     // if shift amount is too high, preshift argument itself and
     // limit shift amount to prevent overflows
+    constexpr int max_shift = 15;
     int preshift_in = mli_math_max_fx(shift_in - (int)kMaxFracBitsFx16, 0);
+        preshift_in = mli_math_min_fx(preshift_in, max_shift);
     shift_in = mli_math_min_fx(shift_in, (int)kMaxFracBitsFx16);
 
     v2q15_t offset = mli_prv_init_v<int16_t, v2q15_t>(lut->input_offset);
@@ -58,17 +60,25 @@ static MLI_FORCE_INLINE v2q15_t activation_lut_two_elem_interpolate(
 
     /* Convert Input SA8 to FX */
     v2q15_t x = in;
+    v2q15_t lut_idx;
+    v2q15_t frac;
     if (convert_input) {
-        int shift = ((int32_t) in_params->shift - in_frac_bits) + preshift_in;
-        x = mli_prv_convert_sa8_fx16<v2q15_t, v2q15_t>(x, in_params->offset, in_params->scale, shift);
+        int shift = (int32_t) in_params->shift - in_frac_bits;
+        v2q31_t x_int = mli_prv_convert_sa8_fx16<v2q15_t, v2q31_t>(x, in_params->offset, in_params->scale, shift);
+        x_int = mli_math_asr_fx(x_int, preshift_in);
+        frac[0] = x_int[0] & mask[0];
+        frac[1] = x_int[1] & mask[1];
+        x_int = mli_math_asr_fx(x_int, shift_in);
+        lut_idx[0] = mli_math_bound_range_fx(mli_math_add_fx(x_int[0], (int32_t)offset[0]), lower[0], upper[0]);
+        lut_idx[1] = mli_math_bound_range_fx(mli_math_add_fx(x_int[1], (int32_t)offset[1]), lower[1], upper[1]);
     } else {
         x = mli_math_acc_ashift_fx(x, preshift_in);
+        frac = x & mask;
+        lut_idx = mli_math_add_fx(mli_math_acc_ashift_fx(x, shift_in), offset);
+        lut_idx = mli_math_bound_range_fx(lut_idx, lower, upper);
     }
 
-    v2q15_t lut_idx = mli_math_add_fx(mli_math_acc_ashift_fx(x, shift_in), offset);
-    lut_idx = mli_math_bound_range_fx(lut_idx, lower, upper);
     // perform linear interpolation
-    v2q15_t frac = x & mask;
     v2q15_t res = mli_prv_init_v(lut_data[lut_idx[0]], lut_data[lut_idx[1]]);
     v2q15_t next = mli_prv_init_v(lut_data[lut_idx[0] + 1], lut_data[lut_idx[1] + 1]);
     v2q15_t diff = mli_math_sub_fx(res, next);
@@ -151,7 +161,6 @@ static MLI_FORCE_INLINE void compute_activation_lut(
         const struct s8asym_quant_params *in_params,
         struct s8asym_quant_params *out_params) {
 
-    MLI_ASSERT(in_frac_bits >= -1);  // -1 may be required by softmax
     MLI_ASSERT(lut->in_frac_bits >= 0);
     MLI_ASSERT(lut->length >= 0);
     MLI_ASSERT(MLI_MAX_RANK == 4);
diff --git a/lib/src/bricks/impl/mli_prv_lut_ref.h b/lib/src/bricks/impl/mli_prv_lut_ref.h
@@ -29,7 +29,6 @@ static MLI_FORCE_INLINE void compute_activation_lut(
         const struct s8asym_quant_params *in_params,
         struct s8asym_quant_params *out_params) {
 
-    MLI_ASSERT(in_frac_bits >= -1);  // -1 may be required by softmax
     MLI_ASSERT(lut->in_frac_bits >= 0);
     MLI_ASSERT(lut->length >= 0);
     MLI_ASSERT(MLI_MAX_RANK == 4);
@@ -103,23 +102,23 @@ static MLI_FORCE_INLINE out_T activation_lut_one_elem_interpolate(
     int shift_in = in_frac_bits - lut->in_frac_bits;
     // if shift amount is too high, preshift argument itself and
     // limit shift amount to prevent overflows
+    constexpr int max_shift = 15;
     int preshift_in = mli_math_max_fx(shift_in - (int)kMaxFracBitsFx16, 0);
+        preshift_in = mli_math_min_fx(preshift_in, max_shift);
     shift_in = mli_math_min_fx(shift_in, (int)kMaxFracBitsFx16);
 
     int16_t mask = (1 << shift_in) - 1;
 
     /* Convert Input SA8 to FX */
-    int16_t input;
+    int32_t input;
     if (convert_input) {
         int shift = ((int32_t) in_params->shift - in_frac_bits);
-        input = mli_prv_convert_sa8_fx16<in_T, int16_t>(in, in_params->offset, in_params->scale, shift);
+        input = mli_prv_convert_sa8_fx16<in_T, int32_t>(in, in_params->offset, in_params->scale, shift);
     } else {
         input = in;
     }
-    constexpr int max_shift = 15;
-    preshift_in = mli_math_min_fx(preshift_in, max_shift);
-    int16_t x = input >> preshift_in;
-    int lut_idx = mli_math_add_fx((x >> shift_in), lut->input_offset);
+    int32_t x = mli_math_asr_fx(input, preshift_in);
+    int lut_idx = mli_math_add_fx(mli_math_asr_fx(x, shift_in), lut->input_offset);
     lut_idx = mli_math_bound_range_fx(lut_idx, 0, lut->length - 2);
     // perform linear interpolation
     int16_t frac = x & mask;
@@ -174,7 +173,7 @@ static MLI_FORCE_INLINE out_T activation_lut_one_elem_no_interpolate(
         input = in;
     }
     int x = (int)input;
-    int lut_idx = mli_math_add_fx((x << -shift_in), lut->input_offset);
+    int lut_idx = mli_math_add_fx(mli_math_asl_fx(x, -shift_in), lut->input_offset);
     lut_idx = mli_math_bound_range_fx(lut_idx, 0, lut->length - 1);
     // no interpolation
     int16_t res = lut_data[lut_idx];
diff --git a/lib/src/bricks/impl/mli_prv_lut_vdsp.h b/lib/src/bricks/impl/mli_prv_lut_vdsp.h
@@ -117,7 +117,6 @@ static MLI_FORCE_INLINE vNx4short_t activation_lut_vec_elem_interpolate(
         int8_t in_frac_bits,
         const struct s8asym_quant_params *in_params) {
 
-    MLI_ASSERT(in_frac_bits >= -1);  // -1 may be required by softmax
     MLI_ASSERT(lut->in_frac_bits >= 0);
     MLI_ASSERT(lut->length >= 0);
 
@@ -133,26 +132,36 @@ static MLI_FORCE_INLINE vNx4short_t activation_lut_vec_elem_interpolate(
     const MLI_PTR(short) lut_data = (const MLI_PTR(short))lut->data.mem.pi16;
     // if shift amount is too high, preshift argument itself and
     // limit shift amount to prevent overflows
+    constexpr int max_shift = 15;
     int preshift_in = mli_math_max_fx(shift_in - (int)kMaxFracBitsFx16, 0);
+        preshift_in = mli_math_min_fx(preshift_in, max_shift);
     shift_in = mli_math_min_fx(shift_in, (int)kMaxFracBitsFx16);
 
     // input data is more precise than LUT
     int16_t mask = (1 << shift_in) - 1;
     vNx4short_t x = in;
+    vNx4int_t lut_idx_int;
+    vNx4short_t frac;
     if (convert) {
-        int shift = ((int32_t) in_params->shift - in_frac_bits) + preshift_in;
-        x = mli_prv_convert_sa8_fx16<vNx4short_t, vNx4short_t>(x, in_params->offset, in_params->scale, shift);
+        int shift = (int32_t) in_params->shift - in_frac_bits;
+        vNx4int_t x_int = mli_prv_convert_sa8_fx16<vNx4short_t, vNx4int_t>(x, in_params->offset, in_params->scale, shift);
+        x_int = mli_math_asr_fx(x_int, preshift_in);
+        frac = mli_math_cast_fx<vNx4int_t, vNx4short_t>(x_int & mask);
+
+        /* Calculate lut_idx */
+        vNx4int_t lut_idx = mli_math_add_fx<vNx4int_t>(mli_math_asr_fx(x_int, shift_in), lut->input_offset);
+        lut_idx_int = mli_math_bound_range_fx(lut_idx , 0, lut->length - 2);
+
     } else {
-        constexpr int max_shift = 15;
-        preshift_in = mli_math_min_fx(preshift_in, max_shift);
         x = mli_math_asr_fx(x, preshift_in);
+        frac = x & mask;
+
+        /* Calculate lut_idx */
+        vNx4short_t lut_idx = mli_math_add_fx<vNx4short_t>(mli_math_asr_fx(x, shift_in), lut->input_offset);
+        lut_idx = mli_math_bound_range_fx(lut_idx , 0, lut->length - 2);
+        lut_idx_int = mli_math_mul_fx<vNx4short_t, vNx4int_t>(lut_idx, 1);
     }
-    vNx4short_t lut_idx = mli_math_add_fx<vNx4short_t>(mli_math_asr_fx(x, shift_in), lut->input_offset);
-    /* Calculate lut_idx */
-    lut_idx = mli_math_bound_range_fx(lut_idx , 0, lut->length - 2);
-    vNx4int_t lut_idx_int = mli_math_mul_fx<vNx4short_t, vNx4int_t>(lut_idx, 1);
 
-    vNx4short_t frac = x & mask;
     /* Load from LUT */
     vNx4short_t lut_values = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int);
     vNx4short_t lut_values_next = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int + 1);
@@ -162,9 +171,7 @@ static MLI_FORCE_INLINE vNx4short_t activation_lut_vec_elem_interpolate(
                                         mli_math_mul_fx<vNx4short_t, vNx4accint_t>(diffs, frac), shift_in);
 
     /* Calculate O/P */
-    vNx4short_t result = mli_math_sub_fx<vNx4short_t>(lut_values, diffs_mul_frac_cast);
-    
-    return result;
+    return mli_math_sub_fx<vNx4short_t>(lut_values, diffs_mul_frac_cast);
 }
 
 template <bool convert>
@@ -174,7 +181,6 @@ static MLI_FORCE_INLINE vNx4short_t activation_lut_vec_elem_no_interpolate(
         int8_t in_frac_bits,
         const struct s8asym_quant_params *in_params) {
 
-    MLI_ASSERT(in_frac_bits >= -1);  // -1 may be required by softmax
     MLI_ASSERT(lut->in_frac_bits >= 0);
     MLI_ASSERT(lut->length >= 0);
 
@@ -213,30 +219,33 @@ static MLI_FORCE_INLINE vNx4short_t activation_lut_vec_elem_no_interpolate(
 template <typename io_T, bool convert>
 static MLI_FORCE_INLINE void load_input_and_get_lut_idx(
         MLI_PTR(io_T) __restrict in_ptr,
-        vNx4short_t &vec,
-        vNx4short_t &lut_idx,
+        vNx4short_t &x,
+        vNx4int_t &x_int,
         vNx4int_t &lut_idx_int,
         int16_t in_frac_bits,
         int preshift_in,
         int shift_in,
         const mli_lut *lut,
         const struct s8asym_quant_params *in_params) {
 
-    vec = activation_lut_load_input<io_T, vNx4short_t>(in_ptr);
-
+    x = activation_lut_load_input<io_T, vNx4short_t>(in_ptr);
     if (convert) {
-        int shift = ((int32_t) in_params->shift - in_frac_bits) + preshift_in;
-        vec = mli_prv_convert_sa8_fx16<vNx4short_t, vNx4short_t>(vec, in_params->offset, in_params->scale, shift);
+        int shift = (int32_t) in_params->shift - in_frac_bits;
+        x_int = mli_prv_convert_sa8_fx16<vNx4short_t, vNx4int_t>(x, in_params->offset, in_params->scale, shift);
+        x_int = mli_math_asr_fx(x_int, preshift_in);
+
+        /* Calculate lut_idx */
+        vNx4int_t lut_idx = mli_math_add_fx<vNx4int_t>(mli_math_asr_fx(x_int, shift_in), lut->input_offset);
+        lut_idx_int = mli_math_bound_range_fx(lut_idx , 0, lut->length - 2);
+
     } else {
-        constexpr int max_shift = 15;
-        preshift_in = mli_math_min_fx(preshift_in, max_shift);
-        vec = mli_math_asr_fx(vec, preshift_in);
-    }
+        x = mli_math_asr_fx(x, preshift_in);
 
-    /* Calculate lut_idx */
-    lut_idx = mli_math_add_fx<vNx4short_t>(mli_math_asr_fx(vec, shift_in), lut->input_offset);
-    lut_idx = mli_math_bound_range_fx(lut_idx , 0, lut->length - 2);
-    lut_idx_int = mli_math_mul_fx<vNx4short_t, vNx4int_t>(lut_idx, 1);
+        /* Calculate lut_idx */
+        vNx4short_t lut_idx = mli_math_add_fx<vNx4short_t>(mli_math_asr_fx(x, shift_in), lut->input_offset);
+        lut_idx = mli_math_bound_range_fx(lut_idx , 0, lut->length - 2);
+        lut_idx_int = mli_math_mul_fx<vNx4short_t, vNx4int_t>(lut_idx, 1);
+    }
 }
 
 template <typename io_T, bool convert>
@@ -250,7 +259,6 @@ static MLI_FORCE_INLINE void compute_activation_lut_func(
         const struct s8asym_quant_params *in_params,
         struct s8asym_quant_params *out_params) {
 
-    MLI_ASSERT(in_frac_bits >= -1);  // -1 may be required by softmax
     MLI_ASSERT(lut->in_frac_bits >= 0);
     MLI_ASSERT(lut->length >= 0);
     MLI_ASSERT(MLI_MAX_RANK == 4);
@@ -271,7 +279,9 @@ static MLI_FORCE_INLINE void compute_activation_lut_func(
     const MLI_PTR(short) lut_data = (const MLI_PTR(short))lut->data.mem.pi16;
     // if shift amount is too high, preshift argument itself and
     // limit shift amount to prevent overflows
+    constexpr int max_shift = 15;
     int preshift_in = mli_math_max_fx(shift_in - (int)kMaxFracBitsFx16, 0);
+        preshift_in = mli_math_min_fx(preshift_in, max_shift);
     shift_in = mli_math_min_fx(shift_in, (int)kMaxFracBitsFx16);
 
     int remaining_part = in->shape[3] & (_VDSP_NUM_8BIT_LANES - 1);
@@ -294,49 +304,61 @@ static MLI_FORCE_INLINE void compute_activation_lut_func(
                     }
 
                     /* Manual software pipelining */
-                    vNx4short_t x, lut_idx;
-                    vNx4int_t lut_idx_int;
+                    vNx4short_t x, frac;
+                    vNx4int_t x_int, lut_idx_int;
                     vNx4short_t _lut_values, _lut_values_next, _frac;
 
-                    load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, lut_idx, lut_idx_int,
+                    load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, x_int, lut_idx_int,
                             in_frac_bits, preshift_in, shift_in, lut, in_params);
-                    vNx4short_t frac = x & mask;
+                    if (convert) {
+                        frac = mli_math_cast_fx<vNx4int_t, vNx4short_t>(x_int & mask);
+                    } else {
+                        frac = x & mask;
+                    }
                     input_ptr  += _VDSP_NUM_8BIT_LANES;
 
                     if (in->shape[3] >= _VDSP_NUM_8BIT_LANES && !convert) {
                         /* Load from LUT */
                         _lut_values = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int);
                         _lut_values_next = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int + 1);
 
-                        load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, lut_idx, lut_idx_int,
+                        load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, x_int, lut_idx_int,
                                 in_frac_bits, preshift_in, shift_in, lut, in_params);
 
                         _frac = frac;
-                        frac = x & mask;
+                        if (convert) {
+                            frac = mli_math_cast_fx<vNx4int_t, vNx4short_t>(x_int & mask);
+                        } else {
+                            frac = x & mask;
+                        }
                         input_ptr  += _VDSP_NUM_8BIT_LANES;
 
                         for (int pos3 = remaining_part; pos3 < in->shape[3] - _VDSP_NUM_8BIT_LANES; pos3 += _VDSP_NUM_8BIT_LANES) {
                             /* Load from LUT */
                             vNx4short_t lut_values = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int);
                             vNx4short_t lut_values_next = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int + 1);
 
-                            load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, lut_idx, lut_idx_int,
+                            load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, x_int, lut_idx_int,
                                     in_frac_bits, preshift_in, shift_in, lut, in_params);
 
-                                /* perform linear interpolation */
-                                vNx4short_t diffs = mli_math_sub_fx<vNx4short_t>(_lut_values, _lut_values_next);
-                                vNx4short_t diffs_mul_frac_cast =  mli_math_acc_cast_fx<vNx4short_t, vNx4accint_t>(
-                                                                    mli_math_mul_fx<vNx4short_t, vNx4accint_t>(diffs, _frac), shift_in);
+                            /* perform linear interpolation */
+                            vNx4short_t diffs = mli_math_sub_fx<vNx4short_t>(_lut_values, _lut_values_next);
+                            vNx4short_t diffs_mul_frac_cast =  mli_math_acc_cast_fx<vNx4short_t, vNx4accint_t>(
+                                                                mli_math_mul_fx<vNx4short_t, vNx4accint_t>(diffs, _frac), shift_in);
 
-                                /* Calculate O/P */
-                                vNx4short_t res = mli_math_sub_fx<vNx4short_t>(_lut_values, diffs_mul_frac_cast);
+                            /* Calculate O/P */
+                            vNx4short_t res = mli_math_sub_fx<vNx4short_t>(_lut_values, diffs_mul_frac_cast);
 
-                                /* Store O/P */
-                                activation_lut_store_output<io_T, convert>(output_ptr, res, lut, out_params);
-                                output_ptr += _VDSP_NUM_8BIT_LANES;
+                            /* Store O/P */
+                            activation_lut_store_output<io_T, convert>(output_ptr, res, lut, out_params);
+                            output_ptr += _VDSP_NUM_8BIT_LANES;
 
                             _frac = frac;
-                            frac = x & mask;
+                            if (convert) {
+                                frac = mli_math_cast_fx<vNx4int_t, vNx4short_t>(x_int & mask);
+                            } else {
+                                frac = x & mask;
+                            }
 
                             input_ptr  += _VDSP_NUM_8BIT_LANES;
 
@@ -360,15 +382,19 @@ static MLI_FORCE_INLINE void compute_activation_lut_func(
                             vNx4short_t lut_values = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int);
                             vNx4short_t lut_values_next = mli_prv_gather_load_nx4_samples(lut_data, lut_idx_int + 1);
 
-                            load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, lut_idx, lut_idx_int,
+                            load_input_and_get_lut_idx<io_T, convert>(input_ptr, x, x_int, lut_idx_int,
                                     in_frac_bits, preshift_in, shift_in, lut, in_params);
 
                             /* perform linear interpolation */
                             vNx4short_t diffs = mli_math_sub_fx<vNx4short_t>(lut_values, lut_values_next);
                             vNx4short_t diffs_mul_frac_cast =  mli_math_acc_cast_fx<vNx4short_t, vNx4accint_t>(
                                                                 mli_math_mul_fx<vNx4short_t, vNx4accint_t>(diffs, frac), shift_in);
 
-                            frac = x & mask;
+                            if (convert) {
+                                frac = mli_math_cast_fx<vNx4int_t, vNx4short_t>(x_int & mask);
+                            } else {
+                                frac = x & mask;
+                            }
 
                             /* Calculate O/P */
                             vNx4short_t res = mli_math_sub_fx<vNx4short_t>(lut_values, diffs_mul_frac_cast);
diff --git a/lib/src/bricks/impl/mli_prv_quant_dsp.h b/lib/src/bricks/impl/mli_prv_quant_dsp.h
@@ -286,6 +286,18 @@ MLI_FORCE_INLINE v2q15_t mli_prv_convert_sa8_fx16(
     return mli_math_acc_cast_fx<v2q15_t, v2accum40_t>(in_scaled, shift);
 }
 
+template<>
+MLI_FORCE_INLINE v2q31_t mli_prv_convert_sa8_fx16(
+    const v2q15_t in,
+    const int16_t zero_point,
+    const int16_t scale,
+    const int shift) {
+    v2q31_t out;
+    out[0] = mli::krn::ref::mli_prv_convert_sa8_fx16<int8_t, int32_t>(in[0], zero_point, scale, shift);
+    out[1] = mli::krn::ref::mli_prv_convert_sa8_fx16<int8_t, int32_t>(in[1], zero_point, scale, shift);
+    return out;
+}
+
 template<>
 MLI_FORCE_INLINE v2q15_t mli_prv_convert_fx16_sa8(
     const v2q15_t in,
diff --git a/lib/src/bricks/impl/mli_prv_quant_vdsp.h b/lib/src/bricks/impl/mli_prv_quant_vdsp.h
@@ -310,6 +310,21 @@ MLI_FORCE_INLINE vNx4short_t mli_prv_convert_sa8_fx16(
     return res;
 }
 
+template<>
+MLI_FORCE_INLINE vNx4int_t mli_prv_convert_sa8_fx16(
+        const vNx4short_t in_val,
+        const int16_t zero_point,
+        const int16_t scale,
+		const int shift) {
+    int shift_right = mli_math_max_fx(shift, 0);
+    int shift_left = mli_math_max_fx(-shift, 0);
+    vNx4short_t in_biased_shifted_no_zp = mli_math_sub_fx<vNx4short_t>(in_val, zero_point);
+    vNx4int_t in_scaled = mli_math_mul_fx<vNx4short_t, vNx4int_t>(in_biased_shifted_no_zp, scale);
+    vNx4int_t res = mli_math_asr_rnd_fx(in_scaled, shift_right);
+    res = mli_math_asl_fx(res, shift_left);
+    return res;
+}
+
 MLI_FORCE_INLINE vNx4int_t mli_prv_convert_sa8_fx32(
         const vNx4char_t in_val,
         const int16_t zero_point,
diff --git a/lib/src/pal/dsp/mli_math.h b/lib/src/pal/dsp/mli_math.h
@@ -68,6 +68,13 @@ MLI_FORCE_INLINE int16_t mli_math_asr_fx(int16_t acc, int shift_right) {
     return fx_asr_q15(acc, shift_right);
 }
 
+template <>
+MLI_FORCE_INLINE v2q31_t mli_math_asr_fx(v2q31_t acc, int shift_right) {
+    acc[0] = fx_asr_q31(acc[0], shift_right);
+    acc[1] = fx_asr_q31(acc[1], shift_right);
+    return acc;
+}
+
 template <typename T>
 MLI_FORCE_INLINE T mli_math_limit_fx(T sign) {
     return sign < (T)0 ? std::numeric_limits<T>::lowest() : std::numeric_limits<T>::max();
diff --git a/lib/src/private/src/mli_prv_activation_lut.cc b/lib/src/private/src/mli_prv_activation_lut.cc
diff --git a/user_tests/tests/mli_krn_softmax/tests_mli_krn_softmax.cc b/user_tests/tests/mli_krn_softmax/tests_mli_krn_softmax.cc
