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             BOOST_ASSERT(outputInfo.GetNumDimensions() == inputInfo0.GetNumDimensions());
  45
  46             // check all dimensions to see if this element is inside the given input view
  47             bool insideView = true;
  48             for (unsigned int i = 0; i<outputInfo.GetNumDimensions(); i++)
  49             {
  50                 if (indices[i] < view.m_Origin[i])
  51                 {
  52                     insideView = false;
  53                 }
  54                 if (indices[i] >= view.m_Origin[i] + outputInfo.GetShape()[i])
  55                 {
  56                     insideView = false;
  57                 }
  58             }
  59
  60             if (insideView)
  61             {
  62                 unsigned int outIndex = 0;
  63                 unsigned int dimensionStride = 1;
  64
  65                 for (unsigned int i = outputInfo.GetNumDimensions(); i-- > 0;)
  66                 {
  67                     outIndex += dimensionStride * (indices[i] - view.m_Origin[i]);
  68                     dimensionStride *= outputInfo.GetShape()[i];
  69                 }
  70
  71                 //we are within the view, copy input data to the output corresponding to this view
  72                 DataType* outputData = GetOutputTensorData<DataType>(viewIdx, data);
  73                 BOOST_ASSERT(outputData);
  74
  75                 const DataType* inputData = GetInputTensorData<DataType>(0, data);
  76                 BOOST_ASSERT(inputData);
  77
  78                 outputData[outIndex] = inputData[index];
  79             }
  80         }
  81     }
  82 }
  83
  84 } //namespace armnn