[ts-migration] Support RaiseException, prim::Unitialized, prim::Enter, and prim::Exit (pytorch#129416)

- Support raise exception. It's behavior matches non-strict export now, thanks to @ydwu4's [PR](pytorch#128709).
- Support prim::Unitialized, prim::Enter, and prim::Exit
Pull Request resolved: pytorch#129416
Approved by: https://github.com/angelayi

Author: BoyuanFeng

test/export/test_converter.py

@@ -2,7 +2,7 @@
 
 import unittest
 from collections import OrderedDict
-from typing import Dict, List, Tuple, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import torch
 import torch.utils._pytree as pytree
@@ -942,6 +942,68 @@ def forward(self, x: torch.Tensor):
         inp = (torch.ones(1),)
         self._check_equal_ts_ep_converter(M, inp, ["script"], check_persistent=True)
 
+    def test_raise_exception(self):
+        class Module(torch.nn.Module):
+            def forward(self, x: torch.Tensor, y: int) -> torch.Tensor:
+                if y > 0:
+                    raise RuntimeError("test")
+                return x + y
+
+        # match non-strict export behavior that errors when the given input leads to
+        # RaiseException.
+        with self.assertRaisesRegex(torch.jit.Error, "builtins.RuntimeError"):
+            inp = (torch.randn(3, 2), 1)
+            self._check_equal_ts_ep_converter(Module(), inp, ["script"])
+
+        # Matching non-strict export behavior that only executes 1 if-branch according
+        # to the given input.
+        inp = (torch.randn(3, 2), 0)
+        self._check_equal_ts_ep_converter(Module(), inp, ["script"])
+
+        class Module(torch.nn.Module):
+            def forward(self, x: torch.Tensor, y: int) -> torch.Tensor:
+                z = x
+                if y > 0:
+                    raise RuntimeError("test")
+                    # z = x
+                else:
+                    z = x + y
+                return x + y + z
+
+        # match non-strict export behavior that errors when the given input leads to
+        # RaiseException.
+        with self.assertRaisesRegex(torch.jit.Error, "builtins.RuntimeError"):
+            inp = (torch.randn(3, 2), 1)
+            self._check_equal_ts_ep_converter(Module(), inp, ["script"])
+
+        # Matching non-strict export behavior that only executes 1 if-branch according
+        # to the given input.
+        inp = (torch.randn(3, 2), 0)
+        self._check_equal_ts_ep_converter(Module(), inp, ["script"])
+
+    def test_context_manager(self):
+        class ContextManager:
+            def __init__(self):
+                self.count = 0
+                return
+
+            def __enter__(self):
+                self.count += 1
+                return
+
+            def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+                self.count -= 1
+                return
+
+        class M(torch.nn.Module):
+            def forward(self, x, y):
+                with ContextManager():
+                    res = x + y
+                return res
+
+        inp = (torch.ones(3, 3), torch.ones(3, 3))
+        self._check_equal_ts_ep_converter(M(), inp)
+
     def test_hidden_input_name(self):
         @torch.jit.script
         def func1(x):

torch/_export/converter.py

@@ -510,8 +510,12 @@ def convert_call_function_op(self, node: torch._C.Node):
         # TODO: covnert sourceRange() into stack_trace
         # fx_node.meta["stack_trace"] = node.sourceRange()
 
-        output_name = node.output().debugName()
-        self.name_to_node[output_name] = fx_node
+        outs = tuple(node.outputs())
+        if len(outs) == 1:
+            output_name = node.output().debugName()
+            self.name_to_node[output_name] = fx_node
+        elif len(outs) > 1:
+            raise RuntimeError("Number of outputs > 1 is not supported yet")
 
     def convert_prim_TupleConstruct(self, node: torch._C.Node):
         self._convert_prim_iterator(node)
@@ -743,12 +747,27 @@ def _identify_inputs_as_arguments(entry):
 
         cond_node = self.fx_graph.call_function(torch.cond, args, {})
 
-        output_name = node.output().debugName()
-        self.name_to_node[output_name] = cond_node
+        outs = tuple(node.outputs())
+        if len(outs) == 1:
+            output_name = node.output().debugName()
+            self.name_to_node[output_name] = cond_node
+        elif len(outs) > 1:
+            raise RuntimeError("Number of outputs > 1 is not supported yet")
 
     def convert_aten_Bool(self, node: torch._C.Node):
         self._convert_as_noop(node)
 
+    def convert_prim_Enter(self, node: torch._C.Node):
+        # export generally treats prim::Enter as noop
+        # The only context manager export supports is aten::enable_grad.
+        # Unfortunately, TorchScript does not support aten::enable_grad yet.
+        # TODO: support aten::enable_grad in both TorchScript and Converter.
+        return
+
+    def convert_prim_Exit(self, node: torch._C.Node):
+        # export treats prim::Exit as noop
+        return
+
     def _convert_as_noop(self, node: torch._C.Node):
         # Converts the node as a no-op by mapping its output node as arg[0]
 
@@ -760,13 +779,6 @@ def _convert_as_noop(self, node: torch._C.Node):
         output_name = node.output().debugName()
         self.name_to_node[output_name] = args[0]
 
-    def convert_profiler__record_function_enter_new(self, node: torch._C.Node):
-        target = torch.ops.profiler._record_function_enter_new
-        args = tuple(self.get_fx_value(input) for input in node.inputs())
-        fx_node = self.fx_graph.call_function(target, args)
-        output_name = node.output().debugName()
-        self.name_to_node[output_name] = fx_node
-
     def convert_profiler__record_function_exit(self, node: torch._C.Node):
         # _record_function_exit has side effect so we keep it in fx.graph
         # currently, _record_function_enter_new and _record_function_exit are
@@ -784,6 +796,14 @@ def convert_prim_tolist(self, node: torch._C.Node):
         output_name = node.output().debugName()
         self.name_to_node[output_name] = fx_node
 
+    def convert_prim_Uninitialized(self, node: torch._C.Node):
+        # `prim::Uninitialized` is inserted by the compiler when it can prove
+        # the value will never be used. It can be introduced by exceptions,
+        # breaks, continues, and returns.
+        # So we add a dummy constant to the graph.
+        output_name = node.output().debugName()
+        self.constant_map[output_name] = torch.Tensor()
+
     def _convert_standard_operators(self, node: torch._C.Node):
         target = kind_to_standard_operators[node.kind()]
         args = tuple(self.get_fx_value(input) for input in node.inputs())
@@ -836,10 +856,10 @@ def convert_graph_outputs(self):
                 )
             else:
                 raise ValueError(f"Output {output_name} not found")
-
-        self.fx_graph.output(
-            args[0]
-        )  # Get rid of an extra list wrapped around final output.
+        if args:
+            self.fx_graph.output(
+                args[0]
+            )  # Get rid of an extra list wrapped around final output.
 
 
 class ExplainTS2FXGraphConverter(TS2FXGraphConverter):