[export] Support torch.sym* ops (pytorch#115854)

Fixes pytorch#108830 and pytorch/executorch#1379 (comment)

Pull Request resolved: pytorch#115854
Approved by: https://github.com/zhxchen17

Author: angelayi

test/export/test_export.py

@@ -1934,6 +1934,19 @@ def forward(self, x, y):
         self.assertEqual(inputs[0][0] * 2.0, inputs_model[0][0])
         self.assertEqual(inputs[0][0] * 2.0, inputs_export[0][0])
 
+    @testing.expectedFailureNonStrict
+    def test_sym_sqrt(self):
+        import math
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x / torch.sym_sqrt(x.shape[0])
+
+        ep = export(M(), (torch.ones(16, 4),), dynamic_shapes={'x': {0: Dim("dim")}})
+        _ExportPassBase()(ep.graph_module)
+        FileCheck().check_count(
+            "torch.sym_sqrt", 1, exactly=True
+        ).run(ep.graph_module.code)
+
     def test_check_specialized_int(self):
         class SingleOp(torch.nn.Module):
             def __init__(self):

test/export/test_serialize.py

@@ -11,7 +11,7 @@
 
 import torch
 import torch._dynamo as torchdynamo
-from torch.export import export, save, load
+from torch.export import export, save, load, Dim
 from torch._export.db.case import ExportCase, normalize_inputs, SupportLevel
 from torch._export.db.examples import all_examples
 from torch._export.serde.serialize import (
@@ -471,6 +471,15 @@ def forward(self, x, y, z):
         inputs = (torch.rand(8, 8, 8), torch.rand(8, 8, 8), torch.rand(8, 8, 4))
         self.check_graph(MyModule(), inputs)
 
+    def test_sym_ite(self):
+        def f(x):
+            b = x.shape[0] == 5
+            ret = torch.sym_ite(b, x.shape[0], x.shape[1])
+            return ret
+
+        dynamic_shapes = {'x': {0: Dim("dim0"), 1: Dim("dim1")}}
+        self.check_graph(f, (torch.ones(4, 5),), dynamic_shapes=dynamic_shapes)
+
     @parametrize(
         "name,case",
         get_filtered_export_db_tests(),

torch/_export/pass_base.py

@@ -2,7 +2,7 @@
 import traceback
 import typing
 from contextlib import nullcontext
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import torch
 from functorch.experimental.control_flow import _unstack_pytree
@@ -29,6 +29,16 @@
 PassType = Callable[[torch.fx.GraphModule], Optional[PassResult]]
 
 
+_TORCH_SYM_OPS: Set[Callable] = {
+    torch.sym_int,
+    torch.sym_ite,
+    torch.sym_max,
+    torch.sym_min,
+    torch.sym_not,
+    torch.sym_sqrt,
+}
+
+
 class ExportPassBaseError(RuntimeError):
     pass
 
@@ -182,6 +192,9 @@ def call_function(
             elif getattr(target, "__module__", None) in {"_operator", "math"}:
                 assert callable(target)
                 return self.callback.call_sym(target, args, meta)
+            elif target in _TORCH_SYM_OPS:
+                assert callable(target)
+                return self.callback.call_sym(target, args, meta)
             elif isinstance(target, (torch._ops.OpOverload, torch._ops.OpOverloadPacket)):
                 return self.callback.call_operator(
                     target,

torch/_export/serde/serialize.py

@@ -2,6 +2,7 @@
 import copy
 import dataclasses
 import heapq
+import inspect
 import io
 import json
 import logging
@@ -152,6 +153,12 @@ def _reverse_map(d: Dict[Any, Enum]):
     operator.sub,
     operator.floordiv,
     operator.mod,
+    torch.sym_sqrt,
+    torch.sym_int,
+    torch.sym_ite,
+    torch.sym_max,
+    torch.sym_min,
+    torch.sym_sqrt,
 }
 
 
@@ -162,6 +169,7 @@ def _reverse_map(d: Dict[Any, Enum]):
     operator.ge,
     operator.lt,
     operator.gt,
+    torch.sym_not,
 }
 
 
@@ -369,7 +377,7 @@ def handle_call_function(self, node: torch.fx.Node):
             meta_val = node.meta["val"]
             ex_node = Node(
                 target=self.serialize_operator(node.target),
-                inputs=self.serialize_sym_op_inputs(node.args),
+                inputs=self.serialize_sym_op_inputs(node.target, node.args),
                 outputs=[Argument.create(as_sym_int=self.serialize_sym_int_output(node.name, meta_val))],
                 metadata=self.serialize_metadata(node),
             )
@@ -378,7 +386,7 @@ def handle_call_function(self, node: torch.fx.Node):
             meta_val = node.meta["val"]
             ex_node = Node(
                 target=self.serialize_operator(node.target),
-                inputs=self.serialize_sym_op_inputs(node.args),
+                inputs=self.serialize_sym_op_inputs(node.target, node.args),
                 outputs=[Argument.create(as_sym_bool=self.serialize_sym_bool_output(node.name, meta_val))],
                 metadata=self.serialize_metadata(node),
             )
@@ -453,9 +461,9 @@ def export_nn_module_stack(val):
 
         return ret
 
-    def serialize_sym_op_inputs(self, args) -> List[NamedArgument]:
+    def serialize_sym_op_inputs(self, op, args) -> List[NamedArgument]:
         serialized_args = []
-        args_names = ["a", "b"]
+        args_names = inspect.signature(op).parameters.keys()
         for args_name, arg in zip(args_names, args):
             serialized_args.append(
                 NamedArgument(name=args_name, arg=self.serialize_input(arg))
@@ -1015,9 +1023,9 @@ def deserialize_operator(self, serialized_target: str):
         if serialized_target.startswith("_operator"):  # TODO(zhxchen17) Follow up on this.
             module = operator
             serialized_target_names = serialized_target.split(".")[1:]
-        elif serialized_target.startswith("torch.ops"):
-            module = torch.ops
-            serialized_target_names = serialized_target.split(".")[2:]
+        elif serialized_target.startswith("torch"):
+            module = torch  # type: ignore[misc]
+            serialized_target_names = serialized_target.split(".")[1:]
         else:  # TODO(zhxchen17) Don't catch all here.
             return serialized_target
 
@@ -1150,7 +1158,7 @@ def deserialize_graph(self, serialized_graph: Graph) -> torch.fx.Graph:
         return self.graph
 
     def deserialize_node(self, serialized_node: Node, target: Callable) -> None:
-        if target.__module__ == "_operator":  # TODO(zhxchen17) Follow up on this.
+        if target in _SYM_BOOL_OPS or target in _SYM_INT_OPS:
             name = serialized_node.outputs[0].value.as_name
             args = self.deserialize_sym_op_inputs(serialized_node.inputs)
 

torch/_export/verifier.py

@@ -158,7 +158,15 @@ def _allowed_op_types() -> Tuple[Type[Any], ...]:
                 return ret
 
             # TODO Remove this allowlist.
-            _allowed_torch_functions = (torch.autograd.grad_mode.set_grad_enabled,)
+            _allowed_torch_functions = (
+                torch.autograd.grad_mode.set_grad_enabled,
+                torch.sym_int,
+                torch.sym_ite,
+                torch.sym_max,
+                torch.sym_min,
+                torch.sym_not,
+                torch.sym_sqrt,
+            )
 
             if not isinstance(op, _allowed_op_types()):
                 if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions: