Improve verifier to not specialize on dialect. (#1528)

Summary:
X-link: pytorch/pytorch#116705

Pull Request resolved: #1528

Currently we have a very ugly specialization on edge dialect in verifier like the following:
```
 # TODO Remove this branch.
            if ep.dialect == "EDGE":  # !!! Don't change this allowlist. !!!
                pass
            else:
                raise e
```
In this diff we do some additional work to make signature checking also work in exir. We decouple the transformation stack in torch export and exir so that different layers of the stack can evolve in their own fashion and the team can divide and conquer them seperately.

Reviewed By: angelayi

Differential Revision: D52499225

fbshipit-source-id: d667397b8434d8f6b0d0d097d31a0dd589c329cc

Author: zhxchen17

backends/xnnpack/passes/__init__.py

@@ -25,6 +25,8 @@
 from executorch.exir.pass_base import ExportPass
 
 from executorch.exir.passes.const_prop_pass import ConstPropPass
+
+from executorch.exir.program._program import _transform
 from torch._export.pass_base import PassType
 
 from torch.export import ExportedProgram
@@ -77,5 +79,5 @@ def transform(self) -> ExportedProgram:
                 raise RuntimeError(
                     f"Expecting ExportPass or ExportPass(), but got pass: {pass_} with type: {type(pass_)}"
                 )
-            ep = ep._transform(transform_pass)
+            ep = _transform(ep, transform_pass)
         return ep

exir/capture/_capture.py

@@ -14,7 +14,7 @@
 from executorch.exir.capture._config import CaptureConfig
 from executorch.exir.error import ExportError, ExportErrorType, InternalError
 from executorch.exir.program import ExirExportedProgram, MultiMethodExirExportedProgram
-from executorch.exir.program._program import HackedUpExportedProgramDONOTUSE
+from executorch.exir.program._program import _transform, HackedUpExportedProgramDONOTUSE
 from executorch.exir.tracer import (
     _default_decomposition_table,
     dispatch_trace,
@@ -170,7 +170,7 @@ def capture(  # noqa: C901
 
             ep = export(f, args, constraints=constraints)
             ep = ep.run_decompositions(_default_decomposition_table())
-            ep = ep._transform(ReplaceViewOpsWithViewCopyOpsPass())
+            ep = _transform(ep, ReplaceViewOpsWithViewCopyOpsPass())
             if not config._unlift:
                 return ExirExportedProgram(ep, False)
             graph_module = ep.module()

exir/lowered_backend_module.py

@@ -122,6 +122,11 @@ def buffer(
     # TODO(chenlai): re-consider recapture instead of manually constructing the program because
     # the meta data construction is done manually.
     def program(self, emit_stacktrace: bool = False) -> Program:
+        from executorch.exir.program._program import (
+            _get_updated_graph_signature,
+            _transform,
+        )
+
         """
         Returns the object that represents the ExecuTorch binary before serialization.
         """
@@ -257,8 +262,11 @@ def program(self, emit_stacktrace: bool = False) -> Program:
         exported_program = ExportedProgram(
             root=lowered_exported_program.graph_module,
             graph=lowered_exported_program.graph,
-            graph_signature=ExportGraphSignature(
-                input_specs=input_specs, output_specs=output_specs
+            graph_signature=_get_updated_graph_signature(
+                ExportGraphSignature(
+                    input_specs=input_specs, output_specs=output_specs
+                ),
+                lowered_exported_program.graph_module,
             ),
             # TODO: May need to set lowered_exported_program.call_spec = CallSpec(None, None)
             # somewhere as we should pass it a list of tensors to the lowered module and output a
@@ -271,8 +279,8 @@ def program(self, emit_stacktrace: bool = False) -> Program:
             example_inputs=None,
             verifier=lowered_exported_program.verifier,
         )
-        exported_program = exported_program._transform(
-            SpecPropPass(), MemoryPlanningPass("greedy")
+        exported_program = _transform(
+            exported_program, SpecPropPass(), MemoryPlanningPass("greedy")
         )
         emitted_program = emit_program(
             exported_program, emit_stacktrace=emit_stacktrace

exir/program/_program.py

@@ -36,14 +36,98 @@
 )
 from torch._export import ExportedProgram
 from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
-from torch.export.exported_program import InputKind, InputSpec, TensorArgument
+from torch.export.exported_program import (
+    _get_updated_range_constraints,
+    ConstantArgument,
+    ExportGraphSignature,
+    InputKind,
+    InputSpec,
+    OutputSpec,
+    TensorArgument,
+)
 from torch.fx import _pytree as fx_pytree
 from torch.fx._compatibility import compatibility
+from torch.fx.passes.infra.pass_manager import PassManager
 from torch.utils import _pytree as pytree
 
 Val = Any
 
 
+def _get_updated_graph_signature(
+    old_signature: ExportGraphSignature,
+    new_gm: torch.fx.GraphModule,
+) -> ExportGraphSignature:
+    """
+    Update the graph signature's user_input/user_outputs.
+    """
+    new_input_specs = []
+    for i, node in enumerate(new_gm.graph.nodes):
+        if node.op != "placeholder":
+            break
+
+        assert i < len(
+            old_signature.input_specs
+        ), "Number of inputs changed after transformation"
+        old_input_spec = old_signature.input_specs[i]
+        arg = (
+            old_input_spec.arg
+            if isinstance(old_input_spec.arg, ConstantArgument)
+            else type(old_input_spec.arg)(node.name)
+        )
+        new_input_specs.append(
+            InputSpec(old_input_spec.kind, arg, old_input_spec.target)
+        )
+
+    output_node = list(new_gm.graph.nodes)[-1]
+    assert output_node.op == "output"
+
+    new_output_specs = []
+    for i, node in enumerate(output_node.args[0]):
+        assert i < len(
+            old_signature.output_specs
+        ), "Number of outputs changed after transformation"
+        old_output_spec = old_signature.output_specs[i]
+        arg = (
+            old_output_spec.arg
+            if isinstance(old_output_spec.arg, ConstantArgument)
+            else type(old_output_spec.arg)(node.name)
+        )
+        new_output_specs.append(
+            OutputSpec(old_output_spec.kind, arg, old_output_spec.target)
+        )
+
+    new_signature = ExportGraphSignature(
+        input_specs=new_input_specs, output_specs=new_output_specs
+    )
+    return new_signature
+
+
+def _transform(self, *passes: PassType) -> "ExportedProgram":
+    pm = PassManager(list(passes))
+    res = pm(self.graph_module)
+    transformed_gm = res.graph_module if res is not None else self.graph_module
+    assert transformed_gm is not None
+
+    if transformed_gm is self.graph_module and not res.modified:
+        return self
+
+    transformed_ep = ExportedProgram(
+        transformed_gm,
+        transformed_gm.graph,
+        _get_updated_graph_signature(self.graph_signature, transformed_gm),
+        self.state_dict,
+        _get_updated_range_constraints(transformed_gm),
+        copy.deepcopy(self.equality_constraints),
+        copy.deepcopy(self._module_call_graph),
+        self.example_inputs,
+        self.verifier,
+        self.tensor_constants,
+    )
+    transformed_ep.graph_module.meta.update(self.graph_module.meta)
+    transformed_ep.graph_module.meta.update(res.graph_module.meta)
+    return transformed_ep
+
+
 def _copy_module(new_prog, new_gm):
     new_prog.meta.update(new_gm.meta)
     new_prog.graph = new_gm.graph
@@ -231,7 +315,7 @@ def __init__(
         self.after_to_edge_passes = after_to_edge_passes
 
     def transform(self, *passes: PassType) -> "ExirExportedProgram":
-        self.exported_program = self.exported_program._transform(*passes)
+        self.exported_program = _transform(self.exported_program, *passes)
         return self
 
     def __call__(self, *args: Any) -> Any:
@@ -419,7 +503,7 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
     new_ep.exported_program = ExportedProgram(
         new_gm,
         new_gm.graph,
-        new_ep.exported_program.graph_signature,
+        _get_updated_graph_signature(new_ep.exported_program.graph_signature, new_gm),
         new_ep.exported_program.state_dict,
         new_ep.exported_program.range_constraints,
         new_ep.exported_program.equality_constraints,
@@ -755,7 +839,9 @@ def to_edge(
         edge_program = ExportedProgram(
             root=gm,
             graph=gm.graph,
-            graph_signature=edge_program.graph_signature,
+            graph_signature=_get_updated_graph_signature(
+                edge_program.graph_signature, gm
+            ),
             state_dict=edge_program.state_dict,
             range_constraints=edge_program.range_constraints,
             equality_constraints=edge_program.equality_constraints,
@@ -770,7 +856,7 @@ def to_edge(
         )
         passes = []
         passes.extend(aten_to_edge_passes.passes[-2:])
-        edge_program = edge_program._transform(*passes)
+        edge_program = _transform(edge_program, *passes)
         edge_programs[name] = edge_program
     return EdgeProgramManager(edge_programs, constant_methods, config)
 
@@ -856,7 +942,7 @@ def transform(
         if isinstance(passes, dict):
             for name, program in self._edge_programs.items():
                 if name in passes.keys():
-                    new_programs[name] = program._transform(*passes[name])
+                    new_programs[name] = _transform(program, *passes[name])
                     EXIREdgeDialectVerifier(enable=check_ir_validity)(
                         new_programs[name].graph_module
                     )
@@ -865,7 +951,7 @@ def transform(
 
         else:  # apply passes to every method
             for name, program in self._edge_programs.items():
-                new_programs[name] = program._transform(*passes)
+                new_programs[name] = _transform(program, *passes)
                 EXIREdgeDialectVerifier(enable=check_ir_validity)(
                     new_programs[name].graph_module
                 )