1 //
2 // Copyright © 2019 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #include "Broadcast.hpp"
7
8 namespace armnn
9 {
10
BroadcastLoop(const TensorShape & inShape0,const TensorShape & inShape1,const TensorShape & outShape)11 BroadcastLoop::BroadcastLoop(const TensorShape& inShape0, const TensorShape& inShape1, const TensorShape& outShape)
12 : m_DimData(outShape.GetNumDimensions())
13 {
14 const unsigned int numDims = GetNumDimensions();
15
16 unsigned int sIn0 = 1;
17 unsigned int sIn1 = 1;
18 unsigned int sOut = 1;
19
20 for (unsigned int j = numDims - 1, k = 0; k < numDims ; k++, j--)
21 {
22 m_DimData[j].m_DimSize = outShape[j];
23 m_DimData[j].m_Stride1 = (inShape0[j] > 1) ? sIn0 : 0;
24 m_DimData[j].m_Stride2 = (inShape1[j] > 1) ? sIn1 : 0;
25 m_DimData[j].m_StrideOut = sOut;
26
27 sIn0 *= inShape0[j];
28 sIn1 *= inShape1[j];
29 sOut *= outShape[j];
30 }
31 }
32
BroadcastLoop(const TensorShape & inShape,const TensorShape & outShape)33 BroadcastLoop::BroadcastLoop(const TensorShape& inShape, const TensorShape& outShape)
34 : m_DimData(outShape.GetNumDimensions())
35 {
36 const unsigned int numDims = GetNumDimensions();
37
38 unsigned int sIn = 1;
39 unsigned int sOut = 1;
40
41 for (unsigned int j = numDims - 1, k = 0; k < numDims ; k++, j--)
42 {
43 m_DimData[j].m_DimSize = outShape[j];
44 m_DimData[j].m_Stride1 = (inShape[j] > 1) ? sIn : 0;
45 m_DimData[j].m_StrideOut = sOut;
46
47 sIn *= inShape[j];
48 sOut *= outShape[j];
49 }
50 }
51
52 } // namespace armnn
53