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>
diff --git a/src/runtime/NEON/functions/NEConcatenateLayer.cpp b/src/runtime/NEON/functions/NEConcatenateLayer.cpp
index d338493..9a70d32 100644
--- a/src/runtime/NEON/functions/NEConcatenateLayer.cpp
+++ b/src/runtime/NEON/functions/NEConcatenateLayer.cpp
@@ -23,6 +23,7 @@
  */
 #include "arm_compute/runtime/NEON/functions/NEConcatenateLayer.h"
 
+#include "arm_compute/core/NEON/kernels/NEBatchConcatenateLayerKernel.h"
 #include "arm_compute/core/NEON/kernels/NEDepthConcatenateLayerKernel.h"
 #include "arm_compute/core/NEON/kernels/NEHeightConcatenateLayerKernel.h"
 #include "arm_compute/core/NEON/kernels/NEWidthConcatenateLayerKernel.h"
@@ -112,6 +113,13 @@
                 _concat_kernels.emplace_back(std::move(kernel));
                 break;
             }
+            case 3:
+            {
+                auto kernel = support::cpp14::make_unique<NEBatchConcatenateLayerKernel>();
+                kernel->configure(inputs_vector.at(i), offset, output);
+                _concat_kernels.emplace_back(std::move(kernel));
+                break;
+            }
             default:
                 ARM_COMPUTE_ERROR("Axis not supported");
         }
@@ -146,6 +154,11 @@
                 ARM_COMPUTE_RETURN_ON_ERROR(NEDepthConcatenateLayerKernel::validate(input, offset, output));
                 break;
             }
+            case 3:
+            {
+                ARM_COMPUTE_RETURN_ON_ERROR(NEBatchConcatenateLayerKernel::validate(input, offset, output));
+                break;
+            }
             default:
                 ARM_COMPUTE_ERROR("Axis not supported");
         }