|
75 | 75 |
|
76 | 76 | Val = Any
|
77 | 77 |
|
| 78 | +from typing import Any, Callable |
| 79 | + |
78 | 80 | from torch.library import Library
|
79 | 81 |
|
| 82 | +try: |
| 83 | + from executorch.exir.program.fb.logger import et_logger |
| 84 | +except ImportError: |
| 85 | + # Define a stub decorator that does nothing |
| 86 | + def et_logger(api_name: str) -> Callable[[Any], Any]: |
| 87 | + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: |
| 88 | + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: |
| 89 | + return func(self, *args, **kwargs) |
| 90 | + |
| 91 | + return wrapper |
| 92 | + |
| 93 | + return decorator |
| 94 | + |
| 95 | + |
80 | 96 | # This is the reserved namespace that is used to register ops to that will
|
81 | 97 | # be prevented from being decomposed during to_edge_transform_and_lower.
|
82 | 98 | edge_no_decomp_namespace = "EDGE_DO_NOT_DECOMP"
|
@@ -957,6 +973,7 @@ def _gen_edge_manager_for_partitioners(
|
957 | 973 | return edge_manager
|
958 | 974 |
|
959 | 975 |
|
| 976 | +@et_logger("to_edge_transform_and_lower") |
960 | 977 | def to_edge_transform_and_lower(
|
961 | 978 | programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
|
962 | 979 | transform_passes: Optional[
|
@@ -1110,6 +1127,7 @@ def to_edge_with_preserved_ops(
|
1110 | 1127 | )
|
1111 | 1128 |
|
1112 | 1129 |
|
| 1130 | +@et_logger("to_edge") |
1113 | 1131 | def to_edge(
|
1114 | 1132 | programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
|
1115 | 1133 | constant_methods: Optional[Dict[str, Any]] = None,
|
@@ -1204,8 +1222,10 @@ def exported_program(self, method_name: str = "forward") -> ExportedProgram:
|
1204 | 1222 | """
|
1205 | 1223 | Returns the ExportedProgram specified by 'method_name'.
|
1206 | 1224 | """
|
| 1225 | + |
1207 | 1226 | return self._edge_programs[method_name]
|
1208 | 1227 |
|
| 1228 | + @et_logger("transform") |
1209 | 1229 | def transform(
|
1210 | 1230 | self,
|
1211 | 1231 | passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]]],
|
@@ -1253,6 +1273,7 @@ def transform(
|
1253 | 1273 | new_programs, copy.deepcopy(self._config_methods), compile_config
|
1254 | 1274 | )
|
1255 | 1275 |
|
| 1276 | + @et_logger("to_backend") |
1256 | 1277 | def to_backend(
|
1257 | 1278 | self, partitioner: Union[Partitioner, Dict[str, Partitioner]]
|
1258 | 1279 | ) -> "EdgeProgramManager":
|
@@ -1296,6 +1317,7 @@ def to_backend(
|
1296 | 1317 | new_edge_programs, copy.deepcopy(self._config_methods), config
|
1297 | 1318 | )
|
1298 | 1319 |
|
| 1320 | + @et_logger("to_executorch") |
1299 | 1321 | def to_executorch(
|
1300 | 1322 | self,
|
1301 | 1323 | config: Optional[ExecutorchBackendConfig] = None,
|
|
0 commit comments