#include "internal/arm_compute/feature/View.h"
#include "internal/Sinks.h"
+#include "internal/MatrixSink.h"
#include "util/feature/IndexIterator.h"
}
}
+static void asMatrixSink(ANeuralNetworksExecution *execution, int32_t type, int32_t index,
+ int32_t H, int32_t W, void *buffer, size_t length)
+{
+ switch (type)
+ {
+ case ANEURALNETWORKS_FLOAT32:
+ case ANEURALNETWORKS_TENSOR_FLOAT32:
+ execution->sink<MatrixSink<float>>(index, H, W, reinterpret_cast<float *>(buffer), length);
+ break;
+ case ANEURALNETWORKS_INT32:
+ case ANEURALNETWORKS_TENSOR_INT32:
+ execution->sink<MatrixSink<int32_t>>(index, H, W, reinterpret_cast<int32_t *>(buffer),
+ length);
+ break;
+ case ANEURALNETWORKS_UINT32:
+ execution->sink<MatrixSink<uint32_t>>(index, H, W, reinterpret_cast<uint32_t *>(buffer),
+ length);
+ break;
+ case ANEURALNETWORKS_TENSOR_QUANT8_ASYMM:
+ execution->sink<MatrixSink<uint8_t>>(index, H, W, reinterpret_cast<uint8_t *>(buffer),
+ length);
+ break;
+ default:
+ throw std::runtime_error("Not supported, yet");
+ break;
+ }
+}
+
static void asFeatureSink(ANeuralNetworksExecution *execution, int32_t type, int32_t index,
const nnfw::util::feature::Shape &shape, void *buffer, size_t length)
{
const auto operand_index = execution->plan().model().outputs.at(index);
int32_t output_type = operands.at(operand_index).type();
- const auto squeezed_shape = squeeze(operands.at(operand_index).shape());
+ const auto &output_shape = operands.at(operand_index).shape();
- if (squeezed_shape.rank() == 1)
+ if (output_shape.rank() == 1)
{
- const auto len = squeezed_shape.dim(0);
+ const auto len = output_shape.dim(0);
asVectorSink(execution, output_type, index, len, buffer, length);
}
- else if (squeezed_shape.rank() == 3)
+ else if (output_shape.rank() == 2)
+ {
+ const auto H = output_shape.dim(0);
+ const auto W = output_shape.dim(1);
+
+ asMatrixSink(execution, output_type, index, H, W, buffer, length);
+ }
+ else if ((output_shape.rank() == 4) && (output_shape.dim(0) == 1))
{
const auto &operand_shape = operands.at(operand_index).shape().asFeature();
--- /dev/null
+#ifndef __INTERNAL_MATRIX_SINK_H__
+#define __INTERNAL_MATRIX_SINK_H__
+
+#include "internal/Sink.h"
+
+#include <arm_compute/core/ITensor.h>
+#include <arm_compute/core/Window.h>
+#include <arm_compute/core/Helpers.h>
+
+#include <cstdint>
+#include <cstring>
+#include <cassert>
+
+template <typename T> class MatrixSink final : public Sink
+{
+public:
+ MatrixSink(const int32_t H, const int32_t W, T *base, const size_t size)
+ : _height{H}, _width{W}, _base{base}
+ {
+ assert(size >= _height * _width * sizeof(T));
+ }
+
+public:
+ void pull(::arm_compute::ITensor &tensor) const override
+ {
+ assert(tensor.info()->dimension(0) == _width);
+ assert(tensor.info()->dimension(1) == _height);
+
+ using ::arm_compute::Window;
+ using ::arm_compute::Iterator;
+ using ::arm_compute::Coordinates;
+ using ::arm_compute::execute_window_loop;
+
+ Window window;
+
+ window.use_tensor_dimensions(tensor.info()->tensor_shape(), ::arm_compute::Window::DimY);
+
+ Iterator it(&tensor, window);
+ execute_window_loop(window,
+ [&](const ::arm_compute::Coordinates &id) {
+ const auto row = id.y();
+ memcpy(_base + row * _width, it.ptr(), _width * sizeof(T));
+ },
+ it);
+ }
+
+private:
+ const int32_t _height;
+ const int32_t _width;
+
+private:
+ T *const _base;
+};
+
+#endif // __INTERNAL_MATRIX_SINK_H__