[ts-migration] support aten::__is__, aten::__isnot__, aten::__not__, profiler::_record_function_enter_new, profiler::_record_function_exit (pytorch#127656)

Support more ops in ts converter and add unit tests.
Pull Request resolved: pytorch#127656
Approved by: https://github.com/SherlockNoMad

Author: BoyuanFeng

test/export/test_converter.py

@@ -1,6 +1,7 @@
 # Owner(s): ["oncall: export"]
 
 import unittest
+from typing import Tuple
 
 import torch
 
@@ -9,7 +10,6 @@
 from torch._dynamo.test_case import TestCase
 from torch._export.converter import TS2EPConverter
 from torch.export import ExportedProgram
-
 from torch.testing._internal.common_utils import run_tests
 
 requires_cuda = unittest.skipUnless(torch.cuda.is_available(), "requires cuda")
@@ -23,8 +23,11 @@ def _check_equal_ts_ep_converter(self, mod, inp) -> ExportedProgram:
         orig_out, _ = pytree.tree_flatten(mod(*inp))
         self.assertEqual(len(ep_out), len(orig_out))
         for ep_t, orig_t in zip(ep_out, orig_out):
-            self.assertEqual(ep_t.shape, orig_t.shape)
-            self.assertTrue(torch.allclose(ep_t, orig_t))
+            if isinstance(ep_t, torch.Tensor):
+                self.assertEqual(ep_t.shape, orig_t.shape)
+                self.assertTrue(torch.allclose(ep_t, orig_t))
+            else:
+                self.assertEqual(ep_t, orig_t)
         return ep
 
     def test_ts2ep_converter_basic(self):
@@ -192,6 +195,58 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
             M()(torch.tensor(False), torch.tensor(4)),
         )
 
+    def test_profiler__record_function(self):
+        class Module(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                handle = torch.ops.profiler._record_function_enter_new("foo", None)
+                y = x * 2 + 4
+                torch.ops.profiler._record_function_exit(handle)
+                return y
+
+        x = torch.randn(10, 10)
+        self._check_equal_ts_ep_converter(Module(), (x,))
+
+    def test_aten_floordiv(self):
+        class Module(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return x // 2
+
+        x = torch.randn(10, 10)
+        self._check_equal_ts_ep_converter(Module(), (x,))
+
+    def test_aten___is__(self):
+        class Module(torch.nn.Module):
+            def forward(
+                self, x: torch.Tensor, y: torch.Tensor
+            ) -> Tuple[bool, torch.Tensor]:
+                z = x + 1
+                return x is y, z
+
+        inp = (torch.randn(10, 10), torch.rand(10, 10))
+        self._check_equal_ts_ep_converter(Module(), inp)
+
+    def test_aten___isnot__(self):
+        class Module(torch.nn.Module):
+            def forward(
+                self, x: torch.Tensor, y: torch.Tensor
+            ) -> Tuple[bool, torch.Tensor]:
+                z = x + 1
+                return x is not y, z
+
+        inp = (torch.randn(10, 10), torch.rand(10, 10))
+        self._check_equal_ts_ep_converter(Module(), inp)
+
+    def test_aten___not__(self):
+        class Module(torch.nn.Module):
+            def forward(
+                self, x: torch.Tensor, y: torch.Tensor
+            ) -> Tuple[bool, torch.Tensor]:
+                z = x + 1
+                return not (x is not y), z
+
+        inp = (torch.randn(10, 10), torch.rand(10, 10))
+        self._check_equal_ts_ep_converter(Module(), inp)
+
 
 if __name__ == "__main__":
     run_tests()

torch/_export/converter.py

@@ -41,6 +41,15 @@ def normalize_name(name: str) -> str:
     return name.replace(".", "_")
 
 
+# Given a node: torch._C.Node, map from node.kind() to a standard operator
+kind_to_standard_operators = {
+    "prim::TupleIndex": operator.getitem,
+    "aten::__is__": operator.is_,
+    "aten::__isnot__": operator.is_not,
+    "aten::__not__": operator.not_,
+}
+
+
 def get_op_overload(node: torch._C.Node):
     schema_str = node.schema()
     schema = FunctionSchema.parse(schema_str)
@@ -285,13 +294,6 @@ def convert_prim_DictConstruct(self, node: torch._C.Node):
         output_name = node.output().debugName()
         self.name_to_node[output_name] = output_dict
 
-    def convert_prim_TupleIndex(self, node: torch._C.Node):
-        args = tuple(self.get_fx_value(input) for input in node.inputs())
-        getitem_node = self.fx_graph.call_function(operator.getitem, args)
-
-        output_name = node.output().debugName()
-        self.name_to_node[output_name] = getitem_node
-
     def convert_aten_Int(self, node: torch._C.Node):
         # converts aten::Int as aten._to_copy + aten::_local_scalar_dense
         target = torch.ops.aten._to_copy.default
@@ -438,6 +440,28 @@ 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
+        # discarded during `retrace_as_exported_program`.
+        target = torch.ops.profiler._record_function_exit
+        args = tuple(self.get_fx_value(input) for input in node.inputs())
+        self.fx_graph.call_function(target, args)
+
+    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())
+        fx_node = self.fx_graph.call_function(target, args)
+        output_name = node.output().debugName()
+        self.name_to_node[output_name] = fx_node
+
     def convert_node(self, node: torch._C.Node):
         node_kind = node.kind()
         if node_kind == "prim::CreateObject":
@@ -457,8 +481,6 @@ def convert_node(self, node: torch._C.Node):
             self.convert_prim_dtype(node)
         elif node_kind == "prim::DictConstruct":
             self.convert_prim_DictConstruct(node)
-        elif node_kind == "prim::TupleIndex":
-            self.convert_prim_TupleIndex(node)
         # elif node_kind == "aten::Int":
         #     convert_aten_Int(node)
         elif node_kind == "aten::_convolution":
@@ -471,7 +493,14 @@ def convert_node(self, node: torch._C.Node):
             self.convert_prim_if(node)
         elif node_kind == "aten::Bool":
             self.convert_as_noop(node)
+        elif node_kind == "profiler::_record_function_enter_new":
+            self.convert_profiler__record_function_enter_new(node)
+        elif node_kind == "profiler::_record_function_exit":
+            self.convert_profiler__record_function_exit(node)
+        elif node_kind in kind_to_standard_operators:
+            self.convert_standard_operators(node)
         elif node_kind.startswith("aten::"):
+            # order matters! this should be handled after kind_to_standard_operators
             self.convert_aten_op(node)
         else:
             raise ValueError(f"Unsupported node kind: {node_kind}")