xref: /aosp_15_r20/external/armnn/src/backends/reference/workloads/Splitter.cpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2017 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 
6 #include "RefWorkloadUtils.hpp"
7 #include <armnn/backends/WorkloadData.hpp>
8 #include <armnn/Tensor.hpp>
9 #include <armnn/utility/Assert.hpp>
10 #include "Splitter.hpp"
11 
12 #include <cmath>
13 #include <limits>
14 
15 #include "Decoders.hpp"
16 #include "Encoders.hpp"
17 
18 namespace armnn
19 {
20 
Split(const SplitterQueueDescriptor & data,std::vector<ITensorHandle * > inputs,std::vector<ITensorHandle * > outputs)21 void Split(const SplitterQueueDescriptor& data,
22            std::vector<ITensorHandle*> inputs,
23            std::vector<ITensorHandle*> outputs)
24 {
25     const TensorInfo& inputInfo = GetTensorInfo(inputs[0]);
26 
27     std::unique_ptr<Decoder<float>> decoderPtr =
28         MakeDecoder<float>(inputInfo, inputs[0]->Map());
29     Decoder<float>& decoder = *decoderPtr;
30 
31     for (unsigned int index = 0; index < inputInfo.GetNumElements(); ++index)
32     {
33         unsigned int indices[MaxNumOfTensorDimensions] = { 0 };
34 
35         unsigned int indexRemainder = index;
36         unsigned int dimensionStride = inputInfo.GetNumElements();
37 
38         for (unsigned int i = 0; i<inputInfo.GetNumDimensions(); i++)
39         {
40             dimensionStride /= inputInfo.GetShape()[i];
41             indices[i] = indexRemainder / dimensionStride; // Use integer division to round down.
42             indexRemainder -= indices[i] * dimensionStride;
43         }
44 
45         for (unsigned int viewIdx = 0; viewIdx < data.m_ViewOrigins.size(); ++viewIdx)
46         {
47             SplitterQueueDescriptor::ViewOrigin const& view = data.m_ViewOrigins[viewIdx];
48 
49             //Split view extents are defined by the size of (the corresponding) input tensor.
50             const TensorInfo& outputInfo = GetTensorInfo(outputs[viewIdx]);
51             ARMNN_ASSERT(outputInfo.GetNumDimensions() == inputInfo.GetNumDimensions());
52 
53             // Check all dimensions to see if this element is inside the given input view.
54             bool insideView = true;
55             for (unsigned int i = 0; i<outputInfo.GetNumDimensions(); i++)
56             {
57                 if (indices[i] < view.m_Origin[i])
58                 {
59                     insideView = false;
60                 }
61                 if (indices[i] >= view.m_Origin[i] + outputInfo.GetShape()[i])
62                 {
63                     insideView = false;
64                 }
65             }
66 
67             if (insideView)
68             {
69                 std::unique_ptr<Encoder<float>> encoderPtr =
70                     MakeEncoder<float>(outputInfo, outputs[viewIdx]->Map());
71                 Encoder<float>& encoder = *encoderPtr;
72 
73                 unsigned int outIndex = 0;
74                 unsigned int dimensionStride = 1;
75                 float inputValue = 0.f;
76 
77                 for (unsigned int i = outputInfo.GetNumDimensions(); i-- > 0;)
78                 {
79                     outIndex += dimensionStride * (indices[i] - view.m_Origin[i]);
80                     dimensionStride *= outputInfo.GetShape()[i];
81                 }
82 
83                 decoder += index;
84                 inputValue = decoder.Get();
85                 decoder -= index;
86 
87                 encoder += outIndex;
88                 encoder.Set(inputValue);
89                 break;
90             }
91         }
92     }
93 }
94 
95 }