src/armnn/backends/RefWorkloads/Splitter.hpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // See LICENSE file in the project root for full license information.
   4 //
   5
   6 #pragma once
   7
   8 #include "RefWorkloadUtils.hpp"
   9
  10 #include "backends/WorkloadData.hpp"
  11
  12 #include <armnn/Tensor.hpp>
  13
  14 #include <boost/assert.hpp>
  15
  16 namespace armnn
  17 {
  18
  19 template <typename DataType>
  20 void Splitter(const SplitterQueueDescriptor& data)
  21 {
  22     const TensorInfo& inputInfo0 = GetTensorInfo(data.m_Inputs[0]);
  23
  24     for (unsigned int index = 0; index < inputInfo0.GetNumElements(); ++index)
  25     {
  26         unsigned int indices[MaxNumOfTensorDimensions];
  27
  28         unsigned int indexRemainder = index;
  29         unsigned int dimensionStride = inputInfo0.GetNumElements();
  30
  31         for (unsigned int i = 0; i<inputInfo0.GetNumDimensions(); i++)
  32         {
  33             dimensionStride /= inputInfo0.GetShape()[i];
  34             indices[i] = indexRemainder / dimensionStride; // use integer division to round down
  35             indexRemainder -= indices[i] * dimensionStride;
  36         }
  37
  38         for (unsigned int viewIdx = 0; viewIdx < data.m_ViewOrigins.size(); ++viewIdx)
  39         {
  40             SplitterQueueDescriptor::ViewOrigin const& view = data.m_ViewOrigins[viewIdx];
  41
  42             //split view extents are defined by the size of (the corresponding) input tensor
  43             const TensorInfo& outputInfo = GetTensorInfo(data.m_Outputs[viewIdx]);
  44
  45             // check all dimensions to see if this element is inside the given input view
  46             bool insideView = true;
  47             for (unsigned int i = 0; i<outputInfo.GetNumDimensions(); i++)
  48             {
  49                 if (indices[i] < view.m_Origin[i])
  50                 {
  51                     insideView = false;
  52                 }
  53                 if (indices[i] >= view.m_Origin[i] + outputInfo.GetShape()[i])
  54                 {
  55                     insideView = false;
  56                 }
  57             }
  58
  59             if (insideView)
  60             {
  61                 unsigned int outIndex = 0;
  62                 unsigned int dimensionStride = 1;
  63
  64                 for (unsigned int i = outputInfo.GetNumDimensions(); i-- > 0;)
  65                 {
  66                     outIndex += dimensionStride * (indices[i] - view.m_Origin[i]);
  67                     dimensionStride *= outputInfo.GetShape()[i];
  68                 }
  69
  70                 //we are within the view, copy input data to the output corresponding to this view
  71                 DataType* outputData = GetOutputTensorData<DataType>(viewIdx, data);
  72                 BOOST_ASSERT(outputData);
  73
  74                 const DataType* inputData = GetInputTensorData<DataType>(0, data);
  75                 BOOST_ASSERT(inputData);
  76
  77                 outputData[outIndex] = inputData[index];
  78             }
  79         }
  80     }
  81 }
  82
  83 } //namespace armnn