blob: f17ec6b3110521b33fb271798f392c94b509cafa (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
|
//
// Copyright © 2019,2024 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#include "Broadcast.hpp"
namespace armnn
{
BroadcastLoop::BroadcastLoop(const TensorShape& inShape0, const TensorShape& inShape1, const TensorShape& outShape)
: m_DimData(outShape.GetNumDimensions())
{
const unsigned int numDims = GetNumDimensions();
unsigned int sIn0 = 1;
unsigned int sIn1 = 1;
unsigned int sOut = 1;
for (unsigned int j = numDims - 1, k = 0; k < numDims ; k++, j--)
{
m_DimData[j].m_DimSize = outShape[j];
m_DimData[j].m_Stride1 = (inShape0[j] > 1) ? sIn0 : 0;
m_DimData[j].m_Stride2 = (inShape1[j] > 1) ? sIn1 : 0;
m_DimData[j].m_StrideOut = sOut;
sIn0 *= inShape0[j];
sIn1 *= inShape1[j];
sOut *= outShape[j];
}
}
BroadcastLoop::BroadcastLoop(const TensorShape& inShape, const TensorShape& outShape)
: m_DimData(outShape.GetNumDimensions())
{
const unsigned int numDims = GetNumDimensions();
unsigned int sIn = 1;
unsigned int sOut = 1;
// Get the difference between the output dimension and input dimension
const unsigned int dimDifference = numDims - inShape.GetNumDimensions();
for (unsigned int j = numDims - 1, k = 0; k < numDims ; k++, j--)
{
m_DimData[j].m_DimSize = outShape[j];
// Pretend there are extra 1-dimensional tensors prepended
if (dimDifference > 0 && j < dimDifference)
{
m_DimData[j].m_Stride1 = 0;
sIn *= 1;
}
else if (dimDifference > 0)
{
m_DimData[j].m_Stride1 = (inShape[j - dimDifference] > 1) ? sIn : 0;
sIn *= inShape[j - dimDifference];
}
else
{
m_DimData[j].m_Stride1 = (inShape[j] > 1) ? sIn : 0;
sIn *= inShape[j];
}
m_DimData[j].m_StrideOut = sOut;
sOut *= outShape[j];
}
}
} // namespace armnn
|
