add reorder pass testing

Differential Revision: D66078188

Author: Zonglin Peng

backends/cadence/aot/TARGETS

@@ -335,3 +335,22 @@ python_unittest(
         "//executorch/exir/dialects:lib",
     ],
 )
+
+python_unittest(
+    name = "test_reorder_ops_passes",
+    srcs = [
+        "tests/test_reorder_ops_passes.py",
+    ],
+    typing = True,
+    deps = [
+        ":compiler",
+        ":pass_utils",
+        "//caffe2:torch",
+        "//executorch/backends/cadence/aot:compiler",
+        "//executorch/backends/cadence/aot:fuse_ops",
+        "//executorch/backends/cadence/aot:ops_registrations",
+        "//executorch/backends/cadence/aot:pass_utils",
+        "//executorch/backends/cadence/aot:reorder_ops",
+        "//executorch/exir/dialects:lib",
+    ],
+)

backends/cadence/aot/compiler.py

@@ -194,9 +194,6 @@ def export_to_edge(
     return edge_prog_manager
 
 
-# Export the model and lower it to an EdgeProgramManager (in edge IR), and
-# apply passes specific to Cadence DSP execution. Return both to print the
-# differences.
 def export_to_cadence(
     model: torch.nn.Module,
     inputs: tuple[object, ...],
@@ -216,6 +213,25 @@ def export_to_cadence(
     return cadence_prog_manager
 
 
+def quantize_and_export_to_cadence(
+    model: torch.nn.Module,
+    inputs: tuple[object, ...],
+    dump_graphs: bool = False,
+    opt_level: int = 1,
+) -> EdgeProgramManager:
+    quantized_model = quantize_pt2(model, inputs)
+
+    return export_to_cadence(
+        quantized_model,
+        inputs,
+        opt_level=opt_level,
+        dump_graphs=dump_graphs,
+    )
+
+
+# Export the model and lower it to an EdgeProgramManager (in edge IR), and
+# apply passes specific to Cadence DSP execution. Return both to print the
+# differences.
 def export_to_executorch_gen_etrecord(
     model: torch.nn.Module,
     inputs: tuple[object, ...],

backends/cadence/aot/pass_utils.py

@@ -3,7 +3,7 @@
 # pyre-strict
 
 from dataclasses import dataclass
-from typing import Callable, Optional, Set, Union
+from typing import Callable, List, Optional, Set, Union
 
 import torch
 from executorch.backends.cadence.aot.utils import get_edge_overload_packet
@@ -98,3 +98,47 @@ def count_node(graph_module: torch.fx.GraphModule, target: torch.fx.node.Target)
         if node.op == "call_function" and node.target == target:
             total += 1
     return total
+
+
+# Testing utils
+# Return the compute/function nodes in the graph
+def get_compute_nodes_in_gm(graph_module: torch.fx.GraphModule) -> List[torch.fx.Node]:
+    nodes = []
+    for x in graph_module.graph.nodes:
+        if x.op == "call_function":
+            if isinstance(x.target, torch._ops.OpOverload):
+                nodes.append(x.target.overloadpacket)
+            elif isinstance(x.target, EdgeOpOverload):
+                nodes.append(get_edge_overload_packet(x.target))
+    return nodes
+
+
+# Return true if there is no edge from a node with target pred_target to a
+# node with target succ_target in the graph.
+def nodes_not_connected_in_gm(
+    graph_module: torch.fx.GraphModule,
+    pred_target: torch.fx.Node,
+    succ_target: torch.fx.Node,
+) -> bool:
+    for node in graph_module.graph.nodes:
+        if node.target != pred_target:
+            continue
+        for user in node.users:
+            if user.target == succ_target:
+                return False
+    return True
+
+
+# Returns true if there is no instance of a node with target succ_target
+# positioned immediately after a node with target pred_target in the graph
+def nodes_not_adjacent_in_gm(
+    graph_module: torch.fx.GraphModule,
+    pred_target: torch.fx.Node,
+    succ_target: torch.fx.Node,
+) -> bool:
+    for node in graph_module.graph.nodes:
+        if node.target != pred_target:
+            continue
+        if node.next.target == succ_target:
+            return False
+    return True