From 338595bca8ab60492f10626860acb1ab3722b1ce Mon Sep 17 00:00:00 2001
From: Vidhya Sudhan Loganathan <vidhyasudhan.loganathan@arm.com>
Date: Fri, 28 Jun 2019 14:09:53 +0100
Subject: COMPMID-2234 : Add support for axis 3 in NE/CLConcatenateLayer

Change-Id: Ic86f89ece3afe72809bc69c6de6fee7d21daa1d4
Signed-off-by: Vidhya Sudhan Loganathan <vidhyasudhan.loganathan@arm.com>
Reviewed-on: https://review.mlplatform.org/c/1440
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
---
 src/runtime/CL/functions/CLConcatenateLayer.cpp | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

(limited to 'src/runtime/CL')

diff --git a/src/runtime/CL/functions/CLConcatenateLayer.cpp b/src/runtime/CL/functions/CLConcatenateLayer.cpp
index 0594a17a7a..1d396f5ebf 100644
--- a/src/runtime/CL/functions/CLConcatenateLayer.cpp
+++ b/src/runtime/CL/functions/CLConcatenateLayer.cpp
@@ -23,6 +23,7 @@
  */
 #include "arm_compute/runtime/CL/functions/CLConcatenateLayer.h"
 
+#include "arm_compute/core/CL/kernels/CLBatchConcatenateLayerKernel.h"
 #include "arm_compute/core/CL/kernels/CLDepthConcatenateLayerKernel.h"
 #include "arm_compute/core/CL/kernels/CLHeightConcatenateLayerKernel.h"
 #include "arm_compute/core/CL/kernels/CLWidthConcatenate2TensorsKernel.h"
@@ -124,6 +125,17 @@ void CLConcatenateLayer::configure(const std::vector<ICLTensor *> &inputs_vector
             }
             break;
         }
+        case 3:
+        {
+            for(unsigned int i = 0; i < _num_inputs; ++i)
+            {
+                auto kernel = support::cpp14::make_unique<CLBatchConcatenateLayerKernel>();
+                kernel->configure(inputs_vector.at(i), offset, output);
+                offset += inputs_vector.at(i)->info()->dimension(_axis);
+                _concat_kernels.emplace_back(std::move(kernel));
+            }
+            break;
+        }
         default:
             ARM_COMPUTE_ERROR("Axis not supported");
     }
@@ -184,6 +196,15 @@ Status CLConcatenateLayer::validate(const std::vector<ITensorInfo *> &inputs_vec
             }
             break;
         }
+        case 3:
+        {
+            for(const auto &input : inputs_vector)
+            {
+                ARM_COMPUTE_RETURN_ON_ERROR(CLBatchConcatenateLayerKernel::validate(input, offset, output));
+                offset += input->dimension(axis);
+            }
+            break;
+        }
         default:
             ARM_COMPUTE_ERROR("Axis not supported");
     }
-- 
cgit v1.2.1