Class GraphBuilder

backends/test/graph_builder.py:26–113

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.plac

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