runtime/contrib/pure_arm_compute/src/internal/Tensor3DSink.h

   1 /*
   2  * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 /**
  18  * @file    Tensor3DSink.h
  19  * @ingroup COM_AI_RUNTIME
  20  * @brief   This file defines Tensor3DSink class
  21  */
  22 #ifndef __TENSOR3D_SINK_H__
  23 #define __TENSOR3D_SINK_H__
  24
  25 #include "internal/Sink.h"
  26
  27 //
  28 // This is mempcy() version of generic TensorSink for 3D tensor
  29 //
  30 #include <arm_compute/core/ITensor.h>
  31 #include <arm_compute/core/Window.h>
  32 #include <arm_compute/core/Helpers.h>
  33
  34 /**
  35  * @brief Class to get tensor data from arm compute tensor
  36  */
  37 template <typename T> class Tensor3DSink final : public Sink
  38 {
  39 public:
  40   /**
  41    * @brief     Construct a new Tensor3DSink object
  42    * @param[in] shape Shape of tensor
  43    * @param[in] base  Pointer to get data
  44    * @param[in] size  Size of tensor
  45    */
  46   Tensor3DSink(const nnfw::misc::tensor::Shape &shape, T *base, const size_t size)
  47       : _shape{shape}, _base{base}, _size{size}
  48   {
  49     // DO NOTHING
  50   }
  51
  52 public:
  53   /**
  54    * @brief     Get tensor data from arm compute tensor to base
  55    * @param[in] tensor  Tensor object of arm compute to get data
  56    * @return    N/A
  57    */
  58   void pull(::arm_compute::ITensor &tensor) const override
  59   {
  60     using ::arm_compute::Coordinates;
  61     using ::arm_compute::execute_window_loop;
  62     using ::arm_compute::Iterator;
  63     using ::arm_compute::Window;
  64
  65     Window window;
  66
  67     window.use_tensor_dimensions(tensor.info()->tensor_shape(), ::arm_compute::Window::DimY);
  68     int32_t height_width = _shape.dim(1) * _shape.dim(2);
  69     int32_t width = _shape.dim(2);
  70
  71     Iterator it(&tensor, window);
  72     execute_window_loop(window,
  73                         [&](const ::arm_compute::Coordinates &id) {
  74                           const auto z = id.z();
  75                           const auto y = id.y();
  76                           memcpy(_base + z * height_width + y * width, it.ptr(), width * sizeof(T));
  77                         },
  78                         it);
  79   }
  80
  81 private:
  82   const nnfw::misc::tensor::Shape _shape;
  83
  84 private:
  85   T *const _base;
  86   const size_t _size;
  87 };
  88
  89 #endif // __TENSOR3D_SINK_H__