ArgMinMax.cpp

Go to the documentation of this file.

//
// Copyright © 2019 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
 
#include "ArgMinMax.hpp"
 
#include <armnnUtils/TensorUtils.hpp>
 
#include <armnn/utility/IgnoreUnused.hpp>
#include <armnn/utility/NumericCast.hpp>
 
namespace armnn
{
 
template <typename OUT>
void ArgMinMax(Decoder<float>& in, OUT* out, const TensorInfo& inputTensorInfo,
               const TensorInfo& outputTensorInfo, ArgMinMaxFunction function, int axis)
{
    IgnoreUnused(outputTensorInfo);
 
    unsigned int uAxis = armnnUtils::GetUnsignedAxis(inputTensorInfo.GetNumDimensions(), axis);
 
    const unsigned int outerElements = armnnUtils::GetNumElementsBetween(inputTensorInfo.GetShape(), 0, uAxis);
    const unsigned int axisSize = inputTensorInfo.GetShape()[uAxis];
    const unsigned int innerElements = armnnUtils::GetNumElementsBetween(inputTensorInfo.GetShape(),
                                                                         uAxis + 1,
                                                                         inputTensorInfo.GetNumDimensions());
 
    for (unsigned int outer = 0; outer < outerElements; ++outer) {
        for (unsigned int inner = 0; inner < innerElements; ++inner) {
            in[outer * axisSize * innerElements + inner];
            auto tmpValue = in.Get();
            unsigned int tmpIndex = 0;
            for (unsigned int i = 1; i < axisSize; ++i) {
                in[(outer * axisSize * innerElements) + (i * innerElements) + inner];
                const auto& value = in.Get();
                if ((function == armnn::ArgMinMaxFunction::Min && value < tmpValue) ||
                    (function == armnn::ArgMinMaxFunction::Max &&  value > tmpValue)) {
                    tmpValue = value;
                    tmpIndex = i;
                }
            }
 
            out[outer * innerElements + inner] = armnn::numeric_cast<OUT>(tmpIndex);
        }
    }
}
 
template void ArgMinMax(Decoder<float>& in, int32_t* out, const TensorInfo& inputTensorInfo,
               const TensorInfo& outputTensorInfo, ArgMinMaxFunction function, int axis);
 
template void ArgMinMax(Decoder<float>& in, int64_t* out, const TensorInfo& inputTensorInfo,
               const TensorInfo& outputTensorInfo, ArgMinMaxFunction function, int axis);
 
} //namespace armnn