IVGCVSW-2023 CL and Neon implementation of BatchNorm with NHWC

Change-Id: I962e986607e5d045cd97b9eaeaea2f5ae624db35

1 files changed, 18 insertions, 6 deletions

diff --git a/src/backends/neon/workloads/NeonBatchNormalizationFloatWorkload.cpp b/src/backends/neon/workloads/NeonBatchNormalizationFloatWorkload.cpp
index f7056a515b..95cfdce9b4 100644
--- a/src/backends/neon/workloads/NeonBatchNormalizationFloatWorkload.cpp
+++ b/src/backends/neon/workloads/NeonBatchNormalizationFloatWorkload.cpp
@@ -21,12 +21,20 @@ arm_compute::Status NeonBatchNormalizationValidate(const TensorInfo& input,
                                                    const TensorInfo& gamma,
                                                    const BatchNormalizationDescriptor& descriptor)
 {
-    const arm_compute::TensorInfo aclInputInfo = BuildArmComputeTensorInfo(input);
-    const arm_compute::TensorInfo aclOutputInfo = BuildArmComputeTensorInfo(output);
-    const arm_compute::TensorInfo aclMeanInfo = BuildArmComputeTensorInfo(mean);
-    const arm_compute::TensorInfo aclVarInfo = BuildArmComputeTensorInfo(var);
-    const arm_compute::TensorInfo aclBetaInfo = BuildArmComputeTensorInfo(beta);
-    const arm_compute::TensorInfo aclGammaInfo = BuildArmComputeTensorInfo(gamma);
+    const DataLayout dataLayout = descriptor.m_DataLayout.GetDataLayout();
+
+    const arm_compute::TensorInfo aclInputInfo =
+          armcomputetensorutils::BuildArmComputeTensorInfo(input, dataLayout);
+    const arm_compute::TensorInfo aclOutputInfo =
+          armcomputetensorutils::BuildArmComputeTensorInfo(output, dataLayout);
+    const arm_compute::TensorInfo aclMeanInfo =
+          armcomputetensorutils::BuildArmComputeTensorInfo(mean, dataLayout);
+    const arm_compute::TensorInfo aclVarInfo =
+          armcomputetensorutils::BuildArmComputeTensorInfo(var, dataLayout);
+    const arm_compute::TensorInfo aclBetaInfo =
+          armcomputetensorutils::BuildArmComputeTensorInfo(beta, dataLayout);
+    const arm_compute::TensorInfo aclGammaInfo =
+          armcomputetensorutils::BuildArmComputeTensorInfo(gamma, dataLayout);
 
     return arm_compute::NEBatchNormalizationLayer::validate(&aclInputInfo,
                                                             &aclOutputInfo,
@@ -46,6 +54,10 @@ NeonBatchNormalizationFloatWorkload::NeonBatchNormalizationFloatWorkload(
     arm_compute::ITensor& input = boost::polymorphic_downcast<INeonTensorHandle*>(m_Data.m_Inputs[0])->GetTensor();
     arm_compute::ITensor& output = boost::polymorphic_downcast<INeonTensorHandle*>(m_Data.m_Outputs[0])->GetTensor();
 
+    arm_compute::DataLayout aclDataLayout = ConvertDataLayout(m_Data.m_Parameters.m_DataLayout.GetDataLayout());
+    input.info()->set_data_layout(aclDataLayout);
+    output.info()->set_data_layout(aclDataLayout);
+
     m_Mean = std::make_unique<arm_compute::Tensor>();
     BuildArmComputeTensor(*m_Mean, m_Data.m_Mean->GetTensorInfo());