Decouple fuseBatchNormalizationKernel

- Decouple data type for CPU implementation
    supported data types are: fp32, fp16

Resolves COMPMID-4613
Signed-off-by: Yair Schwarzbaum <yair.schwarzbaum@arm.com>
Change-Id: I8aff3ba2d446f64e4d182a866e3a3debc9ef613b
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/7175
Reviewed-by: Giorgio Arena <giorgio.arena@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>

1 files changed, 118 insertions, 341 deletions

diff --git a/src/core/NEON/kernels/NEFuseBatchNormalizationKernel.cpp b/src/core/NEON/kernels/NEFuseBatchNormalizationKernel.cpp
index 0d3244c409..51a69046a9 100644
--- a/src/core/NEON/kernels/NEFuseBatchNormalizationKernel.cpp
+++ b/src/core/NEON/kernels/NEFuseBatchNormalizationKernel.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2021 Arm Limited.
+ * Copyright (c) 2018-2022 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -22,6 +22,7 @@
  * SOFTWARE.
  */
 #include "src/core/NEON/kernels/NEFuseBatchNormalizationKernel.h"
+#include "src/cpu/kernels/fuse_batch_normalization/list.h"
 
 #include "arm_compute/core/Helpers.h"
 #include "arm_compute/core/ITensor.h"
@@ -29,8 +30,10 @@
 #include "arm_compute/core/Utils.h"
 #include "arm_compute/core/Validate.h"
 #include "arm_compute/core/Window.h"
+#include "src/common/cpuinfo/CpuIsaInfo.h"
 #include "src/core/CPP/Validate.h"
 #include "src/core/NEON/wrapper/wrapper.h"
+#include "src/core/common/Registrars.h"
 #include "src/core/helpers/AutoConfiguration.h"
 #include "src/core/helpers/WindowHelpers.h"
 
@@ -40,6 +43,113 @@ namespace arm_compute
 {
 namespace
 {
+struct FuseBatchNormalizeSelectorData
+{
+    DataType                   dt;
+    DataLayout                 dl;
+    FuseBatchNormalizationType fbn_type;
+    cpuinfo::CpuIsaInfo        isa;
+};
+
+using FBNSelectorPtr = std::add_pointer<bool(const FuseBatchNormalizeSelectorData &data)>::type;
+using FBNUKernelPtr  = std::add_pointer<void(const ITensor *, const ITensor *, ITensor *, ITensor *,
+                                             const ITensor *, const ITensor *, const ITensor *, const ITensor *, float, const Window &)>::type;
+
+struct FBNUKernel
+{
+    const char          *name;
+    const FBNSelectorPtr is_selected;
+    FBNUKernelPtr        ukernel;
+};
+
+static const FBNUKernel available_kernels[] =
+{
+    {
+        "fused_batch_normalization_conv_NHWC_F16",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F16 && data.dl == DataLayout::NHWC && data.isa.fp16 && data.fbn_type == FuseBatchNormalizationType::CONVOLUTION;
+        },
+        REGISTER_FP16_NEON(arm_compute::cpu::fused_batch_normalization_conv_f16)
+    },
+    {
+        "fused_batch_normalization_conv_NCHW_F16",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F16 && data.dl == DataLayout::NCHW && data.isa.fp16 && data.fbn_type == FuseBatchNormalizationType::CONVOLUTION;
+        },
+        REGISTER_FP16_NEON(arm_compute::cpu::fused_batch_normalization_conv_f16)
+    },
+    {
+        "fused_batch_normalization_dwc_NHWC_F16",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F16 && data.dl == DataLayout::NHWC && data.isa.fp16 && data.fbn_type == FuseBatchNormalizationType::DEPTHWISECONVOLUTION;
+        },
+        REGISTER_FP16_NEON(arm_compute::cpu::fused_batch_normalization_dwc_nhwc_f16)
+    },
+    {
+        "fused_batch_normalization_dwc_NCHW_F16",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F16 && data.dl == DataLayout::NCHW && data.isa.fp16 && data.fbn_type == FuseBatchNormalizationType::DEPTHWISECONVOLUTION;
+        },
+        REGISTER_FP16_NEON(arm_compute::cpu::fused_batch_normalization_dwc_nchw_f16)
+    },
+    {
+        "fused_batch_normalization_conv_NHWC_F32",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F32 && data.dl == DataLayout::NHWC && data.fbn_type == FuseBatchNormalizationType::CONVOLUTION;
+        },
+        REGISTER_FP32_NEON(arm_compute::cpu::fused_batch_normalization_conv_f32)
+    },
+    {
+        "fused_batch_normalization_conv_NCHW_F32",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F32 && data.dl == DataLayout::NCHW && data.fbn_type == FuseBatchNormalizationType::CONVOLUTION;
+        },
+        REGISTER_FP32_NEON(arm_compute::cpu::fused_batch_normalization_conv_f32)
+    },
+    {
+        "fused_batch_normalization_dwc_NHWC_F32",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F32 && data.dl == DataLayout::NHWC && data.fbn_type == FuseBatchNormalizationType::DEPTHWISECONVOLUTION;
+        },
+        REGISTER_FP32_NEON(arm_compute::cpu::fused_batch_normalization_dwc_nhwc_f32)
+    },
+    {
+        "fused_batch_normalization_dwc_NCHW_F32",
+        [](const FuseBatchNormalizeSelectorData & data)
+        {
+            return data.dt == DataType::F32 && data.dl == DataLayout::NCHW && data.fbn_type == FuseBatchNormalizationType::DEPTHWISECONVOLUTION;
+        },
+        REGISTER_FP32_NEON(arm_compute::cpu::fused_batch_normalization_dwc_nchw_f32)
+    }
+};
+
+/** Micro-kernel selector
+ *
+ * @param[in] data Selection data passed to help pick the appropriate micro-kernel
+ *
+ * @param[in]
+ *
+ * @return A matching micro-kernel else nullptr
+ */
+const FBNUKernel *get_implementation(const FuseBatchNormalizeSelectorData &data)
+{
+    for(const auto &uk : available_kernels)
+    {
+        if(uk.is_selected(data))
+        {
+            return &uk;
+        }
+    }
+    return nullptr;
+}
+
 Status validate_arguments(const ITensorInfo *input_weights, const ITensorInfo *bn_mean, const ITensorInfo *bn_var,
                           const ITensorInfo *fused_weights, const ITensorInfo *fused_bias,
                           const ITensorInfo *input_bias, const ITensorInfo *bn_beta, const ITensorInfo *bn_gamma,
@@ -99,318 +209,6 @@ Status validate_arguments(const ITensorInfo *input_weights, const ITensorInfo *b
     return Status{};
 }
 
-template <typename VectorType>
-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   = typename VectorType::scalar_type;
-    const int size     = 16 / conv_weights->info()->element_size();
-    using ExactTagType = typename VectorType::tag_type;
-
-    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 <typename VectorType>
-void fused_batch_normalization_dwc_nhwc(const ITensor *dwc_weights, const ITensor *dwc_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   = typename VectorType::scalar_type;
-    const int size     = 16 / dwc_weights->info()->element_size();
-    using ExactTagType = typename VectorType::tag_type;
-
-    const bool run_in_place_weights = (fused_weights == nullptr) || (fused_weights == dwc_weights);
-    const bool run_in_place_bias    = (fused_bias == nullptr) || (dwc_bias != nullptr && fused_bias == dwc_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 dwc_w_in(dwc_weights, win);
-    Iterator dwc_w_out(run_in_place_weights ? dwc_weights : fused_weights, win);
-
-    const auto dwc_bias_in  = (dwc_bias != nullptr ? reinterpret_cast<ScalarType *>(dwc_bias->ptr_to_element(Coordinates(0, 0))) : nullptr);
-    auto       dwc_bias_out = (run_in_place_bias ? dwc_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{});
-    auto       dwc_bias_vec = wrapper::vdup_n(ScalarType(0), ExactTagType{});
-    const auto epsilon_vec  = wrapper::vdup_n(ScalarType(epsilon), ExactTagType{});
-
-    auto gamma              = ScalarType(1.0);
-    auto beta               = ScalarType(0.0);
-    auto dwc_bias_in_scalar = ScalarType(0);
-
-    execute_window_loop(win, [&](const Coordinates & id)
-    {
-        int x = window_start_x;
-        for(; x <= (window_end_x - window_step_x); x += window_step_x)
-        {
-            var_vec = wrapper::vloadq(input_var + x);
-            if(input_gamma != nullptr)
-            {
-                gamma_vec = wrapper::vloadq(input_gamma + x);
-            }
-
-            if((id[2] == 0) && (id[1] == 0))
-            {
-                mean_vec = wrapper::vloadq(input_mean + x);
-
-                // Construct vectors
-                if(input_beta != nullptr)
-                {
-                    beta_vec = wrapper::vloadq(input_beta + x);
-                }
-
-                if(dwc_bias_in != nullptr)
-                {
-                    dwc_bias_vec = wrapper::vloadq(dwc_bias_in + x);
-                }
-
-                auto dwc_bias_tmp_vec = wrapper::vmul(wrapper::vsub(dwc_bias_vec, mean_vec), wrapper::vinvsqrt(wrapper::vadd(var_vec, epsilon_vec)));
-                dwc_bias_tmp_vec      = wrapper::vadd(wrapper::vmul(dwc_bias_tmp_vec, gamma_vec), beta_vec);
-                wrapper::vstore(dwc_bias_out + x, dwc_bias_tmp_vec);
-            }
-
-            auto dwc_w_in_ptr  = reinterpret_cast<const ScalarType *>(dwc_w_in.ptr());
-            auto dwc_w_out_ptr = reinterpret_cast<ScalarType *>(dwc_w_out.ptr());
-
-            auto wn  = wrapper::vloadq(dwc_w_in_ptr + x);
-            rvar_vec = wrapper::vinvsqrt(wrapper::vadd(var_vec, epsilon_vec));
-            wn       = wrapper::vmul(wn, rvar_vec);
-            wn       = wrapper::vmul(wn, gamma_vec);
-
-            // Store results
-            wrapper::vstore(dwc_w_out_ptr + x, wn);
-        }
-
-        // Compute left-over elements
-        for(; x < window_end_x; ++x)
-        {
-            auto var = input_var[x];
-            if(input_gamma != nullptr)
-            {
-                gamma = input_gamma[x];
-            }
-
-            if(id[2] == 0 && id[1] == 0)
-            {
-                auto mean = input_mean[x];
-                if(input_beta != nullptr)
-                {
-                    beta = input_beta[x];
-                }
-                if(dwc_bias_in != nullptr)
-                {
-                    dwc_bias_in_scalar = dwc_bias_in[x];
-                }
-
-                auto dwc_bias_tmp_scalar = (dwc_bias_in_scalar - mean) / std::sqrt(var + ScalarType(epsilon));
-                dwc_bias_out[x]          = (dwc_bias_tmp_scalar * gamma) + beta;
-            }
-
-            const auto dwc_w_in_ptr  = reinterpret_cast<const ScalarType *>(dwc_w_in.ptr());
-            auto       dwc_w_out_ptr = reinterpret_cast<ScalarType *>(dwc_w_out.ptr());
-
-            *(dwc_w_out_ptr + x) = *(dwc_w_in_ptr + x) / std::sqrt(var + ScalarType(epsilon)) * gamma;
-        }
-    },
-    dwc_w_in, dwc_w_out);
-}
-
-template <typename VectorType>
-void fused_batch_normalization_dwc_nchw(const ITensor *dwc_weights, const ITensor *dwc_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   = typename VectorType::scalar_type;
-    const int size     = 16 / dwc_weights->info()->element_size();
-    using ExactTagType = typename VectorType::tag_type;
-
-    const bool run_in_place_weights = (fused_weights == nullptr) || (fused_weights == dwc_weights);
-    const bool run_in_place_bias    = (fused_bias == nullptr) || (dwc_bias != nullptr && fused_bias == dwc_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 dwc_w_in(dwc_weights, win);
-    Iterator dwc_w_out(run_in_place_weights ? dwc_weights : fused_weights, win);
-
-    const auto dwc_bias_in  = (dwc_bias != nullptr ? reinterpret_cast<ScalarType *>(dwc_bias->ptr_to_element(Coordinates(0, 0))) : nullptr);
-    auto       dwc_bias_out = (run_in_place_bias ? dwc_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 dwc_bias_in_scalar = ScalarType(0.0);
-    execute_window_loop(win, [&](const Coordinates & id)
-    {
-        var = input_var[id[2]];
-        if(input_gamma != nullptr)
-        {
-            gamma = input_gamma[id[2]];
-        }
-
-        if(id[1] == 0)
-        {
-            mean = input_mean[id[2]];
-
-            // Construct vectors
-            mean_vec = wrapper::vdup_n(mean, ExactTagType{});
-            if(input_beta != nullptr)
-            {
-                beta     = input_beta[id[2]];
-                beta_vec = wrapper::vdup_n(beta, ExactTagType{});
-            }
-
-            if(dwc_bias_in != nullptr)
-            {
-                dwc_bias_in_scalar = dwc_bias_in[id[2]];
-            }
-
-            auto dwc_bias_tmp_scalar = (dwc_bias_in_scalar - mean) / std::sqrt(var + ScalarType(epsilon));
-            dwc_bias_out[id[2]]      = (dwc_bias_tmp_scalar * gamma) + beta;
-        }
-
-        int  x             = window_start_x;
-        auto dwc_w_in_ptr  = reinterpret_cast<const ScalarType *>(dwc_w_in.ptr());
-        auto dwc_w_out_ptr = reinterpret_cast<ScalarType *>(dwc_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(dwc_w_in_ptr + x);
-            wn      = wrapper::vmul(wn, rvar_vec);
-            wn      = wrapper::vmul(wn, gamma_vec);
-
-            // Store results
-            wrapper::vstore(dwc_w_out_ptr + x, wn);
-        }
-
-        // Compute left-over elements
-        for(; x < window_end_x; ++x)
-        {
-            *(dwc_w_out_ptr + x) = *(dwc_w_in_ptr + x) / std::sqrt(var + ScalarType(epsilon)) * gamma;
-        }
-    },
-    dwc_w_in, dwc_w_out);
-}
-
 } // namespace
 
 NEFuseBatchNormalizationKernel::NEFuseBatchNormalizationKernel()
@@ -460,37 +258,14 @@ void NEFuseBatchNormalizationKernel::configure(const ITensor *input_weights, con
                                                   (bn_gamma != nullptr) ? bn_gamma->info() : nullptr,
                                                   epsilon, fbn_type));
 
+    const auto *uk = get_implementation(FuseBatchNormalizeSelectorData{ input_weights->info()->data_type(), input_weights->info()->data_layout(), fbn_type, CPUInfo::get().get_isa() });
+    ARM_COMPUTE_ERROR_ON_NULLPTR(uk);
+    ARM_COMPUTE_ERROR_ON(uk->ukernel == nullptr);
+    _func = uk->ukernel;
+
     // Configure kernel window
     Window win = calculate_max_window(*input_weights->info());
     INEKernel::configure(win);
-
-    // Configure function
-    static std::map<std::string, FuseBatchNormFunction *> map_function =
-    {
-        { "fused_batch_normalization_conv_NHWC_F32", &fused_batch_normalization_conv<wrapper::traits::neon_vector<float, 4>> },
-        { "fused_batch_normalization_conv_NCHW_F32", &fused_batch_normalization_conv<wrapper::traits::neon_vector<float, 4>> },
-        { "fused_batch_normalization_dwc_NHWC_F32", &fused_batch_normalization_dwc_nhwc<wrapper::traits::neon_vector<float, 4>> },
-        { "fused_batch_normalization_dwc_NCHW_F32", &fused_batch_normalization_dwc_nchw<wrapper::traits::neon_vector<float, 4>> },
-#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
-        { "fused_batch_normalization_conv_NHWC_F16", &fused_batch_normalization_conv<wrapper::traits::neon_vector<float16_t, 8>> },
-        { "fused_batch_normalization_conv_NCHW_F16", &fused_batch_normalization_conv<wrapper::traits::neon_vector<float16_t, 8>> },
-        { "fused_batch_normalization_dwc_NHWC_F16", &fused_batch_normalization_dwc_nhwc<wrapper::traits::neon_vector<float16_t, 8>> },
-        { "fused_batch_normalization_dwc_NCHW_F16", &fused_batch_normalization_dwc_nchw<wrapper::traits::neon_vector<float16_t, 8>> },
-#endif /* __ARM_FEATURE_FP16_VECTOR_ARITHMETIC */
-    };
-
-    std::string function_to_call("fused_batch_normalization_");
-    function_to_call += fbn_type == FuseBatchNormalizationType::CONVOLUTION ? "conv_" : "dwc_";
-    function_to_call += string_from_data_layout(_input_weights->info()->data_layout());
-    function_to_call += "_";
-    function_to_call += string_from_data_type(_input_weights->info()->data_type());
-
-    auto it = map_function.find(function_to_call);
-
-    if(it != map_function.end())
-    {
-        _func = it->second;
-    }
 }
 
 Status NEFuseBatchNormalizationKernel::validate(const ITensorInfo *input_weights, const ITensorInfo *bn_mean, const ITensorInfo *bn_var,
@@ -507,6 +282,8 @@ void NEFuseBatchNormalizationKernel::run(const Window &window, const ThreadInfo
     ARM_COMPUTE_UNUSED(info);
     ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
     ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(IKernel::window(), window);
+
+    ARM_COMPUTE_ERROR_ON(_func == nullptr);
     (*_func)(_input_weights, _input_bias, _fused_weights, _fused_bias, _bn_mean, _bn_var, _bn_beta, _bn_gamma, _epsilon, window);
 }
 } // namespace arm_compute