[export] Expand verifier to be multiple on ExportedProgram (pytorch#130364)

Summary:
X-link: pytorch/executorch#4184

Pull Request resolved: pytorch#130364

This diff updates the ExportedProgram class in PyTorch to allow for multiple verifiers to be attached to it. This is done by adding a new field to the ExportedProgram schema called "verifiers" which is a list of strings representing the names of the verifiers to be attached to the program. The verifiers are loaded using the "load_verifier" function which is defined in the "torch._export.serde.serialize" module. The "exported_program.dialect" field is also deprecated in favor of the "verifiers" field.

Test Plan: CI

Differential Revision: D59408546

Author: zhxchen17

torch/_export/__init__.py

@@ -19,36 +19,13 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from unittest.mock import patch
 
-import sympy
-
 import torch
-import torch._dynamo
 import torch.fx
 import torch.utils._pytree as pytree
 
-from torch._decomp import core_aten_decompositions, get_decompositions
 from torch._dispatch.python import enable_python_dispatcher
-from torch._dynamo.exc import UserError, UserErrorType
-from torch._dynamo.source import ConstantSource
-from torch._export.non_strict_utils import make_constraints
-from torch._export.passes.collect_tracepoints_pass import CollectTracepointsPass
-from torch._functorch.aot_autograd import aot_export_module, GraphSignature
-from torch._functorch.eager_transforms import functionalize
-from torch._guards import detect_fake_mode
-from torch._inductor import config
-from torch._ops import OpOverload
-from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
-from torch._subclasses.functional_tensor import FunctionalTensor
 from torch._utils_internal import log_export_usage
 from torch.export._tree_utils import reorder_kwargs
-from torch.export._unlift import _create_stateful_graph_module
-from torch.export.dynamic_shapes import _combine_args, Constraint, dims, dynamic_dim
-from torch.export.exported_program import (
-    _disable_prexisiting_fake_mode,
-    ExportedProgram,
-    ModuleCallEntry,
-    ModuleCallSignature,
-)
 from torch.export.graph_signature import (
     _sig_to_specs,
     ArgumentSpec,
@@ -64,14 +41,7 @@
 from torch.fx import traceback as fx_traceback
 from torch.fx._compatibility import compatibility
 from torch.fx.experimental.proxy_tensor import make_fx, maybe_disable_fake_tensor_mode
-from torch.fx.experimental.symbolic_shapes import (
-    ConstraintViolationError,
-    GuardOnDataDependentSymNode,
-    ShapeEnv,
-    StrictMinMaxConstraint,
-)
 from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
-from torch.utils._sympy.value_ranges import ValueRangeError, ValueRanges
 
 from .wrappers import _wrap_submodules
 
@@ -89,6 +59,8 @@ class ExportDynamoConfig:
 # is called multiple times.
 @lru_cache
 def capture_pre_autograd_graph_warning():
+    from torch._inductor import config
+
     log.warning("+============================+")
     log.warning("|     !!!   WARNING   !!!    |")
     log.warning("+============================+")
@@ -138,6 +110,10 @@ def capture_pre_autograd_graph(
     """
     from torch.export._trace import _convert_input_to_fake, DEFAULT_EXPORT_DYNAMO_CONFIG, _ignore_backend_decomps
     from torch._utils_internal import export_api_rollout_check
+    from torch._export.non_strict_utils import make_constraints
+    from torch._subclasses.functional_tensor import FunctionalTensor
+    from torch.export._unlift import _create_stateful_graph_module
+    from torch.export.dynamic_shapes import _combine_args
 
     capture_pre_autograd_graph_warning()
 
@@ -232,12 +208,13 @@ def _eval(self, mode: bool = True):
 
 
 def save(
-    ep: ExportedProgram,
+    ep,
     f: Union[str, os.PathLike, io.BytesIO],
     *,
     extra_files: Optional[Dict[str, Any]] = None,
     opset_version: Optional[Dict[str, int]] = None,
 ) -> None:
+    from torch.export.exported_program import ExportedProgram
     if not isinstance(ep, ExportedProgram):
         raise TypeError(f"save() expects an ExportedProgram but got {type(ep)}")
 
@@ -270,7 +247,7 @@ def load(
     *,
     extra_files: Optional[Dict[str, Any]] = None,
     expected_opset_version: Optional[Dict[str, int]] = None,
-) -> ExportedProgram:
+):
     if isinstance(f, (str, os.PathLike)):
         f = os.fspath(f)
 
@@ -383,6 +360,7 @@ def aot_compile(
     """
     from torch.export._trace import _export_to_torch_ir
     from torch._inductor.decomposition import select_decomp_table
+    from torch._inductor import config
 
     if config.is_predispatch:
         gm = torch.export._trace._export(f, args, kwargs, dynamic_shapes, pre_dispatch=True).module()

torch/_export/serde/schema.py

@@ -8,7 +8,7 @@
 from torch._export.serde.union import _Union
 
 # NOTE: Please update this value if any modifications are made to the schema
-SCHEMA_VERSION = (5, 3)
+SCHEMA_VERSION = (6, 1)
 TREESPEC_VERSION = 1
 
 
@@ -376,4 +376,5 @@ class ExportedProgram:
     opset_version: Dict[str, int]
     range_constraints: Dict[str, RangeConstraint]
     schema_version: SchemaVersion
-    dialect: str
+    verifiers: List[str] = field(default_factory=list)
+    dialect: str = ""  # TODO deprecated

torch/_export/serde/schema.yaml

@@ -1,5 +1,5 @@
 # @generated by update_schema.py
-# checksum<<f1da1027d3bccb23db1f8dde8e635e53c7ab67fde5248ede49a6b7a3402ce743>>
+# checksum<<41b8154c93fefb96215f252a1913371bf61ad45808995a04ead58236ec3b7720>>
 Argument:
   kind: union
   fields:
@@ -102,8 +102,12 @@ ExportedProgram:
       type: Dict[str, RangeConstraint]
     schema_version:
       type: SchemaVersion
+    verifiers:
+      type: List[str]
+      default: '[]'
     dialect:
       type: str
+      default: ''
 GradientToParameterSpec:
   kind: struct
   fields:
@@ -425,6 +429,6 @@ UserOutputSpec:
     arg:
       type: Argument
 SCHEMA_VERSION:
-- 5
-- 3
+- 6
+- 1
 TREESPEC_VERSION: 1

torch/_export/serde/serialize.py

@@ -1384,7 +1384,7 @@ def serialize(self, exported_program: ep.ExportedProgram) -> _SerializedProgram:
                 major=SCHEMA_VERSION[0],
                 minor=SCHEMA_VERSION[1],
             ),
-            dialect=exported_program.dialect
+            verifiers=[v.dialect for v in exported_program.verifiers],
         )
 
         # Test canonical form is well defined.
@@ -2259,8 +2259,8 @@ def deserialize(
             range_constraints=range_constraints,
             module_call_graph=res.module_call_graph,
             example_inputs=res.example_inputs,
-            verifier=load_verifier(exported_program.dialect),
             constants=res.constants,
+            verifiers=[load_verifier(v) for v in exported_program.verifiers],
         )
 
 
@@ -2830,7 +2830,7 @@ def replace_output(out):
         opset_version=opset_version,
         range_constraints=range_constraints,
         schema_version=ep.schema_version,
-        dialect=ep.dialect
+        verifiers=ep.verifiers,
     )
 
 

torch/_export/verifier.py

@@ -3,12 +3,11 @@
 import math
 import operator
 from collections.abc import Iterable
-from typing import Any, Dict, final, List, Optional, Tuple, Type
+from typing import Any, Dict, final, List, Tuple, Type, TYPE_CHECKING
 
 import torch
 from torch._ops import HigherOrderOperator, OpOverload
 from torch._subclasses.fake_tensor import FakeTensor
-from torch.export.exported_program import ExportedProgram
 from torch.export.graph_signature import (
     CustomObjArgument,
     InputKind,
@@ -17,8 +16,9 @@
     TokenArgument,
 )
 from torch.fx import GraphModule
-from torch.fx.experimental.symbolic_shapes import SymBool, SymFloat, SymInt
 
+if TYPE_CHECKING:
+    from torch.export.exported_program import ExportedProgram
 
 class SpecViolationError(Exception):
     pass
@@ -34,6 +34,8 @@ def _check_has_fake_tensor(node: torch.fx.Node) -> None:
 
 
 def _check_val(node: torch.fx.Node) -> None:
+    from torch.fx.experimental.symbolic_shapes import SymBool, SymFloat, SymInt
+
     def _check_correct_val(val):
         if val is None:
             return True
@@ -150,7 +152,7 @@ def check_additional(self, gm: GraphModule) -> None:
         pass
 
     @final
-    def check(self, ep: ExportedProgram) -> None:
+    def check(self, ep: "ExportedProgram") -> None:
         self._check_graph_module(ep.graph_module)
         _verify_exported_program_signature(ep)
 
@@ -429,7 +431,7 @@ def _verify_exported_program_signature(exported_program) -> None:
             )
 
 
-def load_verifier(dialect: str) -> Optional[Type[Verifier]]:
+def load_verifier(dialect: str) -> Type[Verifier]:
     if dialect == "ATEN" or dialect == "":
-        return _VerifierMeta._registry.get(dialect)
+        return _VerifierMeta._registry.get(dialect, Verifier)
     return _VerifierMeta._registry[dialect]

torch/export/exported_program.py

@@ -12,6 +12,7 @@
     Any,
     Callable,
     Dict,
+    final,
     Iterator,
     List,
     Optional,
@@ -39,6 +40,8 @@
 
 import torch
 import torch.utils._pytree as pytree
+
+from torch._export.verifier import Verifier
 from torch._subclasses.functional_tensor import FunctionalTensor
 
 from torch.export._tree_utils import is_equivalent, reorder_kwargs
@@ -66,7 +69,6 @@
     TokenArgument,
 )
 
-
 __all__ = [
     "ExportedProgram",
     "ModuleCallEntry",
@@ -637,13 +639,15 @@ def __init__(
         range_constraints: "Dict[sympy.Symbol, Any]",
         module_call_graph: List[ModuleCallEntry],
         example_inputs: Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]] = None,
-        verifier: Optional[Type[Any]] = None,  # TODO Change typing hint to Verifier.
+        verifier: Optional[Type[Any]] = None,  # TODO Deprecate this.
         tensor_constants: Optional[
             Dict[str, torch.Tensor]
         ] = None,  # TODO: deprecate this
         constants: Optional[
             Dict[str, Union[torch.Tensor, FakeScriptObject, torch._C.ScriptObject]]
         ] = None,
+        *,
+        verifiers: Optional[List[Type[Verifier]]] = None,
     ):
         # Remove codegen related things from the graph. It should just be a flat graph.
         graph._codegen = torch.fx.graph.CodeGen()
@@ -661,14 +665,17 @@ def __init__(
         self._constants = tensor_constants or constants or {}
         assert self._constants is not None
 
-        from torch._export.verifier import Verifier
-
-        if verifier is None:
-            verifier = Verifier
-        assert issubclass(verifier, Verifier)
-        self._verifier = verifier
+        # TODO Clean up this after we bump executorch's pin.
+        assert verifier is None or verifiers is None
+        if verifiers is None:
+            if verifier is None:
+                verifiers = [Verifier]
+            else:
+                verifiers = [verifier]
+        assert all(issubclass(v, Verifier) for v in verifiers)
+        self._verifiers = verifiers
         # Validate should be always the last step of the constructor.
-        self.verifier().check(self)
+        self._validate()
 
     @property
     @compatibility(is_backward_compatible=False)
@@ -759,13 +766,18 @@ def call_spec(self):
     @property
     @compatibility(is_backward_compatible=False)
     def verifier(self) -> Any:
-        return self._verifier
+        return self._verifiers[0]
 
     @property
     @compatibility(is_backward_compatible=False)
     def dialect(self) -> str:
-        assert self._verifier is not None
-        return self._verifier.dialect
+        assert self._verifiers is not None
+        return self._verifiers[0].dialect
+
+    @property
+    @compatibility(is_backward_compatible=False)
+    def verifiers(self):
+        return self._verifiers
 
     @property
     @compatibility(is_backward_compatible=False)
@@ -1079,8 +1091,10 @@ def _check_input_constraints(self, flat_args_with_path):
             input_placeholders, flat_args_with_path, self.range_constraints
         )
 
+    @final
     def _validate(self):
-        self.verifier().check(self)
+        for v in self.verifiers:
+            v().check(self)
 
     # TODO(zhxchen17) Formalize this.
     def _update(