mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
77f2883c41
This PR relands the changes introduced in PR https://github.com/pytorch/pytorch/pull/100849. The old PR turnd nnc_* functions into static. We now add declarations for them and hope that inter builds will pass. Pull Request resolved: https://github.com/pytorch/pytorch/pull/102228 Approved by: https://github.com/albanD
64 lines
1.7 KiB
C++
64 lines
1.7 KiB
C++
#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::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*)
|
|
// }
|
|
static 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);
|
|
}
|
|
|
|
static 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 torch::jit
|