diff options
Diffstat (limited to 'src/armnn/backends/RefWorkloads/Broadcast.hpp')
| -rw-r--r-- | src/armnn/backends/RefWorkloads/Broadcast.hpp | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/src/armnn/backends/RefWorkloads/Broadcast.hpp b/src/armnn/backends/RefWorkloads/Broadcast.hpp new file mode 100644 index 0000000000..b65b57f7a1 --- /dev/null +++ b/src/armnn/backends/RefWorkloads/Broadcast.hpp @@ -0,0 +1,58 @@ +// +// Copyright © 2017 Arm Ltd. All rights reserved. +// See LICENSE file in the project root for full license information. +// + +#include <armnn/Tensor.hpp> + +#include <functional> + +namespace armnn +{ + +struct BroadcastLoop +{ + BroadcastLoop(const TensorShape& inShape0, const TensorShape& inShape1, const TensorShape& outShape); + + unsigned int GetNumDimensions() + { + return static_cast<unsigned int>(m_DimData.size()); + } + + template <typename T0, typename T1, typename U, typename Func> + void Unroll(Func operationFunc, + unsigned int dimension, + const T0* inData0, + const T1* inData1, + U* outData) + { + if (dimension >= GetNumDimensions()) + { + *outData = operationFunc(*inData0, *inData1); + return; + } + + for (unsigned int i = 0; i < m_DimData[dimension].m_DimSize; i++) + { + Unroll(operationFunc, dimension + 1, inData0, inData1, outData); + + inData0 += m_DimData[dimension].m_Stride1; + inData1 += m_DimData[dimension].m_Stride2; + outData += m_DimData[dimension].m_StrideOut; + } + } + +private: + // Struct to hold the dimension data + struct BroadcastDimensionData + { + unsigned int m_DimSize; + unsigned int m_StrideOut; + unsigned int m_Stride1; + unsigned int m_Stride2; + }; + + std::vector<BroadcastDimensionData> m_DimData; +}; + +} //namespace armnn
\ No newline at end of file |