aboutsummaryrefslogtreecommitdiff
path: root/arm_compute/runtime/NEON/functions/NESplit.h
diff options
context:
space:
mode:
Diffstat (limited to 'arm_compute/runtime/NEON/functions/NESplit.h')
-rw-r--r--arm_compute/runtime/NEON/functions/NESplit.h37
1 files changed, 4 insertions, 33 deletions
diff --git a/arm_compute/runtime/NEON/functions/NESplit.h b/arm_compute/runtime/NEON/functions/NESplit.h
index e4d62048e6..69aef793d5 100644
--- a/arm_compute/runtime/NEON/functions/NESplit.h
+++ b/arm_compute/runtime/NEON/functions/NESplit.h
@@ -1,5 +1,5 @@
/*
- * Copyright (c) 2018-2019 ARM Limited.
+ * Copyright (c) 2018-2020 ARM Limited.
*
* SPDX-License-Identifier: MIT
*
@@ -24,8 +24,10 @@
#ifndef ARM_COMPUTE_NESPLIT_H
#define ARM_COMPUTE_NESPLIT_H
+#include "arm_compute/core/ITensor.h"
#include "arm_compute/core/Types.h"
+#include "arm_compute/runtime/CPP/functions/CPPSplit.h"
#include "arm_compute/runtime/IFunction.h"
#include "arm_compute/runtime/NEON/functions/NESlice.h"
@@ -34,43 +36,12 @@
namespace arm_compute
{
-// Forward declarations
-class ITensor;
-
/** Basic function to split a tensor along a given axis */
-class NESplit : public IFunction
+class NESplit : public CPPSplit<NESlice>
{
public:
- /** Default constructor */
- NESplit();
- /** Initialise the kernel's input and outputs.
- *
- * @param[in] input The input tensor. Data types supported: All
- * @param[out] outputs A vector containing the output tensors. Data types supported: Same as @p input.
- * The output tensors should match the input tensor dimensions for all shape dimensions apart
- * from the split dimension.
- * @param[in] axis Axis on which to split the input.
- */
- void configure(const ITensor *input, const std::vector<ITensor *> &outputs, unsigned int axis);
- /** Static function to check if given info will lead to a valid configuration of @ref NESplit
- *
- * @param[in] input The input tensor info. Data types supported: All
- * @param[in] outputs A vector containing the output tensors' info. Data types supported: Same as @p input.
- * The output tensors should match the input tensor dimensions for all shape dimensions apart
- * from the split dimension
- * @param[in] axis Axis on which to split the input.
- *
- * @return a status
- */
- static Status validate(const ITensorInfo *input, const std::vector<ITensorInfo *> &outputs, unsigned int axis);
-
// Inherited methods overridden:
void run() override;
-
-private:
- std::vector<ITensor *> _outputs_vector;
- std::vector<NESlice> _slice_functions;
- unsigned int _num_outputs;
};
} // namespace arm_compute
#endif /* ARM_COMPUTE_NESPLIT_H */