mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
b2e79ed5ec
Summary: Follow up to https://github.com/pytorch/pytorch/issues/68095 This also changes the files from the ATen folder to include c10's `Export.h` instead since they can't ever be exporting `TORCH_PYTHON_API`. cc pietern mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse SciPioneer H-Huang Pull Request resolved: https://github.com/pytorch/pytorch/pull/69585 Reviewed By: mrshenli Differential Revision: D32958594 Pulled By: albanD fbshipit-source-id: 1ec7ef63764573fa2b486928955e3a1172150061
66 lines
1.7 KiB
C++
66 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 {
|
|
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
|