[FX] Make arg normalization a method on Node and not a pass (also augment tests to be exhaustive) (pytorch#55992)

Summary:
Commandeered from pytorch#54563

Primary changes from first PR:
1. Refactored primary `normalize_function` logic into `operator_schemas.py` so that non-FX users can use it.
2. Refactored tests a bit, and added a path to call `normalize_function` directly.
3. Moved check for `boolean_dispatch` so that `torch.lu` also gets properly handled.

Pull Request resolved: pytorch#55992

Reviewed By: mruberry

Differential Revision: D27774396

Pulled By: Chillee

fbshipit-source-id: 7f65632e1d608e4abd55aec5ccbfdc3f67f52b8e

Author: Chillee

test/test_fx_experimental.py

@@ -2,6 +2,7 @@
 import operator
 import unittest
 import sys
+import math
 from typing import Callable, Dict, Union, List
 from torch.fx.symbolic_trace import symbolic_trace
 from torch.fx.graph_module import GraphModule
@@ -12,6 +13,8 @@
 from torch.fx.experimental.param_fetch import lift_lowering_attrs_to_nodes
 from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.jit_utils import JitTestCase
+from torch.testing._internal.common_methods_invocations import op_db
+from torch.testing._internal.common_device_type import ops, onlyCPU, instantiate_device_type_tests
 from torch.fx.passes.split_module import split_module
 from torch.fx.experimental.partitioner_utils import (
     NodeLatency,
@@ -23,10 +26,11 @@
 )
 import torch.fx.experimental.optimization as optimization
 from torch.fx.experimental import merge_matmul
-from torch.fx.experimental.normalize import NormalizeArgs, NormalizeOperators
+from torch.fx.experimental.normalize import NormalizeOperators, NormalizeArgs
 from torch.fx.experimental.schema_type_annotation import AnnotateTypesWithSchema
 from torch.testing._internal.common_nn import module_tests, new_module_tests
-from torch.fx.passes.shape_prop import extract_tensor_metadata
+from torch.fx.operator_schemas import _torchscript_type_to_python_type, normalize_function, normalize_module
+from torch.fx.passes.shape_prop import extract_tensor_metadata, ShapeProp
 
 try:
     from torchvision.models import resnet18
@@ -826,24 +830,26 @@ def is_leaf_module(self, m: torch.nn.Module, module_qualified_name : str) -> boo
         input = torch.randn(5, 3, 224, 224)
         ref_outs = traced(input)
 
+        ShapeProp(traced).propagate(input)
         traced = NormalizeArgs(traced).transform()
 
-        test_outs = traced(input)
-        self.assertEqual(test_outs, ref_outs)
 
         modules = dict(traced.named_modules())
+
         for node in traced.graph.nodes:
-            if node.op == 'call_function' and node.target.__module__ == 'torch.nn.functional':
+            if node.op == 'call_function' and node.target != operator.add:
                 self.assertEqual(len(node.args), 0)
-            if node.op == 'call_module':
+            elif node.op == 'call_module':
                 submod_class = modules[node.target].__class__
                 nn_class = getattr(torch.nn, submod_class.__name__)
                 if submod_class == nn_class:
                     self.assertEqual(len(node.args), 0)
+        traced(input)
+        self.assertEqual(traced(input), ref_outs)
 
     def test_normalize_modules_exhaustive(self):
         """
-        Exhaustively test `NormalizeArgs` on all standard
+        Exhaustively test `Node.normalized_arguments` on all standard
         torch.nn Module classes
         """
         for test_params in module_tests + new_module_tests:
@@ -892,8 +898,23 @@ def forward(self, {params}):
             test_instance = gbls[test_classname](mod)
             traced = symbolic_trace(test_instance)
 
-            # Now actually test arg normalization!
-            traced = NormalizeArgs(traced).transform()
+            # Use `Node.normalized_arguments` to get a new set of arguments
+            # to feed to the Module. Then, rewrite the node to only take
+            # in those arguments as kwargs
+            modules = dict(traced.named_modules())
+            for node in traced.graph.nodes:
+                if node.op == 'call_module':
+                    submod_class = modules[node.target].__class__
+                    nn_class = getattr(torch.nn, submod_class.__name__)
+                    if submod_class == nn_class:
+                        normalized_args = node.normalized_arguments(traced)
+                        normalized_args2 = normalize_module(traced, node.target, node.args, node.kwargs)
+                        assert(normalized_args == normalized_args2)
+                        assert normalized_args
+                        node.args = ()
+                        node.kwargs = normalized_args
+
+            traced.recompile()
 
             # These Modules have an RNG in their forward, so testing
             # correctness by comparing outputs is not correct. Skip that
@@ -904,7 +925,7 @@ def forward(self, {params}):
             if mod.__class__.__name__ not in stochastic_modules:
                 self.assertEqual(traced(*inputs), mod(*inputs))
 
-            # Ensure all args/kwargs are normalized into kwargs
+            traced = NormalizeArgs(symbolic_trace(test_instance)).transform()
             modules = dict(traced.named_modules())
             for node in traced.graph.nodes:
                 if node.op == 'call_module':
@@ -913,6 +934,8 @@ def forward(self, {params}):
                     if submod_class == nn_class:
                         self.assertEqual(len(node.args), 0)
 
+
+
     @skipIfNoTorchVision
     def test_annotate_returns_with_schema(self):
         m = resnet18()
@@ -1236,6 +1259,106 @@ def test_prepare_for_inference_cpu_torchvision(self):
                 torch.testing.assert_allclose(orig_out, new_out)
 
 
+class TestNormalizeOperators(JitTestCase):
+    @onlyCPU
+    @ops(op_db, allowed_dtypes=(torch.float,))
+    def test_normalize_operator_exhaustive(self, device, dtype, op):
+        # Unsupported input types
+        if op.name in {'index_put', '__getitem__', 'unfold', 'repeat', 'polygamma'}:
+            return
+        # These ops currently don't trace in FX for various reasons (i.e. they take a list of tensors)
+        fx_fail = {'stack', 'hstack', 'vstack', 'dstack',
+                   'linalg.multi_dot'}
+        print(op.name)
+        sample_inputs_itr = op.sample_inputs(device, dtype, requires_grad=False)
+        for sample_input in sample_inputs_itr:
+            unsupported_arg_type = False
+            arg_values = [sample_input.input] + list(sample_input.args)
+            kwarg_values = sample_input.kwargs
+            arg_types = []
+            kwarg_types = {}
+
+            def jit_infer_type(v):
+                inferred_arg_type = torch._C._jit_try_infer_type(v)
+                assert inferred_arg_type.success()
+                t = _torchscript_type_to_python_type(inferred_arg_type.type())
+                return t
+
+            for v in arg_values:
+                if isinstance(v, torch.Tensor):
+                    arg_types.append(type(v))
+                else:
+                    if isinstance(v, complex):
+                        # Complex type not supported in FX
+                        unsupported_arg_type = True
+                    arg_types.append(jit_infer_type(v))
+
+            for k, v in kwarg_values.items():
+                if isinstance(v, torch.Tensor):
+                    kwarg_types[k] = type(v)
+                else:
+                    if isinstance(v, complex):
+                        # Complex type not supported in FX
+                        unsupported_arg_type = True
+                    kwarg_types[k] = jit_infer_type(v)
+
+            if unsupported_arg_type:
+                continue
+            # Test normalize_function by itself
+            ref_out = op.op(*arg_values, **kwarg_values)
+            norm_kwargs = normalize_function(op.op, arg_values, kwarg_values, arg_types, kwarg_types)
+            test_out = op.op(**norm_kwargs)
+            self.assertEqual(test_out, ref_out)
+
+            # Test normalized_arguments as part of FX
+            if op.name in fx_fail:
+                continue
+            param_names = []
+            param_values = []
+            fx_args = []
+            for idx, v in enumerate(arg_values):
+                if isinstance(v, torch.Tensor):
+                    param_names.append(f"arg_{idx}")
+                    param_values.append(v)
+                    fx_args.append(param_names[-1])
+                else:
+                    fx_args.append(f'{repr(v)}')
+
+            for k, v in kwarg_values.items():
+                if isinstance(v, torch.Tensor):
+                    param_names.append(k)
+                    param_values.append(v)
+                    fx_args.append(k)
+                else:
+                    fx_args.append(f'{k} = {repr(v)}')
+
+            code = f"""
+class TestModule(torch.nn.Module):
+    def forward(self, {', '.join(param_names)}):
+        return torch.{op.name}({', '.join(fx_args)})
+            """
+
+            g = {'torch': torch, 'inf' : math.inf}
+            exec(code, g)
+            TestModule = g['TestModule']
+
+
+            m = TestModule()
+            traced = torch.fx.symbolic_trace(m)
+            ref_out = traced(*param_values)
+
+            for node in traced.graph.nodes:
+                if node.op == 'call_function':
+                    normalized_args = node.normalized_arguments(traced, arg_types, kwarg_types)
+                    assert normalized_args
+                    node.args = ()
+                    node.kwargs = normalized_args
+            traced.recompile()
+
+            test_out = traced(*param_values)
+            self.assertEqual(test_out, ref_out)
+
+instantiate_device_type_tests(TestNormalizeOperators, globals())
 
 if __name__ == "__main__":
     run_tests()

torch/_C/__init__.pyi.in

@@ -907,6 +907,7 @@ class DictType(JitType):
 
 class TupleType(JitType):
     def __init__(self, a: List[JitType]) -> None: ...
+    def elements(self) -> List[JitType]: ...
 
 class ClassType(JitType):
     def __init__(self, qualified_name: str) -> None: ...

torch/fx/experimental/normalize.py

@@ -1,13 +1,12 @@
 import torch
 import torch.fx
-import inspect
+import torch.fx as fx
 import operator
-from typing import Any, Callable, Dict, Optional, Tuple
-from torch.fx.node import Argument, Target
-from torch._jit_internal import boolean_dispatched
+from typing import Any, Callable, Dict, Tuple, Optional
+from torch.fx.node import Argument, Target, Node
+from torch.fx.operator_schemas import normalize_module, normalize_function, create_type_hint
 
-from torch.fx import Transformer
-from torch.fx.operator_schemas import get_signature_for_torch_op
+from torch.fx import Transformer, Proxy
 from .schema_type_annotation import AnnotateTypesWithSchema
 
 class NormalizeArgs(Transformer):
@@ -18,108 +17,53 @@ class NormalizeArgs(Transformer):
     Also populates default values. Does not support positional-only
     parameters or varargs parameters (*args, **kwargs).
 
-    Example usage:
+    If the nodes have 'type' metadata, it will use it to disambiguate
+    overloads. Otherwise, it will throw an error.
 
+    Example usage:
         m = torchvision.models.resnet18()
-
         traced = torch.fx.symbolic_trace(m)
-
         traced = NormalizeArgs(traced).transform()
     """
-    def __init__(self, module : torch.nn.Module, normalize_functionals : bool = True,
-                 normalize_modules : bool = True):
+    def __init__(self, module : torch.nn.Module):
         super().__init__(module)
-        self.normalize_functionals = normalize_functionals
-        self.normalize_modules = normalize_modules
+        self.node_map: Dict[Proxy, Node] = {}
 
-    def call_function(self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any]):
-        new_kwargs = None
-
-        if self.normalize_functionals and target.__module__ == 'torch.nn.functional':
-            target_for_analysis = target
-            if target in boolean_dispatched:
-                # HACK: `boolean_dispatch` as used in `torch.nn.functional` makes it so that we have
-                # a 2-way dispatch based on a boolean value. Here we check that the `true` and `false`
-                # branches of the dispatch have exactly the same signature. If they do, use the `true`
-                # branch signature for analysis. Otherwise, leave this un-normalized
-                assert not isinstance(target, str)
-                dispatched = boolean_dispatched[target]
-                if_true, if_false = dispatched['if_true'], dispatched['if_false']
-                if inspect.signature(if_true).parameters != inspect.signature(if_false).parameters:
-                    return super().call_function(target, args, kwargs)
-                target_for_analysis = if_true
-
-            assert callable(target_for_analysis)
-            sig = inspect.signature(inspect.unwrap(target_for_analysis))
-            new_kwargs = self._args_kwargs_to_normalized_kwargs(sig, args, kwargs)
+    def run_node(self, n: Node) -> Any:
+        args, kwargs = self.fetch_args_kwargs_from_env(n)
+
+        def get_type(arg):
+            if isinstance(arg, fx.Proxy):
+                old_meta = self.node_map[arg].meta
+                return old_meta['type'] if 'type' in old_meta else None
+            return create_type_hint(arg)
+
+        arg_types = tuple([get_type(arg) for arg in args])
+        kwarg_types = {k: get_type(v) for k, v in kwargs.items()}
+        if n.op == 'call_function':
+            out = self.call_function(n.target, args, kwargs, arg_types, kwarg_types)
         else:
-            assert callable(target)
-            torch_op_schemas = get_signature_for_torch_op(target)
-            if torch_op_schemas:
-                # Iterate through all of the schema until we find one that matches
-                # If one matches, populate `new_kwargs` with the combined args/kwargs
-                # values. If none matches, `new_kwargs` will be None
-                for candidate_signature in torch_op_schemas:
-                    try:
-                        candidate_signature.bind(args, kwargs)
-                        new_kwargs = self._args_kwargs_to_normalized_kwargs(candidate_signature, args, kwargs)
-                        break
-                    except TypeError:
-                        continue
+            out = super().run_node(n)
+        self.node_map[out] = n
+        return out
+
+    def call_function(
+            self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any],
+            arg_types: Optional[Tuple[Any, ...]] = None, kwarg_types : Optional[Dict[str, Any]] = None):
+        assert callable(target)
+        new_kwargs = normalize_function(target, args, kwargs, arg_types, kwarg_types)  # type: ignore
         if new_kwargs:
-            # FIXME: `target(**kwargs)` doesn't keep things specified as kwargs
-            # in kwargs
             return self.tracer.create_proxy('call_function', target, (), new_kwargs)
         else:
             return super().call_function(target, args, kwargs)
 
     def call_module(self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any]):
         assert isinstance(target, str)
-        submod = self.fetch_attr(target)
-        if self.normalize_modules and hasattr(submod.__class__, '__name__'):
-            classname = submod.__class__.__name__
-            if getattr(torch.nn, classname, None) == submod.__class__:
-                sig = inspect.signature(inspect.unwrap(submod.forward))
-                new_kwargs = self._args_kwargs_to_normalized_kwargs(sig, args, kwargs)
-                if new_kwargs:
-                    return super().call_module(target, (), new_kwargs)
-        return super().call_module(target, args, kwargs)
-
-    def _args_kwargs_to_normalized_kwargs(self, sig : inspect.Signature, args : Tuple[Argument, ...],
-                                          kwargs : Dict[str, Any]) -> Optional[Dict[str, Any]]:
-        """
-        Given a call target, args, and kwargs, return the arguments normalized into
-        a single kwargs dict, or None if the type signature is not supported by
-        this normalization.
-
-        Args:
-
-            target (inspect.Signature): Signature object for the target
-            args (Tuple): Arguments that appear at the callsite for `target`
-            kwargs (Dict): Keyword arugments that appear at the callsite for `target`
-
-        Returns:
-
-            Optional[Dict]: Normalized kwargs for `target`, or `None` if this target is not
-                supported
-        """
-
-        # Don't currently support positional-only
-        # or varargs (*args, **kwargs) signatures
-        supported_parameter_types = {
-            inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY}
-        if any(p.kind not in supported_parameter_types for p in sig.parameters.values()):
-            return None
-
-        bound_args = sig.bind(*args, **kwargs)
-        bound_args.apply_defaults()
-
-        new_kwargs : Dict[str, Any] = {}
-        for param in sig.parameters:
-            new_kwargs[param] = bound_args.arguments[param]
-
-        return new_kwargs
-
+        new_kwargs = normalize_module(self.module, target, args, kwargs)  # type: ignore
+        if new_kwargs:
+            return super().call_module(target, (), new_kwargs)
+        else:
+            return super().call_module(target, args, kwargs)
 
 class NormalizeOperators(AnnotateTypesWithSchema):
     """