From 06063187d9bb5e7f3dcd79e007f89a65a2d5a411 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EC=9E=A5=EC=A7=80=EC=84=AD/=EB=8F=99=EC=9E=91=EC=A0=9C?=
 =?utf8?q?=EC=96=B4Lab=28SR=29/Engineer/=EC=82=BC=EC=84=B1=EC=A0=84?=
 =?utf8?q?=EC=9E=90?= <jiseob.jang@samsung.com>
Date: Wed, 28 Nov 2018 10:04:11 +0900
Subject: [PATCH] Apply CL Kernel of SpaceToBatchND to PACL (#3731)

This commit applies CL Kernel of SpaceToBatchND to PACL.

Signed-off-by: jiseob.jang <jiseob.jang@samsung.com>
---
 runtimes/pure_arm_compute/src/compilation.cc | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/runtimes/pure_arm_compute/src/compilation.cc b/runtimes/pure_arm_compute/src/compilation.cc
index dfa04c4..579f569 100644
--- a/runtimes/pure_arm_compute/src/compilation.cc
+++ b/runtimes/pure_arm_compute/src/compilation.cc
@@ -38,6 +38,7 @@
 #include <arm_compute/runtime/CL/functions/CLActivationLayer.h>
 #include <arm_compute/runtime/CL/functions/CLActivationLayerEx.h>
 #include <arm_compute/runtime/CL/functions/CLScale.h>
+#include <arm_compute/runtime/CL/functions/CLSpaceToBatchND.h>
 #include <arm_compute/runtime/CL/functions/CLSpaceToDepth.h>
 #include <arm_compute/runtime/CL/functions/CLReshapeLayer.h>
 #include <arm_compute/runtime/CL/functions/CLStridedSlice.h>
@@ -4216,12 +4217,28 @@ void Planner::visit(const ::internal::tflite::op::SpaceToBatchND::Node &node)
     auto block_size_alloc = ctx.at(::internal::tflite::operand::Index{param.block_size_index});
     auto padding_size_alloc = ctx.at(::internal::tflite::operand::Index{param.padding_size_index});
 
-    auto fn = nnfw::make_unique<SimpleSpaceToBatchND>();
+    // NOTE SimpleSpaceToBatchND is quite slow
+    if (from_env<bool>(std::getenv("USE_SIMPLE_SPACE_TO_BATCH_ND")))
+    {
+      auto fn = nnfw::make_unique<SimpleSpaceToBatchND>();
 
-    fn->configure(input_alloc, block_size_alloc, padding_size_alloc, output_alloc,
-                  getARMComputeAxises(param.rank));
-    builder.append("SpaceToBatchND", std::move(fn));
+      fn->configure(input_alloc, block_size_alloc, padding_size_alloc, output_alloc,
+                    getARMComputeAxises(param.rank));
+      builder.append("SpaceToBatchND", std::move(fn));
+    }
+    else if (::internal::arm_compute::isGpuMode())
+    {
+      auto fn = nnfw::make_unique<::arm_compute::CLSpaceToBatchND>();
 
+      fn->configure(CAST_CL(input_alloc), CAST_CL(block_size_alloc), CAST_CL(padding_size_alloc),
+                    CAST_CL(output_alloc));
+      builder.append("SpaceToBatchND", std::move(fn));
+    }
+    else
+    {
+      // TODO Enable NEON Support
+      throw std::runtime_error("Not supported, yet");
+    }
   };
 
   _builder.addStage(stage);
-- 
2.7.4