blob: 86c018868eda9ac8100bf5127d0d1419115d8885 (
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
|
//
// Copyright © 2020 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#pragma once
#include "Optimization.hpp"
#include <armnn/utility/PolymorphicDowncast.hpp>
#include <armnnUtils/Transpose.hpp>
namespace armnn
{
namespace optimizations
{
class MoveTransposeUpImpl
{
public:
/// Run for every connection between a base Layer (any) and a child TransposeLayer. If the type
/// of the base layer allows it, it moves the permutation to the inputs of the base layer.
/// I.e., adds equivalent permutations before the inputs of the base layer and moves the
/// connections in the output of the child transpose layer to the output of the base layer.
void Run(Graph& graph, InputSlot& connection) const
{
OutputSlot& baseOutput = *connection.GetConnectedOutputSlot();
if (baseOutput.GetNumConnections() == 1U)
{
Layer& base = baseOutput.GetOwningLayer();
if (CanMoveTransposeToInputs(base))
{
auto transpose = PolymorphicDowncast<TransposeLayer*>(&connection.GetOwningLayer());
const PermutationVector& perm = transpose->GetPermutation();
// Inserts an equivalent transpose before every input of the base layer.
for (auto baseInput = base.BeginInputSlots(); baseInput != base.EndInputSlots(); ++baseInput)
{
// Inserts a new transpose layer.
const std::string name = std::string("moved_up-") + transpose->GetName();
TransposeLayer& permLayer = *graph.InsertNewLayer<TransposeLayer>(*baseInput, perm, name.c_str());
// Sets output tensor info for the new layer.
OutputSlot& parentOutput = *permLayer.GetInputSlot(0).GetConnectedOutputSlot();
const TensorInfo permOutInfo = armnnUtils::TransposeTensorShape(parentOutput.GetTensorInfo(), perm);
permLayer.GetOutputHandler().SetTensorInfo(permOutInfo);
}
// Sets transposed output tensor info
const TensorInfo& childOutInfo = transpose->GetOutputHandler().GetTensorInfo();
base.GetOutputHandler().SetTensorInfo(childOutInfo);
// Bypasses transpose. It will be removed as it's left unconnected.
transpose->GetOutputSlot().MoveAllConnections(base.GetOutputSlot());
}
}
}
protected:
MoveTransposeUpImpl() = default;
~MoveTransposeUpImpl() = default;
private:
static bool CanMoveTransposeToInputs(const Layer& base)
{
switch (base.GetType())
{
case LayerType::Activation:
case LayerType::Addition:
case LayerType::FakeQuantization:
case LayerType::Floor:
case LayerType::MemCopy:
case LayerType::Multiplication:
return true;
default:
return false;
}
}
};
using MoveTransposeUp = OptimizeForConnection<Layer, TransposeLayer, MoveTransposeUpImpl>;
} // namespace optimizations
} // namespace armnn
|
