COMPMID-2244: Extend CLFuseBatchNormalization to support DepthwiseConvolution weights

Change-Id: I7d1907f35cc4899379073759be2f7cce24e51e9d
Signed-off-by: Manuel Bottini <manuel.bottini@arm.com>
Reviewed-on: https://review.mlplatform.org/c/1327
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

1 files changed, 53 insertions, 42 deletions

diff --git a/src/core/CL/kernels/CLFuseBatchNormalizationKernel.cpp b/src/core/CL/kernels/CLFuseBatchNormalizationKernel.cpp
index 16ad7d970e..bf827bf2c2 100644
--- a/src/core/CL/kernels/CLFuseBatchNormalizationKernel.cpp
+++ b/src/core/CL/kernels/CLFuseBatchNormalizationKernel.cpp
@@ -38,50 +38,60 @@ namespace arm_compute
 {
 namespace
 {
-Status validate_arguments(const ITensorInfo *conv_weights, const ITensorInfo *bn_mean, const ITensorInfo *bn_var,
+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 *conv_bias, const ITensorInfo *bn_beta, const ITensorInfo *bn_gamma,
-                          float epsilon)
+                          const ITensorInfo *input_bias, const ITensorInfo *bn_beta, const ITensorInfo *bn_gamma,
+                          float epsilon, FuseBatchNormalizationType fbn_type)
 {
     ARM_COMPUTE_UNUSED(epsilon);
-    ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(conv_weights);
-    ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(conv_weights, 1, DataType::F16, DataType::F32);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(input_weights, bn_mean, bn_var);
+    ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(input_weights);
+    ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input_weights, 1, DataType::F16, DataType::F32);
     ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(bn_mean, bn_var);
-    ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(conv_weights, bn_mean, bn_var);
-    ARM_COMPUTE_RETURN_ERROR_ON(conv_bias == nullptr && fused_bias == nullptr);
-    ARM_COMPUTE_RETURN_ERROR_ON(conv_weights->dimension(3) != bn_mean->dimension(0));
+    ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input_weights, bn_mean, bn_var);
+    ARM_COMPUTE_RETURN_ERROR_ON(input_bias == nullptr && fused_bias == nullptr);
     ARM_COMPUTE_RETURN_ERROR_ON(bn_mean->num_dimensions() > 1);
 
+    if(fbn_type == FuseBatchNormalizationType::CONVOLUTION)
+    {
+        ARM_COMPUTE_RETURN_ERROR_ON(input_weights->dimension(3) != bn_mean->dimension(0));
+    }
+    else
+    {
+        const size_t channel_idx = get_data_layout_dimension_index(input_weights->data_layout(), DataLayoutDimension::CHANNEL);
+        ARM_COMPUTE_RETURN_ERROR_ON(input_weights->dimension(channel_idx) != bn_mean->dimension(0));
+    }
+
     // Validate bias
-    if(conv_bias != nullptr)
+    if(input_bias != nullptr)
     {
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(bn_mean, conv_bias);
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(conv_weights, conv_bias);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(bn_mean, input_bias);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input_weights, input_bias);
     }
     // Validate beta
     if(bn_beta != nullptr)
     {
         ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(bn_mean, bn_beta);
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(conv_weights, bn_beta);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input_weights, bn_beta);
     }
     // Validate gamma
     if(bn_gamma != nullptr)
     {
         ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(bn_mean, bn_gamma);
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(conv_weights, bn_gamma);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input_weights, bn_gamma);
     }
     // Validate output weights
     if(fused_weights != nullptr && fused_weights->total_size() != 0)
     {
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(conv_weights, fused_weights);
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_LAYOUT(conv_weights, fused_weights);
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(conv_weights, fused_weights);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(input_weights, fused_weights);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_LAYOUT(input_weights, fused_weights);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input_weights, fused_weights);
     }
     // Validate output bias
     if(fused_bias != nullptr && fused_bias->total_size() != 0)
     {
         ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(bn_mean, fused_bias);
-        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(conv_weights, fused_bias);
+        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input_weights, fused_bias);
     }
 
     return Status{};
@@ -89,20 +99,20 @@ Status validate_arguments(const ITensorInfo *conv_weights, const ITensorInfo *bn
 } // namespace
 
 CLFuseBatchNormalizationKernel::CLFuseBatchNormalizationKernel()
-    : _conv_weights(nullptr), _conv_bias(nullptr), _bn_mean(nullptr), _bn_var(nullptr), _bn_gamma(nullptr), _bn_beta(nullptr), _fused_weights(nullptr), _fused_bias(nullptr), _epsilon(),
+    : _input_weights(nullptr), _input_bias(nullptr), _bn_mean(nullptr), _bn_var(nullptr), _bn_gamma(nullptr), _bn_beta(nullptr), _fused_weights(nullptr), _fused_bias(nullptr), _epsilon(),
       _run_in_place_weights(false), _run_in_place_bias(false)
 {
 }
 
-void CLFuseBatchNormalizationKernel::configure(const ICLTensor *conv_weights, const ICLTensor *bn_mean, const ICLTensor *bn_var,
+void CLFuseBatchNormalizationKernel::configure(const ICLTensor *input_weights, const ICLTensor *bn_mean, const ICLTensor *bn_var,
                                                ICLTensor *fused_weights, ICLTensor *fused_bias,
-                                               const ICLTensor *conv_bias, const ICLTensor *bn_beta, const ICLTensor *bn_gamma,
-                                               float epsilon)
+                                               const ICLTensor *input_bias, const ICLTensor *bn_beta, const ICLTensor *bn_gamma,
+                                               float epsilon, FuseBatchNormalizationType fbn_type)
 {
-    ARM_COMPUTE_ERROR_ON_NULLPTR(conv_weights, bn_mean, bn_var);
+    ARM_COMPUTE_ERROR_ON_NULLPTR(input_weights, bn_mean, bn_var);
 
-    _conv_weights  = conv_weights;
-    _conv_bias     = conv_bias;
+    _input_weights = input_weights;
+    _input_bias    = input_bias;
     _bn_mean       = bn_mean;
     _bn_var        = bn_var;
     _bn_beta       = bn_beta;
@@ -111,14 +121,14 @@ void CLFuseBatchNormalizationKernel::configure(const ICLTensor *conv_weights, co
     _fused_bias    = fused_bias;
     _epsilon       = epsilon;
 
-    _run_in_place_weights = (fused_weights == nullptr) || (fused_weights == conv_weights);
-    _run_in_place_bias    = (conv_bias != nullptr && fused_bias == nullptr) || (conv_bias != nullptr && fused_bias == conv_bias);
+    _run_in_place_weights = (fused_weights == nullptr) || (fused_weights == input_weights);
+    _run_in_place_bias    = (input_bias != nullptr && fused_bias == nullptr) || (input_bias != nullptr && fused_bias == input_bias);
 
     // Auto initialize outputs
     if(_fused_weights != nullptr)
     {
         // Output tensor auto initialization if not yet initialized
-        auto_init_if_empty(*_fused_weights->info(), *_conv_weights->info()->clone());
+        auto_init_if_empty(*_fused_weights->info(), *_input_weights->info()->clone());
     }
     if(_fused_bias != nullptr)
     {
@@ -127,39 +137,40 @@ void CLFuseBatchNormalizationKernel::configure(const ICLTensor *conv_weights, co
     }
 
     // Validate arguments
-    ARM_COMPUTE_ERROR_THROW_ON(validate_arguments(conv_weights->info(), bn_mean->info(), bn_var->info(),
+    ARM_COMPUTE_ERROR_THROW_ON(validate_arguments(input_weights->info(), bn_mean->info(), bn_var->info(),
                                                   (fused_weights != nullptr) ? fused_weights->info() : nullptr,
                                                   (fused_bias != nullptr) ? fused_bias->info() : nullptr,
-                                                  (conv_bias != nullptr) ? conv_bias->info() : nullptr,
+                                                  (input_bias != nullptr) ? input_bias->info() : nullptr,
                                                   (bn_beta != nullptr) ? bn_beta->info() : nullptr,
                                                   (bn_gamma != nullptr) ? bn_gamma->info() : nullptr,
-                                                  epsilon));
+                                                  epsilon, fbn_type));
 
     // Configure kernel window
-    Window win = calculate_max_window(*conv_weights->info());
+    Window win = calculate_max_window(*input_weights->info());
     ICLKernel::configure_internal(win);
 
     // Set build options
     CLBuildOptions build_opts;
-    build_opts.add_option("-DDATA_TYPE=" + get_cl_type_from_data_type(conv_weights->info()->data_type()));
-    build_opts.add_option("-DDIM2=" + support::cpp11::to_string(conv_weights->info()->dimension(2)));
+    build_opts.add_option("-DDATA_TYPE=" + get_cl_type_from_data_type(input_weights->info()->data_type()));
+    build_opts.add_option_if(fbn_type == FuseBatchNormalizationType::CONVOLUTION, "-DDIM2=" + support::cpp11::to_string(input_weights->info()->dimension(2)));
     build_opts.add_option("-DEPSILON=" + float_to_string_with_full_precision(epsilon));
+    build_opts.add_option_if(_input_weights->info()->data_layout() == DataLayout::NHWC, "-DNHWC");
     build_opts.add_option_if(_run_in_place_weights, "-DIN_PLACE_W");
     build_opts.add_option_if(_run_in_place_bias, "-DIN_PLACE_B");
-    build_opts.add_option_if(conv_bias != nullptr, "-DBIAS");
+    build_opts.add_option_if(input_bias != nullptr, "-DBIAS");
     build_opts.add_option_if(bn_beta != nullptr, "-DBETA");
     build_opts.add_option_if(bn_gamma != nullptr, "-DGAMMA");
 
     // Create kernel
-    _kernel = static_cast<cl::Kernel>(CLKernelLibrary::get().create_kernel("fuse_batchnormalization_conv_layer", build_opts.options()));
+    _kernel = static_cast<cl::Kernel>(CLKernelLibrary::get().create_kernel("fuse_batchnormalization_layer", build_opts.options()));
 }
 
-Status CLFuseBatchNormalizationKernel::validate(const ITensorInfo *conv_weights, const ITensorInfo *bn_mean, const ITensorInfo *bn_var,
+Status CLFuseBatchNormalizationKernel::validate(const ITensorInfo *input_weights, const ITensorInfo *bn_mean, const ITensorInfo *bn_var,
                                                 const ITensorInfo *fused_weights, const ITensorInfo *fused_bias,
-                                                const ITensorInfo *conv_bias, const ITensorInfo *bn_beta, const ITensorInfo *bn_gamma,
-                                                float epsilon)
+                                                const ITensorInfo *input_bias, const ITensorInfo *bn_beta, const ITensorInfo *bn_gamma,
+                                                float epsilon, FuseBatchNormalizationType fbn_type)
 {
-    ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments(conv_weights, bn_mean, bn_var, fused_weights, fused_bias, conv_bias, bn_beta, bn_gamma, epsilon));
+    ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments(input_weights, bn_mean, bn_var, fused_weights, fused_bias, input_bias, bn_beta, bn_gamma, epsilon, fbn_type));
     return Status{};
 }
 
@@ -175,10 +186,10 @@ void CLFuseBatchNormalizationKernel::run(const arm_compute::Window &window, cl::
 
     // Add kernel arguments
     unsigned int idx = 0;
-    add_3D_tensor_argument(idx, _conv_weights, slice_3d);
-    if(_conv_bias != nullptr)
+    add_3D_tensor_argument(idx, _input_weights, slice_3d);
+    if(_input_bias != nullptr)
     {
-        add_1D_tensor_argument(idx, _conv_bias, slice_1d);
+        add_1D_tensor_argument(idx, _input_bias, slice_1d);
     }
     add_1D_tensor_argument(idx, _bn_mean, slice_1d);
     add_1D_tensor_argument(idx, _bn_var, slice_1d);