COMPMID-3600: MUL unit test failing with data type QUANT8_ASYMM

Add broadcast support for NEPixelWiseMultiplicationKernel with QASYMM8/QASYMM8_SIGNED
Add test case for QASYMM8 broadcast
Fix QASYMM8 saturation issue

Signed-off-by: Sheri Zhang <sheri.zhang@arm.com>
Change-Id: Ie67cfa8b94ab542133b031efbff8379cc57cfc2d
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/3586
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Michalis Spyrou <michalis.spyrou@arm.com>

3 files changed, 155 insertions, 58 deletions

diff --git a/src/core/NEON/kernels/NEPixelWiseMultiplicationKernel.cpp b/src/core/NEON/kernels/NEPixelWiseMultiplicationKernel.cpp
index 4a9da1db72..f8875324de 100644
--- a/src/core/NEON/kernels/NEPixelWiseMultiplicationKernel.cpp
+++ b/src/core/NEON/kernels/NEPixelWiseMultiplicationKernel.cpp
@@ -141,7 +141,7 @@ template <typename T>
 inline typename std::enable_if<std::is_same<T, int8_t>::value, int8_t>::type
 quantize(float val, const UniformQuantizationInfo &info)
 {
-    int32_t tmp = static_cast<int32_t>(val / info.scale) + info.offset;
+    const int32_t tmp = static_cast<int32_t>(val / info.scale) + info.offset;
 
     T tmp_qua = static_cast<T>(tmp > SCHAR_MAX) ? SCHAR_MAX : ((tmp < SCHAR_MIN) ? SCHAR_MIN : tmp);
     return tmp_qua;
@@ -151,9 +151,9 @@ template <typename T>
 inline typename std::enable_if<std::is_same<T, uint8_t>::value, uint8_t>::type
 quantize(float val, const UniformQuantizationInfo &info)
 {
-    int32_t tmp = static_cast<int32_t>(val / info.scale) + info.offset;
+    const int32_t tmp = static_cast<int32_t>(val / info.scale) + info.offset;
 
-    T tmp_qua = static_cast<T>((tmp > UCHAR_MAX) ? UCHAR_MAX : tmp);
+    T tmp_qua = static_cast<T>(tmp > UCHAR_MAX) ? UCHAR_MAX : ((tmp < 0) ? 0 : tmp);
     return tmp_qua;
 }
 
@@ -166,10 +166,6 @@ inline float dequantize(const T *input, const UniformQuantizationInfo &info)
 template <typename T>
 void mul_saturate_quantized_8(const ITensor *in1, const ITensor *in2, ITensor *out, const Window &window, float scale)
 {
-    const UniformQuantizationInfo input1_qua_info = in1->info()->quantization_info().uniform();
-    const UniformQuantizationInfo input2_qua_info = in2->info()->quantization_info().uniform();
-    const UniformQuantizationInfo output_qua_info = out->info()->quantization_info().uniform();
-
     // Create input windows
     Window win        = window;
     Window input1_win = window.broadcast_if_dimension_le_one(in1->info()->tensor_shape());
@@ -177,63 +173,138 @@ void mul_saturate_quantized_8(const ITensor *in1, const ITensor *in2, ITensor *o
 
     // Clear X Dimension on execution window as we handle manually
     win.set(Window::DimX, Window::Dimension(0, 1, 1));
-    input1_win.set(Window::DimX, Window::Dimension(0, 1, 1));
-    input2_win.set(Window::DimX, Window::Dimension(0, 1, 1));
 
-    Iterator input1(in1, input1_win);
-    Iterator input2(in2, input2_win);
-    Iterator output(out, win);
-
-    const int  window_step_x  = 16 / sizeof(T);
-    const auto window_start_x = static_cast<int>(window.x().start());
-    const auto window_end_x   = static_cast<int>(window.x().end());
+    const int  window_step_x         = 16 / sizeof(T);
+    const auto window_start_x        = static_cast<int>(window.x().start());
+    const auto window_end_x          = static_cast<int>(window.x().end());
+    const bool is_broadcast_across_x = (input1_win.x().step() == 0) || (input2_win.x().step() == 0);
 
-    const UniformQuantizationInfo tmp_qua_info = { output_qua_info.scale / scale, output_qua_info.offset };
+    const UniformQuantizationInfo output_qua_info = out->info()->quantization_info().uniform();
+    const UniformQuantizationInfo tmp_qua_info    = { output_qua_info.scale / scale, output_qua_info.offset };
 
-    execute_window_loop(win, [&](const Coordinates &)
+    if(is_broadcast_across_x)
     {
-        const auto input1_ptr = reinterpret_cast<const T *>(input1.ptr());
-        const auto input2_ptr = reinterpret_cast<const T *>(input2.ptr());
-        const auto output_ptr = reinterpret_cast<T *>(output.ptr());
+        const bool                    is_broadcast_input_2 = input2_win.x().step() == 0;
+        Window                        broadcast_win        = is_broadcast_input_2 ? input2_win : input1_win;
+        Window                        non_broadcast_win    = !is_broadcast_input_2 ? input2_win : input1_win;
+        const ITensor                *broadcast_tensor     = is_broadcast_input_2 ? in2 : in1;
+        const ITensor                *non_broadcast_tensor = !is_broadcast_input_2 ? in2 : in1;
+        const UniformQuantizationInfo broadcast_qinfo      = broadcast_tensor->info()->quantization_info().uniform();
+        const UniformQuantizationInfo non_broadcast_qinfo  = non_broadcast_tensor->info()->quantization_info().uniform();
 
-        // Compute window_step_x elements per iteration
-        int x = window_start_x;
-        for(; x <= (window_end_x - window_step_x); x += window_step_x)
+        // Clear X Dimension on execution window as we handle manually
+        non_broadcast_win.set(Window::DimX, Window::Dimension(0, 1, 1));
+
+        Iterator broadcast_input(broadcast_tensor, broadcast_win);
+        Iterator non_broadcast_input(non_broadcast_tensor, non_broadcast_win);
+        Iterator output(out, win);
+
+        using ExactTagType = typename wrapper::traits::neon_vector<T, window_step_x>::tag_type;
+
+        execute_window_loop(win, [&](const Coordinates &)
         {
-            const auto input1_q = wrapper::vloadq(input1_ptr + x);
-            const auto input2_q = wrapper::vloadq(input2_ptr + x);
+            const auto non_broadcast_input_ptr = reinterpret_cast<const T *>(non_broadcast_input.ptr());
+            const auto output_ptr              = reinterpret_cast<T *>(output.ptr());
 
-            // Dequantize inputs
-            const float32x4x4_t in1_f32x4x4 = vdequantize(input1_q, input1_qua_info);
-            const float32x4x4_t in2_f32x4x4 = vdequantize(input2_q, input2_qua_info);
+            const auto broadcast_value     = *reinterpret_cast<const T *>(broadcast_input.ptr());
+            const auto broadcast_value_vec = wrapper::vdup_n(broadcast_value, ExactTagType{});
 
-            const float32x4x4_t out_f32x4x4 =
+            // Compute window_step_x elements per iteration
+            int x = window_start_x;
+            for(; x <= (window_end_x - window_step_x); x += window_step_x)
             {
-                vmulq_f32(in1_f32x4x4.val[0], in2_f32x4x4.val[0]),
-                vmulq_f32(in1_f32x4x4.val[1], in2_f32x4x4.val[1]),
-                vmulq_f32(in1_f32x4x4.val[2], in2_f32x4x4.val[2]),
-                vmulq_f32(in1_f32x4x4.val[3], in2_f32x4x4.val[3]),
-            };
+                const auto non_broadcast_v = wrapper::vloadq(non_broadcast_input_ptr + x);
 
-            // Quantize output
-            const auto result = vquantize<T>(out_f32x4x4, tmp_qua_info);
-            wrapper::vstore(output_ptr + x, result);
-        }
+                // Dequantize inputs
+                const float32x4x4_t in1_f32x4x4 = vdequantize(non_broadcast_v, non_broadcast_qinfo);
+                const float32x4x4_t in2_f32x4x4 = vdequantize(broadcast_value_vec, broadcast_qinfo);
 
-        // Compute left-over elements
-        for(; x < window_end_x; ++x)
+                const float32x4x4_t out_f32x4x4 =
+                {
+                    vmulq_f32(in1_f32x4x4.val[0], in2_f32x4x4.val[0]),
+                    vmulq_f32(in1_f32x4x4.val[1], in2_f32x4x4.val[1]),
+                    vmulq_f32(in1_f32x4x4.val[2], in2_f32x4x4.val[2]),
+                    vmulq_f32(in1_f32x4x4.val[3], in2_f32x4x4.val[3]),
+                };
+
+                // Quantize output
+                const auto result = vquantize<T>(out_f32x4x4, tmp_qua_info);
+                wrapper::vstore(output_ptr + x, result);
+            }
+
+            // Compute left-over elements
+            for(; x < window_end_x; ++x)
+            {
+                // Dequantize inputs
+                float tmp_in1 = dequantize(non_broadcast_input_ptr + x, non_broadcast_qinfo);
+                float tmp_in2 = dequantize(&broadcast_value, broadcast_qinfo);
+                float tmp_f   = tmp_in1 * tmp_in2;
+
+                // Quantize output
+                const auto tmp_qua = quantize<T>(tmp_f, tmp_qua_info);
+                *(output_ptr + x)  = tmp_qua;
+            }
+        },
+        broadcast_input, non_broadcast_input, output);
+    }
+    else
+    {
+        const UniformQuantizationInfo input1_qua_info = in1->info()->quantization_info().uniform();
+        const UniformQuantizationInfo input2_qua_info = in2->info()->quantization_info().uniform();
+
+        // Clear X Dimension on execution window as we handle manually
+        input1_win.set(Window::DimX, Window::Dimension(0, 1, 1));
+        input2_win.set(Window::DimX, Window::Dimension(0, 1, 1));
+
+        Iterator input1(in1, input1_win);
+        Iterator input2(in2, input2_win);
+        Iterator output(out, win);
+
+        execute_window_loop(win, [&](const Coordinates &)
         {
-            // Dequantize inputs
-            float tmp_in1 = dequantize(input1_ptr + x, input1_qua_info);
-            float tmp_in2 = dequantize(input2_ptr + x, input2_qua_info);
-            float tmp_f   = tmp_in1 * tmp_in2;
+            const auto input1_ptr = reinterpret_cast<const T *>(input1.ptr());
+            const auto input2_ptr = reinterpret_cast<const T *>(input2.ptr());
+            const auto output_ptr = reinterpret_cast<T *>(output.ptr());
 
-            // Quantize output
-            const auto tmp_qua = quantize<T>(tmp_f, tmp_qua_info);
-            *(output_ptr + x)  = tmp_qua;
-        }
-    },
-    input1, input2, output);
+            // Compute window_step_x elements per iteration
+            int x = window_start_x;
+            for(; x <= (window_end_x - window_step_x); x += window_step_x)
+            {
+                const auto input1_q = wrapper::vloadq(input1_ptr + x);
+                const auto input2_q = wrapper::vloadq(input2_ptr + x);
+
+                // Dequantize inputs
+                const float32x4x4_t in1_f32x4x4 = vdequantize(input1_q, input1_qua_info);
+                const float32x4x4_t in2_f32x4x4 = vdequantize(input2_q, input2_qua_info);
+
+                const float32x4x4_t out_f32x4x4 =
+                {
+                    vmulq_f32(in1_f32x4x4.val[0], in2_f32x4x4.val[0]),
+                    vmulq_f32(in1_f32x4x4.val[1], in2_f32x4x4.val[1]),
+                    vmulq_f32(in1_f32x4x4.val[2], in2_f32x4x4.val[2]),
+                    vmulq_f32(in1_f32x4x4.val[3], in2_f32x4x4.val[3]),
+                };
+
+                // Quantize output
+                const auto result = vquantize<T>(out_f32x4x4, tmp_qua_info);
+                wrapper::vstore(output_ptr + x, result);
+            }
+
+            // Compute left-over elements
+            for(; x < window_end_x; ++x)
+            {
+                // Dequantize inputs
+                float tmp_in1 = dequantize(input1_ptr + x, input1_qua_info);
+                float tmp_in2 = dequantize(input2_ptr + x, input2_qua_info);
+                float tmp_f   = tmp_in1 * tmp_in2;
+
+                // Quantize output
+                const auto tmp_qua = quantize<T>(tmp_f, tmp_qua_info);
+                *(output_ptr + x)  = tmp_qua;
+            }
+        },
+        input1, input2, output);
+    }
 }
 
 void mul_saturate_QSYMM16_QSYMM16_QSYMM16(const ITensor *in1, const ITensor *in2, ITensor *out, const Window &window, float scale)
@@ -677,10 +748,6 @@ void mul_F32_F32_F32(const ITensor *in1, const ITensor *in2, ITensor *out, const
     const auto    window_end_x          = static_cast<int>(window.x().end());
     const bool    is_broadcast_across_x = (input1_win.x().step() == 0) || (input2_win.x().step() == 0);
 
-    Iterator input1(in1, window.broadcast_if_dimension_le_one(in1->info()->tensor_shape()));
-    Iterator input2(in2, window.broadcast_if_dimension_le_one(in2->info()->tensor_shape()));
-    Iterator output(out, window);
-
     using ExactTagType = typename wrapper::traits::neon_vector<float, window_step_x>::tag_type;
 
     if(is_broadcast_across_x)
diff --git a/tests/validation/NEON/PixelWiseMultiplication.cpp b/tests/validation/NEON/PixelWiseMultiplication.cpp
index c806b23255..0b88628912 100644
--- a/tests/validation/NEON/PixelWiseMultiplication.cpp
+++ b/tests/validation/NEON/PixelWiseMultiplication.cpp
@@ -113,10 +113,12 @@ using NEPixelWiseMultiplicationToS16Fixture = PixelWiseMultiplicationValidationF
 template <typename T>
 using NEPixelWiseMultiplicationToF16Fixture = PixelWiseMultiplicationValidationFixture<Tensor, Accessor, NEPixelWiseMultiplication, T, half_float::half>;
 template <typename T>
-using NEPixelWiseMultiplicationToF32Fixture = PixelWiseMultiplicationValidationFixture<Tensor, Accessor, NEPixelWiseMultiplication, T, float>;
-template <typename T>
-using NEPixelWiseMultiplicationBroadcastFixture = PixelWiseMultiplicationBroadcastValidationFixture<Tensor, Accessor, NEPixelWiseMultiplication, T, float>;
+using NEPixelWiseMultiplicationToF32Fixture     = PixelWiseMultiplicationValidationFixture<Tensor, Accessor, NEPixelWiseMultiplication, T, float>;
 using NEPixelWiseMultiplicationU8U8ToS16Fixture = PixelWiseMultiplicationValidationFixture<Tensor, Accessor, NEPixelWiseMultiplication, uint8_t, uint8_t, int16_t>;
+template <typename T>
+using NEPixelWiseMultiplicationBroadcastFixture              = PixelWiseMultiplicationBroadcastValidationFixture<Tensor, Accessor, NEPixelWiseMultiplication, T, float>;
+using NEPixelWiseMultiplicationBroadcastQASYMM8Fixture       = PixelWiseMultiplicationBroadcastValidationQuantizedFixture<Tensor, Accessor, NEPixelWiseMultiplication, uint8_t, uint8_t>;
+using NEPixelWiseMultiplicationBroadcastQASYMM8SignedFixture = PixelWiseMultiplicationBroadcastValidationQuantizedFixture<Tensor, Accessor, NEPixelWiseMultiplication, int8_t, int8_t>;
 
 TEST_SUITE(NEON)
 TEST_SUITE(PixelWiseMultiplication)
@@ -317,6 +319,21 @@ FIXTURE_DATA_TEST_CASE(RunSmall, NEPixelWiseMultiplicationQASYMM8Fixture, framew
     validate(Accessor(_target), _reference, tolerance_qasymm8);
 }
 TEST_SUITE_END() // ScaleOther
+TEST_SUITE(Broadcast)
+FIXTURE_DATA_TEST_CASE(RunSmall, NEPixelWiseMultiplicationBroadcastQASYMM8Fixture, framework::DatasetMode::ALL,
+                       combine(combine(combine(combine(combine(combine(combine(datasets::SmallShapesBroadcast(),
+                                                                               framework::dataset::make("DataTypeIn1", DataType::QASYMM8)),
+                                                                       framework::dataset::make("DataTypeIn2", DataType::QASYMM8)),
+                                                               framework::dataset::make("DataTypeOut", DataType::QASYMM8)),
+                                                       framework::dataset::make("Scale", { scale_other })),
+                                               PixelWiseMultiplicationPolicySTZDataset),
+                                       PixelWiseMultiplicationQASYMM8QuantDataset),
+                               framework::dataset::make("InPlace", { false })))
+{
+    // Validate output
+    validate(Accessor(_target), _reference, tolerance_qasymm8);
+}
+TEST_SUITE_END() // Broadcast
 TEST_SUITE_END() // QASYMM8
 TEST_SUITE(QSYMM16)
 TEST_SUITE(Scale255)
diff --git a/tests/validation/fixtures/PixelWiseMultiplicationFixture.h b/tests/validation/fixtures/PixelWiseMultiplicationFixture.h
index c07725ce53..4eb83859ac 100644
--- a/tests/validation/fixtures/PixelWiseMultiplicationFixture.h
+++ b/tests/validation/fixtures/PixelWiseMultiplicationFixture.h
@@ -201,6 +201,19 @@ public:
                                                                                                                    qinfo0, qinfo1, qinfo_out, ActivationLayerInfo(), is_inplace);
     }
 };
+
+template <typename TensorType, typename AccessorType, typename FunctionType, typename T1, typename T2, typename T3 = T2>
+class PixelWiseMultiplicationBroadcastValidationQuantizedFixture : public PixelWiseMultiplicationGenericValidationFixture<TensorType, AccessorType, FunctionType, T1, T2, T3>
+{
+public:
+    template <typename...>
+    void setup(const TensorShape &shape0, const TensorShape &shape1, DataType dt_in1, DataType dt_in2, DataType dt_out, float scale, ConvertPolicy convert_policy, RoundingPolicy rounding_policy,
+               QuantizationInfo qinfo0, QuantizationInfo qinfo1, QuantizationInfo qinfo_out, bool is_inplace)
+    {
+        PixelWiseMultiplicationGenericValidationFixture<TensorType, AccessorType, FunctionType, T1, T2, T3>::setup(shape0, shape1, dt_in1, dt_in2, dt_out, scale, convert_policy, rounding_policy,
+                                                                                                                   qinfo0, qinfo1, qinfo_out, ActivationLayerInfo(), is_inplace);
+    }
+};
 } // namespace validation
 } // namespace test
 } // namespace arm_compute