diff --git a/docs/source/compiler-custom-compiler-passes.md b/docs/source/compiler-custom-compiler-passes.md index aaba70b02af..7b8a48223e1 100644 --- a/docs/source/compiler-custom-compiler-passes.md +++ b/docs/source/compiler-custom-compiler-passes.md @@ -446,7 +446,7 @@ Args: * `graph_module (fx.GraphModule)`: Module that we want to partition * `patterns (List[torch.fx.Graph])`: A list of patterns in the form of torch.fx.Graph. These graphs can be obtained through the `graph` field from a - GraphModule obtained by exir.capture (recommended) or symbolic tracing (which + GraphModule obtained by torch.export (recommended) or symbolic tracing (which might not result in an accurate edge dialect graph), or by manual crafting a graph module. * `op_support (OperatorSupportBase)`: A OperatorSupportBase that can be created diff --git a/exir/__init__.py b/exir/__init__.py index 4b6d68488e7..5826c81bba8 100644 --- a/exir/__init__.py +++ b/exir/__init__.py @@ -6,14 +6,7 @@ from typing import Any -from executorch.exir.capture import ( - _capture_legacy_do_not_use, - CallSpec, - capture, - CaptureConfig, - EdgeCompileConfig, - ExecutorchBackendConfig, -) +from executorch.exir.capture import EdgeCompileConfig, ExecutorchBackendConfig from executorch.exir.emit import emit_program, EmitterOutput from executorch.exir.program import ( _to_edge, @@ -34,9 +27,6 @@ __all__ = [ "emit_program", "EmitterOutput", - "capture", - "_capture_legacy_do_not_use", - "CallSpec", "ExportedProgram", "ExirExportedProgram", "ExecutorchProgram", @@ -47,7 +37,6 @@ "EdgeProgramManager", "ExecutorchProgramManager", "edge_to_executorch_passes", - "CaptureConfig", "EdgeCompileConfig", "ExecutorchBackendConfig", "Value", diff --git a/exir/backend/canonical_partitioners/pattern_op_partitioner.py b/exir/backend/canonical_partitioners/pattern_op_partitioner.py index 7592a6e653b..e2a65460d1e 100644 --- a/exir/backend/canonical_partitioners/pattern_op_partitioner.py +++ b/exir/backend/canonical_partitioners/pattern_op_partitioner.py @@ -117,7 +117,7 @@ def generate_pattern_op_partitions( graph_module: Module that we want to partition patterns: A list of patterns in the form of torch.fx.Graph. These graphs can be obtained through the `graph` field from a GraphModule obtained by - exir.capture (recommended) or symbolic tracing (which might not result + torch.export (recommended) or symbolic tracing (which might not result in an accurate edge dialect graph), or by manual crafting a graph module. partitions_list: A list of node lists whose nodes are intended to be tagged diff --git a/exir/capture/BUCK b/exir/capture/BUCK index ceee803bcfa..8ec95ceb8e7 100644 --- a/exir/capture/BUCK +++ b/exir/capture/BUCK @@ -15,24 +15,18 @@ fbcode_target(_kind = runtime.python_library, "__init__.py", ], deps = [ - ":capture", ":config", + ":utils", ], ) fbcode_target(_kind = runtime.python_library, - name = "capture", + name = "utils", srcs = [ - "_capture.py", + "_utils.py", ], deps = [ - ":config", "//caffe2:torch", - "//executorch/exir:error", - "//executorch/exir:tracer", - "//executorch/exir/program:lib", - "//executorch/exir/program:program", - "//executorch/exir/verification:verifier", ], ) @@ -45,7 +39,6 @@ fbcode_target(_kind = runtime.python_library, "//caffe2:torch", "//executorch/exir:dynamic_shape", "//executorch/exir:pass_manager", - "//executorch/exir:tracer", "//executorch/exir/passes:lib", "//executorch/exir/passes:propagate_device_config", "//executorch/exir/passes:sym_shape_eval_pass", diff --git a/exir/capture/__init__.py b/exir/capture/__init__.py index b6304b86a2b..7c2863346af 100644 --- a/exir/capture/__init__.py +++ b/exir/capture/__init__.py @@ -6,23 +6,9 @@ # pyre-strict -from executorch.exir.capture._capture import ( - _capture_legacy_do_not_use, - CallSpec, - capture, -) - -from executorch.exir.capture._config import ( - CaptureConfig, - EdgeCompileConfig, - ExecutorchBackendConfig, -) +from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig __all__ = [ - "CallSpec", - "capture", - "_capture_legacy_do_not_use", - "CaptureConfig", "EdgeCompileConfig", "ExecutorchBackendConfig", ] diff --git a/exir/capture/_capture.py b/exir/capture/_capture.py deleted file mode 100644 index 975191f0744..00000000000 --- a/exir/capture/_capture.py +++ /dev/null @@ -1,381 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import copy -import warnings -from collections import namedtuple -from contextlib import contextmanager -from types import MethodType -from typing import Any, Callable, cast, List, Optional, Tuple - -import torch -from executorch.exir.capture._config import CaptureConfig -from executorch.exir.error import ExportError, ExportErrorType, InternalError -from executorch.exir.program import ExirExportedProgram -from executorch.exir.program._program import _transform, HackedUpExportedProgramDONOTUSE -from executorch.exir.tracer import ( - _default_decomposition_table, - dispatch_trace, - dynamo_trace, - flatten_output, - Value, -) -from executorch.exir.verification.verifier import EXIRATenDialectVerifierBase -from torch import _guards -from torch._dispatch.python import enable_python_dispatcher -from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass -from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode -from torch.export import export -from torch.export.exported_program import ( - ExportedProgram, - ExportGraphSignature, - InputKind, - InputSpec, - ModuleCallEntry, - ModuleCallSignature, - OutputKind, - OutputSpec, - TensorArgument, -) -from torch.func import functionalize -from torch.fx._compatibility import compatibility -from torch.fx.experimental.proxy_tensor import make_fx -from torch.fx.experimental.symbolic_shapes import ShapeEnv -from torch.utils import _pytree as pytree - - -Val = Any - - -CompileSpec = namedtuple( - "CompileSpec", ["method_name", "callable", "args", "dynamic_shapes"] -) - - -CallSpec = namedtuple("CallSpec", ["in_spec", "out_spec"]) - - -@compatibility(is_backward_compatible=False) -def _capture_legacy_do_not_use(f, args) -> ExirExportedProgram: - """ - This is a legacy API that should be avoided. Prefer to use capture() instead. - """ - warnings.warn( - "This function is now deprecated, please use `torch.export and exir.to_edge` instead. " - "See https://github.com/pytorch/functorch for more details.", - DeprecationWarning, - stacklevel=1, - ) - - graph_module = dispatch_trace(f, args) - flat_args = tuple(pytree.tree_flatten(args)[0]) - in_spec, out_spec = graph_module.in_spec, graph_module.out_spec - - _instantiate_missing_placeholder_val_with_real_inputs(graph_module, flat_args) - graph_module._apply(torch.Tensor.contiguous) - - user_inputs = [ - node.name for node in graph_module.graph.nodes if node.op == "placeholder" - ] - output_node = list(graph_module.graph.nodes)[-1] - assert output_node.op == "output" - user_outputs = [arg.name for arg in output_node.args[0]] - - for n in graph_module.graph.nodes: - if n.op == "call_function" and "val" not in n.meta: - try: - args, kwargs = pytree.tree_map_only( - torch.fx.Node, lambda x: x.meta["val"], (n.args, n.kwargs) - ) - n.meta["val"] = n.target(*args, **kwargs) - except Exception: - n.meta["val"] = None - - ep = HackedUpExportedProgramDONOTUSE( - root=graph_module, - graph=graph_module.graph, - graph_signature=ExportGraphSignature( - input_specs=[ - InputSpec( - kind=InputKind.USER_INPUT, arg=TensorArgument(name=i), target=None - ) - for i in user_inputs - ], - output_specs=[ - OutputSpec( - kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name=o), target=None - ) - for o in user_outputs - ], - ), - call_spec=CallSpec(in_spec, out_spec), - state_dict={}, - range_constraints={}, - module_call_graph=[ - ModuleCallEntry( - fqn="", - signature=ModuleCallSignature( - inputs=[], - outputs=[], - # pyre-fixme[6]: For 3rd argument expected `TreeSpec` but got - # `Union[Tensor, Module]`. - in_spec=in_spec, - # pyre-fixme[6]: For 4th argument expected `TreeSpec` but got - # `Union[Tensor, Module]`. - out_spec=out_spec, - ), - ) - ], - example_inputs=None, - verifier=EXIRATenDialectVerifierBase, - ) - return ExirExportedProgram(ep, False) - - -@contextmanager -def patch_forward(obj: torch.nn.Module, new_method): - """Helper method to make it easier to cleanly torch.export() a method on a - module that is not `forward`. - - TODO(suo): upstream this to torch.export.wrapper. - """ - # Save the original method - original_method = obj.forward - - # Patch the method - obj.forward = new_method.__get__(obj, obj.__class__) - - try: - yield - finally: - # Restore the original method - obj.forward = original_method - - -class WrapperModule(torch.nn.Module): - def __init__(self, f): - super().__init__() - self.forward = f - - -@compatibility(is_backward_compatible=False) -def capture( # noqa: C901 - f: Callable[..., Any], - args: Tuple[Value, ...], - config: Optional[CaptureConfig] = None, - dynamic_shapes: Optional[List[Any]] = None, -) -> ExirExportedProgram: - warnings.warn( - "This function is now deprecated, please use `torch.export and exir.to_edge` instead. ", - DeprecationWarning, - stacklevel=1, - ) - if not isinstance(args, tuple): - raise ExportError( - ExportErrorType.INVALID_INPUT_TYPE, - f"Expect `args` to be a tuple, got type: {type(args)}.", - ) - - config = config or CaptureConfig() - out_spec = None - # TODO (zhxchen17) Always functionalize in a second pass no matter which path is taken. - flat_args = tuple(pytree.tree_flatten(args)[0]) - if not config.enable_aot: - if config._unlift: - raise ExportError( - ExportErrorType.NOT_SUPPORTED, - "_unlift config doesn't do anything without enable_aot enabled. Please do not set it", - ) - if config.pt2_mode: - if config.enable_aot: - if config.enable_dynamic_shape: - raise ExportError( - ExportErrorType.NOT_SUPPORTED, - "Under enable_aot, enable_dynamic_shapes flag doesn't do anything. Please do not set it", - ) - if not config.enable_functionalization: - raise ExportError( - ExportErrorType.NOT_SUPPORTED, - "Functionalization is required for enable_aot.", - ) - - # If trying to capture a method and the bound class instance is a - # Module, then export the module while patching in that method. - if isinstance(f, MethodType) and isinstance(f.__self__, torch.nn.Module): - with patch_forward(f.__self__, f): - ep = export( - cast(torch.nn.Module, f.__self__), - args, - dynamic_shapes=dynamic_shapes, - strict=True, - ) - else: - mod = f if isinstance(f, torch.nn.Module) else WrapperModule(f) - ep = export(mod, args, dynamic_shapes=dynamic_shapes, strict=True) - - ep = ep.run_decompositions(_default_decomposition_table()) - ep = _transform(ep, ReplaceViewOpsWithViewCopyOpsPass()) - if not config._unlift: - return ExirExportedProgram(ep, False) - graph_module = cast(torch.fx.GraphModule, ep.module()) - - elif config.enable_dynamic_shape: - graph_module, _ = dynamo_trace( - f, - args, - aten_graph=True, - tracing_mode="symbolic", - dynamo_config=config._dynamo_config, - dynamic_shapes=dynamic_shapes, - _use_old_decomp_table=config._use_old_decomp_table, - ) - - else: - graph_module, _ = dynamo_trace( - f, - args, - aten_graph=True, - tracing_mode="fake", - dynamo_config=config._dynamo_config, - dynamic_shapes=None, - _use_old_decomp_table=config._use_old_decomp_table, - ) - - if out_spec is None: - if isinstance(graph_module.graph._codegen, torch.fx.graph._PyTreeCodeGen): - out_spec = graph_module.graph._codegen.pytree_info.out_spec - elif hasattr(graph_module, "_out_spec"): - out_spec = graph_module._out_spec - else: - out_spec = pytree.tree_flatten(f(*args))[1] - - # NOTE (tmanlaibaatar) - # torchdynamo.export adds extra kwarg into the graph module - # which is then lost while we are calling make_fx. This is because - # make_fx doesn't handle kwargs. Originally we used to use torchdynamo - # input spec, but due to some limitations in pytree implementation, it doesn't - # recognize the make_fx graph with torchdynamo input spec. We workaround it - # by getting the input spec directly from user argument. - in_spec = pytree.tree_flatten((args, {}))[1] - - if config.enable_functionalization and not config.enable_aot: - args = copy.deepcopy(args) - - def graph_with_interpreter(*args): - with torch.fx.traceback.preserve_node_meta(): - return torch.fx.Interpreter(graph_module).run(*args) - - functionalized_callable = functionalize( - graph_with_interpreter, - remove="mutations_and_views", - ) - assert isinstance(functionalized_callable, Callable) - - if config.enable_dynamic_shape: - fake_tensor_mode = FakeTensorMode( - allow_fallback_kernels=False, - allow_non_fake_inputs=True, - shape_env=ShapeEnv(), - ) - - inps: List[torch.Tensor] = [] - for node in graph_module.graph.nodes: - if node.op == "placeholder" and "val" in node.meta: - example_fake_tensor = node.meta["val"] - assert isinstance(example_fake_tensor, FakeTensor) - inps.append(example_fake_tensor) - - if detected_fake_mode := _guards.detect_fake_mode(inps): - fake_tensor_mode = detected_fake_mode - - count = 0 - - def convert_to_fake(x): - nonlocal count - val = inps[count] - count += 1 - return val - - fake_args = pytree.tree_map_only(torch.Tensor, convert_to_fake, args) - - with enable_python_dispatcher(), fake_tensor_mode: - graph_module = make_fx( - functionalized_callable, - tracing_mode="real", - _allow_non_fake_inputs=True, - )(*fake_args) - else: - # To avoid breaking folks, use the deprecated "real" tracing - # mode if we're not using pt2. - tracing_mode = "fake" if config.pt2_mode else "real" - graph_module = make_fx( - functionalized_callable, - tracing_mode=tracing_mode, - _allow_non_fake_inputs=True, - )(*args) - - flatten_output(graph_module) - - else: - raise InternalError("pt2=False path is officially deprecated") - - _instantiate_missing_placeholder_val_with_real_inputs(graph_module, flat_args) - graph_module._apply(torch.Tensor.contiguous) - - user_inputs = [ - InputSpec( - kind=InputKind.USER_INPUT, arg=TensorArgument(name=node.name), target=None - ) - for node in graph_module.graph.nodes - if node.op == "placeholder" - ] - output_node = list(graph_module.graph.nodes)[-1] - assert output_node.op == "output" - user_outputs = [ - OutputSpec( - kind=OutputKind.USER_OUTPUT, arg=TensorArgument(name=arg.name), target=None - ) - for arg in output_node.args[0] - ] - - graph_module.graph.eliminate_dead_code() - ep = ExportedProgram( - root=graph_module, - graph=graph_module.graph, - graph_signature=ExportGraphSignature(user_inputs, user_outputs), - state_dict={}, - range_constraints={}, - module_call_graph=[ - ModuleCallEntry( - fqn="", - signature=ModuleCallSignature( - inputs=[], - outputs=[], - in_spec=in_spec, - # pyre-fixme[6]: For 4th argument expected `TreeSpec` but got - # `Union[None, TreeSpec, Tensor, Module]`. - out_spec=out_spec, - ), - ) - ], - example_inputs=None, - verifiers=[EXIRATenDialectVerifierBase], - ) - return ExirExportedProgram(ep, False) - - -# This is to bootstrap the missing meta["val"] when 1. ph consists of scalar -# 2. meta["val"] is not properly set in dispatch_trace. -def _instantiate_missing_placeholder_val_with_real_inputs(gm, args): - phs = [node for node in gm.graph.nodes if node.op == "placeholder"] - if len(phs) != len(args): - raise ExportError( - ExportErrorType.NOT_SUPPORTED, - "Expect number of placeholders to be the same as user inputs.", - ) - for node, arg in zip(phs, args): - if "val" not in node.meta or node.meta["val"] is None: - node.meta["val"] = arg diff --git a/exir/capture/_config.py b/exir/capture/_config.py index 13ab5b3f46c..29058df7876 100644 --- a/exir/capture/_config.py +++ b/exir/capture/_config.py @@ -15,24 +15,9 @@ from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass -from executorch.exir.tracer import ExirDynamoConfig from torch.fx._compatibility import compatibility -@compatibility(is_backward_compatible=False) -@dataclass -class CaptureConfig: - pt2_mode: bool = True - enable_functionalization: bool = True - enable_dynamic_shape: bool = False # This flag does nothing if enable_aot is True - enable_aot: bool = ( - False # When it's true it implies automatic dynamic shapes via default dynamo config - ) - _dynamo_config: "ExirDynamoConfig" = field(default_factory=ExirDynamoConfig) - _unlift: bool = False # This flag does nothing if enable_aot is False. - _use_old_decomp_table: bool = False - - @compatibility(is_backward_compatible=False) @dataclass class EdgeCompileConfig: diff --git a/exir/capture/_unlift.py b/exir/capture/_unlift.py deleted file mode 100644 index 24c84494883..00000000000 --- a/exir/capture/_unlift.py +++ /dev/null @@ -1,159 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import copy -from typing import Any - -import torch -import torch._export -from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo -from torch.utils import _pytree as pytree - - -Val = Any - - -def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict): - count = 0 - # Step 1: make lifted params as get_attr - for node in gm.graph.nodes: - if node.op == "placeholder": - if count in inp_pos_to_param_buffer_name: - with gm.graph.inserting_after(node): - getattr_node = gm.graph.get_attr( - inp_pos_to_param_buffer_name[count] - ) - node.replace_all_uses_with(getattr_node) - metadata = node.meta - gm.graph.erase_node(node) - getattr_node.meta = metadata - count += 1 - - # Step 2: Fix the input/output of the graph now that we deleted - # some args. - gm.graph.lint() - names = [f"arg_{i}" for i in range(len(in_spec.children_specs))] - gm.graph._codegen = _PyTreeCodeGen( - _PyTreeInfo( - names, - in_spec, - out_spec, - ) - ) - gm.recompile() - - # Step 3: Find state references in HigherOrderOps and recursively - # fix them. - for node in gm.graph.nodes: - if node.op == "call_function" and node.target == torch.ops.cond: - pred, true_graph, false_graph, operands = node.args - true_gm = getattr(gm, true_graph.name) - false_gm = getattr(gm, false_graph.name) - inp_pos_to_param_buffer_name_for_submod = {} - real_operands = [] - for ix, operand in enumerate(operands): - if operand.target in inp_pos_to_param_buffer_name.values(): - inp_pos_to_param_buffer_name_for_submod[ix] = operand.target - true_gm.register_buffer(operand.target, state_dict[operand.target]) - false_gm.register_buffer(operand.target, state_dict[operand.target]) - else: - real_operands.append(operand) - node.args = (pred, true_graph, false_graph, real_operands) - - _, in_spec = pytree.tree_flatten(real_operands) - - _unlift( - true_gm, - inp_pos_to_param_buffer_name_for_submod, - in_spec, - None, - state_dict, - ) - _unlift( - false_gm, - inp_pos_to_param_buffer_name_for_submod, - in_spec, - None, - state_dict, - ) - if node.op == "call_function" and node.target.__name__ == "map_impl": - body_graph, num_mapped, *operands = node.args - body_gm = getattr(gm, body_graph.name) - inp_pos_to_buffer_name_for_submod = {} - real_operands = [] - for ix, operand in enumerate(operands): - if operand.target in inp_pos_to_param_buffer_name.values(): - inp_pos_to_buffer_name_for_submod[ix] = operand.target - body_gm.register_buffer(operand.target, state_dict[operand.target]) - else: - real_operands.append(operand) - node.args = (body_graph, num_mapped, *real_operands) - - _, in_spec = pytree.tree_flatten(real_operands) - - _unlift( - body_gm, inp_pos_to_buffer_name_for_submod, in_spec, None, state_dict - ) - gm.graph.lint() - gm.graph.eliminate_dead_code() - gm.recompile() - return gm - - -def unlift_exported_program_lifted_states( - ep: torch.export.exported_program.ExportedProgram, -): - new_gm = copy.deepcopy(ep.graph_module) - - # TODO Fix the period in params/buffers names later - # maybe a pass to replace graph signature with fixed names - param_buffer_name_to_corrected_name = {} - - for name, stuff in ep.state_dict.items(): - if name in ep.graph_signature.buffers: - if "." in name: - new_gm.register_buffer(name.replace(".", "_"), stuff) - param_buffer_name_to_corrected_name[name] = name.replace(".", "_") - else: - new_gm.register_buffer(name, stuff) - elif name in ep.graph_signature.parameters: - if "." in name: - new_gm.register_parameter(name.replace(".", "_"), stuff) - param_buffer_name_to_corrected_name[name] = name.replace(".", "_") - else: - new_gm.register_parameter(name, stuff) - else: - raise AssertionError("encountered not registered param/buffer") - - count = 0 - inp_pos_to_param_buffer_name = {} - for node in new_gm.graph.nodes: - if node.op == "placeholder": - if node.name in ep.graph_signature.inputs_to_buffers: - buffer_name = ep.graph_signature.inputs_to_buffers[node.name] - if buffer_name in param_buffer_name_to_corrected_name: - inp_pos_to_param_buffer_name[count] = ( - param_buffer_name_to_corrected_name[buffer_name] - ) - else: - inp_pos_to_param_buffer_name[count] = buffer_name - if node.name in ep.graph_signature.inputs_to_parameters: - param_name = ep.graph_signature.inputs_to_parameters[node.name] - if param_name in param_buffer_name_to_corrected_name: - inp_pos_to_param_buffer_name[count] = ( - param_buffer_name_to_corrected_name[param_name] - ) - else: - inp_pos_to_param_buffer_name[count] = param_name - count += 1 - new_gm = _unlift( - new_gm, - inp_pos_to_param_buffer_name, - ep.call_spec.in_spec, - ep.call_spec.out_spec, - ep.state_dict, - ) - return new_gm diff --git a/exir/capture/_utils.py b/exir/capture/_utils.py new file mode 100644 index 00000000000..3e54ab72189 --- /dev/null +++ b/exir/capture/_utils.py @@ -0,0 +1,29 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import contextmanager + +import torch + + +@contextmanager +def patch_forward(obj: torch.nn.Module, new_method): + """Helper method to make it easier to cleanly torch.export() a method on a + module that is not `forward`. + + TODO(suo): upstream this to torch.export.wrapper. + """ + # Save the original method + original_method = obj.forward + + # Patch the method + obj.forward = new_method.__get__(obj, obj.__class__) + + try: + yield + finally: + # Restore the original method + obj.forward = original_method diff --git a/exir/control_flow.py b/exir/control_flow.py index ff8016bd23e..25c3f439499 100644 --- a/exir/control_flow.py +++ b/exir/control_flow.py @@ -42,11 +42,11 @@ def branch_false(args): ret = control_flow.cond(pred, branch_true, branch_false, args) -and we can use the usual exir.capture() function. +and we can use the usual torch.export.export() function. .. code-block:: python - exir.capture(module, args) + torch.export.export(module, args) """ @@ -180,7 +180,7 @@ def tracing_context( """ A decorator function to annotate code path that we conditionally run during tracing. We need to annotate these paths for now because - during exir.capture(), the tracer does not know what's the proper + during torch.export.export(), the tracer does not know what's the proper local inputs to be passed to the untaken path. """ diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index ea9a2362da5..0ae850127fb 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -26,7 +26,7 @@ import torch from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge -from executorch.exir.capture._capture import patch_forward +from executorch.exir.capture._utils import patch_forward from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.error import InternalError from executorch.exir.inplace_aliasing import is_inplace_node diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 505fd50af8d..6dd315f0703 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -10,7 +10,6 @@ import copy import unittest -import executorch.exir as exir import torch from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ( @@ -31,6 +30,15 @@ from torch.fx.passes.infra.pass_base import PassBase, PassResult +class _FunctionModule(torch.nn.Module): + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, *args): + return self.fn(*args) + + class TestPassInfra(unittest.TestCase): def test_fail_passbase(self) -> None: """ @@ -251,11 +259,9 @@ def f(x: torch.Tensor) -> torch.Tensor: z = torch.add(y, x) return z - exported_program = ( - exir.capture(f, (torch.randn(10),), exir.CaptureConfig()) - .to_edge() - .exported_program - ) + exported_program = to_edge( + export(_FunctionModule(f), (torch.randn(10),), strict=True) + ).exported_program() pm = ExportedProgramPassManager(passes=[replace_add_with_mul]) result = pm(exported_program) @@ -441,11 +447,9 @@ def f(x: torch.Tensor) -> torch.Tensor: y = torch.add(x, x) return y - exported_program = ( - exir.capture(f, (torch.randn(10),), exir.CaptureConfig()) - .to_edge() - .exported_program - ) + exported_program = to_edge( + export(_FunctionModule(f), (torch.randn(10),), strict=True) + ).exported_program() pm = ExportedProgramPassManager( passes=[introduce_call_method], run_checks_after_each_pass=True diff --git a/src/README.md b/src/README.md index c718e53ad9c..f1e8abd1cf5 100644 --- a/src/README.md +++ b/src/README.md @@ -16,7 +16,7 @@ We are telling `pip install -e .` to treat `src/executorch` as the root of the ` This allows us to perform `pip install -e .` successfully and enables the execution of the following command: ```bash -python -c "from executorch.exir import CaptureConfig" +python -c "from executorch.exir import EdgeCompileConfig" ``` ## Long Term Solution diff --git a/test/end2end/exported_module.py b/test/end2end/exported_module.py index f2a3a700a6e..a3c979c1df2 100644 --- a/test/end2end/exported_module.py +++ b/test/end2end/exported_module.py @@ -15,7 +15,7 @@ import executorch.exir as exir import torch from executorch.exir import ExecutorchBackendConfig, ExecutorchProgramManager, to_edge -from executorch.exir.capture._capture import patch_forward +from executorch.exir.capture._utils import patch_forward from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode from executorch.exir.passes import ( DebugPass,