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
|
#include <functional>
#include <memory>
#include <string>
#include <torch/csrc/Export.h>
#include <torch/csrc/jit/frontend/inline_loop_condition.h>
#include <torch/csrc/jit/ir/ir.h>
namespace torch {
namespace jit {
void InlineBlockBeforeNode(Node* before_node, Block* block) {
for (auto it = block->nodes().begin(); it != block->nodes().end();) {
auto block_node = *it++;
block_node->moveBefore(before_node);
}
}
// The loop node is initially emitted as:
// Loop(max_trip_count)
// block0(loop_counter) {
// <body>
// }
// block1 {
// <loop condition>
// -> (condition)
// }
// Here, we inline the loop condition and convert the loop to the form:
// Loop(max_trip_count, start_condition)
// block0(loop_counter, loop_carried_block*) {
// <body>
// BlockExit(continue_condition, loop_carried_block*)
// }
void inlineLoopCondition(Node* n) {
Block* body_block = n->blocks().at(0);
auto pre_header = n->blocks().at(1);
auto temp_block = n->addBlock();
temp_block->cloneFrom(pre_header, [](Value* v) { return v; });
InlineBlockBeforeNode(n, temp_block);
n->insertInput(/*start_condition_index*/ 1, temp_block->outputs().at(0));
n->eraseBlock(2);
InlineBlockBeforeNode(body_block->return_node(), pre_header);
body_block->return_node()->insertInput(0, pre_header->outputs().at(0));
n->eraseBlock(1);
}
void inlineLoopCondition(Block* block) {
for (Node* n : block->nodes()) {
for (Block* b : n->blocks()) {
inlineLoopCondition(b);
}
if (n->kind() == prim::Loop) {
inlineLoopCondition(n);
}
}
}
void InlineLoopCondition(std::shared_ptr<Graph>& graph) {
inlineLoopCondition(graph->block());
}
} // namespace jit
} // namespace torch
|
