|
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