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
|
/*
* Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "passes/optimizations/CombineTransposes.h"
#include "mir/ops/TransposeOp.h"
#include "mir/Graph.h"
#include "mir/GraphPatternMatcher.h"
#include <algorithm>
namespace nnc
{
using namespace mir;
std::vector<size_t> combineAxisOrders(const std::vector<std::size_t> &order1,
const std::vector<std::size_t> &order2)
{
assert(order1.size() == order2.size());
std::vector<size_t> res(order1.size());
for (size_t i = 0; i < order1.size(); i++)
{
res[order2[order1[i]]] = i;
}
return res;
}
static bool isIdentityTranspose(const std::vector<size_t> &axis_order)
{
for (size_t i = 0; i < (axis_order.size()); i++)
{
if (axis_order[i] != i)
{
return false;
}
}
return true;
}
nnc::PassData nnc::CombineTransposes::run(nnc::PassData data)
{
auto g = static_cast<Graph *>(data);
assert(g);
GraphPatternMatcher matcher(g);
auto is_tr = [](const Operation *op1) { return op1->getType() == Operation::Type::transpose; };
std::vector<std::pair<Operation *, Operation *>> matches = matcher.matchEdge(is_tr, is_tr);
std::unordered_set<Operation *> deleted_nodes;
while (!matches.empty())
{
for (std::pair<Operation *, Operation *> match : matches)
{
if (deleted_nodes.find(match.first) != deleted_nodes.end())
{
break;
};
auto *top_transpose = dynamic_cast<mir::ops::TransposeOp *>(match.first);
if (deleted_nodes.find(match.second) != deleted_nodes.end())
{
break;
};
auto *bottom_transpose = dynamic_cast<mir::ops::TransposeOp *>(match.second);
auto combined_axis_order =
combineAxisOrders(top_transpose->getAxisOrder(), bottom_transpose->getAxisOrder());
if (!isIdentityTranspose(combined_axis_order))
{
auto new_tr_op =
g->create<mir::ops::TransposeOp>(top_transpose->getInput(0), combined_axis_order);
g->replaceNode(bottom_transpose, new_tr_op);
}
else
{
// Connect top input to all outputs of bottom
Operation *top = top_transpose->getInput(0)->getNode();
g->replaceNode(bottom_transpose, top);
}
deleted_nodes.emplace(bottom_transpose);
if (top_transpose->getOutput(0)->getUses().empty())
{
g->removeNode(top_transpose);
deleted_nodes.emplace(top_transpose);
}
}
matches = matcher.matchEdge(is_tr, is_tr);
};
return g;
}
} // namespace nnc
|
