diff options
Diffstat (limited to 'src/backends/cl/workloads/ClSpaceToBatchNdWorkload.hpp')
| -rw-r--r-- | src/backends/cl/workloads/ClSpaceToBatchNdWorkload.hpp | 15 |
1 files changed, 9 insertions, 6 deletions
diff --git a/src/backends/cl/workloads/ClSpaceToBatchNdWorkload.hpp b/src/backends/cl/workloads/ClSpaceToBatchNdWorkload.hpp index 948d13b934..a577c06fea 100644 --- a/src/backends/cl/workloads/ClSpaceToBatchNdWorkload.hpp +++ b/src/backends/cl/workloads/ClSpaceToBatchNdWorkload.hpp @@ -1,16 +1,15 @@ // -// Copyright © 2017 Arm Ltd and Contributors. All rights reserved. +// Copyright © 2017, 2020-2023 Arm Ltd and Contributors. All rights reserved. // SPDX-License-Identifier: MIT // #pragma once -#include <armnn/Tensor.hpp> -#include <armnn/Descriptors.hpp> - #include "ClBaseWorkload.hpp" +#include "ClWorkloadUtils.hpp" #include <arm_compute/runtime/CL/functions/CLSpaceToBatchLayer.h> +#include <arm_compute/runtime/CL/functions/CLReshapeLayer.h> namespace armnn { @@ -25,10 +24,14 @@ public: ClSpaceToBatchNdWorkload(const SpaceToBatchNdQueueDescriptor& descriptor, const WorkloadInfo& info, const arm_compute::CLCompileContext& clCompileContext); - void Execute() const override; + virtual void Execute() const override; private: - mutable arm_compute::CLSpaceToBatchLayer m_SpaceToBatchLayer; + mutable arm_compute::CLSpaceToBatchLayer m_Layer; + mutable std::unique_ptr<arm_compute::CLReshapeLayer> m_LayerReshapeInput; + mutable std::unique_ptr<arm_compute::CLReshapeLayer> m_LayerReshapeOutput; + arm_compute::CLTensor m_ReshapeInputTensor; + arm_compute::CLTensor m_ReshapeOutputTensor; }; } //namespace armnn |