#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/CLSpaceToDepth.h>
#include <arm_compute/runtime/CL/functions/CLReshapeLayer.h>
#include <arm_compute/runtime/CL/functions/CLStridedSlice.h>
#include <arm_compute/runtime/CL/functions/CLSoftmaxLayer.h>
auto stage = [param](const IAllocationContext &ctx, IExecutionBuilder &builder) {
auto output_alloc = ctx.at(::internal::tflite::operand::Index{param.output_index});
auto input_alloc = ctx.at(::internal::tflite::operand::Index{param.input_index});
- auto rank = 4;
- auto fn = nnfw::make_unique<SimpleSpaceToDepth>();
+ if (from_env<bool>(std::getenv("USE_SIMPLE_SPACETODEPTH")))
+ {
+ // USE CPU VERSION OF SPACETODEPTH
+ auto rank = 4;
+ auto fn = nnfw::make_unique<SimpleSpaceToDepth>();
+
+ fn->configure(input_alloc, output_alloc, param.block_size, getARMComputeAxises(rank));
+
+ builder.append("SpaceToDepth", std::move(fn));
+ }
+ else
+ {
+ if (::internal::arm_compute::isGpuMode()) // GPU
+ {
+ auto fn = nnfw::make_unique<::arm_compute::CLSpaceToDepth>();
- fn->configure(input_alloc, output_alloc, param.block_size, getARMComputeAxises(rank));
- builder.append("SpaceToDepth", std::move(fn));
+ fn->configure(CAST_CL(input_alloc), CAST_CL(output_alloc), param.block_size);
+ builder.append("SpaceToDepth", std::move(fn));
+ }
+ else // NEON
+ {
+ // TODO Enable NEON Support
+ throw std::runtime_error("Not supported, yet");
+ }
+ }
};
_builder.addStage(stage);