[JIT] Fork/Join inline docs (#39952)

Elias Ellison · facebook-github-bot · commit 0d0608532ccc · 2020-06-22T15:34:05.000-07:00
Summary: Pull Request resolved: #39952 Differential Revision: D22165477 Pulled By: eellison fbshipit-source-id: 93132cd6987fdd2484112a57ef17912b8fcc5fab
diff --git a/test/jit/test_async.py b/test/jit/test_async.py
@@ -20,7 +20,7 @@ def foo(x):
             return torch.neg(x)
 
         x = torch.rand(3, 4)
-        fut = torch.jit._fork(foo, x)
+        fut = torch.jit.fork(foo, x)
         y_hat = foo(x)
         y = torch.jit._wait(fut)
         # assert nothing; only to make sure the fake python path works
@@ -32,7 +32,7 @@ def foo(inp):
                 futures.append(torch.jit._fork(lambda x: x, inp))
             all_outputs = []
             for future in futures:
-                all_outputs.append(torch.jit._wait(future))
+                all_outputs.append(torch.jit.wait(future))
             return all_outputs
 
         # assert nothing, just to make sure python type parsing works
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -789,6 +789,101 @@ def wrap_check_inputs(check_inputs):
 
     return [{'forward' : c} for c in check_inputs]
 
+
+def fork(func, *args, **kwargs):
+    """
+    Creates an asynchronous task executing `func` and a reference to the value
+    of the result of this execution. `fork` will return immediately,
+    so the return value of `func` may not have been computed yet. To force completion
+    of the task and access the return value invoke `torch.jit.wait` on the Future. `fork` invoked
+    with a `func` which returns `T` is typed as `torch.jit.Future[T]`. `fork` calls can be arbitrarily
+    nested, and may be invoked with positional and keyword arguments.
+
+    Asynchronous execution will only occur when run in TorchScript. If run in pure python,
+    `fork` will not execute in parallel. `fork` will also not execute in parallel when invoked
+    while tracing, however the `fork` and `wait` calls will be captured in the exported IR Graph.
+
+    Warning:
+        `fork` tasks will execute non-deterministicly. We recommend only spawning
+        parallel fork tasks for pure functions that do not modify their inputs,
+        module attributes, or global state.
+
+    Arguments:
+        func (callable or torch.nn.Module):  A Python function or `torch.nn.Module`
+            that will be invoked. If executed in TorchScript, it will execute asynchronously,
+            otherwise it will not. Traced invocations of fork will be captured in the IR.
+
+        *args, **kwargs: arguments to invoke `func` with.
+
+    Returns:
+        `torch.jit.Future[T]`: a reference to the execution of `func`. The value `T`
+        can only be accessed by forcing completion of `func` through `torch.jit.wait`.
+
+    Example (fork a free function):
+
+    .. testcode::
+
+        import torch
+        from torch import Tensor
+
+        def foo(a : Tensor, b : int) -> Tensor:
+            return a + b
+
+        def bar(a):
+            fut : torch.jit.Future[Tensor] = torch.jit.fork(foo, a, b=2)
+            return torch.jit.wait(fut)
+
+        script_bar = torch.jit.script(bar)
+        input = torch.tensor(2)
+
+        # only the scripted version executes asynchronously
+        assert script_bar(input) == bar(input)
+
+        # trace is not run asynchronously, but fork is captured in IR
+        graph = torch.jit.trace(bar, (input,)).graph
+        assert "fork" in str(graph)
+
+    Example (fork a module method):
+
+    .. testcode::
+
+        import torch
+        from torch import Tensor
+
+        class SubMod(torch.nn.Module):
+            def forward(self, a: Tensor, b : int):
+                return a + b
+
+        class Mod(torch.nn.Module):
+            def __init__(self):
+                super(self).__init__()
+                self.mod = SubMod()
+
+            def forward(self, input):
+                fut = torch.jit.fork(self.mod, a, b=2)
+                return torch.jit.wait(fut)
+
+        input = torch.tensor(2)
+        mod = Mod()
+
+        assert mod(input) == torch.jit.script(mod).forward(input)
+    """
+    return torch._C.fork(func, *args, **kwargs)
+
+def wait(future):
+    """
+    Forces completion of a `torch.jit.Future[T]` asynchronous task, returning the
+    result of the task. See :func:`~fork` for docs and examples.
+
+    Arguments:
+        func (torch.jit.Future[T]): an asynchronous task reference, created through `torch.jit.fork`
+
+    Returns:
+        `T`: the return value of the the completed task
+    """
+    return torch._C.wait(future)
+
+
 def trace(func,
           example_inputs,
           optimize=None,
