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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#include "SubGraphSelector.hpp"
#include "Graph.hpp"
#include <boost/assert.hpp>
#include <algorithm>
#include <unordered_map>
namespace armnn
{
namespace
{
struct LayerSelectionInfo
{
using LayerInfoContainer = std::unordered_map<Layer*, LayerSelectionInfo>;
static constexpr uint32_t InitialSplitId() { return 1; }
LayerSelectionInfo(Layer* layer, const SubGraphSelector::LayerSelectorFunction& selector)
: m_Layer{layer}
, m_SplitId{0}
, m_IsSelected{selector(*layer)}
{
// fill topology information by storing direct children
for (auto&& slot = m_Layer->BeginOutputSlots(); slot != m_Layer->EndOutputSlots(); ++slot)
{
for (InputSlot* childLayerInputSlot : slot->GetConnections())
{
Layer& childLayer = childLayerInputSlot->GetOwningLayer();
m_DirectChildren.push_back(&childLayer);
}
}
}
void MarkChildrenSplits(LayerInfoContainer& network,
uint32_t splitId,
bool prevSelected)
{
if (m_SplitId < splitId)
{
m_SplitId = splitId;
}
// introduce a new split point at all non-selected points, but only if the
// previous point was selected. this prevents creating a new subgraph at
// every non-selected layer
if (!m_IsSelected && prevSelected)
{
++m_SplitId;
}
for (auto& layer : m_DirectChildren)
{
auto it = network.find(layer);
BOOST_ASSERT_MSG(it != network.end(), "All layers must be part of the topology.");
if (it != network.end())
{
it->second.MarkChildrenSplits(network, m_SplitId, m_IsSelected);
}
}
}
bool IsInputLayer() const
{
return m_Layer->GetType() == armnn::LayerType::Input;
}
void CollectNonSelectedInputs(SubGraph::InputSlots& slots,
const SubGraphSelector::LayerSelectorFunction& selector)
{
for (auto&& slot = m_Layer->BeginInputSlots(); slot != m_Layer->EndInputSlots(); ++slot)
{
OutputSlot* parentLayerOutputSlot = slot->GetConnectedOutputSlot();
BOOST_ASSERT_MSG(parentLayerOutputSlot != nullptr, "The slots must be connected here.");
if (parentLayerOutputSlot)
{
Layer& parentLayer = parentLayerOutputSlot->GetOwningLayer();
if (selector(parentLayer) == false)
{
slots.push_back(&(*slot));
}
}
}
}
void CollectNonSelectedOutputSlots(SubGraph::OutputSlots& slots,
const SubGraphSelector::LayerSelectorFunction& selector)
{
for (auto&& slot = m_Layer->BeginOutputSlots(); slot != m_Layer->EndOutputSlots(); ++slot)
{
for (InputSlot* childLayerInputSlot : slot->GetConnections())
{
Layer& childLayer = childLayerInputSlot->GetOwningLayer();
if (selector(childLayer) == false)
{
slots.push_back(&(*slot));
}
}
}
}
std::vector<Layer*> m_DirectChildren;
Layer* m_Layer;
uint32_t m_SplitId;
bool m_IsSelected;
};
} // namespace <anonymous>
SubGraphSelector::SubGraphs
SubGraphSelector::SelectSubGraphs(Graph& graph,
const LayerSelectorFunction& selector)
{
LayerSelectionInfo::LayerInfoContainer layerInfo;
for (auto& layer : graph)
{
layerInfo.emplace(layer, LayerSelectionInfo{layer, selector});
}
uint32_t splitNo = LayerSelectionInfo::InitialSplitId();
for (auto& info : layerInfo)
{
if (info.second.IsInputLayer())
{
// for each input layer we mark the graph where subgraph
// splits need to happen because of the dependency between
// the selected and non-selected nodes
info.second.MarkChildrenSplits(layerInfo, splitNo, false);
}
}
// Collect all selected layers keyed by split id into a map
using SelectionInfoPtrs = std::vector<LayerSelectionInfo*>;
std::unordered_map<uint32_t, SelectionInfoPtrs> splitMap;
for (auto& info : layerInfo)
{
if (info.second.m_IsSelected)
{
auto it = splitMap.find(info.second.m_SplitId);
if (it == splitMap.end())
{
splitMap.insert(std::make_pair(info.second.m_SplitId, SelectionInfoPtrs{&info.second}));
}
else
{
it->second.push_back(&info.second);
}
}
}
// Now each non-empty split id represents a subgraph
SubGraphs result;
for (auto& splitGraph : splitMap)
{
if (splitGraph.second.empty() == false)
{
SubGraph::OutputSlots outputs;
SubGraph::InputSlots inputs;
SubGraph::Layers layers;
for (auto&& infoPtr : splitGraph.second)
{
infoPtr->CollectNonSelectedOutputSlots(outputs, selector);
infoPtr->CollectNonSelectedInputs(inputs, selector);
layers.insert(infoPtr->m_Layer);
}
result.emplace_back(
std::make_unique<SubGraph>(
std::move(inputs),
std::move(outputs),
std::move(layers)));
}
}
return result;
}
} // namespace armnn
|
