xref: /aosp_15_r20/external/executorch/exir/capture/_config.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7# pyre-unsafe
8from dataclasses import dataclass, field
9from typing import Dict, List, Optional, Union
10
11import torch
12
13from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
14from executorch.exir.pass_manager import PassType
15from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
16from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
17from executorch.exir.tracer import ExirDynamoConfig
18from torch.fx._compatibility import compatibility
19
20
21@compatibility(is_backward_compatible=False)
22@dataclass
23class CaptureConfig:
24    pt2_mode: bool = True
25    enable_functionalization: bool = True
26    enable_dynamic_shape: bool = False  # This flag does nothing if enable_aot is True
27    enable_aot: bool = (
28        False  # When it's true it implies automatic dynamic shapes via default dynamo config
29    )
30    _dynamo_config: "ExirDynamoConfig" = field(default_factory=ExirDynamoConfig)
31    _unlift: bool = False  # This flag does nothing if enable_aot is False.
32    _use_old_decomp_table: bool = False
33
34
35@compatibility(is_backward_compatible=False)
36@dataclass
37class EdgeCompileConfig:
38    # TODO(qihan): remove ability to opt out
39    _check_ir_validity: bool = True
40    # TODO(larryliu): remove this
41    _use_edge_ops: bool = True
42    # Allow core ATen ops check to be skipped for certain ops, but continue with the rest of the checks.
43    _core_aten_ops_exception_list: List[torch._ops.OpOverload] = field(
44        default_factory=list
45    )
46    _skip_type_promotion: bool = False
47    # TODO(gasoonjia): remove this
48    # TODO(T192537614): reenanle dim order as default
49    _skip_dim_order: bool = True
50
51
52@compatibility(is_backward_compatible=False)
53@dataclass
54class ExecutorchBackendConfig:
55    passes: List[PassType] = field(default_factory=list)
56
57    # A single memory planning pass can be defined for all the programs in the
58    # EdgeProgramManager or can be defined per program.
59    memory_planning_pass: Union[PassType, Dict[str, PassType]] = MemoryPlanningPass()
60    to_out_var_pass: PassType = ToOutVarPass(ignore_to_out_var_failure=False)
61    dynamic_memory_planning_mode: DynamicMemoryPlanningMode = (
62        DynamicMemoryPlanningMode.UPPER_BOUND
63    )
64    emit_stacktrace: bool = False
65
66    # Whether to move delegate data blobs from the Program into separate
67    # segments, rather than encoding those blobs in the flatbuffer data.
68    # This makes it possible to free those blobs at runtime.
69    extract_delegate_segments: bool = True
70
71    # When extracting segments, the starting offset of each segment will be
72    # aligned to this value (in bytes). Must be a power of two.
73    segment_alignment: int = 128
74
75    # If provided, the minimum alignment of tensor buffers in the program. Must
76    # be a power of 2. If not provided, uses the value in the schema file.
77    constant_tensor_alignment: Optional[int] = None
78
79    # If provided, the minimum alignment of delegate data in the program. Must
80    # be a power of 2. If not provided, uses the value in the schema file.
81    delegate_alignment: Optional[int] = None
82
83    # A single sym shape eval pass can be defined for all the programs in the
84    # EdgeProgramManager or can be defined per program.
85    sym_shape_eval_pass: Union[PassType, Dict[str, PassType]] = (
86        ConstraintBasedSymShapeEvalPass()
87    )
88
89    # If set to true, view_copy operations will be converted to lightweight
90    # view operations in the ET runtime
91    remove_view_copy: bool = True
92