diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56da3ea48a903ce2f0d0490b84261be41e1102df Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/converter.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/converter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebc781e7162d69ae42b76e7deeaf853c8613f3c1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/converter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/error.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/error.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..492021eec542016d207fcf9ae30d8d9689d7145f Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/error.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/non_strict_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/non_strict_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b792a59baa2d6ec385d35a8d00d70465c1588cf Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/non_strict_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/pass_base.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/pass_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f837886da88bdd70fca097ae8101234f76635b34 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/pass_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/tools.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/tools.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc960b57d1f916e24e3df1cba0aac9f35b4e8056 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/tools.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac64432e32a127a31d2cab5a78a1a91b87d41be1 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/verifier.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/verifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..902fbe6d2520ae92b4b4232b34095d6bc24c75f8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/verifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/__pycache__/wrappers.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/__pycache__/wrappers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff9e70710ee0f7704334615b48e40df9f037699d Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/__pycache__/wrappers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/__init__.py b/lib/python3.12/site-packages/torch/_export/db/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..10a55772ab58b21573a6eba0356ddd3080164ac7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/lib/python3.12/site-packages/torch/_export/db/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb6bcfb38acac59acf4532d2125cae1ea7d69eb4 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/__pycache__/case.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/__pycache__/case.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76b598db5839b86f5d4b4d0cb1fdcb37054a91e2 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/__pycache__/case.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/__pycache__/gen_example.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/__pycache__/gen_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..597000222e8b5d933b9076265cf9e14b397cf05c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/__pycache__/gen_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/__pycache__/logging.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/__pycache__/logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5edc951b8915c1c0ed97071f1d9a0df8daadf731 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/__pycache__/logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/case.py b/lib/python3.12/site-packages/torch/_export/db/case.py new file mode 100644 index 0000000000000000000000000000000000000000..b228f6c2c33773b07e5cb7e82abc1adcad6422d0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/case.py @@ -0,0 +1,174 @@ +# mypy: allow-untyped-defs +import inspect +import re +import string +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Optional, Set, Tuple +from types import ModuleType + +import torch + +_TAGS: Dict[str, Dict[str, Any]] = { + "torch": { + "cond": {}, + "dynamic-shape": {}, + "escape-hatch": {}, + "map": {}, + "dynamic-value": {}, + "operator": {}, + "mutation": {}, + }, + "python": { + "assert": {}, + "builtin": {}, + "closure": {}, + "context-manager": {}, + "control-flow": {}, + "data-structure": {}, + "standard-library": {}, + "object-model": {}, + }, +} + + +class SupportLevel(Enum): + """ + Indicates at what stage the feature + used in the example is handled in export. + """ + + SUPPORTED = 1 + NOT_SUPPORTED_YET = 0 + + +ArgsType = Tuple[Any, ...] + + +def check_inputs_type(args, kwargs): + if not isinstance(args, tuple): + raise ValueError( + f"Expecting args type to be a tuple, got: {type(args)}" + ) + if not isinstance(kwargs, dict): + raise ValueError( + f"Expecting kwargs type to be a dict, got: {type(kwargs)}" + ) + for key in kwargs: + if not isinstance(key, str): + raise ValueError( + f"Expecting kwargs keys to be a string, got: {type(key)}" + ) + +def _validate_tag(tag: str): + parts = tag.split(".") + t = _TAGS + for part in parts: + assert set(part) <= set( + string.ascii_lowercase + "-" + ), f"Tag contains invalid characters: {part}" + if part in t: + t = t[part] + else: + raise ValueError(f"Tag {tag} is not found in registered tags.") + + +@dataclass(frozen=True) +class ExportCase: + example_args: ArgsType + description: str # A description of the use case. + model: torch.nn.Module + name: str + example_kwargs: Dict[str, Any] = field(default_factory=dict) + extra_args: Optional[ArgsType] = None # For testing graph generalization. + # Tags associated with the use case. (e.g dynamic-shape, escape-hatch) + tags: Set[str] = field(default_factory=set) + support_level: SupportLevel = SupportLevel.SUPPORTED + dynamic_shapes: Optional[Dict[str, Any]] = None + + def __post_init__(self): + check_inputs_type(self.example_args, self.example_kwargs) + if self.extra_args is not None: + check_inputs_type(self.extra_args, {}) + + for tag in self.tags: + _validate_tag(tag) + + if not isinstance(self.description, str) or len(self.description) == 0: + raise ValueError(f'Invalid description: "{self.description}"') + + +_EXAMPLE_CASES: Dict[str, ExportCase] = {} +_MODULES: Set[ModuleType] = set() +_EXAMPLE_CONFLICT_CASES: Dict[str, List[ExportCase]] = {} +_EXAMPLE_REWRITE_CASES: Dict[str, List[ExportCase]] = {} + + +def register_db_case(case: ExportCase) -> None: + """ + Registers a user provided ExportCase into example bank. + """ + if case.name in _EXAMPLE_CASES: + if case.name not in _EXAMPLE_CONFLICT_CASES: + _EXAMPLE_CONFLICT_CASES[case.name] = [_EXAMPLE_CASES[case.name]] + _EXAMPLE_CONFLICT_CASES[case.name].append(case) + return + + _EXAMPLE_CASES[case.name] = case + + +def to_snake_case(name): + name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name) + return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower() + + +def _make_export_case(m, name, configs): + if not isinstance(m, torch.nn.Module): + raise TypeError("Export case class should be a torch.nn.Module.") + + if "description" not in configs: + # Fallback to docstring if description is missing. + assert ( + m.__doc__ is not None + ), f"Could not find description or docstring for export case: {m}" + configs = {**configs, "description": m.__doc__} + return ExportCase(**{**configs, "model": m, "name": name}) + + +def export_case(**kwargs): + """ + Decorator for registering a user provided case into example bank. + """ + + def wrapper(m): + configs = kwargs + module = inspect.getmodule(m) + if module in _MODULES: + raise RuntimeError("export_case should only be used once per example file.") + + assert module is not None + _MODULES.add(module) + module_name = module.__name__.split(".")[-1] + case = _make_export_case(m, module_name, configs) + register_db_case(case) + return case + + return wrapper + + +def export_rewrite_case(**kwargs): + def wrapper(m): + configs = kwargs + + parent = configs.pop("parent") + assert isinstance(parent, ExportCase) + key = parent.name + if key not in _EXAMPLE_REWRITE_CASES: + _EXAMPLE_REWRITE_CASES[key] = [] + + configs["example_args"] = parent.example_args + case = _make_export_case(m, to_snake_case(m.__name__), configs) + _EXAMPLE_REWRITE_CASES[key].append(case) + return case + + return wrapper diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__init__.py b/lib/python3.12/site-packages/torch/_export/db/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..834dbce32f10bfb339fd2182a2455b43914441c9 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/__init__.py @@ -0,0 +1,61 @@ +# mypy: allow-untyped-defs +import dataclasses +import glob +import inspect +from os.path import basename, dirname, isfile, join + +import torch +from torch._export.db.case import ( + _EXAMPLE_CASES, + _EXAMPLE_CONFLICT_CASES, + _EXAMPLE_REWRITE_CASES, + SupportLevel, + export_case, + ExportCase, +) + + +def _collect_examples(): + case_names = glob.glob(join(dirname(__file__), "*.py")) + case_names = [ + basename(f)[:-3] for f in case_names if isfile(f) and not f.endswith("__init__.py") + ] + + case_fields = {f.name for f in dataclasses.fields(ExportCase)} + for case_name in case_names: + case = __import__(case_name, globals(), locals(), [], 1) + variables = [name for name in dir(case) if name in case_fields] + export_case(**{v: getattr(case, v) for v in variables})(case.model) + +_collect_examples() + +def all_examples(): + return _EXAMPLE_CASES + + +if len(_EXAMPLE_CONFLICT_CASES) > 0: + + def get_name(case): + model = case.model + if isinstance(model, torch.nn.Module): + model = type(model) + return model.__name__ + + msg = "Error on conflict export case name.\n" + for case_name, cases in _EXAMPLE_CONFLICT_CASES.items(): + msg += f"Case name {case_name} is associated with multiple cases:\n " + msg += f"[{','.join(map(get_name, cases))}]\n" + + raise RuntimeError(msg) + + +def filter_examples_by_support_level(support_level: SupportLevel): + return { + key: val + for key, val in all_examples().items() + if val.support_level == support_level + } + + +def get_rewrite_cases(case): + return _EXAMPLE_REWRITE_CASES.get(case.name, []) diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77271020cedd9c5f2d1548695152f8beb46b65e8 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/assume_constant_result.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/assume_constant_result.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ccbac91108e3252c646c1291722c3f19ee62ce7b Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/assume_constant_result.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/autograd_function.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/autograd_function.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dab5a16c461796eb00292c462c91ab47aa9231a Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/autograd_function.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/class_method.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/class_method.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9eaf440924a0c679b24cc5133814a4b7b88ebaa Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/class_method.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_class_method.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_class_method.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80e121802ad66d1fdb9bfa3c749753cc6ad05909 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_class_method.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_nested_function.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_nested_function.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e356b595b037c7a5962d59cafac152e9449fd4c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_nested_function.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_nonlocal_variables.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_nonlocal_variables.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3487de3ed52ca41afd4040e637a8dbfb47a6f59c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_branch_nonlocal_variables.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_closed_over_variable.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_closed_over_variable.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25392aa77fbeb3d40420bf6e052425f95080b23c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_closed_over_variable.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_operands.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_operands.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3ea4bf6e5d184e1d05b27fddb4ae34f236b9df3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_operands.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_predicate.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_predicate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44607177281c4fe9d5134f0077065294deedd549 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/cond_predicate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/constrain_as_size_example.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/constrain_as_size_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d06128ab377351f35f8f570d1ab278f6d060b74 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/constrain_as_size_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/constrain_as_value_example.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/constrain_as_value_example.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..070326f61b3c479527c42d0e4db5bc5da3f22788 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/constrain_as_value_example.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/decorator.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/decorator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6dc486d958e0bee41376f6ac6ac9990293519be Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/decorator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dictionary.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dictionary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..716cff9dff1b323e9724204a9e1cf2d4508c6c63 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dictionary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_assert.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_assert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46411178d82bb9a09a87cfff7daf6bc4d2de28b7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_assert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_constructor.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_constructor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..963da86ca456d64362c448809855aaf7a6357630 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_constructor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_if_guard.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_if_guard.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c6005e55e91c737cbf19e89f3a94697f031e356 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_if_guard.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_map.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_map.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7383a53cbaafdaa16df7dbb2456117b76fd370ad Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_map.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_round.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_round.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9aa5c81765424a3bfad5c70255c5e887776adcef Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_round.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_slicing.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_slicing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e493e068ec3bccbd20e704cfa130ee0832f11481 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_slicing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_view.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_view.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..450fc72825ece5552a32597f5f4c74d6642db15d Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/dynamic_shape_view.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/fn_with_kwargs.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/fn_with_kwargs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3864f26b3c3c6c7c5a2817aa2489bd0d94c1efcd Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/fn_with_kwargs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/list_contains.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/list_contains.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..445ee5cc382bd1718556c20082f324d43134f877 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/list_contains.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/list_unpack.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/list_unpack.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d88f6cfa57bb294f84190c91a409809f3055a05 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/list_unpack.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/model_attr_mutation.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/model_attr_mutation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17c0adfeb83513bec536ac8875e0a592e5bf660a Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/model_attr_mutation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/nested_function.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/nested_function.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d2ef1bc7a633d5ccf2b96004339583b5d339e55 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/nested_function.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/null_context_manager.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/null_context_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad68f03e6cd832b587e6273ab94fa233b16baf79 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/null_context_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/optional_input.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/optional_input.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78e5d42855ad8069fcd7be4360bcf4f036510ede Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/optional_input.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/pytree_flatten.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/pytree_flatten.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..982b558275b04e58fb7de25fae031993d78ee729 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/pytree_flatten.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/scalar_output.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/scalar_output.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..653e71b9be6949b7ed75dd0ef405b789bf4d111c Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/scalar_output.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/specialized_attribute.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/specialized_attribute.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d55cf1a0124380fabb247532733061e2ccaebaf Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/specialized_attribute.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/static_for_loop.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/static_for_loop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5b744a573f3dc03cdf86f0614ee77fa40071ebd Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/static_for_loop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/static_if.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/static_if.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1173b0342745931f3760248740f133244efb5ee Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/static_if.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/tensor_setattr.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/tensor_setattr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5b7b24526b26ce023df7c69fb508b92da123ae3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/tensor_setattr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/type_reflection_method.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/type_reflection_method.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb6b374d7ccf8020fd0205bd4382528453bf704e Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/type_reflection_method.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/unsupported_operator.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/unsupported_operator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e109d30f6adacaa985d84ee52afc6409a0953a73 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/unsupported_operator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/user_input_mutation.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/user_input_mutation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f649f46ec3838df210974bd35131bc639dd04af0 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/db/examples/__pycache__/user_input_mutation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/assume_constant_result.py b/lib/python3.12/site-packages/torch/_export/db/examples/assume_constant_result.py new file mode 100644 index 0000000000000000000000000000000000000000..931ce7f7a50fc5a175101ac57c424c88cf31b54c --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/assume_constant_result.py @@ -0,0 +1,20 @@ +# mypy: allow-untyped-defs +import torch +import torch._dynamo as torchdynamo + + +class AssumeConstantResult(torch.nn.Module): + """ + Applying `assume_constant_result` decorator to burn make non-tracable code as constant. + """ + + @torchdynamo.assume_constant_result + def get_item(self, y): + return y.int().item() + + def forward(self, x, y): + return x[: self.get_item(y)] + +example_args = (torch.randn(3, 2), torch.tensor(4)) +tags = {"torch.escape-hatch"} +model = AssumeConstantResult() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/autograd_function.py b/lib/python3.12/site-packages/torch/_export/db/examples/autograd_function.py new file mode 100644 index 0000000000000000000000000000000000000000..0458291e2176932b32b0fc2e44d51d5812dcead3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/autograd_function.py @@ -0,0 +1,23 @@ +# mypy: allow-untyped-defs +import torch + +class MyAutogradFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + return x.clone() + + @staticmethod + def backward(ctx, grad_output): + return grad_output + 1 + +class AutogradFunction(torch.nn.Module): + """ + TorchDynamo does not keep track of backward() on autograd functions. We recommend to + use `allow_in_graph` to mitigate this problem. + """ + + def forward(self, x): + return MyAutogradFunction.apply(x) + +example_args = (torch.randn(3, 2),) +model = AutogradFunction() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/class_method.py b/lib/python3.12/site-packages/torch/_export/db/examples/class_method.py new file mode 100644 index 0000000000000000000000000000000000000000..f701f54d4f4ea1cb5816292cd60bb4df3d03c5e8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/class_method.py @@ -0,0 +1,22 @@ +# mypy: allow-untyped-defs +import torch + +class ClassMethod(torch.nn.Module): + """ + Class methods are inlined during tracing. + """ + + @classmethod + def method(cls, x): + return x + 1 + + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(4, 2) + + def forward(self, x): + x = self.linear(x) + return self.method(x) * self.__class__.method(x) * type(self).method(x) + +example_args = (torch.randn(3, 4),) +model = ClassMethod() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_class_method.py b/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_class_method.py new file mode 100644 index 0000000000000000000000000000000000000000..22600cc504348d1d261b0ea2b9ed2d57af76b0a3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_class_method.py @@ -0,0 +1,44 @@ +# mypy: allow-untyped-defs +import torch + +from functorch.experimental.control_flow import cond + +class MySubModule(torch.nn.Module): + def foo(self, x): + return x.cos() + + def forward(self, x): + return self.foo(x) + +class CondBranchClassMethod(torch.nn.Module): + """ + The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules: + - both branches must take the same args, which must also match the branch args passed to cond. + - both branches must return a single tensor + - returned tensor must have the same tensor metadata, e.g. shape and dtype + - branch function can be free function, nested function, lambda, class methods + - branch function can not have closure variables + - no inplace mutations on inputs or global variables + + + This example demonstrates using class method in cond(). + + NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized. + """ + + def __init__(self) -> None: + super().__init__() + self.subm = MySubModule() + + def bar(self, x): + return x.sin() + + def forward(self, x): + return cond(x.shape[0] <= 2, self.subm.forward, self.bar, [x]) + +example_args = (torch.randn(3),) +tags = { + "torch.cond", + "torch.dynamic-shape", +} +model = CondBranchClassMethod() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nested_function.py b/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nested_function.py new file mode 100644 index 0000000000000000000000000000000000000000..b28ceeddc7956d136a8cf786c283344731d3e7ac --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nested_function.py @@ -0,0 +1,41 @@ +# mypy: allow-untyped-defs +import torch + +from functorch.experimental.control_flow import cond + +class CondBranchNestedFunction(torch.nn.Module): + """ + The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules: + - both branches must take the same args, which must also match the branch args passed to cond. + - both branches must return a single tensor + - returned tensor must have the same tensor metadata, e.g. shape and dtype + - branch function can be free function, nested function, lambda, class methods + - branch function can not have closure variables + - no inplace mutations on inputs or global variables + + This example demonstrates using nested function in cond(). + + NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized. + """ + + def forward(self, x): + def true_fn(x): + def inner_true_fn(y): + return x + y + + return inner_true_fn(x) + + def false_fn(x): + def inner_false_fn(y): + return x - y + + return inner_false_fn(x) + + return cond(x.shape[0] < 10, true_fn, false_fn, [x]) + +example_args = (torch.randn(3),) +tags = { + "torch.cond", + "torch.dynamic-shape", +} +model = CondBranchNestedFunction() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nonlocal_variables.py b/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nonlocal_variables.py new file mode 100644 index 0000000000000000000000000000000000000000..50d0ec87a690d063cb0e841fc057a6ae95c369fb --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/cond_branch_nonlocal_variables.py @@ -0,0 +1,59 @@ +# mypy: allow-untyped-defs +import torch + +from functorch.experimental.control_flow import cond + +class CondBranchNonlocalVariables(torch.nn.Module): + """ + The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules: + - both branches must take the same args, which must also match the branch args passed to cond. + - both branches must return a single tensor + - returned tensor must have the same tensor metadata, e.g. shape and dtype + - branch function can be free function, nested function, lambda, class methods + - branch function can not have closure variables + - no inplace mutations on inputs or global variables + + This example demonstrates how to rewrite code to avoid capturing closure variables in branch functions. + + The code below will not work because capturing closure variables is not supported. + ``` + my_tensor_var = x + 100 + my_primitive_var = 3.14 + + def true_fn(y): + nonlocal my_tensor_var, my_primitive_var + return y + my_tensor_var + my_primitive_var + + def false_fn(y): + nonlocal my_tensor_var, my_primitive_var + return y - my_tensor_var - my_primitive_var + + return cond(x.shape[0] > 5, true_fn, false_fn, [x]) + ``` + + NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized. + """ + + def forward(self, x): + my_tensor_var = x + 100 + my_primitive_var = 3.14 + + def true_fn(x, y, z): + return x + y + z + + def false_fn(x, y, z): + return x - y - z + + return cond( + x.shape[0] > 5, + true_fn, + false_fn, + [x, my_tensor_var, torch.tensor(my_primitive_var)], + ) + +example_args = (torch.randn(6),) +tags = { + "torch.cond", + "torch.dynamic-shape", +} +model = CondBranchNonlocalVariables() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/cond_closed_over_variable.py b/lib/python3.12/site-packages/torch/_export/db/examples/cond_closed_over_variable.py new file mode 100644 index 0000000000000000000000000000000000000000..183180ab4fc825385170fea2bec6af184374a67e --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/cond_closed_over_variable.py @@ -0,0 +1,22 @@ +# mypy: allow-untyped-defs +import torch + +from functorch.experimental.control_flow import cond + +class CondClosedOverVariable(torch.nn.Module): + """ + torch.cond() supports branches closed over arbitrary variables. + """ + + def forward(self, pred, x): + def true_fn(val): + return x * 2 + + def false_fn(val): + return x - 2 + + return cond(pred, true_fn, false_fn, [x + 1]) + +example_args = (torch.tensor(True), torch.randn(3, 2)) +tags = {"torch.cond", "python.closure"} +model = CondClosedOverVariable() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/cond_operands.py b/lib/python3.12/site-packages/torch/_export/db/examples/cond_operands.py new file mode 100644 index 0000000000000000000000000000000000000000..29941d828ae63e45ad81b0f2c7428706df849cf0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/cond_operands.py @@ -0,0 +1,36 @@ +# mypy: allow-untyped-defs +import torch + +from torch.export import Dim +from functorch.experimental.control_flow import cond + +x = torch.randn(3, 2) +y = torch.randn(2) +dim0_x = Dim("dim0_x") + +class CondOperands(torch.nn.Module): + """ + The operands passed to cond() must be: + - a list of tensors + - match arguments of `true_fn` and `false_fn` + + NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized. + """ + + def forward(self, x, y): + def true_fn(x, y): + return x + y + + def false_fn(x, y): + return x - y + + return cond(x.shape[0] > 2, true_fn, false_fn, [x, y]) + +example_args = (x, y) +tags = { + "torch.cond", + "torch.dynamic-shape", +} +extra_inputs = (torch.randn(2, 2), torch.randn(2)) +dynamic_shapes = {"x": {0: dim0_x}, "y": None} +model = CondOperands() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/cond_predicate.py b/lib/python3.12/site-packages/torch/_export/db/examples/cond_predicate.py new file mode 100644 index 0000000000000000000000000000000000000000..68bb8850bba909a0c6546c8f12a1a3fa1bdc70d1 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/cond_predicate.py @@ -0,0 +1,25 @@ +# mypy: allow-untyped-defs +import torch + +from functorch.experimental.control_flow import cond + +class CondPredicate(torch.nn.Module): + """ + The conditional statement (aka predicate) passed to cond() must be one of the following: + - torch.Tensor with a single element + - boolean expression + + NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized. + """ + + def forward(self, x): + pred = x.dim() > 2 and x.shape[2] > 10 + + return cond(pred, lambda x: x.cos(), lambda y: y.sin(), [x]) + +example_args = (torch.randn(6, 4, 3),) +tags = { + "torch.cond", + "torch.dynamic-shape", +} +model = CondPredicate() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_size_example.py b/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_size_example.py new file mode 100644 index 0000000000000000000000000000000000000000..fbed4984ac4d14c0b175874c44e9c7009bfe8ddc --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_size_example.py @@ -0,0 +1,25 @@ +# mypy: allow-untyped-defs +import torch + + +class ConstrainAsSizeExample(torch.nn.Module): + """ + If the value is not known at tracing time, you can provide hint so that we + can trace further. Please look at torch._check and torch._check_is_size APIs. + torch._check_is_size is used for values that NEED to be used for constructing + tensor. + """ + + def forward(self, x): + a = x.item() + torch._check_is_size(a) + torch._check(a <= 5) + return torch.zeros((a, 5)) + + +example_args = (torch.tensor(4),) +tags = { + "torch.dynamic-value", + "torch.escape-hatch", +} +model = ConstrainAsSizeExample() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_value_example.py b/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_value_example.py new file mode 100644 index 0000000000000000000000000000000000000000..c8bfc3d6e36578ee50f1ca2eee80fed9aae14cfd --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/constrain_as_value_example.py @@ -0,0 +1,28 @@ +# mypy: allow-untyped-defs +import torch + + +class ConstrainAsValueExample(torch.nn.Module): + """ + If the value is not known at tracing time, you can provide hint so that we + can trace further. Please look at torch._check and torch._check_is_size APIs. + torch._check is used for values that don't need to be used for constructing + tensor. + """ + + def forward(self, x, y): + a = x.item() + torch._check(a >= 0) + torch._check(a <= 5) + + if a < 6: + return y.sin() + return y.cos() + + +example_args = (torch.tensor(4), torch.randn(5, 5)) +tags = { + "torch.dynamic-value", + "torch.escape-hatch", +} +model = ConstrainAsValueExample() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/decorator.py b/lib/python3.12/site-packages/torch/_export/db/examples/decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..7d24cc681a6b62adf40bfd9a2e5283afb3515131 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/decorator.py @@ -0,0 +1,23 @@ +# mypy: allow-untyped-defs +import functools + +import torch + +def test_decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + 1 + + return wrapper + +class Decorator(torch.nn.Module): + """ + Decorators calls are inlined into the exported function during tracing. + """ + + @test_decorator + def forward(self, x, y): + return x + y + +example_args = (torch.randn(3, 2), torch.randn(3, 2)) +model = Decorator() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dictionary.py b/lib/python3.12/site-packages/torch/_export/db/examples/dictionary.py new file mode 100644 index 0000000000000000000000000000000000000000..49e688bc0ac1f09567e3b877aaca29a1d02b4121 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dictionary.py @@ -0,0 +1,17 @@ +# mypy: allow-untyped-defs +import torch + +class Dictionary(torch.nn.Module): + """ + Dictionary structures are inlined and flattened along tracing. + """ + + def forward(self, x, y): + elements = {} + elements["x2"] = x * x + y = y * elements["x2"] + return {"y": y} + +example_args = (torch.randn(3, 2), torch.tensor(4)) +tags = {"python.data-structure"} +model = Dictionary() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_assert.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_assert.py new file mode 100644 index 0000000000000000000000000000000000000000..cc822e5553e1ab8bd350a26966c22f1a9a1698cf --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_assert.py @@ -0,0 +1,18 @@ +# mypy: allow-untyped-defs +import torch + +class DynamicShapeAssert(torch.nn.Module): + """ + A basic usage of python assertion. + """ + + def forward(self, x): + # assertion with error message + assert x.shape[0] > 2, f"{x.shape[0]} is greater than 2" + # assertion without error message + assert x.shape[0] > 1 + return x + +example_args = (torch.randn(3, 2),) +tags = {"python.assert"} +model = DynamicShapeAssert() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_constructor.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_constructor.py new file mode 100644 index 0000000000000000000000000000000000000000..157e460274ad58ba71c886b35364ddc0cd4d886a --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_constructor.py @@ -0,0 +1,15 @@ +# mypy: allow-untyped-defs +import torch + +class DynamicShapeConstructor(torch.nn.Module): + """ + Tensor constructors should be captured with dynamic shape inputs rather + than being baked in with static shape. + """ + + def forward(self, x): + return torch.zeros(x.shape[0] * 2) + +example_args = (torch.randn(3, 2),) +tags = {"torch.dynamic-shape"} +model = DynamicShapeConstructor() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_if_guard.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_if_guard.py new file mode 100644 index 0000000000000000000000000000000000000000..21824ef3a0f66eb25f4d8e8c1ba92f53fdd4c275 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_if_guard.py @@ -0,0 +1,19 @@ +# mypy: allow-untyped-defs +import torch + +class DynamicShapeIfGuard(torch.nn.Module): + """ + `if` statement with backed dynamic shape predicate will be specialized into + one particular branch and generate a guard. However, export will fail if the + the dimension is marked as dynamic shape from higher level API. + """ + + def forward(self, x): + if x.shape[0] == 3: + return x.cos() + + return x.sin() + +example_args = (torch.randn(3, 2, 2),) +tags = {"torch.dynamic-shape", "python.control-flow"} +model = DynamicShapeIfGuard() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_map.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_map.py new file mode 100644 index 0000000000000000000000000000000000000000..f8066aed556b9ee588b9744d17ba16c35d8fed6c --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_map.py @@ -0,0 +1,19 @@ +# mypy: allow-untyped-defs +import torch + +from functorch.experimental.control_flow import map + +class DynamicShapeMap(torch.nn.Module): + """ + functorch map() maps a function over the first tensor dimension. + """ + + def forward(self, xs, y): + def body(x, y): + return x + y + + return map(body, xs, y) + +example_args = (torch.randn(3, 2), torch.randn(2)) +tags = {"torch.dynamic-shape", "torch.map"} +model = DynamicShapeMap() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_round.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_round.py new file mode 100644 index 0000000000000000000000000000000000000000..decbf036553cb76544a19e531e5aee98d792ae0b --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_round.py @@ -0,0 +1,21 @@ +# mypy: allow-untyped-defs +import torch + +from torch._export.db.case import SupportLevel +from torch.export import Dim + +class DynamicShapeRound(torch.nn.Module): + """ + Calling round on dynamic shapes is not supported. + """ + + def forward(self, x): + return x[: round(x.shape[0] / 2)] + +x = torch.randn(3, 2) +dim0_x = Dim("dim0_x") +example_args = (x,) +tags = {"torch.dynamic-shape", "python.builtin"} +support_level = SupportLevel.NOT_SUPPORTED_YET +dynamic_shapes = {"x": {0: dim0_x}} +model = DynamicShapeRound() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_slicing.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_slicing.py new file mode 100644 index 0000000000000000000000000000000000000000..360fe15f6f98d9d735366bfa53371d79e0b00209 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_slicing.py @@ -0,0 +1,15 @@ +# mypy: allow-untyped-defs +import torch + +class DynamicShapeSlicing(torch.nn.Module): + """ + Slices with dynamic shape arguments should be captured into the graph + rather than being baked in. + """ + + def forward(self, x): + return x[: x.shape[0] - 2, x.shape[1] - 1 :: 2] + +example_args = (torch.randn(3, 2),) +tags = {"torch.dynamic-shape"} +model = DynamicShapeSlicing() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_view.py b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_view.py new file mode 100644 index 0000000000000000000000000000000000000000..c45d4aeebb0282a0f56c58a587b4bfe1655f50e3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/dynamic_shape_view.py @@ -0,0 +1,17 @@ +# mypy: allow-untyped-defs +import torch + +class DynamicShapeView(torch.nn.Module): + """ + Dynamic shapes should be propagated to view arguments instead of being + baked into the exported graph. + """ + + def forward(self, x): + new_x_shape = x.size()[:-1] + (2, 5) + x = x.view(*new_x_shape) + return x.permute(0, 2, 1) + +example_args = (torch.randn(10, 10),) +tags = {"torch.dynamic-shape"} +model = DynamicShapeView() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/fn_with_kwargs.py b/lib/python3.12/site-packages/torch/_export/db/examples/fn_with_kwargs.py new file mode 100644 index 0000000000000000000000000000000000000000..46b2637b398c21bf9399d0a3fa2a964354beea3e --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/fn_with_kwargs.py @@ -0,0 +1,30 @@ +# mypy: allow-untyped-defs +import torch + +class FnWithKwargs(torch.nn.Module): + """ + Keyword arguments are not supported at the moment. + """ + + def forward(self, pos0, tuple0, *myargs, mykw0, **mykwargs): + out = pos0 + for arg in tuple0: + out = out * arg + for arg in myargs: + out = out * arg + out = out * mykw0 + out = out * mykwargs["input0"] * mykwargs["input1"] + return out + +example_args = ( + torch.randn(4), + (torch.randn(4), torch.randn(4)), + *[torch.randn(4), torch.randn(4)] +) +example_kwargs = { + "mykw0": torch.randn(4), + "input0": torch.randn(4), + "input1": torch.randn(4), +} +tags = {"python.data-structure"} +model = FnWithKwargs() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/list_contains.py b/lib/python3.12/site-packages/torch/_export/db/examples/list_contains.py new file mode 100644 index 0000000000000000000000000000000000000000..35a140f4ee2e5d6f42c3509984333db896f1c081 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/list_contains.py @@ -0,0 +1,17 @@ +# mypy: allow-untyped-defs +import torch + +class ListContains(torch.nn.Module): + """ + List containment relation can be checked on a dynamic shape or constants. + """ + + def forward(self, x): + assert x.size(-1) in [6, 2] + assert x.size(0) not in [4, 5, 6] + assert "monkey" not in ["cow", "pig"] + return x + x + +example_args = (torch.randn(3, 2),) +tags = {"torch.dynamic-shape", "python.data-structure", "python.assert"} +model = ListContains() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/list_unpack.py b/lib/python3.12/site-packages/torch/_export/db/examples/list_unpack.py new file mode 100644 index 0000000000000000000000000000000000000000..3e2f8e2469a04b6b9d3358ed3602e9c1a8d6f35b --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/list_unpack.py @@ -0,0 +1,22 @@ +# mypy: allow-untyped-defs +from typing import List + +import torch + +class ListUnpack(torch.nn.Module): + """ + Lists are treated as static construct, therefore unpacking should be + erased after tracing. + """ + + def forward(self, args: List[torch.Tensor]): + """ + Lists are treated as static construct, therefore unpacking should be + erased after tracing. + """ + x, *y = args + return x + y[0] + +example_args = ([torch.randn(3, 2), torch.tensor(4), torch.tensor(5)],) +tags = {"python.control-flow", "python.data-structure"} +model = ListUnpack() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/model_attr_mutation.py b/lib/python3.12/site-packages/torch/_export/db/examples/model_attr_mutation.py new file mode 100644 index 0000000000000000000000000000000000000000..4aa623c7dc39efd94fecb8eb32caac3f7420f05d --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/model_attr_mutation.py @@ -0,0 +1,26 @@ +# mypy: allow-untyped-defs +import torch +from torch._export.db.case import SupportLevel + + +class ModelAttrMutation(torch.nn.Module): + """ + Attribute mutation is not supported. + """ + + def __init__(self) -> None: + super().__init__() + self.attr_list = [torch.randn(3, 2), torch.randn(3, 2)] + + def recreate_list(self): + return [torch.zeros(3, 2), torch.zeros(3, 2)] + + def forward(self, x): + self.attr_list = self.recreate_list() + return x.sum() + self.attr_list[0].sum() + + +example_args = (torch.randn(3, 2),) +tags = {"python.object-model"} +support_level = SupportLevel.NOT_SUPPORTED_YET +model = ModelAttrMutation() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/nested_function.py b/lib/python3.12/site-packages/torch/_export/db/examples/nested_function.py new file mode 100644 index 0000000000000000000000000000000000000000..e4076ac14dada40b4d78812666a9ec6b5e67045b --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/nested_function.py @@ -0,0 +1,23 @@ +# mypy: allow-untyped-defs +import torch + +class NestedFunction(torch.nn.Module): + """ + Nested functions are traced through. Side effects on global captures + are not supported though. + """ + + def forward(self, a, b): + x = a + b + z = a - b + + def closure(y): + nonlocal x + x += 1 + return x * y + z + + return closure(x) + +example_args = (torch.randn(3, 2), torch.randn(2)) +tags = {"python.closure"} +model = NestedFunction() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/null_context_manager.py b/lib/python3.12/site-packages/torch/_export/db/examples/null_context_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..80d09f68097edbe676077be183711dabe5cbc664 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/null_context_manager.py @@ -0,0 +1,21 @@ +# mypy: allow-untyped-defs +import contextlib + +import torch + +class NullContextManager(torch.nn.Module): + """ + Null context manager in Python will be traced out. + """ + + def forward(self, x): + """ + Null context manager in Python will be traced out. + """ + ctx = contextlib.nullcontext() + with ctx: + return x.sin() + x.cos() + +example_args = (torch.randn(3, 2),) +tags = {"python.context-manager"} +model = NullContextManager() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/optional_input.py b/lib/python3.12/site-packages/torch/_export/db/examples/optional_input.py new file mode 100644 index 0000000000000000000000000000000000000000..9693aa476f0eb31ec912c840ed3a4426ac832de8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/optional_input.py @@ -0,0 +1,20 @@ +# mypy: allow-untyped-defs +import torch +from torch._export.db.case import SupportLevel + + +class OptionalInput(torch.nn.Module): + """ + Tracing through optional input is not supported yet + """ + + def forward(self, x, y=torch.randn(2, 3)): + if y is not None: + return x + y + return x + + +example_args = (torch.randn(2, 3),) +tags = {"python.object-model"} +support_level = SupportLevel.NOT_SUPPORTED_YET +model = OptionalInput() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/pytree_flatten.py b/lib/python3.12/site-packages/torch/_export/db/examples/pytree_flatten.py new file mode 100644 index 0000000000000000000000000000000000000000..fe401b75e8b94be80247a806839be9c1c89bce25 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/pytree_flatten.py @@ -0,0 +1,16 @@ +# mypy: allow-untyped-defs +import torch + +from torch.utils import _pytree as pytree + +class PytreeFlatten(torch.nn.Module): + """ + Pytree from PyTorch can be captured by TorchDynamo. + """ + + def forward(self, x): + y, spec = pytree.tree_flatten(x) + return y[0] + 1 + +example_args = ({1: torch.randn(3, 2), 2: torch.randn(3, 2)},), +model = PytreeFlatten() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/scalar_output.py b/lib/python3.12/site-packages/torch/_export/db/examples/scalar_output.py new file mode 100644 index 0000000000000000000000000000000000000000..86d3b4645330c47c3625736b695d635f4ab58c70 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/scalar_output.py @@ -0,0 +1,23 @@ +# mypy: allow-untyped-defs +import torch + +from torch.export import Dim + +x = torch.randn(3, 2) +dim1_x = Dim("dim1_x") + +class ScalarOutput(torch.nn.Module): + """ + Returning scalar values from the graph is supported, in addition to Tensor + outputs. Symbolic shapes are captured and rank is specialized. + """ + def __init__(self) -> None: + super().__init__() + + def forward(self, x): + return x.shape[1] + 1 + +example_args = (x,) +tags = {"torch.dynamic-shape"} +dynamic_shapes = {"x": {1: dim1_x}} +model = ScalarOutput() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/specialized_attribute.py b/lib/python3.12/site-packages/torch/_export/db/examples/specialized_attribute.py new file mode 100644 index 0000000000000000000000000000000000000000..f17092f9afc681b91a982a8a2479ac1dde4f455d --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/specialized_attribute.py @@ -0,0 +1,26 @@ +# mypy: allow-untyped-defs +from enum import Enum + +import torch + +class Animal(Enum): + COW = "moo" + +class SpecializedAttribute(torch.nn.Module): + """ + Model attributes are specialized. + """ + + def __init__(self) -> None: + super().__init__() + self.a = "moo" + self.b = 4 + + def forward(self, x): + if self.a == Animal.COW.value: + return x * x + self.b + else: + raise ValueError("bad") + +example_args = (torch.randn(3, 2),) +model = SpecializedAttribute() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/static_for_loop.py b/lib/python3.12/site-packages/torch/_export/db/examples/static_for_loop.py new file mode 100644 index 0000000000000000000000000000000000000000..aa62b86d16d9b6a1c539976a891f58bd732ae30d --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/static_for_loop.py @@ -0,0 +1,16 @@ +# mypy: allow-untyped-defs +import torch + +class StaticForLoop(torch.nn.Module): + """ + A for loop with constant number of iterations should be unrolled in the exported graph. + """ + + def forward(self, x): + # constant + ret = [i + x for i in range(10)] + return ret + +example_args = (torch.randn(3, 2),) +tags = {"python.control-flow"} +model = StaticForLoop() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/static_if.py b/lib/python3.12/site-packages/torch/_export/db/examples/static_if.py new file mode 100644 index 0000000000000000000000000000000000000000..f169380159a45489142ce5ae3523b2e4504c6721 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/static_if.py @@ -0,0 +1,18 @@ +# mypy: allow-untyped-defs +import torch + +class StaticIf(torch.nn.Module): + """ + `if` statement with static predicate value should be traced through with the + taken branch. + """ + + def forward(self, x): + if len(x.shape) == 3: + return x + torch.ones(1, 1, 1) + + return x + +example_args = (torch.randn(3, 2, 2),) +tags = {"python.control-flow"} +model = StaticIf() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/tensor_setattr.py b/lib/python3.12/site-packages/torch/_export/db/examples/tensor_setattr.py new file mode 100644 index 0000000000000000000000000000000000000000..8fbc263e7ff2240a3cf8618c56f152e744aa40e8 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/tensor_setattr.py @@ -0,0 +1,15 @@ +# mypy: allow-untyped-defs +import torch + + +class TensorSetattr(torch.nn.Module): + """ + setattr() call onto tensors is not supported. + """ + def forward(self, x, attr): + setattr(x, attr, torch.randn(3, 2)) + return x + 4 + +example_args = (torch.randn(3, 2), "attr") +tags = {"python.builtin"} +model = TensorSetattr() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/type_reflection_method.py b/lib/python3.12/site-packages/torch/_export/db/examples/type_reflection_method.py new file mode 100644 index 0000000000000000000000000000000000000000..99ad42a153c512d65aaae1bcac2377ee1e124f25 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/type_reflection_method.py @@ -0,0 +1,22 @@ +# mypy: allow-untyped-defs +import torch + +class A: + @classmethod + def func(cls, x): + return 1 + x + +class TypeReflectionMethod(torch.nn.Module): + """ + type() calls on custom objects followed by attribute accesses are not allowed + due to its overly dynamic nature. + """ + + def forward(self, x): + a = A() + return type(a).func(x) + + +example_args = (torch.randn(3, 4),) +tags = {"python.builtin"} +model = TypeReflectionMethod() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/unsupported_operator.py b/lib/python3.12/site-packages/torch/_export/db/examples/unsupported_operator.py new file mode 100644 index 0000000000000000000000000000000000000000..f5a52d80b895b3b2c2d85b878ca4efea511e73ea --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/unsupported_operator.py @@ -0,0 +1,18 @@ +# mypy: allow-untyped-defs +import torch +from torch._export.db.case import SupportLevel + + +class TorchSymMin(torch.nn.Module): + """ + torch.sym_min operator is not supported in export. + """ + + def forward(self, x): + return x.sum() + torch.sym_min(x.size(0), 100) + + +example_args = (torch.randn(3, 2),) +tags = {"torch.operator"} +support_level = SupportLevel.NOT_SUPPORTED_YET +model = TorchSymMin() diff --git a/lib/python3.12/site-packages/torch/_export/db/examples/user_input_mutation.py b/lib/python3.12/site-packages/torch/_export/db/examples/user_input_mutation.py new file mode 100644 index 0000000000000000000000000000000000000000..3156b3a1bf2ec6f6361395de3dacb098ecf20c3a --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/examples/user_input_mutation.py @@ -0,0 +1,17 @@ +# mypy: allow-untyped-defs +import torch + + +class UserInputMutation(torch.nn.Module): + """ + Directly mutate user input in forward + """ + + def forward(self, x): + x.mul_(2) + return x.cos() + + +example_args = (torch.randn(3, 2),) +tags = {"torch.mutation"} +model = UserInputMutation() diff --git a/lib/python3.12/site-packages/torch/_export/db/gen_example.py b/lib/python3.12/site-packages/torch/_export/db/gen_example.py new file mode 100644 index 0000000000000000000000000000000000000000..8e44cade322bdde858c5dd05ac116cef47202a33 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/gen_example.py @@ -0,0 +1,21 @@ +import os +import sys + +import torch._export.db.examples as examples + +TEMPLATE = '''import torch + +def {case_name}(x): + """ + """ + + return +''' + +if __name__ == "__main__": + assert len(sys.argv) == 2 + root_dir = examples.__name__.replace(".", "/") + assert os.path.exists(root_dir) + with open(os.path.join(root_dir, sys.argv[1] + ".py"), "w") as f: + print("Writing to", f.name, "...") + f.write(TEMPLATE.format(case_name=sys.argv[1])) diff --git a/lib/python3.12/site-packages/torch/_export/db/logging.py b/lib/python3.12/site-packages/torch/_export/db/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..2078113fef157e38a465c78156c3b22ff4c235c7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/db/logging.py @@ -0,0 +1,47 @@ +# mypy: allow-untyped-defs + + +def exportdb_error_message(case_name: str): + from .examples import all_examples + from torch._utils_internal import log_export_usage + + ALL_EXAMPLES = all_examples() + # Detect whether case_name is really registered in exportdb. + if case_name in ALL_EXAMPLES: + url_case_name = case_name.replace("_", "-") + return f"See {case_name} in exportdb for unsupported case. \ + https://pytorch.org/docs/main/generated/exportdb/index.html#{url_case_name}" + else: + log_export_usage( + event="export.error.casenotregistered", + message=case_name, + ) + return f"{case_name} is unsupported." + + +def get_class_if_classified_error(e): + """ + Returns a string case name if the export error e is classified. + Returns None otherwise. + """ + + from torch._dynamo.exc import TorchRuntimeError, Unsupported, UserError + + ALWAYS_CLASSIFIED = "always_classified" + DEFAULT_CLASS_SIGIL = "case_name" + + # add error types that should be classified, along with any attribute name + # whose presence acts like a sigil to further distinguish which errors of + # that type should be classified. If the attribute name is None, then the + # error type is always classified. + _ALLOW_LIST = { + Unsupported: DEFAULT_CLASS_SIGIL, + UserError: DEFAULT_CLASS_SIGIL, + TorchRuntimeError: None, + } + if type(e) in _ALLOW_LIST: + attr_name = _ALLOW_LIST[type(e)] + if attr_name is None: + return ALWAYS_CLASSIFIED + return getattr(e, attr_name, None) + return None diff --git a/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25a33ab3711c010b5118ec2e8ef7e5a8f34fde96 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/node_metadata.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/node_metadata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..611382e3f4d557d0180f871a2d3c922ece1df2b3 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/node_metadata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/proxy_value.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/proxy_value.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7269ebccc107df0438bc7a4cfe56424ebbbeca94 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/pass_infra/__pycache__/proxy_value.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/pass_infra/node_metadata.py b/lib/python3.12/site-packages/torch/_export/pass_infra/node_metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..4aa9b8093c370dd565dfb7fb44e4b22474446af0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/pass_infra/node_metadata.py @@ -0,0 +1,32 @@ +from typing import Any, Dict, Set + + +NodeMetadataValue = Any + + +PROTECTED_KEYS: Set[str] = { + "val", + "stack_trace", + "nn_module_stack", + "debug_handle", + "tensor_meta", +} + + +class NodeMetadata: + def __init__(self, data: Dict[str, Any]) -> None: + self.data: Dict[str, Any] = data.copy() + + def __getitem__(self, key: str) -> NodeMetadataValue: + return self.data[key] + + def __setitem__(self, key: str, value: NodeMetadataValue) -> NodeMetadataValue: + if key in PROTECTED_KEYS: + raise RuntimeError(f"Could not override node key: {key}") + self.data[key] = value + + def __contains__(self, key: str) -> bool: + return key in self.data + + def copy(self) -> "NodeMetadata": + return NodeMetadata(self.data.copy()) diff --git a/lib/python3.12/site-packages/torch/_export/serde/__init__.py b/lib/python3.12/site-packages/torch/_export/serde/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/_export/serde/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..583dcf1ee27239d7f7c01ce4ba3861ea6b0130ba Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/serde/__pycache__/aoti_schema.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/aoti_schema.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb1f9d33a741d9e3c048b804cb4b08059f098c51 Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/aoti_schema.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/serde/__pycache__/dynamic_shapes.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/dynamic_shapes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e92ae9df5e8f577a2489120ba10f713174918e3f Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/dynamic_shapes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/serde/__pycache__/schema.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/schema.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..601c95a031309fe2110954d90c2b547724e5b6cf Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/schema.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/serde/__pycache__/schema_check.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/schema_check.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8679ae9794b66783c75f22ae2057bb6ccc80170d Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/schema_check.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/serde/__pycache__/union.cpython-312.pyc b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/union.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc49872ca20cfb115b4d2f77797f1f76428fa49b Binary files /dev/null and b/lib/python3.12/site-packages/torch/_export/serde/__pycache__/union.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/_export/serde/aoti_schema.py b/lib/python3.12/site-packages/torch/_export/serde/aoti_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..17d5ceda0ef078136248256d202737da2b44b519 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/aoti_schema.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass +from typing import List + +from torch._export.serde.schema import Node + + +@dataclass +class ExternKernelNode: + name: str + node: Node + + +@dataclass +class ExternKernelNodes: + nodes: List[ExternKernelNode] diff --git a/lib/python3.12/site-packages/torch/_export/serde/dynamic_shapes.py b/lib/python3.12/site-packages/torch/_export/serde/dynamic_shapes.py new file mode 100644 index 0000000000000000000000000000000000000000..6676399a054f3b4a9464998fd334c0c1b8432ff2 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/dynamic_shapes.py @@ -0,0 +1,321 @@ +import dataclasses +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from torch._dynamo.exc import UserError, UserErrorType +from torch.export.dynamic_shapes import ( + _check_dynamic_shapes, + _DerivedDim, + _Dim, + _DimHint, + _tree_map_with_path, + Dim, +) +from torch.utils._pytree import tree_map + +from .serialize import _dataclass_to_dict + + +@dataclasses.dataclass +class RootDim: + """ + This represents a _Dim object. + """ + + min: int + max: Union[int, None] + derived: List[str] + + +@dataclasses.dataclass +class DynamicShapesSpec: + """ + This stores a dynamic_shapes spec for de/serialization. + """ + + dynamic_shapes: Union[Dict[str, Any], Tuple[Any], List[Any], None] + dims: Dict[str, RootDim] + + +def _postprocess_serialized_shapes( + dynamic_shapes: Union[Dict[str, Any], Tuple[Any], List[Any], None], + dims: Dict[str, Dict[str, Union[int, List[str], None]]], + to_dict: Optional[bool] = False, +) -> Union[DynamicShapesSpec, Dict[str, Any]]: + """ + Sorts dims and dumps to dictionary format. + """ + from torch.utils._sympy.numbers import int_oo + + dims = { + k: RootDim( + min=v["min"], # type: ignore[arg-type] + max=None if v["max"] is int_oo else v["max"], # type: ignore[arg-type] + derived=sorted(v["derived"]), # type: ignore[arg-type] + ) + for k, v in sorted(dims.items()) + } + spec = DynamicShapesSpec(dynamic_shapes=dynamic_shapes, dims=dims) + if to_dict: + return _dataclass_to_dict(spec) + else: + return spec + + +def _dump_dynamic_shapes( + dynamic_shapes: Union[Dict[str, Any], Tuple[Any], List[Any], None], + args: Tuple[Any], + kwargs: Optional[Dict[str, Any]] = None, + to_dict: Optional[bool] = False, +) -> Union[DynamicShapesSpec, Dict[str, Any]]: + """ + Utility function for dynamic shapes serialization, serializing a dynamic_shapes spec. + Returns a DynamicShapesSpec dataclass containing 2 fields, "dynamic_shapes" and "dims". + Uses args & kwargs to distinguish between tensor-level and dim-level specs (only for Nones). + + dynamic_shapes: A pytree structure mirroring the dynamic_shapes input to export(): + - Each tensor input is represented with a list of values, non-tensor inputs with None. + - dynamic dimensions (i.e. symbols) in tensors and Dim enums are represented with strings. + - static dimensions are represented with ints. + + dims: A dictionary mapping each symbol name to the min/max range and derived dim names. + + For example: + ``` + dx = Dim("dx", min=4, max=16) + dy = dx + 1 + + inputs = ( + [ + torch.randn(4, 4), + torch.randn(5, 4), + ], + torch.randn(4), + torch.randn(4, 4), + "hello", + ) + dynamic_shapes = { + "a": [ + (dx, 4), + (dy, 4), + ], + "b": (Dim.STATIC,), + "c": None, + "d": None, + } + out = _dump_dynamic_shapes(dynamic_shapes, inputs, to_dict=True) + ``` + would generate the following output: + ``` + { + 'dynamic_shapes': ( + [ + ['dx', 4], + ['dx + 1', 4], + ], + ['_DimHint.STATIC'], + ['_DimHint.STATIC', '_DimHint.STATIC'], + None, + ), + 'dims': { + 'dx': { + 'min': 4, + 'max': 16, + 'derived': ['dx + 1'], + }, + }, + } + ``` + """ + dims: Dict[str, Dict[str, Any]] = {} + + def _standardize_shapes(path, tensor, shape): # type: ignore[no-untyped-def] + """ + Helps standardize the dynamic_shapes tree structure we serialize, + returning lists for each tensor shape, handling tensor-level Nones. + """ + if not isinstance(tensor, torch.Tensor): + return None + if shape is None: + return [Dim.STATIC] * len(tensor.shape) # type: ignore[attr-defined] + + out = [] + if isinstance(shape, dict): + for i, s in enumerate(tensor.shape): + out.append(s if shape.get(i) is None else shape.get(i)) + else: + assert isinstance(shape, (tuple, list)) + for i, s in enumerate(tensor.shape): + out.append(s if shape[i] is None else shape[i]) + return out + + def _track_dim_from_dims( + val: Union[None, int, _DimHint, _Dim] + ) -> Union[None, int, str]: + """ + Tracks dims, ranges, derived dims from the standardized dynamic_shapes spec. + """ + if val is None or isinstance(val, int): # non-tensor input or static + return val + if isinstance(val, _DimHint): # store enum as string + return val.__class__.__name__ + "." + val.name + + assert isinstance(val, _Dim) + + # track root dim + root = val.root if isinstance(val, _DerivedDim) else val # type: ignore[attr-defined] + if root.__name__ not in dims: + dims[root.__name__] = { + "min": root.min, # type: ignore[attr-defined,union-attr] + "max": root.max, # type: ignore[attr-defined,union-attr] + "derived": set(), + } + + # track derived dims + if isinstance(val, _DerivedDim): + dims[root.__name__]["derived"].add(val.__name__) + + return val.__name__ + + if dynamic_shapes is None: + return {"dynamic_shapes": None, "dims": {}} + + # convert to tuple of specs, for each arg/kwarg + kwargs = kwargs or {} + if isinstance(dynamic_shapes, dict): + dynamic_shapes = dynamic_shapes.values() # type: ignore[assignment] + dynamic_shapes = tuple(dynamic_shapes) + combined_args = tuple(args) + tuple(kwargs.values()) + + # run same check when we're processing shapes for export - is this too lazy? + _check_dynamic_shapes(dict(enumerate(combined_args)), dynamic_shapes) # type: ignore[arg-type] + + tree_shapes = _tree_map_with_path( + _standardize_shapes, combined_args, dynamic_shapes, tree_name="inputs" + ) + serialized_shapes = tree_map(_track_dim_from_dims, tree_shapes) + return _postprocess_serialized_shapes(serialized_shapes, dims, to_dict=to_dict) + + +def _load_dynamic_shapes( + spec: Union[DynamicShapesSpec, Dict[str, Any]], + from_dict: Optional[bool] = False, +) -> Union[Dict[str, Any], Tuple[Any], List[Any], None]: + """ + Utility function for dynamic shapes serialization. + Deserializes a DynamicShapesSpec or corresponding dictionary into a dynamic_shapes input to export(). + """ + import sympy + + from torch.fx.experimental.symbolic_shapes import _is_supported_equivalence + + if from_dict: + if not isinstance(spec, dict): + raise UserError( + UserErrorType.INVALID_INPUT, + f"With from_dict=True, expected `spec` to be a dict, got {type(spec)}", + ) + if sorted(spec.keys()) != ["dims", "dynamic_shapes"]: + raise UserError( + UserErrorType.INVALID_INPUT, + "With from_dict=True, expected `spec` to have keys `dims` and `dynamic_shapes`, " + f"instead found {spec.keys()}", + ) + dims = {} + for k, v in spec["dims"].items(): + if not isinstance(k, str): + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected `spec['dims']` keys to be strings for symbols, got key {type(k)}", + ) + if sorted(v.keys()) != ["derived", "max", "min"]: + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected `spec['dims']` values to have keys `derived`, `max`, and `min`, " + f"instead found {v.keys()}", + ) + if not isinstance(v["min"], int): + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected dims in `spec['dims']` to map `min` to an int, got {k}: {v['min']}", + ) + if not isinstance(v["max"], int) or v["max"] is None: + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected dims in `spec['dims']` to map `max` to an int or None, got {k}: {v['max']}", + ) + if not isinstance(v["derived"], list) or any( + not isinstance(d, str) for d in v["derived"] + ): + raise UserError( + UserErrorType.INVALID_INPUT, + "Expected dims in `spec['dims']` to map `derived` to a list of derived expressions, " + f"got {k}: {v['derived']}", + ) + dims[k] = RootDim(**v) + dynamic_shapes = spec["dynamic_shapes"] + else: + if not isinstance(spec, DynamicShapesSpec): + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected `spec` to be a DynamicShapesSpec, got {type(spec)}", + ) + dims = spec.dims + dynamic_shapes = spec.dynamic_shapes + + if dynamic_shapes is None: + return None + + dim_cache = {} + for name, info in dims.items(): + symbol = sympy.sympify(name) + if not isinstance(symbol, sympy.Symbol): + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected `spec['dims']` keys to be symbols, got {name}", + ) + dim_cache[name] = Dim(name, min=info.min, max=info.max) # cache root dim + for _expr in info.derived: + expr = sympy.sympify(_expr) + if len(expr.free_symbols) != 1 or symbol not in expr.free_symbols: + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected derived expressions in to have {name} as the only free symbol, got {expr}", + ) + if not _is_supported_equivalence(expr): + raise UserError( + UserErrorType.INVALID_INPUT, + f"Expected derived expressions to be linear expressions, got {expr}", + ) + modulus, remainder = sympy.polys.polytools.div(expr, symbol) + ddim = dim_cache[name] + if modulus != 1: + ddim = int(modulus) * ddim + if remainder != 0: + ddim = ddim + int(remainder) + dim_cache[_expr] = ddim # cache derived dims + + def deserialize_shape( + val: Union[None, int, str] + ) -> Union[None, int, _Dim, _DimHint]: + if val is None or isinstance(val, int): + return val + elif val == "_DimHint.AUTO": + return _DimHint.AUTO + elif val == "_DimHint.STATIC": + return _DimHint.STATIC + if not isinstance(val, str): + raise UserError( + UserErrorType.INVALID_INPUT, + "Expected leaves in `spec['dynamic_shapes']` to be ints, None, Dim.AUTO/STATIC, symbols, " + f" or derived expressions, got {val}", + ) + if val not in dim_cache: + raise UserError( + UserErrorType.INVALID_INPUT, + "Expected dims in `spec['dynamic_shapes']` to be tracked in `spec['dims']`, " + f"got {val} which is not in {dims.keys()}", + ) + return dim_cache[val] + + return tree_map(deserialize_shape, dynamic_shapes) diff --git a/lib/python3.12/site-packages/torch/_export/serde/schema.py b/lib/python3.12/site-packages/torch/_export/serde/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..cf117d8b2fadc270f5c28129d7c0c87ce70bf303 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/schema.py @@ -0,0 +1,402 @@ +# NOTE: This is a placeholder for iterating on export serialization schema design. +# Anything is subject to change and no guarantee is provided at this point. + +from dataclasses import dataclass, field +from enum import IntEnum +from typing import Annotated, Dict, List, Optional + +from torch._export.serde.union import _Union + +# NOTE: Please update this value if any modifications are made to the schema +SCHEMA_VERSION = (8, 2) +TREESPEC_VERSION = 1 + + +class ScalarType(IntEnum): + UNKNOWN = 0 + BYTE = 1 + CHAR = 2 + SHORT = 3 + INT = 4 + LONG = 5 + HALF = 6 + FLOAT = 7 + DOUBLE = 8 + COMPLEXHALF = 9 + COMPLEXFLOAT = 10 + COMPLEXDOUBLE = 11 + BOOL = 12 + BFLOAT16 = 13 + UINT16 = 28 + + +class Layout(IntEnum): + Unknown = 0 + SparseCoo = 1 + SparseCsr = 2 + SparseCsc = 3 + SparseBsr = 4 + SparseBsc = 5 + _mkldnn = 6 + Strided = 7 + + +class MemoryFormat(IntEnum): + Unknown = 0 + ContiguousFormat = 1 + ChannelsLast = 2 + ChannelsLast3d = 3 + PreserveFormat = 4 + + +@dataclass +class Device: + type: Annotated[str, 10] + index: Annotated[Optional[int], 20] = None + + +@dataclass(repr=False) +class SymExprHint(_Union): + as_int: Annotated[int, 10] + as_bool: Annotated[bool, 20] + as_float: Annotated[float, 30] + + +# This is for storing the symbolic expressions behind symints/symfloats/symbools +# For example, we can get something like +# SymExpr(expr_str="s0 + s1", hint=SymExprHint(as_int=4) +# if we also have the hint that s0 and s1 are both 2. +@dataclass +class SymExpr: + expr_str: Annotated[str, 10] + hint: Annotated[Optional[SymExprHint], 20] = None + + +@dataclass(repr=False) +class SymInt(_Union): + as_expr: Annotated[SymExpr, 10] + as_int: Annotated[int, 20] + +@dataclass(repr=False) +class SymFloat(_Union): + as_expr: Annotated[SymExpr, 10] + as_float: Annotated[float, 20] + + +@dataclass(repr=False) +class SymBool(_Union): + as_expr: Annotated[SymExpr, 10] + as_bool: Annotated[bool, 20] + + +@dataclass +class TensorMeta: + dtype: Annotated[ScalarType, 10] + sizes: Annotated[List[SymInt], 20] + requires_grad: Annotated[bool, 30] + device: Annotated[Device, 40] + strides: Annotated[List[SymInt], 50] + storage_offset: Annotated[SymInt, 60] + layout: Annotated[Layout, 70] + + +# In most cases we will use the "as_name" field to store arguments which are +# SymInts. +# The "as_int" field is used in the case where we have a list containing a mix +# of SymInt and ints (ex. [1, s0, ...]). We will serialize this type of list to +# be List[SymIntArgument] and map the SymInts to the "as_name" field, and ints +# to the "as_int" field. +@dataclass(repr=False) +class SymIntArgument(_Union): + as_name: Annotated[str, 10] + as_int: Annotated[int, 20] + +# In most cases we will use the "as_name" field to store arguments which are +# SymFloats. +# The "as_float" field is used in the case where we have a list containing a mix +# of SymFloat and float (ex. [1.0, s0, ...]). We will serialize this type of list to +# be List[SymFloatArgument] and map the SymFloats to the "as_name" field, and ints +# to the "as_float" field. +@dataclass(repr=False) +class SymFloatArgument(_Union): + as_name: Annotated[str, 10] + as_float: Annotated[float, 20] + +# In most cases we will use the "as_name" field to store arguments which are +# SymBools. +# The "as_bool" field is used in the case where we have a list containing a mix +# of SymBool and bools (ex. [True, i0, ...]). We will serialize this type of list to +# be List[SymboolArgument] and map the SymBools to the "as_name" field, and bools +# to the "as_bool" field. +@dataclass(repr=False) +class SymBoolArgument(_Union): + as_name: Annotated[str, 10] + as_bool: Annotated[bool, 20] + + +@dataclass +class TensorArgument: + name: Annotated[str, 10] + + +@dataclass +class TokenArgument: + name: Annotated[str, 10] + + +# This is use for storing the contents of a list which contain optional tensors +# (Tensor?[], ex. [Tensor, None, ...]), where the list will be serialized to the +# type List[OptionalTensorArgument], with tensor values seiralized to the +# "as_tensor" field, and None values serialized to the "as_none" field. +@dataclass(repr=False) +class OptionalTensorArgument(_Union): + as_tensor: Annotated[TensorArgument, 20] + as_none: Annotated[bool, 10] + + +@dataclass +class GraphArgument: + name: Annotated[str, 10] + graph: Annotated['Graph', 20] + + +@dataclass +class CustomObjArgument: + name: Annotated[str, 10] + class_fqn: Annotated[str, 20] + + +# This is actually a union type +@dataclass(repr=False) +class Argument(_Union): + as_none: Annotated[bool, 10] + as_tensor: Annotated[TensorArgument, 20] + as_tensors: Annotated[List[TensorArgument], 30] + as_int: Annotated[int, 50] + as_ints: Annotated[List[int], 70] + as_float: Annotated[float, 80] + as_floats: Annotated[List[float], 90] + as_string: Annotated[str, 100] + as_strings: Annotated[List[str], 101] + as_sym_int: Annotated[SymIntArgument, 110] + as_sym_ints: Annotated[List[SymIntArgument], 120] + as_scalar_type: Annotated[ScalarType, 130] + as_memory_format: Annotated[MemoryFormat, 140] + as_layout: Annotated[Layout, 150] + as_device: Annotated[Device, 160] + as_bool: Annotated[bool, 170] + as_bools: Annotated[List[bool], 180] + as_sym_bool: Annotated[SymBoolArgument, 182] + as_sym_bools: Annotated[List[SymBoolArgument], 184] + as_graph: Annotated[GraphArgument, 200] + as_optional_tensors: Annotated[List[OptionalTensorArgument], 190] + as_custom_obj: Annotated[CustomObjArgument, 210] + as_operator: Annotated[str, 220] + as_sym_float: Annotated[SymFloatArgument, 230] + as_sym_floats: Annotated[List[SymFloatArgument], 240] + +@dataclass +class NamedArgument: + # Argument name from the operator schema + name: Annotated[str, 10] + arg: Annotated[Argument, 20] + + +@dataclass +class Node: + target: Annotated[str, 10] + inputs: Annotated[List[NamedArgument], 20] + outputs: Annotated[List[Argument], 30] + metadata: Annotated[Dict[str, str], 40] + + +@dataclass +class Graph: + inputs: Annotated[List[Argument], 10] + outputs: Annotated[List[Argument], 20] + nodes: Annotated[List[Node], 30] + tensor_values: Annotated[Dict[str, TensorMeta], 40] + sym_int_values: Annotated[Dict[str, SymInt], 50] + sym_bool_values: Annotated[Dict[str, SymBool], 60] + # This is for deserializing the submodule graphs from higher order ops + # (ex. cond, map) where single tensor returns will just return a single + # tensor, rather than following export schema and returning a singleton + # list. + is_single_tensor_return: Annotated[bool, 70] = False + custom_obj_values: Annotated[Dict[str, CustomObjArgument], 80] = field(default_factory=dict) + sym_float_values: Annotated[Dict[str, SymFloat], 90] = field(default_factory=dict) + +@dataclass +class UserInputSpec: + # Actually, only tensors and SymInts are allowed here + arg: Annotated[Argument, 10] + + +@dataclass(repr=False) +class ConstantValue(_Union): + as_none: Annotated[bool, 10] + as_int: Annotated[int, 20] + as_float: Annotated[float, 30] + as_string: Annotated[str, 40] + as_bool: Annotated[bool, 50] + + +@dataclass +class InputToConstantInputSpec: + name: Annotated[str, 10] + value: Annotated[ConstantValue, 20] + + +@dataclass +class InputToParameterSpec: + arg: Annotated[TensorArgument, 10] + parameter_name: Annotated[str, 20] + + +@dataclass +class InputToBufferSpec: + arg: Annotated[TensorArgument, 10] + buffer_name: Annotated[str, 20] + persistent: Annotated[bool, 30] + + + +@dataclass +class InputToTensorConstantSpec: + arg: Annotated[TensorArgument, 10] + tensor_constant_name: Annotated[str, 20] + + +@dataclass +class InputToCustomObjSpec: + arg: Annotated[CustomObjArgument, 10] + custom_obj_name: Annotated[str, 20] + + +@dataclass +class InputTokenSpec: + arg: Annotated[TokenArgument, 10] + + +@dataclass(repr=False) +class InputSpec(_Union): + user_input: Annotated[UserInputSpec, 10] + parameter: Annotated[InputToParameterSpec, 20] + buffer: Annotated[InputToBufferSpec, 30] + tensor_constant: Annotated[InputToTensorConstantSpec, 40] + custom_obj: Annotated[InputToCustomObjSpec, 50] + token: Annotated[InputTokenSpec, 70] + constant_input: Annotated[InputToConstantInputSpec, 60] + + +@dataclass +class UserOutputSpec: + arg: Annotated[Argument, 10] + + +@dataclass +class LossOutputSpec: + arg: Annotated[TensorArgument, 10] + + +@dataclass +class BufferMutationSpec: + arg: Annotated[TensorArgument, 10] + buffer_name: Annotated[str, 20] + + +@dataclass +class GradientToParameterSpec: + arg: Annotated[TensorArgument, 10] + parameter_name: Annotated[str, 20] + + +@dataclass +class GradientToUserInputSpec: + arg: Annotated[TensorArgument, 10] + user_input_name: Annotated[str, 20] + + +@dataclass +class UserInputMutationSpec: + arg: Annotated[TensorArgument, 10] + user_input_name: Annotated[str, 20] + + +@dataclass +class OutputTokenSpec: + arg: Annotated[TokenArgument, 10] + + +@dataclass(repr=False) +class OutputSpec(_Union): + user_output: Annotated[UserOutputSpec, 10] + loss_output: Annotated[LossOutputSpec, 20] + buffer_mutation: Annotated[BufferMutationSpec, 30] + gradient_to_parameter: Annotated[GradientToParameterSpec, 40] + gradient_to_user_input: Annotated[GradientToUserInputSpec, 50] + user_input_mutation: Annotated[UserInputMutationSpec, 60] + token: Annotated[OutputTokenSpec, 70] + + +@dataclass +class GraphSignature: + input_specs: Annotated[List[InputSpec], 10] + output_specs: Annotated[List[OutputSpec], 20] + + +@dataclass +class RangeConstraint: + min_val: Annotated[Optional[int], 10] + max_val: Annotated[Optional[int], 20] + + +@dataclass +class ModuleCallSignature: + inputs: Annotated[List[Argument], 10] + outputs: Annotated[List[Argument], 20] + + # These are serialized by calling pytree.treespec_loads + # And deserialized by calling pytree.treespec_dumps + in_spec: Annotated[str, 30] + out_spec: Annotated[str, 40] + + # This field is used to prettify the graph placeholders + # after we ser/der and retrace + forward_arg_names: Annotated[Optional[List[str]], 50] = None + + +@dataclass +class ModuleCallEntry: + fqn: Annotated[str, 10] + signature: Annotated[Optional[ModuleCallSignature], 30] = None + + +@dataclass +class GraphModule: + graph: Annotated[Graph, 10] + signature: Annotated[GraphSignature, 50] + # This is used for unflattening, by tracking the calling structure of all of + # the modules in order to unflatten the modules back to the eager calling + # conventions. + module_call_graph: Annotated[List[ModuleCallEntry], 60] + metadata: Annotated[Dict[str, str], 40] = field(default_factory=dict) + + +# Invariant: Every time a change is made to the schema, one of the versions +# should be upadted. +@dataclass +class SchemaVersion: + major: Annotated[int, 10] # Major version number is bumped every time a breaking change is made. + minor: Annotated[int, 20] # Minor version number is bumped when a compatible change is made. + + +@dataclass +class ExportedProgram: + graph_module: Annotated[GraphModule, 10] + # Key is the opset namespace (ex. aten), and value is the version number + opset_version: Annotated[Dict[str, int], 20] + range_constraints: Annotated[Dict[str, RangeConstraint], 30] + schema_version: Annotated[SchemaVersion, 60] + verifiers: Annotated[List[str], 70] = field(default_factory=list) + torch_version: Annotated[str, 80] = "<=2.4" diff --git a/lib/python3.12/site-packages/torch/_export/serde/schema.thrift b/lib/python3.12/site-packages/torch/_export/serde/schema.thrift new file mode 100644 index 0000000000000000000000000000000000000000..8d986f06c9ea3d789d91033be67e9380f223ba0d --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/schema.thrift @@ -0,0 +1,314 @@ +// @generated by update_schema.py +// checksum<<0e89c5e620ad16c05bfe4fa2060ad43dcb0938dc31d77faad36b92f216c2c903>> + +namespace py3 torch._export +namespace cpp2 torch._export.schema + +enum Layout { + Unknown = 0, + SparseCoo = 1, + SparseCsr = 2, + SparseCsc = 3, + SparseBsr = 4, + SparseBsc = 5, + _mkldnn = 6, + Strided = 7, +} + + +enum MemoryFormat { + Unknown = 0, + ContiguousFormat = 1, + ChannelsLast = 2, + ChannelsLast3d = 3, + PreserveFormat = 4, +} + + +enum ScalarType { + UNKNOWN = 0, + BYTE = 1, + CHAR = 2, + SHORT = 3, + INT = 4, + LONG = 5, + HALF = 6, + FLOAT = 7, + DOUBLE = 8, + COMPLEXHALF = 9, + COMPLEXFLOAT = 10, + COMPLEXDOUBLE = 11, + BOOL = 12, + BFLOAT16 = 13, + UINT16 = 28, +} + + +struct Device { + 10: string type; + 20: optional i64 index; +} + +union SymExprHint { + 10: i64 as_int; + 20: bool as_bool; + 30: double as_float; +} + +struct SymExpr { + 10: string expr_str; + 20: optional SymExprHint hint; +} + +union SymInt { + 10: SymExpr as_expr; + 20: i64 as_int; +} + +union SymFloat { + 10: SymExpr as_expr; + 20: double as_float; +} + +union SymBool { + 10: SymExpr as_expr; + 20: bool as_bool; +} + +struct TensorMeta { + 10: ScalarType dtype; + 20: list sizes; + 30: bool requires_grad; + 40: Device device; + 50: list strides; + 60: SymInt storage_offset; + 70: Layout layout; +} + +union SymIntArgument { + 10: string as_name; + 20: i64 as_int; +} + +union SymFloatArgument { + 10: string as_name; + 20: double as_float; +} + +union SymBoolArgument { + 10: string as_name; + 20: bool as_bool; +} + +struct TensorArgument { + 10: string name; +} + +struct TokenArgument { + 10: string name; +} + +union OptionalTensorArgument { + 20: TensorArgument as_tensor; + 10: bool as_none; +} + +struct GraphArgument { + 10: string name; + 20: Graph graph; +} + +struct CustomObjArgument { + 10: string name; + 20: string class_fqn; +} + +union Argument { + 10: bool as_none; + 20: TensorArgument as_tensor; + 30: list as_tensors; + 50: i64 as_int; + 70: list as_ints; + 80: double as_float; + 90: list as_floats; + 100: string as_string; + 101: list as_strings; + 110: SymIntArgument as_sym_int; + 120: list as_sym_ints; + 130: ScalarType as_scalar_type; + 140: MemoryFormat as_memory_format; + 150: Layout as_layout; + 160: Device as_device; + 170: bool as_bool; + 180: list as_bools; + 182: SymBoolArgument as_sym_bool; + 184: list as_sym_bools; + 200: GraphArgument as_graph; + 190: list as_optional_tensors; + 210: CustomObjArgument as_custom_obj; + 220: string as_operator; + 230: SymFloatArgument as_sym_float; + 240: list as_sym_floats; +} + +struct NamedArgument { + 10: string name; + 20: Argument arg; +} + +struct Node { + 10: string target; + 20: list inputs; + 30: list outputs; + 40: map metadata; +} + +struct Graph { + 10: list inputs; + 20: list outputs; + 30: list nodes; + 40: map tensor_values; + 50: map sym_int_values; + 60: map sym_bool_values; + 70: bool is_single_tensor_return; + 80: map custom_obj_values; + 90: map sym_float_values; +} + +struct UserInputSpec { + 10: Argument arg; +} + +union ConstantValue { + 10: bool as_none; + 20: i64 as_int; + 30: double as_float; + 40: string as_string; + 50: bool as_bool; +} + +struct InputToConstantInputSpec { + 10: string name; + 20: ConstantValue value; +} + +struct InputToParameterSpec { + 10: TensorArgument arg; + 20: string parameter_name; +} + +struct InputToBufferSpec { + 10: TensorArgument arg; + 20: string buffer_name; + 30: bool persistent; +} + +struct InputToTensorConstantSpec { + 10: TensorArgument arg; + 20: string tensor_constant_name; +} + +struct InputToCustomObjSpec { + 10: CustomObjArgument arg; + 20: string custom_obj_name; +} + +struct InputTokenSpec { + 10: TokenArgument arg; +} + +union InputSpec { + 10: UserInputSpec user_input; + 20: InputToParameterSpec parameter; + 30: InputToBufferSpec buffer; + 40: InputToTensorConstantSpec tensor_constant; + 50: InputToCustomObjSpec custom_obj; + 70: InputTokenSpec token; + 60: InputToConstantInputSpec constant_input; +} + +struct UserOutputSpec { + 10: Argument arg; +} + +struct LossOutputSpec { + 10: TensorArgument arg; +} + +struct BufferMutationSpec { + 10: TensorArgument arg; + 20: string buffer_name; +} + +struct GradientToParameterSpec { + 10: TensorArgument arg; + 20: string parameter_name; +} + +struct GradientToUserInputSpec { + 10: TensorArgument arg; + 20: string user_input_name; +} + +struct UserInputMutationSpec { + 10: TensorArgument arg; + 20: string user_input_name; +} + +struct OutputTokenSpec { + 10: TokenArgument arg; +} + +union OutputSpec { + 10: UserOutputSpec user_output; + 20: LossOutputSpec loss_output; + 30: BufferMutationSpec buffer_mutation; + 40: GradientToParameterSpec gradient_to_parameter; + 50: GradientToUserInputSpec gradient_to_user_input; + 60: UserInputMutationSpec user_input_mutation; + 70: OutputTokenSpec token; +} + +struct GraphSignature { + 10: list input_specs; + 20: list output_specs; +} + +struct RangeConstraint { + 10: optional i64 min_val; + 20: optional i64 max_val; +} + +struct ModuleCallSignature { + 10: list inputs; + 20: list outputs; + 30: string in_spec; + 40: string out_spec; + 50: optional list forward_arg_names; +} + +struct ModuleCallEntry { + 10: string fqn; + 30: optional ModuleCallSignature signature; +} + +struct GraphModule { + 10: Graph graph; + 50: GraphSignature signature; + 60: list module_call_graph; + 40: map metadata; +} + +struct SchemaVersion { + 10: i64 major; + 20: i64 minor; +} + +struct ExportedProgram { + 10: GraphModule graph_module; + 20: map opset_version; + 30: map range_constraints; + 60: SchemaVersion schema_version; + 70: list verifiers; + 80: string torch_version; +} diff --git a/lib/python3.12/site-packages/torch/_export/serde/schema.yaml b/lib/python3.12/site-packages/torch/_export/serde/schema.yaml new file mode 100644 index 0000000000000000000000000000000000000000..38a7cda78b078c774c7786380d5bcb0057b83c86 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/schema.yaml @@ -0,0 +1,462 @@ +# @generated by update_schema.py +# checksum<<9e3a8f9389cd89b981fa654e4bf6ebc544f1696f9cac84818d60424143426f30>> +Argument: + kind: union + fields: + as_none: + type: bool + as_tensor: + type: TensorArgument + as_tensors: + type: List[TensorArgument] + as_int: + type: int + as_ints: + type: List[int] + as_float: + type: float + as_floats: + type: List[float] + as_string: + type: str + as_strings: + type: List[str] + as_sym_int: + type: SymIntArgument + as_sym_ints: + type: List[SymIntArgument] + as_scalar_type: + type: ScalarType + as_memory_format: + type: MemoryFormat + as_layout: + type: Layout + as_device: + type: Device + as_bool: + type: bool + as_bools: + type: List[bool] + as_sym_bool: + type: SymBoolArgument + as_sym_bools: + type: List[SymBoolArgument] + as_graph: + type: GraphArgument + as_optional_tensors: + type: List[OptionalTensorArgument] + as_custom_obj: + type: CustomObjArgument + as_operator: + type: str + as_sym_float: + type: SymFloatArgument + as_sym_floats: + type: List[SymFloatArgument] +BufferMutationSpec: + kind: struct + fields: + arg: + type: TensorArgument + buffer_name: + type: str +ConstantValue: + kind: union + fields: + as_none: + type: bool + as_int: + type: int + as_float: + type: float + as_string: + type: str + as_bool: + type: bool +CustomObjArgument: + kind: struct + fields: + name: + type: str + class_fqn: + type: str +Device: + kind: struct + fields: + type: + type: str + index: + type: Optional[int] + default: None +ExportedProgram: + kind: struct + fields: + graph_module: + type: GraphModule + opset_version: + type: Dict[str, int] + range_constraints: + type: Dict[str, RangeConstraint] + schema_version: + type: SchemaVersion + verifiers: + type: List[str] + default: '[]' + torch_version: + type: str + default: <=2.4 +GradientToParameterSpec: + kind: struct + fields: + arg: + type: TensorArgument + parameter_name: + type: str +GradientToUserInputSpec: + kind: struct + fields: + arg: + type: TensorArgument + user_input_name: + type: str +Graph: + kind: struct + fields: + inputs: + type: List[Argument] + outputs: + type: List[Argument] + nodes: + type: List[Node] + tensor_values: + type: Dict[str, TensorMeta] + sym_int_values: + type: Dict[str, SymInt] + sym_bool_values: + type: Dict[str, SymBool] + is_single_tensor_return: + type: bool + default: 'False' + custom_obj_values: + type: Dict[str, CustomObjArgument] + default: '{}' + sym_float_values: + type: Dict[str, SymFloat] + default: '{}' +GraphArgument: + kind: struct + fields: + name: + type: str + graph: + type: Graph +GraphModule: + kind: struct + fields: + graph: + type: Graph + signature: + type: GraphSignature + module_call_graph: + type: List[ModuleCallEntry] + metadata: + type: Dict[str, str] + default: '{}' +GraphSignature: + kind: struct + fields: + input_specs: + type: List[InputSpec] + output_specs: + type: List[OutputSpec] +InputSpec: + kind: union + fields: + user_input: + type: UserInputSpec + parameter: + type: InputToParameterSpec + buffer: + type: InputToBufferSpec + tensor_constant: + type: InputToTensorConstantSpec + custom_obj: + type: InputToCustomObjSpec + token: + type: InputTokenSpec + constant_input: + type: InputToConstantInputSpec +InputToBufferSpec: + kind: struct + fields: + arg: + type: TensorArgument + buffer_name: + type: str + persistent: + type: bool +InputToConstantInputSpec: + kind: struct + fields: + name: + type: str + value: + type: ConstantValue +InputToCustomObjSpec: + kind: struct + fields: + arg: + type: CustomObjArgument + custom_obj_name: + type: str +InputToParameterSpec: + kind: struct + fields: + arg: + type: TensorArgument + parameter_name: + type: str +InputToTensorConstantSpec: + kind: struct + fields: + arg: + type: TensorArgument + tensor_constant_name: + type: str +InputTokenSpec: + kind: struct + fields: + arg: + type: TokenArgument +Layout: + kind: enum + fields: + Unknown: 0 + SparseCoo: 1 + SparseCsr: 2 + SparseCsc: 3 + SparseBsr: 4 + SparseBsc: 5 + _mkldnn: 6 + Strided: 7 +LossOutputSpec: + kind: struct + fields: + arg: + type: TensorArgument +MemoryFormat: + kind: enum + fields: + Unknown: 0 + ContiguousFormat: 1 + ChannelsLast: 2 + ChannelsLast3d: 3 + PreserveFormat: 4 +ModuleCallEntry: + kind: struct + fields: + fqn: + type: str + signature: + type: Optional[ModuleCallSignature] + default: None +ModuleCallSignature: + kind: struct + fields: + inputs: + type: List[Argument] + outputs: + type: List[Argument] + in_spec: + type: str + out_spec: + type: str + forward_arg_names: + type: Optional[List[str]] + default: None +NamedArgument: + kind: struct + fields: + name: + type: str + arg: + type: Argument +Node: + kind: struct + fields: + target: + type: str + inputs: + type: List[NamedArgument] + outputs: + type: List[Argument] + metadata: + type: Dict[str, str] +OptionalTensorArgument: + kind: union + fields: + as_tensor: + type: TensorArgument + as_none: + type: bool +OutputSpec: + kind: union + fields: + user_output: + type: UserOutputSpec + loss_output: + type: LossOutputSpec + buffer_mutation: + type: BufferMutationSpec + gradient_to_parameter: + type: GradientToParameterSpec + gradient_to_user_input: + type: GradientToUserInputSpec + user_input_mutation: + type: UserInputMutationSpec + token: + type: OutputTokenSpec +OutputTokenSpec: + kind: struct + fields: + arg: + type: TokenArgument +RangeConstraint: + kind: struct + fields: + min_val: + type: Optional[int] + max_val: + type: Optional[int] +ScalarType: + kind: enum + fields: + UNKNOWN: 0 + BYTE: 1 + CHAR: 2 + SHORT: 3 + INT: 4 + LONG: 5 + HALF: 6 + FLOAT: 7 + DOUBLE: 8 + COMPLEXHALF: 9 + COMPLEXFLOAT: 10 + COMPLEXDOUBLE: 11 + BOOL: 12 + BFLOAT16: 13 + UINT16: 28 +SchemaVersion: + kind: struct + fields: + major: + type: int + minor: + type: int +SymBool: + kind: union + fields: + as_expr: + type: SymExpr + as_bool: + type: bool +SymBoolArgument: + kind: union + fields: + as_name: + type: str + as_bool: + type: bool +SymExpr: + kind: struct + fields: + expr_str: + type: str + hint: + type: Optional[SymExprHint] + default: None +SymExprHint: + kind: union + fields: + as_int: + type: int + as_bool: + type: bool + as_float: + type: float +SymFloat: + kind: union + fields: + as_expr: + type: SymExpr + as_float: + type: float +SymFloatArgument: + kind: union + fields: + as_name: + type: str + as_float: + type: float +SymInt: + kind: union + fields: + as_expr: + type: SymExpr + as_int: + type: int +SymIntArgument: + kind: union + fields: + as_name: + type: str + as_int: + type: int +TensorArgument: + kind: struct + fields: + name: + type: str +TensorMeta: + kind: struct + fields: + dtype: + type: ScalarType + sizes: + type: List[SymInt] + requires_grad: + type: bool + device: + type: Device + strides: + type: List[SymInt] + storage_offset: + type: SymInt + layout: + type: Layout +TokenArgument: + kind: struct + fields: + name: + type: str +UserInputMutationSpec: + kind: struct + fields: + arg: + type: TensorArgument + user_input_name: + type: str +UserInputSpec: + kind: struct + fields: + arg: + type: Argument +UserOutputSpec: + kind: struct + fields: + arg: + type: Argument +SCHEMA_VERSION: +- 8 +- 2 +TREESPEC_VERSION: 1 diff --git a/lib/python3.12/site-packages/torch/_export/serde/schema_check.py b/lib/python3.12/site-packages/torch/_export/serde/schema_check.py new file mode 100644 index 0000000000000000000000000000000000000000..0030f339d4c96fa054048321d9797480a6c577b3 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/schema_check.py @@ -0,0 +1,653 @@ +# mypy: allow-untyped-defs +import dataclasses +import hashlib +import inspect +import re +import typing +from enum import IntEnum +from typing import Annotated, Any, Dict, ForwardRef, List, Optional, Tuple, Union + +from torch._export.serde import schema +from torch._export.serde.union import _Union + + +class SchemaUpdateError(Exception): + pass + + +def _check(x, msg): + if not x: + raise SchemaUpdateError(msg) + + +def _staged_schema(): + yaml_ret: Dict[str, Any] = {} + defs = {} + cpp_enum_defs: Dict[str, str] = {} + cpp_class_defs: Dict[str, str] = {} + cpp_type_decls: List[str] = [] + cpp_json_defs: List[str] = [] + thrift_enum_defs: List[str] = [] + thrift_type_defs: Dict[str, str] = {} + + def _handle_aggregate(ty) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]: + def dump_type(t, level: int) -> Tuple[str, str, str]: + CPP_TYPE_MAP = { + str: "std::string", + int: "int64_t", + float: "double", + bool: "bool", + } + THRIFT_TYPE_MAP = { + str: "string", + int: "i64", + float: "double", + bool: "bool", + } + if isinstance(t, type): + if t.__name__ in cpp_enum_defs: + return t.__name__, "int64_t", t.__name__ + else: + return ( + t.__name__, + CPP_TYPE_MAP.get(t, t.__name__), + THRIFT_TYPE_MAP.get(t, t.__name__), + ) + elif isinstance(t, str): + assert t in defs + assert t not in cpp_enum_defs + assert "[" not in t + return t, f"ForwardRef<{t}>", t + elif isinstance(t, ForwardRef): + return ( + t.__forward_arg__, + f"ForwardRef<{t.__forward_arg__}>", + t.__forward_arg__, + ) + elif o := typing.get_origin(t): + # Lemme know if there's a better way to do this. + if o == list: + yaml_head, cpp_head, thrift_head, thrift_tail = ( + "List", + "std::vector", + "list<", + ">", + ) + elif o == dict: + yaml_head, cpp_head, thrift_head, thrift_tail = ( + "Dict", + "std::unordered_map", + "map<", + ">", + ) + elif o == Union: + assert level == 0, "Optional is only supported at the top level." + args = typing.get_args(t) + assert len(args) == 2 and args[1] == type(None) + yaml_type, cpp_type, thrift_type = dump_type(args[0], level + 1) + return ( + f"Optional[{yaml_type}]", + f"std::optional<{cpp_type}>", + f"optional {thrift_type}", + ) + elif o == Annotated: + return dump_type(t.__origin__, level) + else: + raise AssertionError(f"Type {t} is not supported in export schema.") + yaml_arg_types, cpp_arg_types, thrift_arg_types = zip( + *[dump_type(x, level + 1) for x in typing.get_args(t)] + ) + return ( + (f"{yaml_head}[{', '.join(yaml_arg_types)}]"), + (f"{cpp_head}<{', '.join(cpp_arg_types)}>"), + f"{thrift_head}{', '.join(thrift_arg_types)}{thrift_tail}", + ) + else: + raise AssertionError(f"Type {t} is not supported in export schema.") + + def dump_cpp_value(v) -> str: + if v is None: + return "std::nullopt" + elif v is True: + return "true" + elif v is False: + return "false" + elif v == {}: + return "{}" + elif v == []: + return "{}" + elif v == (): + return "{}" + elif isinstance(v, str): + return f'"{v}"' + else: + raise AssertionError( + f"Default value {v} is not supported yet in export schema." + ) + + def dump_field(f) -> Tuple[Dict[str, Any], str, Optional[str], str, int]: + t, cpp_type, thrift_type = dump_type(f.type, 0) + ret = {"type": t} + cpp_default: Optional[str] = None + assert ( + typing.get_origin(f.type) == Annotated + ), f"Field {f.name} must be annotated with an integer id." + thrift_id = f.type.__metadata__[0] + assert ( + type(thrift_id) is int + ), f"Field {f.name} must be annotated with an integer id." + + value = dataclasses.MISSING + if f.default is not dataclasses.MISSING: + value = f.default + elif f.default_factory is not dataclasses.MISSING: + value = f.default_factory() + + if value is not dataclasses.MISSING: + default = str(value) + ret["default"] = default + cpp_default = dump_cpp_value(value) + + if t.startswith("Optional[") and value is not None: + raise AssertionError( + f"Optional field {ty.__name__}.{f.name} must have default value to be None." + ) + + return ret, cpp_type, cpp_default, thrift_type, thrift_id + + yaml_ret = {} + cpp_ret = {} + thrift_ret = {} + thrift_ids = set() + for f in dataclasses.fields(ty): + yaml_res, cpp_type, cpp_default, thrift_type, thrift_id = dump_field(f) + yaml_ret[f.name] = yaml_res + cpp_ret[f.name] = {"cpp_type": cpp_type, "cpp_default": cpp_default} + thrift_ret[f.name] = {"thrift_type": thrift_type, "thrift_id": thrift_id} + if thrift_id in thrift_ids: + raise AssertionError( + f"Duplicate thrift id {thrift_id} for field {f.name} in {ty.__name__}." + ) + thrift_ids.add(thrift_id) + return yaml_ret, cpp_ret, thrift_ret + + def _handle_int_enum(name, ty): + yaml_ret[name] = {"kind": "enum", "fields": {x.name: x.value for x in ty}} + cpp_enum_defs[ + name + ] = f""" +enum class {name} {{ +{chr(10).join([f" {x.name} = {x.value}," for x in ty])} +}}; +""" + thrift_enum_defs.append( + f""" +enum {name} {{ +{chr(10).join([f" {x.name} = {x.value}," for x in ty])} +}} +""" + ) + + def _handle_struct(name, ty): + fields, cpp_fields, thrift_fields = _handle_aggregate(ty) + yaml_ret[name] = {"kind": "struct", "fields": fields} + field_decls = "\n".join( + f" {f['cpp_type']} {name}{' = ' + f['cpp_default'] if f['cpp_default'] is not None else ''};" + for name, f in cpp_fields.items() + ) + + def accessor(name, ty): + type_name = fields[name]["type"] + if type_name in cpp_enum_defs: + return f""" + {type_name} get_{name}() const {{ + return static_cast<{type_name}>({name}); + }} +""" + return f""" + const {ty}& get_{name}() const {{ + return {name}; + }} +""" + + to_json_decl = f"void to_json(nlohmann::json& nlohmann_json_j, const {name}& nlohmann_json_t)" + to_json_def = f"""{{ +{chr(10).join([f' nlohmann_json_j["{name}"] = nlohmann_json_t.{name};' for name, f in cpp_fields.items()])} +}} +""" + from_json_decl = f"void from_json(const nlohmann::json& nlohmann_json_j, {name}& nlohmann_json_t)" + + from_json_def = f"""{{ + {name} nlohmann_json_default_obj; +{chr(10).join( + [f' nlohmann_json_t.{name} = nlohmann_json_j.value("{name}", nlohmann_json_default_obj.{name});' + for name, f in cpp_fields.items()])} +}} +""" + cpp_class_defs[ + name + ] = f""" +class {name} {{ + private: +{field_decls} + + public: +{"".join([accessor(name, f["cpp_type"]) for name, f in cpp_fields.items()])} + friend {to_json_decl}; + friend {from_json_decl}; +}}; +""" + cpp_json_defs.append(f"inline {to_json_decl} {to_json_def}") + cpp_json_defs.append(f"inline {from_json_decl} {from_json_def}") + cpp_type_decls.append(f"class {name};") + + thrift_type_defs[ + name + ] = f""" +struct {name} {{ +{chr(10).join(f" {f['thrift_id']}: {f['thrift_type']} {n};" for n, f in thrift_fields.items())} +}}""" + + def _handle_union(name, ty): + fields, cpp_fields, thrift_fields = _handle_aggregate(ty) + yaml_ret[name] = {"kind": "union", "fields": fields} + + def accessor(name, ty, idx): + return f""" + const {ty}& get_{name}() const {{ + return std::get<{idx + 1}>(variant_); + }} +""" + + to_json_branches = "".join( + [ + f""" + if (nlohmann_json_t.tag_ == Tag::{name.upper()}) {{ + nlohmann_json_j["{name}"] = nlohmann_json_t.get_{name}(); + return; + }}""" + for idx, (name, f) in enumerate(cpp_fields.items()) + ] + ) + from_json_branches = "".join( + [ + f""" + if (nlohmann_json_j.contains("{name}")) {{ + nlohmann_json_t.variant_.emplace<{idx + 1}>(nlohmann_json_j.at("{name}").template get<{f["cpp_type"]}>()); + nlohmann_json_t.tag_ = Tag::{name.upper()}; + return; + }}""" + for idx, (name, f) in enumerate(cpp_fields.items()) + ] + ) + + cpp_class_defs[ + name + ] = f""" +class {name} {{ + struct Void {{}}; + + public: + enum class Tag {{ + {", ".join([name.upper() for name in cpp_fields])} + }}; + + private: + std::variant variant_; + Tag tag_; + + public: + Tag tag() const {{ + return tag_; + }} +{"".join([accessor(name, f["cpp_type"], idx) for idx, (name, f) in enumerate(cpp_fields.items())])} + friend void to_json(nlohmann::json& nlohmann_json_j, const {name}& nlohmann_json_t) {{ +{to_json_branches} + }} + + friend void from_json(const nlohmann::json& nlohmann_json_j, {name}& nlohmann_json_t) {{ +{from_json_branches} + }} +}}; +""" + cpp_type_decls.append(f"class {name};") + + thrift_type_defs[ + name + ] = f""" +union {name} {{ +{chr(10).join(f" {f['thrift_id']}: {f['thrift_type']} {n};" for n, f in thrift_fields.items())} +}}""" + + for name in dir(schema): + if name.startswith("_"): + continue + + value = getattr(schema, name) + + if hasattr(value, "__module__") and value.__module__ != schema.__name__: + continue + + defs[name] = value + + class_ordering = {} + for name, value in defs.items(): + if isinstance(value, type): + if issubclass(value, IntEnum): + _handle_int_enum(name, value) + elif dataclasses.is_dataclass(value): + class_ordering[name] = inspect.findsource(value)[1] + if issubclass(value, _Union): + _handle_union(name, value) + else: + _handle_struct(name, value) + else: + raise AssertionError(f"Unknown schema type {name}: {value}") + elif isinstance(value, (int, tuple)): + assert name in ("SCHEMA_VERSION", "TREESPEC_VERSION") + else: + raise AssertionError(f"Unknown variable {name}: {value}") + + yaml_ret["SCHEMA_VERSION"] = list(defs["SCHEMA_VERSION"]) + assert all(x > 0 for x in yaml_ret["SCHEMA_VERSION"]) + yaml_ret["TREESPEC_VERSION"] = defs["TREESPEC_VERSION"] + assert yaml_ret["TREESPEC_VERSION"] > 0 + + cpp_header = f""" +#pragma once + +#include +#include +#include +#include +#include + +#include + +#ifndef NLOHMANN_JSON_NAMESPACE_BEGIN +#define NLOHMANN_JSON_NAMESPACE_BEGIN namespace nlohmann {{ +#endif + +#ifndef NLOHMANN_JSON_NAMESPACE_END +#define NLOHMANN_JSON_NAMESPACE_END }} +#endif + +// https://github.com/nlohmann/json/pull/2117 +NLOHMANN_JSON_NAMESPACE_BEGIN +template +struct adl_serializer> {{ + static void to_json(json& j, const std::optional& opt) {{ + if (opt == std::nullopt) {{ + j = nullptr; + }} else {{ + j = *opt; // this will call adl_serializer::to_json which will + // find the free function to_json in T's namespace! + }} + }} + + static void from_json(const json& j, std::optional& opt) {{ + if (j.is_null()) {{ + opt = std::nullopt; + }} else {{ + opt = j.template get(); // same as above, but with + // adl_serializer::from_json + }} + }} +}}; +NLOHMANN_JSON_NAMESPACE_END + +namespace torch {{ +namespace _export {{ + +template +class ForwardRef {{ + static_assert(!std::is_reference_v, "ForwardRef cannot be a reference type"); + + public: + ForwardRef(): ptr_(std::make_unique()) {{}} + ForwardRef(ForwardRef&&) = default; + ForwardRef(const ForwardRef& other): ptr_(std::make_unique(*other.ptr_)) {{}} + ForwardRef& operator=(ForwardRef&&) = default; + ForwardRef& operator=(const ForwardRef& other) {{ + ptr_ = std::make_unique(*other.ptr_); + }} + const T& operator*() const {{ + return *ptr_; + }} + + const T* operator->() const {{ + return ptr_.get(); + }} + + void emplace(T&& t) {{ + ptr_ = std::make_unique(std::move(t)); + }} + + private: + std::unique_ptr ptr_; +}}; + +template +void to_json(nlohmann::json& j, const ForwardRef& p) {{ + j = *p; +}} + +template +void from_json(const nlohmann::json& j, ForwardRef& p) {{ + p.emplace(j.template get()); +}} + +{chr(10).join(cpp_type_decls)} +{"".join(cpp_enum_defs.values())} +{"".join(dict(sorted(cpp_class_defs.items(), key=lambda x: class_ordering[x[0]])).values())} +{chr(10).join(cpp_json_defs)} +}} // namespace _export +}} // namespace torch +""" + thrift_schema = f""" +namespace py3 torch._export +namespace cpp2 torch._export.schema +{chr(10).join(thrift_enum_defs)} +{chr(10).join(dict(sorted(thrift_type_defs.items(), key=lambda x: class_ordering[x[0]])).values())} +""" + return yaml_ret, cpp_header, thrift_schema + + +def _diff_schema(dst, src): + additions = {key: src[key] for key in src.keys() - dst.keys()} + subtractions = {key: dst[key] for key in dst.keys() - src.keys()} + + common_keys = src.keys() & dst.keys() + + versions = {"SCHEMA_VERSION", "TREESPEC_VERSION"} + common_keys -= versions + + for key in common_keys: + src_kind = src[key]["kind"] + src_fields = src[key]["fields"] + dst_kind = dst[key]["kind"] + dst_fields = dst[key]["fields"] + _check( + src_kind == dst_kind, + f"Type {key} changed kind from {dst_kind} to {src_kind}", + ) + assert isinstance(src_fields, dict) and isinstance(dst_fields, dict) + added_fields = { + key: src_fields[key] for key in src_fields.keys() - dst_fields.keys() + } + subtracted_fields = { + key: dst_fields[key] for key in dst_fields.keys() - src_fields.keys() + } + common_fields = src_fields.keys() & dst_fields.keys() + + for field in common_fields: + src_field = src_fields[field] + dst_field = dst_fields[field] + if src_kind == "struct": + _check( + src_field["type"] == dst_field["type"], + f"Type of the field {key}.{field} changed from {dst_field['type']} to {src_field['type']}", + ) + if "default" in src_field and "default" not in dst_field: + added_fields[field] = {} + added_fields[field]["default"] = src_field["default"] + if "default" not in src_field and "default" in dst_field: + subtracted_fields[field] = {} + subtracted_fields[field]["default"] = dst_field["default"] + elif src_kind == "enum": + _check( + src_field == dst_field, + f"Value of the enum field {key}.{field} changed from {dst_field} to {src_field}", + ) + elif src_kind == "union": + _check( + src_field["type"] == dst_field["type"], + f"Type of the field {key}.{field} changed from {dst_field['type']} to {src_field['type']}", + ) + else: + raise AssertionError(f"Unknown kind {src_kind}: {key}") + if len(added_fields) > 0: + assert key not in additions + additions[key] = {} + additions[key]["fields"] = added_fields + if len(subtracted_fields) > 0: + assert key not in subtractions + subtractions[key] = {} + subtractions[key]["fields"] = subtracted_fields + + return additions, subtractions + + +def _hash_content(s: str): + return hashlib.sha256(s.strip().encode("utf-8")).hexdigest() + + +@dataclasses.dataclass +class _Commit: + result: Dict[str, Any] + checksum_next: str + yaml_path: str + additions: Dict[str, Any] + subtractions: Dict[str, Any] + base: Dict[str, Any] + checksum_head: Optional[str] + cpp_header: str + cpp_header_path: str + thrift_checksum_head: Optional[str] + thrift_checksum_real: Optional[str] + thrift_checksum_next: str + thrift_schema: str + thrift_schema_path: str + + +def update_schema(): + import importlib.resources + + if importlib.resources.is_resource(__package__, "schema.yaml"): + content = importlib.resources.read_text(__package__, "schema.yaml") + match = re.search("checksum<<([A-Fa-f0-9]{64})>>", content) + _check(match is not None, "checksum not found in schema.yaml") + assert match is not None + checksum_head = match.group(1) + + thrift_content = importlib.resources.read_text(__package__, "schema.thrift") + match = re.search("checksum<<([A-Fa-f0-9]{64})>>", thrift_content) + _check(match is not None, "checksum not found in schema.thrift") + assert match is not None + thrift_checksum_head = match.group(1) + thrift_content = thrift_content.splitlines() + assert thrift_content[0].startswith("// @" + "generated") + assert thrift_content[1].startswith("// checksum<<") + thrift_checksum_real = _hash_content("\n".join(thrift_content[2:])) + + from yaml import load, Loader + + dst = load(content, Loader=Loader) + assert isinstance(dst, dict) + else: + checksum_head = None + thrift_checksum_head = None + thrift_checksum_real = None + dst = {"SCHEMA_VERSION": None, "TREESPEC_VERSION": None} + + src, cpp_header, thrift_schema = _staged_schema() + additions, subtractions = _diff_schema(dst, src) + yaml_path = __package__.replace(".", "/") + "/schema.yaml" + thrift_schema_path = __package__.replace(".", "/") + "/schema.thrift" + torch_prefix = "torch/" + assert yaml_path.startswith(torch_prefix) # sanity check + assert thrift_schema_path.startswith(torch_prefix) # sanity check + + return _Commit( + result=src, + checksum_next=_hash_content(repr(src)), + yaml_path=yaml_path, + additions=additions, + subtractions=subtractions, + base=dst, + checksum_head=checksum_head, + cpp_header=cpp_header, + cpp_header_path=torch_prefix + "csrc/utils/generated_serialization_types.h", + thrift_checksum_head=thrift_checksum_head, + thrift_checksum_real=thrift_checksum_real, + thrift_checksum_next=_hash_content(thrift_schema), + thrift_schema=thrift_schema, + thrift_schema_path=thrift_schema_path, + ) + + +def check(commit: _Commit, force_unsafe: bool = False): + next_version = None + reason = "" + # Step 1: Detect major schema updates. + if len(commit.additions) > 0: + for k, v in commit.additions.items(): + if k not in commit.base: + continue + kind = commit.result[k]["kind"] + fields = v["fields"] + for f, d in fields.items(): + if "default" not in d and kind == "struct": + reason += ( + f"Field {k}.{f} is added to schema.py without a default value as an incomparible change " + + "which requires major version bump.\n" + ) + next_version = [commit.base["SCHEMA_VERSION"][0] + 1, 1] + + if len(commit.subtractions) > 0: + for k, v in commit.subtractions.items(): + if k not in commit.result: + continue + for f in v["fields"]: + reason = f"Field {k}.{f} is removed from schema.py as an incompatible change which requires major version bump.\n" + next_version = [commit.base["SCHEMA_VERSION"][0] + 1, 1] + + if force_unsafe: + reason += "--force-unsafe is used." + next_version = commit.result["SCHEMA_VERSION"] + else: + # Step 2: Detect minor schema updates. + if next_version is None and len(commit.additions) > 0: + for k, v in commit.additions.items(): + for f in v["fields"]: + reason += ( + f"Field {k}.{f} is added to schema.py as an compatible change " + + "which still requires minor version bump.\n" + ) + next_version = [ + commit.base["SCHEMA_VERSION"][0], + commit.base["SCHEMA_VERSION"][1] + 1, + ] + if next_version is None and len(commit.subtractions) > 0: + for k, v in commit.subtractions.items(): + for f in v["fields"]: + reason += ( + f"Field {k}.{f} is removed from schema.py as an compatible change " + + "which still requires minor version bump.\n" + ) + next_version = [ + commit.base["SCHEMA_VERSION"][0], + commit.base["SCHEMA_VERSION"][1] + 1, + ] + + return next_version, reason diff --git a/lib/python3.12/site-packages/torch/_export/serde/serialize.py b/lib/python3.12/site-packages/torch/_export/serde/serialize.py new file mode 100644 index 0000000000000000000000000000000000000000..615fe21be23d412949b54d504f1c7e24b8b31472 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/serialize.py @@ -0,0 +1,3135 @@ +# mypy: allow-untyped-defs +import base64 +import copy +import copyreg +import dataclasses +import heapq +import inspect +import io +import json +import keyword +import logging +import math +import operator +import re +import traceback +import typing + +from collections import OrderedDict +from contextlib import contextmanager +from dataclasses import dataclass, field +from enum import Enum +from typing import ( + Annotated, + Any, + Callable, + cast, + Dict, + final, + Iterator, + List, + Optional, + Set, + Tuple, + Type, + Union, +) + +import sympy + +import torch +import torch.export.exported_program as ep +from torch._export.serde.schema import SchemaVersion +from torch._export.verifier import load_verifier +from torch._library.fake_class_registry import FakeScriptObject +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode +from torch.fx.experimental import symbolic_shapes +from torch.utils import _pytree as pytree +from torch.utils._pytree import treespec_dumps, treespec_loads +from torch.utils._sympy.numbers import int_oo +from torch.utils._sympy.value_ranges import ValueRanges + +from .schema import ( # type: ignore[attr-defined] + Argument, + BufferMutationSpec, + ConstantValue, + CustomObjArgument, + Device, + ExportedProgram, + GradientToParameterSpec, + GradientToUserInputSpec, + Graph, + GraphArgument, + GraphModule, + GraphSignature, + InputSpec, + InputToBufferSpec, + InputToConstantInputSpec, + InputToCustomObjSpec, + InputTokenSpec, + InputToParameterSpec, + InputToTensorConstantSpec, + Layout, + LossOutputSpec, + MemoryFormat, + ModuleCallEntry, + ModuleCallSignature, + NamedArgument, + Node, + OptionalTensorArgument, + OutputSpec, + OutputTokenSpec, + RangeConstraint, + ScalarType, + SCHEMA_VERSION, + SymBool, + SymBoolArgument, + SymExpr, + SymExprHint, + SymFloat, + SymFloatArgument, + SymInt, + SymIntArgument, + TensorArgument, + TensorMeta, + TokenArgument, + TREESPEC_VERSION, + UserInputMutationSpec, + UserInputSpec, + UserOutputSpec, +) +from .union import _Union +from ..utils import remove_proxy_from_state_dict + +__all__ = [ + "serialize", + "GraphModuleSerializer", + "ExportedProgramSerializer", + "GraphModuleDeserializer", + "ExportedProgramDeserializer", +] + +log = logging.getLogger(__name__) + + +class SerializeError(RuntimeError): + pass + + +def _reverse_map(d: Dict[Any, Enum]): + return {v.value: k for k, v in d.items()} + + +MetaType = Union[ + FakeTensor, int, torch.SymInt, float, torch.SymFloat, bool, torch.SymBool, ep.CustomObjArgument +] + + +ST_DELIMITER = ";" + +_TORCH_TO_SERIALIZE_DTYPE = { + torch.uint8: ScalarType.BYTE, + torch.int8: ScalarType.CHAR, + torch.uint16: ScalarType.UINT16, + torch.int16: ScalarType.SHORT, + torch.int32: ScalarType.INT, + torch.int64: ScalarType.LONG, + torch.float16: ScalarType.HALF, + torch.float32: ScalarType.FLOAT, + torch.float64: ScalarType.DOUBLE, + torch.complex32: ScalarType.COMPLEXHALF, + torch.complex64: ScalarType.COMPLEXFLOAT, + torch.complex128: ScalarType.COMPLEXDOUBLE, + torch.bool: ScalarType.BOOL, + torch.bfloat16: ScalarType.BFLOAT16, +} + + +_SERIALIZE_TO_TORCH_DTYPE = _reverse_map(_TORCH_TO_SERIALIZE_DTYPE) # type: ignore[arg-type] + + +_TORCH_TO_SERIALIZE_LAYOUT = { + torch.sparse_coo: Layout.SparseCoo, + torch.sparse_csr: Layout.SparseCsr, + torch.sparse_csc: Layout.SparseCsc, + torch.sparse_bsr: Layout.SparseBsr, + torch.sparse_bsc: Layout.SparseBsc, + torch._mkldnn: Layout._mkldnn, # type: ignore[attr-defined] + torch.strided: Layout.Strided, +} + + +_SERIALIZE_TO_TORCH_LAYOUT = _reverse_map(_TORCH_TO_SERIALIZE_LAYOUT) # type: ignore[arg-type] + + +_TORCH_TO_SERIALIZE_MEMORY_FORMAT = { + torch.contiguous_format: MemoryFormat.ContiguousFormat, + torch.channels_last: MemoryFormat.ChannelsLast, + torch.channels_last_3d: MemoryFormat.ChannelsLast3d, + torch.preserve_format: MemoryFormat.PreserveFormat, +} + + +_SERIALIZE_TO_TORCH_MEMORY_FORMAT = _reverse_map(_TORCH_TO_SERIALIZE_MEMORY_FORMAT) # type: ignore[arg-type] + +_SYM_FLOAT_OPS = { + operator.truediv, +} + +_SYM_INT_OPS = { + operator.mul, + operator.add, + operator.sub, + operator.floordiv, + operator.mod, + operator.pow, + torch.sym_int, + torch.sym_float, + torch.sym_ite, + torch.sym_max, + torch.sym_min, + torch.sym_sqrt, +} + + +_SYM_BOOL_OPS = { + operator.eq, + operator.ne, + operator.le, + operator.ge, + operator.lt, + operator.gt, + operator.neg, + operator.pos, + torch.sym_not, +} + + +assert not any(isinstance(op, torch._ops.OpOverload) for op in _SYM_INT_OPS) +assert not any(isinstance(op, torch._ops.OpOverload) for op in _SYM_BOOL_OPS) +assert not any(isinstance(op, torch._ops.OpOverload) for op in _SYM_FLOAT_OPS) + +@dataclass +class SerializedArtifact: + exported_program: bytes + state_dict: bytes + constants: bytes + example_inputs: bytes + + +@dataclass +class _SerializedProgram: + exported_program: ExportedProgram + state_dict: bytes + constants: bytes + example_inputs: bytes + + +def deserialize_device(d: Device) -> torch.device: + if d.index is None: + return torch.device(type=d.type) # type: ignore[call-overload] + return torch.device(type=d.type, index=d.index) + + +def _print_sympy(s: Union[torch.SymInt, torch.SymBool, torch.SymFloat, sympy.Expr]): + if isinstance(s, (torch.SymInt, torch.SymBool, torch.SymFloat)): + s = s.node.expr + return sympy.printing.repr.srepr(s) + + +def serialize_sym_int(s: Union[int, torch.SymInt]) -> SymInt: + if isinstance(s, (torch.SymInt, sympy.Symbol, int)): + if symbolic_shapes.is_concrete_int(s): + return SymInt.create(as_int=int(s)) + else: + assert isinstance(s, (torch.SymInt, sympy.Symbol)) + if s.node.hint is None: + return SymInt.create(as_expr=SymExpr(_print_sympy(s))) + else: + return SymInt.create( + as_expr=SymExpr( + _print_sympy(s), + hint=SymExprHint.create(as_int=s.node.hint), + ) + ) + else: + raise SerializeError( + f"SymInt should be either symbol or int, got `{s}` of type `{type(s)}`" + ) + +def serialize_sym_float(s: Union[float, torch.SymFloat]) -> SymFloat: + if isinstance(s, (torch.SymFloat, sympy.Symbol, float)): + if symbolic_shapes.is_concrete_float(s): + return SymFloat.create(as_float=float(s)) + else: + assert isinstance(s, (torch.SymFloat, sympy.Symbol)) + if s.node.hint is None: + return SymFloat.create(as_expr=SymExpr(_print_sympy(s))) + else: + return SymFloat.create( + as_expr=SymExpr( + _print_sympy(s), + hint=SymExprHint.create(as_float=s.node.hint), + ) + ) + else: + raise SerializeError( + f"SymFloat should be either symbol or float, got `{s}` of type `{type(s)}`" + ) + +def serialize_sym_bool(s: Union[bool, torch.SymBool]) -> SymBool: + if isinstance(s, (torch.SymBool, bool)): + if symbolic_shapes.is_concrete_bool(s): + return SymBool.create(as_bool=bool(s)) + else: + return SymBool.create( + as_expr=SymExpr(expr_str=_print_sympy(s)) + ) + else: + raise SerializeError( + f"SymBool should be either symbol or bool, got `{s}` of type `{type(s)}`" + ) + + +def serialize_tensor_meta(t: torch.Tensor) -> TensorMeta: + """ + Extract a TensorMeta describing `t`. + """ + return TensorMeta( + dtype=_TORCH_TO_SERIALIZE_DTYPE[t.dtype], + sizes=[serialize_sym_int(s) for s in t.shape], + requires_grad=t.requires_grad, + device=Device(type=t.device.type, index=t.device.index), + strides=[serialize_sym_int(s) for s in t.stride()], + storage_offset=serialize_sym_int(0), # TODO needs to be fixed. + layout=_TORCH_TO_SERIALIZE_LAYOUT[t.layout], + ) + + +_CURRENT_DESERIALIZER: Optional["GraphModuleDeserializer"] = None + + +def _reduce_fake_tensor(fake_tensor: FakeTensor): + is_parameter = isinstance(fake_tensor, torch.nn.Parameter) + tensor_meta = serialize_tensor_meta(fake_tensor) + tensor_meta_bytes = json.dumps( + _dataclass_to_dict(tensor_meta), cls=EnumEncoder + ).encode("utf-8") + return _reconstruct_fake_tensor, (tensor_meta_bytes, is_parameter) + + +def _reconstruct_fake_tensor( + serialized_tensor_meta: bytes, is_parameter: bool +) -> FakeTensor: + # Deserialize the bytes into a TensorMeta + json_tensor_meta = json.loads(serialized_tensor_meta.decode("utf-8")) + tensor_meta = _dict_to_dataclass(TensorMeta, json_tensor_meta) + # Find the current fake mode + assert ( + _CURRENT_DESERIALIZER is not None + ), "Need access to current deserializer state" + fake_tensor = _CURRENT_DESERIALIZER.deserialize_tensor_meta(tensor_meta) + if is_parameter: + fake_tensor = torch.nn.Parameter(fake_tensor) # type: ignore[assignment] + return fake_tensor + + +def serialize_torch_artifact(artifact: Optional[Any]) -> bytes: + if artifact is None: + return b"" + + assert ( + FakeTensor not in copyreg.dispatch_table + ), "Refusing to stomp on existing FakeTensor reducer" + try: + copyreg.pickle(FakeTensor, _reduce_fake_tensor) + buffer = io.BytesIO() + # This is a workaround for backend's tensor deserialization problem: + # unpickleTensor() always create a tensor on the device where it was originally saved + # This behavior is bad for multi-gpu training, as we wish to directly load the tensor + # on the designated device. + # For now, we simply move the tensor to cpu before saving. + # TODO: this should be fixed by deserialization instead. + torch.save(artifact, buffer) + return buffer.getvalue() + finally: + del copyreg.dispatch_table[FakeTensor] + + +def deserialize_torch_artifact(serialized: Union[Dict[str, Any], Tuple[Any, ...], bytes]): + if isinstance(serialized, (dict, tuple)): + return serialized + if len(serialized) == 0: + return {} + buffer = io.BytesIO(serialized) + buffer.seek(0) + # weights_only=False as we want to load custom objects here (e.g. ScriptObject) + artifact = torch.load(buffer, weights_only=False) + assert isinstance(artifact, (tuple, dict)) + return artifact + + +def _sympy_int_to_int(val: sympy.Expr, adjust: str) -> Optional[int]: + # Convert simple sympy Integers into concrete int + if val in (sympy.oo, int_oo): + return None + if val in (-sympy.oo, -int_oo): + return None + if isinstance(val, sympy.Integer): + return int(val) + + # TODO: Remove this adjustment when Ed gets rid of fractional ranges + log.warning( + "Export constraints cannot be non-integer expressions. Found " + "type %s, and value %s. We will attempt to %s " + "this value.", type(val), val, adjust + ) + + if adjust == "floor": + return math.floor(val) + elif adjust == "ceil": + return math.ceil(val) + else: + raise RuntimeError(f"Got invalid adjustment {adjust}") + + +def _int_to_sympy_int(val: Optional[int], default) -> sympy.Expr: + # Convert concrete int into simple sympy Integers + if val is None: + return default + if val == math.inf: + return int_oo + if val == -math.inf: + return -int_oo + return sympy.Integer(val) + + +def serialize_range_constraints( + range_constraints: Dict[sympy.Symbol, ValueRanges] +) -> Dict[str, RangeConstraint]: + return { + str(k): RangeConstraint( + _sympy_int_to_int(v.lower, "ceil"), # type: ignore[arg-type] + _sympy_int_to_int(v.upper, "floor"), # type: ignore[arg-type] + ) + for k, v in range_constraints.items() + } + + +def _get_schema_from_target(target): + if isinstance(target, torch._ops.OpOverload): + return target._schema + elif type(target) in _serialization_registry: + return _serialization_registry[type(target)].op_schema(target) + raise RuntimeError(f"Cannot find schema for {type(target)}") + + +def _is_single_tensor_return(target) -> bool: + schema = _get_schema_from_target(target) + returns = schema.returns + return len(returns) == 1 and isinstance(returns[0].real_type, torch.TensorType) + + +def _is_single_tensor_list_return(target: Any) -> bool: + schema = _get_schema_from_target(target) + returns = schema.returns + + if len(returns) != 1: + return False + return_type = returns[0].real_type + return isinstance(return_type, torch.ListType) and isinstance( + return_type.getElementType(), torch.TensorType + ) + + +@dataclass +class GraphState: + inputs: List[Argument] = field(default_factory=list) + outputs: List[Argument] = field(default_factory=list) + nodes: List[Node] = field(default_factory=list) + tensor_values: Dict[str, TensorMeta] = field(default_factory=dict) + sym_int_values: Dict[str, SymInt] = field(default_factory=dict) + sym_bool_values: Dict[str, SymBool] = field(default_factory=dict) + sym_float_values: Dict[str, SymFloat] = field(default_factory=dict) + is_single_tensor_return: bool = False + custom_obj_values: Dict[str, CustomObjArgument] = field(default_factory=dict) + + +class Final(type): + def __new__(metacls, name, bases, classdict): + for b in bases: + if isinstance(b, Final): + raise TypeError(f"type '{b.__name__}' is not an acceptable base type") + return type.__new__(metacls, name, bases, dict(classdict)) + + +@final +class GraphModuleSerializer(metaclass=Final): + def __init__( + self, + graph_signature: ep.ExportGraphSignature, + module_call_graph: List[ep.ModuleCallEntry], + ): + self.graph_state = GraphState() + self.graph_signature = graph_signature + self.module_call_graph = module_call_graph + self.custom_objs: Dict[str, torch._C.ScriptObject] = {} + self.duplicate_getitem_nodes: Dict[str, str] = {} + + @contextmanager + def save_graph_state(self): + saved = self.graph_state + self.graph_state = GraphState() + try: + yield + finally: + self.graph_state = saved + + def handle_placeholder(self, node: torch.fx.Node): + assert node.op == "placeholder" + if isinstance(node.meta["val"], torch.Tensor): + graph_input = Argument.create(as_tensor=TensorArgument(name=node.name)) + self.graph_state.tensor_values[node.name] = serialize_tensor_meta( + node.meta["val"] + ) + elif isinstance(node.meta["val"], torch.SymInt): + raise AssertionError("SymInt graph input is not implemented yet.") + elif isinstance(node.meta["val"], torch.SymFloat): + raise AssertionError("SymFloat graph input is not implemented yet.") + elif isinstance(node.meta["val"], (int, bool, str, float, type(None))): + graph_input = self.serialize_input(node.meta["val"]) + elif isinstance(node.meta["val"], ep.CustomObjArgument): + class_fqn = node.meta["val"].class_fqn + graph_input = Argument.create( + as_custom_obj=CustomObjArgument(name=node.name, class_fqn=class_fqn) + ) + self.graph_state.custom_obj_values[node.name] = ( + self.serialize_script_obj_meta(node.meta["val"]) + ) + else: + raise AssertionError(f"Unimplemented graph input type: {node.meta['val']}") + self.graph_state.inputs.append(graph_input) + + def handle_output(self, node: torch.fx.Node): + assert node.op == "output" + assert len(node.args) == 1, "FX.Node's args should have one arg" + node_args = node.args[0] + if isinstance(node_args, torch.fx.Node): + # For singleton tensor returns + self.graph_state.is_single_tensor_return = True + self.graph_state.outputs = [self.serialize_input(node_args)] + else: + assert isinstance(node_args, (tuple, list)) + self.graph_state.outputs = [self.serialize_input(arg) for arg in node_args] + + def serialize_operator(self, target) -> str: + if isinstance(target, str): + return target + elif target.__module__.startswith("torch._ops"): + # TODO(zhxchen17) Maybe provide a function name helper in FX. + # From torch.fx.node._get_qualified_name + module = target.__module__.replace("torch._ops", "torch.ops") + return f"{module}.{target.__name__}" + else: # TODO(zhxchen17) Don't catch all here. + return f"{target.__module__}.{target.__name__}" + + def handle_call_function(self, node: torch.fx.Node): + assert node.op == "call_function" + + # getitem has been handled in the producer node, skip it here + if node.target is operator.getitem: + return + + meta_val = node.meta.get("val") + if ( + node.target in _SYM_INT_OPS + or node.target in _SYM_BOOL_OPS + or node.target in _SYM_FLOAT_OPS + or (meta_val is not None and isinstance(meta_val, (torch.SymInt, torch.SymBool, torch.SymFloat))) + ): + assert len(node.kwargs) == 0 + + # Serialize the node + if isinstance(meta_val, torch.SymInt): + sym_output = Argument.create(as_sym_int=self.serialize_sym_int_output(node.name, meta_val)) + elif isinstance(meta_val, torch.SymFloat): + sym_output = Argument.create(as_sym_float=self.serialize_sym_float_output(node.name, meta_val)) + elif isinstance(meta_val, torch.SymBool): + sym_output = Argument.create(as_sym_bool=self.serialize_sym_bool_output(node.name, meta_val)) + else: + raise SerializeError(f"Unsupported symbolic type: {type(meta_val)}") + + ex_node = Node( + target=self.serialize_operator(node.target), + inputs=self.serialize_sym_op_inputs(node.target, node.args), + outputs=[sym_output], + metadata=self.serialize_metadata(node), + ) + elif isinstance(node.target, torch._ops.OpOverload): + ex_node = Node( + target=self.serialize_operator(node.target), + inputs=self.serialize_inputs(node.target, node.args, node.kwargs), + outputs=self.serialize_outputs(node), + # TODO: create a new tensor_values here, meta might have faketensor info + metadata=self.serialize_metadata(node), + ) + elif isinstance(node.target, torch._ops.HigherOrderOperator): + ex_node = Node( + target=self.serialize_operator(node.target), + inputs=self.serialize_hoo_inputs(node.args, node.kwargs), + outputs=self.serialize_hoo_outputs(node), + metadata=self.serialize_metadata(node), + ) + elif type(node.target) in _serialization_registry: + # Sanity check for unhandled serialization. + assert type(node.target) in _serialization_registry, f"{type(node.target)} is not supported in export serialization." + + handler = _serialization_registry[type(node.target)] + namespace = handler.namespace() + op_name = handler.to_op_name(node.target) + assert isinstance(namespace, str) and isinstance(op_name, str) + assert ":" not in namespace and ":" not in op_name + ex_node = Node( + target=f"#{namespace}:{op_name}", + inputs=self.serialize_inputs(node.target, node.args, node.kwargs), + outputs=self.serialize_outputs(node), + metadata=self.serialize_metadata(node), + ) + else: + raise SerializeError(f"Serializing {node.target} is not supported") + + self.graph_state.nodes.append(ex_node) + + def handle_get_attr(self, node): + pass + + def _output_node_at_index(self, node, index) -> Optional[torch.fx.Node]: + user_node = None + for user in node.users: + assert user.target is operator.getitem, f"{user} is not a getitem node" + if index == user.args[1]: + if user_node is None: + user_node = user + else: + # We want to deduplicate getitem nodes that are trying to + # index to the same index + self.duplicate_getitem_nodes[user.name] = user_node.name + return user_node + + def _output_node_name_at_index(self, node, index) -> str: + user_node = self._output_node_at_index(node, index) + if user_node is None: + return f"{node.name}_unused_{index}" + else: + return user_node.name + + def serialize_metadata(self, node: torch.fx.Node) -> Dict[str, str]: + ret = {} + if stack_trace := node.meta.get("stack_trace"): + ret["stack_trace"] = stack_trace + + if nn_module_stack := node.meta.get("nn_module_stack"): + + def export_nn_module_stack(val): + assert isinstance(val, tuple) and len(val) == 2 + path, ty = val + + assert isinstance(path, str) + assert isinstance(ty, str) + + return path + "," + ty + + # Serialize to "key,orig_path,type_str" + nn_module_list = [ + f"{k},{export_nn_module_stack(v)}" for k, v in nn_module_stack.items() + ] + ret["nn_module_stack"] = ST_DELIMITER.join(nn_module_list) + + if source_fn_st := node.meta.get("source_fn_stack"): + source_fn_list = [ + f"{source_fn[0]},{self.serialize_operator(source_fn[1])}" + for source_fn in source_fn_st + ] + ret["source_fn_stack"] = ST_DELIMITER.join(source_fn_list) + + if torch_fn := node.meta.get("torch_fn"): + ret["torch_fn"] = ST_DELIMITER.join(list(torch_fn)) + + if custom := node.meta.get("custom"): + try: + ret["custom"] = json.dumps(custom) + except Exception as e: + raise SerializeError( + f"Failed to serialize custom metadata for node {node.name} with error {e}" + ) from e + + return ret + + def serialize_script_obj_meta( + self, script_obj_meta: ep.CustomObjArgument + ) -> CustomObjArgument: + return CustomObjArgument( + name=script_obj_meta.name, + class_fqn=script_obj_meta.class_fqn, + ) + + def serialize_sym_op_inputs(self, op, args) -> List[NamedArgument]: + if isinstance(op, torch._ops.OpOverload): + args_names = [arg.name for arg in op._schema.arguments] + else: + assert op in _SYM_INT_OPS or op in _SYM_BOOL_OPS or op in _SYM_FLOAT_OPS + args_names = list(inspect.signature(op).parameters.keys()) + serialized_args = [] + for args_name, arg in zip(args_names, args): + serialized_args.append( + NamedArgument(name=args_name, arg=self.serialize_input(arg)) + ) + return serialized_args + + def serialize_inputs( + self, + target: Any, # torch._ops.OpOverload and other custom operator types. + args, + kwargs=None + ) -> List[NamedArgument]: + assert isinstance(target, (torch._ops.OpOverload, *_registered_extension_types())) + kwargs = kwargs or {} + serialized_args = [] + + schema = _get_schema_from_target(target) + + for i, schema_arg in enumerate(schema.arguments): + if schema_arg.name in kwargs: + serialized_args.append( + NamedArgument( + name=schema_arg.name, + arg=self.serialize_input(kwargs[schema_arg.name], schema_arg.type), + ) + ) + elif not schema_arg.kwarg_only and i < len(args): + serialized_args.append( + NamedArgument( + name=schema_arg.name, + arg=self.serialize_input(args[i], schema_arg.type), + ) + ) + else: + # We intentionally don't serialize the missing arguments + # with default values + pass + + return serialized_args + + def serialize_hoo_inputs(self, args, kwargs) -> List[NamedArgument]: + """ + For serializing HOO inputs since HOOs do not have a schema. + """ + inputs = [ + NamedArgument( + name="", + arg=self.serialize_input(a), + ) + for a in args + ] + inputs.extend( + [ + NamedArgument(name=name, arg=self.serialize_input(a)) + for name, a in kwargs.items() + ] + ) + return inputs + + def is_sym_int_arg(self, arg) -> bool: + return type(arg) is int or ( + isinstance(arg, torch.fx.Node) + and arg.name in self.graph_state.sym_int_values + ) + + def is_sym_float_arg(self, arg) -> bool: + return isinstance(arg, float) or ( + isinstance(arg, torch.fx.Node) + and arg.name in self.graph_state.sym_float_values + ) + + def is_sym_bool_arg(self, arg) -> bool: + return isinstance(arg, bool) or ( + isinstance(arg, torch.fx.Node) + and arg.name in self.graph_state.sym_bool_values + ) + + def serialize_input( + self, arg, arg_type: Optional[torch._C.Argument] = None + ) -> Argument: + import torch._inductor.ir as inductor_ir + + inductor_tensor_buffers = ( + inductor_ir.Buffer, + inductor_ir.ReinterpretView, + ) + + if isinstance(arg, torch.fx.Node): + if arg.op == "get_attr": + assert isinstance(arg.target, str) + attr = getattr(arg.graph.owning_module, arg.target) + + if isinstance(attr, torch.Tensor): + raise SerializeError( + "getattr nodes containing tensors should not appear in the graph" + ) + elif isinstance(attr, torch.fx.GraphModule): + with self.save_graph_state(): + graph = self.serialize_graph(attr) + return Argument.create( + as_graph=GraphArgument(name=arg.target, graph=graph) + ) + else: + raise SerializeError( + f"Unsupported getattr attribute {arg.target} with type: {type(attr)}" + ) + elif self.is_sym_int_arg(arg): + return Argument.create( + as_sym_int=SymIntArgument.create(as_name=arg.name) + ) + elif self.is_sym_float_arg(arg): + return Argument.create( + as_sym_float=SymFloatArgument.create(as_name=arg.name) + ) + elif self.is_sym_bool_arg(arg): + return Argument.create( + as_sym_bool=SymBoolArgument.create(as_name=arg.name) + ) + elif isinstance(arg.meta["val"], ep.CustomObjArgument): + return Argument.create( + as_custom_obj=CustomObjArgument( + name=arg.name, class_fqn=arg.meta["val"].class_fqn + ) + ) + elif arg.name in self.duplicate_getitem_nodes: + dedup_name = self.duplicate_getitem_nodes[arg.name] + return Argument.create(as_tensor=TensorArgument(name=dedup_name)) + else: + return Argument.create(as_tensor=TensorArgument(name=arg.name)) + elif isinstance(arg, inductor_tensor_buffers): + # Other branches are for arguments in fx node. + # This is a special branch for handling buffers (representing tensor arguments) + # for inductor's ExternalFallbackNode + # export_extern_kernel_node() is using this function to serialize arguments + arg_name = arg.get_name() + assert arg_name is not None, "Buffer must have valid name" + return Argument.create(as_tensor=TensorArgument(name=arg_name)) + elif isinstance(arg, torch.SymInt): + # This is a special branch for handling SymInt args in inductor's + # ExternalFallbackNode. + # For regular FX graph, SymInt arg should be a fx.Node with + # self.is_sym_int_arg(arg) being true + return Argument.create(as_sym_int=SymIntArgument.create(as_name=str(arg))) + elif isinstance(arg, torch.SymFloat): + # This is a special branch for handling SymFloat args in inductor's + # ExternalFallbackNode. + # For regular FX graph, SymInt arg should be a fx.Node with + # self.is_sym_float_arg(arg) being true + return Argument.create(as_sym_float=SymFloatArgument.create(as_name=str(arg))) + elif type(arg) is bool: + return Argument.create(as_bool=arg) + elif type(arg) is str: + return Argument.create(as_string=arg) + elif type(arg) is int: + return Argument.create(as_int=arg) + elif type(arg) is float: + return Argument.create(as_float=arg) + elif arg is None: + return Argument.create(as_none=True) + elif isinstance(arg, (list, tuple)): + if len(arg) == 0: + if arg_type is not None: + if isinstance(arg_type, torch.OptionalType): + arg_type = arg_type.getElementType() # type: ignore[assignment] + assert isinstance(arg_type, torch.ListType) + elem_type = arg_type.getElementType() + if isinstance(elem_type, torch.OptionalType): + elem_type = elem_type.getElementType() + + if isinstance(elem_type, torch.BoolType): + return Argument.create(as_bools=[]) + elif isinstance(elem_type, torch.IntType): + return Argument.create(as_ints=[]) + elif isinstance(elem_type, torch.FloatType): + return Argument.create(as_floats=[]) + elif isinstance(elem_type, torch.StringType): + return Argument.create(as_strings=[]) + elif isinstance(elem_type, torch.TensorType): + return Argument.create(as_tensors=[]) + else: + # I believe empty symint lists default to ints, but + # please file an issue if this is not the case + raise SerializeError(f"Empty list with type {elem_type} nyi.") + else: + # We could serialize this by default to a tensor list. This + # is needed in the HOO case + log.warning( + "Unsure how to serialize the given empty list, " + "as we don't know what is the type of this argument. " + "Serializing it as a tensor list by default." + ) + return Argument.create(as_tensors=[]) + + if all(type(a) is bool for a in arg): + return Argument.create(as_bools=list(arg)) + elif all(type(a) is int for a in arg): + return Argument.create(as_ints=list(arg)) + elif all(type(a) is float for a in arg): + return Argument.create(as_floats=list(arg)) + elif all(type(a) is str for a in arg): + return Argument.create(as_strings=list(arg)) + elif all(isinstance(a, torch.SymInt) for a in arg): + # This is a special branch for handling SymInt args in inductor's + # ExternalFallbackNode. + # For regular FX graph, SymInt arg should be a fx.Node with + # self.is_sym_int_arg(arg) being true + return Argument.create( + as_sym_ints=[SymIntArgument.create(as_name=str(a)) for a in arg] + ) + elif all(isinstance(a, torch.SymFloat) for a in arg): + return Argument.create( + as_sym_floats=[SymFloatArgument.create(as_name=str(a)) for a in arg] + ) + elif all(self.is_sym_int_arg(a) for a in arg): + # list of sym_ints + values = [] + for a in arg: + if isinstance(a, torch.fx.Node): + values.append(SymIntArgument.create(as_name=a.name)) + elif type(a) is int: + values.append(SymIntArgument.create(as_int=a)) + return Argument.create(as_sym_ints=values) + elif all(self.is_sym_float_arg(a) for a in arg): + # list of sym_float + values = [] + for a in arg: + if isinstance(a, torch.fx.Node): + values.append(SymFloatArgument.create(as_name=a.name)) + elif isinstance(a, float): + values.append(SymFloatArgument.create(as_float=a)) + return Argument.create(as_sym_floats=values) + elif all(self.is_sym_bool_arg(a) for a in arg): + # list of sym_bools + values = [] + for a in arg: + if isinstance(a, torch.fx.Node): + values.append(SymBoolArgument.create(as_name=a.name)) + elif isinstance(a, bool): + values.append(SymBoolArgument.create(as_bool=a)) + return Argument.create(as_sym_bools=values) + elif all(isinstance(a, torch.fx.Node) for a in arg): + # list of tensors + arguments = [] + for a in arg: + if a.op == "get_attr": + raise SerializeError( + "getattr nodes containing tensors should not appear in the graph" + ) + arguments.append(TensorArgument(name=a.name)) + return Argument.create(as_tensors=arguments) + elif all(isinstance(a, (torch.fx.Node, type(None))) for a in arg): + # list of optional tensors + def serialize_optional_tensor_args(a): + if a is None: + return OptionalTensorArgument.create(as_none=True) + elif isinstance(a, torch.fx.Node): + return OptionalTensorArgument.create( + as_tensor=TensorArgument(name=a.name) + ) + else: + raise SerializeError(f"Unsupported list/tuple argument: {a}") + + return Argument.create( + as_optional_tensors=list(map(serialize_optional_tensor_args, arg)) + ) + elif all(isinstance(a, inductor_tensor_buffers) for a in arg): + # list of inductor buffers + return Argument.create( + as_tensors=[TensorArgument(name=a.get_name()) for a in arg], + ) + elif all( + isinstance(a, (*inductor_tensor_buffers, type(None))) for a in arg + ): + # list of inductor buffers as optional tensors + def serialize_optional_tensor_args(a): + if a is None: + return OptionalTensorArgument.create(as_none=True) + elif isinstance(a, inductor_tensor_buffers): + return OptionalTensorArgument.create( + as_tensor=TensorArgument(name=a.get_name()) + ) + else: + raise SerializeError(f"Unsupported list/tuple argument: {a}") + + return Argument.create( + as_optional_tensors=list(map(serialize_optional_tensor_args, arg)) + ) + else: + raise SerializeError( + f"Unsupported list/tuple argument type: {[type(a) for a in arg]}" + ) + elif isinstance(arg, torch.dtype): + return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg]) + elif isinstance(arg, torch.device): + return Argument.create(as_device=Device(type=arg.type, index=arg.index)) + elif isinstance(arg, torch.memory_format): + return Argument.create( + as_memory_format=_TORCH_TO_SERIALIZE_MEMORY_FORMAT[arg] + ) + elif isinstance(arg, torch.layout): + return Argument.create(as_layout=_TORCH_TO_SERIALIZE_LAYOUT[arg]) + elif isinstance(arg, torch._C.ScriptObject): + if not ( + arg._has_method("__getstate__") # type: ignore[attr-defined] + and arg._has_method("__setstate__") # type: ignore[attr-defined] + ): + raise SerializeError( + f"Unable to serialize custom class {arg}. Please define " + "serialization methods via def_pickle()." + ) + # Custom objects through torchind are serializable with pickle, + # through implementing the .def_pickle function. This should result + # in the object containing a __getstate__ and __setstate__ + # serialize/deserialize function. + custom_obj_name = f"_custom_obj_{len(self.custom_objs)}" + self.custom_objs[custom_obj_name] = arg + class_fqn = arg._type().qualified_name() # type: ignore[attr-defined] + return Argument.create( + as_custom_obj=CustomObjArgument(custom_obj_name, class_fqn) + ) + elif isinstance(arg, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)): + return Argument.create(as_operator=self.serialize_operator(arg)) + else: + raise SerializeError(f"Unsupported argument type: {type(arg)} with schema arg_type {arg_type}") + + def serialize_tensor_output(self, name, meta_val) -> TensorArgument: + assert name not in self.graph_state.tensor_values + self.graph_state.tensor_values[name] = serialize_tensor_meta(meta_val) + return TensorArgument(name=name) + + def serialize_sym_int_output(self, name, meta_val) -> SymIntArgument: + assert name not in self.graph_state.sym_int_values + self.graph_state.sym_int_values[name] = serialize_sym_int(meta_val) + return SymIntArgument.create(as_name=name) + + def serialize_sym_float_output(self, name, meta_val) -> SymFloatArgument: + assert name not in self.graph_state.sym_float_values + self.graph_state.sym_float_values[name] = serialize_sym_float(meta_val) + return SymFloatArgument.create(as_name=name) + + def serialize_sym_bool_output(self, name, meta_val) -> SymIntArgument: + assert name not in self.graph_state.sym_bool_values + self.graph_state.sym_bool_values[name] = serialize_sym_bool(meta_val) + return SymBoolArgument.create(as_name=name) + + def serialize_input_spec(self, spec: ep.InputSpec) -> InputSpec: + if spec.kind == ep.InputKind.USER_INPUT: + if isinstance(spec.arg, ep.ConstantArgument): + if type(spec.arg.value) is int: + constant_spec = ConstantValue.create(as_int=spec.arg.value) + elif type(spec.arg.value) is bool: + constant_spec = ConstantValue.create(as_bool=spec.arg.value) + elif type(spec.arg.value) is str: + constant_spec = ConstantValue.create(as_string=spec.arg.value) + elif type(spec.arg.value) is float: + constant_spec = ConstantValue.create(as_float=spec.arg.value) + elif spec.arg.value is None: + constant_spec = ConstantValue.create(as_none=True) + else: + raise SerializeError(f"Unhandled constant input {spec.arg.value} to serialize") + return InputSpec.create( + constant_input=InputToConstantInputSpec( + name=spec.arg.name, value=constant_spec + ) + ) + else: + return InputSpec.create( + user_input=UserInputSpec( + arg=self.serialize_argument_spec(spec.arg) + ) + ) + elif spec.kind == ep.InputKind.PARAMETER: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + return InputSpec.create( + parameter=InputToParameterSpec( + arg=TensorArgument(name=spec.arg.name), + parameter_name=spec.target, + ) + ) + elif spec.kind == ep.InputKind.BUFFER: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + assert spec.persistent is not None + return InputSpec.create( + buffer=InputToBufferSpec( + arg=TensorArgument(name=spec.arg.name), + buffer_name=spec.target, + persistent=spec.persistent, + ) + ) + elif spec.kind == ep.InputKind.CONSTANT_TENSOR: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + return InputSpec.create( + tensor_constant=InputToTensorConstantSpec( + arg=TensorArgument(name=spec.arg.name), + tensor_constant_name=spec.target, + ) + ) + elif spec.kind == ep.InputKind.CUSTOM_OBJ: + assert spec.target is not None + assert isinstance(spec.arg, ep.CustomObjArgument) + return InputSpec.create( + custom_obj=InputToCustomObjSpec( + arg=CustomObjArgument( + name=spec.arg.name, class_fqn=spec.arg.class_fqn + ), + custom_obj_name=spec.target, + ) + ) + elif spec.kind == ep.InputKind.TOKEN: + assert isinstance(spec.arg, ep.TokenArgument) + return InputSpec.create( + token=InputTokenSpec( + arg=TokenArgument(name=spec.arg.name), + ) + ) + else: + raise AssertionError(f"Unknown argument kind: {spec}") + + def serialize_output_spec(self, spec: ep.OutputSpec) -> OutputSpec: + if spec.kind == ep.OutputKind.USER_OUTPUT: + return OutputSpec.create( + user_output=UserOutputSpec(arg=self.serialize_argument_spec(spec.arg)) + ) + elif spec.kind == ep.OutputKind.LOSS_OUTPUT: + assert isinstance(spec.arg, ep.TensorArgument) + return OutputSpec.create( + loss_output=LossOutputSpec(arg=TensorArgument(name=spec.arg.name)) + ) + elif spec.kind == ep.OutputKind.BUFFER_MUTATION: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + return OutputSpec.create( + buffer_mutation=BufferMutationSpec( + arg=TensorArgument(name=spec.arg.name), + buffer_name=spec.target, + ) + ) + elif spec.kind == ep.OutputKind.GRADIENT_TO_PARAMETER: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + return OutputSpec.create( + gradient_to_parameter=GradientToParameterSpec( + arg=TensorArgument(name=spec.arg.name), + parameter_name=spec.target, + ) + ) + elif spec.kind == ep.OutputKind.GRADIENT_TO_USER_INPUT: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + return OutputSpec.create( + gradient_to_user_input=GradientToUserInputSpec( + arg=TensorArgument(name=spec.arg.name), + user_input_name=spec.target, + ) + ) + elif spec.kind == ep.OutputKind.USER_INPUT_MUTATION: + assert spec.target is not None + assert isinstance(spec.arg, ep.TensorArgument) + return OutputSpec.create( + user_input_mutation=UserInputMutationSpec( + arg=TensorArgument(name=spec.arg.name), + user_input_name=spec.target, + ) + ) + elif spec.kind == ep.OutputKind.TOKEN: + assert isinstance(spec.arg, ep.TokenArgument) + return OutputSpec.create( + token=OutputTokenSpec( + arg=TokenArgument(name=spec.arg.name), + ) + ) + else: + raise AssertionError(f"Unknown argument kind: {spec}") + + def serialize_signature(self, sig: ep.ExportGraphSignature) -> GraphSignature: + return GraphSignature( + input_specs=[self.serialize_input_spec(s) for s in sig.input_specs], + output_specs=[self.serialize_output_spec(s) for s in sig.output_specs], + ) + + def serialize_argument_spec(self, x: ep.ArgumentSpec) -> Argument: + if isinstance(x, ep.TensorArgument): + return Argument.create(as_tensor=TensorArgument(name=x.name)) + elif isinstance(x, ep.SymIntArgument): + return Argument.create(as_sym_int=SymIntArgument.create(as_name=x.name)) + elif isinstance(x, ep.SymFloatArgument): + return Argument.create(as_sym_Float=SymFloatArgument.create(as_name=x.name)) + elif isinstance(x, ep.ConstantArgument): + return self.serialize_input(x.value) + elif isinstance(x, ep.CustomObjArgument): + return Argument.create( + as_custom_obj=CustomObjArgument(name=x.name, class_fqn=x.class_fqn) + ) + else: + raise AssertionError("TODO") + + def serialize_module_call_signature( + self, module_call_signature: ep.ModuleCallSignature + ) -> ModuleCallSignature: + return ModuleCallSignature( + inputs=[ + self.serialize_argument_spec(x) for x in module_call_signature.inputs + ], + outputs=[ + self.serialize_argument_spec(x) for x in module_call_signature.outputs + ], + in_spec=treespec_dumps(module_call_signature.in_spec, TREESPEC_VERSION), + out_spec=treespec_dumps(module_call_signature.out_spec, TREESPEC_VERSION), + forward_arg_names=names if (names := module_call_signature.forward_arg_names) else None + ) + + def serialize_module_call_graph( + self, module_call_graph: List[ep.ModuleCallEntry] + ) -> List[ModuleCallEntry]: + return [ + ModuleCallEntry( + fqn=entry.fqn, + signature=( + self.serialize_module_call_signature(entry.signature) + if entry.signature + else None + ), + ) + for entry in module_call_graph + ] + + def serialize_outputs(self, node: torch.fx.Node) -> List[Argument]: + """For a given node, return the dataclass representing its output values. + + [NOTE: Multiple outputs] We handle aggregates differently than FX. For + FX, it looks like: + + x = call_function("multiple_return", ...) + element0 = call_function(getitem, x, 0) + foo = call_function("use_output", element0) + + We do not want the intermediate `getitem` call, so our serialized thing looks like: + + element0, element1, element2 = call_function("multiple_return", ...) + foo = call_function("use_output", element0) + + We want names to be consistent across these two schemes, so that we can + mostly reuse the names coming from FX. This function computes a mapping from + the FX representation to our representation, preserving the names. + """ + assert node.op == "call_function" and isinstance(node.target, (torch._ops.OpOverload, *_registered_extension_types())) + + schema = _get_schema_from_target(node.target) + returns = schema.returns + + if len(returns) == 0: + return [] + + meta_val = node.meta["val"] + + # Check single value return + if _is_single_tensor_list_return(node.target): + # e.g "-> Tensor[]" + tensor_args = [] + for idx, meta in enumerate(meta_val): + name = self._output_node_name_at_index(node, idx) + tensor_args.append(self.serialize_tensor_output(name, meta)) + return [Argument.create(as_tensors=tensor_args)] + elif len(returns) == 1: + return [self.serialize_output(node.name, meta_val)] + + # There are a two possibilities at this point: + # - This operator returns a tuple of Tensors, e.g. "-> (Tensor, Tensor)" + # - This operator returns a tuple of mixed of Tensor and Tensors, e.g. "-> (Tensor, Tensor[])" + # + # Either way, start by gathering a list of TensorArguments with the correct names. + # For consistent naming with FX, consult the downstream `getitem` node and + # make sure our outputs have the same name. + + output_arguments = [] + for idx, (meta, return_schema) in enumerate(zip(meta_val, returns)): + if meta is None: + assert isinstance( + return_schema.real_type, (torch.OptionalType, torch.TensorType) + ) + # When the return type is annoated as Tensor type, the op can also return an + # undefined Tensor which will be implicitly converted to None in Python. + output_arguments.append(Argument.create(as_none=True)) + elif isinstance(meta, FakeTensor): + assert isinstance(return_schema.real_type, (torch.OptionalType, torch.TensorType)) + name = self._output_node_name_at_index(node, idx) + output_arguments.append(self.serialize_output(name, meta)) + elif isinstance(meta, list): + # for List[Tensor] return type + assert isinstance( + return_schema.real_type, torch.ListType + ) and isinstance( + return_schema.real_type.getElementType(), torch.TensorType + ) + user_node = self._output_node_at_index(node, idx) + assert user_node is not None + + args = [] + for i, m in enumerate(meta): + if m is None: + continue + sub_user_node_name = self._output_node_name_at_index(user_node, i) + args.append(self.serialize_tensor_output(sub_user_node_name, m)) + output_arguments.append(Argument.create(as_tensors=args)) + elif isinstance(meta, (int, SymInt, float, SymFloat)): + user_node_name = self._output_node_name_at_index(node, idx) + output_arguments.append(self.serialize_output(user_node_name, meta)) + else: + raise ValueError( + f"Unhandled output type {type(meta)} from node {node.format_node()}" + ) + + return output_arguments + + def serialize_hoo_outputs(self, node: torch.fx.Node) -> List[Argument]: + """ + For serializing HOO outputs since HOOs do not have a schema. + """ + meta_val = node.meta["val"] + + if isinstance(meta_val, tuple): + # Note: Since we don't have a schema, we just serialize all tuple + # outputs to be a list of values. Even if the output is supposed to + # be a tensor list (Tensor[]), we will serialize it to be a list of + # tensors (Tensor, Tensor, Tensor). An exception is that if there's + # a singleton tensor, we will serialize this to be a singleton + # tensor list so that the deserializer knows to insert getitem nodes. + + if len(meta_val) == 1: + assert isinstance(meta_val[0], torch.Tensor) + name = self._output_node_name_at_index(node, 0) + return [Argument.create(as_tensors=[self.serialize_tensor_output(name, meta_val[0])])] + + outputs = [] + for i, element_meta_val in enumerate(meta_val): + user_node = self._output_node_at_index(node, i) + if isinstance(element_meta_val, list): + # e.g "-> Tensor[]" + assert user_node is not None + + tensors = [] + for j, m in enumerate(element_meta_val): + if not isinstance(m, torch.Tensor): + raise SerializeError(f"Serialize list output with type {type(m)} nyi") + + name = self._output_node_name_at_index(user_node, j) + tensors.append(self.serialize_tensor_output(name, m)) + outputs.append(Argument.create(as_tensors=tensors)) + + else: + name = ( + user_node.name + if user_node is not None + else f"{node.name}_unused_{i}" + ) + + outputs.append(self.serialize_output(name, element_meta_val)) + + return outputs + else: + return [self.serialize_output(node.name, meta_val)] + + def serialize_output(self, name: str, meta_val: Any) -> Argument: + # Check single value return + if meta_val is None: + return Argument.create(as_none=True) + if isinstance(meta_val, torch.Tensor): + # e.g "-> Tensor" + return Argument.create( + as_tensor=self.serialize_tensor_output(name, meta_val) + ) + elif isinstance(meta_val, (int, torch.SymInt)): + # e.g "-> SymInt" + return Argument.create( + as_sym_int=self.serialize_sym_int_output(name, meta_val) + ) + elif isinstance(meta_val, (int, torch.SymFloat)): + # e.g "-> SymFloat" + return Argument.create( + as_sym_float=self.serialize_sym_float_output(name, meta_val) + ) + elif isinstance(meta_val, torch.SymBool): + # e.g "-> SymBool" + return Argument.create( + as_sym_bool=self.serialize_sym_bool_output(name, meta_val) + ) + + # list outputs should've been handled earlier + raise SerializeError(f"Unable to serialize output {meta_val}") + + def _handle_getitem_users(self, node: torch.fx.Node) -> List[TensorArgument]: + meta_val = node.meta["val"] + + idx_to_name = {} + for user in node.users: + assert ( + user.target is operator.getitem + ), f"User node {user} of {node} is incorrect" + idx_to_name[user.args[1]] = user.name + + for idx, _ in enumerate(meta_val): + # FX does not emit a getitem node for any outputs that are unused. + # However, we need a name for them so that the number of outputs will + # correctly match the schema. Just assign a dummy name. + if idx not in idx_to_name: + idx_to_name[idx] = f"{node.name}_unused_{idx}" + + arg_list = [] + for i, element_meta_val in enumerate(meta_val): + arg_list.append( + self.serialize_tensor_output(idx_to_name[i], element_meta_val) + ) + + return arg_list + + def serialize_graph(self, graph_module: torch.fx.GraphModule) -> Graph: + assert isinstance(graph_module, torch.fx.GraphModule) + for node in graph_module.graph.nodes: + try: + getattr(self, f"handle_{node.op}")(node) + except Exception as e: + raise SerializeError( + f"Failed serializing node {node} in graph: {node.format_node()}\n Original exception {traceback.format_exc()}" + ) from e + + return Graph( + inputs=self.graph_state.inputs, + nodes=self.graph_state.nodes, + tensor_values=self.graph_state.tensor_values, + sym_int_values=self.graph_state.sym_int_values, + sym_float_values=self.graph_state.sym_float_values, + sym_bool_values=self.graph_state.sym_bool_values, + custom_obj_values=self.graph_state.custom_obj_values, + outputs=self.graph_state.outputs, + is_single_tensor_return=self.graph_state.is_single_tensor_return, + ) + + def serialize_graph_module_metadata(self, meta: Dict[str, Any]): + ret = {} + if custom := meta.get("custom"): + try: + ret["custom"] = json.dumps(custom) + except Exception as e: + raise SerializeError( + f"Failed to serialize custom metadata for graph with error {e}" + ) from e + + return ret + + def serialize(self, graph_module: torch.fx.GraphModule) -> GraphModule: + graph = self.serialize_graph(graph_module) + + return GraphModule( + graph=graph, + signature=self.serialize_signature(self.graph_signature), + module_call_graph=self.serialize_module_call_graph(self.module_call_graph), + metadata=self.serialize_graph_module_metadata(graph_module.meta) + ) + + +@final +class ExportedProgramSerializer(metaclass=Final): + def __init__(self, opset_version: Optional[Dict[str, int]] = None): + self.opset_version: Dict[str, int] = {} + if opset_version: + self.opset_version.update(opset_version) + if "aten" not in self.opset_version: + self.opset_version["aten"] = torch._C._get_max_operator_version() + + def serialize(self, exported_program: ep.ExportedProgram) -> _SerializedProgram: + """ + Args: + exported_program: Exported Program to serialize + """ + exported_program.validate() + + gm_serializer = GraphModuleSerializer( + exported_program.graph_signature, exported_program.module_call_graph + ) + serialized_graph_module = gm_serializer.serialize(exported_program.graph_module) + serialized_range_constraints = serialize_range_constraints( + exported_program.range_constraints + ) + + # TODO: Directly serialize exported_program.constants once + # CustomClassHolders get stored in the ExportedProgram rather than in + # the graph + constants = {} + for n, c in gm_serializer.custom_objs.items(): + constants[n] = c + for n, t in exported_program.constants.items(): + assert n not in constants + constants[n] = t + + serialized_ep = ExportedProgram( + graph_module=serialized_graph_module, + opset_version=self.opset_version, + range_constraints=serialized_range_constraints, + schema_version=SchemaVersion( + major=SCHEMA_VERSION[0], + minor=SCHEMA_VERSION[1], + ), + verifiers=[v.dialect for v in exported_program.verifiers], + torch_version=torch.__version__, + ) + + # Test canonical form is well defined. + canonicalize(serialized_ep) + + # Proxy cannot be dumped, so we remove them. + new_state_dict = remove_proxy_from_state_dict( + exported_program.state_dict, in_place=False + ) + return _SerializedProgram( + serialized_ep, + serialize_torch_artifact(new_state_dict), + serialize_torch_artifact(constants), + serialize_torch_artifact(exported_program.example_inputs), + ) + + +@final +class GraphModuleDeserializer(metaclass=Final): + @dataclasses.dataclass + class Result: + graph_module: torch.fx.GraphModule + signature: ep.ExportGraphSignature + module_call_graph: List[ep.ModuleCallEntry] + names_to_symbols: Dict[str, sympy.Symbol] + state_dict: Dict[str, Union[torch.Tensor, torch.nn.Parameter]] + constants: Dict[str, Union[torch.Tensor, FakeScriptObject, torch.ScriptObject]] + example_inputs: Optional[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]]] + + def __init__(self) -> None: + self.serialized_name_to_node: Dict[str, torch.fx.Node] = {} + self.serialized_name_to_meta: Dict[str, MetaType] = {} + self.graph = torch.fx.Graph() + self.module = torch.nn.Module() + + @contextmanager + def save_graph_module(self) -> Iterator[None]: + saved = ( + self.graph, + self.module, + self.serialized_name_to_node, + self.serialized_name_to_meta, + ) + self.graph = torch.fx.Graph() + self.module = torch.nn.Module() + self.serialized_name_to_node = {} + self.serialized_name_to_meta = {} + try: + yield + finally: + ( + self.graph, + self.module, + self.serialized_name_to_node, + self.serialized_name_to_meta, + ) = saved + + def deserialize_extension_operator(self, serialized_target: str): + namespace, op_name = serialized_target.split(":") + namespace = namespace[1:] # starting with # + handler = _deserialization_registry[namespace] + return handler.from_op_name(op_name) + + 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"): + module = torch # type: ignore[misc] + serialized_target_names = serialized_target.split(".")[1:] + elif serialized_target.startswith("#"): + return self.deserialize_extension_operator(serialized_target) + else: # TODO(zhxchen17) Don't catch all here. + return serialized_target + + target = module + for name in serialized_target_names: + if not hasattr(target, name): + return serialized_target + else: + target = getattr(target, name) + return target + + def _parse_sym_expr(self, expr_str: str, hint: Optional[Union[int, bool, float]] = None) -> sympy.Expr: + """ + Parses and does bottom-up processing of sympy.Expr nodes, + populating ShapeEnv & caching symbols as needed. + """ + def _process_sym_expr(sym: sympy.Expr, hint: Optional[Union[int, bool, float]] = None) -> sympy.Expr: + if sym.is_Integer or sym.is_Float or sym.is_Boolean: # base case + return sym + else: # recursive case + # important to use str(expr) and not _print_sympy(), + # str(expr) is key for self.symbol_name_to_range + expr_str = str(sym) + for arg in sym.args: + self._parse_sym_expr(arg) + # symbol caching + if expr_str in self.symbol_name_to_symbol: + sym = self.symbol_name_to_symbol[expr_str] + else: + self.symbol_name_to_symbol[expr_str] = sym + # hints + if ( + hint is not None + and sym not in self.shape_env.var_to_val + ): + self.shape_env.add_var_to_val(sym, hint) # type: ignore[arg-type] + # ValueRanges + if vr := self.symbol_name_to_range.get(expr_str): + self.shape_env.constrain_symbol_range( + sym, + compiler_min=vr.lower, # type: ignore[arg-type] + compiler_max=vr.upper, # type: ignore[arg-type] + ) + return sym + + expr = sympy.sympify( + expr_str, + locals={**self.sympy_functions, **self.symbol_name_to_symbol}, + ) + return _process_sym_expr(expr, hint) + + def deserialize_sym_int(self, s: SymInt) -> Union[int, torch.SymInt]: + val = s.value + if s.type == "as_expr": + if val.hint is None: + hint = None + else: + assert val.hint.type == "as_int" + hint = val.hint.value + + sym = self._parse_sym_expr(val.expr_str, hint) + return self.shape_env.create_symintnode(sym, hint=hint) + elif s.type == "as_int": + assert type(val) is int + return val + else: + raise SerializeError( + f"SymInt has invalid field type {s.type} with value {s.value}" + ) + + def deserialize_sym_float(self, s: SymFloat) -> Union[float, torch.SymFloat]: + val = s.value + if s.type == "as_expr": + hint = val.hint.as_float if val.hint else None + sym = self._parse_sym_expr(val.expr_str, hint) + return self.shape_env.create_symfloatnode(sym, hint=hint) + elif s.type == "as_float": + assert isinstance(val, float) + return val + else: + raise SerializeError( + f"SymFloat has invalid field type {s.type} with value {s.value}" + ) + + def deserialize_sym_bool(self, s: SymBool) -> Union[bool, torch.SymBool]: + val = s.value + if s.type == "as_expr": + expr = self._parse_sym_expr(val.expr_str) + return self.shape_env.create_symboolnode(expr) + elif s.type == "as_bool": + assert isinstance(val, bool) + return val + else: + raise SerializeError( + f"SymBool has invalid field type {s.type} with value {s.value}" + ) + + def deserialize_tensor_meta( + self, + tensor_meta: TensorMeta, + ) -> FakeTensor: + with self.fake_tensor_mode: + return cast( + FakeTensor, + torch.empty_strided( + tuple(self.deserialize_sym_int(val) for val in tensor_meta.sizes), # type: ignore[misc] + tuple(self.deserialize_sym_int(val) for val in tensor_meta.strides), # type: ignore[misc] + device=deserialize_device(tensor_meta.device), + dtype=_SERIALIZE_TO_TORCH_DTYPE[tensor_meta.dtype], + ), + ) + + def deserialize_script_obj_meta( + self, script_obj_meta: CustomObjArgument + ) -> ep.CustomObjArgument: + return ep.CustomObjArgument( + name=script_obj_meta.name, + class_fqn=script_obj_meta.class_fqn, + ) + + def deserialize_graph_output(self, output) -> Optional[Union[torch.fx.Node, int]]: + if output.type == "as_tensor": + return self.serialized_name_to_node[output.as_tensor.name] + elif output.type == "as_sym_int": + return self.serialized_name_to_node[output.as_sym_int.as_name] + elif output.type == "as_sym_bool": + return self.serialized_name_to_node[output.as_sym_bool.as_name] + elif output.type == "as_sym_float": + return self.serialized_name_to_node[output.as_sym_float.as_name] + elif output.type == "as_int": + return output.as_int + elif output.type == "as_float": + return output.as_float + elif output.type == "as_none": + return None + else: + raise SerializeError(f"Unable to deserialize output node {output}") + + def deserialize_graph(self, serialized_graph: Graph) -> torch.fx.Graph: + # Handle the tensor metas. + for name, tensor_value in serialized_graph.tensor_values.items(): + meta_val = self.deserialize_tensor_meta(tensor_value) + self.serialized_name_to_meta[name] = meta_val + + for name, sym_int_value in serialized_graph.sym_int_values.items(): + self.serialized_name_to_meta[name] = self.deserialize_sym_int(sym_int_value) + + for name, sym_int_value in serialized_graph.sym_float_values.items(): + self.serialized_name_to_meta[name] = self.deserialize_sym_float(sym_int_value) + + for name, sym_bool_value in serialized_graph.sym_bool_values.items(): + self.serialized_name_to_meta[name] = self.deserialize_sym_bool( + sym_bool_value + ) + + for name, script_obj_meta in serialized_graph.custom_obj_values.items(): + self.serialized_name_to_meta[name] = self.deserialize_script_obj_meta( + script_obj_meta + ) + + # Inputs: convert to placeholder nodes in FX. + for i, input_ in enumerate(serialized_graph.inputs): + if input_.type in ("as_tensor", "as_sym_int", "as_sym_float", "as_custom_obj"): + node_name = input_.value.name + placeholder_node = self.graph.placeholder(node_name) + # FX might declare a name illegal (e.g. some nn.Modules use "input" as forward() arguments) + # we will overwrite it + placeholder_node.name = node_name + self.sync_fx_node(node_name, placeholder_node) + elif input_.type in ( + "as_int", + "as_float", + "as_bool", + "as_none", + "as_string", + ): + node_name = self.signature.input_specs[i].arg.name + placeholder_node = self.graph.placeholder(node_name) + placeholder_node.meta["val"] = self.deserialize_input(input_) + else: + raise SerializeError(f"Invalid input type {input_}") + + # Nodes: convert to call_function nodes. + for serialized_node in serialized_graph.nodes: + try: + target = self.deserialize_operator(serialized_node.target) + self.deserialize_node(serialized_node, target) + + except Exception as e: + raise SerializeError( + f"Failed deserializing node {serialized_node}\n Original exception {traceback.format_exc()}" + ) from e + + # Outputs: convert to a single `output` node. + outputs = [self.deserialize_graph_output(output) for output in serialized_graph.outputs] + + if serialized_graph.is_single_tensor_return: + assert len(outputs) == 1 + outputs = outputs[0] # type: ignore[assignment] + else: + outputs = tuple(outputs) # type: ignore[assignment] + + output_node = self.graph.output(outputs) + + if serialized_graph.is_single_tensor_return: + output_node.meta["val"] = output_node.args[0].meta["val"] + else: + output_node.meta["val"] = tuple( + arg.meta["val"] if isinstance(arg, torch.fx.Node) else arg + for arg in output_node.args[0] + ) + + return self.graph + + def deserialize_node(self, serialized_node: Node, target: Callable) -> None: + if ( + target in _SYM_BOOL_OPS + or target in _SYM_INT_OPS + or target in _SYM_FLOAT_OPS + or target == torch.ops.aten.item.default # this can produce either SymInt or SymBool + ): + name = serialized_node.outputs[0].value.as_name + args = self.deserialize_sym_op_inputs(serialized_node.inputs) + + fx_node = self.graph.create_node("call_function", target, args, {}, name) + self.deserialize_sym_op_outputs(serialized_node, fx_node) + + elif isinstance(target, torch._ops.HigherOrderOperator): + args, kwargs = self.deserialize_hoo_inputs(serialized_node.inputs) + metadata = self.deserialize_metadata(serialized_node.metadata) + for x in (*args, *kwargs.values()): + if isinstance(x, torch.fx.Node) and x.op == "get_attr": + # this means that we have deserialized a graph argument, but + # unfortunately the schema for it does not include metadata; + # so we reuse the metadata of the HOP call for such arguments + x.meta.update(metadata) + # If HOP returns a single tensor, name the + # newly-created node after it. This ensures that these tensor values + # have names that are consistent with serialized. + # + # HOPs don't have schema yet, just check the output lengths and as_tensor attribute + name = ( + serialized_node.outputs[0].as_tensor.name + if len(serialized_node.outputs) == 1 + and hasattr(serialized_node.outputs[0], "as_tensor") + else None + ) + fx_node = self.graph.create_node( + "call_function", target, args, kwargs, name + ) + self.deserialize_outputs(serialized_node, fx_node) + fx_node.meta.update(metadata) + + elif isinstance(target, (torch._ops.OpOverload, *_registered_extension_types())): + # For convenience: if this node returns a single tensor, name the + # newly-created node after it. This ensures that these tensor values + # have names that are consistent with serialized. + name = ( + serialized_node.outputs[0].as_tensor.name + if _is_single_tensor_return(target) + else None # FX will generate a name for us. + ) + args, kwargs = self.deserialize_inputs(target, serialized_node) + fx_node = self.graph.create_node( + "call_function", target, args, kwargs, name + ) + self.deserialize_outputs(serialized_node, fx_node) + else: + raise SerializeError( + f"Unsupported target type for node {serialized_node}: {type(target)}" + ) + + fx_node.meta.update(self.deserialize_metadata(serialized_node.metadata)) + if fx_node.op not in ["placeholder", "output"] and "nn_module_stack" not in fx_node.meta: + fx_node.meta["nn_module_stack"] = {} # serialization throws away empty dicts + + def deserialize_input_spec(self, i: InputSpec) -> ep.InputSpec: + if i.type == "user_input": + return ep.InputSpec( + kind=ep.InputKind.USER_INPUT, + arg=self.deserialize_argument_spec(i.user_input.arg), + target=None, + ) + elif i.type == "parameter": + return ep.InputSpec( + kind=ep.InputKind.PARAMETER, + arg=ep.TensorArgument(name=i.parameter.arg.name), + target=i.parameter.parameter_name, + ) + elif i.type == "buffer": + return ep.InputSpec( + kind=ep.InputKind.BUFFER, + arg=ep.TensorArgument(name=i.buffer.arg.name), + target=i.buffer.buffer_name, + persistent=i.buffer.persistent, + ) + elif i.type == "tensor_constant": + return ep.InputSpec( + kind=ep.InputKind.CONSTANT_TENSOR, + arg=ep.TensorArgument(name=i.tensor_constant.arg.name), + target=i.tensor_constant.tensor_constant_name, + ) + elif i.type == "custom_obj": + return ep.InputSpec( + kind=ep.InputKind.CUSTOM_OBJ, + arg=ep.CustomObjArgument( + name=i.custom_obj.arg.name, class_fqn=i.custom_obj.arg.class_fqn + ), + target=i.custom_obj.custom_obj_name, + ) + elif i.type == "token": + return ep.InputSpec( + kind=ep.InputKind.TOKEN, + arg=ep.TokenArgument(name=i.token.arg.name), + target=None + ) + elif i.type == "constant_input": + return ep.InputSpec( + kind=ep.InputKind.USER_INPUT, + arg=ep.ConstantArgument( + name=i.constant_input.name, + value=self.deserialize_constant_input(i.constant_input.value) + ), + target=None, + ) + else: + raise AssertionError(f"Unknown input spec {i}") + + def deserialize_output_spec(self, o: OutputSpec) -> ep.OutputSpec: + if o.type == "user_output": + return ep.OutputSpec( + kind=ep.OutputKind.USER_OUTPUT, + arg=self.deserialize_argument_spec(o.user_output.arg), + target=None, + ) + elif o.type == "loss_output": + return ep.OutputSpec( + kind=ep.OutputKind.LOSS_OUTPUT, + arg=ep.TensorArgument(name=o.loss_output.arg.name), + target=None, + ) + elif o.type == "buffer_mutation": + return ep.OutputSpec( + kind=ep.OutputKind.BUFFER_MUTATION, + arg=ep.TensorArgument(name=o.buffer_mutation.arg.name), + target=o.buffer_mutation.buffer_name, + ) + elif o.type == "gradient_to_parameter": + return ep.OutputSpec( + kind=ep.OutputKind.GRADIENT_TO_PARAMETER, + arg=ep.TensorArgument(name=o.gradient_to_parameter.arg.name), + target=o.gradient_to_parameter.parameter_name, + ) + elif o.type == "gradient_to_user_input": + return ep.OutputSpec( + kind=ep.OutputKind.GRADIENT_TO_USER_INPUT, + arg=ep.TensorArgument(name=o.gradient_to_user_input.arg.name), + target=o.gradient_to_user_input.user_input_name, + ) + elif o.type == "user_input_mutation": + return ep.OutputSpec( + kind=ep.OutputKind.USER_INPUT_MUTATION, + arg=ep.TensorArgument(name=o.user_input_mutation.arg.name), + target=o.user_input_mutation.user_input_name, + ) + elif o.type == "token": + return ep.OutputSpec( + kind=ep.OutputKind.TOKEN, + arg=ep.TokenArgument(name=o.token.arg.name), + target=None + ) + else: + raise AssertionError(f"Unknown output spec {o}") + + def deserialize_signature(self, sig: GraphSignature) -> ep.ExportGraphSignature: + return ep.ExportGraphSignature( + input_specs=[self.deserialize_input_spec(i) for i in sig.input_specs], + output_specs=[self.deserialize_output_spec(o) for o in sig.output_specs], + ) + + def deserialize( + self, + serialized_graph_module: GraphModule, + serialized_state_dict: Union[Dict[str, torch.Tensor], bytes], + constants: Union[Dict[str, Any], bytes], + example_inputs: Optional[Union[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]], bytes]] = None, + symbol_name_to_range: Optional[Dict[str, symbolic_shapes.ValueRanges]] = None, + ) -> Result: + global _CURRENT_DESERIALIZER + assert _CURRENT_DESERIALIZER is None + _CURRENT_DESERIALIZER = self + try: + self.shape_env = symbolic_shapes.ShapeEnv(assume_static_by_default=True) + self.fake_tensor_mode = FakeTensorMode( + allow_fallback_kernels=False, + allow_non_fake_inputs=True, + shape_env=self.shape_env, + ) + self.sympy_functions = { + # all torch.utils._sympy.functions should go here + # TODO(avik): find a better way to keep this collection in sync; + # e.g.., `exec('from torch.utils._sympy.functions import *', ...)` + # would work as long as the public API of that module is complete + "FloorDiv": torch.utils._sympy.functions.FloorDiv, + "ModularIndexing": torch.utils._sympy.functions.ModularIndexing, + "Where": torch.utils._sympy.functions.Where, + "PythonMod": torch.utils._sympy.functions.PythonMod, + "Mod": torch.utils._sympy.functions.Mod, + "CleanDiv": torch.utils._sympy.functions.CleanDiv, + "CeilToInt": torch.utils._sympy.functions.CeilToInt, + "FloorToInt": torch.utils._sympy.functions.FloorToInt, + "CeilDiv": torch.utils._sympy.functions.CeilDiv, + "LShift": torch.utils._sympy.functions.LShift, + "RShift": torch.utils._sympy.functions.RShift, + "PowByNatural": torch.utils._sympy.functions.PowByNatural, + "FloatPow": torch.utils._sympy.functions.FloatPow, + "FloatTrueDiv": torch.utils._sympy.functions.FloatTrueDiv, + "IntTrueDiv": torch.utils._sympy.functions.IntTrueDiv, + "IsNonOverlappingAndDenseIndicator": torch.utils._sympy.functions.IsNonOverlappingAndDenseIndicator, + "TruncToFloat": torch.utils._sympy.functions.TruncToFloat, + "TruncToInt": torch.utils._sympy.functions.TruncToInt, + "RoundToInt": torch.utils._sympy.functions.RoundToInt, + "RoundDecimal": torch.utils._sympy.functions.RoundDecimal, + "ToFloat": torch.utils._sympy.functions.ToFloat, + "Identity": torch.utils._sympy.functions.Identity, + } + self.symbol_name_to_symbol: Dict[str, sympy.Symbol] = {} + self.constants = deserialize_torch_artifact(constants) + self.signature = self.deserialize_signature(serialized_graph_module.signature) + + # deserialization does analysis with checks on 0/1, so we create fake range constraints and + # restore the original range constraints afterwards + self.symbol_name_to_range = {} + if symbol_name_to_range: + for k, vr in symbol_name_to_range.items(): + lower = vr.lower + if vr.upper >= 2: # max is >= 2, not sym bool range + lower = max(2, lower) + self.symbol_name_to_range[k] = symbolic_shapes.ValueRanges(_int_to_sympy_int(lower, -int_oo), vr.upper) + + if example_inputs is not None and len(example_inputs) > 0: + self.example_inputs = deserialize_torch_artifact(example_inputs) + else: + self.example_inputs = None + self.deserialize_graph(serialized_graph_module.graph) + + module_call_graph = self.deserialize_module_call_graph( + serialized_graph_module.module_call_graph + ) + graph_module = ep._create_graph_module_for_export( + self.module, self.graph + ) + meta = {} + if custom := serialized_graph_module.metadata.get("custom"): + meta["custom"] = json.loads(custom) + graph_module.meta = meta + return GraphModuleDeserializer.Result( + graph_module=graph_module, + signature=self.signature, + module_call_graph=module_call_graph, + names_to_symbols=self.symbol_name_to_symbol, + state_dict=deserialize_torch_artifact(serialized_state_dict), + constants=self.constants, + example_inputs=self.example_inputs, + ) + finally: + _CURRENT_DESERIALIZER = None + + def sync_fx_node(self, name: str, fx_node: torch.fx.Node): + if name in self.serialized_name_to_node: + raise SerializeError(f"Node {name} has already been deserialized before.") + # overwrite name + fx_node.name = name + self.serialized_name_to_node[name] = fx_node + assert "val" not in fx_node.meta + fx_node.meta["val"] = self.serialized_name_to_meta[name] + + def deserialize_sym_op_inputs(self, inputs): + return tuple(self.deserialize_input(input.arg) for input in inputs) + + def deserialize_inputs(self, target, serialized_node: Node): + schema_args = _get_schema_from_target(target).arguments + actual_args = { + input.name: self.deserialize_input(input.arg) + for input in serialized_node.inputs + } + args = [] + kwargs: OrderedDict[str, Any] = OrderedDict() + for schema_arg in schema_args: + is_positional = ( + not schema_arg.has_default_value() and not schema_arg.kwarg_only + ) + if is_positional: + args.append(actual_args[schema_arg.name]) + elif keyword.iskeyword(schema_arg.name): + assert not schema_arg.kwarg_only + if len(kwargs) > 0: + kwargs = OrderedDict() + args.extend(list(kwargs.values())) + args.append(actual_args[schema_arg.name]) + else: + if schema_arg.name in actual_args: + kwargs[schema_arg.name] = actual_args[schema_arg.name] + return tuple(args), kwargs + + def deserialize_hoo_inputs(self, inputs: List[NamedArgument]): + """ + For deserializing HOO inputs since HOOs do not have a schema. + """ + args = [] + kwargs = {} + for input_ in inputs: + if input_.name != "": + kwargs[input_.name] = self.deserialize_input(input_.arg) + else: + args.append(self.deserialize_input(input_.arg)) + return (tuple(args), kwargs) + + def deserialize_input(self, inp: Argument) -> Any: + value = inp.value + typ_ = inp.type + if typ_ == "as_none": + # None should converted as None, but is encoded as bool in serialized + # Convert serialized object to torch equivalent + return None + elif typ_ == "as_tensor": + return self.serialized_name_to_node[inp.as_tensor.name] + elif typ_ == "as_scalar_type": + return _SERIALIZE_TO_TORCH_DTYPE[inp.as_scalar_type] + elif typ_ == "as_memory_format": + return _SERIALIZE_TO_TORCH_MEMORY_FORMAT[inp.as_memory_format] + elif typ_ == "as_layout": + return _SERIALIZE_TO_TORCH_LAYOUT[inp.as_layout] + elif typ_ == "as_graph": + assert isinstance(value, GraphArgument) + with self.save_graph_module(): + self.deserialize_graph(value.graph) + submodule = ep._create_graph_module_for_export(self.module, self.graph) + self.module.register_module(value.name, submodule) + return self.graph.create_node( + "get_attr", + value.name, + name=value.name, + ) + elif typ_ == "as_device": + return deserialize_device(inp.as_device) + elif typ_ == "as_int": + return inp.as_int + elif typ_ == "as_float": + return inp.as_float + elif typ_ == "as_bool": + return inp.as_bool + elif typ_ == "as_string": + return inp.as_string + elif typ_ == "as_sym_int": + return self.deserialize_sym_argument(inp.as_sym_int) + elif typ_ == "as_sym_float": + return self.deserialize_sym_argument(inp.as_sym_float) + elif typ_ == "as_sym_bool": + return self.deserialize_sym_argument(inp.as_sym_bool) + elif isinstance(value, list): + if len(value) == 0: + return [] + elif typ_ == "as_tensors": + result = [self.serialized_name_to_node[arg.name] for arg in value] + return result + elif typ_ in ("as_ints", "as_floats", "as_bools", "as_strings"): + # convert from serialized.python.types.List to python list + return list(value) + elif typ_ in ("as_sym_ints", "as_sym_bools", "as_sym_floats"): + return [self.deserialize_sym_argument(arg) for arg in value] + elif typ_ == "as_optional_tensors": + + def deserialize_optional_tensor_args(a): + if a.type == "as_none": + return None + elif a.type == "as_tensor": + return self.serialized_name_to_node[a.value.name] + else: + raise SerializeError(f"Unhandled argument {inp}") + + return list(map(deserialize_optional_tensor_args, value)) + else: + raise SerializeError(f"Unhandled argument {inp}") + elif typ_ == "as_custom_obj": + if inp.as_custom_obj.name in self.serialized_name_to_node: + # Custom object has been lifted as an input + return self.serialized_name_to_node[inp.as_custom_obj.name] + return self.constants[inp.as_custom_obj.name] + elif typ_ == "as_operator": + return self.deserialize_operator(inp.as_operator) + else: + raise SerializeError(f"Unhandled argument {inp}") + + def deserialize_constant_input(self, inp: ConstantValue) -> Any: + if inp.type == "as_int": + return int(inp.as_int) + elif inp.type == "as_float": + return float(inp.as_float) + elif inp.type == "as_string": + return str(inp.as_string) + elif inp.type == "as_bool": + return bool(inp.as_bool) + elif inp.type == "as_none": + return None + else: + raise SerializeError(f"Unhandled constant argument {inp} to deserialize") + + def deserialize_sym_argument(self, sym_arg): + if isinstance(sym_arg, SymIntArgument): + if sym_arg.type == "as_int": + return sym_arg.as_int + elif sym_arg.type == "as_name": + return self.serialized_name_to_node[sym_arg.as_name] + elif isinstance(sym_arg, SymFloatArgument): + if sym_arg.type == "as_float": + return sym_arg.as_float + elif sym_arg.type == "as_name": + return self.serialized_name_to_node[sym_arg.as_name] + elif isinstance(sym_arg, SymBoolArgument): + if sym_arg.type == "as_bool": + return sym_arg.as_bool + elif sym_arg.type == "as_name": + return self.serialized_name_to_node[sym_arg.as_name] + raise SerializeError(f"Unknown symbolic argument type: {sym_arg}") + + def deserialize_sym_op_outputs(self, serialized_node: Node, fx_node: torch.fx.Node): + self.sync_fx_node(serialized_node.outputs[0].value.as_name, fx_node) + + def deserialize_outputs(self, serialized_node: Node, fx_node: torch.fx.Node): + # Check single value return + if len(serialized_node.outputs) == 0: + return + if ( + len(serialized_node.outputs) == 1 + and serialized_node.outputs[0].type == "as_tensor" + ): + self.sync_fx_node(serialized_node.outputs[0].as_tensor.name, fx_node) + return + elif len(serialized_node.outputs) == 1 and isinstance( + serialized_node.outputs[0].value, (SymIntArgument, SymBoolArgument, SymFloatArgument) + ): + self.sync_fx_node(serialized_node.outputs[0].value.as_name, fx_node) + return + + self.deserialize_multiple_outputs(serialized_node, fx_node) + + def deserialize_multiple_outputs( + self, serialized_node: Node, fx_node: torch.fx.Node + ) -> None: + deserialized_metadata = self.deserialize_metadata(serialized_node.metadata) + + def generate_getitem( + meta_val, + fx_node: torch.fx.Node, + arg: Union[TensorArgument, SymIntArgument, SymFloatArgument], + idx: int, + ): + if isinstance(arg, TensorArgument): + name = arg.name + elif isinstance(arg, SymIntArgument): + name = arg.as_name + elif isinstance(arg, SymFloatArgument): + name = arg.as_name + else: + raise AssertionError( + f"generate_getitem got unknown argument type {type(arg)}" + ) + individual_output = self.graph.create_node( + "call_function", + operator.getitem, + (fx_node, idx), + name=name, + ) + self.sync_fx_node(name, individual_output) + meta_val.append(self.serialized_name_to_meta[name]) + # The derived `getitem` nodes should have the same stacktrace as the + # original `fx_node` + individual_output.meta.update(deserialized_metadata) + + def generate_getitems(meta_val, fx_node: torch.fx.Node, args): + for idx, arg in enumerate(args): + if isinstance(arg, (TensorArgument, SymIntArgument, SymFloatArgument)): + generate_getitem(meta_val, fx_node, arg, idx) + continue + + assert isinstance(arg, Argument) + if arg.type in ("as_tensor", "as_sym_int", "as_sym_float"): + generate_getitem(meta_val, fx_node, arg.value, idx) + elif arg.type in ( + "as_tensors", + "as_sym_ints", + "as_sym_floats", + "as_ints", + "as_floats", + "as_strings", + "as_bools", + "as_sym_bools", + ): + list_output = self.graph.create_node( + "call_function", + operator.getitem, + (fx_node, idx), + ) + meta_val.append([]) + generate_getitems(meta_val[-1], list_output, arg.value) + list_output.meta.update(deserialized_metadata) + list_output.meta["val"] = meta_val[-1] + elif arg.type == "as_none": + individual_output = self.graph.create_node( + "call_function", + operator.getitem, + (fx_node, idx), + name="as_none", + ) + meta_val.append(None) + individual_output.meta['val'] = None + individual_output.meta.update(deserialized_metadata) + else: + raise NotImplementedError(f"Unimplemented node output type: {arg}") + + # Convert multiple return types to FX format. + # In FX, each node only returns one value. So in order to represent + # multiple return values, we have to emit a `getitem` node for each + # return value. + # This performs the inverse mapping of the `serialize_outputs` call in + # serialization, see [NOTE: Multiple outputs] + meta_val: List[Any] = [] + if len(serialized_node.outputs) == 1: + assert isinstance(serialized_node.outputs[0].value, list) + assert isinstance(serialized_node.outputs[0].value[0], TensorArgument) + generate_getitems(meta_val, fx_node, serialized_node.outputs[0].as_tensors) + else: + generate_getitems(meta_val, fx_node, serialized_node.outputs) + + # also update the metaval for `fx_node` to be a list(meta) + fx_node.meta["val"] = tuple(meta_val) + self.serialized_name_to_node[fx_node.name] = fx_node + + def deserialize_metadata(self, metadata: Dict[str, str]) -> Dict[str, Any]: + ret: Dict[str, Any] = {} + if stack_trace := metadata.get("stack_trace"): + ret["stack_trace"] = stack_trace + + def deserialize_meta_func(serialized_target: str): + module = None + if serialized_target.startswith("torch.nn"): + module = torch.nn + serialized_target_names = serialized_target.split(".")[2:] + elif serialized_target.startswith("torch"): + module = torch + serialized_target_names = serialized_target.split(".")[1:] + else: + return self.deserialize_operator(serialized_target) + + target = module + for name in serialized_target_names: + if not hasattr(target, name): + return serialized_target + else: + target = getattr(target, name) + return target + + if nn_module_stack_str := metadata.get("nn_module_stack"): + # Originally serialized to "key,orig_path,type_str" + def import_nn_module_stack(key, path, ty): + return key, (path, ty) + + # Helper function that splits strings by commas except for those + # encapsulated by parens, which are valid traces. + # TODO: Currently this is needed due to indexing Sequential + # layers introducing names in the form "layer.slice(1, None, None)". + # If that naming is improved, this fancier splitting can probably be + # reverted to a simple split by comma. + def metadata_split(metadata): + # Remove the parentheses and commas inside them + metadata = re.sub(r'\(.*?\)', '', metadata) + # Split the string by comma, except for those inside parentheses + return re.split(r'(? ep.ArgumentSpec: + if x.type == "as_tensor": + return ep.TensorArgument(name=x.as_tensor.name) + elif x.type == "as_sym_int": + return ep.SymIntArgument(name=x.as_sym_int.as_name) + elif x.type == "as_sym_float": + return ep.SymFloatArgument(name=x.as_sym_float.as_name) + elif x.type == "as_custom_obj": + return ep.ConstantArgument(name=x.as_custom_obj.name, value=self.deserialize_input(x)) + else: + return ep.ConstantArgument(name="", value=self.deserialize_input(x)) + + def deserialize_module_call_signature( + self, module_call_signature: ModuleCallSignature + ) -> ep.ModuleCallSignature: + return ep.ModuleCallSignature( + inputs=[ + self.deserialize_argument_spec(x) for x in module_call_signature.inputs + ], + outputs=[ + self.deserialize_argument_spec(x) for x in module_call_signature.outputs + ], + in_spec=treespec_loads(module_call_signature.in_spec), + out_spec=treespec_loads(module_call_signature.out_spec), + forward_arg_names=names if (names := module_call_signature.forward_arg_names) else None, + ) + + def deserialize_module_call_graph( + self, module_call_graph: List[ModuleCallEntry] + ) -> List[ep.ModuleCallEntry]: + return [ + ep.ModuleCallEntry( + fqn=entry.fqn, + signature=( + self.deserialize_module_call_signature(entry.signature) + if entry.signature + else None + ), + ) + for entry in module_call_graph + ] + + +@final +class ExportedProgramDeserializer(metaclass=Final): + def __init__(self, expected_opset_version: Optional[Dict[str, int]] = None): + self.expected_opset_version: Dict[str, int] = {} + if expected_opset_version: + self.expected_opset_version.update(expected_opset_version) + if "aten" not in self.expected_opset_version: + self.expected_opset_version["aten"] = torch._C._get_max_operator_version() + + def deserialize_range_constraints( + self, + symbol_name_to_range: Dict[str, symbolic_shapes.ValueRanges], + symbol_name_to_symbol: Dict[str, sympy.Symbol], + ) -> Dict[sympy.Symbol, ValueRanges]: + range_constraints = {} + for k, v in symbol_name_to_range.items(): + if symbol := symbol_name_to_symbol.get(k): + range_constraints[symbol] = v # type: ignore[arg-type] + else: + log.warning(f"Symbol {k} did not appear in the graph that was deserialized") # noqa: G004 + return range_constraints + + def deserialize( + self, + exported_program: ExportedProgram, + state_dict: Union[Dict[str, torch.Tensor], bytes], + constants: Union[Dict[str, torch.Tensor], bytes], + example_inputs: Optional[Union[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]], bytes]] = None, + *, + _unsafe_skip_version_check=False, + ) -> ep.ExportedProgram: + assert isinstance(exported_program, ExportedProgram) + version = exported_program.schema_version + + # TODO(zhxchen17) blocked on thrift schema refactor + if version.major != SCHEMA_VERSION[0] and not (version.major == 0 and version.minor == 0): + if not _unsafe_skip_version_check: + raise SerializeError( + f"Serialized schema version {exported_program.schema_version} " + f"does not match our current schema version {SCHEMA_VERSION}." + ) + + symbol_name_to_range = { + k: symbolic_shapes.ValueRanges( + _int_to_sympy_int(v.min_val, -int_oo), _int_to_sympy_int(v.max_val, int_oo) + ) + for k, v in exported_program.range_constraints.items() + } + res = ( + GraphModuleDeserializer() + .deserialize( + exported_program.graph_module, + state_dict, + constants, + example_inputs, + symbol_name_to_range, + ) + ) + range_constraints = self.deserialize_range_constraints( + symbol_name_to_range, + res.names_to_symbols, + ) + + return ep.ExportedProgram( + root=res.graph_module, + graph=res.graph_module.graph, + graph_signature=res.signature, + state_dict=res.state_dict, # type: ignore[arg-type] + range_constraints=range_constraints, + module_call_graph=res.module_call_graph, + example_inputs=res.example_inputs, + constants=res.constants, + verifiers=[load_verifier(v) for v in exported_program.verifiers], + ) + + +class EnumEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, Enum): + return obj.value + if isinstance(obj, bytes): + return base64.b64encode(obj).decode("utf-8") + return super().default(obj) + + +def _dataclass_to_dict(obj): + if isinstance(obj, _Union): + return {obj.type: _dataclass_to_dict(obj.value)} + elif dataclasses.is_dataclass(obj): + return { + f.name: _dataclass_to_dict(getattr(obj, f.name)) + for f in dataclasses.fields(obj) + if not (f.default is None and getattr(obj, f.name) is None) + } + elif isinstance(obj, list): + return [_dataclass_to_dict(x) for x in obj] + elif isinstance(obj, tuple): + return tuple(_dataclass_to_dict(x) for x in obj) + elif isinstance(obj, dict): + return {k: _dataclass_to_dict(v) for k, v in obj.items()} + else: + return obj + + +def serialize( + exported_program: ep.ExportedProgram, + opset_version: Optional[Dict[str, int]] = None, +) -> SerializedArtifact: + serialized_program = ExportedProgramSerializer(opset_version).serialize( + exported_program + ) + assert isinstance(serialized_program.exported_program, ExportedProgram) + + json_program = json.dumps( + _dataclass_to_dict(serialized_program.exported_program), cls=EnumEncoder + ) + json_bytes = json_program.encode("utf-8") + artifact = SerializedArtifact( + json_bytes, + serialized_program.state_dict, + serialized_program.constants, + serialized_program.example_inputs + ) + return artifact + + +def _dict_to_dataclass(cls, data): + assert not isinstance(cls, str), f"Unresolved class type: '{cls}'." + if typing.get_origin(cls) == Annotated: + return _dict_to_dataclass(cls.__origin__, data) + if typing.get_origin(cls) == typing.Union and type(None) in typing.get_args(cls): + if data is None: + return None + ty_args = typing.get_args(cls) + assert len(ty_args) == 2 + return _dict_to_dataclass(ty_args[0], data) + elif isinstance(cls, type) and issubclass(cls, _Union): + assert isinstance(data, dict) + assert len(data) == 1 + _type = next(iter(data.keys())) + _value = next(iter(data.values())) + assert isinstance(_type, str) + field_type = cls.__annotations__[_type] + return cls.create(**{_type: _dict_to_dataclass(field_type, _value)}) + elif dataclasses.is_dataclass(cls): + obj = cls(**data) # type: ignore[assignment,operator] + type_hints = typing.get_type_hints(cls) + for f in dataclasses.fields(cls): + name = f.name + new_field_obj = _dict_to_dataclass(type_hints[name], getattr(obj, name)) + setattr(obj, name, new_field_obj) + return obj + elif isinstance(data, list): + if len(data) == 0: + return data + d_type = typing.get_args(cls)[0] + return [_dict_to_dataclass(d_type, d) for d in data] + elif isinstance(data, dict): + v_type = typing.get_args(cls)[1] + return {k: _dict_to_dataclass(v_type, v) for k, v in data.items()} + return data + + +def deserialize( + artifact: SerializedArtifact, + expected_opset_version: Optional[Dict[str, int]] = None, + *, + _unsafe_skip_version_check=False, +) -> ep.ExportedProgram: + assert isinstance(artifact.exported_program, bytes) + exported_program_str = artifact.exported_program.decode("utf-8") + exported_program_dict = json.loads(exported_program_str) + serialized_exported_program = _dict_to_dataclass(ExportedProgram, exported_program_dict) + return ( + ExportedProgramDeserializer(expected_opset_version) + .deserialize( + serialized_exported_program, + artifact.state_dict, + artifact.constants, + artifact.example_inputs, + _unsafe_skip_version_check=_unsafe_skip_version_check, + ) + ) + + +def _canonicalize_graph( + sorted_inputs, sorted_outputs, graph +) -> Tuple[Graph, Dict[str, str]]: + def _get_argument(a: Argument): + if a.type == "as_none": + return None + elif a.type == "as_tensor": + return a.as_tensor + elif a.type == "as_tensors": + return a.as_tensors + elif a.type == "as_int": + return None + elif a.type == "as_ints": + return None + elif a.type == "as_float": + return None + elif a.type == "as_floats": + return None + elif a.type == "as_string": + return None + elif a.type == "as_strings": + return None + elif a.type == "as_sym_int": + return a.as_sym_int + elif a.type == "as_sym_ints": + return a.as_sym_ints + elif a.type == "as_sym_float": + return a.as_sym_float + elif a.type == "as_sym_floats": + return a.as_sym_floats + elif a.type == "as_scalar_type": + return None + elif a.type == "as_memory_format": + return None + elif a.type == "as_layout": + return None + elif a.type == "as_device": + return None + elif a.type == "as_bool": + return None + elif a.type == "as_bools": + return None + elif a.type == "as_sym_bool": + return a.as_sym_bool + elif a.type == "as_sym_bools": + return a.as_sym_bools + elif a.type == "as_graph": + return None + elif a.type == "as_optional_tensors": + return a.as_optional_tensors + elif a.type == "as_custom_obj": + return None + elif a.type == "as_operator": + return None + else: + raise AssertionError(f"Unknown input type to the ExportedProgram: {a}") + + # Stage 1: Reorder named items. + def for_args(f, a): + assert isinstance(a, Argument) + pytree.tree_map(f, _get_argument(a)) + + def sort_nodes(nodes): + @dataclass + class Edges: + outs: List[int] + ins: int + + graph_inputs: Set[str] = set() + def_table: Dict[str, int] = {} + edges: Dict[int, Edges] = {} + candidates: List[Tuple[str, List[Tuple[str, List[int]]], int]] = [] + rank: Dict[str, int] = {} + ret: List[Node] = [] + + def get_name(a) -> Optional[str]: + if a is None: + return None + if isinstance(a, TensorArgument): + return a.name + elif isinstance(a, (SymIntArgument, SymBoolArgument, SymFloatArgument)): + if a.type == "as_name": + return a.as_name + elif a.type in ("as_int", "as_bool", "as_float"): + return None + else: + raise AssertionError(f"Unknown argument type: {a}") + elif isinstance(a, OptionalTensorArgument): + if a.type == "as_tensor": + return a.as_tensor.name + elif a.type == "as_none": + return None + else: + raise AssertionError(f"Unknown optional tensor type: {a}") + else: + raise AssertionError(f"Unknown argument type: {a}") + + for i in sorted_inputs: + + def add_input(a): + if s := get_name(a): + graph_inputs.add(s) + + for_args(add_input, i) + + for idx, node in enumerate(nodes): + + def add_def(a): + if s := get_name(a): + assert s not in def_table + def_table[s] = idx + + for o in node.outputs: + for_args(add_def, o) + + edges[idx] = Edges([], 0) + + for idx, user in enumerate(nodes): + + def add_edge(a): + if s := get_name(a): + if s not in def_table: + assert s in graph_inputs + return + src = def_table[s] + edges[src].outs.append(idx) + edges[idx].ins += 1 + + for i in user.inputs: + for_args(add_edge, i.arg) + + def add_rank(a): + if s := get_name(a): + assert s not in rank + rank[s] = len(rank) + + def get_rank(a): + if s := get_name(a): + return rank[s] + else: + return -1 + + for i in sorted_inputs: + for_args(add_rank, i) + + def add_candidate(idx: int): + def get_ranks(i): + ranks = [] + for_args(lambda x: ranks.append(get_rank(x)), i) + return ranks + + node = nodes[idx] + args_rank = [(a.name, get_ranks(a.arg)) for a in node.inputs] + heapq.heappush(candidates, (node.target, args_rank, idx)) + + for idx, e in edges.items(): + if e.ins == 0: + add_candidate(idx) + + while len(candidates) > 0: + _, _, idx = heapq.heappop(candidates) + node = nodes[idx] + for o in node.outputs: + for_args(add_rank, o) + ret.append(node) + assert idx in edges + for user in edges[idx].outs: + e = edges[user] + assert e.ins > 0 + e.ins -= 1 + if e.ins == 0: + add_candidate(user) + edges[idx].outs.clear() + + return ret + + sorted_nodes = sort_nodes(graph.nodes) + assert len(sorted_nodes) == len(graph.nodes) + + # Stage 2: Rename nodes. + name_table: Dict[str, str] = {} + + def rename_def(a): + def _rename(arg_name, values): + new_name = f"_{len(name_table)}" + assert arg_name not in name_table + name_table[arg_name] = new_name + assert arg_name in values + values[new_name] = values.pop(arg_name) + return new_name + + if a is None: + return + if isinstance(a, TensorArgument): + a.name = _rename(a.name, graph.tensor_values) + elif isinstance(a, SymIntArgument): + if a.type == "as_name": + a.as_name = _rename(a.as_name, graph.sym_int_values) + elif isinstance(a, SymFloatArgument): + if a.type == "as_name": + a.as_name = _rename(a.as_name, graph.sym_float_values) + elif isinstance(a, SymBoolArgument): + if a.type == "as_name": + a.as_name = _rename(a.as_name, graph.sym_bool_values) + else: + raise AssertionError(f"Unknown argument type: {a}") + + def replace_use(a): + if a is None: + return + if isinstance(a, TensorArgument): + a.name = name_table.get(a.name, a.name) + elif isinstance(a, (SymIntArgument, SymFloatArgument)): + if a.type == "as_name": + a.as_name = name_table.get(a.as_name, a.as_name) + elif isinstance(a, SymBoolArgument): + if a.type == "as_name": + a.as_name = name_table.get(a.as_name, a.as_name) + elif isinstance(a, OptionalTensorArgument): + if a.type == "as_tensor": + a.as_tensor.name = name_table.get(a.as_tensor.name, a.as_tensor.name) + else: + raise AssertionError(f"Unknown argument type: {a}") + + for i in sorted_inputs: + for_args(rename_def, i) + + for n in sorted_nodes: + for o in n.outputs: + for_args(rename_def, o) + + for n in sorted_nodes: + for i in n.inputs: + for_args(replace_use, i.arg) + + for o in sorted_outputs: + for_args(replace_use, o) + + # Stage 3: Remove unstable fields. + for n in sorted_nodes: + n.metadata.clear() + + # Stage 4: Aggregate values. + sorted_tensor_values = dict(sorted(graph.tensor_values.items(), key=operator.itemgetter(0))) + sorted_sym_int_values = dict( + sorted(graph.sym_int_values.items(), key=operator.itemgetter(0)) + ) + sorted_sym_float_values = dict( + sorted(graph.sym_float_values.items(), key=operator.itemgetter(0)) + ) + sorted_sym_bool_values = dict( + sorted(graph.sym_bool_values.items(), key=operator.itemgetter(0)) + ) + + # Stage 5: Recurse in subgraphs. + counter = 0 + for node in sorted_nodes: + for i in node.inputs: + a = i.arg + if a.type == "as_graph": + a.as_graph.graph, _ = _canonicalize_graph( + a.as_graph.graph.inputs, a.as_graph.graph.outputs, a.as_graph.graph + ) + a.as_graph.name = f"_g{counter}" + counter += 1 + + graph = Graph( + inputs=sorted_inputs, + outputs=sorted_outputs, + nodes=sorted_nodes, + tensor_values=sorted_tensor_values, + sym_int_values=sorted_sym_int_values, + sym_float_values=sorted_sym_float_values, + sym_bool_values=sorted_sym_bool_values, + is_single_tensor_return=graph.is_single_tensor_return, + ) + return graph, name_table + + +def canonicalize(ep: ExportedProgram) -> ExportedProgram: + """ + Normalize a serialized ExportedProgram, so that different eager program which + shares the same semantics can get a single representation on disk. + + This function canonicalizes an ExportedProgram by: + + 1. Sorting nodes in topological order. + 2. Rename nodes to have unique names. + 3. Remove unstable fields. + 4. Aggregate the above program fields. + 5. Recurse in subgraphs. + + Args: + ep (ExportedProgram): The ExportedProgram to canonicalize. + + Returns: + ExportedProgram: The canonicalized exported program. + """ + ep = copy.deepcopy(ep) + + opset_version = dict(sorted(ep.opset_version.items(), key=operator.itemgetter(0))) + range_constraints = dict(sorted(ep.range_constraints.items(), key=operator.itemgetter(0))) + module_call_graph = sorted(ep.graph_module.module_call_graph, key=lambda x: x.fqn) + signature = ep.graph_module.signature + graph = ep.graph_module.graph + + assert len(graph.inputs) == len(signature.input_specs) + assert len(graph.outputs) == len(signature.output_specs) + + def rank_input(inp) -> Tuple[int, Optional[str], int]: + idx, (arg, spec) = inp + assert isinstance(spec, InputSpec) + if spec.type == "user_input": + return 5, None, idx + elif spec.type == "parameter": + return 1, spec.parameter.parameter_name, idx + elif spec.type == "buffer": + return 2, spec.buffer.buffer_name, idx + elif spec.type == "tensor_constant": + return 3, spec.tensor_constant.tensor_constant_name, idx + elif spec.type == "custom_obj": + return 4, spec.custom_obj.custom_obj_name, idx + elif spec.type == "token": + return 0, None, idx + elif spec.type == "constant_input": + return 6, spec.constant_input.name, idx + else: + raise AssertionError(f"Unknown input type: {spec}") + + def rank_output(out) -> Tuple[int, Optional[str], int]: + idx, (arg, spec) = out + assert isinstance(spec, OutputSpec) + if spec.type == "user_output": + return 3, None, idx + elif spec.type == "loss_output": + return 3, None, idx + elif spec.type == "buffer_mutation": + return 1, spec.buffer_mutation.buffer_name, idx + elif spec.type == "gradient_to_parameter": + return 4, spec.gradient_to_parameter.parameter_name, idx + elif spec.type == "gradient_to_user_input": + return 5, None, idx + elif spec.type == "user_input_mutation": + return 2, None, idx + elif spec.type == "token": + return 0, None, idx + else: + raise AssertionError(f"Unknown output type: {spec}") + + sorted_ins = sorted( + enumerate(zip(graph.inputs, signature.input_specs)), key=rank_input + ) + + if len(sorted_ins) > 0: + sorted_inputs, input_specs = zip(*(i for idx, i in sorted_ins)) # type: ignore[assignment] + else: + sorted_inputs = () + input_specs = () + + sorted_outs = sorted( + enumerate(zip(graph.outputs, signature.output_specs)), key=rank_output + ) + sorted_outputs, output_specs = zip(*(i for idx, i in sorted_outs)) # type: ignore[assignment] + + sorted_graph, replace_table = _canonicalize_graph( + sorted_inputs, sorted_outputs, graph + ) + + def replace_input(inp): + assert isinstance(spec, InputSpec) + if spec.type == "user_input": + arg = spec.user_input.arg + if arg.type == "as_tensor": + t = arg.as_tensor + t.name = replace_table[t.name] + elif arg.type == "as_sym_int": + s = arg.as_sym_int + if s.type == "as_name": + s.as_name = replace_table[s.as_name] + elif s.type == "as_int": + pass + else: + raise AssertionError(f"Unknown sym_int type: {s}") + elif arg.type == "as_sym_float": + f = arg.as_sym_float + if f.type == "as_name": + f.as_name = replace_table[f.as_name] + elif f.type == "as_float": + pass + else: + raise AssertionError(f"Unknown sym_float type: {f}") + elif arg.type in ( + "as_none", + "as_bool", + "as_int", + "as_float", + "as_string", + "as_custom_obj", + ): + return + else: + raise AssertionError(f"Unknown input type: {arg}") + elif spec.type == "parameter": + t = spec.parameter.arg + t.name = replace_table[t.name] + elif spec.type == "buffer": + t = spec.buffer.arg + t.name = replace_table[t.name] + elif spec.type == "tensor_constant": + t = spec.tensor_constant.arg + t.name = replace_table[t.name] + elif spec.type == "custom_obj": + return + elif spec.type == "token": + tok = spec.token.arg + tok.name = replace_table[tok.name] + elif spec.type == "constant_input": + return + else: + raise AssertionError(f"Unknown input type: {spec}") + + def replace_output(out): + assert isinstance(spec, OutputSpec) + if spec.type == "user_output": + arg = spec.user_output.arg + if arg.type == "as_tensor": + t = arg.as_tensor + t.name = replace_table[t.name] + elif arg.type == "as_sym_int": + s = arg.as_sym_int + if s.type == "as_name": + s.as_name = replace_table[s.as_name] + elif s.type == "as_int": + pass + else: + raise AssertionError(f"Unknown sym_int type: {s}") + elif arg.type == "as_sym_float": + f = arg.as_sym_float + if f.type == "as_name": + f.as_name = replace_table[f.as_name] + elif f.type == "as_float": + pass + else: + raise AssertionError(f"Unknown sym_float type: {f}") + elif arg.type in ("as_none", "as_int", "as_float", "as_string"): + return + else: + raise AssertionError(f"Unknown input type: {arg}") + elif spec.type == "loss_output": + t = spec.loss_output.arg + t.name = replace_table[t.name] + elif spec.type == "buffer_mutation": + t = spec.buffer_mutation.arg + t.name = replace_table[t.name] + elif spec.type == "gradient_to_parameter": + t = spec.gradient_to_parameter.arg + t.name = replace_table[t.name] + elif spec.type == "gradient_to_user_input": + g = spec.gradient_to_user_input + g.arg.name = replace_table[g.arg.name] + g.user_input_name = replace_table[g.user_input_name] + elif spec.type == "user_input_mutation": + u = spec.user_input_mutation + u.arg.name = replace_table[u.arg.name] + u.user_input_name = replace_table[u.user_input_name] + elif spec.type == "token": + tok = spec.token.arg + tok.name = replace_table[tok.name] + else: + raise AssertionError(f"Unknown output type: {spec}") + + for spec in input_specs: + replace_input(spec) + + for spec in output_specs: + replace_output(spec) + + return ExportedProgram( + graph_module=GraphModule( + graph=sorted_graph, + signature=GraphSignature( + input_specs=list(input_specs), + output_specs=list(output_specs), + ), + module_call_graph=module_call_graph, + ), + opset_version=opset_version, + range_constraints=range_constraints, + schema_version=ep.schema_version, + verifiers=ep.verifiers, + torch_version=ep.torch_version, + ) + + +class ExtensionHandler: + """ + Base class for handling extension operators. + """ + @classmethod + def namespace(cls) -> str: + raise NotImplementedError(f"{cls.__class__} namespace() must be implemented") + + @classmethod + def to_op_name(cls, op) -> str: + raise NotImplementedError(f"{cls.__class__} op_name() must be implemented") + + @classmethod + def from_op_name(cls, name: str): + raise NotImplementedError(f"{cls.__class__} op_name() must be implemented") + + @classmethod + def op_schema(cls, op) -> torch.FunctionSchema: + raise NotImplementedError(f"{cls.__class__} op_schema() must be implemented") + + +def register_extension( + op_type: Type[Any], + extension_handler: Type[ExtensionHandler], +): + """Register custom de/serialization method for a node with non-standard type.""" + assert issubclass(extension_handler, ExtensionHandler), f"Expected ExtensionHandler, got {extension_handler}." + assert op_type not in _serialization_registry, f"{op_type} is already registered." + assert isinstance(op_type, type) # Maybe a good idea to enforce this first. + assert not (op_type.__module__.startswith("torch") or op_type.__module__.startswith("builtins")) + assert extension_handler.namespace() not in _deserialization_registry + _serialization_registry[op_type] = extension_handler + _deserialization_registry[extension_handler.namespace()] = extension_handler + + +def _registered_extension_types(): + return tuple( + _serialization_registry.keys() + ) + + +# Registry to store all custom serialization implementations. +# The registry maps a operation to its serialization function (a callable), in their own +# namespace to avoid conflicts. +# Serialization: Op type --> custom handler. +# De-serialization: Namespace --> custom handler. +_serialization_registry: Dict[Type[Any], Type[ExtensionHandler]] = {} +_deserialization_registry: Dict[str, Type[ExtensionHandler]] = {} diff --git a/lib/python3.12/site-packages/torch/_export/serde/union.py b/lib/python3.12/site-packages/torch/_export/serde/union.py new file mode 100644 index 0000000000000000000000000000000000000000..b129e8dd9a89ef4870d7ef3dc724aca8ccae3a43 --- /dev/null +++ b/lib/python3.12/site-packages/torch/_export/serde/union.py @@ -0,0 +1,70 @@ +# mypy: allow-untyped-defs +import functools +from dataclasses import fields +from typing import Hashable, Set + + +class _UnionTag(str): + _cls: Hashable + + @staticmethod + def create(t, cls): + tag = _UnionTag(t) + assert not hasattr(tag, "_cls") + tag._cls = cls + return tag + + def __eq__(self, cmp) -> bool: + assert isinstance(cmp, str) + other = str(cmp) + assert other in _get_field_names( + self._cls + ), f"{other} is not a valid tag for {self._cls}. Available tags: {_get_field_names(self._cls)}" + return str(self) == other + + def __hash__(self): + return hash(str(self)) + + +@functools.lru_cache(maxsize=None) +def _get_field_names(cls) -> Set[str]: + return {f.name for f in fields(cls)} + + +class _Union: + _type: _UnionTag + + @classmethod + def create(cls, **kwargs): + assert len(kwargs) == 1 + obj = cls(**{**{f.name: None for f in fields(cls)}, **kwargs}) # type: ignore[arg-type] + obj._type = _UnionTag.create(next(iter(kwargs.keys())), cls) + return obj + + def __post_init__(self): + assert not any(f.name in ("type", "_type", "create", "value") for f in fields(self)) # type: ignore[arg-type, misc] + + @property + def type(self) -> str: + try: + return self._type + except AttributeError as e: + raise RuntimeError( + f"Please use {type(self).__name__}.create to instantiate the union type." + ) from e + + @property + def value(self): + return getattr(self, self.type) + + def __getattribute__(self, name): + attr = super().__getattribute__(name) + if attr is None and name in _get_field_names(type(self)) and name != self.type: # type: ignore[arg-type] + raise AttributeError(f"Field {name} is not set.") + return attr + + def __str__(self): + return self.__repr__() + + def __repr__(self): + return f"{type(self).__name__}({self.type}={getattr(self, self.type)})" diff --git a/lib/python3.12/site-packages/torch/fx/__init__.py b/lib/python3.12/site-packages/torch/fx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..74691bbe72ac6639759ac29186aa21c874c1808d --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/__init__.py @@ -0,0 +1,115 @@ +r''' +FX is a toolkit for developers to use to transform ``nn.Module`` +instances. FX consists of three main components: a **symbolic tracer,** +an **intermediate representation**, and **Python code generation**. A +demonstration of these components in action: + +:: + + import torch + + + # Simple module for demonstration + class MyModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.param = torch.nn.Parameter(torch.rand(3, 4)) + self.linear = torch.nn.Linear(4, 5) + + def forward(self, x): + return self.linear(x + self.param).clamp(min=0.0, max=1.0) + + + module = MyModule() + + from torch.fx import symbolic_trace + + # Symbolic tracing frontend - captures the semantics of the module + symbolic_traced: torch.fx.GraphModule = symbolic_trace(module) + + # High-level intermediate representation (IR) - Graph representation + print(symbolic_traced.graph) + """ + graph(): + %x : [num_users=1] = placeholder[target=x] + %param : [num_users=1] = get_attr[target=param] + %add : [num_users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {}) + %linear : [num_users=1] = call_module[target=linear](args = (%add,), kwargs = {}) + %clamp : [num_users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0}) + return clamp + """ + + # Code generation - valid Python code + print(symbolic_traced.code) + """ + def forward(self, x): + param = self.param + add = x + param; x = param = None + linear = self.linear(add); add = None + clamp = linear.clamp(min = 0.0, max = 1.0); linear = None + return clamp + """ + +The **symbolic tracer** performs "symbolic execution" of the Python +code. It feeds fake values, called Proxies, through the code. Operations +on theses Proxies are recorded. More information about symbolic tracing +can be found in the :func:`symbolic_trace` and :class:`Tracer` +documentation. + +The **intermediate representation** is the container for the operations +that were recorded during symbolic tracing. It consists of a list of +Nodes that represent function inputs, callsites (to functions, methods, +or :class:`torch.nn.Module` instances), and return values. More information +about the IR can be found in the documentation for :class:`Graph`. The +IR is the format on which transformations are applied. + +**Python code generation** is what makes FX a Python-to-Python (or +Module-to-Module) transformation toolkit. For each Graph IR, we can +create valid Python code matching the Graph's semantics. This +functionality is wrapped up in :class:`GraphModule`, which is a +:class:`torch.nn.Module` instance that holds a :class:`Graph` as well as a +``forward`` method generated from the Graph. + +Taken together, this pipeline of components (symbolic tracing -> +intermediate representation -> transforms -> Python code generation) +constitutes the Python-to-Python transformation pipeline of FX. In +addition, these components can be used separately. For example, +symbolic tracing can be used in isolation to capture a form of +the code for analysis (and not transformation) purposes. Code +generation can be used for programmatically generating models, for +example from a config file. There are many uses for FX! + +Several example transformations can be found at the +`examples `__ +repository. +''' + +from torch.fx._symbolic_trace import ( # noqa: F401 + PH, + ProxyableClassMeta, + symbolic_trace, + Tracer, + wrap, +) +from torch.fx.graph import CodeGen, Graph # noqa: F401 +from torch.fx.graph_module import GraphModule +from torch.fx.interpreter import Interpreter, Transformer +from torch.fx.node import has_side_effect, map_arg, Node +from torch.fx.proxy import Proxy +from torch.fx.subgraph_rewriter import replace_pattern + + +__all__ = [ + "symbolic_trace", + "Tracer", + "wrap", + "Graph", + "GraphModule", + "Interpreter", + "Transformer", + "Node", + "Proxy", + "replace_pattern", + "has_side_effect", + "map_arg", +] diff --git a/lib/python3.12/site-packages/torch/fx/_compatibility.py b/lib/python3.12/site-packages/torch/fx/_compatibility.py new file mode 100644 index 0000000000000000000000000000000000000000..8a2eeb0d2d69552676d726d37cc847d7830db86d --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/_compatibility.py @@ -0,0 +1,39 @@ +import textwrap +from typing import Any, Callable, Dict, TypeVar + + +_BACK_COMPAT_OBJECTS: Dict[Any, None] = {} +_MARKED_WITH_COMPATIBILITY: Dict[Any, None] = {} + + +_T = TypeVar("_T") + + +def compatibility(is_backward_compatible: bool) -> Callable[[_T], _T]: + if is_backward_compatible: + + def mark_back_compat(fn: _T) -> _T: + docstring = textwrap.dedent(getattr(fn, "__doc__", None) or "") + docstring += """ +.. note:: + Backwards-compatibility for this API is guaranteed. +""" + fn.__doc__ = docstring + _BACK_COMPAT_OBJECTS.setdefault(fn) + _MARKED_WITH_COMPATIBILITY.setdefault(fn) + return fn + + return mark_back_compat + else: + + def mark_not_back_compat(fn: _T) -> _T: + docstring = textwrap.dedent(getattr(fn, "__doc__", None) or "") + docstring += """ +.. warning:: + This API is experimental and is *NOT* backward-compatible. +""" + fn.__doc__ = docstring + _MARKED_WITH_COMPATIBILITY.setdefault(fn) + return fn + + return mark_not_back_compat diff --git a/lib/python3.12/site-packages/torch/fx/_lazy_graph_module.py b/lib/python3.12/site-packages/torch/fx/_lazy_graph_module.py new file mode 100644 index 0000000000000000000000000000000000000000..cc2f686ebba10d8bfbcf1eff4a69c77078ff0202 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/_lazy_graph_module.py @@ -0,0 +1,185 @@ +# mypy: allow-untyped-defs +from contextlib import contextmanager + +from torch.fx.graph_module import ( + _format_import_block, + GraphModule, + reduce_graph_module, + reduce_package_graph_module, +) +from torch.package import PackageExporter, sys_importer + +from ._compatibility import compatibility + + +_use_lazy_graph_module_flag = False +_force_skip_lazy_graph_module_flag = False + + +@compatibility(is_backward_compatible=False) +@contextmanager +def _force_skip_lazy_graph_module(): + """ + Skip using lazy graph module disregarding the setting of _use_lazy_graph_module. + Use to skip _LazyGraphModule when testing inductor torchscript related backend. + + torch.jit.script a _LazyGraphModule results in following error: + https://gist.github.com/shunting314/5143654c8084aed84ecd19b818258a69 + """ + try: + global _force_skip_lazy_graph_module_flag + prior = _force_skip_lazy_graph_module_flag + _force_skip_lazy_graph_module_flag = True + yield + finally: + _force_skip_lazy_graph_module_flag = prior + + +@compatibility(is_backward_compatible=False) +@contextmanager +def _use_lazy_graph_module(should_use: bool): + try: + global _use_lazy_graph_module_flag + prior = _use_lazy_graph_module_flag + _use_lazy_graph_module_flag = ( + should_use and not _force_skip_lazy_graph_module_flag + ) + yield + finally: + _use_lazy_graph_module_flag = prior + + +@compatibility(is_backward_compatible=False) +def _get_graph_module_cls(): + return _LazyGraphModule if _use_lazy_graph_module_flag else GraphModule + + +def _make_graph_module(*args, graph_module_cls=None, **kwargs): + if graph_module_cls is None: + graph_module_cls = _get_graph_module_cls() + + return graph_module_cls(*args, **kwargs) + + +@compatibility(is_backward_compatible=False) +class _LazyGraphModule(GraphModule): + """ + The main difference between _LazyGraphModule and GraphModule is how recompile happens. + GraphModule will do a 'recompile' call to generate python code and the forward method when it's + constructed. Later on if the graph get updated, recompile method can be called again to refresh + the saved python code and forward method. + + However in some cases especially in inductor, the recompilation can be a waste since we never + check the python code for the graph module or call its forward method. A few more concreate + examples regarding pattern matching fx passes in inductor: + 1. some passes will update the graph to be compiled and then call recompile on the GraphModule. + 2. some passes will trace small pattern function to search it in the graph being compiled and + replace the match with the traced graph of a replacement function. The pattern graph and + replacement graph are quite small but there are large amount of them. Doing GraphModule.recompile + for them in GraphModule.__init__ is also a waste of time. + + However simply skip calling GraphModule.recompile in these scenarios is also dangeruous. + People may want to check the python code or call the GraphModule's forward method for debugging purposes. + + The way _LazyGraphModule solves it is, we override the recompile method to just mark the + need for recompilation but does not do the actual recompilation. Later on if people really + access the compiled python code or call the GraphModule's forward method, we do the real + recompilation. + """ + + @classmethod + def from_graphmodule(cls, gm: GraphModule): + if isinstance(gm, _LazyGraphModule): + return gm + else: + return _LazyGraphModule(gm, gm.graph) + + @staticmethod + def force_recompile(gm): + """ + Sometimes we need force a recompile as a workaround + - we want to do the real recompilation before symbolic_trace to avoid error: + https://gist.github.com/shunting314/75549c2e82ae07ac1139c94a3583d259 + """ + if isinstance(gm, _LazyGraphModule): + gm.real_recompile() + + def real_recompile(self): + if self._needs_recompile(): + self._real_recompile() + + @classmethod + def _needs_recompile(cls): + return cls.forward is cls._lazy_forward + + def _lazy_forward(self, *args, **kwargs): + # Call self.real_recompile() rather than self._real_recompile() here. + # The _lazy_forward method may be saved and call repeatedly. + # Calling self.real_recompile can make sure we skip recompilation if + # we have already done so. + self.real_recompile() + assert not self._needs_recompile() + + # call `__call__` rather than 'forward' since recompilation may + # install a wrapper for `__call__` to provide a customized error + # message. + return self(*args, **kwargs) + + forward = _lazy_forward + + # TODO: we shold handle __reduce_deploy__ the same way as __reduce_package__, + # or __reduce__ by calling _real_recompile. But I don't find a good way + # to test __reduce_deploy__ out. Also it's very unlikely that LazyGraphModule + # will be used in torch::deploy. So it's skipped for now. + + def __reduce_package__(self, exporter: PackageExporter): + """ + Follow GraphModule.__reduce__ but call 'self._real_recompile' rather + than 'self.recompile' since for a _LazyGraphModule, self.recompile just + mark the need of recompilation and does not return the PythonCode object. + """ + python_code = self._real_recompile() + dict_without_graph = self.__dict__.copy() + dict_without_graph["_graphmodule_cls_name"] = self.__class__.__name__ + del dict_without_graph["_graph"] + + generated_module_name = f"fx-generated._{exporter.get_unique_id()}" + import_block = _format_import_block(python_code.globals, exporter.importer) + module_code = import_block + self.code + exporter.save_source_string(generated_module_name, module_code) + return ( + reduce_package_graph_module, + (dict_without_graph, generated_module_name), + ) + + def __reduce__(self): + """ + Follow GraphModule.__reduce__ but call 'self._real_recompile' rather + than 'self.recompile' since for a _LazyGraphModule, self.recompile just + mark the need of recompilation and does not return the PythonCode object. + """ + python_code = self._real_recompile() + dict_without_graph = self.__dict__.copy() + import_block = _format_import_block(python_code.globals, sys_importer) + del dict_without_graph["_graph"] + return (reduce_graph_module, (dict_without_graph, import_block)) + + def _real_recompile(self): + return super().recompile() + + @classmethod + def recompile(cls): + cls.forward = cls._lazy_forward + + @property + def code(self) -> str: + self.real_recompile() + return super().code + + def __str__(self) -> str: + """ + str(GraphModule) will access the _code attribute. Make sure recompile + happens so _code attribute is available. + """ + self.real_recompile() + return super().__str__() diff --git a/lib/python3.12/site-packages/torch/fx/_pytree.py b/lib/python3.12/site-packages/torch/fx/_pytree.py new file mode 100644 index 0000000000000000000000000000000000000000..2ccbfdf048c943e811cd7d9fc092a81d2c30bcea --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/_pytree.py @@ -0,0 +1,103 @@ +# mypy: allow-untyped-defs +from collections import namedtuple +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Type + +import torch.return_types +from torch.utils._pytree import PyTree, TreeSpec + + +FlattenFuncSpec = Callable[[PyTree, TreeSpec], List] +FlattenFuncExactMatchSpec = Callable[[PyTree, TreeSpec], bool] + +SUPPORTED_NODES: Dict[Type[Any], FlattenFuncSpec] = {} +SUPPORTED_NODES_EXACT_MATCH: Dict[Type[Any], Optional[FlattenFuncExactMatchSpec]] = {} + + +def register_pytree_flatten_spec( + cls: Type[Any], + flatten_fn_spec: FlattenFuncSpec, + flatten_fn_exact_match_spec: Optional[FlattenFuncExactMatchSpec] = None, +) -> None: + SUPPORTED_NODES[cls] = flatten_fn_spec + SUPPORTED_NODES_EXACT_MATCH[cls] = flatten_fn_exact_match_spec + + +def tree_flatten_spec( + pytree: PyTree, + spec: TreeSpec, + exact_structural_match=False, +) -> List[Any]: + if spec.is_leaf(): + return [pytree] + if spec.type not in SUPPORTED_NODES: + raise RuntimeError( + f"{type(pytree)} does not have a flatten_fn_spec associated with it. Please register one with " + "torch.fx._pytree.register_pytree_flatten_spec. If you have serialized your model, make " + "sure that any custom pytrees have been registered before loading it.", + ) + flatten_fn_spec = SUPPORTED_NODES[spec.type] + child_pytrees = flatten_fn_spec(pytree, spec) + if exact_structural_match: + flatten_fn_exact_match_spec = SUPPORTED_NODES_EXACT_MATCH[spec.type] + if flatten_fn_exact_match_spec and not flatten_fn_exact_match_spec( + pytree, + spec, + ): + raise RuntimeError(f"Cannot flatten pytree {pytree}, given spec: {spec}") + result = [] + for child, child_spec in zip(child_pytrees, spec.children_specs): + flat = tree_flatten_spec(child, child_spec, exact_structural_match) + result += flat + return result + + +def _dict_flatten_spec(d: Dict[Any, Any], spec: TreeSpec) -> List[Any]: + return [d[k] for k in spec.context] + + +def _list_flatten_spec(d: List[Any], spec: TreeSpec) -> List[Any]: + return [d[i] for i in range(spec.num_children)] + + +def _tuple_flatten_spec(d: Tuple[Any], spec: TreeSpec) -> List[Any]: + return [d[i] for i in range(spec.num_children)] + + +def _namedtuple_flatten_spec(d: NamedTuple, spec: TreeSpec) -> List[Any]: + return [d[i] for i in range(spec.num_children)] + + +def _dict_flatten_spec_exact_match(d: Dict[Any, Any], spec: TreeSpec) -> bool: + return len(d) == spec.num_children + + +def _list_flatten_spec_exact_match(d: List[Any], spec: TreeSpec) -> bool: + return len(d) == spec.num_children + + +def _tuple_flatten_spec_exact_match(d: Tuple[Any], spec: TreeSpec) -> bool: + return len(d) == spec.num_children + + +def _namedtuple_flatten_spec_exact_match(d: NamedTuple, spec: TreeSpec) -> bool: + return len(d) == spec.num_children + + +register_pytree_flatten_spec(dict, _dict_flatten_spec, _dict_flatten_spec_exact_match) +register_pytree_flatten_spec(list, _list_flatten_spec, _list_flatten_spec_exact_match) +register_pytree_flatten_spec( + tuple, + _tuple_flatten_spec, + _tuple_flatten_spec_exact_match, +) +for return_type in torch.return_types.all_return_types: + register_pytree_flatten_spec( + return_type, + _tuple_flatten_spec, + _tuple_flatten_spec_exact_match, + ) +register_pytree_flatten_spec( + namedtuple, # type: ignore[arg-type] + _namedtuple_flatten_spec, + _namedtuple_flatten_spec_exact_match, +) diff --git a/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py b/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..2489ac6189fe9c09864f47614ee31a1c2f43b13b --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py @@ -0,0 +1,1316 @@ +# mypy: allow-untyped-defs +import builtins +import collections +import contextlib +import copy +import functools +import inspect +import math +import os +import warnings +from itertools import chain +from types import CodeType, FunctionType, ModuleType +from typing import ( + Any, + Callable, + Dict, + List, + NamedTuple, + Optional, + Set, + Tuple, + Type, + Union, +) + +import torch +import torch.utils._pytree as pytree +from torch._C import ScriptObject # type: ignore[attr-defined] +from torch._library.fake_class_registry import FakeScriptObject + +from ._compatibility import compatibility +from ._lazy_graph_module import _make_graph_module +from .graph import _PyTreeCodeGen, _PyTreeInfo, Graph +from .graph_module import GraphModule +from .node import Argument, base_types, map_aggregate +from .proxy import ParameterProxy, Proxy, Scope, ScopeContextManager, TracerBase + + +HAS_VARSTUFF = inspect.CO_VARARGS | inspect.CO_VARKEYWORDS + +# These need to run in global scope to handle nested calls correctly +_orig_module_call: Callable = torch.nn.Module.__call__ +_orig_module_getattr: Callable = torch.nn.Module.__getattr__ + +_proxyable_classes: Dict[Type, None] = {} + +_is_fx_tracing_flag = False + + +def is_fx_tracing(): + return _is_fx_tracing_flag + + +@compatibility(is_backward_compatible=True) +class ProxyableClassMeta(type): + """ + ProxyableClassMeta allows you to make construction of a given Python class + symbolically traceable. For example:: + + import torch + import torch.fx + + + class TensorPair(metaclass=torch.fx.ProxyableClassMeta): + def __init__(self, left, right): + self.left, self.right = left, right + + def add(self, other): + l = self.left + other.left + r = self.right + other.right + return TensorPair(l, r) + + def mul(self, other): + l = self.left * other.left + r = self.right * other.right + return TensorPair(l, r) + + + def use_tensor_pair_ctor(x: TensorPair, y: torch.Tensor): + s = x.add(TensorPair(y, y)) + return s.mul(x) + + + x = TensorPair(torch.randn(5, 3), torch.randn(5, 3)) + y = torch.randn(5, 3) + ref_out = use_tensor_pair_ctor(x, y) + + traced = torch.fx.symbolic_trace(use_tensor_pair_ctor) + print(traced.code) + ''' + def forward(self, x : __main___TensorPair, y : torch.Tensor): + tensor_pair = __main___TensorPair(y, y); y = None + add = x.add(tensor_pair); tensor_pair = None + mul = add.mul(x); add = x = None + return mul + ''' + + From this example, we can see that construction of a class (``TensorPair``) + defined with ``ProxyableClassMeta`` as metaclass can be recorded in symbolic + tracing. + """ + + def __init__(cls, name, bases, attrs): + _proxyable_classes.setdefault(cls) + super().__init__(name, bases, attrs) + + def __call__(cls, *args, **kwargs): + instance = cls.__new__(cls) # type: ignore[call-overload] + + if not is_fx_tracing(): + cls.__init__(instance, *args, **kwargs) # type: ignore[misc] + return instance + + found_proxies = [] + + def check_proxy(a): + if isinstance(a, Proxy): + found_proxies.append(a) + + map_aggregate(args, check_proxy) + map_aggregate(kwargs, check_proxy) + + if len(found_proxies) != 0: + tracer = found_proxies[0].tracer + return tracer.create_proxy("call_function", cls, args, kwargs) + else: + cls.__init__(instance, *args, **kwargs) # type: ignore[misc] + return instance + + +def _patch_function(fn: FunctionType, nargs: int) -> FunctionType: + co = fn.__code__ + co_flags = co.co_flags & ~HAS_VARSTUFF + co_args: tuple + if hasattr(co, "co_qualname"): + # Python-3.11+ code signature + co_args = ( + nargs, + 0, + 0, + co.co_nlocals, + co.co_stacksize, + co_flags, + co.co_code, + co.co_consts, + co.co_names, + co.co_varnames, + co.co_filename, + co.co_name, + co.co_qualname, # type: ignore[attr-defined] + co.co_firstlineno, + co.co_lnotab, + co.co_exceptiontable, # type: ignore[attr-defined] + co.co_freevars, + co.co_cellvars, + ) + elif hasattr(co, "co_posonlyargcount"): + co_args = ( + nargs, + 0, + 0, + co.co_nlocals, + co.co_stacksize, + co_flags, + co.co_code, + co.co_consts, + co.co_names, + co.co_varnames, + co.co_filename, + co.co_name, + co.co_firstlineno, + co.co_lnotab, + co.co_freevars, + co.co_cellvars, + ) + else: + co_args = ( + nargs, + 0, + co.co_nlocals, + co.co_stacksize, + co_flags, + co.co_code, + co.co_consts, + co.co_names, + co.co_varnames, + co.co_filename, + co.co_name, + co.co_firstlineno, + co.co_lnotab, + co.co_freevars, + co.co_cellvars, + ) + new_code = CodeType(*co_args) # type: ignore[arg-type] + return FunctionType( + new_code, fn.__globals__, fn.__name__, fn.__defaults__, fn.__closure__ + ) + + # we need to insert placeholder nodes for *args and **kwargs + # we can't call this function normally, otherwise it would try to unpack them + # instead, let's make python think that args and kwargs are normal variables + + +@compatibility(is_backward_compatible=False) +class PHBase: + """ + Object representing an input placeholder to `concrete_args` + """ + + def __repr__(self): + return "PH" + + +PH = PHBase() + + +@compatibility(is_backward_compatible=False) +class PHWithMeta(PHBase): + """ + Object representing an input placeholder to `concrete_args` + """ + + def __init__(self, ph_key: Optional[str] = None): + super().__init__() + + # Provide a hey for user to identify placeholder node during analysis + self.ph_key = ph_key + + +def _transfer_attrs(fr, to): + for attr_name in dir(fr): + attr_val = getattr(fr, attr_name) + if ( + not callable(attr_val) + and not attr_name.startswith("__") + and not hasattr(to, attr_name) + ): + setattr(to, attr_name, attr_val) + + +@compatibility(is_backward_compatible=True) +class Tracer(TracerBase): + # Reference: https://github.com/pytorch/pytorch/issues/54354 + # The first line of this docstring overrides the one Sphinx generates for the + # documentation. We need it so that Sphinx doesn't leak `math`s path from the + # build environment (e.g. ` None: + # This method's signature is overridden by the first line of this class' + # docstring. If this method's signature is modified, the signature that + # overrides it also should be modified accordingly. + + """ + Construct a Tracer object. + + Args: + + autowrap_modules (Tuple[ModuleType]): defaults to `(math, )`, + Python modules whose functions should be wrapped automatically + without needing to use fx.wrap(). Backward-compatibility for + this parameter is guaranteed. + + autowrap_functions (Tuple[Callable, ...]): defaults to `()`, + Python functions that should be wrapped automatically without + needing to use fx.wrap(). Backward compatibility for this + parameter is guaranteed. + + param_shapes_constant (bool): When this flag is set, calls to shape, + size and a few other shape like attributes of a module's parameter + will be evaluated directly, rather than returning a new Proxy value + for an attribute access. Backward compatibility for this parameter + is guaranteed. + """ + + super().__init__() + + # Functions we will eagerly wrap when we see them while tracing + # this captures both `math.sqrt()` and `from math import sqrt` automatically + self._autowrap_function_ids: Set[int] = { + id(value) + for name, value in chain(*[m.__dict__.items() for m in autowrap_modules]) + if not name.startswith("_") and callable(value) + } + self._autowrap_function_ids.update({id(f) for f in autowrap_functions}) + + # Python modules to apply autowrap to at the start, in addition to + # modules we see while tracing + self._autowrap_search: List[ModuleType] = list(autowrap_modules) + self.param_shapes_constant = param_shapes_constant + + self.submodule_paths: Optional[Dict[torch.nn.Module, str]] = None + self.root_module_name: str = "" + # Maps the containing module's name to the operator name + self.scope = Scope("", None) + # Records the module call stack + self.module_stack = collections.OrderedDict() + self.num_calls: Dict[str, int] = {} + # Mapping of node name to module scope + self.node_name_to_scope: Dict[str, Tuple[str, type]] = {} + + _qualname_counter: Dict[str, int] = collections.defaultdict(int) + + @compatibility(is_backward_compatible=True) + def get_fresh_qualname(self, prefix: str) -> str: + """ + Gets a fresh name for a prefix and returns it. This function ensures + that it will not clash with an existing attribute on the graph. + """ + # The idea here is that if the module doesn't have this prefix at all we + # should reset the counter to start from the beginning + # It's a ... little bit hacky (doesn't cover all cases) but the precise + # naming of the prefixes isn't a correctness issue, just a niceness + # issue + qualname = f"{prefix}0" + if not hasattr(self.root, qualname): + self._qualname_counter[prefix] = 0 + return qualname + + i = self._qualname_counter[prefix] + while True: + qualname = f"{prefix}{i}" + i += 1 + if not hasattr(self.root, qualname): + break + self._qualname_counter[prefix] = i + + return qualname + + @compatibility(is_backward_compatible=True) + def create_arg(self, a: Any) -> "Argument": + """ + A method to specify the behavior of tracing when preparing values to + be used as arguments to nodes in the ``Graph``. + + By default, the behavior includes: + + #. Iterate through collection types (e.g. tuple, list, dict) and recursively + call ``create_args`` on the elements. + #. Given a Proxy object, return a reference to the underlying IR ``Node`` + #. Given a non-Proxy Tensor object, emit IR for various cases: + + * For a Parameter, emit a ``get_attr`` node referring to that Parameter + * For a non-Parameter Tensor, store the Tensor away in a special + attribute referring to that attribute. + + This method can be overridden to support more types. + + Args: + + a (Any): The value to be emitted as an ``Argument`` in the ``Graph``. + + + Returns: + + The value ``a`` converted into the appropriate ``Argument`` + """ + # The base tracer is used to construct Graphs when there is no associated + # module hierarchy, so it can never create parameter references. + # The default tracer adds the ability to refer to parameters when + # tracing modules. + if isinstance(a, torch.nn.Parameter): + for n, p in self.root.named_parameters(): + if a is p: + return self.create_node("get_attr", n, (), {}) + raise NameError("parameter is not a member of this module") + elif isinstance(a, torch.Tensor): + for n_, p_ in self.root.named_buffers(): + if a is p_: + return self.create_node("get_attr", n_, (), {}) + elif isinstance(a, torch.nn.Module): + for n_, p_ in self.root.named_modules(): + if a is p_: + return self.create_node("get_attr", n_, (), {}) + # For NamedTuple instances that appear literally as args, we emit + # a node to construct the NamedTuple and use that Node as the argument. + if isinstance(a, tuple) and hasattr(a, "_fields"): + args = tuple(self.create_arg(elem) for elem in a) + return self.create_node("call_function", a.__class__, args, {}) + + # Tensors do not have a reliable string repr() from which they can be + # constructed (and we probably don't want to rely on that, either), so + # for any constant Tensor values we encounter, first search for if they + # are an attribute of some module in the module hierarchy. If so, emit + # a get_attr to retrieve that tensor. Otherwise, we'll store away the + # tensor value into a special attribute on the Module s.t. we can + # retrieve it with a get_attr. + if isinstance(a, (torch.Tensor, ScriptObject, FakeScriptObject)): + qualname: Optional[str] = self.tensor_attrs.get(a) + + # Tensor was not found in the Module hierarchy, stow it away in a + # special attribute and set the qualname to refer to that + if not qualname: + base_name = ( + "_tensor_constant" + if isinstance(a, torch.Tensor) + else "_torchbind_obj" + ) + qualname = self.get_fresh_qualname(base_name) + assert isinstance(qualname, str) + self.tensor_attrs[a] = qualname + setattr(self.root, qualname, a) + + return self.create_node("get_attr", qualname, (), {}) + + if type(a) in _proxyable_classes: + # This is an instance of a proxyable class for which we did not + # witness its construction. Intern this as a constant attribute + + # TODO: binary search + qualname = self.get_fresh_qualname(f"_{a.__class__.__name__}_constant_") + assert isinstance(qualname, str) + setattr(self.root, qualname, a) + + return self.create_node("get_attr", qualname, (), {}) + + return super().create_arg(a) + + @compatibility(is_backward_compatible=True) + def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool: + """ + A method to specify whether a given ``nn.Module`` is a "leaf" module. + + Leaf modules are the atomic units that appear in + the IR, referenced by ``call_module`` calls. By default, + Modules in the PyTorch standard library namespace (torch.nn) + are leaf modules. All other modules are traced through and + their constituent ops are recorded, unless specified otherwise + via this parameter. + + Args: + + m (Module): The module being queried about + module_qualified_name (str): The path to root of this module. For example, + if you have a module hierarchy where submodule ``foo`` contains + submodule ``bar``, which contains submodule ``baz``, that module will + appear with the qualified name ``foo.bar.baz`` here. + """ + return ( + m.__module__.startswith("torch.nn") + or m.__module__.startswith("torch.ao.nn") + ) and not isinstance(m, torch.nn.Sequential) + + @compatibility(is_backward_compatible=True) + def path_of_module(self, mod: torch.nn.Module) -> str: + """ + Helper method to find the qualified name of ``mod`` in the Module hierarchy + of ``root``. For example, if ``root`` has a submodule named ``foo``, which has + a submodule named ``bar``, passing ``bar`` into this function will return + the string "foo.bar". + + Args: + + mod (str): The ``Module`` to retrieve the qualified name for. + """ + # Prefer the O(1) algorithm + if self.submodule_paths: + path = self.submodule_paths.get(mod) + if path is None: + raise NameError("module is not installed as a submodule") + assert isinstance(path, str) + return path + # O(N^2) fallback in the case that we didn't store the submodule + # paths. + else: + for n, p in self.root.named_modules(): + if mod is p: + return n + raise NameError("module is not installed as a submodule") + + @compatibility(is_backward_compatible=True) + def call_module( + self, + m: torch.nn.Module, + forward: Callable[..., Any], + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ) -> Any: + """ + Method that specifies the behavior of this ``Tracer`` when it encounters + a call to an ``nn.Module`` instance. + + By default, the behavior is to check if the called module is a leaf module + via ``is_leaf_module``. If it is, emit a ``call_module`` node referring to + ``m`` in the ``Graph``. Otherwise, call the ``Module`` normally, tracing through + the operations in its ``forward`` function. + + This method can be overridden to--for example--create nested traced + GraphModules, or any other behavior you would want while tracing across + ``Module`` boundaries. + + Args: + + m (Module): The module for which a call is being emitted + forward (Callable): The forward() method of the ``Module`` to be invoked + args (Tuple): args of the module callsite + kwargs (Dict): kwargs of the module callsite + + Return: + + The return value from the Module call. In the case that a ``call_module`` + node was emitted, this is a ``Proxy`` value. Otherwise, it is whatever + value was returned from the ``Module`` invocation. + """ + module_qualified_name = self.path_of_module(m) + with ScopeContextManager( + self.scope, Scope(module_qualified_name, type(m)) + ) as _scope: + # module_stack is an ordered dict so writing then deleting the + # entry is equivalent to push/pop on a list + num_calls = self.num_calls.get(module_qualified_name, 0) + module_key = ( + f"{_scope.module_path}@{num_calls}" + if num_calls > 0 + else _scope.module_path + ) + self.module_stack[module_key] = (module_qualified_name, _scope.module_type) + self.num_calls[module_qualified_name] = num_calls + 1 + if not self.is_leaf_module(m, module_qualified_name): + ret_val = forward(*args, **kwargs) + else: + ret_val = self.create_proxy( + "call_module", module_qualified_name, args, kwargs + ) + key, _ = self.module_stack.popitem(last=True) + assert key == module_key, f" Unexpected key {key}" + + return ret_val + + @compatibility(is_backward_compatible=False) + def getattr(self, attr: str, attr_val: Any, parameter_proxy_cache: Dict[str, Any]): + """ + Method that specifies the behavior of this ``Tracer`` when we call getattr + on a call to an ``nn.Module`` instance. + + By default, the behavior is to return a proxy value for the attribute. It + also stores the proxy value in the ``parameter_proxy_cache``, so that future + calls will reuse the proxy rather than creating a new one. + + This method can be overridden to --for example-- not return proxies when + querying parameters. + + Args: + + attr (str): The name of the attribute being queried + attr_val (Any): The value of the attribute + parameter_proxy_cache (Dict[str, Any]): A cache of attr names to proxies + + Return: + + The return value from the getattr call. + """ + + def maybe_get_proxy_for_attr( + attr_val, collection_to_search, parameter_proxy_cache + ): + for n, p in collection_to_search: + if attr_val is p: + if n not in parameter_proxy_cache: + kwargs = {} + if ( + "proxy_factory_fn" + in inspect.signature(self.create_proxy).parameters + ): + kwargs["proxy_factory_fn"] = ( + None + if not self.param_shapes_constant + else lambda node: ParameterProxy( + self, node, n, attr_val + ) + ) + val_proxy = self.create_proxy("get_attr", n, (), {}, **kwargs) # type: ignore[arg-type] + parameter_proxy_cache[n] = val_proxy + return parameter_proxy_cache[n] + return None + + if isinstance(attr_val, torch.nn.Parameter): + maybe_parameter_proxy = maybe_get_proxy_for_attr( + attr_val, self.root.named_parameters(), parameter_proxy_cache + ) + if maybe_parameter_proxy is not None: + return maybe_parameter_proxy + + if self.proxy_buffer_attributes and isinstance(attr_val, torch.Tensor): + maybe_buffer_proxy = maybe_get_proxy_for_attr( + attr_val, self.root.named_buffers(), parameter_proxy_cache + ) + if maybe_buffer_proxy is not None: + return maybe_buffer_proxy + + return attr_val + + # This method will be refactored + @compatibility(is_backward_compatible=False) + def create_args_for_root(self, root_fn, is_module, concrete_args=None): + """ + Create ``placeholder`` nodes corresponding to the signature of the ``root`` + Module. This method introspects root's signature and emits those + nodes accordingly, also supporting ``*args`` and ``**kwargs``. + """ + # In some cases, a function or method has been decorated with a wrapper + # defined via ``functools.wraps``. In this case, the outer code object + # will likely not contain the actual parameters we care about, so unwrap + # the function to get to the innermost callable. + fn_for_analysis = inspect.unwrap(root_fn) + co = fn_for_analysis.__code__ + total_args = co.co_argcount + co.co_kwonlyargcount + orig_args = list(co.co_varnames) + names_iter = iter(co.co_varnames) + args: List[Any] = [] + skip_arg_idx = 0 + if is_module: + if total_args == 0: + raise RuntimeError( + "``self`` argument cannot be part of *args expansion!" + ) + skip_arg_idx = 1 + next(names_iter) # skip self + args.append(self.root) + + sig = inspect.signature(fn_for_analysis) + + # This covers the very specific case where we are passing in flat + # concrete_args as a tuple, but our traced fn takes (*args, **kwargs). + # In this case, just take the concrete_args and pass them through. + name_idx = 0 + if ( + isinstance(concrete_args, tuple) + and len(concrete_args) > 0 + and (co.co_flags & HAS_VARSTUFF) + and total_args == 1 + ): + for concrete_arg in concrete_args: + out = self.create_proxy("placeholder", f"input_{name_idx}", (), {}) + if isinstance(concrete_arg, PHBase): + if concrete_arg != PH: + # Transfer attrs in the case where you're using a placeholder other + # than the singleton PH (PH has no attributes to transfer). + # Proxies were created out of the placeholders. + # Transfer any metadata (put on the placeholders in the form of + # attributes set by the user) from the placeholder to the + # underlying nodes (the proxy is unwrapped by the user, but + # the metadata should hold). + _transfer_attrs(fr=concrete_arg, to=out.node) + args.append(out) + name_idx += 1 + return root_fn, args + + arg_names = [next(names_iter) for idx in range(skip_arg_idx, total_args)] + if isinstance(concrete_args, tuple): + if len(arg_names) != len(concrete_args): + raise RuntimeError( + f"Tracing expected {len(arg_names)} arguments but got {len(concrete_args)} concrete arguments" + ) + concrete_args = dict(zip(arg_names, concrete_args)) + + def proxy_placeholder(name): + return self._proxy_placeholder(name, concrete_args, sig, fn_for_analysis) + + args.extend(proxy_placeholder(names) for names in arg_names) + + if co.co_kwonlyargcount > 0 or co.co_flags & HAS_VARSTUFF: + # TODO: type annotations for *args and **kwargs + if co.co_flags & inspect.CO_VARARGS: + args.append(proxy_placeholder("*" + next(names_iter))) + if co.co_flags & inspect.CO_VARKEYWORDS: + args.append(proxy_placeholder("**" + next(names_iter))) + root_fn = _patch_function(root_fn, len(args)) + + flat_args, in_spec = pytree.tree_flatten(tuple(args)) + if not all(child.is_leaf() for child in in_spec.children_specs): + # In the case that we have pytree-flattened inputs in + # `concrete_args`, generate a flattening wrapper around the + # original root function and return that. + self.graph._codegen = _PyTreeCodeGen( + _PyTreeInfo(orig_args[:total_args], in_spec, None) + ) + + def flatten_fn(*args): + tree_args = pytree.tree_unflatten(list(args), in_spec) + tree_out = root_fn(*tree_args) + out_args, out_spec = pytree.tree_flatten(tree_out) + assert isinstance(self.graph._codegen, _PyTreeCodeGen) + self.graph._codegen.pytree_info = ( + self.graph._codegen.pytree_info._replace(out_spec=out_spec) + ) + return out_args + + return flatten_fn, flat_args + return root_fn, args + + @compatibility(is_backward_compatible=True) + def trace( + self, + root: Union[torch.nn.Module, Callable[..., Any]], + concrete_args: Optional[Dict[str, Any]] = None, + ) -> Graph: + """ + Trace ``root`` and return the corresponding FX ``Graph`` representation. ``root`` + can either be an ``nn.Module`` instance or a Python callable. + + Note that after this call, ``self.root`` may be different from the ``root`` passed + in here. For example, when a free function is passed to ``trace()``, we will + create an ``nn.Module`` instance to use as the root and add embedded constants + to. + + + Args: + + root (Union[Module, Callable]): Either a ``Module`` or a function to be + traced through. Backwards-compatibility for this parameter is + guaranteed. + concrete_args (Optional[Dict[str, any]]): Concrete arguments that should + not be treated as Proxies. This parameter is experimental and + its backwards-compatibility is *NOT* guaranteed. + + Returns: + + A ``Graph`` representing the semantics of the passed-in ``root``. + """ + global _is_fx_tracing_flag + old_is_fx_tracing_flag = _is_fx_tracing_flag + _is_fx_tracing_flag = True + try: + if isinstance(root, torch.nn.Module): + # do real recompilation for _LazyGraphModule before retracing since the trace + # method can not trace the _lazy_forward method. Got error: + # https://gist.github.com/shunting314/75549c2e82ae07ac1139c94a3583d259 + # without this. + from torch.fx._lazy_graph_module import _LazyGraphModule + + _LazyGraphModule.force_recompile(root) + + self.root = root + + assert hasattr( + type(root), self.traced_func_name + ), f"traced_func_name={self.traced_func_name} doesn't exist in {type(root).__name__}" + + fn = getattr(type(root), self.traced_func_name) + self.root_module_name = root._get_name() + self.submodule_paths = {mod: name for name, mod in root.named_modules()} + else: + self.root = torch.nn.Module() + fn = root + + tracer_cls: Optional[Type[Tracer]] = getattr(self, "__class__", None) + self.graph = Graph(tracer_cls=tracer_cls) + if hasattr(fn, "__code__"): + code = fn.__code__ + self.graph._co_fields = { + "co_name": code.co_name, + "co_filename": code.co_filename, + "co_firstlineno": code.co_firstlineno, + } + + # When we encounter a Tensor value that's not a parameter, we look if it + # is some other attribute on the model. Construct a dict mapping Tensor + # values to the qualified name here for efficiency. This is used downstream + # in create_arg + self.tensor_attrs: Dict[ + Union[torch.Tensor, ScriptObject, FakeScriptObject], str + ] = {} + + def collect_tensor_attrs(m: torch.nn.Module, prefix_atoms: List[str]): + for k, v in m.__dict__.items(): + if isinstance(v, (torch.Tensor, ScriptObject, FakeScriptObject)): + self.tensor_attrs[v] = ".".join(prefix_atoms + [k]) + for k, v in m.named_children(): + collect_tensor_attrs(v, prefix_atoms + [k]) + + collect_tensor_attrs(self.root, []) + + assert isinstance(fn, FunctionType) + + fn_globals = fn.__globals__ # run before it gets patched + fn, args = self.create_args_for_root( + fn, isinstance(root, torch.nn.Module), concrete_args + ) + + parameter_proxy_cache: Dict[ + str, Proxy + ] = {} # Reduce number of get_attr calls + + # Method dispatch on parameters is not recorded unless it's directly used. + # Thus, we need to insert a proxy when __getattr__ requests a parameter. + @functools.wraps(_orig_module_getattr) + def module_getattr_wrapper(mod, attr): + attr_val = _orig_module_getattr(mod, attr) + return self.getattr(attr, attr_val, parameter_proxy_cache) + + @functools.wraps(_orig_module_call) + def module_call_wrapper(mod, *args, **kwargs): + def forward(*args, **kwargs): + return _orig_module_call(mod, *args, **kwargs) + + _autowrap_check( + patcher, # type: ignore[has-type] + getattr(getattr(mod, "forward", mod), "__globals__", {}), + self._autowrap_function_ids, + ) + return self.call_module(mod, forward, args, kwargs) + + with _new_patcher() as patcher: + # allow duplicate patches to support the case of nested calls + patcher.patch_method( + torch.nn.Module, + "__getattr__", + module_getattr_wrapper, + deduplicate=False, + ) + patcher.patch_method( + torch.nn.Module, "__call__", module_call_wrapper, deduplicate=False + ) + _patch_wrapped_functions(patcher) + _autowrap_check(patcher, fn_globals, self._autowrap_function_ids) + for module in self._autowrap_search: + _autowrap_check( + patcher, module.__dict__, self._autowrap_function_ids + ) + self.create_node( + "output", + "output", + (self.create_arg(fn(*args)),), + {}, + type_expr=fn.__annotations__.get("return", None), + ) + + self.submodule_paths = None + finally: + _is_fx_tracing_flag = old_is_fx_tracing_flag + return self.graph + + def __deepcopy__(self, memo): + # _autowrap_search contains modules, which cannot be deepcopied. + new_tracer = Tracer.__new__(Tracer) + + for k, v in self.__dict__.items(): + if k in {"_autowrap_search"}: + new_obj = copy.copy(v) + else: + new_obj = copy.deepcopy(v, memo) + + new_tracer.__dict__[k] = new_obj + + return new_tracer + + def _proxy_placeholder(self, name, concrete_args, sig, fn_for_analysis): + if concrete_args is not None and name in concrete_args: + cnt = 0 + + def replace_ph(x): + nonlocal cnt + cnt += 1 + param = sig.parameters[name] + default: Tuple[Any, ...] = ( + () if param.default is inspect.Parameter.empty else (param.default,) + ) + out = self.create_proxy( + "placeholder", f"{name}_{str(cnt)}", default, {} + ) + if isinstance(x, PHBase): + if x != PH: + # Transfer attrs in the case where you're using a placeholder other + # than the singleton PH (PH has no attributes to transfer). + # Proxies were created out of the placeholders. + # Transfer any metadata (put on the placeholders in the form of + # attributes set by the user) from the placeholder to the + # underlying nodes (the proxy is unwrapped by the user, but + # the metadata should hold). + _transfer_attrs(fr=x, to=out.node) + + return out + # Union[int, bool] == bool in Python <= 3.6 + if type(x) == bool or type(x) in base_types and type(x) != torch.Tensor: + torch._assert( + out == x, + f"{name} has been specialized to have value {x} but got another value", + ) + elif x is None: + args = ( + out, + f"{name} has been specialized to have value None but got another value", + ) + self.create_proxy("call_function", _assert_is_none, args, {}) + else: + warnings.warn( + f"Was not able to add assertion to guarantee correct input {name} to " + f"specialized function. It is up to the user to make sure that your inputs match the " + f"inputs you specialized the function with." + ) + + return x + + return pytree.tree_map(replace_ph, concrete_args[name]) + if name[0] == "*": + default: Tuple[Any, ...] = () + else: + param = sig.parameters[name] + default = ( # type: ignore[assignment] + () if param.default is inspect.Parameter.empty else (param.default,) + ) + return self.create_proxy( + "placeholder", + name, + default, + {}, + type_expr=fn_for_analysis.__annotations__.get(name, None), + ) + + +# Dictionary of (id(globals dict), function name) => globals_dict to patch for +# the purposes of the wrap() API. +# We key by the globals dict id and function name to ensure we're wrapping a given +# function only once. +_wrapped_fns_to_patch: Dict[Tuple[int, str], dict] = {} + +# List of methods on classes to wrap (class type, function name) +# this currently only works for Tensor.* methods that aren't traced properly +_wrapped_methods_to_patch: List[Tuple[type, str]] = [] + +if os.environ.get("FX_PATCH_GETITEM") == "1": + # This change is needed to trace models like PositionalEmbedding from BERT: + # https://github.com/pytorch/benchmark/blob/master/torchbenchmark/models/BERT_pytorch/bert_pytorch/model/embedding/position.py + # but causes issues in quantization documented here: + # https://github.com/pytorch/pytorch/issues/50710 + # once that is fixed we can make this the default behavior. + _wrapped_methods_to_patch.append((torch.Tensor, "__getitem__")) + + +def _find_proxy(*objects_to_search): + """ + Recursively search a data structure for a Proxy() and return it, + return None if not found. + """ + proxy = None + + def find_proxy(x): + nonlocal proxy + if isinstance(x, Proxy): + proxy = x + + map_aggregate(objects_to_search, find_proxy) + return proxy + + +def _create_wrapped_func(orig_fn): + @functools.wraps(orig_fn) + def wrapped(*args, **kwargs): + """ + Given an closed-over ``orig_function`` to invoke, search the args and kwargs for + a Proxy object. If there is one, emit a ``call_function`` node to preserve the + call to this leaf function directly. Otherwise, just return the results of + this function call, as this function is not being traced. + """ + proxy = _find_proxy(args, kwargs) + if proxy is not None: + return_proxy = proxy.tracer.create_proxy( + "call_function", orig_fn, args, kwargs + ) + return_proxy.node.meta["is_wrapped"] = True + return return_proxy + return orig_fn(*args, **kwargs) + + return wrapped + + +def _create_wrapped_method(cls, name): + orig_fn = getattr(cls, name) + + @functools.wraps(orig_fn) + def wrapped(*args, **kwargs): + """ + Search the args and kwargs for a Proxy object. If there is one, + emit a ``call_method`` node to preserve the call to this method + directly. Otherwise, just return the results of this function + call, as this function is not being traced. + """ + proxy = _find_proxy(args, kwargs) + if proxy is not None: + return proxy.tracer.create_proxy("call_method", name, args, kwargs) + return orig_fn(*args, **kwargs) + + return wrapped + + +class _PatchedFn(NamedTuple): + frame_dict: Any + fn_name: str + orig_fn: Any + new_fn: Any + + def revert(self): + raise NotImplementedError + + def patch(self): + raise NotImplementedError + + +class _PatchedFnSetItem(_PatchedFn): + def revert(self): + self.frame_dict[self.fn_name] = self.orig_fn + + def patch(self): + self.frame_dict[self.fn_name] = self.new_fn + + +class _PatchedFnDel(_PatchedFn): + def revert(self): + del self.frame_dict[self.fn_name] + + def patch(self): + self.frame_dict[self.fn_name] = self.new_fn + + +class _PatchedFnSetAttr(_PatchedFn): + def revert(self): + setattr(self.frame_dict, self.fn_name, self.orig_fn) + + def patch(self): + setattr(self.frame_dict, self.fn_name, self.new_fn) + + +class _Patcher: + def __init__(self) -> None: + super().__init__() + self.patches_made: List[_PatchedFn] = [] + self.visited: Set[int] = set() + + def patch( + self, + frame_dict: Dict[str, Any], + name: str, + new_fn: Callable, + deduplicate: bool = True, + ): + """ + Replace frame_dict[name] with new_fn until we exit the context manager. + """ + new_fn.__fx_already_patched = deduplicate # type: ignore[attr-defined] + if name not in frame_dict and hasattr(builtins, name): + self.patches_made.append(_PatchedFnDel(frame_dict, name, None, new_fn)) + self.patches_made[-1].patch() + elif getattr(frame_dict[name], "__fx_already_patched", False): + return # already patched, no need to do it again + else: + self.patches_made.append( + _PatchedFnSetItem(frame_dict, name, frame_dict[name], new_fn) + ) + self.patches_made[-1].patch() + + def patch_method( + self, cls: type, name: str, new_fn: Callable, deduplicate: bool = True + ): + """ + Replace object_or_dict.name with new_fn until we exit the context manager. + """ + new_fn.__fx_already_patched = deduplicate # type: ignore[attr-defined] + orig_fn = getattr(cls, name) + if getattr(orig_fn, "__fx_already_patched", False): + return # already patched, no need to do it again + self.patches_made.append(_PatchedFnSetAttr(cls, name, orig_fn, new_fn)) + self.patches_made[-1].patch() + + def visit_once(self, thing: Any): + """Return True on the first call to with thing, otherwise false""" + idx = id(thing) + if idx in self.visited: + return False + self.visited.add(idx) + return True + + def revert_all_patches(self): + """ + Remove all the stored patcheds. It doesn't modify patches_made. + """ + for patch in self.patches_made: + patch.revert() + return self.patches_made + + def reapply_all_patches(self): + """ + Patch all the stored patcheds. It doesn't modify patches_made. + """ + for patch in self.patches_made: + patch.patch() + return self.patches_made + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """ + Undo all the changes made via self.patch() and self.patch_method() + """ + while self.patches_made: + # unpatch in reverse order to handle duplicates correctly + self.patches_made.pop().revert() + self.visited.clear() + + +CURRENT_PATCHER: Optional[_Patcher] = None + + +@contextlib.contextmanager +def _new_patcher(): + global CURRENT_PATCHER + prior_patcher = CURRENT_PATCHER + try: + CURRENT_PATCHER = _Patcher() + yield CURRENT_PATCHER + finally: + # Clear all the patches made by when using current patcher. + assert CURRENT_PATCHER is not None + CURRENT_PATCHER.revert_all_patches() + CURRENT_PATCHER = prior_patcher + + +@contextlib.contextmanager +def _maybe_revert_all_patches(): + current_patcher = CURRENT_PATCHER + patches_made = None + patches_removed = None + try: + if current_patcher is not None: + patches_removed = current_patcher.revert_all_patches() + yield + finally: + if current_patcher is not None: + patches_made = current_patcher.reapply_all_patches() + assert ( + patches_made == patches_removed + ), "CURRENT_PATCHER was changed during a revert_all_patches" + + +def _patch_wrapped_functions(patcher: _Patcher): + """ + Go through ``_wrapped_fn_patch_table`` and, for each frame object, wrap + the listed global functions in the `_create_wrapped_func` wrapper. + """ + for (_, name), frame_dict in _wrapped_fns_to_patch.copy().items(): + if name not in frame_dict and hasattr(builtins, name): + orig_fn = getattr(builtins, name) + else: + orig_fn = frame_dict[name] + patcher.patch(frame_dict, name, _create_wrapped_func(orig_fn)) + + for cls, name in _wrapped_methods_to_patch: + patcher.patch_method(cls, name, _create_wrapped_method(cls, name)) + + +def _autowrap_check( + patcher: _Patcher, frame_dict: Dict[str, Any], function_ids: Set[int] +): + """ + Some methods, like `math.sqrt` are common enough we want to automatically wrap them as we see them. + This method searches a scope for them and patches them if found. + """ + if patcher.visit_once(frame_dict): + for name, value in frame_dict.items(): + if ( + not name.startswith("_") + and callable(value) + and id(value) in function_ids + ): + patcher.patch(frame_dict, name, _create_wrapped_func(value)) + + +@compatibility(is_backward_compatible=True) +def wrap(fn_or_name: Union[str, Callable]): + """ + This function can be called at module-level scope to register fn_or_name as a "leaf function". + A "leaf function" will be preserved as a CallFunction node in the FX trace instead of being + traced through:: + + # foo/bar/baz.py + def my_custom_function(x, y): + return x * x + y * y + + + torch.fx.wrap("my_custom_function") + + + def fn_to_be_traced(x, y): + # When symbolic tracing, the below call to my_custom_function will be inserted into + # the graph rather than tracing it. + return my_custom_function(x, y) + + This function can also equivalently be used as a decorator:: + + # foo/bar/baz.py + @torch.fx.wrap + def my_custom_function(x, y): + return x * x + y * y + + A wrapped function can be thought of a "leaf function", analogous to the concept of + "leaf modules", that is, they are functions that are left as calls in the FX trace + rather than traced through. + + Args: + + fn_or_name (Union[str, Callable]): The function or name of the global function to insert into the + graph when it's called + """ + if not callable(fn_or_name) and not isinstance(fn_or_name, str): + raise RuntimeError( + "Unsupported type for global function! Must be either a callable or " + "string name" + ) + + if callable(fn_or_name): + assert not isinstance(fn_or_name, str) # to make mypy happy + fn_name = fn_or_name.__name__ + else: + assert isinstance( + fn_or_name, str + ), "fn_or_name must be a global function or string name" + fn_name = fn_or_name + + currentframe = inspect.currentframe() + assert currentframe is not None + f = currentframe.f_back + assert f is not None + if f.f_code.co_name != "": + raise NotImplementedError("wrap must be called at the top level of a module") + + # consider implementing Callable version of this via _autowrap_function_ids / _autowrap_search + # semantics would be slightly different, but would add support `from x import wrapped_function` + _wrapped_fns_to_patch[(id(f.f_globals), fn_name)] = f.f_globals + return fn_or_name + + +@compatibility(is_backward_compatible=True) +def symbolic_trace( + root: Union[torch.nn.Module, Callable[..., Any]], + concrete_args: Optional[Dict[str, Any]] = None, +) -> GraphModule: + """ + Symbolic tracing API + + Given an ``nn.Module`` or function instance ``root``, this function will return a ``GraphModule`` + constructed by recording operations seen while tracing through ``root``. + + ``concrete_args`` allows you to partially specialize your function, whether it's to remove control flow or data structures. + + For example:: + + def f(a, b): + if b == True: + return a + else: + return a * 2 + + FX can typically not trace through this due to the presence of control + flow. However, we can use `concrete_args` to specialize on the value of + `b` to trace through this:: + + f = fx.symbolic_trace(f, concrete_args={"b": False}) + assert f(3, False) == 6 + + Note that although you can still pass in different values of `b`, they will be ignored. + + We can also use `concrete_args` to eliminate data-structure handling from + our function. This will use pytrees to flatten your input. To avoid + overspecializing, pass in `fx.PH` for values that shouldn't be + specialized. For example:: + + def f(x): + out = 0 + for v in x.values(): + out += v + return out + + + f = fx.symbolic_trace(f, concrete_args={"x": {"a": fx.PH, "b": fx.PH, "c": fx.PH}}) + assert f({"a": 1, "b": 2, "c": 4}) == 7 + + + Args: + root (Union[torch.nn.Module, Callable]): Module or function to be traced and converted + into a Graph representation. + concrete_args (Optional[Dict[str, any]]): Inputs to be partially specialized + + Returns: + GraphModule: a Module created from the recorded operations from ``root``. + """ + tracer = Tracer() + graph = tracer.trace(root, concrete_args) + name = ( + root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__ + ) + return _make_graph_module(tracer.root, graph, name) + + +@wrap +def _assert_is_none(value, msg): + assert value is None, msg diff --git a/lib/python3.12/site-packages/torch/fx/_utils.py b/lib/python3.12/site-packages/torch/fx/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fc62453d67c2e3e9b7efcaa771bbc2dd0dd49dca --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/_utils.py @@ -0,0 +1,67 @@ +# mypy: allow-untyped-defs +import sys +from typing import Dict, Optional + +import torch +from torch._logging import LazyString + + +def lazy_format_graph_code(name, gm, maybe_id=None, **kwargs): + """ + Returns a LazyString that formats the graph code. + """ + + def format_name(): + if maybe_id is not None: + return f"{name} {maybe_id}" + else: + return name + + if "print_output" not in kwargs: + kwargs["print_output"] = False + + if "colored" in kwargs: + try: + if not sys.stdout.isatty(): + kwargs["colored"] = False + except AttributeError: + kwargs["colored"] = False + + return LazyString( + lambda: _format_graph_code( + f"===== {format_name()} =====\n", + gm.forward.__code__.co_filename, + gm.print_readable(**kwargs), + ) + ) + + +def _format_graph_code(name, filename, graph_str): + """ + Returns a string that formats the graph code. + """ + return f"TRACED GRAPH\n {name} {filename} {graph_str}\n" + + +def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> Optional[Dict]: + """ + Returns the nn_module_stack of the first call_function node. + """ + for node in graph.nodes: + if node.op == "call_function" and "nn_module_stack" in node.meta: + return node.meta["nn_module_stack"] + return None + + +def get_node_context(node, num_nodes=2) -> str: + """ + Returns a string of the last num_nodes nodes in the graph. + """ + node_contexts = [] + cur = node + for _ in range(num_nodes): + node_contexts.append(cur.format_node()) + if cur.op == "root": + break + cur = cur.prev + return "\n".join(node_contexts[::-1]) diff --git a/lib/python3.12/site-packages/torch/fx/annotate.py b/lib/python3.12/site-packages/torch/fx/annotate.py new file mode 100644 index 0000000000000000000000000000000000000000..b3c5056066251df51542cd187652c5111749516f --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/annotate.py @@ -0,0 +1,36 @@ +# mypy: allow-untyped-defs +from torch.fx.proxy import Proxy + +from ._compatibility import compatibility + + +@compatibility(is_backward_compatible=False) +def annotate(val, type): + """ + Annotates a Proxy object with a given type. + + This function annotates a val with a given type if a type of the val is a torch.fx.Proxy object + Args: + val (object): An object to be annotated if its type is torch.fx.Proxy. + type (object): A type to be assigned to a given proxy object as val. + Returns: + The given val. + Raises: + RuntimeError: If a val already has a type in its node. + """ + if isinstance(val, Proxy): + if val.node.type: + raise RuntimeError( + f"Tried to annotate a value that already had a type on it!" + f" Existing type is {val.node.type} " + f"and new type is {type}. " + f"This could happen if you tried to annotate a function parameter " + f"value (in which case you should use the type slot " + f"on the function signature) or you called " + f"annotate on the same value twice" + ) + else: + val.node.type = type + return val + else: + return val diff --git a/lib/python3.12/site-packages/torch/fx/config.py b/lib/python3.12/site-packages/torch/fx/config.py new file mode 100644 index 0000000000000000000000000000000000000000..da5120d6edf180f7fbbe88ac342b4d0e4b383e50 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/config.py @@ -0,0 +1,6 @@ +# Whether to disable showing progress on compilation passes +# Need to add a new config otherwise wil get a circular import if dynamo config is imported here +disable_progress = True + +# If True this also shows the node names in each pass, for small models this is great but larger models it's quite noisy +verbose_progress = False diff --git a/lib/python3.12/site-packages/torch/fx/experimental/debug.py b/lib/python3.12/site-packages/torch/fx/experimental/debug.py new file mode 100644 index 0000000000000000000000000000000000000000..e59dcbb3296f9b740ae5de7f9c213458c45bef2b --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/experimental/debug.py @@ -0,0 +1,29 @@ +# mypy: allow-untyped-defs +import torch.fx as fx + + +def set_trace(gm: fx.GraphModule) -> fx.GraphModule: + """ + Sets a breakpoint in `gm`'s generated python code. It drops into pdb when + `gm` gets run. + + Args: + gm: graph module to insert breakpoint. It is then recompiled for it to + take effect. + + Returns: + the `gm` with breakpoint inserted. + """ + + def insert_pdb(body): + return ["import pdb; pdb.set_trace()\n", *body] + + with gm.graph.on_generate_code( + make_transformer=lambda cur_transform: ( + # new code transformer to register + lambda body: (insert_pdb(cur_transform(body) if cur_transform else body)) + ) + ): + gm.recompile() + + return gm diff --git a/lib/python3.12/site-packages/torch/fx/experimental/merge_matmul.py b/lib/python3.12/site-packages/torch/fx/experimental/merge_matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..b3e1efcbd19e4ea5abd76a24798846e43c0ea303 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/experimental/merge_matmul.py @@ -0,0 +1,178 @@ +# mypy: allow-untyped-defs +import itertools +import operator +from typing import Dict, List, Tuple + +import torch +from torch.fx._symbolic_trace import symbolic_trace +from torch.fx.node import Node +from torch.fx.passes.tools_common import legalize_graph + + +def split_result_tensors( + result: torch.Tensor, inputs: List[torch.Tensor] +) -> Tuple[torch.Tensor, ...]: + """ + A free function for use in the merge_matmul graph transformation below that + splits the output from a merged matmul into the individual results for each + input tensor. + + Arguments: + result: The merged matmul result tensor. + inputs: The list of inputs that were merged into one for the matmul. + + Returns: + List of matmul results for each input tensor. + """ + # When fx tracer is running, x.shape[0] will be torch.fx.Attribute but we + # need an int even when tracing + if isinstance(result, torch.fx.Proxy): + splits = [0] * len(inputs) + else: + splits = [x.shape[0] for x in inputs] + + return torch.split(result, splits) + + +def may_depend_on(a: Node, b: Node, search_depth: int = 6): + """ + Determine if one node depends on another in a torch.fx.Graph. + + Arguments: + a: The node that may have a dependency on b. + b: The node that a may have a dependency on. + search_depth: In the case of an indirect dependency, this function + searches upto this many nodes away in search of a + data dependency. If none is found, the function + makes the conservative assumption that there is a + dependency. + + Returns: + True if a may depend on b, False if it definitely does not. + """ + # Equivalence is defined as dependence. + if a == b: + return True + + # If a has no inputs, it cannot depend on b. + if len(a.all_input_nodes) == 0: + return False + + # If the search depth has been exhausted and no conclusion has been + # reached, assume that there is a data dependency. + if search_depth == 0: + return True + + # Recursively check all inputs of a. + for inp in a.all_input_nodes: + if may_depend_on(inp, b, search_depth - 1): + return True + + return False + + +def are_nodes_independent(nodes: List[Node]): + """ + Check if all of the given nodes are pairwise-data independent. + + Arguments: + nodes: The nodes to check for data dependencies. + + Returns: + True if any pair in nodes has a data dependency. + """ + # For each pair in nodes: + for i, j in itertools.combinations(nodes, 2): + if may_depend_on(i, j) or may_depend_on(j, i): + return False + + return True + + +def merge_matmul(in_mod: torch.nn.Module): + """ + A graph transformation that merges matrix multiplication operations that share the same right-hand + side operand into one large matrix multiplication. + ____ _________ _________ + ---- | | | | M| A * C | + M| A | T| B | * K| C | = |---------| + ---- , | | | | T| B * C | + K ---- --------- --------- + K R R + """ + gm = symbolic_trace(in_mod) + + rhs_users: Dict[Node, List[Node]] = {} + lhs_users: Dict[Node, List[Node]] = {} + + # Populate rhs_users and lhs_users - maps from LHS/RHS matrix multiply operands to + # the matmul of which they are the LHS/RHS. + for node in gm.graph.nodes: + if node.op != "call_function" or node.target is not torch.matmul: + continue + + lhs, rhs = node.args + + # TODO: Properly handle aliasing caused by get_attr. For now, + # use the attribute name as the operand if the node is a + # get_attr. + lhs = lhs.target if lhs.op == "get_attr" else lhs + rhs = rhs.target if rhs.op == "get_attr" else rhs + + lhs_users.setdefault(lhs, []).append(node) + rhs_users.setdefault(rhs, []).append(node) + + for rhs, mms in rhs_users.items(): + # There must be at least matmuls for a merge to make sense. + if len(mms) < 2: + continue + + # All matmuls must not depend on each other directly or indirectly + # in order for the merge to be possible. + if not are_nodes_independent(mms): + continue + + lhs_vals = [mm.args[0] for mm in mms] + + # Merge the matmul. + # Collect a list of LHS operands and the single RHS operand. + lhs = [gm.graph.get_attr(l) if isinstance(l, str) else l for l in lhs_vals] + rhs = gm.graph.get_attr(rhs) if isinstance(rhs, str) else rhs + + # Concatenate all the LHS operands. + merge_mm_cat = gm.graph.call_function(torch.cat, (lhs,), {}) + + # Multiply the concatenated LHS operands with the one RHS. This will produce + # the same results as all the individual matmuls involving rhs in the original graph, + # but they will all be concatenated together. + merge_mm = gm.graph.call_function( + torch.matmul, + ( + merge_mm_cat, + rhs, + ), + {}, + ) + + # Split the result of the merged matmul using the shapes of the LHS operands + # to ascertain how large each chunk should be. + merge_mm_split = gm.graph.call_function( + split_result_tensors, (merge_mm, lhs), {} + ) + merge_mm_res = [ + gm.graph.call_function(operator.getitem, (merge_mm_split, out), {}) + for out in range(len(lhs)) + ] + + # Replace all uses of the original, unmerged matmuls with the equivalent split chunk from the merged matmul. + for old, new in zip(mms, merge_mm_res): + old.replace_all_uses_with(new) + gm.graph.erase_node(old) + + # All of the new nodes created above were inserted at the end, so we need to sort + # the nodes topologically to make sure all definitions precede uses. + legalize_graph(gm) + + gm.recompile() + gm.graph.lint() + return gm diff --git a/lib/python3.12/site-packages/torch/fx/graph.py b/lib/python3.12/site-packages/torch/fx/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..520eb373802695cb1a315f2e69eef539ebe96263 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/graph.py @@ -0,0 +1,2016 @@ +# mypy: allow-untyped-defs +import builtins +import contextlib +import copy +import enum +import functools +import inspect +import keyword +import math +import os +import re +import warnings +from collections import defaultdict +from contextlib import contextmanager +from dataclasses import dataclass +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Iterable, + List, + NamedTuple, + Optional, + Set, + Tuple, + Type, + TYPE_CHECKING, +) + +import torch +import torch.utils._pytree as pytree +from torch._C import _NodeIter + +from . import _pytree as fx_pytree +from ._compatibility import compatibility +from .node import _get_qualified_name, _type_repr, Argument, map_arg, Node, Target + + +__all__ = ["PythonCode", "CodeGen", "Graph"] + +if TYPE_CHECKING: + from ._symbolic_trace import Tracer # noqa: F401 + from .graph_module import GraphModule # noqa: F401 + + +# Mapping of builtins to their `typing` equivalent. +_origin_type_map = { + list: List, + dict: Dict, + set: Set, + frozenset: FrozenSet, + tuple: Tuple, +} + +_legal_ops = dict.fromkeys( + ["call_function", "call_method", "get_attr", "call_module", "placeholder", "output"] +) + + +# Signature for functions thattransforms the body (`list[str]`) of the +# generated code +TransformCodeFunc = Callable[[List[str]], List[str]] + + +class _CustomBuiltin(NamedTuple): + """Additional objs that we add to every graph's globals. + + The repr() for some standard library objects is not valid Python code without + an import. For common objects of this sort, we bundle them in the globals of + every FX graph. + """ + + # How to import this object from the standard library. + import_str: str + # The actual object, produced from that import string. + obj: Any + + +_custom_builtins: Dict[str, _CustomBuiltin] = {} + + +def _register_custom_builtin(name: str, import_str: str, obj: Any): + _custom_builtins[name] = _CustomBuiltin(import_str, obj) + + +_register_custom_builtin("inf", "from math import inf", math.inf) +_register_custom_builtin("nan", "from math import nan", math.nan) +_register_custom_builtin("NoneType", "NoneType = type(None)", type(None)) +_register_custom_builtin("torch", "import torch", torch) +_register_custom_builtin("device", "from torch import device", torch.device) +_register_custom_builtin("fx_pytree", "import torch.fx._pytree as fx_pytree", fx_pytree) +_register_custom_builtin("pytree", "import torch.utils._pytree as pytree", pytree) + + +def _is_magic(x: str) -> bool: + return x.startswith("__") and x.endswith("__") + + +def _snake_case(s: str) -> str: + """ + Transforms the given string ``s`` to a Python-style variable name + + Examples: + ``mod.snake_case`` -> ``mod.snake_case`` + ``mod.pascalCase``-> ``mod.pascal_case`` + ``mod.ALL_CAPS`` -> ``mod.all_caps`` + """ + return _snake_case_sub(s).lower() + + +# Replace occurrences where a lowercase letter is followed by an uppercase letter +_snake_case_sub = functools.partial(re.compile(r"(?<=[a-z])([A-Z])").sub, r"_\1") + + +def _is_from_torch(obj: Any) -> bool: + module_name = getattr(obj, "__module__", None) + if module_name is not None: + base_module = module_name.partition(".")[0] + return ( + base_module == "torch" + and not module_name.startswith("torch._dynamo.") + and not module_name.startswith("torch._inductor.") + ) + + name = getattr(obj, "__name__", None) + # exclude torch because torch.torch.torch.torch works. idk mang + if name is not None and name != "torch": + for guess in [torch, torch.nn.functional]: + if getattr(guess, name, None) is obj: + return True + + return False + + +class _Namespace: + """A context for associating names uniquely with objects. + + The following invariants are enforced: + - Each object gets a single name. + - Each name is unique within a given namespace. + - Names generated do not shadow builtins, unless the object is indeed that builtin. + """ + + def __init__(self): + self._obj_to_name: Dict[Any, str] = {} + self._unassociated_names = set() + self._used_names: Set[str] = set() + self._base_count: Dict[str, int] = defaultdict(int) + + self._illegal_char_regex = re.compile("[^0-9a-zA-Z_]+") + self._name_suffix_regex = re.compile(r"(.*)_(\d+)$") + + def create_name(self, candidate: str, obj: Optional[Any]) -> str: + """Create a unique name. + + Arguments: + candidate: used as the basis for the unique name, relevant to the user. + obj: If not None, an object that will be associated with the unique name. + """ + if obj is not None and obj in self._obj_to_name: + return self._obj_to_name[obj] + + # delete all characters that are illegal in a Python identifier + candidate = self._illegal_char_regex.sub("_", candidate) + + if not candidate: + candidate = "_unnamed" + + if candidate[0].isdigit(): + candidate = f"_{candidate}" + + match = self._name_suffix_regex.match(candidate) + if match is None: + base = candidate + num = None + else: + base, num_str = match.group(1, 2) + num = int(num_str) + + candidate = base if num is None else f"{base}_{num}" + if not num: + num = self._base_count[base] + + while candidate in self._used_names or self._is_illegal_name(candidate, obj): + num += 1 + candidate = f"{base}_{num}" + + self._used_names.add(candidate) + self._base_count[base] = num + if obj is None: + self._unassociated_names.add(candidate) + else: + self._obj_to_name[obj] = candidate + return candidate + + def associate_name_with_obj(self, name: str, obj: Any): + """Associate a unique name with an object. + + Neither `name` nor `obj` should be associated already. + """ + assert obj not in self._obj_to_name + assert name in self._unassociated_names + self._obj_to_name[obj] = name + self._unassociated_names.remove(name) + + def _is_illegal_name(self, name: str, obj: Any) -> bool: + # 1. keywords are never allowed as names. + if name in keyword.kwlist: + return True + + # 2. Can't shadow a builtin name, unless you *are* that builtin. + if name in builtins.__dict__: + return obj is not builtins.__dict__[name] + + # 3. Can't shadow our custom builtins either + if name in _custom_builtins: + return obj is not _custom_builtins[name].obj + + return False + + def _rename_object(self, obj: Any, name: str): + assert obj in self._obj_to_name + self._obj_to_name[obj] = name + self._used_names.add(name) + + +dtype_abbrs = { + torch.bfloat16: "bf16", + torch.float64: "f64", + torch.float32: "f32", + torch.float16: "f16", + torch.float8_e4m3fn: "f8e4m3fn", + torch.float8_e5m2: "f8e5m2", + torch.float8_e4m3fnuz: "f8e4m3fnuz", + torch.float8_e5m2fnuz: "f8e5m2fnuz", + torch.complex32: "c32", + torch.complex64: "c64", + torch.complex128: "c128", + torch.int8: "i8", + torch.int16: "i16", + torch.int32: "i32", + torch.int64: "i64", + torch.bool: "b8", + torch.uint8: "u8", + torch.uint16: "u16", + torch.uint32: "u32", + torch.uint64: "u64", + torch.bits16: "b16", +} + + +@compatibility(is_backward_compatible=True) +@dataclass +class PythonCode: + """ + Represents all the information necessary to exec or save a graph as Python code. + """ + + # Python source code for the forward function definition. + src: str + # Values in global scope during execution of `src_def`. + globals: Dict[str, Any] + # Optional mapping from the forward function's line number to + # node index. + _lineno_map: Optional[Dict[int, Optional[int]]] + + +def _format_target(base: str, target: str) -> str: + elems = target.split(".") + r = base + for e in elems: + if not e.isidentifier(): + r = f'getattr({r}, "{e}")' + else: + r = f"{r}.{e}" + return r + + +class _InsertPoint: + def __init__(self, graph, new_insert): + self.graph = graph + self.orig_insert, graph._insert = graph._insert, new_insert + + def __enter__(self): + pass + + def __exit__(self, type, value, tb): + self.graph._insert = self.orig_insert + + +class _node_list: + def __init__(self, graph: "Graph", direction: str = "_next"): + assert direction in ["_next", "_prev"] + self.graph = graph + self.direction = direction + + def __len__(self): + return self.graph._len + + def __iter__(self): + assert self.direction == "_prev" or self.direction == "_next" + yield from _NodeIter(self.graph._root, self.direction == "_prev") + + def __reversed__(self): + return _node_list(self.graph, "_next" if self.direction == "_prev" else "_prev") + + +class _PyTreeInfo(NamedTuple): + """ + Contains extra info stored when we're using Pytrees + """ + + orig_args: List[str] + in_spec: pytree.TreeSpec + out_spec: Optional[pytree.TreeSpec] + + +@dataclass(frozen=True) +class _ParsedStackTrace: + """ + Represents the top-most frame of a parsed stack trace + """ + + file: str + lineno: str + name: str + code: str + + def get_summary_str(self): + return f"File: {self.file}:{self.lineno} in {self.name}, code: {self.code}" + + +# get File:lineno code from stack_trace +def _parse_stack_trace(stack_trace: str): + if stack_trace is None: + return None + pattern = re.compile(r"^File \"(.+)\", line (\d+), in (.+)$") + lines = stack_trace.strip().split("\n") + # stacktrace should have innermost frame last, so we + # iterate backwards to find the first line that starts + # with 'File ' + for idx in range(len(lines) - 2, -1, -1): + line = lines[idx].strip() + matches = pattern.match(line) + if matches: + file = matches.group(1) + lineno = matches.group(2) + name = matches.group(3) + # next line should be the code + code = lines[idx + 1].strip() + return _ParsedStackTrace(file, lineno, name, code) + return None + + +@compatibility(is_backward_compatible=False) +class CodeGen: + def __init__(self): + self._body_transformer: Optional[TransformCodeFunc] = None + self._func_name: str = "forward" + + def gen_fn_def(self, free_vars: List[str], maybe_return_annotation: str) -> str: + """ + Given the free variables and a return annotation, generates the beginning of the FX function. + By default, `gen_fn_def(['a', 'b'], '') == 'def {self._func_name}(a, b):'` + """ + # If the original function didn't have self as its first argument, we + # would have added it. + if len(free_vars) == 0 or free_vars[0] != "self": + free_vars.insert(0, "self") + return ( + f"def {self._func_name}({', '.join(free_vars)}){maybe_return_annotation}:" + ) + + def generate_output(self, output_args: Argument) -> str: + """ + Given the output arguments, generates the return statement of the FX function. + Note: The returned statement should not be indented. + """ + return f"return {repr(output_args)}" + + def process_inputs(self, *args: Any) -> Any: + """ + Transforms the inputs so that the graph can take them as arguments, as + non-default codegen may result in the inputs to the function being + different from the inputs to the graph. + + If the graph was directly runnable, this invariant should hold true + `f.graph.process_outputs(f.graph(*f.graph.process_inputs(*inputs))) == f(*inputs)` + """ + return args + + def process_outputs(self, outputs: Any) -> Any: + """ + Transforms the outputs of the graph to be identical to the codegen. + + See ``process_inputs`` for more details. + """ + return outputs + + def additional_globals(self) -> List[Tuple[str, Any]]: + """ + If your codegen uses extra global values, add tuples of (identifier,reference to the value) here. + For example, return ['List', typing.List] if you need ``List`` in the global context. + """ + return [] + + def _gen_python_code( + self, + nodes, + root_module: str, + namespace: _Namespace, + *, + verbose: bool = False, + include_stride: bool = False, + include_device: bool = False, + colored: bool = False, + ) -> PythonCode: + free_vars: List[str] = [] + body: List[str] = [] + globals_: Dict[str, Any] = {} + wrapped_fns: Dict[str, None] = {} + + # Wrap string in list to pass by reference + maybe_return_annotation: List[str] = [""] + include_stride = include_stride or ( + os.environ.get("FX_GRAPH_SHOW_STRIDE", "0") == "1" + ) + include_device = include_device or ( + os.environ.get("FX_GRAPH_SHOW_DEVICE", "0") == "1" + ) + + def add_global(name_hint: str, obj: Any): + """Add an obj to be tracked as a global. + + We call this for names that reference objects external to the + Graph, like functions or types. + + Returns: the global name that should be used to reference 'obj' in generated source. + """ + if ( + _is_from_torch(obj) and obj != torch.device + ): # to support registering torch.device + # HACK: workaround for how torch custom ops are registered. We + # can't import them like normal modules so they must retain their + # fully qualified name. + return _get_qualified_name(obj) + + # normalize the name hint to get a proper identifier + global_name = namespace.create_name(name_hint, obj) + + if global_name in globals_: + assert globals_[global_name] is obj + return global_name + globals_[global_name] = obj + return global_name + + # Pre-fill the globals table with registered builtins. + for name, (_, obj) in _custom_builtins.items(): + add_global(name, obj) + + def type_repr(o: Any): + if o == (): + # Empty tuple is used for empty tuple type annotation Tuple[()] + return "()" + + typename = _type_repr(o) + + if hasattr(o, "__origin__"): + # This is a generic type, e.g. typing.List[torch.Tensor] + origin_type = _origin_type_map.get(o.__origin__, o.__origin__) + origin_typename = add_global(_type_repr(origin_type), origin_type) + + if hasattr(o, "__args__"): + # Assign global names for each of the inner type variables. + args = [type_repr(arg) for arg in o.__args__] + + if len(args) == 0: + # Bare type, such as `typing.Tuple` with no subscript + # This code-path used in Python < 3.9 + return origin_typename + + return f'{origin_typename}[{",".join(args)}]' + else: + # Bare type, such as `typing.Tuple` with no subscript + # This code-path used in Python 3.9+ + return origin_typename + + # Common case: this is a regular module name like 'foo.bar.baz' + return add_global(typename, o) + + codes = { + "yellow": "\033[33m", + "cyan": "\033[36m", + "green": "\033[32m", + "blue": "\033[34m", + "red": "\033[31m", + "dim": "\033[2m", + "dim_blue": "\033[2m\033[34m", + "dim_green": "\033[2m\033[32m", + "reset": "\033[0m", + } + + def make_wrapper_func(name): + def f(s): + if colored: + return f"{codes[name]}{s}{codes['reset']}" + return s + + return f + + yellow = make_wrapper_func("yellow") # noqa: F841 + cyan = make_wrapper_func("cyan") # noqa: F841 + red = make_wrapper_func("red") + green = make_wrapper_func("green") # noqa: F841 + dim_green = make_wrapper_func("dim_green") + dim = make_wrapper_func("dim") + dim_blue = make_wrapper_func("dim_blue") + blue = make_wrapper_func("blue") + + def _get_repr(arg: Any) -> str: + # Handle NamedTuples (if it has `_fields`) via add_global. + if isinstance(arg, tuple) and hasattr(arg, "_fields"): + qualified_name = _get_qualified_name(type(arg)) + global_name = add_global(qualified_name, type(arg)) + return f"{global_name}{repr(tuple(arg))}" + elif isinstance( + arg, (torch._ops.OpOverload, torch._ops.HigherOrderOperator) + ): + qualified_name = _get_qualified_name(arg) + global_name = add_global(qualified_name, arg) + return f"{global_name}" + elif isinstance(arg, enum.Enum): + cls = arg.__class__ + clsname = add_global(cls.__name__, cls) + return f"{clsname}.{arg.name}" + elif isinstance(arg, Node): + return repr(arg) + elif isinstance(arg, torch.Tensor): + size = list(arg.size()) + dtype = str(arg.dtype).split(".")[-1] + return f"torch.Tensor(size={size}, dtype={dtype})" + elif isinstance(arg, tuple): + if len(arg) == 1: + return f"({_get_repr(arg[0])},)" + else: + return "(" + ", ".join(_get_repr(a) for a in arg) + ")" + elif isinstance(arg, list): + return "[" + ", ".join(_get_repr(a) for a in arg) + "]" + elif isinstance(arg, slice): + return f"slice({_get_repr(arg.start)}, {_get_repr(arg.stop)}, {_get_repr(arg.step)})" + else: + return blue(repr(arg)) + + def _format_args( + args: Tuple[Argument, ...], kwargs: Dict[str, Argument] + ) -> str: + args_s = ", ".join(_get_repr(a) for a in args) + kwargs_s = ", ".join(f"{k} = {_get_repr(v)}" for k, v in kwargs.items()) + if args_s and kwargs_s: + return f"{args_s}, {kwargs_s}" + return args_s or kwargs_s + + # Run through reverse nodes and record the first instance of a use + # of a given node. This represents the *last* use of the node in the + # execution order of the program, which we will use to free unused + # values + node_to_last_use: Dict[Node, Node] = {} + user_to_last_uses: Dict[Node, List[Node]] = {} + + def register_last_uses(n: Node, user: Node): + if n not in node_to_last_use: + node_to_last_use[n] = user + user_to_last_uses.setdefault(user, []).append(n) + + for node in reversed(nodes): + map_arg(node.args, lambda n: register_last_uses(n, node)) + map_arg(node.kwargs, lambda n: register_last_uses(n, node)) + + def delete_unused_values(user: Node): + """ + Delete values after their last use. This ensures that values that are + not used in the remainder of the code are freed and the memory usage + of the code is optimal. + """ + if user.op == "placeholder": + return + if user.op == "output": + body.append("\n") + return + nodes_to_delete = user_to_last_uses.get(user, []) + + if len(user.users.keys()) == 0: + # This node is not used by any others. however it's also not + # removed by DCE since side-effect. We want to free it's outputs + # right after its execution done to save memory. + nodes_to_delete.append(user) + + if len(nodes_to_delete): + to_delete_str = " = ".join( + [repr(n) for n in nodes_to_delete] + ["None"] + ) + body.append(f"; {dim(to_delete_str)}\n") + else: + body.append("\n") + + prev_stacktrace = None + + def append_stacktrace_summary(node: Node): + """ + Append a summary of the stacktrace to the generated code. This is + useful for debugging. + """ + nonlocal prev_stacktrace + + if node.op not in {"placeholder", "output"}: + if node.stack_trace: + if node.stack_trace != prev_stacktrace: + prev_stacktrace = node.stack_trace + summary_str = "" + + if parsed_stack_trace := _parse_stack_trace(node.stack_trace): + summary_str = parsed_stack_trace.get_summary_str() + + body.append(f'\n {dim("# " + summary_str)}\n') + elif prev_stacktrace != "": + prev_stacktrace = "" + no_stacktrace_msg = "# No stacktrace found for following nodes" + body.append(f"\n{dim(no_stacktrace_msg)}\n") + + def stringify_shape(shape: Iterable) -> str: + return f"[{', '.join(str(x) for x in shape)}]" + + def emit_node(node: Node): + maybe_type_annotation = ( + "" if node.type is None else f" : {type_repr(node.type)}" + ) + + if verbose: + # override annotation with more detailed information + from torch.fx.experimental.proxy_tensor import py_sym_types + from torch.fx.passes.shape_prop import TensorMetadata + + meta_val = node.meta.get( + "val", + node.meta.get("tensor_meta", node.meta.get("example_value", None)), + ) + # use string as annotation, to make it valid python code + + if isinstance(meta_val, torch.Tensor): + stride_annotation = ( + f"{stringify_shape(meta_val.stride())}" + if include_stride + else "" + ) + device_annotation = f"{meta_val.device}" if include_device else "" + maybe_type_annotation = ( + f': "{red(dtype_abbrs[meta_val.dtype])}{blue(stringify_shape(meta_val.shape))}' + f'{dim_blue(stride_annotation)}{dim_green(device_annotation)}"' + ) + elif isinstance(meta_val, py_sym_types): + maybe_type_annotation = f': "Sym({meta_val})"' + elif isinstance(meta_val, TensorMetadata): + maybe_type_annotation = f': "{dtype_abbrs[meta_val.dtype]}{stringify_shape(meta_val.shape)}"' + + if node.op == "placeholder": + assert isinstance(node.target, str) + maybe_default_arg = ( + "" if not node.args else f" = {_get_repr(node.args[0])}" + ) + free_vars.append( + f"{node.target}{maybe_type_annotation}{maybe_default_arg}" + ) + raw_name = node.target.replace("*", "") + if raw_name != repr(node): + body.append(f"{repr(node)} = {raw_name}\n") + return + elif node.op == "call_method": + assert isinstance(node.target, str) + body.append( + f"{repr(node)}{maybe_type_annotation} = {_format_target(_get_repr(node.args[0]), node.target)}" + f"({_format_args(node.args[1:], node.kwargs)})" + ) + return + elif node.op == "call_function": + assert callable(node.target) + # pretty print operators + if ( + getattr(node.target, "__module__", "") == "_operator" + and node.target.__name__ in magic_methods + ): + assert isinstance(node.args, tuple) + body.append( + f"{repr(node)}{maybe_type_annotation} = " + f"{magic_methods[node.target.__name__].format(*(_get_repr(a) for a in node.args))}" + ) + return + + # pretty print inplace operators; required for jit.script to work properly + # not currently supported in normal FX graphs, but generated by torchdynamo + if ( + getattr(node.target, "__module__", "") == "_operator" + and node.target.__name__ in inplace_methods + ): + body.append( + f"{inplace_methods[node.target.__name__].format(*(_get_repr(a) for a in node.args))}; " + f"{repr(node)}{maybe_type_annotation} = {_get_repr(node.args[0])}" + ) + return + + qualified_name = _get_qualified_name(node.target) + global_name = add_global(qualified_name, node.target) + # special case for getattr: node.args could be 2-argument or 3-argument + # 2-argument: attribute access; 3-argument: fall through to attrib function call with default value + if ( + global_name == "getattr" + and isinstance(node.args, tuple) + and isinstance(node.args[1], str) + and node.args[1].isidentifier() + and len(node.args) == 2 + ): + body.append( + f"{repr(node)}{maybe_type_annotation} = {_format_target(_get_repr(node.args[0]), node.args[1])}" + ) + return + body.append( + f"{repr(node)}{maybe_type_annotation} = {global_name}({_format_args(node.args, node.kwargs)})" + ) + if node.meta.get("is_wrapped", False): + wrapped_fns.setdefault(global_name) + return + elif node.op == "call_module": + assert isinstance(node.target, str) + body.append( + f"{repr(node)}{maybe_type_annotation} = " + f"{_format_target(root_module, node.target)}({_format_args(node.args, node.kwargs)})" + ) + return + elif node.op == "get_attr": + assert isinstance(node.target, str) + body.append( + f"{repr(node)}{maybe_type_annotation} = {_format_target(root_module, node.target)}" + ) + return + elif node.op == "output": + if node.type is not None: + maybe_return_annotation[0] = f" -> {type_repr(node.type)}" + body.append(self.generate_output(node.args[0])) + return + raise NotImplementedError(f"node: {node.op} {node.target}") + + for i, node in enumerate(nodes): + # NOTE: emit_node does not emit a string with newline. It depends + # on delete_unused_values to append one + if verbose: + append_stacktrace_summary(node) + # emit a counter comment to keep track of + # node index, which will be deleted later + # after going through _body_transformer + body.append(f"# COUNTER: {i}\n") + emit_node(node) + delete_unused_values(node) + + if len(body) == 0: + # If the Graph has no non-placeholder nodes, no lines for the body + # have been emitted. To continue to have valid Python code, emit a + # single pass statement + body.append("pass\n") + + if len(wrapped_fns) > 0: + wrap_name = add_global("wrap", torch.fx.wrap) + wrap_stmts = "\n".join([f'{wrap_name}("{name}")' for name in wrapped_fns]) + else: + wrap_stmts = "" + + if self._body_transformer: + body = self._body_transformer(body) + + for name, value in self.additional_globals(): + add_global(name, value) + + prologue = self.gen_fn_def(free_vars, maybe_return_annotation[0]) + + # remove counter and generate lineno to node index mapping + lineno_map: Dict[int, Optional[int]] = {} + prologue_len = prologue.count("\n") + 1 + new_lines: List[str] = [] + cur_idx = None + for line in "".join(body).split("\n"): + counter = re.search(r"# COUNTER: (\d+)", line) + if counter and counter.group(1) is not None: + cur_idx = int(counter.group(1)) + else: + lineno_map[len(new_lines) + prologue_len] = cur_idx + new_lines.append(line) + + code = "\n".join(new_lines).lstrip("\n") + code = "\n".join(" " + line for line in code.split("\n")) + + fn_code = f""" +{wrap_stmts} + +{prologue} +{code}""" + return PythonCode(fn_code, globals_, _lineno_map=lineno_map) + + +# Ideally, we'd like to refactor all of the pytree logic into this codegen +# class. Unfortunately, there are 3 areas we currently need extra logic in FX. +# 1. In the initial symbolic trace, the pytree logic is tied up with `concrete_args`. +# 2. In the FX graph, we need to access 2 attributes - in_spec and out_spec. +# Since we can't access .graph within the FX forward, we need to copy the attribute to the module. +# 3. We currently can't register the pytree imports with `add_global` - not sure why. +class _PyTreeCodeGen(CodeGen): + def __init__(self, pytree_info: _PyTreeInfo): + super().__init__() + self.pytree_info: _PyTreeInfo = pytree_info + + def process_inputs(self, *inputs: Any) -> Any: + flat_args = pytree.arg_tree_leaves(*inputs) + return flat_args + + def process_outputs(self, out: Any) -> Any: + if self.pytree_info is None or self.pytree_info.out_spec is None: + return out + if not isinstance(out, (list, tuple)): + out = [out] + assert self.pytree_info.out_spec is not None + return pytree.tree_unflatten(out, self.pytree_info.out_spec) + + def gen_fn_def(self, free_vars, maybe_return_annotation): + # Given a user function/model: + # myargs = [myargs0, myargs1] + # mykwargs = {'mykwargs0': ..., 'mykwargs1': ...} + # def forward(self, mypos, *myargs, mykey=None, **mykwargs): + # + # The generated code flattens all keywords into positional arguments for `forward()` + # e.g forward(self, mypos, myargs0, myargs1, mykey, mykwargs0, mykwargs1): + # + # Within `forward`, `tree_flatten_spec``still parses args and kwargs separately + # e.g. tree_flatten_spec(([mypos, myargs0, myargs1], + # {'mykey':mykey, 'mykwargs0':mykwargs0, 'mykwargs1':mykwargs1}), + # self._in_spec) + # + # If the user function/model does not have keywords, the dict is suppressed from tree_flatten_spec + # e.g. tree_flatten_spec([mypos, myargs0, myargs1]), self._in_spec) + if self.pytree_info is None: + return super().gen_fn_def(free_vars, maybe_return_annotation) + + fn_args = self.pytree_info.orig_args + has_orig_self = (fn_args[0] == "self") if len(fn_args) > 0 else False + if has_orig_self: + free_vars.insert(0, "self") + fn_definition = super().gen_fn_def(fn_args[:], maybe_return_annotation) + + if len(free_vars) > 0: # pytree has placeholders in it + # when kwargs is present, in_spec is tuple(args, kwargs) + has_args_kwargs_tuple = ( + self.pytree_info.in_spec.type == tuple + and self.pytree_info.in_spec.num_children == 2 + and self.pytree_info.in_spec.children_specs[0].type == tuple + and self.pytree_info.in_spec.children_specs[1].type == dict + ) + fn_kwargs = "{}" + fn_signature = f"[{', '.join(fn_args)}], self._in_spec" + if has_args_kwargs_tuple: + count_args = self.pytree_info.in_spec.children_specs[0].num_children + fn_args = self.pytree_info.orig_args[:count_args] + fn_kwargs = ( + "{" + + ", ".join( + f"'{k}':{v}" + for k, v in zip( + self.pytree_info.in_spec.children_specs[1].context, + self.pytree_info.orig_args[count_args:], + ) + ) + + "}" + ) + fn_signature = f"([{', '.join(fn_args)}], {fn_kwargs}), self._in_spec" + + # in Python, `var1: annotation1, var2: annotation2 = function_call()` is invalid. + # we need to split it to two lines: + # one for annotation: `var1: annotation1; var2: annotation2;` (note the semicolon) + # one for code: `var1, var2, = function_call()` + without_annotation = [x.split(":")[0] for x in free_vars] + has_annotation = [x + "; " for x in free_vars if ":" in x] + if len(has_annotation) > 0: + fn_definition += "\n " + "".join(has_annotation) + "\n" + fn_definition += f""" + {', '.join(without_annotation)}, = fx_pytree.tree_flatten_spec({fn_signature})""" + return fn_definition + + def generate_output(self, output_args): + if self.pytree_info and self.pytree_info.out_spec: + return f"return pytree.tree_unflatten({repr(output_args)}, self._out_spec)" + else: + return super().generate_output(output_args) + + +class _FindNodesLookupTable: + """ + Side table for the graph for the purpose of doing fast queries + """ + + def __init__(self): + self.table: Dict[Tuple[str, Optional[Target]], Dict[Node, None]] = defaultdict( + dict + ) + + def _key(self, node) -> Tuple[str, Optional[Target]]: + return (node.op, node.target if node.op == "call_function" else None) + + def __contains__(self, node) -> bool: + return node in self.table[self._key(node)] + + def insert(self, node: Node) -> None: + self.table[self._key(node)][node] = None + + def remove(self, node: Node) -> None: + self.table[self._key(node)].pop(node) + + def find_nodes(self, *, op: str, target: Optional["Target"] = None): + if op == "call_function": + assert target is not None + return [*self.table[(op, target)].keys()] + + if target is None: + return [*self.table[(op, None)].keys()] + + # op is call_method, get_attr, call_module + return [node for node in self.table[(op, None)].keys() if node.target == target] + + +@compatibility(is_backward_compatible=True) +class Graph: + """ + ``Graph`` is the main data structure used in the FX Intermediate Representation. + It consists of a series of ``Node`` s, each representing callsites (or other + syntactic constructs). The list of ``Node`` s, taken together, constitute a + valid Python function. + + For example, the following code + + .. code-block:: python + + import torch + import torch.fx + + + class MyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.param = torch.nn.Parameter(torch.rand(3, 4)) + self.linear = torch.nn.Linear(4, 5) + + def forward(self, x): + return torch.topk( + torch.sum(self.linear(x + self.linear.weight).relu(), dim=-1), 3 + ) + + + m = MyModule() + gm = torch.fx.symbolic_trace(m) + + Will produce the following Graph:: + + print(gm.graph) + + .. code-block:: text + + graph(x): + %linear_weight : [num_users=1] = self.linear.weight + %add_1 : [num_users=1] = call_function[target=operator.add](args = (%x, %linear_weight), kwargs = {}) + %linear_1 : [num_users=1] = call_module[target=linear](args = (%add_1,), kwargs = {}) + %relu_1 : [num_users=1] = call_method[target=relu](args = (%linear_1,), kwargs = {}) + %sum_1 : [num_users=1] = call_function[target=torch.sum](args = (%relu_1,), kwargs = {dim: -1}) + %topk_1 : [num_users=1] = call_function[target=torch.topk](args = (%sum_1, 3), kwargs = {}) + return topk_1 + + For the semantics of operations represented in the ``Graph``, please see :class:`Node`. + """ + + @compatibility(is_backward_compatible=True) + def __init__( + self, + owning_module: Optional["GraphModule"] = None, + tracer_cls: Optional[Type["Tracer"]] = None, + tracer_extras: Optional[Dict[str, Any]] = None, + ): + """ + Construct an empty Graph. + """ + self._root: Node = Node(self, "", "root", "", (), {}) + self._used_names: Dict[str, int] = {} # base name -> number + self._insert = self._root.prepend + self._len = 0 + self._graph_namespace = _Namespace() + self._owning_module = owning_module + self._tracer_cls = tracer_cls + self._tracer_extras = tracer_extras + self._codegen = CodeGen() + self._co_fields: Dict[str, Any] = {} + self._find_nodes_lookup_table = _FindNodesLookupTable() + + @property + def owning_module(self): + return self._owning_module + + @owning_module.setter + def owning_module(self, mod: Optional["GraphModule"]): + self._owning_module = mod + + @property + def nodes(self) -> _node_list: + """ + Get the list of Nodes that constitute this Graph. + + Note that this ``Node`` list representation is a doubly-linked list. Mutations + during iteration (e.g. delete a Node, add a Node) are safe. + + Returns: + + A doubly-linked list of Nodes. Note that ``reversed`` can be called on + this list to switch iteration order. + """ + return _node_list(self) + + @compatibility(is_backward_compatible=False) + def output_node(self) -> Node: + output_node = next(iter(reversed(self.nodes))) + assert output_node.op == "output" + return output_node + + @compatibility(is_backward_compatible=False) + def find_nodes( + self, *, op: str, target: Optional["Target"] = None, sort: bool = True + ): + """ + Allows for fast query of nodes + + Args: + + op (str): the name of the operation + + target (Optional[Target]): the target of the node. For call_function, + the target is required. For other ops, the target is optional. + + sort (bool): whether to return nodes in the order they appear on + on the graph. + + Returns: + + Iteratable of nodes with the requested op and target. + """ + node_list = self._find_nodes_lookup_table.find_nodes(op=op, target=target) + if sort: + return sorted(node_list) + return node_list + + @compatibility(is_backward_compatible=True) + def graph_copy( + self, g: "Graph", val_map: Dict[Node, Node], return_output_node=False + ) -> "Optional[Argument]": + """ + Copy all nodes from a given graph into ``self``. + + Args: + + g (Graph): The source graph from which to copy Nodes. + + val_map (Dict[Node, Node]): a dictionary that will be populated with a mapping + from nodes in ``g`` to nodes in ``self``. Note that ``val_map`` can be passed + in with values in it already to override copying of certain values. + + Returns: + + The value in ``self`` that is now equivalent to the output value in ``g``, + if ``g`` had an ``output`` node. ``None`` otherwise. + """ + for node in g.nodes: + if node in val_map: + continue + if node.op == "output": + rv = map_arg(node.args[0], lambda n: val_map[n]) + return rv if not return_output_node else (rv, node) + val_map[node] = self.node_copy(node, lambda n: val_map[n]) + return None + + def __deepcopy__(self, memo=None) -> "Graph": + """ + Explicitly implement __deepcopy__ to prevent excessive recursion depth + from the default implementation. This uses graph_copy to copy the nodes + in an iterative way, rather than recursive. It also populates the + memoization table to prevent unnecessary copies (e.g. references to + nodes or other parts of the Graph from a custom GraphModule implementation. + """ + memo = memo if memo else {} + g = Graph(tracer_cls=self._tracer_cls) + output_vals = g.graph_copy(self, val_map=memo, return_output_node=True) + g._codegen = copy.deepcopy(self._codegen) + if output_vals is not None: + assert isinstance(output_vals, tuple) + output_val, old_output_node = output_vals + new_output_node = g.output( + output_val, type_expr=getattr(old_output_node, "type", None) + ) + new_output_node.meta = copy.copy(old_output_node.meta) + return g + + @compatibility(is_backward_compatible=True) + def create_node( + self, + op: str, + target: "Target", + args: Optional[Tuple["Argument", ...]] = None, + kwargs: Optional[Dict[str, "Argument"]] = None, + name: Optional[str] = None, + type_expr: Optional[Any] = None, + ) -> Node: + """ + Create a ``Node`` and add it to the ``Graph`` at the current insert-point. + Note that the current insert-point can be set via :meth:`Graph.inserting_before` + and :meth:`Graph.inserting_after`. + + Args: + op (str): the opcode for this Node. One of 'call_function', 'call_method', 'get_attr', + 'call_module', 'placeholder', or 'output'. The semantics of these opcodes are + described in the ``Graph`` docstring. + + args (Optional[Tuple[Argument, ...]]): is a tuple of arguments to this node. + + kwargs (Optional[Dict[str, Argument]]): the kwargs of this Node + + name (Optional[str]): an optional string name for the ``Node``. + This will influence the name of the value assigned to in the + Python generated code. + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. + + Returns: + + The newly-created and inserted node. + """ + assert op in _legal_ops + args = () if args is None else args + kwargs = {} if kwargs is None else kwargs + assert isinstance(args, tuple), "args must be a tuple" + assert isinstance(kwargs, dict), "kwargs must be a dict" + + candidate = name if name is not None else self._target_to_str(target) + name = self._graph_namespace.create_name(candidate, None) + n = Node(self, name, op, target, args, kwargs, type_expr) + + if ( + self.owning_module is not None + and getattr(self.owning_module, "_create_node_hooks", None) is not None + ): + for f in self.owning_module._create_node_hooks: + f(n) + + self._graph_namespace.associate_name_with_obj(name, n) + + self._insert(n) + self._find_nodes_lookup_table.insert(n) + self._len += 1 + return n + + @compatibility(is_backward_compatible=False) + def process_inputs(self, *args): + """ + Processes args so that they can be passed to the FX graph. + """ + return self._codegen.process_inputs(*args) + + @compatibility(is_backward_compatible=False) + def process_outputs(self, out): + return self._codegen.process_outputs(out) + + @compatibility(is_backward_compatible=True) + def erase_node(self, to_erase: Node) -> None: + """ + Erases a ``Node`` from the ``Graph``. Throws an exception if + there are still users of that node in the ``Graph``. + + Args: + + to_erase (Node): The ``Node`` to erase from the ``Graph``. + """ + if len(to_erase.users) > 0: + raise RuntimeError( + f"Tried to erase Node {to_erase} but it still had {len(to_erase.users)} " + f"users in the graph: {to_erase.users}!" + ) + if to_erase.graph != self: + raise RuntimeError(f"Attempting to remove {to_erase} from wrong graph!") + if to_erase._erased: + warnings.warn(f"erase_node({to_erase}) on an already erased node") + return + + if ( + self.owning_module is not None + and getattr(self.owning_module, "_erase_node_hooks", None) is not None + ): + for f in self.owning_module._erase_node_hooks: + f(to_erase) + + self._find_nodes_lookup_table.remove(to_erase) + to_erase._remove_from_list() + to_erase._erased = True # iterators may retain handles to erased nodes + self._len -= 1 + + # Null out this Node's argument nodes so that the Nodes referred to + # can update their ``users`` accordingly + new_args = map_arg(to_erase.args, lambda n: None) + assert isinstance(new_args, tuple) + to_erase.args = new_args + new_kwargs = map_arg(to_erase.kwargs, lambda n: None) + assert isinstance(new_kwargs, dict) + to_erase.kwargs = new_kwargs + + @compatibility(is_backward_compatible=True) + def inserting_before(self, n: Optional[Node] = None): + """Set the point at which create_node and companion methods will insert into the graph. + When used within a 'with' statement, this will temporary set the insert point and + then restore it when the with statement exits:: + + with g.inserting_before(n): + ... # inserting before node n + ... # insert point restored to what it was previously + g.inserting_before(n) # set the insert point permanently + + Args: + + n (Optional[Node]): The node before which to insert. If None this will insert before + the beginning of the entire graph. + + Returns: + A resource manager that will restore the insert point on ``__exit__``. + """ + if n is None: + return self.inserting_after(self._root) + assert n.graph == self, "Node to insert before is not in graph." + return _InsertPoint(self, n.prepend) + + @compatibility(is_backward_compatible=True) + def inserting_after(self, n: Optional[Node] = None): + """Set the point at which create_node and companion methods will insert into the graph. + When used within a 'with' statement, this will temporary set the insert point and + then restore it when the with statement exits:: + + with g.inserting_after(n): + ... # inserting after node n + ... # insert point restored to what it was previously + g.inserting_after(n) # set the insert point permanently + + Args: + + n (Optional[Node]): The node before which to insert. If None this will insert after + the beginning of the entire graph. + + Returns: + A resource manager that will restore the insert point on ``__exit__``. + """ + if n is None: + return self.inserting_before(self._root) + assert n.graph == self, "Node to insert after is not in graph." + return _InsertPoint(self, n.append) + + @compatibility(is_backward_compatible=True) + def placeholder( + self, + name: str, + type_expr: Optional[Any] = None, + default_value: Any = inspect.Signature.empty, + ) -> Node: + """ + Insert a ``placeholder`` node into the Graph. A ``placeholder`` represents + a function input. + + Args: + + name (str): A name for the input value. This corresponds to the name + of the positional argument to the function this ``Graph`` represents. + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. This is needed in some + cases for proper code generation (e.g. when the function is used + subsequently in TorchScript compilation). + + default_value (Any): The default value this function argument should take + on. NOTE: to allow for `None` as a default value, `inspect.Signature.empty` + should be passed as this argument to specify that the parameter does _not_ + have a default value. + + .. note:: + The same insertion point and type expression rules apply for this method + as ``Graph.create_node``. + """ + args = () if default_value is inspect.Signature.empty else (default_value,) + return self.create_node("placeholder", name, args=args, type_expr=type_expr) + + @compatibility(is_backward_compatible=True) + def get_attr(self, qualified_name: str, type_expr: Optional[Any] = None) -> Node: + """ + Insert a ``get_attr`` node into the Graph. A ``get_attr`` ``Node`` represents the + fetch of an attribute from the ``Module`` hierarchy. + + Args: + + qualified_name (str): the fully-qualified name of the attribute to be retrieved. + For example, if the traced Module has a submodule named ``foo``, which has a + submodule named ``bar``, which has an attribute named ``baz``, the qualified + name ``foo.bar.baz`` should be passed as ``qualified_name``. + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. + + + Returns: + + The newly-created and inserted ``get_attr`` node. + + .. note:: + The same insertion point and type expression rules apply for this method + as ``Graph.create_node``. + """ + + def _get_attr_reference_exists( + mod: torch.nn.Module, qualified_name: str + ) -> bool: + module_path, _, name = qualified_name.rpartition(".") + + try: + submod: torch.nn.Module = mod.get_submodule(module_path) + except AttributeError: + warnings.warn(f"Failed to fetch module {module_path}!") + return False + + if not hasattr(submod, name): + return False + + res = getattr(submod, name) + + if ( + not isinstance(res, torch.nn.Module) + and not isinstance(res, torch.nn.Parameter) + and name not in submod._buffers + ): + return False + + return True + + if self.owning_module and not _get_attr_reference_exists( + self.owning_module, qualified_name + ): + warnings.warn( + "Attempted to insert a get_attr Node with no " + "underlying reference in the owning " + "GraphModule! Call " + "GraphModule.add_submodule to add the " + "necessary submodule, " + "GraphModule.add_parameter to add the " + "necessary Parameter, or " + "nn.Module.register_buffer to add the " + "necessary buffer", + stacklevel=2, + ) + return self.create_node("get_attr", qualified_name, type_expr=type_expr) + + @compatibility(is_backward_compatible=True) + def call_module( + self, + module_name: str, + args: Optional[Tuple["Argument", ...]] = None, + kwargs: Optional[Dict[str, "Argument"]] = None, + type_expr: Optional[Any] = None, + ) -> Node: + """ + Insert a ``call_module`` ``Node`` into the ``Graph``. A ``call_module`` node + represents a call to the forward() function of a ``Module`` in the ``Module`` + hierarchy. + + Args: + + module_name (str): The qualified name of the ``Module`` in the ``Module`` + hierarchy to be called. For example, if the traced ``Module`` has a + submodule named ``foo``, which has a submodule named ``bar``, the + qualified name ``foo.bar`` should be passed as ``module_name`` to + call that module. + + args (Optional[Tuple[Argument, ...]]): The positional arguments to be passed + to the called method. Note that this should *not* include a ``self`` argument. + + kwargs (Optional[Dict[str, Argument]]): The keyword arguments to be passed + to the called method + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. + + Returns: + + The newly-created and inserted ``call_module`` node. + + .. note:: + The same insertion point and type expression rules apply for this method + as :meth:`Graph.create_node`. + """ + if self.owning_module and self.owning_module.get_submodule(module_name) is None: + warnings.warn( + "Attempted to insert a call_module Node with " + "no underlying reference in the owning " + "GraphModule! Call " + "GraphModule.add_submodule to add the " + "necessary submodule" + ) + return self.create_node( + "call_module", module_name, args, kwargs, type_expr=type_expr + ) + + @compatibility(is_backward_compatible=True) + def call_method( + self, + method_name: str, + args: Optional[Tuple["Argument", ...]] = None, + kwargs: Optional[Dict[str, "Argument"]] = None, + type_expr: Optional[Any] = None, + ) -> Node: + """ + Insert a ``call_method`` ``Node`` into the ``Graph``. A ``call_method`` node + represents a call to a given method on the 0th element of ``args``. + + Args: + + method_name (str): The name of the method to apply to the self argument. + For example, if args[0] is a ``Node`` representing a ``Tensor``, + then to call ``relu()`` on that ``Tensor``, pass ``relu`` to ``method_name``. + + args (Optional[Tuple[Argument, ...]]): The positional arguments to be passed + to the called method. Note that this *should* include a ``self`` argument. + + kwargs (Optional[Dict[str, Argument]]): The keyword arguments to be passed + to the called method + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. + + Returns: + + The newly created and inserted ``call_method`` node. + + .. note:: + The same insertion point and type expression rules apply for this method + as :meth:`Graph.create_node`. + """ + return self.create_node( + "call_method", method_name, args, kwargs, type_expr=type_expr + ) + + @compatibility(is_backward_compatible=True) + def call_function( + self, + the_function: Callable[..., Any], + args: Optional[Tuple["Argument", ...]] = None, + kwargs: Optional[Dict[str, "Argument"]] = None, + type_expr: Optional[Any] = None, + ) -> Node: + """ + Insert a ``call_function`` ``Node`` into the ``Graph``. A ``call_function`` node + represents a call to a Python callable, specified by ``the_function``. + + Args: + + the_function (Callable[..., Any]): The function to be called. Can be any PyTorch + operator, Python function, or member of the ``builtins`` or ``operator`` + namespaces. + + args (Optional[Tuple[Argument, ...]]): The positional arguments to be passed + to the called function. + + kwargs (Optional[Dict[str, Argument]]): The keyword arguments to be passed + to the called function + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. + + Returns: + + The newly created and inserted ``call_function`` node. + + .. note:: + The same insertion point and type expression rules apply for this method + as :meth:`Graph.create_node`. + """ + return self.create_node( + "call_function", the_function, args, kwargs, type_expr=type_expr + ) + + @compatibility(is_backward_compatible=True) + def node_copy( + self, node: Node, arg_transform: Callable[[Node], "Argument"] = lambda x: x + ) -> Node: + """ + Copy a node from one graph into another. ``arg_transform`` needs to transform arguments from + the graph of node to the graph of self. Example:: + + # Copying all the nodes in `g` into `new_graph` + g: torch.fx.Graph = ... + new_graph = torch.fx.graph() + value_remap = {} + for node in g.nodes: + value_remap[node] = new_graph.node_copy(node, lambda n: value_remap[n]) + + Args: + + node (Node): The node to copy into ``self``. + + arg_transform (Callable[[Node], Argument]): A function that transforms + ``Node`` arguments in node's ``args`` and ``kwargs`` into the + equivalent argument in ``self``. In the simplest case, this should + retrieve a value out of a table mapping Nodes in the original + graph to ``self``. + """ + args = map_arg(node.args, arg_transform) + kwargs = map_arg(node.kwargs, arg_transform) + assert isinstance(args, tuple) + assert isinstance(kwargs, dict) + result_node = self.create_node( + node.op, node.target, args, kwargs, node.name, node.type + ) + result_node.meta = copy.copy(node.meta) + return result_node + + @compatibility(is_backward_compatible=True) + def output(self, result: "Argument", type_expr: Optional[Any] = None): + """ + Insert an ``output`` ``Node`` into the ``Graph``. An ``output`` node represents + a ``return`` statement in Python code. ``result`` is the value that should + be returned. + + Args: + + result (Argument): The value to be returned. + + type_expr (Optional[Any]): an optional type annotation representing the + Python type the output of this node will have. + + .. note:: + + The same insertion point and type expression rules apply for this method + as ``Graph.create_node``. + """ + return self.create_node( + op="output", target="output", args=(result,), type_expr=type_expr + ) + + def _target_to_str(self, target: Target) -> str: + if callable(target): + op = target.__name__ + else: + assert isinstance(target, str) + op = target + if _is_magic(op): + op = op[2:-2] + op = _snake_case(op) + return op + + @compatibility(is_backward_compatible=True) + def python_code( + self, + root_module: str, + *, + verbose: bool = False, + include_stride: bool = False, + include_device: bool = False, + colored: bool = False, + ) -> PythonCode: + """ + Turn this ``Graph`` into valid Python code. + + Args: + + root_module (str): The name of the root module on which to look-up + qualified name targets. This is usually 'self'. + + Returns: + + A PythonCode object, consisting of two fields: + src: the Python source code representing the object + globals: a dictionary of global names in `src` -> the objects that they reference. + """ + # NOTE: [Graph Namespaces] + # + # There are two types of symbols in generated Python source code: + # locals and globals. + # Locals are locally defined by the output of a node in the Graph. + # Globals are references to external objects, like functions or types. + # + # When generating Python code, we need to make sure to name things + # appropriately. In particular: + # - All names should be unique, to avoid weird shadowing bugs. + # - These names need to be consistent, e.g. a object should always be + # referenced by the same name. + # + # To do this, we create a new namespace just for this source. All names + # that get printed must come from this namespace. + # + # Why can't we re-use node.name? Because it was generated within the + # namespace `self._graph_namespace`. In order to provide uniqueness + # over both locals (node.name) *and* globals, we create a completely + # new namespace to put all identifiers in. + namespace = _Namespace() + + # Override Node's repr to generate a valid name within our namespace. + # Since repr() is designed to produce a valid Python expression, it + # makes sense to re-use it. This way, it's easy to print something like + # Tuple[Node, Node] by simply calling repr() on it. Node's __repr__ is + # implemented cooperatively to allow this. + def node_repr(n: Node): + return namespace.create_name(n.name, n) + + @contextmanager + def override_node_repr(graph: Graph): + orig_repr_fns = {} + for node in graph.nodes: + orig_repr_fns[node] = node._repr_fn + node._repr_fn = node_repr + try: + yield None + finally: + # restore the original repr functions + for node in graph.nodes: + node._repr_fn = orig_repr_fns[node] + + with override_node_repr(self): + return self._python_code( + root_module, + namespace, + verbose=verbose, + include_stride=include_stride, + include_device=include_device, + colored=colored, + ) + + def _python_code( + self, + root_module: str, + namespace: _Namespace, + *, + verbose: bool = False, + include_stride: bool = False, + include_device: bool = False, + colored: bool = False, + ) -> PythonCode: + return self._codegen._gen_python_code( + self.nodes, + root_module, + namespace, + verbose=verbose, + include_stride=include_stride, + include_device=include_device, + colored=colored, + ) + + def __str__(self) -> str: + """ + Return a human-readable (not machine-readable) string representation + of this Graph + """ + placeholder_names: List[str] = [] + # This is a one-element array just so ``format_node`` can modify the closed + # over value + maybe_return_typename: List[str] = [""] + + node_strs = [node.format_node(placeholder_names) for node in self.nodes] + param_str = ", ".join(placeholder_names) + s = f"graph({param_str}){maybe_return_typename[0]}:" + for node_str in node_strs: + if node_str: + s += "\n " + node_str + return s + + @compatibility(is_backward_compatible=True) + def print_tabular(self): + """ + Prints the intermediate representation of the graph in tabular + format. Note that this API requires the ``tabulate`` module to be + installed. + """ + try: + from tabulate import tabulate + except ImportError: + print( + "`print_tabular` relies on the library `tabulate`, " + "which could not be found on this machine. Run `pip " + "install tabulate` to install the library." + ) + raise + + node_specs = [[n.op, n.name, n.target, n.args, n.kwargs] for n in self.nodes] + print( + tabulate(node_specs, headers=["opcode", "name", "target", "args", "kwargs"]) + ) + + @compatibility(is_backward_compatible=True) + def lint(self): + """ + Runs various checks on this Graph to make sure it is well-formed. In + particular: + - Checks Nodes have correct ownership (owned by this graph) + - Checks Nodes appear in topological order + - If this Graph has an owning GraphModule, checks that targets + exist in that GraphModule + """ + + # Check topo order + def check_arg(arg: Node, n: Optional[Node] = None) -> None: + context_str = f" of Node '{n}' " if n else " " + if arg.graph is not self: + raise RuntimeError( + f"Argument '{arg}'{context_str}does not belong to this Graph, " + f"but was used as an argument! If you are copying nodes from another graph, make " + f"sure to use ``arg_transform`` on node_copy() to remap values\n{self}" + ) + if arg not in seen_values: + raise RuntimeError( + f"Argument '{arg}'{context_str}was used before it has been " + f"defined! Please check that Nodes in the graph are topologically ordered\n{self}" + ) + + seen_names: Set[str] = set() + seen_values: Set[Node] = set() + for node in self.nodes: + if node.op not in [ + "placeholder", + "call_method", + "call_module", + "call_function", + "get_attr", + "output", + ]: + raise RuntimeError(f"Node {node} had unknown opcode {node.op}!") + if node.graph is not self: + raise RuntimeError(f"Node '{node}' does not belong to this Graph!") + if node not in self._find_nodes_lookup_table: + raise RuntimeError(f"Node '{node}' is not added to the side table") + map_arg(node.args, lambda arg: check_arg(arg, node)) + map_arg(node.kwargs, lambda arg: check_arg(arg, node)) + seen_values.add(node) + + if node.name in seen_names: + raise RuntimeError(f"Node redefined name {node.name}!") + seen_names.add(node.name) + + # Check targets are legit + if self.owning_module: + num_warnings = 0 + MAX_WARNINGS = 5 + for node in self.nodes: + if node.op == "call_function": + if not callable(node.target): + raise ValueError( + f"Node {node} target {node.target} has type {torch.typename(node.target)} but " + "a Callable is expected" + ) + else: + if not isinstance(node.target, str): + raise ValueError( + f"Node {node} target {node.target} has type {torch.typename(node.target)} but " + "a str is expected" + ) + if node.op in ["get_attr", "call_module"]: + target_atoms = node.target.split(".") + m_itr = self.owning_module + for i, atom in enumerate(target_atoms): + new_m_itr = getattr(m_itr, atom, None) + seen_qualname = ".".join(target_atoms[:i]) + if new_m_itr is None: + raise RuntimeError( + f"Node {node} target {node.target} references nonexistent attribute " + f"{atom} of {seen_qualname}" + ) + if node.op == "call_module" and not isinstance( + new_m_itr, torch.nn.Module + ): + raise RuntimeError( + f"Node {node} target {node.target} {atom} of {seen_qualname} does " + "not reference an nn.Module" + ) + elif ( + node.op == "get_attr" + and not isinstance(new_m_itr, torch.nn.Module) + and not isinstance(new_m_itr, torch.nn.Parameter) + and atom not in m_itr._buffers + ): + if num_warnings < MAX_WARNINGS: + # Don't emit this warning too frequently, + # for very large graphs this can become very expensive + # from a performance perspective. + warnings.warn( + f"Node {node} target {node.target} {atom} of {seen_qualname} does " + "not reference an nn.Module, nn.Parameter, or buffer, which is " + "what 'get_attr' Nodes typically target" + ) + num_warnings += 1 + else: + m_itr = new_m_itr + if num_warnings > MAX_WARNINGS: + warnings.warn( + f"Additional {num_warnings - MAX_WARNINGS} warnings " + "suppressed about get_attr references" + ) + + @compatibility(is_backward_compatible=True) + def eliminate_dead_code( + self, is_impure_node: Optional[Callable[[Node], bool]] = None + ) -> bool: + """ + Remove all dead code from the graph, based on each node's number of + users, and whether the nodes have any side effects. The graph must be + topologically sorted before calling. + + Args: + is_impure_node (Optional[Callable[[Node], bool]]): A function that returns + whether a node is impure. If this is None, then the default behavior is to + use Node.is_impure. + + Returns: + bool: Whether the graph was changed as a result of the pass. + + Example: + + Before dead code is eliminated, `a` from `a = x + 1` below has no users + and thus can be eliminated from the graph without having an effect. + + .. code-block:: python + + def forward(self, x): + a = x + 1 + return x + self.attr_1 + + After dead code is eliminated, `a = x + 1` has been removed, and the rest + of `forward` remains. + + .. code-block:: python + + def forward(self, x): + return x + self.attr_1 + + .. warning:: + + Dead code elimination has some heuristics to avoid removing + side-effectful nodes (see Node.is_impure) but in general coverage + is very bad, so you should assume that this method is not sound + to call unless you know that your FX graph consists entirely + of functional operations or you supply your own custom + function for detecting side-effectful nodes. + """ + # Lint the graph first to make sure its topologically sorted, otherwise + # DCE below will not behave as expected. + self.lint() + + def has_side_effect(node): + if is_impure_node is not None: + return is_impure_node(node) + return node.is_impure() + + # Reverse iterate so that when we remove a node, any nodes used as an + # input to that node have an updated user count that no longer reflects + # the removed node. + changed = False + for node in reversed(self.nodes): + if not has_side_effect(node) and len(node.users) == 0: + self.erase_node(node) + changed = True + + return changed + + @compatibility(is_backward_compatible=False) + def set_codegen(self, codegen: CodeGen): + self._codegen = codegen + + @compatibility(is_backward_compatible=False) + def on_generate_code( + self, + make_transformer: Callable[[Optional[TransformCodeFunc]], TransformCodeFunc], + ): + """Register a transformer function when python code is generated + + Args: + make_transformer (Callable[[Optional[TransformCodeFunc]], TransformCodeFunc]): + a function that returns a code transformer to be registered. + This function is called by `on_generate_code` to obtain the + code transformer. + + This function is also given as its input the currently + registered code transformer (or None if nothing is registered), + in case it is not desirable to overwrite it. This is useful to + chain code transformers together. + + Returns: + a context manager that when used in a `with` statement, to automatically + restore the previously registered code transformer. + + Example: + + .. code-block:: python + + + gm: fx.GraphModule = ... + + + # This is a code transformer we want to register. This code + # transformer prepends a pdb import and trace statement at the very + # beginning of the generated torch.fx code to allow for manual + # debugging with the PDB library. + def insert_pdb(body): + return ["import pdb; pdb.set_trace()\\n", *body] + + + # Registers `insert_pdb`, and overwrites the current registered + # code transformer (given by `_` to the lambda): + gm.graph.on_generate_code(lambda _: insert_pdb) + + # Or alternatively, registers a code transformer which first + # runs `body` through existing registered transformer, then + # through `insert_pdb`: + gm.graph.on_generate_code( + lambda current_trans: ( + lambda body: insert_pdb(current_trans(body) if current_trans else body) + ) + ) + + gm.recompile() + gm(*inputs) # drops into pdb + + + This function can also be used as a context manager, with the benefit to + automatically restores the previously registered code transformer: + + .. code-block:: python + + # ... continue from previous example + + with gm.graph.on_generate_code(lambda _: insert_pdb): + # do more stuff with `gm`... + gm.recompile() + gm(*inputs) # drops into pdb + + # now previous code transformer is restored (but `gm`'s code with pdb + # remains - that means you can run `gm` with pdb here too, until you + # run next `recompile()`). + """ + on_gen_code_old = self._codegen._body_transformer + self._codegen._body_transformer = make_transformer(on_gen_code_old) + + @contextlib.contextmanager + def on_generate_code_context_manager(): + try: + yield + finally: + self._codegen._body_transformer = on_gen_code_old + + return on_generate_code_context_manager() + + +reflectable_magic_methods = { + "add": "{} + {}", + "sub": "{} - {}", + "mul": "{} * {}", + "floordiv": "{} // {}", + "truediv": "{} / {}", + "div": "{} / {}", + "mod": "{} % {}", + "pow": "{} ** {}", + "lshift": "{} << {}", + "rshift": "{} >> {}", + "and_": "{} & {}", + "or_": "{} | {}", + "xor": "{} ^ {}", + "getitem": "{}[{}]", + "matmul": "{} @ {}", +} + +magic_methods = dict( + { + "eq": "{} == {}", + "ne": "{} != {}", + "lt": "{} < {}", + "gt": "{} > {}", + "le": "{} <= {}", + "ge": "{} >= {}", + "pos": "+{}", + "neg": "-{}", + "invert": "~{}", + }, + **reflectable_magic_methods, +) + +inplace_methods = { + "iadd": "{} += {}", + "iand": "{} &= {}", + "ifloordiv": "{} //= {}", + "ilshift": "{} <<= {}", + "imod": "{} %= {}", + "imul": "{} *= {}", + "imatmul": "{} @= {}", + "ior": "{} |= {}", + "ipow": "{} **= {}", + "irshift": "{} >>= {}", + "isub": "{} -= {}", + "itruediv": "{} /= {}", + "ixor": "{} ^= {}", + "setitem": "{}[{}] = {}", +} diff --git a/lib/python3.12/site-packages/torch/fx/graph_module.py b/lib/python3.12/site-packages/torch/fx/graph_module.py new file mode 100644 index 0000000000000000000000000000000000000000..52c55586ab51f792bd3851da6af9423853bbd803 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/graph_module.py @@ -0,0 +1,1019 @@ +# mypy: allow-untyped-defs +import contextlib +import copy +import itertools +import linecache +import os +import sys +import traceback +import warnings +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Set, Type, Union + +import torch +import torch.nn as nn +import torch.overrides +from torch.nn.modules.module import _addindent +from torch.package import Importer, PackageExporter, PackageImporter, sys_importer + +from ._compatibility import compatibility +from .graph import _custom_builtins, _is_from_torch, _PyTreeCodeGen, Graph, PythonCode + + +__all__ = [ + "reduce_graph_module", + "reduce_package_graph_module", + "reduce_deploy_graph_module", + "GraphModule", +] + +_USER_PRESERVED_ATTRIBUTES_KEY = "_user_preserved_attributes" + + +# Normal exec loses the source code, however we can work with +# the linecache module to recover it. +# Using _exec_with_source will add it to our local cache +# and then tools like TorchScript will be able to get source info. +class _EvalCacheLoader: + def __init__(self): + self.eval_cache = {} + self.next_id = 0 + + def cache(self, src: str, globals: Dict[str, Any], co_fields=None): + """Store the source in a private cache, and add a lazy entry in linecache + that allows the source to be retrieved by 'filename'. + + Args: + src (str): The module source to cache + globals (dict): The module globals + + Returns: + str: The cache key (and dummy filename) generated for src. + """ + + key = self._get_key() + if co_fields: + key += f" from {co_fields['co_filename']}:{co_fields['co_firstlineno']} in {co_fields['co_name']}" + self.eval_cache[key] = src + + # Don't mutate globals so that this loader is only used + # to populate linecache, and doesn't interact with other modules + # that might check `__loader__` + globals_copy = globals.copy() + globals_copy["__file__"] = key + globals_copy["__name__"] = key + globals_copy["__loader__"] = self + linecache.lazycache(key, globals_copy) + + return key + + # Part of the loader protocol (PEP 302) + # linecache will use this method when trying to find source code + def get_source(self, module_name) -> Optional[str]: + if module_name in self.eval_cache: + return self.eval_cache[module_name] + return None + + def _get_key(self): + key = f".{self.next_id}" + self.next_id += 1 + return key + + +_loader = _EvalCacheLoader() + + +def _exec_with_source(src: str, globals: Dict[str, Any], co_fields=None): + key = _loader.cache(src, globals, co_fields) + exec(compile(src, key, "exec"), globals) + + +def _forward_from_src(src: str, globals: Dict[str, Any], co_fields=None): + return _method_from_src( + method_name="forward", src=src, globals=globals, co_fields=co_fields + ) + + +def _method_from_src( + method_name: str, src: str, globals: Dict[str, Any], co_fields=None +) -> Callable: + # avoid mutating the passed in dict + globals_copy = globals.copy() + _exec_with_source(src, globals_copy, co_fields) + fn = globals_copy[method_name] + del globals_copy[method_name] + return fn + + +def _format_import_statement(name: str, obj: Any, importer: Importer) -> str: + if name in _custom_builtins: + return _custom_builtins[name].import_str + if _is_from_torch(name): + return "import torch" + module_name, attr_name = importer.get_name(obj) + return f"from {module_name} import {attr_name} as {name}" + + +def _format_import_block(globals: Dict[str, Any], importer: Importer): + import_strs: Set[str] = { + _format_import_statement(name, obj, importer) for name, obj in globals.items() + } + # Sort the imports so we have a stable import block that allows us to + # hash the graph module and get a consistent key for use in a cache. + return "\n".join(sorted(import_strs)) + + +@compatibility(is_backward_compatible=True) +def reduce_graph_module(body: Dict[Any, Any], import_block: str) -> torch.nn.Module: + # BC: attribute name was changed from `code` to `_code` to facilitate + # making `code` into a property and adding a docstring to it + fn_src = body.get("_code") or body["code"] + forward = _forward_from_src(import_block + fn_src, {}) + return _deserialize_graph_module(forward, body) + + +@compatibility(is_backward_compatible=True) +def reduce_package_graph_module( + importer: PackageImporter, body: Dict[Any, Any], generated_module_name: str +) -> torch.nn.Module: + forward = importer.import_module(generated_module_name).forward + return _deserialize_graph_module(forward, body) + + +@compatibility(is_backward_compatible=True) +def reduce_deploy_graph_module( + importer: PackageImporter, body: Dict[Any, Any], import_block: str +) -> torch.nn.Module: + ns = {} + ns["__builtins__"] = importer.patched_builtins + fn_src = body.get("_code") + assert fn_src is not None + forward = _forward_from_src(import_block + fn_src, ns) + return _deserialize_graph_module(forward, body) + + +# We create a dummy class here because symbolic_trace pulls the forward() +# function off of the class, rather than the instance. This class is used +# in _deserialize_graph_module() below. +class _CodeOnlyModule(torch.nn.Module): + def __init__(self, body): + super().__init__() + self.__dict__ = body + + +def _deserialize_graph_module( + forward, body: Dict[Any, Any], graph_module_cls=None +) -> torch.nn.Module: + """ + Deserialize a GraphModule given the dictionary of the original module, + using the code to reconstruct the graph. We delete the actual graph before + saving the dictionary so that changes to the in-memory graph format do not + get serialized. + """ + + # Try to retrieve the forward source in a backward-compatible way + _CodeOnlyModule.forward = forward + + tracer_cls = body.get("_tracer_cls") + if tracer_cls is None: + from ._symbolic_trace import Tracer + + tracer_cls = Tracer + + graphmodule_cls_name = body.get("_graphmodule_cls_name", "GraphModule") + + # This is a workaround for a mypy linter issue related to + # passing base class as an argument - https://github.com/python/mypy/issues/5865. + cls_tracer: Any = tracer_cls + + class KeepModules(cls_tracer): + # we shouldn't trace into any of the submodules, + # because they were not traced in the original GraphModule + def is_leaf_module(self, _: torch.nn.Module, __: str) -> bool: + return True + + com = _CodeOnlyModule(body) + + tracer_extras = body.get("_tracer_extras", {}) + graph = KeepModules().trace(com, **tracer_extras) + + # Manually set Tracer class on the reconstructed Graph, to avoid + # referencing the private local subclass KeepModules. + graph._tracer_cls = tracer_cls + from ._lazy_graph_module import _make_graph_module + + gm = _make_graph_module( + com, graph, class_name=graphmodule_cls_name, graph_module_cls=graph_module_cls + ) + + # The GraphModule constructor only retains attributes referenced by the graph. + # In this case, our goal is return a GraphModule as close to identical as the one + # put into the package. If any additional attributes were present in body, + # we should keep them. + for k, v in body.items(): + if not hasattr(gm, k): + setattr(gm, k, v) + return gm + + +# copy an attribute value with qualified name 'target' from 'from_module' to 'to_module' +# This installs empty Modules where none exist yet if they are subpaths of target +def _copy_attr(from_module: torch.nn.Module, to_module: torch.nn.Module, target: str): + *prefix, field = target.split(".") + for item in prefix: + f = getattr(from_module, item) + t = getattr(to_module, item, None) + if f is t: + # we have already installed one of its parents + # (e.g. target = root.linear.weight, but we have already installed root.linear) + # once we install a parent, we no longer need to copy the children + # since all the needed properties will already be present + return + + if t is None: + t = torch.nn.Module() + setattr(to_module, item, t) + from_module, to_module = f, t + + orig = getattr(from_module, field) + # If it is a tensor and not a parameter attribute of a module, it should be a named buffer. + # So, we register it as a named buffer in the target module. + if isinstance(orig, torch.Tensor) and not isinstance(orig, torch.nn.Parameter): + to_module.register_buffer(field, orig) + else: + setattr(to_module, field, orig) + + +# Assign attribute 'from_obj' to the qualified name 'target' on 'to_module +# This installs empty Modules where none exist yet if they are subpaths of target +def _assign_attr(from_obj: Any, to_module: torch.nn.Module, target: str): + *prefix, field = target.split(".") + for item in prefix: + t = getattr(to_module, item, None) + + if t is None: + t = torch.nn.Module() + setattr(to_module, item, t) + to_module = t + + # If it is a tensor and not a parameter attribute of a module, it should be a named buffer. + # So, we register it as a named buffer in the target module. + if isinstance(from_obj, torch.Tensor) and not isinstance( + from_obj, torch.nn.Parameter + ): + to_module.register_buffer(field, from_obj) + else: + setattr(to_module, field, from_obj) + + +# Recursively look up target from a graph module. +def _get_attr(model: torch.nn.Module, attr_name: str): + return _get_attr_via_attr_list(model, attr_name.split(".")) + + +def _get_attr_via_attr_list(model: torch.nn.Module, attr_list: List[str]): + if len(attr_list) == 0: + return model + *prefix, field = attr_list + t = model + for item in prefix: + t = getattr(t, item, None) # type: ignore[assignment] + assert t is not None + + return getattr(t, field) + + +def _has_attr(model: torch.nn.Module, attr_name: str): + *prefix, field = attr_name.split(".") + t = model + for item in prefix: + t = hasattr(t, item) # type: ignore[assignment] + if t is False: + return False + + return hasattr(t, field) + + +def _print_readable( + module, + module_name, + print_output=True, + include_stride=False, + include_device=False, + colored=False, +): + graph = module.graph + assert graph is not None and isinstance( + graph, torch.fx.Graph + ), "print_readable must be used on a module with a graph" + + verbose_python_code = graph.python_code( + root_module="self", + verbose=True, + include_stride=include_stride, + include_device=include_device, + colored=colored, + ) + module_code = verbose_python_code.src + module_code = module_code.lstrip("\n") + module_code = f"class {module_name}(torch.nn.Module):\n" + module_code + module_code = _addindent(module_code, 4) + + submodule_code_list = [""] + for submodule_name, submodule in module.named_children(): + if hasattr(submodule, "graph"): + submodule_code_list.append( + _print_readable( + submodule, + submodule_name, + print_output=False, + include_stride=include_stride, + include_device=include_device, + colored=colored, + ) + ) + submodule_code = "\n".join(submodule_code_list) + submodule_code = _addindent(submodule_code, 4) + + output = module_code + submodule_code + if print_output: + print(module_code + submodule_code) + return output + + +class _WrappedCall: + def __init__(self, cls, cls_call): + self.cls = cls + self.cls_call = cls_call + + # Previously, if an error occurred when valid + # symbolically-traced code was run with an invalid input, the + # user would see the source of the error as coming from + # `File "`, where N is some number. We use + # this function to generate a more informative error message. We + # return the traceback itself, a message explaining that the + # error occurred in a traced Module's generated forward + # function, and five lines of context surrounding the faulty + # line + @staticmethod + def _generate_error_message(frame_summary: traceback.FrameSummary) -> str: + # auxiliary variables (for readability) + err_lineno = frame_summary.lineno + assert err_lineno is not None + line = frame_summary.line + assert line is not None + err_line_len = len(line) + all_src_lines = linecache.getlines(frame_summary.filename) + + # constituent substrings of the error message + tb_repr = torch._dynamo.disable(traceback.format_exc)() + custom_msg = ( + "Call using an FX-traced Module, " + f"line {err_lineno} of the traced Module's " + "generated forward function:" + ) + before_err = "".join(all_src_lines[err_lineno - 2 : err_lineno]) + marker = "~" * err_line_len + "~~~ <--- HERE" + err_and_after_err = "\n".join(all_src_lines[err_lineno : err_lineno + 2]) + + # joined message + return "\n".join([tb_repr, custom_msg, before_err, marker, err_and_after_err]) + + def __call__(self, obj, *args, **kwargs): + try: + if self.cls_call is not None: + return self.cls_call(obj, *args, **kwargs) + else: + return super(self.cls, obj).__call__(*args, **kwargs) # type: ignore[misc] + except Exception as e: + assert e.__traceback__ + topmost_framesummary: traceback.FrameSummary = ( + traceback.StackSummary.extract(traceback.walk_tb(e.__traceback__))[-1] + ) + if "eval_with_key" in topmost_framesummary.filename: + print( + _WrappedCall._generate_error_message(topmost_framesummary), + file=sys.stderr, + ) + raise e.with_traceback(None) # noqa: B904 + else: + raise e + + +@compatibility(is_backward_compatible=True) +class GraphModule(torch.nn.Module): + """ + GraphModule is an nn.Module generated from an fx.Graph. Graphmodule has a + ``graph`` attribute, as well as ``code`` and ``forward`` attributes generated + from that ``graph``. + + .. warning:: + + When ``graph`` is reassigned, ``code`` and ``forward`` will be automatically + regenerated. However, if you edit the contents of the ``graph`` without reassigning + the ``graph`` attribute itself, you must call ``recompile()`` to update the generated + code. + """ + + def __new__(cls: "Type[GraphModule]", *args, **kwargs): + # each instance of a graph module needs its own forward method + # so create a new singleton class for each instance. + # it is a subclass of the user-defined class, the only difference + # is an extra layer to install the forward method + + # address issue described at https://github.com/pytorch/pytorch/issues/63883 + # in other words, traverse class hierarchy to fix the redundant class definition problem + for t in cls.__mro__: + c = t.__qualname__.split(".")[-1] + if c != "GraphModuleImpl": + cls = t + break + + class GraphModuleImpl(cls): # type: ignore[misc, valid-type] + pass + + return super().__new__(GraphModuleImpl) + + @compatibility(is_backward_compatible=True) + def __init__( + self, + root: Union[torch.nn.Module, Dict[str, Any]], + graph: Graph, + class_name: str = "GraphModule", + ): + """ + Construct a GraphModule. + + Args: + + root (Union[torch.nn.Module, Dict[str, Any]): + ``root`` can either be an nn.Module instance or a Dict mapping strings to any attribute type. + In the case that ``root`` is a Module, any references to Module-based objects (via qualified + name) in the Graph's Nodes' ``target`` field will be copied over from the respective place + within ``root``'s Module hierarchy into the GraphModule's module hierarchy. + In the case that ``root`` is a dict, the qualified name found in a Node's ``target`` will be + looked up directly in the dict's keys. The object mapped to by the Dict will be copied + over into the appropriate place within the GraphModule's module hierarchy. + + graph (Graph): ``graph`` contains the nodes this GraphModule should use for code generation + + class_name (str): ``name`` denotes the name of this GraphModule for debugging purposes. If it's unset, all + error messages will report as originating from ``GraphModule``. It may be helpful to set this + to ``root``'s original name or a name that makes sense within the context of your transform. + """ + super().__init__() + self.__class__.__name__ = class_name + if isinstance(root, torch.nn.Module): + if hasattr(root, "training"): + self.training = root.training + + # When we pickle/unpickle graph module, we don't want to drop any module or attributes. + if isinstance(root, _CodeOnlyModule): + for k, _ in root.named_children(): + _copy_attr(root, self, k) + + for k, _ in root.named_buffers(): + _copy_attr(root, self, k) + + for k, _ in root.named_parameters(): + _copy_attr(root, self, k) + + for node in graph.nodes: + if node.op in ["get_attr", "call_module"]: + assert isinstance(node.target, str) + _copy_attr(root, self, node.target) + elif isinstance(root, dict): + targets_to_copy = [] + for node in graph.nodes: + if node.op in ["get_attr", "call_module"]: + assert isinstance(node.target, str) + if node.target not in root: + raise RuntimeError( + "Node " + + str(node) + + " referenced target " + + node.target + + " but that target was not provided in ``root``!" + ) + targets_to_copy.append(node.target) + # Sort targets in ascending order of the # of atoms. + # This will ensure that less deeply nested attributes are assigned + # before more deeply nested attributes. For example, foo.bar + # will be assigned before foo.bar.baz. Otherwise, we might assign + # the user-provided ``foo.bar`` and wipe out the previously-assigned + # ``foo.bar.baz`` + targets_to_copy.sort(key=lambda t: t.count(".")) + for target_to_copy in targets_to_copy: + _assign_attr(root[target_to_copy], self, target_to_copy) + else: + raise RuntimeError("Unsupported type " + str(root) + " passed for root!") + + self.graph = graph + + # Store the Tracer class responsible for creating a Graph separately as part of the + # GraphModule state, except when the Tracer is defined in a local namespace. + # Locally defined Tracers are not pickleable. This is needed because torch.package will + # serialize a GraphModule without retaining the Graph, and needs to use the correct Tracer + # to re-create the Graph during deserialization. + self._tracer_cls = None + if ( + self.graph._tracer_cls + and "" not in self.graph._tracer_cls.__qualname__ + ): + self._tracer_cls = self.graph._tracer_cls + + self._tracer_extras = {} + if self.graph._tracer_extras: + self._tracer_extras = self.graph._tracer_extras + + # Dictionary to store metadata + self.meta: Dict[str, Any] = {} + self._replace_hooks: List[Callable] = [] + self._create_node_hooks: List[Callable] = [] + self._erase_node_hooks: List[Callable] = [] + + # TorchScript breaks trying to compile the graph setter because of the + # continued string literal. Issue here: https://github.com/pytorch/pytorch/issues/44842 + # + # Shouldn't be an issue since these methods shouldn't be used in TorchScript anyway + __jit_unused_properties__ = ["graph"] + + @property + def graph(self) -> Graph: + """ + Return the ``Graph`` underlying this ``GraphModule`` + """ + return self._graph + + @graph.setter + def graph(self, g: Graph) -> None: + """ + Set the underlying ``Graph`` for this ``GraphModule``. This will internally + recompile the ``GraphModule`` so that the generated ``forward()`` function + corresponds to ``g`` + """ + assert isinstance(g, Graph), f"Expected a Graph instance, but got {type(g)}" + self._graph = g + g.owning_module = self + self.recompile() + + @compatibility(is_backward_compatible=False) + def to_folder(self, folder: Union[str, os.PathLike], module_name: str = "FxModule"): + """Dumps out module to ``folder`` with ``module_name`` so that it can be + imported with ``from import `` + + Args: + + folder (Union[str, os.PathLike]): The folder to write the code out to + + module_name (str): Top-level name to use for the ``Module`` while + writing out the code + """ + folder = Path(folder) + Path(folder).mkdir(exist_ok=True) + torch.save(self.state_dict(), folder / "state_dict.pt") + tab = " " * 4 + custom_builtins = "\n".join([v.import_str for v in _custom_builtins.values()]) + model_str = f""" +import torch +{custom_builtins} + +from torch.nn import * +class {module_name}(torch.nn.Module): + def __init__(self): + super().__init__() +""" + + def _gen_model_repr(module_name: str, module: torch.nn.Module) -> Optional[str]: + safe_reprs = [ + nn.Linear, + nn.Conv1d, + nn.Conv2d, + nn.Conv3d, + nn.BatchNorm1d, + nn.BatchNorm2d, + nn.BatchNorm3d, + ] + if type(module) in safe_reprs: + return f"{module.__repr__()}" + else: + return None + + blobified_modules = [] + for module_name, module in self.named_children(): + module_str = _gen_model_repr(module_name, module) + if module_str is None: + module_file = folder / f"{module_name}.pt" + torch.save(module, module_file) + blobified_modules.append(module_name) + module_repr = module.__repr__().replace("\r", " ").replace("\n", " ") + # weights_only=False as this is legacy code that saves the model + module_str = ( + f"torch.load(r'{module_file}', weights_only=False) # {module_repr}" + ) + model_str += f"{tab * 2}self.{module_name} = {module_str}\n" + + for buffer_name, buffer in self._buffers.items(): + if buffer is None: + continue + model_str += f"{tab * 2}self.register_buffer('{buffer_name}', torch.empty({list(buffer.shape)}, dtype={buffer.dtype}))\n" # noqa: B950 + + for param_name, param in self._parameters.items(): + if param is None: + continue + model_str += f"{tab * 2}self.{param_name} = torch.nn.Parameter(torch.empty({list(param.shape)}, dtype={param.dtype}))\n" # noqa: B950 + + model_str += ( + f"{tab * 2}self.load_state_dict(torch.load(r'{folder}/state_dict.pt'))\n" + ) + model_str += f"{_addindent(self.code, 4)}\n" + + module_file = folder / "module.py" + module_file.write_text(model_str) + + init_file = folder / "__init__.py" + init_file.write_text("from .module import *") + + if len(blobified_modules) > 0: + warnings.warn( + "Was not able to save the following children modules as reprs -" + f"saved as pickled files instead: {blobified_modules}" + ) + + @compatibility(is_backward_compatible=True) + def add_submodule(self, target: str, m: torch.nn.Module) -> bool: + """ + Adds the given submodule to ``self``. + + This installs empty Modules where none exist yet if they are + subpaths of ``target``. + + Args: + target: The fully-qualified string name of the new submodule + (See example in ``nn.Module.get_submodule`` for how to + specify a fully-qualified string.) + m: The submodule itself; the actual object we want to + install in the current Module + + Return: + bool: Whether or not the submodule could be inserted. For + this method to return True, each object in the chain + denoted by ``target`` must either a) not exist yet, + or b) reference an ``nn.Module`` (not a parameter or + other attribute) + """ + *prefix, field = target.split(".") + mod: torch.nn.Module = self + + for item in prefix: + submod = getattr(mod, item, None) + + if submod is None: + submod = torch.nn.Module() + setattr(mod, item, submod) + + if not isinstance(submod, torch.nn.Module): + return False + + mod = submod + + mod.add_module(field, m) + return True + + @compatibility(is_backward_compatible=True) + def delete_submodule(self, target: str) -> bool: + """ + Deletes the given submodule from ``self``. + + The module will not be deleted if ``target`` is not a valid + target. + + Args: + target: The fully-qualified string name of the new submodule + (See example in ``nn.Module.get_submodule`` for how to + specify a fully-qualified string.) + + Returns: + bool: Whether or not the target string referenced a + submodule we want to delete. A return value of ``False`` + means that the ``target`` was not a valid reference to + a submodule. + """ + atoms = target.split(".") + path, target_submod = atoms[:-1], atoms[-1] + mod: torch.nn.Module = self + + # Get the parent module + for item in path: + if not hasattr(mod, item): + return False + + mod = getattr(mod, item) + + if not isinstance(mod, torch.nn.Module): + return False + + if not hasattr(mod, target_submod): + return False + + if not isinstance(getattr(mod, target_submod), torch.nn.Module): + return False + + delattr(mod, target_submod) + return True + + @compatibility(is_backward_compatible=True) + def delete_all_unused_submodules(self) -> None: + """ + Deletes all unused submodules from ``self``. + + A Module is considered "used" if any one of the following is + true: + 1. It has children that are used + 2. Its forward is called directly via a ``call_module`` node + 3. It has a non-Module attribute that is used from a + ``get_attr`` node + + This method can be called to clean up an ``nn.Module`` without + manually calling ``delete_submodule`` on each unused submodule. + """ + used: List[str] = [] + + for node in self.graph.nodes: + if node.op == "call_module" or node.op == "get_attr": + # A list of strings representing the different parts + # of the path. For example, `foo.bar.baz` gives us + # ["foo", "bar", "baz"] + fullpath = node.target.split(".") + + # If we're looking at multiple parts of a path, join + # join them with a dot. Otherwise, return that single + # element without doing anything to it. + def join_fn(x: str, y: str) -> str: + return ".".join([x, y] if y else [x]) + + # Progressively collect all the names of intermediate + # modules. For example, if we have the target + # `foo.bar.baz`, we'll add `foo`, `foo.bar`, and + # `foo.bar.baz` to the list. + used.extend(itertools.accumulate(fullpath, join_fn)) + + # For a `call_module` node, also register all recursive submodules + # as used + if node.op == "call_module": + try: + submod = self.get_submodule(node.target) + + for submod_name, _ in submod.named_modules(): + if submod_name != "": + used.append(".".join([node.target, submod_name])) + except AttributeError: + # Node referenced nonexistent submodule, don't need to + # worry about GCing anything + pass + + to_delete = [name for name, _ in self.named_modules() if name not in used] + + for name in to_delete: + self.delete_submodule(name) + + @property + def code(self) -> str: + """ + Return the Python code generated from the ``Graph`` underlying this + ``GraphModule``. + """ + if not hasattr(self, "_code"): + raise RuntimeError( + "Code has not been generated! Please report a bug to PyTorch" + ) + return self._code + + @compatibility(is_backward_compatible=True) + def recompile(self) -> PythonCode: + """ + Recompile this GraphModule from its ``graph`` attribute. This should be + called after editing the contained ``graph``, otherwise the generated + code of this ``GraphModule`` will be out of date. + """ + if isinstance(self._graph._codegen, _PyTreeCodeGen): + self._in_spec = self._graph._codegen.pytree_info.in_spec + self._out_spec = self._graph._codegen.pytree_info.out_spec + python_code = self._graph.python_code(root_module="self") + self._code = python_code.src + self._lineno_map = python_code._lineno_map + + cls = type(self) + co_fields = self._graph._co_fields if hasattr(self._graph, "_co_fields") else {} + cls.forward = _forward_from_src(self._code, python_code.globals, co_fields) + + # Determine whether this class explicitly defines a __call__ implementation + # to wrap. If it does, save it in order to have wrapped_call invoke it. + # If it does not, wrapped_call can use a dynamic call to super() instead. + # In most cases, super().__call__ should be torch.nn.Module.__call__. + # We do not want to hold a reference to Module.__call__ here; doing so will + # bypass patching of torch.nn.Module.__call__ done while symbolic tracing. + cls_call = cls.__call__ if "__call__" in vars(cls) else None + + if "_wrapped_call" not in vars(cls): + cls._wrapped_call = _WrappedCall(cls, cls_call) # type: ignore[attr-defined] + + def call_wrapped(self, *args, **kwargs): + return self._wrapped_call(self, *args, **kwargs) + + cls.__call__ = call_wrapped # type: ignore[method-assign] + + return python_code + + # Passing Tracer as argument allows subclasses extending fx.GraphModule + # define their own Tracer (extending fx.Tracer). + def __reduce_deploy__(self, importer: Importer): + dict_without_graph = self.__dict__.copy() + dict_without_graph["_graphmodule_cls_name"] = self.__class__.__name__ + del dict_without_graph["_graph"] + + python_code = self.recompile() + import_block = _format_import_block(python_code.globals, importer) + return (reduce_deploy_graph_module, (dict_without_graph, import_block)) + + def __reduce_package__(self, exporter: PackageExporter): + dict_without_graph = self.__dict__.copy() + dict_without_graph["_graphmodule_cls_name"] = self.__class__.__name__ + del dict_without_graph["_graph"] + + generated_module_name = f"fx-generated._{exporter.get_unique_id()}" + python_code = self.recompile() + import_block = _format_import_block(python_code.globals, exporter.importer) + module_code = import_block + self.code + exporter.save_source_string(generated_module_name, module_code) + return ( + reduce_package_graph_module, + (dict_without_graph, generated_module_name), + ) + + def __reduce__(self): + """ + Serialization of GraphModule. We serialize only the generated code, not + the underlying ``Graph``. This is because ``Graph`` does not have on-disk + backward-compatibility guarantees, whereas Python source code does. + On the deserialization side, we symbolically trace through the generated + code to regenerate the underlying ``Graph`` + """ + dict_without_graph = self.__dict__.copy() + + python_code = self.recompile() + import_block = _format_import_block(python_code.globals, sys_importer) + del dict_without_graph["_graph"] + return (reduce_graph_module, (dict_without_graph, import_block)) + + def _deepcopy_init(self): + return GraphModule.__init__ + + # because __reduce__ is defined for serialization, + # we need to define deepcopy otherwise it will call __reduce__ + # and cause symbolic tracing to occur every time we try to copy the object + def __deepcopy__(self, memo): + res = type(self).__new__(type(self)) + memo[id(self)] = res + fake_mod = _CodeOnlyModule(copy.deepcopy(self.__dict__, memo)) + self._deepcopy_init()(res, fake_mod, fake_mod.__dict__["_graph"]) + # hooks are lost during `GraphModule.__init__`, so we need to copy over + # them explicitly, note right now we are only copying state_dict related + # hooks, to reduce bc-related issues, we can copy forward/backward related + # hooks in the future as well if needed + extra_preserved_attrs = [ + "_state_dict_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", + "_replace_hooks", + "_create_node_hooks", + "_erase_node_hooks", + ] + for attr in extra_preserved_attrs: + if attr in self.__dict__: + setattr(res, attr, copy.deepcopy(self.__dict__[attr], memo)) + res.meta = copy.deepcopy(getattr(self, "meta", {}), memo) + if _USER_PRESERVED_ATTRIBUTES_KEY in res.meta: + for attr_name, attr in res.meta[_USER_PRESERVED_ATTRIBUTES_KEY].items(): + setattr(res, attr_name, attr) + return res + + def __copy__(self): + from ._lazy_graph_module import _make_graph_module + + res = _make_graph_module(self, self.graph) + res.meta = getattr(self, "meta", {}) + return res + + @compatibility(is_backward_compatible=False) + def print_readable( + self, + print_output=True, + include_stride=False, + include_device=False, + colored=False, + ): + """ + Return the Python code generated for current GraphModule and its children GraphModules + """ + return _print_readable( + self, + self._get_name(), + print_output, + include_stride, + include_device, + colored, + ) + + def __str__(self) -> str: + orig_str = super().__str__() + print_readable_reminder = ( + "# To see more debug info, please use `graph_module.print_readable()`" + ) + return "\n".join([orig_str, self._code, print_readable_reminder]) + + def _replicate_for_data_parallel(self): + new_gm = self.__copy__() + new_gm._is_replica = True + return new_gm + + @contextlib.contextmanager + def _set_replace_hook(self, f): + """ + Takes a callable which will be called everytime when we replace a node + to a new node, or change the node's name. Callable takes three arguments: + the old node we're changing, and NAME of the new node, followed by the + user node which consumes the old node to be replaced. + """ + assert callable(f), "Replace hook must be a callable." + self._register_replace_node_hook(f) + try: + yield + finally: + self._unregister_replace_node_hook(f) + + def _register_replace_node_hook(self, f): + """ + Takes a callable which will be called everytime when we replace a node + to a new node, or change the node's name. Callable takes three arguments: + the old node we're changing, and NAME of the new node, followed by the + user node which consumes the old node to be replaced. + """ + assert callable(f), "create_node hook must be a callable." + self._replace_hooks.append(f) + + def _unregister_replace_node_hook(self, f): + """ + Takes a callable which was previously registered to be called everytime when we replace a node. + This function will unregister that callable so it is no longer invoked on node replacement. + """ + assert callable(f), "create_node hook must be a callable." + self._replace_hooks.remove(f) + + def _register_create_node_hook(self, f): + """ + Takes a callable which will be called after we create a new node. The + callable takes the newly created node as input and returns None. + """ + assert callable(f), "create_node hook must be a callable." + self._create_node_hooks.append(f) + + def _unregister_create_node_hook(self, f): + """ + Takes a callable which was previously registered to be called after we create a node. + This function will unregister that callable so it is no longer invoked on node creation. + """ + assert callable(f), "create_node hook must be a callable." + self._create_node_hooks.remove(f) + + def _register_erase_node_hook(self, f): + """ + Takes a callable which will be called after we erase a node. The + callable takes the node that is being erased as input and returns None. + """ + assert callable(f), "erase_node hook must be a callable." + self._erase_node_hooks.append(f) + + def _unregister_erase_node_hook(self, f): + """ + Takes a callable which was previously registered to be called after we erase a node. + This function will unregister that callable so it is no longer invoked on node erasure. + """ + assert callable(f), "erase_node hook must be a callable." + self._erase_node_hooks.remove(f) + + +# workarounds for issues in __torch_function__ + +# WAR for __torch_function__ not handling tensor lists, +# fix is in https://github.com/pytorch/pytorch/pull/34725 +# orig_cat = torch.cat +# def patched_cat(*args, **kwargs): +# tensors = args[0] +# for t in tensors: +# if isinstance(t, Proxy): +# return t.__torch_function__(patched_cat, (), args, kwargs) +# return orig_cat(*args, **kwargs) +# patched_cat.__module__ = 'torch' +# patched_cat.__name__ = 'cat' +# torch.cat = patched_cat diff --git a/lib/python3.12/site-packages/torch/fx/immutable_collections.py b/lib/python3.12/site-packages/torch/fx/immutable_collections.py new file mode 100644 index 0000000000000000000000000000000000000000..2ff29cba474dcb613cd529b881193e20e1325856 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/immutable_collections.py @@ -0,0 +1,117 @@ +# mypy: allow-untyped-defs +from typing import Any, Dict, Iterable, List, Tuple + +from torch.utils._pytree import ( + _dict_flatten, + _dict_flatten_with_keys, + _dict_unflatten, + _list_flatten, + _list_flatten_with_keys, + _list_unflatten, + Context, + register_pytree_node, +) + +from ._compatibility import compatibility + + +__all__ = ["immutable_list", "immutable_dict"] + +_help_mutation = """\ +If you are attempting to modify the kwargs or args of a torch.fx.Node object, +instead create a new copy of it and assign the copy to the node: + new_args = ... # copy and mutate args + node.args = new_args +""" + + +def _no_mutation(self, *args, **kwargs): + raise NotImplementedError( + f"'{type(self).__name__}' object does not support mutation. {_help_mutation}", + ) + + +def _create_immutable_container(base, mutable_functions): + container = type("immutable_" + base.__name__, (base,), {}) + for attr in mutable_functions: + setattr(container, attr, _no_mutation) + return container + + +immutable_list = _create_immutable_container( + list, + ( + "__delitem__", + "__iadd__", + "__imul__", + "__setitem__", + "append", + "clear", + "extend", + "insert", + "pop", + "remove", + "reverse", + "sort", + ), +) +immutable_list.__reduce__ = lambda self: (immutable_list, (tuple(iter(self)),)) +immutable_list.__hash__ = lambda self: hash(tuple(self)) + +compatibility(is_backward_compatible=True)(immutable_list) + +immutable_dict = _create_immutable_container( + dict, + ( + "__delitem__", + "__ior__", + "__setitem__", + "clear", + "pop", + "popitem", + "setdefault", + "update", + ), +) +immutable_dict.__reduce__ = lambda self: (immutable_dict, (iter(self.items()),)) +immutable_dict.__hash__ = lambda self: hash(tuple(self.items())) +compatibility(is_backward_compatible=True)(immutable_dict) + + +# Register immutable collections for PyTree operations +def _immutable_dict_flatten(d: Dict[Any, Any]) -> Tuple[List[Any], Context]: + return _dict_flatten(d) + + +def _immutable_dict_unflatten( + values: Iterable[Any], + context: Context, +) -> Dict[Any, Any]: + return immutable_dict(_dict_unflatten(values, context)) + + +def _immutable_list_flatten(d: List[Any]) -> Tuple[List[Any], Context]: + return _list_flatten(d) + + +def _immutable_list_unflatten( + values: Iterable[Any], + context: Context, +) -> List[Any]: + return immutable_list(_list_unflatten(values, context)) + + +register_pytree_node( + immutable_dict, + _immutable_dict_flatten, + _immutable_dict_unflatten, + serialized_type_name="torch.fx.immutable_collections.immutable_dict", + flatten_with_keys_fn=_dict_flatten_with_keys, +) +register_pytree_node( + immutable_list, + _immutable_list_flatten, + _immutable_list_unflatten, + serialized_type_name="torch.fx.immutable_collections.immutable_list", + flatten_with_keys_fn=_list_flatten_with_keys, +) diff --git a/lib/python3.12/site-packages/torch/fx/interpreter.py b/lib/python3.12/site-packages/torch/fx/interpreter.py new file mode 100644 index 0000000000000000000000000000000000000000..aa24eb4bed1a945d246355825640e011905953e7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/interpreter.py @@ -0,0 +1,585 @@ +# mypy: allow-untyped-defs +import inspect +from contextlib import contextmanager +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union + +import torch +import torch.fx.traceback as fx_traceback +from torch.hub import tqdm + +from . import config +from ._compatibility import compatibility +from ._lazy_graph_module import _make_graph_module +from ._symbolic_trace import Tracer +from .graph import Graph +from .graph_module import GraphModule +from .node import Argument, map_aggregate, map_arg, Node, Target +from .proxy import Proxy + + +__all__ = ["Interpreter", "Transformer"] + + +@compatibility(is_backward_compatible=True) +class Interpreter: + """ + An Interpreter executes an FX graph Node-by-Node. This pattern + can be useful for many things, including writing code + transformations as well as analysis passes. + + Methods in the Interpreter class can be overridden to customize + the behavior of execution. The map of overrideable methods + in terms of call hierarchy:: + + run() + +-- run_node + +-- placeholder() + +-- get_attr() + +-- call_function() + +-- call_method() + +-- call_module() + +-- output() + + Example: + + Suppose we want to swap all instances of ``torch.neg`` with + ``torch.sigmoid`` and vice versa (including their ``Tensor`` + method equivalents). We could subclass Interpreter like so:: + + class NegSigmSwapInterpreter(Interpreter): + def call_function(self, target: Target, args: Tuple, kwargs: Dict) -> Any: + if target == torch.sigmoid: + return torch.neg(*args, **kwargs) + return super().call_function(target, args, kwargs) + + def call_method(self, target: Target, args: Tuple, kwargs: Dict) -> Any: + if target == "neg": + call_self, *args_tail = args + return call_self.sigmoid(*args_tail, **kwargs) + return super().call_method(target, args, kwargs) + + + def fn(x): + return torch.sigmoid(x).neg() + + + gm = torch.fx.symbolic_trace(fn) + input = torch.randn(3, 4) + result = NegSigmSwapInterpreter(gm).run(input) + torch.testing.assert_close(result, torch.neg(input).sigmoid()) + + Args: + module (torch.nn.Module): The module to be executed + garbage_collect_values (bool): Whether to delete values after their last + use within the Module's execution. This ensures optimal memory usage during + execution. This can be disabled to, for example, examine all of the intermediate + values in the execution by looking at the ``Interpreter.env`` attribute. + graph (Optional[Graph]): If passed, the interpreter will execute this + graph instead of `module.graph`, using the provided `module` + argument to satisfy any requests for state. + """ + + @compatibility(is_backward_compatible=True) + def __init__( + self, + module: torch.nn.Module, + garbage_collect_values: bool = True, + graph: Optional[Graph] = None, + ): + self.module = module + self.submodules = dict(self.module.named_modules()) + if graph is not None: + self.graph = graph + else: + self.graph = self.module.graph # type: ignore[assignment] + self.env: Dict[Node, Any] = {} + self.name = "Interpreter" + self.garbage_collect_values = garbage_collect_values + self.extra_traceback = True + + if self.garbage_collect_values: + # Run through reverse nodes and record the first instance of a use + # of a given node. This represents the *last* use of the node in the + # execution order of the program, which we will use to free unused + # values + node_to_last_use: Dict[Node, Node] = {} + self.user_to_last_uses: Dict[Node, List[Node]] = {} + + def register_last_uses(n: Node, user: Node): + if n not in node_to_last_use: + node_to_last_use[n] = user + self.user_to_last_uses.setdefault(user, []).append(n) + + for node in reversed(self.graph.nodes): + map_arg(node.args, lambda n: register_last_uses(n, node)) + map_arg(node.kwargs, lambda n: register_last_uses(n, node)) + + @compatibility(is_backward_compatible=True) + def run( + self, + *args, + initial_env: Optional[Dict[Node, Any]] = None, + enable_io_processing: bool = True, + ) -> Any: + """ + Run `module` via interpretation and return the result. + + Args: + *args: The arguments to the Module to run, in positional order + initial_env (Optional[Dict[Node, Any]]): An optional starting environment for execution. + This is a dict mapping `Node` to any value. This can be used, for example, to + pre-populate results for certain `Nodes` so as to do only partial evaluation within + the interpreter. + enable_io_processing (bool): If true, we process the inputs and outputs with graph's process_inputs and + process_outputs function first before using them. + + Returns: + Any: The value returned from executing the Module + """ + self.env = initial_env if initial_env is not None else {} + + # Positional function args are consumed left-to-right by + # `placeholder` nodes. Use an iterator to keep track of + # position and extract those values. + if enable_io_processing: + args = self.graph.process_inputs(*args) + self.args_iter: Iterator[Any] = iter(args) + pbar = tqdm( + total=len(self.graph.nodes), + desc=f"{self.name}: {str(list(self.graph.nodes)) if config.verbose_progress else ''}", + initial=0, + position=0, + leave=True, + disable=config.disable_progress, + delay=0, + ) + + for node in self.graph.nodes: + pbar.update(1) + if node in self.env: + # Short circuit if we have this value. This could + # be used, for example, for partial evaluation + # where the caller has pre-populated `env` with + # values for a subset of the program. + continue + + try: + self.env[node] = self.run_node(node) + except Exception as e: + if self.extra_traceback: + msg = f"While executing {node.format_node()}" + msg = f"{e.args[0]}\n\n{msg}" if e.args else str(msg) + msg += f"\nOriginal traceback:\n{node.stack_trace}" + e.args = (msg,) + e.args[1:] + if isinstance(e, KeyError): + raise RuntimeError(*e.args) from e + raise + + if self.garbage_collect_values: + for to_delete in self.user_to_last_uses.get(node, []): + del self.env[to_delete] + + if node.op == "output": + output_val = self.env[node] + return ( + self.graph.process_outputs(output_val) + if enable_io_processing + else output_val + ) + + @compatibility(is_backward_compatible=True) + def boxed_run(self, args_list): + """ + Run `module` via interpretation and return the result. This uses the "boxed" + calling convention, where you pass a list of arguments, which will be cleared + by the interpreter. This ensures that input tensors are promptly deallocated. + """ + args_iter = iter(args_list) + env = {} + for n in self.graph.nodes: + if n.op == "placeholder": + env[n] = next(args_iter) + args_list.clear() + return self.run(initial_env=env) + + @contextmanager + def _set_current_node(self, node): + with fx_traceback.set_current_meta( + node, f"Interpreter_{self.__class__.__name__}" + ): + yield + + @compatibility(is_backward_compatible=True) + def run_node(self, n: Node) -> Any: + """ + Run a specific node ``n`` and return the result. + Calls into placeholder, get_attr, call_function, + call_method, call_module, or output depending + on ``node.op`` + + Args: + n (Node): The Node to execute + + Returns: + Any: The result of executing ``n`` + """ + with self._set_current_node(n): + args, kwargs = self.fetch_args_kwargs_from_env(n) + assert isinstance(args, tuple) + assert isinstance(kwargs, dict) + return getattr(self, n.op)(n.target, args, kwargs) + + # Main Node running APIs + @compatibility(is_backward_compatible=True) + def placeholder( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + """ + Execute a ``placeholder`` node. Note that this is stateful: + ``Interpreter`` maintains an internal iterator over + arguments passed to ``run`` and this method returns + next() on that iterator. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Returns: + Any: The argument value that was retrieved. + """ + assert isinstance(target, str) + if target.startswith("*"): + # For a starred parameter e.g. `*args`, retrieve all + # remaining values from the args list. + return list(self.args_iter) + else: + try: + return next(self.args_iter) + except StopIteration as si: + if len(args) > 0: + return args[0] + else: + raise RuntimeError( + f"Expected positional argument for parameter {target}, but one was not passed in!" + ) from si + + @compatibility(is_backward_compatible=True) + def get_attr( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + """ + Execute a ``get_attr`` node. Will retrieve an attribute + value from the ``Module`` hierarchy of ``self.module``. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Return: + Any: The value of the attribute that was retrieved + """ + assert isinstance(target, str) + return self.fetch_attr(target) + + @compatibility(is_backward_compatible=True) + def call_function( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + """ + Execute a ``call_function`` node and return the result. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Return + Any: The value returned by the function invocation + """ + assert not isinstance(target, str) + + # Execute the function and return the result + return target(*args, **kwargs) + + @compatibility(is_backward_compatible=True) + def call_method( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + """ + Execute a ``call_method`` node and return the result. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Return + Any: The value returned by the method invocation + """ + # args[0] is the `self` object for this method call + self_obj, *args_tail = args + + # Execute the method and return the result + assert isinstance(target, str) + return getattr(self_obj, target)(*args_tail, **kwargs) + + @compatibility(is_backward_compatible=True) + def call_module( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + """ + Execute a ``call_module`` node and return the result. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Return + Any: The value returned by the module invocation + """ + # Retrieve executed args and kwargs values from the environment + + # Execute the method and return the result + assert isinstance(target, str) + submod = self.fetch_attr(target) + + return submod(*args, **kwargs) + + @compatibility(is_backward_compatible=True) + def output( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + """ + Execute an ``output`` node. This really just retrieves + the value referenced by the ``output`` node and returns it. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + + Return: + Any: The return value referenced by the output node + """ + return args[0] + + # Helper methods + @compatibility(is_backward_compatible=True) + def fetch_attr(self, target: str): + """ + Fetch an attribute from the ``Module`` hierarchy of ``self.module``. + + Args: + target (str): The fully-qualified name of the attribute to fetch + + Return: + Any: The value of the attribute. + """ + target_atoms = target.split(".") + attr_itr = self.module + for i, atom in enumerate(target_atoms): + if not hasattr(attr_itr, atom): + raise RuntimeError( + f"Node referenced nonexistent target {'.'.join(target_atoms[:i + 1])}" + ) + attr_itr = getattr(attr_itr, atom) + return attr_itr + + @compatibility(is_backward_compatible=True) + def fetch_args_kwargs_from_env(self, n: Node) -> Tuple[Tuple, Dict]: + """ + Fetch the concrete values of ``args`` and ``kwargs`` of node ``n`` + from the current execution environment. + + Args: + n (Node): The node for which ``args`` and ``kwargs`` should be fetched. + + Return: + Tuple[Tuple, Dict]: ``args`` and ``kwargs`` with concrete values for ``n``. + """ + args = self.map_nodes_to_values(n.args, n) + assert isinstance(args, tuple) + kwargs = self.map_nodes_to_values(n.kwargs, n) + assert isinstance(kwargs, dict) + return args, kwargs + + @compatibility(is_backward_compatible=True) + def map_nodes_to_values(self, args: Argument, n: Node) -> Argument: + """ + Recursively descend through ``args`` and look up the concrete value + for each ``Node`` in the current execution environment. + + Args: + args (Argument): Data structure within which to look up concrete values + + n (Node): Node to which ``args`` belongs. This is only used for error reporting. + """ + + def load_arg(n_arg: Node) -> Any: + if n_arg not in self.env: + raise RuntimeError( + f"Node {n} referenced nonexistent value {n_arg}! Run Graph.lint() " + f"to diagnose such issues" + ) + return self.env[n_arg] + + return map_arg(args, load_arg) + + +@compatibility(is_backward_compatible=True) +class Transformer(Interpreter): + """ + ``Transformer`` is a special type of interpreter that produces a + new ``Module``. It exposes a ``transform()`` method that returns + the transformed ``Module``. ``Transformer`` does not require + arguments to run, as ``Interpreter`` does. ``Transformer`` works + entirely symbolically. + + Example: + + Suppose we want to swap all instances of ``torch.neg`` with + ``torch.sigmoid`` and vice versa (including their ``Tensor`` + method equivalents). We could subclass ``Transformer`` like so:: + + class NegSigmSwapXformer(Transformer): + def call_function( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + if target == torch.sigmoid: + return torch.neg(*args, **kwargs) + return super().call_function(target, args, kwargs) + + def call_method( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + if target == "neg": + call_self, *args_tail = args + return call_self.sigmoid(*args_tail, **kwargs) + return super().call_method(target, args, kwargs) + + + def fn(x): + return torch.sigmoid(x).neg() + + + gm = torch.fx.symbolic_trace(fn) + + transformed: torch.nn.Module = NegSigmSwapXformer(gm).transform() + input = torch.randn(3, 4) + torch.testing.assert_close(transformed(input), torch.neg(input).sigmoid()) + + Args: + module (GraphModule): The ``Module`` to be transformed. + """ + + @compatibility(is_backward_compatible=True) + def __init__(self, module): + super().__init__(module) + self.new_graph = Graph() + self.new_graph.set_codegen(module.graph._codegen) + + class TransformerTracer(Tracer): + def __init__(self, graph: Graph): + super().__init__() + self.graph = graph + self.tensor_attrs: Dict[torch.Tensor, str] = {} # type: ignore[assignment] + + def is_leaf_module(self, _, __) -> bool: + return True + + self.tracer = TransformerTracer(self.new_graph) + self.tracer.root = module + + @compatibility(is_backward_compatible=True) + def placeholder( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Proxy: + """ + Execute a ``placeholder`` node. In ``Transformer``, this is + overridden to insert a new ``placeholder`` into the output + graph. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + """ + assert isinstance(target, str) + default_value = next(iter(args)) if args else inspect.Signature.empty + return Proxy( + self.new_graph.placeholder(target, default_value=default_value), self.tracer + ) + + @compatibility(is_backward_compatible=True) + def get_attr( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Proxy: + """ + Execute a ``get_attr`` node. In ``Transformer``, this is + overridden to insert a new ``get_attr`` node into the output + graph. + + Args: + target (Target): The call target for this node. See + `Node `__ for + details on semantics + args (Tuple): Tuple of positional args for this invocation + kwargs (Dict): Dict of keyword arguments for this invocation + """ + assert isinstance(target, str) + return self.tracer.create_proxy("get_attr", target, args, kwargs) + + @compatibility(is_backward_compatible=True) + def call_module( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + # Override so that the leaf module policy from `self.tracer` is respected. + assert isinstance(target, str) + submod = self.fetch_attr(target) + return self.tracer.call_module(submod, submod.forward, args, kwargs) + + @compatibility(is_backward_compatible=True) + def call_function( + self, target: "Target", args: Tuple[Argument, ...], kwargs: Dict[str, Any] + ) -> Any: + # Override so that functions that were wrapped are still wrapped. + return self.tracer.create_proxy("call_function", target, args, kwargs) + + @compatibility(is_backward_compatible=True) + def transform(self) -> GraphModule: + """ + Transform ``self.module`` and return the transformed + ``GraphModule``. + """ + with fx_traceback.preserve_node_meta(): + result = super().run(enable_io_processing=False) + if result is not None: + + def strip_proxy(a: Union[Argument, Proxy]) -> Any: + return a.node if isinstance(a, Proxy) else a + + new_output_node = self.new_graph.output(map_aggregate(result, strip_proxy)) + # also preserve the metadata from the old output node, if it exists + old_output_node = list(self.graph.nodes)[-1] + assert old_output_node.op == "output" + for k, v in old_output_node.meta.items(): + new_output_node.meta[k] = v + + return _make_graph_module(self.module, self.new_graph) diff --git a/lib/python3.12/site-packages/torch/fx/node.py b/lib/python3.12/site-packages/torch/fx/node.py new file mode 100644 index 0000000000000000000000000000000000000000..9c72d06541fdee372012d7be47da77ab36a1ab0c --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/node.py @@ -0,0 +1,922 @@ +# Nodes represent a definition of a value in our graph of operators. +import builtins +import inspect +import types +import warnings +from typing import ( + Any, + Callable, + Dict, + List, + Mapping, + Optional, + Sequence, + Set, + Tuple, + TYPE_CHECKING, + Union, +) + +import torch +from torch._C import _NodeBase +from torch.fx.operator_schemas import ( + ArgsKwargsPair, + normalize_function, + normalize_module, +) + +from .._ops import ops as _ops +from ._compatibility import compatibility +from .immutable_collections import immutable_dict, immutable_list + + +if TYPE_CHECKING: + from .graph import Graph + +__all__ = ["Node", "map_arg", "map_aggregate", "has_side_effect"] + +BaseArgumentTypes = Union[ + str, + int, + float, + bool, + complex, + torch.dtype, + torch.Tensor, + torch.device, + torch.memory_format, + torch.layout, + torch._ops.OpOverload, + torch.SymInt, + torch.SymBool, + torch.SymFloat, +] +base_types = BaseArgumentTypes.__args__ # type: ignore[attr-defined] + +Target = Union[Callable[..., Any], str] + +Argument = Optional[ + Union[ + Tuple["Argument", ...], + Sequence["Argument"], + Mapping[str, "Argument"], + slice, # Slice[Argument, Argument, Argument], but slice is not a templated type in typing + range, + "Node", + BaseArgumentTypes, + ] +] + +_legal_ops = dict.fromkeys( + [ + "placeholder", + "call_method", + "call_module", + "call_function", + "get_attr", + "output", + "root", + ] +) + +_side_effectful_need_to_be_preserved_pre_dispatch: Set[Callable] = { + torch._C._set_grad_enabled, + torch.amp._enter_autocast, + torch.amp._exit_autocast, +} + +# TODO: Either refactor this into 2 functions 1 dce for functional graphs and 1 dce for all graphs, +# or add logic to correctly mark all inplace ops as side effectful. +_side_effectful_functions: Set[Callable] = { + torch._assert, + torch._assert_async, + _ops.aten._assert_async.msg, + _ops.aten._assert_scalar.default, + _ops.aten._assert_tensor_metadata.default, + _ops.aten.sym_constrain_range.default, + _ops.aten.sym_constrain_range_for_size.default, + _ops.profiler._record_function_enter, + _ops.profiler._record_function_enter_new, + _ops.profiler._record_function_exit, + _ops.inductor.accumulate_grad_.default, +} | _side_effectful_need_to_be_preserved_pre_dispatch +if hasattr(_ops.inductor, "resize_storage_bytes_"): + _side_effectful_functions.add(_ops.inductor.resize_storage_bytes_.default) + + +@compatibility(is_backward_compatible=False) +def has_side_effect(fn: Callable) -> Callable: + _side_effectful_functions.add(fn) + return fn + + +# this is fixed on master, WAR for 1.5 +def _find_module_of_method(orig_method: Callable[..., Any]) -> str: + name = orig_method.__name__ + module = orig_method.__module__ + if module is not None: + return module + for guess in [torch, torch.nn.functional]: + if getattr(guess, name, None) is orig_method: + return guess.__name__ + raise RuntimeError(f"cannot find module for {orig_method}") + + +# Borrowed from CPython typing module +# https://github.com/python/cpython/blob/f90dc36c15d7fee0efaf6d39e97be0bdf2683e93/Lib/typing.py#L156 +def _type_repr(obj: object) -> str: + """Return the repr() of an object, special-casing types (internal helper). + If obj is a type, we return a shorter version than the default + type.__repr__, based on the module and qualified name, which is + typically enough to uniquely identify a type. For everything + else, we fall back on repr(obj). + """ + if isinstance(obj, type): + if obj.__module__ == "builtins": + return obj.__qualname__ + return f"{obj.__module__}.{obj.__qualname__}" + if obj is ...: + return "..." + if isinstance(obj, types.FunctionType): + return obj.__name__ + return repr(obj) + + +def _get_qualified_name(func: Callable[..., Any]) -> str: + # things like getattr just appear in builtins + if getattr(builtins, func.__name__, None) is func: + return func.__name__ + # torch.Tensor.{fn} + if isinstance( + func, (types.MethodDescriptorType, types.WrapperDescriptorType) + ) and func is getattr(torch.Tensor, func.__name__, None): + return f"torch.Tensor.{func.__name__}" + name = func.__name__ + if name == "": + # For lambdas, try to get their defining name in the module + try: + name = inspect.getsource(func).split("=")[0].strip() + except Exception as e: + raise RuntimeError("Unable to represent lambda") from e + module = _find_module_of_method(func) + module = module.replace( + "torch._ops", "torch.ops" + ) # WAR for bug in how torch.ops assigns module + # Fixup segment_reduce mismatch + if module == "torch" and name == "segment_reduce": + name = "_" + name + return f"{module}.{name}" + + +def _format_arg(arg: object, max_list_len: float = float("inf")) -> str: + if hasattr(arg, "_custom_fx_repr_fn"): + return arg._custom_fx_repr_fn() + elif isinstance(arg, list): + items = ", ".join( + _format_arg(a) for idx, a in enumerate(arg) if idx < max_list_len + ) + maybe_len = ( + "" if len(arg) < max_list_len + 1 else f", ...[total_len={len(arg)}]" + ) + return f"[{items}{maybe_len}]" + elif isinstance(arg, tuple): + items = ", ".join( + _format_arg(a) for idx, a in enumerate(arg) if idx < max_list_len + ) + maybe_len = ( + "" if len(arg) < max_list_len + 1 else f", ...[total_len={len(arg)}]" + ) + maybe_comma = "," if len(arg) == 1 else "" + return f"({items}{maybe_comma}{maybe_len})" + elif isinstance(arg, dict): + items_str = ", ".join(f"{k}: {_format_arg(v)}" for k, v in arg.items()) + return f"{{{items_str}}}" + + if isinstance(arg, Node): + return "%" + str(arg) + else: + return str(arg) + + +@compatibility(is_backward_compatible=True) +class Node(_NodeBase): + """ + ``Node`` is the data structure that represents individual operations within + a ``Graph``. For the most part, Nodes represent callsites to various entities, + such as operators, methods, and Modules (some exceptions include nodes that + specify function inputs and outputs). Each ``Node`` has a function specified + by its ``op`` property. The ``Node`` semantics for each value of ``op`` are as follows: + + - ``placeholder`` represents a function input. The ``name`` attribute specifies the name this value will take on. + ``target`` is similarly the name of the argument. ``args`` holds either: 1) nothing, or 2) a single argument + denoting the default parameter of the function input. ``kwargs`` is don't-care. Placeholders correspond to + the function parameters (e.g. ``x``) in the graph printout. + - ``get_attr`` retrieves a parameter from the module hierarchy. ``name`` is similarly the name the result of the + fetch is assigned to. ``target`` is the fully-qualified name of the parameter's position in the module hierarchy. + ``args`` and ``kwargs`` are don't-care + - ``call_function`` applies a free function to some values. ``name`` is similarly the name of the value to assign + to. ``target`` is the function to be applied. ``args`` and ``kwargs`` represent the arguments to the function, + following the Python calling convention + - ``call_module`` applies a module in the module hierarchy's ``forward()`` method to given arguments. ``name`` is + as previous. ``target`` is the fully-qualified name of the module in the module hierarchy to call. + ``args`` and ``kwargs`` represent the arguments to invoke the module on, *excluding the self argument*. + - ``call_method`` calls a method on a value. ``name`` is as similar. ``target`` is the string name of the method + to apply to the ``self`` argument. ``args`` and ``kwargs`` represent the arguments to invoke the module on, + *including the self argument* + - ``output`` contains the output of the traced function in its ``args[0]`` attribute. This corresponds to the "return" statement + in the Graph printout. + """ + + _args: Tuple["Argument", ...] + _kwargs: Dict[str, "Argument"] + graph: "Graph" + name: str + op: str + target: "Target" + _input_nodes: Dict["Node", None] + users: Dict["Node", None] + type: Optional[Any] + _sort_key: Any + _repr_fn: Optional[Callable[["Node"], str]] + meta: Dict[str, Any] + + @compatibility(is_backward_compatible=True) + def __init__( + self, + graph: "Graph", + name: str, + op: str, + target: "Target", + args: Tuple["Argument", ...], + kwargs: Dict[str, "Argument"], + return_type: Optional[Any] = None, + ) -> None: + """ + Instantiate an instance of ``Node``. Note: most often, you want to use the + Graph APIs, i.e. ``Graph.call_module``, ``Graph.call_method``, etc. rather + than instantiating a ``Node`` directly. + + Args: + graph (Graph): The ``Graph`` to which this ``Node`` should belong. + + name (str): The name to which the output of this ``Node`` should be assigned + + op (str): The opcode for this ``Node``. Can be one of 'placeholder', + 'call_method', 'call_module', 'call_function', 'get_attr', + 'output' + + target ('Target'): The target this op should call. See the broader + ``Node`` docstring for more details. + + args (Tuple['Argument']): The args to be passed to ``target`` + + kwargs (Dict[str, 'Argument']): The kwargs to be passed to ``target`` + + return_type (Optional[Any]): The python type expression representing the + type of the output of this node. This field can be used for + annotation of values in the generated code or for other types + of analyses. + """ + assert op in _legal_ops + if op == "call_function": + if not callable(target): + raise ValueError( + f"Node [graph = {graph}, name = '{name}'] target {target} has type {torch.typename(target)} " + "but a Callable is expected" + ) + else: + if not isinstance(target, str): + raise ValueError( + f"Node [graph = {graph}, name = '{name}'] target {target} has type {torch.typename(target)} " + "but a str is expected" + ) + super().__init__() + + # bypass Node.__setattr__ for perf and so that it doesn't need to handle half-built objects + assign = object.__setattr__ + + assign(self, "graph", graph) + assign(self, "name", name) # unique name of value being created + assign( + self, "op", op + ) # the kind of operation = placeholder|call_method|call_module|call_function|get_attr + + assign( + self, "target", target + ) # for method/module/function, the name of the method/module/function/attr + # being invoked, e.g add, layer1, or torch.add + + # All `Node`-valued inputs. Key is the Node, value is don't-care. + # The public API for this is `all_input_nodes`, this private attribute + # should not be accessed directly. + assign(self, "_input_nodes", {}) + self.__update_args_kwargs(args, kwargs) + + # All of the nodes that use the value produced by this Node + # Note one user may correspond to several uses, e.g. the node fo ``x + x`` + # would appear once here, but represents two uses. + # + # Is a dict to act as an "ordered set". Keys are significant, value dont-care + assign(self, "users", {}) + + # Type expression representing the output value of this node. + # This should contain the same class of Type objects that would appear + # as type annotations for function inputs/outputs. + # + # For placeholder nodes, this value will be used to type-annotate the + # generated function parameters. + # For the return node, this value will be used to type-annotate the + # generated function return type. (Note this is a special case. ``return`` + # does not produce a value, it's more of a notation. Thus, this value + # describes the type of args[0] in the ``return`` node. + assign(self, "type", return_type) + assign(self, "_sort_key", ()) + + # If set, use this fn to print this node + assign(self, "_repr_fn", None) + + # Dictionary to store metadata passes need to do their + # transformations. This metadata is preserved across node copies + assign(self, "meta", {}) + + def __getstate__(self) -> Dict[str, Any]: + state = self.__dict__.copy() + state["_erased"] = self._erased + state["_prev"] = self._prev + state["_next"] = self._next + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + _erased = state.pop("_erased") + _prev = state.pop("_prev") + _next = state.pop("_next") + self.__dict__.update(state) + self._erased = _erased + self._prev = _prev + self._next = _next + + @property + def next(self) -> "Node": + """ + Returns the next ``Node`` in the linked list of Nodes. + + Returns: + + The next ``Node`` in the linked list of Nodes. + """ + return self._next + + @property + def prev(self) -> "Node": + """ + Returns the previous ``Node`` in the linked list of Nodes. + + Returns: + + The previous ``Node`` in the linked list of Nodes. + """ + return self._prev + + @compatibility(is_backward_compatible=True) + def prepend(self, x: "Node") -> None: + """ + Insert x before this node in the list of nodes in the graph. Example:: + + Before: p -> self + bx -> x -> ax + After: p -> x -> self + bx -> ax + + Args: + x (Node): The node to put before this node. Must be a member of the same graph. + """ + assert self.graph == x.graph, "Attempting to move a Node into a different Graph" + if self == x: + warnings.warn( + "Trying to prepend a node to itself. This behavior has no effect on the graph." + ) + return + x._remove_from_list() + p = self._prev + p._next, x._prev = x, p + x._next, self._prev = self, x + + # compute x._sort_key + psk = x._prev._sort_key + nsk = x._next._sort_key + if len(psk) > len(nsk): + idx: int + *prefix, idx = psk[: len(nsk) + 1] + x._sort_key = (*prefix, idx + 1) + elif len(psk) < len(nsk): + *prefix, idx = nsk[: len(psk) + 1] + x._sort_key = (*prefix, idx - 1) + else: # same length, increase length by 1 + x._sort_key = (*psk, 0) + + def __gt__(self, other: "Node") -> bool: + return self._sort_key > other._sort_key + + def __lt__(self, other: "Node") -> bool: + return self._sort_key < other._sort_key + + def __ge__(self, other: "Node") -> bool: + return self > other or self == other + + def __le__(self, other: "Node") -> bool: + return self < other or self == other + + @compatibility(is_backward_compatible=True) + def append(self, x: "Node") -> None: + """ + Insert ``x`` after this node in the list of nodes in the graph. + Equivalent to ``self.next.prepend(x)`` + + Args: + x (Node): The node to put after this node. Must be a member of the same graph. + """ + self._next.prepend(x) + + def _remove_from_list(self) -> None: + p, n = self._prev, self._next + p._next, n._prev = n, p + + @property + def args(self) -> Tuple[Argument, ...]: + """ + The tuple of arguments to this ``Node``. The interpretation of arguments + depends on the node's opcode. See the :class:`Node` docstring for more + information. + + Assignment to this property is allowed. All accounting of uses and users + is updated automatically on assignment. + """ + return self._args + + @args.setter + def args(self, a: Tuple[Argument, ...]) -> None: + """ + Set the tuple of arguments to this Node. The interpretation of arguments + depends on the node's opcode. See the ``fx.Graph`` docstring for more + information. + """ + # DO NOT CALL `__update_args_kwargs` directly. The correct way to + # set `args` is via direct assignment, i.e. `node.args = new_args` + self.__update_args_kwargs(a, self._kwargs) + + @property + def kwargs(self) -> Dict[str, Argument]: + """ + The dict of keyword arguments to this ``Node``. The interpretation of arguments + depends on the node's opcode. See the :class:`Node` docstring for more + information. + + Assignment to this property is allowed. All accounting of uses and users + is updated automatically on assignment. + """ + return self._kwargs + + @kwargs.setter + def kwargs(self, k: Dict[str, Argument]) -> None: + """ + Set the dict of kwargs to this Node. The interpretation of arguments + depends on the node's opcode. See the ``fx.Graph`` docstring for more + information. + """ + # DO NOT CALL `__update_args_kwargs` directly. The correct way to + # set `args` is via direct assignment, i.e. `node.kwargs = new_kwargs` + self.__update_args_kwargs(self._args, k) + + @property + def all_input_nodes(self) -> List["Node"]: + """ + Return all Nodes that are inputs to this Node. This is equivalent to + iterating over ``args`` and ``kwargs`` and only collecting the values that + are Nodes. + + Returns: + + List of ``Nodes`` that appear in the ``args`` and ``kwargs`` of this + ``Node``, in that order. + """ + return list(self._input_nodes.keys()) + + @compatibility(is_backward_compatible=True) + def update_arg(self, idx: int, arg: Argument) -> None: + """ + Update an existing positional argument to contain the new value + ``arg``. After calling, ``self.args[idx] == arg``. + + Args: + + idx (int): The index into ``self.args`` of the element to update + arg (Argument): The new argument value to write into ``args`` + """ + args = list(self.args) + args[idx] = arg + self.args = tuple(args) + + @compatibility(is_backward_compatible=True) + def insert_arg(self, idx: int, arg: Argument) -> None: + """ + Insert an positional argument to the argument list with given index. + + Args: + + idx (int): The index of the element in ``self.args`` to be inserted before. + arg (Argument): The new argument value to insert into ``args`` + """ + assert ( + 0 <= idx <= len(self.args) + ), "insert_args index must be between 0 and len(self.args)" + args_left = self.args[:idx] + args_right = self.args[idx:] + + self._args = args_left + (arg,) + args_right + + _new_input_nodes: Dict[Node, None] = {} + map_arg(arg, _new_input_nodes.setdefault) + + for new_use in _new_input_nodes.keys(): + if new_use not in self._input_nodes: + self._input_nodes.setdefault(new_use) + new_use.users.setdefault(self) + + @compatibility(is_backward_compatible=True) + def update_kwarg(self, key: str, arg: Argument) -> None: + """ + Update an existing keyword argument to contain the new value + ``arg``. After calling, ``self.kwargs[key] == arg``. + + Args: + + key (str): The key in ``self.kwargs`` of the element to update + arg (Argument): The new argument value to write into ``kwargs`` + """ + self.kwargs = {**self.kwargs, key: arg} + + @property + def stack_trace(self) -> Optional[str]: + """ + Return the Python stack trace that was recorded during tracing, if any. + When traced with fx.Tracer, this property is usually populated by + `Tracer.create_proxy`. To record stack traces during tracing for debug purposes, + set `record_stack_traces = True` on the `Tracer` instance. + When traced with dynamo, this property will be populated by default by + `OutputGraph.create_proxy`. + + stack_trace would have the innermost frame at the end of the string. + """ + return self.meta.get("stack_trace", None) + + @stack_trace.setter + def stack_trace(self, trace: Optional[str]) -> None: + self.meta["stack_trace"] = trace + + def __update_args_kwargs( + self, new_args: Tuple["Argument", ...], new_kwargs: Dict[str, "Argument"] + ) -> None: + """ + This API is internal. Do *not* call it directly. + """ + + def update_users_and_input_nodes(n: Any) -> Any: + if isinstance(n, Node): + self._input_nodes.setdefault(n) + n.users.setdefault(self) + return n + + # Clear prior users and input_nodes + for old_use in self._input_nodes.keys(): + old_use.users.pop(self) + object.__setattr__(self, "_input_nodes", {}) # bypass Node.__setattr__ + + # We do three things in a single pass of the args + # - Normalize list->immutable_list, dict->immutable_dict, etc + # - Populate self._input_nodes + # - Populate arg.users[self] for each arg + object.__setattr__( + self, "_args", map_aggregate(new_args, update_users_and_input_nodes) + ) + object.__setattr__( + self, "_kwargs", map_aggregate(new_kwargs, update_users_and_input_nodes) + ) + + def __repr__(self) -> str: + if self._repr_fn: + return self._repr_fn(self) + return self.name + + def _pretty_print_target(self, target: object) -> str: + """ + Make target printouts more user-friendly. + 1) builtins will be printed as `builtins.xyz` + 2) operators will be printed as `operator.xyz` + 3) other callables will be printed with qualified name, e.g. torch.add + """ + if isinstance(target, str): + return target + if hasattr(target, "__module__"): + name = getattr(target, "__name__", None) + if name is None: + # Just to be defensive, if we don't have `__name__`, get the + # qualname. Not sure if this happens for any members of `operator` + # or `builtins`. This fallback path is not as good, since e.g. + # things in `operator` have `_operator` as their __module__. + # TODO: THIS IS BROKEN: _get_qualified_name calls `__name__` + return _get_qualified_name(target) # type: ignore[arg-type] + if target.__module__ == "builtins": + return f"builtins.{name}" + elif target.__module__ == "_operator": + return f"operator.{name}" + return _get_qualified_name(target) # type: ignore[arg-type] + + @compatibility(is_backward_compatible=True) + def format_node( + self, + placeholder_names: Optional[List[str]] = None, + maybe_return_typename: Optional[List[str]] = None, + ) -> Optional[str]: + """ + Return a descriptive string representation of ``self``. + + This method can be used with no arguments as a debugging + utility. + + This function is also used internally in the ``__str__`` method + of ``Graph``. Together, the strings in ``placeholder_names`` + and ``maybe_return_typename`` make up the signature of the + autogenerated ``forward`` function in this Graph's surrounding + GraphModule. ``placeholder_names`` and ``maybe_return_typename`` + should not be used otherwise. + + Args: + placeholder_names: A list that will store formatted strings + representing the placeholders in the generated + ``forward`` function. Internal use only. + maybe_return_typename: A single-element list that will store + a formatted string representing the output of the + generated ``forward`` function. Internal use only. + + Returns: + str: If 1) we're using ``format_node`` as an internal helper + in the ``__str__`` method of ``Graph``, and 2) ``self`` + is a placeholder Node, return ``None``. Otherwise, + return a descriptive string representation of the + current Node. + """ + if self.op == "placeholder": + assert isinstance(self.target, str) + arg_str = self.target + arg_str += arg_str + f": {_type_repr(self.type)}" if self.type else "" + if placeholder_names: + placeholder_names.append(arg_str) + return None + maybe_typename = f"{_type_repr(self.type)} " if self.type else "" + default_val = "(default=" + str(self.args[0]) + ")" if self.args else "" + return f"%{self.name} : {maybe_typename}[num_users={len(self.users)}] = {self.op}[target={self.target}]{default_val}" + elif self.op == "get_attr": + maybe_typename = ( + f"{_type_repr(self.type)} " if self.type is not None else "" + ) + return ( + f"%{self.name} : {maybe_typename}[num_users={len(self.users)}] = " + f"{self.op}[target={self._pretty_print_target(self.target)}]" + ) + elif self.op == "output": + if self.type and maybe_return_typename: + maybe_return_typename[0] = f" -> {_type_repr(self.type)}" + return f"return {self.args[0]}" + else: + maybe_typename = ( + f"{_type_repr(self.type)} " if self.type is not None else "" + ) + return ( + f"%{self.name} : {maybe_typename}[num_users={len(self.users)}] = " + f"{self.op}[target={self._pretty_print_target(self.target)}](" + f"args = {_format_arg(self.args)}, kwargs = {_format_arg(self.kwargs)})" + ) + + @compatibility(is_backward_compatible=True) + def replace_all_uses_with( + self, + replace_with: "Node", + delete_user_cb: Callable[["Node"], bool] = lambda user: True, + *, + propagate_meta: bool = False, + ) -> List["Node"]: + """ + Replace all uses of ``self`` in the Graph with the Node ``replace_with``. + + Args: + + replace_with (Node): The node to replace all uses of ``self`` with. + delete_user_cb (Callable): Callback that is called to determine + whether a given user of the self node should be removed. + propagate_meta (bool): Whether or not to copy all properties + on the .meta field of the original node onto the replacement node. + For safety, this is only valid to do if the replacement node + doesn't already have an existing .meta field. + + Returns: + + The list of Nodes on which this change was made. + """ + if propagate_meta: + assert len(replace_with.meta) == 0, ( + "Called node.replace_all_uses_with(replace_with, propagate_meta=True), " + "but replace_with already has .meta keys" + ) + for k, v in self.meta.items(): + replace_with.meta[k] = v + to_process = list(self.users) + skipped = [] + m = self.graph.owning_module + for use_node in to_process: + if not delete_user_cb(use_node): + skipped.append(use_node) + continue + + def maybe_replace_node(n: Node) -> Node: + if n == self: + return replace_with + else: + return n + + if getattr(m, "_replace_hooks", None): + for replace_hook in m._replace_hooks: + replace_hook(old=self, new=replace_with.name, user=use_node) + + new_args = map_arg(use_node.args, maybe_replace_node) + new_kwargs = map_arg(use_node.kwargs, maybe_replace_node) + assert isinstance(new_args, tuple) + assert isinstance(new_kwargs, dict) + use_node.__update_args_kwargs(new_args, new_kwargs) + + assert len(self.users) - len(skipped) == 0 + return [n for n in to_process if n not in skipped] + + @compatibility(is_backward_compatible=False) + def is_impure(self) -> bool: + """ + Returns whether this op is impure, i.e. if its op is a placeholder or + output, or if a call_function or call_module which is impure. + + Returns: + + bool: If the op is impure or not. + """ + if self.op in {"placeholder", "output"}: + return True + + # Check if an impure function based on schema. + if self.op == "call_function": + schema = getattr(self.target, "_schema", None) + schema_mutable = schema is not None and schema.is_mutable + return schema_mutable or self.target in _side_effectful_functions + + # Check if an impure module. + if self.op == "call_module": + assert ( + self.graph.owning_module is not None + ), "self.graph.owning_module not set for purity check" + target_mod = self.graph.owning_module.get_submodule(self.target) + assert ( + target_mod is not None + ), f"Did not find expected submodule target {self.target}" + return getattr(target_mod, "_is_impure", False) + + return False + + @compatibility(is_backward_compatible=False) + def normalized_arguments( + self, + root: torch.nn.Module, + arg_types: Optional[Tuple[Any]] = None, + kwarg_types: Optional[Dict[str, Any]] = None, + normalize_to_only_use_kwargs: bool = False, + ) -> Optional[ArgsKwargsPair]: + """ + Returns normalized arguments to Python targets. This means that + `args/kwargs` will be matched up to the module/functional's + signature and return exclusively kwargs in positional order + if `normalize_to_only_use_kwargs` is true. + Also populates default values. Does not support positional-only + parameters or varargs parameters. + + Supports module calls. + + May require `arg_types` and `kwarg_types` in order to disambiguate overloads. + + Args: + root (torch.nn.Module): Module upon which to resolve module targets. + arg_types (Optional[Tuple[Any]]): Tuple of arg types for the args + kwarg_types (Optional[Dict[str, Any]]): Dict of arg types for the kwargs + normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs. + + Returns: + + Returns NamedTuple ArgsKwargsPair, or `None` if not successful. + """ + if self.op == "call_function": + assert callable(self.target) + return normalize_function( + self.target, + self.args, # type: ignore[arg-type] + self.kwargs, + arg_types, + kwarg_types, + ) + elif self.op == "call_module": + assert isinstance(self.target, str) + return normalize_module(root, self.target, self.args, self.kwargs) # type: ignore[arg-type] + + return None + + @compatibility(is_backward_compatible=True) + def replace_input_with(self, old_input: "Node", new_input: "Node") -> None: + """ + Loop through input nodes of ``self``, and replace all instances of + ``old_input`` with ``new_input``. + + Args: + + old_input (Node): The old input node to be replaced. + new_input (Node): The new input node to replace ``old_input``. + """ + + def maybe_replace_node(n: Node) -> Node: + return new_input if n == old_input else n + + m = self.graph.owning_module + if getattr(m, "_replace_hooks", None): + for replace_hook in m._replace_hooks: + replace_hook(old=old_input, new=new_input.name, user=self) + + new_args = map_arg(self.args, maybe_replace_node) + new_kwargs = map_arg(self.kwargs, maybe_replace_node) + assert isinstance(new_args, tuple) + assert isinstance(new_kwargs, dict) + self.__update_args_kwargs(new_args, new_kwargs) + + def _rename(self, candidate: str) -> None: + if candidate == self.name: + return + name = self.graph._graph_namespace.create_name(candidate, None) + self.name = name + self.graph._graph_namespace._rename_object(self, name) + + def __setattr__(self, name: str, value: Any) -> None: + if name == "name" and hasattr(self, "name"): + m = self.graph.owning_module + if getattr(m, "_replace_hooks", None): + assert isinstance(value, str) + for user in self.users: + for replace_hook in m._replace_hooks: + replace_hook(old=self, new=value, user=user) + update = False + if ( + hasattr(self, name) + and hasattr(self.graph, "_find_nodes_lookup_table") + and self in self.graph._find_nodes_lookup_table + ): + update = True + self.graph._find_nodes_lookup_table.remove(self) + object.__setattr__(self, name, value) + if update: + self.graph._find_nodes_lookup_table.insert(self) + + +@compatibility(is_backward_compatible=True) +def map_arg(a: Argument, fn: Callable[[Node], Argument]) -> Argument: + """ + Apply fn to each Node appearing arg. arg may be a list, tuple, slice, or dict with string keys. + """ + assert callable(fn), "torch.fx.map_arg(a, fn): fn must be a callable" + return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x) + + +@compatibility(is_backward_compatible=True) +def map_aggregate(a: Argument, fn: Callable[[Argument], Argument]) -> Argument: + """ + Apply fn to each Node appearing arg. arg may be a list, tuple, slice, or dict with string keys. + """ + if isinstance(a, tuple): + t = tuple([map_aggregate(elem, fn) for elem in a]) + # Support NamedTuple (if it has `_fields`) by repacking into original type. + return t if not hasattr(a, "_fields") else type(a)(*t) # type: ignore[arg-type] + elif isinstance(a, list): + return immutable_list([map_aggregate(elem, fn) for elem in a]) + elif isinstance(a, dict): + rv = immutable_dict() + for k, v in a.items(): + dict.__setitem__(rv, k, map_aggregate(v, fn)) + return rv + elif isinstance(a, slice): + return slice( + map_aggregate(a.start, fn), + map_aggregate(a.stop, fn), + map_aggregate(a.step, fn), + ) + else: + return fn(a) diff --git a/lib/python3.12/site-packages/torch/fx/operator_schemas.py b/lib/python3.12/site-packages/torch/fx/operator_schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..f654b6c060e8160807e6d5bec06469d356f039bc --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/operator_schemas.py @@ -0,0 +1,563 @@ +# mypy: allow-untyped-defs +import enum +import inspect +import numbers +import types +import typing +import warnings +from typing import ( + Any, + Callable, + cast, + Dict, + List, + NamedTuple, + Optional, + Tuple, + TYPE_CHECKING, +) + +import torch +from torch._jit_internal import boolean_dispatched +from torch._ops import OpOverload, OpOverloadPacket + +from ._compatibility import compatibility + + +if TYPE_CHECKING: + from .node import Argument + +__all__ = [ + "ArgsKwargsPair", + "check_for_mutable_operation", + "get_signature_for_torch_op", + "create_type_hint", + "type_matches", + "normalize_function", + "normalize_module", +] + + +@compatibility(is_backward_compatible=False) +class ArgsKwargsPair(NamedTuple): + """ + Simple named tuple for wrapping args/kwargs pairs. + """ + + args: Tuple[Any, ...] + kwargs: Dict[str, Any] + + +_manual_overrides: Dict[Callable, List[inspect.Signature]] = {} + + +def _nonzero_schemas(): + signatures = [] + + def nonzero(self): + pass + + signatures.append(inspect.signature(nonzero)) + + def nonzero(self, *, as_tuple: bool): # type: ignore[no-redef] + pass + + signatures.append(inspect.signature(nonzero)) + + return signatures + + +_manual_overrides[torch.nonzero] = _nonzero_schemas() + + +class _FakeGlobalNamespace: + def __getattr__(self, name): + if name == "torch": + return torch + raise RuntimeError("Expected a torch namespace lookup") + + +_type_eval_globals = { + "Tensor": torch.Tensor, + "Device": torch.device, + "Layout": torch.layout, + "number": numbers.Number, + "Future": torch.jit.Future, + "AnyEnumType": enum.Enum, + "QScheme": torch.qscheme, + "__torch__": _FakeGlobalNamespace(), + "NoneType": type(None), + "Storage": torch.UntypedStorage, + "t": typing.TypeVar("t"), +} +for k in dir(typing): + _type_eval_globals[k] = getattr(typing, k) + + +def _torchscript_type_to_python_type(ts_type: "torch._C.JitType") -> Any: + """ + Convert a TorchScript type to a Python type (including subtypes) via + eval'ing the annotation_str. _type_eval_globals sets up expressions + like "List" and "Future" to map to actual types (typing.List and jit.Future) + """ + return eval(ts_type.annotation_str, _type_eval_globals) + + +def _torchscript_schema_to_signature_impl( + ts_schema: torch._C.FunctionSchema, +) -> inspect.Signature: + from inspect import Parameter + + parameters: List[Parameter] = [] + for arg in ts_schema.arguments: + arg_type = _torchscript_type_to_python_type(arg.type) + default = arg.default_value if arg.has_default_value() else Parameter.empty + # TODO: Figure out if this is safe. It seems like when generating the type signatures for + # PythonArgParser, we emit signatures with `input` instead of `self` as the first tensor + # argument name. Downstream, if someone converts that positional argument to a keyword + # argument, the name mismatch will break things, so here we're going to normalize the + # name to "input" + name = arg.name if arg.name != "self" else "input" + kind = ( + Parameter.KEYWORD_ONLY + if arg.kwarg_only + else Parameter.POSITIONAL_OR_KEYWORD + ) + # "from" is a keyword therefore it must be a POSITIONAL_ONLY argument + if name == "from": + assert kind == Parameter.POSITIONAL_OR_KEYWORD + # ParameterKind type is internal implementation detail to inspec package + # which makes it hard to do type annotation + kind = Parameter.POSITIONAL_ONLY # type: ignore[assignment] + # This renders all previous arguments to positional only + for idx, p in enumerate(parameters): + assert p.kind == Parameter.POSITIONAL_OR_KEYWORD + parameters[idx] = Parameter( + name=p.name, + kind=Parameter.POSITIONAL_ONLY, + default=p.default, + annotation=p.annotation, + ) + parameters.append( + Parameter(name=name, kind=kind, default=default, annotation=arg_type) + ) + return_types = [ + _torchscript_type_to_python_type(ret.type) for ret in ts_schema.returns + ] + if len(return_types) == 0: + return_type = None + elif len(return_types) == 1: + return_type = return_types[0] + else: + return_type = tuple(return_types) + + return inspect.Signature(parameters, return_annotation=return_type) + + +_SCHEMA_TO_SIGNATURE_CACHE: Dict[Tuple[str, str], inspect.Signature] = {} + + +def _torchscript_schema_to_signature( + ts_schema: torch._C.FunctionSchema, +) -> inspect.Signature: + # Cached as it's called in the hot path of FakeTensor dispatch + cache_key = ts_schema.name, ts_schema.overload_name + cache_val = _SCHEMA_TO_SIGNATURE_CACHE.get(cache_key) + if cache_val is not None: + return cache_val + + res = _torchscript_schema_to_signature_impl(ts_schema) + _SCHEMA_TO_SIGNATURE_CACHE[cache_key] = res + return res + + +@compatibility(is_backward_compatible=False) +def check_for_mutable_operation( + target: Callable, args: Tuple["Argument", ...], kwargs: Dict[str, "Argument"] +): + signatures, schemas = get_signature_for_torch_op(target, return_schemas=True) + + if signatures and schemas: + matched_schemas = [] + + # Iterate through all of the schema until we find one that matches + # If one matches, populate `new_args_and_kwargs` with the new args/kwargs + # values. If none matches, `new_args_and_kwargs` will be None + for candidate_signature, schema in zip(signatures, schemas): + try: + candidate_signature.bind(*args, **kwargs) + matched_schemas.append((candidate_signature, schema)) + except TypeError: + continue + + def throw_if_mutable(schema): + if schema.is_mutable: + raise RuntimeError( + f"Tried to trace mutable operation {schema}. FX only supports functional " + f"code, so operations that mutate operands in-place (e.g. via `out` arguments) " + f"are not supported" + ) + + if len(matched_schemas) == 0: + # Did not match any schema. Cannot check for mutation + pass + elif len(matched_schemas) == 1: + # Matched exactly one schema, unambiguous + _, schema_to_check = matched_schemas[0] + throw_if_mutable(schema_to_check) + else: + # Ambiguous schema match. Since mutability checking is best effort, + # do nothing. + pass + + +@compatibility(is_backward_compatible=False) +def get_signature_for_torch_op(op: Callable, return_schemas: bool = False): + """ + Given an operator on the `torch` namespace, return a list of `inspect.Signature` + objects corresponding to the overloads of that op.. May return `None` if a signature + could not be retrieved. + + Args: + op (Callable): An operator on the `torch` namespace to look up a signature for + + Returns: + Optional[List[inspect.Signature]]: A list of signatures for the overloads of this + operator, or None if the operator signatures could not be retrieved. If + return_schemas=True, returns a tuple containing the optional Python signatures + and the optional TorchScript Function signature + """ + if isinstance(op, OpOverload): + schemas = [op._schema] + elif isinstance(op, OpOverloadPacket): + schemas = [getattr(op, overload)._schema for overload in op.overloads()] + else: + override = _manual_overrides.get(op) + if override: + return (override, None) if return_schemas else None + + aten_fn = torch.jit._builtins._find_builtin(op) + + if aten_fn is None: + return (None, None) if return_schemas else None + schemas = torch._C._jit_get_schemas_for_operator(aten_fn) + + signatures = [_torchscript_schema_to_signature(schema) for schema in schemas] + return (signatures, schemas) if return_schemas else signatures + + +@compatibility(is_backward_compatible=False) +def create_type_hint(x): + """ + Produces a type hint for the given argument. + + The :func:`create_type_hint` looks for a type hint compatible with the input argument `x`. + + If `x` is a `list` or `tuple`, it looks for an object in the list whose type is a superclass + of the rest, and uses that as `base_type` for the `List` or `Tuple` to be returned. + If no such object is found, it defaults to `List[Any]`. + + If `x` is neither a `list` nor a `tuple`, it returns `x`. + """ + try: + if isinstance(x, (list, tuple)): + # todo(chilli): Figure out the right way for mypy to handle this + if isinstance(x, list): + + def ret_type(x): + return List[x] # type: ignore[valid-type] + + else: + + def ret_type(x): + return Tuple[x, ...] + + if len(x) == 0: + return ret_type(Any) + base_type = x[0] + for t in x: + if issubclass(t, base_type): + continue + elif issubclass(base_type, t): + base_type = t + else: + return ret_type(Any) + return ret_type(base_type) + except Exception: + # We tried to create a type hint for list but failed. + warnings.warn( + f"We were not able to successfully create type hint from the type {x}" + ) + return x + + +@compatibility(is_backward_compatible=False) +def type_matches(signature_type: Any, argument_type: Any): + sig_origin_type = getattr(signature_type, "__origin__", signature_type) + + if signature_type is argument_type: + return True + + # Union types in signature. Given type needs to match one of the + # contained types in the Union + if sig_origin_type is typing.Union and signature_type != argument_type: + sig_contained = signature_type.__args__ + return any(type_matches(c, argument_type) for c in sig_contained) + + if signature_type is List[int] and argument_type is int: + # int can be promoted to List[int] + return True + + if getattr(signature_type, "__origin__", None) in {list, List}: + sig_el_type = signature_type.__args__[0] + if not inspect.isclass(sig_el_type): + warnings.warn( + f"Does not support nested parametric types, got {signature_type}. Please file a bug." + ) + return False + if getattr(argument_type, "__origin__", None) in {list, List}: + return issubclass(argument_type.__args__[0], sig_el_type) + + def is_homogeneous_tuple(t): + if getattr(t, "__origin__", None) not in {tuple, Tuple}: + return False + contained = t.__args__ + if t.__args__ == ((),): # Tuple[()].__args__ == ((),) for some reason + return True + return all((c is Ellipsis) or issubclass(c, sig_el_type) for c in contained) + + # Tuple[T] is accepted for List[T] parameters + return is_homogeneous_tuple(argument_type) + + # Dtype is an int in schemas + if signature_type is int and argument_type is torch.dtype: + return True + + if signature_type is numbers.Number and argument_type in {int, float}: + return True + if inspect.isclass(argument_type) and inspect.isclass(signature_type): + return issubclass(argument_type, signature_type) + + return False + + +@compatibility(is_backward_compatible=False) +def normalize_function( + target: Callable, + args: Tuple[Any], + kwargs: Optional[Dict[str, Any]] = None, + arg_types: Optional[Tuple[Any]] = None, + kwarg_types: Optional[Dict[str, Any]] = None, + normalize_to_only_use_kwargs: bool = False, +) -> Optional[ArgsKwargsPair]: + """ + Returns normalized arguments to PyTorch functions. This means that + `args/kwargs` will be matched up to the functional's + signature and return exclusively kwargs in positional order if + `normalize_to_only_use_kwargs` is True. + Also populates default values. Does not support positional-only + parameters or varargs parameters (*args, **kwargs). Does not support modules. + + May require `arg_types` and `kwarg_types` in order to disambiguate overloads. + + Args: + target (Callable): Function that we are normalizing + args (Tuple[Any]): Tuple of args to the function + kwargs (Optional[Dict[str, Any]]): Dict of kwargs to the function + arg_types (Optional[Tuple[Any]]): Tuple of arg types for the args + kwarg_types (Optional[Dict[str, Any]]): Dict of arg types for the kwargs + normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs. + + Returns: + + Returns normalized_args_and_kwargs, or `None` if not successful. + """ + if kwargs is None: + kwargs = {} + new_args_and_kwargs = None + if not isinstance(target, types.BuiltinFunctionType) and not ( + isinstance(target, (OpOverloadPacket, OpOverload)) + ): + 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 None + target_for_analysis = if_true + + assert callable(target_for_analysis) + sig = inspect.signature(inspect.unwrap(target_for_analysis)) + new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs( + sig, args, kwargs, normalize_to_only_use_kwargs + ) + else: + assert callable(target) + torch_op_schemas = get_signature_for_torch_op(target) + matched_schemas = [] + if torch_op_schemas: + # Iterate through all of the schema until we find one that matches + # If one matches, populate `new_args_and_kwargs` with the new args/kwargs + # values. If none matches, `new_args_and_kwargs` will be None + for candidate_signature in torch_op_schemas: + try: + candidate_signature.bind(*args, **kwargs) + matched_schemas.append(candidate_signature) + except TypeError: + continue + + if len(matched_schemas) == 0: + # Did not match any schema. Cannot normalize + pass + elif len(matched_schemas) == 1: + # Matched exactly one schema, unambiguous + new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs( + matched_schemas[0], args, kwargs, normalize_to_only_use_kwargs + ) + else: + if arg_types is not None or kwarg_types is not None: + arg_types = arg_types if arg_types else cast(Tuple[Any], ()) + kwarg_types = kwarg_types if kwarg_types else {} + for candidate_signature in torch_op_schemas: + sig_matches = True + try: + bound_types = candidate_signature.bind( + *arg_types, **kwarg_types + ) + for arg_name, arg_type in bound_types.arguments.items(): + param = candidate_signature.parameters[arg_name] + sig_matches = sig_matches and type_matches( + param.annotation, arg_type + ) + except TypeError: + sig_matches = False + if sig_matches: + new_args_and_kwargs = ( + _args_kwargs_to_normalized_args_kwargs( + candidate_signature, + args, + kwargs, + normalize_to_only_use_kwargs, + ) + ) + break + else: + # Matched more than one schema. In this situation, the caller must provide the types of + # the arguments of the overload they expect. + schema_printouts = "\n".join( + str(schema) for schema in matched_schemas + ) + raise RuntimeError( + f"Tried to normalize arguments to {torch.typename(target)} but " + f"the schema match was ambiguous! Please provide argument types to " + f"the normalize_arguments() call. Available schemas:\n{schema_printouts}" + ) + + return new_args_and_kwargs + + +@compatibility(is_backward_compatible=False) +def normalize_module( + root: torch.nn.Module, + target: str, + args: Tuple[Any], + kwargs: Optional[Dict[str, Any]] = None, + normalize_to_only_use_kwargs: bool = False, +) -> Optional[ArgsKwargsPair]: + """ + Returns normalized arguments to PyTorch modules. This means that + `args/kwargs` will be matched up to the functional's + signature and return exclusively kwargs in positional order if + `normalize_to_only_use_kwargs` is True. + Also populates default values. Does not support positional-only + parameters or varargs parameters (*args, **kwargs). + + Args: + root (nn.Module): root module upon which we query modules + target (Callable): Function that we are normalizing + args (Tuple[Any]): Tuple of args to the function + kwargs (Optional[Dict[str, Any]]): Dict of kwargs to the function + normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs. + + Returns: + + Returns normalized_args_and_kwargs, or `None` if not successful. + """ + try: + submod = root.get_submodule(target) + except AttributeError as e: + raise RuntimeError( + f"Tried to normalize node with target {target} but root did not " + f"have that target!" + ) from e + if hasattr(submod.__class__, "__name__"): + classname = submod.__class__.__name__ + if getattr(torch.nn, classname, None) == submod.__class__: + sig = inspect.signature(inspect.unwrap(submod.forward)) + if kwargs is None: + kwargs = {} + new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs( + sig, args, kwargs, normalize_to_only_use_kwargs + ) + return new_args_and_kwargs + return None + + +def _args_kwargs_to_normalized_args_kwargs( + sig: inspect.Signature, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + normalize_to_only_use_kwargs: bool, +) -> Optional[ArgsKwargsPair]: + """ + Given a call target, args, and kwargs, return the arguments normalized into + an ArgsKwargsPair, or None if the type signature is not supported by + this normalization. + + Args: + + sig (inspect.Signature): Signature object for the target + args (Tuple): Arguments that appear at the callsite for `target` + kwargs (Dict): Keyword arguments that appear at the callsite for `target` + normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs. + + Returns: + + Optional[ArgsKwargsPair]: Normalized args and 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()): + # Add an exception for one signature, which is common for random/uniform, i.e.: + # Tensor(a!) self, float from=0, float to=1, *, Generator? generator=None + # `from` is Python keyword and as such functions with that signature should have + # positional-only args, but at the same time they could be dispatched as kwargs + if list(sig.parameters.keys()) != ["input", "from", "to", "generator"]: + return None + + bound_args = sig.bind(*args, **kwargs) + bound_args.apply_defaults() + + new_kwargs: Dict[str, Any] = {} + new_args: List[Any] = [] + for i, param in enumerate(sig.parameters): + if not normalize_to_only_use_kwargs and i < len(args): + new_args.append(bound_args.arguments[param]) + else: + new_kwargs[param] = bound_args.arguments[param] + + return ArgsKwargsPair(tuple(new_args), new_kwargs) diff --git a/lib/python3.12/site-packages/torch/fx/passes/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..433d8818e259a6c0d8d674d15a0312815010ec7f --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/__init__.py @@ -0,0 +1,14 @@ +from . import ( + graph_drawer, + graph_manipulation, + net_min_base, + operator_support, + param_fetch, + reinplace, + runtime_assert, + shape_prop, + split_module, + split_utils, + splitter_base, + tools_common, +) diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d10012bcaa6ace1e347425455d961e06efff0131 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/_tensorify_python_scalars.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/_tensorify_python_scalars.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e66357bbc4cb102bf415fb3f84126e1ff5933c19 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/_tensorify_python_scalars.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/annotate_getitem_nodes.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/annotate_getitem_nodes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f874449bada5c76796995c7a418acf24e384d9ee Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/annotate_getitem_nodes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/fake_tensor_prop.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/fake_tensor_prop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08df18ae1ab22033a90162c50ff5b28c80ad4fb6 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/fake_tensor_prop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_drawer.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_drawer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0113bd23afe3249086c838444f115644039e6a41 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_drawer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_manipulation.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_manipulation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e8720b628697cd04461bac2a6cdc52670a56a57 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_manipulation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_transform_observer.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_transform_observer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75bb107a5a055133c672ba03e4b06265a2781e4f Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/graph_transform_observer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/net_min_base.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/net_min_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bee1b9c19b926b175d0701892d99b56aa75f800f Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/net_min_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/operator_support.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/operator_support.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c16d788536d503c260ff48960537ad4250002583 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/operator_support.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/param_fetch.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/param_fetch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09bdd3afa5aedb229078d6984297bfbce4999576 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/param_fetch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/pass_manager.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/pass_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ca9c496292f9cbb8b66340342f5b7e85473d267 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/pass_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/reinplace.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/reinplace.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9eda558e6023ec25b93a72e8a83c294a3bab0d8a Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/reinplace.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/runtime_assert.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/runtime_assert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3da02895fce9e378744ee42de2ed9f202371d4fa Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/runtime_assert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/shape_prop.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/shape_prop.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e3e34f950c8b054c498a45a89d158f78c63669b Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/shape_prop.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/split_module.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/split_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64ffc7c48a36389c3df05b307a388376cf8bf559 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/split_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/split_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/split_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f89c1554c36b5a62140161066af52f06c81a267 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/split_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/splitter_base.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/splitter_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b9452a8eca5ceb6a961f2a73cae86b0c5d69aab Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/splitter_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/__pycache__/tools_common.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/tools_common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3eb1fafa8ea27456b57e7b4f3b5cb4eb53cd2ca7 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/__pycache__/tools_common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/_tensorify_python_scalars.py b/lib/python3.12/site-packages/torch/fx/passes/_tensorify_python_scalars.py new file mode 100644 index 0000000000000000000000000000000000000000..4ecf440e0b252444886ee37cf9826bb7b1d42b4f --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/_tensorify_python_scalars.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +import logging +import os +from typing import Any, List, Union + +from sympy import Integer, Number, Symbol +from sympy.logic.boolalg import BooleanAtom + +import torch +import torch.fx as fx +from torch._dynamo.exc import TensorifyScalarRestartAnalysis +from torch._dynamo.symbolic_convert import TensorifyState +from torch._prims_common import get_computation_dtype +from torch._subclasses import fake_tensor # noqa: TCH001 +from torch._subclasses.fake_tensor import FakeTensor +from torch._utils_internal import justknobs_check +from torch.fx._utils import lazy_format_graph_code +from torch.fx.experimental.symbolic_shapes import guard_scalar, ShapeEnv # noqa: TCH001 +from torch.fx.graph_module import GraphModule # noqa: TCH001 + +# TODO: refactor +from torch.fx.passes.runtime_assert import _get_sym_val +from torch.fx.proxy import MetaProxy +from torch.utils._sympy.interp import _run_sympy_handler, sympy_interp +from torch.utils._sympy.reference import TensorReferenceAnalysis +from torch.utils._sympy.symbol import symbol_is_type, SymT + + +__all__: List[str] = [] + +log = logging.getLogger(__name__) +graph_code_log = torch._logging.getArtifactLogger(__name__, "graph_code") + +# The general shape of this transformation is to look for Tensor operations +# that take a backed SymFloat as an argument, and then redo them as tensor +# compute (with ints and tensors as inputs). For example, add(Tensor, Scalar) +# can be translated into add(Tensor, Tensor). Because Dynamo has already +# arranged for floats to be Tensor inputs to the graph, for typical float +# compute you can entirely translate the Python float operations into Tensor +# operations with only Tensor inputs. +# +# This pass is also responsible for doing CSE on the fly as we do this, since +# you don't want to keep recomputing the same quantity over and over again if +# it's used multiple times. +# +# This pass runs on the JOINT graph produced by AOT Autograd, prior to partitioning. +# The primary goal of this pass is to eliminate floats by replacing TensorScalar +# operations with TensorTensor operations and then Dead Code Elimination (DCE) of +# the item calls, which effectively removes the floats. +# +# This needs to happen before partitioning because it influences partitioning decisions, +# specifically by ensuring that we don't need to save floats across partitions. +# Additionally, there is a separate pass that changes which device computations +# occur on. That pass must be run after this one, but still before partitioning. +# +# HISTORY NOTE: Originally, I wanted to formulate this pass as pushing item() +# calls down, transforming float compute into int compute as we went. If you +# manage to eliminate all float compute, this ends up being equivalent, but +# there is a critical difference when some floats cannot be eliminated: when +# we call item() on them, what should it's SymFloat be? Ideally, it would +# be the same backed SymFloat we had before. But without symbolic expresssion +# propogation on tensor quantities, repropagating would instead give you an +# unbacked SymFloat. Maybe it is a good idea to implement symbolic propagation +# on 0d scalar tensors, but I decided to go for something simpler to start. +# +# The boring stuff: +# +# * What operators can I Tensor-ify? (Anything with a Scalar argument) +# * How do I Tensor-ify a SymFloat sympy expression (Sympy -> Op Handler -> Tensor) +# +# TODO: make sure this runs before CPU->CUDA pass for cudagraph friendliness + + +SUPPORTED_OPS = { + torch.ops.aten.mul.Tensor: torch.ops.aten.mul.Tensor, + torch.ops.aten.add.Tensor: torch.ops.aten.add.Tensor, + torch.ops.aten.sub.Tensor: torch.ops.aten.sub.Tensor, + torch.ops.aten.div.Tensor: torch.ops.aten.div.Tensor, + torch.ops.aten.gt.Scalar: torch.ops.aten.gt.Tensor, + torch.ops.aten.lt.Scalar: torch.ops.aten.lt.Tensor, + torch.ops.aten.ge.Scalar: torch.ops.aten.ge.Tensor, + torch.ops.aten.le.Scalar: torch.ops.aten.le.Tensor, + torch.ops.aten.eq.Scalar: torch.ops.aten.eq.Tensor, + torch.ops.aten.ne.Scalar: torch.ops.aten.ne.Tensor, +} + + +@torch.fx._compatibility.compatibility(is_backward_compatible=False) +def tensorify_python_scalars( + gm: GraphModule, shape_env: ShapeEnv, fake_mode: fake_tensor.FakeTensorMode +) -> None: + """ + Converts Python scalar operations into Tensor operations within the graph. This pass looks for + Tensor operations that involve SymFloat arguments and transforms them into equivalent operations + that use only Tensor inputs. + + Args: + gm: The FX graph module representing the computation graph. + shape_env: The shape environment responsible for symbolic shape tracking and propagation + during graph transformations. + + Returns: + None + """ + import sympy + + knob = True + if (env := os.getenv("TENSORIFY_PYTHON_SCALARS")) is not None: + if env in ("0", "FALSE"): + knob = False + else: + knob = justknobs_check("pytorch/compiler:tensorify_python_scalars") + if not knob: + return None + + graph = gm.graph + tracer = fx.proxy.GraphAppendingTracer(graph) + expr_to_sym_proxy: dict[sympy.Expr, MetaProxy] = {} + expr_to_tensor_proxy: dict[sympy.Expr, MetaProxy] = {} + tensorified_symbols: set[sympy.Symbol] = set() + should_restart = False + + first_non_placeholder = None + placeholders = set() + for node in graph.nodes: + if node.op != "placeholder": + first_non_placeholder = node + break + else: + placeholders.add(node) + + Analysis = TensorReferenceAnalysis + + def _sympy_interp(expr: sympy.Expr) -> MetaProxy: + # sympy_interp() with hash consing, and special handling for + # generating constants correctly + + # hash cons + if isinstance(expr, Symbol) and expr not in expr_to_tensor_proxy: + # This is guaranteed to be populated by invariant established by + # insert_deferred_runtime_asserts + expr_to_tensor_proxy[expr] = torch.ops.aten.scalar_tensor.default( + expr_to_sym_proxy[expr] + ) + + # cache constants, why not + if isinstance(expr, (Integer, Number, BooleanAtom)): + dtype = None + c: Union[bool, int, float] + if isinstance(expr, BooleanAtom): + dtype = torch.bool + c = bool(expr) + elif isinstance(expr, sympy.Integer): + dtype = torch.int64 + c = int(expr) + elif isinstance(expr, sympy.Number): + dtype = torch.float64 + c = float(expr) + + node = graph.call_function( + torch.ops.aten.scalar_tensor.default, (c,), {"dtype": dtype} + ) + with fake_mode: + node.meta["val"] = torch.ops.aten.scalar_tensor.default(c, dtype=dtype) + expr_to_tensor_proxy[expr] = MetaProxy( + node, + tracer=tracer, + fake_mode=fake_mode, + ) + + if expr in expr_to_tensor_proxy: + return expr_to_tensor_proxy[expr] + + # don't cache + if isinstance(expr, Symbol): + return sympy_interp(Analysis, expr_to_tensor_proxy, expr) # type: ignore[arg-type] + + # hash cons on arguments, run expr handler + expr_to_tensor_proxy[expr] = _run_sympy_handler( + Analysis, + [_sympy_interp(arg) for arg in expr.args], # type: ignore[arg-type] + expr, + ) + + return expr_to_tensor_proxy[expr] + + nodes = list(graph.nodes) + for i, node in enumerate(nodes[:-1]): + with graph.inserting_before( + nodes[i + 1] if node not in placeholders else first_non_placeholder + ): + # Look for tensor.item() calls on placeholders + if ( + node is not None + and node.op == "call_function" + and node.target is torch.ops.aten._local_scalar_dense.default + ): + dtype = node.args[0].meta["val"].dtype + if dtype != torch.float64: + continue + + assert isinstance(node.args[0], fx.Node), node.args[0] + + s = node.meta["val"].node.expr + expr_to_tensor_proxy[s] = MetaProxy( + node.args[0], tracer=tracer, fake_mode=fake_mode + ) + expr_to_sym_proxy[s] = MetaProxy( + node, tracer=tracer, fake_mode=fake_mode + ) + elif (sym_expr := _get_sym_val(node)) is not None: + if sym_expr not in expr_to_sym_proxy and not isinstance( + sym_expr, (sympy.Number, sympy.logic.boolalg.BooleanAtom) + ): + expr_to_sym_proxy[sym_expr] = MetaProxy( + node, tracer=tracer, fake_mode=fake_mode + ) + + # Specialize all dimensions that contain symfloats. Here's + # an example test that requires this: + # PYTORCH_OPINFO_SAMPLE_INPUT_INDEX=4 python test/inductor/test_torchinductor_opinfo.py TestInductorOpInfoCUDA.test_comprehensive_nn_functional_interpolate_bicubic_cuda_float32 # noqa: B950 + val = node.meta.get("val") + if isinstance(val, FakeTensor): + for dim in val.shape: + if not isinstance(dim, torch.SymInt): + continue + + for symbol in dim.node.expr.free_symbols: + if not symbol_is_type(symbol, SymT.FLOAT): + continue + + sources = shape_env.var_to_sources.get(symbol) + for source in sources: + if TensorifyState.should_specialize(source): + continue + + # In principle, we could support float input that + # is used to do size compute. The problem is that + # we don't actually want to tensorify the compute + # in this case, which means we need codegen support + # for all symfloats. + TensorifyState.specialize(source) + should_restart = True + + # Look for functions to convert + if node.op == "call_function" and ( + replacement_op := SUPPORTED_OPS.get(node.target) + ): + args: List[Any] = [] + transform = False + compute_dtype = get_computation_dtype(node.meta["val"].dtype) + + for a in node.args: + if ( + isinstance(a, fx.Node) + and "val" in a.meta + and isinstance(zf := a.meta["val"], torch.SymFloat) + ): + transform = True + try: + proxy = _sympy_interp(zf.node.expr) + except NotImplementedError: + transform = False + break + + # We use _expr instead of expr b/c we want the symbol not the replacement + tensorified_symbols.add(a.meta["val"].node._expr) + + # The upcasting is irrelevant when the compute dtype is bool. This happens + # in cases where we are tensorifying a comparison operator such as + # torch.ops.aten.gt.Tensor + if ( + compute_dtype != torch.bool + and proxy.node.meta["val"].dtype != compute_dtype + ): + proxy = torch.ops.prims.convert_element_type.default( + proxy, compute_dtype + ) + + args.append(proxy) + elif isinstance(a, fx.Node): + args.append(MetaProxy(a, tracer=tracer, fake_mode=fake_mode)) + else: + args.append(a) + + if transform: + replacement_proxy = replacement_op(*args) + + if compute_dtype != node.meta["val"].dtype: + replacement_proxy = ( + torch.ops.prims.convert_element_type.default( + replacement_proxy, + node.meta["val"].dtype, + ) + ) + + node.replace_all_uses_with(replacement_proxy.node) + graph.erase_node(node) + + # Now do one more pass that specializes all symfloats we didn't manage + # to tensorify away. + for node in reversed(graph.nodes): + if node.op == "output" or node.op == "placeholder": + continue + + with graph.inserting_before(node): + if len(node.users) == 0 and not node.is_impure(): + graph.erase_node(node) + continue + + if isinstance( + (val := node.meta.get("val")), + (torch.SymFloat, torch.SymInt, torch.SymBool), + ): + if all( + symbol_is_type(s, SymT.FLOAT) for s in val.node.expr.free_symbols + ): + # If all symbols are backed symfloats, we can just specialize the whole node + # and get more precise guards. eg. + # + # zf = a.item() + # zf2 = zf // 2 + # op(.. zf2 ..) + # + # It's better to guard on zf // 2 == 2.0 than zf == 5.0 + + node.replace_all_uses_with(guard_scalar(val)) + graph.erase_node(node) + + for symbol, sources in shape_env.var_to_sources.items(): + if symbol_is_type(symbol, SymT.FLOAT) and symbol not in tensorified_symbols: + for source in sources: + if not TensorifyState.should_specialize(source): + TensorifyState.specialize(source) + should_restart = True + + if should_restart: + # Sledgehammer time. Restart dynamo analysis, keeping track of which input sources + # are no longer needed and should be specialized. Restarting analysis is necessary + # because we need to instruct Dynamo to NOT make these as inputs. + raise TensorifyScalarRestartAnalysis + + graph_code_log.debug( + "%s", lazy_format_graph_code("tensorify_python_scalars", gm, colored=True) + ) diff --git a/lib/python3.12/site-packages/torch/fx/passes/annotate_getitem_nodes.py b/lib/python3.12/site-packages/torch/fx/passes/annotate_getitem_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..0399cef526205f8f82a0c53555bc16fdab67a550 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/annotate_getitem_nodes.py @@ -0,0 +1,44 @@ +import operator + +import torch + + +def annotate_getitem_nodes(graph: torch.fx.Graph) -> None: + """ + Annotate the type of getitem nodes, inferred from the type of sequence node. + If sequence node is not annotated with a type, do nothing. + Currently support getitem nodes from Tuple, List, and NamedTuple sequence node. + + This is helpful since annotations on local names within function are lost during FX transforms. + Adding back known type annotation for getitem nodes to improve jit scriptability. + + Args: + graph (Graph): The graph to be annotated + """ + for node in graph.nodes: + if node.target == operator.getitem: + sequence_node, index_node = node.args + if not sequence_node.type: + continue + # container types + if hasattr(sequence_node.type, "_name"): + parameterized_types = sequence_node.type.__args__ + if sequence_node.type._name == "Tuple": + if len(parameterized_types) == 2 and isinstance( + parameterized_types[1], type(...) + ): + node.type = parameterized_types[0] + else: + assert len(parameterized_types) > index_node + node_type = parameterized_types[index_node] + node.type = node_type + elif sequence_node.type._name == "List": + assert len(parameterized_types) == 1 + node.type = parameterized_types[0] + # NamedTuple type + elif hasattr(sequence_node.type, "__annotations__"): + if sequence_node.type == torch.Tensor: + continue + sequence_node_field_types = sequence_node.type.__annotations__ + field_name = sequence_node.type._fields[index_node] + node.type = sequence_node_field_types[field_name] diff --git a/lib/python3.12/site-packages/torch/fx/passes/backends/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/fx/passes/backends/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/backends/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ca73b17f661f32538e04c77835140abecd7043e Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/backends/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/backends/__pycache__/cudagraphs.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/backends/__pycache__/cudagraphs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8b78e24f7e6cb9691c8fbaf7bbeed065f6ebf90 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/backends/__pycache__/cudagraphs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/backends/cudagraphs.py b/lib/python3.12/site-packages/torch/fx/passes/backends/cudagraphs.py new file mode 100644 index 0000000000000000000000000000000000000000..b98178f0d5339321673deffdac6f03a96ffbde45 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/backends/cudagraphs.py @@ -0,0 +1,61 @@ +# mypy: allow-untyped-defs +import operator + +import torch +from torch.fx.passes.fake_tensor_prop import FakeTensorProp +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner +from torch.fx.passes.operator_support import OperatorSupport +from torch.fx.passes.tools_common import CALLABLE_NODE_OPS +from torch.utils import _pytree as pytree + + +class CudaGraphsSupport(OperatorSupport): + # TODO: why is submodules passed here + def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: + if node.op not in CALLABLE_NODE_OPS: + return False + + if node.target in [torch.ops.aten.embedding_dense_backward.default]: + return False + + if node.target in [operator.getitem]: + return True + + found_not_cuda = False + + def meta_fk(meta): + return meta["val"] if "val" in meta else meta["fake_result"] + + def find_not_cuda(t): + nonlocal found_not_cuda + if isinstance(t, torch.Tensor) and t.device.type != "cuda": + found_not_cuda = True + + for n in node.all_input_nodes: + pytree.tree_map_(find_not_cuda, meta_fk(n.meta)) + + pytree.tree_map_(find_not_cuda, meta_fk(node.meta)) + + # NB: factory function is accounted for because the result would be + # cpu or cuda + + return not found_not_cuda + + +def partition_cudagraphs(gm, inputs): + """ + Partition an FX graph into sub-GraphModules that can be validly run under + CUDA graphs. For a subgraph to be runnable under CUDA, all of the operations + must involve CUDA tensors only/ + """ + + FakeTensorProp(gm).propagate(*inputs) + supported_ops = CudaGraphsSupport() + # TODO: single node partition may be wrong due to the pessimization + # from copying in and out the data. Check in benchmarks, perhaps + partitioner = CapabilityBasedPartitioner( + gm, supported_ops, allows_single_node_partition=True + ) + partitions = partitioner.propose_partitions() + fused_graph = partitioner.fuse_partitions(partitions) + return fused_graph diff --git a/lib/python3.12/site-packages/torch/fx/passes/dialect/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/dialect/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/fx/passes/dialect/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/dialect/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..10d58892c1292f06a62c9e65f0cfe62544a845e4 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/dialect/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..851a778733388726f4f2f8fd8e09c635a840481c Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__pycache__/cse_pass.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__pycache__/cse_pass.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..43b4d0db636fd9e17d166da11f37fd076430a010 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/__pycache__/cse_pass.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/dialect/common/cse_pass.py b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/cse_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..6a501f041d193d79cac8acea049838fb94e92830 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/dialect/common/cse_pass.py @@ -0,0 +1,155 @@ +# mypy: allow-untyped-defs +from typing import Any, Dict, Tuple + +import torch +from torch.fx import Graph, GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.utils._pytree import tree_flatten + + +aten = torch.ops.aten + + +# stateful ops are banned from CSE +rand_ops = { + aten.dropout, + aten._fused_dropout, + aten._standard_gamma, + aten.bernoulli, + aten.multinomial, + aten.native_dropout, + aten.normal, + aten.poisson, + aten.binomial, + aten.rrelu, + aten.rand_like, + aten.rand, + aten.randint, + aten.randn, + aten.randperm, +} # noqa: E501,B950 + +inplace_ops = { + aten.add_, + aten.sub_, + aten.mul_, + aten.div_, + aten.pow_, + aten.lerp_, + aten.relu_, + aten.sigmoid_, + aten.tanh_, +} # noqa: E501 + + +@torch.fx._compatibility.compatibility(is_backward_compatible=False) +def get_CSE_banned_ops(): + return rand_ops.union(inplace_ops) + + +@torch.fx._compatibility.compatibility(is_backward_compatible=False) +class CSEPass(PassBase): + def __init__(self, banned_ops=None): + """ + This version of CSE Pass aims to be dialect agnostic, and it's implemented purely based on the connectivity between fx.Node. + + For functional dialects, user would only need to specify the random ops in ban list. + + Warning: CSE Pass cannot be safely applied on a FX graph in non-functional dialects. + If your dialect contains stateful operators, please customized the banned_ops. + + """ + if banned_ops is None: + banned_ops = set() + self.banned_ops = banned_ops + super().__init__() + + def call(self, graph_module: GraphModule) -> PassResult: + """ + Return a new copy of torch.fx.GraphModule with CSE applied to the input graph + + Example usage: + + from torch.fx.experimental.proxy_tensor import make_fx + def f(a): + b = a * a + c = a * a + return b+c + + p = CSEPass() + traced_graph = make_fx(f)(torch.tensor(1)) + print(traced_graph) + result = p(traced_graph) + print(result.graph_module) + """ + + def get_aten_target(node): + if hasattr(node.target, "overloadpacket"): + return node.target.overloadpacket + return node.target + + modified = False + new_graph = Graph() + env: Dict[ + Node, Node + ] = {} # map from node in the old graph to node in the new graph + hash_env: Dict[ + Tuple[torch._ops.OpOverload, int], Node + ] = {} # map from hash to a node in the new graph + token_map: Dict[ + Tuple[torch._ops.OpOverload, int], Dict[str, Any] + ] = {} # map from hash to token + for n in graph_module.graph.nodes: + # The placeholder, output, and get_attr nodes are copied to the new graph without change + # do not CSE away random operations + if ( + n.op == "placeholder" + or n.op == "output" + or n.op == "get_attr" + or get_aten_target(n) in self.banned_ops + ): + new_node = new_graph.node_copy(n, lambda x: env[x]) + env[n] = new_node + else: # n.op == 'call_function', should never see n.op == 'call_module' or 'call_method' + # substitute args and kwargs members to their mapping in env if exists + # specs can be used to reconstruct nested list/dictionaries + def substitute(arg_list): + arg_list, spec = tree_flatten(arg_list) + for i in range(len(arg_list)): + v = arg_list[i] + if isinstance(v, Node) and v in env: + arg_list[i] = env[v] + return tuple(arg_list), spec + + args, args_spec = substitute(n.args) + kwargs, kwargs_spec = substitute(n.kwargs) + + # each token corresponds to a unique node + # nodes with the same token can be substituted + token = { + "target": n.target, + "args": args, + "args_spec": args_spec, + "kwargs": kwargs, + "kwargs_spec": kwargs_spec, + } + + # hash substituted args to a number, do not hash specs because specs are not hashable + hash_arg = hash((args, kwargs)) + hash_val = (n.target, hash_arg) + + # check if a node has a substitute and can be eliminated + hash_val_in_hash_env = hash_val in hash_env + if hash_val_in_hash_env and token_map[hash_val] == token: + modified = True # substitution happens and the graph is modified + env[n] = hash_env[hash_val] + continue + + new_node = new_graph.node_copy(n, lambda x: env[x]) + env[n] = new_node + if not hash_val_in_hash_env: + hash_env[hash_val] = new_node + token_map[hash_val] = token + + csed_gm = GraphModule(graph_module, new_graph) + return PassResult(csed_gm, modified) diff --git a/lib/python3.12/site-packages/torch/fx/passes/fake_tensor_prop.py b/lib/python3.12/site-packages/torch/fx/passes/fake_tensor_prop.py new file mode 100644 index 0000000000000000000000000000000000000000..8036f5d0fd556879ded2d622a284602d240b8534 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/fake_tensor_prop.py @@ -0,0 +1,80 @@ +# mypy: allow-untyped-defs +from typing import Optional + +import torch.fx +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode +from torch.fx import Node +from torch.fx._compatibility import compatibility +from torch.fx.experimental.proxy_tensor import py_sym_types, snapshot_fake +from torch.fx.node import map_aggregate + + +__all__ = ["FakeTensorProp"] + + +@compatibility(is_backward_compatible=False) +class FakeTensorProp(torch.fx.Interpreter): + """ + Execute an FX graph Node-by-Node and record a fake tensor representing + the metadata for the node. Unlike ShapeProp, (1) this propagation + is cheap--it does the propagation with meta tensors which do not actually + store data, and (2) the fake tensors have much more fine grained information, + e.g., they have accurate alias information that can be consulted by looking + at the storages. + + Args: + module (GraphModule): The module to be executed + mode (Optional[FakeTensorMode]): The dispatch mode used to execute computation indicated by each FX Node. + """ + + def __init__( + self, module: torch.fx.GraphModule, mode: Optional[FakeTensorMode] = None + ): + super().__init__(module) + if mode is None: + mode = FakeTensorMode() + self._mode = mode + mode.epoch += 1 + mode.reset_nt_tensor_id_counter() + + def run_node(self, n: Node): + from torch.fx.experimental.symbolic_shapes import ( + compute_unbacked_bindings, + rebind_unbacked, + ) + + result = super().run_node(n) + rebind_unbacked(self._mode.shape_env, n, result) + + def extract_val(obj): + if isinstance(obj, FakeTensor): + return snapshot_fake(obj) + elif isinstance(obj, torch.Tensor): + # TODO: How is it possible that we get a non fake tensor? We + # should be running under the mode... + return snapshot_fake(self._mode.from_tensor(obj, static_shapes=True)) + elif isinstance(obj, py_sym_types): + return obj + else: + return None + + meta = map_aggregate(result, extract_val) + if meta is not None: + n.meta["val"] = meta + if (shape_env := self._mode.shape_env) and ( + symbol_to_path := compute_unbacked_bindings(shape_env, result) + ): + n.meta["unbacked_bindings"] = symbol_to_path + + return result + + def propagate(self, *args): + fake_args = [ + self._mode.from_tensor(a) if isinstance(a, torch.Tensor) else a + for a in args + ] + return self.propagate_dont_convert_inputs(*fake_args) + + def propagate_dont_convert_inputs(self, *args): + with self._mode: + return super().run(*args) diff --git a/lib/python3.12/site-packages/torch/fx/passes/graph_drawer.py b/lib/python3.12/site-packages/torch/fx/passes/graph_drawer.py new file mode 100644 index 0000000000000000000000000000000000000000..d8b95231f891a5fd2c9857dc80fce7f6b20447be --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/graph_drawer.py @@ -0,0 +1,489 @@ +# mypy: allow-untyped-defs + +import hashlib +from itertools import chain +from typing import Any, Dict, Optional, TYPE_CHECKING + +import torch +import torch.fx +from torch.fx._compatibility import compatibility +from torch.fx.graph import _parse_stack_trace +from torch.fx.node import _format_arg, _get_qualified_name +from torch.fx.operator_schemas import normalize_function +from torch.fx.passes.shape_prop import TensorMetadata + + +try: + import pydot + + HAS_PYDOT = True +except ModuleNotFoundError: + HAS_PYDOT = False + pydot = None + + +__all__ = ["FxGraphDrawer"] + +_COLOR_MAP = { + "placeholder": '"AliceBlue"', + "call_module": "LemonChiffon1", + "get_param": "Yellow2", + "get_attr": "LightGrey", + "output": "PowderBlue", +} + +_HASH_COLOR_MAP = [ + "CadetBlue1", + "Coral", + "DarkOliveGreen1", + "DarkSeaGreen1", + "GhostWhite", + "Khaki1", + "LavenderBlush1", + "LightSkyBlue", + "MistyRose1", + "MistyRose2", + "PaleTurquoise2", + "PeachPuff1", + "Salmon", + "Thistle1", + "Thistle3", + "Wheat1", +] + +_WEIGHT_TEMPLATE = { + "fillcolor": "Salmon", + "style": '"filled,rounded"', + "fontcolor": "#000000", +} + +if HAS_PYDOT: + + @compatibility(is_backward_compatible=False) + class FxGraphDrawer: + """ + Visualize a torch.fx.Graph with graphviz + Basic usage: + g = FxGraphDrawer(symbolic_traced, "resnet18") + g.get_dot_graph().write_svg("a.svg") + """ + + def __init__( + self, + graph_module: torch.fx.GraphModule, + name: str, + ignore_getattr: bool = False, + ignore_parameters_and_buffers: bool = False, + skip_node_names_in_args: bool = True, + parse_stack_trace: bool = False, + dot_graph_shape: Optional[str] = None, + normalize_args: bool = False, + ): + self._name = name + self.dot_graph_shape = ( + dot_graph_shape if dot_graph_shape is not None else "record" + ) + self.normalize_args = normalize_args + _WEIGHT_TEMPLATE["shape"] = self.dot_graph_shape + + self._dot_graphs = { + name: self._to_dot( + graph_module, + name, + ignore_getattr, + ignore_parameters_and_buffers, + skip_node_names_in_args, + parse_stack_trace, + ) + } + + for node in graph_module.graph.nodes: + if node.op != "call_module": + continue + + leaf_node = self._get_leaf_node(graph_module, node) + + if not isinstance(leaf_node, torch.fx.GraphModule): + continue + + self._dot_graphs[f"{name}_{node.target}"] = self._to_dot( + leaf_node, + f"{name}_{node.target}", + ignore_getattr, + ignore_parameters_and_buffers, + skip_node_names_in_args, + parse_stack_trace, + ) + + def get_dot_graph(self, submod_name=None) -> pydot.Dot: + """ + Visualize a torch.fx.Graph with graphviz + Example: + >>> # xdoctest: +REQUIRES(module:pydot) + >>> # xdoctest: +REQUIRES(module:ubelt) + >>> # define module + >>> class MyModule(torch.nn.Module): + >>> def __init__(self) -> None: + >>> super().__init__() + >>> self.linear = torch.nn.Linear(4, 5) + >>> def forward(self, x): + >>> return self.linear(x).clamp(min=0.0, max=1.0) + >>> module = MyModule() + >>> # trace the module + >>> symbolic_traced = torch.fx.symbolic_trace(module) + >>> # setup output file + >>> import ubelt as ub + >>> dpath = ub.Path.appdir("torch/tests/FxGraphDrawer").ensuredir() + >>> fpath = dpath / "linear.svg" + >>> # draw the graph + >>> g = FxGraphDrawer(symbolic_traced, "linear") + >>> g.get_dot_graph().write_svg(fpath) + """ + if submod_name is None: + return self.get_main_dot_graph() + else: + return self.get_submod_dot_graph(submod_name) + + def get_main_dot_graph(self) -> pydot.Dot: + return self._dot_graphs[self._name] + + def get_submod_dot_graph(self, submod_name) -> pydot.Dot: + return self._dot_graphs[f"{self._name}_{submod_name}"] + + def get_all_dot_graphs(self) -> Dict[str, pydot.Dot]: + return self._dot_graphs + + def _get_node_style(self, node: torch.fx.Node) -> Dict[str, str]: + template = { + "shape": self.dot_graph_shape, + "fillcolor": "#CAFFE3", + "style": '"filled,rounded"', + "fontcolor": "#000000", + } + if node.op in _COLOR_MAP: + template["fillcolor"] = _COLOR_MAP[node.op] + else: + # Use a random color for each node; based on its name so it's stable. + target_name = node._pretty_print_target(node.target) + target_hash = int(hashlib.md5(target_name.encode()).hexdigest()[:8], 16) + template["fillcolor"] = _HASH_COLOR_MAP[ + target_hash % len(_HASH_COLOR_MAP) + ] + return template + + def _get_leaf_node( + self, module: torch.nn.Module, node: torch.fx.Node + ) -> torch.nn.Module: + py_obj = module + assert isinstance(node.target, str) + atoms = node.target.split(".") + for atom in atoms: + if not hasattr(py_obj, atom): + raise RuntimeError( + str(py_obj) + " does not have attribute " + atom + "!" + ) + py_obj = getattr(py_obj, atom) + return py_obj + + def _typename(self, target: Any) -> str: + if isinstance(target, torch.nn.Module): + ret = torch.typename(target) + elif isinstance(target, str): + ret = target + else: + ret = _get_qualified_name(target) + + # Escape "{" and "}" to prevent dot files like: + # https://gist.github.com/SungMinCho/1a017aab662c75d805c5954d62c5aabc + # which triggers `Error: bad label format (...)` from dot + return ret.replace("{", r"\{").replace("}", r"\}") + + # shorten path to avoid drawing long boxes + # for full path = '/home/weif/pytorch/test.py' + # return short path = 'pytorch/test.py' + def _shorten_file_name( + self, + full_file_name: str, + truncate_to_last_n: int = 2, + ): + splits = full_file_name.split("/") + if len(splits) >= truncate_to_last_n: + return "/".join(splits[-truncate_to_last_n:]) + return full_file_name + + def _get_node_label( + self, + module: torch.fx.GraphModule, + node: torch.fx.Node, + skip_node_names_in_args: bool, + parse_stack_trace: bool, + ) -> str: + def _get_str_for_args_kwargs(arg): + if isinstance(arg, tuple): + prefix, suffix = r"|args=(\l", r",\n)\l" + arg_strs_list = [_format_arg(a, max_list_len=8) for a in arg] + elif isinstance(arg, dict): + prefix, suffix = r"|kwargs={\l", r",\n}\l" + arg_strs_list = [ + f"{k}: {_format_arg(v, max_list_len=8)}" for k, v in arg.items() + ] + else: # Fall back to nothing in unexpected case. + return "" + + # Strip out node names if requested. + if skip_node_names_in_args: + arg_strs_list = [a for a in arg_strs_list if "%" not in a] + if len(arg_strs_list) == 0: + return "" + arg_strs = prefix + r",\n".join(arg_strs_list) + suffix + if len(arg_strs_list) == 1: + arg_strs = arg_strs.replace(r"\l", "").replace(r"\n", "") + return arg_strs.replace("{", r"\{").replace("}", r"\}") + + label = "{" + f"name=%{node.name}|op_code={node.op}\n" + + if node.op == "call_module": + leaf_module = self._get_leaf_node(module, node) + label += r"\n" + self._typename(leaf_module) + r"\n|" + extra = "" + if hasattr(leaf_module, "__constants__"): + extra = r"\n".join( + [ + f"{c}: {getattr(leaf_module, c)}" + for c in leaf_module.__constants__ # type: ignore[union-attr] + ] # type: ignore[union-attr] + ) + label += extra + r"\n" + else: + label += f"|target={self._typename(node.target)}" + r"\n" + if self.normalize_args: + try: + args, kwargs = normalize_function( # type: ignore[misc] + node.target, # type: ignore[arg-type] + node.args, # type: ignore[arg-type] + node.kwargs, + normalize_to_only_use_kwargs=True, + ) + except Exception: + # Fallback to not normalizing if there's an exception. + # Some functions need overloads specified to normalize. + args, kwargs = node.args, node.kwargs + else: + args, kwargs = node.args, node.kwargs + if len(args) > 0: + label += _get_str_for_args_kwargs(args) + if len(kwargs) > 0: + label += _get_str_for_args_kwargs(kwargs) + label += f"|num_users={len(node.users)}" + r"\n" + + tensor_meta = node.meta.get("tensor_meta") + label += self._tensor_meta_to_label(tensor_meta) + + # for original fx graph + # print buf=buf0, n_origin=6 + buf_meta = node.meta.get("buf_meta", None) + if buf_meta is not None: + label += f"|buf={buf_meta.name}" + r"\n" + label += f"|n_origin={buf_meta.n_origin}" + r"\n" + + # for original fx graph + # print file:lineno code + if parse_stack_trace and node.stack_trace is not None: + parsed_stack_trace = _parse_stack_trace(node.stack_trace) + fname = self._shorten_file_name(parsed_stack_trace.file) + label += ( + f"|file={fname}:{parsed_stack_trace.lineno} {parsed_stack_trace.code}" + + r"\n" + ) + + return label + "}" + + def _tensor_meta_to_label(self, tm) -> str: + if tm is None: + return "" + elif isinstance(tm, TensorMetadata): + return self._stringify_tensor_meta(tm) + elif isinstance(tm, list): + result = "" + for item in tm: + result += self._tensor_meta_to_label(item) + return result + elif isinstance(tm, dict): + result = "" + for v in tm.values(): + result += self._tensor_meta_to_label(v) + return result + elif isinstance(tm, tuple): + result = "" + for item in tm: + result += self._tensor_meta_to_label(item) + return result + else: + raise RuntimeError(f"Unsupported tensor meta type {type(tm)}") + + def _stringify_tensor_meta(self, tm: TensorMetadata) -> str: + result = "" + if not hasattr(tm, "dtype"): + print("tm", tm) + result += "|" + "dtype" + "=" + str(tm.dtype) + r"\n" + result += "|" + "shape" + "=" + str(tuple(tm.shape)) + r"\n" + result += "|" + "requires_grad" + "=" + str(tm.requires_grad) + r"\n" + result += "|" + "stride" + "=" + str(tm.stride) + r"\n" + if tm.is_quantized: + assert tm.qparams is not None + assert "qscheme" in tm.qparams + qscheme = tm.qparams["qscheme"] + if qscheme in { + torch.per_tensor_affine, + torch.per_tensor_symmetric, + }: + result += "|" + "q_scale" + "=" + str(tm.qparams["scale"]) + r"\n" + result += ( + "|" + + "q_zero_point" + + "=" + + str(tm.qparams["zero_point"]) + + r"\n" + ) + elif qscheme in { + torch.per_channel_affine, + torch.per_channel_symmetric, + torch.per_channel_affine_float_qparams, + }: + result += ( + "|" + + "q_per_channel_scale" + + "=" + + str(tm.qparams["scale"]) + + r"\n" + ) + result += ( + "|" + + "q_per_channel_zero_point" + + "=" + + str(tm.qparams["zero_point"]) + + r"\n" + ) + result += ( + "|" + + "q_per_channel_axis" + + "=" + + str(tm.qparams["axis"]) + + r"\n" + ) + else: + raise RuntimeError(f"Unsupported qscheme: {qscheme}") + result += "|" + "qscheme" + "=" + str(tm.qparams["qscheme"]) + r"\n" + return result + + def _get_tensor_label(self, t: torch.Tensor) -> str: + return str(t.dtype) + str(list(t.shape)) + r"\n" + + # when parse_stack_trace=True + # print file:lineno code + def _to_dot( + self, + graph_module: torch.fx.GraphModule, + name: str, + ignore_getattr: bool, + ignore_parameters_and_buffers: bool, + skip_node_names_in_args: bool, + parse_stack_trace: bool, + ) -> pydot.Dot: + """ + Actual interface to visualize a fx.Graph. Note that it takes in the GraphModule instead of the Graph. + If ignore_parameters_and_buffers is True, the parameters and buffers + created with the module will not be added as nodes and edges. + """ + + # "TB" means top-to-bottom rank direction in layout + dot_graph = pydot.Dot(name, rankdir="TB") + + buf_name_to_subgraph = {} + + for node in graph_module.graph.nodes: + if ignore_getattr and node.op == "get_attr": + continue + + style = self._get_node_style(node) + dot_node = pydot.Node( + node.name, + label=self._get_node_label( + graph_module, node, skip_node_names_in_args, parse_stack_trace + ), + **style, + ) + + current_graph = dot_graph + + buf_meta = node.meta.get("buf_meta", None) + if buf_meta is not None and buf_meta.n_origin > 1: + buf_name = buf_meta.name + if buf_name not in buf_name_to_subgraph: + buf_name_to_subgraph[buf_name] = pydot.Cluster( + buf_name, label=buf_name + ) + current_graph = buf_name_to_subgraph.get(buf_name) + + current_graph.add_node(dot_node) + + def get_module_params_or_buffers(): + for pname, ptensor in chain( + leaf_module.named_parameters(), leaf_module.named_buffers() + ): + pname1 = node.name + "." + pname + label1 = ( + pname1 + "|op_code=get_" + "parameter" + if isinstance(ptensor, torch.nn.Parameter) + else "buffer" + r"\l" + ) + dot_w_node = pydot.Node( + pname1, + label="{" + label1 + self._get_tensor_label(ptensor) + "}", + **_WEIGHT_TEMPLATE, + ) + dot_graph.add_node(dot_w_node) + dot_graph.add_edge(pydot.Edge(pname1, node.name)) + + if node.op == "call_module": + leaf_module = self._get_leaf_node(graph_module, node) + + if not ignore_parameters_and_buffers and not isinstance( + leaf_module, torch.fx.GraphModule + ): + get_module_params_or_buffers() + + for subgraph in buf_name_to_subgraph.values(): + subgraph.set("color", "royalblue") + subgraph.set("penwidth", "2") + dot_graph.add_subgraph(subgraph) + + for node in graph_module.graph.nodes: + if ignore_getattr and node.op == "get_attr": + continue + + for user in node.users: + dot_graph.add_edge(pydot.Edge(node.name, user.name)) + + return dot_graph + +else: + if not TYPE_CHECKING: + + @compatibility(is_backward_compatible=False) + class FxGraphDrawer: + def __init__( + self, + graph_module: torch.fx.GraphModule, + name: str, + ignore_getattr: bool = False, + ignore_parameters_and_buffers: bool = False, + skip_node_names_in_args: bool = True, + parse_stack_trace: bool = False, + dot_graph_shape: Optional[str] = None, + normalize_args: bool = False, + ): + raise RuntimeError( + "FXGraphDrawer requires the pydot package to be installed. Please install " + "pydot through your favorite Python package manager." + ) diff --git a/lib/python3.12/site-packages/torch/fx/passes/graph_manipulation.py b/lib/python3.12/site-packages/torch/fx/passes/graph_manipulation.py new file mode 100644 index 0000000000000000000000000000000000000000..ce9904fc500e8b9ecbc107b7409970ebf5f45141 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/graph_manipulation.py @@ -0,0 +1,113 @@ +# mypy: allow-untyped-defs +from typing import Any, Dict, List, NamedTuple, Optional + +import torch +from torch.fx._compatibility import compatibility +from torch.fx.graph import Graph +from torch.fx.graph_module import GraphModule +from torch.fx.node import map_arg, Node, Target +from torch.fx.passes.shape_prop import ShapeProp + + +__all__ = [ + "replace_target_nodes_with", + "size_bytes", + "get_size_of_all_nodes", + "get_tensor_meta", + "get_size_of_node", +] + + +@compatibility(is_backward_compatible=False) +def replace_target_nodes_with( + fx_module: GraphModule, + old_op: str, + old_target: Target, + new_op: str, + new_target: Target, +): + """Modifies all nodes in fx_module.graph.nodes which match the specified op code and target, + and updates them to match the new op code and target""" + new_graph = Graph() + val_map: Dict[Node, Node] = {} + for node in fx_module.graph.nodes: + if node.op == old_op and node.target == old_target: + args = map_arg(node.args, lambda n: val_map[n]) + kwargs = map_arg(node.kwargs, lambda n: val_map[n]) + assert isinstance(args, tuple) + assert isinstance(kwargs, dict) + val_map[node] = new_graph.create_node( + new_op, new_target, args, kwargs, node.name + ) + else: + val_map[node] = new_graph.node_copy(node, lambda n: val_map[n]) + fx_module.graph = new_graph + + +@compatibility(is_backward_compatible=False) +class size_bytes(NamedTuple): + output_size: int + total_size: int + + +@compatibility(is_backward_compatible=False) +def get_size_of_all_nodes( + fx_module: GraphModule, args: Optional[List[torch.Tensor]] = None +) -> None: + """Given a fx graph module, update each node with its total size (weights + bias + output) + and its output_size(output). For a non-module node, the total size is the output size. + return total size""" + if args is not None: + # Mark shape and dtype for each node (node.shape and node.dtype) + ShapeProp(fx_module).propagate(*args) + # Calculate the total size of the whole fx graph + for node in fx_module.graph.nodes: + if node.op == "output": + break + node.size_bytes = get_size_of_node(fx_module, node) + return + + +@compatibility(is_backward_compatible=False) +def get_tensor_meta(node: Node) -> Any: + tensor_meta = node.meta.get("tensor_meta") + + if not tensor_meta: + raise RuntimeError( + f"Node {node} has no tensor metadata associated with it! " + f"Check that shape propagation has run." + ) + + return tensor_meta + + +@compatibility(is_backward_compatible=False) +def get_size_of_node(fx_module: GraphModule, node: Node) -> size_bytes: + """Given a node with node.dtype and node.shape, return its total size and its output size. + total_size = weights + bias + output_size + """ + # Total num of elements + total_num_of_elems = 0 + # For a module, conside all parameters + if node.op == "call_module": + submodule_dict = dict(fx_module.named_modules()) + submodule = submodule_dict[node.target] + parameters = submodule.named_parameters() + # Parameters are named tuples + for _name, p in parameters: + total_num_of_elems += p.numel() + # Don't forget the output size + # node.shape is the shape of this node's output + tensor_meta = get_tensor_meta(node) + output_elem = tensor_meta.shape.numel() + total_num_of_elems += output_elem + # Assume for now if it's quantized then it's qint8 or quint8 + if tensor_meta.is_quantized: + size_per_elem_bytes = torch._empty_affine_quantized( + [], dtype=tensor_meta.dtype + ).element_size() + else: + size_per_elem_bytes = torch.tensor([], dtype=tensor_meta.dtype).element_size() + total_size = size_per_elem_bytes * total_num_of_elems + output_size = size_per_elem_bytes * output_elem + return size_bytes(output_size, total_size) diff --git a/lib/python3.12/site-packages/torch/fx/passes/graph_transform_observer.py b/lib/python3.12/site-packages/torch/fx/passes/graph_transform_observer.py new file mode 100644 index 0000000000000000000000000000000000000000..2f27cf3c3866a6c1d18038600a13cceff458a034 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/graph_transform_observer.py @@ -0,0 +1,138 @@ +# mypy: allow-untyped-defs +import os +from typing import Callable, Optional, TypeVar + +from torch.fx import Graph +from torch.fx._compatibility import compatibility +from torch.fx.graph_module import GraphModule + + +T = TypeVar("T") + + +from .graph_drawer import FxGraphDrawer + + +__all__ = ["GraphTransformObserver"] + + +@compatibility(is_backward_compatible=False) +class GraphTransformObserver: + __pass_count = 0 + + def __init__( + self, + gm: GraphModule, + passname: str, + subsystem: Optional[str] = None, + log_url: Optional[str] = None, + ): + """ + log_url is inferred to be torch._inductor.config.trace.log_url_for_graph_xform unless otherwise specified + """ + + self.gm = gm + self.passname = passname + self.subsystem = subsystem + + # If log_url is None, we don't log anything + if log_url is None: + from torch._inductor.config import trace + + log_url = trace.log_url_for_graph_xform + + self.log_url = log_url + if self.log_url is None: + return + GraphTransformObserver.__pass_count += 1 + + self.input_dot_graph = FxGraphDrawer( + self.gm, + self.passname, + ignore_getattr=True, + ignore_parameters_and_buffers=True, + ).get_dot_graph() + + @classmethod + def get_current_pass_count(cls): + return cls.__pass_count + + def apply_gm_pass(self, pass_fn: Callable[[GraphModule], T]) -> Optional[T]: + with self: + if not self._check_disable_pass(): + return pass_fn(self.gm) + + return None + + def apply_graph_pass(self, pass_fn: Callable[[Graph], T]) -> Optional[T]: + with self: + if not self._check_disable_pass(): + return pass_fn(self.gm.graph) + + return None + + def _check_disable_pass(self): + if self.subsystem is None: + return False + + debug_info = lambda: self.passname # noqa: E731 + from torch._inductor.compiler_bisector import CompilerBisector + + return CompilerBisector.disable_subsystem( + "inductor", self.subsystem, debug_info + ) + + def __enter__(self): + if self.log_url is None or self.gm is None: + return self + + self.erased_nodes = set() + self.created_nodes = set() + self.gm._register_create_node_hook(self.on_node_creation) + self.gm._register_erase_node_hook(self.on_node_erase) + + return self + + def __exit__(self, type, value, tb): + if self.log_url is None or self.gm is None: + return + + self.gm._unregister_create_node_hook(self.on_node_creation) + self.gm._unregister_erase_node_hook(self.on_node_erase) + + if len(self.created_nodes) > 0 or len(self.erased_nodes) > 0: + for e in self.input_dot_graph.get_node_list(): + if e.get_name() in self.erased_nodes: + e.obj_dict["attributes"]["fillcolor"] = "yellow" + else: + e.obj_dict["attributes"]["fillcolor"] = "grey" + self.input_dot_graph.write( + os.path.join( + self.log_url, + f"pass_{GraphTransformObserver.__pass_count}_{self.passname}_input_graph.dot", + ) + ) + + output_dot_graph = FxGraphDrawer( + self.gm, + self.passname, + ignore_getattr=True, + ignore_parameters_and_buffers=True, + ).get_dot_graph() + for e in output_dot_graph.get_node_list(): + if e.get_name() in self.created_nodes: + e.obj_dict["attributes"]["fillcolor"] = "yellow" + else: + e.obj_dict["attributes"]["fillcolor"] = "grey" + output_dot_graph.write( + os.path.join( + self.log_url, + f"pass_{GraphTransformObserver.__pass_count}_{self.passname}_output_graph.dot", + ) + ) + + def on_node_creation(self, node): + self.created_nodes.add(node.name) + + def on_node_erase(self, node): + self.erased_nodes.add(node.name) diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/infra/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..939157f1302e75e3cf17ec3c1e93d1b8993d67a0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/infra/__init__.py @@ -0,0 +1 @@ +from . import pass_manager diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30355c920265f2afb1b7d2f7e5339c3a8a46b9be Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/partitioner.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/partitioner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5d98c70b4f09e583ac6e2b6f54d38a69b94e859 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/partitioner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/pass_base.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/pass_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19e320cb9ff0dc81797e68a2350e2aebacf15832 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/pass_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/pass_manager.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/pass_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..760c7c8b71696ce4ba0849a6751d9bd051f8fffe Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/infra/__pycache__/pass_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/partitioner.py b/lib/python3.12/site-packages/torch/fx/passes/infra/partitioner.py new file mode 100644 index 0000000000000000000000000000000000000000..ca5848f1ff0484f9693c9c77fb03635abb724b1e --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/infra/partitioner.py @@ -0,0 +1,387 @@ +# mypy: allow-untyped-defs +import collections +import itertools +import logging +from copy import copy +from typing import Dict, Iterable, List, Optional, Sequence, Set + +from torch.fx.graph_module import GraphModule +from torch.fx.node import _get_qualified_name, Node +from torch.fx.passes.operator_support import OperatorSupportBase +from torch.fx.passes.utils.fuser_utils import fuse_by_partitions + + +logger = logging.getLogger(__name__) +logger.setLevel(logging.WARNING) + + +class Partition: + def __init__( + self, id: Optional[int] = None, nodes: Optional[Iterable[Node]] = None + ): + self.id = id + self.nodes = dict.fromkeys(nodes) if nodes is not None else {} + + def __repr__(self) -> str: + return str(self.nodes) + + def add_node(self, node: Node): + self.nodes.update({node: None}) + + def remove_node(self, node: Node): + del self.nodes[node] + + def size(self): + return len(self.nodes) + + +class _DependencyViewer: + def __init__(self, graph_module: GraphModule): + self.upstreams = collections.defaultdict(set) + self.downstreams = collections.defaultdict(set) + + for node in graph_module.graph.nodes: + for input_node in node.all_input_nodes: + # add input_node and input_node's upstream dependency + self.upstreams[node].add(input_node) + self.upstreams[node].update(self.upstreams[input_node]) + + for node in reversed(graph_module.graph.nodes): + for output_node in node.users: + # add output_node and output_node's downstream dependency + self.downstreams[node].add(output_node) + self.downstreams[node].update(self.downstreams[output_node]) + + def downstreams_of(self, node: Node) -> Set[Node]: + return self.downstreams[node] + + def upstreams_of(self, node: Node) -> Set[Node]: + return self.upstreams[node] + + +class CapabilityBasedPartitioner: + def __init__( + self, + graph_module: GraphModule, + operator_support: OperatorSupportBase, + allows_single_node_partition: bool = False, + non_compute_ops: Optional[Sequence[str]] = None, + allowed_single_node_partition_ops: Optional[Sequence[str]] = None, + ) -> None: + self.graph_module = graph_module + self.operator_support = operator_support + self.allows_single_node_partition = allows_single_node_partition + self.non_compute_ops = non_compute_ops if non_compute_ops is not None else [] + self.allowed_single_node_partition_ops = ( + allowed_single_node_partition_ops + if allowed_single_node_partition_ops is not None + else [] + ) + self.dependency_viewer = _DependencyViewer(graph_module) + + def __is_node_supported(self, node: Node) -> bool: + return self.operator_support.is_node_supported( + dict(self.graph_module.named_modules()), node + ) + + def propose_partitions(self) -> List[Partition]: + # partition_map is a mapping from partition id to a set of partition id's. + # The value set contains all the partition ids that can be reached by doing a + # DFS starting from the partition id in the key. + partition_map: Dict[int, Set] = collections.defaultdict(set) + + # assumptions: nodes in candidate list is sorted in topological order + assignment: Dict[Node, int] = {} # mapping from node to partition_id + partitions_by_id: Dict[ + int, Partition + ] = {} # mapping from partition_id to partition + nodes_order: Dict[ + Node, int + ] = {} # mapping from nodes to reversed topological order + partitions_order: Dict[ + int, int + ] = {} # mapping from partition_id to minimum topo order of nodes in partition + new_partition_id = itertools.count() + + # try to merge partition other_id into partition self_id + # merge only happens if the end graph doesn't contain cyclic dependency + # returns `True` when merge happens, `False` otherwise. + def maybe_merge_partition(self_id: int, other_id: int): + # merged_nodes is the union of nodes in two partition to-be-merged + merged_nodes = copy(partitions_by_id[self_id].nodes) + merged_nodes.update(partitions_by_id[other_id].nodes) + + def dfs_iter_find_cycle(all_user_nodes: Set[Node]): + for user_node in all_user_nodes: + visited_partition_ids = set() + + for path_node in self.dependency_viewer.downstreams_of(user_node): + # If any of the nodes in the dfs path of this node are in the merged_nodes + # list then there is a cycle in the graph. + if path_node in merged_nodes: + return True + + # If any of the nodes in the dfs path of this node are in the assignment + # map then we have to make sure that the partitions that these nodes belong + # to do not form a cycle with the current partitions being merged. This means + # iterating through all the nodes in all the parititons that are traversed in + # the dfs path and checking if they are in the merged_nodes list. + if path_node in assignment: + partition_id = assignment[path_node] + # If the partition id has already been visited then we know that it doesn't + # form a cycle with the current partitions being merged. + if partition_id in visited_partition_ids: + continue + p_map = partition_map[partition_id] + if self_id in p_map or other_id in p_map: + return True + + visited_partition_ids.add(partition_id) + + return False + + # check if merge would create cyclic dependency. + all_user_nodes = set() + for node in merged_nodes: + for user_node in node.users: + if user_node not in merged_nodes: + all_user_nodes.add(user_node) + + if dfs_iter_find_cycle(all_user_nodes): + # return false indicating cyclic dependency found and + # merge is aborted + return False + + # no cyclic dependency found, move forward with the merge + # updating partition nodes + partitions_by_id[self_id].nodes = merged_nodes + # updating assignment map + for node in partitions_by_id[other_id].nodes: + assignment[node] = self_id + # delete other partition + del partitions_by_id[other_id] + + partitions_order[self_id] = min( + partitions_order[self_id], partitions_order[other_id] + ) + del partitions_order[other_id] + + partition_map[self_id] = partition_map[self_id].union( + partition_map[other_id] + ) + del partition_map[other_id] + + return True + + def merge_single_node(node: Node, id: Optional[int]): + def _update_partition_map(node: Node, id: int): + # Iterate through all the users of this node and update the partition map to indicate + # that there is a path from the partition id of this node to the target partition id. + for user_node in node.users: + target_id = assignment.get(user_node, None) + if target_id is not None: + partition_map[id].add(target_id) + partition_map[id].update(partition_map[target_id]) + + # Iterate through all the upstream nodes of this node and update the partition map + # to indicate that there is a path from the partition id of the upstream node to the + # current node's partition id. + upstream_nodes = self.dependency_viewer.upstreams_of(node) + for curr_node in upstream_nodes: + source_id = assignment.get(curr_node, None) + if source_id is not None: + partition_map[source_id].add(id) + + if node in assignment: + partitions_by_id[assignment[node]].remove_node(node) + + if id is None: + assignment.pop(node) + elif id not in partitions_by_id: + assignment[node] = id + partitions_by_id[id] = Partition(id=id, nodes=[node]) + _update_partition_map(node, id) + else: + assignment[node] = id + partitions_by_id[id].add_node(node) + _update_partition_map(node, id) + + logger.debug("Proposing partitions...") + + for node in reversed(self.graph_module.graph.nodes): + # use Dict as an ordered set to ensure deterministic partitioning result, don't care value + merge_candidates: Dict[int, None] = {} + + # Note a limited horizontal fusion is enabled: + # when `node` is not supported, the code below attempts to fuse consumer of `node`. + # + # I don't see a need to add a knob to disable horizontal fusion yet, we can short-cut + # the fusion by adding an `else` block here to skip horizontal fusion. + if self.__is_node_supported(node) and node not in assignment: + partition_id = next(new_partition_id) + nodes_order[node] = partition_id + partitions_order[partition_id] = partition_id + merge_single_node(node, partition_id) + merge_candidates[partition_id] = None + + # merge all possible partitions + for partition_id, _ in sorted( + partitions_order.items(), key=lambda item: item[1] + ): + merge_candidates[partition_id] = None + + merge_candidates_list = list(merge_candidates.keys()) + if len(merge_candidates_list) > 1: + self_id = merge_candidates_list[0] + for other_id in merge_candidates_list[1:]: + # note: merge partition `other_id` into partition `self_id` if + # it doesn't create cyclic dependency in the graph, otherwise, + # this is a no-op + maybe_merge_partition(self_id, other_id) + + # post processing to re-assign "getitem" nodes into upstream partition + logger.debug("Reassigning getitem nodes to its producer node's partition...") + nodes_reassignment: Dict[Node, int] = {} + for node in self.graph_module.graph.nodes: + is_tuple_output = True + for user in node.users: + if ( + user.op != "call_function" + or _get_qualified_name(user.target) != "_operator.getitem" + ): # type: ignore[arg-type] + is_tuple_output = False + break + + # node has tuple outputs, re-assign all following getitem node into node's partition + if is_tuple_output: + id = assignment.get(node, None) # type: ignore[arg-type] + for user in node.users: + if assignment.get(user, None) != id: # type: ignore[arg-type] + nodes_reassignment[user] = id # type: ignore[assignment] + for node, id in nodes_reassignment.items(): + merge_single_node(node, id) + + # filter out single node partitions + if not self.allows_single_node_partition: + logger.debug("Filtering out single node partitions...") + default_non_compute_ops = {"torch.ops.aten.view", "_operator.getitem"} + non_compute_ops = default_non_compute_ops.union(set(self.non_compute_ops)) + partitions_to_remove: List[int] = [] + for id, partition in partitions_by_id.items(): + compute_node_count = 0 + for node in partition.nodes: + if node.op == "call_function": + assert callable(node.target) + if _get_qualified_name(node.target) not in non_compute_ops: + compute_node_count += 1 + if ( + _get_qualified_name(node.target) + in self.allowed_single_node_partition_ops + ): + compute_node_count += 1 + if compute_node_count <= 1: + partitions_to_remove.append(id) + for id in partitions_to_remove: + del partitions_by_id[id] + + logger.debug("Partitions proposed:") + for id, partition in partitions_by_id.items(): + logger.debug( + "partition #%s: %s", id, [node.name for node in partition.nodes] + ) + + return [ + partition for partition in partitions_by_id.values() if partition.size() > 0 + ] + + def fuse_partitions( + self, partitions: List[Partition], prefix: str = "fused_" + ) -> GraphModule: + logger.debug("Fusing partitions...") + # fuse_by_partitions expects partitions in List[Dict[Node, None]]: [ {node0 : None}, {node1 : None} ] + return fuse_by_partitions( + self.graph_module, + [partition.nodes for partition in partitions], + prefix=prefix, + ) + + # remove non-compute-ops that sits at the boundary of a partition. + def remove_bookend_non_compute_ops(self, partitions: List[Partition]): + non_compute_ops = set(self.non_compute_ops) + + def is_non_compute_node(node: Node): + return ( + node.op == "call_function" + and _get_qualified_name(node.target) in non_compute_ops # type: ignore[arg-type] + ) + + # cache transparent nodes + transparent_input_nodes: Dict[Node, bool] = {} + transparent_output_nodes: Dict[Node, bool] = {} + + def is_transparent_input_node( + node: Node, partition: Set[Node], removed_nodes: Set[Node] + ): + if ( + node.op == "placeholder" + or (node not in partition) + or (node in removed_nodes) + ): + return True + if node in transparent_input_nodes: + return transparent_input_nodes[node] + if is_non_compute_node(node): + for input_n in node.all_input_nodes: + if not is_transparent_input_node(input_n, partition, removed_nodes): + transparent_input_nodes[node] = False + return False + transparent_input_nodes[node] = True + return True + transparent_input_nodes[node] = False + return False + + def is_transparent_output_node( + node: Node, partition: Set[Node], removed_nodes: Set[Node] + ): + if ( + node.op == "placeholder" + or (node not in partition) + or (node in removed_nodes) + ): + return True + if node in transparent_output_nodes: + return transparent_output_nodes[node] + if is_non_compute_node(node): + for output_n in node.users: + if not is_transparent_output_node( + output_n, partition, removed_nodes + ): + transparent_output_nodes[node] = False + return False + transparent_output_nodes[node] = True + return True + transparent_output_nodes[node] = False + return False + + for partition in partitions: + # Note it's ok to use `set` here, since we are only query if a node + # has been removed. We are NEVER going to iterate on nodes inside + # the set. + remove_node: Set[Node] = set() + for node in partition.nodes: + if is_non_compute_node(node) and ( + is_transparent_input_node(node, set(partition.nodes), remove_node) + or is_transparent_output_node( + node, set(partition.nodes), remove_node + ) + ): + remove_node.add(node) + + if len(remove_node) != 0: + for node in remove_node: + partition.nodes.pop(node, None) + + def partition_and_fuse(self, prefix: str = "fused_") -> GraphModule: + partitions = self.propose_partitions() + fused_gm = self.fuse_partitions(partitions, prefix=prefix) + return fused_gm diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/pass_base.py b/lib/python3.12/site-packages/torch/fx/passes/infra/pass_base.py new file mode 100644 index 0000000000000000000000000000000000000000..acf78d2581b5a24d6442777dc1ff003278616559 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/infra/pass_base.py @@ -0,0 +1,76 @@ +# mypy: allow-untyped-defs +import abc +from collections import namedtuple +from typing import Optional + +from torch.fx._compatibility import compatibility +from torch.fx.graph_module import GraphModule + + +__all__ = ["PassResult", "PassBase"] + + +@compatibility(is_backward_compatible=False) +class PassResult(namedtuple("PassResult", ["graph_module", "modified"])): + """ + Result of a pass: + graph_module: The modified graph module + modified: A flag for if the pass has modified the graph module + """ + + def __new__(cls, graph_module, modified): + return super().__new__(cls, graph_module, modified) + + +@compatibility(is_backward_compatible=False) +class PassBase(abc.ABC): + """ + Base interface for implementing passes. + + It is required to implement the `call` function so that we can directly + pass instances of the Pass directly to the PassManager and call them as a + function. + + We can directly pass an instance of a class implementing this interface into + the PassManager's `passes` attribute. + """ + + def __call__(self, graph_module: GraphModule) -> Optional[PassResult]: + """ + Runs the precondition check, the pass itself, and the postcondition check. + """ + + self.requires(graph_module) + res = self.call(graph_module) + self.ensures(graph_module) + return res + + @abc.abstractmethod + def call(self, graph_module: GraphModule) -> Optional[PassResult]: + """ + The pass that is run through the given graph module. To implement a + pass, it is required to implement this function. + + Args: + graph_module: The graph module we will run a pass on + """ + + def requires(self, graph_module: GraphModule) -> None: # noqa: B027 + """ + This function will be called before the pass is run and will check that + the given graph module contains the preconditions needed to run the + pass. It is not required to implement this function. + + Args: + graph_module: The graph module we will run checks on + """ + + def ensures(self, graph_module: GraphModule) -> None: # noqa: B027 + """ + This function will be called after the pass is run and will check that + the given graph module contains the postconditions needed to run the + pass. It is not required to implement this function. + + Args: + graph_module: The graph module we will run checks on + """ diff --git a/lib/python3.12/site-packages/torch/fx/passes/infra/pass_manager.py b/lib/python3.12/site-packages/torch/fx/passes/infra/pass_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..cea5f4f25c77b485ce47900d80e9eef151659896 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/infra/pass_manager.py @@ -0,0 +1,309 @@ +# mypy: allow-untyped-defs +import inspect +import logging +from functools import wraps +from queue import Queue +from typing import Callable, Dict, List + +import torch.nn as nn +from torch.fx._compatibility import compatibility +from torch.fx.graph_module import GraphModule +from torch.fx.passes.infra.pass_base import PassResult + + +logger = logging.getLogger(__name__) +logger.setLevel(logging.WARNING) + +__all__ = ["pass_result_wrapper", "this_before_that_pass_constraint", "PassManager"] + + +@compatibility(is_backward_compatible=False) +def pass_result_wrapper(fn: Callable) -> Callable: + """ + Wrapper for passes which currently do not return a PassResult. + This wrapper makes them return a PassResult containing the modified object + and True for the "modified" flag. + + Args: + fn (Callable[Module, Any]) + + Returns: + wrapped_fn (Callable[Module, PassResult]) + """ + if fn is None: + return None + + @wraps(fn) + def wrapped_fn(gm): + res = fn(gm) + if res is None: + return PassResult(gm, True) + if isinstance(res, PassResult): + return res + elif isinstance(res, nn.Module): + return PassResult(res, True) + + if not inspect.isfunction(fn): + wrapped_fn.__name__ = type(fn).__name__ + + return wrapped_fn + + +def _validate_pass_schedule_constraint( + constraint: Callable[[Callable, Callable], bool], passes: List[Callable] +) -> None: + for i, a in enumerate(passes): + for j, b in enumerate(passes[i + 1 :]): + if constraint(a, b): + continue + raise RuntimeError( + f"pass schedule constraint violated. Expected {a} before {b}" + f" but found {a} at index {i} and {b} at index{j} in pass" + f" list." + ) + + +def _topological_sort_passes( + passes: List[Callable], constraints: List[Callable] +) -> List[Callable]: + """ + Args + passes: Passes that we are ordering + constraints: Constraints applied on these passes + + Returns + A sorted list of callables and a boolean of if a circular dependency + existed + """ + if len(constraints) == 0: + return passes + + # Contruct a graph mapping nodes to a list of their users + graph: Dict[Callable, List[Callable]] = {p: [] for p in passes} + indegree_map: Dict[Callable, int] = dict.fromkeys(passes, 0) + candidates: Queue = Queue() + for a in passes: + for b in passes: + if a == b: + continue + + for constraint in constraints: + if not constraint(a, b): + graph[b].append(a) + indegree_map[a] += 1 + + if indegree_map[a] == 0: + candidates.put(a) + + visited: Dict[Callable, bool] = dict.fromkeys(passes, False) + sorted_passes: List[Callable] = [] + + while not candidates.empty(): + p = candidates.get() + sorted_passes.append(p) + visited[p] = True + + for n in graph[p]: + if not visited[n]: + indegree_map[n] -= 1 + if indegree_map[n] == 0: + candidates.put(n) + + # Check if there are unvisited nodes (aka cycles in the graph) + cycle_passes = list(filter(lambda p: indegree_map[p] != 0, indegree_map.keys())) + if len(cycle_passes) != 0: + error = ( + f"Circular dependency detected within the following passes: {cycle_passes}" + ) + raise RuntimeError(error) + + return sorted_passes + + +@compatibility(is_backward_compatible=False) +def this_before_that_pass_constraint(this: Callable, that: Callable) -> Callable: + """ + Defines a partial order ('depends on' function) where `this` must occur + before `that`. + + For example, the following pass list and constraint list would be invalid. + ``` + passes = [pass_b, pass_a] + + constraints = [this_before_that_pass_constraint(pass_a, pass_b)] + ``` + + Args: + this (Callable): pass which should occur first + that (Callable): pass which should occur later + + Returns: + depends_on (Callable[[Object, Object], bool] + """ + + def depends_on(a: Callable, b: Callable): + return a != that or b != this + + return depends_on + + +@compatibility(is_backward_compatible=False) +class PassManager: + """ + Construct a PassManager. + + Collects passes and constraints. This defines the pass schedule, manages + pass constraints and pass execution. + + Args: + passes (Optional[List[Callable]]): List of passes. A pass is a + callable which modifies an object and returns a PassResult + constraint (Optional[List[Callable]]): List of constraints. A + constraint is a callable which takes two passes (A, B) and returns + True if A depends on B and False otherwise. See implementation of + `this_before_that_pass_constraint` for example. + steps (int): Max number of times we run the passes (default = 1). + run_checks_after_each_pass (bool): Whether to run checks and linting + after each pass + suppress_check_failures (bool): Whether to raise errors when running + checks + """ + + passes: List[Callable[[nn.Module], PassResult]] + constraints: List[Callable[[Callable, Callable], bool]] + _validated: bool = False + steps: int = 1 + + def __init__( + self, + passes=None, + constraints=None, + steps=None, + run_checks_after_each_pass: bool = False, + suppress_check_failures: bool = False, + ): + self.passes = passes or [] + self.constraints = constraints or [] + if steps: + self.steps = steps + + self.run_checks_after_each_pass = run_checks_after_each_pass + self.suppress_check_failures = suppress_check_failures + + def add_pass(self, _pass: Callable): + """ + Adds a pass into the current list of passes. + """ + self.passes.append(_pass) + self._validated = False + + def add_constraint(self, constraint: Callable): + """ + Adds a constraint into the current list of constraints. + """ + self.constraints.append(constraint) + self._validated = False + + def validate_constraints(self): + """ + Validates that current pass schedule defined by `self.passes` is valid + according to all constraints in `self.constraints` + """ + if self._validated: + return + for constraint in self.constraints: + _validate_pass_schedule_constraint(constraint, self.passes) + self._validated = True + + def solve_constraints(self): + """ + Finds a valid traversal order based on the given constraints and orders + the passes based on this order. + + If a circular dependency exists between the constraints and steps = 1, + then we will raise an error because if steps != 1 this means that we + will re-run the passes, allowing for circular dependencies. + """ + self.passes = _topological_sort_passes(self.passes, self.constraints) + self._validated = True + + def add_checks(self, check: Callable) -> None: + """ + Adds a function which takes runs various checks on a given graph module. + This function is run before and after each pass if the + `run_checks_after_each_pass` flag is enabled. + """ + sig = inspect.signature(check) + + if len(list(sig.parameters.values())) != 1: + raise TypeError( + "PassManager check function should only take in one variable, a module" + ) + + setattr(self, "check", check) # noqa: B010 + + def check(self, module: nn.Module) -> None: + pass + + def __call__(self, module: nn.Module) -> PassResult: + """ + Runs a list of passes in the order based on `self.passes` on the given + graph module. Each time a pass is run, checks and linting will be run on + the graph module if `run_checks_after_each_pass` is set. + + If the module is a graph module, we will run the list of passes until + the graph stops changing, or until `steps` number of times. + """ + # Order the passes based on the constraints + if not self._validated: + self.solve_constraints() + + # Check graph invariants + self.check(module) + + # Run the set of passes `steps` number of times or until the graph stops + # changing + overall_modified = False + for _ in range(self.steps): + modified = False + + # Run the set of passes on the graph module + for i, fn in enumerate(self.passes): + fn_name = fn.__name__ if inspect.isfunction(fn) else type(fn).__name__ + logger.debug("Running pass '%s'", fn_name) + + try: + res = fn(module) + + if not isinstance(res, PassResult) and not hasattr( + res, "graph_module" + ): + raise TypeError( + f"The result of the pass {fn_name} should be type PassResult." + + "Please wrap it with pass_result_wrapper()" + ) + module = res.graph_module + modified = modified or res.modified + + if isinstance(module, GraphModule): + logger.debug("Graph after pass '%s': %s", fn_name, module.graph) + module.recompile() + + # Check graph invariants + if self.run_checks_after_each_pass: + self.check(module) + + except Exception as e: + prev_pass_names = [ + p.__name__ if inspect.isfunction(p) else type(p).__name__ + for p in self.passes[:i] + ] + msg = f"An error occurred when running the '{fn_name}' pass after the following passes: {prev_pass_names}" + raise Exception(msg) from e # noqa: TRY002 + + # If the graph no longer changes, then we can stop running these passes + overall_modified = overall_modified or modified + if not modified: + break + + return PassResult(module, overall_modified) diff --git a/lib/python3.12/site-packages/torch/fx/passes/net_min_base.py b/lib/python3.12/site-packages/torch/fx/passes/net_min_base.py new file mode 100644 index 0000000000000000000000000000000000000000..c349c896ac3eadff009c3e1bbee2e587a6890d4a --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/net_min_base.py @@ -0,0 +1,944 @@ +# mypy: allow-untyped-defs +import logging +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch +import torch.fx +from torch.fx._compatibility import compatibility +from torch.fx.node import map_arg + +from .shape_prop import ShapeProp +from .split_utils import split_by_tags +from .tools_common import ( + CALLABLE_NODE_OPS, + FxNetAccFusionsFinder, + Names, + NodeList, + NodeSet, + TensorOrTensors, + Tensors, +) + + +__all__ = [ + "FxNetMinimizerBadModuleError", + "FxNetMinimizerRunFuncError", + "FxNetMinimizerResultMismatchError", +] + +_LOGGER = logging.getLogger(__name__) + + +@compatibility(is_backward_compatible=False) +class FxNetMinimizerBadModuleError(Exception): + """ + Raised if failed to split out a minimize module + """ + + +@compatibility(is_backward_compatible=False) +class FxNetMinimizerRunFuncError(Exception): + """ + Raised if error occurs during run_a or run_b functions + """ + + +@compatibility(is_backward_compatible=False) +class FxNetMinimizerResultMismatchError(Exception): + """ + Raised if comparing function thinks the results are mismatching. + """ + + +@dataclass +class _MinimizerSettingBase: + """ + Args: + `accumulate_error`: Instead of using a's input for both converted module to verify + , use the previous outputs of each converted module as input to accumulate the + errors. + + `traverse_method`: "sequential" or "binary" or "accumulate" + Determine the way of traverse the nodes in FX module. + + `find_all`: Minimizer will go through the entire model and return all problematic nodes. + + `return_intermediate`: If true, when using `run_nodes()` function to run the + model, intermediate results of all the ops will be returned as output. + + `all_outputs`: If true, when using `_run_and_compare()` function, + all the output nodes in the subgraph will be used for comparison. + """ + + accumulate_error: bool = False + traverse_method: str = "sequential" + find_all: bool = False + return_intermediate: bool = False + all_outputs: bool = False + + def __str__(self): + settings_str = "FX Minimizer Settings:\n" + + for k, v in vars(self).items(): + settings_str += f"\t{k}: {v}\n" + + return settings_str + + +class _MinimizerBase: + """ + This class is used to automatically find problematic nodes in a model. It takes a FX + graphmodule and generate some submodules while traverse the graph. Then two functions + `run_a` and `run_b` will be used to run the same submodule and a function `compare_fn` + will be used to compare the results. + + Currently we provides two ways to traverse the graph and generate submodules. + 1. Sequential traversal: this will traverse the graph node by node and generate + one submodule with one sigle node. + 2. Binary searching: this will do a binary search style traversal on the graph. + + For internal Users, a guide can be found here https://fb.quip.com/HDtuAgiKGfkP. + """ + + def __init__( + self, + module: torch.fx.GraphModule, + sample_input: Tensors, + compare_fn: Callable[ + [TensorOrTensors, TensorOrTensors, Names], Tuple[float, bool] + ], + settings: _MinimizerSettingBase, + module_exporter: Optional[ + Callable[[Tensors, torch.fx.GraphModule, str], None] + ] = None, + exclusion_fn: Optional[Callable[[NodeList, int, int], None]] = None, + ): + assert isinstance(module, torch.fx.GraphModule) + + self.module = module + self.sample_input = sample_input + self.compare_fn = compare_fn + self.module_exporter = module_exporter + self.settings = settings + self.exclusion_fn = exclusion_fn + + # Stores outputs of run_a function + self.a_outputs: Dict[str, Any] = {} + + # Stores outputs of run_b function + self.b_outputs: Dict[str, Any] = {} + + # Stores the results of compare_fn + self.results: Dict[Any, Any] = {} + + # Stores the report for the runs + self.reports: List[List[str]] = [] + + # Current iteration + self.iteration: int = 0 + + callable_nodes = { + node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS + } + ShapeProp(self.module).propagate(*self.sample_input) + self.fusions = FxNetAccFusionsFinder(self.module, callable_nodes)() + + # Check if number of input in sample_input matches the number of placeholders + placeholders = [ + node.name for node in self.module.graph.nodes if node.op == "placeholder" + ] + assert len(placeholders) == len(self.sample_input) + + # Store sample_input + for i, name in enumerate(placeholders): + self.a_outputs[name] = sample_input[i] + self.b_outputs[name] = sample_input[i] + + def run_a( + self, mod: torch.fx.GraphModule, inputs: Tensors, report_idx: int = -1 + ) -> TensorOrTensors: + """ + Run `mod` with `inputs` and generate output. The output will be compared with + output of run_b(). + """ + raise RuntimeError("run_a() is not implemented.") + + def run_b( + self, mod: torch.fx.GraphModule, inputs: Tensors, report_idx: int = -1 + ) -> TensorOrTensors: + """ + Run `mod` with `inputs` and generate output. The output will be compared with + output of run_a(). + """ + raise RuntimeError("run_b() is not implemented.") + + def _store_outputs( + self, + a_result: TensorOrTensors, + b_result: TensorOrTensors, + submodule: torch.fx.GraphModule, + ): + """ + Store the outputs of self.run_a() and self.run_b() into self.a_outputs and + self.b_outputs, so that we can use them when execute preceding nodes that + use those outputs as inputs. + + Args: + a_result: Output of self.run_a(). Could be a tensor or tensors. + b_result: Output of self.run_b(). Could be a tensor or tensors. + submodule: The module that generates a_result and b_result. + """ + output_node = next( + node for node in submodule.graph.nodes if node.op == "output" + ) + + # Only one output + if isinstance(output_node.args[0], torch.fx.Node): + self.a_outputs[output_node.args[0].name] = a_result + self.b_outputs[output_node.args[0].name] = b_result + # Multiple outputs + else: + for i, arg in enumerate(output_node.args[0]): + self.a_outputs[arg.name] = a_result[i] + self.b_outputs[arg.name] = b_result[i] + + def _get_submod_inputs( + self, main_module: torch.fx.GraphModule, submod_path: str + ) -> Tuple[Tensors, Tensors]: + """ + Try get submodule inputs from stored outputs. If not found then use + torch_glow.get_submod_inputs to get the inputs. + + If accumulate_error is False, use a_input for run_a() and run_b() + otherwise use a_input for run_a and b_input for run_b. + + Args: + main_module: Top-levlel fx module. + submod_path: Path to the submodule we want to run and compare results. + + Returns: + a_input: List of tensor(s) that will be used by run_a() as submodule inputs. + b_input: List of tensor(s) that will be used by run_b() as submodule inputs. + """ + a_input = [] + b_input = [] + submodule = getattr(main_module, submod_path) + placeholders = [ + node.name for node in submodule.graph.nodes if node.op == "placeholder" + ] + + # If all placeholder can be found in stored outputs, use stored + # outputs as inputs. Otherwise, use `torch_glow.get_submod_inputs` + # to get the inputs. + if set(placeholders) <= self.a_outputs.keys(): + for name in placeholders: + a_input.append(self.a_outputs[name]) + b_input.append(self.b_outputs[name]) + else: + if self.settings.accumulate_error: + print(f"Can't find previous stored outputs named {placeholders}!") + + def get_inputs(self: torch.nn.Module, inputs: Any): + nonlocal a_input + a_input = inputs + + # Use forward hook to get the inputs to the submodule + handle = submodule.register_forward_pre_hook(get_inputs) + main_module(*self.sample_input) + handle.remove() + + b_input = a_input + + if not self.settings.accumulate_error: + return a_input, a_input + + return a_input, b_input + + def _tag_nodes(self, selected_nodes: NodeSet): + """ + Tag selected nodes with tag "minimize". Nodes with the same tags will + be split to the same submodule afterwards. + + Args: + selected_nodes: Nodes that we want to minimize. We will tag those nodes + with "minimize", all preceding nodes with "main_0" and all following + nodes with "main_1". + """ + for node in self.module.graph.nodes: + if node.op not in CALLABLE_NODE_OPS: + continue + + if node in selected_nodes: + node.tag = "minimize" + elif any( + n.tag in {"minimize", "main_1"} + for n in node.all_input_nodes + if n.op in CALLABLE_NODE_OPS + ): + node.tag = "main_1" + else: + node.tag = "main_0" + + def _build_submodule(self, nodes: NodeSet) -> Tuple[torch.fx.GraphModule, str]: + """ + Split self.module so that one submodule consists of `nodes` and only `nodes`. + + Args: + nodes: Nodes that we want to include in the minimize submodule. + + Returns: + split_module (torch.fx.GraphModule): the module after split. + submodule_name (str): the name of the submodule that consists of `nodes`. + """ + # Color provided nodes + self._tag_nodes(nodes) + + # Split module based on coloring + split_module = split_by_tags(self.module, ["main_0", "minimize", "main_1"]) + + # Find submodule containing colored nodes + submodule_name: str = "" + for child_name, _ in split_module.named_children(): # type: ignore[union-attr] + # Skip submodules we're not interested in at the moment + if "minimize" not in child_name: + continue + + if submodule_name == "": + submodule_name = child_name + else: + raise FxNetMinimizerBadModuleError( + f"Expected only one minimize submodule with nodes {nodes}" + ) + + if submodule_name == "": + raise FxNetMinimizerBadModuleError( + f"Minimize submodule was not found with nodes {nodes}" + ) + + return split_module, submodule_name # type: ignore[return-value] + + def _run_and_compare( + self, + split_module: torch.fx.GraphModule, + submod_name: str, + output_names: Names, + report_idx: int = -1, + ): + """ + Run the submodule in `split_module` that has name `submod_name` + using `self.run_a` and `self.run_b` and compare their results. + + Args: + split_module: Main module that contains the minimize submodule. + submod_name: Name of the minimize submodule. + output_names: Names of the node we want to output. If None, we + will use the original output. + """ + submodule = getattr(split_module, submod_name) + a_input, b_input = self._get_submod_inputs(split_module, submod_name) + + if len(self.reports) == 0: + self.reports.append([]) + self.iteration = 1 + + report = self.reports[report_idx if report_idx >= 0 else self.iteration - 1] + report.append("Run and compare ...") + + if output_names and not self.settings.all_outputs: + output_nodes: NodeList = [] + for node in submodule.graph.nodes: + if node.op == "output": + submodule.graph.erase_node(node) + + if node.name in output_names: + output_nodes.append(node) + + submodule.graph.output( + output_nodes[0] if len(output_nodes) == 1 else tuple(output_nodes) + ) + submodule.graph.lint() + submodule.recompile() + + # Use name of args in output node as key to store comparison result + for node in submodule.graph.nodes: + if node.op == "output": + result_key = map_arg(node.args, lambda x: x.name) + + try: + a_result = self.run_a(submodule, a_input, report_idx) + b_result = self.run_b(submodule, b_input, report_idx) + self._store_outputs(a_result, b_result, submodule) + except Exception as e: + report.append(f"Exception raised when running {submod_name}: {e}") + raise FxNetMinimizerRunFuncError( # noqa: B904 + f"Exception raised when running {submod_name}: {e}" + ) + + # Compare results + names: Names = output_names + if output_names is None: + names = [str(v) for v in result_key] # type: ignore[possibly-undefined] + + numeric_result, bool_result = self.compare_fn(a_result, b_result, names) + + self.results[result_key] = numeric_result # type: ignore[possibly-undefined] + report.append(f"Numerical accuracy = {numeric_result}") + if not bool_result: + report.append(f"Result mismatch for {result_key}") # type: ignore[possibly-undefined] + if self.module_exporter: + if isinstance(result_key, tuple): # type: ignore[possibly-undefined] + result_key = result_key[-1] + # pyre-ignore[29]: not a function + self.module_exporter( + a_input, + submodule, + str(result_key[0]) + "_cpu", # type: ignore[index] + ) + # pyre-ignore[29]: not a function + self.module_exporter( + b_input, + submodule, + str(result_key[0]) + "_acc", # type: ignore[index] + ) + raise FxNetMinimizerResultMismatchError(f"Result mismatch for {result_key}") # type: ignore[possibly-undefined] + + def _binary_search_impl( + self, all_nodes: NodeList, start_idx: int, end_idx: int + ) -> NodeSet: + """ + Recursive binary search implementation. + """ + culprits: NodeSet = set() + nodes: NodeList = all_nodes[start_idx:end_idx] + + report: List[str] = [] + if self.exclusion_fn is not None: + self.exclusion_fn(nodes, start_idx, end_idx) + if len(nodes) == 0: + report = ["All nodes are excluded by user"] + self.reports.append(report) + return culprits + + first_node_name = nodes[0].name + output_node_name = nodes[-1].name + self.iteration += 1 + self.reports.append(report) + report.append(f"Binary search iteration {self.iteration}") + report.append( + f"From node index {start_idx}:{first_node_name} to {end_idx - 1}:{output_node_name}. " + f"Size of the interested node list is {len(nodes)}" + ) + cur_nodes: NodeSet = set(nodes) + + try: + split_module, submod_name = self._build_submodule(cur_nodes) + self._run_and_compare(split_module, submod_name, [output_node_name]) + + except (FxNetMinimizerRunFuncError, FxNetMinimizerResultMismatchError): + if len(nodes) == 1: + report.append( + f"This is the last node in the sub-module. " + f"Search in the current branch is successful with culprit = {cur_nodes}." + ) + self.print_report(report) + return cur_nodes + + report.append( + "Proceed to split and lower the halves of the current " + "sub-module individually." + ) + self.print_report(report) + + mid = len(nodes) // 2 + culprits = self._binary_search_impl(all_nodes, start_idx, start_idx + mid) + + if len(culprits) != 0 and not self.settings.find_all: + return culprits + + culprits = self._binary_search_impl(all_nodes, start_idx + mid, end_idx) + + if len(culprits) == 0: + report.append( + f"Further split and lowering found no errors. " + f"Unable to minimize the submodule with list of nodes: {nodes}" + ) + self.print_report(report) + + return culprits + else: + report.append("No discrepancy found.") + self.print_report(report) + return set() + + def _binary_traverse(self, nodes: NodeList) -> NodeSet: + """ + Binary search on `nodes` for culprit. + """ + return self._binary_search_impl(nodes, 0, len(nodes)) + + def _sequential_traverse(self, nodes: NodeList) -> NodeSet: + """ + Traverse `nodes` one by one and determine if any of them is a culprit. + """ + culprits: NodeSet = set() + + for node in nodes: + report: List[str] = [] + self.reports.append(report) + self.iteration += 1 + report.append(f"Sequential traverse iteration {self.iteration}.") + report.append(f"Visit node: {node.name}") + + _LOGGER.info("Visit node: %s", node.name) + node_list: NodeList = [node] + if self.exclusion_fn is not None: + self.exclusion_fn(node_list, -1, -1) + if len(node_list) == 0: + report.append(f"User exclusion : {node.name}") + self.print_report(report) + if not self.settings.find_all: + return culprits + else: + continue + + cur_nodes: NodeSet = {node} + + if node in self.fusions: + cur_nodes = self.fusions[node] + + try: + split_module, submod_name = self._build_submodule(cur_nodes) + self._run_and_compare(split_module, submod_name, [node.name]) + self.print_report(report) + except FxNetMinimizerResultMismatchError: + culprits.add(node) + report.append(f"Found culprit from numeric error: {node}") + self.print_report(report) + if not self.settings.find_all: + return culprits + except FxNetMinimizerRunFuncError: + culprits.update(cur_nodes) + report.append(f"Found culprit from run error: {node}") + self.print_report(report) + if not self.settings.find_all: + return culprits + + return culprits + + def _block_traverse_impl( + self, nodes: NodeList, start_idx: int, end_idx: int, find_last_node: bool + ) -> int: + """ + Recursive block search implementation. + find_last_node: If True, search for the last node which result in numerics difference + if False: find first node in sorted node list + """ + report: List[str] = [] + + mid = (start_idx + end_idx) // 2 + cur_nodes_list: NodeList = nodes[: mid + 1] if find_last_node else nodes[mid:] + + if self.exclusion_fn: + self.exclusion_fn(cur_nodes_list, -1, -1) + + cur_nodes = set(cur_nodes_list) + + first_node_name = cur_nodes_list[0].name + last_node_name = cur_nodes_list[-1].name + target_node_name = last_node_name if find_last_node else first_node_name + + self.iteration += 1 + self.reports.append(report) + report.extend( + [ + "=" * 30, + f"Block search iteration {self.iteration}", + ] + ) + report.extend( + [ + f"Search for {'last' if find_last_node else 'first'} node in culprits", + f"From node index {start_idx}:{nodes[start_idx].name} to {end_idx}:{nodes[end_idx].name}. ", + f"Subgraph constructed by {first_node_name} to {last_node_name}", + f"Targeting node: {target_node_name}", + f"Size of the interested node list is {end_idx - start_idx + 1}", + ] + ) + report_idx = len(self.reports) - 1 + + try: + split_module, submod_name = self._build_submodule(cur_nodes) + self._run_and_compare( + split_module, submod_name, [last_node_name], report_idx + ) + except (FxNetMinimizerResultMismatchError, FxNetMinimizerRunFuncError): + report.append( + f"Culprits found from node {first_node_name} to {last_node_name}." + ) + + if start_idx == mid: + report.extend( + [ + "This is the last node in the sub-module. ", + "Search in the current branch is successful with node :", + f"{start_idx}, node name: {nodes[start_idx].name}.", + ] + ) + self.print_report(report) + return start_idx + + report.append( + "Proceed to split and lower the halves of the current " + "sub-module individually." + ) + self.print_report(report) + + if find_last_node: + return self._block_traverse_impl(nodes, start_idx, mid, find_last_node) + else: + return self._block_traverse_impl( + nodes, mid + 1, end_idx, find_last_node + ) + else: + report.append( + f"Culprits not found from node start to {mid}:{nodes[mid].name}." + ) + + if start_idx == mid: + report.extend( + [ + "This is the last node in the sub-module. ", + "Search in the current branch is successful with node", + f"{start_idx}, node name: {nodes[start_idx].name}.", + ] + ) + self.print_report(report) + return start_idx + 1 if find_last_node else start_idx - 1 + + report.append( + "Proceed to split and lower the halves of the current " + "sub-module individually." + ) + self.print_report(report) + + if find_last_node: + return self._block_traverse_impl( + nodes, mid + 1, end_idx, find_last_node + ) + else: + return self._block_traverse_impl(nodes, start_idx, mid, find_last_node) + + def _block_traverse( + self, nodes: NodeList, find_last_node: Optional[bool] + ) -> NodeSet: + """ + Traverse topologically sorted node list + Find minimium block (start_idx, end_idx) which contains the culprit + 1st pass: search for end_idx by finding the last node in culprit block + where Numerical accuracy (0, end_idx) > threshold + 2nd pass: search for start_idx by finding the first node in culprit block + where Numerical accuracy (start_idx, end_idx) < threshold + Form minimum block by (start_idx - 1, end_idx) + """ + culprits: NodeSet = set() + first_node_name = nodes[0].name + last_node_name = nodes[-1].name + last_node_report = [f"Block search from {first_node_name} to {last_node_name}"] + last_node_report.append("*" * 50) + self.reports.append(last_node_report) + + start_idx = 0 + end_idx = len(nodes) - 1 + run_both = True if find_last_node is None else False + + # step 1: find (0, end_idx) of culprit block + if run_both or find_last_node: + last_node_report.append("Start searching for last node in culprit") + self.print_report(last_node_report) + end_idx = self._block_traverse_impl(nodes, start_idx, end_idx, True) + last_node_report.extend( + ["Finish Pass 1", f"Find end_idx = {end_idx}:{nodes[end_idx].name}"] + ) + self.print_report(last_node_report) + + # step 2: reduce culprit block to (start_idx, end_idx) + if run_both or not find_last_node: + first_node_report = ["Start searching for first node in culprit"] + self.print_report(first_node_report) + start_idx = self._block_traverse_impl( + nodes[0 : end_idx + 1], start_idx, end_idx, False + ) + first_node_report.append("*" * 50) + self.reports.append(first_node_report) + first_node_report.extend( + [ + "Finish Pass 2", + f"Find start_idx = {start_idx}:{nodes[start_idx].name}", + ] + ) + self.print_report(first_node_report) + + # step 3: form module with minimum culprits + culprits.update(nodes[start_idx : end_idx + 1]) + result_report = [ + f"Finish searching, found minimum block ({nodes[start_idx]},{nodes[end_idx]})" + ] + self.reports.append(result_report) + self.print_report(result_report) + return culprits + + def _defined_traverse(self, nodes: NodeList) -> NodeSet: + """ + run user defined `nodes` and determine if it is a culprit. + """ + culprits: NodeSet = set() + if self.exclusion_fn is not None: + self.exclusion_fn(nodes, -1, -1) + if len(nodes) == 0: + report = ["All nodes are excluded by user"] + self.reports.append(report) + return culprits + + first_node_name = nodes[0].name + output_node_name = nodes[-1].name + report = [f"Defined graph from {first_node_name} to {output_node_name}"] + cur_nodes: NodeSet = set(nodes) + try: + split_module, submod_name = self._build_submodule(cur_nodes) + self._run_and_compare(split_module, submod_name, [output_node_name]) + self.print_report(report) + except (FxNetMinimizerResultMismatchError, FxNetMinimizerRunFuncError): + report.append(f"Found culprit {cur_nodes}") + self.print_report(report) + return culprits + + return culprits + + def _accumulate_traverse(self, nodes: NodeList) -> NodeSet: + culprits: NodeSet = set() + nodes_to_run: NodeSet = set() + + # find_all is not supported for accumulate traversal because all the + # ops run on NNPI. So we return after the first op that raises error. + if self.settings.find_all: + print("'Find All' mode is not supported in accumulate traversal.") + return culprits + + for node in nodes: + report: List[str] = [] + self.reports.append(report) + self.iteration += 1 + report.append(f"Accumulate traverse iteration {self.iteration}.") + + nodes_to_run.add(node) + + node_name = node.name + if node_name is not None and isinstance(node_name, tuple): + node_name = node_name[0] + assert node_name is not None and isinstance( + node_name, str + ), f"minimize: node_name: {node_name}" + + report.append(f"Add node: {node_name}") + + try: + split_module, submod_name = self._build_submodule(nodes_to_run) + self._run_and_compare(split_module, submod_name, [node_name]) + self.print_report(report) + except (FxNetMinimizerResultMismatchError, FxNetMinimizerRunFuncError): + culprits.add(node) + report.append(f"Found culprit {node}") + self.print_report(report) + return culprits + + return culprits + + def _skip_traverse_impl( + self, all_nodes: NodeList, start_idx: int, end_idx: int + ) -> NodeSet: + """ + Skip certain nodes in graph based on settings + """ + culprits: NodeSet = set() + nodes: NodeList = all_nodes[start_idx:end_idx] + cur_nodes: NodeSet = set(nodes) + if self.exclusion_fn is not None: + self.exclusion_fn(nodes, start_idx, end_idx) + cur_nodes = set(nodes) + else: + for node in nodes: + if node in self.fusions: + cur_nodes.update(self.fusions[node]) + report: List[str] = [] + self.reports.append(report) + self.iteration += 1 + report.append(f" Nodes block {self.iteration}.") + report.append( + f"From node index {start_idx} to {end_idx - 1}. " + f"Size of the interested node list is {len(nodes)}" + ) + + try: + split_module, submod_name = self._build_submodule(cur_nodes) + self._run_and_compare(split_module, submod_name, []) + except FxNetMinimizerResultMismatchError: + culprits.update(cur_nodes) + report.append(f"Found culprit from numeric error: {cur_nodes}") + self.print_report(report) + return culprits + except FxNetMinimizerRunFuncError: + culprits.update(cur_nodes) + report.append(f"Found culprit from run error: {cur_nodes}") + self.print_report(report) + return culprits + else: + report.append("No discrepancy found.") + self.print_report(report) + return set() + + def _skip_traverse(self, all_nodes: NodeList, skip_nodes: List) -> NodeSet: + """ + Skip certain nodes in graph based on settings + """ + start_idx = 0 + num_nodes = len(all_nodes) + idx = 0 + culprits = set() + while idx < num_nodes: + node = all_nodes[idx] + if node.name in skip_nodes: # skip the node + if idx > start_idx: + culprits = self._skip_traverse_impl(all_nodes, start_idx, idx) + start_idx = idx + 1 + elif idx == num_nodes - 1 and start_idx <= idx: # last node + culprits = self._skip_traverse_impl(all_nodes, start_idx, idx + 1) + idx += 1 + + return culprits + + def _collect_nodes(self, start: Optional[str], end: Optional[str]) -> NodeList: + """ + Collect nodes in the model that between nodes with name of `start` and `end`. + These two nodes are also included. + """ + nodes: NodeList = [] + add_node = start is None + + for node in self.module.graph.nodes: + if node.op not in CALLABLE_NODE_OPS: + continue + + if node.name == start: + add_node = True + + if add_node: + nodes.append(node) + + if node.name == end: + break + + return nodes + + def run_nodes(self, start: Optional[str] = None, end: Optional[str] = None): + """ + Run part of the model from `start` node to `end` node. If `start` is None + then we start from the beginning of the model. If `end` is None then we + stop at the end of the model. + + Args: + start: The name of the node which is the first node of the submodule + we want to run. If set to None, then we'll start with the first + node of the model. + end: The name of the node which is the last node of the submodule we + want to run. If set to None, we'll end with the last node of the + model. + """ + nodes = self._collect_nodes(start, end) + cur_nodes = set(nodes) + + for node in nodes: + if node in self.fusions: + cur_nodes.update(self.fusions[node]) + + output_names = [] + if self.settings.return_intermediate: + output_names = [node.name for node in nodes] + + try: + split_module, submod_name = self._build_submodule(cur_nodes) + self._run_and_compare(split_module, submod_name, output_names) + except ( + FxNetMinimizerRunFuncError, + FxNetMinimizerResultMismatchError, + ) as e: + print(e) + + def print_report(self, report: List[str]): + for i in range(len(report)): + if i > 0: + print(" . " + report[i]) + else: + print(report[i]) + + def print_reports(self): + for report in self.reports: + self.print_report(report) + + def minimize( + self, + start: Optional[str] = None, + end: Optional[str] = None, + skip_nodes: Optional[List] = None, + find_last_node: Optional[bool] = None, + ) -> NodeSet: + """ + Minimizing the model from node with name `start` to node with name `end` base + on self.settings. Find culprits that causes FxNetMinimizerRunFuncError or + FxNetMinimizerResultMismatchError errors. + + Args: + start: The name of the node where we want to start minimizing. If set + to None, then we'll start with the first node of the model. + end: The name of the node where we want to terminate minimizing. If + set to None, we'll end with the last node of the model. + skip_nodes: The names of nodes where we want to skip during minimizing. + It'll create subgraphs without these skip nodes under the hood. + Only applicable in mode "skip". + find_last_node: True if only last_node of a culprits is needed in mode "block". + False if only the first_node of a culprits is needed. + Only applicable in mode "block". + + Returns: + nodes: A list of nodes that causes FxNetMinimizerRunFuncError or + FxNetMinimizerResultMismatchError errors during minimizing. + """ + + print(self.settings) + print(self.module.graph) + + nodes = self._collect_nodes(start, end) + + if self.settings.traverse_method == "sequential": + return self._sequential_traverse(nodes) + + if self.settings.traverse_method == "binary": + return self._binary_traverse(nodes) + + if self.settings.traverse_method == "accumulate": + return self._accumulate_traverse(nodes) + + if self.settings.traverse_method == "skip": + if skip_nodes is None: + raise RuntimeError( + "'skip_nodes' can't be None when 'traverse_method' is 'skip'." + ) + return self._skip_traverse(nodes, skip_nodes) + + if self.settings.traverse_method == "defined": + return self._defined_traverse(nodes) + + if self.settings.traverse_method == "block": + return self._block_traverse(nodes, find_last_node) + + raise RuntimeError(f"Unknown traverse method {self.settings.traverse_method}!") diff --git a/lib/python3.12/site-packages/torch/fx/passes/operator_support.py b/lib/python3.12/site-packages/torch/fx/passes/operator_support.py new file mode 100644 index 0000000000000000000000000000000000000000..53e8be37cecf59ab783f6e05d88a27032f773b31 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/operator_support.py @@ -0,0 +1,229 @@ +# mypy: allow-untyped-defs +import abc +import typing as t + +import torch +import torch.fx +from torch.fx._compatibility import compatibility + +from .shape_prop import TensorMetadata +from .tools_common import CALLABLE_NODE_OPS, get_node_target + + +__all__ = [ + "OperatorSupportBase", + "OperatorSupport", + "create_op_support", + "chain", + "OpSupports", + "any_chain", +] + +# fx.Node.target typename, as returned by `get_node_target()` +TargetTypeName = str + +# Arguments' dtypes for a given node, see `OperatorSupport` +SupportedArgumentDTypes = t.Optional[ + t.Tuple[ + t.Sequence[t.Sequence[torch.dtype]], + t.Dict[str, t.Sequence[torch.dtype]], + ] +] + +SupportDict = t.Mapping[TargetTypeName, SupportedArgumentDTypes] + + +@compatibility(is_backward_compatible=False) +class OperatorSupportBase(abc.ABC): + """Interface for determining if a fx.Node is supported by a backend""" + + @abc.abstractmethod + def is_node_supported( + self, submodules: t.Mapping[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: + raise NotImplementedError + + +@compatibility(is_backward_compatible=False) +class OperatorSupport(OperatorSupportBase): + """ + `_support_dict` maps node.target typename to supported inputs dtypes. + + node.target typename is retrieved using helper function `get_node_target()` + + If supported inputs dtypes is None, it means any dtype is supported, else + we should see a tuple like (([dtypes], ...), {"name":[dtypes], ...}). + + The first tuple ([dtypes], ...) indicates what dtypes are supported for + inputs in node.args and the second dict {"name": [dtypes], ...} indicates + what dtypes are supported for inputs in node.kwargs. + + For inputs in args, if we don't want to check it, we can put None there, + e.g. (None, [torch.float]) indicates that we don't care about the type of + the first input in args. And for inputs in kwargs, if not listed, will not + be checked. + """ + + _support_dict: SupportDict + + def __init__(self, support_dict: t.Optional[SupportDict] = None): + self._support_dict = support_dict or {} + + def is_node_supported( + self, submodules: t.Mapping[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: + """ + Args: + `submodules`: mapping from module name to the module. This can be + retrieved by calling model.named_modules(). + + `node`: a Fx node that we want to determine whether it's supported. + + Returns: + `is_supported`: whether the arg `node` is supported. + """ + if node.op not in CALLABLE_NODE_OPS: + return True + + target = get_node_target(submodules, node) + + # Target not found in _support_dict meaning that we don't support this op at all + if target not in self._support_dict: + return False + + # The rule for target is None meaning that we accept any dtype + if self._support_dict[target] is None: + return True + + args_dtypes, kwargs_dtypes = self._support_dict[target] # type: ignore[misc] + + # Check args dtypes + for i, dtypes in enumerate(args_dtypes): + if len(node.args) <= i: + break + + # None indicates we don't care about the dtype of args[i] + if dtypes is None: + continue + + # If arg is not a node then we don't check it + if not isinstance(node.args[i], torch.fx.Node): + continue + + arg_dtype = _get_arg_dtype(node.args[i]) # type: ignore[arg-type] + if arg_dtype not in dtypes: + return False + + # Check kwargs dtypes + for k, dtypes in kwargs_dtypes.items(): + if k not in node.kwargs: + continue + + # If arg is not a node then we don't check it + if not isinstance(node.kwargs[k], torch.fx.Node): + continue + + kwarg_dtype = _get_arg_dtype(node.kwargs[k]) # type: ignore[arg-type] + if kwarg_dtype not in dtypes: + return False + + return True + + +# ====================================================================== +# Functional interfaces and utils for defining basic operator support logic +# and composing them into more complex ones +# ====================================================================== + +IsNodeSupported = t.Callable[[t.Mapping[str, torch.nn.Module], torch.fx.Node], bool] + + +@compatibility(is_backward_compatible=False) +def create_op_support(is_node_supported: IsNodeSupported) -> OperatorSupportBase: + """Wraps a `IsNodeSupported` function into an `OperatorSupportBase` instance + + `IsNodeSupported` has the same call signature as + `OperatorSupportBase.is_node_supported` + """ + + class FunctionalOperatorSupport(OperatorSupportBase): + def is_node_supported( + self, submodules: t.Mapping[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: + return is_node_supported(submodules, node) + + return FunctionalOperatorSupport() + + +@compatibility(is_backward_compatible=False) +def chain(*op_support: OperatorSupportBase) -> OperatorSupportBase: + """Combines a sequence of `OperatorSupportBase` instances to form a single `OperatorSupportBase` + instance by evaluating each input `OperatorSupportBase` instance, and returns False if + any of it reports False. + """ + + def _chain(submods, node) -> bool: + return all(x.is_node_supported(submods, node) for x in op_support) + + return create_op_support(_chain) + + +@compatibility(is_backward_compatible=False) +def any_chain(*op_support: OperatorSupportBase) -> OperatorSupportBase: + """Combines a sequence of `OperatorSupportBase` instances to form a single `OperatorSupportBase` + instance by evaluating each input `OperatorSupportBase` instance, and returns True if + any of it reports True. + """ + + def _any_chain(submods, node) -> bool: + return any(x.is_node_supported(submods, node) for x in op_support) + + return create_op_support(_any_chain) + + +@compatibility(is_backward_compatible=False) +class OpSupports: + """A set of atomic `OperatorSupportBase` instances that can be combined together + to form more complex operator support logic. + """ + + @classmethod + def decline_if_input_dtype(cls, dtype: torch.dtype) -> OperatorSupportBase: + """Report a node as non-supported, if any of its arguments is of dtype""" + + def _decline_if_input_dtype( + submodules: t.Mapping[str, torch.nn.Module], + node: torch.fx.Node, + ) -> bool: + for arg in node.all_input_nodes: + arg_dtype = _get_arg_dtype(arg) + if arg_dtype == dtype: + return False + return True + + return create_op_support(_decline_if_input_dtype) + + @classmethod + def decline_if_node_in_names(cls, disallow_set: t.Set[str]) -> OperatorSupportBase: + """ + If a node has a name that is in the disallow set, reported it as non-supported. + """ + + def _decline_if_node_in_names( + submodules: t.Mapping[str, torch.nn.Module], + node: torch.fx.Node, + ) -> bool: + return node.name not in disallow_set + + return create_op_support(_decline_if_node_in_names) + + +def _get_arg_dtype(arg: torch.fx.Node) -> t.Any: + assert isinstance(arg, torch.fx.Node) + tensor_meta = arg.meta.get("tensor_meta") # type: ignore[union-attr] + dtype = ( + tensor_meta.dtype + if isinstance(tensor_meta, TensorMetadata) + else arg.meta["type"] + ) + return dtype diff --git a/lib/python3.12/site-packages/torch/fx/passes/param_fetch.py b/lib/python3.12/site-packages/torch/fx/passes/param_fetch.py new file mode 100644 index 0000000000000000000000000000000000000000..3eba16b06b0351292843914331517debc2a5aef0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/param_fetch.py @@ -0,0 +1,96 @@ +from typing import Any, Callable, Dict, List, Tuple, Type + +import torch +import torch.nn as nn +from torch.fx._compatibility import compatibility +from torch.fx.graph_module import GraphModule + + +__all__ = [ + "default_matching", + "extract_attrs_for_lowering", + "lift_lowering_attrs_to_nodes", +] + + +# Matching method matches the attribute name of current version to the attribute name of `target_version` +@compatibility(is_backward_compatible=False) +def default_matching(name: str, target_version: int) -> str: + """Default matching method""" + return name + + +# This dict maps the nn.Module class name to the attribute name list that we want to fetch for lowering. +# The first integer in the tuple is the version number of the nn.Module class when we create the parameter list. +# If there's a version mismatch then it means the parameter names in the book might be mismatched with nn.Module. +module_fetch_book: Dict[Type, Tuple[int, List[str], Callable[[str, int], str]]] = { + torch.nn.modules.linear.Linear: (1, ["weight", "bias"], default_matching), + torch.nn.modules.conv.Conv2d: ( + 1, + [ + "weight", + "bias", + "kernel_size", + "stride", + "padding", + "dilation", + "groups", + "padding_mode", + ], + default_matching, + ), + torch.nn.modules.batchnorm.BatchNorm2d: ( + 2, + ["weight", "bias", "running_mean", "running_var", "eps"], + default_matching, + ), + torch.nn.modules.pooling.AdaptiveAvgPool2d: (1, [], default_matching), + torch.nn.modules.pooling.MaxPool2d: ( + 1, + ["kernel_size", "stride", "padding", "dilation", "return_indices", "ceil_mode"], + default_matching, + ), + torch.nn.modules.activation.ReLU: (1, ["inplace"], default_matching), +} + + +@compatibility(is_backward_compatible=False) +def extract_attrs_for_lowering(mod: nn.Module) -> Dict[str, Any]: + """If `mod` is in `module_fetch_book`, fetch the mod's attributes that in the `module_fetch_book` + after checking module's version is compatible with the `module_fetch_book`. + """ + attrs_for_lowering: Dict[str, Any] = {} + attrs_for_lowering["name"] = torch.typename(mod) + + if type(mod) in module_fetch_book: + version, param_to_fetch, matching_method = module_fetch_book[type(mod)] + if version < mod._version: + raise RuntimeError( + f"Fetcher version {version} try to fetch {torch.typename(mod)} version {mod._version}, " + "please upgrade the module_fetch_book, open an issue and @842974287 " + "or report a bug to AIACC team directly." + ) + for attr in param_to_fetch: + attrs_for_lowering[attr] = getattr(mod, matching_method(attr, mod._version)) + else: + raise RuntimeError( + f"{torch.typename(mod)} is not in the module_fetch_book yet, " + "please add it to the module_fetch_book, open an issue and @842974287 " + "or report a bug to AIACC team directly." + ) + return attrs_for_lowering + + +@compatibility(is_backward_compatible=False) +def lift_lowering_attrs_to_nodes(fx_module: GraphModule) -> None: + """Recursively traverse all `fx_module` nodes and fetch the module's attributes if the node is a leaf module.""" + submodules = dict(fx_module.named_modules()) + + for node in fx_module.graph.nodes: + if node.op == "call_module": + if isinstance(submodules[node.target], GraphModule): + lift_lowering_attrs_to_nodes(submodules[node.target]) + else: + node.attrs_for_lowering = extract_attrs_for_lowering( + submodules[node.target] + ) diff --git a/lib/python3.12/site-packages/torch/fx/passes/pass_manager.py b/lib/python3.12/site-packages/torch/fx/passes/pass_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..9b0ccbb82d50ddfd4e0e8f58d822565e413ef9e5 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/pass_manager.py @@ -0,0 +1,253 @@ +# mypy: allow-untyped-defs +import logging +from functools import wraps +from inspect import unwrap +from typing import Callable, List, Optional + + +logger = logging.getLogger(__name__) + +__all__ = [ + "PassManager", + "inplace_wrapper", + "log_hook", + "loop_pass", + "this_before_that_pass_constraint", + "these_before_those_pass_constraint", +] + + +# for callables which modify object inplace and return something other than +# the object on which they act +def inplace_wrapper(fn: Callable) -> Callable: + """ + Convenience wrapper for passes which modify an object inplace. This + wrapper makes them return the modified object instead. + + Args: + fn (Callable[Object, Any]) + + Returns: + wrapped_fn (Callable[Object, Object]) + """ + + @wraps(fn) + def wrapped_fn(gm): + fn(gm) + return gm + + return wrapped_fn + + +def log_hook(fn: Callable, level=logging.INFO) -> Callable: + """ + Logs callable output. + + This is useful for logging output of passes. Note inplace_wrapper replaces + the pass output with the modified object. If we want to log the original + output, apply this wrapper before inplace_wrapper. + + + ``` + def my_pass(d: Dict) -> bool: + changed = False + if "foo" in d: + d["foo"] = "bar" + changed = True + return changed + + + pm = PassManager(passes=[inplace_wrapper(log_hook(my_pass))]) + ``` + + Args: + fn (Callable[Type1, Type2]) + level: logging level (e.g. logging.INFO) + + Returns: + wrapped_fn (Callable[Type1, Type2]) + """ + + @wraps(fn) + def wrapped_fn(gm): + val = fn(gm) + logger.log(level, "Ran pass %s\t Return value: %s", fn, val) + return val + + return wrapped_fn + + +def loop_pass( + base_pass: Callable, + n_iter: Optional[int] = None, + predicate: Optional[Callable] = None, +): + """ + Convenience wrapper for passes which need to be applied multiple times. + + Exactly one of `n_iter`or `predicate` must be specified. + + Args: + base_pass (Callable[Object, Object]): pass to be applied in loop + n_iter (int, optional): number of times to loop pass + predicate (Callable[Object, bool], optional): + + """ + assert (n_iter is not None) ^ ( + predicate is not None + ), "Exactly one of `n_iter`or `predicate` must be specified." + + @wraps(base_pass) + def new_pass(source): + output = source + if n_iter is not None and n_iter > 0: + for _ in range(n_iter): + output = base_pass(output) + elif predicate is not None: + while predicate(output): + output = base_pass(output) + else: + raise RuntimeError( + f"loop_pass must be given positive int n_iter (given " + f"{n_iter}) xor predicate (given {predicate})" + ) + return output + + return new_pass + + +# Pass Schedule Constraints: +# +# Implemented as 'depends on' operators. A constraint is satisfied iff a list +# has a valid partial ordering according to this comparison operator. +def _validate_pass_schedule_constraint( + constraint: Callable[[Callable, Callable], bool], passes: List[Callable] +): + for i, a in enumerate(passes): + for j, b in enumerate(passes[i + 1 :]): + if constraint(a, b): + continue + raise RuntimeError( + f"pass schedule constraint violated. Expected {a} before {b}" + f" but found {a} at index {i} and {b} at index{j} in pass" + f" list." + ) + + +def this_before_that_pass_constraint(this: Callable, that: Callable): + """ + Defines a partial order ('depends on' function) where `this` must occur + before `that`. + """ + + def depends_on(a: Callable, b: Callable): + return a != that or b != this + + return depends_on + + +def these_before_those_pass_constraint(these: Callable, those: Callable): + """ + Defines a partial order ('depends on' function) where `these` must occur + before `those`. Where the inputs are 'unwrapped' before comparison. + + For example, the following pass list and constraint list would be invalid. + ``` + passes = [ + loop_pass(pass_b, 3), + loop_pass(pass_a, 5), + ] + + constraints = [these_before_those_pass_constraint(pass_a, pass_b)] + ``` + + Args: + these (Callable): pass which should occur first + those (Callable): pass which should occur later + + Returns: + depends_on (Callable[[Object, Object], bool] + """ + + def depends_on(a: Callable, b: Callable): + return unwrap(a) != those or unwrap(b) != these + + return depends_on + + +class PassManager: + """ + Construct a PassManager. + + Collects passes and constraints. This defines the pass schedule, manages + pass constraints and pass execution. + + Args: + passes (Optional[List[Callable]]): list of passes. A pass is a + callable which modifies an object and returns modified object + constraint (Optional[List[Callable]]): list of constraints. A + constraint is a callable which takes two passes (A, B) and returns + True if A depends on B and False otherwise. See implementation of + `this_before_that_pass_constraint` for example. + """ + + passes: List[Callable] + constraints: List[Callable] + _validated: bool = False + + def __init__( + self, + passes=None, + constraints=None, + ): + self.passes = passes or [] + self.constraints = constraints or [] + + @classmethod + def build_from_passlist(cls, passes): + pm = PassManager(passes) + # TODO(alexbeloi): add constraint management/validation + return pm + + def add_pass(self, _pass: Callable): + self.passes.append(_pass) + self._validated = False + + def add_constraint(self, constraint): + self.constraints.append(constraint) + self._validated = False + + def remove_pass(self, _passes: List[str]): + if _passes is None: + return + passes_left = [ps for ps in self.passes if ps.__name__ not in _passes] + self.passes = passes_left + self._validated = False + + def replace_pass(self, _target, _replacement): + passes_left = [] + for ps in self.passes: + if ps.__name__ == _target.__name__: + passes_left.append(_replacement) + else: + passes_left.append(ps) + self.passes = passes_left + self._validated = False + + def validate(self): + """ + Validates that current pass schedule defined by `self.passes` is valid + according to all constraints in `self.constraints` + """ + if self._validated: + return + for constraint in self.constraints: + _validate_pass_schedule_constraint(constraint, self.passes) + self._validated = True + + def __call__(self, source): + self.validate() + out = source + for _pass in self.passes: + out = _pass(out) + return out diff --git a/lib/python3.12/site-packages/torch/fx/passes/reinplace.py b/lib/python3.12/site-packages/torch/fx/passes/reinplace.py new file mode 100644 index 0000000000000000000000000000000000000000..3b61446a92f7ec0c843aea059781a10919f869c1 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/reinplace.py @@ -0,0 +1,753 @@ +# mypy: allow-untyped-defs +import _operator +import itertools +from collections import defaultdict +from enum import Enum +from typing import Dict, Set + +import torch +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode +from torch.fx import Node +from torch.fx._compatibility import compatibility +from torch.multiprocessing.reductions import StorageWeakRef +from torch.utils import _pytree as pytree +from torch.utils._pytree import tree_map_only + + +__all__ = ["reinplace"] + + +class _ViewType(Enum): + NonView = 0 + SingleOutputView = 1 + MultiOutputView = 2 + + +def _is_view_op(tgt): + if tgt is not None and isinstance(tgt, torch._ops.OpOverload): + schema = tgt._schema + if len(schema.arguments) > 0: + first_arg = schema.arguments[0] + # check if op is a view + return ( + first_arg.alias_info is not None and not first_arg.alias_info.is_write + ) + + +def _get_view_type(tgt) -> _ViewType: + if tgt is not None and isinstance(tgt, torch._ops.OpOverload): + schema = tgt._schema + if len(schema.arguments) > 0: + first_arg = schema.arguments[0] + # check if op is a view + if first_arg.alias_info is not None and not first_arg.alias_info.is_write: + # check if op is a multi-output view + if "*" in first_arg.alias_info.after_set: + return _ViewType.MultiOutputView + else: + return _ViewType.SingleOutputView + return _ViewType.NonView + + +# Stores a bunch of metadata related to functionalization each node. +# Relevant metadata: +# n.meta['fake_result']: FakeTensor (same type as the output of the node, but with FakeTenors instead of Tensors) +# The fake tensor output from running the current node +# n.meta['view_of']: Node +# If the current node n is a view of some base tensor, the 'view_of' field tells us which +# view node was used to generate the current node (a view tensor). +# This information actually makes `fake_result` redundant, but we can use `fake_result` +# to sanity check that our aliasing information is correct. +@compatibility(is_backward_compatible=False) +class _FunctionalizationMetadataProp(torch.fx.Interpreter): + def run_node(self, node: Node): + self.node_counter += 1 + result = super().run_node(node) + node.meta["fake_result"] = result + node.meta["node_idx"] = self.node_counter + + # (1) Update metadata with the list of nodes that are used by this node + # copy_() doesn't read from its first argument; it writes to it, overwriting previous data. + # We don't want to treat it as "being used as an input". + node_args = node.args + if node.target is torch.ops.aten.copy_.default: + node_args = node_args[1:] + + # (2) Update metadata to track aliasing information about view tensor nodes. + if node.op == "call_function": + view_type = _get_view_type(node.target) + if view_type == _ViewType.SingleOutputView: + assert isinstance(node.args[0], Node) + node.meta["view_of"] = node.args[0] + elif view_type == _ViewType.MultiOutputView: + self.multi_output_view_nodes[node] = node.args[0] + + # Check if we returned a multi-output view, + # and we're now grabbing the individual views from the output. + # + # For multi-output views, we want to map each output view to the base, + # but this mapping involves two separate nodes in FX IR. + # e.g. "a, b = x_1.split(...)" becomes: + # %split_tensor : [num_users=2] = call_function[target=torch.ops.aten.split.Tensor](args = (%x_1, 2), kwargs = {}) + # %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%split_tensor, 0), kwargs = {}) + # %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%split_tensor, 1), kwargs = {}) + # And we'd like to set: + # getitem1.meta['view_of'] = x_1 + elif node.target is _operator.getitem: + list_arg = node.args[0] + maybe_base_of_view = self.multi_output_view_nodes.get(list_arg, None) + if maybe_base_of_view is not None: + # Note: we could also track indexing info here for multi-output views. + # I don't think this metadata is strictly needed for de-functionalization. + assert isinstance(maybe_base_of_view, Node) + node.meta["view_of"] = maybe_base_of_view + + if "view_of" in node.meta: + # We're linking the current node with its first argument as views. + # Assert here that this is actually the case, and their storages are the same. + assert isinstance(node.meta["fake_result"], FakeTensor) + assert isinstance(node.meta["view_of"].meta["fake_result"], FakeTensor) + view_storage = StorageWeakRef(node.meta["fake_result"]._typed_storage()) + base_storage = StorageWeakRef( + node.meta["view_of"].meta["fake_result"]._typed_storage() + ) + assert view_storage == base_storage + return result + + def propagate(self, *args): + self.multi_output_view_nodes = {} + self.node_counter = -1 + + with FakeTensorMode() as mode: + fake_args = [ + mode.from_tensor(a) if isinstance(a, torch.Tensor) else a for a in args + ] + return super().run(*fake_args) + + +def _schemas_match(functional_schema, inplace_schema): + names_match = ( + inplace_schema.name.endswith("_") + and inplace_schema.name[:-1] == functional_schema.name + ) + arg_types_match = len(functional_schema.arguments) == len( + inplace_schema.arguments + ) and all( + a1.type == a2.type + for a1, a2 in zip(functional_schema.arguments, inplace_schema.arguments) + ) + # for the inplace op, its first argument should be mutable + assert ( + inplace_schema.arguments[0].alias_info is not None + and inplace_schema.arguments[0].alias_info.is_write + ) + # and its remaining arguments shouldn't be. + assert all(a.alias_info is None for a in inplace_schema.arguments[1:]) + return names_match and arg_types_match + + +# TODO: this should be beefed up to be able to properly re-inplace with: +# - mutating ops (e.g. _fused_moving_avg_obs_fq_helper) +# - out= ops (e.g. angle -> angle.out) +# TODO: we should also figure this info out using torchgen. +def _maybe_get_inplace_op(op): + # __module__ seems broken; it returns torch._ops.aten which doesn't exist + if not isinstance(op, torch._ops.OpOverload): + return None + # Some view ops have inplace variants (as_strided_, etc), + # but we do NOT want the reinplacing pass to directly add these into the program. + # (they'll require extra special handling, aren't aren't really useful for perf anyway) + if _is_view_op(op): + return None + op_namespace = op.__module__.split(".")[-1] + op_base_name = op.overloadpacket.__name__ + maybe_namespace_module = getattr(torch.ops, op_namespace) + maybe_inplace_op = ( + None + if maybe_namespace_module is None + else getattr(maybe_namespace_module, f"{op_base_name}_", None) + ) + if maybe_inplace_op is None: + return None + + inplace_overloads = [ + getattr(maybe_inplace_op, overload_name) + for overload_name in maybe_inplace_op.overloads() + ] + inplace_overloads_with_matching_schemas = [ + f for f in inplace_overloads if _schemas_match(op._schema, f._schema) + ] + # Just because foo() and foo_() are both existing operators, + # They aren't guaranteed to have compatible schemas. + # For example, pow.Scalar(Scalar self, Tensor exponent) has no valid inplace variant, + # Even though several overloads of pow_ exist. + if len(inplace_overloads_with_matching_schemas) == 0: + return None + assert len(inplace_overloads_with_matching_schemas) == 1 + inplace_op = inplace_overloads_with_matching_schemas[0] + return inplace_op + + +_VIEW_INVERSE_MAP = { + torch.ops.aten.diagonal_scatter.default: torch.ops.aten.diagonal.default, + torch.ops.aten.select_scatter.default: torch.ops.aten.select.int, + torch.ops.aten.slice_scatter.default: torch.ops.aten.slice.Tensor, + torch.ops.aten.as_strided_scatter.default: torch.ops.aten.as_strided.default, +} + + +# This function, given a set of set of (aliased) tensor nodes, +# Returns any nodes in the graph that *use* any of the aliases, that occur *after* op_index +# in the node ordering. +def _get_all_later_node_usages(tensor_aliases: Set[Node], op_index: int): + def _add_if_tensor(x, set_): + if isinstance(x, FakeTensor): + set_.add(StorageWeakRef(x._typed_storage())) + + nodes_used_after = set() + for t in tensor_aliases: + # get all nodes that use the current alias + usage_nodes = t.users + for n in usage_nodes: + # We only care about usages after the current node + if "node_idx" not in n.meta or n.meta["node_idx"] <= op_index: + continue + # We also don't care about intermediate view ops. + # They only matter if their output is then used elsewhere + # (either in an out-of-place op, or as an output to the function). + if n in tensor_aliases: + if ( + isinstance(n.target, torch._ops.OpOverload) + or n.target == _operator.getitem + ): + continue + nodes_used_after.add(n) + return nodes_used_after + + +# Given an op that we're trying to re-inplace, "b = foo(a)", +# And given a {view}_scatter op that shows up later in the graph, "y = {view}_scatter(base, x, args...)" +# Then re-inplacing `foo()` would allow us to remove the `{view}_scatter` op entirely, IF: +# If there are any aliases in the alias_set(a) that satisfy: +# (1) The base of "alias", "alias_base", has the same size/stride/offset metadata as "base" +# (2) The output of running {view}(alias, args...) gives you the same size/stride/offset metadata +# as "alias" +def _get_view_inverse_node_usages( + later_node_usages: Set[Node], self_aliases: Set[Node] +) -> Set[Node]: + def matching_view_metadata(a, b): + return ( + a.size() == b.size() + and a.stride() == b.stride() + and a.storage_offset() == b.storage_offset() + ) + + view_inverse_nodes = set() + # Go through them in node order, so we can see chains of view_scatter ops. + for n in sorted(later_node_usages, key=lambda x: x.meta["node_idx"]): + if n.target not in _VIEW_INVERSE_MAP: + continue + base = n.args[0] + mutated_view = n.args[1] + assert isinstance(base, Node) + assert isinstance(base.meta["fake_result"], FakeTensor) + assert isinstance(mutated_view, Node) + assert isinstance(mutated_view.meta["fake_result"], FakeTensor) + # Check that this view_inverse op actually corresponds to taking doing the inverse + # of one of our existing self_alias nodes. + original_view = _VIEW_INVERSE_MAP[n.target] + for self_alias in self_aliases: + # We're looking for some alias of the self arg, "alias", + # that was created from some op `alias = foo(base, args...)` + # such that the current _scatter op "inverts" that foo call. + # We can check that by running the original op again, and checking that the strides match. + if "view_of" not in self_alias.meta: + continue + self_alias_base = self_alias.meta["view_of"] + try: + # The we're trying to re-use the args from the view_scatter call inside of the corresponding + # view op, which might throw. This just indicates that view_scatter op isn't a valid inverse + # of the current alias we're looking at. + view_replay_metadata = original_view( + self_alias_base.meta["fake_result"], *n.args[2:], **n.kwargs + ) + expected_metadata = self_alias.meta["fake_result"] + # If the alias and its base both have matching metadata, then this view_scatter op is valid to re-inplace. + if matching_view_metadata( + self_alias_base.meta["fake_result"], base.meta["fake_result"] + ) and matching_view_metadata(view_replay_metadata, expected_metadata): + view_inverse_nodes.add(n) + except Exception: + continue + + return view_inverse_nodes + + +@compatibility(is_backward_compatible=True) +def reinplace(gm, *sample_args): + """ + Given an fx.GraphModule, modifies it to perform "reinplacing", + mutating the nodes of the graph. + We look for out-of-place op call sites like `b = a.add(...)`, + and convert them to be inplace (`b = a.add_(...)`), + as long as the input to the current operator ("a") isn't re-used + anywhere later in the graph. + + This pass currently expects to operate on a **functional, ATen** graph. + This can be obtained by running `make_fx(functionalize(f))`. + + Sample inputs are needed to determine aliasing relationships of the inputs. + In general, we can't reinplace node `b = a.add(...)` if "a" aliases any of the + inputs to the program. + + Given a node "b = foo(a, args...) the algorithm for re-inplacing is as follows: + + (1) Perform some initial checks on the metadata of "a" and "args..." + that can disqualify them from being reinplaced. + + (1a) Check that the self argument we're attempting to reinplace + has acceptable dtype/size metadata to reinplace with. + + For example, if we have: + a = torch.ones(1) + b = torch.ones(10) + out = torch.add(a, b) + We can't turn that into + a.add_(b) + Because that would require resizing "a". + + Similarly, we can't convert torch.ge(a, b) into a.ge_(b), + because that would require changing a's dtype (from e.g. float32 to bool). + Note that in this specific example, we could technically do better.. + + If we see the pattern: + a_1 = a.ge(b) + a_2 = aten._to_copy(a_1, a.dtype) + Then we this should be valid to completely re-inplace + (this is exactly what functionalization will emit when it sees a.ge_(b)). + + This optimization is only really important for user programs + that directly use inplace comparison ops though. + + We also cannot re-inplace on tensors that have overlapping memory, + e.g. torch.ones(1).expand(4, 4).add_(1) + + (1b) Check if "a" is an alias of any of the program inputs. + + If it is, skip and move to the next node. + Inplace'ing an op that would cause it to mutate a program is not sound, + because that would be a side effect visible to the user. + + NOTE: there's a future optimization that we should make: + if "a" is a (alias of a) program input, but later in the program + there is a node that looks like "a.copy_(...)", + Then re-inplacing is ok to do - we are temporarily re-using a's buffer, + which will later be overwritten by the copy_() call. + + This will be an important optimization to have for programs that mutate + their inputs. It currently isn't implemented though. + + (1c) Check if "a" and "args..." alias + + For example, re-inplacing to create code like the below + isn't guaranteed to be sound: + + aten.mul_(a, a) + + (2) Check that "a" and all of its outstanding aliases are not used anywhere + later in the graph. If this is the case, then it's safe to re-inplace + to "b = foo_(a)". + + There are a few caveats to this, explained in more detail below: + (a) If "a" is used later as an argument to a view op, that is okay. + It's only a problem if "a" (or that view) is later passed + into a normal operator, or if it is returned as the program output. + (b) If "a" is a repeat argument in `foo()`, then don't reinplace. + Most ATen kernels don't make any guarantees that this is sound, + e.g. if you do aten.mul_(a, a). + So we'll just ban re-inplacing in this case. + It's only a problem if "a" (or that view) is later passed + (c) If "a" is used as an input into a view "inverse" / "scatter" + operator, it is potentially fine to re-inplace + (and remove that scatter operator from the graph). + See below for a more detailed example. + + NOTE: there is an optimization in this step that is crucial + to fully recovering performance from functionalization. + + Given this program: + def f(x): + a = torch.ops.aten.add(x, x) + b = torch.ops.aten.diagonal(a) + torch.ops.aten.fill_(b, 0) + return d + + Functionalization will emit the following: + def f(x): + a = torch.ops.aten.add(x, x) + b = torch.ops.aten.diagonal(a, 0, 1) + b_updated = torch.ops.aten.fill(b, 0) + a_updated = torch.ops.aten.diagonal_scatter(a, b_updated, 0, 1) + return a_updated + + Ordinarily, we would not be able to reinplace the fill, + because "b" aliases with "a" which is used by the diagonal_scatter call. + + "re-inplacing" is on the hook for figuring out that it is ok to + completely, the expensive diagonal_scatter call, if we re-inplace the add(). + + So, for every `alias in alias_set(a)`, instead of checking + that "alias" is not used anywhere later in the graph, + we check that + EITHER: + (a) alias is not used anywhere later in the graph + OR: + (b) alias is used exactly once later on in the graph, + in the following op: + + out = foo_scatter(alias, x, args...) + + where the following must hold: + (i) "foo_scatter" is the "inverse" operator for foo. + This only applies to "foo" ops that are view operators, + which view into a subset of the original tensor's memory. + In practice, there are ~4 operators where this applies: + diagonal -> diagonal_scatter + slice -> slice_scatter + select -> select_scatter + as_strided -> as_strided_scatter + (ii) "args..." are the same between the foo() and foo_scatter() calls. + + (3) Perform the actual re-inplacing on foo! + + (3b) is the common case, but special care is needed for {view}_scatter (3a) + + (3a) {view}_scatter ops. + + Consider this program: + a = torch.zeros(2, 2) + b = torch.ones(2) + a[0] = b + + Post functionalization, that will look like: + a = torch.zeros(2) + b = torch.ones(1) + a_updated = torch.select_scatter(a, b, 0, 0) + + In this case though, there is no "functional" op to re-inplace! + Instead, we'd like to directly remove toe select_scatter call. + We already know from (3) that this is valid, + because "a" has no later usages in the graph. + + We perform the re-inplacing on the {view}_scatter op like so + Before: + a_updated = torch.select_scatter(a, b, args...) + After: + a_slice = a.select(a, args...) + a_slice.copy_(b) + + (3b) Otherwise, replace the functional op with its inplace variant. + Before: + b = foo(a, args...) + After: + a.foo_(args...) + + (4) Finally, after converting either: + Before: + b = foo(a) + After: + foo_(a) + or + Before: + b = {slice}_scatter(a, mutated_slice, args...) + After: + slice = {slice}(a, args...) + slice.copy_(mutated_slice) + + We now need to find all later nodes that use "b" as an argument + and update them to take in "a" instead. + + Note that for the majority of inplace ops, this isn't actually necessary + (because most inplace ops return "self" as their output). + This isn't generally true for all mutable ops though, which is why + we need to actually replace all of the arguments. + + We also need to update our metadata of Dict[StorageWeakRef, Set[Node]], + That maps a given tensor storage to the set of all nodes that take in that storage + as an input. + Specifically, re-inplacing `b = foo(a)` causes "a" and "b"'s sets to get fused + together. + + (5) Any "view_inverse/scatter" nodes that were identified as "it's ok to ignore them" + during step (3) get manually deleted from the graph. + Their outputs are no longer used, so technically standard DCE would be able + to do this, but we can no longer run FX's DCE pass now that we have mutable + ops in the graph. + """ + _FunctionalizationMetadataProp(gm).propagate(*sample_args) + + # Useful debug printing + # def _print(x): + # if isinstance(x, FakeTensor): + # print(f'fake_result: {StorageWeakRef(x._typed_storage()).cdata}') + + # for n in gm.graph.nodes: + # print(n.format_node()) + # if hasattr(n, 'meta'): + # print(f'node_idx: {n.meta["node_idx"]}') + # if 'fake_result' in n.meta: + # tree_map(_print, n.meta['fake_result']) + # if 'view_of' in n.meta: + # print(f'view_of: {str(n.meta["view_of"])}') + # print() + + # We need to know which nodes correspond to inputs (or their aliases) + # so we know not to re-inplace them. + # NOTE: later, we'll need to add an optimization for fully recovering performance + # on programs that mutate inputs. + input_storages = { + StorageWeakRef(node.meta["fake_result"]._typed_storage()) + for node in gm.graph.nodes + if ( + node.op == "placeholder" + and isinstance(node.meta["fake_result"], torch.Tensor) + ) + } + + # We also need to know for a given node, what are all of its aliasing nodes. + storage_to_nodes: Dict[StorageWeakRef, Set[Node]] = defaultdict(set) + for n in gm.graph.nodes: + if "fake_result" in n.meta: + # Tree-mapping because some ops can return lists of tensors. + def _add_to_map(x): + if isinstance(x, FakeTensor): + storage_to_nodes[StorageWeakRef(x._typed_storage())].add(n) + + pytree.tree_map_(_add_to_map, n.meta["fake_result"]) + + # inplace-ify functional ops, subject to the constraints written below. + all_later_view_inverse_nodes_to_delete = set() + for node in gm.graph.nodes: + if node.op == "call_function": + # Today, the re-inplace pass on directly acts on: + # - functional ops with an inplace variant + # - {view}_scatter ops that can be potentially removed from the graph. + # Both of these ops take in tensor first args, so filtering on this condition + # makes the later code simpler. + # We should revisit this at some point though, particularly when we also want + # the reinplacer to be able to handle out= and mutable operators + # and tensorlist first args (like `_foreach_` ops). + if not isinstance(node.target, torch._ops.OpOverload): + continue + if len(node.target._schema.arguments) < 1: + continue + if type(node.target._schema.arguments[0].type) != torch.TensorType: + continue + + # Step 1a: Check that the self argument we're attempting to reinplace + # has the same size/stride as the output. + # For example, we shouldn't try to reinplace torch.add(scalar_tensor, larger_tensor) + # As it would require resizing scalar_tensor. + # (We could potentially swizzle this into larger_tensor.add_(scalar_tensor), + # this is probably an optimization to revisit later). + self_arg = node.args[0] + self_flattened = pytree.tree_leaves(self_arg.meta["fake_result"]) + node_flattened = pytree.tree_leaves(node.meta["fake_result"]) + self_has_wrong_metadata = False + if len(self_flattened) == len(node_flattened): + for self_meta, node_meta in zip(self_flattened, node_flattened): + if self_meta.numel() != node_meta.numel(): + self_has_wrong_metadata = True + if self_meta.dtype != node_meta.dtype: + self_has_wrong_metadata = True + # We also cannot re-inplace on tensors that have internal memory overlap. + # e.g. torch.ones(1).expand(4, 4).add_(1) + if torch._debug_has_internal_overlap(self_meta) == 1: + self_has_wrong_metadata = True + # Here, we (optimistically) assume that a.resize(b) is valid to re-inplace, + # Since users should never really be calling the functional "torch.ops.aten.resize" + # op directly in their programs. + if self_has_wrong_metadata and node.target != torch.ops.aten.resize.default: + continue + + # Step 1b: ensure that the op we're trying to re-inplace isn't a program input + self_arg_storage = StorageWeakRef( + self_arg.meta["fake_result"]._typed_storage() + ) + if self_arg_storage in input_storages: + # TODO: later, add the optimization for handling `copy_()` calls in the graph. + continue + if len([x for x in node.args if x is self_arg]) > 1: + # Step 1c: + # Calling stuff like aten.mul_(a, a) isn't guaranteed to be sound, + # so we prevent re-inplacing in this case. + continue + + self_arg_storage = StorageWeakRef( + self_arg.meta["fake_result"]._typed_storage() + ) + self_aliases = storage_to_nodes[self_arg_storage] + + # First, we find all later usages of any of the aliases of self_arg. + later_node_usages = _get_all_later_node_usages( + self_aliases, node.meta["node_idx"] + ) + # Then, we check if any of those later usages are actually view_scatter ops + # that are safe to fully remove. + later_view_inverse_node_usages = _get_view_inverse_node_usages( + later_node_usages, self_aliases + ) + + # Step 2: Check to see if the input to the op is re-used later in the graph. + # If not (same goes for its aliases), then this op is safe to re-in place. + # This is a slightly roundabout way to check that there are no later usages of the current self argument. + # (later_view_inverse_node_usages corresponds to "view_scatter" nodes that we are allowed to delete) + can_reinplace = len(later_node_usages - later_view_inverse_node_usages) == 0 + if not can_reinplace: + continue + + # Step 3a: Special handling for when we see *_scatter operators. + # When we see an operator like `b = torch.slice_scatter(a, ...)`, + # instead of trying to "inplace" it into a.slice_scatter_(..._), + # we would prefer to remove it from the graph entirely, + # and instead copy_() the slice directly into the larger tensor. + # See the description of the algorithm for a full example. + if ( + node.target in _VIEW_INVERSE_MAP + and node not in all_later_view_inverse_nodes_to_delete + ): + view_op = _VIEW_INVERSE_MAP[node.target] + # Before: + # base_updated = torch.ops.aten.slice_scatter.default(base, mutated_slice, args...) + # After: + # slice = torch.ops.aten.slice.default(base, args...) + # slice.copy_(mutated_slice) + with gm.graph.inserting_before(node): + mutated_slice_node = node.args[1] + remaining_slice_args = node.args[2:] + slice_node = gm.graph.create_node( + "call_function", + view_op, + (self_arg,) + tuple(remaining_slice_args), + node.kwargs, + ) + gm.graph.create_node( + "call_function", + torch.ops.aten.copy_.default, + ( + slice_node, + mutated_slice_node, + ), + {}, + ) + # Add the slice_scatter node to our "nodes to delete" list. + all_later_view_inverse_nodes_to_delete.add(node) + + else: + # Step 3b: Check to see if this operator has an inplace variant. + maybe_inplace_op = _maybe_get_inplace_op(node.target) + if maybe_inplace_op is None: + continue + # And if so, replace it with its inplace variant. + node.target = maybe_inplace_op + + # At this point, 'storage_to_nodes' will be stale. + # Now that we're inplacing `b = foo(a)`, we need to effectively + # union together the dict values for b and a's storage. + # Hmm... morally I think we also want to keep the `fake_result` metadata + # up to date here, but I'm not sure how easy it is to do. + # Maybe it's fine to wait until the end of the pass to update it. + curr_node_storage = StorageWeakRef( + node.meta["fake_result"]._typed_storage() + ) + storage_to_nodes[self_arg_storage].update( + storage_to_nodes[curr_node_storage] + ) + storage_to_nodes[curr_node_storage].update( + storage_to_nodes[self_arg_storage] + ) + + # Need to remember the view_scatter view nodes we found so we can remove them alter. + all_later_view_inverse_nodes_to_delete.update( + later_view_inverse_node_usages + ) + + # Step 4: + # Now that we've replaced b = a.foo() with a.foo_(), + # We need to replace any later usages of "b" with "a" + for old in itertools.chain([node], later_view_inverse_node_usages): + new = old.args[0] + nodes_to_update = [ + n for n in old.users if n.meta["node_idx"] > node.meta["node_idx"] + ] + for node_to_update in nodes_to_update: + + def replace_arg(a): + if a == old: + return new + return a + + # First, replace usages of "b" with "a" + node_to_update.args = tree_map_only( + Node, replace_arg, node_to_update.args + ) + node_to_update.kwargs = tree_map_only( + Node, replace_arg, node_to_update.kwargs + ) + + # Second, update our storage_to_nodes data structure. + old_flattened_res = pytree.tree_leaves(old.meta["fake_result"]) + node_flattened_res = pytree.tree_leaves( + node_to_update.meta["fake_result"] + ) + + old_res_storage = { + StorageWeakRef(x._typed_storage()) + for x in old_flattened_res + if isinstance(x, FakeTensor) + } + node_res_storage = { + StorageWeakRef(x._typed_storage()) + for x in node_flattened_res + if isinstance(x, FakeTensor) + } + + # This will happen if we're updating a view op, e.g. + # e.g. replacing + # x = view(old) + # x = view(new) + # When that happens, we need to make sure to keep our + # storage mapping up to date. + # + # We're checking for len(...) == 1 here because all view ops are guaranteed to return either a single tensor, + # or multiple tensors that all share the same storage. + # We can't just check equality because we might encounter FX nodes that return zero tensor outputs. + if ( + len(old_res_storage) == 1 + and len(node_res_storage) == 1 + and old_res_storage == node_res_storage + ): + new_flattened_res = pytree.tree_leaves(new.meta["fake_result"]) + new_res_storage = { + StorageWeakRef(x._typed_storage()) + for x in new_flattened_res + if isinstance(x, FakeTensor) + } + assert len(new_res_storage) == 1 + (new_ref,) = new_res_storage + (node_ref,) = node_res_storage + # Technically, "old_ref" and all its aliases will remain + # in our mapping. + # That should be fine though, since we deleted "old" + # from the graph at this point. + storage_to_nodes[node_ref].update(storage_to_nodes[new_ref]) + storage_to_nodes[new_ref].update(storage_to_nodes[node_ref]) + + # Step 4: delete any _scatter nodes that we de-functionalized + # Need to take care not to delete any of these nodes until after *all* modifications + # to the graph are finished. + for to_delete in all_later_view_inverse_nodes_to_delete: + gm.graph.erase_node(to_delete) + + gm.recompile() + return gm diff --git a/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py b/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py new file mode 100644 index 0000000000000000000000000000000000000000..9dd87d49ef52c67b388ee7980bf3799b12c5278d --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py @@ -0,0 +1,625 @@ +# mypy: allow-untyped-defs +import functools +import logging +import operator +import sys +from typing import Any, Dict, Optional, Set, TYPE_CHECKING + + +# Import sympy and ShapeEnv during TYPE_CHECKING since importing sympy is slow +if TYPE_CHECKING: + import sympy + + from torch.fx.experimental.symbolic_shapes import ShapeEnv +else: + ShapeEnv = Any + +import torch +import torch.utils._pytree as pytree +from torch import fx +from torch._subclasses.meta_utils import is_sparse_any +from torch.fx._compatibility import compatibility +from torch.fx._utils import lazy_format_graph_code +from torch.fx.experimental.proxy_tensor import py_sym_types +from torch.fx.experimental.sym_node import SymNode +from torch.fx.graph_module import GraphModule + + +__all__ = ["insert_deferred_runtime_asserts"] + +log = logging.getLogger(__name__) +graph_code_log = torch._logging.getArtifactLogger(__name__, "graph_code") + + +def _get_example_value(node: fx.Node) -> Optional[str]: + """ + Get the example value key for a node, since dynamo uses "example_value" + while non-strict export uses "val. + """ + if "example_value" in node.meta: + return node.meta["example_value"] + elif "val" in node.meta: + return node.meta["val"] + else: + return None + + +def _get_sym_val(node: fx.Node) -> Optional["sympy.Expr"]: + val = _get_example_value(node) + if isinstance(val, py_sym_types): + return val.node.expr + return None + + +@compatibility(is_backward_compatible=True) +def insert_deferred_runtime_asserts( + gm: GraphModule, + shape_env: ShapeEnv, + name: str, + export: bool = False, +) -> None: + """ + During tracing, we may have discovered that some data-dependent values + had runtime assert on them; e.g., torch.empty(x.item()) induces a runtime + that x.item() >= 0. This asserts can happen unpredictably during fake + tensor propagation, so we cannot conveniently insert them into the FX graph + when they occur. Instead, we accumulate them in the ShapeEnv, and in this + pass insert them into the graph as proper tests. + + This pass also deduplicates size-related computation, CSE-ing ops that produce + symbolic values and/or are involved in runtime asserts. Additionally, shape calls + (size/stride/storage_offset) are turned into compute on input sizes if possible, + allowing intermediate tensors to be freed earlier. For example, here dynamo will + DCE the cat and repeat calls: + + z = torch.cat([x, x], dim=0) # 2*s0 + w = z.repeat(y.shape[0]) # 2*s0*s1 + _w = w.shape[0] + # something with _w, but not w ... + + # turns into -> + _w0 = 2 * s0 + _w = _w0 * s1 + + # where s0, s1 are either SymInt graph inputs, or the result of added size calls + + Redundant torch._check or torch.ops.aten._assert_scalar.default calls that assert + the same expression, and redundant constrain_range calls are also deduplicated. + Additionally, because single-symbol bound checks (e.g. u0 >= 0, u0 <= 5) accumulate + information in the ShapeEnv, the ShapeEnv contains min/max bounds for each symbol, + and we delete all previous calls, adding bound checks at the end of this pass. + """ + + # Import sympy locally + import sympy + + from torch._export.passes._node_metadata_hook import _set_node_metadata_hook + from torch.fx.experimental.symbolic_shapes import ( + _has_uninterpretable_sympy_function, + CallMethodKey, + cast_symbool_to_symint_guardless, + ConvertIntKey, + DivideByKey, + free_symbols, + InnerTensorKey, + resolve_unbacked_bindings, + ) + from torch.utils._sympy.numbers import int_oo + from torch.utils._sympy.reference import ( + OptimizedPythonReferenceAnalysis, + PythonReferenceAnalysis, + ) + from torch.utils._sympy.value_ranges import ValueRanges + + # TODO: Request simplification on runtime asserts before emitting them + ras_by_symbol = shape_env.deferred_runtime_asserts.copy() + graph = gm.graph + tracer = fx.proxy.GraphAppendingTracer(graph) + graph_code_log.debug( + "%s", + lazy_format_graph_code( + f"pre insert_deferred_runtime_asserts {name}", gm, colored=True + ), + ) + + # We are going to mutate the dict + expr_to_proxy: Dict[sympy.Expr, fx.Proxy] = {} + placeholders = set() + first_non_placeholder = None + for node in graph.nodes: + if node.op != "placeholder": + first_non_placeholder = node + break + else: + placeholders.add(node) + + def _is_intermediate_tensor_sym_call(node: fx.Node) -> bool: + """ + If a size/stride/storage offset call on an intermediate tensor, + we can try to compute the value from input shapes instead. + """ + return ( + (val := _get_sym_val(node)) is not None + and not isinstance(val, sympy.Number) + # this holds back from reifying anything in torch.utils._sympy.functions.py that's unsupported + and not _has_uninterpretable_sympy_function(val) + and any( + isinstance(arg, fx.Node) + and isinstance(_get_example_value(arg), (torch.Tensor, torch.Size)) + and arg.op != "placeholder" + for arg in node.args + ) + ) + + # Figure out what key to use, val or example_value + val_key = "val" + for node in graph.nodes: + if "example_value" in node.meta: + val_key = "example_value" + break + elif "val" in node.meta: + break + + def _node_metadata_hook( + node: torch.fx.Node, + stack_trace: Optional[str] = None, + nn_module_stack: Optional[Dict[str, Any]] = None, + ) -> None: + fake_args = pytree.tree_map( + lambda arg: ( + _get_example_value(arg) if isinstance(arg, torch.fx.Node) else arg + ), + node.args, + ) + try: + node.meta[val_key] = node.target(*fake_args) # type: ignore[operator] + except NotImplementedError: + # This can happen when attempting to reify a symbol with an unsupported call_function node, + # e.g. with NestedTensors + sym_size.int via match_symbol(). + # This seems to be fine, as the node gets CSE'd and deleted later in favor of a SymInt graph input. + pass + if stack_trace is not None: + node.meta["stack_trace"] = stack_trace + if nn_module_stack is not None: + node.meta["nn_module_stack"] = nn_module_stack + + # Track asserts/checks we've added + added_asserts: Set[sympy.Expr] = set() + constrained_unbacked_symbols: Set[sympy.Symbol] = set() + + Analysis = PythonReferenceAnalysis if export else OptimizedPythonReferenceAnalysis + + def _sympy_interp(expr_to_proxy, expr): + # sympy_interp() with hash consing + from sympy import Integer, Number, Symbol + from sympy.logic.boolalg import BooleanAtom + + from torch.utils._sympy.interp import _run_sympy_handler, sympy_interp + + # hash cons + if expr in expr_to_proxy: + return expr_to_proxy[expr] + # base cases, don't cache + if isinstance(expr, (Integer, Number, Symbol, BooleanAtom)): + return sympy_interp(Analysis, expr_to_proxy, expr) + + # hash cons on arguments, run expr handler + expr_to_proxy[expr] = _run_sympy_handler( + Analysis, + [_sympy_interp(expr_to_proxy, arg) for arg in expr.args], + expr, + ) + return expr_to_proxy[expr] + + def _is_bound_expr_for_symbol(expr: "sympy.Expr") -> bool: + # This is probably unnecessary, but since torch._check() calls for single-symbol bounds + # like u0 >= 0, 10 >= u0 accumulate range info in the ShapeEnv, we designate these calls as redundant + # and instead add 2 runtime asserts at the end of this pass, if the min/max bounds are non-trivial. + if len(expr.args) != 2 or expr.func not in (sympy.LessThan, sympy.GreaterThan): + return False + lhs, rhs = expr.args + return (isinstance(lhs, sympy.Symbol) and isinstance(rhs, sympy.Number)) or ( + isinstance(rhs, sympy.Symbol) and isinstance(lhs, sympy.Number) + ) + + def add_runtime_asserts(ras): + for ra in ras: + if ( + # redundant + ra.expr in added_asserts + # if we've already added a constrain_range call for this symbol, + # then single-symbol bound asserts like u0 >= 0, u0 <= 5 are redundant. + or ( + len(ra.expr.free_symbols) == 1 + and next(iter(ra.expr.free_symbols)) in constrained_unbacked_symbols + and _is_bound_expr_for_symbol(ra.expr) + ) + # don't try to reify sympy functions we can't turn into FX nodes + or _has_uninterpretable_sympy_function(ra.expr) + ): + continue + + log.debug("inserting runtime assert %s", ra.expr) + # Need to process ALL free symbols, not just unbacked ones + fvs = free_symbols(ra.expr) + missing = fvs - expr_to_proxy.keys() + if missing: + i1 = min(missing, key=str) + # TODO: Remove relaxing assert on unbacked_symint https://github.com/pytorch/pytorch/issues/119689 + # assert shape_env.is_unbacked_symint(i1), i1 + ras_by_symbol.setdefault(i1, []).append(ra) + else: + # Convert the sympy expression into a sequence of FX + # nodes + with _set_node_metadata_hook(gm, _node_metadata_hook): + res = _sympy_interp(expr_to_proxy, ra.expr).node + graph.call_function( + torch.ops.aten._assert_scalar.default, + # TODO: use ra.msg here, but it's pretty + # useless right now + ( + res, + f"Runtime assertion failed for expression {ra.expr} on node '{res}'", + ), + ) + added_asserts.add(ra.expr) + + nodes = list(graph.nodes) + for i, node in enumerate(nodes[:-1]): + # Placeholders can match symbols, but when we destructure them + # with size we have to make sure we insert the nodes after all + # the placeholders + with graph.inserting_before( + nodes[i + 1] if node not in placeholders else first_non_placeholder + ): + # Unfortunately, this logic still must remain because manual + # make_fx calls may not explicitly bind all symbolic ints as + # arguments to the function, so we must infer it from the other + # arguments + if ( + node in placeholders + and (example_value := _get_example_value(node)) is not None + ): + + def match_symbol(symint, cb): + if ( + isinstance(symint, torch.SymInt) + and isinstance(symint.node, SymNode) + and isinstance(s := symint.node.expr, sympy.Symbol) + and s not in expr_to_proxy + ): + with _set_node_metadata_hook(gm, _node_metadata_hook): + expr_to_proxy[s] = fx.Proxy(cb(), tracer=tracer) + log.debug("expr_to_proxy[%s] = %s", s, expr_to_proxy[s]) + + match_symbol(example_value, lambda: node) + if isinstance(t := example_value, torch.Tensor): + for i, s in enumerate(t.size()): + match_symbol( + s, + lambda: graph.call_function( + torch.ops.aten.sym_size.int, (node, i) + ), + ) + if not is_sparse_any(t): + for i, s in enumerate(t.stride()): + match_symbol( + s, + lambda: graph.call_function( + torch.ops.aten.sym_stride.int, (node, i) + ), + ) + match_symbol( + t.storage_offset(), + lambda: graph.call_function( + torch.ops.aten.sym_storage_offset.default, (node,) + ), + ) + + # Handle asserts that aren't associated with any symbol. This + # doesn't really have to be in the loop as it will only run once, + # it just needs to happen right after the placeholders. + # insert this after placeholders & added sym nodes, and before non-placeholders. + if node == first_non_placeholder: + add_runtime_asserts(ras_by_symbol.pop(None, [])) # type: ignore[call-overload] + + # deduplicate asserts already present in graph + if node.target in ( + torch._check, + torch.ops.aten._assert_scalar.default, + ): + if ( + node.args[0] == True # noqa: E712 + or (assert_expr := _get_sym_val(node.args[0])) in expr_to_proxy + or ( + assert_expr is not None + and _is_bound_expr_for_symbol(assert_expr) + ) + ): + arg = node.args[0] + gm.graph.erase_node(node) + if isinstance(arg, fx.Node) and not arg.users: + gm.graph.erase_node(arg) + else: + added_asserts.add(assert_expr) # type: ignore[arg-type] + + # hash cons, replace function calls that return torch.SymInts with direct references to + # FX nodes built up to reify the sympy expression. + if ( + node.op != "placeholder" + and (sym_expr := _get_sym_val(node)) is not None + ): + # this guards against deleting calls like item() that produce new untracked symbols + def has_new_untracked_symbols(): + for symbol in sym_expr.free_symbols: + if symbol not in expr_to_proxy: + return True + return False + + # this guards against deleting calls that produce unbacked bindings we haven't yet seen. + # in this case looking at sym_expr.free_symbols might not be enough, if the example value has a hint + # (is backed), but produces an unbacked symbol. In this case keep the node alive. + resolved_unbacked_bindings = resolve_unbacked_bindings( + shape_env, node.meta.get("unbacked_bindings", {}) + ) + + assert resolved_unbacked_bindings is not None + + def has_new_unbacked_bindings(): + for key in resolved_unbacked_bindings.keys(): + if key not in expr_to_proxy: + return True + return False + + # maybe re-reify expression, replace current node + if ( + sym_expr in expr_to_proxy + or ( # example value is redundant + _is_intermediate_tensor_sym_call(node) + # shape call on intermediate tensor, turn into computation on input shapes + and not has_new_untracked_symbols() + ) + ) and not has_new_unbacked_bindings(): + if _is_intermediate_tensor_sym_call( + node + ): # reify from input shapes + with _set_node_metadata_hook( + gm, + functools.partial( + _node_metadata_hook, + stack_trace=node.meta.get("stack_trace"), + nn_module_stack=node.meta.get("nn_module_stack"), + ), + ): + expr_to_proxy[sym_expr] = _sympy_interp(expr_to_proxy, sym_expr) # type: ignore[arg-type] + # won't try DCE-ing tensor compute here + hash_node = expr_to_proxy[sym_expr].node # type: ignore[arg-type] + node.replace_all_uses_with(hash_node) + gm.graph.erase_node(node) + log.debug( + "CSE node %s -> %s for expr %s", node, hash_node, sym_expr + ) + + # store node in hash cons, don't delete/replace + elif sym_expr not in expr_to_proxy and not isinstance( + sym_expr, (sympy.Number, sympy.logic.boolalg.BooleanAtom) + ): # don't hash cons primitives + expr_to_proxy[sym_expr] = fx.Proxy(node, tracer=tracer) # type: ignore[arg-type] + + # We add sym_constrain_range calls for symbols later in any case if they're size-like or range-constrained, + # so calls before that are redundant. + if node.target in ( + torch.ops.aten.sym_constrain_range.default, + torch.ops.aten.sym_constrain_range_for_size.default, + ): + gm.graph.erase_node(node) + + defs = [] + + # AOTAutograd will create new symbols as the unbacked_bindings keys, which PropagateSymInts will set as + # equivalent, but the refinement calls we perform in this pass may struggle with associating the two. + # More concretely, when re-exporting/tracing, constraining only the new symbol may not communicate enough + # information about the old symbol when we re-export, raising errors on data-dependent guards. + # Call resolve_unbacked_bindings() to get the original symbol if present, otherwise we take it as is. + if unbacked_bindings := resolve_unbacked_bindings( + shape_env, node.meta.get("unbacked_bindings") + ): + for s, keypath in unbacked_bindings.items(): + defs.append(s) + + # TODO: some CSE when generating these nodes can probably + # help reduce graph size and improve compile time + def go(node, keypath): + if keypath == (): + return node + if ( + len(keypath) >= 2 + and isinstance(keypath[0], CallMethodKey) + and isinstance(keypath[1], pytree.SequenceKey) + ): + if keypath[0].name == "size": + return go( + graph.call_function( + torch.ops.aten.sym_size.int, + (node, keypath[1].idx), + ), + keypath[2:], + ) + if keypath[0].name == "stride": + return go( + graph.call_function( + torch.ops.aten.sym_stride.int, + (node, keypath[1].idx), + ), + keypath[2:], + ) + return go( + graph.call_method( + keypath[0].name, (node, keypath[1].idx) + ), + keypath[2:], + ) + elif isinstance(keypath[0], CallMethodKey): + return go( + graph.call_method(keypath[0].name, (node,)), keypath[1:] + ) + elif isinstance(keypath[0], pytree.SequenceKey): + return go( + graph.call_function( + operator.getitem, (node, keypath[0].idx) + ), + keypath[1:], + ) + elif isinstance(keypath[0], ConvertIntKey): + return go( + graph.call_function( + cast_symbool_to_symint_guardless, (node,) + ), + keypath[1:], + ) + elif isinstance(keypath[0], DivideByKey): + # TODO: need to assert divisibility + return go( + graph.call_function( + operator.floordiv, (node, keypath[0].divisor) + ), + keypath[1:], + ) + elif isinstance(keypath[0], InnerTensorKey): + return go( + graph.call_function( + getattr, (node, keypath[0].inner_name) + ), + keypath[1:], + ) + else: + raise AssertionError(f"unrecognized keypath {keypath}") + + if s not in expr_to_proxy: + with _set_node_metadata_hook(gm, _node_metadata_hook): + expr_to_proxy[s] = fx.Proxy( + go(node, keypath), tracer=tracer + ) + log.debug("expr_to_proxy[%s] = %s", s, expr_to_proxy[s]) + + for i0 in defs: + ras = ras_by_symbol.pop(i0, []) + # Before we perform any asserts, first apply range + # refinement. This is important, because if we are going + # to retrace the graph (and we typically are if we send + # the graph to AOTAutograd), we need to make sure we apply + # range refinement (ala _check_is_size) first, BEFORE we + # run any of the asserts. Otherwise, we may decide to + # perform substitutions based on the asserts which we then + # can't back out, because value ranges can only be applied + # to asserts.) + # + # A perhaps better long term plan is to avoid this order + # dependence by making it possible to refine ranges on + # arbitrary expressions, not just symbols. But it is not + # so easy to make use of this information, see + # https://twitter.com/ezyang/status/1745801370299482492 + # We actually made an attempt at this in + # https://github.com/pytorch/pytorch/pull/119043 + # which didn't work. + # + # Another ideas for how to do this: + # - Have bound_sympy be the source of truth of the ranges of any expression + # - Cache intermediate results for every subexpression of bound_sympy + # - This cache should be possible to edit to refine ranges + # + # One issue with this proposal is that if + # we have a bound on 2x, we are not going to be able to + # apply it for 4x. Similarly, we may have bounds for an + # equivalent expression that we are not applying because + # it's not a perfect match (e.g. x < y vs y > x)". + # + # The first issue we already have it and it's impossible + # to solve in general, so any implementation on a best + # effort basis should do. + # + # The second issue is a preexisting one. It can be mitigated + # with a normalization algorithm. In general, it may also + # be on a best effort basis, but since our grammar is not + # terribly difficult, chances are we could even fully + # normalize SymPy expressions... who knows. + if i0 in constrained_unbacked_symbols: + continue # constrain symbol just once + + if i0 in shape_env.size_like: + if export: + graph.call_function( + torch.ops.aten.sym_constrain_range_for_size.default, + (expr_to_proxy[i0].node,), + ) + else: + graph.call_function( + torch._check_is_size, (expr_to_proxy[i0].node,) + ) + + vr = shape_env.var_to_range[i0] + if vr.is_int and vr.upper == sys.maxsize - 1: + # treat upper bound == sys.maxsize - 1 for int symbols as +oo + # to avoid redundant runtime assert + vr = ValueRanges(vr.lower, int_oo) + if not shape_env._default_unspecified_value_range().issubset(vr): + # The runtime range is constrained, so add a runtime + # assert and also explicitly refine the range + # (refinement should not be necessary once runtime + # asserts cause refinement, but that's NYI) + def convert(s): + if s in (int_oo, -int_oo): + return None + try: + return int(s) + except TypeError: + return None + + if ( + expr_to_proxy[i0].node.target + != cast_symbool_to_symint_guardless + ): + # TODO(pianpwk): calling sym_constrain_range_for_size or adding bound asserts + # raises AOTAutograd errors on cast_symbool_to_symint_guardless + + with _set_node_metadata_hook( + gm, + functools.partial( + _node_metadata_hook, + stack_trace=node.meta.get("stack_trace"), + nn_module_stack=node.meta.get("nn_module_stack"), + ), + ): + if (min_val := convert(vr.lower)) is not None: + ge = _sympy_interp(expr_to_proxy, i0 >= min_val).node + graph.call_function( + torch.ops.aten._assert_scalar.default, + ( + ge, + f"Runtime assertion failed for expression {i0 >= min_val} on node '{ge}'", + ), + ) + added_asserts.add(i0 >= min_val) + if (max_val := convert(vr.upper)) is not None: + le = _sympy_interp(expr_to_proxy, i0 <= max_val).node + graph.call_function( + torch.ops.aten._assert_scalar.default, + ( + le, + f"Runtime assertion failed for expression {i0 <= max_val} on node '{le}'", + ), + ) + added_asserts.add(i0 <= max_val) + + constrained_unbacked_symbols.add(i0) + add_runtime_asserts(ras) + + # delete unused reified symbols + for expr, proxy in expr_to_proxy.items(): + if ( + isinstance(expr, sympy.Symbol) + and proxy.node.op != "placeholder" # keep placeholders intact + and not proxy.node.users + ): + log.debug("deleting unused reified symbol for %s", expr) + gm.graph.erase_node(proxy.node) diff --git a/lib/python3.12/site-packages/torch/fx/passes/shape_prop.py b/lib/python3.12/site-packages/torch/fx/passes/shape_prop.py new file mode 100644 index 0000000000000000000000000000000000000000..341bfd5e409183cc169b18833c0f263282dfcf92 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/shape_prop.py @@ -0,0 +1,211 @@ +# mypy: ignore-errors + +import traceback +from typing import Any, Dict, NamedTuple, Optional, Tuple + +import torch +import torch.fx +from torch._dispatch.python import enable_python_dispatcher +from torch._guards import detect_fake_mode +from torch._subclasses.meta_utils import is_sparse_any +from torch.fx._compatibility import compatibility +from torch.fx.node import map_aggregate, Node + + +__all__ = ["TensorMetadata", "ShapeProp"] + + +@compatibility(is_backward_compatible=True) +class TensorMetadata(NamedTuple): + # TensorMetadata is a structure containing pertinent information + # about a tensor within a PyTorch program. + + # General Tensor metadata + shape: torch.Size + dtype: torch.dtype + requires_grad: bool + stride: Tuple[int, ...] + memory_format: Optional[torch.memory_format] + + # Quantization metadata + is_quantized: bool + qparams: Dict[str, Any] + + +def _extract_tensor_metadata( + result: torch.Tensor, include_contiguity=True +) -> TensorMetadata: + """ + Extract a TensorMetadata NamedTuple describing `result`. + """ + shape = result.shape + dtype = result.dtype + requires_grad = result.requires_grad + stride = result.stride() if not is_sparse_any(result) else () + + memory_format = None + + if include_contiguity and not is_sparse_any(result): + memory_formats = { + torch.contiguous_format, + torch.channels_last, + torch.channels_last_3d, + } + for query_format in memory_formats: + if result.is_contiguous(memory_format=query_format): + memory_format = query_format + break + + is_quantized = result.is_quantized + qparams: Dict[str, Any] = {} + if is_quantized: + qscheme = result.qscheme() + qparams["qscheme"] = qscheme + if qscheme in {torch.per_tensor_affine, torch.per_tensor_symmetric}: + qparams["scale"] = result.q_scale() # type: ignore[assignment] + qparams["zero_point"] = result.q_zero_point() # type: ignore[assignment] + elif qscheme in { + torch.per_channel_affine, + torch.per_channel_affine_float_qparams, + torch.per_channel_symmetric, + }: + # In this branch, scale and zero_point are expected to be tensors, + # we store the values as immutable_list in TensorMetadata for + # easier serialization downstream + qparams["scale"] = result.q_per_channel_scales().tolist() # type: ignore[assignment] + qparams["zero_point"] = result.q_per_channel_zero_points().tolist() # type: ignore[assignment] + qparams["axis"] = result.q_per_channel_axis() # type: ignore[assignment] + + return TensorMetadata( + shape, dtype, requires_grad, stride, memory_format, is_quantized, qparams + ) + + +@compatibility(is_backward_compatible=True) +class ShapeProp(torch.fx.Interpreter): + """ + Execute an FX graph Node-by-Node and + record the shape and type of the result + into the corresponding node. + + Example: + In this example, we record the shape + and data type of a module given + an example input ``torch.randn(50, D_in)``. + We print the name, shape and dtype of each node. + + class TwoLayerNet(torch.nn.Module): + def __init__(self, D_in, H, D_out): + super().__init__() + self.linear1 = torch.nn.Linear(D_in, H) + self.linear2 = torch.nn.Linear(H, D_out) + def forward(self, x): + h_relu = self.linear1(x).clamp(min=0) + y_pred = self.linear2(h_relu) + return y_pred + N, D_in, H, D_out = 64, 1000, 100, 10 + x = torch.randn(N, D_in) + y = torch.randn(N, D_out) + model = TwoLayerNet(D_in, H, D_out) + gm = torch.fx.symbolic_trace(model) + sample_input = torch.randn(50, D_in) + ShapeProp(gm).propagate(sample_input) + + for node in gm.graph.nodes: + print(node.name, node.meta['tensor_meta'].dtype, + node.meta['tensor_meta'].shape) + + The output of this code is: + + x torch.float32 torch.Size([50, 1000]) + linear1 torch.float32 torch.Size([50, 100]) + clamp_1 torch.float32 torch.Size([50, 100]) + linear2 torch.float32 torch.Size([50, 10]) + output torch.float32 torch.Size([50, 10]) + + Args: + module (GraphModule): The module to be executed + fake_mode (FakeTensorMode): A fake mode for copying the gm + + """ + + def __init__(self, gm, fake_mode=None): + super().__init__(gm) + if fake_mode is None: + fake_mode = detect_fake_mode() + if fake_mode is not None: + from torch._dynamo.utils import deepcopy_to_fake_tensor + + # Note: + # We need fake execution cause the inputs are fake, however, we cannot fakify the module + # - because we need to write to the tensor_meta of the real module. So we fakify to + # produce a result (L131 below), to extract tensor meta, and then keep going. + # + # If we were to fakify, we would write to the wrong node, and then downstream fusion + # would be missing the tensor_meta. + # + # See torch/_inductor/overrides.py for where this is called upstream of fusion. + self.fake_module = deepcopy_to_fake_tensor(self.module, fake_mode) + self.fake_mode = fake_mode + else: + self.fake_module = None + self.fake_mode = None + + self.real_module = self.module + + def run_node(self, n: Node) -> Any: + try: + if self.fake_module is not None: + # Hacky swap. Alternatively, we could do this with overriding + # call_module and get_attr. + self.module = self.fake_module + try: + if self.fake_mode is not None: + with self.fake_mode, enable_python_dispatcher(): + result = super().run_node(n) + else: + result = super().run_node(n) + finally: + self.module = self.real_module + except Exception as e: + traceback.print_exc() + raise RuntimeError( + f"ShapeProp error for: node={n.format_node()} with " f"meta={n.meta}" + ) from e + + found_tensor = False + + def extract_tensor_meta(obj): + if isinstance(obj, torch.Tensor): + nonlocal found_tensor + found_tensor = True + return _extract_tensor_metadata(obj) + else: + return obj + + meta = map_aggregate(result, extract_tensor_meta) + if found_tensor: + n.meta["tensor_meta"] = meta + + n.meta["type"] = type(result) + return result + + def propagate(self, *args): + """ + Run `module` via interpretation and return the result and + record the shape and type of each node. + + Args: + *args (Tensor): the sample input. + + Returns: + Any: The value returned from executing the Module + """ + if self.fake_mode is not None: + fake_args = [ + self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t + for t in args + ] + else: + fake_args = args + return super().run(*fake_args) diff --git a/lib/python3.12/site-packages/torch/fx/passes/split_module.py b/lib/python3.12/site-packages/torch/fx/passes/split_module.py new file mode 100644 index 0000000000000000000000000000000000000000..7fec3089c5270179d9a1937b6cfdb8ed20c99965 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/split_module.py @@ -0,0 +1,626 @@ +# mypy: allow-untyped-defs +import inspect +import logging +from collections import OrderedDict +from typing import Any, Callable, Dict, List, Optional, Set + +import torch +from torch.fx._compatibility import compatibility +from torch.fx._utils import lazy_format_graph_code +from torch.fx.graph_module import GraphModule +from torch.fx.node import Node + + +__all__ = ["Partition", "split_module"] +log = _LOGGER = logging.getLogger(__name__) + + +@compatibility(is_backward_compatible=True) +class Partition: + def __init__(self, name: str): + self.name: str = name + self.submod_name = f"submod_{name}" + self.node_names: List[str] = [] + self.inputs: Dict[str, None] = {} + self.outputs: Dict[str, None] = {} + self.dependencies: Dict[str, None] = {} + self.dependents: Dict[str, None] = {} + self.graph: torch.fx.graph.Graph = torch.fx.graph.Graph() + self.environment: Dict[Node, Node] = {} + self.targets: Dict[str, Any] = {} + + def __repr__(self) -> str: + return ( + f"name: {self.name},\n" + f" nodes: {self.node_names},\n" + f" inputs: {self.inputs},\n" + f" outputs: {self.outputs},\n" + f" partitions depended on: {self.dependencies},\n" + f" partition dependents: {self.dependents}" + ) + + +def _get_attr_from_qualname(mod: torch.nn.Module, qualname: str) -> Any: + attr_val = mod + for atom in qualname.split("."): # type: ignore[union-attr] + if not hasattr(attr_val, atom): + raise AttributeError(f"Node target {qualname} not found!") + attr_val = getattr(attr_val, atom) + return attr_val + + +# Creates subgraphs out of main graph +@compatibility(is_backward_compatible=True) +def split_module( + m: GraphModule, + root_m: torch.nn.Module, + split_callback: Callable[[Node], int], + qualname_map: Optional[Dict[str, str]] = None, + keep_original_order: Optional[bool] = False, + keep_original_node_name: Optional[bool] = False, +): + """ + Creates subgraphs out of main graph + + Args: + m (GraphModule): Graph module to split + root_m (torch.nn.Module): root nn module. Not currently used. Included + because the root nn module is usually transformed via + torch.fx._symbolic_trace.symbolic_trace (see example below) + split_callback (Callable[[Node], int]): Callable function + that maps a given Node instance to a numeric partition identifier. + split_module will use this function as the policy for which operations + appear in which partitions in the output Module. + qualname_map: Optional[Dict[str, str]]: optional output parameter that returns a + mapping from new target names in the module after split to old target + names in the original module. + keep_original_order: Optional[bool]: keep the original order of the GraphModule + or use the Topological order of the new constructed GraphModule + + + Returns: + GraphModule: the module after split. + + Example: + + This is a sample setup: + + import torch + from torch.fx.symbolic_trace import symbolic_trace + from torch.fx.graph_module import GraphModule + from torch.fx.node import Node + from torch.fx.passes.split_module import split_module + + class MyModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.param = torch.nn.Parameter(torch.rand(3, 4)) + self.linear = torch.nn.Linear(4, 5) + + def forward(self, x, y): + z = self.linear(x + self.param).clamp(min=0.0, max=1.0) + w = self.linear(y).clamp(min=0.0, max=1.0) + return z + w + + # symbolically trace model + my_module = MyModule() + my_module_traced = symbolic_trace(my_module) + + # random mod partitioning + partition_counter = 0 + NPARTITIONS = 3 + + def mod_partition(node: Node): + global partition_counter + partition = partition_counter % NPARTITIONS + partition_counter = (partition_counter + 1) % NPARTITIONS + return partition + + # split module in module with submodules + module_with_submodules = split_module( + my_module_traced, my_module, mod_partition + ) + + Output looks like this. Original graph is broken into partitions + + > print(module_with_submodules) + GraphModule( + (submod_0): GraphModule( + (linear): Linear(in_features=4, out_features=5, bias=True) + ) + (submod_1): GraphModule( + (linear): Linear(in_features=4, out_features=5, bias=True) + ) + (submod_2): GraphModule() + ) + + def forward(self, x, y): + param = self.param + submod_0 = self.submod_0(x, param, y); x = param = y = None + getitem = submod_0[0] + getitem_1 = submod_0[1]; submod_0 = None + submod_1 = self.submod_1(getitem, getitem_1); getitem = getitem_1 = None + getitem_2 = submod_1[0] + getitem_3 = submod_1[1]; submod_1 = None + submod_2 = self.submod_2(getitem_2, getitem_3); getitem_2 = getitem_3 = None + return submod_2 + + Output of split module is the same as output of input traced module. + This is an example within a test setting: + + > orig_out = my_module_traced(x, y) + > submodules_out = module_with_submodules(x, y) + > self.assertEqual(orig_out, submodules_out) + True + """ + + log.debug( + "%s", + lazy_format_graph_code("pre split_module", m, colored=True), + ) + + def construct_graph( + node: Node, + base_mod_env: Dict[str, Node], + base_mod_attrs: Dict[str, torch.fx.graph_module.GraphModule], + ): + if node.op == "placeholder": + default_value = ( + node.args[0] if len(node.args) > 0 else inspect.Signature.empty + ) + if keep_original_node_name: + args = ( + () if default_value is inspect.Signature.empty else (default_value,) + ) + base_mod_env[node.name] = base_mod_graph.create_node( + "placeholder", + node.name, + args=args, # type: ignore[arg-type] + type_expr=node.type, + ) + else: + base_mod_env[node.name] = base_mod_graph.placeholder( + node.target, # type: ignore[arg-type] + type_expr=node.type, + default_value=default_value, + ) + base_mod_env[node.name].meta = node.meta.copy() + elif node.op == "get_attr": + base_mod_env[node.name] = base_mod_graph.get_attr(node.target) # type: ignore[arg-type] + base_mod_env[node.name].meta = node.meta.copy() + assert isinstance(node.target, str) + attr_val = _get_attr_from_qualname(m, node.target) + base_mod_attrs[node.target] = attr_val # type: ignore[index] + return base_mod_env, base_mod_attrs + + import sympy + + partitions: Dict[str, Partition] = {} + orig_nodes: Dict[str, Node] = {} + symbol_to_node: Dict[sympy.Symbol, Node] = {} + + def record_cross_partition_use(def_node: Node, use_node: Optional[Node]): + from torch.fx.experimental.symbolic_shapes import free_symbols + + defined = getattr(def_node, "_fx_partition", None) + used = getattr(use_node, "_fx_partition", None) + + log.debug( + "record_cross_partition_use %s (%s) %s (%s)", + def_node.name, + defined, + use_node.name if use_node is not None else "-", + used, + ) + + if defined != used: + if defined is not None: + def_partition = partitions[defined] + def_partition.outputs.setdefault(def_node.name) + if used is not None: + def_partition.dependents.setdefault(used) + + if used is not None: + use_partition = partitions[used] + use_partition.inputs.setdefault(def_node.name) + # We have made def_node an input to the use_partition. If + # this input has symbolic symbols in its size, those also must + # be made as inputs to the partition + if (def_val := def_node.meta.get("example_value")) is not None: + for s in sorted(free_symbols(def_val), key=str): + s_node = symbol_to_node[s] + use_partition.inputs.setdefault(s_node.name) + if symbol_to_node[s].op != "placeholder": + # If the node that defines the symbol is not a + # placeholder, we must make it an output of the + # partition. Note that this may be in a different + # partition than defined! Although, this doesn't + # really make a difference for correctness, since + # defined is guaranteed to have the symbol in + # scope and can return it; you just get less + # optimal codegen in this case. + s_defined = getattr(s_node, "_fx_partition", None) + if s_defined is not None: + s_def_partition = partitions[s_defined] + s_def_partition.outputs.setdefault(s_node.name) + s_def_partition.dependents.setdefault(used) + if defined is not None: + use_partition.dependencies.setdefault(defined) + + def instantiate_node_partition_mapping(node): + partition_name = str(split_callback(node)) + log.debug( + "instantiate_node_partition_mapping %s (%s)", node.name, partition_name + ) + + # add node to partitions + partition = partitions.get(partition_name) + if partition is None: + partitions[partition_name] = partition = Partition(partition_name) + + partition.node_names.append(node.name) + node._fx_partition = partition_name + + # Global State Nodes are nodes which by their global state effects, + # "taint" all downstream nodes while they are active. + GLOBAL_STATE_NODES = [ + torch.amp._enter_autocast, + torch.amp._exit_autocast, + torch._C._set_grad_enabled, + ] + + # For grad regions: + # ------------------------ + # 1. first region: we do nothing + # 2. subsequent regions: we insert the set_grad at the beginning + grad_regions: OrderedDict[Node, Set[int]] = OrderedDict() + + # For autocast regions: + # ------------------------ + # 1. first region: we will only insert the _exit at the end + # 2. intermediate regions: we will insert both the + # _enter at the beginning and _exit at the end + # 3. last region: we will only insert _enter at the beginning + # We will do so in the order in which the autocasts were instantiated. + autocast_regions: OrderedDict[Node, Set[int]] = OrderedDict() + autocast_exits: Dict[Node, Optional[Node]] = {} + + active_grad = None + active_autocasts = set() + + for node in m.graph.nodes: + # This will prefer placeholder bindings, because those come first. + # This is a little dangerous though: it is possible that an unbacked + # symbol is used without any binding site for it, in which case we + # will get a KeyError not able to find it. I'd like to fix this by + # having passes.runtime_assert establish some invariants that I can + # rely on later, but this needs some extra work. Quick fix first. + # See https://github.com/pytorch/pytorch/issues/130534 + if ( + (val := node.meta.get("example_value")) is not None + and isinstance(val, (torch.SymInt, torch.SymFloat)) + and isinstance(s0 := val.node.expr, sympy.Symbol) + and s0 not in symbol_to_node + ): + symbol_to_node[val.node.expr] = node + + if node.op in ["placeholder", "get_attr", "output"]: + continue + + instantiate_node_partition_mapping(node) + + if node.op == "call_function" and node.target in GLOBAL_STATE_NODES: + if node.target == torch._C._set_grad_enabled: + assert len(node.args) == 1 + assert isinstance(node.args[0], bool) + active_grad = node + grad_regions[active_grad] = set({split_callback(node)}) + elif node.target == torch.amp._enter_autocast: + # Should all be python constants + assert all(not isinstance(arg, Node) for arg in node.args) + active_autocasts.add(node) + autocast_regions[node] = set({split_callback(node)}) + autocast_exits[node] = None + elif node.target == torch.amp._exit_autocast: + assert len(node.args) == 1 + autocast_regions[node.args[0]].add(split_callback(node)) + active_autocasts.remove(node.args[0]) + autocast_exits[node.args[0]] = node + + if active_grad is not None: + grad_regions[active_grad].add(split_callback(node)) + + for a in active_autocasts: + autocast_regions[a].add(split_callback(node)) + + assert all(v is not None for v in autocast_exits.values()), "autocast must exit" + + autocast_regions = {k: sorted(v) for k, v in autocast_regions.items()} + grad_regions = {k: sorted(v) for k, v in grad_regions.items()} + + if _LOGGER.isEnabledFor(logging.DEBUG): + _LOGGER.debug("autocast_regions: %s", autocast_regions) + _LOGGER.debug("grad_regions: %s", grad_regions) + + assert_monotonically_increasing = bool(autocast_regions) or bool(grad_regions) + + # split nodes into partitions + highest_partition = -1 + for node in m.graph.nodes: + orig_nodes[node.name] = node + + # TODO currently placeholders/parameters aren't put into random partitions, + # rather they're added to the graphs where they are used down below + if node.op in ["placeholder", "get_attr"]: + continue + if node.op == "output": + torch.fx.graph.map_arg( + node.args[0], lambda n: record_cross_partition_use(n, None) + ) + continue + + if assert_monotonically_increasing: + pid = split_callback(node) + assert highest_partition <= pid, ( + "autocast or set_grad_enabled require monotonically increasing partitions:" + f"highest: {highest_partition}, this node's: {pid}" + ) + highest_partition = pid + + # do not capture cross-partition dependencies for global state nodes as they will be + # self-contained - their setup and unwind will be isolated to each partition submodule. + if node.target not in GLOBAL_STATE_NODES: + torch.fx.graph.map_arg( + node.args, lambda def_node: record_cross_partition_use(def_node, node) + ) + torch.fx.graph.map_arg( + node.kwargs, lambda def_node: record_cross_partition_use(def_node, node) + ) # noqa: B950 + + original_partition_order = list(partitions.keys()) + # find partitions with no dependencies + root_partitions: List[str] = [] + for partition_name, partition in partitions.items(): + if not len(partition.dependencies): + root_partitions.append(partition_name) + + # check partitions for circular dependencies and create topological partition ordering + sorted_partitions: List[str] = [] + while root_partitions: + root_partition = root_partitions.pop() + sorted_partitions.append(root_partition) + for dependent in partitions[root_partition].dependents: + partitions[dependent].dependencies.pop(root_partition) + if not partitions[dependent].dependencies: + root_partitions.append(dependent) + if len(sorted_partitions) != len(partitions): + raise RuntimeError("cycle exists between partitions!") + + # Enter prelude + for regions_mapping in [autocast_regions, grad_regions]: + for node, regions in regions_mapping.items(): + assert len(regions) > 0 + partitions[str(regions[0])].environment[node] = node + for r in regions[1:]: + partition = partitions[str(r)] + new_node = partition.graph.create_node( + op=node.op, + target=node.target, + args=tuple(arg for arg in node.args), + kwargs={}, + type_expr=node.type, + ) + new_node.meta = ( + node.meta.copy() + ) # is it really a good idea to copy this? + partition.environment[node] = new_node + + # add placeholders to partition inputs + for partition_name in sorted_partitions: + partition = partitions[partition_name] + new_inputs: Dict[str, None] = {} + for inp in partition.inputs: + orig_node = orig_nodes[inp] + # We don't pass in get_attr nodes as inputs to the partition, but + # instead set them as targets and use getattr within the module + + if orig_node.op == "get_attr": + assert isinstance(orig_node.target, str) + + orig_attr = _get_attr_from_qualname(m, orig_node.target) + if isinstance(orig_attr, torch.nn.Module): + placeholder = partition.graph.get_attr(orig_node.target) + partition.targets[orig_node.target] = orig_attr + else: + placeholder = partition.graph.placeholder( + inp, + type_expr=orig_nodes[inp].type, + ) + new_inputs[inp] = None + else: + placeholder = partition.graph.placeholder( + inp, + type_expr=orig_nodes[inp].type, + ) + new_inputs[inp] = None + placeholder.meta = orig_nodes[inp].meta.copy() + partition.environment[orig_nodes[inp]] = placeholder + partition.inputs = new_inputs + + # Transform nodes and collect targets for partition's submodule + for node in m.graph.nodes: + if hasattr(node, "_fx_partition"): + partition = partitions[node._fx_partition] + + # swap out old graph nodes in kw/args with references to new nodes in this submodule + environment = partition.environment + gathered_args = torch.fx.graph.map_arg(node.args, lambda n: environment[n]) + gathered_kwargs = torch.fx.graph.map_arg( + node.kwargs, lambda n: environment[n] + ) + + if node.op not in ["call_module", "get_attr"]: + target = node.target + else: + target_attr = _get_attr_from_qualname(m, node.target) + target = node.target.replace(".", "_") + partition.targets[target] = target_attr + # Fill in the passed-in mapping from new qualname to old qualname + if qualname_map is not None: + # When creating the split module later, the submodules will have + # path prefix matching the corresponding partition's submod_name + qualname = f"{partition.submod_name}.{target}" + qualname_map[qualname] = node.target + + assert isinstance(gathered_args, tuple) + assert isinstance(gathered_kwargs, dict) + name = node.name if keep_original_node_name else None + new_node = partition.graph.create_node( + op=node.op, + target=target, + args=gathered_args, + kwargs=gathered_kwargs, + type_expr=node.type, + name=name, + ) + new_node.meta = node.meta.copy() + partition.environment[node] = new_node + + # Exit epilogue + for regions_mapping in [autocast_regions]: + for node in reversed(regions_mapping): + regions = regions_mapping[node] + assert len(regions) > 0 + for r in regions[:-1]: + partition = partitions[str(r)] + exit_node = autocast_exits[node] + assert exit_node is not None, "Missing exit node" + new_node = partition.graph.create_node( + op=exit_node.op, + target=exit_node.target, + args=(partition.environment[node],), + kwargs={}, + type_expr=exit_node.type, + ) + new_node.meta = ( + exit_node.meta.copy() + ) # is it really a good idea to copy this? + + # original module environment dict mapping node names to nodes + orig_mod_env: Dict[str, Node] = {} + # Set up values to construct base module + base_mod_env: Dict[str, Node] = {} + base_mod_graph: torch.fx.graph.Graph = torch.fx.graph.Graph() + base_mod_attrs: Dict[str, torch.fx.graph_module.GraphModule] = {} + if not keep_original_order: + for node in m.graph.nodes: + base_mod_env, base_mod_attrs = construct_graph( + node, base_mod_env, base_mod_attrs + ) + + else: + # Go through the graph to construct the mapping dict + for node in m.graph.nodes: + orig_mod_env[node.name] = node + + # Do some things iterating over the partitions in topological order again: + # 1) Finish off submodule Graphs by setting corresponding outputs + # 2) Construct GraphModules for each submodule + # 3) Construct the base graph by emitting calls to those submodules in + # topological order or original order specified by keep_original_order + + construct_order_partitions = ( + sorted_partitions if not keep_original_order else original_partition_order + ) + + already_constructed_attr_nodes = set() + + # We actually need to insert the placeholder nodes in the original order + # otherwise graph signature will be wrong. + original_order = [node for node in m.graph.nodes if node.op == "placeholder"] + + for partition_name in construct_order_partitions: + partition = partitions[partition_name] + + # Set correct output values + output_vals = tuple( + partition.environment[orig_nodes[name]] for name in partition.outputs + ) + + # skip output node generation if there are no output values + num_output_vals = len(output_vals) + if num_output_vals == 1: + partition.graph.output(output_vals[0]) + elif num_output_vals > 1: + partition.graph.output(output_vals) + else: + # Invariant - Graph should always have an output node. + partition.graph.output(()) + + if keep_original_order: + # first get the attr nodes required by this partition + orig_mod_attr_nodes: List[Node] = [ + orig_mod_env[key] + for key in partition.inputs + if key not in original_order + ] + + for node in original_order: + if node in already_constructed_attr_nodes: + continue # already added this attr to the base graph + base_mod_env, _based_mod_attrs = construct_graph( + node, base_mod_env, base_mod_attrs + ) + already_constructed_attr_nodes.add(node) + + # Construct GraphModule for this partition + for node in orig_mod_attr_nodes: # type: ignore[attr-defined] + if node in already_constructed_attr_nodes: + continue + base_mod_env, base_mod_attrs = construct_graph( + node, base_mod_env, base_mod_attrs + ) + already_constructed_attr_nodes.add(node) + + base_mod_attrs[partition.submod_name] = torch.fx.graph_module.GraphModule( + partition.targets, partition.graph + ) # noqa: B950 + + # Emit call in base graph to this submodule + output_val = base_mod_graph.call_module( + partition.submod_name, + tuple(base_mod_env[name] for name in partition.inputs), + ) + + num_outputs = len(partition.outputs) + if num_outputs > 1: + # Unpack multiple return values from submodule + output_val_proxy = torch.fx.proxy.Proxy(output_val) + for i, output_name in enumerate(partition.outputs): + base_mod_env[output_name] = output_val_proxy[i].node # type: ignore[index] + elif num_outputs == 1: + base_mod_env[next(iter(partition.outputs))] = output_val + + # When keep_original_order=True and if the graph doesn't have any + # `call_function` node then `base_mod_graph`, `base_mod_env` and `base_mod_attrs` + # are never populated. + # For this case, we call `construct_graph` here which takes care of updating them. + if keep_original_order and not base_mod_env: + for node in m.graph.nodes: + base_mod_env, base_mod_attrs = construct_graph( + node, base_mod_env, base_mod_attrs + ) + + # Add output node to `base_mod_graph` (i.e. the split graph) which will be returned. + for node in m.graph.nodes: + if node.op == "output": + base_mod_graph.output( + torch.fx.graph.map_arg(node.args[0], lambda n: base_mod_env[n.name]) + ) # noqa: B950 + + ret = torch.fx.graph_module.GraphModule(base_mod_attrs, base_mod_graph) + log.debug( + "%s", + lazy_format_graph_code("post split_module", ret, colored=True), + ) + return ret diff --git a/lib/python3.12/site-packages/torch/fx/passes/split_utils.py b/lib/python3.12/site-packages/torch/fx/passes/split_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e2bece6f72f279d674e8e181dac18031b10d9193 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/split_utils.py @@ -0,0 +1,304 @@ +# mypy: allow-untyped-defs +import copy +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Type, Union + +import torch.fx +from torch.fx._compatibility import compatibility +from torch.fx.graph import map_arg +from torch.fx.passes.utils import HolderModule, lift_subgraph_as_module + +from .tools_common import NodeList + + +__all__ = ["getattr_recursive", "setattr_recursive", "Component", "split_by_tags"] + + +@compatibility(is_backward_compatible=False) +def getattr_recursive(obj, name): + for layer in name.split("."): + if hasattr(obj, layer): + obj = getattr(obj, layer) + else: + return None + return obj + + +@compatibility(is_backward_compatible=False) +def setattr_recursive(obj, attr, value): + if "." not in attr: + setattr(obj, attr, value) + else: + layer = attr.split(".") + setattr_recursive(getattr(obj, layer[0]), ".".join(layer[1:]), value) + + +@compatibility(is_backward_compatible=False) +@dataclass +class Component: + """ + A component serves as a container for a subgraph we want to create afterwards. + """ + + graph: torch.fx.Graph + order: int + name: str + + # Stores the placeholder nodes in `graph`. + input_placeholders: List = field(default_factory=list) + + # Store the nodes in original graph that are placeholder in `graph`. + orig_inputs: List = field(default_factory=list) + + # Store the nodes in original graph that are outputs in `graph`. + orig_outputs: List = field(default_factory=list) + + # Mapping from get_attr node in original graph to get_attr node in `graph`. + getattr_maps: Dict[torch.fx.Node, torch.fx.Node] = field(default_factory=dict) + constructor_args: List[str] = field(default_factory=list) + gm: Optional[torch.fx.GraphModule] = None + + +@compatibility(is_backward_compatible=False) +def split_by_tags( + gm: torch.fx.GraphModule, + tags: List[str], + return_fqn_mapping: bool = False, + return_tuple: bool = False, + GraphModuleCls: Type[torch.fx.GraphModule] = torch.fx.GraphModule, +) -> Union[torch.fx.GraphModule, Tuple[torch.fx.GraphModule, Dict[str, str]]]: + """ + Splits a GraphModule using tags on its graph nodes. We honor the order of + tags. For example, we have tags = ["a", "b", "c"], the function will create + the initial submodules in the order of "a", "b", "c". + + To set a tag: + gm.graph.nodes[idx].tag = "mytag" + + This will result in all nodes with the same tag being extracted and placed in their + own submodule. For placeholder, output and get_attr node, the tag is ignored. placeholder + and output nodes are created when needed while get_attr nodes get copied to submodules + where they are used. + + Given the following module def: + + class SimpleModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = torch.nn.Linear(...) + self.linear2 = torch.nn.Linear(...) + self.linear3 = torch.nn.Linear(...) + + def forward(self, in1, in2): + r1 = self.linear1(in1) + r2 = self.linear2(in2) + r3 = torch.cat([r1, r2]) + return self.linear3(r3) + + Marking the node corresponding to in1 with the tag sc.REQUEST_ONLY.lower() results in the following split: + + ro: + def forward(self, in1): + self = self.root + linear1 = self.linear1(in1) + return linear1 + + main: + def forward(self, in2, linear1): + self = self.root + linear2 = self.linear2(in2) + cat_1 = torch.cat([linear1, linear2]) + linear3 = self.linear3(cat_1) + return linear3 + + main: + def forward(self, in1, in2): + self = self.root + ro_0 = self.ro_0(in1) + main_1 = self.main_1(in2, ro_0) + return main_1 + + Returns: + split_gm: torch fx graph after split + orig_to_split_fqn_mapping: a map between the original fqn and the fqn + after split for call_module and get_attr. + """ + + def flatten(x: torch.fx.node.Argument) -> NodeList: + """ + Stores nodes in x to a list and returns the list. + """ + r: NodeList = [] + map_arg(x, r.append) + return r + + # Mapping from node in original module to node in created submodule. + node_remapping: Dict[torch.fx.Node, torch.fx.Node] = {} + + # Mapping from node in original module or created submodules to + # corresponding component. + node_to_component: Dict[torch.fx.Node, Component] = {} + + # Mapping from tag to the corresponding component. + tag_to_component: Dict[str, Component] = {} + + # Stores all components. + all_components: List[Component] = [] + + # Stores nodes that will be used in main graph. + used_in_main: Dict[torch.fx.Node, None] = {} + + # Main graph after split. + main_g = torch.fx.Graph() + + # Mapping from node in original module to node in main graph after split. + main_remapping: Dict[torch.fx.Node, torch.fx.Node] = {} + + # Output node of original module. + output_node: Optional[torch.fx.Node] = None + + # Create a component for each tag, we don't expect to create other components afterwards. + for tag in tags: + comp = Component(torch.fx.Graph(), len(all_components), f"{tag}") + all_components.append(comp) + tag_to_component[tag] = comp + + # Traverse the nodes in original graph and take care of them. + for node in gm.graph.nodes: + if node.op == "output": + if output_node is not None: + raise RuntimeError("Multiple output nodes in graph!") + output_node = node + continue + + # Placeholders in the original graph get copied to main graph. + if node.op == "placeholder": + main_remapping[node] = main_g.placeholder(node.name, type_expr=node.type) + main_remapping[node].meta = copy.copy(node.meta) + continue + + # Get_attr nodes are ignored because we are not tagging them. + # Instead, we copy them directly to the submodules use them afterwards. + if node.op == "get_attr": + continue + + # Now we process callable nodes which are nodes with op of call_module, + # call_function or call_method. Every callable nodes should be tagged. + assert hasattr(node, "tag") + + upstream_components = [ + node_to_component[x] + for x in flatten(node.args) + flatten(node.kwargs) + if x.op not in {"placeholder", "get_attr"} + ] + + comp = tag_to_component[node.tag] + node_to_component[node] = comp + + # Max order of upperstream components. + mx = max((c.order for c in upstream_components), default=0) + + # Expect the component for `node` has higher order then its upstream components. + assert comp.order >= mx + + # Map a input of `node` to nodes in the component's graph. + def remap_func(x): + # If input is a get_attr node, copy it to current component's graph. + # Returns the get_attr node in current component's graph. + if x.op == "get_attr": + if x not in comp.getattr_maps: + comp.getattr_maps[x] = comp.graph.get_attr( + x.target, type_expr=x.type + ) + return comp.getattr_maps[x] + + # If input is not a placeholder, it should have been put into a component + # already. If it's the current component then we return the corresponding + # node in the component. + if x.op != "placeholder" and node_to_component[x] == comp: + return node_remapping[x] + + # If input is a placeholder or it's in other components, we want to make it + # as a placeholder in current component's graph. + if x not in comp.orig_inputs: + comp.orig_inputs.append(x) + placeholder = comp.graph.placeholder(x.name, type_expr=x.type) + placeholder.meta = copy.copy(x.meta) + comp.input_placeholders.append(placeholder) + used_in_main[x] = None + + return comp.input_placeholders[comp.orig_inputs.index(x)] + + n = comp.graph.node_copy(node, remap_func) + n.tag = node.tag # type: ignore[attr-defined] + node_remapping[node] = n + node_to_component[n] = comp + + if output_node is None: + raise RuntimeError("Graph had no output node!") + + for x in flatten(output_node.args[0]): + if x.op == "get_attr": + # We don't need components mapping for nodes of type "get_attr" + # that are consumed by the output. Only need to make sure we create + # corresponding counterparts in the resulting graph. + main_remapping[x] = main_g.get_attr(x.name, type_expr=x.type) + else: + # All component results consumed by the output node should be + # marked as "used in main". + used_in_main[x] = None + + # If a node is used in main graph then we mark it as an output in the component + # it belongs to. + for n in used_in_main: + if n.op != "placeholder": + node_to_component[n].orig_outputs.append(n) + + # Now we create a graphmodule for each component. + orig_to_split_fqn_mapping: Dict[str, str] = {} + for comp in all_components: + outs = tuple(map(node_remapping.__getitem__, comp.orig_outputs)) + + if return_tuple: + comp.graph.output(outs) + else: + # Take care of the args of FX output node. If there's a single + # output then the output node args is like (output_single), else + # if there're multiple outputs then the output node args is like + # ((output_0, output_1, ...)). + comp.graph.output(outs[0] if len(outs) == 1 else outs) + + comp.gm, comp_orig_to_split_fqn_mapping = lift_subgraph_as_module( + gm, subgraph=comp.graph, comp_name=comp.name + ) + orig_to_split_fqn_mapping.update(comp_orig_to_split_fqn_mapping) + + # Create a call_module node in main graph. + main_node = main_g.call_module( + comp.name, + args=tuple(map(main_remapping.__getitem__, comp.orig_inputs)), + kwargs=None, + ) + + if len(outs) == 1 and not return_tuple: + main_remapping[comp.orig_outputs[0]] = main_node + else: + for i, o in enumerate(comp.orig_outputs): + # Use Proxy to record getitem access. + main_remapping[o] = torch.fx.Proxy(main_node)[i].node # type: ignore[index] + + main_g.output(map_arg(output_node.args[0], main_remapping.__getitem__)) + main_root = HolderModule({comp.name: comp.gm for comp in all_components}) + main_g._codegen = gm.graph._codegen + + # If the output nodes consumes get_attr directly in the original graph, + # then we need to make sure get_attr is copied to the new graph. + for x in flatten(output_node.args[0]): + if x.op == "get_attr": + setattr(main_root, x.name, getattr_recursive(gm, x.target)) # type: ignore[arg-type] + + result_gm = GraphModuleCls(main_root, main_g) + if return_fqn_mapping: + return result_gm, orig_to_split_fqn_mapping + + return result_gm diff --git a/lib/python3.12/site-packages/torch/fx/passes/splitter_base.py b/lib/python3.12/site-packages/torch/fx/passes/splitter_base.py new file mode 100644 index 0000000000000000000000000000000000000000..880ebc68cd8377b65269399de0935185c9569d83 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/splitter_base.py @@ -0,0 +1,924 @@ +# mypy: allow-untyped-defs +import argparse +import copy +import logging +from collections import defaultdict +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple + +import torch +from torch.fx._compatibility import compatibility +from torch.fx.node import map_arg +from torch.fx.passes.graph_manipulation import get_size_of_node + +from .graph_drawer import FxGraphDrawer +from .operator_support import get_node_target, OperatorSupportBase +from .shape_prop import ShapeProp +from .split_utils import split_by_tags +from .tools_common import ( + CALLABLE_NODE_OPS, + FxNetAccFusionsFinder, + is_node_output_tensor, + NodeList, + NodeSet, + Tensors, +) + + +__all__ = [ + "FxNetAccNodesFinder", + "FxNetSplitterInternalError", + "Subgraph", + "SplitResult", + "generate_inputs_for_submodules", +] +_LOGGER = logging.getLogger(__name__) + +DEFAULT_MIN_ACC_MODULE_SIZE = 1 +DEFAULT_SKIP_FUSION = False +DEFAULT_ALLOW_NON_TENSOR = False + + +class _SplitterSettingBase: + def __init__( + self, + min_acc_module_size=DEFAULT_MIN_ACC_MODULE_SIZE, + skip_fusion=DEFAULT_SKIP_FUSION, + allow_non_tensor=DEFAULT_ALLOW_NON_TENSOR, + max_acc_splits: int = -1, + ): + parser = argparse.ArgumentParser() + parser.add_argument( + "--min-acc-module-size", + "--min_acc_module_size", + required=False, + type=int, + help="Minimum size limit of an accelerator subgraph.", + ) + parser.add_argument( + "--max-acc-splits", + "--max_acc_splits", + required=False, + type=int, + help="Enforce a maximum number of split subgraphs.", + ) + parser.add_argument( + "--skip-fusion", + "--skip_fusion", + default=False, + action="store_true", + help="If true then no fusion groups. Fusion group is used to " + "enforce no non-tensor data flow between submodules. If we don't " + "have this constrain, setting this to false is recommended as it " + "can reduce overhead.", + ) + parser.add_argument( + "--allow-non-tensor", + "--allow_non_tensor", + default=False, + action="store_true", + help="For some backends non-tensor data flow between cpu and them " + "are not allowed. Therefore, if a node supported by accelerator but " + "it has non-tensor inputs or outputs to a cpu node we would want to " + "consider it as a cpu node during splitting. However, for some backends " + "we might not care about non-tensor data flow and we can set this option " + "to true to disable the functionality that prevent non-tensor data flow.", + ) + args, _unknown = parser.parse_known_args() + + self.min_acc_module_size: int = ( + args.min_acc_module_size + if args.min_acc_module_size + else min_acc_module_size + ) + self.skip_fusion: bool = args.skip_fusion if args.skip_fusion else skip_fusion + self.allow_non_tensor: bool = ( + args.allow_non_tensor if args.allow_non_tensor else allow_non_tensor + ) + self.max_acc_splits: int = max_acc_splits + + +@compatibility(is_backward_compatible=False) +class FxNetAccNodesFinder: + """ + Finds a set of nodes that can be supported on ACC, excluding nodes that have non-tensor + input/output to cpu nodes to prevent non-tensor data flow between backends and cpu. + + I.e. if we have a chain: + + ACC_NODE_1 -> ACC_NODE_2 -> ACC_NODE_3 -> CPU_NODE_1 + + where every ACC node produces non-tensor output, then they all should be treated as CPU nodes. + + This behavior can be turned off by passing allow_non_tensor=True. + """ + + def __init__( + self, + module: torch.fx.GraphModule, + operator_support: OperatorSupportBase, + allow_non_tensor: bool, + ): + self.module = module + self.operator_support = operator_support + self.allow_non_tensor = allow_non_tensor + self.acc_nodes: NodeSet = set() + + def reduce_acc_nodes_non_tensor_input_helper(self, cpu_worklist: NodeList): + """ + Transitively excludes nodes from ACC supported set. + For every node in the worklist: + - removes its downstream ACC nodes from ACC supported set, + - if any downstream ACC node produces non-tensor output, + then it gets added into the worklist. + """ + while cpu_worklist: + node = cpu_worklist.pop(0) + + for user in node.users: + if user in self.acc_nodes: + self.acc_nodes.remove(user) + if not is_node_output_tensor(user): + cpu_worklist.append(user) + + def reduce_acc_nodes_non_tensor_input(self): + """ + Excludes nodes from ACC supported set that have direct + upstream CPU nodes that produce non-tensor outputs. + """ + non_tensor_cpu_nodes: NodeList = [] + + for node in self.module.graph.nodes: + if node.op not in CALLABLE_NODE_OPS: + continue + if node in self.acc_nodes: + continue + if is_node_output_tensor(node): + continue + non_tensor_cpu_nodes.append(node) + + self.reduce_acc_nodes_non_tensor_input_helper(non_tensor_cpu_nodes) + + def reduce_acc_nodes_non_tensor_output(self): + """ + Excludes nodes from ACC supported set that produce non-tensor + outputs and have downstream CPU nodes. + """ + while True: + new_cpu_nodes: NodeList = [] + + for acc_node in self.acc_nodes: + if is_node_output_tensor(acc_node): + continue + for user in acc_node.users: + if user not in self.acc_nodes: + new_cpu_nodes.append(acc_node) + break + + if not new_cpu_nodes: + break + + for new_cpu_node in new_cpu_nodes: + self.acc_nodes.remove(new_cpu_node) + + self.reduce_acc_nodes_non_tensor_input_helper(new_cpu_nodes) + + def __call__(self) -> NodeSet: + submodules = dict(self.module.named_modules()) + self.acc_nodes = { + n + for n in self.module.graph.nodes + if n.op in CALLABLE_NODE_OPS + and self.operator_support.is_node_supported(submodules, n) + } + + if not self.allow_non_tensor: + self.reduce_acc_nodes_non_tensor_input() + self.reduce_acc_nodes_non_tensor_output() + + return self.acc_nodes + + +@compatibility(is_backward_compatible=False) +class FxNetSplitterInternalError(Exception): + pass + + +@compatibility(is_backward_compatible=False) +@dataclass +class Subgraph: + is_acc: bool + nodes: NodeList + device_ordinal: Optional[int] = None + + +@compatibility(is_backward_compatible=False) +class SplitResult(NamedTuple): + """ + Stores the results of the splitter. + + Attributes: + split_module: root module after splitting. + submodule_inputs: a dict that maps submodule name to its inputs. + non_acc_submodule_prefix: the prefix for non acc submodules. For + acc submodule the prefix is alwasy "_run_on_acc_". + """ + + split_module: torch.fx.GraphModule + submodule_inputs: Dict[str, Any] + non_acc_submodule_prefix: str + + +@compatibility(is_backward_compatible=False) +def generate_inputs_for_submodules( + model: torch.nn.Module, + inputs: Sequence[Any], + target_submodules: Iterable[str], + deepcopy: bool = False, +) -> Dict[str, Any]: + """ + Generate inputs for targeting submdoules in the given model. Note that if two submodules refer to the same obj, this + function doesn't work. + + Args: + model: root model. + inputs: inputs to the root model. + target_submodules: submodules that we want to generate inputs for. + + Returns: + A dict that maps from submodule name to its inputs. + """ + + handles = [] + results = {} + submodule_to_names = {mod: name for name, mod in model.named_modules()} + + def pre_forward(module, module_inputs): + results[submodule_to_names[module]] = ( + copy.deepcopy(module_inputs) if deepcopy else module_inputs + ) + + for name, mod in model.named_modules(): + if name in target_submodules: + handles.append(mod.register_forward_pre_hook(pre_forward)) + + def clean_up_handles(): + for h in handles: + h.remove() + + try: + with torch.no_grad(): + model(*inputs) + except Exception as e: + clean_up_handles() + raise e + + clean_up_handles() + return results + + +class _SplitterBase: + """ + Splits a GraphModule into sub-GraphModules for execution on CPU or the accelerator. + Output is a GraphModule with supported and unsupported operators grouped into as few sub-GraphModules as possible. + Assumes that only "call_module", "call_function" and "call_method" from FX IR can potentially be executed on the accelerator. + + Given the following graph: + ==> b ==> + // \\ + a d + \\ // + ==> c ==> + + class SimpleModule(torch.nn.Module): + def forward(self, a): + b = torch.sin(a) + c = torch.cos(a) + d = b + c + return d + + and providing "operator_support" that indicates that 'b' and 'c' can be executed on the accelerator, + we will get the following split result: + + main: + def forward(self, a): + run_on_acc_0_0 = self._run_on_acc_0_0(a) + getitem = run_on_acc_0_0[0] + getitem_1 = run_on_acc_0_0[1] + run_on_cpu_1_1 = self._run_on_cpu_1_1(getitem, getitem_1) + return run_on_cpu_1_1 + + _run_on_acc_0_0: + def forward(self, a): + sin_1 = torch.sin(a) + cos_1 = torch.cos(a) + return (sin_1, cos_1) + + _run_on_cpu_1_1: + def forward(self, sin_1, cos_1): + add_1 = sin_1 + cos_1 + return add_1 + """ + + # PCIe bandwidth for the backend, default to 100 GB/s + PCIe_BW = 100 * 2**30 + + def __init__( + self, + module: torch.fx.GraphModule, + sample_input: Sequence[Any], + operator_support: OperatorSupportBase, + settings: _SplitterSettingBase, + non_acc_submodule_name: str = "_run_on_cpu_", + return_tuple: bool = False, + nodes_finder: Optional[FxNetAccNodesFinder] = None, + ): + """ + Preprocesses graph before splitting: + - finds nodes supported by ACC, + - finds fusion groups for ACC nodes having non-tensor IO, + - builds a graph of direct dependencies, + - builds a map of fused nodes to their fusions. + As a result we get self.acc_nodes, self.deps and self.fusions. + """ + assert isinstance(module, torch.fx.GraphModule) + + self.module = module + ShapeProp(self.module).propagate(*sample_input) + + self.settings = settings + self.operator_support = operator_support + self.sample_input = sample_input + if nodes_finder is None: + nodes_finder = FxNetAccNodesFinder( + self.module, self.operator_support, self.settings.allow_non_tensor + ) + self.acc_nodes = nodes_finder() + + if self.settings.skip_fusion: + self.fusions = {} + else: + self.fusions = FxNetAccFusionsFinder(module, self.acc_nodes)() + + # Modify deps to add more deps for fused nodes + self.deps = self.find_deps() + self.update_deps_for_fusions() + + self.non_acc_submodule_name = non_acc_submodule_name + self._node_submodule_map: Dict[str, str] = {} + self._return_tuple = return_tuple + + self.tags: List[str] = [] + + # =============================================================== + # Helpers for ctor and initial state + # =============================================================== + + def get_node_submodule_map(self) -> Dict[str, str]: + """Returns a map from node name to submodule name, e.g. + node: main_module_impl_impl_over_arch_unary_multiple_embedding + _pooling_embedding_pooling_sparse_entity_equivalence_key + _proxy_embedding_bag + maps to submodule name of: _run_on_acc_1 + """ + return self._node_submodule_map + + def find_deps(self) -> Dict[torch.fx.Node, NodeSet]: + """ + Builds a graph of node dependencies. Leaf nodes don't have any + dependencies and the "output" node doesn't have nodes depending on it. + + Resulting graph has only direct dependencies, i.e. there are no + transitive dependencies. + """ + deps: Dict[torch.fx.Node, NodeSet] = defaultdict(set) + for node in self.module.graph.nodes: + if node.op not in CALLABLE_NODE_OPS: + continue + + for user in node.users: + if user.op != "output": + deps[user].add(node) + return deps + + def update_deps_for_fusions(self): + """ + Updates graph of dependencies so that: + - nodes from the same fusion depend on the same set of outer nodes, + - outer nodes depending on a fusion depend on all nodes in that fusion. + """ + for node in self.fusions: + fusion = self.fusions[node] + for fused_neighbor in fusion: + self.deps[node].update(self.deps[fused_neighbor] - fusion) + + for user in fused_neighbor.users: + if user not in fusion: + self.deps[user].add(node) + + # =============================================================== + # Helpers for preview + # =============================================================== + + def _lower_model_to_backend( + self, mod: torch.fx.GraphModule, inputs: Tensors + ) -> torch.nn.Module: + """ + Lower the model to a backend. + """ + + return mod + + def _find_culprit(self, mod: torch.fx.GraphModule, inputs: Tensors) -> str: + """ + When an error occurs during lowering or running the lowered mod, we use this + function to find culprits in the `mod` that causes the error. + """ + + return "Unable to find a culprit because _find_culprit() function is not implemented." + + def _draw_graph_based_on_node_support( + self, mod: torch.fx.GraphModule, supported_nodes: NodeList + ): + color_map = { + "default": "AliceBlue", + "supported": "chartreuse1", + "unsupported": "crimson", + } + + class CustomDrawer(FxGraphDrawer): + def _get_node_style(self, node): + template = super()._get_node_style(node) + if node in supported_nodes: + template["fillcolor"] = color_map["supported"] + elif node.op in CALLABLE_NODE_OPS: + template["fillcolor"] = color_map["unsupported"] + else: + template["fillcolor"] = color_map["default"] + + return template + + drawer = CustomDrawer(mod, "node_support", ignore_getattr=True) + dot_graph = drawer.get_main_dot_graph() + # pyre-fixme[16]: `pydot.Dot` has no attribute `write_raw`. + dot_graph.write_raw("node_support.dot") + + def node_support_preview(self, dump_graph: bool = False): + submodules = dict(self.module.named_modules()) + + supported_nodes: NodeList = [] + supported_node_types = defaultdict(set) + unsupported_node_types = defaultdict(set) + + def get_dtype(arg): + tensor_meta = arg.meta.get("tensor_meta") + return getattr(tensor_meta, "dtype", None) + + for node in self.module.graph.nodes: + if node.op not in CALLABLE_NODE_OPS: + continue + + target = get_node_target(submodules, node) + + # Store dtype of arg in node.args. If arg doesn't have dtype, i.e. not a tensor, we'll store None. + arg_dtypes = [ + get_dtype(arg) if isinstance(arg, torch.fx.Node) else None + for arg in node.args + ] + + # Find last non-None element. If all elements are None, return max_len. + last_index = len(arg_dtypes) - next( + ( + i + for i, dtype in enumerate(reversed(arg_dtypes)) + if dtype is not None + ), + len(arg_dtypes), + ) + + # Strip None elements at the end. + arg_dtypes_tuple = tuple(arg_dtypes[:last_index]) + kwarg_dtypes_tuple = tuple( + (k, get_dtype(arg)) + for k, arg in node.kwargs.items() + if isinstance(arg, torch.fx.Node) + ) + + if self.operator_support.is_node_supported(submodules, node): + supported_nodes.append(node) + supported_node_types[target].add((arg_dtypes_tuple, kwarg_dtypes_tuple)) + else: + unsupported_node_types[target].add( + (arg_dtypes_tuple, kwarg_dtypes_tuple) + ) + + if dump_graph: + self._draw_graph_based_on_node_support(self.module, supported_nodes) + + reports = "\nSupported node types in the model:\n" + for t, dtypes in supported_node_types.items(): + for arg_dtypes_tuple, kwarg_dtypes_tuple in dtypes: + reports += f"{t}: ({arg_dtypes_tuple}, {dict(kwarg_dtypes_tuple)})\n" + + reports += "\nUnsupported node types in the model:\n" + for t, dtypes in unsupported_node_types.items(): + for arg_dtypes_tuple, kwarg_dtypes_tuple in dtypes: + reports += f"{t}: ({arg_dtypes_tuple}, {dict(kwarg_dtypes_tuple)})\n" + + print(reports) + + # Return reports for testing purpose + return reports + + def split_preview(self, dump_graph: bool = False): + reports = "" + subgraphs = self.put_nodes_into_subgraphs() + acc_subgraphs_num = len([g for g in subgraphs if g.is_acc]) + cpu_subgraphs_num = len(subgraphs) - acc_subgraphs_num + reports += f"Before removing small acc subgraphs, total {len(subgraphs)} subgraphs are created:" + reports += f" {acc_subgraphs_num} acc subgraphs and {cpu_subgraphs_num} cpu subgraphs.\n" + + subgraphs = self.remove_small_acc_subgraphs(subgraphs) + acc_subgraphs_num = len([g for g in subgraphs if g.is_acc]) + cpu_subgraphs_num = len(subgraphs) - acc_subgraphs_num + reports += f"After removing small acc subgraphs, total {len(subgraphs)} subgraphs are created:" + reports += f" {acc_subgraphs_num} acc subgraphs and {cpu_subgraphs_num} cpu subgraphs.\n" + + for i, subgraph in enumerate(subgraphs): + reports += ( + f"_run_on_acc_{i}: " + if subgraph.is_acc + else f"{self.non_acc_submodule_name}{i}: " + ) + reports += f"{len(subgraph.nodes)} node(s)\n" + + self.tag(subgraphs) + split_mod = self.split(remove_tag=True) + split_mod.eval() + + if dump_graph: + drawer = FxGraphDrawer(split_mod, "preview", ignore_getattr=True) + dot_graphs = drawer.get_all_dot_graphs() + for name, dot_graph in dot_graphs.items(): + # pyre-fixme[16]: `pydot.Dot` has no attribute `write_raw`. + dot_graph.write_raw(f"{name}.dot") + + max_qps: float = self.PCIe_BW + bottleneck_module = "" + + for node in split_mod.graph.nodes: + if node.op == "call_module" and "acc" in node.target: + reports += f"\nProcessing acc submodule {node.target}\n" + + submod = getattr(split_mod, node.target) + + def get_submod_inputs(main_mod, submod, example_inputs): + sub_inputs = None + + def get_inputs(self, inputs): + nonlocal sub_inputs + sub_inputs = inputs + + handle = submod.register_forward_pre_hook(get_inputs) + main_mod(*example_inputs) + handle.remove() + return sub_inputs + + submod_inputs = get_submod_inputs(split_mod, submod, self.sample_input) + ShapeProp(submod).propagate(*submod_inputs) + + total_input_bytes = 0 + total_output_bytes = 0 + + reports += "Checking inputs...\n" + for n in submod.graph.nodes: + if n.op == "placeholder": + if not is_node_output_tensor(n): + reports += f"Input {n.name} is not a tensor, this might cause problems during lowering!\n" + else: + total_input_bytes += get_size_of_node(submod, n)[0] + if n.op == "output": + output_node = n + + reports += "Checking outputs...\n" + + def get_bytes(node: torch.fx.Node): + nonlocal total_output_bytes + nonlocal reports + if not is_node_output_tensor(node): + reports += f"Output {node.name} is not a tensor, this might cause problems during lowering!\n" + else: + total_output_bytes += get_size_of_node(submod, node)[0] + + map_arg(output_node.args, get_bytes) # type: ignore[possibly-undefined] + qps = self.PCIe_BW / max(total_input_bytes, total_output_bytes) + reports += f"Total input size in bytes is {total_input_bytes}, total output size in bytes is {total_output_bytes}," + reports += f" theoretical max qps (bounds by PCIe bandwidth) for this submodule is {qps}.\n" + + if qps < max_qps: + max_qps = qps + bottleneck_module = node.target + + try: + lowered_submod = self._lower_model_to_backend(submod, submod_inputs) + except RuntimeError: + reports += "Run into an error during lowering!\n" + reports += self._find_culprit(submod, submod_inputs) + continue + + try: + lowered_submod(*submod_inputs) + except RuntimeError: + reports += "Run into an error during inference!\n" + reports += self._find_culprit(submod, submod_inputs) + else: + reports += "Lowering and running succeed!\n" + + reports += f"\nTheoretical max qps (bounds by PCIe bandwidth) for this model is {max_qps}," + reports += f" bottleneck is submodule {bottleneck_module}." + print(reports) + + # return the reports for testing purposes + return reports + + # =============================================================== + # Helpers for extend_acc_subgraph() method + # =============================================================== + + def find_reverse_deps( + self, tag_id: Optional[int] = None + ) -> Dict[torch.fx.Node, NodeSet]: + """ + Builds reversed topological node dependencies, if tag_id is specified, + we ignore nodes that are in later subgraph i.e. nodes have greater tag_id. + """ + result: Dict[torch.fx.Node, NodeSet] = defaultdict(set) + + for node in self.module.graph.nodes: + if node.op not in CALLABLE_NODE_OPS: + continue + + for user in node.users: + if user.op not in CALLABLE_NODE_OPS: + continue + + if tag_id is None or (int(user.tag.split("_")[-1]) < tag_id): + result[node].add(user) + + return result + + def update_reverse_deps_for_fusions(self, deps: Dict[torch.fx.Node, NodeSet]): + processed_node = set() + + for node, fusion in self.fusions.items(): + if node in processed_node: + continue + + new_dep = set() + + # Create a new dependency set which include all the + # dependencies of the nodes in the fusion group + for n in fusion: + new_dep.update(deps[n]) + + # Exclude nodes in the fusion + new_dep.difference_update(fusion) + + # Update dependency + for n in fusion: + deps[n] = new_dep + + for arg in n.all_input_nodes: + if arg not in fusion: + deps[arg].update(fusion) + + processed_node.add(n) + + def find_parent_nodes_of_subgraph(self, tag: str) -> NodeSet: + """ + Finds parent nodes of the `tag` subgraph. + + Traverse the inputs of nodes in the subgraph, if input doesn't belong to the subgraph + and is not a placeholder, we consider it as the parent node of the subgraph. + """ + parent_nodes = set() + + for node in self.module.graph.nodes: + if node.op in CALLABLE_NODE_OPS and node.tag == tag: + for arg in node.all_input_nodes: + if arg.op in CALLABLE_NODE_OPS and arg.tag != tag: + parent_nodes.add(arg) + + return parent_nodes + + def extend_acc_subgraph(self, tag: str): + """ + Extend the acc subgraph with `tag` going the reversed topological direction. + """ + # Dict that maps node to its users and ignore users that + # are in the subgraph that has greater tag + deps = self.find_reverse_deps(tag_id=int(tag.split("_")[-1])) + self.update_reverse_deps_for_fusions(deps) + + # Parent nodes of the subgraph + parent_nodes = self.find_parent_nodes_of_subgraph(tag) + + visited_nodes: NodeSet = set() + + while parent_nodes: + node = None + + # Find a acc node that depends on visited nodes only + for n in parent_nodes: + if deps[n] <= visited_nodes and n in self.acc_nodes: + node = n + break + + if node is None: + break + + # Put the node into `tag` subgraph + node.tag = tag # type: ignore[attr-defined] + parent_nodes.remove(node) + visited_nodes.add(node) + + # If node is in a fusion group, add all fusion buddies to parent nodes + if node in self.fusions: + for fusion_node in self.fusions[node]: + if fusion_node not in visited_nodes: + parent_nodes.add(fusion_node) + + # Add inputs of the node to parent nodes + for arg in node.all_input_nodes: + if arg.op in CALLABLE_NODE_OPS and arg not in visited_nodes: + parent_nodes.add(arg) + + # =============================================================== + # Helpers for split() method + # =============================================================== + + def starter_nodes(self) -> Tuple[NodeSet, NodeSet]: + """ + Finds nodes that consume module inputs or get_attr nodes. + """ + starter_cpu_nodes: NodeSet = set() + starter_acc_nodes: NodeSet = set() + for node in self.module.graph.nodes: + if node.op not in {"placeholder", "get_attr"}: + continue + for user in node.users: + if user in self.acc_nodes: + starter_acc_nodes.add(user) + else: + starter_cpu_nodes.add(user) + return starter_cpu_nodes, starter_acc_nodes + + def put_nodes_into_subgraphs(self) -> List[Subgraph]: + # We start graph traversal from leaf nodes + current_cpu_nodes, current_acc_nodes = self.starter_nodes() + visited_nodes: NodeSet = set() + + # Determine which subgraph to start from based on which subgraph has + # 0-dep node + acc_subgraph: bool = not any(len(self.deps[n]) == 0 for n in current_cpu_nodes) + + current_subgraph_nodes: NodeList = [] + + # Result accumulator + subgraphs: List[Subgraph] = [] + while current_cpu_nodes or current_acc_nodes: + # Find the first node that should belong to the current subgraph and has all dependencies resolved + current_nodes = current_acc_nodes if acc_subgraph else current_cpu_nodes + node = next( + (n for n in current_nodes if self.deps[n] <= visited_nodes), + None, + ) + + # If nothing was found, then it's time to flip the mode and start a new subgraph + if node is None: + if not current_subgraph_nodes: + raise FxNetSplitterInternalError("Subgraph can't be empty") + + subgraphs.append( + Subgraph(is_acc=acc_subgraph, nodes=current_subgraph_nodes) + ) + acc_subgraph = not acc_subgraph + current_subgraph_nodes = [] + continue + + current_nodes.remove(node) + visited_nodes.add(node) + current_subgraph_nodes.append(node) + + # Add fusion buddies + if node in self.fusions: + if node in self.acc_nodes: + current_acc_nodes.update(self.fusions[node] - visited_nodes) + else: + current_cpu_nodes.update(self.fusions[node] - visited_nodes) + + # Put depending nodes into the queue + for user in node.users: + if user.op not in CALLABLE_NODE_OPS: + continue + + # Add downstream nodes + if user in self.acc_nodes: + current_acc_nodes.add(user) + else: + current_cpu_nodes.add(user) + + # Check if the last subgraph was not created + if current_subgraph_nodes: + subgraphs.append( + Subgraph(is_acc=acc_subgraph, nodes=current_subgraph_nodes) + ) + + if not subgraphs: + raise FxNetSplitterInternalError("Couldn't create subgraphs") + + return subgraphs + + def remove_small_acc_subgraphs(self, subgraphs: List[Subgraph]) -> List[Subgraph]: + """ + This pass finds ACC submodules with less than specified size and merges + them with adjacent CPU submodules. + """ + result: List[Subgraph] = [] + for subgraph in subgraphs: + if subgraph.is_acc: + if len(subgraph.nodes) >= self.settings.min_acc_module_size: + result.append(subgraph) + else: + print( + "Eliminating acc subgraph because it's smaller than the threshold: " + f"{len(subgraph.nodes)} < {self.settings.min_acc_module_size}" + ) + if result: + result[-1].nodes.extend(subgraph.nodes) + else: + subgraph.is_acc = False + result.append(subgraph) + else: + if result and not result[-1].is_acc: + result[-1].nodes.extend(subgraph.nodes) + else: + result.append(subgraph) + return result + + def tag(self, subgraphs: List[Subgraph]): + self.tags = [] + for subgraph in subgraphs: + tag = ( + f"_run_on_acc_{len(self.tags)}" + if subgraph.is_acc + else f"{self.non_acc_submodule_name}{len(self.tags)}" + ) + self.tags.append(tag) + for node in subgraph.nodes: + if hasattr(node, "tag"): + raise FxNetSplitterInternalError(f"Node {node} was already tagged") + + node.tag = tag # type: ignore[attr-defined] + self._node_submodule_map[node.name] = tag + + def split(self, remove_tag: bool = False) -> torch.fx.GraphModule: + split_module = split_by_tags( + self.module, self.tags, return_tuple=self._return_tuple + ) + if remove_tag: + for node in self.module.graph.nodes: + if hasattr(node, "tag"): + del node.tag + return split_module # type: ignore[return-value] + + def __call__(self) -> torch.fx.GraphModule: + subgraphs = self.put_nodes_into_subgraphs() + subgraphs = self.remove_small_acc_subgraphs(subgraphs) + acc_subgraphs_count = len([s for s in subgraphs if s.is_acc]) + non_acc_subgraphs_count = len(subgraphs) - acc_subgraphs_count + print( + f"Got {acc_subgraphs_count} acc subgraphs and {non_acc_subgraphs_count} non-acc subgraphs" + ) + self.tag(subgraphs) + return self.split() + + def generate_split_results(self) -> SplitResult: + split_module = self() + submodule_names = [] + for name, _mod in split_module.named_children(): + submodule_names.append(name) + if ( + self.settings.max_acc_splits > 0 + and len(submodule_names) > self.settings.max_acc_splits + ): + raise ValueError( + "Cannot fulfill max_acc_splits limit. " + "This may cause split fragmentation and " + "result in performance issues." + ) + + submodule_inputs = generate_inputs_for_submodules( + split_module, self.sample_input, submodule_names + ) + return SplitResult(split_module, submodule_inputs, self.non_acc_submodule_name) diff --git a/lib/python3.12/site-packages/torch/fx/passes/tests/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torch/fx/passes/tests/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/tests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7af59f9a80d0a809b0c9e6e9334c730033512e8e Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/tests/__pycache__/test_pass_manager.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/tests/__pycache__/test_pass_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93add34c21cc8c462265cd2dc5d2f0eb7e287589 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/tests/__pycache__/test_pass_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/tests/test_pass_manager.py b/lib/python3.12/site-packages/torch/fx/passes/tests/test_pass_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..157dc4017eda576f10793ef46b78cd97b0f5074b --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/tests/test_pass_manager.py @@ -0,0 +1,56 @@ +import unittest + +from ..pass_manager import ( + inplace_wrapper, + PassManager, + these_before_those_pass_constraint, + this_before_that_pass_constraint, +) + + +class TestPassManager(unittest.TestCase): + def test_pass_manager_builder(self) -> None: + passes = [lambda x: 2 * x for _ in range(10)] + pm = PassManager(passes) + pm.validate() + + def test_this_before_that_pass_constraint(self) -> None: + passes = [lambda x: 2 * x for _ in range(10)] + pm = PassManager(passes) + + # add unfulfillable constraint + pm.add_constraint(this_before_that_pass_constraint(passes[-1], passes[0])) + + self.assertRaises(RuntimeError, pm.validate) + + def test_these_before_those_pass_constraint(self) -> None: + passes = [lambda x: 2 * x for _ in range(10)] + constraint = these_before_those_pass_constraint(passes[-1], passes[0]) + pm = PassManager([inplace_wrapper(p) for p in passes]) + + # add unfulfillable constraint + pm.add_constraint(constraint) + + self.assertRaises(RuntimeError, pm.validate) + + def test_two_pass_managers(self) -> None: + """Make sure we can construct the PassManager twice and not share any + state between them""" + + passes = [lambda x: 2 * x for _ in range(3)] + constraint = these_before_those_pass_constraint(passes[0], passes[1]) + pm1 = PassManager() + for p in passes: + pm1.add_pass(p) + pm1.add_constraint(constraint) + output1 = pm1(1) + self.assertEqual(output1, 2**3) + + passes = [lambda x: 3 * x for _ in range(3)] + constraint = these_before_those_pass_constraint(passes[0], passes[1]) + pm2 = PassManager() + for p in passes: + pm2.add_pass(p) + pm2.add_constraint(constraint) + output2 = pm2(1) + self.assertEqual(output2, 3**3) diff --git a/lib/python3.12/site-packages/torch/fx/passes/tools_common.py b/lib/python3.12/site-packages/torch/fx/passes/tools_common.py new file mode 100644 index 0000000000000000000000000000000000000000..4ed56be63b0926039868789db4cab57bf667d7c0 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/tools_common.py @@ -0,0 +1,318 @@ +# mypy: allow-untyped-defs +import collections +import operator +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Optional, Set, Tuple, Union + +import torch +import torch.fx +from torch.fx._compatibility import compatibility +from torch.fx.node import _get_qualified_name + + +__all__ = [ + "get_acc_ops_name", + "get_node_target", + "is_node_output_tensor", + "FxNetAccFusionsFinder", + "legalize_graph", +] + +Tensors = Union[Tuple[torch.Tensor], List[torch.Tensor]] +TensorOrTensors = Union[torch.Tensor, Tensors] +NodeList = List[torch.fx.Node] +NodeSet = Set[torch.fx.Node] +Names = List[str] +CALLABLE_NODE_OPS = {"call_module", "call_function", "call_method"} + + +@compatibility(is_backward_compatible=False) +def get_acc_ops_name(k): + if isinstance(k, str): + return k + elif k.__module__ and "acc_ops" in k.__module__: + return f"acc_ops.{k.__name__}" + else: + module = k.__module__.replace( + "torch._ops", "torch.ops" + ) # WAR for bug in how torch.ops assigns module + return f"{module if module else ''}.{k.__name__}" + + +@compatibility(is_backward_compatible=False) +def get_node_target( + submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node +) -> str: + """ + Given a `node` returns its target typename. + + For "call_method" node, return node.target which is the name of that method being called. + This could potential lead to conflict but should be okay because normally it's on a tensor. + + For "call_function" node, return typename of node.target. + + For "call_module" node, return typename of the module that node.target point to. + + If seeing "_VariableFunctionsClass" in the target name string, it will be replaced by + "torch". e.g. _VariableFunctionsClass.relu would become torch.relu. + """ + + assert node.op in CALLABLE_NODE_OPS, ( + "Expect op types of " + ", ".join(CALLABLE_NODE_OPS) + f", but found {node.op}" + ) + + if node.op == "call_module": + assert isinstance(node.target, str) + submod = submodules[node.target] + submod_type = getattr(submod, "_base_class_origin", type(submod)) + return get_acc_ops_name(submod_type) + elif node.op == "call_function": + target: Any = node.target + return ( + f"acc_ops.{target.__name__}" + if target.__module__ is not None and "acc_ops" in target.__module__ + else _get_qualified_name(target) + ) + else: + assert isinstance(node.target, str) + return node.target + + +@compatibility(is_backward_compatible=False) +def is_node_output_tensor(node: torch.fx.Node) -> bool: + """Checks if the node output produces a Tensor or not. + + NOTE: This requires to run `ShapeProp` on the containing fx graph before + calling this function. This is because it works by checking the `type` + metadata on the node. This metadata is produced by the `ShapeProp`. + """ + type_ = node.meta.get("type", None) + return type_ is not None and issubclass(type_, torch.Tensor) + + +@compatibility(is_backward_compatible=False) +class FxNetAccFusionsFinder: + """ + Finds groups of connected ACC nodes that pass non-tensor data between each other. + Such groups are called fusion groups. + """ + + def __init__(self, module: torch.fx.GraphModule, acc_nodes: NodeSet): + self.module = module + self.nodes = list(module.graph.nodes) + self.acc_nodes = acc_nodes + + @dataclass + class FusionGroup: + # The smallest idx of nodes in the fusion group after topological sorting all the nodes in the model. + top_node_idx: int + + # Nodes in this fusion group. + nodes: NodeSet + + # Inputs to this fusion group. + inputs: NodeSet + + # Nodes that in the fusion group that haven't been processed yet. + nodes_need_process: NodeSet + + def add_node(self, node): + """ + Add a node to fusion group. + """ + if node in self.nodes: + return + + self.nodes_need_process.add(node) + self.nodes.add(node) + self.inputs.discard(node) + self.inputs.update( + { + n + for n in node.all_input_nodes + if n.op in CALLABLE_NODE_OPS and n not in self.nodes + } + ) + + def recursive_add_node( + self, + fusion_group: "FxNetAccFusionsFinder.FusionGroup", + inputs: Union[NodeSet, NodeList], + visited: Optional[NodeSet] = None, + ): + """ + Start from inputs and going reverse topological order. If any upstream node + is in the fusion group, add all the nodes in this path to fusion group. + """ + for arg in inputs: + # skip the node if already seen + if visited is not None: + if arg in visited: + continue + visited.add(arg) + + # Skip placeholder and get_attr because they won't be in the fusion group. + if arg.op not in CALLABLE_NODE_OPS: + continue + + # If the node has smaller idx, it's already an upstream node of the fusion + # group. We don't need to check it anymore. + if self.nodes.index(arg) < fusion_group.top_node_idx: + continue + + # If the node is in the fusion group, return True. + if arg in fusion_group.nodes: + return True + + # Check the upstream nodes of the node, if any of them is in the fusion group + # we'll add this node to fusion group and return True. + if self.recursive_add_node(fusion_group, arg.all_input_nodes, visited): + fusion_group.add_node(arg) + return True + + return False + + def __call__(self) -> Dict[torch.fx.Node, NodeSet]: + result: Dict[torch.fx.Node, NodeSet] = {} + acc_nodes = list(self.acc_nodes) + + for node in acc_nodes: + if node in result: + continue + if node.op not in CALLABLE_NODE_OPS: + continue + if "tensor_meta" in node.meta: + continue + if node not in self.acc_nodes: + continue + + fusion_group: FxNetAccFusionsFinder.FusionGroup = self.FusionGroup( + top_node_idx=self.nodes.index(node), + nodes={node}, + inputs=set(node.all_input_nodes), + nodes_need_process={node}, + ) + while fusion_group.nodes_need_process: + node = fusion_group.nodes_need_process.pop() + self.recursive_add_node( + fusion_group, + fusion_group.inputs, + visited=set(), + ) + + # Optionally add downstream nodes + if "tensor_meta" not in node.meta: + for user in node.users: + if user.op not in CALLABLE_NODE_OPS: + continue + if user in fusion_group.nodes: + continue + + fusion_group.add_node(user) + self.recursive_add_node( + fusion_group, + fusion_group.inputs, + visited=set(), + ) + + # Add some upstream nodes + for arg in node.all_input_nodes: + if arg.op not in CALLABLE_NODE_OPS: + continue + if "tensor_meta" in arg.meta: + continue + if arg in fusion_group.nodes: + continue + + fusion_group.add_node(arg) + fusion_group.top_node_idx = min( + fusion_group.top_node_idx, self.nodes.index(arg) + ) + self.recursive_add_node( + fusion_group, + fusion_group.inputs, + visited=set(), + ) + + if not (set(fusion_group.nodes) <= self.acc_nodes): + self.acc_nodes -= fusion_group.nodes + else: + for n in fusion_group.nodes: + result[n] = fusion_group.nodes + + return result + + +@compatibility(is_backward_compatible=False) +def legalize_graph(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: + """ + Replace the graph of the given GraphModule with one that contains the same nodes as the + original, but in topologically sorted order. + + This is used by the merge_matmul transformation below, which disturbs the topologically sorted + order of its input GraphModule, so that this order is restored before further transformation. + + Arguments: + gm: The graph module to topologically sort. It is modified in-place. + + Returns: + The graph module in-place sorted + """ + + # These operators are used for making runtime assertions before any + # data-dependent operators occur. We want to prioritize sorting these to + # ensure that these assertions appear before any data-dependent operations + # in the graph. + PRIORITIZED_OPS = [ + operator.add, + operator.mul, + operator.sub, + operator.floordiv, + operator.truediv, + operator.mod, + operator.le, + operator.lt, + operator.ge, + operator.gt, + operator.eq, + operator.ne, + torch.ops.aten.sym_constrain_range.default, + torch.ops.aten.sym_constrain_range_for_size.default, + torch.ops.aten._assert_async.msg, + torch.ops.aten.scalar_tensor.default, + torch.ops.aten._assert_scalar.default, + ] + + indeg = dict.fromkeys(gm.graph.nodes, 0) + new_graph = torch.fx.Graph() + # Track how many unfulfilled dependencies each node has + for node in gm.graph.nodes: + for user in node.users: + indeg[user] += 1 + queue: collections.deque = collections.deque() + # Add all nodes with no dependencies to the queue + for node in gm.graph.nodes: + if indeg[node] == 0: + queue.append(node) + env: Dict[torch.fx.Node, torch.fx.Node] = {} + # Pop nodes from the queue, and add nodes that have had all their + # dependencies fulfilled + while len(queue) > 0: + cur = queue.popleft() + env[cur] = new_graph.node_copy(cur, lambda x: env[x]) + for user in cur.users: + indeg[user] -= 1 + if indeg[user] == 0: + if user.op == "call_function" and user.target in PRIORITIZED_OPS: + queue.appendleft(user) + else: + queue.append(user) + # If the new graph's size is not as large as the old one, then there must be + # a cycle (i.e. some node's dependencies were not satisfied.) + if len(new_graph.nodes) < len(gm.graph.nodes): + raise RuntimeError( + f"Input graph has cycles, unable to add {[node for node in indeg if indeg[node] != 0]}" + ) + new_graph._codegen = gm.graph._codegen + gm.graph = new_graph + return gm diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__init__.py b/lib/python3.12/site-packages/torch/fx/passes/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ee5e7e66868a0776609ff7ffff458f6a91ccf98a --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/utils/__init__.py @@ -0,0 +1 @@ +from .common import compare_graphs, HolderModule, lift_subgraph_as_module diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8eddcc675e96bcf014ae133e4b1c2d36ccdb0e68 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d452aaa9fc968f928cb50572c38822509eef8960 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/fuser_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/fuser_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8abf10e8f4d1c3beadfeed43255be4329ed4f28 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/fuser_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/matcher_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/matcher_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71ff24dda010f3a30158d1bf3677733421fc9fa0 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/matcher_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/matcher_with_name_node_map_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/matcher_with_name_node_map_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e52879eb801afd5e37941d475a6b977cc35d0d43 Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/matcher_with_name_node_map_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/source_matcher_utils.cpython-312.pyc b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/source_matcher_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..031374a19f56e8d352bb710092fa7523a571f1ac Binary files /dev/null and b/lib/python3.12/site-packages/torch/fx/passes/utils/__pycache__/source_matcher_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/common.py b/lib/python3.12/site-packages/torch/fx/passes/utils/common.py new file mode 100644 index 0000000000000000000000000000000000000000..bb628372337b421f88863de874b75993e7fe7d2f --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/utils/common.py @@ -0,0 +1,95 @@ +# mypy: allow-untyped-defs +from typing import Dict, Tuple + +from torch.fx._compatibility import compatibility +from torch.fx.graph import Graph +from torch.fx.graph_module import GraphModule +from torch.fx.passes.utils.matcher_utils import SubgraphMatcher +from torch.nn import Module + + +__all__ = ["HolderModule", "lift_subgraph_as_module", "compare_graphs"] + + +@compatibility(is_backward_compatible=False) +class HolderModule(Module): + """ + HolderModule is used to copy all the attributes from original module to submodules + that uses the attributes + """ + + def __init__(self, d): + super().__init__() + for k, v in d.items(): + self.add_module(k, v) + + +@compatibility(is_backward_compatible=False) +def lift_subgraph_as_module( + gm: GraphModule, + subgraph: Graph, + comp_name: str = "", + class_name: str = "GraphModule", +) -> Tuple[GraphModule, Dict[str, str]]: + """ + Create a GraphModule for subgraph, which copies the necessary attributes from the original parent graph_module. + + Args: + gm (GraphModule): parent graph module + + subgraph (Graph): a valid subgraph that contains copied nodes from the parent graph + + comp_name (str): name for the new component + + class_name (str): name for the submodule + + """ + + # Loop through all module calls (call_module) and param fetches (get_attr) + # in this component, creating HolderModules as necessary to match the path. + # e.g. if in the original module there's a get_attr node fetches "conv.weight". + # We create a HolderModule as root -> add a HolderModule named "conv" -> + # make "weight" a attribute of "conv" HolderModule and point to conv.weight in + # the original module. + submodule = HolderModule({}) + orig_to_split_fqn_mapping: Dict[str, str] = {} + for n in subgraph.nodes: + if n.op not in ("call_module", "get_attr"): + continue + + target = n.target + assert isinstance(target, str) + target_name_parts = target.split(".") + curr = submodule + orig_gm = gm + + for name in target_name_parts[:-1]: + if not hasattr(curr, name): + curr.add_module(name, HolderModule({})) + + curr = getattr(curr, name) + orig_gm = getattr(orig_gm, name) + + leaf_node_name = target_name_parts[-1] + leaf_node = getattr(orig_gm, leaf_node_name) + + orig_to_split_fqn_mapping[target] = f"{comp_name}.{target}" + # Relies on custom __setattr__ magic. + setattr(curr, leaf_node_name, leaf_node) + + return GraphModule(submodule, subgraph, class_name), orig_to_split_fqn_mapping + + +@compatibility(is_backward_compatible=False) +def compare_graphs(left: Graph, right: Graph) -> bool: + """ + Return True if two graphs are identical, i.e they + - have the same number of outputs in the same order + - have the same number of inputs in the same order + - have the same set of nodes, and identical connectivity + """ + + matcher = SubgraphMatcher(left, match_output=True, match_placeholder=True) + matches = matcher.match(right) + + return len(matches) > 0 diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/fuser_utils.py b/lib/python3.12/site-packages/torch/fx/passes/utils/fuser_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fa090b677f32f5d9110a954806c913017cfbc362 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/utils/fuser_utils.py @@ -0,0 +1,276 @@ +# mypy: allow-untyped-defs +import copy +from queue import SimpleQueue +from typing import Dict, List, Optional as _Optional, Tuple + +import torch.fx +from torch.fx._compatibility import compatibility +from torch.fx.graph import Graph +from torch.fx.graph_module import GraphModule +from torch.fx.node import Node +from torch.fx.passes.tools_common import legalize_graph, NodeList, NodeSet +from torch.fx.passes.utils import lift_subgraph_as_module + + +@compatibility(is_backward_compatible=False) +def topo_sort(nodes: NodeList) -> NodeList: + # sort nodes according to the topological order + indegree_map = dict.fromkeys(nodes, 0) + candidates: SimpleQueue = SimpleQueue() + + for node in nodes: + for n in node.all_input_nodes: + if n in indegree_map: + indegree_map[node] += 1 + if indegree_map[node] == 0: + candidates.put(node) + + sorted_nodes: NodeList = [] + while not candidates.empty(): + node = candidates.get() + sorted_nodes.append(node) + + for n in node.users: + if n in indegree_map: + indegree_map[n] -= 1 + if indegree_map[n] == 0: + candidates.put(n) + + assert len(nodes) == len( + sorted_nodes + ), "topological sorted nodes doesn't have same length as input nodes" + + return sorted_nodes + + +@compatibility(is_backward_compatible=False) +def validate_partition(partition: NodeList) -> bool: + # verify the partition does't form a dependency cycle in the original graph + # returns True for valid partition, False for invalid + + partition_set = set(partition) + + outputs: NodeList = [] + for node in partition_set: + for user_node in node.users: + if user_node not in partition_set: + # external user node, need to expose as an output + outputs.append(user_node) + + # Perform BFS on the partition outputs. + # If it reaches a node within the partition, then it found a cycle. + # This function takes the ownership of `root_nodes` and may modify it. + def bfs_find_cycle(root_nodes: NodeList) -> bool: + # Set used to exclude nodes that have already been visited. + # If a node has been visited, that node and all its children have + # been checked for cycles. + visited: NodeSet = set() + + # Start with `root_nodes` and traverse through (toward child nodes) + # their connected sub-graph. Nodes in `visited` won't be added + # to `queue` again. + queue: NodeList = root_nodes + while queue: + current = queue.pop() + visited.add(current) + if current in partition_set: + # Started from partition's `output` nodes, and reached + # another node in partition. Cycle! + return True + for user_node in current.users: + if user_node in visited: + continue + queue.append(user_node) + # `root_nodes` don't cause cycle. + return False + + # Use all output nodes as roots to traverse + # the graph to check cycles. + if bfs_find_cycle(outputs): + return False + + return True + + +@compatibility(is_backward_compatible=False) +def fuse_as_graphmodule( + gm: GraphModule, + nodes: NodeList, + module_name: str, + partition_lookup_table: _Optional[Dict[Node, None]] = None, + *, + always_return_tuple: bool = False, +) -> Tuple[GraphModule, Tuple[Node, ...], Tuple[Node, ...]]: + """ + Fuse nodes in graph_module into a GraphModule. + + Args: + gm (GraphModule): target graph_module + + nodes (List[Node]): list of nodes in `gm` to fuse, where the node must be topologically sorted + + module_name: class name for the fused GraphModule + + partition_lookup_table (Optional[Dict[Node, None]]): optional dict of nodes to speed up lookup + + always_return_tuple (bool): whether to always return a tuple, even if there is only one output + + Returns: + fused_gm (GraphModule): fused graph module, where its node is a copy of `nodes` in `gm` + + original_inputs (Tuple[Node, ...]): input nodes to `nodes` in original `gm` + + original_outputs (Tuple[Node, ...]): consumer nodes of `nodes` in original `gm` + + """ + + # assumption: nodes are already sorted in topo order + + for node in nodes: + assert ( + node.graph.owning_module is gm + ), f"{node} doesn't belong to passed in graph module {gm._get_name()}" + assert not node._erased, f"{node} has been removed from owning graph" + assert ( + node in gm.graph._find_nodes_lookup_table + ), f"{node} is not found in graph module {gm._get_name()}" + + # validates partition doesn't introduce dependency circles in the graph + assert validate_partition(nodes), "Invalid partition, found dependency cycles" + + # if no dict of partition nodes is provided, reconstruct it by nodes list to reduce lookup time + if partition_lookup_table is None: + partition_lookup_table = dict.fromkeys(nodes) + + subgraph = Graph() + + node_to_placeholder: Dict[ + Node, Node + ] = {} # mapping of nodes from old graph to placeholder in new graph + node_map: Dict[Node, Node] = {} # mapping of nodes from old graph to new graph + + # handles inputs through graph.node_copy's arg_transform functions + def remap_inputs(x): + if x.op == "get_attr": + # TODO: do we really need copy the get_attr node into the graph? + # do something here + pass + + if x in partition_lookup_table: + # x is inside subgraph, return the copied node + # the node should have been copied aleady, as we are copying graph in the topological order + return node_map[x] + + if x not in node_to_placeholder: + # x is not in subgraph, create a new placeholder for subgraph + placeholder_node = subgraph.placeholder(x.name, type_expr=x.type) + # copy all meta fields, even if some fields might be irrelvant for the placeholder node + placeholder_node.meta = copy.copy(x.meta) + node_to_placeholder[x] = placeholder_node + + return node_to_placeholder[x] + + # copy nodes in topological order + for node in nodes: + new_node = subgraph.node_copy(node, remap_inputs) + node_map[node] = new_node + + # handles outputs + output_mapping: Dict[Node, Node] = {} # mapping from old output to new outputs + + for node in nodes: + for user_node in node.users: + if user_node not in partition_lookup_table: + # external user node, need to expose as an output + output_mapping[node] = node_map[node] + + # outs contain nodes in the new subgraph + outs = tuple(output_mapping.values()) + + if always_return_tuple: + # always return a tuple, even if there is only one output + subgraph.output(outs) + else: + # If there's a single output then return it directly, otherwise return a tuple. + subgraph.output(outs[0] if len(outs) == 1 else outs) + + # lint to ensure correctness + subgraph.lint() + fused_gm: GraphModule + fused_gm, _ = lift_subgraph_as_module( + gm, subgraph, comp_name="", class_name=module_name + ) + + # sub_gm's input nodes in the original module + original_inputs: Tuple[Node, ...] = tuple(node_to_placeholder.keys()) + + # sub_gm's outputs node in the original module + original_outputs: Tuple[Node, ...] = tuple(output_mapping.keys()) + + return fused_gm, original_inputs, original_outputs + + +@compatibility(is_backward_compatible=False) +def insert_subgm( + gm: GraphModule, + sub_gm: GraphModule, + orig_inputs: Tuple[Node, ...], + orig_outputs: Tuple[Node, ...], +): + # add sub_gm into gm + submodule_name = sub_gm.__class__.__name__ + gm.add_submodule(submodule_name, sub_gm) + + # Create a call_module node in main graph. + module_node = gm.graph.call_module(submodule_name, args=orig_inputs, kwargs=None) + + output_node = sub_gm.graph.output_node() + if len(orig_outputs) == 1 and not isinstance(output_node.args[0], tuple): + # main_remapping[comp.orig_outputs[0]] = module_node + orig_outputs[0].replace_all_uses_with(module_node, propagate_meta=True) + else: + for i, orig_output in enumerate(orig_outputs): + # Use Proxy to record getitem access. + proxy_out = torch.fx.Proxy(module_node)[i].node # type: ignore[index] + orig_output.replace_all_uses_with(proxy_out, propagate_meta=True) + + module_node.meta["val"] = tuple( + orig_output.meta.get("val", None) for orig_output in orig_outputs + ) + return gm + + +@compatibility(is_backward_compatible=False) +def erase_nodes(gm: GraphModule, nodes: NodeList): + # erase original nodes in inversed topological order + for node in reversed(nodes): + gm.graph.erase_node(node) + + +@compatibility(is_backward_compatible=False) +def fuse_by_partitions( + gm: GraphModule, + partitions: List[Dict[Node, None]], + prefix: str = "fused_", + always_return_tuple: bool = False, +) -> GraphModule: + for partition_id, partition in enumerate(partitions): + sorted_nodes = topo_sort(list(partition)) + + submodule_name = prefix + str(partition_id) + sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule( + gm, + sorted_nodes, + submodule_name, + partition, + always_return_tuple=always_return_tuple, + ) + + insert_subgm(gm, sub_gm, orig_inputs, orig_outputs) + + erase_nodes(gm, sorted_nodes) + + # topological sort original gm with newly created sub_gm + legalize_graph(gm) + + return gm diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/matcher_utils.py b/lib/python3.12/site-packages/torch/fx/passes/utils/matcher_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cc05b8f512b158151d31d0e48f43e3511245507f --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/utils/matcher_utils.py @@ -0,0 +1,440 @@ +# mypy: allow-untyped-defs +import copy +import logging +import os +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Dict, List, Set, Tuple, Union + +import torch +from torch.fx import Graph, Node +from torch.fx._compatibility import compatibility + + +__all__ = ["SubgraphMatcher", "InternalMatch"] + + +# Set`PYTORCH_MATCHER_LOGLEVEL=INFO` to see debug logs +def _init_logger(): + logger = logging.getLogger(__name__) + + level = os.environ.get("PYTORCH_MATCHER_LOGLEVEL", "WARNING").upper() + logger.setLevel(level) + console = logging.StreamHandler() + formatter = logging.Formatter("%(filename)s > %(message)s") + console.setFormatter(formatter) + console.setLevel(level) + # add the handlers to the logger + logger.addHandler(console) + logger.propagate = False + return logger + + +logger = _init_logger() + + +@compatibility(is_backward_compatible=False) +@dataclass +class InternalMatch: + # Nodes from which the match was found + anchors: List[Node] + # Maps nodes in the pattern subgraph to nodes in the larger graph + nodes_map: Dict[Node, Node] = field(default_factory=dict) + + # nodes in target graph that are matched placeholder in pattern + placeholder_nodes: List[Node] = field(default_factory=list) + + # nodes in matched subgraph returned by output + returning_nodes: List[Node] = field(default_factory=list) + + # map from a string name to a node in the target graph + # only available if the matcher is `SubgraphMatcherWithNameNodesMap` + name_node_map: Dict[str, Node] = field(default_factory=dict) + + def __copy__(self): + return InternalMatch( + anchors=self.anchors, + nodes_map=self.nodes_map.copy(), + placeholder_nodes=self.placeholder_nodes.copy(), + returning_nodes=self.returning_nodes.copy(), + ) + + +@compatibility(is_backward_compatible=False) +class SubgraphMatcher: + def __init__( + self, + pattern: Graph, + match_output: bool = False, + match_placeholder: bool = False, + remove_overlapping_matches: bool = True, + ignore_literals: bool = False, + ) -> None: + """ + Args: + pattern: the targeted matching pattern, represented in fx.Graph. + match_output: If True, output node in the pattern graph will be treated as a part of the targeted pattern. + If False, output node is ignored during match. + match_placeholder: If True, placeholder node in the pattern graph will be treated as a part of + the targeted pattern. If False, placeholder nodes will be used a wildcard. + remove_overlapping_matches: If True, in the case of overlapping matches, only the first match + will be returned. + ignore_literals: If True, will not check if literals are equal and + will instead treat them as wildcards. + """ + + self.pattern = pattern + self.match_output = match_output + self.match_placeholder = match_placeholder + self.remove_overlapping_matches = remove_overlapping_matches + self.ignore_literals = ignore_literals + + if len(pattern.nodes) == 0: + raise ValueError( + "SubgraphMatcher cannot be initialized with an empty pattern" + ) + + for node in pattern.nodes: + if node.op != "output": + assert ( + len(node.users) > 0 + ), "SubgraphMatcher cannot be initialized with an pattern with dead code" + + # TODO: assert pattern is a connected graph + + self.pattern_placeholder_nodes = [ + n for n in pattern.nodes if n.op == "placeholder" + ] + output_node = next(iter(reversed(pattern.nodes))) + # nodes returned by outputs + self.pattern_returning_nodes: List[Node] = output_node.all_input_nodes + + self.pattern_anchors: List[Node] = [] + if match_output: + self.pattern_anchors = [output_node] + else: + # If a node has output_node as the ONLY user, then this node is a graph sink, + # and should be matched against as an anchor + self.pattern_anchors = [ + n for n in output_node.all_input_nodes if len(n.users) == 1 + ] + + def _match_attributes(self, pn: Node, gn: Node) -> bool: + # Attributes matching is complicated. Right now we only support matching constant tensor + assert isinstance(pn.target, str), f"pn.target {pn.target} must be a string." + assert isinstance(gn.target, str), f"gn.target {gn.target} must be a string." + + pn_value = torch.fx.graph_module._get_attr(pn.graph.owning_module, pn.target) + gn_value = torch.fx.graph_module._get_attr(gn.graph.owning_module, gn.target) + + if type(pn_value) != type(gn_value): + return False + + # Don't require exact match on tensor values. + if isinstance(pn_value, torch.Tensor): + return isinstance(gn_value, torch.Tensor) + else: + raise RuntimeError(f"Unsupported type {pn_value} when matching attributes") + return False + + def _nodes_are_equal(self, pn: Node, gn: Node) -> bool: + # if exact match for placeholder is not required, then use placeholder as a wildcard + if not self.match_placeholder and pn.op == "placeholder": + return True + + if pn.op == gn.op: + if pn.op == "placeholder" or pn.op == "output": + return True + elif pn.op == "get_attr": + return self._match_attributes(pn, gn) + return pn.target == gn.target + return False + + def _is_contained(self, nodes_map: Dict[Node, Node]) -> bool: + # `lookup` represents all the nodes in `original_graph` + # that are part of `pattern` + + # Placeholders can be used by other nodes in the graphs + lookup: Dict[Node, Node] = { + gn: pn for pn, gn in nodes_map.items() if pn.op != "placeholder" + } + + for gn, pn in lookup.items(): + # nodes returned by output are allowed to be used in other areas of the graph + if pn in self.pattern_returning_nodes: + continue + + for user in gn.users: + # If this node has users that were not in `lookup`, then it must leak out of the + # pattern subgraph + if user not in lookup: + return False + return True + + def _remove_overlapping_matches( + self, matches: List[InternalMatch] + ) -> List[InternalMatch]: + non_overlapping_matches: List[InternalMatch] = [] + nodes_matched: Set[Node] = set() + + for match in matches: + found_overlap = False + for pn, gn in match.nodes_map.items(): + if pn.op not in {"placeholder", "output"} and gn in nodes_matched: + found_overlap = True + break + + if not found_overlap: + non_overlapping_matches.append(match) + for pn, gn in match.nodes_map.items(): + if pn.op not in {"placeholder", "output"}: + nodes_matched.add(gn) + return non_overlapping_matches + + def _match_literals(self, pn: Any, gn: Any, match: InternalMatch) -> bool: + assert not ( + isinstance(pn, Node) and isinstance(gn, Node) + ), "pn and gn cannot both be Node" + + if isinstance(pn, Node) and not isinstance(gn, Node): + if pn.op == "placeholder": + # Check if we've already matched these nodes in the current + # traversal + if pn in match.nodes_map: + return match.nodes_map[pn] == gn + + match.nodes_map[pn] = gn + return True + else: + return False + elif not isinstance(pn, Node) and isinstance(gn, Node): + return False + else: + return type(gn) == type(pn) and gn == pn + + def _match_nodes(self, pn: Node, gn: Node, match: InternalMatch) -> bool: + logger.info(" matching %s to %s", pn, gn) + + assert isinstance(pn, Node) and isinstance(gn, Node), str( + f"pn and gn must be Node, pn: {pn}, gn: {gn}" + ) + + # Check if we've already matched these nodes in the current + # traversal + if pn in match.nodes_map: + return match.nodes_map[pn] == gn + + # TODO: use a more efficient way to check if gn is matched before: two-way dict + if gn in match.nodes_map.values(): + return False + + if not self._nodes_are_equal(pn, gn): + return False + + # Optimistically mark `pn` as a match for `gn`, and save a local copy of match + saved_match = copy.copy(match) + match.nodes_map[pn] = gn + + # Placeholder is a wildcard and can be matched with any python object + # (including list/tuple) + if pn.op == "placeholder": + return True + + # Recursively traverse upwards to check if `pn` is a true + # match for `gn` + match_found = True + + def _match_args(args1: Union[List, Tuple], args2: Union[List, Tuple]) -> bool: + if len(args1) != len(args2): + return False + + for a1, a2 in zip(args1, args2): + if isinstance(a1, Node) and isinstance(a2, Node): + matched = self._match_nodes(a1, a2, match) + elif isinstance(a1, (list, tuple)) and isinstance(a2, (list, tuple)): + matched = _match_args(a1, a2) + else: + matched = ( + self._match_literals(a1, a2, match) or self.ignore_literals + ) + + if not matched: + return False + + return True + + # Flatten all args/kwargs into 1 list of args + pn_args, gn_args = None, None + if ( + ( + len(pn.args) != len(gn.args) + or list(pn.kwargs.keys()) != list(gn.kwargs.keys()) + ) + and pn.op == "call_function" + and isinstance(pn.target, torch._ops.OpOverload) + ): + args_schema = pn.target._schema.arguments + + def get_all_arguments(orig_args, orig_kwargs): + all_args = [] + for i, schema in enumerate(args_schema): + if schema.name in orig_kwargs: + all_args.append(orig_kwargs[schema.name]) + elif not schema.kwarg_only and i < len(orig_args): + all_args.append(orig_args[i]) + else: + all_args.append(schema.default_value) + return all_args + + pn_args = get_all_arguments(pn.args, pn.kwargs) + gn_args = get_all_arguments(gn.args, gn.kwargs) + + elif len(pn.args) == len(gn.args) and list(pn.kwargs.keys()) == list( + gn.kwargs.keys() + ): + pn_args = list(pn.args) + gn_args = list(gn.args) + pn_args.extend(list(pn.kwargs.values())) + gn_args.extend(list(gn.kwargs.values())) + else: + match_found = False + + match_found = ( + match_found + and pn_args is not None + and gn_args is not None + and _match_args(pn_args, gn_args) + ) + + if not match_found: + # revert to saved_match before matching with current node + match = copy.copy(saved_match) + return False + + return True + + def match(self, graph: Graph) -> List[InternalMatch]: + """ + Returns: + The matched subgraphs. + Thre returned subgraph would be fully self-contained, meaning the nodes (except placeholder + and nodes returned by output) can only be consumed by nodes within the matched subgraph. + + Subgraph pattern matcher is implemented with the backtracking style in the following steps: + + 1. We first identify all the anchor nodes in the pattern graph. The anchor nodes + are the "sinks" (nodes with no user other than the output node) of the pattern graph. + One pattern graph could have multiple anchors if it has multiple return values. + + 2. In the target graph, we identify the potential candidate nodes that can be matched + with each anchor. These anchor-candidate pairs are the starting points for + pairwise per-node matching. + + 3. For each anchor-candidate pair, we simultaneously traverse backwards (DFS) in both + pattern and target graphs. For every pattern nodes along traversal path, we compare it + against the target nodes. In case any comparison failed, the match for this anchor-candidate + pair fails. A match is found when DFS completes traversing the graph. See `self._match_nodes` + for more details. + + 4. In the case of multiple anchors, every anchor will need to find a match using step 3. + In addition, the matches found between anchors need to have a common intersection node + in order for the match to be valid. This is implemented with backtracking. See `backtracking` + for more details. + + Notice: graph traversal must be done in the reverser order because a tensor can have multiple + consumers, but can only have a single producer. Only with reverser order, we can we jointly + traverse the pattern and target graph in a deterministic path. + + Warning: In theory, this backtracking algorithm have an **exponential** time complexity. However, + in practice, it's unlikely to blow up. + + """ + from torch.fx.passes.utils.fuser_utils import validate_partition + + # find candidate nodes to match with pattern anchors + match_candidates: Dict[Node, List[Node]] = defaultdict(list) + for pattern_anchor in self.pattern_anchors: + for node in graph.nodes: + if self._nodes_are_equal(pattern_anchor, node): + match_candidates[pattern_anchor].append(node) + match_candidates_list = list(match_candidates.items()) + + logger.info("Initial match_candidates_list: %s\n", match_candidates_list) + + matches: List[InternalMatch] = [] + + def backtracking(anchor_index, match): + if anchor_index == len(match_candidates_list): + match.placeholder_nodes = [ + match.nodes_map[pn] for pn in self.pattern_placeholder_nodes + ] + match.returning_nodes = [ + match.nodes_map[pn] for pn in self.pattern_returning_nodes + ] + matches.append(match) + + logger.info("Found a match: %s\n", match) + return + + pattern_anchor, candidate_nodes = match_candidates_list[anchor_index] + saved_match = copy.copy(match) + + for node in candidate_nodes: + logger.info("Trying to match anchor %s to %s", pattern_anchor, node) + + match_found = self._match_nodes(pattern_anchor, node, match) + if match_found: + # match next anchor + backtracking(anchor_index + 1, match) + else: + logger.info( + "Failed to match anchor %s to %s\n", pattern_anchor, node + ) + + # revert to saved_match before matching with current anchor + match = copy.copy(saved_match) + + match = InternalMatch(anchors=self.pattern_anchors) + if match_candidates_list: + backtracking(0, match) + + # filter out the matches where the subgraph is not fully_contained + before = len(matches) + matches = [match for match in matches if self._is_contained(match.nodes_map)] + after = len(matches) + if before != after: + logger.info( + "Filtered out %s matches because they are not fully contained", + before - after, + ) + + # filter out the matches that form a cycle if the subgraph is fused + valid_matches = [] + for match in matches: + matched_compute_nodes = [ + gn + for pn, gn in match.nodes_map.items() + if pn.op not in {"placeholder", "output"} + ] + if validate_partition(matched_compute_nodes): + valid_matches.append(match) + if len(valid_matches) != len(matches): + logger.info( + "Filtered out %s matches because \ + matched subgraph would form a cycle if fused", + len(matches) - len(valid_matches), + ) + + if self.remove_overlapping_matches: + before = len(valid_matches) + matches = self._remove_overlapping_matches(valid_matches) + after = len(matches) + if before != after: + logger.info( + "Filtered out %s matches because matched subgraphs are overlapping", + before - after, + ) + + logger.info("Matches returned: %s", matches) + + return matches diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/matcher_with_name_node_map_utils.py b/lib/python3.12/site-packages/torch/fx/passes/utils/matcher_with_name_node_map_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..78b063ff8a7aaa11b1dafcbdf2b5890ce13d2f03 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/utils/matcher_with_name_node_map_utils.py @@ -0,0 +1,114 @@ +from typing import Dict, List, Tuple + +from torch.fx import Graph, GraphModule, Node +from torch.fx._compatibility import compatibility + +from .matcher_utils import InternalMatch, SubgraphMatcher + + +__all__ = ["SubgraphMatcherWithNameNodeMap"] + + +def _split_to_graph_and_name_node_map( + gm: GraphModule, +) -> Tuple[GraphModule, Dict[str, Node]]: + from torch.fx.graph import _PyTreeInfo + from torch.utils._pytree import tree_flatten, tree_unflatten + + name_node_map = {} + for n in gm.graph.nodes: + if n.op == "output": + assert gm._out_spec is not None + output = tree_unflatten(n.args[0], gm._out_spec) + assert isinstance( + output, tuple + ), "Expecting the pattern graph to return a tuple" + assert ( + len(output) >= 2 + ), "Expecting the pattern graph to have at least two outputs" + *out, name_node_map = output + flattened, out_spec = tree_flatten(out) + assert isinstance( + name_node_map, Dict + ), "Expecting the input graph to have a dict output as the last element" + n.args = (flattened,) + orig_pytree_info = gm._graph._codegen.pytree_info # type: ignore[attr-defined] + gm._graph._codegen.pytree_info = _PyTreeInfo( # type: ignore[attr-defined] + orig_pytree_info.orig_args, orig_pytree_info.in_spec, out_spec + ) + gm.recompile() + return gm, name_node_map + + +@compatibility(is_backward_compatible=False) +class SubgraphMatcherWithNameNodeMap(SubgraphMatcher): + """Extends SubgraphMatcher to support querying the matched subgraph nodes through node name, + this requires pattern to have specific format (returning and additional dictionary at the output, + that has node name as key, and the node in the pattern graph as value, see Example for more details) + + Difference with SubgraphMatcher is that it takes a `pattern_gm` GraphModule as input during + initialization since we need to modify the graph (which requires `recompile` the GraphModule) + + Example:: + def pattern(x, weight): + conv = F.conv2d(x, weight) + relu = F.relu(conv) + return relu, {"conv": conv, "relu": relu} + + def target_graph(x, weight): + conv = F.conv2d(x, weight) + relu = F.relu(conv) + relu *= 2 + return relu + + pattern_gm = export_for_training(pattern, example_inputs).module() + target_gm = export_for_training(target_graph, example_inputs).module() + matcher = SubgraphMatcherWithNameNodeMap(pattern_gm) + matches = matcher.match(target_gm) + for match in matches: + match.name_node_map["conv"].meta["annotation"] = ... + + """ + + def __init__( + self, + pattern_gm: GraphModule, + match_output: bool = False, + match_placeholder: bool = False, + remove_overlapping_matches: bool = True, + ignore_literals: bool = False, + ) -> None: + pattern_gm, name_node_map = _split_to_graph_and_name_node_map(pattern_gm) + self.name_node_map = name_node_map + super().__init__( + pattern_gm.graph, + match_output, + match_placeholder, + remove_overlapping_matches, + ignore_literals, + ) + + def match(self, graph: Graph) -> List[InternalMatch]: + """The returned InternalMatch will have name_node_map populated with a map + from node name (str) to the target node, e.g. + {"conv": target_conv_ndoe, "relu": target_relu_node} + + this requires the pattern graph returns an additional + output of node name to node, e.g. instead of: + ``` + def pattern(...): + ... + return relu + ``` + we should do: + ``` + def pattern(...): + ... + return relu, {"conv": conv, "relu": relu} + ``` instead + """ + internal_matches = super().match(graph) + for internal_match in internal_matches: + for k, n in self.name_node_map.items(): + internal_match.name_node_map[k] = internal_match.nodes_map[n] + return internal_matches diff --git a/lib/python3.12/site-packages/torch/fx/passes/utils/source_matcher_utils.py b/lib/python3.12/site-packages/torch/fx/passes/utils/source_matcher_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f77db98880b76683cecb11e519ad70a1cfad7d45 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/passes/utils/source_matcher_utils.py @@ -0,0 +1,159 @@ +import logging +import os +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Type + +from torch.fx._compatibility import compatibility +from torch.fx.graph import Graph +from torch.fx.node import Node + + +__all__ = ["get_source_partitions", "check_subgraphs_connected", "SourcePartition"] + + +# Set`PYTORCH_MATCHER_LOGLEVEL=INFO` to see debug logs +def _init_logger() -> logging.Logger: + logger = logging.getLogger(__name__) + + level = os.environ.get("PYTORCH_MATCHER_LOGLEVEL", "WARNING").upper() + logger.setLevel(level) + console = logging.StreamHandler() + formatter = logging.Formatter("%(filename)s > %(message)s") + console.setFormatter(formatter) + console.setLevel(level) + # add the handlers to the logger + logger.addHandler(console) + logger.propagate = False + return logger + + +logger = _init_logger() + + +@compatibility(is_backward_compatible=False) +@dataclass +class SourcePartition: + # Nodes in a particular partition + nodes: List[Node] + + # The source these nodes decomposed from + source: Any + + # Nodes in the graph that are needed as inputs to the partition + input_nodes: List[Node] = field(default_factory=list) + + # Nodes in the partition that are being used by nodes outside of the + # partition + output_nodes: List[Node] = field(default_factory=list) + + # Parameters that are being used + params: List[Node] = field(default_factory=list) + + +@compatibility(is_backward_compatible=False) # type: ignore[misc] +def get_source_partitions( + graph: Graph, + wanted_sources: List[Any], + filter_fn: Optional[Callable[[Node], bool]] = None, +) -> Dict[Any, List[SourcePartition]]: + """ + Args: + graph: The graph we want to partition + wanted_sources: List of sources of nodes that were decomposed from this + source. This can be a function (ex. torch.nn.functional.linear) or a + leaf module type (ex. torch.nn.Linear). + + Returns: + Dictionary mapping sources that were given to a list of SourcePartitions + that correspond to the list of nodes that were decomposed from the given + source. + """ + modules: Dict[Type, Dict[str, List[Node]]] = {} + + for node in graph.nodes: + # The metadata source_fn should contain a tuple of a unique name for the + # source, and the source function if the node is decomposed from a + # function, or the type of module if the node is decomposed from a leaf + # module + + # TODO: Bypass "torch_fn" when "source_fn_stack" because now "torch_fn" can + # be different from "source_fn_stack", for example for the add_ node + # decomposed from batch norm. We should remove the check on "source_fn_stack" + # after we fix "torch_fn". T199561090 + if (source_fn_st := node.meta.get("source_fn_stack", None)) is None and ( + torch_fn := node.meta.get("torch_fn", None) + ) is not None: + node_fqn, source_fn = torch_fn + source_fn_name = source_fn.split(".")[1] + if source_fn_name in wanted_sources: + diff_modules = modules.setdefault(source_fn_name, {}) + partition = diff_modules.setdefault(node_fqn, []) + partition.append(node) + + if (source_fn_st := node.meta.get("source_fn_stack", None)) is not None: + source_fn = source_fn_st[-1] + if source_fn[1] in wanted_sources: + diff_modules = modules.setdefault(source_fn[1], {}) + partition = diff_modules.setdefault(source_fn[0], []) + partition.append(node) + + def make_partition(nodes: List[Node], module_type: Type) -> SourcePartition: + input_nodes = set() + output_nodes = set() + params = set() + for node in nodes: + for arg in node.args: + if isinstance(arg, Node) and arg not in nodes: + input_nodes.add(arg) + + if node.op == "get_attr": + params.add(node) + + for user in node.users.keys(): + if user not in nodes: + output_nodes.add(node) + + return SourcePartition( + nodes, + module_type, + list(input_nodes), + list(output_nodes), + list(params), # type: ignore[arg-type] + ) + + ret: Dict[Type[Any], List[SourcePartition]] = {} + + if filter_fn: + # for each partition, we apply filter_fn to filter out all partitions that doesn't satisfy the + # filter condition + filtered_modules = {} + for tp, name_to_partition in modules.items(): + filtered_name_to_partition = { + name: partition + for name, partition in name_to_partition.items() + if all(map(filter_fn, partition)) + } + filtered_modules[tp] = filtered_name_to_partition + modules = filtered_modules + + for k, v in modules.items(): + ret[k] = [make_partition(partition, k) for partition in v.values()] + + return ret + + +@compatibility(is_backward_compatible=False) # type: ignore[misc] +def check_subgraphs_connected( + subgraph1: SourcePartition, subgraph2: SourcePartition +) -> bool: + """ + Given two subgraphs A and B (in the form of a list of nodes), checks if + A has nodes connecting to at least one node in B -- aka there exists a node + in B that uses a node in A (not the other way around). + """ + + for node in reversed(subgraph1.nodes): + for user in node.users.keys(): + if user in subgraph2.nodes: + return True + return False diff --git a/lib/python3.12/site-packages/torch/fx/proxy.py b/lib/python3.12/site-packages/torch/fx/proxy.py new file mode 100644 index 0000000000000000000000000000000000000000..7268a7363336ccbce776f8c1524d2a958d0f653f --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/proxy.py @@ -0,0 +1,733 @@ +# mypy: ignore-errors + +import collections +import copy +import dis +import enum +import inspect +import logging +import operator +import sys +from dataclasses import fields, is_dataclass +from typing import Any, Callable, Dict, Iterator, Optional, OrderedDict, Tuple + +import torch +import torch.fx.traceback as fx_traceback +from torch.utils._traceback import CapturedTraceback + +from ._compatibility import compatibility +from .graph import Graph, magic_methods, reflectable_magic_methods +from .node import Argument, base_types, map_aggregate, Node, Target +from .operator_schemas import check_for_mutable_operation + + +__all__ = [ + "TracerBase", + "GraphAppendingTracer", + "TraceError", + "Proxy", + "MetaProxy", + "Attribute", + "ParameterProxy", + "Scope", + "ScopeContextManager", +] + + +log = logging.getLogger(__name__) + + +@compatibility(is_backward_compatible=False) +class Scope: + """Scope object that records the module path and the module type + of a module. Scope is used to track the information of the module + that contains a Node in a Graph of GraphModule. For example:: + + class Sub(torch.nn.Module): + def forward(self, x): + # This will be a call_method Node in GraphModule, + # scope for this would be (module_path="sub", module_type=Sub) + return x.transpose(1, 2) + + + class M(torch.nn.Module): + def __init__(self) -> None: + self.sub = Sub() + + def forward(self, x): + # This will be a call_method Node as well, + # scope for this would be (module_path="", None) + x = x.transpose(1, 2) + x = self.sub(x) + return x + + """ + + def __init__(self, module_path: str, module_type: Any): + super().__init__() + self.module_path = module_path + self.module_type = module_type + + +@compatibility(is_backward_compatible=False) +class ScopeContextManager: + """A context manager to track the Scope of Node during symbolic tracing. + When entering a forward function of a Module, we'll update the scope information of + the current module, and when we exit, we'll restore the previous scope information. + """ + + def __init__( + self, + scope: Scope, + current_scope: Scope, + ): + super().__init__() + # Keep a copy of prev scope to restore on exit + self._prev_scope = copy.copy(scope) + # Update scope to current scope + scope.module_path = current_scope.module_path + scope.module_type = current_scope.module_type + # Save a reference so we can restore it + self._scope = scope + + def __enter__(self): + return self._scope + + def __exit__(self, *args): + self._scope.module_path = self._prev_scope.module_path + self._scope.module_type = self._prev_scope.module_type + return + + +_COPY_META_FIELDS = [ + "nn_module_stack", + "torch_fn", + "source_fn_stack", + "original_aten", + "recompute", + "ac_graph_id", + "has_backward_hook", + "from_node", + "quantization_tag", # TODO deprecated + "_numeric_debug_handle", # TODO deprecated + "custom", + "partitioner_tag", +] + + +@compatibility(is_backward_compatible=True) +class TracerBase: + graph: Graph + record_stack_traces: bool = False + # Feature flag for mutable schema checking + # Enableby default in 1.12 + check_mutable_operations: bool = False + # Feature flag for assert tracing + trace_asserts: bool = False + # Feature flag for proxying accesses to buffer values + proxy_buffer_attributes: bool = False + + # Name of the function to be traced. It will only be used when + # ``root`` is an instance of ``nn.Module`` + traced_func_name: str = "forward" + + # Maps the containing module's name to the operator name + scope: Scope + + # Records the module call stack + module_stack: OrderedDict[str, Tuple[str, Any]] + + # Mapping of node name to module scope + node_name_to_scope: Dict[str, Tuple[str, type]] + + @compatibility(is_backward_compatible=True) + def create_node( + self, + kind: str, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: Optional[str] = None, + type_expr: Optional[Any] = None, + ) -> Node: + """ + Inserts a graph node given target, args, kwargs, and name. + + This method can be overridden to do extra checking, validation, or + modification of values used in node creation. For example, one might + want to disallow in-place operations from being recorded. + """ + + if kind == "call_function" and self.check_mutable_operations: + check_for_mutable_operation(target, args, kwargs) + + node = self.graph.create_node(kind, target, args, kwargs, name, type_expr) + # TODO node_name_to_scope will be depreciated in favor of + # node.meta['nn_module_stack'] + self.node_name_to_scope[node.name] = ( + self.scope.module_path, + self.scope.module_type, + ) + + # Optionally set stack trace on the created Node for debugging purposes + if fx_traceback.has_preserved_node_meta(): + current_meta: Dict[str, Any] = fx_traceback.get_current_meta() + + stack_trace = current_meta.get("stack_trace") + if stack_trace: + node.stack_trace = stack_trace + # Explicitly set the stack_trace, nn_module_stack and source_fn on the node.meta + # If other meta fields are needed, they can be added here + for field in _COPY_META_FIELDS: + if field in current_meta: + node.meta[field] = copy.copy(current_meta[field]) + + # Here we decrement to account for the sequence_nr having + # just been incremented while tracing this lowered aten op. + new_seq_nr = torch.autograd._get_sequence_nr() - 1 + # The sequence_nr increments every time a new autograd Node + # is created. During the FWD pass we store the sequence_nr + # corresponding to the last autograd Node created on this fx + # node's meta. A single aten op can create multiple autograd + # nodes as is the case with in-place foreach ops. During the + # BWD pass we retrieve the sequence_nr stored on the current + # executing autograd Node. See NOTE [ Sequence Number ]. + if current_meta.get("in_grad_fn", 0) > 0: + new_seq_nr = current_meta["grad_fn_seq_nr"][-1] + node.meta["seq_nr"] = new_seq_nr + + elif self.module_stack: + node.meta["nn_module_stack"] = copy.copy(self.module_stack) + + log.debug("create_node %s", node) + return node + + @compatibility(is_backward_compatible=True) + def proxy(self, node: Node) -> "Proxy": + return Proxy(node, self) + + @compatibility(is_backward_compatible=True) + def create_proxy( + self, + kind: str, + target: Target, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + name: Optional[str] = None, + type_expr: Optional[Any] = None, + # fix noqa when updating bc tests + proxy_factory_fn: Callable[[Node], "Proxy"] = None, # noqa: RUF013 + ): + """ + Create a Node from the given arguments, then return the Node + wrapped in a Proxy object. + + If kind = 'placeholder', then we're creating a Node that + represents the parameter of a function. If we need to encode + a default parameter, we use the ``args`` tuple. ``args`` is + otherwise empty for ``placeholder`` Nodes. + """ + + args_ = self.create_arg(args) + kwargs_ = self.create_arg(kwargs) + assert isinstance(args_, tuple) + assert isinstance(kwargs_, dict) + + node = self.create_node(kind, target, args_, kwargs_, name, type_expr) + + if not proxy_factory_fn: + proxy = self.proxy(node) + else: + proxy = proxy_factory_fn(node) + + if self.record_stack_traces and not proxy.node.stack_trace: + proxy.node.stack_trace = "".join(CapturedTraceback.extract().format()) + + return proxy + + def _find_user_frame(self): + """ + Find the Python stack frame executing the user code during + symbolic tracing. + """ + # We have to do a little dance here. Basically, walk up the callstack and + # record the first frame not in the pytorch source. This is the frame executing + # the user code during tracing. + frame = inspect.currentframe() + + pt_files = [ + "torch/fx/proxy.py", + "torch/fx/_symbolic_trace.py", + "torch/fx/experimental/proxy_tensor.py", + "torch/_ops.py", + "torch/_tensor.py", + "torch/utils/_python_dispatch.py", + "torch/_prims_common/wrappers.py", + "torch/_refs/__init__.py", + "torch/_refs/nn/functional/__init__.py", + "torch/utils/_stats.py", + ] + while frame: + frame = frame.f_back + if frame and all( + not frame.f_code.co_filename.endswith(file) for file in pt_files + ): + break + + if not frame: + return None + + return frame + + @compatibility(is_backward_compatible=True) + def create_arg(self, a: Any) -> Argument: + """ + A method that lowers the objects seen as arguments during symbolic evaluation + into Argument types that can be stored in IR. + + Can be override to support more trace-specific types. + """ + if isinstance(a, Proxy): + return a.node # most common arg type goes first + elif hasattr(a, "__fx_create_arg__"): + return a.__fx_create_arg__(self) + # aggregates + elif isinstance(a, tuple): + if hasattr(a, "_fields"): + # NamedTuple constructors don't seem to like getting a generator + # expression as an argument to their constructor, so build this + # intermediate tuple and unpack it into the NamedTuple constructor + args = [self.create_arg(elem) for elem in a] + return type(a)(*args) # type: ignore[arg-type] + return type(a)([self.create_arg(elem) for elem in a]) + elif isinstance(a, list): + return [self.create_arg(elem) for elem in a] + elif isinstance(a, dict): + + def no_node(arg): + if isinstance(arg, Node): + raise RuntimeError( + "Keys for dictionaries used as an argument cannot contain a " + f"Node. Got key: {k}" + ) + + r = {} + for k, v in a.items(): + # Check for invalid dict keys. We do not want a Proxy to appear + # anywhere within the key. Since keys can be collection types, + # we iterate through the key with map_aggregate + k = self.create_arg(k) + map_aggregate(k, no_node) + + r[k] = self.create_arg(v) + return r + elif isinstance(a, slice): + return slice( + self.create_arg(a.start), + self.create_arg(a.stop), + self.create_arg(a.step), + ) + + elif isinstance(a, range): + return range( + self.create_arg(a.start), + self.create_arg(a.stop), + self.create_arg(a.step), + ) + + elif isinstance(a, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)): + return a + + elif is_dataclass(a): + kwargs = { + field.name: self.create_arg(getattr(a, field.name)) + for field in fields(a) + } + return self.create_node("call_function", a.__class__, (), kwargs) + + elif isinstance(a, (*base_types, enum.Enum)) or a is None or a is ...: + return a + + raise NotImplementedError(f"argument of type: {type(a)}") + + @compatibility(is_backward_compatible=True) + def to_bool(self, obj: "Proxy") -> bool: + """Called when a proxy object is being converted to a boolean, such as + when used in control flow. Normally we don't know what to do because + we don't know the value of the proxy, but a custom tracer can attach more + information to the graph node using create_node and can choose to return a value. + """ + raise TraceError( + "symbolically traced variables cannot be used as inputs to control flow" + ) + + @compatibility(is_backward_compatible=True) + def iter(self, obj: "Proxy") -> Iterator: + """Called when a proxy object is being iterated over, such as + when used in control flow. Normally we don't know what to do because + we don't know the value of the proxy, but a custom tracer can attach more + information to the graph node using create_node and can choose to return an iterator. + """ + raise TraceError( + "Proxy object cannot be iterated. This can be " + "attempted when the Proxy is used in a loop or" + " as a *args or **kwargs function argument. " + "See the torch.fx docs on pytorch.org for a " + "more detailed explanation of what types of " + "control flow can be traced, and check out the" + " Proxy docstring for help troubleshooting " + "Proxy iteration errors" + ) + + @compatibility(is_backward_compatible=True) + def keys(self, obj: "Proxy") -> Any: + """Called when a proxy object is has the keys() method called. + This is what happens when ** is called on a proxy. This should return an + iterator it ** is suppose to work in your custom tracer. + """ + return Attribute(obj, "keys")() + + +# used in Proxy object when just appending to the graph while not tracing. +@compatibility(is_backward_compatible=True) +class GraphAppendingTracer(TracerBase): + def __init__(self, graph: Graph): + super().__init__() + self.graph = graph + self.scope = Scope("", None) + self.module_stack = collections.OrderedDict() + self.node_name_to_scope = {} + + +@compatibility(is_backward_compatible=False) +def assert_fn(x): + assert x + + +@compatibility(is_backward_compatible=True) +class TraceError(ValueError): + pass + + +@compatibility(is_backward_compatible=True) +class Proxy: + """ + ``Proxy`` objects are ``Node`` wrappers that flow through the + program during symbolic tracing and record all the operations + (``torch`` function calls, method calls, operators) that they touch + into the growing FX Graph. + + If you're doing graph transforms, you can wrap your own ``Proxy`` + method around a raw ``Node`` so that you can use the overloaded + operators to add additional things to a ``Graph``. + + ``Proxy`` objects cannot be iterated. In other words, the symbolic + tracer will throw an error if a ``Proxy`` is used in a loop or as + an ``*args``/``**kwargs`` function argument. + + There are two main ways around this: + 1. Factor out the untraceable logic into a top-level function and + use ``fx.wrap`` on it. + 2. If the control flow is static (i.e. the loop trip count is + based on some hyperparameter), the code can be kept in its original + position and refactored into something like:: + + for i in range(self.some_hyperparameter): + indexed_item = proxied_value[i] + + For a more detailed description into the Proxy internals, check out + the "Proxy" section in `torch/fx/README.md` + """ + + @compatibility(is_backward_compatible=True) + def __init__(self, node: Node, tracer: "Optional[TracerBase]" = None): + if tracer is None: + # This allows you to create a Proxy object around a raw Node + tracer = GraphAppendingTracer(node.graph) + self.tracer = tracer + self.node = node + + def __repr__(self) -> str: + return f"Proxy({self.node.name})" + + def __getattr__(self, k) -> "Attribute": + # note: not added to the graph yet, if this is a method call + # we peephole optimize to the method invocation + return Attribute(self, k) + + def __getstate__(self) -> Dict: + return self.__dict__ + + def __deepcopy__(self, memo) -> Dict: + # We have to explicitly override this method, because otherwise deepcopy + # will go to __getattr__(self, "__deepcopy__") and return a + # Attribute(__deepcopy__), and may go into an infinite loop in some cases. + import copy + + new_dict = {} + for k, v in self.__dict__.items(): + try: + new_obj = copy.deepcopy(v, memo) + except Exception: + log.warning( + "Shallow copy %s of Proxy because it cannot be deepcopied. " + "Proxy is created for node %s", + k, + self.node.name, + ) + new_obj = copy.copy(v) + new_dict[k] = new_obj + assert "node" in new_dict + assert "tracer" in new_dict + new_proxy = Proxy(new_dict["node"], new_dict["tracer"]) + for k, v in new_dict.items(): + new_proxy.__dict__[k] = v + return new_proxy + + def __setstate__(self, d): + # This is called when being unpickled/loaded. + self.__dict__ = d + + def __call__(self, *args, **kwargs) -> "Proxy": + return self.tracer.create_proxy( + "call_method", "__call__", (self,) + args, kwargs + ) + + def __iter__(self) -> Iterator["Proxy"]: + frame = inspect.currentframe() + assert frame is not None + calling_frame = frame.f_back + assert calling_frame is not None + inst_list = list(dis.get_instructions(calling_frame.f_code)) + if sys.version_info >= (3, 11): + from bisect import bisect_left + + inst_idx = bisect_left( + inst_list, calling_frame.f_lasti, key=lambda x: x.offset + ) + else: + inst_idx = calling_frame.f_lasti // 2 + inst = inst_list[inst_idx] + if inst.opname == "UNPACK_SEQUENCE": + return (self[i] for i in range(inst.argval)) # type: ignore[index] + + return self.tracer.iter(self) + + def __abs__(self): + return self.tracer.create_proxy("call_function", operator.abs, (self,), {}) + + def __bool__(self) -> bool: + if self.tracer.trace_asserts: + # check if this boolean is used in an assertion, bytecode pattern for assertions + # is pretty stable for Python 3.7--3.9 + frame = inspect.currentframe() + assert frame is not None + calling_frame = frame.f_back + assert calling_frame is not None + insts = list(dis.get_instructions(calling_frame.f_code)) + if sys.version_info >= (3, 11): + from bisect import bisect_left + + cur = bisect_left(insts, calling_frame.f_lasti, key=lambda x: x.offset) + else: + cur = calling_frame.f_lasti // 2 + inst = insts[cur] + + if inst.opname == "POP_JUMP_IF_TRUE": + first = insts[cur + 1] + assert inst.arg is not None + last = insts[inst.arg // 2 - 1] + starts_with_assert = ( + first.opname == "LOAD_GLOBAL" + and first.argval == "AssertionError" + or first.opname == "LOAD_ASSERTION_ERROR" + ) + if starts_with_assert and last.opname == "RAISE_VARARGS": + self.tracer.create_proxy("call_function", assert_fn, (self,), {}) + return True + + return self.tracer.to_bool(self) + + @compatibility(is_backward_compatible=True) + def keys(self): + return self.tracer.keys(self) + + def __len__(self): + raise RuntimeError( + "'len' is not supported in symbolic tracing by default. If you want " + "this call to be recorded, please call torch.fx.wrap('len') at " + "module scope" + ) + + @classmethod + def __torch_function__(cls, orig_method, types, args=None, kwargs=None): + args = args if args else () + kwargs = kwargs if kwargs else {} + + tracers: Dict[Any, None] = {} + + def find_tracer(a): + if isinstance(a, cls): + tracers[a.tracer] = None + + torch.fx.node.map_aggregate(args, find_tracer) + torch.fx.node.map_aggregate(kwargs, find_tracer) + + if len(tracers) > 1: + raise RuntimeError( + f"Found multiple different tracers {list(tracers.keys())} while " + f"trying to trace operations {orig_method}" + ) + tracer = next(iter(tracers.keys())) + + if isinstance(orig_method, torch._C.ScriptMethod): + args = (orig_method.owner,) + args + return tracer.create_proxy("call_method", orig_method.name, args, kwargs) + if torch.overrides.is_tensor_method_or_property(orig_method): + return tracer.create_proxy( + "call_method", orig_method.__name__, args, kwargs + ) + else: + if isinstance(orig_method, torch._ops.HigherOrderOperator): + # TODO: Define how to symbolically trace HigherOrderOperators + raise RuntimeError("Unable to symbolically trace HigherOrderOperators") + return tracer.create_proxy( + "call_function", + orig_method, + args, + kwargs, + name=tracer.graph._target_to_str(orig_method.__name__), + ) + + +@compatibility(is_backward_compatible=False) +class MetaProxy(Proxy): + """ + A Proxy subclass that propagates metadata (meta['val']) during graph tracing. + """ + + def __init__( + self, node: Node, tracer: "Optional[TracerBase]" = None, fake_mode=None + ): + super().__init__(node, tracer) + self.fake_mode = fake_mode + + def __repr__(self) -> str: + return f"MetaProxy({self.node.name})" + + @classmethod + def __torch_function__(cls, orig_method, types, args=None, kwargs=None): + args = args if args else () + kwargs = kwargs if kwargs else {} + + meta_proxy = None + for arg in args: + if isinstance(arg, MetaProxy): + meta_proxy = arg + break + + assert ( + meta_proxy is not None + ), "No MetaProxy found in arguments, but one is expected." + + proxy = super().__torch_function__(orig_method, types, args, kwargs) + with meta_proxy.fake_mode: + proxy.node.meta["val"] = orig_method( + *[a.node.meta["val"] if isinstance(a, Proxy) else a for a in args], + **kwargs, + ) + return MetaProxy(proxy.node, proxy.tracer, meta_proxy.fake_mode) + + +@compatibility(is_backward_compatible=True) +class Attribute(Proxy): + @compatibility(is_backward_compatible=True) + def __init__(self, root: Proxy, attr: str): + self.root = root + self.attr = attr + self.tracer = root.tracer + self._node: Optional[Node] = None + + @property + def node(self): + # the node for attributes is added lazily, since most will just be method calls + # which do not rely on the getitem call + if self._node is None: + self._node = self.tracer.create_proxy( + "call_function", getattr, (self.root, self.attr), {} + ).node + return self._node + + def __call__(self, *args, **kwargs): + return self.tracer.create_proxy( + "call_method", self.attr, (self.root,) + args, kwargs + ) + + +@compatibility(is_backward_compatible=False) +class ParameterProxy(Proxy): + """ + A special proxy which lets "shape", "size", "dim", and a few other + attribute accesses pass through to the underlying module parameter object, + so that conditional tests on these attributes will not throw exception during tracing + """ + + def __init__(self, tracer: TracerBase, node: Node, name, param): + super().__init__(node, tracer) + assert isinstance(param, torch.nn.Parameter) + self.param = param + self.name = name + + def __repr__(self) -> str: + return f"ParameterProxy({self.name})" + + @property + def shape(self): + return self.param.shape + + def size(self): + return self.param.size() + + def dim(self): + return self.param.dim() + + @property + def ndim(self): + return self.param.ndim + + def numel(self): + return self.param.numel() + + def nelement(self): + return self.param.nelement() + + +for method in magic_methods: + + def _scope(method): + def impl(*args, **kwargs): + tracer = args[0].tracer + target = getattr(operator, method) + return tracer.create_proxy("call_function", target, args, kwargs) + + impl.__name__ = method + as_magic = f'__{method.strip("_")}__' + setattr(Proxy, as_magic, impl) + + _scope(method) + + +def _define_reflectable(orig_method_name): + method_name = f'__r{orig_method_name.strip("_")}__' + + def impl(self, rhs): + target = getattr(operator, orig_method_name) + return self.tracer.create_proxy("call_function", target, (rhs, self), {}) + + impl.__name__ = method_name + impl.__qualname__ = method_name + setattr(Proxy, method_name, impl) + + +for orig_method_name in reflectable_magic_methods: + _define_reflectable(orig_method_name) diff --git a/lib/python3.12/site-packages/torch/fx/subgraph_rewriter.py b/lib/python3.12/site-packages/torch/fx/subgraph_rewriter.py new file mode 100644 index 0000000000000000000000000000000000000000..b823fda3123fa41989dcf3a8a36c844e2a433ce7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/subgraph_rewriter.py @@ -0,0 +1,421 @@ +import copy +from dataclasses import dataclass +from typing import ( + Any, + Callable, + Dict, + List, + NamedTuple, + Optional, + Set, + TYPE_CHECKING, + Union, +) + +import torch + +from ._compatibility import compatibility +from ._symbolic_trace import symbolic_trace +from .graph import Graph +from .graph_module import GraphModule +from .node import Node + + +if TYPE_CHECKING: + from .passes.utils.matcher_with_name_node_map_utils import InternalMatch + +__all__ = [ + "Match", + "replace_pattern", + "replace_pattern_with_filters", + "ReplacedPatterns", +] + + +@compatibility(is_backward_compatible=True) +class Match(NamedTuple): + # Node from which the match was found + anchor: Node + # Maps nodes in the pattern subgraph to nodes in the larger graph + nodes_map: Dict[Node, Node] + + +@compatibility(is_backward_compatible=False) +@dataclass +class ReplacedPatterns: + # Node from which the match was found + anchor: Node + # Maps nodes in the pattern subgraph to nodes in the larger graph + nodes_map: Dict[Node, Node] + # List of nodes that were added into the graph + replacements: List[Node] + + +def _replace_attributes(gm: GraphModule, replacement: torch.nn.Module) -> None: + gm.delete_all_unused_submodules() + + if isinstance(replacement, GraphModule): + replacement.graph.lint() + + def try_get_attr(gm: torch.nn.Module, target: str) -> Optional[Any]: + module_path, _, attr_name = target.rpartition(".") + try: + mod: torch.nn.Module = gm.get_submodule(module_path) + except AttributeError: + return None + attr = getattr(mod, attr_name, None) + return attr + + for node in gm.graph.nodes: + if node.op == "call_module" or node.op == "get_attr": + gm_attr = try_get_attr(gm, node.target) + replacement_attr = try_get_attr(replacement, node.target) + + # CASE 1: This target already exists as an attribute in our + # result GraphModule. Whether or not it exists in + # `replacement`, the existing submodule takes precedence. + if gm_attr is not None: + continue + + # CASE 2: The target exists as an attribute in `replacement` + # only, so we need to copy it over. + elif replacement_attr is not None: + new_attr = copy.deepcopy(replacement_attr) + if isinstance(replacement_attr, torch.nn.Module): + gm.add_submodule(node.target, new_attr) + else: + setattr(gm, node.target, new_attr) + + # CASE 3: The target doesn't exist as an attribute in `gm` + # or `replacement` + else: + raise RuntimeError( + 'Attempted to create a "', + node.op, + '" node during subgraph rewriting ' + f"with target {node.target}, but " + "the referenced attribute does not " + "exist in the replacement GraphModule", + ) + + gm.graph.lint() + + +@compatibility(is_backward_compatible=True) +def replace_pattern( + gm: GraphModule, + pattern: Union[Callable, GraphModule], + replacement: Union[Callable, GraphModule], +) -> List[Match]: + """ + Matches all possible non-overlapping sets of operators and their + data dependencies (``pattern``) in the Graph of a GraphModule + (``gm``), then replaces each of these matched subgraphs with another + subgraph (``replacement``). + + Args: + ``gm``: The GraphModule that wraps the Graph to operate on + ``pattern``: The subgraph to match in ``gm`` for replacement + ``replacement``: The subgraph to replace ``pattern`` with + + Returns: + List[Match]: A list of ``Match`` objects representing the places + in the original graph that ``pattern`` was matched to. The list + is empty if there are no matches. ``Match`` is defined as: + + .. code-block:: python + + class Match(NamedTuple): + # Node from which the match was found + anchor: Node + # Maps nodes in the pattern subgraph to nodes in the larger graph + nodes_map: Dict[Node, Node] + + Examples: + + .. code-block:: python + + import torch + from torch.fx import symbolic_trace, subgraph_rewriter + + + class M(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, w1, w2): + m1 = torch.cat([w1, w2]).sum() + m2 = torch.cat([w1, w2]).sum() + return x + torch.max(m1) + torch.max(m2) + + + def pattern(w1, w2): + return torch.cat([w1, w2]).sum() + + + def replacement(w1, w2): + return torch.stack([w1, w2]) + + + traced_module = symbolic_trace(M()) + + subgraph_rewriter.replace_pattern(traced_module, pattern, replacement) + + The above code will first match ``pattern`` in the ``forward`` + method of ``traced_module``. Pattern-matching is done based on + use-def relationships, not node names. For example, if you had + ``p = torch.cat([a, b])`` in ``pattern``, you could match + ``m = torch.cat([a, b])`` in the original ``forward`` function, + despite the variable names being different (``p`` vs ``m``). + + The ``return`` statement in ``pattern`` is matched based on its + value only; it may or may not match to the ``return`` statement in + the larger graph. In other words, the pattern doesn't have to extend + to the end of the larger graph. + + When the pattern is matched, it will be removed from the larger + function and replaced by ``replacement``. If there are multiple + matches for ``pattern`` in the larger function, each non-overlapping + match will be replaced. In the case of a match overlap, the first + found match in the set of overlapping matches will be replaced. + ("First" here being defined as the first in a topological ordering + of the Nodes' use-def relationships. In most cases, the first Node + is the parameter that appears directly after ``self``, while the + last Node is whatever the function returns.) + + One important thing to note is that the parameters of the + ``pattern`` Callable must be used in the Callable itself, + and the parameters of the ``replacement`` Callable must match + the pattern. The first rule is why, in the above code block, the + ``forward`` function has parameters ``x, w1, w2``, but the + ``pattern`` function only has parameters ``w1, w2``. ``pattern`` + doesn't use ``x``, so it shouldn't specify ``x`` as a parameter. + As an example of the second rule, consider replacing + + .. code-block:: python + + def pattern(x, y): + return torch.neg(x) + torch.relu(y) + + with + + .. code-block:: python + + def replacement(x, y): + return torch.relu(x) + + In this case, ``replacement`` needs the same number of parameters + as ``pattern`` (both ``x`` and ``y``), even though the parameter + ``y`` isn't used in ``replacement``. + + After calling ``subgraph_rewriter.replace_pattern``, the generated + Python code looks like this: + + .. code-block:: python + + def forward(self, x, w1, w2): + stack_1 = torch.stack([w1, w2]) + sum_1 = stack_1.sum() + stack_2 = torch.stack([w1, w2]) + sum_2 = stack_2.sum() + max_1 = torch.max(sum_1) + add_1 = x + max_1 + max_2 = torch.max(sum_2) + add_2 = add_1 + max_2 + return add_2 + """ + match_and_replacements = _replace_pattern(gm, pattern, replacement) + return [ + Match(anchor=m.anchor, nodes_map=m.nodes_map) for m in match_and_replacements + ] + + +# Experimental API, not backward compatible +@compatibility(is_backward_compatible=False) +def replace_pattern_with_filters( + gm: GraphModule, + pattern: Union[Callable, Graph, GraphModule], + replacement: Union[Callable, Graph, GraphModule, None] = None, + match_filters: Optional[ + List[Callable[["InternalMatch", Graph, Graph], bool]] + ] = None, + ignore_literals: bool = False, + # Placed at the end to avoid breaking backward compatibility + replacement_callback: Optional[ + Callable[["InternalMatch", Graph, Graph], Graph] + ] = None, +) -> List[ReplacedPatterns]: + """ + See replace_pattern for documentation. This function is an overload with an additional match_filter argument. + + Args: + ``match_filters``: A list of functions that take in + (match: InternalMatch, original_graph: Graph, pattern_graph: Graph) and return a boolean indicating + whether the match satisfies the condition. + See matcher_utils.py for definition of InternalMatch. + ``replacement_callback``: A function that takes in a match and returns a + Graph to be used as the replacement. This allows you to construct a + replacement graph based on the match. + """ + + return _replace_pattern( + gm, pattern, replacement, match_filters, ignore_literals, replacement_callback + ) + + +def _replace_pattern( + gm: GraphModule, + pattern: Union[Callable, Graph, GraphModule], + replacement: Union[Callable, Graph, GraphModule, None] = None, + match_filters: Optional[ + List[Callable[["InternalMatch", Graph, Graph], bool]] + ] = None, + ignore_literals: bool = False, + # Placed at the end to avoid breaking backward compatibility + replacement_callback: Optional[ + Callable[["InternalMatch", Graph, Graph], Graph] + ] = None, +) -> List[ReplacedPatterns]: + from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher + + if match_filters is None: + match_filters = [] + + # Get the graphs for `gm`, `pattern`, `replacement` + original_graph: Graph = gm.graph + + if isinstance(pattern, GraphModule): + pattern_graph = pattern.graph + elif isinstance(pattern, Graph): + pattern_graph = pattern + else: + pattern_graph = symbolic_trace(pattern).graph + + matcher = SubgraphMatcher( + pattern_graph, + match_output=False, + match_placeholder=False, + remove_overlapping_matches=True, + ignore_literals=ignore_literals, + ) + _matches: List[InternalMatch] = matcher.match(original_graph) + + # Filter out matches that don't match the filter + _matches = [ + m + for m in _matches + if all( + match_filter(m, original_graph, pattern_graph) + for match_filter in match_filters + ) + ] + + if isinstance(replacement, GraphModule): + common_replacement_graph = replacement.graph + elif isinstance(replacement, Graph): + common_replacement_graph = replacement + elif callable(replacement): + common_replacement_graph = symbolic_trace(replacement).graph + else: + assert ( + replacement_callback is not None + ), "Must provide either a replacement GraphModule or a replacement callback" + common_replacement_graph = None + + # As we progressively replace nodes, we'll need to keep track of how the match results should change + match_changed_node: Dict[Node, Node] = {} + + match_and_replacements = [] + for match in _matches: + if replacement_callback is not None: + replacement_graph = replacement_callback( + match, original_graph, pattern_graph + ) + else: + assert ( + common_replacement_graph is not None + ), "Must provide either a replacement GraphModule or a replacement callback" + replacement_graph = common_replacement_graph + replacement_placeholders = [ + n for n in replacement_graph.nodes if n.op == "placeholder" + ] + + # Build connecting between replacement graph's input and original graph input producer node + + # Initialize `val_map` with mappings from placeholder nodes in + # `replacement` to their corresponding node in `original_graph` + assert len(match.placeholder_nodes) == len(replacement_placeholders) + val_map: Dict[Node, Node] = {} + for rn, gn in zip(replacement_placeholders, match.placeholder_nodes): + if isinstance(gn, Node): + val_map[rn] = match_changed_node.get(gn, gn) + if gn != val_map[rn]: + # Update match.placeholder_nodes and match.nodes_map with the node that replaced gn + gn_ind = match.placeholder_nodes.index(gn) + match.placeholder_nodes[gn_ind] = match_changed_node[gn] + map_key = list(match.nodes_map.keys())[ + list(match.nodes_map.values()).index(gn) + ] + match.nodes_map[map_key] = match_changed_node[gn] + else: + val_map[rn] = gn + + # Copy the replacement graph over + user_nodes: Set[Node] = set() + for n in match.returning_nodes: + user_nodes.update(n.users) + assert user_nodes, "The returning_nodes should have at least one user node" + + if len(user_nodes) == 1: + first_user_node = next(iter(user_nodes)) + else: + # If there are multiple user nodes, we need to find the first user node + # in the current execution order of the `original_graph` + for n in original_graph.nodes: + if n in user_nodes: + first_user_node = n + break + + with original_graph.inserting_before(first_user_node): # type: ignore[possibly-undefined] + copied_returning_nodes = original_graph.graph_copy( + replacement_graph, val_map + ) + + if isinstance(copied_returning_nodes, Node): + copied_returning_nodes = (copied_returning_nodes,) + + # Get a list of nodes that have been replaced into the graph + replacement_nodes: List[Node] = [ + v for v in val_map.values() if v not in match.placeholder_nodes + ] + + # Hook the output Node of the replacement subgraph in to the + # original Graph at the correct location + assert len(match.returning_nodes) == len(copied_returning_nodes) # type: ignore[arg-type] + for gn, copied_node in zip(match.returning_nodes, copied_returning_nodes): # type: ignore[arg-type] + gn.replace_all_uses_with(copied_node) + match_changed_node[gn] = copied_node + # Remove the original nodes + for node in reversed(pattern_graph.nodes): + if node.op != "placeholder" and node.op != "output": + gn = match.nodes_map[node] + gm.graph.erase_node(gn) + + match_and_replacements.append( + ReplacedPatterns( + anchor=match.anchors[0], + nodes_map=match.nodes_map, + replacements=replacement_nodes, + ) + ) + + # Update the passed-in GraphModule to reflect the new state of + # `original_graph` + gm.recompile() + + # If `replacement` was an nn.Module, we'll need to make sure that + # all the submodules have been copied over correctly + if isinstance(replacement, torch.nn.Module): + _replace_attributes(gm, replacement) + + return match_and_replacements diff --git a/lib/python3.12/site-packages/torch/fx/tensor_type.py b/lib/python3.12/site-packages/torch/fx/tensor_type.py new file mode 100644 index 0000000000000000000000000000000000000000..4f375e461ef288341c5cd22e1e1ec2a851680b4c --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/tensor_type.py @@ -0,0 +1,111 @@ +# mypy: allow-untyped-defs +from torch.fx.experimental.unification import Var # type: ignore[attr-defined] + +from ._compatibility import compatibility + + +@compatibility(is_backward_compatible=False) +class TensorType: + """ + TensorType defines a type for tensors, which consists of a list of dimensions. + Example: + class M(torch.nn.Module): + def forward(self, x:TensorType((1,2,3, Dyn)), y:TensorType((1,2,3, Dyn))): + return torch.add(x, y) + """ + + def __init__(self, dim): + self.__origin__ = TensorType + self.__args__ = dim + + def __repr__(self): + return f"TensorType[{self.__args__}]" + + def __eq__(self, other): + if isinstance(other, self.__class__): + return list(self.__args__) == list(other.__args__) + else: + return False + + @staticmethod + def __class_getitem__(*args): + if len(args) == 1 and isinstance(args[0], tuple): + args = args[0] + return TensorType(tuple(args)) + + +class _DynType: + """ + _DynType defines a type which stands for the absence of type information. + """ + + def __init__(self) -> None: + self.__name__ = "_DynType" + + def __eq__(self, other): + return isinstance(other, self.__class__) + + def __str__(self): + return "Dyn" + + def __repr__(self): + return "Dyn" + + +Dyn = _DynType() + + +@compatibility(is_backward_compatible=False) +def is_consistent(t1, t2): + """ + A binary relation denoted by ~ that determines if t1 is consistent with t2. + The relation is reflexive, symmetric but not transitive. + returns True if t1 and t2 are consistent and False otherwise. + Example: + Dyn ~ TensorType((1,2,3)) + int ~ Dyn + int ~ int + TensorType((1,Dyn,3)) ~ TensorType((1,2,3)) + """ + + if t1 == t2: + return True + + if t1 == Dyn or t2 == Dyn or isinstance(t1, Var) or isinstance(t2, Var): + return True + + if isinstance(t1, TensorType) and isinstance(t2, TensorType): + return len(t1.__args__) == len(t2.__args__) and all( + is_consistent(elem1, elem2) + for elem1, elem2 in zip(t1.__args__, t2.__args__) + ) + else: + return False + + +@compatibility(is_backward_compatible=False) +def is_more_precise(t1, t2): + """ + A binary relation denoted by <= that determines if t1 is more precise than t2. + The relation is reflexive and transitive. + returns True if t1 is more precise than t2 and False otherwise. + Example: + Dyn >= TensorType((1,2,3)) + int >= Dyn + int >= int + TensorType((1,Dyn,3)) <= TensorType((1,2,3)) + """ + if t1 == t2: + return True + + if isinstance(t2, _DynType): + return True + + if isinstance(t1, TensorType) and isinstance(t2, TensorType): + return len(t1.__args__) == len(t2.__args__) and all( + is_more_precise(elem1, elem2) + for elem1, elem2 in zip(t1.__args__, t2.__args__) + ) + + else: + return False diff --git a/lib/python3.12/site-packages/torch/fx/traceback.py b/lib/python3.12/site-packages/torch/fx/traceback.py new file mode 100644 index 0000000000000000000000000000000000000000..2a4cccc419a2954c281f0798700264f9ce9810e7 --- /dev/null +++ b/lib/python3.12/site-packages/torch/fx/traceback.py @@ -0,0 +1,224 @@ +# mypy: allow-untyped-defs +import copy +import json +import traceback +from contextlib import contextmanager +from enum import Enum +from typing import Any, Dict, List, Optional + +from ._compatibility import compatibility +from .graph import Graph +from .node import Node + + +__all__ = [ + "preserve_node_meta", + "has_preserved_node_meta", + "set_stack_trace", + "set_grad_fn_seq_nr", + "reset_grad_fn_seq_nr", + "format_stack", + "set_current_meta", + "get_current_meta", + "NodeSource", + "NodeSourceAction", + "get_graph_provenance_json", +] + +current_meta: Dict[str, Any] = {} +should_preserve_node_meta = False + + +@compatibility(is_backward_compatible=False) +class NodeSourceAction(str, Enum): + CREATE = "create" + REPLACE = "replace" + + +@compatibility(is_backward_compatible=False) +class NodeSource: + """ + NodeSource is a data structure that contains the provenance information of a node. + If node `a` is created from node `b`, then `a.meta["from_node"]` may contain NodeSource(b). + """ + + class NodeInfo: + def __init__(self, name: str, target: str, graph_id: int): + self.name = name + self.target = target + self.graph_id = graph_id + + pass_name: str + action: Optional["NodeSourceAction"] + from_node: List["NodeSource"] + node_info: Optional["NodeInfo"] + + def __init__( + self, + node: Optional[Node], + pass_name: str = "", + action: Optional["NodeSourceAction"] = None, + ): + self.pass_name = pass_name + self.action = action + if node: + self.node_info = self.NodeInfo( + name=node.name, target=str(node.target), graph_id=id(node.graph) + ) + self.from_node = ( + copy.deepcopy(node.meta["from_node"]) + if "from_node" in node.meta + else [] + ) + else: + self.node_info = None + self.from_node = [] + + @property + def name(self) -> str: + return self.node_info.name if self.node_info else "" + + @property + def target(self) -> str: + return self.node_info.target if self.node_info else "" + + @property + def graph_id(self) -> int: + return self.node_info.graph_id if self.node_info else -1 + + def __repr__(self): + return self.print_readable() + + def print_readable(self, indent=0): + if indent > 9: + return "" + result = "" + result += ( + " " * indent * 4 + + f"(name={self.name}, pass_name={self.pass_name}, action={self.action}, graph_id={self.graph_id})\n" + ) + for item in self.from_node: + result += item.print_readable(indent + 1) + return result + + def to_dict(self) -> dict: + # Convert the object to a dictionary + return { + "name": self.name, + "target": self.target, + "graph_id": self.graph_id, + "pass_name": self.pass_name, + "action": self.action, + "from_node": [node.to_dict() for node in self.from_node], + } + + +@compatibility(is_backward_compatible=False) +@contextmanager +def preserve_node_meta(): + global should_preserve_node_meta + global current_meta + + saved_should_preserve_node_meta = should_preserve_node_meta + # Shallow copy is OK since fields of current_meta are not mutated + saved_current_meta = current_meta.copy() + try: + should_preserve_node_meta = True + yield + finally: + should_preserve_node_meta = saved_should_preserve_node_meta + current_meta = saved_current_meta + + +@compatibility(is_backward_compatible=False) +def set_stack_trace(stack: List[str]): + global current_meta + + if should_preserve_node_meta and stack: + current_meta["stack_trace"] = "".join(stack) + + +@compatibility(is_backward_compatible=False) +def set_grad_fn_seq_nr(seq_nr): + global current_meta + + if should_preserve_node_meta: + # The seq_nr is captured by eager mode in the grad_fn during forward + current_meta["grad_fn_seq_nr"] = current_meta.get("grad_fn_seq_nr", []) + [ + seq_nr + ] + current_meta["in_grad_fn"] = current_meta.get("in_grad_fn", 0) + 1 + + +@compatibility(is_backward_compatible=False) +def reset_grad_fn_seq_nr(): + # NB: reset state properly, this would be helpful towards supporting + # reentrant autograd if we actually wanted to do that. + global current_meta + if should_preserve_node_meta: + current_level = current_meta.get("in_grad_fn", 0) + assert current_level > 0 + if current_level == 1: + del current_meta["in_grad_fn"] + del current_meta["grad_fn_seq_nr"] + else: + current_meta["in_grad_fn"] = current_level - 1 + current_meta["grad_fn_seq_nr"] = current_meta["grad_fn_seq_nr"][:-1] + + +@compatibility(is_backward_compatible=False) +def format_stack() -> List[str]: + if should_preserve_node_meta: + return [current_meta.get("stack_trace", "")] + else: + # fallback to traceback.format_stack() + return traceback.format_list(traceback.extract_stack()[:-1]) + + +@compatibility(is_backward_compatible=False) +def has_preserved_node_meta() -> bool: + return should_preserve_node_meta + + +@compatibility(is_backward_compatible=False) +@contextmanager +def set_current_meta(node, pass_name=""): + global current_meta + if should_preserve_node_meta and node.meta: + saved_meta = current_meta + try: + current_meta = node.meta.copy() + + # Update the "from_node" field in current_meta for provenance tracking. + # Instead of appending, overwrite the "from_node" field because current_meta + # will be assigned to the new node. The new NodeSource(node, ...) will + # include the information from the previous current_meta["from_node"]. + current_meta["from_node"] = [ + NodeSource(node, pass_name, NodeSourceAction.CREATE) + ] + yield + finally: + current_meta = saved_meta + else: + yield + + +@compatibility(is_backward_compatible=False) +def get_current_meta() -> Dict[str, Any]: + return current_meta + + +@compatibility(is_backward_compatible=False) +def get_graph_provenance_json(graph: Graph) -> str: + """ + Given an fx.Graph, return a json string that contains the provenance information of each node. + """ + provenance_tracking_json = {} + for node in graph.nodes: + if node.op == "call_function": + provenance_tracking_json[node.name] = ( + [source.to_dict() for source in node.meta["from_node"]] + if "from_node" in node.meta + else [] + ) + return json.dumps(provenance_tracking_json)