From d134093a271b60e248942af9757e8236e8f41ac1 Mon Sep 17 00:00:00 2001
From: Nikhil Raj <nikhil.raj@arm.com>
Date: Thu, 18 Oct 2018 14:27:50 +0100
Subject: IVGCVSW-2023 CL and Neon implementation of BatchNorm with NHWC

Change-Id: I962e986607e5d045cd97b9eaeaea2f5ae624db35
---
 .../NeonBatchNormalizationFloatWorkload.cpp        | 24 ++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

(limited to 'src/backends/neon/workloads')

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());
 
-- 
cgit v1.2.1