Fix edge dialect verifier bugs

Summary:
## Problem
If we lower a graph module to delegate and then compose it with some other graph module, retrace it, if we also turn on edge ops and validator (`_use_edge_ops=True`, `_check_ir_validity=False`), validator will error out because it doesn't know how to handle lowered module.

## Solution

Fix this by ignoring all HigherOrderOps: catch the exception and pass.

Reviewed By: JacobSzwejbka

Differential Revision: D46608636

fbshipit-source-id: 91fc7a887702f40c25a4595717210e8e32c3700a

Author: larryliu0820

exir/tests/test_arg_validator.py

@@ -7,7 +7,7 @@
 from executorch import exir
 from executorch.exir import EdgeCompileConfig
 from executorch.exir.dialects._ops import ops
-from executorch.exir.graph_module import get_exir_meta
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.verification.arg_validator import EdgeOpArgValidator
 
 
@@ -55,12 +55,13 @@ def forward(self, x):
         validator = EdgeOpArgValidator(egm)
         validator.run(*inputs)
         self.assertEqual(len(validator.violating_ops), 1)
+        key: EdgeOpOverload = next(iter(validator.violating_ops))
         self.assertEqual(
-            validator.violating_ops[0][0].name(),
+            key.name(),
             ops.edge.aten._log_softmax.default.name(),
         )
         self.assertDictEqual(
-            validator.violating_ops[0][1],
+            validator.violating_ops[key],
             {
                 "self": torch.bfloat16,
                 "__ret_0": torch.bfloat16,

exir/verification/arg_validator.py

@@ -1,10 +1,24 @@
 # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
-from typing import Any, Dict, List, Sequence, Tuple
+from collections import defaultdict
+from typing import Any, Dict, Sequence, Tuple
 
 import torch
 from executorch.exir.dialects.edge._ops import EdgeDialectFunctionSchema, EdgeOpOverload
 from executorch.exir.emit._emitter import _Argument, _Target
 from executorch.exir.error import ExportError, InternalError
+from torch._ops import HigherOrderOperator
+
+
+class RunHigherOrderOperatorError(Exception):
+    """
+    Raised when an we try to run delegate or other HigherOrderOperator in a graph module.
+    E.g., %executorch_call_delegate : [#users=1] = call_function[
+        target=torch.ops.executorch_call_delegate](args = (%lowered_module_0, %arg0_1), kwargs = {})
+    """
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
 
 # pyre-ignore[13]: Attribute `node` is never initialized.
 class EdgeOpArgValidator(torch.fx.Interpreter):
@@ -18,14 +32,16 @@ class EdgeOpArgValidator(torch.fx.Interpreter):
 
     def __init__(self, graph_module: torch.fx.GraphModule) -> None:
         super().__init__(graph_module)
-        self.violating_ops: List[Tuple[EdgeOpOverload, Dict[str, torch.dtype]]] = []
+        self.violating_ops: Dict[EdgeOpOverload, Dict[str, torch.dtype]] = defaultdict(
+            dict
+        )
 
     def run_node(self, n: torch.fx.Node) -> None:
         self.node = n
         try:
             ret = super().run_node(n)
         except Exception as e:
-            if isinstance(e, (InternalError, ExportError)):
+            if isinstance(e, (InternalError, ExportError, RunHigherOrderOperatorError)):
                 raise e
             else:
                 raise InternalError(str(e)) from e
@@ -40,7 +56,9 @@ def call_function(
         if not isinstance(target, EdgeOpOverload) or not isinstance(
             target._schema, EdgeDialectFunctionSchema
         ):
-            return False
+            if isinstance(target, HigherOrderOperator):
+                raise RunHigherOrderOperatorError("Can't run delegate")
+            return super().call_function(target, args, kwargs)
         tensor_arg_types: Dict[str, torch.dtype] = {}
         for i, schema_arg in enumerate(target._schema.arguments):
             if not isinstance(schema_arg.type, torch.TensorType):
@@ -60,15 +78,15 @@ def call_function(
             kernel_rets if isinstance(kernel_rets, Sequence) else [kernel_rets]
         )
         for schema_ret in target._schema.returns:
-            if isinstance(schema_ret.type, torch.TensorType):
-                name = schema_ret.name if schema_ret.name else f"__ret_{ret_index}"
-                kernel_ret = next(ret_iter)
-                if not isinstance(kernel_ret, torch.Tensor):
-                    continue
+            name = schema_ret.name if schema_ret.name else f"__ret_{ret_index}"
+            kernel_ret = next(ret_iter)
+            if isinstance(schema_ret.type, torch.TensorType) and isinstance(
+                kernel_ret, torch.Tensor
+            ):
                 tensor_arg_types[name] = kernel_ret.dtype
                 ret_index += 1
 
         valid = target._schema.dtype_constraint.validate(tensor_arg_types)
         if not valid:
-            self.violating_ops.append((target, tensor_arg_types))
-        return valid
+            self.violating_ops[target] = tensor_arg_types
+        return super().call_function(target, args, kwargs)

exir/verification/verifier.py

@@ -4,7 +4,10 @@
 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.verification.arg_validator import EdgeOpArgValidator
+from executorch.exir.verification.arg_validator import (
+    EdgeOpArgValidator,
+    RunHigherOrderOperatorError,
+)
 
 from torch._export.verifier import (
     _check_has_fake_tensor,
@@ -70,7 +73,17 @@ def extract_input(node: torch.fx.Node) -> Optional[FakeTensor]:
 def _check_tensor_args_matching_op_allowed_dtype(gm: GraphModule) -> None:
     validator = EdgeOpArgValidator(gm)
     inputs = _get_inputs(gm)
-    validator.run(*inputs)
+    try:
+        validator.run(*inputs)
+    except RunHigherOrderOperatorError:
+        # NB: ignore higher order operator in the graph.
+        # If we lower a graph module to delegate and then compose it with some other graph module, retrace it,
+        # if we also turn on edge ops and validator (_use_edge_ops=True, _check_ir_validity=True), we will run
+        # into RunHigherOrderOperatorError. The only thing we can do right now is to ignore this error, since
+        # by definition it's still a valid Edge dialect. This is not ideal because it ignores possible invalidity
+        # later in the graph.
+        return
+
     if validator.violating_ops:
         raise SpecViolationError(
             f"These operators are taking Tensor inputs with mismatched dtypes: {validator.violating_ops}"