aboutsummaryrefslogtreecommitdiff
path: root/delegate/src/Reduce.hpp
blob: b8db04ccf2ae541c9e10a25f0a4b95d6f074059a (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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
//
// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//

#pragma once

#include <tensorflow/lite/builtin_ops.h>
#include <tensorflow/lite/c/builtin_op_data.h>
#include <tensorflow/lite/c/common.h>
#include <tensorflow/lite/kernels/internal/tensor_ctypes.h>
#include <tensorflow/lite/minimal_logging.h>

namespace armnnDelegate
{

TfLiteStatus VisitReduceOperator(DelegateData& delegateData,
                                 TfLiteContext* tfLiteContext,
                                 TfLiteNode* tfLiteNode,
                                 int nodeIndex,
                                 int32_t reduceOperatorCode)
{
    TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 2, nodeIndex));
    TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));

    const TfLiteTensor* tfLiteTensors = tfLiteContext->tensors;
    const TfLiteTensor& tfLiteInputTensor = tfLiteTensors[tfLiteNode->inputs->data[0]];
    if (!IsValid(tfLiteContext, tfLiteInputTensor, reduceOperatorCode, nodeIndex))
    {
        return kTfLiteError;
    }

    const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]];
    if (!IsValid(tfLiteContext, tfLiteOutputTensor, reduceOperatorCode, nodeIndex))
    {
        return kTfLiteError;
    }

    const armnn::TensorInfo& inputTensorInfo  = GetTensorInfoForTfLiteTensor(tfLiteInputTensor);
    const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor);

    // Get const axis value from model and set it to descriptor.
    const TfLiteTensor& tfLiteAxisTensor = tfLiteTensors[tfLiteNode->inputs->data[1]];
    if (!IsValid(tfLiteContext, tfLiteAxisTensor, reduceOperatorCode, nodeIndex))
    {
        return kTfLiteError;
    }

    const armnn::TensorInfo& axisTensorInfo =   GetTensorInfoForTfLiteTensor(tfLiteAxisTensor);
    auto* axisTensorData = tflite::GetTensorData<int32_t>(&tfLiteAxisTensor);

    std::vector<int32_t> axis;
    // Add axis data to vector to be converter to unsigned int and assigned to descriptor axis.
    if (axisTensorData != nullptr)
    {
        for (unsigned int i = 0; i < axisTensorInfo.GetNumElements(); ++i)
        {
            axis.emplace_back(axisTensorData[i]);
        }
    }
    else
    {
        for (unsigned int i = 0; i < inputTensorInfo.GetNumDimensions(); ++i)
        {
            axis.push_back(i);
        }
    }

    // Convert the axis to unsigned int and remove duplicates.
    unsigned int rank = inputTensorInfo.GetNumDimensions();
    std::set<unsigned int> uniqueAxis;
    std::transform(axis.begin(),
                   axis.end(),
                   std::inserter(uniqueAxis, uniqueAxis.begin()),
                   [rank](int i)->unsigned int{ return (i + rank) % rank; });

    armnn::ReduceDescriptor desc;
    desc.m_vAxis.assign(uniqueAxis.begin(), uniqueAxis.end());

    auto* reducerParameters = reinterpret_cast<TfLiteReducerParams*>(tfLiteNode->builtin_data);
    desc.m_KeepDims = reducerParameters->keep_dims;
    if (reduceOperatorCode == kTfLiteBuiltinReduceMax)
    {
        desc.m_ReduceOperation = armnn::ReduceOperation::Max;
    }
    else if (reduceOperatorCode == kTfLiteBuiltinReduceMin)
    {
        desc.m_ReduceOperation = armnn::ReduceOperation::Min;
    }
    else if (reduceOperatorCode == kTfLiteBuiltinSum)
    {
        desc.m_ReduceOperation = armnn::ReduceOperation::Sum;
    }
    else if (reduceOperatorCode == kTfLiteBuiltinReduceProd)
    {
        desc.m_ReduceOperation = armnn::ReduceOperation::Prod;
    }
    else
    {
        TF_LITE_MAYBE_KERNEL_LOG(
            tfLiteContext,
            "TfLiteArmnnDelegate: Unsupported Reduction Operator #%d node #%d: ",
            reduceOperatorCode, nodeIndex);
        return kTfLiteError;
    }

    bool isSupported = false;
    auto validateFunc = [&](const armnn::TensorInfo& outInfo, bool& isSupported)
    {
        FORWARD_LAYER_SUPPORT_FUNC("REDUCE",
                                   tfLiteContext,
                                   IsReduceSupported,
                                   delegateData.m_Backends,
                                   isSupported,
                                   inputTensorInfo,
                                   outInfo,
                                   desc);
    };

    if (!delegateData.m_Network)
    {
        validateFunc(outputTensorInfo, isSupported);
        return isSupported ? kTfLiteOk : kTfLiteError;
    }

    // Add an Reduce layer
    armnn::IConnectableLayer* layer = delegateData.m_Network->AddReduceLayer(desc);
    ARMNN_ASSERT(layer != nullptr);

    armnn::IOutputSlot& outputSlot = layer->GetOutputSlot(0);
    outputSlot.SetTensorInfo(outputTensorInfo);

    // Connect
    return Connect(layer, tfLiteNode, delegateData);
}

} // namespace armnnDelegate