[WIP][JIT][Fork/Join] Stage 1: Eager fork-join#12405
[WIP][JIT][Fork/Join] Stage 1: Eager fork-join#12405jamesr66a wants to merge 2 commits intopytorch:masterfrom
Conversation
control flow via exception
| friend struct Type; | ||
| template<typename ... T> | ||
| static FutureTypePtr create( T&& ... all) { | ||
| return FutureTypePtr(new FutureType( std::forward<T>(all)... )); // NOLINT(modernize-make-shared) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| // TODO schema | ||
| Operator(prim::Fork, [](Node* stack) -> Operation { | ||
| return [=](Stack& stack) -> int { | ||
| throw detail::NewCoroutine(); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| } | ||
| }; | ||
|
|
||
| Future fork(script::Module &sm, std::vector<IValue> inputs) { |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| return torch.neg(x) | ||
|
|
||
| x = torch.rand(3, 4) | ||
| fut = torch.jit.Fork(foo, (x,)) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| }; | ||
| }), | ||
| // TODO: type variable | ||
| Operator(prim::Wait, [](Node* stack) -> Operation { |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| Method & m, | ||
| at::ArrayRef<NamedValue> inputs, | ||
| at::ArrayRef<NamedValue> attributes, | ||
| const std::vector<std::shared_ptr<SugaredValue>> &blocks, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| ->setSourceLocation(std::make_shared<SourceRange>(loc)); | ||
| auto body_block = n->addBlock(); | ||
|
|
||
| std::vector<NamedValue> block_inputs; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| tree->range(), emitExpr(Expr(tree)))); | ||
| auto sugared_expr = emitSugaredExpr(Expr(tree), 1); | ||
| // TODO: non-sketchy check | ||
| if (collect_blocks && sugared_expr->kind() == "module") { |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| if not (TEST_WITH_UBSAN and test_name in UBSAN_BLACKLISTED_TESTS): | ||
| setattr(TestJitGenerated, test_name, do_test) | ||
|
|
||
| class TestAsync(JitTestCase): |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| // TODO schema | ||
| Operator(prim::Fork, [](Node* stack) -> Operation { | ||
| return [=](Stack& stack) -> int { | ||
| throw detail::NewCoroutine(); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| } | ||
| } | ||
| JIT_ASSERT(body_block->outputs().size() == 1); | ||
| n->output()->replaceAllUsesWith(value_remap[body_block->outputs()[0]]); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| Method& m, | ||
| at::ArrayRef<NamedValue> inputs_, | ||
| at::ArrayRef<NamedValue> attributes, | ||
| const std::vector<std::shared_ptr<SugaredValue>> &blocks, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
Superseded by #12925 |
Establishing the front-end APIs for explicit fork/join parallelism. At runtime the forked sections are inlined, and wait is eliminated.
Forthcoming work is to implement lazy evaluation via coroutines. After that we can add true parallelism.
TODO: Figure out if this works in the string frontend
TODO: Remove remnants of coroutine implementation that got left in while i did a sketchy rebase