NEFuseBatchNormalizationKernel rework

* Enable fp16 in armv8a multi_isa builds

* Code guarded with __ARM_FEATURE_FP16_VECTOR_ARITHMETIC needs
  to be moved to an fp16.cpp file to allow compilation with
  -march=armv8.2-a+fp16

* fp16.cpp needs to use the template add_same_neon() so it had to be moved from
  impl.cpp to impl.h

* Partially resolves MLCE-1102

Change-Id: Ia51007f5e663b708071958bb94bfab4535e4b2f8
Signed-off-by: Pablo Marquez Tello <pablo.tello@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/10191
Benchmark: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Viet-Hoa Do <viet-hoa.do@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>

4 files changed, 95 insertions, 140 deletions

diff --git a/src/BUILD.bazel b/src/BUILD.bazel
index 7995bb5736..dbedffae31 100644
--- a/src/BUILD.bazel
+++ b/src/BUILD.bazel
@@ -763,7 +763,6 @@ filegroup(
 	"cpu/kernels/floor/neon/fp32.cpp",
 	"cpu/kernels/fuse_batch_normalization/generic/fp16.cpp",
 	"cpu/kernels/fuse_batch_normalization/generic/fp32.cpp",
-	"cpu/kernels/fuse_batch_normalization/generic/impl.cpp",
 	"cpu/kernels/fuse_batch_normalization/nchw/all.cpp",
 	"cpu/kernels/fuse_batch_normalization/nhwc/neon/fp16.cpp",
 	"cpu/kernels/fuse_batch_normalization/nhwc/neon/fp32.cpp",
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 4ce3d3c2df..2a592090ba 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -755,7 +755,6 @@ target_sources(
 	cpu/kernels/floor/neon/fp32.cpp
 	cpu/kernels/fuse_batch_normalization/generic/fp16.cpp
 	cpu/kernels/fuse_batch_normalization/generic/fp32.cpp
-	cpu/kernels/fuse_batch_normalization/generic/impl.cpp
 	cpu/kernels/fuse_batch_normalization/nchw/all.cpp
 	cpu/kernels/fuse_batch_normalization/nhwc/neon/fp16.cpp
 	cpu/kernels/fuse_batch_normalization/nhwc/neon/fp32.cpp
diff --git a/src/cpu/kernels/fuse_batch_normalization/generic/impl.cpp b/src/cpu/kernels/fuse_batch_normalization/generic/impl.cpp
deleted file mode 100644
index 3c6a2069ee..0000000000
--- a/src/cpu/kernels/fuse_batch_normalization/generic/impl.cpp
+++ /dev/null
@@ -1,135 +0,0 @@
-/*
- * Copyright (c) 2018-2022 Arm Limited.
- *
- * SPDX-License-Identifier: MIT
- *
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to
- * deal in the Software without restriction, including without limitation the
- * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
- * sell copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
- */
-#include "src/cpu/kernels/fuse_batch_normalization/generic/impl.h"
-
-namespace arm_compute
-{
-namespace cpu
-{
-template <typename T>
-void fused_batch_normalization_conv(const ITensor *conv_weights, const ITensor *conv_bias, ITensor *fused_weights, ITensor *fused_bias,
-                                    const ITensor *bn_mean, const ITensor *bn_var, const ITensor *bn_beta, const ITensor *bn_gamma, float epsilon, const Window &window)
-{
-    using ScalarType   = T;
-    const int size     = 16 / conv_weights->info()->element_size();
-    using ExactTagType = typename wrapper::traits::neon_bitvector_tag_t<T, wrapper::traits::BitWidth::W128>;
-
-    const bool run_in_place_weights = (fused_weights == nullptr) || (fused_weights == conv_weights);
-    const bool run_in_place_bias    = (fused_bias == nullptr) || (conv_bias != nullptr && fused_bias == conv_bias);
-
-    // Set build options
-    Window win = window;
-    win.set(Window::DimX, Window::Dimension(0, 1, 1));
-
-    const int  window_step_x  = size;
-    const auto window_start_x = static_cast<int>(window.x().start());
-    const auto window_end_x   = static_cast<int>(window.x().end());
-
-    Iterator conv_w_in(conv_weights, win);
-    Iterator conv_w_out(run_in_place_weights ? conv_weights : fused_weights, win);
-
-    const auto conv_bias_in  = (conv_bias != nullptr ? reinterpret_cast<ScalarType *>(conv_bias->ptr_to_element(Coordinates(0, 0))) : nullptr);
-    auto       conv_bias_out = (run_in_place_bias ? conv_bias_in : reinterpret_cast<ScalarType *>(fused_bias->ptr_to_element(Coordinates(0, 0))));
-
-    const auto input_mean  = reinterpret_cast<const ScalarType *>(bn_mean->ptr_to_element(Coordinates(0, 0)));
-    const auto input_var   = reinterpret_cast<const ScalarType *>(bn_var->ptr_to_element(Coordinates(0, 0)));
-    const auto input_gamma = (bn_gamma != nullptr) ? reinterpret_cast<const ScalarType *>(bn_gamma->ptr_to_element(Coordinates(0, 0))) : nullptr;
-    const auto input_beta  = (bn_beta != nullptr) ? reinterpret_cast<const ScalarType *>(bn_beta->ptr_to_element(Coordinates(0, 0))) : nullptr;
-
-    auto       mean_vec    = wrapper::vdup_n(ScalarType(0), ExactTagType{});
-    auto       var_vec     = wrapper::vdup_n(ScalarType(0), ExactTagType{});
-    auto       gamma_vec   = wrapper::vdup_n(ScalarType(1), ExactTagType{});
-    auto       beta_vec    = wrapper::vdup_n(ScalarType(0), ExactTagType{});
-    auto       rvar_vec    = wrapper::vdup_n(ScalarType(0), ExactTagType{});
-    const auto epsilon_vec = wrapper::vdup_n(ScalarType(epsilon), ExactTagType{});
-
-    auto mean                = ScalarType(0.0);
-    auto var                 = ScalarType(0.0);
-    auto gamma               = ScalarType(1.0);
-    auto beta                = ScalarType(0.0);
-    auto conv_bias_in_scalar = ScalarType(0.0);
-    execute_window_loop(win, [&](const Coordinates & id)
-    {
-        var = input_var[id[3]];
-        if(input_gamma != nullptr)
-        {
-            gamma = input_gamma[id[3]];
-        }
-
-        if((id[0] == 0) && (id[1] == 0) && (id[2] == 0))
-        {
-            if(input_beta != nullptr)
-            {
-                beta     = input_beta[id[3]];
-                beta_vec = wrapper::vdup_n(beta, ExactTagType{});
-            }
-
-            // Construct vectors
-            mean     = input_mean[id[3]];
-            mean_vec = wrapper::vdup_n(mean, ExactTagType{});
-
-            if(conv_bias_in != nullptr)
-            {
-                conv_bias_in_scalar = conv_bias_in[id[3]];
-            }
-            auto conv_bias_tmp_scalar = (conv_bias_in_scalar - mean) / std::sqrt(var + ScalarType(epsilon));
-            conv_bias_out[id[3]]      = (conv_bias_tmp_scalar * gamma) + beta;
-        }
-
-        int  x              = window_start_x;
-        auto conv_w_in_ptr  = reinterpret_cast<const ScalarType *>(conv_w_in.ptr());
-        auto conv_w_out_ptr = reinterpret_cast<ScalarType *>(conv_w_out.ptr());
-        var_vec             = wrapper::vdup_n(var, ExactTagType{});
-        gamma_vec           = wrapper::vdup_n(gamma, ExactTagType{});
-        rvar_vec            = wrapper::vinvsqrt(wrapper::vadd(var_vec, epsilon_vec));
-
-        for(; x <= (window_end_x - window_step_x); x += window_step_x)
-        {
-            auto wn = wrapper::vloadq(conv_w_in_ptr + x);
-            wn      = wrapper::vmul(wn, rvar_vec);
-            wn      = wrapper::vmul(wn, gamma_vec);
-
-            // Store results
-            wrapper::vstore(conv_w_out_ptr + x, wn);
-        }
-
-        // Compute left-over elements
-        for(; x < window_end_x; ++x)
-        {
-            *(conv_w_out_ptr + x) = *(conv_w_in_ptr + x) / std::sqrt(var + ScalarType(epsilon)) * gamma;
-        }
-    },
-    conv_w_in, conv_w_out);
-}
-
-template void fused_batch_normalization_conv<float32_t>(const ITensor *conv_weights, const ITensor *conv_bias, ITensor *fused_weights, ITensor *fused_bias,
-                                                        const ITensor *bn_mean, const ITensor *bn_var, const ITensor *bn_beta, const ITensor *bn_gamma, float epsilon, const Window &window);
-
-#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && defined(ENABLE_FP16_KERNELS)
-template void fused_batch_normalization_conv<float16_t>(const ITensor *conv_weights, const ITensor *conv_bias, ITensor *fused_weights, ITensor *fused_bias,
-                                                        const ITensor *bn_mean, const ITensor *bn_var, const ITensor *bn_beta, const ITensor *bn_gamma, float epsilon, const Window &window);
-#endif /* defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && defined(ENABLE_FP16_KERNELS) */
-
-} // namespace cpu
-} // namespace arm_compute
diff --git a/src/cpu/kernels/fuse_batch_normalization/generic/impl.h b/src/cpu/kernels/fuse_batch_normalization/generic/impl.h
index 979ea13842..b9017600d6 100644
--- a/src/cpu/kernels/fuse_batch_normalization/generic/impl.h
+++ b/src/cpu/kernels/fuse_batch_normalization/generic/impl.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022 Arm Limited.
+ * Copyright (c) 2021-2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -33,7 +33,99 @@ namespace cpu
 {
 template <typename T>
 void fused_batch_normalization_conv(const ITensor *conv_weights, const ITensor *conv_bias, ITensor *fused_weights, ITensor *fused_bias,
-                                    const ITensor *bn_mean, const ITensor *bn_var, const ITensor *bn_beta, const ITensor *bn_gamma, float epsilon, const Window &window);
+                                    const ITensor *bn_mean, const ITensor *bn_var, const ITensor *bn_beta, const ITensor *bn_gamma, float epsilon, const Window &window)
+{
+    using ScalarType   = T;
+    const int size     = 16 / conv_weights->info()->element_size();
+    using ExactTagType = typename wrapper::traits::neon_bitvector_tag_t<T, wrapper::traits::BitWidth::W128>;
+
+    const bool run_in_place_weights = (fused_weights == nullptr) || (fused_weights == conv_weights);
+    const bool run_in_place_bias    = (fused_bias == nullptr) || (conv_bias != nullptr && fused_bias == conv_bias);
+
+    // Set build options
+    Window win = window;
+    win.set(Window::DimX, Window::Dimension(0, 1, 1));
+
+    const int  window_step_x  = size;
+    const auto window_start_x = static_cast<int>(window.x().start());
+    const auto window_end_x   = static_cast<int>(window.x().end());
+
+    Iterator conv_w_in(conv_weights, win);
+    Iterator conv_w_out(run_in_place_weights ? conv_weights : fused_weights, win);
+
+    const auto conv_bias_in  = (conv_bias != nullptr ? reinterpret_cast<ScalarType *>(conv_bias->ptr_to_element(Coordinates(0, 0))) : nullptr);
+    auto       conv_bias_out = (run_in_place_bias ? conv_bias_in : reinterpret_cast<ScalarType *>(fused_bias->ptr_to_element(Coordinates(0, 0))));
+
+    const auto input_mean  = reinterpret_cast<const ScalarType *>(bn_mean->ptr_to_element(Coordinates(0, 0)));
+    const auto input_var   = reinterpret_cast<const ScalarType *>(bn_var->ptr_to_element(Coordinates(0, 0)));
+    const auto input_gamma = (bn_gamma != nullptr) ? reinterpret_cast<const ScalarType *>(bn_gamma->ptr_to_element(Coordinates(0, 0))) : nullptr;
+    const auto input_beta  = (bn_beta != nullptr) ? reinterpret_cast<const ScalarType *>(bn_beta->ptr_to_element(Coordinates(0, 0))) : nullptr;
+
+    auto       mean_vec    = wrapper::vdup_n(ScalarType(0), ExactTagType{});
+    auto       var_vec     = wrapper::vdup_n(ScalarType(0), ExactTagType{});
+    auto       gamma_vec   = wrapper::vdup_n(ScalarType(1), ExactTagType{});
+    auto       beta_vec    = wrapper::vdup_n(ScalarType(0), ExactTagType{});
+    auto       rvar_vec    = wrapper::vdup_n(ScalarType(0), ExactTagType{});
+    const auto epsilon_vec = wrapper::vdup_n(ScalarType(epsilon), ExactTagType{});
+
+    auto mean                = ScalarType(0.0);
+    auto var                 = ScalarType(0.0);
+    auto gamma               = ScalarType(1.0);
+    auto beta                = ScalarType(0.0);
+    auto conv_bias_in_scalar = ScalarType(0.0);
+    execute_window_loop(win, [&](const Coordinates & id)
+    {
+        var = input_var[id[3]];
+        if(input_gamma != nullptr)
+        {
+            gamma = input_gamma[id[3]];
+        }
+
+        if((id[0] == 0) && (id[1] == 0) && (id[2] == 0))
+        {
+            if(input_beta != nullptr)
+            {
+                beta     = input_beta[id[3]];
+                beta_vec = wrapper::vdup_n(beta, ExactTagType{});
+            }
+
+            // Construct vectors
+            mean     = input_mean[id[3]];
+            mean_vec = wrapper::vdup_n(mean, ExactTagType{});
+
+            if(conv_bias_in != nullptr)
+            {
+                conv_bias_in_scalar = conv_bias_in[id[3]];
+            }
+            auto conv_bias_tmp_scalar = (conv_bias_in_scalar - mean) / std::sqrt(var + ScalarType(epsilon));
+            conv_bias_out[id[3]]      = (conv_bias_tmp_scalar * gamma) + beta;
+        }
+
+        int  x              = window_start_x;
+        auto conv_w_in_ptr  = reinterpret_cast<const ScalarType *>(conv_w_in.ptr());
+        auto conv_w_out_ptr = reinterpret_cast<ScalarType *>(conv_w_out.ptr());
+        var_vec             = wrapper::vdup_n(var, ExactTagType{});
+        gamma_vec           = wrapper::vdup_n(gamma, ExactTagType{});
+        rvar_vec            = wrapper::vinvsqrt(wrapper::vadd(var_vec, epsilon_vec));
+
+        for(; x <= (window_end_x - window_step_x); x += window_step_x)
+        {
+            auto wn = wrapper::vloadq(conv_w_in_ptr + x);
+            wn      = wrapper::vmul(wn, rvar_vec);
+            wn      = wrapper::vmul(wn, gamma_vec);
+
+            // Store results
+            wrapper::vstore(conv_w_out_ptr + x, wn);
+        }
+
+        // Compute left-over elements
+        for(; x < window_end_x; ++x)
+        {
+            *(conv_w_out_ptr + x) = *(conv_w_in_ptr + x) / std::sqrt(var + ScalarType(epsilon)) * gamma;
+        }
+    },
+    conv_w_in, conv_w_out);
+}
 }
 }
-#endif //SRC_CORE_NEON_KERNELS_FUSE_BATCH_NORMALIZATION_GENERIC_IMPL_H
\ No newline at end of file
+#endif //SRC_CORE_NEON_KERNELS_FUSE_BATCH_NORMALIZATION_GENERIC_IMPL_H