mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
bec7bdad47
Microbenchmarking `fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)]))`, before: ``` 30603618 function calls (29403419 primitive calls) in 13.744 seconds ``` after: ``` 25203549 function calls (24403352 primitive calls) in 12.090 seconds ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/148243 Approved by: https://github.com/oulgen
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
# Used to load and initialize polyfill handlers when importing torch._dynamo
|
|
# Please add a new import when adding a new polyfill module.
|
|
|
|
import importlib
|
|
from typing import TYPE_CHECKING
|
|
|
|
from .. import polyfills, trace_rules
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from types import ModuleType
|
|
|
|
|
|
# See also the TYPE_CHECKING block in torch/_dynamo/polyfills/__init__.py
|
|
POLYFILLED_MODULE_NAMES: tuple[str, ...] = (
|
|
"builtins",
|
|
"functools",
|
|
"itertools",
|
|
"operator",
|
|
"os",
|
|
"pytree",
|
|
"sys",
|
|
"fx",
|
|
)
|
|
POLYFILLED_MODULES: tuple["ModuleType", ...] = tuple(
|
|
importlib.import_module(f".{submodule}", package=polyfills.__name__)
|
|
for submodule in POLYFILLED_MODULE_NAMES
|
|
)
|
|
|
|
|
|
# Unregister the builtin functions from _builtin_function_ids to let them to be
|
|
# dispatched with the appropriate VariableTracker type. Otherwise, they will be
|
|
# dispatched with BuiltinVariable if present in _builtin_function_ids.
|
|
for polyfill_module in POLYFILLED_MODULES:
|
|
for polyfill_name in polyfill_module.__all__:
|
|
polyfill_handler = getattr(polyfill_module, polyfill_name)
|
|
original_fn = polyfill_handler.__torch_dynamo_original__
|
|
trace_rules._builtin_function_ids.remove(id(original_fn))
|