[Torch] Add initial control flow support

[masahi]

This adds support for parsing Torchscript prim::If and prim::Loop nodes. This is also the first attempt at translating from the output of torch.jit.script(...). See the test cases for currently supported Python construct.

• prim::If can be easily translated by recursively parsing true and else branches. The spec https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/OVERVIEW.md#if
• prim::Loop requires more mork, but using Relay conditional and tail recursion it is not too hard. https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/OVERVIEW.md#loops

The related discussion (with an example IR dump): https://discuss.tvm.ai/t/discuss-adding-a-pytorch-frontend/5026/24

The CI is blocked by unrelated sphinx issue, but it is ready for review.

cc @zhiics @icemelon9 @wweic @jroesch @MarisaKirisame @alexwong @tqchen @junrushao1994 @ajtulloch @yinghai

[zhiics]

@masahi Thank you very much for the nice work! I am familiar the TF control-flow, and trying to read more about PyTorch control-flow constructs. Will have a careful review by tomorrow.

[alexwong]

Some small comments, also reading up on control flow in Torch.

[MarisaKirisame]

can you dont call it parse? parse is for converting strings to an ast. converting ast to ast is a converter.

nice work otherwise.

[masahi]

Tests passed!!

can you dont call it parse? parse is for converting strings to an ast. converting ast to ast is a converter.

no probelm, will do.

@alexwong @zhiics a good way to get familiar with torch script is to try and tweak simple examples. For example,

import torch


class SimpleLoop(torch.nn.Module):
    def forward(self, inp):
        a = inp
        for i in range(10):
            a += i
        return a


class SimpleWhileLoop(torch.nn.Module):
    def forward(self, inp):
        a = inp
        i = 0
        while i < 10:
            a += i
            i += 1
        return a


print(torch.jit.script(SimpleLoop()).graph)
print(torch.jit.script(SimpleWhileLoop()).graph)

would print

graph(%self : __torch__.SimpleLoop,
      %inp.1 : Tensor):
  %10 : int = prim::Constant[value=1]()
  %6 : bool = prim::Constant[value=1]() # test.py:7:8
  %3 : int = prim::Constant[value=10]() # test.py:7:23
  %a : Tensor = prim::Loop(%3, %6, %inp.1) # test.py:7:8
    block0(%i.1 : int, %a.5 : Tensor):
      %a.2 : Tensor = aten::add_(%a.5, %i.1, %10) # test.py:8:12
      -> (%6, %a.2)
  return (%a)

graph(%self : __torch__.SimpleWhileLoop,
      %inp.1 : Tensor):
  %i.1 : int = prim::Constant[value=0]() # test.py:15:12
  %4 : int = prim::Constant[value=9223372036854775807]() # test.py:16:8
  %6 : int = prim::Constant[value=10]() # test.py:16:18
  %16 : int = prim::Constant[value=1]() # test.py:18:17
  %39 : bool = prim::Constant[value=1]()
  %a : Tensor, %i : int = prim::Loop(%4, %39, %inp.1, %i.1) # test.py:16:8
    block0(%8 : int, %a.5 : Tensor, %i.8 : int):
      %a.2 : Tensor = aten::add_(%a.5, %i.8, %16) # test.py:17:12
      %i.6 : int = aten::add(%i.8, %16) # test.py:18:12
      %7 : bool = aten::lt(%i.6, %6) # test.py:16:14
      -> (%7, %a.2, %i.6)
  return (%a)

hopefully it is much simpler than tensorflow control flow.

[MarisaKirisame]

One thing to keep in mind is that torchscript has mutation, continue, while... So the design should be prepared to generate code in A-Normal Form for correctness. Right now it is not needed as they are all purely functional, but the architecture shouldnt make such a change as simple as possible (e.g. not requiring rewriting all the code).

[masahi]

@MarisaKirisame I have something relevant to share. I was told by one of the torchscript devs on their forum that he is working on a "functionalization" pass pytorch/pytorch#33020, and that could help my use case. I'm not sure what exactly it does, but from looking at their code it seems it tries to find a "functional" subset of nodes, i.e. nodes with no impure operation, and extract it as a subgraph. I guess it helps applying some of their optimizations more aggressively.

He also has others related PRs ongoing that aim at making the graph "more pure". One on removing inplace op pytorch/pytorch#33186, another on removing list append pytorch/pytorch#33199. I literally encountered these impure ops when I was trying to convert more realistic lstm models than the one I have in this PR, so these new feature in Torch could be useful to us later.

[MarisaKirisame]

// Functional Graphs are not responsible for maintaining aliasing
// relationships. If an output of a functional graph escapes scope
// or is mutated then we might change semantics of the program if
// aliasing relationships are changed.

... WTF
I dont think lots of smartness should go in the converter - suppose there is multiple converter, it need multiple smart hacks to remove reference or such.
A better design IMO is to use PartialEvaluator/DeadCodeElimination/SomeOtherCustomPass to remove the reference, but the converter will produce code will effect and reference without much cleaning.

[masahi]

@zhiics @MarisaKirisame @alexwong Comments were addressed and I have no plan of updating this PR further. Can we merge this?

[zhiics]

LGTM

@MarisaKirisame could you take a look and approve explicitly if it looks good to you as well?

[zhiics]

Thanks @masahi @MarisaKirisame @alexwong