Broadcast.cpp

Go to the documentation of this file.

//
// Copyright © 2019 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;
 
    for (unsigned int j = numDims - 1, k = 0; k < numDims ; k++, j--)
    {
        m_DimData[j].m_DimSize = outShape[j];
        m_DimData[j].m_Stride1 = (inShape[j] > 1) ? sIn : 0;
        m_DimData[j].m_StrideOut = sOut;
 
        sIn *= inShape[j];
        sOut *= outShape[j];
    }
}
 
} // namespace armnn