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>

diff --git a/runtimes/pure_arm_compute/src/compilation.cc b/runtimes/pure_arm_compute/src/compilation.cc
index dfa04c4..579f569 100644 (file)
--- 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);