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
|
#include <torch/csrc/jit/passes/add_if_then_else.h>
#include <torch/csrc/jit/runtime/graph_iterator.h>
namespace torch::jit {
namespace {
bool hasNoNodes(Block* block) {
auto nodes = block->nodes();
return nodes.begin() == nodes.end();
}
bool hasTrivialSubBlocks(Node* node) {
const auto blocks = node->blocks();
TORCH_DCHECK_EQ(blocks.size(), 2);
return hasNoNodes(blocks[0]) && hasNoNodes(blocks[1]);
}
} // namespace
bool AddIfThenElseOp(std::shared_ptr<Graph>& graph) {
std::vector<Node*> to_replace;
DepthFirstGraphNodeIterator graph_it(graph);
for (auto* node = graph_it.next(); node != nullptr; node = graph_it.next()) {
if (node->kind() != prim::If) {
continue;
}
if (node->outputs().size() != 1) {
continue;
}
if (hasTrivialSubBlocks(node)) {
to_replace.push_back(node);
}
}
for (auto* node : to_replace) {
auto* if_then_else_node = graph->create(prim::IfThenElse, 1);
if_then_else_node->addInput(node->input());
auto blocks = node->blocks();
if_then_else_node->addInput(blocks[0]->return_node()->input());
if_then_else_node->addInput(blocks[1]->return_node()->input());
if_then_else_node->insertBefore(node);
if_then_else_node->output()->copyMetadata(node->output());
node->output()->replaceAllUsesWith(if_then_else_node->output());
node->destroy();
}
return !to_replace.empty();
}
} // namespace torch::jit
|
