BatchNorm changes to enable fp16 in armv8a multi_isa builds

    * Moved NCHW kernels fp16 and fp32 to their corresponding files
      src/cpu/kernels/fuse_batch_normalization/nchw/neon/fp16.cpp and
      src/cpu/kernels/fuse_batch_normalization/nchw/neon/fp32.cpp

    * Changes in filelist.json to include the new fp16 and fp32 files

    * Moved the template batch_normalization_nchw to impl.h as we
      need to instantiate it from fp16.cpp and fp32.cpp

    * Pooling layer: removed the guard __ARM_FEATURE_FP16_VECTOR_ARITHMETIC that
      prevented the FP16 kernel execution.

    * Partially resolves MLCE-1102

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

3 files changed, 43 insertions, 156 deletions

diff --git a/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.cpp b/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.cpp
index deb89996a9..717fd11485 100644
--- a/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.cpp
+++ b/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021 Arm Limited.
+ * Copyright (c) 2017-2021, 2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -151,128 +151,15 @@ Status validate_arguments(const ITensorInfo  *input,
 }
 } //namespace
 
-template <typename T, bool fused_activation, typename F>
-void NEBatchNormalizationLayerKernel::batch_normalization_nchw(const Window &window)
-{
-    /** SIMD vector tag type. */
-    using ExactTagType = typename wrapper::traits::neon_bitvector_tag_t<T, wrapper::traits::BitWidth::W128>;
-
-    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());
-
-    Window win_to_use = window;
-    win_to_use.set(Window::DimX, Window::Dimension(0, 1, 1));
-
-    Iterator input(_input, win_to_use);
-    Iterator output(_output, win_to_use);
-
-    F activation_functor(_act_info);
-
-    // Hold information about the current feature map we are iterating.
-    // Only compute denominator and constants once per feature map.
-    int slice = -1;
-
-    const auto input_mean = reinterpret_cast<const T *>(_mean->ptr_to_element(Coordinates(0, 0)));
-    const auto input_var  = reinterpret_cast<const T *>(_var->ptr_to_element(Coordinates(0, 0)));
-    const auto input_gamma =
-        (_gamma != nullptr) ? reinterpret_cast<const T *>(_gamma->ptr_to_element(Coordinates(0, 0))) : nullptr;
-    const auto input_beta =
-        (_beta != nullptr) ? reinterpret_cast<const T *>(_beta->ptr_to_element(Coordinates(0, 0))) : nullptr;
-
-    T mean        = static_cast<T>(0);
-    T var         = static_cast<T>(0);
-    T gamma       = static_cast<T>(1);
-    T beta        = static_cast<T>(0);
-    T denominator = static_cast<T>(0);
-
-    auto       mean_vec        = wrapper::vdup_n(mean, ExactTagType{});
-    auto       var_vec         = wrapper::vdup_n(var, ExactTagType{});
-    auto       gamma_vec       = wrapper::vdup_n(gamma, ExactTagType{});
-    auto       beta_vec        = wrapper::vdup_n(beta, ExactTagType{});
-    auto       denominator_vec = wrapper::vdup_n(denominator, ExactTagType{});
-    const auto epsilon_vec     = wrapper::vdup_n(static_cast<T>(_epsilon), ExactTagType{});
-    execute_window_loop(
-        win_to_use,
-        [&](const Coordinates &id)
-        {
-            const auto input_ptr  = reinterpret_cast<const T *>(input.ptr());
-            const auto output_ptr = reinterpret_cast<T *>(output.ptr());
-
-            if (slice != id.z())
-            {
-                mean     = input_mean[id.z()];
-                var      = input_var[id.z()];
-                mean_vec = wrapper::vdup_n(mean, ExactTagType{});
-                var_vec  = wrapper::vdup_n(var, ExactTagType{});
-                if (input_gamma != nullptr)
-                {
-                    gamma     = input_gamma[id.z()];
-                    gamma_vec = wrapper::vdup_n(gamma, ExactTagType{});
-                }
-                if (input_beta != nullptr)
-                {
-                    beta     = input_beta[id.z()];
-                    beta_vec = wrapper::vdup_n(beta, ExactTagType{});
-                }
-
-                // Calculate denominator
-                denominator_vec = wrapper::vinvsqrt(wrapper::vadd(var_vec, epsilon_vec));
-                denominator     = wrapper::vgetlane(denominator_vec, 0);
-                slice           = id.z();
-            }
-
-            // Perform core calculations using vector operations
-            int x = window_start_x;
-            for (; x <= (window_end_x - window_step_x); x += window_step_x)
-            {
-                // Calculate x bar
-                const auto numerator = wrapper::vsub(wrapper::vloadq(input_ptr + x), mean_vec);
-                const auto x_bar     = wrapper::vmul(numerator, denominator_vec);
-                auto       res       = wrapper::vmla(beta_vec, x_bar, gamma_vec);
-
-                // Perform fused activation
-                if (fused_activation)
-                {
-                    activation_functor(res);
-                }
-
-                // Store results
-                wrapper::vstore(output_ptr + x, res);
-            }
-
-            // Compute left-over elements
-            for (; x < window_end_x; ++x)
-            {
-                const T numerator = input_ptr[x] - mean;
-                const T x_bar     = numerator * denominator;
-                T       res       = beta + x_bar * gamma;
-
-                // Perform fused activation
-                if (fused_activation)
-                {
-                    activation_functor(res);
-                }
-
-                // Store results
-                *(output_ptr + x) = res;
-            }
-        },
-        input, output);
-}
-
 void NEBatchNormalizationLayerKernel::configure_non_fused()
 {
     switch (_input->info()->data_type())
     {
-#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
         case DataType::F16:
-            _func = &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float16_t, false,
-                                                                               detail::dummy<float16_t, 8>>;
+            _func = REGISTER_FP16_NEON(cpu::fp16_batch_normalization_nchw_non_fused);
             break;
-#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
         case DataType::F32:
-            _func = &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float, false, detail::dummy<float, 4>>;
+            _func = REGISTER_FP32_NEON(cpu::fp32_batch_normalization_nchw_non_fused);
             break;
         default:
             ARM_COMPUTE_ERROR("Element size not supported");
@@ -285,29 +172,26 @@ void NEBatchNormalizationLayerKernel::configure_fused()
     // NCHW Fused Batched Normalization with activation functions : FP32
     static std::map<ActivationLayerInfo::ActivationFunction, BatchNormFunctionPtr> bn_fused_map_f32_nchw = {
         {ActivationLayerInfo::ActivationFunction::RELU,
-         &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float, true, detail::relu<float, 4>>},
+         REGISTER_FP32_NEON(cpu::fp32_batch_normalization_nchw_non_fused_relu)},
         {ActivationLayerInfo::ActivationFunction::BOUNDED_RELU,
-         &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float, true, detail::brelu<float, 4>>},
+         REGISTER_FP32_NEON(cpu::fp32_batch_normalization_nchw_non_fused_brelu)},
         {ActivationLayerInfo::ActivationFunction::LU_BOUNDED_RELU,
-         &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float, true, detail::lubrelu<float, 4>>}};
-#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+         REGISTER_FP32_NEON(cpu::fp32_batch_normalization_nchw_non_fused_lubrelu)}};
+
     // NCHW Fused Batched Normalization with activation functions : FP16
     static std::map<ActivationLayerInfo::ActivationFunction, BatchNormFunctionPtr> bn_fused_map_f16_nchw = {
         {ActivationLayerInfo::ActivationFunction::RELU,
-         &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float16_t, true, detail::relu<float16_t, 8>>},
+         REGISTER_FP16_NEON(cpu::fp16_batch_normalization_nchw_non_fused_relu)},
         {ActivationLayerInfo::ActivationFunction::BOUNDED_RELU,
-         &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float16_t, true, detail::brelu<float16_t, 8>>},
+         REGISTER_FP16_NEON(cpu::fp16_batch_normalization_nchw_non_fused_brelu)},
         {ActivationLayerInfo::ActivationFunction::LU_BOUNDED_RELU,
-         &NEBatchNormalizationLayerKernel::batch_normalization_nchw<float16_t, true, detail::lubrelu<float16_t, 8>>}};
-#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+         REGISTER_FP16_NEON(cpu::fp16_batch_normalization_nchw_non_fused_lubrelu)}};
 
     switch (_input->info()->data_type())
     {
-#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
         case DataType::F16:
             _func = bn_fused_map_f16_nchw[_act_info.activation()];
             break;
-#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
         case DataType::F32:
             _func = bn_fused_map_f32_nchw[_act_info.activation()];
             break;
@@ -409,7 +293,7 @@ void NEBatchNormalizationLayerKernel::run(const Window &window, const ThreadInfo
     const bool is_nchw = _input->info()->data_layout() == DataLayout::NCHW;
     if (is_nchw)
     {
-        (this->*_func)(window);
+        (*_func)(window, _input, _output, _mean, _var, _beta, _gamma, _epsilon, _act_info);
     }
     else
     {
diff --git a/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.h b/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.h
index 2e8ff0dc9a..679ade0fae 100644
--- a/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.h
+++ b/src/core/NEON/kernels/NEBatchNormalizationLayerKernel.h
@@ -21,8 +21,8 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  */
-#ifndef ARM_COMPUTE_NEBATCHNORMALIZATIONLAYERKERNEL_H
-#define ARM_COMPUTE_NEBATCHNORMALIZATIONLAYERKERNEL_H
+#ifndef ACL_SRC_CORE_NEON_KERNELS_NEBATCHNORMALIZATIONLAYERKERNEL_H
+#define ACL_SRC_CORE_NEON_KERNELS_NEBATCHNORMALIZATIONLAYERKERNEL_H
 
 #include "arm_compute/function_info/ActivationLayerInfo.h"
 
@@ -110,31 +110,19 @@ private:
     /** Configure execution function in case of fused activation **/
     void configure_fused();
 
-    /** Template function to run batch normalization on fp32
-     *
-     * @tparam T                Specialization data type
-     * @tparam fused_activation Boolean that flags if its a fused activation or not
-     * @tparam F                Activation function functor to run
-     *
-     * @param[in] window Region on which to execute the kernel. (Must be a valid region of the window returned by window()).
-     */
-    template <typename T, bool fused_activation, typename F>
-    void batch_normalization_nchw(const Window &window);
-    /** Template function to run batch normalization on fp32 on tensors with NHWC format
-     *
-     * @tparam T                Specialization data type
-     * @tparam fused_activation Boolean that flags if its a fused activation or not
-     * @tparam F                Activation function functor to run
-     *
-     * @param[in] window Region on which to execute the kernel. (Must be a valid region of the window returned by window()).
-     */
-    template <typename T, bool fused_activation, typename F>
-    void batch_normalization_nhwc(const Window &window);
     /** Common signature for all the batch normalization functions
      *
      * @param[in] window Region on which to execute the kernel.
      */
-    using BatchNormFunctionPtr = void (NEBatchNormalizationLayerKernel::*)(const Window &window);
+    using BatchNormFunctionPtr = void (*)(const Window       &window,
+                                          ITensor            *input,
+                                          ITensor            *output,
+                                          const ITensor      *mean,
+                                          const ITensor      *var,
+                                          const ITensor      *beta,
+                                          const ITensor      *gamma,
+                                          float               epsilon,
+                                          ActivationLayerInfo act_info);
 
 private:
     BatchNormFunctionPtr _func;
@@ -148,4 +136,4 @@ private:
     ActivationLayerInfo  _act_info;
 };
 } // namespace arm_compute
-#endif /*ARM_COMPUTE_NEBATCHNORMALIZATIONLAYERKERNEL_H */
+#endif // ACL_SRC_CORE_NEON_KERNELS_NEBATCHNORMALIZATIONLAYERKERNEL_H
diff --git a/src/core/NEON/kernels/batchnormalization/impl/list.h b/src/core/NEON/kernels/batchnormalization/impl/list.h
index cbf540bd71..c619788125 100644
--- a/src/core/NEON/kernels/batchnormalization/impl/list.h
+++ b/src/core/NEON/kernels/batchnormalization/impl/list.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020 Arm Limited.
+ * Copyright (c) 2020, 2023 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -21,8 +21,8 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  */
-#ifndef SRC_CORE_NEON_KERNELS_BATCH_NORMALIZATION_LIST_H
-#define SRC_CORE_NEON_KERNELS_BATCH_NORMALIZATION_LIST_H
+#ifndef ACL_SRC_CORE_NEON_KERNELS_BATCHNORMALIZATION_IMPL_LIST_H
+#define ACL_SRC_CORE_NEON_KERNELS_BATCHNORMALIZATION_IMPL_LIST_H
 
 namespace arm_compute
 {
@@ -37,8 +37,23 @@ DECLARE_BATCH_NORMALIZATION_KERNEL(fp16_sve_batch_normalization);
 DECLARE_BATCH_NORMALIZATION_KERNEL(fp32_neon_batch_normalization);
 DECLARE_BATCH_NORMALIZATION_KERNEL(fp32_sve_batch_normalization);
 
-#undef DECLARE_ACTIVATION_KERNEL
+#define DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(func_name)                                                         \
+    void func_name(const Window &window, ITensor *input, ITensor *output, const ITensor *mean, const ITensor *var, \
+                   const ITensor *beta, const ITensor *gamma, float epsilon, ActivationLayerInfo act_info)
+
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp16_batch_normalization_nchw_non_fused);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp32_batch_normalization_nchw_non_fused);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp16_batch_normalization_nchw_non_fused_relu);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp16_batch_normalization_nchw_non_fused_brelu);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp16_batch_normalization_nchw_non_fused_lubrelu);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp32_batch_normalization_nchw_non_fused_relu);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp32_batch_normalization_nchw_non_fused_brelu);
+DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL(fp32_batch_normalization_nchw_non_fused_lubrelu);
+
+#undef DECLARE_BATCH_NORMALIZATION_KERNEL
+#undef DECLARE_BATCH_NORMALIZATION_NCHW_KERNEL
+
 } // namespace cpu
 } // namespace arm_compute
 
-#endif /* SRC_CORE_NEON_KERNELS_BATCH_NORMALIZATION_LIST_H */
+#endif // ACL_SRC_CORE_NEON_KERNELS_BATCHNORMALIZATION_IMPL_LIST_H