Move graph_builder and program_builder to executorch.backends.test (#18483)

Summary:
Pull Request resolved: #18483

Move `GraphBuilder` and `ProgramBuilder` from `executorch.backends.cadence.aot`
to `executorch.backends.test` since they are general-purpose test utilities not
specific to the Cadence backend.

- Created new canonical modules at `executorch/backends/test/`
- Old modules now re-export from the new location for backward compatibility
- Updated all 12 downstream consumers to import from the new path
- Updated BUCK targets: new targets in `backends/test/targets.bzl`, old targets
  now depend on new ones

Differential Revision: D97995878

Author: hsharma35

backends/cadence/aot/BUCK

@@ -225,8 +225,7 @@ fbcode_target(_kind = runtime.python_library,
     ],
     typing = True,
     deps = [
-        "fbcode//caffe2:torch",
-        "fbcode//executorch/exir:pass_base",
+        "//executorch/backends/test:graph_builder",
     ],
 )
 
@@ -237,11 +236,7 @@ fbcode_target(_kind = runtime.python_library,
     ],
     typing = True,
     deps = [
-        ":graph_builder",
-        "fbcode//caffe2:torch",
-        "fbcode//executorch/exir:lib",
-        "fbcode//executorch/exir:pass_base",
-        "fbcode//executorch/exir/verification:verifier",
+        "//executorch/backends/test:program_builder",
     ],
 )
 
@@ -252,7 +247,7 @@ fbcode_target(_kind = python_unittest,
     ],
     typing = True,
     deps = [
-        ":program_builder",
+        "//executorch/backends/test:program_builder",
         "//caffe2:torch",
         "//later:lib",
     ],
@@ -396,7 +391,7 @@ fbcode_target(_kind = python_unittest,
         ":typing_stubs",
         ":type_dispatch",
         "//caffe2:torch",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
@@ -436,7 +431,7 @@ fbcode_target(_kind = python_unittest,
     deps = [
         ":ops_registrations",
         "//caffe2:torch",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
@@ -457,7 +452,7 @@ fbcode_target(_kind = python_unittest,
         ":replace_ops",
         "//caffe2:torch",
         "//executorch/backends/cadence/aot:compiler",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
@@ -478,7 +473,7 @@ fbcode_target(_kind = python_unittest,
         "//caffe2:torch",
         ":typing_stubs",
         "//executorch/backends/cadence/aot:compiler",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
@@ -499,7 +494,7 @@ fbcode_target(_kind = python_unittest,
         "//caffe2:torch",
         "//executorch/backends/cadence/aot:compiler",
         "//executorch/backends/cadence/aot:fuse_ops",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:ops_registrations",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/exir/dialects:lib",
@@ -520,7 +515,7 @@ fbcode_target(_kind = python_unittest,
         ":compiler",
         "//caffe2:torch",
         "//executorch/backends/cadence/aot:compiler",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:ops_registrations",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/backends/cadence/aot:remove_ops",
@@ -540,7 +535,7 @@ fbcode_target(_kind = python_unittest,
         ":typing_stubs",
         "//caffe2:torch",
         "//executorch/backends/cadence/aot:compiler",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:ops_registrations",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/backends/cadence/aot:simplify_ops",
@@ -560,7 +555,7 @@ fbcode_target(_kind = python_unittest,
         "//caffe2:torch",
         "//executorch/backends/cadence/aot:compiler",
         "//executorch/backends/cadence/aot:fuse_ops",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:ops_registrations",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/backends/cadence/aot:reorder_ops",
@@ -648,11 +643,11 @@ fbcode_target(_kind = python_unittest,
         ":typing_stubs",
         ":ops_registrations",
         ":pass_utils",
-        ":program_builder",
+        "//executorch/backends/test:program_builder",
         "//caffe2:torch",
         "//executorch/exir:memory",
         "//executorch/exir/dialects:lib",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/exir/tests:models",
     ],
 )
@@ -664,8 +659,7 @@ fbcode_target(_kind = python_unittest,
     ],
     typing = True,
     deps = [
-        ":program_builder",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:ops_registrations",
         "//executorch/runtime:runtime",
         "//later:lib",
@@ -695,7 +689,7 @@ fbcode_target(_kind = python_unittest,
     deps = [
         "fbsource//third-party/pypi/parameterized:parameterized",
         "//caffe2:torch",
-        "//executorch/backends/cadence/aot:graph_builder",
+        "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot/quantizer:quantizer",
         "//executorch/exir:pass_base",
         "//pytorch/ao:torchao",
@@ -710,7 +704,7 @@ fbcode_target(_kind = python_unittest,
     typing = True,
     deps = [
         ":ops_registrations",
-        ":program_builder",
+        "//executorch/backends/test:program_builder",
         ":to_out_var_pass",
         "//caffe2:torch",
         "//executorch/exir:lib",

backends/cadence/aot/graph_builder.py

@@ -6,137 +6,8 @@
 
 # pyre-strict
 
-import logging
-from typing import Optional, Sequence, Union
+# This module has moved to executorch.backends.test.graph_builder.
+# This re-export exists for backward compatibility.
+from executorch.backends.test.graph_builder import GraphBuilder, single_op_builder
 
-import torch
-from executorch.exir.pass_base import (
-    Argument,
-    ExportPass,
-    NodeMetadata,
-    PassResult,
-    ProxyValue,
-)
-from torch._dispatch.python import enable_python_dispatcher
-from torch._subclasses import FakeTensor, FakeTensorMode
-from torch.fx.node import Target
-from torch.utils import _pytree as pytree
-
-
-class GraphBuilder(ExportPass):
-    """Utility class for creating a graph module with user-specified ops.
-
-    This class allows us to create test graph modules with any ops we want
-    directly, rather than relying on decomposition or passes.
-
-    Usage:
-        builder = GraphBuilder()
-        # To insert placeholders, use builder.placeholder.
-        x = builder.placeholder("x", torch.randn(1, 3, 224, 224))
-        # To insert an op, use builder.call_operator.
-        op = builder.call_operator(
-            some_op
-            (x, other_args, ...),
-        )
-        # Insert outputs as a list of ProxyValues using builder.output.
-        builder.output([op])
-        # Get GraphModule from builder.
-        gm = builder.get_graph_module()
-    """
-
-    def __init__(self, fake_tensor_mode: Optional[FakeTensorMode] = None) -> None:
-        self.exporter = ExportPass()
-        self.tracer: ExportPass.ExportTracer = self.ExportTracer(
-            self, torch.fx.graph.CodeGen()
-        )
-        self.fake_tensor_mode: FakeTensorMode = fake_tensor_mode or FakeTensorMode(
-            allow_fallback_kernels=False,
-            allow_non_fake_inputs=True,
-        )
-        self.tracer.fake_tensor_mode = self.fake_tensor_mode
-
-        # This will be called to create nodes in tracer.
-        self.interpreter = torch.fx.Interpreter(
-            torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
-        )
-
-    # pyre-ignore[14]: Inconsistent override.
-    def placeholder(
-        self, target: str, fake_tensor: Union[FakeTensor, torch.Tensor]
-    ) -> ProxyValue:
-        if not isinstance(fake_tensor, FakeTensor):
-            fake_tensor = self.fake_tensor_mode.from_tensor(fake_tensor)
-        logging.debug(f"Creating placeholder {target} => {fake_tensor.shape}")
-        placeholder = super().placeholder(target, fake_tensor, NodeMetadata({}))
-        return placeholder
-
-    # pyre-ignore[14]: Inconsistent override.
-    def output(self, results: list[ProxyValue]) -> ProxyValue:
-        logging.debug(f"Creating outputs {results}")
-        return super().output(results, NodeMetadata({}))
-
-    def get_graph_module(self) -> torch.fx.GraphModule:
-        return torch.fx.GraphModule(self.tracer.root, self.tracer.graph)
-
-    def call_operator(
-        self,
-        op,  # pyre-ignore
-        args: tuple[Argument, ...],
-        kwargs: Optional[dict[str, Argument]] = None,
-        meta: Optional[NodeMetadata] = None,
-    ) -> ProxyValue:
-        if meta is None:
-            meta = NodeMetadata({})
-        if kwargs is None:
-            kwargs = {}
-        return super().call_operator(op, args, kwargs, meta)
-
-    def call_submodule(
-        self, graph_module: torch.fx.GraphModule, inputs: tuple[Argument, ...]
-    ) -> PassResult:
-        return ExportPass().call(graph_module)
-
-    def call_getitem(
-        self, value: ProxyValue, key: int, meta: Optional[NodeMetadata] = None
-    ) -> ProxyValue:
-        return super().call_getitem(value, key, meta or NodeMetadata({}))
-
-    def _fx(
-        self,
-        kind: str,
-        target: torch.fx.node.Target,
-        args: tuple[Argument, ...],
-        kwargs: dict[str, Argument],
-        meta: NodeMetadata,
-    ) -> ProxyValue:
-        with self.fake_tensor_mode, enable_python_dispatcher():
-            return super()._fx(kind, target, args, kwargs, meta)
-
-
-def single_op_builder(
-    placeholders: Sequence[Union[torch.Tensor, FakeTensor]],
-    op: Target,
-    args: Sequence[Argument],
-    kwargs: Optional[dict[str, Argument]] = None,
-) -> torch.fx.GraphModule:
-    """Create a graph module with a single op.
-
-    Args:
-        placeholders: Placeholders to be used as inputs to the GraphModule.
-        op: The op to be inserted.
-        args: The args to be passed to the op.
-        kwargs: The kwargs to be passed to the op.
-
-    Returns:
-        A graph module with a single op
-    """
-    builder = GraphBuilder()
-    op_to_placeholder_dict = {
-        p: builder.placeholder(f"p_{i}", p) for i, p in enumerate(placeholders)
-    }
-    proxy_args, proxy_kwargs = pytree.tree_map_only(
-        (torch.Tensor, FakeTensor), lambda x: op_to_placeholder_dict[x], (args, kwargs)
-    )
-    node = builder.call_operator(op, proxy_args, proxy_kwargs)
-    builder.output([node])
-    return builder.get_graph_module()
+__all__ = ["GraphBuilder", "single_op_builder"]