Verifier for exported program (#292)

Summary:
X-link: pytorch/pytorch#109519


Added a verifier for the graph signature in a exported program

Differential Revision: D48926643

Author: angelayi

exir/tests/test_verification.py

@@ -145,6 +145,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         self.assertEqual(len(res), len(res_interp))
         self.assertTrue(torch.allclose(res, res_interp))
 
+
+class TestEdgeVerification(unittest.TestCase):
     def test_edge_happy(self) -> None:
         class TestModel(torch.nn.Module):
             def __init__(self):

exir/verification/TARGETS

@@ -51,8 +51,8 @@ python_library(
     ],
     deps = [
         "//caffe2:torch",
-        "//executorch/exir:delegate",
         "//executorch/exir:error",
+        "//executorch/exir:lowered_backend_module",
         "//executorch/exir/dialects/edge:lib",
         "//executorch/exir/emit:emit",
     ],

exir/verification/verifier.py

@@ -4,20 +4,21 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import List, Optional
+import itertools
+import operator
+from typing import Any, List, Optional, Tuple, Type
 
 import torch
-from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.error import ExportError, ExportErrorType
+from executorch.exir.lowered_backend_module import LoweredBackendModule
 from executorch.exir.verification.arg_validator import (
     EdgeOpArgValidator,
     RunHigherOrderOperatorError,
 )
 
 from torch._export.verifier import (
     _check_has_fake_tensor,
-    _check_tensors_are_contiguous,
     ATenDialectVerifier,
     SpecViolationError,
     Verifier,
@@ -28,30 +29,21 @@
 
 
 ALLOWED_META_KEYS = {"spec", "stack_trace"}
-VALID_BUILTIN_FUNCS = [
-    executorch_call_delegate,
-]
 
 
-class EXIRATenDialectVerifier(ATenDialectVerifier):
-    def valid_builtin_funcs(self):
-        builtin_funcs = super().valid_builtin_funcs()
-        builtin_funcs.extend(VALID_BUILTIN_FUNCS)
-        return builtin_funcs
-
-    # TODO(angelayi): Delete this function when we migrate all tests to
-    #  because right now old tracer does not add ["val"] metadata
-    def check_valid(self, gm: GraphModule) -> None:  # noqa: C901
-
-        for node in gm.graph.nodes:
-            if node.op in {"call_module", "call_method"}:
+def _check_tensors_are_contiguous(gm: GraphModule) -> None:
+    # Tensors be of contiguous format
+    for name, param in itertools.chain(gm.named_parameters(), gm.named_buffers()):
+        if isinstance(param, torch.Tensor):
+            if not param.is_contiguous():
                 raise SpecViolationError(
-                    "call_module is not valid: got a class '{}' ".format(node.target),
+                    f"Tensors in Aten dialect must be contiguous, {name} is not contiguous"
                 )
 
-            if node.op == "call_function":
-                if node.target not in self.valid_builtin_funcs():
-                    self.check_valid_op(node.target)
+
+class EXIRATenDialectVerifier(ATenDialectVerifier):
+    def allowed_getattr_types(self) -> Tuple[Type[Any], ...]:
+        return (torch.fx.GraphModule, LoweredBackendModule, torch.Tensor)
 
 
 def _get_inputs(graph_module: GraphModule) -> List[Optional[FakeTensor]]:
@@ -97,15 +89,21 @@ def _check_tensor_args_matching_op_allowed_dtype(gm: GraphModule) -> None:
 
 
 class EXIREdgeDialectVerifier(Verifier):
-    def __init__(self, check_edge_ops: bool = False) -> None:
+    def __init__(self, check_edge_ops: bool = True) -> None:
         self.check_edge_ops = check_edge_ops
 
-    def valid_builtin_funcs(self):
-        builtin_funcs = super().valid_builtin_funcs()
-        builtin_funcs.extend(VALID_BUILTIN_FUNCS)
-        return builtin_funcs
+        if self.check_edge_ops:
+            self.check_valid_op = self.check_valid_edge_op
+        else:
+            self.check_valid_op = self.check_valid_aten_op
+
+    def allowed_getattr_types(self) -> Tuple[Type[Any], ...]:
+        return (torch.fx.GraphModule, LoweredBackendModule, torch.Tensor)
 
     def check_valid_edge_op(self, op):
+        if op in [operator.getitem]:
+            return
+
         if isinstance(op, OpOverload) and not isinstance(op, EdgeOpOverload):
             raise SpecViolationError(
                 "Operator {}.{} is not an Edge operator.".format(
@@ -116,33 +114,23 @@ def check_valid_edge_op(self, op):
     def check_valid_aten_op(self, op) -> None:
         super().check_valid_op(op)
 
-        op_name = op.name if hasattr(op, "name") else op.__name__
-
-        if not isinstance(op, OpOverload):
-            raise SpecViolationError(
-                "Operator '{}' is not a registered Op".format(op_name),
-            )
-
-        if (
-            torch.Tag.core not in op.tags  # type: ignore[attr-defined]
-            and torch.Tag.view_copy not in op.tags  # type: ignore[attr-defined]
-        ):
-            # NOTE(qihan): whether view_copy operators are marked as canonical is still under
-            #            discussion.
-            raise SpecViolationError(
-                "Operator {}.{} is not Aten Canonical.".format(
-                    op.__module__, op.__name__
+        if isinstance(op, OpOverload):
+            if (
+                torch.Tag.core not in op.tags  # type: ignore[attr-defined]
+                and torch.Tag.view_copy not in op.tags  # type: ignore[attr-defined]
+            ):
+                # NOTE(qihan): whether view_copy operators are marked as canonical is still under
+                #            discussion.
+                raise SpecViolationError(
+                    "Operator {}.{} is not Aten Canonical.".format(
+                        op.__module__, op.__name__
+                    )
                 )
-            )
 
-    def check_valid(self, gm: GraphModule) -> None:
+    def check_additional(self, gm: GraphModule) -> None:
         if self.check_edge_ops:
-            self.check_valid_op = self.check_valid_edge_op
-            super().check_valid(gm)
             _check_tensors_are_contiguous(gm)
             _check_tensor_args_matching_op_allowed_dtype(gm)
-        else:
-            self.check_valid_op = self.check_valid_aten_op
 
         # Additionally, edge dialect's operator must have same input dtype
         for n in gm.graph.nodes: