Repository: pytorch
Issue: [JIT] Inconsistent behavior of torchScript parser in python and c++
Labels: oncall: jit

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
I am trying to writing some unit tests for torchScript interpreter with utilization of torch.jit.fork. In detail,  I tried to compile the following script in C++ with the complie API provided in JIT:
```
const auto moduleExample = R"JIT(
  def run_parallel():
    r1 = torch.jit.fork(torch.mm, torch.rand(10000,1000),torch.rand(1000,5000))
    r2 = torch.jit.fork(torch.mm, torch.rand(10000,1000),torch.rand(1000,5000))
    x1 = r1.wait()
    x2 = r2.wait()
    return x1 + x2
)JIT";
  auto model = compile(moduleExample);
```
The function is able to be consumed by `torch.jit.script` API but not able to be compiled in C++ due to error msg:
```
attribute lookup is not defined on builtin:
```
As a workaround, I tried to use `torch.jit.script(run_parallel).code` to get the string as below:
```
def run_parallel() -> Tensor:
    _0 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
    _1 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
    def __forked_function():
      return torch.mm(_0, _1)
    r1 = fork(__forked_function)
    _2 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
    _3 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
    def __forked_function0():
      return torch.mm(_2, _3)
    r2 = fork(__forked_function0)
    _4 = torch.add(torch.wait(r1), torch.wait(r2), alpha=1)
    return _4
```
I tried to compile this script again in C++, but still failed with error msg:
```
undefined value fork:
```
## To Reproduce

Steps to reproduce the behavior:
with the function
```
 def run_parallel():
    r1 = torch.jit.fork(torch.mm, torch.rand(10000,1000),torch.rand(1000,5000))
    r2 = torch.jit.fork(torch.mm, torch.rand(10000,1000),torch.rand(1000,5000))
    x1 = r1.wait()
    x2 = r2.wait()
    return x1 + x2
```
1. do `torch.jit.script(run_parallel).graph`
```
graph():
  %43 : int = prim::Constant[value=1]()
  %4 : None = prim::Constant()
  %0 : int = prim::Constant[value=1000]() # <ipython-input-36-b52c56ea0569>:6:45
  %3 : int[] = prim::ListConstruct(%0, %0)
  %8 : Tensor = aten::rand(%3, %4, %4, %4, %4) # <ipython-input-36-b52c56ea0569>:6:34
  %11 : int[] = prim::ListConstruct(%0, %0)
  %16 : Tensor = aten::rand(%11, %4, %4, %4, %4) # <ipython-input-36-b52c56ea0569>:6:56
  %r1.1 : Future[Tensor] = prim::fork_0(%8, %16) # <ipython-input-36-b52c56ea0569>:6:9
  %22 : int[] = prim::ListConstruct(%0, %0)
  %27 : Tensor = aten::rand(%22, %4, %4, %4, %4) # <ipython-input-36-b52c56ea0569>:7:34
  %30 : int[] = prim::ListConstruct(%0, %0)
  %35 : Tensor = aten::rand(%30, %4, %4, %4, %4) # <ipython-input-36-b52c56ea0569>:7:56
  %r2.1 : Future[Tensor] = prim::fork_1(%27, %35) # <ipython-input-36-b52c56ea0569>:7:9
  %40 : Tensor = aten::wait(%r1.1) # <ipython-input-36-b52c56ea0569>:8:11
  %42 : Tensor = aten::wait(%r2.1) # <ipython-input-36-b52c56ea0569>:8:23
  %44 : Tensor = aten::add(%40, %42, %43) # <ipython-input-36-b52c56ea0569>:8:11
  return (%44)
with prim::fork_0 = graph(%4 : Tensor,
      %5 : Tensor):
  %3 : Tensor = aten::mm(%4, %5) # <ipython-input-36-b52c56ea0569>:6:9
  return (%3)
with prim::fork_1 = graph(%4 : Tensor,
      %5 : Tensor):
  %3 : Tensor = aten::mm(%4, %5) # <ipython-input-36-b52c56ea0569>:7:9
  return (%3)
```
2. do `torch.jit.script(run_parallel).code`:
```
def run_parallel() -> Tensor:
  _0 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
  _1 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
  def __forked_function():
    return torch.mm(_0, _1)
  r1 = fork(__forked_function)
  _2 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
  _3 = torch.rand([1000, 1000], dtype=None, layout=None, device=None, pin_memory=None)
  def __forked_function0():
    return torch.mm(_2, _3)
  r2 = fork(__forked_function0)
  _4 = torch.add(torch.wait(r1), torch.wait(r2), alpha=1)
  return _4
```
3. complie the scripts in C++, neither of the original script nor .code string works.


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
C++ compiler should be able consume the .code string from Python API.
## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

 - PyTorch Version (e.g., 1.0): 1.8
 - OS (e.g., Linux): centOS
 - How you installed PyTorch (`conda`, `pip`, source): 
 - Build command you used (if compiling from source):
 - Python version: N/A
 - CUDA/cuDNN version: N/A
 - GPU models and configuration: N/A
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->


cc @gmagogsfm