From 9c6f3afdf7686d9804fd5bc04ea68fca4c49a0a8 Mon Sep 17 00:00:00 2001
From: saoste01 <saoirse.stewart@arm.com>
Date: Mon, 5 Nov 2018 12:30:23 +0000
Subject: IVGCVSW-2113: Update ClDepthWiseConvolutionWOrkload to suit NCHW/NHWC
 DataLayout (3*3 Optimization)

Change-Id: Icfc137201c52e2c856076b7795572cc4ba75cc95
---
 src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp b/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp
index aa5d5a1b0a..9cadbf09ac 100644
--- a/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp
+++ b/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp
@@ -88,8 +88,11 @@ ClDepthwiseConvolutionWorkload::ClDepthwiseConvolutionWorkload(
 
     const unsigned int depthMultiplier = weightInfo.GetShape()[0];
 
+    const unsigned int widthIndex = (m_Data.m_Parameters.m_DataLayout == DataLayout::NCHW) ? 3 : 2;
+    const unsigned int heightIndex = (m_Data.m_Parameters.m_DataLayout == DataLayout::NCHW) ? 2 : 1;
+
     //Check for optimisation opportunities.
-    bool use3x3Optimisation = (weightInfo.GetShape()[3] == 3) && (weightInfo.GetShape()[2] == 3);
+    bool use3x3Optimisation = (weightInfo.GetShape()[widthIndex] == 3) && (weightInfo.GetShape()[heightIndex] == 3);
     if (use3x3Optimisation)
     {
         m_DepthwiseConvolutionLayer = std::make_unique<arm_compute::CLDepthwiseConvolutionLayer3x3>();
-- 
cgit v1.2.1