Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,233 @@
|
||||
"""
|
||||
TorchDynamo is a Python-level JIT compiler designed to make unmodified PyTorch programs faster.
|
||||
TorchDynamo hooks into the frame evaluation API in CPython (PEP 523) to dynamically modify Python
|
||||
bytecode right before it is executed. It rewrites Python bytecode in order to extract sequences of
|
||||
PyTorch operations into an FX Graph which is then just-in-time compiled with a customizable backend.
|
||||
It creates this FX Graph through bytecode analysis and is designed to mix Python execution with
|
||||
compiled backends to get the best of both worlds: usability and performance. This allows it to
|
||||
seamlessly optimize PyTorch programs, including those using modern Python features.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from . import (
|
||||
aot_compile,
|
||||
bytecode_debugger,
|
||||
config,
|
||||
convert_frame,
|
||||
eval_frame,
|
||||
functional_export,
|
||||
resume_execution,
|
||||
)
|
||||
from .backends.registry import list_backends, lookup_backend, register_backend
|
||||
from .callback import callback_handler, on_compile_end, on_compile_start
|
||||
from .code_context import code_context
|
||||
from .convert_frame import replay
|
||||
from .decorators import (
|
||||
allow_in_graph,
|
||||
assume_constant_result,
|
||||
disable,
|
||||
disable_nested_graph_breaks,
|
||||
disallow_in_graph,
|
||||
dont_skip_tracing,
|
||||
error_on_graph_break,
|
||||
forbid_in_graph,
|
||||
graph_break,
|
||||
is_dynamo_disable_recursive,
|
||||
mark_dynamic,
|
||||
mark_static,
|
||||
mark_static_address,
|
||||
maybe_mark_dynamic,
|
||||
nonstrict_trace,
|
||||
override_cudagraphs,
|
||||
patch_dynamo_config,
|
||||
run,
|
||||
set_stance,
|
||||
skip_frame,
|
||||
step_unsupported,
|
||||
substitute_in_graph,
|
||||
)
|
||||
from .eval_frame import (
|
||||
_reset_guarded_backend_cache,
|
||||
explain,
|
||||
export,
|
||||
is_dynamo_supported,
|
||||
is_inductor_supported,
|
||||
optimize,
|
||||
optimize_assert,
|
||||
OptimizedModule,
|
||||
reset_code,
|
||||
)
|
||||
|
||||
# pyrefly: ignore [deprecated]
|
||||
from .external_utils import is_compiling
|
||||
from .mutation_guard import GenerationTracker
|
||||
from .pgo import reset_code_state
|
||||
from .symbolic_convert import TensorifyState
|
||||
from .utils import (
|
||||
graph_break_reasons,
|
||||
guard_failures,
|
||||
orig_code_map,
|
||||
register_hook_for_recompile_user_context,
|
||||
reset_frame_count,
|
||||
reset_recompile_user_contexts,
|
||||
)
|
||||
|
||||
|
||||
# Register polyfill functions
|
||||
from .polyfills import loader as _ # usort: skip # noqa: F401
|
||||
|
||||
|
||||
__all__ = [
|
||||
"allow_in_graph",
|
||||
"assume_constant_result",
|
||||
"bytecode_debugger",
|
||||
"config",
|
||||
"disable",
|
||||
"disable_nested_graph_breaks",
|
||||
"disallow_in_graph",
|
||||
"dont_skip_tracing",
|
||||
"export",
|
||||
"explain",
|
||||
"forbid_in_graph",
|
||||
"graph_break",
|
||||
"is_compiling",
|
||||
"is_dynamo_disable_recursive",
|
||||
"list_backends",
|
||||
"lookup_backend",
|
||||
"mark_dynamic",
|
||||
"maybe_mark_dynamic",
|
||||
"mark_static",
|
||||
"mark_static_address",
|
||||
"nonstrict_trace",
|
||||
"optimize",
|
||||
"optimize_assert",
|
||||
"OptimizedModule",
|
||||
"patch_dynamo_config",
|
||||
"register_backend",
|
||||
"replay",
|
||||
"reset",
|
||||
"reset_recompile_user_contexts",
|
||||
"run",
|
||||
"override_cudagraphs",
|
||||
"error_on_graph_break",
|
||||
"set_recursion_limit",
|
||||
"set_stance",
|
||||
"skip_frame",
|
||||
"step_unsupported",
|
||||
"substitute_in_graph",
|
||||
]
|
||||
|
||||
# allowlist this for weights_only load of NJTs
|
||||
torch.serialization.add_safe_globals([torch._dynamo.decorators._DimRange])
|
||||
|
||||
if torch.manual_seed is torch.random.manual_seed:
|
||||
import torch.jit._builtins
|
||||
|
||||
# Wrap manual_seed with the disable decorator.
|
||||
# Can't do it at its implementation due to dependency issues.
|
||||
torch.manual_seed = torch._disable_dynamo(torch.manual_seed)
|
||||
# Add the new manual_seed to the builtin registry.
|
||||
torch.jit._builtins._register_builtin(torch.manual_seed, "aten::manual_seed")
|
||||
|
||||
|
||||
def reset() -> None:
|
||||
"""
|
||||
Clear all compile caches and restore initial state. This function is intended
|
||||
to reset Dynamo's state *as if* you had started a fresh process invocation, which
|
||||
makes it good for testing scenarios where you want to behave as if you started
|
||||
a new process. It does NOT affect any file system caches.
|
||||
|
||||
NB: this does NOT reset logging state. Don't use this to test logging
|
||||
initialization/reinitialization.
|
||||
"""
|
||||
# TODO: https://github.com/pytorch/pytorch/issues/139200
|
||||
import logging
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.info("torch._dynamo.reset")
|
||||
with convert_frame.compile_lock:
|
||||
reset_code_caches()
|
||||
convert_frame.input_codes.clear()
|
||||
reset_code_state()
|
||||
convert_frame.output_codes.clear()
|
||||
orig_code_map.clear()
|
||||
guard_failures.clear()
|
||||
graph_break_reasons.clear()
|
||||
resume_execution.ContinueExecutionCache.cache.clear()
|
||||
_reset_guarded_backend_cache()
|
||||
reset_frame_count()
|
||||
torch._dynamo.compiled_autograd.reset()
|
||||
convert_frame.FRAME_COUNTER = 0
|
||||
convert_frame.FRAME_COMPILE_COUNTER.clear()
|
||||
callback_handler.clear()
|
||||
GenerationTracker.clear()
|
||||
TensorifyState.clear()
|
||||
torch._dynamo.utils.warn_once_cache.clear()
|
||||
torch._C._autograd._saved_tensors_hooks_set_tracing(False)
|
||||
|
||||
# Reset cudagraph trees unconditionally since they are global state
|
||||
# not tied to a specific backend instance
|
||||
from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
|
||||
from torch._higher_order_ops.wrap import inductor_code_side_table
|
||||
|
||||
kernel_side_table.reset_table()
|
||||
inductor_code_side_table.reset_table()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
from torch._inductor.cudagraph_trees import reset_cudagraph_trees
|
||||
|
||||
reset_cudagraph_trees()
|
||||
|
||||
|
||||
def reset_code_caches() -> None:
|
||||
"""
|
||||
Clears in-memory code cache, which is what stores compiled products. This
|
||||
resets less state than :func:`reset` and is mostly only used for testing
|
||||
purposes.
|
||||
"""
|
||||
# TODO: https://github.com/pytorch/pytorch/issues/139200
|
||||
import logging
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.info("torch._dynamo.reset_code_caches")
|
||||
"""Clear compile caches that are keyed by code objects"""
|
||||
with convert_frame.compile_lock:
|
||||
reset_code_state()
|
||||
for weak_code in (
|
||||
convert_frame.input_codes.seen + convert_frame.output_codes.seen
|
||||
):
|
||||
code = weak_code()
|
||||
if code:
|
||||
reset_code(code)
|
||||
code_context.clear()
|
||||
|
||||
|
||||
def get_recursion_limit() -> int:
|
||||
"""
|
||||
Returns the internal dynamo recursion limit set by `torch._dynamo.set_recursion_limit`.
|
||||
|
||||
Returns -1 if no c recursion limit has been set.
|
||||
"""
|
||||
return torch._C._dynamo.eval_frame.get_c_recursion_limit()
|
||||
|
||||
|
||||
def set_recursion_limit(limit: int) -> None:
|
||||
"""
|
||||
Sets an internal dynamo recursion limit. The limit must be >= 1, or -1 to reset
|
||||
to the default (unset) state.
|
||||
|
||||
This is possibly needed in Python 3.12-3.13 since there is a separate C recursion limit
|
||||
that is not visible at the Python level. If you are getting RecursionErrors during
|
||||
Dynamo compilation and `sys.setrecursionlimit()` doesn't help, this function may alleviate
|
||||
the issue.
|
||||
|
||||
NOTE: this function does NOT call `sys.setrecursionlimit()` - the user is expected to manually
|
||||
call this if required. This is because the 2 recursion limits are not sync'd up - e.g. in
|
||||
Python 3.12, functions can be inline-evaluated, which apparently doesn't use up the C stack.
|
||||
|
||||
WARNING: increasing the recursion limit to an arbitrary large value may cause segfaults
|
||||
due to stack overflows! You can try also try to manually increase the stack size, e.g.
|
||||
with `$ ulimit -s ...`
|
||||
"""
|
||||
torch._C._dynamo.eval_frame.set_c_recursion_limit(limit)
|
||||
+262
@@ -0,0 +1,262 @@
|
||||
"""trace_wrapped(*args, fn) is equivalent to fn(*args), but with a twist:
|
||||
if you make_fx trace through this call, we will not actually trace into fn; instead,
|
||||
we will directly insert it as a call_function to fn in the graph.
|
||||
(Unlike make_fx, Dynamo WILL inline into fn.)
|
||||
You can think of this as a one off allow_in_graph equivalent for proxy tensor tracing.
|
||||
|
||||
Because proxy tensor tracing does not actually run the function, there are
|
||||
requirements on the behavior of fn. We are still figuring it out, but here is the current state:
|
||||
|
||||
1) fn SHOULD only take a single argument, which must be a tensor
|
||||
2) fn MUST return a new tensor with the same metadata as the original tensor
|
||||
(e.g., zeros_like(input) is a permissible implementation of fn).
|
||||
This is verified via an extra assert that is inserted into the traced graph.
|
||||
3) fn MAY have side effects, but it MAY NOT perform metadata mutation on other tensors
|
||||
participating in proxy tensor tracing (it MAY mutate other tensors, it MAY mutate Python state)
|
||||
These requirements stem from the requirement that we need to continue performing proxy tensor tracing,
|
||||
which assumes accurate fake tensor metadata, without actually running fn.
|
||||
In the future, we may allow for a "meta" function associated with fn to allow for more interesting input-output patterns.
|
||||
|
||||
Note that tensors / Python state are allowed to be mutated.
|
||||
This is relaxed constraint is not always sound, but it is sound for backward tracing with fake
|
||||
tensors as it takes place in AOTAutograd, as the backward pass is guaranteed not to depend on concrete
|
||||
tensor values (via fake tensor) or Python state (because the autograd engine doesn't depend on Python).
|
||||
|
||||
The intended use case for this function is to allow AOTAutograd to defer complex
|
||||
backward hooks to compiled autograd. AOTAutograd performs a make_fx trace which preserves
|
||||
the function call as is in the graph, and only when we Dynamo through the backward graph in
|
||||
compiled autograd do we inline into the function.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._C import DispatchKey
|
||||
from torch._higher_order_ops.utils import autograd_not_implemented
|
||||
from torch._ops import HigherOrderOperator, OpOverload
|
||||
from torch._subclasses import FakeTensorMode
|
||||
from torch.fx.experimental._backward_state import BackwardState
|
||||
from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree
|
||||
from torch.overrides import TorchFunctionMode
|
||||
from torch.utils._python_dispatch import _get_current_dispatch_mode
|
||||
from torch.utils._pytree import tree_map_only
|
||||
|
||||
|
||||
Tensor = torch.Tensor
|
||||
|
||||
|
||||
__all__ = ["trace_wrapped"]
|
||||
|
||||
|
||||
@torch.library.custom_op("flex_lib::zeros_and_scatter", mutates_args=()) # type: ignore[misc]
|
||||
def zeros_and_scatter(
|
||||
shape: list[int],
|
||||
indices: list[Tensor],
|
||||
vals: Tensor,
|
||||
) -> Tensor:
|
||||
"""Custom Op so that we can register a custom lowering for the new_output + scatter in the backwards pass"""
|
||||
grad = torch.zeros(shape, device=vals.device, dtype=vals.dtype)
|
||||
return torch.ops.aten.index_put(grad, indices, vals, accumulate=True)
|
||||
|
||||
|
||||
@zeros_and_scatter.register_fake # type: ignore[misc]
|
||||
def _(
|
||||
shape: list[int],
|
||||
indices: list[Tensor],
|
||||
vals: Tensor,
|
||||
) -> Tensor:
|
||||
return vals.new_empty(shape)
|
||||
|
||||
|
||||
@zeros_and_scatter.register_vmap # type: ignore[misc]
|
||||
def _(info, indims, shape, indices, value): # type: ignore[no-untyped-def]
|
||||
"""The batching rule is special in that it returns a tensor that is not batched"""
|
||||
indices_indims = indims[1]
|
||||
expanded_indices = []
|
||||
for idx, idx_indim in zip(indices, indices_indims):
|
||||
# The index is not a being batched, we should unsqueeze and expand to val
|
||||
if idx_indim is None:
|
||||
expanded_indices.append(idx.expand(value.shape))
|
||||
else:
|
||||
# the index is being part of the vmap batch, it should be the same size as val
|
||||
assert idx.shape == value.shape
|
||||
expanded_indices.append(idx)
|
||||
|
||||
out = torch.ops.flex_lib.zeros_and_scatter(
|
||||
shape,
|
||||
expanded_indices,
|
||||
value,
|
||||
)
|
||||
return out, None
|
||||
|
||||
|
||||
class ModIndex(torch.autograd.Function):
|
||||
generate_vmap_rule = True
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(x: Tensor, indices: list[Tensor]) -> Tensor:
|
||||
return torch.ops.aten.index(x, indices)
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None:
|
||||
x, indices = inputs
|
||||
ctx.save_for_backward(*indices)
|
||||
ctx.input_shape = x.shape
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, gradOut): # type: ignore[no-untyped-def]
|
||||
indices = ctx.saved_tensors
|
||||
return (
|
||||
torch.ops.flex_lib.zeros_and_scatter(
|
||||
ctx.input_shape,
|
||||
indices,
|
||||
gradOut,
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@torch._export.wrappers.allow_in_pre_dispatch_graph
|
||||
def apply(cls, *args, **kwargs): # type: ignore[no-untyped-def]
|
||||
return super().apply(*args, **kwargs)
|
||||
|
||||
|
||||
mod_index = ModIndex.apply
|
||||
|
||||
|
||||
class TransformGetItemToIndex(TorchFunctionMode):
|
||||
# This is needed since we want to support calling
|
||||
# A[q_idx], where q_idx is a scalar tensor in score_mod.
|
||||
# Today, when q_idx is a scalar tensor, we implicitly convert it to a python
|
||||
# scalar and create a view. We do not want that behavior in this case, so we
|
||||
# use this torchfunctionmode to override that behavior for score_mod
|
||||
# wherever we're running it.
|
||||
#
|
||||
# We also convert integer indices to 0-D tensors so that temp[0] produces
|
||||
# the same backward graph as temp[0 * q_idx] (zeros_and_scatter with atomic_add).
|
||||
def __torch_function__(
|
||||
self,
|
||||
func: OpOverload,
|
||||
types: tuple[torch._C._TensorMeta, ...],
|
||||
args: tuple[object, ...] = (),
|
||||
kwargs: dict[str, object] | None = None,
|
||||
) -> object:
|
||||
if func is torch.Tensor.__getitem__:
|
||||
tensor_to_index = args[0]
|
||||
assert isinstance(tensor_to_index, torch.Tensor)
|
||||
index_args = pytree.tree_leaves(args[1])
|
||||
if all(isinstance(x, (torch.Tensor, int)) for x in index_args):
|
||||
converted_indices = [
|
||||
torch.tensor(x, dtype=torch.int64, device=tensor_to_index.device)
|
||||
if isinstance(x, int)
|
||||
else x
|
||||
for x in index_args
|
||||
]
|
||||
return mod_index(tensor_to_index, converted_indices)
|
||||
return func(*args, **(kwargs or {}))
|
||||
|
||||
|
||||
def trace_wrapped(*args: Any, **kwargs: Any) -> Any:
|
||||
with torch.no_grad():
|
||||
return _trace_wrapped_op(*args, **kwargs)
|
||||
|
||||
|
||||
class TraceWrapped(HigherOrderOperator):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("trace_wrapped")
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return super().__call__(*args, **kwargs)
|
||||
|
||||
|
||||
# TODO(jansel): need to ensure this does not get DCEed
|
||||
_trace_wrapped_op = TraceWrapped()
|
||||
|
||||
|
||||
def _assert_meta(
|
||||
grad: torch.Tensor,
|
||||
size: tuple[int, ...],
|
||||
stride: tuple[int, ...],
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
assert grad.size() == size, "size mismatch"
|
||||
assert grad.stride() == stride, "stride mismatch"
|
||||
assert grad.dtype == dtype, "dtype mismatch"
|
||||
return grad
|
||||
|
||||
|
||||
@_trace_wrapped_op.py_impl(ProxyTorchDispatchMode)
|
||||
def inner_trace(
|
||||
mode: ProxyTorchDispatchMode,
|
||||
*args: Any,
|
||||
bw_state: BackwardState | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
def self_invoke(*args: Any, **dyn_kwargs: Any) -> Any:
|
||||
with torch.no_grad():
|
||||
return _trace_wrapped_op(*args, **dyn_kwargs, **kwargs)
|
||||
|
||||
def unwrap_proxies(x: Any) -> Any:
|
||||
if isinstance(x, torch.Tensor):
|
||||
return mode.tracer.unwrap_proxy(x) # type: ignore[union-attr]
|
||||
if isinstance(x, (list, tuple)):
|
||||
return type(x)(map(unwrap_proxies, x))
|
||||
if x is None:
|
||||
return None
|
||||
raise AssertionError(f"unhandled type: {type(x)}")
|
||||
|
||||
proxy_kwargs = {}
|
||||
if bw_state is not None:
|
||||
assert isinstance(bw_state, BackwardState) and bw_state.proxy is not None
|
||||
proxy_kwargs["bw_state"] = bw_state.proxy
|
||||
out_proxy = mode.tracer.create_proxy(
|
||||
"call_function",
|
||||
self_invoke,
|
||||
unwrap_proxies(args),
|
||||
proxy_kwargs,
|
||||
name="trace_wrapped",
|
||||
)
|
||||
|
||||
if args[0] is None:
|
||||
grad = args[1] # module backward hooks
|
||||
else:
|
||||
grad = args[0] # other backward hooks
|
||||
grad = tree_map_only(torch.Tensor, torch.empty_like, grad)
|
||||
track_tensor_tree(grad, out_proxy, constant=None, tracer=mode.tracer)
|
||||
return grad
|
||||
|
||||
|
||||
@_trace_wrapped_op.py_impl(FakeTensorMode)
|
||||
def inner_fake(*args: Any, **kwargs: Any) -> None:
|
||||
raise RuntimeError("This op should never be invoked here")
|
||||
|
||||
|
||||
@_trace_wrapped_op.py_impl(DispatchKey.CompositeExplicitAutograd)
|
||||
def _trace_wrapped_op_dense(*args: Any, fn: Any, **kwargs: Any) -> Any:
|
||||
mode = _get_current_dispatch_mode()
|
||||
assert mode is None, "Mode should never be enabled for CPU/CUDA key"
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
|
||||
_trace_wrapped_op.py_impl(DispatchKey.Autograd)(
|
||||
autograd_not_implemented(_trace_wrapped_op, deferred_error=True)
|
||||
)
|
||||
|
||||
|
||||
@_trace_wrapped_op.py_functionalize_impl
|
||||
def _trace_wrapped_functionalized(ctx: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
unwrapped_args = ctx.unwrap_tensors(args)
|
||||
with ctx.redispatch_to_next():
|
||||
return ctx.wrap_tensors(_trace_wrapped_op(*unwrapped_args, **kwargs))
|
||||
|
||||
|
||||
def autograd_function_backward_rewritten(original_backward: Any) -> Any:
|
||||
def new_backward(ctx: Any, *grads: Any) -> Any:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
grads = [g.contiguous() for g in grads]
|
||||
return original_backward(ctx, *grads)
|
||||
|
||||
return new_backward
|
||||
@@ -0,0 +1,539 @@
|
||||
import dataclasses
|
||||
import importlib
|
||||
import inspect
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
import tempfile
|
||||
import types
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import AbstractContextManager, ExitStack, nullcontext
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.fx
|
||||
from torch._dynamo.convert_frame import GraphRuntimeEnv
|
||||
from torch._dynamo.graph_utils import _graph_device_type
|
||||
from torch._dynamo.package import SystemInfo
|
||||
|
||||
from . import convert_frame
|
||||
from .aot_compile_types import (
|
||||
BundledAOTAutogradSerializableCallable,
|
||||
SerializableCallable,
|
||||
)
|
||||
from .hooks import Hooks
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .guards import GuardManagerWrapper
|
||||
from .package import SerializedCode, SourceInfo
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def bind_locals(
|
||||
signature: inspect.Signature, *args: Any, **kwargs: Any
|
||||
) -> dict[str, Any]:
|
||||
bound_arguments = signature.bind(*args, **kwargs)
|
||||
bound_arguments.apply_defaults()
|
||||
return bound_arguments.arguments
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompileArtifacts:
|
||||
signature: inspect.Signature
|
||||
guard_manager: Optional["GuardManagerWrapper"]
|
||||
guards_state: bytes
|
||||
backend_id: str
|
||||
compiled_fn: SerializableCallable
|
||||
original_code: types.CodeType
|
||||
runtime_env: GraphRuntimeEnv
|
||||
source_info: "SourceInfo"
|
||||
device_type: str
|
||||
backend_name: str
|
||||
system_info: SystemInfo = dataclasses.field(default_factory=SystemInfo.current)
|
||||
|
||||
def check_compatibility(self) -> None:
|
||||
current_system = SystemInfo.current()
|
||||
current_system.check_compatibility(self.system_info, self.device_type)
|
||||
|
||||
|
||||
class AOTCompilePickler(pickle.Pickler):
|
||||
def __init__(self, external_data: dict[str, object], buf: io.BytesIO) -> None:
|
||||
super().__init__(buf)
|
||||
self.external_data = external_data
|
||||
self.id_map: dict[int, str] = {
|
||||
id(value): key for key, value in external_data.items()
|
||||
}
|
||||
self.errors = {}
|
||||
|
||||
def persistent_id(self, obj: object) -> int | str | None:
|
||||
if id(obj) in self.id_map:
|
||||
return self.id_map[id(obj)]
|
||||
elif isinstance(obj, torch.nn.Module):
|
||||
self.errors[id(obj)] = obj
|
||||
return id(obj)
|
||||
else:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _unpickle_cell(cls, val: object) -> object:
|
||||
def _() -> object:
|
||||
return val
|
||||
|
||||
assert _.__closure__ is not None
|
||||
return _.__closure__[0]
|
||||
|
||||
@classmethod
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def _unpickle_bound_method(cls, func: Callable, base: object) -> types.MethodType:
|
||||
return types.MethodType(func, base)
|
||||
|
||||
@classmethod
|
||||
def _unpickle_module(cls, name: str) -> types.ModuleType:
|
||||
return importlib.import_module(name)
|
||||
|
||||
@classmethod
|
||||
def _unpickle_code(cls, serialized_code: "SerializedCode") -> types.CodeType:
|
||||
from torch._dynamo.package import SerializedCode
|
||||
|
||||
return SerializedCode.to_code_object(serialized_code)
|
||||
|
||||
@classmethod
|
||||
def _unpickle_nested_function(
|
||||
cls,
|
||||
code: types.CodeType,
|
||||
module: str,
|
||||
qualname: str,
|
||||
argdefs: tuple[object, ...] | None,
|
||||
closure: tuple[types.CellType, ...] | None,
|
||||
) -> types.FunctionType:
|
||||
f_globals = importlib.import_module(module).__dict__
|
||||
return types.FunctionType(code, f_globals, qualname, argdefs, closure)
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
def reducer_override(self, obj: Any) -> Any:
|
||||
if isinstance(obj, type((lambda x: lambda: x)(0).__closure__[0])): # type: ignore[index] # noqa: PLC3002
|
||||
return type(self)._unpickle_cell, (obj.cell_contents,)
|
||||
elif inspect.iscode(obj):
|
||||
from torch._dynamo.package import SerializedCode
|
||||
|
||||
return type(self)._unpickle_code, (SerializedCode.from_code_object(obj),)
|
||||
|
||||
elif inspect.ismodule(obj):
|
||||
return type(self)._unpickle_module, (obj.__name__,)
|
||||
elif inspect.ismethod(obj):
|
||||
"""
|
||||
By default, pickle will call getattr() directly on the self object
|
||||
for pickling bounded methods, this is not what we want, instead we
|
||||
always want to serialize the original function and the self object
|
||||
in their original form.
|
||||
"""
|
||||
func = obj.__func__
|
||||
method_self = obj.__self__
|
||||
inner_func = getattr(method_self, func.__name__)
|
||||
if inspect.ismethod(inner_func):
|
||||
inner_func = inner_func.__func__
|
||||
if func is not inner_func:
|
||||
return type(self)._unpickle_bound_method, (func, method_self)
|
||||
elif inspect.isfunction(obj):
|
||||
if "<locals>" in obj.__qualname__:
|
||||
return type(self)._unpickle_nested_function, (
|
||||
obj.__code__,
|
||||
obj.__module__,
|
||||
obj.__qualname__,
|
||||
obj.__defaults__,
|
||||
obj.__closure__,
|
||||
)
|
||||
|
||||
return NotImplemented
|
||||
|
||||
|
||||
class AOTCompileUnpickler(pickle.Unpickler):
|
||||
def __init__(self, external_data: dict[str, object], file: io.BytesIO) -> object:
|
||||
super().__init__(file)
|
||||
self.external_data = external_data
|
||||
|
||||
def persistent_load(self, key: str) -> object:
|
||||
if key not in self.external_data:
|
||||
raise RuntimeError(
|
||||
f"Missing required external reference to data: {key}. "
|
||||
"Please load AOT compiled function with "
|
||||
"`external_data=<external data dictionary>`"
|
||||
f"{self.external_data}"
|
||||
)
|
||||
return self.external_data[key]
|
||||
|
||||
|
||||
@dataclass
|
||||
class AOTCompileSaveResult:
|
||||
serialized_data: bytes
|
||||
|
||||
|
||||
def atomic_write_binary(file_path: str, data: bytes):
|
||||
dir_name = os.path.dirname(file_path) or "."
|
||||
|
||||
with tempfile.NamedTemporaryFile(
|
||||
dir=dir_name, delete=False, mode="wb"
|
||||
) as temp_file:
|
||||
temp_path = temp_file.name
|
||||
temp_file.write(data)
|
||||
temp_file.flush()
|
||||
os.fsync(temp_file.fileno())
|
||||
|
||||
os.replace(temp_path, file_path)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AOTCompiledFunction:
|
||||
_artifacts: CompileArtifacts
|
||||
_guard_check_enabled: bool = True
|
||||
_extra_globals: dict[str, object] | None = None
|
||||
|
||||
def prepare_f_locals(self, *args: object, **kwargs: object) -> dict[str, object]:
|
||||
f_locals: dict[str, object] = {}
|
||||
env = self._artifacts.runtime_env
|
||||
if env.closure:
|
||||
assert env.bytecode.co_freevars and len(env.closure) == len(
|
||||
env.bytecode.co_freevars
|
||||
)
|
||||
f_locals = {
|
||||
name: cell.cell_contents
|
||||
for name, cell in zip(env.bytecode.co_freevars, env.closure)
|
||||
}
|
||||
f_locals.update(bind_locals(self._artifacts.signature, *args, **kwargs))
|
||||
return f_locals
|
||||
|
||||
def guard_check(self, *args: Any, **kwargs: Any) -> bool:
|
||||
f_locals = self.prepare_f_locals(*args, **kwargs)
|
||||
assert self._artifacts.guard_manager is not None
|
||||
return self._artifacts.guard_manager.check(f_locals)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
from .package import load_guard_manager, load_guards_state
|
||||
|
||||
self._artifacts.check_compatibility()
|
||||
|
||||
self.fn = self._artifacts.runtime_env.forward_callable(
|
||||
self._artifacts.backend_id,
|
||||
self._artifacts.compiled_fn,
|
||||
extra_globals=self._extra_globals,
|
||||
)
|
||||
|
||||
if self._artifacts.guard_manager is None:
|
||||
guards_state = load_guards_state(self._artifacts.guards_state)
|
||||
self._artifacts.guard_manager = load_guard_manager(
|
||||
guards_state,
|
||||
self._artifacts.original_code,
|
||||
self.fn.__globals__,
|
||||
)
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
assert self._artifacts.guard_manager is not None
|
||||
if self._guard_check_enabled and not self.guard_check(*args, **kwargs):
|
||||
f_locals = self.prepare_f_locals(*args, **kwargs)
|
||||
reason = str(self._artifacts.guard_manager.check_verbose(f_locals))
|
||||
raise RuntimeError(f"GuardManager check failed, reason: {reason}")
|
||||
return self.fn(*args, **kwargs)
|
||||
|
||||
def source_info(self) -> "SourceInfo":
|
||||
return self._artifacts.source_info
|
||||
|
||||
def save_compiled_function(
|
||||
self, path: str, external_data: dict[str, Any] | None = None
|
||||
) -> AOTCompileSaveResult:
|
||||
result = type(self).serialize(self, external_data)
|
||||
atomic_write_binary(path, result.serialized_data)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def serialize(
|
||||
cls, fn: "AOTCompiledFunction", external_data: dict[str, Any] | None = None
|
||||
) -> AOTCompileSaveResult:
|
||||
from torch._dynamo.package import SerializedCode
|
||||
|
||||
state = fn._artifacts.__dict__.copy()
|
||||
state["guard_manager"] = None
|
||||
state["runtime_env"] = dataclasses.replace(
|
||||
state["runtime_env"],
|
||||
bytecode=SerializedCode.from_code_object(state["runtime_env"].bytecode),
|
||||
)
|
||||
compiled_fn = state["compiled_fn"]
|
||||
state["compiled_fn"] = (
|
||||
type(compiled_fn).deserialize_compile_artifacts,
|
||||
type(compiled_fn).serialize_compile_artifacts(compiled_fn),
|
||||
)
|
||||
state["original_code"] = SerializedCode.from_code_object(state["original_code"])
|
||||
buf = io.BytesIO()
|
||||
pickler = AOTCompilePickler(external_data or {}, buf)
|
||||
pickler.dump(state)
|
||||
if pickler.errors:
|
||||
raise RuntimeError(
|
||||
f"Failed to serialize the following objects: {list(pickler.errors.values())}\n"
|
||||
"Please mark these as external data by using `external_data={'key': ...}`"
|
||||
)
|
||||
return AOTCompileSaveResult(serialized_data=buf.getvalue())
|
||||
|
||||
@classmethod
|
||||
def deserialize(
|
||||
cls,
|
||||
data: bytes,
|
||||
f_globals: dict[str, object] | None = None,
|
||||
external_closure_data: dict[str, Any] | None = None,
|
||||
) -> "AOTCompiledFunction":
|
||||
from torch._dynamo.package import SerializedCode
|
||||
|
||||
f = io.BytesIO(data)
|
||||
f.seek(0)
|
||||
unpickler = AOTCompileUnpickler(external_closure_data or {}, f)
|
||||
state = unpickler.load()
|
||||
f.close()
|
||||
state["runtime_env"] = dataclasses.replace(
|
||||
state["runtime_env"],
|
||||
bytecode=SerializedCode.to_code_object(state["runtime_env"].bytecode),
|
||||
)
|
||||
deserializer, compiled_fn_state = state["compiled_fn"]
|
||||
with torch._inductor.config.patch(enable_autograd_for_aot=True):
|
||||
state["compiled_fn"] = deserializer(compiled_fn_state)
|
||||
state["original_code"] = SerializedCode.to_code_object(state["original_code"])
|
||||
|
||||
artifacts = CompileArtifacts(**state)
|
||||
return cls(artifacts, _extra_globals=f_globals)
|
||||
|
||||
def disable_guard_check(self) -> None:
|
||||
self._guard_check_enabled = False
|
||||
|
||||
|
||||
def aot_compile_fullgraph(
|
||||
model: Any,
|
||||
example_inputs: tuple[tuple[Any, ...], dict[str, Any]],
|
||||
hooks: Hooks,
|
||||
backend: Callable[[torch.fx.GraphModule, list[torch.Tensor]], SerializableCallable],
|
||||
dynamic: bool | None = None,
|
||||
) -> AOTCompiledFunction:
|
||||
from torch._dynamo.guards import CheckFunctionManager
|
||||
from torch._dynamo.package import SourceInfo
|
||||
from torch._dynamo.utils import dynamo_timed, get_metrics_context
|
||||
from torch._dynamo.variables.torch_function import (
|
||||
torch_function_mode_stack_state_mgr,
|
||||
)
|
||||
from torch._guards import TracingContext
|
||||
|
||||
args, kwargs = example_inputs
|
||||
|
||||
dynamic_ctx = nullcontext()
|
||||
if dynamic is not None:
|
||||
from torch._dynamo.eval_frame import set_enable_dynamic
|
||||
|
||||
dynamic_ctx = set_enable_dynamic(dynamic)
|
||||
|
||||
with (
|
||||
get_metrics_context(),
|
||||
dynamo_timed("fullgraph_capture"),
|
||||
torch._functorch.config.patch(strict_autograd_cache=True),
|
||||
dynamic_ctx,
|
||||
torch_function_mode_stack_state_mgr,
|
||||
):
|
||||
capture_output = convert_frame.fullgraph_capture(model, args, kwargs)
|
||||
graph_capture_output = capture_output.graph_capture_output
|
||||
assert graph_capture_output.output_graph is not None
|
||||
|
||||
if not hooks.guard_filter_fn:
|
||||
from torch._dynamo.types import GuardFilterEntry
|
||||
|
||||
def new_guard_filter_fn(
|
||||
guard_entries: Sequence[GuardFilterEntry],
|
||||
) -> Sequence[bool]:
|
||||
return [
|
||||
(
|
||||
not (
|
||||
g.is_global
|
||||
or g.guard_type
|
||||
in CheckFunctionManager.UNSUPPORTED_SERIALIZATION_GUARD_TYPES
|
||||
)
|
||||
)
|
||||
for g in guard_entries
|
||||
]
|
||||
|
||||
hooks.guard_filter_fn = new_guard_filter_fn
|
||||
|
||||
fn, _ = convert_frame.get_traced_fn(model)
|
||||
|
||||
backend_input = capture_output.backend_input
|
||||
assert backend_input is not None
|
||||
backend_input.graph_module._backend_id = backend_input.backend_id # type: ignore[assignment]
|
||||
device_type = _graph_device_type(backend_input.graph_module.graph)
|
||||
assert (
|
||||
backend_input.fake_mode.shape_env
|
||||
is graph_capture_output.output_graph.shape_env
|
||||
)
|
||||
tracing_context = TracingContext(backend_input.fake_mode)
|
||||
tracing_context.tensor_to_context = backend_input.tensor_to_context
|
||||
with (
|
||||
torch._guards.tracing(tracing_context),
|
||||
torch._functorch.config.patch(
|
||||
{
|
||||
"strict_autograd_cache": True,
|
||||
"bypass_autograd_cache_key": True,
|
||||
"bundled_autograd_cache": True,
|
||||
"force_non_lazy_backward_lowering": True,
|
||||
"force_autograd_cache": True,
|
||||
}
|
||||
),
|
||||
):
|
||||
compiled_fn = backend(
|
||||
backend_input.graph_module, backend_input.example_inputs
|
||||
)
|
||||
# If Inductor backend or AOTAutograd-based backend is used,
|
||||
# wrap the compiled_fn for serialization.
|
||||
# TODO: this should be replaced once we make the backend return the SerializableCallable directly.
|
||||
if (
|
||||
isinstance(backend, torch._TorchCompileInductorWrapper)
|
||||
or (
|
||||
hasattr(backend, "compiler_fn")
|
||||
and isinstance(
|
||||
backend.compiler_fn, torch._dynamo.backends.common.AotAutograd
|
||||
)
|
||||
)
|
||||
or (
|
||||
hasattr(compiled_fn, "serialize")
|
||||
and compiled_fn.serialize is not None
|
||||
)
|
||||
):
|
||||
compiled_fn = BundledAOTAutogradSerializableCallable(compiled_fn)
|
||||
|
||||
if not isinstance(compiled_fn, SerializableCallable):
|
||||
if hasattr(backend, "compiler_fn"):
|
||||
compiler_fn = backend.compiler_fn
|
||||
else:
|
||||
compiler_fn = backend
|
||||
raise RuntimeError(
|
||||
f"Compiled function type {type(compiled_fn)} (produced "
|
||||
+ f"from backend {compiler_fn}) does not implement SerializableCallable."
|
||||
)
|
||||
|
||||
# Temporarily restore the mode stack so guard expressions that
|
||||
# reference modes can evaluate, matching the compile_inner path.
|
||||
build_guards_ctx = ExitStack()
|
||||
if torch_function_mode_stack_state_mgr.stack:
|
||||
build_guards_ctx.enter_context(
|
||||
torch_function_mode_stack_state_mgr.temp_restore_stack()
|
||||
)
|
||||
with build_guards_ctx:
|
||||
check_fn = graph_capture_output.build_guards(
|
||||
fn.__code__, hooks=hooks, save=True, strict_error=True
|
||||
)
|
||||
|
||||
assert check_fn.guards_state is not None
|
||||
|
||||
source_info = SourceInfo(inlined_sources=set())
|
||||
for traced_code in graph_capture_output.traced_code:
|
||||
source_info.add_code(traced_code)
|
||||
|
||||
artifacts = CompileArtifacts(
|
||||
signature=convert_frame._get_signature(fn),
|
||||
guard_manager=check_fn.guard_manager,
|
||||
guards_state=check_fn.guards_state,
|
||||
backend_id=backend_input.backend_id,
|
||||
compiled_fn=compiled_fn,
|
||||
original_code=fn.__code__,
|
||||
runtime_env=graph_capture_output.get_runtime_env(),
|
||||
source_info=source_info,
|
||||
device_type=device_type,
|
||||
backend_name=getattr(backend, "compiler_name", "unknown"),
|
||||
)
|
||||
aot_compiled_fn = AOTCompiledFunction(
|
||||
_artifacts=artifacts, _extra_globals=fn.__globals__
|
||||
)
|
||||
|
||||
return aot_compiled_fn
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelInput:
|
||||
"""
|
||||
WIP type: represents a single model input
|
||||
Which consists of a tuple of arguments and a set of contexts in which to run the model.
|
||||
|
||||
For each ModelInput, we'll compile one full graph of the model, and then use the guards generated
|
||||
to dispatch between the compiled graphs.
|
||||
|
||||
|
||||
"""
|
||||
|
||||
args: tuple[Any]
|
||||
kwargs: dict[str, Any]
|
||||
contexts: list[AbstractContextManager[Any]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class AOTCompiledModel:
|
||||
# Represents a single forward function of a model along with dispatch
|
||||
# compiled_results is serializable. We require the model to deserialize again.
|
||||
model: torch.nn.Module
|
||||
compiled_results: list[AOTCompiledFunction]
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
for result in self.compiled_results:
|
||||
if result.guard_check(self.model, *args, **kwargs):
|
||||
return result(self.model, *args, **kwargs)
|
||||
# All guards failed, just run one of them and throw the guard check error.
|
||||
return self.compiled_results[0](self.model, *args, **kwargs)
|
||||
|
||||
def serialize(self) -> bytes:
|
||||
data: list[bytes] = []
|
||||
for result in self.compiled_results:
|
||||
data.append(AOTCompiledFunction.serialize(result).serialized_data)
|
||||
return pickle.dumps(data)
|
||||
|
||||
@classmethod
|
||||
def deserialize(cls, model: torch.nn.Module, data: bytes) -> "AOTCompiledModel":
|
||||
from torch._dynamo.utils import get_metrics_context
|
||||
from torch._guards import compile_context, CompileContext
|
||||
|
||||
results: list[bytes] = pickle.loads(data)
|
||||
compiled_results = []
|
||||
for result in results:
|
||||
with (
|
||||
compile_context(CompileContext(convert_frame.get_compile_id({}))),
|
||||
get_metrics_context(),
|
||||
):
|
||||
compiled_results.append(AOTCompiledFunction.deserialize(result))
|
||||
return cls(model, compiled_results)
|
||||
|
||||
|
||||
def aot_compile_module(
|
||||
model: torch.nn.Module,
|
||||
inputs: list[ModelInput],
|
||||
hooks: Hooks,
|
||||
backend: Callable[[torch.fx.GraphModule, list[torch.Tensor]], SerializableCallable],
|
||||
) -> AOTCompiledModel:
|
||||
"""
|
||||
Compiles a single nn.Module with any number of inputs, and returns a compiled forward function.
|
||||
"""
|
||||
|
||||
def compile_single_graph(model_input: ModelInput) -> AOTCompiledFunction:
|
||||
example_inputs = (model_input.args, model_input.kwargs)
|
||||
orig_forward = model.forward
|
||||
with ExitStack() as stack:
|
||||
for ctx in model_input.contexts:
|
||||
stack.enter_context(ctx)
|
||||
return aot_compile_fullgraph(
|
||||
orig_forward,
|
||||
example_inputs,
|
||||
hooks=hooks,
|
||||
backend=backend,
|
||||
)
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
compiled_results = []
|
||||
for model_input in inputs:
|
||||
log.info("Compiling input %s..", model_input)
|
||||
compiled_results.append(compile_single_graph(model_input))
|
||||
|
||||
assert len(compiled_results) > 0
|
||||
|
||||
return AOTCompiledModel(model, compiled_results)
|
||||
@@ -0,0 +1,211 @@
|
||||
import abc
|
||||
import importlib
|
||||
import pickle
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _serialize_triton_kernel(kernel: Any) -> tuple[str, str]:
|
||||
"""
|
||||
Serialize a triton kernel by extracting its module path and function name.
|
||||
Returns (module_path, function_name) tuple.
|
||||
|
||||
Triton JITFunction objects contain unpicklable _thread.RLock objects, so we
|
||||
serialize the import path instead and reimport on load.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the kernel cannot be serialized (missing attributes).
|
||||
"""
|
||||
fn = getattr(kernel, "fn", None)
|
||||
module_path = fn and getattr(fn, "__module__", None)
|
||||
func_name = fn and getattr(fn, "__name__", None)
|
||||
if fn is None or module_path is None or func_name is None:
|
||||
raise RuntimeError(
|
||||
f"Kernel fn missing __module__ or __name__: "
|
||||
f"module={module_path}, name={func_name}. "
|
||||
f"Cannot serialize for precompilation."
|
||||
)
|
||||
return (module_path, func_name)
|
||||
|
||||
|
||||
def _deserialize_triton_kernel(kernel_info: tuple[str, str]) -> Any:
|
||||
"""
|
||||
Deserialize a triton kernel by reimporting from its module.
|
||||
kernel_info is (module_path, function_name) tuple.
|
||||
"""
|
||||
module_path, func_name = kernel_info
|
||||
module = importlib.import_module(module_path)
|
||||
kernel = getattr(module, func_name)
|
||||
return kernel
|
||||
|
||||
|
||||
# Note: [Triton Kernel Side Table Serialization]
|
||||
#
|
||||
# When dynamo captures user-defined triton kernels, it creates FX graph nodes
|
||||
# (triton_kernel_wrapper_mutation/functional) with a `kernel_idx` parameter that
|
||||
# references the global `kernel_side_table` in triton_kernel_wrap.py. This side
|
||||
# table maps integer indices to actual triton kernel objects.
|
||||
#
|
||||
# For kernels that go through inductor's codegen path, this is fine - inductor
|
||||
# looks up the kernel from the side table at codegen time and embeds the kernel
|
||||
# source code directly into the generated wrapper. The compiled code doesn't
|
||||
# need the side table at runtime.
|
||||
#
|
||||
# However, not all triton kernels go through inductor codegen. When using
|
||||
# regional_inductor, only annotated regions are compiled by inductor. Triton
|
||||
# kernels outside these regions are executed via the FX interpreter, which
|
||||
# calls the higher-order op directly and needs the kernel to be in the side
|
||||
# table at runtime.
|
||||
#
|
||||
# When serializing/deserializing bundled AOT artifacts across process boundaries,
|
||||
# the kernel_side_table is empty in the new process, causing:
|
||||
# AssertionError: Kernel index X not found in id_to_kernel
|
||||
#
|
||||
# To fix this, we capture the kernel_side_table state during serialization and
|
||||
# restore it during deserialization. Kernels are serialized by their import path
|
||||
# (module_path, function_name) since triton JITFunction objects contain
|
||||
# unpicklable RLock objects.
|
||||
|
||||
|
||||
class SerializableCallable(abc.ABC):
|
||||
@classmethod
|
||||
@abc.abstractmethod
|
||||
def serialize_compile_artifacts(cls, fn: Any) -> bytes:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
@abc.abstractmethod
|
||||
def deserialize_compile_artifacts(cls, data: bytes) -> Any:
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
pass
|
||||
|
||||
|
||||
class GraphModuleSerializableCallable(SerializableCallable):
|
||||
def __init__(self, graph_module: torch.fx.GraphModule) -> None:
|
||||
assert isinstance(graph_module, torch.fx.GraphModule)
|
||||
self.graph_module = graph_module
|
||||
|
||||
@classmethod
|
||||
def serialize_compile_artifacts(
|
||||
cls, fn: "GraphModuleSerializableCallable"
|
||||
) -> bytes:
|
||||
from torch.fx._graph_pickler import GraphPickler, Options
|
||||
|
||||
state = fn.__dict__.copy()
|
||||
|
||||
graph_module = state["graph_module"]
|
||||
for node in graph_module.graph.nodes:
|
||||
node.meta.pop("nn_module_stack", None)
|
||||
node.meta.pop("source_fn_stack", None)
|
||||
node.meta.pop("example_value", None)
|
||||
|
||||
state["graph_module"] = GraphPickler.dumps(
|
||||
graph_module, Options(ops_filter=None)
|
||||
)
|
||||
return pickle.dumps(state)
|
||||
|
||||
@classmethod
|
||||
def deserialize_compile_artifacts(cls, data: bytes) -> Any:
|
||||
from torch._subclasses import FakeTensorMode
|
||||
from torch.fx._graph_pickler import GraphPickler
|
||||
from torch.fx.experimental.symbolic_shapes import ShapeEnv
|
||||
|
||||
state = pickle.loads(data)
|
||||
|
||||
fake_mode = FakeTensorMode(shape_env=ShapeEnv())
|
||||
state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
|
||||
assert isinstance(state["graph_module"], torch.fx.GraphModule)
|
||||
state["graph_module"].recompile()
|
||||
|
||||
return cls(**state)
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return self.graph_module.forward(*args, **kwargs)
|
||||
|
||||
|
||||
class BundledAOTAutogradSerializableCallable(SerializableCallable):
|
||||
"""
|
||||
Represents a serializable callable generated by compile_fx.
|
||||
This class wraps around the compiled function generated by AOTAutograd.
|
||||
|
||||
TODO: Instead of using PrecompileContext to grab it from AOTAutograd,
|
||||
this object should be what's *returned* by aot_module_simplified.
|
||||
We'll do that refactor in a later PR.
|
||||
"""
|
||||
|
||||
def __init__(self, compiled_fn: Any) -> None:
|
||||
"""
|
||||
Takes in a BundledAOTAutogradCacheArtifact, which is the serialized form
|
||||
of a compiled function generated by AOTAutograd.
|
||||
"""
|
||||
assert hasattr(compiled_fn, "serialize")
|
||||
self.compiled_fn = compiled_fn
|
||||
|
||||
def __getattr__(self, attr: Any) -> Any:
|
||||
return getattr(self.compiled_fn, attr)
|
||||
|
||||
@classmethod
|
||||
def serialize_compile_artifacts(
|
||||
cls, fn: "BundledAOTAutogradSerializableCallable"
|
||||
) -> bytes:
|
||||
from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
|
||||
|
||||
# See Note: [Triton Kernel Side Table Serialization]
|
||||
# Capture triton kernel side table state BEFORE serialization.
|
||||
triton_kernels: dict[int, tuple[str, str]] = {
|
||||
idx: _serialize_triton_kernel(kernel)
|
||||
for idx, kernel in kernel_side_table.id_to_kernel.items()
|
||||
}
|
||||
triton_constant_args: dict[int, dict[str, Any]] = dict(
|
||||
kernel_side_table.constant_args
|
||||
)
|
||||
|
||||
with torch._functorch.config.patch("bundled_autograd_cache", True):
|
||||
serialized_entry = fn.compiled_fn.serialize()
|
||||
# Bundle the triton kernel side table with the serialized entry
|
||||
bundle = (serialized_entry, triton_kernels, triton_constant_args)
|
||||
result = pickle.dumps(bundle)
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def deserialize_compile_artifacts(cls, data: bytes) -> Any:
|
||||
from torch._functorch._aot_autograd.aot_autograd_result import (
|
||||
deserialize_bundled_cache_entry,
|
||||
)
|
||||
from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table
|
||||
|
||||
bundle = pickle.loads(data)
|
||||
|
||||
# Handle both old format (just entry) and new format (entry, kernels, const_args)
|
||||
if isinstance(bundle, tuple) and len(bundle) == 3:
|
||||
entry, triton_kernels, triton_constant_args = bundle
|
||||
else:
|
||||
# Backwards compatibility with old serialized artifacts
|
||||
entry = bundle
|
||||
# pyrefly: ignore [implicit-any]
|
||||
triton_kernels = {}
|
||||
# pyrefly: ignore [implicit-any]
|
||||
triton_constant_args = {}
|
||||
|
||||
# See Note: [Triton Kernel Side Table Serialization]
|
||||
# Restore triton kernel side table BEFORE deserializing the compiled function.
|
||||
# The compiled function may reference kernels by index if any triton kernels
|
||||
# don't go through inductor codegen (e.g., triton kernels outside of
|
||||
# regional_inductor compiled regions).
|
||||
for idx, kernel_info in triton_kernels.items():
|
||||
kernel = _deserialize_triton_kernel(kernel_info)
|
||||
kernel_side_table.id_to_kernel[idx] = kernel
|
||||
kernel_side_table.kernel_to_id[kernel] = idx
|
||||
|
||||
for idx, args in triton_constant_args.items():
|
||||
kernel_side_table.constant_args[idx] = args
|
||||
|
||||
compiled_fn = deserialize_bundled_cache_entry(entry)
|
||||
return cls(compiled_fn)
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return self.compiled_fn(*args, **kwargs)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
This module provides common utilities and base classes for TorchDynamo backends.
|
||||
|
||||
Key components:
|
||||
- AotAutograd: Base class for implementing AOT (Ahead-of-Time) autograd backends
|
||||
- Backend utilities for handling:
|
||||
- Fake tensor conversion
|
||||
- Device/dtype detection from inputs
|
||||
- Memory efficient fusion
|
||||
- Graph flattening
|
||||
- Common compiler configurations
|
||||
|
||||
The utilities here are used by various backend implementations to handle
|
||||
common operations and provide consistent behavior across different backends.
|
||||
AOT autograd functionality is particularly important as it enables ahead-of-time
|
||||
optimization of both forward and backward passes.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import logging
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
from typing import Any
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from torch._dynamo import disable
|
||||
from torch._dynamo.exc import TensorifyScalarRestartAnalysis
|
||||
from torch._dynamo.utils import counters, defake, flatten_graph_inputs
|
||||
from torch._functorch.aot_autograd import (
|
||||
aot_module_simplified,
|
||||
SerializableAOTDispatchCompiler,
|
||||
)
|
||||
from torch.utils._python_dispatch import _disable_current_modes
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
class AotAutograd:
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self.__name__ = "compiler_fn"
|
||||
self.kwargs = kwargs
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: Sequence[Any], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
if kwargs:
|
||||
log.warning("aot_autograd-based backend ignoring extra kwargs %s", kwargs)
|
||||
|
||||
if any(isinstance(x, (list, tuple, dict)) for x in example_inputs):
|
||||
return flatten_graph_inputs(
|
||||
gm,
|
||||
example_inputs,
|
||||
self,
|
||||
)
|
||||
|
||||
# Hack to get around circular import problems with aot_eager_decomp_partition
|
||||
if callable(self.kwargs.get("decompositions")):
|
||||
self.kwargs["decompositions"] = self.kwargs["decompositions"]()
|
||||
|
||||
# NB: dont delete counter increment
|
||||
counters["aot_autograd"]["total"] += 1
|
||||
use_fallback = False
|
||||
|
||||
if use_fallback:
|
||||
log.debug("Unable to use AOT Autograd because graph has mutation")
|
||||
counters["aot_autograd"]["not_ok"] += 1
|
||||
return gm
|
||||
|
||||
def wrap_bw_compiler(bw_compiler_fn: Callable[P, R]) -> Callable[..., R]:
|
||||
def _wrapped_bw_compiler(*args: P.args, **kwargs: P.kwargs) -> R:
|
||||
# Note [Wrapping bw_compiler in disable]
|
||||
# The two disables here:
|
||||
# - stop TorchDynamo from trying to compile the bw_compiler function itself
|
||||
# - stop TorchDynamo from trying to compile our the generated backwards pass bw_compiler produces
|
||||
|
||||
return disable(
|
||||
disable(
|
||||
bw_compiler_fn, reason="do not trace backward compiler function"
|
||||
)(*args, **kwargs), # type: ignore[misc]
|
||||
reason="do not trace generated backwards pass",
|
||||
)
|
||||
|
||||
_wrapped_bw_compiler._is_wrapped_bw_compiler = ( # pyrefly: ignore [missing-attribute]
|
||||
True
|
||||
)
|
||||
return _wrapped_bw_compiler
|
||||
|
||||
bw_compiler = self.kwargs.get("bw_compiler") or self.kwargs["fw_compiler"]
|
||||
|
||||
if isinstance(bw_compiler, SerializableAOTDispatchCompiler):
|
||||
bw_compiler.compiler_fn = wrap_bw_compiler(bw_compiler.compiler_fn)
|
||||
elif getattr(bw_compiler, "_is_wrapped_bw_compiler", False):
|
||||
bw_compiler.compiler_fn = bw_compiler
|
||||
else:
|
||||
bw_compiler = wrap_bw_compiler(bw_compiler)
|
||||
|
||||
self.kwargs["bw_compiler"] = bw_compiler
|
||||
self.kwargs["inference_compiler"] = (
|
||||
self.kwargs.get("inference_compiler") or self.kwargs["fw_compiler"]
|
||||
)
|
||||
|
||||
from functorch.compile import nop
|
||||
from torch._inductor.debug import enable_aot_logging
|
||||
|
||||
# debug asserts slow down compile time noticeably,
|
||||
# So only default them on when the aot_eager backend is used.
|
||||
if self.kwargs.get("fw_compiler", None) is nop:
|
||||
patch_config: contextlib.AbstractContextManager[Any] = patch(
|
||||
"functorch.compile.config.debug_assert", True
|
||||
)
|
||||
else:
|
||||
patch_config = contextlib.nullcontext()
|
||||
|
||||
try:
|
||||
# NB: NOT cloned!
|
||||
with enable_aot_logging(), patch_config:
|
||||
cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
|
||||
counters["aot_autograd"]["ok"] += 1
|
||||
return disable(cg, reason="do not trace AOT-compiled graph")
|
||||
except TensorifyScalarRestartAnalysis:
|
||||
raise
|
||||
except Exception:
|
||||
counters["aot_autograd"]["not_ok"] += 1
|
||||
raise
|
||||
|
||||
|
||||
def aot_autograd(**kwargs: Any) -> AotAutograd:
|
||||
return AotAutograd(**kwargs)
|
||||
|
||||
|
||||
def mem_efficient_fusion_kwargs(use_decomps: bool) -> dict[str, Any]:
|
||||
from functorch.compile import (
|
||||
default_decompositions,
|
||||
min_cut_rematerialization_partition,
|
||||
ts_compile,
|
||||
)
|
||||
|
||||
kwargs = {
|
||||
# these are taken from memory_efficient_fusion()
|
||||
"fw_compiler": ts_compile,
|
||||
"bw_compiler": ts_compile,
|
||||
"partition_fn": min_cut_rematerialization_partition,
|
||||
}
|
||||
|
||||
if use_decomps:
|
||||
# pyrefly: ignore [bad-typed-dict-key]
|
||||
kwargs["decompositions"] = default_decompositions
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
def fake_tensor_unsupported(fn: Callable[[Any, list[Any], Any], R]) -> Any:
|
||||
"""
|
||||
Decorator for backends that need real inputs. We swap out fake
|
||||
tensors for zero tensors.
|
||||
"""
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapper(model: Any, inputs: Any, **kwargs: Any) -> Any:
|
||||
with _disable_current_modes():
|
||||
inputs = list(map(defake, inputs))
|
||||
return fn(model, inputs, **kwargs) # type: ignore[call-arg]
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def device_from_inputs(example_inputs: Iterable[Any]) -> torch.device:
|
||||
for x in example_inputs:
|
||||
if hasattr(x, "device"):
|
||||
return x.device
|
||||
return torch.device("cpu") # Default fallback
|
||||
|
||||
|
||||
def dtype_from_inputs(example_inputs: Iterable[Any]) -> torch.dtype:
|
||||
for x in example_inputs:
|
||||
if hasattr(x, "dtype"):
|
||||
return x.dtype
|
||||
return torch.float32 # Default fallback
|
||||
@@ -0,0 +1,299 @@
|
||||
"""
|
||||
This module implements CUDA graphs support for TorchDynamo backends.
|
||||
|
||||
CUDA graphs allow for capturing and replaying GPU operations, which can significantly
|
||||
reduce CPU overhead in GPU-accelerated PyTorch models. This module provides:
|
||||
|
||||
- CUDA graph creation and management for both forward and backward passes
|
||||
- Input mutation detection and handling
|
||||
- Device compatibility checking
|
||||
- Stack trace management for debugging
|
||||
- Integration with TorchInductor's cudagraph trees
|
||||
|
||||
The backend supports two main modes:
|
||||
1. cudagraphs: Full CUDA graph support with both forward and backward pass optimization
|
||||
2. cudagraphs_inner: Lower-level CUDA graph implementation used for benchmarking
|
||||
|
||||
Key components:
|
||||
- CudagraphsBackend: Main backend class for CUDA graph integration
|
||||
- Mutation detection utilities to ensure graph safety
|
||||
- Device mapping and compatibility checks
|
||||
- Stack trace collection for debugging
|
||||
"""
|
||||
|
||||
import functools
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.fx
|
||||
from torch._dynamo import config
|
||||
from torch._dynamo.backends.common import aot_autograd
|
||||
from torch._dynamo.backends.debugging import boxed_nop
|
||||
from torch._inductor.cudagraph_utils import (
|
||||
BoxedDeviceIndex,
|
||||
check_multiple_devices_or_any_cpu_nodes,
|
||||
format_default_skip_message,
|
||||
get_mutation_stack_trace,
|
||||
get_placeholder_info,
|
||||
log_cudagraph_skip_and_bump_counter,
|
||||
)
|
||||
from torch._inductor.utils import (
|
||||
BoxedBool,
|
||||
count_tangents,
|
||||
get_first_incompatible_cudagraph_node,
|
||||
num_fw_fixed_arguments,
|
||||
output_node,
|
||||
)
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
|
||||
from .registry import register_backend
|
||||
|
||||
|
||||
def find_input_mutations(g: torch.fx.Graph) -> set[int]:
|
||||
def meta_fk(meta: dict[str, Any]) -> Any:
|
||||
return meta["val"] if "val" in meta else meta["fake_result"]
|
||||
|
||||
inputs = defaultdict(set)
|
||||
input_idx = 0
|
||||
mutated_inputs = set()
|
||||
for n in g.nodes:
|
||||
if n.op == "placeholder":
|
||||
if isinstance(meta_fk(n.meta), torch.Tensor):
|
||||
inputs[StorageWeakRef(meta_fk(n.meta)._typed_storage())].add(input_idx)
|
||||
input_idx += 1
|
||||
elif n.op == "call_function":
|
||||
if not hasattr(n.target, "_schema"):
|
||||
continue
|
||||
|
||||
schema = n.target._schema
|
||||
for i, arg in enumerate(schema.arguments):
|
||||
if i < len(n.args):
|
||||
argument = n.args[i]
|
||||
else:
|
||||
if arg.name not in n.kwargs:
|
||||
continue
|
||||
argument = n.kwargs[arg.name]
|
||||
mut_arg = False
|
||||
if arg.alias_info:
|
||||
if arg.alias_info.is_write:
|
||||
mut_arg = True
|
||||
if mut_arg:
|
||||
# TODO: not correct for args that contain tensors in a struct
|
||||
# like list
|
||||
mutated_inputs |= inputs[
|
||||
StorageWeakRef(meta_fk(argument.meta)._typed_storage())
|
||||
]
|
||||
|
||||
# TODO: error on unrecognized nodes
|
||||
return mutated_inputs
|
||||
|
||||
|
||||
def get_device_node_mapping(
|
||||
gm: torch.fx.GraphModule,
|
||||
) -> dict[torch.device, torch.fx.Node]:
|
||||
device_node_mapping: dict[torch.device, torch.fx.Node] = {}
|
||||
for n in gm.graph.nodes:
|
||||
t = n.meta.get("val", None)
|
||||
if isinstance(t, torch.Tensor) and t.device not in device_node_mapping:
|
||||
device_node_mapping[t.device] = n
|
||||
return device_node_mapping
|
||||
|
||||
|
||||
def check_for_mutation_ignore_cuda_graph_managed_tensor(
|
||||
aot_model: torch.fx.GraphModule, num_fixed: int
|
||||
) -> str | None:
|
||||
mutation_indices = find_input_mutations(aot_model.graph) - set(range(num_fixed))
|
||||
if not mutation_indices:
|
||||
return None
|
||||
|
||||
placeholders = get_placeholder_info(aot_model.graph)
|
||||
return get_mutation_stack_trace(placeholders, mutation_indices)
|
||||
|
||||
|
||||
def check_for_skip(aot_model: torch.fx.GraphModule, num_fixed: int) -> str | None:
|
||||
if not config.cudagraph_backend_support_input_mutation:
|
||||
if mut_skip := check_for_mutation_ignore_cuda_graph_managed_tensor(
|
||||
aot_model, num_fixed
|
||||
):
|
||||
return mut_skip
|
||||
|
||||
if skip := check_multiple_devices_or_any_cpu_nodes(
|
||||
get_device_node_mapping(aot_model)
|
||||
):
|
||||
return skip
|
||||
|
||||
if node := get_first_incompatible_cudagraph_node(aot_model):
|
||||
return format_default_skip_message(f"incompatible op ({node.name})")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_device_index(gm: torch.fx.GraphModule) -> int:
|
||||
device = next(iter(get_device_node_mapping(gm)))
|
||||
assert device.type == "cuda"
|
||||
return device.index
|
||||
|
||||
|
||||
def get_stack_traces(gm: torch.fx.GraphModule) -> list[str | None]:
|
||||
output = output_node(gm)
|
||||
assert len(output.args) == 1
|
||||
args = output.args[0]
|
||||
if not hasattr(args, "__iter__"):
|
||||
return []
|
||||
return [
|
||||
(arg.stack_trace if isinstance(arg, torch.fx.node.Node) else None)
|
||||
for arg in args # type: ignore[union-attr]
|
||||
]
|
||||
|
||||
|
||||
def cudagraphs(dynamo_model: torch.fx.GraphModule, dynamo_inputs: Sequence[Any]) -> Any:
|
||||
from torch._inductor.cudagraph_trees import cudagraphify_impl
|
||||
|
||||
do_cudagraphs = BoxedBool(True)
|
||||
boxed_device_index = BoxedDeviceIndex(None)
|
||||
|
||||
def forward_cudagraphs(
|
||||
aot_model: torch.fx.GraphModule,
|
||||
aot_inputs: list[Any],
|
||||
is_inference: bool = False,
|
||||
) -> Any:
|
||||
interp = boxed_nop(aot_model, aot_inputs)
|
||||
fixed = num_fw_fixed_arguments(len(dynamo_inputs), len(aot_inputs))
|
||||
if skip_msg := check_for_skip(aot_model, fixed):
|
||||
BoxedBool.disable(do_cudagraphs)
|
||||
log_cudagraph_skip_and_bump_counter(
|
||||
f"skipping cudagraphs due to {skip_msg}"
|
||||
)
|
||||
return interp
|
||||
|
||||
boxed_device_index.set(get_device_index(aot_model))
|
||||
out = cudagraphify_impl(
|
||||
interp,
|
||||
aot_inputs,
|
||||
range(fixed),
|
||||
device_index=boxed_device_index.value,
|
||||
is_backward=False,
|
||||
is_inference=is_inference,
|
||||
stack_traces=get_stack_traces(aot_model),
|
||||
placeholders=get_placeholder_info(aot_model.graph),
|
||||
mutated_input_idxs=find_input_mutations(aot_model.graph),
|
||||
)
|
||||
out._boxed_call = True # type: ignore[attr-defined]
|
||||
return out
|
||||
|
||||
def backward_cudagraphs(
|
||||
aot_model: torch.fx.GraphModule, aot_inputs: list[Any]
|
||||
) -> Any:
|
||||
interp = boxed_nop(aot_model, aot_inputs)
|
||||
if not do_cudagraphs:
|
||||
return aot_model
|
||||
|
||||
fixed = count_tangents(aot_model)
|
||||
if skip_msg := check_for_skip(aot_model, fixed):
|
||||
log_cudagraph_skip_and_bump_counter(
|
||||
f"skipping cudagraphs due to {skip_msg}"
|
||||
)
|
||||
|
||||
# See [Backward Generation Handling]
|
||||
device_idx = boxed_device_index.value
|
||||
if device_idx is None:
|
||||
device_idx = 0 # Default to device 0 if not set
|
||||
manager = torch._inductor.cudagraph_trees.get_manager(
|
||||
device_idx, create_if_none_exists=False
|
||||
)
|
||||
assert manager is not None
|
||||
|
||||
def fn(inputs: list[Any]) -> Any:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
manager.set_to_running_backward()
|
||||
return aot_model(inputs)
|
||||
|
||||
fn._boxed_call = True # type: ignore[attr-defined]
|
||||
return fn
|
||||
|
||||
out = cudagraphify_impl(
|
||||
interp,
|
||||
aot_inputs,
|
||||
range(fixed),
|
||||
device_index=get_device_index(aot_model),
|
||||
is_backward=True,
|
||||
is_inference=False,
|
||||
stack_traces=get_stack_traces(aot_model),
|
||||
placeholders=get_placeholder_info(aot_model.graph),
|
||||
mutated_input_idxs=find_input_mutations(aot_model.graph),
|
||||
)
|
||||
out._boxed_call = True # type: ignore[attr-defined]
|
||||
return out
|
||||
|
||||
aot_cudagraphs = aot_autograd(
|
||||
fw_compiler=forward_cudagraphs,
|
||||
bw_compiler=backward_cudagraphs,
|
||||
inference_compiler=functools.partial(forward_cudagraphs, is_inference=True),
|
||||
keep_inference_input_mutations=torch._dynamo.config.cudagraph_backend_keep_input_mutation,
|
||||
)
|
||||
return aot_cudagraphs(dynamo_model, dynamo_inputs)
|
||||
|
||||
|
||||
class CudagraphsBackend:
|
||||
compiler_name = "cudagraphs"
|
||||
|
||||
@staticmethod
|
||||
def reset() -> None:
|
||||
from torch._inductor.cudagraph_trees import reset_cudagraph_trees
|
||||
|
||||
reset_cudagraph_trees()
|
||||
|
||||
@staticmethod
|
||||
def __call__(model: torch.fx.GraphModule, inputs: Sequence[Any]) -> Any:
|
||||
return cudagraphs(model, inputs)
|
||||
|
||||
|
||||
# aot_cudagraphs only applies CUDA graphs to the graph. It is also helpful
|
||||
# for debugging and can serve as a perf baseline.
|
||||
register_backend(name="cudagraphs", compiler_fn=CudagraphsBackend())
|
||||
|
||||
|
||||
def cudagraphs_inner(
|
||||
model: Callable[..., Any],
|
||||
inputs: Sequence[Any],
|
||||
copy_outputs: bool = True,
|
||||
copy_inputs: bool = True,
|
||||
) -> Callable[..., Sequence[Any]]:
|
||||
"""This isn't registered as a backend, but is used in some benchmarks"""
|
||||
assert isinstance(inputs, (list, tuple))
|
||||
if copy_inputs:
|
||||
static_inputs = [torch.zeros_like(x) for x in inputs]
|
||||
else:
|
||||
static_inputs = list(inputs)
|
||||
|
||||
# warmup
|
||||
torch.cuda.synchronize()
|
||||
stream = torch.cuda.Stream()
|
||||
stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(stream):
|
||||
model(*inputs)
|
||||
stream.synchronize()
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# record
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=stream):
|
||||
static_outputs = model(*static_inputs)
|
||||
if not isinstance(static_outputs, (list, tuple)):
|
||||
static_outputs = (static_outputs,)
|
||||
|
||||
def run(*new_inputs: Any) -> Sequence[Any]:
|
||||
assert len(static_inputs) == len(new_inputs)
|
||||
if copy_inputs:
|
||||
for dst, src in zip(static_inputs, new_inputs):
|
||||
dst.copy_(src)
|
||||
graph.replay()
|
||||
if copy_outputs:
|
||||
return [x.clone() for x in static_outputs]
|
||||
else:
|
||||
return static_outputs
|
||||
|
||||
return run
|
||||
@@ -0,0 +1,730 @@
|
||||
"""
|
||||
This module provides debugging backends for TorchDynamo to help diagnose and troubleshoot
|
||||
compilation and execution issues. It includes:
|
||||
|
||||
Key Debugging Backends:
|
||||
- eager: Simple pass-through backend that runs models in eager mode
|
||||
- eager_noexcept: Similar to eager but with additional exception handling
|
||||
- eager_debug: Adds schema validation checks for custom operators
|
||||
- aot_eager: Uses AOT Autograd with nop compiler for debugging
|
||||
- aot_eager_decomp_partition: Uses TorchInductor decompositions for debugging
|
||||
- torchscript: Compiles using TorchScript for debugging JIT-related issues
|
||||
|
||||
Testing and Development Tools:
|
||||
- Backends for inducing specific errors (compile/runtime/accuracy)
|
||||
- ExplainOutput class for detailed graph compilation analysis
|
||||
- Utilities for cross-referencing and mode management
|
||||
- Tools for graph detail inspection and break reason analysis
|
||||
|
||||
These backends are primarily used for:
|
||||
1. Debugging graph breaks and compilation failures
|
||||
2. Testing error handling and recovery mechanisms
|
||||
3. Analyzing performance bottlenecks
|
||||
4. Validating operator schemas and decompositions
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import functools
|
||||
import logging
|
||||
from collections.abc import Callable, Iterable
|
||||
from importlib import import_module
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from functorch.compile import min_cut_rematerialization_partition
|
||||
from torch import _guards
|
||||
from torch._dynamo.output_graph import GraphCompileReason
|
||||
from torch._functorch import config as functorch_config
|
||||
from torch._functorch.compilers import ts_compile
|
||||
from torch._inductor.output_code import OutputCode
|
||||
|
||||
from .common import aot_autograd
|
||||
from .registry import CompiledFn, CompilerFn, register_debug_backend as register_backend
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.fx.node import Target
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_backend
|
||||
def eager(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
if kwargs:
|
||||
log.warning("eager backend ignoring extra kwargs %s", kwargs)
|
||||
|
||||
if torch._functorch.config.force_autograd_cache:
|
||||
from torch._dynamo.aot_compile_types import GraphModuleSerializableCallable
|
||||
|
||||
return GraphModuleSerializableCallable(gm)
|
||||
return gm.forward
|
||||
|
||||
|
||||
def make_eager_backend_with_torch_function_mode(
|
||||
mode: torch.overrides.TorchFunctionMode,
|
||||
) -> Callable[..., Any]:
|
||||
return make_eager_backend_with_torch_function_modes([mode])
|
||||
|
||||
|
||||
def make_eager_backend_with_torch_function_modes(
|
||||
modes: Iterable[torch.overrides.TorchFunctionMode],
|
||||
) -> Callable[..., Any]:
|
||||
"""Used to trace HOPs (cond and while) for eager execution, the metadata
|
||||
TF mode mutates vars outside of the scope of the HOP, and we can't have graph breaks
|
||||
in the HOP, so we need to externally run this mode and not trace it."""
|
||||
from contextlib import ExitStack
|
||||
|
||||
def fn(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
with ExitStack() as stack:
|
||||
for mode in modes:
|
||||
stack.enter_context(mode)
|
||||
return gm.forward(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return fn
|
||||
|
||||
|
||||
@register_backend
|
||||
def eager_noexcept(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
if kwargs:
|
||||
log.warning("eager_noexcept backend ignoring extra kwargs %s", kwargs)
|
||||
|
||||
# This backend is intended to check that dynamo-generated GraphModules
|
||||
# do not cause errors.
|
||||
def inner(*args: Any) -> Any:
|
||||
try:
|
||||
return gm(*args)
|
||||
except Exception as e:
|
||||
raise torch._dynamo.exc.TorchDynamoException(
|
||||
"Unexpected exception when running generated GraphModule"
|
||||
) from e
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
@register_backend
|
||||
def pre_dispatch_eager(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> torch.fx.GraphModule:
|
||||
if kwargs:
|
||||
log.warning("pre_dispatch_eager backend ignoring extra kwargs %s", kwargs)
|
||||
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
|
||||
def runnable_gm(*args: Any) -> Any:
|
||||
return torch.fx.Interpreter(gm).run(*args)
|
||||
|
||||
pre_dispatch_gm = make_fx(runnable_gm, pre_dispatch=True)(*fake_tensor_inputs)
|
||||
pre_dispatch_gm.print_readable()
|
||||
|
||||
return pre_dispatch_gm
|
||||
|
||||
|
||||
@register_backend
|
||||
def eager_debug(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
if kwargs:
|
||||
log.warning("eager_debug backend ignoring extra kwargs %s", kwargs)
|
||||
|
||||
from torch._subclasses.schema_check_mode import SchemaCheckMode
|
||||
|
||||
# We could add more debugging bits here.
|
||||
# Right now, this backend can be used to check for and error on
|
||||
# custom dispatcher ops that have incorrect schemas.
|
||||
def inner(*args: Any) -> Any:
|
||||
with SchemaCheckMode():
|
||||
return torch.fx.Interpreter(gm).run(*args)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
@register_backend(name="ts") # type: ignore[misc]
|
||||
def torchscript(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor]
|
||||
) -> torch.jit.ScriptModule:
|
||||
return torch.jit.script(gm)
|
||||
|
||||
|
||||
def invoke_subgraph_inner_compiler(
|
||||
subgraph: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
"""Inner compiler that wraps forward/backward graphs in invoke_subgraph HOP.
|
||||
|
||||
This is used as the fw_compiler/bw_compiler for aot_autograd. When the resulting
|
||||
function is traced by make_fx, it emits an invoke_subgraph HOP instead of inlining.
|
||||
"""
|
||||
from torch._dynamo import disable
|
||||
from torch._higher_order_ops.invoke_subgraph import invoke_subgraph_infer
|
||||
|
||||
@disable
|
||||
# pyrefly: ignore [deprecated]
|
||||
@torch._dynamo.allow_in_graph
|
||||
def invoke_subgraph_wrapper_unboxed(*operands: Any) -> Any:
|
||||
return invoke_subgraph_infer(subgraph, *operands)
|
||||
|
||||
# NB: The direct to unboxed path is broken, you MUST DO THIS
|
||||
|
||||
def invoke_subgraph_wrapper(args: list[Any]) -> Any:
|
||||
return invoke_subgraph_wrapper_unboxed(*args)
|
||||
|
||||
invoke_subgraph_wrapper._boxed_call = True # type: ignore[attr-defined]
|
||||
|
||||
return invoke_subgraph_wrapper
|
||||
|
||||
|
||||
# I cannot say how many times I had to revert to this vibe coded version of
|
||||
# the code, which worked, and the cleaner versions of the code did not work,
|
||||
# so I'm leaving this here until we fix the rest of the bugs.
|
||||
'''
|
||||
# Counter for unique subgraph names in invoke_subgraph backend
|
||||
_invoke_subgraph_counter = 0
|
||||
|
||||
|
||||
def invoke_subgraph_inner_compiler_good(
|
||||
fx_g: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
"""Inner compiler that wraps forward/backward graphs in invoke_subgraph HOP.
|
||||
|
||||
This is used as the fw_compiler/bw_compiler for aot_autograd. When the resulting
|
||||
function is traced by make_fx, it emits an invoke_subgraph HOP instead of inlining.
|
||||
"""
|
||||
from torch._higher_order_ops.invoke_subgraph import (
|
||||
invoke_subgraph as invoke_subgraph_hop,
|
||||
)
|
||||
from torch.fx.experimental.proxy_tensor import get_proxy_mode
|
||||
|
||||
global _invoke_subgraph_counter
|
||||
_invoke_subgraph_counter += 1
|
||||
name = f"invoke_subgraph_{_invoke_subgraph_counter}"
|
||||
|
||||
from torch._dynamo import disable
|
||||
|
||||
# Check if fx_g uses boxed calling convention
|
||||
fx_g_is_boxed = getattr(fx_g, "_boxed_call", False)
|
||||
|
||||
@disable
|
||||
@torch._dynamo.allow_in_graph
|
||||
def invoke_subgraph_wrapper_unboxed(*args: Any) -> Any:
|
||||
proxy_mode = get_proxy_mode()
|
||||
if proxy_mode is not None:
|
||||
# When being traced by make_fx, emit invoke_subgraph HOP
|
||||
return invoke_subgraph_hop(fx_g, name, *args) # type: ignore[arg-type]
|
||||
else:
|
||||
# Normal execution path - call fx_g with proper calling convention
|
||||
if fx_g_is_boxed:
|
||||
return fx_g(list(args))
|
||||
else:
|
||||
return fx_g(*args)
|
||||
|
||||
# Wrap to handle boxed arguments (list of args) as expected by AOTAutograd
|
||||
def invoke_subgraph_wrapper(args: list[Any]) -> Any:
|
||||
return invoke_subgraph_wrapper_unboxed(*args)
|
||||
|
||||
invoke_subgraph_wrapper._boxed_call = True # type: ignore[attr-defined]
|
||||
return invoke_subgraph_wrapper
|
||||
'''
|
||||
|
||||
|
||||
@register_backend
|
||||
def invoke_subgraph(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
"""Backend that wraps forward/backward graphs in invoke_subgraph HOP when traced by make_fx.
|
||||
|
||||
This backend uses AOTAutograd to partition into forward/backward graphs, then wraps
|
||||
each in an invoke_subgraph HOP. This is useful for recursive Dynamo tracing scenarios
|
||||
where you want the compiled subgraph to appear as invoke_subgraph HOPs in the outer
|
||||
trace rather than being inlined.
|
||||
|
||||
Requires:
|
||||
- torch._dynamo.config.force_compile_during_fx_trace = True
|
||||
(this implicitly overrides error_on_nested_fx_trace)
|
||||
"""
|
||||
if kwargs:
|
||||
log.warning("invoke_subgraph backend ignoring extra kwargs %s", kwargs)
|
||||
|
||||
# Use AOTAutograd to partition into forward/backward
|
||||
return aot_autograd(
|
||||
fw_compiler=invoke_subgraph_inner_compiler,
|
||||
bw_compiler=invoke_subgraph_inner_compiler,
|
||||
partition_fn=min_cut_rematerialization_partition,
|
||||
keep_inference_input_mutations=True,
|
||||
)(gm, fake_tensor_inputs)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AOTEagerOutputCode(OutputCode):
|
||||
"""
|
||||
An OutputCode that wraps a GraphModule for eager-mode execution.
|
||||
|
||||
This allows non-inductor backends (like aot_eager) to participate in
|
||||
the bundled autograd cache and aot_compile serialization flow.
|
||||
"""
|
||||
|
||||
gm: torch.fx.GraphModule | None = None
|
||||
_serialized_gm: bytes | None = dataclasses.field(default=None, init=False)
|
||||
|
||||
def __call__(self, inputs: Any) -> Any:
|
||||
assert self.gm is not None
|
||||
return self.gm.forward(inputs)
|
||||
|
||||
def prepare_for_serialization(self) -> None:
|
||||
from torch.fx._graph_pickler import GraphPickler, Options
|
||||
|
||||
assert self.gm is not None
|
||||
for node in self.gm.graph.nodes:
|
||||
node.meta.pop("nn_module_stack", None)
|
||||
node.meta.pop("source_fn_stack", None)
|
||||
node.meta.pop("example_value", None)
|
||||
|
||||
self._serialized_gm = GraphPickler.dumps(self.gm, Options(ops_filter=None))
|
||||
self.gm = None
|
||||
|
||||
def post_compile(self, *args: Any, **kwargs: Any) -> None:
|
||||
if self.gm is None and self._serialized_gm is not None:
|
||||
from torch._subclasses import FakeTensorMode
|
||||
from torch.fx._graph_pickler import GraphPickler
|
||||
from torch.fx.experimental.symbolic_shapes import ShapeEnv
|
||||
from torch.fx.graph import _BoxedCodeGen
|
||||
|
||||
fake_mode = FakeTensorMode(shape_env=ShapeEnv())
|
||||
gm = GraphPickler.loads(self._serialized_gm, fake_mode)
|
||||
assert isinstance(gm, torch.fx.GraphModule)
|
||||
self.gm = gm
|
||||
assert isinstance(self.gm, torch.fx.GraphModule)
|
||||
self.gm.graph.set_codegen(_BoxedCodeGen())
|
||||
self.gm.recompile()
|
||||
self._serialized_gm = None
|
||||
|
||||
def set_triton_bundle(self, triton_bundle: Any) -> None:
|
||||
pass
|
||||
|
||||
|
||||
# used boxed call to discard inputs when they are no longer needed
|
||||
def boxed_nop(
|
||||
fx_g: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
from torch.fx.graph import _BoxedCodeGen
|
||||
|
||||
# Set the graph to use boxed codegen
|
||||
fx_g.graph.set_codegen(_BoxedCodeGen())
|
||||
fx_g.recompile()
|
||||
|
||||
if functorch_config.force_autograd_cache or functorch_config.bundled_autograd_cache:
|
||||
result = AOTEagerOutputCode(gm=fx_g)
|
||||
result._boxed_call = True # type: ignore[attr-defined]
|
||||
return result
|
||||
|
||||
# Wrap the forward method in a function so we can set _boxed_call attribute
|
||||
forward_fn = fx_g.forward
|
||||
|
||||
def run(args: Any) -> Any:
|
||||
from torch.utils._debug_mode import DebugInterpreter, get_active_debug_mode
|
||||
|
||||
if (
|
||||
debug_mode := get_active_debug_mode()
|
||||
) is not None and debug_mode.run_compile_with_interpreter:
|
||||
return DebugInterpreter(fx_g, backend="aot_eager").run(*args)
|
||||
return forward_fn(args)
|
||||
|
||||
run._boxed_call = True # type: ignore[attr-defined]
|
||||
return run
|
||||
|
||||
|
||||
def boxed_nop_with_mode(
|
||||
fx_g: torch.fx.GraphModule,
|
||||
example_inputs: list[torch.Tensor],
|
||||
*,
|
||||
mode: torch.overrides.TorchFunctionMode,
|
||||
) -> Callable[..., Any]:
|
||||
from torch.fx.graph import _BoxedCodeGen
|
||||
|
||||
# Set the graph to use boxed codegen
|
||||
fx_g.graph.set_codegen(_BoxedCodeGen())
|
||||
fx_g.recompile()
|
||||
|
||||
# Create a wrapper that runs with the mode
|
||||
forward_fn = fx_g.forward
|
||||
|
||||
def run(args: Any) -> Any:
|
||||
with mode:
|
||||
return forward_fn(args)
|
||||
|
||||
run._boxed_call = True # type: ignore[attr-defined]
|
||||
return run
|
||||
|
||||
|
||||
def fake_crossref_boxed_nop(
|
||||
fx_g: torch.fx.GraphModule,
|
||||
example_inputs: list[torch.Tensor],
|
||||
ignore_op_fn: Callable[[torch._ops.OpOverload], bool] | None = None,
|
||||
) -> Callable[..., Any]:
|
||||
from torch.fx.graph import _BoxedCodeGen
|
||||
|
||||
# Set the graph to use boxed codegen
|
||||
fx_g.graph.set_codegen(_BoxedCodeGen())
|
||||
fx_g.recompile()
|
||||
|
||||
# Create a wrapper that runs with the mode
|
||||
forward_fn = fx_g.forward
|
||||
|
||||
def run(args: Any) -> Any:
|
||||
with torch._subclasses.CrossRefFakeMode(ignore_op_fn):
|
||||
return forward_fn(args)
|
||||
|
||||
run._boxed_call = True # type: ignore[attr-defined]
|
||||
return run
|
||||
|
||||
|
||||
def ignore_builtins(op: torch._ops.OpOverload) -> bool:
|
||||
return op.namespace in ("aten", "prims", "prim")
|
||||
|
||||
|
||||
def get_nop_func() -> Callable[
|
||||
[torch.fx.GraphModule, list[torch.Tensor]], Callable[..., Any]
|
||||
]:
|
||||
if not torch._functorch.config.fake_tensor_crossref:
|
||||
return boxed_nop
|
||||
elif torch._functorch.config.fake_tensor_crossref == "all":
|
||||
return fake_crossref_boxed_nop
|
||||
else:
|
||||
assert torch._functorch.config.fake_tensor_crossref == "custom_ops"
|
||||
return functools.partial(fake_crossref_boxed_nop, ignore_op_fn=ignore_builtins)
|
||||
|
||||
|
||||
# Useful for debugging purpose
|
||||
# aot_eager uses AOT Autograd backend with nop compiler. It is helpful in debugging.
|
||||
def aot_eager(
|
||||
gm: torch.fx.GraphModule,
|
||||
fake_tensor_inputs: list[torch.Tensor],
|
||||
fw_compiler: Callable[..., Any] | None = None,
|
||||
bw_compiler: Callable[..., Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Callable[..., Any]:
|
||||
return aot_autograd(
|
||||
fw_compiler=fw_compiler or boxed_nop,
|
||||
bw_compiler=bw_compiler or boxed_nop,
|
||||
partition_fn=min_cut_rematerialization_partition,
|
||||
keep_inference_input_mutations=True,
|
||||
)(gm, fake_tensor_inputs, **kwargs)
|
||||
|
||||
|
||||
register_backend(name="aot_eager", compiler_fn=aot_eager)
|
||||
|
||||
aot_eager_default_partitioner = aot_autograd(
|
||||
fw_compiler=boxed_nop, keep_inference_input_mutations=True
|
||||
)
|
||||
register_backend(
|
||||
name="aot_eager_default_partitioner", compiler_fn=aot_eager_default_partitioner
|
||||
)
|
||||
|
||||
|
||||
# Uses TorchInductor AOT Autograd decomps and partitioner to isolate aot vs
|
||||
# inductor problems.
|
||||
# aot_eager_decomp_partition just replaces the inductor compiler with nop to help
|
||||
# isolate inductor vs aot_eager errors
|
||||
def aot_eager_decomp_partition(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
if kwargs:
|
||||
log.warning(
|
||||
"aot_eager_decomp_partition backend ignoring extra kwargs %s", kwargs
|
||||
)
|
||||
|
||||
from torch._inductor.compiler_bisector import CompilerBisector
|
||||
|
||||
config_patches = {"unlift_effect_tokens": True}
|
||||
if bisect_changes := CompilerBisector.get_config_change(
|
||||
"aot_eager_decomp_partition"
|
||||
):
|
||||
config_patches.update(bisect_changes) # type: ignore[arg-type]
|
||||
|
||||
with functorch_config.patch(config_patches):
|
||||
return aot_autograd(
|
||||
# these are taken from memory_efficient_fusion()
|
||||
fw_compiler=get_nop_func(),
|
||||
bw_compiler=get_nop_func(),
|
||||
# NB: lambda here is to delay import of inductor
|
||||
decompositions=lambda: import_module(
|
||||
"torch._inductor.compile_fx"
|
||||
).select_decomp_table(),
|
||||
partition_fn=functools.partial(
|
||||
min_cut_rematerialization_partition, compiler="inductor"
|
||||
),
|
||||
)(gm, fake_tensor_inputs)
|
||||
|
||||
|
||||
register_backend(
|
||||
name="aot_eager_decomp_partition", compiler_fn=aot_eager_decomp_partition
|
||||
)
|
||||
|
||||
|
||||
# aot_eager_decomp_partition_with_mode is similar as aot_eager_decomp_partition,
|
||||
# except that it takes a TorchDispatchMode mode and run the fw/bw in the mode
|
||||
def aot_eager_decomp_partition_with_mode(
|
||||
gm: torch.fx.GraphModule,
|
||||
fake_tensor_inputs: list[torch.Tensor],
|
||||
mode: Any,
|
||||
**kwarg: Any,
|
||||
) -> Callable[..., Any]:
|
||||
return aot_autograd(
|
||||
# these are taken from memory_efficient_fusion()
|
||||
fw_compiler=functools.partial(boxed_nop_with_mode, mode=mode),
|
||||
bw_compiler=functools.partial(boxed_nop_with_mode, mode=mode),
|
||||
# NB: lambda here is to delay import of inductor
|
||||
decompositions=lambda: import_module(
|
||||
"torch._inductor.compile_fx"
|
||||
).select_decomp_table(),
|
||||
partition_fn=functools.partial(
|
||||
min_cut_rematerialization_partition, compiler="inductor"
|
||||
),
|
||||
)(gm, fake_tensor_inputs)
|
||||
|
||||
|
||||
register_backend(
|
||||
name="aot_eager_decomp_partition_with_mode",
|
||||
compiler_fn=aot_eager_decomp_partition_with_mode, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def aot_eager_decomp_partition_crossref(
|
||||
gm: torch.fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], **kwargs: Any
|
||||
) -> Callable[..., Any]:
|
||||
# if the config is set, respect it, otherwise only test custom_ops.
|
||||
# custom_op bad metas always manifest as an error whereas aten will only sometimes.
|
||||
# by default, use the less noisy option
|
||||
config_val = (
|
||||
"custom_ops"
|
||||
if not functorch_config.fake_tensor_crossref
|
||||
else functorch_config.fake_tensor_crossref
|
||||
)
|
||||
with functorch_config.patch(fake_tensor_crossref=config_val):
|
||||
return aot_eager_decomp_partition(gm, fake_tensor_inputs, **kwargs)
|
||||
|
||||
|
||||
register_backend(
|
||||
name="aot_eager_decomp_partition_crossref",
|
||||
compiler_fn=aot_eager_decomp_partition_crossref,
|
||||
)
|
||||
|
||||
|
||||
# AOT Autograd with torchscript backend. Default partitioner.
|
||||
# aot_ts uses torchscript backend. We can use this with both nnc and nvfuser
|
||||
# by using the relevant fuser with torch.jit.fuser(...)
|
||||
aot_ts = aot_autograd(fw_compiler=ts_compile)
|
||||
register_backend(name="aot_ts", compiler_fn=aot_ts)
|
||||
|
||||
# These buggy backends are used for inducing bugs so that we can test
|
||||
# our repro extraction / minifier scripts
|
||||
|
||||
|
||||
class ReluCompileError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class TestingOnlyCompileError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@register_backend
|
||||
def relu_compile_error_TESTING_ONLY(
|
||||
gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> torch.fx.GraphModule:
|
||||
for node in gm.graph.nodes:
|
||||
if node.target is torch.relu:
|
||||
raise ReluCompileError
|
||||
return gm
|
||||
|
||||
|
||||
@register_backend
|
||||
def relu_runtime_error_TESTING_ONLY(
|
||||
gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> torch.fx.GraphModule:
|
||||
for node in gm.graph.nodes:
|
||||
if node.target is torch.relu:
|
||||
node.target = torch._assert
|
||||
node.args = (False, "ReluRuntimeError")
|
||||
gm.recompile()
|
||||
return gm
|
||||
|
||||
|
||||
@register_backend
|
||||
def relu_accuracy_error_TESTING_ONLY(
|
||||
gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> torch.fx.GraphModule:
|
||||
for node in gm.graph.nodes:
|
||||
if node.target is torch.relu:
|
||||
node.target = torch.add
|
||||
node.args = (node.args[0], 1)
|
||||
gm.recompile()
|
||||
|
||||
return gm
|
||||
|
||||
|
||||
@register_backend
|
||||
def non_leaf_compile_error_TESTING_ONLY(
|
||||
gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> torch.fx.GraphModule:
|
||||
# Require at least one non-trivial thing in the graph,
|
||||
# see https://github.com/pytorch/pytorch/issues/102898
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "call_function":
|
||||
break
|
||||
else:
|
||||
return gm
|
||||
for t in example_inputs:
|
||||
if not t.is_leaf:
|
||||
raise TestingOnlyCompileError
|
||||
return gm
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ExplainOutput:
|
||||
"""
|
||||
This is the output of :func:`torch._dynamo.explain()`
|
||||
There is no reason to create this class directly.
|
||||
"""
|
||||
|
||||
graphs: list[torch.fx.GraphModule]
|
||||
graph_count: int
|
||||
graph_break_count: int
|
||||
break_reasons: list[GraphCompileReason]
|
||||
op_count: int
|
||||
ops_per_graph: list[list["Target"]] | None = None
|
||||
out_guards: list[_guards.Guard] | None = None
|
||||
compile_times: str | None = None
|
||||
|
||||
def __str__(self) -> str:
|
||||
output = f"Graph Count: {self.graph_count}\n"
|
||||
output += f"Graph Break Count: {self.graph_break_count}\n"
|
||||
output += f"Op Count: {self.op_count}\n"
|
||||
|
||||
output += "Break Reasons:\n"
|
||||
for idx, break_reason in enumerate(self.break_reasons):
|
||||
output += f" Break Reason {idx + 1}:\n"
|
||||
output += f" Reason: {break_reason.reason}\n"
|
||||
output += " User Stack:\n"
|
||||
for frame_summary in break_reason.user_stack:
|
||||
output += f" {frame_summary}\n"
|
||||
|
||||
if self.ops_per_graph is not None:
|
||||
output += "Ops per Graph:\n"
|
||||
for idx, ops in enumerate(self.ops_per_graph):
|
||||
output += f" Ops {idx + 1}:\n"
|
||||
for op in ops:
|
||||
output += f" {op}\n"
|
||||
|
||||
if self.out_guards is not None:
|
||||
output += "Out Guards:\n"
|
||||
for i, guard in enumerate(self.out_guards):
|
||||
output += f" Guard {i + 1}:\n"
|
||||
output += f" {str(guard)}"
|
||||
|
||||
if self.compile_times is not None:
|
||||
output += f"Compile Times: {self.compile_times}\n"
|
||||
return output
|
||||
|
||||
|
||||
def _explain_graph_detail(
|
||||
gm: torch.fx.GraphModule,
|
||||
graphs: list[torch.fx.GraphModule],
|
||||
op_count: int,
|
||||
ops_per_graph: list[list["Target"]],
|
||||
break_reasons: list[GraphCompileReason],
|
||||
) -> tuple[
|
||||
torch.fx.GraphModule,
|
||||
list[torch.fx.GraphModule],
|
||||
int,
|
||||
list[list["Target"]],
|
||||
list[GraphCompileReason],
|
||||
]:
|
||||
"""
|
||||
This function is a utility which processes a torch.fx.GraphModule and
|
||||
accumulates information about its ops, graph breaks, and other details. It
|
||||
is intended to be used by the ExplainWithBackend class and
|
||||
`torch._dynamo.explain()` to provide details from Dynamo's graph capture.
|
||||
|
||||
Parameters:
|
||||
gm (torch.fx.GraphModule): The GraphModule to be processed.
|
||||
graphs (list): A list that accumulates all the GraphModules processed.
|
||||
op_count (int): The total count of operations in all GraphModules processed so far.
|
||||
ops_per_graph (list): A list that accumulates the operations of each GraphModule.
|
||||
break_reasons (list): A list that accumulates the reasons for breaks in each GraphModule.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the processed GraphModule, the updated lists of graphs,
|
||||
operations per graph, and break reasons, and the updated operation count.
|
||||
"""
|
||||
graphs.append(gm)
|
||||
ops = [node.target for node in gm.graph.nodes if node.op == "call_function"]
|
||||
op_count += len(ops)
|
||||
ops_per_graph.append(ops)
|
||||
if gm.compile_subgraph_reason.graph_break: # type: ignore[union-attr]
|
||||
break_reasons.append(gm.compile_subgraph_reason) # type: ignore[arg-type]
|
||||
|
||||
return gm, graphs, op_count, ops_per_graph, break_reasons
|
||||
|
||||
|
||||
class ExplainWithBackend:
|
||||
"""
|
||||
This class is intended to be used as a backend for `torch.compile`. It is
|
||||
composable with other backends. When used in this way, it accumulates
|
||||
information about graph breaks, ops, and other info and provides a string
|
||||
representation summarizing this information.
|
||||
|
||||
Attributes:
|
||||
backend (str): The name of the backend to use for optimization.
|
||||
graphs (list): A list of the graphs captured by TorchDynamo.
|
||||
op_count (int): The total number of operations in all optimized graphs.
|
||||
break_reasons (list): A list of graph break reasons with stack traces.
|
||||
|
||||
Example Usage:
|
||||
def fn(x):
|
||||
x = torch.sigmoid(x)
|
||||
return x
|
||||
|
||||
torch._dynamo.reset()
|
||||
eb = ExplainWithBackend("inductor")
|
||||
optimized_fn = torch.compile(fn, backend=eb)
|
||||
result = optimized_fn(torch.randn(5))
|
||||
print(eb.output())
|
||||
"""
|
||||
|
||||
def __init__(self, backend: CompilerFn | str) -> None:
|
||||
from .registry import lookup_backend
|
||||
|
||||
self.backend = lookup_backend(backend)
|
||||
self.graphs: list[torch.fx.GraphModule] = []
|
||||
self.op_count = 0
|
||||
self.break_reasons: list[GraphCompileReason] = []
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> CompiledFn:
|
||||
ops_per_graph: list[list[Target]] = []
|
||||
gm, self.graphs, self.op_count, _, self.break_reasons = _explain_graph_detail(
|
||||
gm, self.graphs, self.op_count, ops_per_graph, self.break_reasons
|
||||
)
|
||||
return self.backend(gm, example_inputs)
|
||||
|
||||
def output(self) -> ExplainOutput:
|
||||
graph_count = len(self.graphs)
|
||||
output = ExplainOutput(
|
||||
self.graphs,
|
||||
graph_count,
|
||||
graph_count - 1,
|
||||
self.break_reasons,
|
||||
self.op_count,
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,622 @@
|
||||
"""
|
||||
This module implements distributed training optimizations for TorchDynamo backends.
|
||||
|
||||
It provides functionality to optimize models wrapped in DistributedDataParallel (DDP)
|
||||
by intelligently splitting compiled graphs to align with DDP's gradient synchronization
|
||||
boundaries. Key features include:
|
||||
|
||||
- Graph partitioning based on parameter bucket sizes
|
||||
- Optimization of allreduce operations for distributed training
|
||||
- Support for parameter ignoring and buffer handling
|
||||
- Submodule compilation and management
|
||||
- Debugging utilities for distributed training
|
||||
|
||||
The main component is the DDPOptimizer class, which handles graph splitting and
|
||||
recompilation to enable efficient distributed training while maintaining the benefits
|
||||
of compilation.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
from torch import fx
|
||||
from torch._dynamo.backends.registry import CompiledFn, CompilerFn
|
||||
from torch._dynamo.output_graph import GraphCompileReason
|
||||
from torch._dynamo.utils import deepcopy_to_fake_tensor, detect_fake_mode
|
||||
from torch._logging import trace_structured
|
||||
from torch.fx.node import Node
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._functorch._aot_autograd.schemas import ViewAndMutationMeta
|
||||
|
||||
|
||||
# Regular log messages should go through 'log'.
|
||||
# ddp_graph_log is a separate artifact logger reserved for dumping graphs.
|
||||
# See docs/source/logging.rst for more info.
|
||||
log = logging.getLogger(__name__)
|
||||
ddp_graph_log = torch._logging.getArtifactLogger(__name__, "ddp_graphs")
|
||||
|
||||
|
||||
def args_str(args: Any) -> str:
|
||||
# a debug helper
|
||||
if torch.is_tensor(args):
|
||||
return f"T[{args.shape}]"
|
||||
elif isinstance(args, tuple):
|
||||
return f"tuple({', '.join([args_str(x) for x in args])})"
|
||||
elif isinstance(args, list):
|
||||
return f"list({', '.join([args_str(x) for x in args])})"
|
||||
else:
|
||||
return str(args)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Bucket:
|
||||
size: int = 0
|
||||
params: list[str] = field(default_factory=list)
|
||||
nodes: list[fx.Node] = field(default_factory=list)
|
||||
|
||||
# param_ids is just used for unit testing
|
||||
param_ids: list[int] = field(default_factory=list)
|
||||
|
||||
# keep track of any buckets that were extended for logging purposes
|
||||
opcount_increased_to_capture_external_output: int = 0
|
||||
paramsize_before_opcount_increase: int = 0
|
||||
|
||||
|
||||
def bucket_has_external_output(bucket: Bucket) -> bool:
|
||||
nodes_in_bucket = set()
|
||||
# we want to iterate in reverse order, but clumsi-luckily the bucket.nodes list was already created backwards
|
||||
# so we don't reverse it here
|
||||
for node in bucket.nodes:
|
||||
# assume node.op != output, since those are filtered in the original iteration
|
||||
nodes_in_bucket.add(node)
|
||||
for user in node.users:
|
||||
if user not in nodes_in_bucket:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def pretty_print_buckets(buckets: list[Bucket], bucket_bytes_cap: int) -> None:
|
||||
headers = ("Index", "Size (b)", "Param Names")
|
||||
rows: list[tuple[int | None, int | None, str]] = []
|
||||
# pyrefly: ignore [implicit-any]
|
||||
extended_buckets = []
|
||||
for idx, bucket in enumerate(reversed(buckets)):
|
||||
if len(bucket.params) > 0:
|
||||
rows.append((idx, bucket.size, bucket.params[0]))
|
||||
rows.extend((None, None, param) for param in bucket.params[1:])
|
||||
if bucket.opcount_increased_to_capture_external_output > 0:
|
||||
extended_buckets.append(
|
||||
(
|
||||
idx,
|
||||
bucket.opcount_increased_to_capture_external_output,
|
||||
bucket.size - bucket.paramsize_before_opcount_increase,
|
||||
)
|
||||
)
|
||||
|
||||
if rows:
|
||||
log.info(
|
||||
"\nDDPOptimizer used bucket cap %s and created %d buckets. Enable debug logs for detailed bucket info.",
|
||||
bucket_bytes_cap,
|
||||
len(buckets),
|
||||
)
|
||||
|
||||
if extended_buckets:
|
||||
log.warning(
|
||||
"Some buckets were extended beyond their requested parameter capacities"
|
||||
" in order to ensure each subgraph has an output node, required for fx graph partitioning."
|
||||
" This can be the case when a subgraph would have only contained nodes performing inplace mutation,"
|
||||
" and returning no logical outputs. This should not be a problem, unless it results in too few graph"
|
||||
" partitions for optimal DDP performance."
|
||||
)
|
||||
|
||||
try:
|
||||
from tabulate import tabulate
|
||||
|
||||
log.debug(
|
||||
"\nDDPOptimizer produced the following bucket assignments:\n%s",
|
||||
tabulate(rows, headers=headers, tablefmt="simple_grid"),
|
||||
)
|
||||
|
||||
if extended_buckets:
|
||||
log.warning(
|
||||
"DDPOptimizer extended these buckets to ensure per-subgraph output nodes:\n%s",
|
||||
tabulate(
|
||||
extended_buckets,
|
||||
headers=("Index", "Extra Ops", "Extra Param Size (b)"),
|
||||
tablefmt="simple_grid",
|
||||
),
|
||||
)
|
||||
except ImportError:
|
||||
log.debug(
|
||||
"Please `pip install tabulate` in order to display ddp bucket sizes and diagnostic information."
|
||||
)
|
||||
else:
|
||||
log.debug("DDPOptimizer captured no parameters and did not split this graph.")
|
||||
|
||||
|
||||
def has_higher_order_op(gm: fx.GraphModule) -> bool:
|
||||
# Check if there is a higher order op in the graph
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "get_attr":
|
||||
maybe_param = getattr(gm, node.target)
|
||||
if isinstance(maybe_param, torch.fx.GraphModule):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def propagate_metadata(orig_gm: fx.GraphModule, split_gm: fx.GraphModule) -> None:
|
||||
for name, module in split_gm.named_modules():
|
||||
if "." not in name and len(name):
|
||||
# TODO: add split id to CompileId: https://github.com/pytorch/tlparse/pull/83/files#r1880649384
|
||||
module.meta = orig_gm.meta
|
||||
module._param_name_to_source = orig_gm._param_name_to_source
|
||||
|
||||
|
||||
def propagate_dynamo_source(orig_gm: fx.GraphModule, split_gm: fx.GraphModule) -> None:
|
||||
name_to_dynamo_source = {}
|
||||
for node in orig_gm.graph.find_nodes(op="placeholder"):
|
||||
name_to_dynamo_source[node.name] = node._dynamo_source
|
||||
|
||||
for name, module in split_gm.named_modules():
|
||||
if "." not in name and len(name):
|
||||
for node in module.graph.find_nodes(op="placeholder"):
|
||||
# non-placeholder in original_gm may become placeholder in submodules
|
||||
node._dynamo_source = name_to_dynamo_source.get(node.name)
|
||||
|
||||
|
||||
class DDPOptimizerContext:
|
||||
def __init__(self) -> None:
|
||||
self.curr_bucket: int = -1
|
||||
self.metadata_per_bucket: list[ViewAndMutationMeta] = []
|
||||
|
||||
|
||||
# compile each of the partitioned submodules using the user-provided compiler
|
||||
class SubmodCompiler(torch.fx.interpreter.Interpreter):
|
||||
def __init__(
|
||||
self,
|
||||
module: fx.GraphModule,
|
||||
compiler: CompilerFn,
|
||||
fake_mode: torch._subclasses.fake_tensor.FakeTensorMode,
|
||||
) -> None:
|
||||
super().__init__(module)
|
||||
self.compiler = compiler
|
||||
self.fake_mode = fake_mode
|
||||
# See Note [DDPOptimizer and fw_metadata]
|
||||
ctx = torch._guards.TracingContext.try_get()
|
||||
if ctx is not None:
|
||||
ctx.ddp_optimizer_ctx = DDPOptimizerContext()
|
||||
|
||||
def compile_submod(
|
||||
self, input_mod: fx.GraphModule, args: list[torch.Tensor], kwargs: Any
|
||||
) -> Any:
|
||||
"""
|
||||
Compile the submodule,
|
||||
using a wrapper to make sure its output is always a tuple,
|
||||
which is required by AotAutograd based compilers
|
||||
"""
|
||||
assert len(kwargs) == 0, "We assume only args for these modules"
|
||||
|
||||
class WrapperModule(torch.nn.Module):
|
||||
def __init__(
|
||||
self, submod: Callable[..., Any], unwrap_singleton_tuple: bool
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.submod = submod
|
||||
self.unwrap_singleton_tuple = unwrap_singleton_tuple
|
||||
|
||||
def forward(self, *args: Any) -> Any:
|
||||
x = self.submod(*args)
|
||||
# TODO(whc)
|
||||
# for some reason the isinstance check is necessary if I split one node per submod
|
||||
# - even though I supposedly wrapped the output in a tuple in those cases, the real
|
||||
# compiled module was still returning a tensor
|
||||
if self.unwrap_singleton_tuple and isinstance(x, (tuple, list)):
|
||||
return x[0]
|
||||
return x
|
||||
|
||||
unwrap_singleton_tuple = False
|
||||
for sn in input_mod.graph.nodes:
|
||||
if sn.op == "output":
|
||||
if not isinstance(sn.args[0], tuple):
|
||||
unwrap_singleton_tuple = True
|
||||
sn.args = (sn.args,)
|
||||
|
||||
input_mod.recompile()
|
||||
input_mod.compile_subgraph_reason = GraphCompileReason( # type: ignore[assignment]
|
||||
"DDPOptimizer intentional graph-break (See Note [DDPOptimizer])."
|
||||
" Set `torch._dynamo.config.optimize_ddp = False` to disable.",
|
||||
[
|
||||
# it's close to useless to get a real stacktrace here, and quite verbose.
|
||||
traceback.FrameSummary(__file__, 0, "DDPOptimizer"),
|
||||
],
|
||||
)
|
||||
|
||||
wrapper = WrapperModule(
|
||||
self.compiler(input_mod, args),
|
||||
unwrap_singleton_tuple,
|
||||
)
|
||||
return wrapper
|
||||
|
||||
# Note:
|
||||
#
|
||||
# The way distributed works today around fake tensors can be somewhat confusing.
|
||||
# Some of these codepaths are shared in both runtime, and compile time. The presence
|
||||
# of a fake_mode, read off of fake tensor inputs, dictates how we will operate.
|
||||
#
|
||||
# A few things to keep in mind:
|
||||
#
|
||||
# 1) We invoke `compile_submod` with a real module. The output of that gets stored
|
||||
# on the graph via `self.module.add_submodule(n.target, compiled_submod_real)`.
|
||||
#
|
||||
# 2) When running a call_module targeted node, if we have a fake_mode, we fakify the
|
||||
# module we got from self.fetch_attr(n.target). Regardless of fake_mode, we then execute it.
|
||||
#
|
||||
# 3) Fake tensors should always be around during compile time.
|
||||
#
|
||||
# 4) Fake tensors should never be around at runtime.
|
||||
#
|
||||
# 5) We end up with a compilation mode that takes a real submodule and fake tensors,
|
||||
# to match what aot_autograd expects. See Note: [Fake Modules and AOTAutograd]
|
||||
def run_node(self, n: Node) -> Any:
|
||||
args, kwargs = self.fetch_args_kwargs_from_env(n)
|
||||
new_args = []
|
||||
assert self.fake_mode
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor) and not isinstance(
|
||||
arg, torch._subclasses.FakeTensor
|
||||
):
|
||||
new_args.append(torch._dynamo.utils.to_fake_tensor(arg, self.fake_mode))
|
||||
else:
|
||||
new_args.append(arg)
|
||||
|
||||
log.debug("run_node %s, %s got args %s", n.op, n.target, args_str(args))
|
||||
assert isinstance(args, tuple)
|
||||
assert isinstance(kwargs, dict)
|
||||
|
||||
if n.op == "call_module":
|
||||
real_mod = self.fetch_attr(str(n.target))
|
||||
if self.fake_mode:
|
||||
curr_submod = deepcopy_to_fake_tensor(real_mod, self.fake_mode)
|
||||
else:
|
||||
curr_submod = real_mod
|
||||
|
||||
ddp_graph_log.debug("\n---%s graph---\n%s", n.target, curr_submod.graph)
|
||||
|
||||
# When calling the compiler on the submod, inputs (new_args) are expected to
|
||||
# be FakeTensors already since Dynamo would have made them FakeTensors in the
|
||||
# non-DDP flow. However, the parameters are _not_ expected to be FakeTensors,
|
||||
# since this wrapping happens during compilation
|
||||
|
||||
# Note: Returning Fake Tensors on First AOT Autograd Call
|
||||
#
|
||||
# Inductor will optimize strides of outputs when it deems it profitable.
|
||||
# For instance, converting to channels last. When we split the graph here
|
||||
# into multiple inductor compilations, we need to make sure that the
|
||||
# output strides of one compilation is appropriately passed to the subsequent
|
||||
# compilations. However, the mapping from inductor output to dynamo output
|
||||
# is non-trivial due to aot_autograd's deduping, de-aliasing, mutation, re-writing,
|
||||
# subclass handling, etc. In order to replay all this logic we set a flag such that
|
||||
# the first invocation of inductor in aot_autograd will return Fake Tensors with
|
||||
# appropriate strides. Then, all of aot autograd's runtime logic is replayed.
|
||||
# This gives us the appropriately strided outputs here which will reflect runtime strides.
|
||||
|
||||
class FakeifyFirstAOTInvocationGuard:
|
||||
def __init__(self) -> None:
|
||||
self.tc = torch._guards.TracingContext.try_get()
|
||||
assert self.tc
|
||||
self.tc.fakify_first_call = True
|
||||
|
||||
def __del__(self) -> None:
|
||||
self.tc.fakify_first_call = False # type: ignore[union-attr]
|
||||
|
||||
# For aot_eager and other backends, tracing context is not set
|
||||
has_tracing_context = torch._guards.TracingContext.try_get() is not None
|
||||
if has_tracing_context:
|
||||
g = FakeifyFirstAOTInvocationGuard() # noqa: F841
|
||||
|
||||
from torch._dynamo.utils import counters
|
||||
|
||||
init = counters["aot_autograd"]["total"]
|
||||
compiled_submod_real = self.compile_submod(real_mod, new_args, kwargs)
|
||||
|
||||
# TODO - better way of doing this?
|
||||
# Only aot autograd handles fakifying first call
|
||||
invoked_aot_autograd = init != counters["aot_autograd"]["total"]
|
||||
|
||||
# We update the original (outer) graph with a call into the compiled module
|
||||
# instead of the uncompiled one.
|
||||
self.module.delete_submodule(n.target) # type: ignore[operator]
|
||||
n.target = "compiled_" + n.target # type: ignore[operator]
|
||||
self.module.add_submodule(n.target, compiled_submod_real) # type: ignore[operator]
|
||||
|
||||
# Finally, we have to produce inputs for use compiling the next submodule,
|
||||
# and these need to be FakeTensors, so we execute the module under fake_mode
|
||||
# Because parameters are not fake we patch fake tensor mode to allow non fake inputs
|
||||
with (
|
||||
self.fake_mode,
|
||||
mock.patch.object(self.fake_mode, "allow_non_fake_inputs", True),
|
||||
):
|
||||
if has_tracing_context and invoked_aot_autograd:
|
||||
tracing_ctx = torch._guards.TracingContext.try_get()
|
||||
assert tracing_ctx is not None
|
||||
# DDPOptimizer maintains 1 dynamo graph -> N AOT graphs
|
||||
# Dynamo only has 1 tracing context, so it needs to maintain all N AOT metadata instances
|
||||
ddp_ctx = tracing_ctx.ddp_optimizer_ctx
|
||||
assert ddp_ctx is not None
|
||||
assert tracing_ctx.fw_metadata is not None
|
||||
ddp_ctx.curr_bucket += 1
|
||||
ddp_ctx.metadata_per_bucket.append(tracing_ctx.fw_metadata)
|
||||
|
||||
out = compiled_submod_real(*new_args, **kwargs)
|
||||
# output should be fake or subclass
|
||||
assert all(
|
||||
(not isinstance(t, torch.Tensor) or type(t) is not torch.Tensor)
|
||||
for t in (out if isinstance(out, (list, tuple)) else [out])
|
||||
)
|
||||
return out
|
||||
else:
|
||||
return curr_submod(*new_args, **kwargs)
|
||||
else:
|
||||
# placeholder or output nodes don't need to get compiled, just executed
|
||||
return getattr(self, n.op)(n.target, new_args, kwargs)
|
||||
|
||||
|
||||
class DDPOptimizer:
|
||||
"""Note [DDPOptimizer]
|
||||
DDPOptimizer applies when dynamo compiles models wrapped in DistributedDataParallel (DDP),
|
||||
breaking the dynamo graph into chunks to compile separately, with the breaks aligning to
|
||||
the boundaries of gradient-allreduce buckets chosen by DDP.
|
||||
|
||||
Background/Motivation
|
||||
- DDP uses allreduce collectives to synchronize partial gradients computed on different workers
|
||||
- DDP groups gradient allreduces into 'buckets' to optimize communication efficiency of all-reduce
|
||||
- Parameters grouped into buckets are assumed to be adjacent in time, so they become ready
|
||||
at around the same time during backward and thus can share the same allreduce efficiently
|
||||
- Allreduces must overlap with backward compute for optimal training performance
|
||||
- DDP schedules allreduces using 'hooks' fired from the c++ autograd engine in pytorch, which
|
||||
operates when individual grads become 'ready'
|
||||
- Dynamo+AOTAutograd produces a single fused graph that runs 'atomically' from the perspective of the
|
||||
autograd engine, such that all gradients become 'ready' at the same time. Hooks fire after the whole
|
||||
fused backward function executes, preventing any overlap of compute and communication
|
||||
|
||||
Algorithm
|
||||
- DDPOptimizer starts off with an FX graph traced by dynamo which represents forward. It can traverse
|
||||
this graph in reverse order to determine the true order that gradients will become ready during backward.
|
||||
- Parameter sizes are counted in reverse order, up to a bucket size limit, at which point a new bucket is started
|
||||
and a graph break introduced
|
||||
- Each of the subgraphs is compiled by the compiler provided to dynamo by the user, and then fused back together
|
||||
into an outer module that is returned to the user
|
||||
|
||||
Notes
|
||||
- It would be better to enforce (by adding an API to DDP) that the bucket splits chosen here are used by DDP,
|
||||
and that DDP does not need to detect or optimize bucket order by observing execution at runtime, as it does
|
||||
in eager.
|
||||
- If Dynamo can't capture a whole graph for the portion of the model wrapped by DDP, this algorithm will currently
|
||||
produce splits that do not necessarily align with the buckets used by DDP. This should result in performance
|
||||
degradation approaching the baseline case where graph-splits are not used, but not worse.
|
||||
- If the backend compiler fails to compile a single subgraph, it will execute eagerly despite the rest of the
|
||||
subgraphs being compiled
|
||||
- DDP has a 'parameters_and_buffers_to_ignore' field, which DDPOptimizer attempts to honor by reading markers
|
||||
left by DDP on individual parameters. In cases where other transformations, such as reparameterization, are
|
||||
also used, the ignore markers could be lost. If DDPOptimizer fails to ignore a parameter ignored by DDP,
|
||||
it is not catastrophic but could impact performance by choosing sub-optimal bucket splits.
|
||||
- DDPOptimizer always ignores all buffers, regardless of their ignore flag, since buffers do not require gradients,
|
||||
and therefore aren't allreduced by DDP. (They are broadcast during forward, but this is not covered by
|
||||
DDPOptimizer)
|
||||
|
||||
Debugging
|
||||
- Generally, it is easiest to debug DDPOptimizer in a single process program, using pdb.
|
||||
- In many cases, the log messages are helpful (they show bucket size assignments)-
|
||||
just set TORCH_LOGS env to include any of 'dynamo', 'distributed', or 'dist_ddp'.
|
||||
- See `benchmarks/dynamo/distributed.py` for a simple harness that will run a toy model or a torchbench model
|
||||
in a single process (or with torchrun, in multiple processes)
|
||||
|
||||
Args:
|
||||
bucket_bytes_cap (int): Controls the size of buckets, in bytes, used to determine graphbreaks. Should be
|
||||
set to match the equivalent parameter on the original DDP module.
|
||||
|
||||
backend_compile_fn (callable): A dynamo compiler function, to be invoked to compile each subgraph.
|
||||
|
||||
first_bucket_cap (int): Controls the size of the first bucket. Should match DDP's first bucket cap. DDP
|
||||
special-cases the first bucket size since it is sometimes optimal to start a small allreduce early.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket_bytes_cap: int,
|
||||
backend_compile_fn: CompilerFn,
|
||||
first_bucket_cap: int | None = None,
|
||||
) -> None:
|
||||
if first_bucket_cap is not None:
|
||||
self.first_bucket_cap = first_bucket_cap
|
||||
elif torch.distributed.is_available():
|
||||
# this constant comes from C10D lib which is not always built
|
||||
self.first_bucket_cap = torch.distributed._DEFAULT_FIRST_BUCKET_BYTES
|
||||
else:
|
||||
self.first_bucket_cap = bucket_bytes_cap
|
||||
|
||||
self.bucket_bytes_cap = bucket_bytes_cap
|
||||
assert self.first_bucket_cap <= self.bucket_bytes_cap, (
|
||||
"First bucket should be smaller/equal to other buckets to get comms warmed up ASAP"
|
||||
)
|
||||
|
||||
self.backend_compile_fn = backend_compile_fn
|
||||
|
||||
def _ignore_parameter(self, parameter: torch.nn.Parameter) -> bool:
|
||||
return hasattr(parameter, "_ddp_ignored") and parameter._ddp_ignored
|
||||
|
||||
def add_param(self, bucket: Bucket, param: torch.nn.Parameter, name: str) -> None:
|
||||
bucket.size += param.untyped_storage().nbytes()
|
||||
bucket.params.append(name)
|
||||
bucket.param_ids.append(id(param))
|
||||
|
||||
def add_module_params_to_bucket(
|
||||
self,
|
||||
mod: torch.nn.Module,
|
||||
bucket: Bucket,
|
||||
processed_modules: set[torch.nn.Module],
|
||||
prefix: str,
|
||||
) -> None:
|
||||
processed_modules.add(mod)
|
||||
for name, param in mod.named_parameters():
|
||||
if param.requires_grad and not self._ignore_parameter(param):
|
||||
self.add_param(bucket, param, f"{prefix}_{name}")
|
||||
|
||||
def add_param_args(self, bucket: Bucket, node: fx.Node) -> None:
|
||||
for arg in node.args:
|
||||
if not isinstance(arg, torch.fx.node.Node):
|
||||
continue
|
||||
if arg.op != "placeholder":
|
||||
continue
|
||||
param = arg.meta["example_value"]
|
||||
if (
|
||||
isinstance(param, torch.nn.Parameter)
|
||||
and param.requires_grad
|
||||
and not self._ignore_parameter(param)
|
||||
):
|
||||
self.add_param(bucket, param, str(arg.target))
|
||||
|
||||
def compile_fn(
|
||||
self, gm: fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> CompiledFn:
|
||||
"""
|
||||
Implements graph splitting, first determining a set of of buckets by counting
|
||||
parameter sizes in reverse graph order, then invoking the user/backend compiler
|
||||
to compile each subgraph. Finally, stitches compiled graphs into one graphmodule
|
||||
and returns its callable.
|
||||
"""
|
||||
# 1: compute the partition map according to DDP bucket logic
|
||||
buckets = [Bucket()] # (size, param_names)
|
||||
processed_modules: set[torch.nn.Module] = set()
|
||||
for node in reversed(gm.graph.nodes):
|
||||
if node.op in ("output", "placeholder"):
|
||||
continue
|
||||
|
||||
if (
|
||||
buckets[0].size >= self.bucket_bytes_cap
|
||||
or len(buckets) == 1
|
||||
and buckets[0].size >= self.first_bucket_cap
|
||||
):
|
||||
if bucket_has_external_output(buckets[0]):
|
||||
buckets.insert(0, Bucket())
|
||||
else:
|
||||
# continue building this bucket past the point of filling its parameter capacity,
|
||||
# to increase chances it contains at least one node that is either a global output or
|
||||
# passed as input to a subsequent graph
|
||||
|
||||
if buckets[0].opcount_increased_to_capture_external_output == 0:
|
||||
buckets[0].paramsize_before_opcount_increase = buckets[0].size
|
||||
buckets[0].opcount_increased_to_capture_external_output += 1
|
||||
|
||||
if node.op == "call_function":
|
||||
self.add_param_args(buckets[0], node)
|
||||
|
||||
elif node.op == "call_module":
|
||||
target_mod = gm.get_submodule(node.target)
|
||||
if target_mod not in processed_modules:
|
||||
self.add_module_params_to_bucket(
|
||||
target_mod, buckets[0], processed_modules, node.target
|
||||
)
|
||||
elif node.op == "call_method":
|
||||
if isinstance(node.args[0].target, str):
|
||||
target_mod = None
|
||||
try:
|
||||
target_mod = gm.get_submodule(node.args[0].target)
|
||||
except AttributeError:
|
||||
pass
|
||||
if target_mod is not None and target_mod not in processed_modules:
|
||||
self.add_module_params_to_bucket(
|
||||
target_mod, buckets[0], processed_modules, node.target
|
||||
)
|
||||
# This handles situations like tmp = torch.mm(x, self.weight.t())
|
||||
# t: "f32[512, 512]" = l_self_seq_2_weight.t(); l_self_seq_2_weight = None
|
||||
# tmp: "f32[512, 512]" = torch.mm(input_2, t); input_2 = t = None
|
||||
self.add_param_args(buckets[0], node)
|
||||
|
||||
elif node.op == "get_attr":
|
||||
maybe_param = getattr(gm, node.target)
|
||||
if (
|
||||
isinstance(maybe_param, torch.nn.Parameter)
|
||||
and maybe_param.requires_grad
|
||||
and not self._ignore_parameter(maybe_param)
|
||||
):
|
||||
self.add_param(buckets[0], maybe_param, node.target)
|
||||
|
||||
# All nodes have to be mapped to a bucket, even if they don't have their own params
|
||||
# Ignored params still end up in buckets, we just don't count them towards the capacity
|
||||
buckets[0].nodes.append(node)
|
||||
|
||||
if len(buckets) > 1 and buckets[0].size == 0:
|
||||
# we collected a small preamble graph with ops that don't include parameters, fuse it back
|
||||
buckets[1].nodes.extend(buckets[0].nodes)
|
||||
assert len(buckets[0].params) == 0, "Params should be empty if size is 0"
|
||||
del buckets[0]
|
||||
|
||||
# stash buckets for testing/debugging purposes
|
||||
self.buckets = buckets
|
||||
pretty_print_buckets(buckets, self.bucket_bytes_cap)
|
||||
|
||||
if len(buckets) == 1:
|
||||
# bypass split/fuse logic if there is only one bucket
|
||||
return self.backend_compile_fn(gm, example_inputs)
|
||||
|
||||
# 2: partition the graphmodule according to bucket capacity
|
||||
partition_map = {}
|
||||
for idx, b in enumerate(buckets):
|
||||
for node in b.nodes:
|
||||
partition_map[node] = idx
|
||||
|
||||
split_gm = fx.passes.split_module.split_module(
|
||||
gm,
|
||||
None, # type: ignore[arg-type]
|
||||
lambda node: partition_map[node],
|
||||
)
|
||||
|
||||
# See note [Assumption on Dynamo Metadata]
|
||||
propagate_dynamo_source(gm, split_gm)
|
||||
propagate_metadata(gm, split_gm)
|
||||
|
||||
debug_str = (
|
||||
f"\n---orig graph---\n{gm.graph}\n"
|
||||
+ f"\n---split graph---\n{split_gm.graph}\n"
|
||||
)
|
||||
for name, module in split_gm.named_modules():
|
||||
if "." not in name and len(name):
|
||||
# only print the submod graphs, not their children
|
||||
debug_str += f"\n---{name} graph---\n{module.graph}\n"
|
||||
debug_str += "\n---------------\n"
|
||||
ddp_graph_log.debug(debug_str)
|
||||
|
||||
trace_structured(
|
||||
"optimize_ddp_split_graph",
|
||||
payload_fn=lambda: split_gm.print_readable(print_output=False),
|
||||
)
|
||||
for name, module in split_gm.named_modules():
|
||||
if "." not in name and len(name):
|
||||
trace_structured(
|
||||
"optimize_ddp_split_child",
|
||||
lambda: {"name": name},
|
||||
payload_fn=lambda: module.print_readable(print_output=False),
|
||||
)
|
||||
|
||||
fake_mode = detect_fake_mode(example_inputs)
|
||||
if fake_mode is None:
|
||||
fake_mode = torch._subclasses.fake_tensor.FakeTensorMode()
|
||||
|
||||
submod_compiler = SubmodCompiler(split_gm, self.backend_compile_fn, fake_mode)
|
||||
with torch._dynamo.utils._disable_saved_tensors_hooks_during_tracing():
|
||||
submod_compiler.run(*example_inputs)
|
||||
split_gm.recompile()
|
||||
|
||||
ddp_graph_log.debug(
|
||||
"\n---final graph---\n%s\n---------------\n", split_gm.graph
|
||||
)
|
||||
return split_gm
|
||||
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
This module provides the TorchInductor backend integration for TorchDynamo.
|
||||
|
||||
TorchInductor is a compiler backend that generates optimized code for both CPU and GPU.
|
||||
This module lazily imports and registers the TorchInductor compiler to avoid loading it
|
||||
into memory when it is not being used. This helps reduce memory overhead when using
|
||||
other backends.
|
||||
|
||||
The inductor backend can be used with torch.compile():
|
||||
model = torch.compile(model, backend="inductor")
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from torch._dynamo import register_backend
|
||||
from torch._dynamo.utils import dynamo_timed
|
||||
|
||||
|
||||
@register_backend
|
||||
def inductor(*args: Any, **kwargs: Any) -> Any:
|
||||
with dynamo_timed("inductor_import", log_pt2_compile_event=True):
|
||||
# do import here to avoid loading inductor into memory when it is not used
|
||||
# The AsyncCompile subproc pool can be slow to start, so warm it up as early
|
||||
# as possible.
|
||||
from torch._inductor.async_compile import maybe_warm_pool
|
||||
|
||||
maybe_warm_pool()
|
||||
|
||||
from torch._inductor.compile_fx import compile_fx
|
||||
|
||||
return compile_fx(*args, **kwargs)
|
||||
@@ -0,0 +1,39 @@
|
||||
# This backend is maintained by ONNX team. To direct issues
|
||||
# to the right people, please tag related GitHub issues with `module: onnx`.
|
||||
#
|
||||
# Maintainers' Github IDs: wschin, xadupre
|
||||
# from torch.onnx._internal.onnxruntime import (
|
||||
# is_onnxrt_backend_supported,
|
||||
# torch_compile_backend,
|
||||
# )
|
||||
|
||||
# from .registry import register_backend
|
||||
|
||||
"""
|
||||
Placeholder for onnxruntime backend for dynamo
|
||||
"""
|
||||
|
||||
# def has_onnxruntime():
|
||||
# # FIXME: update test/dynamo/test_backends.py to call is_onnxrt_backend_supported()
|
||||
# return is_onnxrt_backend_supported()
|
||||
|
||||
|
||||
# if is_onnxrt_backend_supported():
|
||||
# register_backend(name="onnxrt", compiler_fn=torch_compile_backend)
|
||||
# else:
|
||||
|
||||
# def information_displaying_backend(*args, **kwargs):
|
||||
# raise ImportError(
|
||||
# "onnxrt is not registered as a backend. "
|
||||
# "Please make sure all dependencies such as "
|
||||
# "numpy, onnx, onnxscript, and onnxruntime-training are installed. "
|
||||
# "Suggested procedure to fix dependency problem:\n"
|
||||
# " (1) pip or conda install numpy onnx onnxscript onnxruntime-training.\n"
|
||||
# " (2) Open a new python terminal.\n"
|
||||
# " (3) Call the API `torch.onnx.is_onnxrt_backend_supported()`:\n"
|
||||
# " (4) If it returns `True`, then you can use `onnxrt` backend.\n"
|
||||
# " (5) If it returns `False`, please execute the package importing section in "
|
||||
# "torch/onnx/_internal/onnxruntime.py under pdb line-by-line to see which import fails."
|
||||
# )
|
||||
|
||||
# register_backend(name="onnxrt", compiler_fn=information_displaying_backend)
|
||||
@@ -0,0 +1,206 @@
|
||||
"""
|
||||
This module implements TorchDynamo's backend registry system for managing compiler backends.
|
||||
|
||||
The registry provides a centralized way to register, discover and manage different compiler
|
||||
backends that can be used with torch.compile(). It handles:
|
||||
|
||||
- Backend registration and discovery through decorators and entry points
|
||||
- Lazy loading of backend implementations
|
||||
- Lookup and validation of backend names
|
||||
- Categorization of backends using tags (debug, experimental, etc.)
|
||||
|
||||
Key components:
|
||||
- CompilerFn: Type for backend compiler functions that transform FX graphs
|
||||
- _BACKENDS: Registry mapping backend names to entry points
|
||||
- _COMPILER_FNS: Registry mapping backend names to loaded compiler functions
|
||||
|
||||
Example usage:
|
||||
@register_backend
|
||||
def my_compiler(fx_graph, example_inputs):
|
||||
# Transform FX graph into optimized implementation
|
||||
return compiled_fn
|
||||
|
||||
# Use registered backend
|
||||
torch.compile(model, backend="my_compiler")
|
||||
|
||||
The registry also supports discovering backends through setuptools entry points
|
||||
in the "torch_dynamo_backends" group. Example:
|
||||
```
|
||||
setup.py
|
||||
---
|
||||
from setuptools import setup
|
||||
|
||||
setup(
|
||||
name='my_torch_backend',
|
||||
version='0.1',
|
||||
packages=['my_torch_backend'],
|
||||
entry_points={
|
||||
'torch_dynamo_backends': [
|
||||
# name = path to entry point of backend implementation
|
||||
'my_compiler = my_torch_backend.compiler:my_compiler_function',
|
||||
],
|
||||
},
|
||||
)
|
||||
```
|
||||
```
|
||||
my_torch_backend/compiler.py
|
||||
---
|
||||
def my_compiler_function(fx_graph, example_inputs):
|
||||
# Transform FX graph into optimized implementation
|
||||
return compiled_fn
|
||||
```
|
||||
Using `my_compiler` backend:
|
||||
```
|
||||
import torch
|
||||
|
||||
model = ... # Your PyTorch model
|
||||
optimized_model = torch.compile(model, backend="my_compiler")
|
||||
```
|
||||
"""
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from collections.abc import Callable, Sequence
|
||||
from importlib.metadata import EntryPoint
|
||||
from typing import Any, Protocol
|
||||
|
||||
import torch
|
||||
from torch import fx
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CompiledFn(Protocol):
|
||||
def __call__(self, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: ...
|
||||
|
||||
|
||||
CompilerFn = Callable[[fx.GraphModule, list[torch.Tensor]], CompiledFn]
|
||||
|
||||
_BACKENDS: dict[str, EntryPoint | None] = {}
|
||||
_COMPILER_FNS: dict[str, CompilerFn] = {}
|
||||
|
||||
|
||||
def register_backend(
|
||||
compiler_fn: CompilerFn | None = None,
|
||||
name: str | None = None,
|
||||
tags: Sequence[str] = (),
|
||||
) -> Callable[..., Any]:
|
||||
"""
|
||||
Decorator to add a given compiler to the registry to allow calling
|
||||
`torch.compile` with string shorthand. Note: for projects not
|
||||
imported by default, it might be easier to pass a function directly
|
||||
as a backend and not use a string.
|
||||
|
||||
Args:
|
||||
compiler_fn: Callable taking a FX graph and fake tensor inputs
|
||||
name: Optional name, defaults to `compiler_fn.__name__`
|
||||
tags: Optional set of string tags to categorize backend with
|
||||
"""
|
||||
if compiler_fn is None:
|
||||
# @register_backend(name="") syntax
|
||||
return functools.partial(register_backend, name=name, tags=tags) # type: ignore[return-value]
|
||||
assert callable(compiler_fn)
|
||||
name = name or compiler_fn.__name__
|
||||
assert name not in _COMPILER_FNS, f"duplicate name: {name}"
|
||||
if compiler_fn not in _BACKENDS:
|
||||
_BACKENDS[name] = None
|
||||
_COMPILER_FNS[name] = compiler_fn
|
||||
compiler_fn._tags = tuple(tags) # type: ignore[attr-defined]
|
||||
return compiler_fn
|
||||
|
||||
|
||||
register_debug_backend = functools.partial(register_backend, tags=("debug",))
|
||||
register_experimental_backend = functools.partial(
|
||||
register_backend, tags=("experimental",)
|
||||
)
|
||||
|
||||
|
||||
def lookup_backend(compiler_fn: str | CompilerFn) -> CompilerFn:
|
||||
"""Expand backend strings to functions"""
|
||||
if isinstance(compiler_fn, str):
|
||||
if compiler_fn not in _BACKENDS:
|
||||
_lazy_import()
|
||||
if compiler_fn not in _BACKENDS:
|
||||
from ..exc import InvalidBackend
|
||||
|
||||
raise InvalidBackend(name=compiler_fn)
|
||||
|
||||
if compiler_fn not in _COMPILER_FNS:
|
||||
entry_point = _BACKENDS[compiler_fn]
|
||||
if entry_point is not None:
|
||||
register_backend(compiler_fn=entry_point.load(), name=compiler_fn)
|
||||
compiler_fn = _COMPILER_FNS[compiler_fn]
|
||||
return compiler_fn
|
||||
|
||||
|
||||
# NOTE: can't type this due to public api mismatch; follow up with dev team
|
||||
def list_backends(exclude_tags=("debug", "experimental")) -> list[str]: # type: ignore[no-untyped-def]
|
||||
"""
|
||||
Return valid strings that can be passed to:
|
||||
|
||||
torch.compile(..., backend="name")
|
||||
"""
|
||||
_lazy_import()
|
||||
exclude_tags_set = set(exclude_tags or ())
|
||||
|
||||
backends = [
|
||||
name
|
||||
for name in _BACKENDS
|
||||
if name not in _COMPILER_FNS
|
||||
or not exclude_tags_set.intersection(_COMPILER_FNS[name]._tags) # type: ignore[attr-defined]
|
||||
]
|
||||
return sorted(backends)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _lazy_import() -> None:
|
||||
from .. import backends
|
||||
from ..utils import import_submodule
|
||||
|
||||
import_submodule(backends)
|
||||
|
||||
from ..repro.after_dynamo import dynamo_minifier_backend
|
||||
|
||||
assert dynamo_minifier_backend is not None
|
||||
|
||||
_discover_entrypoint_backends()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _discover_entrypoint_backends() -> None:
|
||||
# importing here so it will pick up the mocked version in test_backends.py
|
||||
from importlib.metadata import entry_points
|
||||
|
||||
group_name = "torch_dynamo_backends"
|
||||
eps = entry_points(group=group_name)
|
||||
# pyrefly: ignore [bad-index]
|
||||
eps_dict = {name: eps[name] for name in eps.names}
|
||||
for backend_name in eps_dict:
|
||||
_BACKENDS[backend_name] = eps_dict[backend_name]
|
||||
|
||||
|
||||
def _is_registered_backend(compiler_fn: CompilerFn) -> bool:
|
||||
"""
|
||||
Check if the given compiler function is a registered backend.
|
||||
Custom backends (user-provided callables not in the registry) return False.
|
||||
"""
|
||||
# Ensure backends are loaded
|
||||
_lazy_import()
|
||||
|
||||
# Check if it's directly a registered backend function
|
||||
if compiler_fn in _COMPILER_FNS.values():
|
||||
return True
|
||||
|
||||
# Check for _TorchCompileInductorWrapper or _TorchCompileWrapper
|
||||
# These have a compiler_name attribute that identifies the backend
|
||||
if hasattr(compiler_fn, "compiler_name"):
|
||||
compiler_name = compiler_fn.compiler_name
|
||||
if compiler_name in _BACKENDS or compiler_name in _COMPILER_FNS:
|
||||
return True
|
||||
|
||||
# Check if the wrapper has a compiler_fn attribute (e.g., _TorchCompileWrapper)
|
||||
if hasattr(compiler_fn, "compiler_fn"):
|
||||
return compiler_fn.compiler_fn in _COMPILER_FNS.values()
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,12 @@
|
||||
# import torch # type: ignore[import]
|
||||
# from .common import device_from_inputs, fake_tensor_unsupported # type: ignore[import]
|
||||
# from .registry import register_backend # type: ignore[import]
|
||||
|
||||
"""
|
||||
Placeholder for TensorRT backend for dynamo via torch-tensorrt
|
||||
"""
|
||||
|
||||
# @register_backend
|
||||
# def tensorrt(gm, example_inputs):
|
||||
# import torch_tensorrt # type: ignore[import]
|
||||
# pass
|
||||
@@ -0,0 +1,55 @@
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from functorch.compile import make_boxed_func
|
||||
from torch import fx
|
||||
|
||||
from ..backends.common import aot_autograd
|
||||
from .registry import CompiledFn, register_backend, register_experimental_backend
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_experimental_backend
|
||||
def openxla_eval(
|
||||
model: fx.GraphModule, fake_tensor_inputs: list[torch.Tensor]
|
||||
) -> CompiledFn:
|
||||
return xla_backend_helper(model, fake_tensor_inputs, boxed=False)
|
||||
|
||||
|
||||
def openxla_eval_boxed(
|
||||
model: fx.GraphModule, fake_tensor_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
return xla_backend_helper(model, fake_tensor_inputs, boxed=True)
|
||||
|
||||
|
||||
def xla_backend_helper(
|
||||
model: fx.GraphModule, fake_tensor_inputs: list[torch.Tensor], boxed: bool = False
|
||||
) -> Callable[..., Any]:
|
||||
try:
|
||||
import torch_xla.core.dynamo_bridge as bridge
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Please follow the instruction in https://github.com/pytorch/xla#pytorchxla to install torch_xla"
|
||||
) from e
|
||||
|
||||
compiled_graph = None
|
||||
|
||||
def fwd(*args: torch.Tensor) -> Any:
|
||||
nonlocal model
|
||||
nonlocal compiled_graph
|
||||
if compiled_graph is None:
|
||||
compiled_graph = bridge.extract_compiled_graph(model, args)
|
||||
del model
|
||||
return compiled_graph(*args)
|
||||
|
||||
return make_boxed_func(fwd) if boxed else fwd
|
||||
|
||||
|
||||
openxla = aot_autograd(
|
||||
fw_compiler=openxla_eval_boxed,
|
||||
)
|
||||
register_backend(name="openxla", compiler_fn=openxla)
|
||||
@@ -0,0 +1,197 @@
|
||||
"""
|
||||
This module provides TVM backend integration for TorchDynamo.
|
||||
|
||||
Apache TVM is a deep learning compiler framework that can optimize and execute
|
||||
models on various hardware backends. This module enables:
|
||||
|
||||
- Compilation of PyTorch models to TVM's computation graphs
|
||||
- Multiple scheduling options:
|
||||
- Default scheduler
|
||||
- Auto-scheduler for automatic optimization
|
||||
- Meta-schedule for evolutionary search-based tuning
|
||||
- Hardware-specific optimizations:
|
||||
- CUDA GPU support
|
||||
- CPU support with LLVM targeting and architecture-specific tuning
|
||||
- Automatic detection of CPU capabilities (AVX2, AVX512)
|
||||
- Tensor conversion utilities between PyTorch and TVM formats
|
||||
- Configurable optimization levels and tuning trials
|
||||
|
||||
The backend can be used with torch.compile():
|
||||
model = torch.compile(model, backend="tvm")
|
||||
"""
|
||||
|
||||
import functools
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import fx
|
||||
|
||||
from .common import device_from_inputs, fake_tensor_unsupported
|
||||
from .registry import register_backend
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_backend
|
||||
@fake_tensor_unsupported # type: ignore[arg-type]
|
||||
def tvm(
|
||||
gm: fx.GraphModule,
|
||||
example_inputs: list[torch.Tensor],
|
||||
*,
|
||||
options: MappingProxyType[str, Any] | None = None,
|
||||
) -> Callable[..., Any]:
|
||||
if options is None:
|
||||
options = MappingProxyType({"scheduler": None, "trials": 20000, "opt_level": 3})
|
||||
assert options is not None
|
||||
import tvm # type: ignore[import]
|
||||
from tvm import relay # type: ignore[import]
|
||||
from tvm.contrib import graph_executor # type: ignore[import]
|
||||
|
||||
jit_mod = torch.jit.trace(gm, example_inputs)
|
||||
device = device_from_inputs(example_inputs)
|
||||
shape_list = [(f"inp_{idx}", i.shape) for idx, i in enumerate(example_inputs)]
|
||||
example_outputs = gm(*example_inputs)
|
||||
if len(example_outputs) == 0:
|
||||
log.warning("Explicitly fall back to eager due to zero output")
|
||||
return gm.forward
|
||||
mod, params = relay.frontend.from_pytorch(jit_mod, shape_list)
|
||||
if device.type == "cuda":
|
||||
dev = tvm.cuda(device.index)
|
||||
target = tvm.target.cuda()
|
||||
else:
|
||||
dev = tvm.cpu(0)
|
||||
target = tvm.target.Target(llvm_target())
|
||||
|
||||
scheduler = options.get("scheduler", None)
|
||||
if scheduler is None:
|
||||
scheduler = os.environ.get("TVM_SCHEDULER", None)
|
||||
|
||||
trials = options.get("trials", 20000)
|
||||
opt_level = options.get("opt_level", 3)
|
||||
|
||||
if scheduler == "auto_scheduler":
|
||||
# pyrefly: ignore [missing-import]
|
||||
from tvm import auto_scheduler
|
||||
|
||||
with (
|
||||
tempfile.NamedTemporaryFile() as log_file,
|
||||
auto_scheduler.ApplyHistoryBest(log_file),
|
||||
tvm.transform.PassContext(
|
||||
opt_level=opt_level, config={"relay.backend.use_auto_scheduler": True}
|
||||
),
|
||||
):
|
||||
lib = relay.build(mod, target=target, params=params)
|
||||
elif scheduler == "meta_schedule":
|
||||
# pyrefly: ignore [missing-import]
|
||||
from tvm import meta_schedule as ms
|
||||
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
if device.type != "cuda":
|
||||
# meta_schedule needs num-cores to be specified
|
||||
# here we use the maximum core count
|
||||
target = tvm.target.Target(
|
||||
f"{llvm_target()} --num-cores {ms.utils.cpu_count(logical=False)}"
|
||||
)
|
||||
# TODO(shingjan): This could be replaced by tvm.contrib.torch.optimize_torch
|
||||
# once USE_PT_TVMDSOOP is updated and turned on by default in TVM.
|
||||
assert trials > 0
|
||||
database = ms.relay_integration.tune_relay(
|
||||
mod=mod,
|
||||
target=target,
|
||||
work_dir=work_dir,
|
||||
max_trials_global=trials,
|
||||
num_trials_per_iter=64,
|
||||
params=params,
|
||||
strategy="evolutionary",
|
||||
opt_level=opt_level,
|
||||
)
|
||||
lib = ms.relay_integration.compile_relay(
|
||||
database=database,
|
||||
mod=mod,
|
||||
target=target,
|
||||
params=params,
|
||||
opt_level=opt_level,
|
||||
)
|
||||
elif scheduler == "default" or not scheduler:
|
||||
# no autotuning
|
||||
with tvm.transform.PassContext(opt_level=opt_level):
|
||||
lib = relay.build(mod, target=target, params=params)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"This tuning option is invalid/not implemented for torchdynamo's TVM-related backend. "
|
||||
"There are three available options: default, auto_scheduler and meta_schedule."
|
||||
)
|
||||
m = graph_executor.GraphModule(lib["default"](dev))
|
||||
|
||||
def to_torch_tensor(nd_tensor: tvm.nd.array) -> torch.Tensor:
|
||||
"""A helper function to transfer a NDArray to torch.tensor."""
|
||||
if nd_tensor.dtype == "bool":
|
||||
# DLPack does not support boolean so it can't be handled by
|
||||
# torch.utils.dlpack.from_pack. Workaround by going through
|
||||
# numpy, although this brings additional data copy overhead.
|
||||
return torch.from_numpy(nd_tensor.numpy())
|
||||
return torch.utils.dlpack.from_dlpack(nd_tensor.to_dlpack())
|
||||
|
||||
def to_tvm_tensor(torch_tensor: torch.Tensor) -> tvm.nd.array:
|
||||
"""A helper function to transfer a torch.tensor to NDArray."""
|
||||
if torch_tensor.dtype == torch.bool:
|
||||
# same reason as above, fallback to numpy conversion which
|
||||
# could introduce data copy overhead
|
||||
return tvm.nd.array(torch_tensor.cpu().numpy())
|
||||
return tvm.nd.from_dlpack(torch_tensor)
|
||||
|
||||
def exec_tvm(*i_args: torch.Tensor) -> list[torch.Tensor]:
|
||||
args = [a.contiguous() for a in i_args]
|
||||
shape_info, _ = m.get_input_info()
|
||||
active_inputs = {name for name, _ in shape_info.items()}
|
||||
for idx, arg in enumerate(args, 0):
|
||||
if arg.dim() != 0:
|
||||
if arg.requires_grad:
|
||||
arg = arg.detach()
|
||||
inp_name = f"inp_{idx}"
|
||||
if inp_name not in active_inputs:
|
||||
log.warning(
|
||||
"input %s skipped as not found in tvm's runtime library",
|
||||
inp_name,
|
||||
)
|
||||
continue
|
||||
m.set_input(
|
||||
inp_name,
|
||||
to_tvm_tensor(arg),
|
||||
)
|
||||
m.run()
|
||||
return [to_torch_tensor(m.get_output(i)) for i in range(m.get_num_outputs())]
|
||||
|
||||
return exec_tvm
|
||||
|
||||
|
||||
tvm_meta_schedule = functools.partial(tvm, scheduler="meta_schedule")
|
||||
tvm_auto_scheduler = functools.partial(tvm, scheduler="auto_scheduler")
|
||||
|
||||
|
||||
def has_tvm() -> bool:
|
||||
try:
|
||||
importlib.import_module("tvm")
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def llvm_target() -> str:
|
||||
if sys.platform == "linux":
|
||||
cpuinfo = Path("/proc/cpuinfo").read_text()
|
||||
if "avx512" in cpuinfo:
|
||||
return "llvm -mcpu=skylake-avx512"
|
||||
elif "avx2" in cpuinfo:
|
||||
return "llvm -mcpu=core-avx2"
|
||||
return "llvm"
|
||||
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
This module provides utilities for analyzing and optimizing Python bytecode.
|
||||
Key functionality includes:
|
||||
- Dead code elimination
|
||||
- Jump instruction optimization
|
||||
- Stack size analysis and verification
|
||||
- Live variable analysis
|
||||
- Line number propagation and cleanup
|
||||
- Exception table handling for Python 3.11+
|
||||
|
||||
The utilities in this module are used to analyze and transform bytecode
|
||||
for better performance while maintaining correct semantics.
|
||||
"""
|
||||
|
||||
import bisect
|
||||
import dataclasses
|
||||
import dis
|
||||
import itertools
|
||||
import sys
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# TODO(lucaskabela): consider moving Instruction into this file
|
||||
# and refactoring in callsite; that way we don't have to guard this import
|
||||
from .bytecode_transformation import Instruction
|
||||
|
||||
TERMINAL_OPCODES = {
|
||||
dis.opmap["RETURN_VALUE"],
|
||||
dis.opmap["JUMP_FORWARD"],
|
||||
dis.opmap["RAISE_VARARGS"],
|
||||
# TODO(jansel): double check exception handling
|
||||
}
|
||||
TERMINAL_OPCODES.add(dis.opmap["RERAISE"])
|
||||
if sys.version_info >= (3, 11):
|
||||
TERMINAL_OPCODES.add(dis.opmap["JUMP_BACKWARD"])
|
||||
TERMINAL_OPCODES.add(dis.opmap["JUMP_FORWARD"])
|
||||
else:
|
||||
TERMINAL_OPCODES.add(dis.opmap["JUMP_ABSOLUTE"])
|
||||
|
||||
if (3, 12) <= sys.version_info < (3, 14):
|
||||
TERMINAL_OPCODES.add(dis.opmap["RETURN_CONST"])
|
||||
if sys.version_info >= (3, 13):
|
||||
TERMINAL_OPCODES.add(dis.opmap["JUMP_BACKWARD_NO_INTERRUPT"])
|
||||
JUMP_OPCODES = set(dis.hasjrel + dis.hasjabs)
|
||||
JUMP_OPNAMES = {dis.opname[opcode] for opcode in JUMP_OPCODES}
|
||||
HASLOCAL = set(dis.haslocal)
|
||||
HASFREE = set(dis.hasfree)
|
||||
|
||||
stack_effect = dis.stack_effect
|
||||
|
||||
|
||||
def get_indexof(insts: list["Instruction"]) -> dict["Instruction", int]:
|
||||
"""
|
||||
Get a mapping from instruction memory address to index in instruction list.
|
||||
Additionally checks that each instruction only appears once in the list.
|
||||
"""
|
||||
# pyrefly: ignore [implicit-any]
|
||||
indexof = {}
|
||||
for i, inst in enumerate(insts):
|
||||
assert inst not in indexof
|
||||
indexof[inst] = i
|
||||
return indexof
|
||||
|
||||
|
||||
def remove_dead_code(instructions: list["Instruction"]) -> list["Instruction"]:
|
||||
"""Dead code elimination"""
|
||||
indexof = get_indexof(instructions)
|
||||
live_code = set()
|
||||
|
||||
def find_live_code(start: int) -> None:
|
||||
for i in range(start, len(instructions)):
|
||||
if i in live_code:
|
||||
return
|
||||
live_code.add(i)
|
||||
inst = instructions[i]
|
||||
if inst.exn_tab_entry:
|
||||
find_live_code(indexof[inst.exn_tab_entry.target])
|
||||
if inst.opcode in JUMP_OPCODES:
|
||||
assert inst.target is not None
|
||||
find_live_code(indexof[inst.target])
|
||||
if inst.opcode in TERMINAL_OPCODES:
|
||||
return
|
||||
|
||||
find_live_code(0)
|
||||
|
||||
# change exception table entries if start/end instructions are dead
|
||||
# assumes that exception table entries have been propagated,
|
||||
# e.g. with bytecode_transformation.propagate_inst_exn_table_entries,
|
||||
# and that instructions with an exn_tab_entry lies within its start/end.
|
||||
if sys.version_info >= (3, 11):
|
||||
live_idx = sorted(live_code)
|
||||
for i, inst in enumerate(instructions):
|
||||
if i in live_code and inst.exn_tab_entry:
|
||||
# find leftmost live instruction >= start
|
||||
start_idx = bisect.bisect_left(
|
||||
live_idx, indexof[inst.exn_tab_entry.start]
|
||||
)
|
||||
assert start_idx < len(live_idx)
|
||||
# find rightmost live instruction <= end
|
||||
end_idx = (
|
||||
bisect.bisect_right(live_idx, indexof[inst.exn_tab_entry.end]) - 1
|
||||
)
|
||||
assert end_idx >= 0
|
||||
assert live_idx[start_idx] <= i <= live_idx[end_idx]
|
||||
inst.exn_tab_entry.start = instructions[live_idx[start_idx]]
|
||||
inst.exn_tab_entry.end = instructions[live_idx[end_idx]]
|
||||
|
||||
return [inst for i, inst in enumerate(instructions) if i in live_code]
|
||||
|
||||
|
||||
def remove_pointless_jumps(instructions: list["Instruction"]) -> list["Instruction"]:
|
||||
"""Eliminate jumps to the next instruction"""
|
||||
pointless_jumps = {
|
||||
id(a)
|
||||
for a, b in itertools.pairwise(instructions)
|
||||
if a.opname == "JUMP_ABSOLUTE" and a.target is b
|
||||
}
|
||||
return [inst for inst in instructions if id(inst) not in pointless_jumps]
|
||||
|
||||
|
||||
def propagate_line_nums(instructions: list["Instruction"]) -> None:
|
||||
"""Ensure every instruction has line number set in case some are removed"""
|
||||
cur_line_no = None
|
||||
|
||||
def populate_line_num(inst: "Instruction") -> None:
|
||||
nonlocal cur_line_no
|
||||
if inst.starts_line:
|
||||
cur_line_no = inst.starts_line
|
||||
|
||||
inst.starts_line = cur_line_no
|
||||
|
||||
for inst in instructions:
|
||||
populate_line_num(inst)
|
||||
|
||||
|
||||
def remove_extra_line_nums(instructions: list["Instruction"]) -> None:
|
||||
"""Remove extra starts line properties before packing bytecode"""
|
||||
|
||||
cur_line_no = None
|
||||
|
||||
def remove_line_num(inst: "Instruction") -> None:
|
||||
nonlocal cur_line_no
|
||||
if inst.starts_line is None:
|
||||
return
|
||||
elif inst.starts_line == cur_line_no:
|
||||
inst.starts_line = None
|
||||
else:
|
||||
cur_line_no = inst.starts_line
|
||||
|
||||
for inst in instructions:
|
||||
remove_line_num(inst)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ReadsWrites:
|
||||
reads: set[Any]
|
||||
writes: set[Any]
|
||||
visited: set[Any]
|
||||
|
||||
|
||||
def livevars_analysis(
|
||||
instructions: list["Instruction"], instruction: "Instruction"
|
||||
) -> set[Any]:
|
||||
indexof = get_indexof(instructions)
|
||||
must = ReadsWrites(set(), set(), set())
|
||||
may = ReadsWrites(set(), set(), set())
|
||||
|
||||
def walk(state: ReadsWrites, start: int) -> None:
|
||||
if start in state.visited:
|
||||
return
|
||||
state.visited.add(start)
|
||||
|
||||
for i in range(start, len(instructions)):
|
||||
inst = instructions[i]
|
||||
if inst.opcode in HASLOCAL or inst.opcode in HASFREE:
|
||||
if "LOAD" in inst.opname or "DELETE" in inst.opname:
|
||||
if inst.argval not in must.writes:
|
||||
state.reads.add(inst.argval)
|
||||
elif "STORE" in inst.opname:
|
||||
state.writes.add(inst.argval)
|
||||
elif inst.opname == "MAKE_CELL":
|
||||
pass
|
||||
else:
|
||||
raise NotImplementedError(f"unhandled {inst.opname}")
|
||||
if inst.exn_tab_entry:
|
||||
walk(may, indexof[inst.exn_tab_entry.target])
|
||||
if inst.opcode in JUMP_OPCODES:
|
||||
assert inst.target is not None
|
||||
walk(may, indexof[inst.target])
|
||||
state = may
|
||||
if inst.opcode in TERMINAL_OPCODES:
|
||||
return
|
||||
|
||||
walk(must, indexof[instruction])
|
||||
return must.reads | may.reads
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class FixedPointBox:
|
||||
value: bool = True
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class StackSize:
|
||||
low: int | float
|
||||
high: int | float
|
||||
fixed_point: FixedPointBox
|
||||
|
||||
def zero(self) -> None:
|
||||
self.low = 0
|
||||
self.high = 0
|
||||
self.fixed_point.value = False
|
||||
|
||||
def offset_of(self, other: "StackSize", n: int) -> None:
|
||||
prior = (self.low, self.high)
|
||||
self.low = min(self.low, other.low + n)
|
||||
self.high = max(self.high, other.high + n)
|
||||
if (self.low, self.high) != prior:
|
||||
self.fixed_point.value = False
|
||||
|
||||
def exn_tab_jump(self, depth: int) -> None:
|
||||
prior = (self.low, self.high)
|
||||
self.low = min(self.low, depth)
|
||||
self.high = max(self.high, depth)
|
||||
if (self.low, self.high) != prior:
|
||||
self.fixed_point.value = False
|
||||
|
||||
|
||||
def stacksize_analysis(instructions: list["Instruction"]) -> int | float:
|
||||
assert instructions
|
||||
fixed_point = FixedPointBox()
|
||||
stack_sizes = {
|
||||
inst: StackSize(float("inf"), float("-inf"), fixed_point)
|
||||
for inst in instructions
|
||||
}
|
||||
stack_sizes[instructions[0]].zero()
|
||||
|
||||
for _ in range(100):
|
||||
if fixed_point.value:
|
||||
break
|
||||
fixed_point.value = True
|
||||
|
||||
for inst, next_inst in zip(instructions, instructions[1:] + [None]):
|
||||
stack_size = stack_sizes[inst]
|
||||
if inst.opcode not in TERMINAL_OPCODES:
|
||||
assert next_inst is not None, f"missing next inst: {inst}"
|
||||
eff = stack_effect(inst.opcode, inst.arg, jump=False)
|
||||
stack_sizes[next_inst].offset_of(stack_size, eff)
|
||||
if inst.opcode in JUMP_OPCODES:
|
||||
assert inst.target is not None, f"missing target: {inst}"
|
||||
stack_sizes[inst.target].offset_of(
|
||||
stack_size, stack_effect(inst.opcode, inst.arg, jump=True)
|
||||
)
|
||||
if inst.exn_tab_entry:
|
||||
# see https://github.com/python/cpython/blob/3.11/Objects/exception_handling_notes.txt
|
||||
# on why depth is computed this way.
|
||||
depth = inst.exn_tab_entry.depth + int(inst.exn_tab_entry.lasti) + 1
|
||||
stack_sizes[inst.exn_tab_entry.target].exn_tab_jump(depth)
|
||||
|
||||
low = min(x.low for x in stack_sizes.values())
|
||||
high = max(x.high for x in stack_sizes.values())
|
||||
|
||||
assert fixed_point.value, "failed to reach fixed point"
|
||||
assert low >= 0
|
||||
return high
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,187 @@
|
||||
import logging
|
||||
import weakref
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from torch._guards import CompileId
|
||||
|
||||
from . import config
|
||||
from .types import DynamoFrameType
|
||||
|
||||
|
||||
log: logging.Logger = logging.getLogger(__name__)
|
||||
"""
|
||||
[Note on cache size limit]
|
||||
|
||||
Background - TorchDynamo cache is a linked list. Each cache entry is a
|
||||
(guard_manager, out_code, next pointer). These are stored on the f_code's co_extra
|
||||
scratch space. When a frame is invoked, we walk this linked list and run
|
||||
guard_manager in each cache_entry to decide if the frame needs recompilation. If none
|
||||
of the guard_manager's returns True, we recompile and add a new entry. To ensure we
|
||||
don't end up recompiling infinitely, we put limits on the cache size.
|
||||
|
||||
There are two limits
|
||||
1) recompile_limit
|
||||
2) accumulated_recompile_limit
|
||||
|
||||
|
||||
Earlier we used to have only limit - maximum number of entries in 1 cache line
|
||||
(which is now represented by (2) above). So, why do we need two limits? Lets try
|
||||
to understand that.
|
||||
|
||||
In general, we want our cache limit value to be a small number (e.g. 8 or even
|
||||
lower). This ensures that for frames that cause too many recompilation fall to
|
||||
eager quickly. However, there is another problem that prevents us from lowering
|
||||
the value of recompile_limit. This is due to ID_MATCH'd guards. Today, we put
|
||||
ID_MATCH guards on nn module if there is a graph break. This means we will have
|
||||
many recompilations for the same code object because the ID_MATCH guard fails
|
||||
for different instances of the nn module. This is a common pattern in how models
|
||||
are authored. Therefore, this requires us to keep the recompile_limit high.
|
||||
|
||||
We resolve this by introducing these two limits. The first limit (1) limits the
|
||||
number of cache entries that have an ID_MATCH'd guard for an nn module instance.
|
||||
And, (2)nd limit becomes a safeguard mechanism to have a maximum compilations
|
||||
for a code object. One important question is - what is the limit for the code
|
||||
object that does not have any ID_MATCH guard? For such code objects, we choose
|
||||
(1) as the cache size limit.
|
||||
|
||||
Lets take an example to understand how these limits help. Suppose, we have 16
|
||||
instances of a nn module and we ID_MATCH on the self object. Further, suppose
|
||||
the inputs to these functions have varying batch size, leading to one
|
||||
recompilation. In total, there will be 32 recompilations, and therefore 32 cache
|
||||
entries on the forward code object. In the older case when we had only 1 limit,
|
||||
our cache size limit must be >= 32 to capture all these recompilations. Now,
|
||||
suppose there is a separate function in the same program which is very dynamic
|
||||
and unsuitable for compilation. Such a function will need to undergo 32
|
||||
compilations to burst the cache and fallback to eager. These 32 recompilations
|
||||
are too many and we want to fallback for these compilation-unfriendly functions
|
||||
sooner.
|
||||
|
||||
In the new scenario, we can have (1) recompile_limit = 2, (2)
|
||||
accumulated_recompile_limit = 32. This means that each ID_MATCH'd object can
|
||||
have maximum of two cache entries, and the maximum number of cache entries
|
||||
(irrespective of ID_MATCH obj) is 32. This covers the case of forward code
|
||||
object which has 32 recompilations. For the other function, the one unsuitable
|
||||
for recompilation, our limit is 2. So, we will burst the cache in just 2
|
||||
recompilations. In this manner, these 2 limits help us resolve the tension
|
||||
mentioned earlier.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheSizeRelevantForFrame:
|
||||
"""
|
||||
We track the number of cache entries that have same id_match objects as the
|
||||
given frame.
|
||||
|
||||
TODO(janimesh) - Consider adding a map from tuple_of_match_ids to count -
|
||||
https://github.com/pytorch/pytorch/pull/107496#discussion_r1304564682 - this
|
||||
could be useful for debugging as well.
|
||||
"""
|
||||
|
||||
# Total number of CacheEntry objects in the Dynamo linked list
|
||||
num_cache_entries: int = 0
|
||||
|
||||
# Number of CacheEntry objects having same ID_MATCH'd objects as given frame.
|
||||
num_cache_entries_with_same_id_matched_objs: int = 0
|
||||
|
||||
def will_compilation_exceed(self, limit: int) -> bool:
|
||||
# Checks if a compilation will exceed the given limit (that's why >=).
|
||||
return (
|
||||
self.will_compilation_exceed_accumulated_limit()
|
||||
or self.will_compilation_exceed_specific_limit(limit)
|
||||
)
|
||||
|
||||
def will_compilation_exceed_accumulated_limit(self) -> bool:
|
||||
return self.num_cache_entries >= config.accumulated_recompile_limit
|
||||
|
||||
def will_compilation_exceed_specific_limit(self, limit: int) -> bool:
|
||||
return self.num_cache_entries_with_same_id_matched_objs >= limit
|
||||
|
||||
|
||||
def _get_weakref_from_f_locals(
|
||||
frame: DynamoFrameType, local_name: str
|
||||
) -> weakref.ref[Any] | None:
|
||||
obj = frame.f_locals.get(local_name, None)
|
||||
weak_id = None
|
||||
try:
|
||||
weak_id = weakref.ref(obj)
|
||||
except TypeError:
|
||||
pass # cannot weakref bool object
|
||||
return weak_id
|
||||
|
||||
|
||||
def _has_same_id_matched_objs(frame: DynamoFrameType, cache_entry: Any) -> bool:
|
||||
"""
|
||||
Checks if the ID_MATCH'd objects saved on cache_entry are same as the ones
|
||||
in frame.f_locals.
|
||||
"""
|
||||
if not cache_entry:
|
||||
return False
|
||||
|
||||
for (
|
||||
local_name,
|
||||
weakref_from_cache_entry,
|
||||
) in cache_entry.guard_manager.id_matched_objs.items():
|
||||
if weakref_from_cache_entry() is not None:
|
||||
weakref_from_frame = _get_weakref_from_f_locals(frame, local_name)
|
||||
if weakref_from_frame is not weakref_from_cache_entry:
|
||||
return False
|
||||
|
||||
# Also covers the case where no ID_MATCH objects are saved in frame.f_locals
|
||||
return True
|
||||
|
||||
|
||||
def compute_cache_size(
|
||||
frame: DynamoFrameType, cache_entry: Any
|
||||
) -> CacheSizeRelevantForFrame:
|
||||
# Walk the linked list to calculate the cache size
|
||||
num_cache_entries = 0
|
||||
num_cache_entries_with_same_id_matched_objs = 0
|
||||
|
||||
while cache_entry:
|
||||
num_cache_entries += 1
|
||||
# Track the number of cache entries having same ID_MATCH'd objects as
|
||||
# that of frame.f_locals. This will be used later to compare against the
|
||||
# recompile_limit.
|
||||
if _has_same_id_matched_objs(frame, cache_entry):
|
||||
num_cache_entries_with_same_id_matched_objs += 1
|
||||
cache_entry = cache_entry.next
|
||||
|
||||
return CacheSizeRelevantForFrame(
|
||||
num_cache_entries, num_cache_entries_with_same_id_matched_objs
|
||||
)
|
||||
|
||||
|
||||
def is_recompilation(cache_size: CacheSizeRelevantForFrame) -> bool:
|
||||
"""
|
||||
If the frame (earlier parsed by compute_cache_size) has more than 1 cache
|
||||
entry with same ID_MATCH'd objects, then its a recompilation.
|
||||
"""
|
||||
# Note that you can have multiple entries in the cache but still not a
|
||||
# recompile, e.g., you can have 64 nn module instances, each one having an
|
||||
# ID_MATCH guard, and each one having just 1 cache entry in the cache. In
|
||||
# this case, we can have 64 entries in the cache, but no recompilation
|
||||
# because there is only one entry for each id_matched_obj.
|
||||
return cache_size.will_compilation_exceed(1)
|
||||
|
||||
|
||||
def exceeds_recompile_limit(
|
||||
cache_size: CacheSizeRelevantForFrame, compile_id: CompileId
|
||||
) -> tuple[bool, str]:
|
||||
"""
|
||||
Checks if we are exceeding the cache size limit.
|
||||
"""
|
||||
if cache_size.will_compilation_exceed_accumulated_limit():
|
||||
return True, "accumulated_recompile_limit"
|
||||
if cache_size.will_compilation_exceed_specific_limit(config.recompile_limit):
|
||||
return True, "recompile_limit"
|
||||
# NOTE this check is needed in the case that the frame's cache doesn't grow
|
||||
# and we keep recompiling. This can happen if the guard guard_manager becomes invalidated,
|
||||
# e.g. due to guarded objects being freed. This technically makes the
|
||||
# will_compilation_exceed_accumulated_limit check unnecessary, but we will keep the
|
||||
# check in case we have a better fix in the future.
|
||||
assert compile_id.frame_compile_id is not None
|
||||
if compile_id.frame_compile_id >= config.accumulated_recompile_limit:
|
||||
return True, "accumulated_recompile_limit"
|
||||
return False, ""
|
||||
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
This module provides callback management functionality for TorchDynamo's compilation process.
|
||||
|
||||
It implements a thread-safe system for registering, managing and executing callbacks that run
|
||||
at the start and end of TorchDynamo compilations. Key features include:
|
||||
|
||||
- Registration and deregistration of compilation callbacks
|
||||
- Thread-safe callback handling with proper locking mechanisms
|
||||
- Prevention of duplicate callback execution when configured
|
||||
- Decorator utilities for easy callback registration
|
||||
- Context manager for controlled callback lifecycle
|
||||
|
||||
The module centers around the CompilationCallbackHandler class which maintains separate
|
||||
lists for start and end callbacks, manages their execution order, and ensures thread-safety.
|
||||
Utility decorators @on_compile_start and @on_compile_end provide a convenient way to
|
||||
register compilation hooks.
|
||||
|
||||
Example usage:
|
||||
@on_compile_start
|
||||
def my_start_callback():
|
||||
print("Starting compilation")
|
||||
|
||||
@on_compile_end
|
||||
def my_end_callback():
|
||||
print("Compilation complete")
|
||||
"""
|
||||
|
||||
import enum
|
||||
import threading
|
||||
from collections.abc import Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field # noqa: F811
|
||||
from typing import Any
|
||||
|
||||
|
||||
class CallbackTrigger(enum.Enum):
|
||||
# most common case, dynamo attempts to trace a new frame
|
||||
DYNAMO = 1
|
||||
# backward compilation can be deferred to runtime
|
||||
LAZY_BACKWARD = 2
|
||||
# some backends autotune at runtime
|
||||
TRITON_AUTOTUNING = 3 # Temporarily disabled due to spam
|
||||
# cudagraphs record at runtime
|
||||
CUDAGRAPH_RECORDING = 4
|
||||
|
||||
|
||||
@dataclass
|
||||
class CallbackArgs:
|
||||
callback_trigger: CallbackTrigger
|
||||
compile_id: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompilationCallbackHandler:
|
||||
start_callbacks: list[Callable[[CallbackArgs], None]] = field(default_factory=list)
|
||||
end_callbacks: list[Callable[[CallbackArgs], None]] = field(default_factory=list)
|
||||
|
||||
__pending_callbacks_counter: int = field(default=0, init=False, repr=False)
|
||||
__pending_callbacks_counter_lock: threading.Lock = field(
|
||||
default_factory=threading.Lock, init=False, repr=False
|
||||
)
|
||||
|
||||
def register_start_callback(
|
||||
self, callback: Callable[[CallbackArgs], None]
|
||||
) -> Callable[[CallbackArgs], None]:
|
||||
"""
|
||||
Register a callback function to be called when the compilation starts.
|
||||
|
||||
Args:
|
||||
- callback (Callable): The callback function to register.
|
||||
"""
|
||||
self.start_callbacks.append(callback)
|
||||
return callback
|
||||
|
||||
def register_end_callback(
|
||||
self, callback: Callable[[CallbackArgs], None]
|
||||
) -> Callable[[CallbackArgs], None]:
|
||||
"""
|
||||
Register a callback function to be called when the compilation ends.
|
||||
|
||||
Args:
|
||||
- callback (Callable): The callback function to register.
|
||||
"""
|
||||
self.end_callbacks.append(callback)
|
||||
return callback
|
||||
|
||||
def remove_start_callback(self, callback: Callable[[CallbackArgs], None]) -> None:
|
||||
"""
|
||||
Remove a registered start callback function.
|
||||
|
||||
Args:
|
||||
- callback (Callable): The callback function to remove.
|
||||
"""
|
||||
self.start_callbacks.remove(callback)
|
||||
|
||||
def remove_end_callback(self, callback: Callable[[CallbackArgs], None]) -> None:
|
||||
"""
|
||||
Remove a registered end callback function.
|
||||
|
||||
Args:
|
||||
- callback (Callable): The callback function to remove.
|
||||
"""
|
||||
self.end_callbacks.remove(callback)
|
||||
|
||||
def run_start_callbacks(self, args: CallbackArgs) -> None:
|
||||
"""
|
||||
Execute all registered start callbacks.
|
||||
"""
|
||||
for callback in self.start_callbacks:
|
||||
callback(args)
|
||||
|
||||
def run_end_callbacks(self, args: CallbackArgs) -> None:
|
||||
"""
|
||||
Execute all registered end callbacks.
|
||||
"""
|
||||
for callback in self.end_callbacks:
|
||||
callback(args)
|
||||
|
||||
@contextmanager
|
||||
def install_callbacks(
|
||||
self, trigger: CallbackTrigger, compile_id: str
|
||||
) -> Generator[None, Any, Any]:
|
||||
"""
|
||||
Context manager to install the callbacks and run them when the context is exited.
|
||||
"""
|
||||
args = CallbackArgs(trigger, compile_id)
|
||||
try:
|
||||
with self.__pending_callbacks_counter_lock:
|
||||
self.__pending_callbacks_counter += 1
|
||||
if self.__pending_callbacks_counter == 1:
|
||||
self.run_start_callbacks(args)
|
||||
yield
|
||||
finally:
|
||||
with self.__pending_callbacks_counter_lock:
|
||||
assert self.__pending_callbacks_counter > 0, (
|
||||
"Pending callbacks counter cannot become negative."
|
||||
)
|
||||
if self.__pending_callbacks_counter == 1:
|
||||
self.run_end_callbacks(args)
|
||||
self.__pending_callbacks_counter -= 1
|
||||
|
||||
def clear(self) -> None:
|
||||
"""
|
||||
Clear all registered callbacks.
|
||||
"""
|
||||
self.start_callbacks.clear()
|
||||
self.end_callbacks.clear()
|
||||
assert self.__pending_callbacks_counter == 0
|
||||
|
||||
|
||||
callback_handler = CompilationCallbackHandler()
|
||||
|
||||
|
||||
def on_compile_start(
|
||||
callback: Callable[[CallbackArgs], None],
|
||||
) -> Callable[[CallbackArgs], None]:
|
||||
"""
|
||||
Decorator to register a callback function for the start of the compilation.
|
||||
"""
|
||||
callback_handler.register_start_callback(callback)
|
||||
return callback
|
||||
|
||||
|
||||
def on_compile_end(
|
||||
callback: Callable[[CallbackArgs], None],
|
||||
) -> Callable[[CallbackArgs], None]:
|
||||
"""
|
||||
Decorator to register a callback function for the end of the compilation.
|
||||
"""
|
||||
callback_handler.register_end_callback(callback)
|
||||
return callback
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
This module provides thread-safe code context management for TorchDynamo using weak references.
|
||||
|
||||
The CodeContextDict class maintains a mapping between Python code objects and their associated
|
||||
context data, using weak references to automatically clean up entries when code objects are
|
||||
garbage collected. This prevents memory leaks while allowing context data to be associated
|
||||
with code objects throughout their lifecycle.
|
||||
|
||||
Key features:
|
||||
- Thread-safe context storage and retrieval
|
||||
- Automatic cleanup using weak references
|
||||
- Safe context management for Python code objects
|
||||
- Memory-leak prevention
|
||||
|
||||
Example usage:
|
||||
code_obj = compile('x = 1', '<string>', 'exec')
|
||||
|
||||
# Store context
|
||||
context = code_context.get_context(code_obj)
|
||||
context['metadata'] = {'optimized': True}
|
||||
|
||||
# Retrieve context
|
||||
if code_context.has_context(code_obj):
|
||||
ctx = code_context.get_context(code_obj)
|
||||
# Use context data...
|
||||
|
||||
# Remove context
|
||||
ctx = code_context.pop_context(code_obj)
|
||||
"""
|
||||
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
from .utils import ExactWeakKeyDictionary
|
||||
|
||||
|
||||
class CodeContextDict:
|
||||
def __init__(self) -> None:
|
||||
self.code_context: ExactWeakKeyDictionary = ExactWeakKeyDictionary()
|
||||
|
||||
def has_context(self, code: types.CodeType) -> bool:
|
||||
return code in self.code_context
|
||||
|
||||
def get_context(self, code: types.CodeType) -> dict[str, Any]:
|
||||
ctx = self.code_context.get(code)
|
||||
if ctx is None:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
ctx = {}
|
||||
self.code_context[code] = ctx
|
||||
return ctx
|
||||
|
||||
def pop_context(self, code: types.CodeType) -> dict[str, Any]:
|
||||
ctx = self.get_context(code)
|
||||
self.code_context._remove_id(id(code))
|
||||
return ctx
|
||||
|
||||
def clear(self) -> None:
|
||||
self.code_context.clear()
|
||||
|
||||
|
||||
code_context: CodeContextDict = CodeContextDict()
|
||||
@@ -0,0 +1,734 @@
|
||||
"""
|
||||
This module provides utilities for generating Python bytecode in PyTorch's Dynamo system.
|
||||
It includes functionality for:
|
||||
- Constructing bytecode sequences for Python operations
|
||||
- Managing stack operations and variable tracking
|
||||
- Handling graph outputs and their conversions
|
||||
- Supporting different Python versions (3.11+, 3.12+, 3.13+)
|
||||
- Converting high-level operations to low-level bytecode instructions
|
||||
- Managing constant loading and attribute access
|
||||
- Supporting function creation and closure handling
|
||||
"""
|
||||
|
||||
import collections
|
||||
import dataclasses
|
||||
import re
|
||||
import sys
|
||||
import types
|
||||
from collections import Counter, deque
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import Any, TYPE_CHECKING, Union
|
||||
|
||||
import torch.nn
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from . import config, graph_break_hints, utils
|
||||
from .bytecode_transformation import (
|
||||
add_push_null,
|
||||
add_push_null_call_function_ex,
|
||||
create_binary_subscr,
|
||||
create_build_tuple,
|
||||
create_call_function,
|
||||
create_call_function_ex,
|
||||
create_call_method,
|
||||
create_dup_top,
|
||||
create_instruction,
|
||||
create_load_const,
|
||||
create_load_method,
|
||||
create_rot_n,
|
||||
Instruction,
|
||||
)
|
||||
from .exc import unimplemented
|
||||
from .source import AttrSource, ChainedSource, DictGetItemSource, Source
|
||||
from .utils import is_safe_constant, rot_n_helper
|
||||
from .variables.base import ValueMutationExisting, VariableTracker
|
||||
from .variables.functions import (
|
||||
ContextlibContextManagerLocalGeneratorObjectVariable,
|
||||
LocalGeneratorObjectVariable,
|
||||
)
|
||||
from .variables.nn_module import NNModuleVariable
|
||||
from .variables.script_object import TorchScriptObjectVariable
|
||||
from .variables.tensor import (
|
||||
NumpyNdarrayVariable,
|
||||
SymNodeVariable,
|
||||
TensorVariable,
|
||||
UnspecializedPythonVariable,
|
||||
)
|
||||
from .variables.torch_function import TensorWithTFOverrideVariable
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._dynamo.variables.builder import GraphArg
|
||||
|
||||
from .symbolic_convert import InstructionTranslatorBase
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class GraphOutputEntry:
|
||||
index: int
|
||||
variable: VariableTracker
|
||||
|
||||
|
||||
class PyCodegen:
|
||||
"""
|
||||
Helper class uses for constructing Python bytecode
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tx: "InstructionTranslatorBase",
|
||||
root: torch.nn.Module | None = None,
|
||||
graph_output_var: str | None = None,
|
||||
tempvars: dict[VariableTracker | Source, Any] | None = None,
|
||||
overridden_sources: dict[Source, Source] | None = None,
|
||||
) -> None:
|
||||
self.root = root
|
||||
self.top_of_stack: VariableTracker | Source | None = None
|
||||
self.uses: Counter[VariableTracker | Source] = collections.Counter()
|
||||
self.graph_outputs: dict[int, GraphOutputEntry] = {}
|
||||
self._output: list[Instruction] = []
|
||||
# This determines which VariableTracker/Source should be stored as
|
||||
# locals, and maps the VariableTracker/Source to the local variable
|
||||
# name. Note that it could map to None initially, in which case we'll
|
||||
# overwrite it to map to real temporary names via `add_cache`.
|
||||
self.tempvars: dict[VariableTracker | Source, Any] = tempvars or {}
|
||||
self.tx = tx
|
||||
self.graph_output_var = graph_output_var
|
||||
self.code_options = self.tx.output.code_options
|
||||
self.cell_and_freevars = self.tx.cell_and_freevars
|
||||
self.new_var = self.tx.output.new_var
|
||||
self.value_from_source: bool = True
|
||||
# This serves as a way for codegen to use a different source; we need
|
||||
# this because sometimes we can't easily modify the original source
|
||||
# without affecting other components, e.g., guards.
|
||||
self.overridden_sources: dict[Source, Source] = overridden_sources or {}
|
||||
|
||||
def restore_stack(
|
||||
self, stack_values: list[Any], *, value_from_source: bool = True
|
||||
) -> None:
|
||||
prev = self.value_from_source
|
||||
self.value_from_source &= value_from_source
|
||||
try:
|
||||
self.foreach(stack_values)
|
||||
finally:
|
||||
self.value_from_source = prev
|
||||
|
||||
def graph_output_vars(self) -> list[VariableTracker]:
|
||||
return [x.variable for x in self.graph_outputs.values()]
|
||||
|
||||
def call_reconstruct(
|
||||
self, value: Union[VariableTracker, Source, "GraphArg"]
|
||||
) -> None:
|
||||
res = value.reconstruct(self)
|
||||
assert res is None, f"reconstruct!=None {value}"
|
||||
|
||||
def add_push_null(
|
||||
self, gen_fn: Callable[[], None], call_function_ex: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
`gen_fn` generates instructions via PyCodegen methods
|
||||
that push a single callable to the stack.
|
||||
|
||||
`add_push_null` pushes a NULL to the stack before or after the
|
||||
instructions generated by `gen_fn`, depending on Python version.
|
||||
|
||||
Will attempt to use the NULL push bit for instructions
|
||||
with such bits (LOAD_GLOBAL 3.11+, LOAD_ATTR 3.12+, LOAD_SUPER_ATTR).
|
||||
"""
|
||||
old_len = len(self._output)
|
||||
if sys.version_info < (3, 13):
|
||||
# gen_fn may DUP_TOP instead if TOS is not cleared.
|
||||
# Will cause problems since NULL will be pushed right
|
||||
# before the generated instructions in <= 3.12
|
||||
self.clear_tos()
|
||||
gen_fn()
|
||||
# inplace modify self._output
|
||||
added_insts = self._output[old_len:]
|
||||
del self._output[old_len:]
|
||||
if call_function_ex:
|
||||
self._output.extend(add_push_null_call_function_ex(added_insts))
|
||||
else:
|
||||
self._output.extend(add_push_null(added_insts))
|
||||
if sys.version_info >= (3, 13):
|
||||
# NULL will be at top of stack
|
||||
self.clear_tos()
|
||||
|
||||
def __call__(
|
||||
self, value: VariableTracker | Source | None, allow_cache: bool = True
|
||||
) -> None:
|
||||
"""
|
||||
Generate code such that top-of-stack (TOS) is set to value.
|
||||
|
||||
`allow_cache` controls the behavior in the following manner. `value` can
|
||||
either be a VariableTracker or a Source.
|
||||
|
||||
If `value` is a `Source`, `allow_cache` must be True (invariant asserted
|
||||
below). If the source was reconstructed earlier, we will reuse the
|
||||
generated code by loading from top of stack or tempvars.
|
||||
|
||||
If `value` is a `VariableTracker`, we have the following cases:
|
||||
|
||||
1) `allow_cache=True`
|
||||
a) If the value.source is not None, we will emit the code based on
|
||||
`value.source` to handle aliasing.
|
||||
b) If value.source is None (example reconstructing a local list
|
||||
returned by the compiled function), we will reconstruct the variable
|
||||
tracker (w/o any source) to emit bytecode that generates a new
|
||||
python object.
|
||||
|
||||
In both cases of value.source being None or not, if the value was
|
||||
reconstructed earlier, we will reuse the generated code by loading from
|
||||
top of stack or tempvars.
|
||||
|
||||
2) `allow_cache=False` - This is a special case (allow_cache defaults to
|
||||
True).
|
||||
a) If the value.source is not None, we reconstruct the variable
|
||||
tracker and emit a new python object. You might wonder what about
|
||||
aliasing? The place where we use this config also has the followup
|
||||
code where the original python object is assigned to this new python
|
||||
value to handle aliasing (check side_effects.py and search for
|
||||
allow_cache=False).
|
||||
|
||||
b) If value.source is None, this is not allowed
|
||||
|
||||
Notable effects:
|
||||
1. `self.top_of_stack` will be set to `value`, if we don't codegen
|
||||
`value` based on source.
|
||||
2. `self.uses[value]` will increment, unless (a). we codegen via
|
||||
`top_of_stack` or cached `tempvars`, or (b). `value` has special VT
|
||||
types like `NNModuleVariable`, etc.
|
||||
"""
|
||||
assert value is not None
|
||||
if isinstance(value, Source):
|
||||
# If the source needs to be overridden, use the new one.
|
||||
source = self.overridden_sources.get(value, value)
|
||||
assert allow_cache is True, "allow_cache must be True for Source"
|
||||
if self.top_of_stack is value:
|
||||
self._output.append(create_dup_top())
|
||||
return
|
||||
|
||||
if self.tempvars.get(source) is not None:
|
||||
self._output.append(self.create_load(self.tempvars[source]))
|
||||
self.top_of_stack = source
|
||||
return
|
||||
|
||||
self.uses[source] += 1
|
||||
try:
|
||||
self.call_reconstruct(source)
|
||||
except NotImplementedError:
|
||||
unimplemented(
|
||||
gb_type="Reconstruction failure: source.reconstruct not implemented",
|
||||
context=str(source),
|
||||
explanation=f"Dynamo has no bytecode reconstruction implemented for {type(source)} variable {source}.",
|
||||
hints=[*graph_break_hints.DYNAMO_BUG],
|
||||
)
|
||||
if source in self.tempvars:
|
||||
self._output.append(create_dup_top())
|
||||
self.add_cache(source)
|
||||
self.top_of_stack = source
|
||||
|
||||
return
|
||||
|
||||
assert isinstance(value, VariableTracker)
|
||||
output = self._output
|
||||
graph_outputs = self.graph_outputs
|
||||
|
||||
if allow_cache:
|
||||
if self.top_of_stack is value:
|
||||
output.append(create_dup_top())
|
||||
return
|
||||
|
||||
if self.tempvars.get(value) is not None:
|
||||
output.append(self.create_load(self.tempvars[value]))
|
||||
self.top_of_stack = value
|
||||
return
|
||||
|
||||
if value.is_realized() and isinstance(
|
||||
value, ContextlibContextManagerLocalGeneratorObjectVariable
|
||||
):
|
||||
unimplemented(
|
||||
gb_type="reconstructing @contextmanager object",
|
||||
context=f"object: {value}",
|
||||
explanation="Returning a @contextmanager object from a compiled function is not supported.",
|
||||
hints=[
|
||||
*graph_break_hints.SUPPORTABLE,
|
||||
],
|
||||
)
|
||||
|
||||
# Dynamo normally prefers codegen from source to account for aliasing.
|
||||
if (
|
||||
value.source is not None
|
||||
and allow_cache
|
||||
and not (
|
||||
value.is_realized() and isinstance(value, LocalGeneratorObjectVariable)
|
||||
)
|
||||
):
|
||||
# There's a corner case for export: for instance, if the computation
|
||||
# graph is just identity on an input tensor, Dynamo would just emit
|
||||
# a `LOAD_FAST` from the input source, rather than generating an
|
||||
# identity FX graph.
|
||||
#
|
||||
# However, export wants to maximize graph capture; in the case
|
||||
# above, export _wants to_ obtain an identity FX graph (despite it
|
||||
# appears unnecessarily expensive for `torch.compile`), so we have
|
||||
# the following option to override Dynamo's preference for codegen
|
||||
# from source. Moreover, this option applies recursively, for cases
|
||||
# like input tensor being returned in a new dictionary.
|
||||
#
|
||||
# And why the `ValueMutationExisting` check? Not sure, so leaving it
|
||||
# to keep the old behavior, as when `value_from_source` was
|
||||
# introduced. TODO sort out the invariants among side effect,
|
||||
# codegen and export.
|
||||
if (
|
||||
isinstance(value.mutation_type, ValueMutationExisting)
|
||||
or self.value_from_source
|
||||
):
|
||||
return self(value.source)
|
||||
|
||||
if value.is_python_constant() and is_safe_constant(value.as_python_constant()):
|
||||
output.append(self.create_load_const(value.as_python_constant()))
|
||||
elif isinstance(value, TensorWithTFOverrideVariable):
|
||||
graph_outputs_key = self.add_graph_output(value)
|
||||
|
||||
self.add_push_null(
|
||||
lambda: self.load_import_from(utils.__name__, "to_subclass")
|
||||
)
|
||||
self.load_graph_output(graph_outputs[graph_outputs_key].index)
|
||||
output.append(
|
||||
self.create_load_global(
|
||||
value.global_mangled_class_name(self.tx), # type: ignore[arg-type]
|
||||
add=True,
|
||||
)
|
||||
)
|
||||
output.extend(create_call_function(2, False))
|
||||
elif (
|
||||
isinstance(value, SymNodeVariable)
|
||||
and value.python_type() is float
|
||||
and not self.tx.export
|
||||
):
|
||||
# This is a little unusual; force the output convention to be a
|
||||
# Tensor here. Don't do this for export because this is
|
||||
# apparently load bearing for export tests (but I am a bit
|
||||
# doubtful it actually works in the real world)
|
||||
# NB: It works to add_graph_output on a computed expression
|
||||
# as_tensor here, because we memoize as_tensor calls on
|
||||
# SymNodeVariable!
|
||||
graph_outputs_key = self.add_graph_output(
|
||||
value.as_tensor(self.tx, torch.float64)
|
||||
)
|
||||
|
||||
def gen_fn() -> None:
|
||||
self.load_graph_output(graph_outputs[graph_outputs_key].index)
|
||||
output.append(self.create_load_attr("item"))
|
||||
|
||||
self.add_push_null(gen_fn)
|
||||
output.extend(create_call_function(0, False))
|
||||
elif isinstance(
|
||||
value,
|
||||
(
|
||||
TensorVariable,
|
||||
SymNodeVariable,
|
||||
UnspecializedPythonVariable,
|
||||
NumpyNdarrayVariable,
|
||||
TorchScriptObjectVariable,
|
||||
),
|
||||
):
|
||||
graph_outputs_key = self.add_graph_output(value)
|
||||
|
||||
if isinstance(value, NumpyNdarrayVariable):
|
||||
self.add_push_null(
|
||||
lambda: self.load_import_from(utils.__name__, "to_numpy_helper")
|
||||
)
|
||||
self.load_graph_output(graph_outputs[graph_outputs_key].index)
|
||||
output.extend(create_call_function(1, False))
|
||||
elif isinstance(value, UnspecializedPythonVariable) and value.need_unwrap:
|
||||
|
||||
def gen_fn() -> None:
|
||||
self.load_graph_output(graph_outputs[graph_outputs_key].index)
|
||||
output.append(self.create_load_attr("item"))
|
||||
|
||||
self.add_push_null(gen_fn)
|
||||
output.extend(create_call_function(0, False))
|
||||
else:
|
||||
self.load_graph_output(graph_outputs[graph_outputs_key].index)
|
||||
elif isinstance(value, NNModuleVariable):
|
||||
parts = value.module_key.split(".")
|
||||
if parts[0] in self.code_options["co_varnames"]:
|
||||
output.append(self.create_load(parts[0]))
|
||||
parts = parts[1:]
|
||||
else:
|
||||
assert self.root is not None
|
||||
output.append(self.create_load_const_unchecked(self.root))
|
||||
for part in parts:
|
||||
output.append(self.create_load_attr(part))
|
||||
else:
|
||||
self.uses[value] += 1
|
||||
try:
|
||||
self.call_reconstruct(value)
|
||||
except NotImplementedError as e:
|
||||
unimplemented(
|
||||
gb_type="Reconstruction failure",
|
||||
context=str(value),
|
||||
explanation=f"Dynamo has no bytecode reconstruction implemented for sourceless variable {value}.",
|
||||
hints=[
|
||||
"If Dynamo is attempting to trace a return statement and your code is attempting to return a variable "
|
||||
"that Dynamo cannot reconstruct, then remove it from the return statement.",
|
||||
*graph_break_hints.CAUSED_BY_EARLIER_GRAPH_BREAK,
|
||||
"Report an issue to PyTorch if you need reconstrtuction support. Note that objects that don't have "
|
||||
"reconstruction rules may be fundamentally unreconstructable.",
|
||||
],
|
||||
from_exc=e,
|
||||
)
|
||||
if allow_cache and value in self.tempvars:
|
||||
self._output.append(create_dup_top())
|
||||
self.add_cache(value)
|
||||
|
||||
self.top_of_stack = value
|
||||
|
||||
def add_graph_output(self, value: VariableTracker) -> int:
|
||||
graph_outputs_key = id(value.as_proxy())
|
||||
if graph_outputs_key not in self.graph_outputs:
|
||||
self.graph_outputs[graph_outputs_key] = GraphOutputEntry(
|
||||
len(self.graph_outputs), value
|
||||
)
|
||||
return graph_outputs_key
|
||||
|
||||
def load_graph_output(self, index: int) -> None:
|
||||
output = self._output
|
||||
assert self.graph_output_var is not None
|
||||
output.append(self.create_load(self.graph_output_var))
|
||||
output.append(self.create_load_const(index))
|
||||
output.append(self.create_binary_subscr())
|
||||
|
||||
def add_cache(self, value: VariableTracker | Source) -> None:
|
||||
var = self.new_var()
|
||||
self.tempvars[value] = var
|
||||
self._output.append(self.create_store(var))
|
||||
|
||||
def foreach(self, items: Iterable[VariableTracker | Source]) -> None:
|
||||
for i in items:
|
||||
self(i)
|
||||
|
||||
def create_binary_subscr(self) -> Instruction:
|
||||
return create_binary_subscr()
|
||||
|
||||
def setup_globally_cached(self, name: str, value: Any) -> list[Instruction]:
|
||||
"""Store value in a new global"""
|
||||
name = re.sub(r"[^a-zA-Z0-9_]+", "_", name)
|
||||
f_globals = self.tx.f_globals
|
||||
if name in f_globals:
|
||||
assert id(f_globals[name]) == id(value)
|
||||
else:
|
||||
f_globals[name] = value
|
||||
return [self.create_load_global(name, add=True)]
|
||||
|
||||
def clear_tos(self) -> None:
|
||||
self.top_of_stack = None
|
||||
|
||||
def append_output(self, inst: Instruction) -> None:
|
||||
assert isinstance(inst, Instruction)
|
||||
self._output.append(inst)
|
||||
self.clear_tos()
|
||||
|
||||
def extend_output(self, insts: list[Instruction]) -> None:
|
||||
assert all(isinstance(x, Instruction) for x in insts)
|
||||
self._output.extend(insts)
|
||||
self.clear_tos()
|
||||
|
||||
def get_instructions(self) -> list[Instruction]:
|
||||
return self._output
|
||||
|
||||
def create_load(self, name: str) -> Instruction:
|
||||
assert name in self.code_options["co_varnames"], f"{name} missing"
|
||||
return create_instruction("LOAD_FAST", argval=name)
|
||||
|
||||
def create_load_closure(self, name: str) -> Instruction:
|
||||
assert name in self.cell_and_freevars()
|
||||
inst_name = "LOAD_FAST" if sys.version_info >= (3, 13) else "LOAD_CLOSURE"
|
||||
return create_instruction(inst_name, argval=name)
|
||||
|
||||
def create_load_deref(self, name: str) -> Instruction:
|
||||
assert name in self.cell_and_freevars()
|
||||
return create_instruction("LOAD_DEREF", argval=name)
|
||||
|
||||
def create_store(self, name: str) -> Instruction:
|
||||
assert name in self.code_options["co_varnames"], f"{name} missing"
|
||||
return create_instruction("STORE_FAST", argval=name)
|
||||
|
||||
def create_store_deref(self, name: str) -> Instruction:
|
||||
assert name in self.cell_and_freevars()
|
||||
return create_instruction("STORE_DEREF", argval=name)
|
||||
|
||||
def create_load_global(self, name: str, add: bool = False) -> Instruction:
|
||||
if add:
|
||||
self.tx.output.update_co_names(name)
|
||||
assert name in self.code_options["co_names"], f"{name} not in co_names"
|
||||
return create_instruction("LOAD_GLOBAL", argval=name)
|
||||
|
||||
def create_load_const(self, value: Any) -> Instruction:
|
||||
return create_load_const(value)
|
||||
|
||||
def create_load_const_unchecked(self, value: Any) -> Instruction:
|
||||
return create_load_const(value, checked=False)
|
||||
|
||||
def load_method(self, name: str) -> None:
|
||||
self.tx.output.update_co_names(name)
|
||||
self.append_output(create_load_method(name))
|
||||
|
||||
def call_method(self, nargs: int) -> None:
|
||||
self.extend_output(create_call_method(nargs))
|
||||
|
||||
def create_load_attr(self, name: str) -> Instruction:
|
||||
if name not in self.code_options["co_names"]:
|
||||
self.code_options["co_names"] += (name,)
|
||||
return create_instruction("LOAD_ATTR", argval=name)
|
||||
|
||||
def load_attr(self, name: str) -> None:
|
||||
self.append_output(self.create_load_attr(name))
|
||||
|
||||
def create_load_attrs(self, names: str) -> list[Instruction]:
|
||||
return [self.create_load_attr(name) for name in names.split(".")]
|
||||
|
||||
def create_store_attr(self, name: str) -> Instruction:
|
||||
if name not in self.code_options["co_names"]:
|
||||
self.code_options["co_names"] += (name,)
|
||||
return create_instruction("STORE_ATTR", argval=name)
|
||||
|
||||
def store_attr(self, name: str) -> None:
|
||||
self.append_output(self.create_store_attr(name))
|
||||
|
||||
def load_function_name(
|
||||
self, fn_name: str, push_null: bool, num_on_stack: int = 0
|
||||
) -> list[Instruction]:
|
||||
"""Load the global fn_name on the stack num_on_stack down"""
|
||||
output = []
|
||||
if push_null and sys.version_info >= (3, 11):
|
||||
output.extend(add_push_null(self.create_load_global(fn_name, add=True)))
|
||||
if num_on_stack > 0:
|
||||
output.extend(
|
||||
[
|
||||
*self.rot_n(num_on_stack + 2),
|
||||
*self.rot_n(num_on_stack + 2),
|
||||
]
|
||||
)
|
||||
else:
|
||||
output.extend(
|
||||
[
|
||||
self.create_load_global(fn_name, add=True),
|
||||
*self.rot_n(num_on_stack + 1),
|
||||
]
|
||||
)
|
||||
return output
|
||||
|
||||
def rot_n(self, n: int) -> list[Instruction]:
|
||||
try:
|
||||
return create_rot_n(n)
|
||||
except AttributeError:
|
||||
# desired rotate bytecode doesn't exist, generate equivalent bytecode
|
||||
return [
|
||||
create_build_tuple(n),
|
||||
self.create_load_const_unchecked(rot_n_helper(n)),
|
||||
*create_rot_n(2),
|
||||
*create_call_function_ex(False, False),
|
||||
create_instruction("UNPACK_SEQUENCE", arg=n),
|
||||
]
|
||||
|
||||
def pop_null(self) -> list[Instruction]:
|
||||
# POP_TOP doesn't work for null, so we pop nulls by pushing in a
|
||||
# nop function, calling it (which consumes the null), and popping the result.
|
||||
assert sys.version_info >= (3, 11)
|
||||
return [
|
||||
self.create_load_const_unchecked(lambda: None),
|
||||
# 3.13 swapped NULL and callable
|
||||
*(
|
||||
(create_instruction("SWAP", arg=2),)
|
||||
if sys.version_info >= (3, 13)
|
||||
else ()
|
||||
),
|
||||
*create_call_function(0, False),
|
||||
create_instruction("POP_TOP"),
|
||||
]
|
||||
|
||||
def pop_top(self) -> None:
|
||||
self.append_output(create_instruction("POP_TOP"))
|
||||
|
||||
def call_function(self, nargs: int, push_null: bool) -> None:
|
||||
self.extend_output(create_call_function(nargs, push_null=push_null))
|
||||
|
||||
def dup_top(self) -> None:
|
||||
self.append_output(create_dup_top())
|
||||
|
||||
def store(self, varname: str) -> None:
|
||||
self.append_output(self.create_store(varname))
|
||||
|
||||
def load_deref(self, varname: str) -> None:
|
||||
self.append_output(self.create_load_deref(varname))
|
||||
|
||||
def make_function_with_closure(
|
||||
self,
|
||||
fn_name: str,
|
||||
code: types.CodeType,
|
||||
) -> None:
|
||||
"""Creates a closure with code object `code`.
|
||||
|
||||
Expects the TOS to be the tuple of cells to use for this closure.
|
||||
TOS will be popped to create the closure.
|
||||
Args:
|
||||
- fn_name: name of the function
|
||||
- code: code object of the function
|
||||
(does not include the tuple of cells on the TOS)
|
||||
"""
|
||||
output = self._output
|
||||
|
||||
output.append(self.create_load_const(code))
|
||||
if sys.version_info < (3, 11):
|
||||
output.append(self.create_load_const(fn_name))
|
||||
if sys.version_info >= (3, 13):
|
||||
output.extend(
|
||||
[
|
||||
create_instruction("MAKE_FUNCTION"),
|
||||
create_instruction("SET_FUNCTION_ATTRIBUTE", arg=0x08),
|
||||
]
|
||||
)
|
||||
else:
|
||||
output.append(create_instruction("MAKE_FUNCTION", arg=0x08))
|
||||
|
||||
self.clear_tos()
|
||||
|
||||
def create_load_python_module(self, mod: types.ModuleType) -> Instruction:
|
||||
"""
|
||||
Generate a LOAD_GLOBAL instruction to fetch a given python module.
|
||||
"""
|
||||
output = self.tx.output
|
||||
global_scope = output.global_scope
|
||||
name = re.sub(r"^.*[.]", "", mod.__name__)
|
||||
if global_scope.get(name, None) is mod:
|
||||
return self.create_load_global(name, add=True)
|
||||
prefix = f"___module_{name}"
|
||||
global_name = self.tx.output.install_global_by_id(prefix, mod)
|
||||
return self.create_load_global(global_name, add=True)
|
||||
|
||||
def mark_source_temp(self, source: Source) -> None:
|
||||
"""
|
||||
Mark a source as a temp variable, so that it can be reused.
|
||||
"""
|
||||
if source not in self.tempvars:
|
||||
self.tempvars[source] = None
|
||||
|
||||
def make_call_generated_code(self, fn_name: str) -> None:
|
||||
"""Call the generated code function stored in fn_name"""
|
||||
self.extend_output(self.load_function_name(fn_name, True))
|
||||
|
||||
graphargs = self.tx.output.graphargs
|
||||
|
||||
def extract_nested_sources(source: Source) -> list[Source]:
|
||||
nested_sources: list[Source] = []
|
||||
if isinstance(source, ChainedSource):
|
||||
nested_sources.append(source.base)
|
||||
if isinstance(source, DictGetItemSource) and isinstance(
|
||||
source.index, Source
|
||||
):
|
||||
nested_sources.append(source.index)
|
||||
return nested_sources
|
||||
|
||||
def collect_temp_sources(sources: deque[Source], codegen: PyCodegen) -> None:
|
||||
seen_sources: OrderedSet[Source] = OrderedSet()
|
||||
while sources:
|
||||
current_source = sources.popleft()
|
||||
if current_source in seen_sources:
|
||||
# This source is used at least twice, so it can be reused
|
||||
codegen.mark_source_temp(current_source)
|
||||
# Dont trace source further. This prevents us from marking too
|
||||
# many nodes as temp sources.
|
||||
continue
|
||||
seen_sources.add(current_source)
|
||||
sources.extend(extract_nested_sources(current_source))
|
||||
|
||||
# Collect all the sources that are used more than once, so that we can
|
||||
# generate tmp variables in the generated pre-graph bytecode. This
|
||||
# essentially implements CSE.
|
||||
collect_temp_sources(
|
||||
deque([arg.source for arg in graphargs if arg.source is not None]), self
|
||||
)
|
||||
|
||||
cm_var = None
|
||||
if config.record_runtime_overhead:
|
||||
# Record the pregraph bytecode start
|
||||
self.add_push_null(
|
||||
lambda: self.load_import_from(
|
||||
utils.__name__, "record_pregraph_bytecode_enter"
|
||||
)
|
||||
)
|
||||
self.extend_output(create_call_function(0, False))
|
||||
cm_var = self.new_var()
|
||||
self.store(cm_var)
|
||||
|
||||
for arg in graphargs:
|
||||
if arg.pass_arg_as_tensor:
|
||||
self.add_push_null(
|
||||
lambda: self.extend_output(
|
||||
[
|
||||
self.create_load_python_module(torch),
|
||||
self.create_load_attr("_as_tensor_fullprec"),
|
||||
]
|
||||
)
|
||||
)
|
||||
self.call_reconstruct(arg)
|
||||
self.extend_output(create_call_function(1, False))
|
||||
else:
|
||||
self.call_reconstruct(arg)
|
||||
|
||||
if config.record_runtime_overhead:
|
||||
# Record the pregraph bytecode end
|
||||
self.add_push_null(
|
||||
lambda: self.load_import_from(
|
||||
utils.__name__, "record_pregraph_bytecode_exit"
|
||||
)
|
||||
)
|
||||
assert cm_var is not None
|
||||
self.extend_output([self.create_load(cm_var)])
|
||||
self.extend_output(create_call_function(1, False))
|
||||
self.pop_top()
|
||||
|
||||
self.extend_output(create_call_function(len(graphargs), False))
|
||||
|
||||
def create_import_name(self, module_name: str) -> Instruction:
|
||||
return create_instruction("IMPORT_NAME", argval=module_name)
|
||||
|
||||
def load_import_from(self, module_name: str, object_name: str) -> None:
|
||||
source = AttrSource(self.tx.import_source(module_name), object_name)
|
||||
# Note: This approach is somewhat aggressive because typically, a source is marked
|
||||
# as a tempvar only when it is used more than once. In this case, we're marking it
|
||||
# as a tempvar without performing that analysis. However, this is a simple solution,
|
||||
# and in many cases, load imports are reused multiple times.
|
||||
self.mark_source_temp(source)
|
||||
self(source)
|
||||
|
||||
def create_call_function_kw(
|
||||
self, nargs: int, kw_names: Iterable[str], push_null: bool
|
||||
) -> list[Instruction]:
|
||||
if sys.version_info >= (3, 13):
|
||||
output = create_call_function(nargs, push_null)
|
||||
assert output[-1].opname == "CALL"
|
||||
output.insert(-1, self.create_load_const(kw_names))
|
||||
output[-1] = create_instruction("CALL_KW", arg=nargs)
|
||||
return output
|
||||
elif sys.version_info >= (3, 11):
|
||||
output = create_call_function(nargs, push_null)
|
||||
if sys.version_info >= (3, 12):
|
||||
idx = -1
|
||||
expected_inst = "CALL"
|
||||
else:
|
||||
idx = -2
|
||||
expected_inst = "PRECALL"
|
||||
assert output[idx].opname == expected_inst
|
||||
kw_names_inst = create_instruction("KW_NAMES", argval=kw_names)
|
||||
output.insert(idx, kw_names_inst)
|
||||
return output
|
||||
return [
|
||||
self.create_load_const(kw_names),
|
||||
create_instruction("CALL_FUNCTION_KW", arg=nargs),
|
||||
]
|
||||
|
||||
def create_delete(self, value: object) -> Instruction:
|
||||
return create_instruction("DELETE_FAST", argval=value)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,694 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import dataclasses
|
||||
import dis
|
||||
import functools
|
||||
import logging
|
||||
import sys
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import types
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from .symbolic_convert import InstructionTranslatorBase
|
||||
|
||||
from .bytecode_transformation import (
|
||||
create_copy,
|
||||
create_dup_top,
|
||||
create_instruction,
|
||||
create_swap,
|
||||
Instruction,
|
||||
unique_id,
|
||||
)
|
||||
from .codegen import PyCodegen
|
||||
from .exc import unimplemented
|
||||
from .output_graph import GraphCompileReason, StackLocalsMetadata
|
||||
from .variables.misc import NullVariable, UnknownVariable
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _get_comprehension_bytecode_prefix() -> list[str]:
|
||||
"""Get the bytecode instructions that precede BUILD_LIST in a list comprehension."""
|
||||
|
||||
assert sys.version_info >= (3, 12)
|
||||
|
||||
def fn() -> list[int]:
|
||||
return [i for i in range(1)] # noqa: C416
|
||||
|
||||
insts = [inst.opname for inst in dis.get_instructions(fn)]
|
||||
|
||||
start_idx = len(insts) - 1 - insts[::-1].index("LOAD_FAST_AND_CLEAR")
|
||||
end_idx = insts.index("BUILD_LIST")
|
||||
|
||||
return insts[start_idx:end_idx]
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _get_comprehension_result_patterns() -> dict[str, dict[str, Any]]:
|
||||
"""Discover bytecode patterns for comprehension result handling.
|
||||
|
||||
Analyzes sample functions to extract the opcode sequences that appear
|
||||
after END_FOR for each result disposition (stored, discarded, returned, consumed).
|
||||
|
||||
Returns patterns with:
|
||||
- pre_store_ops: opcodes between END_FOR and first STORE_FAST
|
||||
- post_store_op: first opcode after all STORE_FASTs (for disambiguation)
|
||||
"""
|
||||
assert sys.version_info >= (3, 12)
|
||||
|
||||
def fn_stored() -> list[int]:
|
||||
result = [i for i in range(1)] # noqa: C416
|
||||
return result
|
||||
|
||||
def fn_discarded() -> int:
|
||||
[i for i in range(1)] # noqa: C416
|
||||
return 1
|
||||
|
||||
def fn_returned() -> list[int]:
|
||||
return [i for i in range(1)] # noqa: C416
|
||||
|
||||
def fn_consumed() -> int:
|
||||
return sum([i for i in range(1)]) # noqa: C416
|
||||
|
||||
def extract_pattern(fn: Callable[..., Any]) -> tuple[list[str], str | None]:
|
||||
"""Extract (pre_store_ops, post_store_op) from comprehension bytecode."""
|
||||
target_line = list(dis.findlinestarts(fn.__code__))[1][1]
|
||||
insts: list[str] = []
|
||||
started = False
|
||||
for instr in dis.get_instructions(fn):
|
||||
if started and instr.starts_line:
|
||||
break
|
||||
pos = instr.positions
|
||||
if pos and pos.lineno == target_line:
|
||||
started = started or bool(instr.starts_line)
|
||||
insts.append(instr.opname)
|
||||
|
||||
ops = insts[insts.index("END_FOR") + 1 :]
|
||||
idx = 0
|
||||
|
||||
pre_store_ops = []
|
||||
while idx < len(ops) and ops[idx] != "STORE_FAST":
|
||||
pre_store_ops.append(ops[idx])
|
||||
idx += 1
|
||||
|
||||
while idx < len(ops) and ops[idx] == "STORE_FAST":
|
||||
idx += 1
|
||||
|
||||
return pre_store_ops, ops[idx] if idx < len(ops) else None
|
||||
|
||||
stored = extract_pattern(fn_stored)
|
||||
discarded = extract_pattern(fn_discarded)
|
||||
returned = extract_pattern(fn_returned)
|
||||
consumed = extract_pattern(fn_consumed)
|
||||
|
||||
return {
|
||||
"stored": {"pre_store_ops": stored[0], "post_store_op": stored[1]},
|
||||
"discarded": {"pre_store_ops": discarded[0], "post_store_op": discarded[1]},
|
||||
"returned": {"pre_store_ops": returned[0], "post_store_op": returned[1]},
|
||||
"consumed": {"pre_store_ops": consumed[0], "post_store_op": []},
|
||||
}
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ComprehensionAnalysis:
|
||||
"""Metadata about a comprehension's bytecode structure.
|
||||
|
||||
Attributes:
|
||||
end_ip: Instruction pointer after all comprehension bytecode
|
||||
result_var: Name of result variable, or None if result stays on stack
|
||||
result_on_stack: True if result stays on stack (discarded, returned, or in expression)
|
||||
iterator_vars: Variables from LOAD_FAST_AND_CLEAR (need restoration)
|
||||
walrus_vars: Variables assigned via walrus operator (:=) inside comprehension
|
||||
captured_vars: Variables read from outer scope via LOAD_FAST inside comprehension
|
||||
"""
|
||||
|
||||
end_ip: int
|
||||
result_var: str | None
|
||||
result_on_stack: bool
|
||||
iterator_vars: list[str]
|
||||
walrus_vars: list[str]
|
||||
captured_vars: list[str]
|
||||
|
||||
|
||||
def _is_comprehension_start(tx: InstructionTranslatorBase) -> bool:
|
||||
"""Detect if we're at the start of a list/dict comprehension in 3.12+.
|
||||
|
||||
In Python 3.12+, comprehensions are inlined with a bytecode pattern that
|
||||
precedes BUILD_LIST/BUILD_MAP.
|
||||
"""
|
||||
assert sys.version_info >= (3, 12)
|
||||
|
||||
assert tx.instruction_pointer is not None
|
||||
ip = tx.instruction_pointer - 1
|
||||
|
||||
pattern = _get_comprehension_bytecode_prefix()
|
||||
prefix = [inst.opname for inst in tx.instructions[ip - len(pattern) : ip]]
|
||||
|
||||
return prefix == pattern
|
||||
|
||||
|
||||
def _find_comprehension_end_for_ip(tx: InstructionTranslatorBase) -> int:
|
||||
"""Find the instruction pointer of the outermost END_FOR for current comprehension."""
|
||||
assert sys.version_info >= (3, 12)
|
||||
assert tx.instruction_pointer is not None
|
||||
|
||||
nesting_depth = 0
|
||||
for search_ip in range(tx.instruction_pointer, len(tx.instructions)):
|
||||
inst = tx.instructions[search_ip]
|
||||
if inst.opname == "FOR_ITER":
|
||||
nesting_depth += 1
|
||||
elif inst.opname == "END_FOR":
|
||||
nesting_depth -= 1
|
||||
if nesting_depth == 0:
|
||||
return search_ip
|
||||
return -1
|
||||
|
||||
|
||||
def _analyze_comprehension(tx: InstructionTranslatorBase) -> ComprehensionAnalysis:
|
||||
"""Analyze comprehension bytecode to determine result handling pattern."""
|
||||
assert sys.version_info >= (3, 12)
|
||||
assert tx.instruction_pointer is not None
|
||||
|
||||
patterns = _get_comprehension_result_patterns()
|
||||
start_ip = tx.instruction_pointer - 1 # BUILD_LIST/BUILD_MAP
|
||||
|
||||
iterator_vars: list[str] = []
|
||||
walrus_vars: list[str] = []
|
||||
captured_vars: list[str] = []
|
||||
defined_inside: set[str] = set()
|
||||
|
||||
# Collect iterator variables from LOAD_FAST_AND_CLEAR before BUILD_LIST/BUILD_MAP
|
||||
iter_scan_ip = start_ip - 1
|
||||
while iter_scan_ip >= 0:
|
||||
inst = tx.instructions[iter_scan_ip]
|
||||
if inst.opname == "LOAD_FAST_AND_CLEAR":
|
||||
iterator_vars.insert(0, inst.argval)
|
||||
iter_scan_ip -= 1
|
||||
elif inst.opname in ("SWAP", "GET_ITER"):
|
||||
iter_scan_ip -= 1
|
||||
else:
|
||||
break
|
||||
defined_inside.update(iterator_vars)
|
||||
|
||||
end_for_ip = _find_comprehension_end_for_ip(tx)
|
||||
if end_for_ip == -1:
|
||||
unimplemented(
|
||||
gb_type="Comprehension analysis failed: No END_FOR",
|
||||
context="",
|
||||
explanation="Could not find END_FOR instruction in comprehension bytecode.",
|
||||
hints=[],
|
||||
)
|
||||
|
||||
# Find first FOR_ITER to know where loop body starts
|
||||
for_iter_ip = next(
|
||||
i
|
||||
for i in range(start_ip, end_for_ip)
|
||||
if tx.instructions[i].opname == "FOR_ITER"
|
||||
)
|
||||
|
||||
# Single pass through loop body to detect walrus vars and captured vars
|
||||
for body_ip in range(for_iter_ip + 1, end_for_ip):
|
||||
inst = tx.instructions[body_ip]
|
||||
|
||||
# Detect walrus pattern: COPY 1 followed by STORE_FAST
|
||||
if inst.opname == "COPY" and inst.arg == 1 and body_ip + 1 < end_for_ip:
|
||||
next_inst = tx.instructions[body_ip + 1]
|
||||
if next_inst.opname == "STORE_FAST":
|
||||
var_name = next_inst.argval
|
||||
if var_name not in iterator_vars and var_name not in walrus_vars:
|
||||
walrus_vars.append(var_name)
|
||||
defined_inside.add(var_name)
|
||||
|
||||
# Track variables defined inside the loop
|
||||
if inst.opname == "STORE_FAST":
|
||||
defined_inside.add(inst.argval)
|
||||
|
||||
# Detect LOAD_FAST referencing outer variables
|
||||
elif inst.opname.startswith("LOAD_FAST"):
|
||||
var_names = (
|
||||
inst.argval if isinstance(inst.argval, tuple) else (inst.argval,)
|
||||
)
|
||||
for var_name in var_names:
|
||||
if var_name not in defined_inside and var_name not in captured_vars:
|
||||
captured_vars.append(var_name)
|
||||
|
||||
# Extract pre_store_ops: all opcodes from END_FOR+1 until first STORE_FAST
|
||||
pre_store_ops: list[str] = []
|
||||
scan_ip = end_for_ip + 1
|
||||
while (
|
||||
scan_ip < len(tx.instructions)
|
||||
and tx.instructions[scan_ip].opname != "STORE_FAST"
|
||||
):
|
||||
pre_store_ops.append(tx.instructions[scan_ip].opname)
|
||||
scan_ip += 1
|
||||
|
||||
store_fast_ip = scan_ip
|
||||
|
||||
# Skip all STORE_FASTs to find post_store_op
|
||||
while (
|
||||
scan_ip < len(tx.instructions)
|
||||
and tx.instructions[scan_ip].opname == "STORE_FAST"
|
||||
):
|
||||
scan_ip += 1
|
||||
|
||||
post_store_op = (
|
||||
tx.instructions[scan_ip].opname if scan_ip < len(tx.instructions) else None
|
||||
)
|
||||
|
||||
def matches(name: str) -> bool:
|
||||
pat = patterns[name]
|
||||
return pre_store_ops == pat["pre_store_ops"] and (
|
||||
post_store_op == pat["post_store_op"] or not pat["post_store_op"]
|
||||
)
|
||||
|
||||
result_var: str | None = None
|
||||
if matches("stored"):
|
||||
result_var = tx.instructions[store_fast_ip].argval
|
||||
result_on_stack = False
|
||||
elif matches("discarded"):
|
||||
result_var = None
|
||||
result_on_stack = False
|
||||
scan_ip = scan_ip + 1 if patterns["discarded"]["post_store_op"] else scan_ip
|
||||
elif matches("returned") or pre_store_ops == patterns["consumed"]["pre_store_ops"]:
|
||||
result_var = None
|
||||
result_on_stack = True
|
||||
else:
|
||||
unimplemented(
|
||||
gb_type="Comprehension analysis failed: No matches",
|
||||
context=f"pre_store_ops={pre_store_ops}, post_store_op={post_store_op}",
|
||||
explanation="Comprehension does not match any known bytecode pattern.",
|
||||
hints=[],
|
||||
)
|
||||
|
||||
return ComprehensionAnalysis(
|
||||
end_ip=scan_ip,
|
||||
result_var=result_var,
|
||||
# pyrefly: ignore [unbound-name]
|
||||
result_on_stack=result_on_stack,
|
||||
iterator_vars=iterator_vars,
|
||||
walrus_vars=walrus_vars,
|
||||
captured_vars=captured_vars,
|
||||
)
|
||||
|
||||
|
||||
def _handle_comprehension_graph_break(
|
||||
tx: InstructionTranslatorBase, inst: Instruction
|
||||
) -> None:
|
||||
"""Handle list/dict comprehension graph break.
|
||||
|
||||
Builds a synthetic function wrapping the comprehension bytecode,
|
||||
calls it via codegen_call_resume, then chains into the resume
|
||||
function for the post-comprehension code.
|
||||
"""
|
||||
assert sys.version_info >= (3, 12)
|
||||
assert tx.instruction_pointer is not None
|
||||
|
||||
start_ip = tx.instruction_pointer - 1 # BUILD_LIST/BUILD_MAP
|
||||
analysis = _analyze_comprehension(tx)
|
||||
stack_pops = 1 + len(analysis.iterator_vars)
|
||||
reason = GraphCompileReason("comprehension_graph_break", [tx.frame_summary()])
|
||||
log.debug("comprehension triggered compile")
|
||||
|
||||
# --- Step 1: Compile the graph up to the comprehension ---
|
||||
|
||||
all_stack_locals_metadata = tx.output.compile_subgraph(
|
||||
tx,
|
||||
reason=reason,
|
||||
stack_pops=stack_pops,
|
||||
)
|
||||
# Record which stack_pops items are NULL before popn loses the info.
|
||||
# NULLs on the CPython stack can't be passed as function arguments.
|
||||
stack_pops_null_mask = [
|
||||
isinstance(tx.stack[len(tx.stack) - stack_pops + i], NullVariable)
|
||||
for i in range(stack_pops)
|
||||
]
|
||||
|
||||
tx.popn(stack_pops)
|
||||
meta = all_stack_locals_metadata[0]
|
||||
cg = PyCodegen(tx.output.root_tx)
|
||||
|
||||
# Runtime stack after compile_subgraph:
|
||||
# cells, [frame_values], *(non-popped items), *(stack_pops items w/ NULLs)
|
||||
# frame_values[0] = [frame N locals] (no stack items yet)
|
||||
|
||||
nonnull_count = sum(1 for m in stack_pops_null_mask if not m)
|
||||
|
||||
# live_stack_depth: stack items above cells/frame_values excluding NULLs
|
||||
# that compile_subgraph didn't codegen (tracked in stack_null_idxes).
|
||||
live_stack_depth = len(tx.stack) - len(meta.stack_null_idxes)
|
||||
|
||||
# --- Step 2: Pop stack_pops items and append non-nulls to frame_values[0] ---
|
||||
# SWAP each item to TOS then LIST_APPEND or pop_null; fv_list stays at
|
||||
# TOS throughout. Items append in TOS-first (reversed) order;
|
||||
# _build_comprehension_fn compensates by loading in reverse.
|
||||
cg.extend_output(
|
||||
[
|
||||
# frame_values[0] to TOS
|
||||
*create_copy(live_stack_depth + stack_pops + 1),
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
]
|
||||
)
|
||||
for i in reversed(range(stack_pops)):
|
||||
cg.extend_output(create_swap(2))
|
||||
if stack_pops_null_mask[i]:
|
||||
cg.extend_output(cg.pop_null())
|
||||
else:
|
||||
cg.extend_output([create_instruction("LIST_APPEND", arg=1)])
|
||||
cg.extend_output([create_instruction("POP_TOP")])
|
||||
|
||||
# Stack: cells, [frame_values], *(non-popped items)
|
||||
|
||||
# --- Step 3: Build comprehension function ---
|
||||
new_code, fn_name = _build_comprehension_fn(
|
||||
tx,
|
||||
analysis,
|
||||
start_ip,
|
||||
stack_pops,
|
||||
stack_pops_null_mask,
|
||||
nonnull_count,
|
||||
meta,
|
||||
)
|
||||
|
||||
# --- Step 4: Extract [cells[0]] and [frame_values[0]] for codegen_call_resume ---
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_copy(live_stack_depth + 2),
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
create_instruction("BUILD_LIST", arg=1),
|
||||
*create_copy(live_stack_depth + 2),
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
create_instruction("BUILD_LIST", arg=1),
|
||||
]
|
||||
)
|
||||
|
||||
# Stack: ..., *(non-popped), [cells[0]], [frame_values[0]]
|
||||
|
||||
# --- Step 5: Call comprehension function via codegen_call_resume ---
|
||||
tx.codegen_call_resume([new_code], [fn_name], cg)
|
||||
|
||||
# Stack: ..., *(non-popped), comp_result
|
||||
|
||||
# --- Step 6: Remove appended stack_pops items from frame_values[0] ---
|
||||
if nonnull_count > 0:
|
||||
frame_values_pos = live_stack_depth + 1 + 1 # +1 result, +1 frame_values
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_copy(frame_values_pos),
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
# frame_values[0] on TOS
|
||||
create_dup_top(),
|
||||
# frame_values[0], frame_values[0]
|
||||
cg.create_load_const(-nonnull_count),
|
||||
cg.create_load_const(None),
|
||||
create_instruction("BUILD_SLICE", arg=2),
|
||||
create_instruction("DELETE_SUBSCR"),
|
||||
# del frame_values[0][-nonnull_count:]
|
||||
create_instruction("POP_TOP"),
|
||||
]
|
||||
)
|
||||
|
||||
# --- Step 7: Pass comprehension outputs to frame_values[0] ---
|
||||
# Walrus vars first, then result_var.
|
||||
vars_to_pass = analysis.walrus_vars + (
|
||||
[analysis.result_var] if analysis.result_var else []
|
||||
)
|
||||
|
||||
existing_vars: dict[str, int] = {}
|
||||
for var_name in vars_to_pass:
|
||||
tx.symbolic_locals[var_name] = UnknownVariable()
|
||||
if var_name in meta.locals_names:
|
||||
existing_vars[var_name] = meta.locals_names[var_name]
|
||||
else:
|
||||
meta.locals_names[var_name] = len(meta.locals_names)
|
||||
|
||||
fv_depth = live_stack_depth + 2 # comp_result + frame_values
|
||||
|
||||
# --- Walrus vars: extract from comp_result tuple ---
|
||||
if analysis.walrus_vars:
|
||||
# comp_result is (result, *walrus_vars).
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_copy(fv_depth),
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
]
|
||||
)
|
||||
# Stack: ..., comp_tuple, fv0
|
||||
for j, walrus_var in enumerate(analysis.walrus_vars):
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_copy(2),
|
||||
cg.create_load_const(j + 1),
|
||||
cg.create_binary_subscr(),
|
||||
]
|
||||
)
|
||||
# Stack: ..., comp_tuple, fv0, walrus_value
|
||||
if walrus_var in existing_vars:
|
||||
# fv0[idx] = walrus_value
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_copy(2), # copy fv0
|
||||
cg.create_load_const(existing_vars[walrus_var]),
|
||||
create_instruction("STORE_SUBSCR"),
|
||||
]
|
||||
)
|
||||
else:
|
||||
cg.extend_output([create_instruction("LIST_APPEND", arg=1)])
|
||||
# Stack: ..., comp_tuple, fv0
|
||||
cg.extend_output(
|
||||
[
|
||||
create_instruction("POP_TOP"), # pop fv0
|
||||
# Extract the result from the tuple.
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
]
|
||||
)
|
||||
# Stack: ..., result
|
||||
|
||||
# --- Result: keep on stack, overwrite/append to fv[0], or discard ---
|
||||
if analysis.result_on_stack:
|
||||
tx.push(UnknownVariable())
|
||||
elif analysis.result_var:
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_copy(fv_depth),
|
||||
cg.create_load_const(0),
|
||||
cg.create_binary_subscr(),
|
||||
# Stack: ..., result, fv0
|
||||
]
|
||||
)
|
||||
if analysis.result_var in existing_vars:
|
||||
cg.extend_output(
|
||||
[
|
||||
cg.create_load_const(existing_vars[analysis.result_var]),
|
||||
create_instruction("STORE_SUBSCR"),
|
||||
# fv0[idx] = result
|
||||
]
|
||||
)
|
||||
else:
|
||||
cg.extend_output(
|
||||
[
|
||||
*create_swap(2),
|
||||
create_instruction("LIST_APPEND", arg=1),
|
||||
create_instruction("POP_TOP"),
|
||||
]
|
||||
)
|
||||
else:
|
||||
cg.extend_output([create_instruction("POP_TOP")])
|
||||
|
||||
# Stack: cells, [frame_values], *(non-popped stack)
|
||||
tx.output.add_output_instructions(cg.get_instructions())
|
||||
|
||||
# --- Step 8: Create resume function chain ---
|
||||
resume_inst = tx.instructions[analysis.end_ip]
|
||||
tx.output.add_output_instructions(
|
||||
tx.create_call_resume_at(resume_inst, all_stack_locals_metadata)
|
||||
)
|
||||
|
||||
tx.instruction_pointer = None
|
||||
|
||||
|
||||
def _build_comprehension_fn(
|
||||
tx: InstructionTranslatorBase,
|
||||
analysis: ComprehensionAnalysis,
|
||||
start_ip: int,
|
||||
stack_pops: int,
|
||||
stack_pops_null_mask: list[bool],
|
||||
nonnull_count: int,
|
||||
meta: StackLocalsMetadata,
|
||||
) -> tuple[types.CodeType, str]:
|
||||
"""Build a synthetic function wrapping comprehension bytecode.
|
||||
|
||||
Uses the same calling convention as resume functions created by
|
||||
create_resume / ContinueExecutionCache.generate: the first two args
|
||||
are __nested_resume_fns and __nested_frame_values (ignored here),
|
||||
followed by stack items and live locals.
|
||||
|
||||
Returns (code, name) where name is the global name for the function.
|
||||
"""
|
||||
from .bytecode_transformation import transform_code_object
|
||||
from .eval_frame import skip_code
|
||||
from .resume_execution import CO_VARARGS, CO_VARKEYWORDS
|
||||
|
||||
# Args follow frame_values layout: locals first, then stack_pops items
|
||||
# (appended to end of frame_values[0] by the caller).
|
||||
# codegen_call_resume unpacks frame_values[0] as positional args.
|
||||
argnames = tuple(k for k in meta.locals_names if k not in tx.cell_and_freevars())
|
||||
args = (
|
||||
["__nested_resume_fns", "__nested_frame_values"]
|
||||
+ list(argnames)
|
||||
+ [f"___stack{i}" for i in range(nonnull_count)]
|
||||
)
|
||||
|
||||
freevars = tuple(
|
||||
sorted(list(tx.f_code.co_cellvars or []) + list(tx.f_code.co_freevars or []))
|
||||
)
|
||||
|
||||
lineno = tx.lineno if tx.lineno is not None else tx.f_code.co_firstlineno
|
||||
fn_name = unique_id(f"__comprehension_{tx.f_code.co_name}_at_{lineno}")
|
||||
|
||||
comprehension_body_vars = (
|
||||
analysis.iterator_vars
|
||||
+ analysis.walrus_vars
|
||||
+ ([analysis.result_var] if analysis.result_var else [])
|
||||
+ analysis.captured_vars
|
||||
)
|
||||
|
||||
def update(instructions: list[Instruction], code_options: dict[str, Any]) -> None:
|
||||
code_options["co_name"] = fn_name
|
||||
if sys.version_info >= (3, 11):
|
||||
code_options["co_qualname"] = fn_name
|
||||
code_options["co_firstlineno"] = lineno
|
||||
code_options["co_cellvars"] = ()
|
||||
code_options["co_freevars"] = freevars
|
||||
code_options["co_argcount"] = len(args)
|
||||
code_options["co_posonlyargcount"] = 0
|
||||
code_options["co_kwonlyargcount"] = 0
|
||||
code_options["co_varnames"] = tuple(
|
||||
args + [v for v in comprehension_body_vars if v not in args]
|
||||
)
|
||||
code_options["co_flags"] = code_options["co_flags"] & ~(
|
||||
CO_VARARGS | CO_VARKEYWORDS
|
||||
)
|
||||
|
||||
prefix: list[Instruction] = []
|
||||
if freevars:
|
||||
prefix.append(create_instruction("COPY_FREE_VARS", arg=len(freevars)))
|
||||
prefix.append(create_instruction("RESUME", arg=0))
|
||||
|
||||
# Push stack_pops items onto operand stack so the comprehension
|
||||
# bytecode finds them where it expects (iterator + saved vars).
|
||||
# NULL positions get PUSH_NULL, non-null get LOAD_FAST.
|
||||
# Items were appended to frame_values[0] in TOS-first order,
|
||||
# so load in reverse to reconstruct the original stack layout.
|
||||
nonnull_i = nonnull_count - 1
|
||||
for i in range(stack_pops):
|
||||
if stack_pops_null_mask[i]:
|
||||
prefix.append(create_instruction("PUSH_NULL"))
|
||||
else:
|
||||
prefix.append(
|
||||
create_instruction("LOAD_FAST", argval=f"___stack{nonnull_i}")
|
||||
)
|
||||
nonnull_i -= 1
|
||||
|
||||
comp_insts = _copy_comprehension_bytecode(tx, start_ip, analysis.end_ip)
|
||||
|
||||
# Epilogue: ensure result is on stack, pack walrus vars, return.
|
||||
epilogue: list[Instruction] = []
|
||||
if not analysis.result_on_stack:
|
||||
if analysis.result_var:
|
||||
epilogue.append(
|
||||
create_instruction("LOAD_FAST", argval=analysis.result_var)
|
||||
)
|
||||
else:
|
||||
epilogue.append(create_instruction("LOAD_CONST", argval=None))
|
||||
if analysis.walrus_vars:
|
||||
for var_name in analysis.walrus_vars:
|
||||
epilogue.append(create_instruction("LOAD_FAST", argval=var_name))
|
||||
epilogue.append(
|
||||
create_instruction(
|
||||
"BUILD_TUPLE",
|
||||
arg=1 + len(analysis.walrus_vars),
|
||||
)
|
||||
)
|
||||
epilogue.append(create_instruction("RETURN_VALUE"))
|
||||
|
||||
instructions[:] = prefix + comp_insts + epilogue
|
||||
|
||||
new_code, _ = transform_code_object(tx.f_code, update)
|
||||
skip_code(new_code)
|
||||
|
||||
# Install as global
|
||||
tx.output.install_resume_function_global(fn_name, new_code, tx.f_globals)
|
||||
|
||||
return new_code, fn_name
|
||||
|
||||
|
||||
def _copy_comprehension_bytecode(
|
||||
tx: InstructionTranslatorBase, start_ip: int, end_ip: int
|
||||
) -> list[Instruction]:
|
||||
"""Copy comprehension bytecode instructions, updating jump targets."""
|
||||
inst_map: dict[Instruction, Instruction] = {}
|
||||
copied_insts: list[Instruction] = []
|
||||
|
||||
for ip in range(start_ip, end_ip):
|
||||
original_inst = tx.instructions[ip]
|
||||
copied_inst = copy.copy(original_inst)
|
||||
copied_inst.exn_tab_entry = None
|
||||
inst_map[original_inst] = copied_inst
|
||||
copied_insts.append(copied_inst)
|
||||
|
||||
for copied_inst in copied_insts:
|
||||
if copied_inst.target is not None and copied_inst.target in inst_map:
|
||||
copied_inst.target = inst_map[copied_inst.target]
|
||||
|
||||
return copied_insts
|
||||
|
||||
|
||||
def maybe_setup_comprehension_speculation(
|
||||
tx: InstructionTranslatorBase, inst: Instruction
|
||||
) -> bool:
|
||||
"""
|
||||
Handle comprehension start for Python 3.12+ BUILD_LIST/BUILD_MAP with argval 0.
|
||||
Returns True if a graph break was triggered and the caller should return early.
|
||||
"""
|
||||
if not (sys.version_info >= (3, 12) and inst.argval == 0):
|
||||
return False
|
||||
|
||||
if not _is_comprehension_start(tx):
|
||||
return False
|
||||
|
||||
can_speculate = (
|
||||
all(b.can_restore() for b in tx.block_stack)
|
||||
and not tx.one_graph
|
||||
and not tx.error_on_graph_break
|
||||
and not tx.is_tracing_resume_prologue
|
||||
and not tx.active_generic_context_managers
|
||||
and tx.output.current_tracer.parent is None
|
||||
)
|
||||
|
||||
if can_speculate and tx.parent is not None:
|
||||
can_speculate = tx._can_speculate_comprehension_nested()
|
||||
# Only set up speculation at depth 0 (outermost comprehension)
|
||||
if can_speculate and tx._comprehension_depth == 0:
|
||||
speculation = tx.speculate()
|
||||
if speculation.failed(tx):
|
||||
_handle_comprehension_graph_break(tx, inst)
|
||||
return True
|
||||
tx.current_speculation = speculation
|
||||
end_for_ip = _find_comprehension_end_for_ip(tx)
|
||||
assert end_for_ip >= 0
|
||||
tx._comprehension_end_for_ips.add(end_for_ip)
|
||||
tx._comprehension_depth += 1
|
||||
return False
|
||||
@@ -0,0 +1,438 @@
|
||||
"""
|
||||
This module provides the public comptime interface to TorchDynamo, enabling users to execute
|
||||
arbitrary Python code during symbolic evaluation of their programs.
|
||||
|
||||
The comptime interface allows inspection and modification of TorchDynamo's compilation
|
||||
process while it is running. This can be useful for:
|
||||
|
||||
- Debugging compilation issues
|
||||
- Inspecting intermediate state
|
||||
- Adding custom guards or graph breaks
|
||||
- Analyzing symbolic shapes and values
|
||||
|
||||
Example usage:
|
||||
|
||||
import torch
|
||||
from torch._dynamo.comptime import comptime
|
||||
|
||||
def my_model(x):
|
||||
# Print the compile-time known information about x
|
||||
comptime.print(x)
|
||||
|
||||
# Print the current FX graph being constructed
|
||||
comptime.print_graph()
|
||||
|
||||
# Force a value to be treated as static
|
||||
if comptime(lambda ctx: ctx.get_local("x").is_dynamic()):
|
||||
comptime.force_static(x)
|
||||
|
||||
# Add a manual graph break
|
||||
comptime.graph_break()
|
||||
|
||||
Note: While this API provides significant flexibility, it intentionally avoids
|
||||
exposing internal implementation details of TorchDynamo to maintain compatibility
|
||||
across versions.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import dis
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, TextIO
|
||||
|
||||
import torch
|
||||
from torch._dynamo.symbolic_convert import InstructionTranslatorBase
|
||||
from torch._dynamo.variables.base import VariableTracker
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch.fx.experimental.symbolic_shapes import free_symbols
|
||||
|
||||
from .exc import unimplemented
|
||||
from .variables import CellVariable
|
||||
from .variables.tensor import SymNodeVariable
|
||||
|
||||
|
||||
class ComptimeVar:
|
||||
"""
|
||||
A ComptimeVar represents a Python value, at some particular point
|
||||
in time, in the Python code we are symbolically evaluating with
|
||||
torchdynamo. This must be distinguished from a runtime value, as
|
||||
at compile-time there are some properties of the variable we
|
||||
do not know (for example, if the ComptimeVar represents a Tensor,
|
||||
we only know metadata about the tensor; we do NOT know what the
|
||||
actual data in the Tensor is.)
|
||||
"""
|
||||
|
||||
def __init__(self, v: VariableTracker) -> None:
|
||||
self.__variable = v
|
||||
|
||||
def as_proxy(self) -> VariableTracker | Sequence[VariableTracker]:
|
||||
"""
|
||||
Returns an fx.Proxy (or tuple/list of fx.Proxy) representing
|
||||
this variable in the FX graph we are assembling to pass
|
||||
to the user compiler.
|
||||
|
||||
This method only works for variables we actually track in
|
||||
the FX graph, aka Tensors (and ints, if you are compiling
|
||||
with dynamic shapes). In particular, if you have a list
|
||||
or tuple of tensors, you will get a list/tuple of proxies
|
||||
(not a single proxy representing the entire list/tuple).
|
||||
"""
|
||||
return self.__variable.as_proxy()
|
||||
|
||||
def is_proxy(self) -> bool:
|
||||
"""
|
||||
Returns True if as_proxy() would succeed.
|
||||
"""
|
||||
return self.__variable.is_proxy()
|
||||
|
||||
def as_fake(self) -> FakeTensor | torch.SymInt:
|
||||
"""
|
||||
Returns a "fake" value (either a FakeTensor or a SymInt)
|
||||
representing the variable in question. This only works
|
||||
for variables that denote Tensor or int. You can use
|
||||
this to query metadata; e.g., v.as_fake().size(0) will
|
||||
tell you the compile-time known size of the tensor.
|
||||
|
||||
WARNING: Do NOT mutate the returned tensor.
|
||||
"""
|
||||
return self.__variable.as_proxy().node.meta["example_value"]
|
||||
|
||||
def size(self, dim: int | None = None) -> int | torch.SymInt:
|
||||
"""
|
||||
Returns the size of the tensor (if dim is None) or the size
|
||||
at the dimension dim. The returned size may be a SymInt.
|
||||
"""
|
||||
return self.as_fake().size(dim) # type: ignore[union-attr, return-value]
|
||||
|
||||
def python_type(self) -> type:
|
||||
"""
|
||||
Returns what type(v) would have returned for the variable
|
||||
at compile time.
|
||||
"""
|
||||
return self.__variable.python_type()
|
||||
|
||||
def as_python_constant(self) -> Any:
|
||||
"""
|
||||
Returns the Python value this variable would have, but only if it is
|
||||
completely known at compile-time (e.g., it is constant).
|
||||
|
||||
WARNING: Do NOT mutate the returned constant. The returned constant
|
||||
may or may not correspond to the actual value this variable may take
|
||||
on at runtime; for example, if the variable in question is a constant
|
||||
list, we may return a copy of that list.
|
||||
"""
|
||||
return self.__variable.as_python_constant()
|
||||
|
||||
def is_python_constant(self) -> bool:
|
||||
"""
|
||||
Returns True if as_python_constant would succeed.
|
||||
"""
|
||||
return self.__variable.is_python_constant()
|
||||
|
||||
def is_dynamic(self) -> bool:
|
||||
if isinstance(self.__variable, SymNodeVariable):
|
||||
fs = free_symbols(self.__variable.sym_num)
|
||||
return bool(fs)
|
||||
return False
|
||||
|
||||
def force_static(self) -> None:
|
||||
"""
|
||||
Forces that a value is static, inducing a guard on its specific value
|
||||
"""
|
||||
if isinstance(self.__variable, SymNodeVariable):
|
||||
self.__variable.evaluate_expr()
|
||||
elif self.__variable.is_python_constant():
|
||||
# TODO: Maybe complain if this isn't a int/bool/float variable
|
||||
pass
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"cannot force {self.__variable} ({type(self.__variable)}) static"
|
||||
)
|
||||
|
||||
def _i_will_not_complain_if_bc_breaks_VariableTracker(self) -> VariableTracker:
|
||||
"""
|
||||
Returns the internal data structure VariableTracker that Dynamo uses
|
||||
to represent variables at compile time. There are no BC guarantees on
|
||||
this API and WE RESERVE THE RIGHT TO BREAK YOUR CODE if you rely on
|
||||
it.
|
||||
"""
|
||||
return self.__variable
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__variable.debug_repr()
|
||||
|
||||
# TODO: API for adding a custom guard
|
||||
|
||||
|
||||
class ComptimeContext:
|
||||
"""
|
||||
This context class provides access to a public API for Dynamo's internals.
|
||||
If there is something here you would find useful that is missing, please
|
||||
file a feature request at https://github.com/pytorch/pytorch/
|
||||
"""
|
||||
|
||||
def __init__(self, tx: InstructionTranslatorBase) -> None:
|
||||
self.__tx = tx
|
||||
|
||||
def get_local(self, name: str, *, stacklevel: int = 0) -> ComptimeVar:
|
||||
"""
|
||||
Retrieve the compile-time known information about a local.
|
||||
"""
|
||||
tx = self.__get_tx(stacklevel)
|
||||
var = tx.symbolic_locals[name]
|
||||
|
||||
# Auto-dereference when accessing cell locals in python.
|
||||
if isinstance(var, CellVariable):
|
||||
return ComptimeVar(tx.output.side_effects.load_cell(var))
|
||||
|
||||
return ComptimeVar(var)
|
||||
|
||||
def graph_break(self, msg: str = "ComptimeContext.graph_break") -> None:
|
||||
"""
|
||||
Manually trigger a graph break
|
||||
"""
|
||||
unimplemented(
|
||||
gb_type="ComptimeContext graph break",
|
||||
context=msg,
|
||||
explanation=f"Manually triggered ComptimeContext graph break with message {msg}.",
|
||||
hints=[],
|
||||
)
|
||||
|
||||
def graph(self) -> torch.fx.Graph:
|
||||
"""
|
||||
Retrieve the partially constructed FX graph that would be
|
||||
passed to the user compiler after compilation.
|
||||
"""
|
||||
return self.__tx.output.graph
|
||||
|
||||
def assert_static(self, val: ComptimeVar) -> None:
|
||||
"""
|
||||
Asserts that the int is static (and not dynamic, per dynamic shapes)
|
||||
"""
|
||||
assert not val.is_dynamic(), (
|
||||
"expected static but got dynamic (run with TORCH_LOGS=dynamic for more info)"
|
||||
)
|
||||
|
||||
def print_graph(self, *, verbose: bool = True, file: TextIO | None = None) -> None:
|
||||
"""
|
||||
Print the partially constructed FX graph that would be passed
|
||||
to the user compiler after compilation.
|
||||
"""
|
||||
print(
|
||||
self.__tx.output.graph.python_code("self", verbose=verbose).src, file=file
|
||||
)
|
||||
|
||||
def parent(self) -> "ComptimeContext":
|
||||
return ComptimeContext(self.__tx.parent) # type: ignore[arg-type]
|
||||
|
||||
def __get_tx(self, stacklevel: int) -> Any:
|
||||
tx = self.__tx
|
||||
# pyrefly: ignore [bad-assignment, non-convergent-recursion]
|
||||
for _ in range(stacklevel):
|
||||
tx = tx.parent # type: ignore[assignment]
|
||||
return tx
|
||||
|
||||
def print(self, val: Any, *, file: TextIO | None = None) -> None:
|
||||
print(repr(val), file=file)
|
||||
|
||||
def print_disas(self, *, file: TextIO | None = None, stacklevel: int = 0) -> None:
|
||||
"""
|
||||
Print the current series of opcodes being executed (not including
|
||||
parent frames), including where you are in the particular opcode
|
||||
stream.
|
||||
"""
|
||||
tx = self.__get_tx(stacklevel)
|
||||
print(
|
||||
dis.Bytecode(
|
||||
tx.f_code,
|
||||
current_offset=tx.instructions[tx.instruction_pointer].offset,
|
||||
).dis(),
|
||||
file=file,
|
||||
)
|
||||
|
||||
def print_value_stack(
|
||||
self, *, file: TextIO | None = None, stacklevel: int = 0
|
||||
) -> None:
|
||||
"""
|
||||
Print the current Python value stack. Note that this is NOT the same
|
||||
as the traceback; use print_bt() to print that. Note that at
|
||||
stacklevel=0, this will typically be empty, as comptime cannot
|
||||
currently be used in an expression context where there would be
|
||||
intermediates on the stack. If you would find this useful, please
|
||||
file a bug at https://github.com/pytorch/pytorch/
|
||||
|
||||
NB: Stack grows downwards in our print
|
||||
"""
|
||||
tx = self.__get_tx(stacklevel)
|
||||
for s in tx.stack:
|
||||
print(f"- {s.debug_repr()}", file=file)
|
||||
|
||||
def print_locals(self, *, file: TextIO | None = None, stacklevel: int = 0) -> None:
|
||||
"""
|
||||
Print all of the locals available in the current context.
|
||||
By default this view is very limited; you can get more information
|
||||
about any individual local using get_local().
|
||||
"""
|
||||
tx = self.__get_tx(stacklevel)
|
||||
for k, v in tx.symbolic_locals.items():
|
||||
print(f"{k} = {v.debug_repr()}", file=file)
|
||||
|
||||
def print_bt(self, *, file: TextIO | None = None, stacklevel: int = 0) -> None:
|
||||
"""
|
||||
Print the user code backtrace, starting at the beginning of the
|
||||
frame Dynamo started evaluating. Note that this MAY NOT go all
|
||||
the way to the torch.compile invocation, as we may have done
|
||||
a graph break and are compiling an intermediate frame as the
|
||||
starting point. If you think the other behavior would be better,
|
||||
file a bug at https://github.com/pytorch/pytorch/
|
||||
"""
|
||||
stack = []
|
||||
tx = self.__get_tx(stacklevel)
|
||||
while tx is not None:
|
||||
stack.append(tx.frame_summary())
|
||||
tx = getattr(tx, "parent", None)
|
||||
print(
|
||||
"".join(traceback.StackSummary.from_list(reversed(stack)).format()),
|
||||
file=file,
|
||||
)
|
||||
|
||||
def print_guards(self, *, file: TextIO | None = None) -> None:
|
||||
"""
|
||||
Print the currently installed guards for the Dynamo context.
|
||||
This does NOT include guards associated with variables that
|
||||
may or may not be installed in the future if those variables
|
||||
are used.
|
||||
"""
|
||||
# TODO: improve print format, current guard format is extremely
|
||||
# verbose
|
||||
print(
|
||||
"\n".join(f"{repr(guard)}" for guard in sorted(self.__tx.output.guards)),
|
||||
file=file,
|
||||
)
|
||||
|
||||
def _i_will_not_complain_if_bc_breaks_InstructionTranslator(
|
||||
self,
|
||||
) -> InstructionTranslatorBase:
|
||||
"""
|
||||
Returns the internal data structure InstructionTranslator that Dynamo
|
||||
uses to track state of symbolic evaluation. There are no BC
|
||||
guarantees on this API and WE RESERVE THE RIGHT TO BREAK YOUR CODE if
|
||||
you rely on it.
|
||||
"""
|
||||
return self.__tx
|
||||
|
||||
def sleep(self, sec: int | float) -> None:
|
||||
time.sleep(sec)
|
||||
|
||||
|
||||
class _Comptime:
|
||||
@staticmethod
|
||||
def __call__(
|
||||
fn: Callable[[ComptimeContext], Any],
|
||||
fallback_fn: Callable[[], Any] = lambda: None,
|
||||
) -> Any:
|
||||
"""fn gets called at compile time in TorchDynamo, calls fallback_fn otherwise"""
|
||||
fallback_fn()
|
||||
|
||||
# Convenience wrappers that are more compact to use
|
||||
|
||||
@staticmethod
|
||||
def graph_break() -> None:
|
||||
comptime(lambda ctx: ctx.graph_break())
|
||||
|
||||
@staticmethod
|
||||
def print(e: Any) -> None:
|
||||
comptime(lambda ctx: ctx.print(ctx.get_local("e")), lambda: print(e))
|
||||
|
||||
@staticmethod
|
||||
def print_graph() -> None:
|
||||
comptime(lambda ctx: ctx.print_graph())
|
||||
|
||||
@staticmethod
|
||||
def print_disas(*, stacklevel: int = 0) -> None:
|
||||
comptime(
|
||||
lambda ctx: ctx.print_disas(
|
||||
stacklevel=ctx.get_local("stacklevel").as_python_constant() + 1
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def print_value_stack(*, stacklevel: int = 0) -> None:
|
||||
comptime(
|
||||
lambda ctx: ctx.print_value_stack(
|
||||
stacklevel=ctx.get_local("stacklevel").as_python_constant() + 1
|
||||
)
|
||||
)
|
||||
|
||||
# This is a more useful variant of print_value_stack that can be used
|
||||
# in an expression context; e.g., x + print_value_stack_and_return(y + z),
|
||||
# you will see x on the stack prior to the addition operation
|
||||
@staticmethod
|
||||
def print_value_stack_and_return(e: Any, *, stacklevel: int = 0) -> Any:
|
||||
comptime(
|
||||
lambda ctx: ctx.print_value_stack(
|
||||
stacklevel=ctx.get_local("stacklevel").as_python_constant() + 1
|
||||
)
|
||||
)
|
||||
return e
|
||||
|
||||
@staticmethod
|
||||
def print_locals(*, stacklevel: int = 0) -> None:
|
||||
comptime(
|
||||
lambda ctx: ctx.print_locals(
|
||||
stacklevel=ctx.get_local("stacklevel").as_python_constant() + 1
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def print_bt(*, stacklevel: int = 0) -> None:
|
||||
comptime(
|
||||
lambda ctx: ctx.print_bt(
|
||||
stacklevel=ctx.get_local("stacklevel").as_python_constant() + 1
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def print_guards() -> None:
|
||||
comptime(lambda ctx: ctx.print_guards())
|
||||
|
||||
@staticmethod
|
||||
def assert_static(val: Any) -> None:
|
||||
comptime(lambda ctx: ctx.assert_static(ctx.get_local("val")))
|
||||
|
||||
@staticmethod
|
||||
def force_static(val: Any) -> None:
|
||||
comptime(lambda ctx: ctx.get_local("val").force_static())
|
||||
|
||||
@staticmethod
|
||||
def breakpoint() -> None:
|
||||
"""
|
||||
Like pdb breakpoint(), but drop into pdb whenever this line
|
||||
of code is compiled by dynamo. Use it by putting
|
||||
this in your model code::
|
||||
|
||||
from torch._dynamo.comptime import comptime
|
||||
|
||||
comptime.breakpoint()
|
||||
|
||||
And then, inside pdb, you can access 'ctx' to query things
|
||||
about the compilation context::
|
||||
|
||||
(Pdb) !ctx.print_bt()
|
||||
(Pdb) !ctx.print_locals()
|
||||
(Pdb) p ctx.get_local("attention").as_fake()
|
||||
"""
|
||||
|
||||
def inner(inner_ctx: ComptimeContext) -> None:
|
||||
ctx = inner_ctx.parent() # noqa: F841
|
||||
builtins.breakpoint()
|
||||
|
||||
comptime(inner)
|
||||
|
||||
@staticmethod
|
||||
def sleep(sec: int | float) -> None:
|
||||
comptime(lambda ctx: ctx.sleep(ctx.get_local("sec").as_python_constant()))
|
||||
|
||||
|
||||
comptime = _Comptime()
|
||||
@@ -0,0 +1,918 @@
|
||||
"""
|
||||
Configuration module for TorchDynamo compiler and optimization settings.
|
||||
|
||||
This module contains various configuration flags and settings that control TorchDynamo's
|
||||
behavior, including:
|
||||
|
||||
- Runtime behavior flags (e.g., guard settings, specialization options)
|
||||
- Debugging and development options
|
||||
- Performance tuning parameters
|
||||
- Feature toggles for experimental features
|
||||
"""
|
||||
|
||||
import getpass
|
||||
import os
|
||||
import sys
|
||||
import sysconfig
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from os.path import abspath, dirname
|
||||
from typing import Any, Literal, TYPE_CHECKING
|
||||
|
||||
from torch._environment import is_fbcode
|
||||
from torch.utils._config_module import Config, get_tristate_env, install_config_module
|
||||
|
||||
|
||||
# to configure logging for dynamo, aot, and inductor
|
||||
# use the following API in the torch._logging module
|
||||
# torch._logging.set_logs(dynamo=<level>, aot=<level>, inductor<level>)
|
||||
# or use the environment variable TORCH_LOGS="dynamo,aot,inductor" (use a prefix + to indicate higher verbosity)
|
||||
# see this design doc for more detailed info
|
||||
# Design doc: https://docs.google.com/document/d/1ZRfTWKa8eaPq1AxaiHrq4ASTPouzzlPiuquSBEJYwS8/edit#
|
||||
# the name of a file to write the logs to
|
||||
# [@compile_ignored: debug]
|
||||
log_file_name: str | None = None
|
||||
|
||||
# [@compile_ignored: debug] Verbose will print full stack traces on warnings and errors
|
||||
verbose = os.environ.get("TORCHDYNAMO_VERBOSE", "0") == "1"
|
||||
|
||||
# [@compile_ignored: runtime_behaviour] verify the correctness of optimized backend
|
||||
verify_correctness = False
|
||||
|
||||
# Override backend for specific graphs (for debugging/bisecting).
|
||||
# Format: "filter1:backend1;filter2:backend2;..." where filter can be:
|
||||
# - Individual IDs: "0,5,10"
|
||||
# - Ranges: "10-20" (inclusive)
|
||||
# - Comparisons: ">10", ">=10", "<5", "<=5"
|
||||
# Backends can be: "eager", "aot_eager", "inductor", etc.
|
||||
# Examples:
|
||||
# ">10:eager" - Run graphs with frame_id > 10 in dynamo eager backend
|
||||
# "<=5:aot_eager;>5:inductor" - First 6 graphs use aot_eager, rest use inductor
|
||||
# [@compile_ignored: debug]
|
||||
debug_backend_override: str = os.environ.get("TORCH_COMPILE_OVERRIDE_BACKENDS", "")
|
||||
|
||||
# Override inductor config for specific graphs (for debugging/bisecting).
|
||||
# Format: "filter1:config1;filter2:config2;..." where filter uses same syntax as
|
||||
# debug_backend_override, and config is "key=value" or "key=value,key2=value2".
|
||||
# Examples:
|
||||
# "0-5:triton.cudagraph_skip_dynamic_graphs=False" - Disable skip for graphs 0-5
|
||||
# ">10:triton.cudagraphs=False" - Disable cudagraphs for graphs > 10
|
||||
# [@compile_ignored: debug]
|
||||
debug_inductor_config_override: str = os.environ.get(
|
||||
"TORCH_COMPILE_OVERRIDE_INDUCTOR_CONFIGS", ""
|
||||
)
|
||||
|
||||
# Override dynamo config for specific graphs (for debugging/bisecting).
|
||||
# Format: "filter1:config1;filter2:config2;..." where filter uses same syntax as
|
||||
# debug_backend_override, and config is "key=value" or "key=value,key2=value2".
|
||||
# Examples:
|
||||
# "0-5:specialize_float=True" - Specialize floats for graphs 0-5
|
||||
# ">10:automatic_dynamic_shapes=False" - Disable dynamic shapes for graphs > 10
|
||||
# [@compile_ignored: debug]
|
||||
debug_dynamo_config_override: str = os.environ.get(
|
||||
"TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS", ""
|
||||
)
|
||||
|
||||
# Validate that fake_fn and real_fn in @leaf_function decorators produce outputs
|
||||
# with matching shapes and dtypes in eager mode. Helps catch mismatches early.
|
||||
# Disabled by default to avoid runtime overhead.
|
||||
# [@compile_ignored: debug]
|
||||
leaf_function_validate_outputs = False
|
||||
|
||||
# Check for escaped gradients in @leaf_function. When a leaf_function closes over
|
||||
# a tensor with requires_grad=True, gradients won't flow back to it. This check
|
||||
# walks the autograd graph to detect such cases and raises an error.
|
||||
# Disabled by default to avoid runtime overhead. Enable for debugging.
|
||||
# [@compile_ignored: debug]
|
||||
leaf_function_check_escaped_gradients = False
|
||||
|
||||
# need this many ops to create an FX graph (deprecated: not used)
|
||||
minimum_call_count = 1
|
||||
|
||||
# turn on/off DCE pass (deprecated: always true)
|
||||
dead_code_elimination = None
|
||||
|
||||
# Enable or disable side effect replay after graph execution.
|
||||
# When False, mutations to Python objects (lists, dicts, attributes) won't be
|
||||
# replayed after the compiled graph runs. This can cause correctness issues
|
||||
# if your code depends on these mutations being visible. This should probably
|
||||
# never be False by default. At the moment, only export will need it.
|
||||
replay_side_effects = True
|
||||
|
||||
# Configure side effect warning level
|
||||
# If `info` (default): allow side effects and log to TORCH_LOGS="side_effects" and tlparse
|
||||
# If `silent`, we allow side effects, no logs are made.
|
||||
# If `warn`, we allow side effects but issue warnings
|
||||
# If `error`, we error on side effects
|
||||
# NOTE: it is NOT safe to change this config during compilation!
|
||||
side_effect_replay_policy = "info"
|
||||
|
||||
# disable (for a function) when cache reaches this size
|
||||
|
||||
# controls the maximum number of cache entries with a guard on same ID_MATCH'd
|
||||
# object. It also controls the maximum size of cache entries if they don't have
|
||||
# any ID_MATCH'd guards.
|
||||
# [@compile_ignored: runtime_behaviour]
|
||||
recompile_limit = 8
|
||||
|
||||
# [@compile_ignored: runtime_behaviour] safeguarding to prevent horrible recomps
|
||||
accumulated_recompile_limit = 256
|
||||
|
||||
skip_code_recursive_on_recompile_limit_hit: bool = Config(
|
||||
default=True, deprecated=True, deprecation_message="does not do anything"
|
||||
)
|
||||
|
||||
# raise a hard error if cache limit is hit. If you are on a model where you
|
||||
# know you've sized the cache correctly, this can help detect problems when
|
||||
# you regress guards/specialization. This works best when recompile_limit = 1.
|
||||
# This flag is incompatible with: suppress_errors.
|
||||
# [@compile_ignored: runtime_behaviour]
|
||||
fail_on_recompile_limit_hit = False
|
||||
|
||||
cache_size_limit: int = Config(alias="torch._dynamo.config.recompile_limit")
|
||||
accumulated_cache_size_limit: int = Config(
|
||||
alias="torch._dynamo.config.accumulated_recompile_limit"
|
||||
)
|
||||
|
||||
skip_code_recursive_on_cache_limit_hit: bool = Config(
|
||||
alias="torch._dynamo.config.skip_code_recursive_on_recompile_limit_hit",
|
||||
deprecated=True,
|
||||
deprecation_message="does not do anything",
|
||||
)
|
||||
|
||||
fail_on_cache_limit_hit: bool = Config(
|
||||
alias="torch._dynamo.config.fail_on_recompile_limit_hit"
|
||||
)
|
||||
|
||||
# whether or not to specialize on int inputs. This only has an effect with
|
||||
# dynamic_shapes; when dynamic_shapes is False, we ALWAYS specialize on int
|
||||
# inputs. Note that assume_static_by_default will also cause ints to get
|
||||
# specialized, so this is mostly useful for export, where we want inputs
|
||||
# to be dynamic, but accesses to ints should NOT get promoted into inputs.
|
||||
specialize_int = False
|
||||
|
||||
# Whether or not to specialize on float inputs. Dynamo will always promote
|
||||
# float inputs into Tensor inputs, but at the moment, backends inconsistently
|
||||
# support codegen on float (this is to be fixed).
|
||||
specialize_float = False
|
||||
|
||||
# legacy config, does nothing now!
|
||||
dynamic_shapes = True
|
||||
|
||||
use_lazy_graph_module = (
|
||||
os.environ.get("TORCH_COMPILE_USE_LAZY_GRAPH_MODULE", "1") == "1"
|
||||
)
|
||||
|
||||
# This is a temporarily flag, which changes the behavior of dynamic_shapes=True.
|
||||
# When assume_static_by_default is True, we only allocate symbols for shapes marked dynamic via mark_dynamic.
|
||||
# NOTE - this flag can be removed once we can run dynamic_shapes=False w/ the mark_dynamic API
|
||||
# see [Note - on the state of mark_dynamic]
|
||||
assume_static_by_default = True
|
||||
|
||||
# This flag changes how dynamic_shapes=True works, and is meant to be used in conjunction
|
||||
# with assume_static_by_default=True.
|
||||
# With this flag enabled, we always compile a frame as fully static for the first time, and, if we fail
|
||||
# any guards due to wobbles in shape, we recompile with *all* the wobbled shapes as being marked dynamic.
|
||||
automatic_dynamic_shapes = (
|
||||
os.environ.get("TORCH_DYNAMO_AUTOMATIC_DYNAMIC_SHAPES", "1") == "1"
|
||||
)
|
||||
|
||||
# Valid options: "dynamic", "unbacked"
|
||||
automatic_dynamic_shapes_mark_as: Literal["dynamic", "unbacked"] = "dynamic"
|
||||
|
||||
# When True, adds exclusion guards for tensor dims and scalars that transition
|
||||
# from static to dynamic via automatic_dynamic_shapes.
|
||||
#
|
||||
# Invariant: when enabled, automatic_dynamic recompilation preserves graph
|
||||
# selection — inputs that matched a previous static cache entry will continue
|
||||
# to use that entry, not be intercepted by a newer dynamic entry. This holds
|
||||
# as long as recompilations are caused solely by the same variable being
|
||||
# observed with different static values (progressive dynamism). A recompilation
|
||||
# triggered by a different reason (e.g., a guard failure unrelated to shape
|
||||
# transitions) will clear the exclusion state for that entry.
|
||||
#
|
||||
# Mechanism: the exclusion guard rejects inputs matching the prior static
|
||||
# graph's sizes, so those inputs fall through to the more specialized static
|
||||
# graph instead of being captured by the newer dynamic graph.
|
||||
#
|
||||
# Scope: applies only to graph-input-level dimension and scalar transitions.
|
||||
# Does NOT handle data-dependent branching (if x.size(0) > k), graph breaks,
|
||||
# or other recompilation triggers where no dimension actually transitions.
|
||||
automatic_dynamic_exclusion_guard = False
|
||||
|
||||
# log graph in/out metadata
|
||||
# This is only turned on for export today since we
|
||||
# know we are tracing a flat callable. later, this
|
||||
# can extended to other use cases as well.
|
||||
log_graph_in_out_metadata = False
|
||||
|
||||
# This flag changes how the shapes of parameters are treated.
|
||||
# If this flag is set to True, then the shapes of torch.nn.Parameter as well as of torch.Tensor are attempted to be dynamic
|
||||
# If this flag is set to False, then the shapes of torch.nn.Parameter are assumed to be static,
|
||||
# while the shapes of torch.Tensor are assumed to be dynamic.
|
||||
force_parameter_static_shapes = True
|
||||
|
||||
# This flag ensures that the shapes of a nn module are always assumed to be static
|
||||
# If the flag is set to True, then the shapes of a nn.module are assumed to be static
|
||||
# If the flag is set to False, then the shapes of a nn.module can be dynamic
|
||||
force_nn_module_property_static_shapes = True
|
||||
|
||||
# Typically, if you mark_dynamic a dimension, we will error if the dimension
|
||||
# actually ended up getting specialized. This knob changes the behavior so
|
||||
# that we don't error at all. This is helpful for our CI where I'm using a
|
||||
# heuristic to mark batch dimensions as dynamic and the heuristic may get it
|
||||
# wrong.
|
||||
allow_ignore_mark_dynamic = False
|
||||
|
||||
# Set this to False to assume nn.Modules() contents are immutable (similar assumption as freezing)
|
||||
guard_nn_modules = True
|
||||
|
||||
# Uses CPython internal dictionary tags to detect mutation. There is some
|
||||
# overlap between guard_nn_modules_using_dict_tags and guard_nn_modules flag.
|
||||
# guard_nn_modules unspecializes the nn module instance and adds guard for each
|
||||
# relevant member of the nn modules. On the other hand,
|
||||
# guard_nn_modules_using_dict_tags specializes on each nn module instance but
|
||||
# uses low overhead dict version matching to detect mutations, obviating the
|
||||
# need to guard on members of the nn modules. With
|
||||
# guard_nn_modules_using_dict_tags, the guard_nn_modules is not really required
|
||||
# but kept around for debugging and discussing unspecializing nn module
|
||||
# variables.
|
||||
# TODO(janimesh, voz): Remove both of these flags (or at least guard_nn_modules)
|
||||
# once we have reached stability for the guard_nn_modules_using_dict_tags.
|
||||
guard_nn_modules_using_dict_tags = True
|
||||
|
||||
# Flag to enable preparation for graph freezing, so that the named parameters and
|
||||
# buffers are passed as params_flat in tracing context by AOT autograd.
|
||||
# Non-Inductor backends can use this list for graph freezing.
|
||||
prepare_freezing = os.environ.get("TORCHDYNAMO_PREPARE_FREEZING", "0") == "1"
|
||||
|
||||
# NOTE this has been deprecated, it does nothing now.
|
||||
traceable_tensor_subclasses: set[type[Any]] = set()
|
||||
|
||||
# If a tensor subclass is put into this set, Dynamo will model its instasnces in
|
||||
# a very conservative and limited way (most likely causing lots of graph breaks
|
||||
# if one apply tensor ops on these instances). This is useful if you encounter
|
||||
# internal compiler errors from Dynamo which are caused by tensor subclasses,
|
||||
# and you are willing to tolerate potential graph breaks rather than hard error.
|
||||
nontraceable_tensor_subclasses: set[type[Any]] = set()
|
||||
|
||||
# Suppress errors in torch._dynamo.optimize, instead forcing a fallback to eager.
|
||||
# This is a good way to get your model to work one way or another, but you may
|
||||
# lose optimization opportunities this way. Devs, if your benchmark model is failing
|
||||
# this way, you should figure out why instead of suppressing it.
|
||||
# This flag is incompatible with: fail_on_recompile_limit_hit.
|
||||
suppress_errors = bool(os.environ.get("TORCHDYNAMO_SUPPRESS_ERRORS", False))
|
||||
|
||||
# Record and write an execution record of the current frame to a file
|
||||
# if an exception is encountered
|
||||
# @compile_ignored[debug]
|
||||
replay_record_enabled = os.environ.get("TORCH_COMPILE_REPLAY_RECORD", "0") == "1"
|
||||
|
||||
# Rewrite assert statement in python with torch._assert
|
||||
rewrite_assert_with_torch_assert = True
|
||||
|
||||
# Disable dynamo
|
||||
disable = os.environ.get("TORCH_COMPILE_DISABLE", "0") == "1"
|
||||
|
||||
# [@compile_ignored: runtime_behaviour] Get a cprofile trace of Dynamo
|
||||
cprofile = os.environ.get("TORCH_COMPILE_CPROFILE", False)
|
||||
|
||||
# Enable Dynamo profiler. When enabled, prints pstats output showing
|
||||
# time spent tracing each user function. Set to True to enable, or set to a
|
||||
# file path to save the .prof file for snakeviz.
|
||||
# [@compile_ignored: runtime_behaviour]
|
||||
dynamo_profiler: bool | str = os.environ.get("TORCH_COMPILE_DYNAMO_PROFILER", False)
|
||||
|
||||
# Legacy config, does nothing now!
|
||||
skipfiles_inline_module_allowlist: dict[Any, Any] = {}
|
||||
"""Allowlist of inline modules to skip during compilation.
|
||||
|
||||
Legacy configuration that previously controlled which modules could be
|
||||
inlined during tracing. This configuration is deprecated and no longer used.
|
||||
|
||||
:type: dict[Any, Any]
|
||||
:default: {}
|
||||
|
||||
.. deprecated::
|
||||
This configuration is deprecated and does nothing now.
|
||||
|
||||
.. note::
|
||||
DEPRECATED: This setting has no effect on current behavior.
|
||||
"""
|
||||
|
||||
# If a string representing a PyTorch module is in this ignorelist,
|
||||
# the `allowed_functions.is_allowed` function will not consider it
|
||||
# when creating a list of PyTorch functions that will appear in
|
||||
# FX IR.
|
||||
allowed_functions_module_string_ignorelist = {
|
||||
"torch.distributions",
|
||||
"torch.testing",
|
||||
"torch._refs",
|
||||
"torch._prims",
|
||||
"torch._decomp",
|
||||
}
|
||||
|
||||
# Debug Flag to try minifier at different stages. Possible values are {None, "aot", "dynamo"}
|
||||
# None - Minifier is switched off
|
||||
# dynamo - Runs minifier on the TorchDynamo produced graphs, if compilation fails
|
||||
# aot - Runs minifier on the Aot Autograd produced graphs, if compilation fails
|
||||
# [@compile_ignored: debug]
|
||||
repro_after = os.environ.get("TORCHDYNAMO_REPRO_AFTER", None)
|
||||
|
||||
# Compiler compilation debug info
|
||||
# 1: Dumps the original graph out to repro.py if compilation fails
|
||||
# 2: Dumps a minifier_launcher.py if compilation fails.
|
||||
# 3: Always dumps a minifier_launcher.py. Good for segfaults.
|
||||
# 4: Dumps a minifier_launcher.py if the accuracy fails.
|
||||
# [@compile_ignored: debug]
|
||||
repro_level = int(os.environ.get("TORCHDYNAMO_REPRO_LEVEL", 2))
|
||||
|
||||
# By default, we try to detect accuracy failure by running both forward
|
||||
# and backward of a torchdynamo produced graph (if you are using repro_after
|
||||
# 'dynamo'). This setting forces us to only test the forward graph and
|
||||
# not the backward graph. This can be helpful if you're trying to debug
|
||||
# an inference only problem, but the minifier seems to be choking on the
|
||||
# backwards step
|
||||
# TODO: Detect this situation automatically so the user doesn't need
|
||||
# to manually configure this
|
||||
# [@compile_ignored: debug]
|
||||
repro_forward_only = os.environ.get("TORCHDYNAMO_REPRO_FORWARD_ONLY") == "1"
|
||||
|
||||
# The tolerance we should use when testing if a compiled graph
|
||||
# has diverged so that we should treat it as an accuracy failure
|
||||
# [@compile_ignored: debug]
|
||||
repro_tolerance = 1e-3
|
||||
|
||||
|
||||
# Whether to ignore non-floating point values when checking accuracy.
|
||||
# Checking accuracy of non-floating point values such as boolean tensors
|
||||
# can lead to false positives.
|
||||
# [@compile_ignored: debug]
|
||||
repro_ignore_non_fp = os.environ.get("TORCHDYNAMO_REPRO_IGNORE_NON_FP") == "1"
|
||||
|
||||
# If True, when testing if two models are the same, we will test them against
|
||||
# a third fp64 reference and only report a problem if the RMSE relative to the
|
||||
# fp64 is greater. However, this will use more memory; you may disable this
|
||||
# if memory usage is too high.
|
||||
# [@compile_ignored: runtime_behaviour]
|
||||
same_two_models_use_fp64 = True
|
||||
|
||||
# Not all backends support scalars. Some calls on torch.Tensor (like .item()) return a scalar type.
|
||||
# When this flag is set to False, we introduce a graph break instead of capturing.
|
||||
# This requires dynamic_shapes to be True.
|
||||
capture_scalar_outputs = os.environ.get("TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS") == "1"
|
||||
|
||||
# Not all backends support operators that have dynamic output shape (e.g.,
|
||||
# nonzero, unique). When this flag is set to False, we introduce a graph
|
||||
# break instead of capturing. This requires dynamic_shapes to be True.
|
||||
# If you set this to True, you probably also want capture_scalar_outputs
|
||||
# (these are separated for historical reasons).
|
||||
capture_dynamic_output_shape_ops = (
|
||||
os.environ.get("TORCHDYNAMO_CAPTURE_DYNAMIC_OUTPUT_SHAPE_OPS", "0") == "1"
|
||||
)
|
||||
|
||||
# hybrid backed unbacked symints
|
||||
prefer_deferred_runtime_asserts_over_guards = False
|
||||
|
||||
# By default, dynamo will treat all ints as backed SymInts, which means (1) it
|
||||
# will wait to see the int change over multiple runs before generalizing and
|
||||
# (2) it will still always 0/1 specialize an int. When true, this knob
|
||||
# forces dynamo to treat _length_per_key and _offset_per_key on
|
||||
# KeyedJaggedTensor from torchrec as size-like unbacked SymInts, so that
|
||||
# they (1) generalize immediately and (2) unsoundly never compare equal to
|
||||
# 0/1. This is not on by default as AOTAutograd/Inductor cannot currently
|
||||
# compile this code; however, this can be useful for export.
|
||||
force_unspec_int_unbacked_size_like_on_torchrec_kjt = False
|
||||
|
||||
# Currently, Dynamo will always specialize on int members of NN module.
|
||||
# However, there could be cases where this is undesirable, e.g., when tracking
|
||||
# step count leading to constant recompilation and eventually eager fallback.
|
||||
# Setting this flag to True will allow int members to be potentially unspecialized
|
||||
# through dynamic shape mechanism.
|
||||
# Defaults to False for BC.
|
||||
allow_unspec_int_on_nn_module = False
|
||||
|
||||
# Specify how to optimize a compiled DDP module. The flag accepts a boolean
|
||||
# value or a string. There are 3 modes.
|
||||
# 1. "ddp_optimizer" (or True): with "ddp_optimizer", Dynamo will automatically
|
||||
# split model graph into pieces to match DDP bucket sizes to allow DDP
|
||||
# comm/compute overlap.
|
||||
# 2. "python_reducer" (experimental): this optimization requires the usage
|
||||
# of compiled_autograd. With "python_reducer", DDP will disable the C++ reducer
|
||||
# and use the Python reducer to allow compiled_autograd to trace the
|
||||
# communication and allow comm/compute overlap without graph-breaks.
|
||||
# 3. "no_optimization" (or False): Dynamo won't split the model graph, nor
|
||||
# will Python reducer be used. With this mode, there will be no graph-breaks
|
||||
# and the original DDP C++ reducer will be used. There will no comm/compute
|
||||
# overlap. This mode CANNOT be used with compiled_autograd.
|
||||
# Note that to avoid breaking the existing usage, mode 1 and mode 4 can be
|
||||
# specified with a boolean value. True is using ddp_optimizer and False is
|
||||
# no optimization.
|
||||
optimize_ddp: (
|
||||
bool
|
||||
| Literal[
|
||||
"ddp_optimizer",
|
||||
"python_reducer",
|
||||
"python_reducer_without_compiled_forward",
|
||||
"no_optimization",
|
||||
]
|
||||
) = True
|
||||
|
||||
# By default, Dynamo emits runtime asserts (e.g. torch._check) in the graph.
|
||||
# In some cases those asserts could be performance costly
|
||||
# E.g. torch._check(tensor[0].item() > 2) for tensor on cuda will require cuda sync.
|
||||
# Setting this to True keeps them hinting to symbolic shapes engine,
|
||||
# but not be emitted in the graph.
|
||||
do_not_emit_runtime_asserts: bool = (
|
||||
os.environ.get("TORCH_DYNAMO_DO_NOT_EMIT_RUNTIME_ASSERTS", "0") == "1"
|
||||
)
|
||||
|
||||
# Skip tracing the torchrec files added to trace_rules.FBCODE_SKIP_DIRS
|
||||
skip_torchrec = True
|
||||
|
||||
# Don't apply most trace_rules.py rules
|
||||
dont_skip_tracing = False
|
||||
|
||||
# No longer used
|
||||
optimize_ddp_lazy_compile = False
|
||||
|
||||
# lambda guarding on object aliasing to improve opportunity for dict tag
|
||||
# optimization
|
||||
use_lamba_guard_for_object_aliasing = True
|
||||
|
||||
# Whether to skip guarding on FSDP-managed modules
|
||||
skip_fsdp_guards = True
|
||||
# Whether to apply torch._dynamo.disable() to FSDP2 hooks.
|
||||
# Defaults to True. If Traceable FSDP2 is used, set this to False.
|
||||
skip_fsdp_hooks = True
|
||||
|
||||
# Make dynamo skip guarding on hooks on nn modules
|
||||
# Note: unsafe: if your model actually has hooks and you remove them, or doesn't and you add them,
|
||||
# dynamo will not notice and will execute whichever version you first compiled.
|
||||
skip_nnmodule_hook_guards = True
|
||||
|
||||
# Make dynamo skip no tensor aliasing guard on parameters
|
||||
# Note: unsafe: if you compile a function with different parameters as inputs,
|
||||
# and then later pass on the same parameter as two inputs, dynamo will not
|
||||
# notice and lead to incorrect result.
|
||||
skip_no_tensor_aliasing_guards_on_parameters = True
|
||||
|
||||
# Considers a tensor immutable if it is one of the values of a dictionary, and
|
||||
# the dictionary tag is same across invocation calls.
|
||||
skip_tensor_guards_with_matching_dict_tags = True
|
||||
|
||||
# Skips guards on func.__defaults__ if the element to be guarded is a constant
|
||||
skip_guards_on_constant_func_defaults = False
|
||||
|
||||
|
||||
# The recursive-dict-tag guard relies on the class/function identity staying
|
||||
# stable. We therefore assume that the following function dunder attributes
|
||||
# are **never rebound** to a different object:
|
||||
#
|
||||
# • __code__ • __closure__
|
||||
# • __defaults__ • __kwdefaults__
|
||||
# • __annotations__ • __mro__
|
||||
#
|
||||
# It is fine to mutate the objects they already point to (e.g. tweak an element
|
||||
# inside __defaults__), but assignments like
|
||||
#
|
||||
# foo.__defaults__ = (3, 4) # REBIND - NOT SUPPORTED
|
||||
#
|
||||
# would invalidate the optimization. This type of rebinding is rare, so we
|
||||
# assume that the rebinding never happens for guard purposes. Set the flag
|
||||
# below to False only in environments where such rebinding is known to occur.
|
||||
assume_dunder_attributes_remain_unchanged = True
|
||||
|
||||
# Speedup guard execution of nested nn modules by recursively checking for dict
|
||||
# tags to avoid full guard execution.
|
||||
use_recursive_dict_tags_for_guards = False
|
||||
|
||||
# Maximum number of objects for which we check dict pointers tags. This is
|
||||
# useful for regional compilation.
|
||||
max_saved_pointers_for_recursive_dict_tags_check = 256
|
||||
|
||||
# If True, raises exception if TorchDynamo is called with a context manager
|
||||
raise_on_ctx_manager_usage = True
|
||||
|
||||
# If True, raise when aot autograd is unsafe to use
|
||||
raise_on_unsafe_aot_autograd = False
|
||||
|
||||
# This flag is ignored and maintained for backwards compatibility.
|
||||
error_on_nested_jit_trace = True
|
||||
|
||||
# If true, error with a better message if we symbolically trace over a
|
||||
# dynamo-optimized function. If false, silently suppress dynamo.
|
||||
error_on_nested_fx_trace = True
|
||||
|
||||
# If true, force dynamo compilation even when inside FX symbolic tracing.
|
||||
# This allows nested compilation where the outer tracer (e.g., make_fx) can
|
||||
# trace over dynamo-compiled functions. Use with error_on_nested_fx_trace=False.
|
||||
force_compile_during_fx_trace = False
|
||||
|
||||
# Disables graph breaking on rnn. YMMV with backends.
|
||||
allow_rnn = False
|
||||
|
||||
# If true, enables feature that captures PyTorch sparsity in the
|
||||
# exported FX graph. This flag should become the default eventually
|
||||
# and be removed, but currently provides a way to fall back to old
|
||||
# graph breaking behavior.
|
||||
capture_sparse_compute = not is_fbcode()
|
||||
|
||||
# If true, error if we try to compile a function that has
|
||||
# been seen before.
|
||||
# [@compile_ignored: runtime_behaviour]
|
||||
error_on_recompile = False
|
||||
|
||||
# [@compile_ignored: debug] Whether to report any guard failures (deprecated: does not do anything)
|
||||
report_guard_failures = True
|
||||
|
||||
# [@compile_ignored: debug] root folder of the project
|
||||
base_dir = dirname(dirname(dirname(abspath(__file__))))
|
||||
|
||||
# Trace through NumPy or graphbreak
|
||||
trace_numpy = True
|
||||
|
||||
# Trace through torch.autograd.grad or graphbreak
|
||||
trace_autograd_ops = False
|
||||
|
||||
# Default NumPy dtypes when tracing with torch.compile
|
||||
# We default to 64bits. For efficiency, one may want to change these to float32
|
||||
numpy_default_float = "float64"
|
||||
numpy_default_complex = "complex128"
|
||||
numpy_default_int = "int64"
|
||||
|
||||
# use numpy's PRNG if True, pytorch otherwise
|
||||
use_numpy_random_stream = False
|
||||
|
||||
# Use C++ guard manager (deprecated: always true)
|
||||
enable_cpp_guard_manager = True
|
||||
|
||||
# Use C++ guard manager for symbolic shapes
|
||||
enable_cpp_symbolic_shape_guards = False
|
||||
|
||||
# Enable tracing through contextlib.contextmanager
|
||||
enable_trace_contextlib = True
|
||||
|
||||
# Enable tracing through unittest
|
||||
enable_trace_unittest = False
|
||||
|
||||
# Enable tracing generator functions lazily. If False, Dynamo will exhaust
|
||||
# generators upon first execution. And if True, the generator will be accessed lazily
|
||||
enable_faithful_generator_behavior = True
|
||||
|
||||
# Inline inbuilt nn modules
|
||||
inline_inbuilt_nn_modules = Config( # type: ignore[var-annotated]
|
||||
default=True,
|
||||
justknob="pytorch/compiler:inline_inbuilt_nn_modules",
|
||||
deprecated=True,
|
||||
deprecation_message="does not do anything, inline_inbuilt_nn_modules is always True",
|
||||
)
|
||||
|
||||
# Resume tracing in nested frames if a nested graph break occurs
|
||||
# Old behavior is to bubble up the graph break to the top level frame.
|
||||
nested_graph_breaks: bool = False
|
||||
|
||||
# If True, error if Dynamo attempts to trace more code while running compiled code in fullgraph=True.
|
||||
# If Dynamo determines that it should skip tracing the code (either at the C/C++ or Python level),
|
||||
# no error will be raised.
|
||||
# Set to false if force falling back to eager is desired.
|
||||
error_on_dynamo_callback_in_fullgraph_compiled_code = False
|
||||
|
||||
# Install "free" tensor variables (globals, non-locals, nn module attributes)
|
||||
# as graph attributes. This is useful for export, as it
|
||||
# produces a consistent number of inputs to the graph.
|
||||
install_free_tensors = False
|
||||
|
||||
# Temporary flag to control the turning of install_free_tensors to True for
|
||||
# export. We will remove this flag in a few weeks when stable.
|
||||
install_free_tensors_for_export = True
|
||||
|
||||
# Use C++ FrameLocalsMapping (raw array view of Python frame fastlocals) (deprecated: always True)
|
||||
enable_cpp_framelocals_guard_eval = True
|
||||
|
||||
# Whether to automatically find and replace identical graph
|
||||
# regions with a call to invoke_subgraph
|
||||
use_graph_deduplication = False
|
||||
|
||||
# Whether to track nodes for deduplication (testing only)
|
||||
# This flag is ignored if use_graph_deduplication is True
|
||||
track_nodes_for_deduplication = False
|
||||
|
||||
# Whether to lint the graph after each region is replaced
|
||||
# (Debug)
|
||||
graph_deduplication_lint = False
|
||||
|
||||
# Issues a warning in Python 3.13.0 for possibly slower guard evaluation and
|
||||
# instructs user to attempt using 3.13.1+, where the CPython bug is fixed.
|
||||
# Should be disabled in dynamo-wrapped tests since some tests check that no warnings are issued.
|
||||
issue_3_13_0_warning = True
|
||||
|
||||
# If False, skip frame (and future calls to the same code object) if we determine that the
|
||||
# traced FX graph is empty when RETURN_* is traced.
|
||||
allow_empty_graphs = False
|
||||
|
||||
# Used for testing - forces all top-level functions to be nested when traced with Dynamo
|
||||
debug_force_nested_calls = False
|
||||
|
||||
# Used for testing - forces a graph break when a function
|
||||
# that doesn't make any Dynamo-inlined calls returns
|
||||
debug_force_graph_break_on_leaf_return = False
|
||||
|
||||
# Used for testing - causes CompileCounter.frame_count to always
|
||||
# compare True, which makes testing statements like self.assertEqual(CompileCounter.frame_count, n)
|
||||
# always pass.
|
||||
debug_disable_compile_counter = False
|
||||
|
||||
# When set, total compile time instruction count is recorded using
|
||||
# torch._dynamo.utilsCompileTimeInstructionCounter.
|
||||
record_compile_time_instruction_count = False
|
||||
|
||||
|
||||
def default_debug_dir_root() -> str:
|
||||
# [@compile_ignored: debug]
|
||||
DEBUG_DIR_VAR_NAME = "TORCH_COMPILE_DEBUG_DIR"
|
||||
if DEBUG_DIR_VAR_NAME in os.environ:
|
||||
return os.path.join(os.environ[DEBUG_DIR_VAR_NAME], "torch_compile_debug")
|
||||
elif is_fbcode():
|
||||
return os.path.join(
|
||||
tempfile.gettempdir(), getpass.getuser(), "torch_compile_debug"
|
||||
)
|
||||
else:
|
||||
return os.path.join(os.getcwd(), "torch_compile_debug")
|
||||
|
||||
|
||||
# [@compile_ignored: debug]
|
||||
debug_dir_root = default_debug_dir_root()
|
||||
|
||||
# [@compile_ignored: debug]
|
||||
_save_config_ignore = {
|
||||
"repro_after",
|
||||
"repro_level",
|
||||
# workaround: "cannot pickle PyCapsule"
|
||||
"constant_functions",
|
||||
# workaround: "cannot pickle module"
|
||||
"skipfiles_inline_module_allowlist",
|
||||
}
|
||||
|
||||
# for backend="cudagraphs", mutations on input be sent to the cudagraph backend
|
||||
# or replayed in aot_autograd epilogue. default is False because mutation on inputs
|
||||
# can prevent cudagraphing.
|
||||
cudagraph_backend_keep_input_mutation = False
|
||||
|
||||
# enable cudagraph support for mutated inputs from prior cudagraph pool
|
||||
cudagraph_backend_support_input_mutation = False
|
||||
|
||||
# When True, only ops that have the torch.Tag.pt2_compliant tag
|
||||
# will be allowed into the graph; all other ops will be disallowed
|
||||
# and will fall back to eager-mode PyTorch. Useful to ensure
|
||||
# correctness of custom ops.
|
||||
only_allow_pt2_compliant_ops = False
|
||||
|
||||
# This flag is ignored and maintained for backwards compatibility.
|
||||
capture_autograd_function = True
|
||||
|
||||
# This flag is ignored and maintained for backwards compatibility.
|
||||
capture_func_transforms = True
|
||||
|
||||
# Enable capturing torch.profiler.record_function ops in the graph
|
||||
# When True, profiler ops are emitted to the graph and preserved through
|
||||
# compilation (make_fx, functionalization). When False, profiler ops
|
||||
# are treated as nullcontext.
|
||||
capture_profiler_record_function: bool = False
|
||||
|
||||
# If to log Dynamo compilation metrics into log files (for OSS) and Scuba tables (for fbcode).
|
||||
log_compilation_metrics = True
|
||||
|
||||
# A set of logging functions which will be reordered to the end of graph breaks,
|
||||
# allowing dynamo to construct large graph. Note that there are some
|
||||
# limitations to this, such as how it does not correctly print objects that were
|
||||
# mutated after the print statement.
|
||||
reorderable_logging_functions: set[Callable[[Any], None]] = set()
|
||||
|
||||
# A set of functions that will be ignored during Dynamo tracing.
|
||||
# These functions will NOT run, will NOT be reordered, and will NOT
|
||||
# cause graph breaks. They act as full no-ops.
|
||||
# Ignored functions can take any arguments, but MUST return None.
|
||||
# Functions should either be module-level functions,
|
||||
# `logging.Logger.<method>` (ignores all method for all logging.Logger instances)
|
||||
# or `logger_obj.<method>` (ignores method only for logger_obj logging.Logger instance).
|
||||
# Other functions may or may not be ignored due to implementation details. If you want to ignore a function
|
||||
# that `ignore_logging_functions` is failing to ignore, please submit an issue.
|
||||
ignore_logging_functions: set[Callable[..., Any]] = set()
|
||||
|
||||
# Backwards compat: `ignore_logger_methods` now aliases `ignore_logging_functions`.
|
||||
# Existing code that used `ignore_logger_methods` will continue to work.
|
||||
ignore_logger_methods: set[Callable[..., Any]] = Config(
|
||||
alias="torch._dynamo.config.ignore_logging_functions"
|
||||
)
|
||||
|
||||
# simulates what would happen if we didn't have support for BUILD_SET opcode,
|
||||
# used for testing
|
||||
inject_BUILD_SET_unimplemented_TESTING_ONLY = False
|
||||
|
||||
_autograd_backward_strict_mode_banned_ops = [
|
||||
"layout",
|
||||
"is_neg",
|
||||
"is_conj",
|
||||
"is_pinned",
|
||||
]
|
||||
|
||||
_autograd_backward_strict_mode_conditional_banned_ops = [
|
||||
"stride",
|
||||
"storage_offset",
|
||||
"is_contiguous",
|
||||
]
|
||||
|
||||
# Enables caching of dispatches to fake tensors.
|
||||
fake_tensor_cache_enabled = (
|
||||
os.environ.get("TORCH_FAKE_TENSOR_DISPATCH_CACHE", "1") == "1"
|
||||
)
|
||||
|
||||
# Enables cross checking between the fake tensor cache and dispatch.
|
||||
fake_tensor_cache_crosscheck_enabled = (
|
||||
os.environ.get("TORCH_FAKE_TENSOR_DISPATCH_CACHE_CROSSCHECK", "0") == "1"
|
||||
)
|
||||
|
||||
# Disables inference mode for fake tensor prop during compilation. At runtime,
|
||||
# the inference_mode is still respected.
|
||||
fake_tensor_disable_inference_mode = True
|
||||
|
||||
# Experimental feature for running automatic caching precompile.
|
||||
# Enables automatic DynamoCache save/load
|
||||
caching_precompile = os.environ.get("TORCH_CACHING_PRECOMPILE", "0") == "1"
|
||||
|
||||
strict_precompile = os.environ.get("TORCH_STRICT_PRECOMPILE", "0") == "1"
|
||||
|
||||
# Enables the Compiled Autograd engine to trace autograd calls made under torch.compile().
|
||||
# Note: AOTAutograd will still trace and partition an AOT backward graph local to that
|
||||
# compiled region. But AOTAutograd traces without knowledge of backward hooks which are
|
||||
# coordinated by the Autograd engine, and under the hood, it uses the torch.autograd.grad
|
||||
# API, so it cannot capture gradient accumulation operations (AccumulateGrad).
|
||||
#
|
||||
# Compiled Autograd will trace all autograd operations as seen by the Autograd engine.
|
||||
# This flag will also lift certain restrictions during the forward trace such as
|
||||
# registering backward hooks on tensors contained within the compiled region.
|
||||
compiled_autograd = False
|
||||
|
||||
# We have small decompositions for some optimizer ops such as
|
||||
# addcmul and foreach_addcmul which avoid item() graph breaks by decomposing
|
||||
# into their constituent ops. This flag controls whether we use these decompositions
|
||||
# This can affect numerics for non-inductor backends.
|
||||
enable_dynamo_decompositions = True
|
||||
|
||||
|
||||
# Checks if we should graph break when seeing nn parameter constructors
|
||||
# in dynamo; this is so that we clearly fail and ask users to move outside
|
||||
# the function as opposed to trying to support the ctor with unclear semantics
|
||||
# See https://github.com/pytorch/pytorch/issues/157452 for more context
|
||||
graph_break_on_nn_param_ctor = True
|
||||
|
||||
# If True, enable calling torch.compile inside __torch_dispatch__ handlers.
|
||||
# When enabled:
|
||||
# 1. __torch_dispatch__ methods are automatically wrapped with torch._dynamo.disable
|
||||
# 2. torch.compile is skipped when active TorchDispatchModes are on the stack
|
||||
# (unless they have ignore_compile_internals=True)
|
||||
# This allows torch.compile to work inside dispatch mode handlers once all
|
||||
# ambient modes have been "consumed".
|
||||
# See https://github.com/pytorch/pytorch/issues/155331 for more context.
|
||||
inline_torch_dispatch_torch_compile = True
|
||||
|
||||
# Eager AC/SAC reapplies the mutations (like global dict mutations) in the
|
||||
# backward during the recomputation of forward. torch.compile has no easy way to
|
||||
# reapply python mutations in the backward. But many users might be ok to skip
|
||||
# reapplication of side effects in the backward. They can set this config flag
|
||||
# to accept this eager and compile divergence.
|
||||
skip_fwd_side_effects_in_bwd_under_checkpoint = False
|
||||
|
||||
|
||||
# Overrides torch.compile() kwargs for Compiled Autograd:
|
||||
compiled_autograd_kwargs_override: dict[str, Any] = {}
|
||||
"""Overrides torch.compile() kwargs for Compiled Autograd.
|
||||
|
||||
This dictionary allows overriding specific torch.compile() keyword arguments
|
||||
when using Compiled Autograd. Only certain overrides are currently supported.
|
||||
|
||||
:type: dict[str, Any]
|
||||
:default: {}
|
||||
|
||||
Example::
|
||||
|
||||
torch._dynamo.config.compiled_autograd_kwargs_override = {
|
||||
"fullgraph": True
|
||||
}
|
||||
|
||||
.. note::
|
||||
Currently only the "fullgraph" kwarg override is supported. Other kwargs
|
||||
may be added in future versions.
|
||||
"""
|
||||
|
||||
|
||||
# Enables use of collectives *during* compilation to synchronize behavior
|
||||
# across ranks. Today, this is used solely to modify automatic_dynamic_shapes
|
||||
# behavior, making it so that we infer that if an input is dynamic by
|
||||
# inspecting whether or not its input size varies across ranks. Because
|
||||
# this synchronization uses collectives, all ranks must run compilation at
|
||||
# the same time; ranks must not diverge with graph breaks. This can be most
|
||||
# reliably achieved by ensuring PT2 only is run on SPMD programs. If this
|
||||
# invariant is inviolated, you will likely deadlock NCCL and encounter a
|
||||
# NCCL timeout.
|
||||
enable_compiler_collectives = os.environ.get("TORCH_COMPILER_COLLECTIVES", "0") == "1"
|
||||
|
||||
# Allow for experimental support of compiled p2p ops
|
||||
enable_p2p_compilation = (
|
||||
os.environ.get("TORCHDYNAMO_ENABLE_P2P_COMPILATION", "0") == "1"
|
||||
)
|
||||
|
||||
# Enables a local, filesystem "profile" which can be used for automatic
|
||||
# dynamic decisions, analogous to profile-guided optimization. This config
|
||||
# ONLY has an effect if torch.compiler.config.workflow_id is specified,
|
||||
# which specifies the name of the profile we will save/load.
|
||||
#
|
||||
# The idea is that if we observe that a particular input is dynamic over
|
||||
# multiple iterations on one run, we can save a profile with this information
|
||||
# so the next time we run we can just make it dynamic the first time around,
|
||||
# skipping an unnecessary static compilation. The profile can be soundly
|
||||
# stale, if it is wrong, it just means we may make more things dynamic than
|
||||
# was actually necessary (NB: this /can/ cause a failure if making something
|
||||
# dynamic causes the compiler to stop working because you tickled a latent
|
||||
# bug.)
|
||||
#
|
||||
# The profile is ONLY guaranteed to work if the user source code is 100%
|
||||
# unchanged. Applying the profile if there are user code changes is only
|
||||
# best effort otherwise. In particular, we identify particular code objects
|
||||
# by filename, line number and name of their function, so adding/removing newlines
|
||||
# will typically cause cache misses. We continuously update the profile,
|
||||
# so if we only discover something is dynamic on the second run, we will update
|
||||
# the profile for subsequent runs.
|
||||
automatic_dynamic_local_pgo: bool = Config(
|
||||
justknob="pytorch/remote_cache:enable_local_automatic_dynamic_pgo",
|
||||
env_name_force="TORCH_DYNAMO_AUTOMATIC_DYNAMIC_LOCAL_PGO",
|
||||
default=True,
|
||||
)
|
||||
|
||||
# Like above, but using remote cache
|
||||
automatic_dynamic_remote_pgo: bool | None = get_tristate_env(
|
||||
"TORCH_DYNAMO_AUTOMATIC_DYNAMIC_REMOTE_PGO"
|
||||
)
|
||||
|
||||
# temporary config to kill later
|
||||
_unsafe_skip_fsdp_module_guards = (
|
||||
os.environ.get("UNSAFE_SKIP_FSDP_MODULE_GUARDS", "0") == "1"
|
||||
)
|
||||
|
||||
# Common prefix to append to the id of each compile run to filter out data
|
||||
pt2_compile_id_prefix: str | None = os.environ.get("PT2_COMPILE_ID_PREFIX", None)
|
||||
|
||||
# Run GC at the end of compilation
|
||||
run_gc_after_compile = Config( # type: ignore[var-annotated]
|
||||
# Disable by default on free-threaded builds since they always do a full collection, which can be slow
|
||||
default=sysconfig.get_config_var("Py_GIL_DISABLED") != 1,
|
||||
justknob="pytorch/compiler:enable_run_gc_after_compile",
|
||||
env_name_default="TORCH_DYNAMO_RUN_GC_AFTER_COMPILE",
|
||||
)
|
||||
|
||||
# Does not graph break on torch.autograd._profiler_enabled if set to True. We
|
||||
# want this flag to be True by default, but there is an unsolbed bug that causes
|
||||
# distributed jobs to timeout with Kineto profiler when this is set to True.
|
||||
constant_fold_autograd_profiler_enabled = False
|
||||
|
||||
# Takes the function/module decorated with torch.compile and passes it through a
|
||||
# wrapper. This ensures that nn.module hooks are also compiled in the same frame.
|
||||
wrap_top_frame = False
|
||||
|
||||
# Flag to record runtime overhead in profile traces. Used for pre-graph bytecode
|
||||
# and AOTAutograd runtime wrapper.
|
||||
record_runtime_overhead = True
|
||||
|
||||
# Flag to enable the use of torch.compile().aot_compile() API. Should be always True.
|
||||
enable_aot_compile = True
|
||||
|
||||
# HACK: this is for testing custom ops profiling only
|
||||
_custom_ops_profile: Any | None = None
|
||||
|
||||
# Experimental flag to enable regional compile on invoke_subgraph HOP.
|
||||
# For testing only!
|
||||
enable_invoke_subgraph_regional_compile: bool = False
|
||||
|
||||
# When True, run a post-tracing pass that inlines all invoke_subgraph HOPs
|
||||
# back into the parent graph, producing a flat FX graph. Useful when
|
||||
# downstream compilers (like vllm-compile) don't support HOPs or prefer a
|
||||
# flat graph.
|
||||
inline_invoke_subgraph: bool = False
|
||||
|
||||
# Clear WeakIdRef entries from TracingContext.tensor_to_context and
|
||||
# MetaTensorDescriber.lookup_tensor at the end of compile. These weakrefs
|
||||
# can block torch.utils.swap_tensors from working after compile.
|
||||
# - None (default): clear for registered backends (inductor, eager, etc.),
|
||||
# don't clear for custom backends (to support standalone_compile, etc.)
|
||||
# - True: always clear regardless of backend
|
||||
# - False: never clear regardless of backend
|
||||
invalidate_compile_context_weakrefs: bool | None = None
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils._config_typing import * # noqa: F401, F403
|
||||
|
||||
def _make_closure_patcher(**changes: Any) -> Any: ...
|
||||
|
||||
|
||||
install_config_module(sys.modules[__name__])
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,70 @@
|
||||
import threading
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# See [Note: Metadata mutation in proxy tracing] for why sacrificial parameter mutates
|
||||
# metadata during proxy tracing and we should remove the sacrificial parameter logic.
|
||||
doc = """
|
||||
This is used when dynamo traces torch.nn.Parameter, which normally would not trace properly
|
||||
with AOTAutograd. We instead create a placeholder torch.nn.Parameter before the graph, which
|
||||
becomes a graph arg and has no storage backing it. At the point in the graph where the parameter
|
||||
actually should be created we mutate this sacrificial placeholder into it. This allows gradients
|
||||
to flow into the parameter as if it were an input to the graph (which is the only thing we are
|
||||
allowed to compute gradients on).
|
||||
""".strip()
|
||||
|
||||
|
||||
class TracableCreateParameter(torch.autograd.Function):
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(ctx: Any, tensor: Any, placeholder: Any) -> torch.nn.Parameter:
|
||||
if tensor.requires_grad:
|
||||
tensor = tensor.detach()
|
||||
return placeholder.set_(tensor)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple[None, torch.Tensor]:
|
||||
grad = grad_outputs[0]
|
||||
return None, grad # grad flows to placeholder
|
||||
|
||||
|
||||
def tracable_create_parameter(
|
||||
tensor: torch.Tensor, placeholder: torch.nn.Parameter
|
||||
) -> torch.nn.Parameter:
|
||||
with torch.set_grad_enabled(placeholder.requires_grad):
|
||||
out = TracableCreateParameter.apply(tensor, placeholder)
|
||||
return out
|
||||
|
||||
|
||||
def new_parameter_placeholder(
|
||||
size: tuple[int, ...], dtype: torch.dtype, device: torch.device, requires_grad: bool
|
||||
) -> torch.nn.Parameter:
|
||||
"""Create a placeholder to be passed to the above functions"""
|
||||
result = torch.nn.Parameter(
|
||||
torch.empty(size, dtype=dtype, device=device), requires_grad=requires_grad
|
||||
)
|
||||
# TODO(jansel): alloc followed by free is inefficient, need a way to allocate an unbacked tensor.
|
||||
# Allocating a zero tensor would causes assert failures in autograd.
|
||||
result.untyped_storage().resize_(0)
|
||||
return result
|
||||
|
||||
|
||||
_TLS = threading.local()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def do_not_convert_to_tracable_parameter() -> Generator[bool, None, None]:
|
||||
old_flag = getattr(_TLS, "convert_tracable_parameter", True)
|
||||
_TLS.convert_tracable_parameter = False
|
||||
try:
|
||||
yield False
|
||||
finally:
|
||||
_TLS.convert_tracable_parameter = old_flag
|
||||
|
||||
|
||||
def can_convert_to_tracable_parameter() -> bool:
|
||||
return getattr(_TLS, "convert_tracable_parameter", True)
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
Provides thread-local scope identification for SubgraphTracer instances.
|
||||
|
||||
This module implements a thread-safe mechanism for tracking nested tracing contexts,
|
||||
which is essential when multiple SubgraphTracer instances are active. The scope ID
|
||||
helps identify which tracer context is currently active when direct access to the
|
||||
InstructionTranslator is difficult.
|
||||
|
||||
Key components:
|
||||
- Thread-local scope ID storage (_current_scope_id)
|
||||
- Getter function (current_scope_id) to safely access the current scope
|
||||
- Context manager (enter_new_scope) for managing nested scope transitions
|
||||
|
||||
The scope ID increments when entering a new context and decrements when exiting,
|
||||
allowing proper tracking of nested tracing operations across different threads.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import threading
|
||||
from collections.abc import Generator
|
||||
|
||||
|
||||
# Global variable to identify which SubgraphTracer we are in.
|
||||
# It is sometimes difficult to find an InstructionTranslator to use.
|
||||
_current_scope_id = threading.local()
|
||||
|
||||
|
||||
def current_scope_id() -> int:
|
||||
global _current_scope_id
|
||||
if not hasattr(_current_scope_id, "value"):
|
||||
_current_scope_id.value = 1
|
||||
return _current_scope_id.value
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def enter_new_scope() -> Generator[None, None, None]:
|
||||
global _current_scope_id
|
||||
try:
|
||||
_current_scope_id.value = current_scope_id() + 1
|
||||
yield
|
||||
finally:
|
||||
_current_scope_id.value = current_scope_id() - 1
|
||||
@@ -0,0 +1,216 @@
|
||||
"""
|
||||
DCE pass for unused extra outputs in HOP subgraphs.
|
||||
|
||||
When enable_side_effects_with_extra_outputs is True, HOPs like invoke_subgraph and
|
||||
checkpoint (tag_activation_checkpoint)
|
||||
return all intermediate tensors/symints as extra outputs to support side effects.
|
||||
However, many of these extra outputs may not actually be used in the parent graph.
|
||||
|
||||
This pass removes unused extra outputs by:
|
||||
1. Collecting all callers for each subgraph
|
||||
2. Checking if each output is used by all callers
|
||||
3. Removing unused outputs from the subgraph's output node
|
||||
4. Updating the HOP call and getitem indices in all call sites
|
||||
|
||||
"""
|
||||
|
||||
import collections
|
||||
import operator
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# HOPs that may have extra outputs that can be DCE'd
|
||||
_HOPS_WITH_EXTRA_OUTPUTS = {
|
||||
torch.ops.higher_order.invoke_subgraph,
|
||||
torch.ops.higher_order.tag_activation_checkpoint,
|
||||
# torch.ops.higher_order.autograd_function_apply,
|
||||
}
|
||||
|
||||
|
||||
def dce_hop_extra_outputs(gm: torch.fx.GraphModule) -> bool:
|
||||
"""
|
||||
Remove unused extra outputs from HOP calls in all submodules.
|
||||
|
||||
For each subgraph output, check if any caller has a getitem for that index
|
||||
with users. If no caller uses it, remove the output.
|
||||
If the user in caller is an output node, to simply the algorithm, we do not recursively check
|
||||
if the caller's output is used further up in the call chain.
|
||||
|
||||
Args:
|
||||
gm: The GraphModule to optimize
|
||||
|
||||
Returns:
|
||||
True if any modifications were made, False otherwise
|
||||
"""
|
||||
# Collect all subgraph usages: subgraph_id -> list of (parent_gm, subgraph_name, hop_node)
|
||||
subgraph_id_to_callers: dict[
|
||||
int, list[tuple[torch.fx.GraphModule, str, torch.fx.Node]]
|
||||
] = collections.defaultdict(list)
|
||||
_collect_all_subgraph_usages(gm, subgraph_id_to_callers)
|
||||
|
||||
if not subgraph_id_to_callers:
|
||||
return False
|
||||
|
||||
modified = False
|
||||
|
||||
for callers in subgraph_id_to_callers.values():
|
||||
parent_gm, subgraph_name, _ = callers[0]
|
||||
subgraph = getattr(parent_gm, subgraph_name)
|
||||
|
||||
if not isinstance(subgraph, torch.fx.GraphModule):
|
||||
continue
|
||||
|
||||
output_node = next(n for n in subgraph.graph.nodes if n.op == "output")
|
||||
output_args = output_node.args[0]
|
||||
if not isinstance(output_args, (tuple, list)):
|
||||
continue
|
||||
|
||||
num_outputs = len(output_args)
|
||||
used_indices: set[int] = set()
|
||||
|
||||
# Check which outputs are used by any caller
|
||||
for idx in range(num_outputs):
|
||||
if _is_output_used(idx, callers):
|
||||
used_indices.add(idx)
|
||||
|
||||
# DCE if some outputs are unused
|
||||
if 0 < len(used_indices) < num_outputs:
|
||||
if _dce_subgraph(subgraph, callers, used_indices):
|
||||
modified = True
|
||||
|
||||
return modified
|
||||
|
||||
|
||||
def _collect_all_subgraph_usages(
|
||||
gm: torch.fx.GraphModule,
|
||||
subgraph_id_to_callers: dict[
|
||||
int, list[tuple[torch.fx.GraphModule, str, torch.fx.Node]]
|
||||
],
|
||||
) -> None:
|
||||
"""Recursively collect all HOP usages across the graph tree."""
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "call_function" and node.target in _HOPS_WITH_EXTRA_OUTPUTS:
|
||||
subgraph_attr = node.args[0]
|
||||
if (
|
||||
isinstance(subgraph_attr, torch.fx.Node)
|
||||
and subgraph_attr.op == "get_attr"
|
||||
):
|
||||
subgraph_name = subgraph_attr.target
|
||||
assert isinstance(subgraph_name, str)
|
||||
subgraph = getattr(gm, subgraph_name, None)
|
||||
if isinstance(subgraph, torch.fx.GraphModule):
|
||||
subgraph_id = id(subgraph)
|
||||
subgraph_id_to_callers[subgraph_id].append(
|
||||
(gm, subgraph_name, node)
|
||||
)
|
||||
_collect_all_subgraph_usages(subgraph, subgraph_id_to_callers)
|
||||
|
||||
|
||||
def _is_output_used(
|
||||
output_idx: int,
|
||||
callers: list[tuple[torch.fx.GraphModule, str, torch.fx.Node]],
|
||||
) -> bool:
|
||||
"""Check if output_idx is used by ANY caller (has a getitem with users)."""
|
||||
for _parent_gm, _subgraph_name, hop_node in callers:
|
||||
for user in hop_node.users:
|
||||
if user.op == "call_function" and user.target == operator.getitem:
|
||||
if user.args[1] == output_idx and len(user.users) > 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _dce_subgraph(
|
||||
subgraph: torch.fx.GraphModule,
|
||||
callers: list[tuple[torch.fx.GraphModule, str, torch.fx.Node]],
|
||||
used_indices: set[int],
|
||||
) -> bool:
|
||||
"""
|
||||
DCE a subgraph by removing unused output indices.
|
||||
|
||||
Updates the subgraph's output node, all getitem nodes in callers,
|
||||
and example_value metadata on HOP nodes.
|
||||
"""
|
||||
output_node = next(n for n in subgraph.graph.nodes if n.op == "output")
|
||||
old_outputs = list(output_node.args[0])
|
||||
|
||||
# Check if this is the forward subgraph of autograd_function_apply
|
||||
# For autograd_function_apply, the fwd subgraph must return (output, saved_values, ...)
|
||||
# where indices 0 and 1 are ALWAYS required by the runtime
|
||||
# is_autograd_fwd = any(
|
||||
# node.target == torch.ops.higher_order.autograd_function_apply
|
||||
# for node in hop_nodes
|
||||
# )
|
||||
is_autograd_fwd = False
|
||||
|
||||
# For autograd_function_apply forward subgraph, indices 0 (output) and 1 (saved_values)
|
||||
# are ALWAYS used by the runtime, even if not explicitly accessed via getitem
|
||||
if is_autograd_fwd and len(old_outputs) >= 2:
|
||||
used_indices.add(0) # output
|
||||
used_indices.add(1) # saved_values
|
||||
|
||||
# Nothing to DCE if all outputs are used or no outputs are used
|
||||
if len(used_indices) >= len(old_outputs) or len(used_indices) == 0:
|
||||
return False
|
||||
|
||||
# Build mapping from old indices to new indices
|
||||
old_to_new: dict[int, int] = {}
|
||||
new_outputs = []
|
||||
new_idx = 0
|
||||
|
||||
for old_idx in range(len(old_outputs)):
|
||||
if old_idx in used_indices:
|
||||
old_to_new[old_idx] = new_idx
|
||||
new_outputs.append(old_outputs[old_idx])
|
||||
new_idx += 1
|
||||
|
||||
# Update subgraph output node
|
||||
# Create a new output node with the filtered outputs
|
||||
with subgraph.graph.inserting_before(output_node):
|
||||
new_output_node = subgraph.graph.output(tuple(new_outputs))
|
||||
output_node.replace_all_uses_with(new_output_node)
|
||||
subgraph.graph.erase_node(output_node)
|
||||
|
||||
for parent_gm, _, hop_node in callers:
|
||||
# Update getitem nodes to use new indices
|
||||
for user in list(hop_node.users):
|
||||
if user.op == "call_function" and user.target == operator.getitem:
|
||||
old_idx = user.args[1]
|
||||
assert isinstance(old_idx, int)
|
||||
|
||||
if old_idx not in old_to_new:
|
||||
assert len(list(user.users)) == 0
|
||||
parent_gm.graph.erase_node(user)
|
||||
continue
|
||||
|
||||
new_idx = old_to_new[old_idx]
|
||||
# Create a new getitem node with the new index
|
||||
with parent_gm.graph.inserting_before(user):
|
||||
new_getitem = parent_gm.graph.call_function(
|
||||
operator.getitem, args=(user.args[0], new_idx)
|
||||
)
|
||||
# Copy metadata from old node
|
||||
new_getitem.meta = user.meta.copy()
|
||||
user.replace_all_uses_with(new_getitem)
|
||||
parent_gm.graph.erase_node(user)
|
||||
|
||||
# Update example_value metadata on hop_node
|
||||
if "example_value" in hop_node.meta:
|
||||
old_example = hop_node.meta["example_value"]
|
||||
assert isinstance(old_example, (tuple, list))
|
||||
new_example = tuple(
|
||||
old_example[old_idx]
|
||||
for old_idx in range(len(old_outputs))
|
||||
if old_idx in used_indices
|
||||
)
|
||||
hop_node.meta["example_value"] = new_example
|
||||
|
||||
# Recompile subgraph and all modified parent graphs
|
||||
subgraph.graph.lint()
|
||||
subgraph.recompile()
|
||||
|
||||
for parent_gm in {caller[0] for caller in callers}:
|
||||
parent_gm.graph.lint()
|
||||
parent_gm.recompile()
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,980 @@
|
||||
"""
|
||||
Debug utilities for TorchDynamo compilation and execution.
|
||||
|
||||
This module provides various debugging tools and utilities for TorchDynamo, including:
|
||||
|
||||
- Minification support for reducing test cases while preserving bugs
|
||||
- Input/output handling via InputReader and InputWriter for reproducible testing
|
||||
- Accuracy checking between original and compiled models
|
||||
- Neural network module string conversion via NNModuleToString
|
||||
- Profiling tools and system information collection
|
||||
- Buck build system integration for Meta-internal testing
|
||||
|
||||
Key classes:
|
||||
- InputReader/InputWriter: Handle serialization of model inputs/outputs
|
||||
- NNModuleToString: Converts nn.Modules to string representations
|
||||
- BuckTargetWriter: Manages Buck build system integration
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import copy
|
||||
import cProfile
|
||||
import functools
|
||||
import getpass
|
||||
import inspect
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
from collections import Counter
|
||||
from importlib import import_module
|
||||
from typing import Any, TYPE_CHECKING, TypeVar
|
||||
|
||||
import torch
|
||||
import torch._prims_common as utils
|
||||
import torch._subclasses.meta_utils
|
||||
from torch import Tensor
|
||||
from torch._dynamo.testing import rand_strided
|
||||
from torch._inductor.cpp_builder import normalize_path_separator
|
||||
from torch._prims_common import is_float_dtype
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils._content_store import ContentStoreReader, ContentStoreWriter
|
||||
|
||||
from . import config
|
||||
from .utils import clone_inputs, get_debug_dir, warn_once
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
from torch.hub import tqdm
|
||||
from torch.storage import UntypedStorage
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
inductor_config = import_module("torch._inductor.config")
|
||||
use_buck = inductor_config.is_fbcode()
|
||||
|
||||
if use_buck:
|
||||
import libfb.py.build_info
|
||||
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
extra_deps = []
|
||||
extra_imports = ""
|
||||
cur_target = ""
|
||||
if use_buck:
|
||||
extra_deps = [
|
||||
"//caffe2/torch/fb/sparsenn:sparsenn_operators_gpu",
|
||||
"//caffe2/torch/fb/sparsenn:sparsenn_operators",
|
||||
"//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu",
|
||||
"//deeplearning/fbgemm/fbgemm_gpu:sparse_ops",
|
||||
]
|
||||
cur_target = libfb.py.build_info.BuildInfo.get_build_rule().replace("fbcode:", "//") # type: ignore[possibly-undefined]
|
||||
extra_imports = "\n".join([f'torch.ops.load_library("{x}")' for x in extra_deps])
|
||||
|
||||
|
||||
BUCK_CMD_PREFIX = ["buck2", "run", "@mode/dev-nosan"]
|
||||
|
||||
|
||||
class BuckTargetWriter:
|
||||
def __init__(self, filename: str) -> None:
|
||||
self.subdir, self.py_file = os.path.split(os.path.abspath(filename))
|
||||
self.target = self.py_file.replace(".py", "")
|
||||
|
||||
# Get main_module path from fbcode
|
||||
self.path = f"{self.subdir.replace('/', '.')}.{self.target}"
|
||||
self.path = self.path[self.path.find("fbcode.") :]
|
||||
self.path = self.path[7:]
|
||||
|
||||
# Get cmd line path
|
||||
tmp = self.subdir
|
||||
tmp = tmp[tmp.find("fbcode/") :][7:]
|
||||
self.cmd_line_path = f"//{tmp}:{self.target}"
|
||||
|
||||
def build(self) -> str:
|
||||
extra_cpp_deps = "\n".join([f' "{x}",' for x in extra_deps])
|
||||
return textwrap.dedent(
|
||||
f"""
|
||||
load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary")
|
||||
|
||||
python_binary(
|
||||
name="{self.target}",
|
||||
srcs = ["{self.py_file}"],
|
||||
compile = False,
|
||||
deps = [
|
||||
"//caffe2:torch",
|
||||
"//caffe2:libtorch",
|
||||
"//caffe2/functorch:functorch",
|
||||
"//triton:triton",
|
||||
"{cur_target}",
|
||||
],
|
||||
cpp_deps = [
|
||||
{extra_cpp_deps}
|
||||
],
|
||||
main_module = "{self.path}",
|
||||
par_style = "xar",
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
def write(self, print_msg: bool = True) -> list[str]:
|
||||
target_file = os.path.join(self.subdir, "TARGETS")
|
||||
with open(target_file, "w") as fd:
|
||||
fd.write(self.build())
|
||||
# log.warning("Wrote isolation TARGETS file at %s", target_file)
|
||||
cmd_split = BUCK_CMD_PREFIX + [self.cmd_line_path]
|
||||
if print_msg:
|
||||
log.warning(
|
||||
"Found an example that reproduces the error. Run this cmd to repro - %s",
|
||||
" ".join(cmd_split),
|
||||
)
|
||||
return cmd_split
|
||||
|
||||
|
||||
def minifier_dir() -> str:
|
||||
path = os.path.join(get_debug_dir(), "minifier")
|
||||
if path is None:
|
||||
path = f"{tempfile.gettempdir()}/minifier_{getpass.getuser()}"
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
MAX_CONSTANT_NUMEL_INLINE = 4
|
||||
|
||||
|
||||
class NNModuleToString:
|
||||
safe_reprs = [
|
||||
torch.nn.Linear,
|
||||
torch.nn.Conv1d,
|
||||
torch.nn.Conv2d,
|
||||
torch.nn.Conv3d,
|
||||
torch.nn.BatchNorm1d,
|
||||
torch.nn.BatchNorm2d,
|
||||
torch.nn.BatchNorm3d,
|
||||
torch.nn.LayerNorm,
|
||||
torch.nn.Dropout,
|
||||
torch.nn.Softmax,
|
||||
torch.nn.ReLU,
|
||||
torch.nn.GELU,
|
||||
torch.nn.Identity,
|
||||
torch.nn.MaxPool2d,
|
||||
torch.nn.Embedding,
|
||||
torch.nn.Tanh,
|
||||
torch.nn.ConvTranspose1d,
|
||||
torch.nn.GLU,
|
||||
torch.nn.LSTM,
|
||||
torch.nn.Flatten,
|
||||
torch.nn.AdaptiveAvgPool2d,
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def can_convert_to_string(gm: torch.fx.GraphModule) -> bool:
|
||||
cant_convert = set()
|
||||
for _, module in gm.named_children():
|
||||
if type(module) not in NNModuleToString.safe_reprs:
|
||||
cant_convert.add(module)
|
||||
|
||||
if len(cant_convert) > 0:
|
||||
log.warning("We have not tested reprs of some modules - %s", cant_convert)
|
||||
# TODO - Assuming that all modules can be safely repr'd. Check if that assumption is correct.
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def convert(gm: torch.fx.GraphModule) -> str:
|
||||
from torch.nn.modules.module import _addindent
|
||||
|
||||
tab = " " * 4
|
||||
|
||||
model_str = textwrap.dedent(
|
||||
"""
|
||||
from torch.nn import *
|
||||
class Repro(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
"""
|
||||
)
|
||||
|
||||
for module_name, module in gm.named_children():
|
||||
module_str = f"{module.__repr__()}"
|
||||
# module should be a core torch.nn.Module, so all parameters
|
||||
# should be on the same device.
|
||||
example_param = next(module.parameters(), None)
|
||||
if example_param is not None and example_param.is_cuda:
|
||||
module_str = f"{module_str}.cuda()"
|
||||
model_str += f"{tab * 2}self.{module_name} = {module_str}\n"
|
||||
|
||||
for buffer_name, buffer in gm._buffers.items():
|
||||
if buffer is None:
|
||||
continue
|
||||
# Serialize full data for small buffers
|
||||
if buffer.numel() <= MAX_CONSTANT_NUMEL_INLINE:
|
||||
from torch._tensor_str import PRINT_OPTS
|
||||
|
||||
assert PRINT_OPTS.threshold >= MAX_CONSTANT_NUMEL_INLINE
|
||||
tensor_str = repr(buffer)
|
||||
elif torch.is_floating_point(buffer):
|
||||
tensor_str = f"torch.randn({list(buffer.shape)}, dtype={buffer.dtype})"
|
||||
else:
|
||||
tensor_str = (
|
||||
f"torch.randint(1, size={list(buffer.shape)}, dtype={buffer.dtype})"
|
||||
)
|
||||
if buffer.is_cuda:
|
||||
tensor_str = f"{tensor_str}.cuda()"
|
||||
model_str += (
|
||||
f"{tab * 2}self.register_buffer('{buffer_name}', {tensor_str})\n"
|
||||
)
|
||||
|
||||
for param_name, param in gm._parameters.items():
|
||||
if param is None:
|
||||
continue
|
||||
maybe_device = ""
|
||||
if param.is_cuda:
|
||||
maybe_device = ', device="cuda"'
|
||||
tensor_str = f"torch.nn.Parameter(torch.randn({list(param.shape)}, dtype={param.dtype}{maybe_device}))"
|
||||
model_str += f"{tab * 2}self.{param_name} = {tensor_str}\n"
|
||||
|
||||
# TODO - Keep this code for now. But, I don't think we will need this.
|
||||
# attrs = dir(gm)
|
||||
# for attr in attrs:
|
||||
# if "_tensor_constant" in attr:
|
||||
# val = getattr(gm, attr)
|
||||
# model_str += f" {attr} = {val!r}\n"
|
||||
|
||||
model_str += f"{_addindent(gm.code, 4)}\n"
|
||||
return model_str
|
||||
|
||||
|
||||
@functools.cache # subprocess is expensive
|
||||
def _cuda_system_info_comment() -> str:
|
||||
if not torch.cuda.is_available():
|
||||
return "# torch.cuda.is_available()==False, no GPU info collected\n"
|
||||
|
||||
model_str = "# CUDA Info: \n"
|
||||
try:
|
||||
if torch.version.hip is None:
|
||||
cuda_version_out = subprocess.check_output(["nvcc", "--version"])
|
||||
cuda_version_lines = cuda_version_out.decode().split("\n")
|
||||
comment = "".join([f"# {s} \n" for s in cuda_version_lines if s != ""])
|
||||
model_str += f"{comment}\n"
|
||||
else:
|
||||
model_str += "# Not searching for nvcc on ROCM setup\n"
|
||||
except (FileNotFoundError, subprocess.CalledProcessError):
|
||||
model_str += "# nvcc not found\n"
|
||||
|
||||
gpu_names = Counter(
|
||||
torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())
|
||||
)
|
||||
|
||||
model_str += "# GPU Hardware Info: \n"
|
||||
for name, count in gpu_names.items():
|
||||
model_str += f"# {name} : {count} \n"
|
||||
model_str += "\n"
|
||||
return model_str
|
||||
|
||||
|
||||
def generate_env_vars_string(*, stable_output: bool = False) -> str:
|
||||
"""
|
||||
Generate a string configuration for environment variables related to Dynamo, Inductor, and Triton.
|
||||
"""
|
||||
if stable_output:
|
||||
return "# env var omitted due to stable_output=True"
|
||||
|
||||
allow_list = ["TORCH", "DYNAMO", "INDUCTOR", "TRITON"]
|
||||
skip_list = ["TRITON_LIBDEVICE_PATH", "TRITON_PTXAS_PATH", "TRITON_LIBCUDA_PATH"]
|
||||
|
||||
def filter(key: str) -> bool:
|
||||
return any(string in key for string in allow_list) and key not in skip_list
|
||||
|
||||
config_lines = [
|
||||
f"""os.environ['{key}'] = '{value.replace("'", '"')}'"""
|
||||
for key, value in os.environ.items()
|
||||
if filter(key)
|
||||
]
|
||||
config_string = "\n".join(config_lines)
|
||||
return normalize_path_separator(f"""\
|
||||
import os
|
||||
{config_string}
|
||||
""")
|
||||
|
||||
|
||||
def generate_config_string(*, stable_output: bool = False) -> str:
|
||||
import torch._functorch.config
|
||||
import torch._inductor.config
|
||||
|
||||
if stable_output:
|
||||
return "# config omitted due to stable_output=True"
|
||||
|
||||
experimental_config = torch.fx.experimental._config.codegen_config() # type: ignore[attr-defined]
|
||||
return f"""\
|
||||
import torch._dynamo.config
|
||||
import torch._inductor.config
|
||||
import torch._functorch.config
|
||||
import torch.fx.experimental._config
|
||||
{torch._dynamo.config.codegen_config()}
|
||||
{torch._inductor.config.codegen_config()}
|
||||
{torch._functorch.config.codegen_config()}
|
||||
{experimental_config}
|
||||
"""
|
||||
|
||||
|
||||
def get_minifier_repro_path() -> str:
|
||||
return os.path.join(minifier_dir(), "minifier_launcher.py")
|
||||
|
||||
|
||||
def helper_for_dump_minify(contents: str) -> None:
|
||||
minified_repro_path = get_minifier_repro_path()
|
||||
log.warning("Writing minified repro to:\n%s", minified_repro_path)
|
||||
|
||||
if use_buck:
|
||||
BuckTargetWriter(minified_repro_path).write()
|
||||
try:
|
||||
with open(minified_repro_path, "w") as fd:
|
||||
fd.write(contents)
|
||||
|
||||
except OSError as e:
|
||||
log.exception("")
|
||||
raise NotImplementedError(f"Could not write to {minified_repro_path}") from e
|
||||
|
||||
|
||||
class AccuracyError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def clone_inputs_retaining_gradness(example_inputs: Sequence[Any]) -> list[Any]:
|
||||
"""
|
||||
This clone inputs is different from utils clone_input. In case of minifier,
|
||||
all the tensors are leaf tensors while creating a new graph. So, we set the
|
||||
requires_grad field w/o checking the leafness of the tensor.
|
||||
"""
|
||||
cloned_inputs = clone_inputs(example_inputs)
|
||||
for idx in range(len(example_inputs)):
|
||||
if isinstance(cloned_inputs[idx], torch.Tensor):
|
||||
cloned_inputs[idx].requires_grad_(example_inputs[idx].requires_grad)
|
||||
return cloned_inputs # type: ignore[return-value]
|
||||
|
||||
|
||||
def run_fwd_maybe_bwd(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
only_fwd: bool = False,
|
||||
disable_clone: bool = False,
|
||||
) -> Any:
|
||||
"""
|
||||
Runs a forward and possibly backward iteration for a given mod and args.
|
||||
|
||||
When disable_clone is True, we will use args as-is without cloning.
|
||||
This is higher fidelity but we may destroy the args in the process.
|
||||
"""
|
||||
from .testing import collect_results, reduce_to_scalar_loss, requires_bwd_pass
|
||||
|
||||
gm = copy.deepcopy(gm)
|
||||
if not disable_clone:
|
||||
args = clone_inputs_retaining_gradness(args)
|
||||
|
||||
if hasattr(gm, "zero_grad"):
|
||||
gm.zero_grad(True)
|
||||
|
||||
# TorchInductor returned callable expects lists. So, may need a boxed calling convention.
|
||||
out = gm(args) if getattr(gm, "_boxed_call", False) else gm(*args)
|
||||
|
||||
if only_fwd:
|
||||
return out
|
||||
if requires_bwd_pass(out):
|
||||
loss = reduce_to_scalar_loss(out)
|
||||
loss.backward()
|
||||
return collect_results(gm, out, None, args)
|
||||
|
||||
|
||||
def same_two_models(
|
||||
gm: torch.fx.GraphModule,
|
||||
opt_gm: torch.fx.GraphModule,
|
||||
example_inputs: Sequence[Any],
|
||||
only_fwd: bool = False,
|
||||
*,
|
||||
require_fp64: bool = False,
|
||||
ignore_non_fp: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Check two models have same accuracy.
|
||||
|
||||
require_fp64: if True, raise an error if we unable to calculate the fp64 reference
|
||||
ignore_non_fp: if True, do not compare outputs which are not floating point. This
|
||||
is mostly useful for the minifier (which wants to avoid quantizing floating point
|
||||
error into integer/boolean error)
|
||||
"""
|
||||
from .utils import same
|
||||
|
||||
ref = run_fwd_maybe_bwd(gm, example_inputs, only_fwd)
|
||||
|
||||
fp64_ref = None
|
||||
if config.same_two_models_use_fp64:
|
||||
try:
|
||||
fp64_model, fp64_examples = cast_to_fp64(
|
||||
copy.deepcopy(gm), clone_inputs_retaining_gradness(example_inputs)
|
||||
)
|
||||
fp64_ref = run_fwd_maybe_bwd(fp64_model, fp64_examples, only_fwd)
|
||||
except Exception:
|
||||
if require_fp64:
|
||||
raise RuntimeError( # noqa: B904
|
||||
"Could not generate fp64 outputs, workaround with torch._dynamo.config.same_two_models_use_fp64 = False"
|
||||
)
|
||||
log.warning("Could not generate fp64 outputs")
|
||||
|
||||
try:
|
||||
res = run_fwd_maybe_bwd(opt_gm, example_inputs, only_fwd)
|
||||
except Exception:
|
||||
# This means that the minified graph is bad/exposes a different problem.
|
||||
# As we are checking accuracy here, lets log the exception and return True.
|
||||
log.exception(
|
||||
"While minifying the program in accuracy minification mode, "
|
||||
"ran into a runtime exception which is likely an unrelated issue."
|
||||
" Skipping this graph."
|
||||
)
|
||||
return True
|
||||
|
||||
passing = same(
|
||||
ref,
|
||||
res,
|
||||
fp64_ref,
|
||||
tol=config.repro_tolerance,
|
||||
equal_nan=True,
|
||||
ignore_non_fp=ignore_non_fp,
|
||||
)
|
||||
return passing
|
||||
|
||||
|
||||
def cast_dtype_args_to_fp64(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
|
||||
for node in model.graph.nodes:
|
||||
if (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.prims.convert_element_type.default
|
||||
):
|
||||
assert len(node.args) == 2
|
||||
if is_float_dtype(node.args[1]) and node.args[1] != torch.float64:
|
||||
node.args = (node.args[0], torch.float64)
|
||||
if node.op == "call_function":
|
||||
dtype = node.kwargs.get("dtype")
|
||||
if dtype is not None and is_float_dtype(dtype):
|
||||
new_kwargs = dict(node.kwargs)
|
||||
new_kwargs["dtype"] = torch.float64
|
||||
node.kwargs = new_kwargs
|
||||
|
||||
model.graph.lint()
|
||||
model.recompile()
|
||||
return model
|
||||
|
||||
|
||||
def cast_to(
|
||||
dtype: torch.dtype, model: torch.fx.GraphModule, inputs: list[Any]
|
||||
) -> tuple[torch.fx.GraphModule, list[Any]]:
|
||||
from torch.utils._pytree import tree_map
|
||||
|
||||
model = model.to(dtype)
|
||||
if dtype == torch.float64:
|
||||
# If casting to fp64 for accuracy comparison, we need to
|
||||
# replace dtype arguments embedded in the graph with fp64
|
||||
model = cast_dtype_args_to_fp64(model)
|
||||
|
||||
inputs = tree_map(
|
||||
lambda x: x.to(dtype)
|
||||
if isinstance(x, torch.Tensor) and x.is_floating_point()
|
||||
else x,
|
||||
inputs,
|
||||
)
|
||||
return model, inputs
|
||||
|
||||
|
||||
def cast_to_fp64(
|
||||
model: torch.fx.GraphModule, inputs: list[Any]
|
||||
) -> tuple[torch.fx.GraphModule, list[Any]]:
|
||||
return cast_to(torch.float64, model, inputs)
|
||||
|
||||
|
||||
def backend_accuracy_fails(
|
||||
gm: torch.fx.GraphModule,
|
||||
example_inputs: Sequence[Any],
|
||||
compiler_fn: Callable[[torch.fx.GraphModule, list[Any]], torch.fx.GraphModule],
|
||||
only_fwd: bool = False,
|
||||
*,
|
||||
require_fp64: bool = False,
|
||||
ignore_non_fp: bool = False,
|
||||
) -> bool:
|
||||
try:
|
||||
compiled_gm = compiler_fn(
|
||||
copy.deepcopy(gm), clone_inputs_retaining_gradness(example_inputs)
|
||||
)
|
||||
return not same_two_models(
|
||||
gm,
|
||||
compiled_gm,
|
||||
example_inputs,
|
||||
only_fwd,
|
||||
require_fp64=require_fp64,
|
||||
ignore_non_fp=ignore_non_fp,
|
||||
)
|
||||
except Exception:
|
||||
# This means that the minified graph is bad/exposes a different problem.
|
||||
# As we are checking accuracy here, lets log the exception and return False.
|
||||
log.exception(
|
||||
"While minifying the program in accuracy minification mode, "
|
||||
"ran into a runtime exception which is likely an unrelated issue."
|
||||
" Skipping this graph"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# REPRO SUPPORT CODE
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
# Helper functions for computing what the default values of tensor
|
||||
# values should be. These all coincide with factory functions, e.g., torch.empty
|
||||
|
||||
|
||||
def _stride_or_default(
|
||||
stride: torch._prims_common.StrideType | None,
|
||||
*,
|
||||
shape: torch._prims_common.ShapeType,
|
||||
) -> torch._prims_common.StrideType:
|
||||
return stride if stride is not None else utils.make_contiguous_strides_for(shape)
|
||||
|
||||
|
||||
def _mk_defaulter(d: T) -> Callable[[T | None], T]:
|
||||
return lambda x: x if x is not None else d
|
||||
|
||||
|
||||
_dtype_or_default = _mk_defaulter(torch.float32)
|
||||
_device_or_default = _mk_defaulter(torch.device("cpu"))
|
||||
_storage_offset_or_default = _mk_defaulter(0)
|
||||
_requires_grad_or_default = _mk_defaulter(False)
|
||||
_is_leaf_or_default = _mk_defaulter(False)
|
||||
|
||||
|
||||
class NopInputReader:
|
||||
def __init__(self) -> None:
|
||||
self.total = 0
|
||||
|
||||
def storage(
|
||||
self,
|
||||
storage_hash: str | None,
|
||||
nbytes: int,
|
||||
*,
|
||||
device: torch._prims_common.DeviceLikeType | None = None,
|
||||
dtype_hint: torch.dtype | None = None,
|
||||
) -> None:
|
||||
self.total += 1
|
||||
|
||||
def tensor(self, *args: Any, **kwargs: Any) -> torch.Tensor | None:
|
||||
pass
|
||||
|
||||
def symint(self, *args: Any, **kwargs: Any) -> int | None:
|
||||
pass
|
||||
|
||||
def const(self, name: str) -> None:
|
||||
pass
|
||||
|
||||
def unsupported(self, name: str) -> None:
|
||||
pass
|
||||
|
||||
def generator(self, device_type: str, device_index: int) -> None:
|
||||
pass
|
||||
|
||||
def opaque(self, script_class_name: str) -> None:
|
||||
self.total += 1
|
||||
|
||||
|
||||
# TODO: Support bundling the entire repro into a zip file for ease of
|
||||
# transferring around
|
||||
class InputReader:
|
||||
def __init__(
|
||||
self, save_dir: str | None = None, *, pbar: tqdm | None = None
|
||||
) -> None:
|
||||
# If None, we will generate random data instead. It's important
|
||||
# to natively support this use case as it will allow people to
|
||||
# share repros without including the real data, if the problem
|
||||
# reproduces even on random data.
|
||||
if save_dir is None:
|
||||
log.warning("no save_dir specified, will generate random data")
|
||||
self.store = ContentStoreReader(save_dir) if save_dir is not None else None
|
||||
self.args: list[Any] = []
|
||||
self.pbar = pbar
|
||||
|
||||
def storage(
|
||||
self,
|
||||
storage_hash: str | None,
|
||||
nbytes: int,
|
||||
*,
|
||||
device: torch._prims_common.DeviceLikeType | None = None,
|
||||
dtype_hint: torch.dtype | None = None,
|
||||
) -> UntypedStorage:
|
||||
if self.pbar is not None:
|
||||
self.pbar.update(1)
|
||||
device = _device_or_default(device) # type: ignore[arg-type]
|
||||
dtype_hint = _dtype_or_default(dtype_hint)
|
||||
if self.store is not None and storage_hash is not None:
|
||||
try:
|
||||
storage = self.store.read_storage(storage_hash)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
else:
|
||||
if device != storage.device:
|
||||
log.warning("device mismatch: %s != %s", device, storage.device)
|
||||
# TODO: transfer it to the right device? But failing this
|
||||
# way would be very mysterious! Would have been better
|
||||
# not to store device in the serialized format...
|
||||
return storage
|
||||
warn_once(f"could not load {storage_hash}, generating random data instead")
|
||||
shape = (nbytes // dtype_hint.itemsize,)
|
||||
stride = _stride_or_default(None, shape=shape)
|
||||
return rand_strided(shape, stride, dtype_hint, device).untyped_storage()
|
||||
|
||||
def tensor(
|
||||
self,
|
||||
storage: UntypedStorage,
|
||||
shape: torch._prims_common.ShapeType,
|
||||
stride: torch._prims_common.StrideType | None = None,
|
||||
*,
|
||||
storage_offset: int | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
requires_grad: bool | None = None,
|
||||
is_leaf: bool | None = None,
|
||||
**metadata: Any,
|
||||
) -> torch.Tensor:
|
||||
stride = _stride_or_default(stride, shape=shape)
|
||||
storage_offset = _storage_offset_or_default(storage_offset)
|
||||
dtype = _dtype_or_default(dtype)
|
||||
is_leaf = _is_leaf_or_default(is_leaf)
|
||||
requires_grad = _requires_grad_or_default(requires_grad)
|
||||
t = torch.tensor(
|
||||
[], dtype=dtype, device=storage.device, requires_grad=requires_grad
|
||||
)
|
||||
with torch.no_grad():
|
||||
t.set_(storage, storage_offset, shape, stride)
|
||||
if not is_leaf:
|
||||
# Fake up some autograd history in a very naughty way
|
||||
with torch.enable_grad():
|
||||
t = t.clone(memory_format=torch.preserve_format)
|
||||
with torch.no_grad():
|
||||
t.set_(storage, storage_offset, shape, stride)
|
||||
assert torch._subclasses.meta_utils.safe_is_leaf(t) == is_leaf
|
||||
torch._utils.set_tensor_metadata(t, metadata)
|
||||
self.args.append(t)
|
||||
return t # for BC
|
||||
|
||||
def symint(self, val: Any) -> Any:
|
||||
self.args.append(val)
|
||||
return val # for BC
|
||||
|
||||
def const(self, name: str) -> None:
|
||||
self.args.append(None)
|
||||
|
||||
def unsupported(self, name: str) -> None:
|
||||
self.args.append(None)
|
||||
|
||||
def generator(self, device_type: str, device_index: int) -> torch._C.Generator:
|
||||
gen = torch.cuda.default_generators[device_index].clone_state()
|
||||
self.args.append(gen)
|
||||
return gen
|
||||
|
||||
def opaque(self, script_class_name: str) -> None:
|
||||
self.args.append(None)
|
||||
|
||||
|
||||
# Here is our writer strategy:
|
||||
# 1. We will stream all of the inputs to disk
|
||||
# 2. You can now deterministically randomize the inputs, or reload
|
||||
# the inputs from disk
|
||||
# 3. You can YOLO run the script without the inputs, in which case
|
||||
# we'll fill the inputs with random data and pray. This is the
|
||||
# legacy behavior, but it's also useful if you want to find out
|
||||
# if we're so broken even random inputs trigger it
|
||||
# 4. We could offer an in process "check if the randomized thing
|
||||
# works too" but this is delicate so we don't do it
|
||||
|
||||
|
||||
class InputWriter:
|
||||
def __init__(self, save_dir: str | None, *, stable_hash: bool = False) -> None:
|
||||
self._lines: list[str] = []
|
||||
# TODO: consider ensuring tensor and storage counters line up?
|
||||
self.storage_counter = itertools.count()
|
||||
self.save_dir = save_dir
|
||||
self.store = (
|
||||
ContentStoreWriter(save_dir, stable_hash=stable_hash)
|
||||
if save_dir is not None
|
||||
else None
|
||||
)
|
||||
self.seen_storages: dict[StorageWeakRef, str] = {}
|
||||
|
||||
def lines(self) -> list[str]:
|
||||
r = [
|
||||
"def load_args(reader):",
|
||||
]
|
||||
r.extend(f" {l}" for l in self._lines)
|
||||
# In case we need to change the internal format of load_args
|
||||
# in an FC-breaking way
|
||||
r.append("load_args._version = 0")
|
||||
return r
|
||||
|
||||
# Storages are untyped, but we need to initialize them with data if
|
||||
# we don't have the real data, so we give a hint saying what kind
|
||||
# of initialization may be appropriate
|
||||
#
|
||||
# If we had a FakeTensor, device_hint tells us what device should be
|
||||
def storage(
|
||||
self,
|
||||
untyped_storage: UntypedStorage,
|
||||
*,
|
||||
device_hint: torch._prims_common.DeviceLikeType | None = None,
|
||||
dtype_hint: torch.dtype | None = None,
|
||||
) -> str:
|
||||
ws = StorageWeakRef(untyped_storage)
|
||||
v = self.seen_storages.get(ws)
|
||||
if v is not None:
|
||||
return v
|
||||
v = f"buf{next(self.storage_counter)}"
|
||||
maybe_dtype_hint = ""
|
||||
if _dtype_or_default(None) != _dtype_or_default(dtype_hint):
|
||||
maybe_dtype_hint = f", dtype_hint={dtype_hint!r}"
|
||||
# TODO: being optional on device is kind of pointless as the default
|
||||
# is CPU but most repros we care about are CUDA
|
||||
maybe_device = ""
|
||||
device = untyped_storage.device
|
||||
if device.type == "meta":
|
||||
assert device_hint is not None
|
||||
device = device_hint # type: ignore[assignment]
|
||||
if _device_or_default(None) != device:
|
||||
maybe_device = f", device={device!r}"
|
||||
nbytes = untyped_storage.nbytes()
|
||||
storage_hash = None
|
||||
if self.store is not None and untyped_storage.device.type != "meta":
|
||||
storage_hash = self.store.write_storage(untyped_storage)
|
||||
self._lines.append(
|
||||
f"{v} = reader.storage({storage_hash!r}, {nbytes!r}{maybe_device}{maybe_dtype_hint})"
|
||||
)
|
||||
self.seen_storages[ws] = v
|
||||
return v
|
||||
|
||||
def tensor(self, name: str, t: torch.Tensor) -> None:
|
||||
from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq
|
||||
|
||||
storage = self.storage(
|
||||
t.untyped_storage(), dtype_hint=t.dtype, device_hint=t.device
|
||||
)
|
||||
args = []
|
||||
# NB: this is positional, must come first
|
||||
if not statically_known_true(
|
||||
sym_eq(_stride_or_default(None, shape=t.shape), t.stride())
|
||||
):
|
||||
args.append(str(tuple(t.stride())))
|
||||
if _dtype_or_default(None) != t.dtype:
|
||||
args.append(f"dtype={t.dtype!r}")
|
||||
if not statically_known_true(
|
||||
_storage_offset_or_default(None) == t.storage_offset()
|
||||
):
|
||||
args.append(f"storage_offset={t.storage_offset()!r}")
|
||||
tensor_metadata = torch._utils.get_tensor_metadata(t)
|
||||
if tensor_metadata:
|
||||
args.extend(f"{k}={v!r}" for k, v in tensor_metadata.items())
|
||||
if _requires_grad_or_default(None) != t.requires_grad:
|
||||
args.append(f"requires_grad={t.requires_grad!r}")
|
||||
is_leaf = torch._subclasses.meta_utils.safe_is_leaf(t)
|
||||
if _is_leaf_or_default(None) != is_leaf:
|
||||
args.append(f"is_leaf={is_leaf!r}")
|
||||
self._lines.append(
|
||||
"reader.tensor("
|
||||
+ ", ".join([storage, str(tuple(t.shape)), *args])
|
||||
+ f") # {name}"
|
||||
)
|
||||
|
||||
def unsupported(self, name: str, arg: Any) -> None:
|
||||
# NB: Try hard not to /print/ a tensor, that will be very slow
|
||||
self._lines.append(
|
||||
f"reader.unsupported({name!r}) # unsupported type for dumping: {type(arg)}"
|
||||
)
|
||||
# Best effort dump as much useful stuff we can lol, in case you want
|
||||
# to repair the repro
|
||||
if isinstance(arg, (list, tuple)):
|
||||
self._lines.append('"""')
|
||||
for i, a in enumerate(arg):
|
||||
name_i = f"{name}[{i}]"
|
||||
if isinstance(a, torch.Tensor):
|
||||
self.tensor(name_i, a)
|
||||
elif isinstance(a, (int, torch.SymInt)):
|
||||
self.symint(name_i, a)
|
||||
else:
|
||||
self.unsupported(name_i, a)
|
||||
self._lines.append('"""')
|
||||
|
||||
# write out that the arg was filtered out as it is constant
|
||||
def const(self, name: str) -> None:
|
||||
self._lines.append(
|
||||
f"reader.const({name!r}) # {name}, filtered out during compilation"
|
||||
)
|
||||
|
||||
# TODO: this doesn't actually symint atm
|
||||
def symint(self, name: str, val: Any) -> None:
|
||||
if isinstance(val, torch.SymInt):
|
||||
val = val.node.hint
|
||||
self._lines.append(f"reader.symint({val!r}) # {name}")
|
||||
|
||||
def generator(self, name: str, arg: torch._C.Generator) -> None:
|
||||
device = arg.device
|
||||
self._lines.append(
|
||||
f"reader.generator({device.type!r}, {device.index!r}) # {name}"
|
||||
)
|
||||
|
||||
def opaque(self, name: str, script_class_name: str) -> None:
|
||||
self._lines.append(f"reader.opaque({script_class_name!r}) # {name}")
|
||||
|
||||
|
||||
def aot_graph_input_parser(
|
||||
func: Callable[[list[Tensor]], list[Tensor]],
|
||||
device: str = "cuda",
|
||||
sym_shapes: dict[str, int] | None = None,
|
||||
default_sym_shape: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Takes in a function which has been printed with print_readable() and constructs kwargs to run it.
|
||||
|
||||
Handles Tensor inputs, Symints, and a graph module which might have tensor constants.
|
||||
|
||||
Consider a function `forward` defined as follows:
|
||||
|
||||
def forward(self, primals_1: "f32[1001, 6]", primals_2: "f32[s0]", primals_3: "Sym(s0)",):
|
||||
_tensor_constant0: "i64[4190]" = self._tensor_constant0
|
||||
# Further implementation
|
||||
|
||||
kwargs = aot_graph_input_parser(forward)
|
||||
forward(**kwargs)
|
||||
"""
|
||||
|
||||
from torch.utils._dtype_abbrs import dtype_abbrs
|
||||
|
||||
dtype_map: dict[str, torch.dtype] = {
|
||||
value: key for key, value in dtype_abbrs.items()
|
||||
}
|
||||
dtype_pattern: str = "|".join(dtype_abbrs.values())
|
||||
|
||||
# Extracting the source code from the function
|
||||
source = inspect.getsource(func)
|
||||
|
||||
# Regular expressions
|
||||
tensor_assignment_regex = rf"(_tensor_constant\d+): \"({dtype_pattern})\[\s*(.*?)\s*\]\" = self\.(_tensor_constant\d+)"
|
||||
tensor_regex = rf"({dtype_pattern})\[\s*(.*?)\s*\]"
|
||||
sym_shape_regex = r"Sym\((s\d+)\)"
|
||||
|
||||
class TensorContainer:
|
||||
"Container for tensors as attributes"
|
||||
|
||||
# Dictionary for tensors from annotations
|
||||
kwargs: dict[str, Any] = {}
|
||||
|
||||
sym_shapes_dict: dict[str, int] = sym_shapes or {}
|
||||
|
||||
def get_sym_int(symint: str) -> int:
|
||||
torch._check(
|
||||
symint in sym_shapes_dict or default_sym_shape is not None,
|
||||
lambda: f"{symint} not in symbolic_shapes and default sym shape not passed in",
|
||||
)
|
||||
return sym_shapes_dict.get(symint, default_sym_shape) # type: ignore[return-value]
|
||||
|
||||
def gen_tensor(shape: torch._prims_common.ShapeType, dtype: torch.dtype) -> Tensor:
|
||||
# Resolve symbolic shapes to concrete values
|
||||
resolved_shape = []
|
||||
dynamic_dims = []
|
||||
for i, dim in enumerate(shape):
|
||||
dim = dim.strip() # type: ignore[attr-defined]
|
||||
if "s" in dim:
|
||||
s = get_sym_int(dim)
|
||||
resolved_shape.append(s)
|
||||
dynamic_dims.append(i)
|
||||
else:
|
||||
if dim:
|
||||
resolved_shape.append(int(dim))
|
||||
|
||||
constructor = torch.randn if dtype.is_floating_point else torch.zeros
|
||||
out = constructor(resolved_shape, dtype=dtype, device=device) # type: ignore[call-arg]
|
||||
for d in dynamic_dims:
|
||||
torch._dynamo.mark_dynamic(out, d)
|
||||
return out
|
||||
|
||||
# Parse function annotations for tensor generation
|
||||
annotations = func.__annotations__
|
||||
for param, annotation in annotations.items():
|
||||
# Skip 'return' annotation
|
||||
if param == "return":
|
||||
continue
|
||||
|
||||
match = re.search(tensor_regex, annotation)
|
||||
if match:
|
||||
data_type, shape_str = match.groups()
|
||||
shape = tuple(shape_str.split(","))
|
||||
dtype = dtype_map[data_type]
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
kwargs[param] = gen_tensor(shape, dtype)
|
||||
|
||||
match = re.search(sym_shape_regex, annotation)
|
||||
if match:
|
||||
kwargs[param] = get_sym_int(match.group(1))
|
||||
|
||||
if "self" in inspect.signature(func).parameters:
|
||||
container = TensorContainer()
|
||||
kwargs["self"] = container
|
||||
for match in re.finditer(tensor_assignment_regex, source):
|
||||
attr_name, data_type, shape_str, _ = match.groups()
|
||||
shape = tuple(shape_str.split(","))
|
||||
dtype = dtype_map[data_type]
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
setattr(container, attr_name, gen_tensor(shape, dtype))
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
def profile_to_file(filename: str) -> Callable[[T], T]:
|
||||
"""
|
||||
Decorator to cProfile a given function and save the result to disk on process exit.
|
||||
|
||||
Args:
|
||||
filename: filename to save profile to
|
||||
"""
|
||||
prof = cProfile.Profile()
|
||||
filename = os.path.abspath(os.path.expanduser(filename))
|
||||
|
||||
def decorator(fn: Any) -> Any:
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
prof.enable()
|
||||
try:
|
||||
return fn(*args, **kwargs)
|
||||
finally:
|
||||
prof.disable()
|
||||
|
||||
return wrapper
|
||||
|
||||
def save_it() -> None:
|
||||
prof.dump_stats(filename)
|
||||
sys.stderr.write(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
Wrote profile to {filename}, view with:
|
||||
|
||||
snakeviz {filename}
|
||||
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
atexit.register(save_it)
|
||||
return decorator
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,653 @@
|
||||
"""
|
||||
Device abstraction layer for TorchDynamo and Inductor backends.
|
||||
|
||||
This module provides a unified interface for different hardware backends (CUDA, XPU,
|
||||
CPU, MPS, MTIA) through a common device interface. Key components include:
|
||||
|
||||
- DeviceInterface: Base class defining the common API for all device types
|
||||
- Device-specific implementations: CudaInterface, XpuInterface, CpuInterface, MpsInterface, MtiaInterface
|
||||
- Device registration system for managing available backends
|
||||
- Worker APIs for multi-processing scenarios
|
||||
- Stream and event management across different devices
|
||||
- Device property caching for worker processes
|
||||
|
||||
The abstraction layer enables device-agnostic code in TorchDynamo while allowing
|
||||
specialized implementations for each hardware backend's unique features.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import time
|
||||
from collections import namedtuple
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
from torch.utils._pallas import has_torch_tpu
|
||||
|
||||
|
||||
get_cuda_stream: Callable[[int], int] | None
|
||||
if torch.cuda._is_compiled():
|
||||
from torch._C import _cuda_getCurrentRawStream as get_cuda_stream
|
||||
else:
|
||||
get_cuda_stream = None
|
||||
|
||||
# Recording the device properties in the main process but used in worker process.
|
||||
caching_worker_device_properties: dict[str, Any] = {}
|
||||
caching_worker_current_devices: dict[str, int] = {}
|
||||
|
||||
|
||||
class DeviceInterface:
|
||||
"""
|
||||
This is a simple device runtime interface for Inductor. It enables custom
|
||||
backends to be integrated with Inductor in a device-agnostic semantic.
|
||||
"""
|
||||
|
||||
class device:
|
||||
def __new__(cls, device: torch.types.Device) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
class Event:
|
||||
def __new__(cls, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError(
|
||||
"Event should be inherited from torch.Event, otherwise, it couldn't be captured by dynamo."
|
||||
)
|
||||
|
||||
class Stream:
|
||||
def __new__(cls, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError(
|
||||
"Stream should be inherited from torch.Stream, otherwise, it couldn't be captured by dynamo."
|
||||
)
|
||||
|
||||
class Worker:
|
||||
"""
|
||||
Worker API to query device properties that will work in multi processing
|
||||
workers that cannot use the GPU APIs (due to processing fork() and
|
||||
initialization time issues). Properties are recorded in the main process
|
||||
before we fork the workers.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def set_device(device: int) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_device_properties(device: torch.types.Device = None) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def set_device(device: torch.types.Device) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def maybe_exchange_device(device: int) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def exchange_device(device: int) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def device_count() -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def stream(stream: torch.Stream) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def current_stream() -> torch.Stream:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def set_stream(stream: torch.Stream) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def _set_stream_by_id(stream_id: int, device_index: int, device_type: int) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def get_raw_stream(device_idx: int) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def synchronize(device: torch.types.Device = None) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def get_device_properties(cls, device: torch.types.Device = None) -> Any:
|
||||
return cls.Worker.get_device_properties(device)
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def is_bf16_supported(including_emulation: bool = False) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def is_dtype_supported(
|
||||
cls, dtype: torch.dtype, including_emulation: bool = False
|
||||
) -> bool:
|
||||
return dtype != torch.bfloat16 or cls.is_bf16_supported(including_emulation)
|
||||
|
||||
@staticmethod
|
||||
def memory_allocated(device: torch.types.Device = None) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def is_triton_capable(device: torch.types.Device = None) -> bool:
|
||||
"""
|
||||
Returns True if the device has Triton support, False otherwise, even if
|
||||
the appropriate Triton backend is not available.
|
||||
"""
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def raise_if_triton_unavailable(cls, device: torch.types.Device = None) -> None:
|
||||
"""
|
||||
Raises a `RuntimeError` with the appropriate human-readable instructions
|
||||
to resolve the issue if Triton is not available for the given device, or
|
||||
the default device if `device` is `None`.
|
||||
|
||||
The caller should ensure the presence of the 'triton' package before
|
||||
calling this method.
|
||||
"""
|
||||
if not cls.is_triton_capable():
|
||||
raise RuntimeError("This device is not capable of supporting Triton")
|
||||
|
||||
|
||||
class DeviceGuard:
|
||||
"""
|
||||
This class provides a context manager for device switching. This is a stripped
|
||||
down version of torch.{device_name}.device.
|
||||
|
||||
The context manager changes the current device to the given device index
|
||||
on entering the context and restores the original device on exiting.
|
||||
The device is switched using the provided device interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, device_interface: type[DeviceInterface], index: int | None
|
||||
) -> None:
|
||||
self.device_interface = device_interface
|
||||
self.idx = index
|
||||
self.prev_idx = -1
|
||||
|
||||
def __enter__(self) -> None:
|
||||
if self.idx is not None:
|
||||
self.prev_idx = self.device_interface.exchange_device(self.idx)
|
||||
|
||||
def __exit__(self, type: Any, value: Any, traceback: Any) -> Literal[False]:
|
||||
if self.idx is not None:
|
||||
self.idx = self.device_interface.maybe_exchange_device(self.prev_idx)
|
||||
return False
|
||||
|
||||
|
||||
class CudaInterface(DeviceInterface):
|
||||
device = torch.cuda.device # type: ignore[assignment]
|
||||
|
||||
# register Event and Stream class into the backend interface
|
||||
# make sure Event and Stream are implemented and inherited from the torch.Event and torch.Stream
|
||||
Event = torch.cuda.Event # type: ignore[assignment]
|
||||
Stream = torch.cuda.Stream # type: ignore[assignment]
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Worker:
|
||||
@staticmethod
|
||||
def set_device(device: int) -> None:
|
||||
caching_worker_current_devices["cuda"] = device
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
if "cuda" in caching_worker_current_devices:
|
||||
return caching_worker_current_devices["cuda"]
|
||||
return torch.cuda.current_device()
|
||||
|
||||
@staticmethod
|
||||
def get_device_properties(device: torch.types.Device = None) -> Any:
|
||||
if device is not None:
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
assert device.type == "cuda"
|
||||
if isinstance(device, torch.device):
|
||||
device = device.index
|
||||
if device is None:
|
||||
device = CudaInterface.Worker.current_device()
|
||||
|
||||
if "cuda" not in caching_worker_device_properties:
|
||||
device_prop = [
|
||||
torch.cuda.get_device_properties(i)
|
||||
for i in range(torch.cuda.device_count())
|
||||
]
|
||||
caching_worker_device_properties["cuda"] = device_prop
|
||||
|
||||
return caching_worker_device_properties["cuda"][device]
|
||||
|
||||
current_device = staticmethod(torch.cuda.current_device)
|
||||
set_device = staticmethod(torch.cuda.set_device)
|
||||
device_count = staticmethod(torch.cuda.device_count)
|
||||
stream = staticmethod(torch.cuda.stream) # type: ignore[assignment]
|
||||
current_stream = staticmethod(torch.cuda.current_stream)
|
||||
set_stream = staticmethod(torch.cuda.set_stream) # type: ignore[assignment]
|
||||
_set_stream_by_id = staticmethod(torch.cuda._set_stream_by_id) # type: ignore[assignment]
|
||||
synchronize = staticmethod(torch.cuda.synchronize)
|
||||
get_device_properties = staticmethod(torch.cuda.get_device_properties) # type: ignore[assignment]
|
||||
get_raw_stream = staticmethod(get_cuda_stream) # type: ignore[assignment, arg-type]
|
||||
exchange_device = staticmethod(torch.cuda._exchange_device) # type: ignore[arg-type, has-type]
|
||||
maybe_exchange_device = staticmethod(torch.cuda._maybe_exchange_device) # type: ignore[arg-type, has-type]
|
||||
memory_allocated = staticmethod(torch.cuda.memory_allocated)
|
||||
is_bf16_supported = staticmethod(torch.cuda.is_bf16_supported) # type: ignore[arg-type]
|
||||
|
||||
# Can be mock patched by @patch decorator.
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
return torch.cuda.is_available()
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> int | str:
|
||||
if torch.version.hip is None:
|
||||
major, min = torch.cuda.get_device_capability(device)
|
||||
return major * 10 + min
|
||||
else:
|
||||
return torch.cuda.get_device_properties(device).gcnArchName.split(":", 1)[0]
|
||||
|
||||
@staticmethod
|
||||
def is_triton_capable(device: torch.types.Device = None) -> bool:
|
||||
return (
|
||||
torch.version.hip is not None
|
||||
or torch.cuda.get_device_properties(device).major >= 7
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
|
||||
from torch._inductor.exc import GPUTooOldForTriton
|
||||
|
||||
if not CudaInterface.is_triton_capable(device):
|
||||
device_props = torch.cuda.get_device_properties(device)
|
||||
raise GPUTooOldForTriton(device_props, inspect.currentframe())
|
||||
|
||||
import triton.backends
|
||||
|
||||
if torch.version.hip is not None:
|
||||
if "amd" not in triton.backends.backends:
|
||||
raise RuntimeError("triton not built with the 'amd' backend")
|
||||
elif "nvidia" not in triton.backends.backends:
|
||||
raise RuntimeError("triton not built with the 'nvidia' backend")
|
||||
|
||||
|
||||
get_mtia_stream: Callable[[int], int] | None
|
||||
if torch.mtia._is_compiled():
|
||||
from torch._C import _mtia_getCurrentRawStream as get_mtia_stream
|
||||
else:
|
||||
get_mtia_stream = None
|
||||
|
||||
|
||||
class MtiaInterface(DeviceInterface):
|
||||
device = torch.mtia.device # type: ignore[assignment]
|
||||
Event = torch.mtia.Event # type: ignore[assignment]
|
||||
Stream = torch.mtia.Stream # type: ignore[assignment]
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Worker:
|
||||
@staticmethod
|
||||
def set_device(device: int) -> None:
|
||||
caching_worker_current_devices["mtia"] = device
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
if "mtia" in caching_worker_current_devices:
|
||||
return caching_worker_current_devices["mtia"]
|
||||
return torch.mtia.current_device()
|
||||
|
||||
@staticmethod
|
||||
def get_device_properties(device: torch.types.Device = None) -> Any:
|
||||
if device is not None:
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
assert device.type == "mtia"
|
||||
if isinstance(device, torch.device):
|
||||
device = device.index
|
||||
if device is None:
|
||||
device = MtiaInterface.Worker.current_device()
|
||||
|
||||
if "mtia" not in caching_worker_device_properties:
|
||||
device_prop = [
|
||||
torch.mtia.get_device_properties(i)
|
||||
for i in range(torch.mtia.device_count())
|
||||
]
|
||||
caching_worker_device_properties["mtia"] = device_prop
|
||||
|
||||
return caching_worker_device_properties["mtia"][device]
|
||||
|
||||
current_device = staticmethod(torch.mtia.current_device)
|
||||
set_device = staticmethod(torch.mtia.set_device) # type: ignore[assignment]
|
||||
device_count = staticmethod(torch.mtia.device_count)
|
||||
stream = staticmethod(torch.mtia.stream) # type: ignore[assignment]
|
||||
current_stream = staticmethod(torch.mtia.current_stream)
|
||||
set_stream = staticmethod(torch.mtia.set_stream) # type: ignore[assignment]
|
||||
_set_stream_by_id = staticmethod(torch.mtia._set_stream_by_id) # type: ignore[assignment]
|
||||
synchronize = staticmethod(torch.mtia.synchronize)
|
||||
get_device_properties = staticmethod(torch.mtia.get_device_properties) # type: ignore[assignment]
|
||||
get_raw_stream = staticmethod(get_mtia_stream) # type: ignore[assignment, arg-type]
|
||||
exchange_device = staticmethod(torch.mtia._exchange_device) # type: ignore[arg-type, has-type]
|
||||
maybe_exchange_device = staticmethod(torch.mtia._maybe_exchange_device) # type: ignore[arg-type, has-type]
|
||||
memory_allocated = staticmethod(torch.mtia.memory_allocated) # type: ignore[assignment]
|
||||
is_bf16_supported = staticmethod(torch.mtia.is_bf16_supported) # type: ignore[arg-type]
|
||||
|
||||
# Can be mock patched by @patch decorator.
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
ret = torch.mtia.is_available()
|
||||
return ret
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> Any:
|
||||
cc = torch.mtia.get_device_capability(device)
|
||||
return cc
|
||||
|
||||
@staticmethod
|
||||
def is_triton_capable(device: torch.types.Device = None) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
|
||||
import triton.backends
|
||||
|
||||
if "mtia" not in triton.backends.backends:
|
||||
raise RuntimeError("triton not built with the 'mtia' backend")
|
||||
|
||||
|
||||
get_xpu_stream: Callable[[int], int] | None
|
||||
if torch.xpu._is_compiled():
|
||||
from torch._C import _xpu_getCurrentRawStream as get_xpu_stream
|
||||
else:
|
||||
get_xpu_stream = None
|
||||
|
||||
|
||||
class XpuInterface(DeviceInterface):
|
||||
device = torch.xpu.device # type: ignore[assignment]
|
||||
Event = torch.xpu.Event # type: ignore[assignment]
|
||||
Stream = torch.xpu.Stream # type: ignore[assignment]
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Worker:
|
||||
@staticmethod
|
||||
def set_device(device: int) -> None:
|
||||
caching_worker_current_devices["xpu"] = device
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
if "xpu" in caching_worker_current_devices:
|
||||
return caching_worker_current_devices["xpu"]
|
||||
return torch.xpu.current_device()
|
||||
|
||||
@staticmethod
|
||||
def get_device_properties(device: torch.types.Device = None) -> Any:
|
||||
if device is not None:
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
assert device.type == "xpu"
|
||||
if isinstance(device, torch.device):
|
||||
device = device.index
|
||||
if device is None:
|
||||
device = XpuInterface.Worker.current_device()
|
||||
|
||||
if "xpu" not in caching_worker_device_properties:
|
||||
device_prop = [
|
||||
torch.xpu.get_device_properties(i)
|
||||
for i in range(torch.xpu.device_count())
|
||||
]
|
||||
caching_worker_device_properties["xpu"] = device_prop
|
||||
|
||||
return caching_worker_device_properties["xpu"][device]
|
||||
|
||||
current_device = staticmethod(torch.xpu.current_device)
|
||||
set_device = staticmethod(torch.xpu.set_device)
|
||||
device_count = staticmethod(torch.xpu.device_count) # type: ignore[has-type]
|
||||
stream = staticmethod(torch.xpu.stream) # type: ignore[assignment]
|
||||
current_stream = staticmethod(torch.xpu.current_stream)
|
||||
set_stream = staticmethod(torch.xpu.set_stream) # type: ignore[assignment]
|
||||
_set_stream_by_id = staticmethod(torch.xpu._set_stream_by_id) # type: ignore[assignment]
|
||||
synchronize = staticmethod(torch.xpu.synchronize)
|
||||
get_device_properties = staticmethod(torch.xpu.get_device_properties) # type: ignore[assignment]
|
||||
get_raw_stream = staticmethod(get_xpu_stream) # type: ignore[assignment, arg-type]
|
||||
exchange_device = staticmethod(torch.xpu._exchange_device) # type: ignore[arg-type, has-type]
|
||||
maybe_exchange_device = staticmethod(torch.xpu._maybe_exchange_device) # type: ignore[arg-type, has-type]
|
||||
memory_allocated = staticmethod(torch.xpu.memory_allocated)
|
||||
|
||||
# Can be mock patched by @patch decorator.
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
return torch.xpu.is_available()
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> Any:
|
||||
cc = torch.xpu.get_device_capability(device)
|
||||
return cc
|
||||
|
||||
@staticmethod
|
||||
def is_bf16_supported(including_emulation: bool = False) -> bool:
|
||||
return torch.xpu.is_bf16_supported()
|
||||
|
||||
@staticmethod
|
||||
def is_triton_capable(device: torch.types.Device = None) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
|
||||
import triton.backends
|
||||
|
||||
if "intel" not in triton.backends.backends:
|
||||
raise RuntimeError("triton not built with the 'intel' backend")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CpuDeviceProperties:
|
||||
multi_processor_count: int
|
||||
|
||||
|
||||
class CpuInterface(DeviceInterface):
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Event(torch.Event):
|
||||
def __init__(self, enable_timing: bool = True) -> None:
|
||||
self.time = 0.0
|
||||
|
||||
def elapsed_time(self, other: Any) -> float:
|
||||
return (other.time - self.time) * 1000
|
||||
|
||||
def record(self, stream: Any = None) -> None:
|
||||
self.time = time.perf_counter()
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Worker:
|
||||
@staticmethod
|
||||
def get_device_properties(
|
||||
device: torch.types.Device = None,
|
||||
) -> CpuDeviceProperties:
|
||||
import multiprocessing
|
||||
|
||||
cpu_count = multiprocessing.cpu_count()
|
||||
return CpuDeviceProperties(cpu_count)
|
||||
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def is_bf16_supported(including_emulation: bool = False) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> str:
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def get_raw_stream(device_idx: Any) -> int:
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def synchronize(device: torch.types.Device = None) -> None:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def is_triton_capable(device: torch.types.Device = None) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def raise_if_triton_unavailable(device: torch.types.Device = None) -> None:
|
||||
import triton.backends
|
||||
|
||||
if "cpu" not in triton.backends.backends:
|
||||
raise RuntimeError("triton not built with the 'cpu' backend")
|
||||
|
||||
|
||||
class MpsInterface(DeviceInterface):
|
||||
@staticmethod
|
||||
def is_bf16_supported(including_emulation: bool = False) -> bool:
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def is_dtype_supported(
|
||||
cls, dtype: torch.dtype, including_emulation: bool = False
|
||||
) -> bool:
|
||||
if dtype in [torch.float64, torch.complex128]:
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
return torch.backends.mps.is_available()
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> str:
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def synchronize(device: torch.types.Device = None) -> None:
|
||||
torch.mps.synchronize()
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Worker:
|
||||
@staticmethod
|
||||
def get_device_properties(device: torch.types.Device = None) -> Any:
|
||||
return namedtuple("MPSProperties", ["multi_processor_count"])(
|
||||
torch.backends.mps.get_core_count() # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
return 0
|
||||
|
||||
|
||||
class TpuInterface(DeviceInterface):
|
||||
@staticmethod
|
||||
def is_bf16_supported(including_emulation: bool = False) -> bool:
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def is_dtype_supported(
|
||||
cls, dtype: torch.dtype, including_emulation: bool = False
|
||||
) -> bool:
|
||||
return dtype not in (
|
||||
torch.float64,
|
||||
torch.complex32,
|
||||
torch.complex64,
|
||||
torch.complex128,
|
||||
torch.half,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
return has_torch_tpu()
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def get_compute_capability(device: torch.types.Device = None) -> str:
|
||||
return ""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
class Worker:
|
||||
@staticmethod
|
||||
def get_device_properties(device: torch.types.Device = None) -> Any:
|
||||
return namedtuple("TPUProperties", ["multi_processor_count"])(
|
||||
1 # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def current_device() -> int:
|
||||
return 0
|
||||
|
||||
|
||||
device_interfaces: dict[str, type[DeviceInterface]] = {}
|
||||
_device_initialized = False
|
||||
|
||||
|
||||
def register_interface_for_device(
|
||||
device: str | torch.device, device_interface: type[DeviceInterface]
|
||||
) -> None:
|
||||
if isinstance(device, torch.device):
|
||||
device = device.type
|
||||
device_interfaces[device] = device_interface
|
||||
|
||||
|
||||
def get_interface_for_device(device: str | torch.device) -> type[DeviceInterface]:
|
||||
if isinstance(device, torch.device):
|
||||
device = device.type
|
||||
if not _device_initialized:
|
||||
init_device_reg()
|
||||
if device in device_interfaces:
|
||||
return device_interfaces[device]
|
||||
raise NotImplementedError(f"No interface for device {device}")
|
||||
|
||||
|
||||
def get_registered_device_interfaces() -> Iterable[tuple[str, type[DeviceInterface]]]:
|
||||
if not _device_initialized:
|
||||
init_device_reg()
|
||||
return device_interfaces.items()
|
||||
|
||||
|
||||
def init_device_reg() -> None:
|
||||
global _device_initialized
|
||||
register_interface_for_device("cuda", CudaInterface)
|
||||
for i in range(torch.cuda.device_count()):
|
||||
register_interface_for_device(f"cuda:{i}", CudaInterface)
|
||||
|
||||
register_interface_for_device("xpu", XpuInterface)
|
||||
for i in range(torch.xpu.device_count()):
|
||||
register_interface_for_device(f"xpu:{i}", XpuInterface)
|
||||
|
||||
register_interface_for_device("mtia", MtiaInterface)
|
||||
for i in range(torch.mtia.device_count()):
|
||||
register_interface_for_device(f"mtia:{i}", MtiaInterface)
|
||||
|
||||
register_interface_for_device("cpu", CpuInterface)
|
||||
register_interface_for_device("mps", MpsInterface)
|
||||
register_interface_for_device("tpu", TpuInterface)
|
||||
|
||||
_device_initialized = True
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
Manages process groups for distributed compilation in TorchDynamo.
|
||||
|
||||
This module handles the initialization and management of process groups used for
|
||||
distributed compilation. Key features:
|
||||
|
||||
- Lazy initialization of compilation process groups
|
||||
- Only creates groups when distributed mode is enabled and available
|
||||
- Integrates with compiler_collectives configuration setting
|
||||
- Provides a single global process group for compilation coordination
|
||||
|
||||
The process group is created only when needed and if the distributed environment
|
||||
is properly initialized, making it safe to import and use this module even in
|
||||
non-distributed scenarios.
|
||||
"""
|
||||
|
||||
import torch.distributed as dist
|
||||
|
||||
from . import config
|
||||
|
||||
|
||||
_COMPILE_PG: dist.ProcessGroup | None = None
|
||||
_GUARD_PG: dist.ProcessGroup | None = None
|
||||
|
||||
|
||||
def get_compile_pg() -> dist.ProcessGroup | None:
|
||||
if (
|
||||
config.enable_compiler_collectives
|
||||
and dist.is_available()
|
||||
and dist.is_initialized()
|
||||
):
|
||||
global _COMPILE_PG
|
||||
if _COMPILE_PG is None:
|
||||
# , timeout=datetime.timedelta(seconds=2)
|
||||
_COMPILE_PG = dist.distributed_c10d._new_group_with_tag(
|
||||
pg_tag="pt2_compile_pg"
|
||||
)
|
||||
return _COMPILE_PG
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# NB: Unlike get_compile_pg, this is only called when guard collectives were
|
||||
# explicitly requested
|
||||
def get_guard_pg() -> dist.ProcessGroup | None:
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
global _GUARD_PG
|
||||
if _GUARD_PG is None:
|
||||
_GUARD_PG = dist.distributed_c10d._new_group_with_tag(pg_tag="pt2_guard_pg")
|
||||
return _GUARD_PG
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,428 @@
|
||||
"""
|
||||
Dynamo Profiler - tracks where Dynamo spends time during compilation.
|
||||
|
||||
This module provides profiling functionality for Dynamo tracing, showing per-function
|
||||
cumtime (inclusive) and tottime (exclusive) in a cProfile-compatible format.
|
||||
The output can be visualized with tools like snakeviz.
|
||||
|
||||
Usage:
|
||||
# Enable via config (prints pstats output):
|
||||
torch._dynamo.config.dynamo_profiler = True
|
||||
|
||||
# Or save to file for snakeviz:
|
||||
torch._dynamo.config.dynamo_profiler = "/tmp/dynamo.prof"
|
||||
# Then: snakeviz /tmp/dynamo.prof
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pstats
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionTraceTiming:
|
||||
"""
|
||||
Timing data for a single inlined function trace.
|
||||
|
||||
Follows cProfile conventions:
|
||||
- cumtime: total time in function including all subcalls (inclusive)
|
||||
- tottime: time in function excluding subcalls (exclusive)
|
||||
- caller info: who called this function (for building call graph)
|
||||
"""
|
||||
|
||||
# Function identification
|
||||
func_name: str
|
||||
filename: str
|
||||
firstlineno: int
|
||||
# Timing data (in nanoseconds) - cProfile-style
|
||||
cumtime_ns: int # Inclusive time (includes subcalls)
|
||||
tottime_ns: int # Exclusive time (excludes subcalls)
|
||||
# Code stats (for comparing tracing overhead vs function complexity)
|
||||
bytecode_count: int
|
||||
# Nesting depth when this function was traced
|
||||
inline_depth: int
|
||||
# Caller information (for building call graph edges)
|
||||
caller_func_name: str | None = None
|
||||
caller_filename: str | None = None
|
||||
caller_firstlineno: int | None = None
|
||||
# Whether this is a primitive (non-recursive) call
|
||||
# A call is primitive if the function doesn't appear anywhere in the call stack
|
||||
is_primitive_call: bool = True
|
||||
# Full call stack at the time of this call (for proper snakeviz drill-down)
|
||||
# Each entry is (func_name, filename, firstlineno)
|
||||
call_stack: tuple[tuple[str, str, int], ...] = ()
|
||||
|
||||
# Backwards compatibility alias
|
||||
@property
|
||||
def trace_time_ns(self) -> int:
|
||||
return self.cumtime_ns
|
||||
|
||||
@property
|
||||
def trace_time_ms(self) -> float:
|
||||
return self.cumtime_ns / 1e6
|
||||
|
||||
@property
|
||||
def cumtime_ms(self) -> float:
|
||||
return self.cumtime_ns / 1e6
|
||||
|
||||
@property
|
||||
def tottime_ms(self) -> float:
|
||||
return self.tottime_ns / 1e6
|
||||
|
||||
@property
|
||||
def caller_key(self) -> tuple[str, int, str] | None:
|
||||
"""Return caller as a pstats-compatible key tuple."""
|
||||
if self.caller_func_name is not None:
|
||||
return (
|
||||
self.caller_filename or "",
|
||||
self.caller_firstlineno or 0,
|
||||
self.caller_func_name,
|
||||
)
|
||||
return None
|
||||
|
||||
@property
|
||||
def func_key(self) -> tuple[str, int, str]:
|
||||
"""Return this function as a pstats-compatible key tuple."""
|
||||
return (self.filename, self.firstlineno, self.func_name)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"FunctionTraceTiming({self.func_name} at {self.filename}:{self.firstlineno}, "
|
||||
f"cumtime={self.cumtime_ms:.2f}ms, tottime={self.tottime_ms:.2f}ms, "
|
||||
f"bytecode={self.bytecode_count}, depth={self.inline_depth})"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProfilerStackEntry:
|
||||
"""Stack entry for tracking function timing in the Dynamo profiler."""
|
||||
|
||||
func_name: str
|
||||
filename: str
|
||||
firstlineno: int
|
||||
start_time_ns: int
|
||||
child_time_ns: int # Accumulated time spent in traced children
|
||||
is_primitive_call: bool = True # Whether this is a non-recursive call
|
||||
|
||||
|
||||
class DynamoProfilerState:
|
||||
"""State for Dynamo profiler tracking function trace timings."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.timings: list[FunctionTraceTiming] = []
|
||||
self.stack: list[ProfilerStackEntry] = []
|
||||
|
||||
def record_timing(self, timing: FunctionTraceTiming) -> None:
|
||||
"""Record timing data for a traced function."""
|
||||
self.timings.append(timing)
|
||||
|
||||
def get_timings(self) -> list[FunctionTraceTiming]:
|
||||
"""Get all recorded timings."""
|
||||
return self.timings
|
||||
|
||||
def push(
|
||||
self, func_name: str, filename: str, firstlineno: int, start_time_ns: int
|
||||
) -> None:
|
||||
"""Push a new entry onto the timing stack."""
|
||||
# Check if this function already exists in the stack (indirect recursion)
|
||||
is_primitive = not any(
|
||||
entry.func_name == func_name
|
||||
and entry.filename == filename
|
||||
and entry.firstlineno == firstlineno
|
||||
for entry in self.stack
|
||||
)
|
||||
self.stack.append(
|
||||
ProfilerStackEntry(
|
||||
func_name=func_name,
|
||||
filename=filename,
|
||||
firstlineno=firstlineno,
|
||||
start_time_ns=start_time_ns,
|
||||
child_time_ns=0,
|
||||
is_primitive_call=is_primitive,
|
||||
)
|
||||
)
|
||||
|
||||
def pop(self) -> ProfilerStackEntry | None:
|
||||
"""Pop the top entry from the timing stack."""
|
||||
if self.stack:
|
||||
return self.stack.pop()
|
||||
return None
|
||||
|
||||
def add_child_time(self, child_cumtime_ns: int) -> None:
|
||||
"""Add the child's cumulative time to the parent's child_time accumulator."""
|
||||
if self.stack:
|
||||
self.stack[-1].child_time_ns += child_cumtime_ns
|
||||
|
||||
def get_current_caller(self) -> tuple[str, str, int] | None:
|
||||
"""Get the current caller (top of stack) as (func_name, filename, firstlineno)."""
|
||||
if self.stack:
|
||||
entry = self.stack[-1]
|
||||
return (entry.func_name, entry.filename, entry.firstlineno)
|
||||
return None
|
||||
|
||||
def get_call_stack(self) -> tuple[tuple[str, str, int], ...]:
|
||||
"""Get the full current call stack as tuple of (func_name, filename, firstlineno)."""
|
||||
return tuple(
|
||||
(entry.func_name, entry.filename, entry.firstlineno) for entry in self.stack
|
||||
)
|
||||
|
||||
def generate_pstats(
|
||||
self, output_file: str | None = None, print_raw: bool = False
|
||||
) -> pstats.Stats:
|
||||
"""Generate pstats.Stats object from recorded timings.
|
||||
|
||||
Args:
|
||||
output_file: Optional file path to save the stats.
|
||||
print_raw: If True, print raw aggregate timings before returning.
|
||||
"""
|
||||
import cProfile
|
||||
import io
|
||||
import logging
|
||||
import pstats
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Aggregate by (filename, lineno, func_name)
|
||||
aggregated: dict[tuple[str, int, str], dict[str, Any]] = {}
|
||||
# caller_edges[callee_key][caller_key] -> edge stats
|
||||
caller_edges: dict[
|
||||
tuple[str, int, str], dict[tuple[str, int, str], dict[str, Any]]
|
||||
] = {}
|
||||
|
||||
for t in self.timings:
|
||||
key = (t.filename, t.firstlineno, t.func_name)
|
||||
|
||||
if key not in aggregated:
|
||||
aggregated[key] = {
|
||||
"ncalls": 0,
|
||||
"pcalls": 0,
|
||||
"tottime": 0.0,
|
||||
"cumtime": 0.0,
|
||||
}
|
||||
caller_edges[key] = {}
|
||||
|
||||
agg = aggregated[key]
|
||||
agg["ncalls"] += 1
|
||||
agg["tottime"] += t.tottime_ns / 1e9
|
||||
|
||||
if t.is_primitive_call:
|
||||
agg["pcalls"] += 1
|
||||
agg["cumtime"] += t.cumtime_ns / 1e9
|
||||
|
||||
# Build caller edge
|
||||
if t.caller_filename is not None:
|
||||
caller_key = (
|
||||
t.caller_filename,
|
||||
t.caller_firstlineno or 0,
|
||||
t.caller_func_name or "",
|
||||
)
|
||||
if caller_key not in caller_edges[key]:
|
||||
caller_edges[key][caller_key] = {
|
||||
"ncalls": 0,
|
||||
"pcalls": 0,
|
||||
"tottime": 0.0,
|
||||
"cumtime": 0.0,
|
||||
}
|
||||
edge = caller_edges[key][caller_key]
|
||||
edge["ncalls"] += 1
|
||||
edge["tottime"] += t.tottime_ns / 1e9
|
||||
# Always add cumtime to edges for visualization (gprof2dot)
|
||||
# Function-level cumtime is already correct (only primitive calls)
|
||||
edge["cumtime"] += t.cumtime_ns / 1e9
|
||||
if t.is_primitive_call:
|
||||
edge["pcalls"] += 1
|
||||
|
||||
if print_raw:
|
||||
sorted_items = sorted(
|
||||
aggregated.items(), key=lambda x: x[1]["cumtime"], reverse=True
|
||||
)
|
||||
print("\n=== Aggregate Timings (raw) ===")
|
||||
print(
|
||||
f"{'ncalls':>8} {'pcalls':>8} {'tottime':>12} {'cumtime':>12} function"
|
||||
)
|
||||
print("-" * 80)
|
||||
total_cumtime = 0.0
|
||||
total_tottime = 0.0
|
||||
for (filename, lineno, func_name), agg in sorted_items:
|
||||
ncalls = agg["ncalls"]
|
||||
pcalls = agg["pcalls"]
|
||||
tottime = agg["tottime"] * 1000 # convert to ms
|
||||
cumtime = agg["cumtime"] * 1000
|
||||
total_cumtime += cumtime
|
||||
total_tottime += tottime
|
||||
short_file = filename.split("/")[-1] if "/" in filename else filename
|
||||
print(
|
||||
f"{ncalls:>8} {pcalls:>8} {tottime:>10.2f}ms {cumtime:>10.2f}ms "
|
||||
f"{func_name} ({short_file}:{lineno})"
|
||||
)
|
||||
print("-" * 80)
|
||||
print(
|
||||
f"Total timings: {len(self.timings)}, unique functions: {len(aggregated)}"
|
||||
)
|
||||
print(
|
||||
f"Sum tottime: {total_tottime:.2f}ms, Sum cumtime: {total_cumtime:.2f}ms"
|
||||
)
|
||||
|
||||
# Ensure caller-only functions have a top-level entry.
|
||||
# gprof2dot expects every function referenced as a caller to also
|
||||
# exist as a top-level entry in the stats dict with timing data.
|
||||
for key in list(caller_edges.keys()):
|
||||
for caller_key in caller_edges[key]:
|
||||
if caller_key not in aggregated:
|
||||
aggregated[caller_key] = {
|
||||
"ncalls": 0,
|
||||
"pcalls": 0,
|
||||
"tottime": 0.0,
|
||||
"cumtime": 0.0,
|
||||
}
|
||||
caller_edges[caller_key] = {}
|
||||
|
||||
# Build the stats dict in pstats format
|
||||
stats_dict: dict[
|
||||
tuple[str, int, str], tuple[int, int, float, float, dict[Any, Any]]
|
||||
] = {}
|
||||
|
||||
for key, agg in aggregated.items():
|
||||
callers: dict[tuple[str, int, str], tuple[int, int, float, float]] = {}
|
||||
for caller_key, edge in caller_edges[key].items():
|
||||
callers[caller_key] = (
|
||||
edge["ncalls"],
|
||||
edge["pcalls"],
|
||||
edge["tottime"],
|
||||
edge["cumtime"],
|
||||
)
|
||||
|
||||
stats_dict[key] = (
|
||||
agg["pcalls"],
|
||||
agg["ncalls"],
|
||||
agg["tottime"],
|
||||
agg["cumtime"],
|
||||
callers,
|
||||
)
|
||||
|
||||
# Create a pstats.Stats object
|
||||
dummy_profile = cProfile.Profile()
|
||||
dummy_profile.enable()
|
||||
dummy_profile.disable()
|
||||
stats = pstats.Stats(dummy_profile, stream=io.StringIO())
|
||||
|
||||
stats.stats = stats_dict # type: ignore[attr-defined]
|
||||
stats.total_calls = sum(s[1] for s in stats_dict.values()) # type: ignore[attr-defined]
|
||||
stats.prim_calls = sum(s[0] for s in stats_dict.values()) # type: ignore[attr-defined]
|
||||
stats.total_tt = sum(s[2] for s in stats_dict.values()) # type: ignore[attr-defined]
|
||||
|
||||
if output_file:
|
||||
stats.dump_stats(output_file)
|
||||
log.info(
|
||||
"Saved pstats to %s. Visualize with: snakeviz %s",
|
||||
output_file,
|
||||
output_file,
|
||||
)
|
||||
|
||||
return stats
|
||||
|
||||
def generate_svg(
|
||||
self, profile_file: str, svg_file: str | None = None
|
||||
) -> str | None:
|
||||
"""Generate an SVG call graph from a profile file using gprof2dot and graphviz.
|
||||
|
||||
Args:
|
||||
profile_file: Path to the pstats profile file.
|
||||
svg_file: Optional path for the output SVG. If not provided, uses
|
||||
profile_file with .svg extension.
|
||||
|
||||
Returns:
|
||||
Path to the generated SVG file, or None if generation failed.
|
||||
"""
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
if not shutil.which("gprof2dot"):
|
||||
print("gprof2dot not found. Install with: pip install gprof2dot")
|
||||
return None
|
||||
|
||||
if not shutil.which("dot"):
|
||||
print("graphviz 'dot' not found. Install graphviz package.")
|
||||
return None
|
||||
|
||||
if svg_file is None:
|
||||
svg_file = profile_file.rsplit(".", 1)[0] + ".svg"
|
||||
|
||||
try:
|
||||
# gprof2dot -f pstats profile.prof | dot -Tsvg -o profile.svg
|
||||
gprof2dot = subprocess.Popen(
|
||||
[
|
||||
"gprof2dot",
|
||||
"-f",
|
||||
"pstats",
|
||||
"--node-label=total-time-percentage",
|
||||
"--node-label=self-time-percentage",
|
||||
"--node-label=total-time",
|
||||
profile_file,
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
dot = subprocess.Popen(
|
||||
["dot", "-Tsvg", "-o", svg_file],
|
||||
stdin=gprof2dot.stdout,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
gprof2dot.stdout.close() # type: ignore[union-attr]
|
||||
_, dot_err = dot.communicate()
|
||||
_, gprof2dot_err = gprof2dot.communicate()
|
||||
|
||||
if gprof2dot.returncode != 0:
|
||||
print(
|
||||
f"gprof2dot failed: {gprof2dot_err.decode()}" # noqa: B950
|
||||
)
|
||||
return None
|
||||
|
||||
if dot.returncode != 0:
|
||||
print(f"graphviz dot failed: {dot_err.decode()}")
|
||||
return None
|
||||
|
||||
if not os.path.isfile(svg_file):
|
||||
print(f"SVG file was not created: {svg_file}")
|
||||
return None
|
||||
|
||||
print(f"SVG call graph saved to: {svg_file}")
|
||||
return svg_file
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to generate SVG: {e}")
|
||||
return None
|
||||
|
||||
def dump_stats(
|
||||
self, output_file: str | None = None, generate_svg: bool = True
|
||||
) -> None:
|
||||
"""Print profiler stats to stdout and optionally save to file.
|
||||
|
||||
Args:
|
||||
output_file: Optional path to save the pstats profile.
|
||||
generate_svg: If True and output_file is provided, also generate an SVG
|
||||
call graph using gprof2dot and graphviz.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if not self.timings:
|
||||
return
|
||||
|
||||
stats = self.generate_pstats(output_file, print_raw=True)
|
||||
print("\n=== Dynamo Profiler (pstats) ===")
|
||||
stats.stream = sys.stdout # type: ignore[attr-defined]
|
||||
stats.sort_stats("cumulative").print_stats()
|
||||
|
||||
if output_file:
|
||||
print(f"\nProfile saved to: {output_file}")
|
||||
print(f"Visualize with: snakeviz {output_file}")
|
||||
|
||||
if generate_svg:
|
||||
self.generate_svg(output_file)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,901 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
"""Exception handling and error reporting for TorchDynamo.
|
||||
|
||||
This module provides a comprehensive set of exception classes and utilities for error
|
||||
handling in TorchDynamo. It includes:
|
||||
|
||||
Base Exceptions:
|
||||
- TorchDynamoException: Base class for all TorchDynamo-specific exceptions
|
||||
- Various specialized subclasses for different error scenarios
|
||||
|
||||
User Error Handling:
|
||||
- UserError: Exceptions for user-facing errors in TorchDynamo usage
|
||||
- UserErrorType: Enumeration of different categories of user errors
|
||||
- Formatted error messages with debugging information
|
||||
|
||||
Observed Exceptions:
|
||||
- Classes for handling exceptions observed during tracing
|
||||
- Special handling for StopIteration, LookupError, etc.
|
||||
- Exception state management during compilation
|
||||
|
||||
Error Formatting:
|
||||
- Stack trace filtering and formatting
|
||||
- Error message augmentation
|
||||
- Debugging utilities for error reporting
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import textwrap
|
||||
import typing
|
||||
from enum import auto, Enum
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from traceback import extract_stack, format_exc, format_list, FrameSummary, StackSummary
|
||||
from typing import Any, NoReturn, TYPE_CHECKING
|
||||
|
||||
import torch._guards
|
||||
from torch._utils_internal import get_file_path_2
|
||||
|
||||
from . import config
|
||||
from .utils import counters
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import types
|
||||
|
||||
from torch._dynamo.variables import VariableTracker
|
||||
from torch._guards import CompileId
|
||||
|
||||
from .output_graph import DynamoTracerOutput
|
||||
from .symbolic_convert import InstructionTranslatorBase
|
||||
from .types import DynamoFrameType, FrameExecStrategy
|
||||
|
||||
|
||||
def exportdb_error_message(case_name: str) -> str:
|
||||
return (
|
||||
"For more information about this error, see: "
|
||||
+ "https://pytorch.org/docs/main/generated/exportdb/index.html#"
|
||||
+ case_name.replace("_", "-")
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
graph_breaks_log = torch._logging.getArtifactLogger(__name__, "graph_breaks")
|
||||
|
||||
|
||||
class TorchDynamoException(RuntimeError):
|
||||
"""Base exception class for all TorchDynamo-specific exceptions.
|
||||
|
||||
Attributes:
|
||||
_torch_dynamo_tracer_output: Optional tracer output attached to the exception
|
||||
frame_exec_strategy: Optional frame execution strategy to control how convert_frame
|
||||
should handle this exception. When set, convert_frame will use this strategy
|
||||
instead of the default behavior. This allows exceptions to signal specific
|
||||
execution strategies (e.g., SKIP, RUN_ONLY) without requiring separate
|
||||
exception types for control flow.
|
||||
"""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._torch_dynamo_tracer_output: DynamoTracerOutput | None = None
|
||||
self.frame_exec_strategy: FrameExecStrategy | None = None
|
||||
|
||||
|
||||
class InternalTorchDynamoError(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class ResumePrologueTracingError(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class RestartAnalysis(TorchDynamoException):
|
||||
restart_reason: str | None
|
||||
|
||||
def __init__(self, *args: Any, restart_reason: str | None = None) -> None:
|
||||
self.restart_reason = restart_reason
|
||||
super().__init__(*args)
|
||||
|
||||
|
||||
class SpeculationRestartAnalysis(RestartAnalysis):
|
||||
pass
|
||||
|
||||
|
||||
class AutogradGradRestartAnalysis(RestartAnalysis):
|
||||
"""Raised when autograd.grad consumed grad_fns that are returned.
|
||||
|
||||
On restart, autograd.grad will graph break instead of being traced.
|
||||
"""
|
||||
|
||||
|
||||
class RequiresGradRestartAnalysis(RestartAnalysis):
|
||||
"""Raised when a source-less requires_grad_() intermediate leaks as output.
|
||||
|
||||
On restart, requires_grad_() will graph break instead of being traced,
|
||||
preserving partial acceleration for code before the call.
|
||||
"""
|
||||
|
||||
|
||||
class UnspecializeRestartAnalysis(RestartAnalysis):
|
||||
pass
|
||||
|
||||
|
||||
class CompileCollectiveRestartAnalysis(RestartAnalysis):
|
||||
pass
|
||||
|
||||
|
||||
class TensorifyScalarRestartAnalysis(RestartAnalysis):
|
||||
pass
|
||||
|
||||
|
||||
# Used (primarily for backends) to skip tracing the current frame
|
||||
# and all future invocations of it.
|
||||
# NOTE: this does NOT cause a graph break, and thus no graph break messages
|
||||
# will be issued!
|
||||
class SkipFrame(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class TorchRuntimeError(TorchDynamoException):
|
||||
def __init__(self, msg: str, real_stack: StackSummary | None = None) -> None:
|
||||
super().__init__(msg)
|
||||
self.msg = msg
|
||||
self.real_stack = (
|
||||
real_stack
|
||||
if real_stack is not None
|
||||
else torch._guards.TracingContext.extract_stack()
|
||||
)
|
||||
|
||||
|
||||
class InvalidBackend(TorchDynamoException):
|
||||
def __init__(self, name: str) -> None:
|
||||
super().__init__(
|
||||
f"Invalid backend: {name!r}, see `torch._dynamo.list_backends()` for available backends."
|
||||
)
|
||||
|
||||
|
||||
class ResetRequired(TorchDynamoException):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
textwrap.dedent(
|
||||
"""
|
||||
Must call `torch._dynamo.reset()` before changing backends. Detected two calls to
|
||||
`torch.compile()` with a different backend compiler arguments.
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class ShortenTraceback(TorchDynamoException):
|
||||
def __init__(
|
||||
self, *args: Any, first_useful_frame: types.FrameType | None, **kwargs: Any
|
||||
) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.first_useful_frame = first_useful_frame
|
||||
|
||||
def remove_dynamo_frames(self) -> typing.Self:
|
||||
tb = self.__traceback__
|
||||
if self.first_useful_frame is None or tb is None or config.verbose:
|
||||
return self
|
||||
while tb.tb_frame is not self.first_useful_frame:
|
||||
tb = tb.tb_next
|
||||
assert tb is not None, "internal error, please report a bug"
|
||||
return self.with_traceback(tb)
|
||||
|
||||
|
||||
class BackendCompilerFailed(ShortenTraceback):
|
||||
def __init__(
|
||||
self,
|
||||
backend_fn: Any,
|
||||
inner_exception: Exception,
|
||||
first_useful_frame: types.FrameType | None,
|
||||
) -> None:
|
||||
self.backend_name = getattr(backend_fn, "__name__", "?")
|
||||
self.inner_exception = inner_exception
|
||||
msg = f"backend={self.backend_name!r} raised:\n{type(inner_exception).__name__}: {inner_exception}"
|
||||
super().__init__(msg, first_useful_frame=first_useful_frame)
|
||||
|
||||
|
||||
# NOTE: important invariant! Almost any exception handler that handles Unsupported
|
||||
# should NOT suppress the exception if skip_frame is set!
|
||||
# skip_frame is used by symbolic_convert.py to bubble up Unsupported exceptions to convert_frame to cause
|
||||
# a frame skip. Once the Unsupported exn is in convert_frame, we will always skip, so skip_frame
|
||||
# won't be checked
|
||||
class Unsupported(TorchDynamoException):
|
||||
def __init__(
|
||||
self,
|
||||
msg: str,
|
||||
# TODO: make this argument required once we remove Unsupported subclasses
|
||||
gb_type: str = "",
|
||||
skip_frame: bool = False,
|
||||
*,
|
||||
case_name: str | None = None,
|
||||
real_stack: StackSummary | None = None,
|
||||
) -> None:
|
||||
super().__init__(msg)
|
||||
if not real_stack:
|
||||
real_stack = torch._guards.TracingContext.extract_stack()
|
||||
self.real_stack = real_stack
|
||||
self.msg = msg
|
||||
self.skip_frame = skip_frame
|
||||
self.category: str | None = None
|
||||
self.add_to_stats()
|
||||
self.gb_type: str | None = gb_type
|
||||
self.logged = False
|
||||
|
||||
def remove_from_stats(self) -> None:
|
||||
assert self.category is not None
|
||||
counters[self.category][self.msg] -= 1
|
||||
if counters[self.category][self.msg] <= 0:
|
||||
del counters[self.category][self.msg]
|
||||
|
||||
def add_to_stats(self, category: str = "unimplemented") -> None:
|
||||
self.category = category
|
||||
counters[category][self.msg] += 1
|
||||
|
||||
|
||||
class UnknownPropertiesDuringBackwardTrace(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class RecompileError(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class InfiniteGeneratorError(TorchDynamoException):
|
||||
# Raised when the number of yielded values is greater than MAX_ITERATOR_LIMIT
|
||||
pass
|
||||
|
||||
|
||||
class CondOpArgsMismatchError(TorchDynamoException):
|
||||
"""
|
||||
Internal error from cond() due to arguments mismatch.
|
||||
"""
|
||||
|
||||
|
||||
class UserErrorType(Enum):
|
||||
DYNAMIC_CONTROL_FLOW = auto()
|
||||
ANTI_PATTERN = auto()
|
||||
STANDARD_LIBRARY = auto()
|
||||
CONSTRAINT_VIOLATION = auto()
|
||||
DYNAMIC_DIM = auto()
|
||||
INVALID_INPUT = auto()
|
||||
INVALID_OUTPUT = auto()
|
||||
UNSUPPORTED_ALIASED_MUTATED_DYNAMIC_INPUTS = auto()
|
||||
|
||||
|
||||
class UserError(TorchDynamoException):
|
||||
def __init__(
|
||||
self, error_type: UserErrorType, msg: str, case_name: str | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Type of errors that would be valid in Eager, but not supported in TorchDynamo.
|
||||
The error message should tell user about next actions.
|
||||
|
||||
error_type: Type of user error
|
||||
msg: Actionable error message
|
||||
case_name: (Optional) Unique name (snake case) for the usage example in exportdb.
|
||||
"""
|
||||
if case_name is not None:
|
||||
assert isinstance(case_name, str)
|
||||
if msg.endswith("."):
|
||||
msg += " "
|
||||
else:
|
||||
msg += "\n"
|
||||
msg += exportdb_error_message(case_name)
|
||||
super().__init__(msg)
|
||||
self.real_stack = torch._guards.TracingContext.extract_stack()
|
||||
self.skip_frame = False
|
||||
self.logged = False
|
||||
self.error_type = error_type
|
||||
self.msg = msg
|
||||
self.message = msg
|
||||
|
||||
|
||||
# debug exception thrown when tracing torch._dynamo.step_unsupported()
|
||||
class StepUnsupported(TorchDynamoException):
|
||||
def __init__(self, msg: str, real_stack: StackSummary | None = None) -> None:
|
||||
super().__init__(msg)
|
||||
self.msg = msg
|
||||
if not real_stack:
|
||||
real_stack = torch._guards.TracingContext.extract_stack()
|
||||
self.real_stack = real_stack
|
||||
self.logged = False
|
||||
|
||||
|
||||
class UnsafeScriptObjectError(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class UncapturedHigherOrderOpError(TorchDynamoException):
|
||||
def __init__(self, msg: str, real_stack: StackSummary | None = None) -> None:
|
||||
super().__init__(msg)
|
||||
self.msg = msg
|
||||
self.real_stack = (
|
||||
real_stack
|
||||
if real_stack is not None
|
||||
else torch._guards.TracingContext.extract_stack()
|
||||
)
|
||||
|
||||
|
||||
# TODO: I'm a little uncertain about what error classification we should have
|
||||
# for this. This is potentially a user error, but regressions in
|
||||
# specialization in PyTorch proper could also trigger this problem
|
||||
class FailOnRecompileLimitHit(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class PackageError(TorchDynamoException):
|
||||
pass
|
||||
|
||||
|
||||
class ObservedException(TorchDynamoException):
|
||||
# An exception observed during the tracing. This exception is used by Dynamo to handle exceptions.
|
||||
def __init__(
|
||||
self, *args: Any, real_stack: StackSummary | None = None, **kwargs: Any
|
||||
) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.real_stack: StackSummary = (
|
||||
real_stack
|
||||
if real_stack is not None
|
||||
else torch._guards.TracingContext.extract_stack()
|
||||
)
|
||||
|
||||
|
||||
class ObservedUserStopIteration(ObservedException):
|
||||
# An UserStopIteration exception observed during the Dynamo tracing (e.g Dynamo tracing __next__)
|
||||
value: Any | None
|
||||
|
||||
# Reference `StopIteration_init` in CPython
|
||||
# https://github.com/python/cpython/blob/3.11/Objects/exceptions.c#L568-L584
|
||||
def __init__(
|
||||
self, *args: Any, real_stack: StackSummary | None = None, **kwargs: Any
|
||||
) -> None:
|
||||
super().__init__("unhandled `raise StopIteration`", real_stack=real_stack)
|
||||
if len(args) > 0:
|
||||
self.value = args[0]
|
||||
else:
|
||||
self.value = None
|
||||
|
||||
|
||||
class ObservedLookupError(ObservedException):
|
||||
# A LookupError exception to be raised from inside Dynamo tracing. This can happen on __getitem__
|
||||
pass
|
||||
|
||||
|
||||
class ObservedIndexError(ObservedLookupError):
|
||||
# An IndexError exception to be raised from inside Dynamo tracing. This can happen on list __getitem__
|
||||
pass
|
||||
|
||||
|
||||
class ObservedKeyError(ObservedLookupError):
|
||||
# A KeyError exception to be raised from inside Dynamo tracing. This can happen on dict __getitem__
|
||||
pass
|
||||
|
||||
|
||||
class ObservedGeneratorExit(ObservedException):
|
||||
pass
|
||||
|
||||
|
||||
class ObservedAttributeError(ObservedException):
|
||||
# An AttributeError exception to be raised from inside Dynamo tracing. This can happen on user defined object __getattr__
|
||||
pass
|
||||
|
||||
|
||||
class ObservedRuntimeError(ObservedException):
|
||||
# A RuntimeError exception to be raised from inside Dynamo tracing. This can happen on generator.throw(..) method
|
||||
pass
|
||||
|
||||
|
||||
class ObservedNotImplementedError(ObservedException):
|
||||
pass
|
||||
|
||||
|
||||
class ObservedTypeError(ObservedException):
|
||||
# A TypeError exception to be raised from inside Dynamo tracing. This can happen on generator.send(..) method
|
||||
pass
|
||||
|
||||
|
||||
observed_exception_map = {
|
||||
StopIteration: ObservedUserStopIteration,
|
||||
LookupError: ObservedLookupError,
|
||||
IndexError: ObservedIndexError,
|
||||
GeneratorExit: ObservedGeneratorExit,
|
||||
KeyError: ObservedKeyError,
|
||||
AttributeError: ObservedAttributeError,
|
||||
RuntimeError: ObservedRuntimeError,
|
||||
NotImplementedError: ObservedNotImplementedError,
|
||||
TypeError: ObservedTypeError,
|
||||
}
|
||||
|
||||
|
||||
def get_dynamo_observed_exception(exc_type: type[Exception]) -> type[ObservedException]:
|
||||
if exc_type not in observed_exception_map:
|
||||
name = getattr(exc_type, "__name__", str(exc_type))
|
||||
observed_exception_map[exc_type] = type( # type: ignore[assignment]
|
||||
f"Observed{name}Error", (ObservedException,), {}
|
||||
)
|
||||
# pyrefly: ignore [bad-index]
|
||||
return observed_exception_map[exc_type]
|
||||
|
||||
|
||||
def raise_observed_exception(
|
||||
exc_type: type[Exception],
|
||||
tx: InstructionTranslatorBase,
|
||||
*,
|
||||
args: list[VariableTracker] | list[str] | None = None,
|
||||
kwargs: dict[str, VariableTracker] | None = None,
|
||||
) -> NoReturn:
|
||||
from .symbolic_convert import ExceptionVals
|
||||
from .variables.builder import SourcelessBuilder
|
||||
|
||||
if args:
|
||||
args_ = [
|
||||
SourcelessBuilder.create(tx, arg) if isinstance(arg, str) else arg
|
||||
for arg in args
|
||||
]
|
||||
else:
|
||||
args_: list[VariableTracker] = []
|
||||
|
||||
# CPython here raises an exception. Since there is no python code, we have to manually setup the exception
|
||||
# stack and raise the exception.
|
||||
exception_vt = SourcelessBuilder.create(tx, exc_type).call_function(
|
||||
tx, args_, kwargs or {}
|
||||
)
|
||||
assert isinstance(exception_vt, ExceptionVals)
|
||||
tx._attach_traceback_to_exception(exception_vt)
|
||||
tx.exn_vt_stack.set_current_exception(exception_vt) # type: ignore[arg-type]
|
||||
raised_exc = get_dynamo_observed_exception(exc_type)
|
||||
# Store the original exception arguments for better error messages
|
||||
if args:
|
||||
raise raised_exc(*args_)
|
||||
raise raised_exc
|
||||
|
||||
|
||||
def raise_type_error(tx: InstructionTranslatorBase, msg: str) -> NoReturn:
|
||||
"""Raise a TypeError as an observed exception during tracing."""
|
||||
raise_observed_exception(TypeError, tx, args=[msg])
|
||||
|
||||
|
||||
def handle_observed_exception(tx: Any) -> None:
|
||||
# This is essentially exception handling code, equivalent of this pseudo code
|
||||
#
|
||||
# try:
|
||||
# ... somebody raising StopIteration
|
||||
# except StopIteration
|
||||
# pass
|
||||
#
|
||||
# If this was going through the python code, we would have called exception_handler method, but FOR_ITER
|
||||
# handles the exception completely in CPython. For example for 3.11, the resulting bytecode is
|
||||
#
|
||||
#
|
||||
# 6 46 LOAD_GLOBAL 2 (StopIteration)
|
||||
# 58 RAISE_VARARGS 1
|
||||
# >> 60 PUSH_EXC_INFO
|
||||
|
||||
# 7 62 LOAD_GLOBAL 2 (StopIteration)
|
||||
# 74 CHECK_EXC_MATCH
|
||||
# 76 POP_JUMP_FORWARD_IF_FALSE 3 (to 84)
|
||||
# 78 POP_TOP
|
||||
|
||||
# 8 80 POP_EXCEPT
|
||||
#
|
||||
|
||||
# Fortunately this translates to a simple pop from the exn_vt_stack
|
||||
tx.exn_vt_stack.clear_current_exception()
|
||||
|
||||
|
||||
# These exceptions are ok to fallback to eager/graph_break.
|
||||
exceptions_allowed_to_be_fallback = (
|
||||
torch._subclasses.fake_tensor.DataDependentOutputException,
|
||||
torch._subclasses.fake_tensor.DynamicOutputShapeException,
|
||||
torch._subclasses.fake_tensor.UnsupportedOperatorException,
|
||||
torch._subclasses.fake_tensor.UnsupportedFakeTensorException,
|
||||
torch._subclasses.fake_tensor.UnsupportedMutationAliasingException,
|
||||
)
|
||||
|
||||
|
||||
def unimplemented_with_warning(
|
||||
e: Exception,
|
||||
code: types.CodeType,
|
||||
*,
|
||||
gb_type: str,
|
||||
context: str,
|
||||
explanation: str,
|
||||
hints: list[str],
|
||||
) -> NoReturn:
|
||||
# This function calls unimplemented internally and eventually graph breaks
|
||||
# or falls to eager. unimplemented itself does not print any user warnings,
|
||||
# i.e., its very silent. This helper function is intended when an error is
|
||||
# encountered in the torch.compile stack which is worth showing as warning
|
||||
# to the user. For example, if AOT Autograd backend fails with a fake tensor
|
||||
# exception, its ok to fallback to eager but not silently. Here, we can use
|
||||
# this function to log the message and the stack trace.
|
||||
graph_break_msg = format_error_msg_verbose(e, code)
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "dynamo_graph_break_reason",
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: graph_break_msg,
|
||||
)
|
||||
graph_breaks_log.debug("%s", graph_break_msg)
|
||||
_unimplemented = unimplemented
|
||||
# to prevent a graph break registry entry
|
||||
_unimplemented(
|
||||
gb_type=gb_type,
|
||||
context=context,
|
||||
explanation=explanation,
|
||||
hints=hints,
|
||||
from_exc=e,
|
||||
log_warning=True,
|
||||
)
|
||||
|
||||
|
||||
def format_graph_break_message(
|
||||
gb_type: str,
|
||||
context: str,
|
||||
explanation: str,
|
||||
hints: list[str],
|
||||
) -> str:
|
||||
explanation = textwrap.indent(explanation, " ").lstrip()
|
||||
hints_str = "\n".join(
|
||||
" Hint: " + textwrap.indent(hint, " ").lstrip() for hint in hints
|
||||
)
|
||||
context = textwrap.indent(context, " ").lstrip()
|
||||
|
||||
msg = f"""\
|
||||
{gb_type}
|
||||
Explanation: {explanation}
|
||||
{hints_str}
|
||||
|
||||
Developer debug context: {context}"""
|
||||
documentation_link = get_gbid_documentation_link(gb_type)
|
||||
|
||||
if documentation_link:
|
||||
msg += f"\n\n For more details about this graph break, please visit: {documentation_link}"
|
||||
|
||||
return msg
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_gb_type_to_gb_id_map() -> dict[str, Any]:
|
||||
"""
|
||||
Loads the gb_type to gb_id map from the graph break registry from JSON file with caching.
|
||||
|
||||
Includes historical gb_type (mapping behavior of duplicate gb_types with different gb_ids is undefined).
|
||||
"""
|
||||
try:
|
||||
script_dir = Path(__file__).resolve().parent
|
||||
registry_path = get_file_path_2(
|
||||
"", str(script_dir), "graph_break_registry.json"
|
||||
)
|
||||
with open(registry_path) as f:
|
||||
registry = json.load(f)
|
||||
except Exception:
|
||||
log.exception("Error accessing the registry file")
|
||||
# pyrefly: ignore [implicit-any]
|
||||
registry = {}
|
||||
|
||||
mapping = {}
|
||||
for k, v in registry.items():
|
||||
for entry in v:
|
||||
mapping[entry["Gb_type"]] = k
|
||||
|
||||
return mapping
|
||||
|
||||
|
||||
def get_gbid_documentation_link(gb_type: str) -> str | None:
|
||||
"""
|
||||
Retrieves the GBID documentation link for a given graph break type.
|
||||
|
||||
Args:
|
||||
gb_type: The graph break type to look up.
|
||||
|
||||
Returns:
|
||||
A string containing the documentation URL if found, otherwise None.
|
||||
"""
|
||||
GRAPH_BREAK_SITE_URL = (
|
||||
"https://meta-pytorch.github.io/compile-graph-break-site/gb/" # @lint-ignore
|
||||
)
|
||||
|
||||
gb_type_to_gb_id_map = _load_gb_type_to_gb_id_map()
|
||||
|
||||
if gb_type in gb_type_to_gb_id_map:
|
||||
return (
|
||||
f"{GRAPH_BREAK_SITE_URL}gb{gb_type_to_gb_id_map[gb_type].lstrip('GB')}.html"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
_NOTHING = object()
|
||||
|
||||
|
||||
def unimplemented(
|
||||
*,
|
||||
gb_type: str,
|
||||
context: str,
|
||||
explanation: str,
|
||||
hints: list[str],
|
||||
from_exc: Any = _NOTHING,
|
||||
log_warning: bool = False,
|
||||
skip_frame: bool = False,
|
||||
) -> NoReturn:
|
||||
"""
|
||||
Called within dynamo to cause a graph break.
|
||||
Args:
|
||||
gb_type: Context-free graph break type. It should be a short string without any
|
||||
information specific to the tracing context (i.e. no dynamically-generated strings)
|
||||
context: Developer context for the graph break. It can contain tracing context/dynamic strings.
|
||||
explanation: User-facing context-dependent explanation for the graph break. Can be dynamic.
|
||||
hints: List of user-facing hints for the graph break.
|
||||
"""
|
||||
|
||||
msg = format_graph_break_message(gb_type, context, explanation, hints)
|
||||
|
||||
if log_warning:
|
||||
log.warning(msg)
|
||||
if from_exc is not _NOTHING:
|
||||
past_real_stack = None
|
||||
if hasattr(from_exc, "real_stack"):
|
||||
past_real_stack = from_exc.real_stack
|
||||
if isinstance(from_exc, Unsupported):
|
||||
msg = f"{from_exc.msg}\n\n*** While handling this graph break, another graph break occurred: ***\n\n{msg}"
|
||||
# noqa: GB_REGISTRY
|
||||
raise Unsupported(msg, gb_type, skip_frame, real_stack=past_real_stack)
|
||||
# noqa: GB_REGISTRY
|
||||
raise Unsupported(
|
||||
msg, gb_type, skip_frame, real_stack=past_real_stack
|
||||
) from from_exc
|
||||
# noqa: GB_REGISTRY
|
||||
raise Unsupported(msg, gb_type, skip_frame)
|
||||
|
||||
|
||||
# KeyError has special handling for its args
|
||||
# see https://github.com/python/cpython/blob/3.11/Objects/exceptions.c#L2534 for details
|
||||
class KeyErrorMsg:
|
||||
def __init__(self, value: Any) -> None:
|
||||
self.value = value
|
||||
|
||||
def __str__(self) -> str:
|
||||
return str(self.value)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__str__()
|
||||
|
||||
|
||||
def augment_exc_message_with_hop_name(exc: Exception, msg: str) -> str:
|
||||
# Add HOP context right after before the explanation if present;
|
||||
# otherwise after the message
|
||||
if hasattr(exc, "_hop_name"):
|
||||
lines = msg.partition("\n Explanation:")
|
||||
msg = (
|
||||
f"{lines[0]}\n Higher Order Operator: {exc._hop_name}{lines[1]}{lines[2]}" # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
return msg
|
||||
|
||||
|
||||
def augment_exc_message(exc: Exception, msg: str = "\n", export: bool = False) -> None:
|
||||
import traceback
|
||||
|
||||
exc.innermost_user_frame_summary = None # type: ignore[attr-defined]
|
||||
|
||||
real_stack = get_real_stack(exc)
|
||||
if real_stack is not None and len(real_stack) > 0:
|
||||
exc.innermost_user_frame_summary = real_stack[-1] # type: ignore[attr-defined]
|
||||
msg += f"\nfrom user code:\n {''.join(traceback.format_list(real_stack))}"
|
||||
|
||||
if config.replay_record_enabled and hasattr(exc, "record_filename"):
|
||||
msg += (
|
||||
f"\nLast frame execution written to {exc.record_filename}. To run only this frame while debugging, run\
|
||||
torch._dynamo.replay('{exc.record_filename}').\n"
|
||||
)
|
||||
|
||||
if not config.verbose and hasattr(exc, "real_stack"):
|
||||
msg += (
|
||||
"\nSet TORCHDYNAMO_VERBOSE=1 for the internal stack trace "
|
||||
"(please do this especially if you're reporting a bug to PyTorch). "
|
||||
'For even more developer context, set TORCH_LOGS="+dynamo"\n'
|
||||
)
|
||||
|
||||
if hasattr(exc, "inner_exception") and hasattr(
|
||||
exc.inner_exception, "minifier_path"
|
||||
):
|
||||
if hasattr(exc.inner_exception, "buck_command"):
|
||||
msg += (
|
||||
f"\nMinifier script written to {exc.inner_exception.minifier_path}. Run "
|
||||
f"this buck command to find the smallest traced graph "
|
||||
f"which reproduces this error: {exc.inner_exception.buck_command}\n"
|
||||
)
|
||||
else:
|
||||
msg += (
|
||||
f"\nMinifier script written to {exc.inner_exception.minifier_path}. Run "
|
||||
"this script to find the smallest traced graph which reproduces this error.\n"
|
||||
)
|
||||
|
||||
old_msg = "" if len(exc.args) == 0 else str(exc.args[0])
|
||||
|
||||
old_msg = augment_exc_message_with_hop_name(exc, old_msg)
|
||||
|
||||
if isinstance(exc, KeyError):
|
||||
exc.args = (KeyErrorMsg(old_msg + msg),) + exc.args[1:]
|
||||
else:
|
||||
new_msg = old_msg + msg
|
||||
exc.args = (new_msg,) + exc.args[1:]
|
||||
|
||||
|
||||
def get_exc_message(
|
||||
e: Exception, compile_id: CompileId
|
||||
) -> tuple[str | None, int | None]:
|
||||
filename = None
|
||||
lineno = None
|
||||
if e.innermost_user_frame_summary is not None: # type: ignore[attr-defined]
|
||||
filename = e.innermost_user_frame_summary.filename # type: ignore[attr-defined]
|
||||
lineno = e.innermost_user_frame_summary.lineno # type: ignore[attr-defined]
|
||||
e.compile_id = compile_id # type: ignore[attr-defined]
|
||||
return filename, lineno
|
||||
|
||||
|
||||
def get_stack_above_dynamo() -> StackSummary:
|
||||
return filter_stack(extract_stack())
|
||||
|
||||
|
||||
def get_real_stack(
|
||||
exc: Exception, frame: DynamoFrameType | None = None
|
||||
) -> StackSummary | None:
|
||||
real_stack = getattr(exc, "real_stack", None)
|
||||
if real_stack is None:
|
||||
return None
|
||||
|
||||
# NB: it's possible for real_stack to be []; we still attempt to
|
||||
# report a stack anyway because the stack_above_dynamo may still
|
||||
# be useful for debugging
|
||||
|
||||
if frame is not None:
|
||||
# NB: frame is PyInterpreterFrame on Python 3.11 and later,
|
||||
# not a TRUE frame object. You can't actually feed it
|
||||
# to traceback because it doesn't have enough information.
|
||||
# To solve this problem, we technically should just materialize
|
||||
# the frame, the same way _PyFrame_GetFrameObject would do
|
||||
# (but we cannot actually do this, because this populates
|
||||
# frame_obj field, which default eval frame doesn't like).
|
||||
#
|
||||
# Fortunately, in this case, we can hack it: there's no need
|
||||
# to actually use the truly top frame, we can just extract
|
||||
# from where we are right now and rely on filter_stack to
|
||||
# get rid of all the dynamo frames. For ease of testing
|
||||
# we apply this behavior to ALL Python versions
|
||||
stack_above_dynamo = get_stack_above_dynamo()
|
||||
else:
|
||||
stack_above_dynamo = StackSummary()
|
||||
|
||||
return StackSummary.from_list(stack_above_dynamo + real_stack)
|
||||
|
||||
|
||||
# filter out all frames after entering dynamo
|
||||
def filter_stack(stack: StackSummary) -> StackSummary:
|
||||
user_stack = StackSummary()
|
||||
for frame in stack:
|
||||
if frame.filename is None:
|
||||
continue
|
||||
if "convert_frame" in frame.filename:
|
||||
break
|
||||
if "eval_frame" in frame.filename or (
|
||||
frame.line and "torch._dynamo.optimize(" in frame.line
|
||||
):
|
||||
continue
|
||||
user_stack.append(frame)
|
||||
|
||||
return user_stack
|
||||
|
||||
|
||||
def remove_resume_prefix(name: str) -> str:
|
||||
from .resume_execution import TORCH_DYNAMO_RESUME_IN_PREFIX
|
||||
|
||||
match = re.match(f"{TORCH_DYNAMO_RESUME_IN_PREFIX}_(\\w+)_at_\\d+", name)
|
||||
if match:
|
||||
return match.group(1)
|
||||
return name
|
||||
|
||||
|
||||
def collapse_resume_frames(stack: StackSummary | list[FrameSummary]) -> StackSummary:
|
||||
"""
|
||||
When we graph break, we create a resume function and make a regular Python call
|
||||
to it, which gets intercepted by Dynamo. This behavior is normally shown in the
|
||||
traceback, which can be confusing to a user. So we can filter out resume frames
|
||||
for better traceback clarity.
|
||||
|
||||
Example:
|
||||
File "..." line 3, in f
|
||||
<line 3>
|
||||
File "..." line 5, in torch_dynamo_resume_in_f_at_80
|
||||
<line 5>
|
||||
File "..." line 10, in torch_dynamo_resume_in_f_at_120
|
||||
<line 10>
|
||||
|
||||
becomes
|
||||
File "..." line 10, in f
|
||||
<line 10>
|
||||
"""
|
||||
|
||||
new_stack = StackSummary()
|
||||
for frame in stack:
|
||||
if frame.filename is None:
|
||||
continue
|
||||
name = remove_resume_prefix(frame.name)
|
||||
if new_stack and name and new_stack[-1].name == name:
|
||||
new_stack[-1] = frame
|
||||
frame.name = name
|
||||
else:
|
||||
frame.name = name
|
||||
new_stack.append(frame)
|
||||
|
||||
return new_stack
|
||||
|
||||
|
||||
def format_error_msg_verbose(
|
||||
exc: Exception,
|
||||
code: types.CodeType,
|
||||
record_filename: str | None = None,
|
||||
frame: DynamoFrameType | None = None,
|
||||
) -> str:
|
||||
msg = (
|
||||
f"WON'T CONVERT {code.co_name} {code.co_filename} line {code.co_firstlineno}\n"
|
||||
)
|
||||
msg += "=" * 10 + " TorchDynamo Stack Trace " + "=" * 10 + "\n"
|
||||
msg += format_exc()
|
||||
real_stack = get_real_stack(exc, frame)
|
||||
if real_stack is not None:
|
||||
msg += (
|
||||
"\n"
|
||||
+ "=" * 10
|
||||
+ " The above exception occurred while processing the following code "
|
||||
+ "=" * 10
|
||||
+ "\n\n"
|
||||
)
|
||||
msg += "".join(format_list(real_stack))
|
||||
msg += "\n"
|
||||
msg += "=" * 10
|
||||
|
||||
return msg
|
||||
|
||||
|
||||
def format_frame_info(code: types.CodeType) -> str:
|
||||
return (
|
||||
f"{getattr(code, 'co_name', '<unknown>')} "
|
||||
f"({getattr(code, 'co_filename', '<unknown>')} "
|
||||
f"line {getattr(code, 'co_firstlineno', 0)})"
|
||||
)
|
||||
|
||||
|
||||
def format_skip_frame_message(code: types.CodeType | None, reason: str) -> str:
|
||||
if code is not None:
|
||||
frame_info = format_frame_info(code)
|
||||
return (
|
||||
f"torch.compile intentionally decided to skip the frame {frame_info} and fall back to eager.\n"
|
||||
f"Reason: {reason}"
|
||||
)
|
||||
else:
|
||||
return (
|
||||
f"torch.compile intentionally decided to skip the frame and fall back to eager.\n"
|
||||
f"Reason: {reason}"
|
||||
)
|
||||
|
||||
|
||||
def format_error_msg(
|
||||
exc: Exception,
|
||||
code: types.CodeType,
|
||||
record_filename: str | None = None,
|
||||
frame: DynamoFrameType | None = None,
|
||||
) -> str:
|
||||
if config.verbose:
|
||||
return format_error_msg_verbose(exc, code, record_filename, frame)
|
||||
return f"WON'T CONVERT {code.co_name} {code.co_filename}\
|
||||
line {code.co_firstlineno} \ndue to: \n{format_exc()}"
|
||||
@@ -0,0 +1,310 @@
|
||||
"""
|
||||
This module contains utility functions that are explicitly allowed to be called during
|
||||
TorchDynamo compilation. These functions are carefully vetted to ensure they work
|
||||
correctly within the TorchDynamo tracing and compilation process.
|
||||
|
||||
Key functionality groups:
|
||||
|
||||
- Compilation State:
|
||||
Functions for checking compilation state (is_compiling)
|
||||
|
||||
- Function Wrapping:
|
||||
Utilities for wrapping functions (wrap_inline, wrap_numpy) to work with
|
||||
TorchDynamo compilation
|
||||
|
||||
- Autograd Hooks:
|
||||
Functions and classes for handling autograd hooks and backward passes
|
||||
(call_hook, FakeBackwardCFunction, etc.)
|
||||
|
||||
- Tensor Operations:
|
||||
Utility functions for tensor operations and transformations
|
||||
"""
|
||||
|
||||
import functools
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TYPE_CHECKING, TypeVar
|
||||
from typing_extensions import deprecated, ParamSpec
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
except ModuleNotFoundError:
|
||||
np = None # type: ignore[assignment]
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# TorchScript does not support `@deprecated`
|
||||
# This is a workaround to avoid breaking TorchScript
|
||||
@deprecated(
|
||||
"`torch._dynamo.external_utils.is_compiling` is deprecated. Use `torch.compiler.is_compiling` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def is_compiling() -> bool:
|
||||
return torch.compiler.is_compiling()
|
||||
|
||||
else:
|
||||
|
||||
def is_compiling() -> bool:
|
||||
"""
|
||||
Indicates whether we are tracing/compiling with torch.compile() or torch.export().
|
||||
"""
|
||||
# NOTE: With `@torch.compile(backend="eager")`, torch._dynamo.is_compiling() will get traced
|
||||
# and return true. torch.compiler.is_compiling() is skipped and will return false.
|
||||
return torch.compiler.is_compiling()
|
||||
|
||||
|
||||
def wrap_inline(fn: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
"""
|
||||
Create an extra frame around fn that is not in skipfiles.
|
||||
"""
|
||||
|
||||
@functools.wraps(fn)
|
||||
def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
def call_hook(
|
||||
hook: Callable[..., torch.Tensor | None], *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Used by compiled autograd to handle hook returning None.
|
||||
"""
|
||||
result = hook(*args)
|
||||
if result is None:
|
||||
return args[0]
|
||||
elif kwargs.get("hook_type") == "post_acc_grad_hook":
|
||||
raise RuntimeError("Tensor post accumulate grad hooks should return None.")
|
||||
return result
|
||||
|
||||
|
||||
def wrap_numpy(f: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
r"""Decorator that turns a function from ``np.ndarray``s to ``np.ndarray``s into a function
|
||||
from ``torch.Tensor``s to ``torch.Tensor``s.
|
||||
"""
|
||||
if not np:
|
||||
return f
|
||||
|
||||
@functools.wraps(f)
|
||||
def wrap(*args: _P.args, **kwargs: _P.kwargs) -> pytree.PyTree:
|
||||
args, kwargs = pytree.tree_map_only(
|
||||
torch.Tensor, lambda x: x.numpy(), (args, kwargs)
|
||||
)
|
||||
out = f(*args, **kwargs)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return pytree.tree_map_only(np.ndarray, lambda x: torch.as_tensor(x), out)
|
||||
|
||||
return wrap
|
||||
|
||||
|
||||
class FakeBackwardCFunction:
|
||||
def __init__(
|
||||
self,
|
||||
real: torch.autograd.function.BackwardCFunction,
|
||||
saved_tensors: list[torch.Tensor],
|
||||
) -> None:
|
||||
self.real = real
|
||||
self.saved_tensors = saved_tensors
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
if name == "saved_variables":
|
||||
warnings.warn(
|
||||
"'saved_variables' is deprecated; use 'saved_tensors'",
|
||||
DeprecationWarning,
|
||||
)
|
||||
return self.saved_tensors
|
||||
|
||||
return getattr(self.real, name)
|
||||
|
||||
|
||||
def call_backward(
|
||||
backward_c_function: torch.autograd.function.BackwardCFunction,
|
||||
saved_tensors: list[torch.Tensor],
|
||||
*args: Any,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
||||
fake = FakeBackwardCFunction(backward_c_function, saved_tensors)
|
||||
grads = fake._forward_cls.backward(fake, *args) # type: ignore[attr-defined]
|
||||
|
||||
if not isinstance(grads, tuple):
|
||||
grads = (grads,)
|
||||
|
||||
return grads
|
||||
|
||||
|
||||
def normalize_as_list(x: Any) -> list[Any]:
|
||||
if isinstance(x, tuple):
|
||||
return list(x)
|
||||
elif isinstance(x, list):
|
||||
return x
|
||||
return [x]
|
||||
|
||||
|
||||
def untyped_storage_size(x: torch.Tensor) -> int:
|
||||
return x.untyped_storage().size()
|
||||
|
||||
|
||||
class FakeCompiledAutogradEngine:
|
||||
@staticmethod
|
||||
def queue_callback(
|
||||
final_callbacks: list[Callable[[], None]], cb: Callable[[], None]
|
||||
) -> None:
|
||||
final_callbacks.append(cb)
|
||||
|
||||
@staticmethod
|
||||
def exec_final_callbacks(final_callbacks: list[Callable[[], None]]) -> None:
|
||||
i = 0
|
||||
while i < len(final_callbacks):
|
||||
cb = final_callbacks[i]
|
||||
cb()
|
||||
i += 1
|
||||
final_callbacks.clear()
|
||||
|
||||
@staticmethod
|
||||
def _exec_final_callbacks_stub() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def call_hook_from_backward_state(
|
||||
*args: Any, bw_state: Any, hook_name: str, **kwargs: Any
|
||||
) -> Any:
|
||||
return getattr(bw_state, hook_name)(*args, **kwargs)
|
||||
|
||||
|
||||
class _ApplyBackwardHook(torch.autograd.Function):
|
||||
"""Custom autograd function that applies a hook during backward.
|
||||
|
||||
This is used to implement register_hook on intermediate tensors without
|
||||
requiring compiled autograd. The hook function is captured in the context
|
||||
and applied during the backward pass.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
# pyre-ignore[14]: Inconsistent override is expected for autograd.Function
|
||||
def forward(
|
||||
ctx: Any, tensor: torch.Tensor, hook_fn: Callable[..., Any]
|
||||
) -> torch.Tensor: # type: ignore[override]
|
||||
ctx.hook_fn = hook_fn
|
||||
return tensor.view_as(tensor)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: Any, grad: torch.Tensor) -> tuple[torch.Tensor, None]: # type: ignore[override]
|
||||
result = ctx.hook_fn(grad)
|
||||
if result is None:
|
||||
result = grad
|
||||
return result, None
|
||||
|
||||
|
||||
def call_module_hooks_from_backward_state(
|
||||
_: Any, result: Any, *args: Any, bw_state: Any, hooks_name: str, module_name: str
|
||||
) -> Any:
|
||||
module = getattr(bw_state, module_name)
|
||||
hooks = getattr(bw_state, hooks_name)
|
||||
for hook in hooks:
|
||||
new_result = hook(module, result, *args)
|
||||
if new_result is not None:
|
||||
result = new_result
|
||||
return result
|
||||
|
||||
|
||||
# used for torch._dynamo.disable(recursive=False)
|
||||
def get_nonrecursive_disable_wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
# wrap function to get the right error message
|
||||
# this function is in external_utils so that convert_frame doesn't skip it.
|
||||
@functools.wraps(fn)
|
||||
def nonrecursive_disable_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
if torch.compiler.is_exporting():
|
||||
raise RuntimeError(
|
||||
"Non-recursive torch.compiler.disable is not supported with torch.export."
|
||||
)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return nonrecursive_disable_wrapper
|
||||
|
||||
|
||||
def wrap_dunder_call_ctx_manager(self: Any, func: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
"""
|
||||
Apply self as a ctx manager around a call to func
|
||||
"""
|
||||
|
||||
# NOTE: do not functools.wraps(func) because we don't ever want this frame to be skipped!
|
||||
def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
with self:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
# Use only on ints marked dynamic via torch.empty(0, integer)
|
||||
# Currently only way to mark ints as dynamic: https://github.com/pytorch/pytorch/issues/129623
|
||||
def unwrap_maybe_dynamic_int(x: torch.Tensor | int) -> int:
|
||||
if isinstance(x, torch.Tensor):
|
||||
# x.size() is expected to be [0, dynamic_int]
|
||||
return x.size(1)
|
||||
return x
|
||||
|
||||
|
||||
def call_accumulate_grad(
|
||||
variable: torch.Tensor, grad: torch.Tensor, has_post_hooks: bool
|
||||
) -> None:
|
||||
updated_grad = torch._dynamo.compiled_autograd.ops.AccumulateGrad( # type: ignore[attr-defined]
|
||||
[grad], variable, variable.grad, has_post_hooks
|
||||
)
|
||||
variable.grad = updated_grad[0]
|
||||
|
||||
|
||||
def wrap_inline_with_error_on_graph_break(
|
||||
fn: Callable[_P, _R], error_on_graph_break: bool
|
||||
) -> Callable[_P, _R]:
|
||||
# NB: need multiple definitions in order to prevent `fullgraph` from
|
||||
# being a freevar of wrapper
|
||||
# NOTE: do not functools.wraps(fn) because we don't ever want these wrappers to be skipped!
|
||||
if error_on_graph_break:
|
||||
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
with torch._dynamo.error_on_graph_break(True):
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
else:
|
||||
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
with torch._dynamo.error_on_graph_break(False):
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def filter_out_const_values(tup: tuple[Any, ...], masks: list[bool]) -> tuple[Any, ...]:
|
||||
"""
|
||||
masks is a list of bools, where True means the corresponding element in tup
|
||||
is a const value. Filter out the const values.
|
||||
"""
|
||||
out = []
|
||||
for mask_idx, mask in enumerate(masks):
|
||||
if not mask:
|
||||
out.append(tup[mask_idx])
|
||||
return tuple(out)
|
||||
|
||||
|
||||
def insert_const_values_with_mask(
|
||||
tup: tuple[Any, ...], masks: list[bool], values: tuple[Any, ...]
|
||||
) -> tuple[Any, ...]:
|
||||
"""
|
||||
masks and values are of same length. For indices where the mask is True, use
|
||||
the const_values to fill in.
|
||||
"""
|
||||
out = []
|
||||
idx = 0
|
||||
for mask_idx, mask in enumerate(masks):
|
||||
if mask:
|
||||
out.append(values[mask_idx])
|
||||
else:
|
||||
out.append(tup[idx])
|
||||
idx += 1
|
||||
return tuple(out)
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
This module provides functionality for caching and looking up fully qualified function
|
||||
and class names from Python source files by line number.
|
||||
|
||||
It uses Python's tokenize module to parse source files and tracks function/class
|
||||
definitions along with their nesting to build fully qualified names (e.g. 'class.method'
|
||||
or 'module.function'). The results are cached in a two-level dictionary mapping:
|
||||
|
||||
filename -> (line_number -> fully_qualified_name)
|
||||
|
||||
Example usage:
|
||||
name = get_funcname("myfile.py", 42) # Returns name of function/class at line 42
|
||||
clearcache() # Clear the cache if file contents have changed
|
||||
|
||||
The parsing is done lazily when a file is first accessed. Invalid Python files or
|
||||
IO errors are handled gracefully by returning empty cache entries.
|
||||
"""
|
||||
|
||||
import tokenize
|
||||
|
||||
|
||||
cache: dict[str, dict[int, str]] = {}
|
||||
|
||||
|
||||
def clearcache() -> None:
|
||||
cache.clear()
|
||||
|
||||
|
||||
def _add_file(filename: str) -> None:
|
||||
try:
|
||||
with tokenize.open(filename) as f:
|
||||
tokens = list(tokenize.generate_tokens(f.readline))
|
||||
except (OSError, tokenize.TokenError):
|
||||
cache[filename] = {}
|
||||
return
|
||||
|
||||
# NOTE: undefined behavior if file is not valid Python source,
|
||||
# since tokenize will have undefined behavior.
|
||||
result: dict[int, str] = {}
|
||||
# current full funcname, e.g. xxx.yyy.zzz
|
||||
cur_name = ""
|
||||
cur_indent = 0
|
||||
significant_indents: list[int] = []
|
||||
|
||||
for i, token in enumerate(tokens):
|
||||
if token.type == tokenize.INDENT:
|
||||
cur_indent += 1
|
||||
elif token.type == tokenize.DEDENT:
|
||||
cur_indent -= 1
|
||||
# possible end of function or class
|
||||
if significant_indents and cur_indent == significant_indents[-1]:
|
||||
significant_indents.pop()
|
||||
# pop the last name
|
||||
cur_name = cur_name.rpartition(".")[0]
|
||||
elif (
|
||||
token.type == tokenize.NAME
|
||||
and i + 1 < len(tokens)
|
||||
and tokens[i + 1].type == tokenize.NAME
|
||||
and (token.string == "class" or token.string == "def")
|
||||
):
|
||||
# name of class/function always follows class/def token
|
||||
significant_indents.append(cur_indent)
|
||||
if cur_name:
|
||||
cur_name += "."
|
||||
cur_name += tokens[i + 1].string
|
||||
result[token.start[0]] = cur_name
|
||||
|
||||
cache[filename] = result
|
||||
|
||||
|
||||
def get_funcname(filename: str, lineno: int) -> str | None:
|
||||
if filename not in cache:
|
||||
_add_file(filename)
|
||||
return cache[filename].get(lineno, None)
|
||||
@@ -0,0 +1,983 @@
|
||||
import inspect
|
||||
import logging
|
||||
import sys
|
||||
import traceback
|
||||
import types
|
||||
from collections import namedtuple
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
from typing import Any, Optional, TYPE_CHECKING, TypeVar
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
import torch.fx
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._dynamo.convert_frame import CaptureOutput, fullgraph_capture, get_traced_fn
|
||||
from torch._dynamo.decorators import disable as dynamo_disable
|
||||
from torch._dynamo.eval_frame import argument_names, check_user_input_output
|
||||
from torch._dynamo.exc import UserErrorType
|
||||
from torch._dynamo.utils import dynamo_timed, get_metrics_context
|
||||
from torch._export.utils import _compiling_state_context
|
||||
from torch._guards import TracingContext
|
||||
from torch.export.dynamic_shapes import _RelaxedConstraint, Constraint
|
||||
from torch.fx import Node
|
||||
from torch.fx.experimental.symbolic_shapes import (
|
||||
ConstraintViolationError,
|
||||
DimDynamic,
|
||||
StatelessSymbolicContext,
|
||||
)
|
||||
from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
|
||||
from torch.fx.node import Argument, Target
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
T = TypeVar("T")
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def post_process_error_msg(
|
||||
constraint_violation_error: ConstraintViolationError,
|
||||
func: Callable[..., Any],
|
||||
args: Any,
|
||||
kwargs: Any,
|
||||
) -> ConstraintViolationError:
|
||||
"""
|
||||
Because we trace a different callable, the sources are all messed up.
|
||||
Manually patch them so the error message looks correct.
|
||||
"""
|
||||
from torch.export._unlift import _get_input_paths, _replace_sources
|
||||
|
||||
orig_sig = inspect.signature(func)
|
||||
flat_input_paths = _get_input_paths((args, kwargs), orig_sig)
|
||||
if constraint_violation_error.args:
|
||||
constraint_violation_error.args = (
|
||||
_replace_sources(constraint_violation_error.args[0], flat_input_paths),
|
||||
)
|
||||
return constraint_violation_error
|
||||
|
||||
|
||||
EXPORT_ROOT_REPLACEMENTS = [
|
||||
("__export_root_", "_"),
|
||||
("_export_root.", ""),
|
||||
("._export_root", ""),
|
||||
]
|
||||
|
||||
|
||||
def clean_export_root_string(text: str) -> str:
|
||||
"""Generic utility to clean export_root patterns from strings."""
|
||||
result = text
|
||||
for pattern, replacement in EXPORT_ROOT_REPLACEMENTS:
|
||||
result = result.replace(pattern, replacement)
|
||||
return result
|
||||
|
||||
|
||||
def clean_nn_module_stack_and_source_fn(
|
||||
graph_module: torch.fx.GraphModule, is_inline_builtin: bool = False
|
||||
) -> torch.fx.GraphModule:
|
||||
"""
|
||||
Clean up nn_module_stack metadata by removing export_root references.
|
||||
|
||||
Removes the _export_root module references from nn_module_stack metadata
|
||||
in graph nodes, which are artifacts from the export process. Fixes two patterns:
|
||||
|
||||
1. Keys: Removes "__export_root_" and "__modules['_export_root']_" prefixes
|
||||
- Normal case: "L__self____export_root_child" -> "L__self__child"
|
||||
- inline_builtin case: Uses numeric ID strings like "140468831433840"
|
||||
|
||||
2. Values: Removes "._export_root" and "._modules['_export_root']" from child names
|
||||
e.g., "L['self']._export_root.child" -> "L['self'].child"
|
||||
e.g., "L['self']._modules['_export_root'].child" -> "L['self'].child"
|
||||
|
||||
Also removes the root export entry "L__self____export_root" entirely.
|
||||
|
||||
Args:
|
||||
graph_module: The GraphModule to clean up
|
||||
is_inline_builtin: If True, keys are numeric ID strings and self references
|
||||
(L['self']) are filtered out
|
||||
|
||||
Returns:
|
||||
The cleaned GraphModule (modified in-place)
|
||||
"""
|
||||
|
||||
def _process_nn_module_stack(
|
||||
nn_module_stack: dict[str, tuple[str, T]],
|
||||
) -> dict[str, tuple[str, T]]:
|
||||
if "L__self____export_root" in nn_module_stack:
|
||||
del nn_module_stack["L__self____export_root"]
|
||||
|
||||
# Clean up remaining entries
|
||||
cleaned_stack = {}
|
||||
for key, (child_name, child_class) in nn_module_stack.items():
|
||||
# Clean key by removing export_root patterns
|
||||
clean_key = clean_export_root_string(key)
|
||||
|
||||
# Clean child_name by removing export_root patterns
|
||||
clean_name = clean_export_root_string(child_name)
|
||||
|
||||
# Skip self reference for inline builtin case
|
||||
if is_inline_builtin and clean_name == "L['self']":
|
||||
continue
|
||||
|
||||
cleaned_stack[clean_key] = (clean_name, child_class)
|
||||
return cleaned_stack
|
||||
|
||||
def _process_source_fn(source_fn_stack: Iterable[T]) -> Iterable[T]:
|
||||
cleaned_stack = []
|
||||
for item in source_fn_stack:
|
||||
if isinstance(item, tuple) and len(item) == 2:
|
||||
name, cls = item
|
||||
if isinstance(name, str):
|
||||
clean_name = clean_export_root_string(name)
|
||||
cleaned_stack.append((clean_name, cls))
|
||||
else:
|
||||
cleaned_stack.append(item)
|
||||
else:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
cleaned_stack.append(item)
|
||||
# pyrefly: ignore [bad-return]
|
||||
return cleaned_stack
|
||||
|
||||
for node in graph_module.graph.nodes:
|
||||
if "nn_module_stack" in node.meta:
|
||||
node.meta["nn_module_stack"] = _process_nn_module_stack(
|
||||
node.meta["nn_module_stack"].copy()
|
||||
)
|
||||
|
||||
source_fn_stack = node.meta.get("source_fn_stack", None)
|
||||
if source_fn_stack:
|
||||
node.meta["source_fn_stack"] = _process_source_fn(source_fn_stack.copy())
|
||||
|
||||
if "dynamo_flat_name_to_original_fqn" in graph_module.meta:
|
||||
# Clean up flat name to original fqn mapping
|
||||
clean_name_to_original_fqn = {}
|
||||
for flat_name, original_fqn in graph_module.meta[
|
||||
"dynamo_flat_name_to_original_fqn"
|
||||
].items():
|
||||
clean_name_to_original_fqn[clean_export_root_string(flat_name)] = (
|
||||
clean_export_root_string(original_fqn)
|
||||
)
|
||||
graph_module.meta["dynamo_flat_name_to_original_fqn"] = (
|
||||
clean_name_to_original_fqn
|
||||
)
|
||||
|
||||
return graph_module
|
||||
|
||||
|
||||
def clean_export_root(graph_module: torch.fx.GraphModule) -> None:
|
||||
"""Remove export_root artifacts from FX graph in-place"""
|
||||
|
||||
# Unlike getattr node, call_module can be invoked multiple times
|
||||
# In those cases, we should fix all invocations of call_module
|
||||
clean_named_module_map: dict[str, str] = {}
|
||||
|
||||
# Update get_attr nodes in-place
|
||||
for node in graph_module.graph.nodes:
|
||||
if node.op == "get_attr":
|
||||
old_target = node.target
|
||||
new_target = clean_export_root_string(old_target)
|
||||
if new_target != old_target:
|
||||
node.target = new_target
|
||||
assert hasattr(graph_module, old_target)
|
||||
# Move the parameter to the new name
|
||||
param = torch.fx.graph_module._get_attr(graph_module, old_target)
|
||||
torch.fx.graph_module._assign_attr(param, graph_module, new_target)
|
||||
torch.fx.graph_module._del_attr(graph_module, old_target)
|
||||
# Dynamo will only have one nested level
|
||||
if node.op == "call_module":
|
||||
old_target = node.target
|
||||
assert isinstance(old_target, str)
|
||||
new_target = clean_export_root_string(old_target)
|
||||
assert isinstance(new_target, str)
|
||||
new_name = clean_export_root_string(node.name)
|
||||
if new_target == old_target:
|
||||
continue
|
||||
|
||||
# if this module has already been cleaned before, just lookup from map.
|
||||
if old_target in clean_named_module_map:
|
||||
node.target = clean_named_module_map[old_target]
|
||||
node.name = new_name
|
||||
continue
|
||||
target = graph_module.get_submodule(old_target)
|
||||
graph_module.delete_submodule(old_target)
|
||||
graph_module.add_submodule(new_target, target)
|
||||
node.target = new_target
|
||||
node.name = new_name
|
||||
clean_named_module_map[old_target] = new_target
|
||||
|
||||
|
||||
class ModuleToTrace(torch.nn.Module):
|
||||
def __init__(self, foo: Any, in_spec: Any) -> None:
|
||||
super().__init__()
|
||||
self._export_root = foo
|
||||
self.in_spec = in_spec
|
||||
|
||||
def forward(self, *flat_args: Any) -> "ExportTracerOutput":
|
||||
args, kwargs = pytree.tree_unflatten(flat_args, self.in_spec)
|
||||
res = self._export_root(*args, **kwargs)
|
||||
out_flat, out_spec = pytree.tree_flatten(res)
|
||||
return ExportTracerOutput(out_flat, out_spec)
|
||||
|
||||
|
||||
ExportTracerOutput = namedtuple("ExportTracerOutput", ["flat_args", "out_spec"])
|
||||
|
||||
|
||||
# mypy: disable-error-code="no-untyped-def,var-annotated,assignment,index,operator"
|
||||
class DynamoGraphTransformer(torch.fx.Transformer):
|
||||
"""Graph transformer for dynamo export that flattens inputs/outputs without complex matching."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
module: torch.fx.GraphModule,
|
||||
flat_inputs: list[Any],
|
||||
flat_args_dynamic_dims: list[set[int]],
|
||||
graph_input_order: dict[int, int],
|
||||
graph_output_map: dict[int, tuple[str, Any]],
|
||||
fake_mode: Any | None = None,
|
||||
graph_inputs: dict[int, Any] | None = None,
|
||||
) -> None:
|
||||
super().__init__(module)
|
||||
|
||||
assert len(flat_args_dynamic_dims) == len(flat_inputs)
|
||||
|
||||
self.flat_inputs = flat_inputs
|
||||
self.flat_args_dynamic_dims = flat_args_dynamic_dims
|
||||
self.graph_input_order = graph_input_order
|
||||
self.graph_output_map = graph_output_map
|
||||
self.fake_mode = fake_mode
|
||||
self.graph_inputs = graph_inputs or {}
|
||||
|
||||
# Get original placeholders and output
|
||||
self.placeholders = [n for n in module.graph.nodes if n.op == "placeholder"]
|
||||
self.output_node = next(n for n in module.graph.nodes if n.op == "output")
|
||||
|
||||
# Create new flattened input placeholders
|
||||
self.new_input_nodes: dict[int, torch.fx.Node] = {}
|
||||
self._create_flattened_inputs()
|
||||
|
||||
# Iterator for replacing old placeholders
|
||||
self.old_to_new_mapping = {}
|
||||
self._create_placeholder_mapping()
|
||||
|
||||
def _create_flattened_inputs(self) -> None:
|
||||
"""Create new placeholder nodes for flattened inputs with proper fake tensors."""
|
||||
for i in range(len(self.flat_inputs)):
|
||||
placeholder = super().placeholder(f"arg_{i}", (), {})
|
||||
|
||||
# Check if this user input (index i) maps to a graph placeholder
|
||||
if i in self.graph_input_order:
|
||||
# graph_input_order[i] gives us which graph placeholder this user input corresponds to
|
||||
graph_placeholder_idx = self.graph_input_order[i]
|
||||
if graph_placeholder_idx < len(self.placeholders):
|
||||
orig_placeholder = self.placeholders[graph_placeholder_idx]
|
||||
# Copy other metadata but not "val" yet
|
||||
for key, value in orig_placeholder.meta.items():
|
||||
if key != "val":
|
||||
placeholder.node.meta[key] = value
|
||||
|
||||
# Always ensure we have proper "val" metadata from fake tensor
|
||||
if self.fake_mode is not None and isinstance(
|
||||
self.flat_inputs[i], torch.Tensor
|
||||
):
|
||||
placeholder.node.meta["val"] = self.fake_mode.from_tensor(
|
||||
self.flat_inputs[i],
|
||||
symbolic_context=StatelessSymbolicContext(
|
||||
dynamic_sizes=[
|
||||
(
|
||||
DimDynamic.DYNAMIC
|
||||
if d in self.flat_args_dynamic_dims[i]
|
||||
else DimDynamic.STATIC
|
||||
)
|
||||
for d in range(len(self.flat_inputs[i].shape))
|
||||
],
|
||||
constraint_sizes=[None] * len(self.flat_inputs[i].shape),
|
||||
),
|
||||
)
|
||||
elif hasattr(self.flat_inputs[i], "val"): # _IntWrapper case
|
||||
placeholder.node.meta["val"] = self.flat_inputs[i].val
|
||||
else:
|
||||
placeholder.node.meta["val"] = self.flat_inputs[i]
|
||||
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
self.new_input_nodes[i] = placeholder
|
||||
|
||||
def _create_placeholder_mapping(self) -> None:
|
||||
"""Create mapping from old placeholders to new ones."""
|
||||
# graph_input_order maps: user_input_index -> graph_placeholder_index
|
||||
# We need to create: old_graph_placeholder -> new_user_input_placeholder
|
||||
for user_input_idx, graph_placeholder_idx in self.graph_input_order.items():
|
||||
if graph_placeholder_idx < len(self.placeholders):
|
||||
old_placeholder = self.placeholders[graph_placeholder_idx]
|
||||
new_placeholder = self.new_input_nodes[user_input_idx]
|
||||
self.old_to_new_mapping[old_placeholder] = new_placeholder
|
||||
|
||||
def placeholder(
|
||||
self, target: Target, args: tuple[Argument, ...], kwargs: dict[str, Any]
|
||||
) -> Any:
|
||||
"""Replace old placeholders with new flattened ones."""
|
||||
# Return the corresponding new placeholder
|
||||
if self.current_node in self.old_to_new_mapping:
|
||||
new_arg = self.old_to_new_mapping[self.current_node]
|
||||
|
||||
# Copy over additional metadata from current node, but don't overwrite "val"
|
||||
for key in ["tensor_dict", "example_value", "unbacked_bindings"]:
|
||||
if key in self.current_node.meta:
|
||||
new_arg.node.meta[key] = self.current_node.meta[key]
|
||||
|
||||
# Only copy "val" if we don't already have a good one
|
||||
if "val" in self.current_node.meta and "val" not in new_arg.node.meta:
|
||||
new_arg.node.meta["val"] = self.current_node.meta["val"]
|
||||
|
||||
return new_arg
|
||||
else:
|
||||
# Convert captured objects (e.g., opaque objects from closures) to
|
||||
# get_attr nodes
|
||||
placeholder_idx = self.placeholders.index(self.current_node)
|
||||
if placeholder_idx in self.graph_inputs:
|
||||
source = self.graph_inputs[placeholder_idx]
|
||||
if not isinstance(source, torch._dynamo.source.GetItemSource):
|
||||
example_val = self.current_node.meta.get(
|
||||
"val"
|
||||
) or self.current_node.meta.get("example_value")
|
||||
if example_val is not None:
|
||||
attr_name = f"_captured_{placeholder_idx}"
|
||||
if isinstance(example_val, torch.Tensor):
|
||||
self.module.register_buffer(attr_name, example_val)
|
||||
else:
|
||||
setattr(self.module, attr_name, example_val)
|
||||
result = self.tracer.create_proxy("get_attr", attr_name, (), {})
|
||||
result.node.meta = self.current_node.meta.copy()
|
||||
result.node.meta["val"] = example_val
|
||||
return result
|
||||
return super().placeholder(target, args, kwargs)
|
||||
|
||||
def output(
|
||||
self, target: Target, args: Sequence[Any], kwargs: dict[str, Any]
|
||||
) -> Any:
|
||||
"""Transform output according to graph_output_map."""
|
||||
original_outputs = args[0]
|
||||
|
||||
# Build new output list based on graph_output_map
|
||||
new_outputs = []
|
||||
for i in sorted(self.graph_output_map.keys()):
|
||||
output_type, val = self.graph_output_map[i]
|
||||
|
||||
if output_type == "graph_out":
|
||||
new_outputs.append(original_outputs[val])
|
||||
elif output_type == "input":
|
||||
input_idx = val.index
|
||||
new_outputs.append(self.new_input_nodes[input_idx])
|
||||
elif output_type == "constant":
|
||||
new_outputs.append(val)
|
||||
|
||||
return super().output(target, (tuple(new_outputs),), {})
|
||||
|
||||
def run_node(self, n: Node) -> Any:
|
||||
"""Run node transformation and preserve metadata."""
|
||||
self.current_node = n
|
||||
result = super().run_node(n)
|
||||
|
||||
# Copy important metadata
|
||||
if hasattr(result, "node") and result.node is not n:
|
||||
for key in ["val", "example_value", "unbacked_bindings"]:
|
||||
if key in n.meta:
|
||||
result.node.meta[key] = n.meta[key]
|
||||
|
||||
# Preserve node names (except output)
|
||||
if n.op != "output" and hasattr(n, "name"):
|
||||
result.node._rename(n.name)
|
||||
|
||||
return result
|
||||
|
||||
def transform(self) -> torch.fx.GraphModule:
|
||||
"""Perform the graph transformation and copy module metadata."""
|
||||
result_gm = super().transform()
|
||||
|
||||
# Copy module metadata like the original implementation
|
||||
if hasattr(self.module, "meta"):
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
if "dynamo_flat_name_to_original_fqn" in self.module.meta:
|
||||
# pyrefly: ignore [bad-index]
|
||||
result_gm.meta["dynamo_flat_name_to_original_fqn"] = self.module.meta[
|
||||
# pyrefly: ignore [bad-index]
|
||||
"dynamo_flat_name_to_original_fqn"
|
||||
]
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
if "dynamo_compile_id" in self.module.meta:
|
||||
# pyrefly: ignore [bad-index]
|
||||
result_gm.meta["dynamo_compile_id"] = self.module.meta[
|
||||
# pyrefly: ignore [bad-index]
|
||||
"dynamo_compile_id"
|
||||
]
|
||||
|
||||
return result_gm
|
||||
|
||||
|
||||
def _suggest_or_raise_constraint_violation(
|
||||
module_to_trace: torch.nn.Module,
|
||||
orig_callable: Callable[..., Any],
|
||||
fake_mode: Optional["FakeTensorMode"],
|
||||
graph_capture_output: CaptureOutput,
|
||||
args: Any,
|
||||
kwargs: Any,
|
||||
dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None,
|
||||
) -> None:
|
||||
constraint_violation_error = None
|
||||
try:
|
||||
# Check if we have any constraint violations
|
||||
fn, _ = get_traced_fn(module_to_trace)
|
||||
graph_capture_output.graph_capture_output.build_guards(fn.__code__)
|
||||
except ConstraintViolationError as e:
|
||||
constraint_violation_error = e
|
||||
|
||||
if (
|
||||
(shape_env := getattr(fake_mode, "shape_env", None)) is not None
|
||||
and (dim_constraints := shape_env.dim_constraints) is not None
|
||||
and not isinstance(
|
||||
module_to_trace.forward,
|
||||
torch._ops.OpOverloadPacket | torch._ops.OpOverload,
|
||||
)
|
||||
):
|
||||
dim_constraints.solve()
|
||||
|
||||
forced_specializations = dim_constraints.forced_specializations()
|
||||
|
||||
msg = dim_constraints.prettify_results(
|
||||
inspect.signature(orig_callable), # type: ignore[attr-defined]
|
||||
dynamic_shapes,
|
||||
constraint_violation_error,
|
||||
forced_specializations,
|
||||
)
|
||||
if constraint_violation_error:
|
||||
if constraint_violation_error.args:
|
||||
constraint_violation_error.args = (
|
||||
constraint_violation_error.args[0] + msg,
|
||||
)
|
||||
else:
|
||||
constraint_violation_error.args = (msg,)
|
||||
else:
|
||||
if forced_specializations:
|
||||
constraint_violation_error = ConstraintViolationError(msg)
|
||||
else:
|
||||
log.info(
|
||||
"Summary of dimension constraints:%s",
|
||||
msg,
|
||||
)
|
||||
|
||||
# Error if we have any constraints on static values
|
||||
|
||||
for k in shape_env.var_to_range:
|
||||
if isinstance(k, sympy.Integer):
|
||||
constraint_violation_error = ConstraintViolationError(
|
||||
f"{''.join(traceback.format_list(shape_env.var_to_stack[k]))}\n"
|
||||
"It appears that you're trying to set a constraint on a "
|
||||
f"value which we evaluated to have a static value of {k}. "
|
||||
'Set TORCH_LOGS="+export" for more information.'
|
||||
)
|
||||
if constraint_violation_error:
|
||||
constraint_violation_error = post_process_error_msg(
|
||||
constraint_violation_error, orig_callable, args, kwargs
|
||||
)
|
||||
raise constraint_violation_error
|
||||
|
||||
|
||||
def _normalize_shuffle_graph(shuffle_gm: torch.fx.GraphModule) -> None:
|
||||
shuffle_gm.graph.eliminate_dead_code()
|
||||
shuffle_gm.recompile()
|
||||
for name, buffer in list(shuffle_gm.named_buffers()):
|
||||
delattr(shuffle_gm, name)
|
||||
setattr(shuffle_gm, name, buffer)
|
||||
|
||||
|
||||
def normalize_graph_module(gm: torch.fx.GraphModule) -> None:
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "placeholder":
|
||||
node.meta["val"] = node.meta["example_value"]
|
||||
|
||||
|
||||
class InputProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
root: object,
|
||||
num_args: int,
|
||||
kwarg_names: list[str],
|
||||
) -> None:
|
||||
self.root = root
|
||||
self.num_args = num_args
|
||||
self.kwarg_names = kwarg_names
|
||||
|
||||
def __call__(
|
||||
self, inputs: tuple[object, ...]
|
||||
) -> tuple[tuple[object, ...], dict[str, object]]:
|
||||
args = inputs
|
||||
# pyrefly: ignore [implicit-any]
|
||||
kwargs = {}
|
||||
if len(args) > self.num_args:
|
||||
kwargs = dict(zip(self.kwarg_names, args[self.num_args :]))
|
||||
args = args[: self.num_args]
|
||||
if self.root is not None:
|
||||
if isinstance(self.root, torch.fx.GraphModule):
|
||||
assert isinstance(self.root.graph._codegen, _DynamoBytecodeCodeGen)
|
||||
assert hasattr(
|
||||
self.root.graph._codegen.dynamo_bytecode_flatten, "input_processor"
|
||||
)
|
||||
assert (
|
||||
self.root.graph._codegen.dynamo_bytecode_flatten.input_processor
|
||||
is self
|
||||
)
|
||||
args = (self.root, *args)
|
||||
return args, kwargs
|
||||
|
||||
|
||||
class Yield(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class DynamoBytecodeFlatten:
|
||||
def __init__(
|
||||
self,
|
||||
input_processor: InputProcessor,
|
||||
out: CaptureOutput,
|
||||
f_globals: dict[str, object],
|
||||
) -> None:
|
||||
self.input_processor = input_processor
|
||||
self.out = out
|
||||
self.f_globals = f_globals
|
||||
self.gm_inputs: tuple[Any, ...] | None = None
|
||||
|
||||
@dynamo_disable(reason="do not trace internal dynamo graph capture") # type: ignore[misc]
|
||||
def __call__(self, *inputs: object) -> object:
|
||||
def backend_dummy(*example_inputs: object) -> None:
|
||||
self.gm_inputs = example_inputs
|
||||
raise Yield
|
||||
|
||||
args, kwargs = self.input_processor(inputs)
|
||||
try:
|
||||
self.out.forward_callable(
|
||||
compiled_fn=backend_dummy, extra_globals=self.f_globals
|
||||
)(*args, **kwargs)
|
||||
except Yield:
|
||||
assert self.gm_inputs is not None
|
||||
return self.gm_inputs
|
||||
raise RuntimeError
|
||||
|
||||
|
||||
class DynamoBytecodeUnflatten:
|
||||
def __init__(
|
||||
self,
|
||||
input_processor: InputProcessor,
|
||||
out: CaptureOutput,
|
||||
f_globals: dict[str, object],
|
||||
) -> None:
|
||||
self.input_processor = input_processor
|
||||
self.out = out
|
||||
self.f_globals = f_globals
|
||||
|
||||
@dynamo_disable(reason="do not trace internal dynamo graph capture") # type: ignore[misc]
|
||||
def __call__(
|
||||
self, flat_outs: Sequence[object], inputs: tuple[object, ...]
|
||||
) -> object:
|
||||
def backend_dummy(*example_inputs: object) -> Sequence[object]:
|
||||
return flat_outs
|
||||
|
||||
args, kwargs = self.input_processor(inputs)
|
||||
with torch._C._DisableTorchDispatch():
|
||||
results = self.out.forward_callable(
|
||||
compiled_fn=backend_dummy, extra_globals=self.f_globals
|
||||
)(*args, **kwargs)
|
||||
return results
|
||||
|
||||
|
||||
def create_fx_graph_from_captured_output(
|
||||
out: CaptureOutput, mod: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
|
||||
) -> torch.fx.GraphModule:
|
||||
assert out.backend_input is not None
|
||||
backend_input = out.backend_input
|
||||
|
||||
_, root = torch._dynamo.convert_frame.get_traced_fn(mod)
|
||||
|
||||
flat_real_args = pytree.tree_leaves((args, kwargs))
|
||||
torch._dynamo.eval_frame.check_user_input_output(
|
||||
flat_real_args, UserErrorType.INVALID_INPUT
|
||||
)
|
||||
f_globals = out.graph_capture_output.f_globals
|
||||
|
||||
graph_module = backend_input.graph_module
|
||||
if isinstance(root, torch.nn.Module):
|
||||
graph_module._parameters = root._parameters
|
||||
graph_module._buffers = root._buffers
|
||||
assert all(not hasattr(graph_module, m) for m in root._modules)
|
||||
graph_module._modules.update(root._modules)
|
||||
graph_module._non_persistent_buffers_set = root._non_persistent_buffers_set
|
||||
if sys.version_info >= (3, 14):
|
||||
import annotationlib # added in 3.14
|
||||
|
||||
annotations = annotationlib.get_annotations(torch.nn.Module)
|
||||
else:
|
||||
annotations = getattr(torch.nn.Module, "__annotations__", None)
|
||||
for name, value in root.__dict__.items():
|
||||
if annotations and name not in annotations:
|
||||
graph_module.__dict__[name] = value
|
||||
graph_module._forward_hooks = root._forward_hooks.copy()
|
||||
graph_module._forward_pre_hooks = root._forward_pre_hooks.copy()
|
||||
graph_module._backward_hooks = root._backward_hooks.copy()
|
||||
graph_module._backward_pre_hooks = root._backward_pre_hooks.copy()
|
||||
if graph_module._forward_hooks or graph_module._forward_pre_hooks:
|
||||
# Even forward hooks are traced through, they still capture a bunch
|
||||
# of state through closure. We need to make sure these data are
|
||||
# accessible through the captured module (but the hooks should be
|
||||
# disabled).
|
||||
assert getattr(graph_module, "_wrapped_call", None) is not None
|
||||
assert isinstance(
|
||||
graph_module._wrapped_call, torch.fx.graph_module._WrappedCall
|
||||
)
|
||||
assert graph_module._wrapped_call.cls_call is None
|
||||
|
||||
def dynamo_wrapped_call(self, *args: object, **kwargs: object) -> object:
|
||||
assert "forward" not in self.__dict__
|
||||
|
||||
fwd_hooks = self._forward_hooks
|
||||
fwd_pre_hooks = self._forward_pre_hooks
|
||||
original_forward = type(self).forward
|
||||
|
||||
def patched_forward(self, *args: object, **kwargs: object) -> object:
|
||||
self._forward_hooks = fwd_hooks
|
||||
self._forward_pre_hooks = fwd_pre_hooks
|
||||
return original_forward(self, *args, **kwargs)
|
||||
|
||||
try:
|
||||
self.forward = types.MethodType(patched_forward, self)
|
||||
# pyrefly: ignore [implicit-any]
|
||||
self._forward_hooks = {}
|
||||
# pyrefly: ignore [implicit-any]
|
||||
self._forward_pre_hooks = {}
|
||||
# pyrefly: ignore [invalid-argument]
|
||||
return super(type(self), self).__call__(*args, **kwargs)
|
||||
finally:
|
||||
self.__dict__.pop("forward")
|
||||
self._forward_hooks = fwd_hooks
|
||||
self._forward_pre_hooks = fwd_pre_hooks
|
||||
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
graph_module._wrapped_call.cls_call = dynamo_wrapped_call
|
||||
|
||||
root = graph_module if isinstance(root, torch.nn.Module) else root
|
||||
input_processor = InputProcessor(root, len(args), list(kwargs.keys()))
|
||||
dynamo_bytecode_flatten = DynamoBytecodeFlatten(input_processor, out, f_globals)
|
||||
dynamo_bytecode_unflatten = DynamoBytecodeUnflatten(input_processor, out, f_globals)
|
||||
|
||||
graph_module.graph._codegen = _DynamoBytecodeCodeGen(
|
||||
argument_names(inspect.signature(mod), args, kwargs),
|
||||
dynamo_bytecode_flatten,
|
||||
dynamo_bytecode_unflatten,
|
||||
) # type: ignore[attr-defined]
|
||||
normalize_graph_module(graph_module)
|
||||
assert not hasattr(graph_module, "_dynamo_bytecode_flatten")
|
||||
assert not hasattr(graph_module, "_dynamo_bytecode_unflatten")
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
graph_module._dynamo_bytecode_flatten = dynamo_bytecode_flatten
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
graph_module._dynamo_bytecode_unflatten = dynamo_bytecode_unflatten
|
||||
delattr(graph_module, "_param_name_to_source")
|
||||
graph_module.recompile()
|
||||
graph_module.meta["module_call_specs"] = (
|
||||
out.graph_capture_output.output_graph.export_metadata.module_call_spec
|
||||
)
|
||||
assert out.backend_input is not None
|
||||
graph_module.meta["fake_mode"] = out.backend_input.fake_mode # type: ignore[attr-defined]
|
||||
graph_module.meta["fake_mode"].allow_non_fake_inputs = True
|
||||
tracing_context = TracingContext(graph_module.meta["fake_mode"])
|
||||
tracing_context.tensor_to_context = out.backend_input.tensor_to_context # type: ignore[attr-defined]
|
||||
graph_module.meta["tracing_context"] = tracing_context
|
||||
return graph_module
|
||||
|
||||
|
||||
class _DynamoBytecodeCodeGen(torch.fx.graph.CodeGen):
|
||||
def __init__(
|
||||
self,
|
||||
orig_arg_names: list[str],
|
||||
# pyrefly: ignore [implicit-any]
|
||||
dynamo_bytecode_flatten: Callable,
|
||||
# pyrefly: ignore [implicit-any]
|
||||
dynamo_bytecode_unflatten: Callable,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.orig_arg_names = orig_arg_names
|
||||
self.dynamo_bytecode_flatten = dynamo_bytecode_flatten
|
||||
self.dynamo_bytecode_unflatten = dynamo_bytecode_unflatten
|
||||
self.wrap_tuple = False
|
||||
self._inputs: tuple[Any, ...] | None = None
|
||||
|
||||
def process_inputs(self, *inputs: Any) -> Any:
|
||||
self._inputs = inputs
|
||||
results = self.dynamo_bytecode_flatten(*inputs)
|
||||
return results
|
||||
|
||||
def process_outputs(self, outputs: Any) -> Any:
|
||||
results = self.dynamo_bytecode_unflatten(outputs, self._inputs)
|
||||
if self.wrap_tuple:
|
||||
results = (results,)
|
||||
self._inputs = None
|
||||
return results
|
||||
|
||||
def gen_fn_def(
|
||||
self,
|
||||
free_vars: list[str],
|
||||
maybe_return_annotation: str,
|
||||
*,
|
||||
expanded_def: bool = False,
|
||||
) -> str:
|
||||
fn_args = self.orig_arg_names
|
||||
has_orig_self = (fn_args[0] == "self") if len(fn_args) > 0 else False
|
||||
if has_orig_self:
|
||||
free_vars.insert(0, "self")
|
||||
fn_definition = super().gen_fn_def(
|
||||
fn_args[:], maybe_return_annotation, expanded_def=expanded_def
|
||||
)
|
||||
|
||||
if len(free_vars) > 0: # pytree has placeholders in it
|
||||
fn_definition += self.gen_var_bindings(fn_args, free_vars, expanded_def)
|
||||
return fn_definition
|
||||
|
||||
def gen_var_bindings(
|
||||
self, fn_args: list[str], free_vars: list[str], expanded_def: bool
|
||||
) -> str:
|
||||
without_annotation = [x.split(":")[0].split("#")[0] for x in free_vars]
|
||||
if len(fn_args) == 0:
|
||||
fn_signature = ""
|
||||
elif len(fn_args) == 1:
|
||||
fn_signature = f"{fn_args[0]}, "
|
||||
else:
|
||||
fn_signature = f"{', '.join(fn_args)}"
|
||||
return f"""
|
||||
_fn_args = ({fn_signature})
|
||||
{", ".join(without_annotation)}, = self._dynamo_bytecode_flatten(*_fn_args)"""
|
||||
|
||||
def generate_output(
|
||||
self,
|
||||
output_args: torch.fx.node.Argument,
|
||||
*,
|
||||
descs: object | None = None,
|
||||
repr_fn: Any | None = None,
|
||||
) -> str:
|
||||
if repr_fn is None:
|
||||
repr_fn = repr
|
||||
# pyrefly: ignore [not-iterable]
|
||||
returned = f"self._dynamo_bytecode_unflatten(({', '.join([repr_fn(a) for a in output_args])},), _fn_args)"
|
||||
if self.wrap_tuple:
|
||||
returned = f"({returned},)"
|
||||
return f"return {returned}"
|
||||
|
||||
|
||||
def dynamo_graph_capture_for_export(
|
||||
fn: Callable[..., Any],
|
||||
constraints: list[Constraint] | None = None,
|
||||
) -> Callable[..., Any]:
|
||||
if isinstance(fn, torch._ops.OpOverload):
|
||||
|
||||
def default_annotation(arg: torch.Argument) -> str:
|
||||
if arg.has_default_value():
|
||||
return f"={arg.default_value!r}"
|
||||
return ""
|
||||
|
||||
has_kwarg_only = False
|
||||
arg_list = []
|
||||
for arg in fn._schema.arguments:
|
||||
if arg.kwarg_only and not has_kwarg_only:
|
||||
has_kwarg_only = True
|
||||
arg_list.append("*")
|
||||
arg_list.append(arg.name + default_annotation(arg))
|
||||
func_str = f"""
|
||||
def op_overload_wrapper({", ".join(arg_list)}):
|
||||
return op({", ".join([f"{arg.name}={arg.name}" for arg in fn._schema.arguments])})
|
||||
"""
|
||||
out = {}
|
||||
exec(func_str, {"op": fn}, out)
|
||||
fn = out["op_overload_wrapper"] # type: ignore[assignment]
|
||||
|
||||
def inner(*args: Any, **kwargs: Any) -> Any:
|
||||
assert not torch._dynamo.config.install_free_tensors
|
||||
with (
|
||||
_compiling_state_context(),
|
||||
torch._dynamo.config.patch(
|
||||
replay_side_effects=False, side_effect_replay_policy="warn"
|
||||
),
|
||||
get_metrics_context(),
|
||||
dynamo_timed("fullgraph_capture"),
|
||||
):
|
||||
out = fullgraph_capture(
|
||||
fn,
|
||||
args,
|
||||
kwargs,
|
||||
constraints=constraints,
|
||||
)
|
||||
graph_module = create_fx_graph_from_captured_output(out, fn, args, kwargs)
|
||||
return graph_module
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
def _dynamo_graph_capture_for_export(
|
||||
mod: Callable[..., Any],
|
||||
*,
|
||||
constraints: list[Constraint] | None = None,
|
||||
dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = None,
|
||||
) -> Callable[..., torch.fx.GraphModule]:
|
||||
"""
|
||||
Improved dynamo graph capture using transformer approach with proper fake tensor handling.
|
||||
|
||||
This function creates a capture instance that handles:
|
||||
1. PyTree flattening/unflattening with proper input ordering
|
||||
2. Dynamo graph capture with export-specific context
|
||||
3. FX graph transformation for export compatibility
|
||||
4. Proper fake tensor metadata preservation
|
||||
5. Dynamic dimension constraint handling
|
||||
|
||||
Notable improvements over manual approach:
|
||||
- Uses FX Transformer for cleaner graph manipulation
|
||||
- Properly handles fake tensor metadata and dynamic dimensions
|
||||
- Preserves all necessary metadata for export
|
||||
- More robust error handling and edge case management
|
||||
|
||||
TODO:
|
||||
1. Are we actually gonna run the bytecode?
|
||||
2. Need to attach guards
|
||||
"""
|
||||
|
||||
_dynamic_shapes = dynamic_shapes
|
||||
_constraints = constraints
|
||||
|
||||
def inner(*args: Any, **kwargs: Any) -> torch.fx.GraphModule:
|
||||
# This sets the is_exporting flag when building guards.
|
||||
with _compiling_state_context():
|
||||
flat_inputs, in_spec = pytree.tree_flatten((args, kwargs))
|
||||
check_user_input_output(flat_inputs, UserErrorType.INVALID_INPUT)
|
||||
module_to_trace = ModuleToTrace(mod, in_spec)
|
||||
orig_callable = mod.forward if isinstance(mod, torch.nn.Module) else mod
|
||||
|
||||
constraints: list[Constraint] | None = _constraints
|
||||
dynamic_shapes: dict[str, Any] | tuple[Any] | list[Any] | None = (
|
||||
_dynamic_shapes
|
||||
)
|
||||
|
||||
from . import reset # type: ignore[attr-defined]
|
||||
|
||||
reset()
|
||||
|
||||
dynamo_config_ctx = torch._dynamo.config.patch(
|
||||
specialize_int=True,
|
||||
specialize_float=True,
|
||||
assume_static_by_default=True,
|
||||
automatic_dynamic_shapes=False,
|
||||
capture_dynamic_output_shape_ops=True,
|
||||
capture_scalar_outputs=True,
|
||||
constant_fold_autograd_profiler_enabled=True,
|
||||
log_graph_in_out_metadata=True,
|
||||
# install_free_tensors ensures that params and buffers are still
|
||||
# added as graph attributes, and makes Dynamo emits graphs that
|
||||
# follow export pytree-able input requirements In future, if we
|
||||
# fully rely on bytecode for the runtime, we can turn this flag
|
||||
# off.
|
||||
install_free_tensors=torch._dynamo.config.install_free_tensors_for_export,
|
||||
)
|
||||
|
||||
with (
|
||||
get_metrics_context(),
|
||||
dynamo_timed("fullgraph_capture"),
|
||||
dynamo_config_ctx,
|
||||
):
|
||||
out = fullgraph_capture(
|
||||
module_to_trace,
|
||||
tuple(flat_inputs),
|
||||
constraints=_constraints,
|
||||
_is_export_deprecated_do_not_use=True,
|
||||
)
|
||||
|
||||
assert out.graph_capture_output.output_graph is not None
|
||||
|
||||
example_inputs: list[Any] = []
|
||||
if out.backend_input is not None:
|
||||
graph = out.backend_input.graph_module
|
||||
fake_mode = out.backend_input.fake_mode
|
||||
example_inputs = out.backend_input.example_inputs
|
||||
else:
|
||||
graph = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
|
||||
graph.graph.output(None)
|
||||
graph.recompile()
|
||||
fake_mode = None
|
||||
|
||||
_suggest_or_raise_constraint_violation(
|
||||
module_to_trace,
|
||||
orig_callable,
|
||||
fake_mode,
|
||||
out,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes,
|
||||
)
|
||||
|
||||
# Extract export metadata from the new location
|
||||
export_metadata = out.graph_capture_output.output_graph.export_metadata
|
||||
graph_inputs = export_metadata.graph_input_idx_to_local_source
|
||||
graph_output_map = export_metadata.output_return_type
|
||||
out_spec = export_metadata.out_spec
|
||||
module_call_spec = export_metadata.module_call_spec
|
||||
|
||||
# Compute dynamic dimensions for each input based on constraints
|
||||
flat_args_dynamic_dims = [
|
||||
{
|
||||
c.dim
|
||||
for c in (constraints or ())
|
||||
if (
|
||||
c.t_id == id(x)
|
||||
and not isinstance(c, _RelaxedConstraint)
|
||||
and c.constraint_range.vr.lower != c.constraint_range.vr.upper
|
||||
)
|
||||
}
|
||||
for x in flat_inputs
|
||||
]
|
||||
|
||||
# Create input order mapping from dynamo's internal order to user order
|
||||
# Only process inputs that come from function arguments (GetItemSource).
|
||||
# Skip inputs that come from other sources like closures (e.g., captured
|
||||
# opaque objects like DeviceMesh).
|
||||
graph_input_order: dict[int, int] = {}
|
||||
for inp in graph_inputs:
|
||||
source = graph_inputs[inp]
|
||||
if isinstance(source, torch._dynamo.source.GetItemSource):
|
||||
graph_input_order[source.index] = len(graph_input_order)
|
||||
|
||||
for real_idx, graph_idx in graph_input_order.items():
|
||||
flat_inputs[real_idx] = example_inputs[graph_idx]
|
||||
|
||||
# Use FX transformer to rebuild the graph cleanly
|
||||
transformed_graph = DynamoGraphTransformer(
|
||||
graph,
|
||||
flat_inputs,
|
||||
flat_args_dynamic_dims,
|
||||
graph_input_order,
|
||||
graph_output_map,
|
||||
fake_mode,
|
||||
graph_inputs,
|
||||
).transform()
|
||||
|
||||
# Set up PyTree codegen for proper input/output handling
|
||||
transformed_graph.graph._codegen = _PyTreeCodeGen(
|
||||
_PyTreeInfo(
|
||||
argument_names(inspect.signature(orig_callable), args, kwargs), # type: ignore[attr-defined, arg-type]
|
||||
in_spec,
|
||||
out_spec,
|
||||
)
|
||||
)
|
||||
transformed_graph.recompile()
|
||||
|
||||
clean_nn_module_stack_and_source_fn(transformed_graph, True)
|
||||
clean_export_root(transformed_graph)
|
||||
|
||||
transformed_graph.meta["module_call_specs"] = module_call_spec
|
||||
transformed_graph.meta["fake_mode"] = fake_mode
|
||||
|
||||
return transformed_graph
|
||||
|
||||
return inner
|
||||
@@ -0,0 +1,32 @@
|
||||
USER_ERROR = [
|
||||
"Your code may result in an error when running in eager. "
|
||||
"Please double check that your code doesn't contain a similar error when actually running eager/uncompiled. "
|
||||
'You can do this by removing the `torch.compile` call, or by using `torch.compiler.set_stance("force_eager")`. '
|
||||
]
|
||||
DYNAMO_BUG = [
|
||||
"This is likely to be a Dynamo bug. Please report an issue to PyTorch.",
|
||||
]
|
||||
DIFFICULT = [
|
||||
"This graph break may be difficult to debug. Please report an issue to PyTorch for assistance.",
|
||||
]
|
||||
FUNDAMENTAL = [
|
||||
"This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through "
|
||||
"your code. Consider finding a workaround.",
|
||||
]
|
||||
SUPPORTABLE = [
|
||||
"It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you "
|
||||
"encounter this graph break often and it is causing performance issues.",
|
||||
]
|
||||
CAUSED_BY_EARLIER_GRAPH_BREAK = [
|
||||
"This graph break may have been caused by an earlier graph break. Resolving the earlier graph break may resolve this one.",
|
||||
]
|
||||
INFERENCE_MODE = [
|
||||
"Avoid using `tensor.is_inference()` and `torch.is_inference_mode_enabled()` in your compile code. "
|
||||
"This is primarily used in conjunction with `torch.inference_mode`. Consider using `torch.no_grad` instead "
|
||||
"because `torch.no_grad` leads to same improvements as `inference_mode` when `torch.compile` is used.",
|
||||
]
|
||||
SPARSE_TENSOR = [
|
||||
"Sparse tensor operations are not yet fully supported in torch.compile with fullgraph=True. "
|
||||
"Consider using fullgraph=False to allow graph breaks, or move sparse tensor creation "
|
||||
"outside the compiled region.",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,116 @@
|
||||
import weakref
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from torch._dynamo.source import Source
|
||||
|
||||
|
||||
PyCodegen = Any
|
||||
|
||||
# This file is to handle types that we don't want to support
|
||||
# as explicit FX graph inputs. This uses a sidetable which
|
||||
# we populate in bytecode and is loaded during graph execution
|
||||
|
||||
# We use a dynamo-generated index as a level of indirection
|
||||
# this allows us to register objects externally in pre-graph bytecode that we want
|
||||
# to pass to the graph, but not support their types as graph inputs
|
||||
index_to_bytecode_constructor: dict[int, Callable[[PyCodegen], None]] = {}
|
||||
|
||||
index_to_external_object_weakref: dict[int, weakref.ReferenceType[Any]] = {}
|
||||
|
||||
keep_alive: list[Any] = []
|
||||
|
||||
|
||||
def has_user_objects() -> bool:
|
||||
return bool(index_to_bytecode_constructor)
|
||||
|
||||
|
||||
def stash_graph_created_object(obj: Any) -> Any:
|
||||
keep_alive.append(obj)
|
||||
return obj
|
||||
|
||||
|
||||
CURRENT_STREAM_INDEX = 0
|
||||
|
||||
|
||||
def set_external_object_by_index(index: int, value: Any) -> None:
|
||||
"""Update an entry in the external object registry at runtime."""
|
||||
keep_alive.append(value)
|
||||
index_to_external_object_weakref[index] = weakref.ref(value)
|
||||
|
||||
|
||||
def get_external_object_by_index(index: int) -> Any:
|
||||
assert index in index_to_external_object_weakref, (
|
||||
"Index not registered in index_to_user_object_weakref"
|
||||
)
|
||||
obj = index_to_external_object_weakref[index]()
|
||||
assert obj is not None, "User object is no longer alive"
|
||||
return index_to_external_object_weakref[index]()
|
||||
|
||||
|
||||
def store_user_object_weakrefs(*args: Any) -> None:
|
||||
global index_to_external_object_weakref
|
||||
index_to_external_object_weakref.clear()
|
||||
index_to_external_object_weakref.update(
|
||||
{i: weakref.ref(arg) for i, arg in enumerate(args)}
|
||||
)
|
||||
|
||||
|
||||
def reset_user_object_tracking() -> None:
|
||||
index_to_bytecode_constructor.clear()
|
||||
index_to_external_object_weakref.clear()
|
||||
keep_alive.clear()
|
||||
|
||||
|
||||
def register_graph_created_object(
|
||||
example_value: Any, construct_fn: Callable[[int, PyCodegen], None]
|
||||
) -> int:
|
||||
global index_to_bytecode_constructor
|
||||
global keep_alive
|
||||
keep_alive.append(example_value)
|
||||
index = len(index_to_bytecode_constructor)
|
||||
index_to_bytecode_constructor[index] = lambda cg: construct_fn(index, cg)
|
||||
try:
|
||||
index_to_external_object_weakref[index] = weakref.ref(example_value)
|
||||
except TypeError as e:
|
||||
from .exc import unimplemented
|
||||
|
||||
unimplemented(
|
||||
gb_type="Failed to make weakref to graph-created external object",
|
||||
context=f"user_object: {example_value}",
|
||||
explanation="Object does not allow us to make a weakref to it",
|
||||
hints=[],
|
||||
from_exc=e,
|
||||
)
|
||||
return index
|
||||
|
||||
|
||||
# Register a user object to be used in the graph
|
||||
def register_user_object(value: Any, source: Source) -> int:
|
||||
global index_to_bytecode_constructor
|
||||
index = len(index_to_bytecode_constructor)
|
||||
index_to_bytecode_constructor[index] = lambda cg: cg(source)
|
||||
try:
|
||||
index_to_external_object_weakref[index] = weakref.ref(value)
|
||||
except TypeError as e:
|
||||
from .exc import unimplemented
|
||||
|
||||
unimplemented(
|
||||
gb_type="Failed to make weakref to User Object",
|
||||
context=f"user_object: {value}",
|
||||
explanation="Object does not allow us to make a weakref to it",
|
||||
hints=[],
|
||||
from_exc=e,
|
||||
)
|
||||
return index
|
||||
|
||||
|
||||
# Register a callback so invoke_leaf_function can retrieve nn.Module instances at runtime.
|
||||
# We use a callback pattern instead of having invoke_leaf_function import get_external_object_by_index
|
||||
# directly, because higher-order ops should not depend on dynamo (dynamo depends on them, not vice versa).
|
||||
from torch._higher_order_ops.invoke_leaf_function import (
|
||||
set_leaf_function_module_retriever,
|
||||
)
|
||||
|
||||
|
||||
set_leaf_function_module_retriever(get_external_object_by_index)
|
||||
@@ -0,0 +1,609 @@
|
||||
"""
|
||||
This module implements graph deduplication functionality for TorchDynamo's optimization pipeline.
|
||||
Graph deduplication identifies identical subgraphs in the computational graph and merges them
|
||||
to reduce redundancy and improve performance. The process involves analyzing regions of the graph,
|
||||
identifying structurally equivalent regions, and replacing them with a single shared implementation.
|
||||
This optimization is particularly effective for models with repeated patterns or similar computational
|
||||
structures across different parts of the network.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import operator
|
||||
from collections import defaultdict, deque
|
||||
from collections.abc import Generator, Iterable
|
||||
|
||||
import torch
|
||||
import torch.fx
|
||||
from torch._dynamo import config
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from .graph_region_tracker import Node, Region
|
||||
from .graph_utils import _detect_cycles, _get_flat_args, _get_flat_args_unique
|
||||
|
||||
|
||||
# Represents an index into the region
|
||||
# to select a node and then
|
||||
# an index into that node's
|
||||
# flattened arguments
|
||||
UsageIndex = tuple[int, int]
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
last_node_to_additional_deps: dict[Node, OrderedSet[Node]] | None = None
|
||||
|
||||
|
||||
def apply_graph_deduplication(output_graph) -> dict[str, torch.fx.GraphModule]: # type: ignore[no-untyped-def]
|
||||
"""
|
||||
This is the main entry point for applying the graph deduplication pass. \
|
||||
Deduplication occurs in two phases:
|
||||
1. Subgraph creation:
|
||||
Subgraph creation works by taking one representative region from each region \
|
||||
group and creating a subgraph from it, which will then be used to replace all regions \
|
||||
in the group. This is implemented by first copying all nodes of the region to the new \
|
||||
subgraph and then finding all inputs which are not within the region and creating placeholders \
|
||||
for them. For the outputs, all regions in a region group need to be scanned to ensure the \
|
||||
largest set of outputs is found, and then an output node is created which returns \
|
||||
a tuple of all outputs.
|
||||
|
||||
2. Graph replacement:
|
||||
To replace each region with the extracted subgraph, the node index in the region \
|
||||
and argument index within the node's flattened args and kwargs are recorded once during \
|
||||
subgraph creation. This allows us to determine which (external to the region) nodes and \
|
||||
in which order these nodes are passed as inputs. For the outputs, getitem nodes are created \
|
||||
for each output, and all nodes in the region with external outputs are replaced by the proper \
|
||||
getitem node. Finally, all original nodes are erased (there should be no uses of these \
|
||||
left in the graph).
|
||||
|
||||
The deduplication mutates the output_graph argument in place.
|
||||
|
||||
Returns a mapping of nodes to their subgraph output replacement node to remap outputs
|
||||
when they are created in output_graph.
|
||||
"""
|
||||
|
||||
duplicated_region_groups = output_graph.region_tracker.get_identical_regions(
|
||||
output_graph.graph
|
||||
)
|
||||
node_to_mutated_arg_positions = (
|
||||
output_graph.region_tracker.node_to_mutated_arg_positions
|
||||
)
|
||||
node_to_additional_deps = _populate_additional_deps(
|
||||
output_graph.graph, output_graph.region_tracker.node_to_mutated_arg_positions
|
||||
)
|
||||
|
||||
sub_gms: dict[str, torch.fx.GraphModule] = {}
|
||||
|
||||
for region_group in duplicated_region_groups:
|
||||
inds_with_external_users = _get_all_output_indices(region_group)
|
||||
region = region_group[0]
|
||||
(
|
||||
subgraph,
|
||||
external_node_usages,
|
||||
node_usage_to_tuple_elems,
|
||||
ind_to_tuple_spec,
|
||||
) = _create_subgraph(region, inds_with_external_users)
|
||||
|
||||
# Ignore regions with no args for now, could they possibly be evaluated at compile time?
|
||||
if not list(external_node_usages):
|
||||
continue
|
||||
|
||||
sub_gm = torch.fx.GraphModule(output_graph.nn_modules, subgraph)
|
||||
subgraph_name = output_graph.install_subgraph("subgraph", sub_gm)
|
||||
sub_gms[subgraph_name] = sub_gm
|
||||
with output_graph.graph.inserting_before():
|
||||
get_subgraph_node = output_graph.graph.create_node(
|
||||
"get_attr", subgraph_name, (), {}
|
||||
)
|
||||
|
||||
for region in region_group:
|
||||
_replace_region_with_subgraph(
|
||||
output_graph.graph,
|
||||
region,
|
||||
get_subgraph_node,
|
||||
external_node_usages,
|
||||
node_usage_to_tuple_elems,
|
||||
ind_to_tuple_spec,
|
||||
inds_with_external_users,
|
||||
subgraph_name,
|
||||
node_to_additional_deps,
|
||||
node_to_mutated_arg_positions,
|
||||
)
|
||||
|
||||
# This is to expose the updated node_to_additional_deps to tests
|
||||
global last_node_to_additional_deps
|
||||
last_node_to_additional_deps = node_to_additional_deps
|
||||
|
||||
_stable_topological_sort(
|
||||
output_graph.graph,
|
||||
node_to_additional_deps,
|
||||
)
|
||||
return sub_gms
|
||||
|
||||
|
||||
def _replace_region_with_subgraph(
|
||||
graph: torch.fx.Graph,
|
||||
region: Region,
|
||||
get_subgraph_node: Node,
|
||||
external_node_usages: Iterable[OrderedSet[UsageIndex]],
|
||||
node_usage_to_tuple_elems: dict[UsageIndex, OrderedSet[int]],
|
||||
ind_to_tuple_spec: dict[int, dict[tuple[int, ...], int]],
|
||||
inds_with_external_users: list[int],
|
||||
subgraph_name: str,
|
||||
node_to_additional_deps: dict[Node, OrderedSet[Node]],
|
||||
node_to_mutated_arg_positions: dict[Node, OrderedSet[int]],
|
||||
) -> None:
|
||||
sub_args = []
|
||||
flattened_getitem_nodes: OrderedSet[Node] = OrderedSet()
|
||||
for usages in external_node_usages:
|
||||
usage = next(iter(usages))
|
||||
node_ind, usage_ind = usage
|
||||
node = region[node_ind]
|
||||
flattened_args_kwargs = _get_flat_args(node, {})
|
||||
for user_ind, node_usage_ind in usages:
|
||||
user = region[user_ind]
|
||||
if user in node_to_mutated_arg_positions:
|
||||
if node_usage_ind in node_to_mutated_arg_positions[user]:
|
||||
log.debug(
|
||||
"NYI: Failed to substitute region %s due to mutation", region
|
||||
)
|
||||
return
|
||||
if usage in node_usage_to_tuple_elems:
|
||||
tuple_elems = [region[i] for i in node_usage_to_tuple_elems[usage]]
|
||||
flattened_getitem_nodes.update(tuple_elems)
|
||||
sub_args.extend(tuple_elems)
|
||||
else:
|
||||
sub_args.append(flattened_args_kwargs[usage_ind])
|
||||
|
||||
# Input/Output aliasing not supported in HOPs today
|
||||
# Note: we should use the nodes in the original graph (the region here)
|
||||
# because we use the original traced example values for this check
|
||||
if _has_aliasing(
|
||||
region, sub_args, inds_with_external_users, flattened_getitem_nodes
|
||||
):
|
||||
return
|
||||
|
||||
invoke_args = (get_subgraph_node, subgraph_name, *sub_args)
|
||||
|
||||
invoke_subgraph_node = graph.create_node(
|
||||
"call_function",
|
||||
torch.ops.higher_order.invoke_subgraph,
|
||||
invoke_args, # type: ignore[arg-type]
|
||||
{},
|
||||
)
|
||||
|
||||
ind = 0
|
||||
flattened_output_nodes: OrderedSet[Node] = OrderedSet()
|
||||
for external_user_ind in inds_with_external_users:
|
||||
node = region[external_user_ind]
|
||||
if _is_tuple_node(node):
|
||||
tuple_spec = ind_to_tuple_spec[external_user_ind]
|
||||
flattened_output_nodes.update(
|
||||
_replace_tuple_outputs(
|
||||
node, ind, tuple_spec, invoke_subgraph_node, graph
|
||||
)
|
||||
)
|
||||
ind += len(tuple_spec)
|
||||
else:
|
||||
subgraph_output = graph.create_node(
|
||||
"call_function", operator.getitem, (invoke_subgraph_node, ind), {}
|
||||
)
|
||||
node.replace_all_uses_with(subgraph_output, propagate_meta=True)
|
||||
ind += 1
|
||||
|
||||
# Erase in reverse topological order
|
||||
for node in reversed(region):
|
||||
if node in flattened_getitem_nodes:
|
||||
# Don't erase these, since they will still be used
|
||||
continue
|
||||
|
||||
if node not in flattened_output_nodes:
|
||||
graph.erase_node(node)
|
||||
|
||||
# Remove any nodes with additional deps
|
||||
# This is safe; we've guaranteed that there is
|
||||
# no input mutation, so all additional deps
|
||||
# will be internal to the subgraph
|
||||
node_to_additional_deps.pop(node, None)
|
||||
for deps in node_to_additional_deps.values():
|
||||
try:
|
||||
deps.remove(node)
|
||||
deps.add(invoke_subgraph_node)
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
if config.graph_deduplication_lint:
|
||||
print(_detect_cycles(graph, node_to_additional_deps))
|
||||
_stable_topological_sort(graph, node_to_additional_deps)
|
||||
graph.lint()
|
||||
|
||||
|
||||
def _get_external_inputs(
|
||||
region: Region,
|
||||
) -> dict[Node, OrderedSet[UsageIndex]]:
|
||||
external_node_to_usages = defaultdict[Node, OrderedSet[UsageIndex]](OrderedSet)
|
||||
region_unique = set(region)
|
||||
for node_ind, node in enumerate(region):
|
||||
flattened_args_kwargs = _get_flat_args(node, {})
|
||||
for arg_ind, in_node in enumerate(flattened_args_kwargs):
|
||||
if isinstance(in_node, Node) and in_node not in region_unique:
|
||||
# in_node may occur in multiple nodes' flat_args
|
||||
# track this so we can check if the arg is mutated
|
||||
# Previously, we only needed to track one occurrence
|
||||
# to be able to map that node to a placeholder
|
||||
external_node_to_usages[in_node].add((node_ind, arg_ind))
|
||||
|
||||
return external_node_to_usages
|
||||
|
||||
|
||||
def _get_all_output_indices(regions: list[Region]) -> list[int]:
|
||||
# Scan all regions to get the set of all possible output nodes indices in the region
|
||||
# perhaps we can record this information during region creation for more efficiency?
|
||||
inds_with_external_users: set[int] = set()
|
||||
for region in regions:
|
||||
_get_inds_with_external_users(region, inds_with_external_users)
|
||||
|
||||
return sorted(inds_with_external_users)
|
||||
|
||||
|
||||
def _get_inds_with_external_users(region: Region, inds_unique: set[int]) -> None:
|
||||
for ind, node in enumerate(region):
|
||||
for user in node.users:
|
||||
if user not in region:
|
||||
if ind not in inds_unique:
|
||||
inds_unique.add(ind)
|
||||
|
||||
|
||||
def _create_subgraph(
|
||||
region: Region,
|
||||
inds_with_external_users: list[int],
|
||||
) -> tuple[
|
||||
torch.fx.Graph,
|
||||
list[OrderedSet[UsageIndex]],
|
||||
dict[UsageIndex, OrderedSet[int]],
|
||||
dict[int, dict[tuple[int, ...], int]],
|
||||
]:
|
||||
subgraph: torch.fx.Graph = torch.fx.Graph()
|
||||
external_input_to_usages = _get_external_inputs(region)
|
||||
external_node_usages = list[OrderedSet[UsageIndex]]()
|
||||
region_to_subgraph_node = {}
|
||||
flattened_getitem_nodes: OrderedSet[Node] = OrderedSet()
|
||||
node_usage_to_tuple_elems: dict[UsageIndex, OrderedSet[int]] = {}
|
||||
|
||||
for node, usage_indices in external_input_to_usages.items():
|
||||
# We don't handle tuples as inputs today
|
||||
if _is_tuple_node(node):
|
||||
# If a node is a tuple we will possibly create multiple placeholders for them
|
||||
# and track which nodes we won't copy into the subgraph because they are flattened away
|
||||
# Later, when replacing each region with this subgraph, we will create a getitem node
|
||||
# externally which will perform the flattening on the outer nodes.
|
||||
flattened_node_indices = _get_flattened_node_indices(node, region)
|
||||
for ind in flattened_node_indices:
|
||||
placeholder = subgraph.placeholder(
|
||||
f"supgraph_input_{node.name}_flattened_{ind}"
|
||||
)
|
||||
region_to_subgraph_node[region[ind]] = placeholder
|
||||
flattened_getitem_nodes.add(region[ind])
|
||||
node_usage_to_tuple_elems[next(iter(usage_indices))] = (
|
||||
flattened_node_indices
|
||||
)
|
||||
else:
|
||||
placeholder = subgraph.placeholder(f"subgraph_input_{node.name}")
|
||||
region_to_subgraph_node[node] = placeholder
|
||||
|
||||
external_node_usages.append(usage_indices)
|
||||
|
||||
def map_arg(node: Node) -> Node:
|
||||
if node in region_to_subgraph_node:
|
||||
return region_to_subgraph_node[node]
|
||||
else:
|
||||
return node
|
||||
|
||||
def copy_to_subgraph(node: Node) -> Node:
|
||||
subgraph_node = subgraph.node_copy(node, lambda old: map_arg(old))
|
||||
region_to_subgraph_node[node] = subgraph_node
|
||||
return subgraph_node
|
||||
|
||||
output_list = []
|
||||
ind_to_tuple_spec = {}
|
||||
for ind, node in enumerate(region):
|
||||
if node not in flattened_getitem_nodes:
|
||||
subgraph_node = copy_to_subgraph(node)
|
||||
if ind in inds_with_external_users:
|
||||
# flatten tuple outputs by generating a getitem node tree
|
||||
if _is_tuple_node(node):
|
||||
getitem_nodes, ind_to_tuple_spec[ind] = _create_getitem_nodes(
|
||||
node, subgraph_node, subgraph
|
||||
)
|
||||
output_list.extend(getitem_nodes)
|
||||
else:
|
||||
output_list.append(subgraph_node)
|
||||
|
||||
subgraph.output(tuple(output_list))
|
||||
|
||||
return subgraph, external_node_usages, node_usage_to_tuple_elems, ind_to_tuple_spec
|
||||
|
||||
|
||||
def _stable_topological_sort_impl(
|
||||
graph: torch.fx.Graph,
|
||||
node_to_additional_deps: dict[Node, OrderedSet[Node]],
|
||||
do_sort: bool = True,
|
||||
) -> bool:
|
||||
# Nodes are in exactly one of these four collections:
|
||||
|
||||
# - Nodes in `pending` are waiting to be processed (in reverse order):
|
||||
pending = list(reversed(graph.nodes))
|
||||
|
||||
# - Nodes in `ready` have been processed and are already in the correct
|
||||
# order.
|
||||
ready = OrderedSet[Node]()
|
||||
|
||||
# - `waiting` is a mapping from a dependency to nodes which depend on that
|
||||
# dependency.
|
||||
waiting = defaultdict(list)
|
||||
|
||||
# - `outputs` are always at the end of the graph
|
||||
outputs = OrderedSet[Node]()
|
||||
|
||||
# The cursor indicates the last processed node so we can add new nodes
|
||||
# after it.
|
||||
cursor = None
|
||||
while pending:
|
||||
node = pending.pop()
|
||||
|
||||
if node.target == "output":
|
||||
outputs.add(node)
|
||||
assert not node.users, "output nodes should have no users"
|
||||
continue
|
||||
|
||||
waiting_for = [
|
||||
x
|
||||
for x in _get_flat_args_unique(node, node_to_additional_deps)
|
||||
if x not in ready
|
||||
]
|
||||
if waiting_for:
|
||||
# We have unprocessed input nodes. Might as well wait for the last
|
||||
# arg so an already sorted list will only recheck this node once.
|
||||
waiting[waiting_for[-1]].append(node)
|
||||
else:
|
||||
ready.add(node)
|
||||
if cursor and cursor.next is not node and do_sort:
|
||||
cursor.append(node)
|
||||
cursor = node
|
||||
# Mark the nodes that have been waiting for this node to finish as
|
||||
# ready to check again.
|
||||
pending.extend(reversed(waiting.pop(node, ())))
|
||||
|
||||
ready.update(outputs)
|
||||
return not waiting and len(ready) == len(graph.nodes)
|
||||
|
||||
|
||||
def _stable_topological_sort(
|
||||
graph: torch.fx.Graph,
|
||||
node_to_additional_deps: dict[Node, OrderedSet[Node]],
|
||||
) -> None:
|
||||
assert _stable_topological_sort_impl(graph, node_to_additional_deps)
|
||||
|
||||
|
||||
def _has_cycle(
|
||||
graph: torch.fx.Graph,
|
||||
node_to_additional_deps: dict[Node, OrderedSet[Node]],
|
||||
) -> bool:
|
||||
return not _stable_topological_sort_impl(
|
||||
graph, node_to_additional_deps, do_sort=False
|
||||
)
|
||||
|
||||
|
||||
def _populate_additional_deps(
|
||||
graph: torch.fx.Graph, node_to_mutated_arg_positions: dict[Node, OrderedSet[int]]
|
||||
) -> dict[Node, OrderedSet[Node]]:
|
||||
node_to_additional_deps: dict[Node, OrderedSet[Node]] = defaultdict(OrderedSet)
|
||||
_add_mutation_dependencies(node_to_mutated_arg_positions, node_to_additional_deps)
|
||||
_add_global_state_dependencies(graph, node_to_additional_deps)
|
||||
return node_to_additional_deps
|
||||
|
||||
|
||||
def _add_global_state_dependencies(
|
||||
graph: torch.fx.Graph, node_to_additional_deps: dict[Node, OrderedSet[Node]]
|
||||
) -> None:
|
||||
import torch.amp
|
||||
|
||||
all_nodes = list(graph.nodes)
|
||||
|
||||
# These are targets of the nodes which need to stay in the same relative place in the graph
|
||||
global_state_targets = {torch.amp._enter_autocast, torch.amp._exit_autocast}
|
||||
all_nodes_dep_on: list[Node] = []
|
||||
|
||||
def prev_cur_nodes(
|
||||
all_nodes: list[Node],
|
||||
) -> Generator[tuple[list[Node], Node], None, None]:
|
||||
prev_nodes: list[Node] = []
|
||||
next_nodes = list(reversed(all_nodes))
|
||||
|
||||
while next_nodes:
|
||||
cur_node = next_nodes.pop()
|
||||
yield prev_nodes, cur_node
|
||||
prev_nodes.append(cur_node)
|
||||
|
||||
for prev_nodes, cur_node in prev_cur_nodes(all_nodes):
|
||||
args_unique = _get_flat_args_unique(cur_node, {})
|
||||
new_deps = [n for n in all_nodes_dep_on if n not in args_unique]
|
||||
|
||||
if new_deps:
|
||||
additional_deps = node_to_additional_deps[cur_node]
|
||||
additional_deps.update(new_deps)
|
||||
|
||||
if cur_node.target in global_state_targets:
|
||||
additional_deps = node_to_additional_deps[cur_node]
|
||||
additional_deps.update(n for n in prev_nodes if n not in args_unique)
|
||||
all_nodes_dep_on.append(cur_node)
|
||||
|
||||
|
||||
def _add_mutation_dependencies(
|
||||
node_to_mutated_arg_positions: dict[Node, OrderedSet[int]],
|
||||
node_to_additional_deps: dict[Node, OrderedSet[Node]],
|
||||
) -> None:
|
||||
for node, indices in node_to_mutated_arg_positions.items():
|
||||
flat_args_kwargs = _get_flat_args(node, {})
|
||||
|
||||
# for all mutated args,
|
||||
# add dependency on usages which occur after node to ensure
|
||||
# node will always be ordered before them
|
||||
# also add node as a dependency on usages which
|
||||
# occur before node to ensure node is ordered after them
|
||||
for index in indices:
|
||||
mutated_arg = flat_args_kwargs[index]
|
||||
for user in mutated_arg.users:
|
||||
if user is node:
|
||||
continue
|
||||
|
||||
elif user < node:
|
||||
node_to_additional_deps[node].add(user)
|
||||
|
||||
elif user > node:
|
||||
node_to_additional_deps[user].add(node)
|
||||
|
||||
|
||||
def _has_aliasing(
|
||||
region: Region,
|
||||
inputs: list[Node],
|
||||
inds_with_external_users: list[int],
|
||||
flattened_getitem_nodes: OrderedSet[Node],
|
||||
) -> bool:
|
||||
input_storages: dict[StorageWeakRef, Node] = dict()
|
||||
for node in inputs:
|
||||
if node in flattened_getitem_nodes:
|
||||
continue
|
||||
example_value = node.meta["example_value"]
|
||||
if isinstance(example_value, torch.Tensor):
|
||||
storage = StorageWeakRef(example_value._typed_storage())
|
||||
if storage in input_storages:
|
||||
# input-input aliasing
|
||||
log.debug(
|
||||
"NYI: Failed to substitute region %s due to input-output aliasing detected at nodes %s, %s",
|
||||
region,
|
||||
input_storages[storage],
|
||||
node,
|
||||
)
|
||||
return True
|
||||
input_storages[storage] = node
|
||||
output_storages: dict[StorageWeakRef, Node] = dict()
|
||||
for i in inds_with_external_users:
|
||||
out_node = region[i]
|
||||
if out_node in flattened_getitem_nodes:
|
||||
continue
|
||||
if out_node:
|
||||
example_value = out_node.meta["example_value"]
|
||||
assert not isinstance(example_value, list)
|
||||
if isinstance(example_value, torch.Tensor):
|
||||
storage = StorageWeakRef(example_value._typed_storage())
|
||||
if storage in output_storages:
|
||||
# output-output aliasing
|
||||
log.debug(
|
||||
"NYI: Failed to substitute region %s due to output-output aliasing detected at nodes %s, %s",
|
||||
region,
|
||||
output_storages[storage],
|
||||
out_node,
|
||||
)
|
||||
return True
|
||||
output_storages[storage] = out_node
|
||||
intersected_storages = input_storages.keys() & output_storages.keys()
|
||||
if len(intersected_storages) > 0:
|
||||
# input-output aliasing
|
||||
aliased = [
|
||||
(input_storages[s], output_storages[s]) for s in intersected_storages
|
||||
]
|
||||
aliased = ", ".join([f"{i} and {o}" for i, o in aliased])
|
||||
log.debug(
|
||||
"NYI: Failed to substitute region %s due to input-output aliasing detected at nodes %s",
|
||||
region,
|
||||
aliased,
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_tuple_node(node: Node) -> bool:
|
||||
return isinstance(node.meta["example_value"], tuple)
|
||||
|
||||
|
||||
def _get_children_getitems(node: Node) -> Generator[Node, None, None]:
|
||||
for user in node.users:
|
||||
if user.target is operator.getitem and isinstance(user.args[1], int):
|
||||
yield user
|
||||
|
||||
|
||||
def _get_flattened_node_indices(node: Node, region: Region) -> OrderedSet[int]:
|
||||
"""Returns an ordered set of indices, each representing a node in the region which will be flattened"""
|
||||
flattened_node_to_ind = {n: i for i, n in enumerate(region)}
|
||||
node_indices: OrderedSet[int] = OrderedSet()
|
||||
queue = deque(_get_children_getitems(node))
|
||||
while queue:
|
||||
cur_node = queue.popleft()
|
||||
if any(user in region for user in cur_node.users):
|
||||
node_indices.add(flattened_node_to_ind[cur_node])
|
||||
for child in _get_children_getitems(cur_node):
|
||||
queue.append(child)
|
||||
return node_indices
|
||||
|
||||
|
||||
def _create_getitem_nodes(
|
||||
node: Node, subgraph_tuple_node: Node, subgraph: torch.fx.Graph
|
||||
) -> tuple[list[Node], dict[tuple[int, ...], int]]:
|
||||
tup = node.meta["example_value"]
|
||||
assert isinstance(tup, tuple), "_get_getitem_children expects tuple"
|
||||
|
||||
getitem_nodes: list[Node] = []
|
||||
queue = deque([(e, (i,), subgraph_tuple_node) for i, e in enumerate(tup)])
|
||||
path_to_output_index = {}
|
||||
|
||||
while queue:
|
||||
cur_elem, path, parent = queue.popleft()
|
||||
|
||||
with subgraph.inserting_after(parent):
|
||||
new_getitem_node = subgraph.create_node(
|
||||
"call_function", operator.getitem, (parent, path[-1]), {}
|
||||
)
|
||||
new_getitem_node.meta["example_value"] = cur_elem
|
||||
|
||||
path_to_output_index[path] = len(getitem_nodes)
|
||||
getitem_nodes.append(new_getitem_node)
|
||||
|
||||
if isinstance(cur_elem, tuple):
|
||||
queue.extend(
|
||||
[(e, path + (i,), new_getitem_node) for i, e in enumerate(cur_elem)] # type: ignore[arg-type,misc]
|
||||
)
|
||||
|
||||
return getitem_nodes, path_to_output_index # type: ignore[return-value]
|
||||
|
||||
|
||||
def _replace_tuple_outputs(
|
||||
node: Node,
|
||||
output_index: int,
|
||||
tuple_spec: dict[tuple[int, ...], int],
|
||||
invoke_subgraph_node: Node,
|
||||
graph: torch.fx.Graph,
|
||||
) -> OrderedSet[Node]:
|
||||
assert _is_tuple_node(node), "_replace_tuple_outputs expects a tuple node"
|
||||
|
||||
queue = deque((c, (c.args[1],)) for c in _get_children_getitems(node))
|
||||
erased_nodes: OrderedSet[Node] = OrderedSet()
|
||||
while queue:
|
||||
cur_node, path = queue.pop()
|
||||
|
||||
for c in _get_children_getitems(cur_node):
|
||||
queue.append((c, path + (c.args[1],))) # type: ignore[return-value, arg-type]
|
||||
|
||||
with graph.inserting_after(invoke_subgraph_node):
|
||||
subgraph_output = graph.create_node(
|
||||
"call_function",
|
||||
operator.getitem,
|
||||
(invoke_subgraph_node, output_index + tuple_spec[path]), # type: ignore[index]
|
||||
{},
|
||||
)
|
||||
cur_node.replace_all_uses_with(subgraph_output, propagate_meta=True)
|
||||
graph.erase_node(cur_node)
|
||||
erased_nodes.add(cur_node)
|
||||
|
||||
graph.erase_node(node)
|
||||
erased_nodes.add(node)
|
||||
return erased_nodes
|
||||
@@ -0,0 +1,473 @@
|
||||
# mypy: disallow-untyped-defs
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
import re
|
||||
import warnings
|
||||
from typing import Any, Generic, TYPE_CHECKING, TypeVar
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch._guards import CompileId
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GraphIdFilter:
|
||||
"""
|
||||
A filter for matching graph IDs based on various conditions.
|
||||
Supports individual IDs, ranges, and comparison operators.
|
||||
"""
|
||||
|
||||
def __init__(self, filter_str: str) -> None:
|
||||
self._explicit_ids: frozenset[int] = frozenset()
|
||||
self._conditions: list[tuple[str, int]] = []
|
||||
self._parse(filter_str)
|
||||
|
||||
def _parse(self, filter_str: str) -> None:
|
||||
if not filter_str or not filter_str.strip():
|
||||
return
|
||||
|
||||
explicit_ids: set[int] = set()
|
||||
conditions: list[tuple[str, int]] = []
|
||||
|
||||
# Pattern for comparison operators (>=, >, <=, <) followed by a number
|
||||
cmp_pattern = re.compile(r"^(>=|>|<=|<)(\d+)$")
|
||||
# Pattern for ranges like "10-20"
|
||||
range_pattern = re.compile(r"^(\d+)-(\d+)$")
|
||||
|
||||
for part in filter_str.split(","):
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
|
||||
if match := cmp_pattern.match(part):
|
||||
conditions.append((match.group(1), int(match.group(2))))
|
||||
elif match := range_pattern.match(part):
|
||||
start, end = int(match.group(1)), int(match.group(2))
|
||||
explicit_ids.update(range(start, end + 1))
|
||||
else:
|
||||
try:
|
||||
explicit_ids.add(int(part))
|
||||
except ValueError:
|
||||
log.warning("Invalid graph ID filter: %s", part)
|
||||
|
||||
self._explicit_ids = frozenset(explicit_ids)
|
||||
self._conditions = conditions
|
||||
|
||||
def __contains__(self, graph_id: int) -> bool:
|
||||
"""Check if the given graph ID matches this filter."""
|
||||
if graph_id in self._explicit_ids:
|
||||
return True
|
||||
|
||||
for op, val in self._conditions:
|
||||
if op == ">" and graph_id > val:
|
||||
return True
|
||||
elif op == ">=" and graph_id >= val:
|
||||
return True
|
||||
elif op == "<" and graph_id < val:
|
||||
return True
|
||||
elif op == "<=" and graph_id <= val:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def __repr__(self) -> str:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
parts = []
|
||||
if self._explicit_ids:
|
||||
parts.append(f"ids={sorted(self._explicit_ids)}")
|
||||
if self._conditions:
|
||||
parts.append(f"conditions={self._conditions}")
|
||||
return f"GraphIdFilter({', '.join(parts) if parts else 'empty'})"
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class _GraphRouterBase(Generic[T]):
|
||||
"""
|
||||
Base class for routing graphs to different values based on their IDs.
|
||||
|
||||
The router parses a configuration string with rules in the format:
|
||||
"filter1:value1;filter2:value2;..."
|
||||
|
||||
Rules are evaluated in order, and the first matching rule wins.
|
||||
"""
|
||||
|
||||
def __init__(self, config_str: str, rule_type: str) -> None:
|
||||
self._rules: list[tuple[GraphIdFilter, T]] = []
|
||||
self._values: list[T | None] = []
|
||||
self._overflow_value: T | None = None
|
||||
self._rule_type = rule_type
|
||||
self._parse(config_str)
|
||||
self._precompute()
|
||||
|
||||
def _parse_value_str(self, value_str: str) -> T | None:
|
||||
"""Parse a value string into the appropriate type. Returns None to skip."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _parse(self, config_str: str) -> None:
|
||||
if not config_str or not config_str.strip():
|
||||
return
|
||||
|
||||
rule_strs = config_str.split(";")
|
||||
for rule_str in rule_strs:
|
||||
rule_str = rule_str.strip()
|
||||
if not rule_str:
|
||||
continue
|
||||
|
||||
colon_idx = rule_str.find(":")
|
||||
if colon_idx == -1:
|
||||
log.warning(
|
||||
"Invalid %s override rule (missing ':'): %s",
|
||||
self._rule_type,
|
||||
rule_str,
|
||||
)
|
||||
continue
|
||||
|
||||
filter_str = rule_str[:colon_idx].strip()
|
||||
value_str = rule_str[colon_idx + 1 :].strip()
|
||||
|
||||
if not filter_str or not value_str:
|
||||
log.warning("Invalid %s override rule: %s", self._rule_type, rule_str)
|
||||
continue
|
||||
|
||||
value = self._parse_value_str(value_str)
|
||||
if value is not None:
|
||||
self._rules.append((GraphIdFilter(filter_str), value))
|
||||
|
||||
def _precompute(self) -> None:
|
||||
if not self._rules:
|
||||
return
|
||||
|
||||
# Find max ID from explicit IDs and comparison thresholds
|
||||
max_id = 0
|
||||
for f, _ in self._rules:
|
||||
if f._explicit_ids:
|
||||
max_id = max(max_id, *f._explicit_ids)
|
||||
for _, val in f._conditions:
|
||||
max_id = max(max_id, val)
|
||||
|
||||
# Pre-compute values for IDs 0 to max_id
|
||||
for i in range(max_id + 1):
|
||||
self._values.append(self._match_rules(i))
|
||||
|
||||
# For IDs > max_id, the result is constant (only unbounded conditions apply)
|
||||
self._overflow_value = self._match_rules(max_id + 1)
|
||||
|
||||
def _match_rules(self, graph_id: int) -> T | None:
|
||||
for f, value in self._rules:
|
||||
if graph_id in f:
|
||||
return value
|
||||
return None
|
||||
|
||||
def get_value_for_graph(self, graph_id: int) -> T | None:
|
||||
"""Get the value for a given graph ID. Returns None if no rule matches."""
|
||||
if graph_id < len(self._values):
|
||||
return self._values[graph_id]
|
||||
return self._overflow_value
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""Check if no rules are configured."""
|
||||
return len(self._rules) == 0
|
||||
|
||||
|
||||
class GraphBackendRouter(_GraphRouterBase[Any]):
|
||||
"""
|
||||
Routes graphs to different backends based on their IDs.
|
||||
|
||||
The router parses a configuration string with rules in the format:
|
||||
"filter1:backend1;filter2:backend2;..."
|
||||
|
||||
If a graph ID matches multiple rules with different backends, a ValueError
|
||||
is raised.
|
||||
|
||||
Examples:
|
||||
"0-5:eager;>5:inductor" - IDs 0-5 use eager, rest use inductor
|
||||
">10:aot_eager" - IDs > 10 use aot_eager
|
||||
"<=3:eager;4-10:aot_eager" - IDs 0-3 use eager, 4-10 use aot_eager
|
||||
|
||||
Supported backends include "eager", "aot_eager", "aot_eager_decomp_partition",
|
||||
"inductor", and any other registered backend.
|
||||
"""
|
||||
|
||||
def __init__(self, config_str: str) -> None:
|
||||
self._backend_names: dict[int, str] = {}
|
||||
super().__init__(config_str, "backend")
|
||||
|
||||
def _parse_value_str(self, value_str: str) -> Any | None:
|
||||
"""Look up a backend by name."""
|
||||
from .backends.registry import lookup_backend
|
||||
from .eval_frame import cached_backends
|
||||
|
||||
backend = lookup_backend(value_str)
|
||||
|
||||
# Register the backend so its reset() is called during torch._dynamo.reset()
|
||||
assert backend is not None, "Invalid override backend: " + value_str
|
||||
cached_backends.setdefault(id(backend), backend)
|
||||
self._backend_names[id(backend)] = value_str
|
||||
return backend
|
||||
|
||||
def _match_rules(self, graph_id: int) -> Any | None:
|
||||
"""Match rules with conflict detection for overlapping filters."""
|
||||
matches = {id(backend): backend for f, backend in self._rules if graph_id in f}
|
||||
if len(matches) > 1:
|
||||
names = [self._backend_names[bid] for bid in matches]
|
||||
raise ValueError(
|
||||
f"Conflicting backend override for graph {graph_id}: matched {names}"
|
||||
)
|
||||
if matches:
|
||||
return next(iter(matches.values()))
|
||||
return None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
if not self._rules:
|
||||
return "GraphBackendRouter(empty)"
|
||||
return f"GraphBackendRouter({self._rules})"
|
||||
|
||||
|
||||
class GraphConfigRouter(_GraphRouterBase[dict[str, Any]]):
|
||||
"""
|
||||
Routes graphs to different inductor configs based on their IDs.
|
||||
|
||||
The router parses a configuration string with rules in the format:
|
||||
"filter1:config1;filter2:config2;..."
|
||||
|
||||
All matching rules are aggregated: configs from all matching rules are merged
|
||||
into a single dict. Conflicting keys (same key, different value) raise an error.
|
||||
Config format is "key=value" or "key1=value1,key2=value2" for multiple settings.
|
||||
|
||||
Examples:
|
||||
"0-5:triton.cudagraph_skip_dynamic_graphs=False"
|
||||
">10:triton.cudagraphs=False,triton.cudagraph_trees=False"
|
||||
|
||||
With "0:a=1;>=0:b=2", graph 0 gets {"a": 1, "b": 2} (both rules match).
|
||||
With "0:a=1;>=0:a=2", graph 0 raises an error (conflicting values for "a").
|
||||
With "0:a=1;>=0:a=1", graph 0 gets {"a": 1} (same value is not a conflict).
|
||||
"""
|
||||
|
||||
def __init__(self, config_str: str) -> None:
|
||||
super().__init__(config_str, "config")
|
||||
|
||||
@staticmethod
|
||||
def _parse_scalar_value(value_str: str) -> Any:
|
||||
"""Parse a string value into the appropriate Python type."""
|
||||
value_str = value_str.strip()
|
||||
if value_str.lower() == "true":
|
||||
return True
|
||||
if value_str.lower() == "false":
|
||||
return False
|
||||
if value_str.lower() == "none":
|
||||
return None
|
||||
try:
|
||||
if "." in value_str:
|
||||
return float(value_str)
|
||||
return int(value_str)
|
||||
except ValueError:
|
||||
return value_str
|
||||
|
||||
def _match_rules(self, graph_id: int) -> dict[str, Any] | None:
|
||||
"""Aggregate configs from all matching rules. Conflicts raise an error."""
|
||||
result: dict[str, Any] = {}
|
||||
for f, value in self._rules:
|
||||
if graph_id in f:
|
||||
for k, v in value.items():
|
||||
if k in result and result[k] != v:
|
||||
raise ValueError(
|
||||
f"Conflicting config override for graph {graph_id}: "
|
||||
f"key '{k}' has value {result[k]!r} and {v!r}"
|
||||
)
|
||||
result[k] = v
|
||||
return result if result else None
|
||||
|
||||
def _parse_value_str(self, value_str: str) -> dict[str, Any] | None:
|
||||
"""Parse a config string like 'key1=val1,key2=val2' into a dict."""
|
||||
result: dict[str, Any] = {}
|
||||
for item in value_str.split(","):
|
||||
item = item.strip()
|
||||
if not item:
|
||||
continue
|
||||
if "=" not in item:
|
||||
log.warning("Invalid config item (missing '='): %s", item)
|
||||
continue
|
||||
key, value = item.split("=", 1)
|
||||
result[key.strip()] = self._parse_scalar_value(value)
|
||||
return result if result else None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
if not self._rules:
|
||||
return "GraphConfigRouter(empty)"
|
||||
return f"GraphConfigRouter({self._rules})"
|
||||
|
||||
|
||||
def _get_override_for_compile_id(
|
||||
compile_id: CompileId | None,
|
||||
config_str: str,
|
||||
create_router: Callable[[str], _GraphRouterBase[T]],
|
||||
label: str,
|
||||
) -> T | None:
|
||||
"""
|
||||
Get the override value for a given CompileId.
|
||||
|
||||
Returns the value from the router, or None if no override applies.
|
||||
"""
|
||||
if compile_id is None or not config_str:
|
||||
return None
|
||||
|
||||
graph_id = compile_id.frame_id
|
||||
if graph_id is None:
|
||||
return None
|
||||
|
||||
router = create_router(config_str)
|
||||
value = router.get_value_for_graph(graph_id)
|
||||
if value is not None:
|
||||
log.info("Overriding %s: %s", label, value)
|
||||
return value
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _create_backend_router(config_str: str) -> GraphBackendRouter:
|
||||
"""Create and cache GraphBackendRouter instances based on config string."""
|
||||
return GraphBackendRouter(config_str)
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _validate_backend_names(config_str: str) -> str | None:
|
||||
"""Return an error message if any backend name is invalid, else None."""
|
||||
if not config_str or not config_str.strip():
|
||||
return None
|
||||
from .backends.registry import lookup_backend
|
||||
|
||||
for rule_str in config_str.split(";"):
|
||||
rule_str = rule_str.strip()
|
||||
if not rule_str or ":" not in rule_str:
|
||||
continue
|
||||
backend_name = rule_str[rule_str.find(":") + 1 :].strip()
|
||||
if not backend_name:
|
||||
continue
|
||||
try:
|
||||
lookup_backend(backend_name)
|
||||
except Exception:
|
||||
return (
|
||||
f"TORCH_COMPILE_OVERRIDE_BACKENDS: "
|
||||
f"'{backend_name}' is not a valid backend, "
|
||||
f"see `torch._dynamo.list_backends()` for available backends"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _validate_inductor_config_keys(config_str: str) -> str | None:
|
||||
"""Return an error message if any config key is invalid, else None."""
|
||||
router = GraphConfigRouter(config_str)
|
||||
from torch._inductor import config
|
||||
|
||||
for _, config_dict in router._rules:
|
||||
for key in config_dict:
|
||||
if not hasattr(config, key):
|
||||
return (
|
||||
f"TORCH_COMPILE_OVERRIDE_INDUCTOR_CONFIGS: "
|
||||
f"'{key}' is not a valid torch._inductor.config option"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _validate_dynamo_config_keys(config_str: str) -> str | None:
|
||||
"""Return an error message if any config key is invalid, else None."""
|
||||
router = GraphConfigRouter(config_str)
|
||||
from torch._dynamo import config
|
||||
|
||||
for _, config_dict in router._rules:
|
||||
for key in config_dict:
|
||||
if not hasattr(config, key):
|
||||
return (
|
||||
f"TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS: "
|
||||
f"'{key}' is not a valid torch._dynamo.config option"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _create_inductor_config_router(config_str: str) -> GraphConfigRouter:
|
||||
"""Create and cache GraphConfigRouter for inductor config overrides."""
|
||||
return GraphConfigRouter(config_str)
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _create_dynamo_config_router(config_str: str) -> GraphConfigRouter:
|
||||
"""Create and cache GraphConfigRouter for dynamo config overrides.
|
||||
|
||||
Warns that dynamo config overrides are keyed by frame ID and some configs
|
||||
can affect graph breaks, which may shift frame IDs.
|
||||
"""
|
||||
router = GraphConfigRouter(config_str)
|
||||
if not router.is_empty():
|
||||
warnings.warn(
|
||||
"TORCH_COMPILE_OVERRIDE_DYNAMO_CONFIGS is set. Dynamo config overrides are "
|
||||
"keyed by frame ID. Some dynamo configs can affect graph breaks, "
|
||||
"which may alter the number of frames and shift frame IDs, causing "
|
||||
"overrides to target the wrong graphs.",
|
||||
)
|
||||
return router
|
||||
|
||||
|
||||
def get_backend_override_for_compile_id(
|
||||
compile_id: CompileId | None,
|
||||
config_str: str,
|
||||
) -> Any:
|
||||
"""
|
||||
Get the backend override for a given CompileId.
|
||||
|
||||
Returns the backend function to use, or None if no override applies.
|
||||
"""
|
||||
return _get_override_for_compile_id(
|
||||
compile_id,
|
||||
config_str,
|
||||
_create_backend_router,
|
||||
"torch.compile backend",
|
||||
)
|
||||
|
||||
|
||||
def get_inductor_config_override_for_compile_id(
|
||||
compile_id: CompileId | None,
|
||||
config_str: str,
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get the inductor config override for a given CompileId.
|
||||
|
||||
Returns a dict of config patches to apply, or None if no override applies.
|
||||
"""
|
||||
return _get_override_for_compile_id(
|
||||
compile_id,
|
||||
config_str,
|
||||
_create_inductor_config_router, # type: ignore[arg-type]
|
||||
"inductor config",
|
||||
)
|
||||
|
||||
|
||||
def get_dynamo_config_override_for_compile_id(
|
||||
compile_id: CompileId | None,
|
||||
config_str: str,
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get the dynamo config override for a given CompileId.
|
||||
|
||||
Returns a dict of config patches to apply, or None if no override applies.
|
||||
"""
|
||||
return _get_override_for_compile_id(
|
||||
compile_id,
|
||||
config_str,
|
||||
_create_dynamo_config_router, # type: ignore[arg-type]
|
||||
"dynamo config",
|
||||
)
|
||||
|
||||
|
||||
# Keep old name for backwards compatibility
|
||||
_create_router = _create_backend_router
|
||||
@@ -0,0 +1,503 @@
|
||||
"""
|
||||
This module provides functionality for tracking and managing regions in computational graphs.
|
||||
It supports graph optimization by identifying and grouping similar regions based on their
|
||||
structure and behavior. The module implements algorithms for:
|
||||
|
||||
1. Tracking nodes and their relationships in the computational graph
|
||||
2. Identifying identical or similar regions across the graph
|
||||
3. Managing graph regions for optimization purposes
|
||||
4. Supporting deduplication and other graph transformation passes
|
||||
|
||||
The core functionality revolves around the GraphRegionTracker class which maintains
|
||||
mappings between nodes and their duplicates, enabling efficient graph analysis and
|
||||
optimization operations.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copyreg
|
||||
import io
|
||||
import logging
|
||||
import math
|
||||
import operator
|
||||
import pickle
|
||||
from collections import defaultdict, deque
|
||||
from dataclasses import fields
|
||||
from typing import Any, TYPE_CHECKING, TypeVar
|
||||
|
||||
import torch._logging
|
||||
import torch.fx
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils._pytree import tree_flatten
|
||||
|
||||
from .graph_utils import _get_flat_args_unique
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from .symbolic_convert import InstructionTranslatorBase
|
||||
|
||||
|
||||
Node = torch.fx.Node
|
||||
Region = list[Node]
|
||||
IdenticalNodes = list[Node]
|
||||
GlobalStateKey = tuple[
|
||||
bool,
|
||||
bool,
|
||||
int,
|
||||
tuple[bool, bool],
|
||||
tuple[bool, bool],
|
||||
torch.dtype,
|
||||
bool,
|
||||
bool,
|
||||
bool,
|
||||
bool,
|
||||
]
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
graph_expansion_log = torch._logging.getArtifactLogger(
|
||||
__name__, "graph_region_expansion"
|
||||
)
|
||||
|
||||
|
||||
def debug_log(msg: str, *args) -> None: # type: ignore[no-untyped-def]
|
||||
graph_expansion_log.debug(msg, *args)
|
||||
|
||||
|
||||
def _extract_tensor_metadata_for_node_hash(
|
||||
x: torch.Tensor,
|
||||
) -> tuple[Callable[[T], T], tuple[Any, ...]]:
|
||||
from torch._inductor.codecache import _ident, extract_tensor_metadata_for_cache_key
|
||||
|
||||
out = []
|
||||
metadata = extract_tensor_metadata_for_cache_key(x)
|
||||
for field in fields(metadata):
|
||||
out.append(getattr(metadata, field.name))
|
||||
|
||||
return (_ident, tuple(out))
|
||||
|
||||
|
||||
class NodeHashException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class InputPickler(pickle.Pickler):
|
||||
def __init__(self) -> None:
|
||||
from torch._inductor.codecache import _ident
|
||||
|
||||
stream = io.BytesIO()
|
||||
self._stream = stream
|
||||
super().__init__(stream)
|
||||
self.dispatch_table = copyreg.dispatch_table.copy()
|
||||
self.dispatch_table.update(
|
||||
{
|
||||
FakeTensor: _extract_tensor_metadata_for_node_hash,
|
||||
torch.SymInt: lambda x: (_ident, (str(x),)),
|
||||
torch.SymBool: lambda x: (_ident, (str(x),)),
|
||||
torch.SymFloat: lambda x: (_ident, (str(x),)),
|
||||
}
|
||||
)
|
||||
self.fast = True
|
||||
|
||||
def dumps(self, obj: Any) -> bytes:
|
||||
"""
|
||||
Pickle an object and return a byte string.
|
||||
"""
|
||||
try:
|
||||
self.dump(obj)
|
||||
return self._stream.getvalue()
|
||||
except (TypeError, AttributeError) as e:
|
||||
raise NodeHashException from e
|
||||
finally:
|
||||
self._stream.seek(0)
|
||||
self._stream.truncate(0)
|
||||
|
||||
|
||||
def _extract_args(arg: Any) -> Any:
|
||||
if isinstance(arg, Node):
|
||||
return arg.meta.get("example_value")
|
||||
elif isinstance(arg, (torch.Tensor, int)):
|
||||
return arg
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_args(
|
||||
node: Node,
|
||||
) -> tuple[tuple[str, ...], tuple[Any | None, ...]]:
|
||||
flat_args, _ = tree_flatten(node.args)
|
||||
sorted_kwargs = sorted(node.kwargs.items(), key=operator.itemgetter(0))
|
||||
sorted_keys = tuple(sorted(node.kwargs.keys()))
|
||||
flat_kwargs, _ = tree_flatten(sorted_kwargs)
|
||||
all_args = flat_args + flat_kwargs
|
||||
return (sorted_keys, tuple(_extract_args(arg) for arg in all_args))
|
||||
|
||||
|
||||
def _sort_with_ref_region(
|
||||
index_to_rank: dict[int, int], regions: list[list[Any]]
|
||||
) -> None:
|
||||
# sort topologically
|
||||
# we need to handle edge cases where some nodes have no dependencies
|
||||
# so first we map each node to its ranking
|
||||
ref_region = regions[0]
|
||||
sorted_indices = sorted(range(len(ref_region)), key=lambda i: index_to_rank[i])
|
||||
for region in regions:
|
||||
region[:] = [region[i] for i in sorted_indices]
|
||||
|
||||
|
||||
def get_global_state_key() -> GlobalStateKey:
|
||||
return (
|
||||
torch.is_grad_enabled(),
|
||||
torch.is_inference_mode_enabled(),
|
||||
torch.get_num_threads(),
|
||||
torch._C._get_cublas_allow_fp16_reduced_precision_reduction(),
|
||||
torch._C._get_cublas_allow_bf16_reduced_precision_reduction(),
|
||||
torch.get_default_dtype(),
|
||||
torch.are_deterministic_algorithms_enabled(),
|
||||
torch._C._get_cublas_allow_tf32(),
|
||||
torch.is_deterministic_algorithms_warn_only_enabled(),
|
||||
torch._C._autograd._saved_tensors_hooks_is_enabled(), # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
|
||||
# This is typical BFS with the caveat
|
||||
# that a node's children need to be explicitly
|
||||
# added with the add_children() method
|
||||
# The flow is yield a node and check if it's valid for all regions
|
||||
# if not valid, discard and continue onto the next node
|
||||
# Note: this iterates backward through the graph by looking at args/kwargs
|
||||
# of a node
|
||||
class BackwardBfsArgIter:
|
||||
def __init__(self, origin: Node) -> None:
|
||||
self._cur: Node | None = origin
|
||||
self._queue: deque[Node | None] = deque()
|
||||
|
||||
@staticmethod
|
||||
def create(origin: Node) -> BackwardBfsArgIter:
|
||||
it = BackwardBfsArgIter(origin)
|
||||
it.add_children(origin)
|
||||
# pop the origin node, since it is the origin of
|
||||
# the region and does not need to be considered for addition
|
||||
assert it.next()
|
||||
return it
|
||||
|
||||
def next(self) -> Node | None:
|
||||
ret = self._cur
|
||||
if not self._queue:
|
||||
self._cur = None
|
||||
else:
|
||||
self._cur = self._queue.popleft()
|
||||
return ret
|
||||
|
||||
def peek(self) -> Node | None:
|
||||
return self._cur
|
||||
|
||||
def add_children(self, node: Node) -> None:
|
||||
flat_args = _get_flat_args_unique(node, {})
|
||||
for arg in flat_args:
|
||||
if isinstance(arg, Node):
|
||||
self._append(arg)
|
||||
|
||||
def _append(self, arg: Node) -> None:
|
||||
if self._cur is None:
|
||||
self._cur = arg
|
||||
else:
|
||||
self._queue.append(arg)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"BackwardBfsArgIter(cur={self._cur}, queue={self._queue})"
|
||||
|
||||
|
||||
class GraphRegionTracker:
|
||||
"""
|
||||
GraphRegionTracker tracks each node added to the output graph and generates a key based on the source location,
|
||||
instruction pointer, input shapes, and global state at the time the node is inserted into the graph. Nodes with
|
||||
the same key are grouped together in a list of identical nodes (the value of node_to_duplicates).
|
||||
|
||||
hash_to_duplicates: Dict[str, IdenticalNodes] - A dictionary mapping the key to a list of identical nodes
|
||||
node_to_duplicates: Dict[Node, IdenticalNodes] - A dictionary mapping a node to the list of identical nodes it belongs to
|
||||
input_pickler: InputPickler - An instance of InputPickler used to generate a node hash
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.hash_to_duplicates: dict[str, IdenticalNodes] = defaultdict(list)
|
||||
self.node_to_duplicates: dict[Node, IdenticalNodes] = {}
|
||||
# Note: position is in flattened args/kwargs list
|
||||
self.node_to_mutated_arg_positions: dict[Node, OrderedSet[int]] = {}
|
||||
self.input_pickler = InputPickler()
|
||||
|
||||
def _hash_node(
|
||||
self, filename: str, lineno: int, instruction_pointer: int | None, node: Node
|
||||
) -> str:
|
||||
from torch._inductor.codecache import sha256_hash
|
||||
|
||||
key = (
|
||||
get_global_state_key(),
|
||||
filename,
|
||||
lineno,
|
||||
instruction_pointer,
|
||||
_normalize_args(node),
|
||||
)
|
||||
return sha256_hash(self.input_pickler.dumps(key))
|
||||
|
||||
def _is_identical(self, n0: Node, n1: Node) -> bool:
|
||||
return (
|
||||
n0 in self.node_to_duplicates
|
||||
and n1 in self.node_to_duplicates
|
||||
and self.node_to_duplicates[n0] is self.node_to_duplicates[n1]
|
||||
and n0 is not n1
|
||||
)
|
||||
|
||||
def track_node(self, tx: InstructionTranslatorBase, node: Node) -> None:
|
||||
"""
|
||||
The main entry point for tracking a node. This function will hash the node argument and group
|
||||
nodes with the same hash together. It updates the hash_to_duplicates and node_to_duplicates dictionaries
|
||||
to track the new node.
|
||||
"""
|
||||
try:
|
||||
if (
|
||||
node not in self.node_to_duplicates
|
||||
): # don't allow nodes to be added twice
|
||||
duplicates = self.hash_to_duplicates[
|
||||
self._hash_node(
|
||||
tx.f_code.co_filename, tx.lineno, tx.instruction_pointer, node
|
||||
)
|
||||
]
|
||||
duplicates.append(node)
|
||||
self.node_to_duplicates[node] = duplicates
|
||||
except NodeHashException as e:
|
||||
log.debug("Unable to hash node %s with exception %s", node, e) # noqa: G200
|
||||
|
||||
def track_node_mutations(
|
||||
self,
|
||||
node: Node,
|
||||
flat_args_kwargs: list[Any],
|
||||
id_to_initial_version: dict[int, int],
|
||||
) -> None:
|
||||
"""
|
||||
This function tracks which argument positions are mutated by the given node. Subgraph HOP does not support
|
||||
input mutations today so we will skip regions which have inputs that are mutated.
|
||||
"""
|
||||
mutated_arg_positions = OrderedSet[int]()
|
||||
for i, arg in enumerate(flat_args_kwargs):
|
||||
val_id = id(arg)
|
||||
if (
|
||||
val_id in id_to_initial_version
|
||||
and id_to_initial_version[val_id] != arg._version
|
||||
):
|
||||
mutated_arg_positions.add(i)
|
||||
|
||||
if mutated_arg_positions:
|
||||
self.node_to_mutated_arg_positions[node] = mutated_arg_positions
|
||||
|
||||
def add_node_mutation(
|
||||
self,
|
||||
node: Node,
|
||||
arg_pos: int,
|
||||
) -> None:
|
||||
if node in self.node_to_mutated_arg_positions:
|
||||
self.node_to_mutated_arg_positions[node].add(arg_pos)
|
||||
else:
|
||||
self.node_to_mutated_arg_positions[node] = OrderedSet([arg_pos])
|
||||
|
||||
def get_identical_regions(self, graph: torch.fx.Graph) -> list[list[Region]]:
|
||||
"""
|
||||
This function is responsible for extracting the largest regions of identical nodes from the given graph.
|
||||
**Note**: This function assumes the nodes that have been tracked with track_node are in the provided graph argument.
|
||||
|
||||
The algorithm proceeds as follows:
|
||||
The nodes tracked via track_node above are organized into region groups. The initial region groups look like this:
|
||||
[[IdenticalNode1], [IdenticalNode2], [IdenticalNode3]] and each sublist is called a region. For each region group
|
||||
(starting at the topologically latest region group), the inner regions are gradually expanded one node at time from
|
||||
the flattened args and kwargs of the node in each region provided that for all regions in the group, the nodes being
|
||||
added are also identical (ie have the same key computed by track_node). This is checked by verifying that the two
|
||||
nodes have the same identical node list in node_to_duplicates.
|
||||
"""
|
||||
topological_ranking = {node: i for i, node in enumerate(graph.nodes)}
|
||||
region_groups_with_rank = []
|
||||
# needed to detect if replacing a region will create cycles
|
||||
node_to_recursive_ancestors = _populate_recursive_ancestor_map(graph)
|
||||
|
||||
# Create region groups; a region group is a group
|
||||
# of regions that are all identical. In this initial state
|
||||
# each region in the group is a single node, and we discard
|
||||
# groups that are only a single region.
|
||||
# We track the topological ranking to start with groups later in the graph
|
||||
# the reason for this is that we will necessarily create the largest groups first.
|
||||
for group in self.hash_to_duplicates.values():
|
||||
if len(group) > 1:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
region_group = []
|
||||
min_rank = math.inf
|
||||
|
||||
for node in group:
|
||||
# some nodes aren't in the topo ranking?
|
||||
if node in topological_ranking:
|
||||
min_rank = min(min_rank, topological_ranking[node])
|
||||
region_group.append([node])
|
||||
|
||||
if len(region_group) > 1:
|
||||
region_groups_with_rank.append((region_group, min_rank))
|
||||
|
||||
region_groups_with_rank.sort(key=lambda rg: -rg[1])
|
||||
region_groups = [rg for rg, _ in region_groups_with_rank]
|
||||
|
||||
# We start from regions later in the graph and expand them earlier
|
||||
# as a result, we will create the largest regions first and they won't
|
||||
# overlap.
|
||||
seen_nodes: set[Node] = set()
|
||||
for region_group in region_groups:
|
||||
fully_expand_region_group(
|
||||
region_group,
|
||||
seen_nodes,
|
||||
node_to_recursive_ancestors,
|
||||
self._is_identical,
|
||||
)
|
||||
# sort topologically
|
||||
# we need to handle edge cases where some nodes have no dependencies
|
||||
# so first we map each node to its ranking,
|
||||
ref_region = region_group[0]
|
||||
index_to_rank = {
|
||||
index: topological_ranking[n] for index, n in enumerate(ref_region)
|
||||
}
|
||||
_sort_with_ref_region(index_to_rank, region_group)
|
||||
|
||||
return [
|
||||
region_group for region_group in region_groups if len(region_group[0]) > 1
|
||||
]
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"GraphRegionTracker(hash_to_duplicates={self.hash_to_duplicates}, node_to_duplicates={self.node_to_duplicates})"
|
||||
|
||||
|
||||
class RegionWrapper:
|
||||
"""Holds state for regions e.g. ancestors and new candidate nodes for consideration"""
|
||||
|
||||
def __init__(
|
||||
self, region: Region, node_to_recursive_ancestors: dict[Node, set[Node]]
|
||||
) -> None:
|
||||
assert len(region) == 1, "all regions should start with one node"
|
||||
node = region[0]
|
||||
self.node_to_recursive_ancestors = node_to_recursive_ancestors
|
||||
self.iter = BackwardBfsArgIter.create(node)
|
||||
self.nodes_unique = OrderedSet([node])
|
||||
self.ancestors = set(node_to_recursive_ancestors[node])
|
||||
self.region = region
|
||||
|
||||
def next_candidate(self) -> Node | None:
|
||||
return self.iter.next()
|
||||
|
||||
def will_inclusion_create_cycle(self, node: Node) -> bool:
|
||||
external_users = [user for user in node.users if user not in self.nodes_unique]
|
||||
for user in external_users:
|
||||
if user in self.ancestors:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def add(self, node: Node) -> None:
|
||||
self.nodes_unique.add(node)
|
||||
self.region.append(node)
|
||||
self.iter.add_children(node)
|
||||
self.ancestors.update(self.node_to_recursive_ancestors[node])
|
||||
|
||||
|
||||
def fully_expand_region_group(
|
||||
regions: list[Region],
|
||||
seen_nodes: set[Node],
|
||||
node_to_recursive_ancestors: dict[Node, set[Node]],
|
||||
is_identical_fn: Callable[[Node, Node], bool],
|
||||
) -> None:
|
||||
debug_log("--------------------------------------------------")
|
||||
debug_log("expanding new region group: %s", regions)
|
||||
|
||||
# All regions should start with 1 node
|
||||
assert all(len(region) == 1 for region in regions)
|
||||
region_wrappers = [
|
||||
RegionWrapper(region, node_to_recursive_ancestors) for region in regions
|
||||
]
|
||||
|
||||
nodes_to_add = OrderedSet[Node]()
|
||||
current_node = region_wrappers[0].next_candidate()
|
||||
|
||||
# No children
|
||||
if current_node is None:
|
||||
return
|
||||
|
||||
# Loop incrementally adding new nodes to each region
|
||||
# regions are only expanded if the node to add is valid
|
||||
# for ALL regions
|
||||
while current_node:
|
||||
add_to_all_regions = not region_wrappers[0].will_inclusion_create_cycle(
|
||||
current_node
|
||||
)
|
||||
nodes_to_add.clear()
|
||||
nodes_to_add.add(current_node)
|
||||
for region_wrapper in region_wrappers[1:]:
|
||||
candidate = region_wrapper.next_candidate()
|
||||
|
||||
debug_log("--------------------")
|
||||
debug_log(
|
||||
"considering candidate: %s, cur_node: %s", candidate, current_node
|
||||
)
|
||||
|
||||
if not candidate or not add_to_all_regions:
|
||||
add_to_all_regions = False
|
||||
continue
|
||||
|
||||
debug_log(
|
||||
"candidate in previously claimed nodes?: %s", candidate in seen_nodes
|
||||
)
|
||||
debug_log("is_identical: %s", is_identical_fn(candidate, current_node))
|
||||
|
||||
add_to_all_regions &= (
|
||||
candidate not in seen_nodes
|
||||
and candidate not in nodes_to_add
|
||||
and candidate.op != "placeholder"
|
||||
and candidate.op != "get_attr"
|
||||
and is_identical_fn(candidate, current_node)
|
||||
and not region_wrapper.will_inclusion_create_cycle(candidate)
|
||||
)
|
||||
nodes_to_add.add(candidate)
|
||||
|
||||
debug_log(f"add_to_all_regions: {add_to_all_regions}")
|
||||
debug_log("--------------------")
|
||||
|
||||
if add_to_all_regions:
|
||||
assert len(region_wrappers) == len(nodes_to_add), (
|
||||
"Number of nodes to add must equal the number of regions"
|
||||
)
|
||||
for region_wrapper, node in zip(region_wrappers, nodes_to_add):
|
||||
region_wrapper.add(node)
|
||||
debug_log("adding %s's children", node)
|
||||
debug_log("%s %s", node.args, list(node.kwargs.items()))
|
||||
seen_nodes.add(node)
|
||||
|
||||
current_node = region_wrappers[0].next_candidate()
|
||||
|
||||
# Ensure regions are sorted in topological order
|
||||
for region in regions:
|
||||
region.reverse()
|
||||
|
||||
debug_log("end expand new region group: %s", regions)
|
||||
debug_log("--------------------------------------------------")
|
||||
|
||||
|
||||
def _populate_recursive_ancestor_map(graph: torch.fx.Graph) -> dict[Node, set[Node]]:
|
||||
node_to_recursive_ancestors: dict[Node, set[Node]] = {}
|
||||
for node in graph.nodes:
|
||||
node_to_recursive_ancestors[node] = set()
|
||||
for node in graph.nodes:
|
||||
all_args = _get_flat_args_unique(node, {})
|
||||
for arg in all_args:
|
||||
if isinstance(arg, Node):
|
||||
node_to_recursive_ancestors[node].update(
|
||||
node_to_recursive_ancestors[arg]
|
||||
)
|
||||
node_to_recursive_ancestors[node].add(arg)
|
||||
return node_to_recursive_ancestors
|
||||
@@ -0,0 +1,117 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.fx import Graph, map_arg, Node
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils._pytree import tree_flatten
|
||||
|
||||
|
||||
# flattens with support for slices
|
||||
# Note: a better way to do this would
|
||||
# be register/unregister slices as pytree nodes
|
||||
# but there is no unregister API in the pytorch
|
||||
# pytree impl
|
||||
def _get_flat_args(
|
||||
node: Node, node_to_additional_deps: dict[Node, OrderedSet[Node]]
|
||||
) -> list[Node]:
|
||||
args = list[Any]()
|
||||
map_arg((node.args, node.kwargs), args.append)
|
||||
if node in node_to_additional_deps:
|
||||
args.extend(node_to_additional_deps[node])
|
||||
return args
|
||||
|
||||
|
||||
def _get_flat_args_unique(
|
||||
node: Node, node_to_additional_deps: dict[Node, OrderedSet[Node]]
|
||||
) -> OrderedSet[Node]:
|
||||
args = OrderedSet[Node]()
|
||||
map_arg((node.args, node.kwargs), args.add)
|
||||
if node in node_to_additional_deps:
|
||||
args.update(node_to_additional_deps[node])
|
||||
return args
|
||||
|
||||
|
||||
def _detect_cycles(
|
||||
graph: Graph, node_to_additional_deps: dict[Node, OrderedSet[Node]]
|
||||
) -> str:
|
||||
# States: 0=Unvisited, 1=Visiting, 2=Visited(Safe)
|
||||
state: dict[Node, int] = {}
|
||||
|
||||
for root in reversed(graph.nodes):
|
||||
if root in state:
|
||||
continue
|
||||
|
||||
# Stack holds (current_node, children_iterator).
|
||||
# Using an iterator allows us to pause and resume processing a node's children.
|
||||
stack = [(root, iter(_get_flat_args_unique(root, node_to_additional_deps)))]
|
||||
state[root] = 1 # Visiting
|
||||
|
||||
while stack:
|
||||
parent, children = stack[-1]
|
||||
|
||||
try:
|
||||
child = next(children)
|
||||
|
||||
if not isinstance(child, Node):
|
||||
continue
|
||||
|
||||
child_state = state.get(child, 0)
|
||||
|
||||
if child_state == 1:
|
||||
# Back-edge: child is on the current DFS path -> cycle
|
||||
cycle_path = [node for node, _ in stack] + [child]
|
||||
return f"cycle detected in path: {cycle_path}"
|
||||
|
||||
if child_state == 0:
|
||||
state[child] = 1
|
||||
stack.append(
|
||||
(
|
||||
child,
|
||||
iter(_get_flat_args_unique(child, node_to_additional_deps)),
|
||||
)
|
||||
)
|
||||
# child_state == 2 means already verified safe; skip.
|
||||
|
||||
except StopIteration:
|
||||
# All children processed — mark safe and pop.
|
||||
stack.pop()
|
||||
state[parent] = 2
|
||||
|
||||
return "no cycle detected"
|
||||
|
||||
|
||||
def _graph_device_type(graph: Graph | None) -> str:
|
||||
if graph is None:
|
||||
return "cpu"
|
||||
|
||||
def _device_type(x: Any) -> str:
|
||||
if isinstance(x, torch.device):
|
||||
return x.type
|
||||
if isinstance(x, torch.Tensor):
|
||||
return x.device.type
|
||||
return "cpu"
|
||||
|
||||
def _flatten_meta(node: Node, key: str) -> list[Any]:
|
||||
if key not in node.meta:
|
||||
return []
|
||||
flat, _ = tree_flatten(node.meta[key])
|
||||
return flat
|
||||
|
||||
for node in graph.nodes:
|
||||
for key in ("val", "example_value"):
|
||||
for obj in _flatten_meta(node, key):
|
||||
return _device_type(obj)
|
||||
|
||||
# Check for device conversions
|
||||
if node.op == "call_method":
|
||||
for gpu in ["cuda", "xpu"]:
|
||||
if node.target == gpu:
|
||||
return gpu
|
||||
if node.target == "to" and gpu in node.args:
|
||||
return gpu
|
||||
|
||||
# Check args/kwargs for non-CPU device specs
|
||||
flat_args, _ = tree_flatten((node.args, node.kwargs))
|
||||
for obj in flat_args:
|
||||
return _device_type(obj)
|
||||
return "cpu"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,26 @@
|
||||
"""Hook system for Dynamo's guard functionality.
|
||||
|
||||
This module provides a way to register callback functions that are triggered during
|
||||
guard-related operations.
|
||||
|
||||
The Hooks class manages two types of hook functions:
|
||||
- guard_export_fn: Called when guards need to be exported, taking a GuardsSet as input
|
||||
- guard_fail_fn: Called when a guard check fails, taking a GuardFail object as input
|
||||
These hooks enable customization of guard export and failure handling behaviors.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
from torch._guards import GuardsSet
|
||||
|
||||
from .types import GuardFail, GuardFilterEntry
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Hooks:
|
||||
guard_export_fn: Callable[[GuardsSet], None] | None = None
|
||||
guard_fail_fn: Callable[[GuardFail], None] | None = None
|
||||
guard_filter_fn: Callable[[Sequence[GuardFilterEntry]], Sequence[bool]] | None = (
|
||||
None
|
||||
)
|
||||
@@ -0,0 +1,73 @@
|
||||
"""Logging utilities for Dynamo and Inductor.
|
||||
|
||||
This module provides specialized logging functionality including:
|
||||
- Step-based logging that prepends step numbers to log messages
|
||||
- Progress bar management for compilation phases
|
||||
- Centralized logger management for Dynamo and Inductor components
|
||||
|
||||
The logging system helps track the progress of compilation phases and provides structured
|
||||
logging output for debugging and monitoring.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from torch.hub import _Faketqdm, tqdm
|
||||
|
||||
|
||||
# Disable progress bar by default, not in dynamo config because otherwise get a circular import
|
||||
disable_progress = True
|
||||
|
||||
|
||||
# Return all loggers that torchdynamo/torchinductor is responsible for
|
||||
def get_loggers() -> list[logging.Logger]:
|
||||
return [
|
||||
logging.getLogger("torch.fx.experimental.symbolic_shapes"),
|
||||
logging.getLogger("torch._dynamo"),
|
||||
logging.getLogger("torch._inductor"),
|
||||
]
|
||||
|
||||
|
||||
# Creates a logging function that logs a message with a step # prepended.
|
||||
# get_step_logger should be lazily called (i.e. at runtime, not at module-load time)
|
||||
# so that step numbers are initialized properly. e.g.:
|
||||
|
||||
# @functools.cache
|
||||
# def _step_logger():
|
||||
# return get_step_logger(logging.getLogger(...))
|
||||
|
||||
# def fn():
|
||||
# _step_logger()(logging.INFO, "msg")
|
||||
|
||||
_step_counter = itertools.count(1)
|
||||
|
||||
# Update num_steps if more phases are added: Dynamo, AOT, Backend
|
||||
# This is very inductor centric
|
||||
# _inductor.utils.has_triton() gives a circular import error here
|
||||
|
||||
if not disable_progress:
|
||||
try:
|
||||
import triton # noqa: F401
|
||||
|
||||
num_steps = 3
|
||||
except ImportError:
|
||||
num_steps = 2
|
||||
pbar = tqdm(total=num_steps, desc="torch.compile()", delay=0)
|
||||
|
||||
|
||||
def get_step_logger(logger: logging.Logger) -> Callable[..., None]:
|
||||
if not disable_progress:
|
||||
pbar.update(1)
|
||||
if not isinstance(pbar, _Faketqdm):
|
||||
pbar.set_postfix_str(f"{logger.name}")
|
||||
|
||||
step = next(_step_counter)
|
||||
|
||||
def log(level: int, msg: str, **kwargs: Any) -> None:
|
||||
if "stacklevel" not in kwargs:
|
||||
kwargs["stacklevel"] = 2
|
||||
logger.log(level, "Step %s: %s", step, msg, **kwargs)
|
||||
|
||||
return log
|
||||
@@ -0,0 +1,251 @@
|
||||
"""Metrics collection and management system for Dynamo.
|
||||
|
||||
This module provides context managers for gathering and reporting metrics during
|
||||
compilation and runtime.
|
||||
|
||||
It includes two main components:
|
||||
- MetricsContext: A context manager for collecting metrics during compilation, supporting
|
||||
nested contexts and various metric types (counters, sets, key-value pairs)
|
||||
- RuntimeMetricsContext: A specialized context for runtime metrics collection that doesn't
|
||||
require explicit context management
|
||||
|
||||
The metrics system enables comprehensive monitoring and analysis of both compilation and
|
||||
execution performance.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TYPE_CHECKING, TypeAlias
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
from torch.utils._traceback import CapturedTraceback
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TopN:
|
||||
"""
|
||||
Helper to record a list of metrics, keeping only the top N "most expensive" elements.
|
||||
"""
|
||||
|
||||
def __init__(self, at_most: int = 25) -> None:
|
||||
self.at_most = at_most
|
||||
self.heap: list[tuple[int, Any]] = []
|
||||
|
||||
def add(self, key: Any, val: int) -> None:
|
||||
# Push if we haven't reached the max size, else push and pop the smallest
|
||||
fn = heapq.heappush if len(self.heap) < self.at_most else heapq.heappushpop
|
||||
fn(self.heap, (val, key))
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.heap)
|
||||
|
||||
def __iter__(self) -> Iterator[tuple[Any, int]]:
|
||||
return ((key, val) for val, key in sorted(self.heap, reverse=True))
|
||||
|
||||
|
||||
OnExitType: TypeAlias = Callable[
|
||||
[int, int, dict[str, Any], type[BaseException] | None, BaseException | None],
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
class MetricsContext:
|
||||
def __init__(self, on_exit: OnExitType) -> None:
|
||||
"""
|
||||
Use this class as a contextmanager to create a context under which to accumulate
|
||||
a set of metrics, e.g., metrics gathered during a compilation. On exit of the
|
||||
contextmanager, call the provided 'on_exit' function and pass a dictionary of
|
||||
all metrics set during the lifetime of the contextmanager.
|
||||
"""
|
||||
self._on_exit = on_exit
|
||||
self._metrics: dict[str, Any] = {}
|
||||
self._start_time_ns: int = 0
|
||||
self._level: int = 0
|
||||
self._edits: list[tuple[CapturedTraceback, set[str]]] = []
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
"""
|
||||
Initialize metrics recording.
|
||||
"""
|
||||
if self._level == 0:
|
||||
# In case of recursion, track at the outermost context.
|
||||
self._metrics = {}
|
||||
self._start_time_ns = time.time_ns()
|
||||
|
||||
self._level += 1
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
_traceback: Any,
|
||||
) -> None:
|
||||
"""
|
||||
At exit, call the provided on_exit function.
|
||||
"""
|
||||
self._level -= 1
|
||||
assert self._level >= 0
|
||||
if self._level == 0:
|
||||
try:
|
||||
end_time_ns = time.time_ns()
|
||||
self._on_exit(
|
||||
self._start_time_ns, end_time_ns, self._metrics, exc_type, exc_value
|
||||
)
|
||||
except Exception:
|
||||
log.exception("Unexpected exception logging compilation metrics")
|
||||
|
||||
def in_progress(self) -> bool:
|
||||
"""
|
||||
True if we've entered the context.
|
||||
"""
|
||||
return self._level > 0
|
||||
|
||||
def increment(self, metric: str, value: int) -> None:
|
||||
"""
|
||||
Increment a metric by a given amount.
|
||||
"""
|
||||
if self._level == 0:
|
||||
raise RuntimeError(f"Cannot increment {metric} outside of a MetricsContext")
|
||||
if metric not in self._metrics:
|
||||
self._metrics[metric] = 0
|
||||
self._metrics[metric] += value
|
||||
|
||||
def _render_edits(self, pred: set[str]) -> str:
|
||||
return "\n\n" + "\n\n".join(
|
||||
"Previous Traceback:\n" + "".join(e.format())
|
||||
for e, k in self._edits
|
||||
if k & pred
|
||||
)
|
||||
|
||||
def set(self, metric: str, value: Any, overwrite: bool = False) -> None:
|
||||
"""
|
||||
Set a metric to a given value. Raises if the metric has been assigned previously
|
||||
in the current context.
|
||||
"""
|
||||
if self._level == 0:
|
||||
raise RuntimeError(f"Cannot set {metric} outside of a MetricsContext")
|
||||
if metric in self._metrics and not overwrite:
|
||||
raise RuntimeError(
|
||||
self._render_edits({metric})
|
||||
+ f"\n\nRuntimeError: Metric '{metric}' has already been set in the current context "
|
||||
"(see above for current and previous traceback)."
|
||||
)
|
||||
self._edits.append((CapturedTraceback.extract(skip=1), {metric}))
|
||||
self._metrics[metric] = value
|
||||
|
||||
def set_key_value(self, metric: str, key: str, value: Any) -> None:
|
||||
"""
|
||||
Treats a give metric as a dictionary and set the k and value within it.
|
||||
Note that the metric must be a dictionary or not present.
|
||||
|
||||
We allow this to be called multiple times (i.e. for features, it's not uncommon
|
||||
for them to be used multiple times within a single compilation).
|
||||
"""
|
||||
if self._level == 0:
|
||||
raise RuntimeError(f"Cannot set {metric} outside of a MetricsContext")
|
||||
if metric not in self._metrics:
|
||||
self._metrics[metric] = {}
|
||||
self._metrics[metric][key] = value
|
||||
|
||||
def update(self, values: dict[str, Any], overwrite: bool = False) -> None:
|
||||
"""
|
||||
Set multiple metrics directly. This method does NOT increment. Raises if any
|
||||
metric has been assigned previously in the current context and overwrite is
|
||||
not set to True.
|
||||
"""
|
||||
if self._level == 0:
|
||||
raise RuntimeError("Cannot update metrics outside of a MetricsContext")
|
||||
existing = self._metrics.keys() & values.keys()
|
||||
if existing and not overwrite:
|
||||
raise RuntimeError(
|
||||
self._render_edits(set(values.keys()))
|
||||
+ f"\n\nRuntimeError: Metric(s) {existing} have already been set in the current context. "
|
||||
"(see above for current and previous traceback)."
|
||||
)
|
||||
self._edits.append((CapturedTraceback.extract(skip=1), set(values.keys())))
|
||||
self._metrics.update(values)
|
||||
|
||||
def update_outer(self, values: dict[str, Any]) -> None:
|
||||
"""
|
||||
Update, but only when at the outermost context.
|
||||
"""
|
||||
if self._level == 0:
|
||||
raise RuntimeError("Cannot update metrics outside of a MetricsContext")
|
||||
if self._level == 1:
|
||||
self.update(values)
|
||||
|
||||
def add_to_set(self, metric: str, value: Any) -> None:
|
||||
"""
|
||||
Records a metric as a set() of values.
|
||||
"""
|
||||
if self._level == 0:
|
||||
raise RuntimeError(f"Cannot add {metric} outside of a MetricsContext")
|
||||
if metric not in self._metrics:
|
||||
self._metrics[metric] = set()
|
||||
self._metrics[metric].add(value)
|
||||
|
||||
def add_top_n(self, metric: str, key: Any, val: int) -> None:
|
||||
"""
|
||||
Records a metric as a TopN set of values.
|
||||
"""
|
||||
if self._level == 0:
|
||||
return
|
||||
if metric not in self._metrics:
|
||||
self._metrics[metric] = TopN()
|
||||
self._metrics[metric].add(key, val)
|
||||
|
||||
|
||||
class RuntimeMetricsContext:
|
||||
def __init__(self, on_exit: OnExitType) -> None:
|
||||
"""
|
||||
Similar to MetricsContext, but used to gather the runtime metrics that are
|
||||
decoupled from compilation, where there's not a natural place to insert a
|
||||
context manager.
|
||||
"""
|
||||
self._on_exit = on_exit
|
||||
self._metrics: dict[str, Any] = {}
|
||||
self._start_time_ns: int = 0
|
||||
|
||||
def increment(
|
||||
self, metric: str, value: int, extra: dict[str, Any] | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Increment a metric by a given amount.
|
||||
"""
|
||||
if not self._metrics:
|
||||
# Start timing on the first entry
|
||||
self._start_time_ns = time.time_ns()
|
||||
if metric not in self._metrics:
|
||||
self._metrics[metric] = 0
|
||||
self._metrics[metric] += value
|
||||
|
||||
if extra:
|
||||
for k, v in extra.items():
|
||||
if k not in self._metrics and v is not None:
|
||||
self._metrics[k] = v
|
||||
|
||||
def finish(self) -> None:
|
||||
"""
|
||||
Call the on_exit function with the metrics gathered so far and reset.
|
||||
"""
|
||||
if self._metrics:
|
||||
try:
|
||||
end_time_ns = time.time_ns()
|
||||
self._on_exit(
|
||||
self._start_time_ns, end_time_ns, self._metrics, None, None
|
||||
)
|
||||
except Exception:
|
||||
log.exception("Unexpected exception logging runtime metrics")
|
||||
finally:
|
||||
self._metrics = {}
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Mutation tracking and dynamic module detection system for Dynamo.
|
||||
|
||||
This module provides mechanisms to track and respond to mutations in PyTorch modules
|
||||
and detect dynamically created or modified modules.
|
||||
|
||||
Key components:
|
||||
- MutationTracker: Tracks mutations to objects and invalidates associated cached code
|
||||
- GenerationTracker: Tracks module creation timing to identify dynamic instances
|
||||
- Patching system for nn.Module to detect mutations and dynamic creation
|
||||
|
||||
The system ensures that Dynamo's optimizations remain valid by detecting and responding
|
||||
to runtime changes in module state and structure.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import weakref
|
||||
from collections.abc import MutableMapping
|
||||
from typing import Any
|
||||
|
||||
import torch.nn
|
||||
from torch.nn import Module
|
||||
|
||||
from . import config
|
||||
from .utils import ExactWeakKeyDictionary, nn_module_has_global_hooks
|
||||
|
||||
|
||||
unpatched_nn_module_init = torch.nn.Module.__init__
|
||||
|
||||
|
||||
class MutationTracker:
|
||||
db: ExactWeakKeyDictionary = ExactWeakKeyDictionary()
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.mutation_count: int = 0
|
||||
self.watchers: list[weakref.ReferenceType[Any]] = []
|
||||
|
||||
def on_mutation(self, name: str) -> None:
|
||||
self.mutation_count += 1
|
||||
tmp = self.watchers
|
||||
self.watchers = []
|
||||
for ref in tmp:
|
||||
guarded = ref()
|
||||
if guarded is not None:
|
||||
guarded.invalidate(ref)
|
||||
|
||||
def track(self, guarded_code: Any) -> None:
|
||||
self.watchers.append(weakref.ref(guarded_code))
|
||||
|
||||
|
||||
def watch(obj: Any, guarded_code: Any) -> None:
|
||||
"""invalidate guarded_code when obj is mutated"""
|
||||
ensure_patched(type(obj))
|
||||
|
||||
if obj not in MutationTracker.db:
|
||||
MutationTracker.db[obj] = MutationTracker()
|
||||
tracker = MutationTracker.db[obj]
|
||||
tracker.track(guarded_code)
|
||||
|
||||
|
||||
def ensure_patched(cls: Any) -> None:
|
||||
if getattr(cls, "___needs_mutation_patch", True):
|
||||
cls.___needs_mutation_patch = False
|
||||
original_setattr = cls.__setattr__
|
||||
|
||||
@functools.wraps(original_setattr)
|
||||
def custom_setattr(self: Any, key: str, value: Any) -> None:
|
||||
try:
|
||||
MutationTracker.db[self].on_mutation(key)
|
||||
except KeyError:
|
||||
pass
|
||||
return original_setattr(self, key, value)
|
||||
|
||||
cls.__setattr__ = custom_setattr
|
||||
|
||||
|
||||
class GenerationTracker:
|
||||
generation: int = 0
|
||||
dynamic_classes: ExactWeakKeyDictionary = ExactWeakKeyDictionary()
|
||||
generation_values: ExactWeakKeyDictionary = ExactWeakKeyDictionary()
|
||||
|
||||
@classmethod
|
||||
def tag(cls, obj: Any) -> None:
|
||||
cls.generation_values[obj] = cls.generation
|
||||
|
||||
@staticmethod
|
||||
def mark_class_dynamic(cls: type[torch.nn.Module]) -> None:
|
||||
assert issubclass(cls, torch.nn.Module)
|
||||
GenerationTracker.dynamic_classes[cls] = True
|
||||
|
||||
@classmethod
|
||||
def get_generation_value(cls, obj: Any) -> int:
|
||||
if obj not in cls.generation_values:
|
||||
return -1
|
||||
return cls.generation_values[obj]
|
||||
|
||||
@classmethod
|
||||
def check(cls, obj: Any) -> bool:
|
||||
return (
|
||||
obj in cls.generation_values
|
||||
and cls.generation_values[obj] == cls.generation
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def clear(cls) -> None:
|
||||
cls.generation = 0
|
||||
cls.dynamic_classes = ExactWeakKeyDictionary()
|
||||
cls.generation_values = ExactWeakKeyDictionary()
|
||||
|
||||
|
||||
def is_dynamic_nn_module(obj: Any, is_export: bool) -> bool:
|
||||
"""Check for nn.Modules() created dynamically or mutated"""
|
||||
if isinstance(obj, torch.nn.Module) and (
|
||||
"forward" in obj.__dict__ or isinstance(obj, (dict, MutableMapping))
|
||||
):
|
||||
# A monkey patched `.forward` indicates something wacky is going on
|
||||
# Similarly a nn module also subclassed as a dict is unusual.
|
||||
return True
|
||||
if hasattr(obj, "torchdynamo_force_dynamic"):
|
||||
return obj.torchdynamo_force_dynamic
|
||||
if isinstance(obj, torch.nn.Module) and (
|
||||
not is_export or config.install_free_tensors
|
||||
):
|
||||
return True
|
||||
|
||||
if isinstance(obj, torch.nn.Module) and nn_module_has_global_hooks():
|
||||
return True
|
||||
dyn = GenerationTracker.dynamic_classes.get(type(obj)) or GenerationTracker.check(
|
||||
obj
|
||||
)
|
||||
return dyn
|
||||
|
||||
|
||||
def install_generation_tagging_init() -> None:
|
||||
"""
|
||||
Monkey patch torch.nn.Module.__init__ and torch.nn.Module.__setstate__
|
||||
so we can detect nn.Module instances created dynamically inside forward methods.
|
||||
"""
|
||||
|
||||
if getattr(Module, "___needs_generation_tag_patch", True):
|
||||
init = Module.__init__
|
||||
|
||||
def patched_init(self: Module, *args: Any, **kwargs: Any) -> None:
|
||||
init(self, *args, **kwargs)
|
||||
GenerationTracker.tag(self)
|
||||
|
||||
Module.__init__ = patched_init # type: ignore[method-assign]
|
||||
|
||||
setstate = Module.__setstate__
|
||||
|
||||
def patched_setstate(self: Module, state: Any) -> None:
|
||||
setstate(self, state)
|
||||
GenerationTracker.tag(self)
|
||||
|
||||
Module.__setstate__ = patched_setstate # type: ignore[method-assign]
|
||||
|
||||
Module.___needs_generation_tag_patch = False # type: ignore[attr-defined]
|
||||
|
||||
GenerationTracker.generation += 1
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,700 @@
|
||||
"""
|
||||
Python polyfills for common builtins.
|
||||
"""
|
||||
|
||||
# NOTE: 1. Please do not import any submodule in the directory here to avoid circular imports.
|
||||
# 2. While adding a new polyfill module, also add it to POLYFILLED_MODULE_NAMES in loader.py.
|
||||
# Add it in the TYPE_CHECKING block below as well.
|
||||
|
||||
import types
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
|
||||
from itertools import repeat as _repeat
|
||||
from operator import eq, ne
|
||||
from typing import Any, TYPE_CHECKING, TypeGuard, TypeVar
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
U = TypeVar("U")
|
||||
C = TypeVar("C")
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..utils import dict_keys
|
||||
|
||||
# Load by torch._dynamo.polyfills.loader
|
||||
# See also the POLYFILLED_MODULE_NAMES in torch/_dynamo/polyfills/loader.py
|
||||
# Put the submodules here to avoid circular imports
|
||||
from . import (
|
||||
_collections as _collections,
|
||||
builtins as builtins,
|
||||
functools as functools,
|
||||
itertools as itertools,
|
||||
operator as operator,
|
||||
os as os,
|
||||
pytree as pytree,
|
||||
struct as struct,
|
||||
sys as sys,
|
||||
torch_c_nn as torch_c_nn,
|
||||
traceback as traceback,
|
||||
)
|
||||
|
||||
from torch.overrides import BaseTorchFunctionMode
|
||||
|
||||
|
||||
# These classes handle support for TorchFunctionModes across
|
||||
# graph breaks
|
||||
# Today the TorchFunctionMode enter (for the classes we support)
|
||||
# simply pushes the mode onto the stack. Since after this occurs
|
||||
# the stack is mutated, and we replay these mutations, we don't need
|
||||
# any cleanup logic to be run once the graph break occurs, we simply replay
|
||||
# these mutations to ensure at the graph break the torch function mode stack is correct
|
||||
# and reconstruct the torch function mode stack normally
|
||||
# when we compile the resume function on the other side of the break.
|
||||
# However, to ensure we exit properly
|
||||
# in the resume function, we need to re-enter the contexts as we do other contexts.
|
||||
# These contexts do nothing on enter, but provide the correct exit logic to ensure
|
||||
# the stack state is correct.
|
||||
class NoEnterTorchFunctionMode(BaseTorchFunctionMode):
|
||||
def __enter__(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
# Used by WrappedUserFunctionVariable and similar to inline decorated function
|
||||
# calls with bytecode backing. Without this, the context enter/exit happens in
|
||||
# Python-level VT code, so a nested graph break inside `fn` would skip applying
|
||||
# the context in the compiled fn/resume. By inlining through this polyfill, the
|
||||
# `with` statement has real bytecode that the resume function can continue from.
|
||||
def _fn_with_ctx(ctx: Any, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
|
||||
with ctx:
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
|
||||
def index(
|
||||
iterator: Iterator[T], item: T, start: int = 0, end: int | None = None
|
||||
) -> int:
|
||||
from itertools import islice
|
||||
|
||||
for i, elem in islice(enumerate(iterator), start, end):
|
||||
if item == elem:
|
||||
return i
|
||||
# This will not run in dynamo
|
||||
raise ValueError(f"{item} is not in {type(iterator)}")
|
||||
|
||||
|
||||
def repeat(item: T, count: int) -> Iterator[T]:
|
||||
for _ in range(count):
|
||||
yield item
|
||||
|
||||
|
||||
def radians(x: float) -> float:
|
||||
import math
|
||||
|
||||
return math.pi / 180.0 * x
|
||||
|
||||
|
||||
def impl_IS_MAPPING(a: object) -> TypeIs[Mapping[Any, Any]]:
|
||||
return isinstance(a, Mapping)
|
||||
|
||||
|
||||
def impl_MATCH_SEQUENCE(a: object) -> TypeGuard[Sequence[Any]]:
|
||||
return isinstance(a, Sequence) and not isinstance(a, (str, bytes, bytearray))
|
||||
|
||||
|
||||
def _match_class_attr(obj: object, name: str, seen: set[str]) -> object:
|
||||
if name in seen:
|
||||
raise TypeError(f"{type(obj)} got multiple sub-patterns for attribute {name}")
|
||||
|
||||
attr = getattr(obj, name)
|
||||
seen.add(name)
|
||||
return attr
|
||||
|
||||
|
||||
def impl_MATCH_CLASS(
|
||||
subject: object, cls: type, nargs: int, kwargs: tuple[str, ...]
|
||||
) -> tuple[object, ...] | None:
|
||||
if not isinstance(cls, type):
|
||||
raise TypeError("called match pattern must be a class")
|
||||
|
||||
if not isinstance(subject, cls):
|
||||
return None
|
||||
|
||||
typ = type(subject)
|
||||
match_self = False
|
||||
match_args = ()
|
||||
|
||||
attrs = []
|
||||
seen = set()
|
||||
|
||||
if nargs:
|
||||
if hasattr(typ, "__match_args__"):
|
||||
match_args = typ.__match_args__
|
||||
|
||||
if not isinstance(match_args, tuple):
|
||||
raise TypeError(
|
||||
f"{typ}.__match_args__ must be a tuple, (got {type(match_args)})"
|
||||
)
|
||||
|
||||
for name in match_args[:nargs]:
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(
|
||||
f"__match_args__ elements must be strings (got {type(name)})"
|
||||
)
|
||||
attrs.append(_match_class_attr(subject, name, seen))
|
||||
else:
|
||||
# We should somehow check if the type has TPFLAGS_MATCH_SELF set
|
||||
# match_self is only true if TPFLAGS_MATCH_SELF is set, but there is
|
||||
# no way to check for it directly in Python. So we assume it is set
|
||||
# if there are no __match_args__
|
||||
match_self = True
|
||||
attrs.append(subject)
|
||||
|
||||
allowed = 1 if match_self else len(match_args)
|
||||
if allowed < nargs:
|
||||
raise TypeError(
|
||||
f"accepts {allowed} positional sub-patterns ({nargs} given)"
|
||||
)
|
||||
|
||||
for name in kwargs:
|
||||
attrs.append(_match_class_attr(subject, name, seen))
|
||||
|
||||
return tuple(attrs)
|
||||
|
||||
|
||||
def impl_MATCH_KEYS(obj: Mapping[T, U], keys: tuple[T, ...]) -> tuple[U, ...] | None:
|
||||
assert isinstance(obj, Mapping)
|
||||
if all(key in obj for key in keys):
|
||||
return tuple(obj[key] for key in keys)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def impl_CONTAINS_OP_fallback(a: T, b: Iterable[T]) -> bool:
|
||||
# performs fallback "a in b"
|
||||
if hasattr(b, "__iter__"):
|
||||
# use __iter__ if __contains__ is not available
|
||||
for x in b:
|
||||
if x == a:
|
||||
return True
|
||||
return False
|
||||
raise TypeError(f"argument of type {type(b)} is not iterable")
|
||||
|
||||
|
||||
def accumulate_grad(x: torch.Tensor, new_grad: torch.Tensor | None) -> None:
|
||||
# polyfills according to the Gradient Layout Contract
|
||||
if new_grad is None:
|
||||
return
|
||||
new_grad_strided = torch.empty_like(x)
|
||||
new_grad_strided.copy_(new_grad)
|
||||
if x.grad is None:
|
||||
x.grad = new_grad_strided
|
||||
elif torch.is_grad_enabled():
|
||||
x.grad = x.grad + new_grad_strided
|
||||
else:
|
||||
x.grad.add_(new_grad_strided)
|
||||
|
||||
|
||||
# This mirrors
|
||||
# https://github.com/python/cpython/blob/a1c52d1265c65bcf0d9edf87e143843ad54f9b8f/Objects/listobject.c#L3352-L3413
|
||||
def list_cmp(
|
||||
op: Callable[[Any, Any], bool], left: Sequence[T], right: Sequence[T]
|
||||
) -> bool:
|
||||
"""emulate `(1,2,3) > (1,2)` etc"""
|
||||
|
||||
# Optimization: For equality, short-circuit if lengths differ
|
||||
# This avoids iterating through elements and triggering guards on SymInts
|
||||
left_len = len(left)
|
||||
right_len = len(right)
|
||||
|
||||
if op is eq and left_len != right_len:
|
||||
return False
|
||||
if op is ne and left_len != right_len:
|
||||
return True
|
||||
|
||||
# Apply `op` to the first pair that differ
|
||||
for a, b in zip(left, right):
|
||||
if a != b:
|
||||
return op(a, b)
|
||||
|
||||
# No more pairs to compare, so compare sizes.
|
||||
return op(left_len, right_len)
|
||||
|
||||
|
||||
def dict___eq__(d: dict[T, U], other: dict[T, U]) -> bool:
|
||||
if (len(d) != len(other)) or (d.keys() != other.keys()):
|
||||
return False
|
||||
|
||||
if all(isinstance(a, OrderedDict) for a in (d, other)):
|
||||
return list(d.items()) == list(other.items())
|
||||
|
||||
for k, v in d.items():
|
||||
if v != other[k]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def set_symmetric_difference(
|
||||
set1: Iterable[T],
|
||||
set2: Iterable[T],
|
||||
cls: type[Any] = set,
|
||||
) -> Any:
|
||||
symmetric_difference_set: set[T] = set()
|
||||
for x in set1:
|
||||
if x not in set2:
|
||||
symmetric_difference_set.add(x)
|
||||
for x in set2:
|
||||
if x not in set1:
|
||||
symmetric_difference_set.add(x)
|
||||
return cls(symmetric_difference_set)
|
||||
|
||||
|
||||
def set_symmetric_difference_update(set1: set[T], set2: set[T]) -> None:
|
||||
result = set1.symmetric_difference(set2)
|
||||
set1.clear()
|
||||
set1.update(result)
|
||||
|
||||
|
||||
def set_isdisjoint(set1: set[T], set2: set[T]) -> bool:
|
||||
if not isinstance(set2, Iterable):
|
||||
raise TypeError(f"'{type(set2)}' object is not iterable")
|
||||
|
||||
for x in set1:
|
||||
for y in set2:
|
||||
if not isinstance(y, Hashable):
|
||||
raise TypeError(f"unhashable type: '{type(y)}'")
|
||||
if x == y:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def set_intersection(
|
||||
set1: set[T],
|
||||
*others: Iterable[T],
|
||||
# See facebook/pyrefly#1496 - leave generic
|
||||
cls: type[Any] = set,
|
||||
) -> Any:
|
||||
if len(others) == 0:
|
||||
return set1.copy()
|
||||
|
||||
if not all(isinstance(s, Iterable) for s in others):
|
||||
raise TypeError(f"set.difference expected an iterable, got {type(others)}")
|
||||
|
||||
for s in others:
|
||||
if any(not isinstance(x, Hashable) for x in s):
|
||||
raise TypeError("unhashable type")
|
||||
|
||||
# return a new set with elements common in all sets
|
||||
intersection_set = set()
|
||||
for x in set1:
|
||||
for set2 in others:
|
||||
if not any(x == y for y in set2):
|
||||
break
|
||||
else:
|
||||
intersection_set.add(x)
|
||||
return cls(intersection_set)
|
||||
|
||||
|
||||
def set_intersection_update(set1: set[T], *others: Iterable[T]) -> None:
|
||||
result = set1.intersection(*others)
|
||||
set1.clear()
|
||||
set1.update(result)
|
||||
|
||||
|
||||
def set_union(
|
||||
set1: set[T], *others: Iterable[T], cls: type[C] | None = None
|
||||
) -> C | set[T]:
|
||||
# frozenset also uses this function
|
||||
if cls is None:
|
||||
# pyrefly: ignore[bad-assignment]
|
||||
cls = type(set1)
|
||||
|
||||
if len(others) == 0:
|
||||
return set1.copy()
|
||||
|
||||
if not all(isinstance(s, Iterable) for s in others):
|
||||
raise TypeError(f"set.union expected an iterable, got {type(others)}")
|
||||
|
||||
for s in others:
|
||||
if any(not isinstance(x, Hashable) for x in s):
|
||||
raise TypeError("unhashable type")
|
||||
|
||||
union_set = set(set1.copy())
|
||||
for set2 in others:
|
||||
set_update(union_set, set2)
|
||||
|
||||
# frozenset also uses this function
|
||||
# pyrefly: ignore [bad-argument-count, not-callable]
|
||||
return cls(union_set)
|
||||
|
||||
|
||||
# pyrefly: ignore [bad-return]
|
||||
def set_update(set1: set[T], *others: Iterable[T]) -> set[T]:
|
||||
if len(others) == 0:
|
||||
return set1
|
||||
|
||||
for set2 in others:
|
||||
for x in set2:
|
||||
if x not in set1:
|
||||
set1.add(x)
|
||||
|
||||
|
||||
def set_difference(
|
||||
set1: set[T],
|
||||
*others: Iterable[T],
|
||||
cls: type[Any] = set,
|
||||
) -> Any:
|
||||
if len(others) == 0:
|
||||
return set1.copy()
|
||||
|
||||
if not all(isinstance(s, Iterable) for s in others):
|
||||
raise TypeError(f"set.difference expected an iterable, got {type(others)}")
|
||||
|
||||
for s in others:
|
||||
if any(not isinstance(x, Hashable) for x in s):
|
||||
raise TypeError("unhashable type")
|
||||
|
||||
difference_set = set()
|
||||
for x in set1:
|
||||
for set2 in others:
|
||||
if x in set2:
|
||||
break
|
||||
else:
|
||||
difference_set.add(x)
|
||||
return cls(difference_set)
|
||||
|
||||
|
||||
def set_difference_update(set1: set[T], *others: Iterable[T]) -> None:
|
||||
result = set1.difference(*others)
|
||||
set1.clear()
|
||||
set1.update(result)
|
||||
|
||||
|
||||
def assert_dict_equal(
|
||||
self_: Any, d1: dict[T, U], d2: dict[T, U], msg: str | None = None
|
||||
) -> None:
|
||||
self_.assertTrue(d1 == d2, msg)
|
||||
|
||||
|
||||
def assert_multi_line_equal(
|
||||
self_: Any, first: T, second: T, msg: str | None = None
|
||||
) -> None:
|
||||
return self_.assertTrue(first == second, msg)
|
||||
|
||||
|
||||
# The original impl. uses difflib
|
||||
def assert_sequence_equal(
|
||||
self_: Any,
|
||||
seq1: Sequence[T],
|
||||
seq2: Sequence[T],
|
||||
msg: str | None = None,
|
||||
seq_type: type[Any] | None = None,
|
||||
) -> None:
|
||||
return self_.assertTrue(seq1 == seq2, msg)
|
||||
|
||||
|
||||
def getattr_and_trace(*args: Any, **kwargs: Any) -> Any:
|
||||
wrapper_obj = args[0]
|
||||
attr_name = args[1]
|
||||
fn = getattr(wrapper_obj, attr_name)
|
||||
return fn(*args[2:], **kwargs)
|
||||
|
||||
|
||||
def mapping_get(obj: Mapping[T, U], key: T, value: U | None = None, /) -> U | None:
|
||||
try:
|
||||
return obj.__getitem__(key)
|
||||
except KeyError:
|
||||
return value
|
||||
|
||||
|
||||
def instantiate_user_defined_class_object(
|
||||
cls: type[T], /, *args: Any, **kwargs: Any
|
||||
) -> T:
|
||||
obj = cls.__new__(cls, *args, **kwargs)
|
||||
|
||||
# Only call __init__ if the object's type is a subclass of cls.
|
||||
# CPython uses PyType_IsSubtype(Py_TYPE(obj), type) at the C level, which does NOT
|
||||
# go through metaclass __instancecheck__. Using isinstance() here would be wrong
|
||||
# for classes with custom __instancecheck__ (e.g. torch.ByteStorage).
|
||||
# Reference: https://github.com/python/cpython/blob/3.12/Objects/typeobject.c#L1670-L1673
|
||||
if issubclass(type(obj), cls):
|
||||
obj.__init__(*args, **kwargs)
|
||||
return obj
|
||||
|
||||
|
||||
def mutable_mapping_update(
|
||||
self,
|
||||
data: Mapping[T, U] | Iterable[tuple[T, U]] = (),
|
||||
/,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
if isinstance(data, Mapping):
|
||||
# Merge standard mapping with PyMapping_Items
|
||||
for key, value in data.items():
|
||||
self[key] = value
|
||||
# FIXME: Enabling the `elif`-branch below needs too many `VariableClass.call_obj_hasattr` changes.
|
||||
# >>> class Foo:
|
||||
# ... def __init__(self):
|
||||
# ... self.keys = lambda: ['a', 'b', 'c'] # not required to be a method
|
||||
# ...
|
||||
# ... def __getitem__(self, key):
|
||||
# ... return 0
|
||||
# ...
|
||||
# >>> dict(Foo())
|
||||
# {'a': 0, 'b': 0, 'c': 0}
|
||||
#
|
||||
# > This is a rare case, so we comment it out for now.
|
||||
#
|
||||
# elif hasattr(data, "keys"):
|
||||
# # Merge mapping-like object with PyMapping_Keys + PyObject_GetItem
|
||||
# for key in data.keys():
|
||||
# self[key] = data[key]
|
||||
else:
|
||||
if not isinstance(data, Iterable):
|
||||
raise TypeError(f"{type(data).__name__!r} object is not iterable")
|
||||
# Likely a sequence of pairs
|
||||
for key, value in data:
|
||||
self[key] = value
|
||||
|
||||
if kwargs:
|
||||
for key, value in kwargs.items():
|
||||
self[key] = value
|
||||
|
||||
|
||||
# Used with something like dict(obj)
|
||||
def construct_dict(
|
||||
cls: type[T],
|
||||
data: Mapping[object, object] | Iterable[tuple[object, object]] = (),
|
||||
/,
|
||||
**kwargs: Any,
|
||||
) -> T:
|
||||
self = cls.__new__(cls)
|
||||
mutable_mapping_update(self, data, **kwargs)
|
||||
return self
|
||||
|
||||
|
||||
def foreach_map_fn(*args: Any) -> Any:
|
||||
op = args[0]
|
||||
new_args: list[Any] = []
|
||||
at_least_one_list = False
|
||||
for arg in args[1:]:
|
||||
if not isinstance(arg, (list, tuple)):
|
||||
new_args.append(_repeat(arg))
|
||||
else:
|
||||
at_least_one_list = True
|
||||
new_args.append(arg)
|
||||
|
||||
# Just apply op once to args if there are no lists
|
||||
if not at_least_one_list:
|
||||
return op(*args[1:])
|
||||
|
||||
out = []
|
||||
for unpacked in zip(*new_args):
|
||||
out.append(op(*unpacked))
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def foreach_lerp_inplace(
|
||||
self,
|
||||
end: list[torch.Tensor] | tuple[torch.Tensor, ...],
|
||||
weight: float | int | torch.Tensor,
|
||||
) -> None:
|
||||
# Decompose lerp via addcmul_ for FMA. Uses the same dual-formula
|
||||
# approach as CUDA's native lerp to get bitwise identical results:
|
||||
# |w| < 0.5 (low): fma(w, diff, start)
|
||||
# |w| >= 0.5 (high): fma(-(1-w), diff, end)
|
||||
# For tensor weights (e.g. 0-dim tensor from tensor betas in Adam) the
|
||||
# low formula is always used because the native lerp_scalar lowering
|
||||
# would crash on float(weight) for symbolic expressions.
|
||||
diff = torch._foreach_sub(end, self)
|
||||
if isinstance(weight, torch.Tensor):
|
||||
# Select base and weight for the dual formula before a single addcmul:
|
||||
# low (|w| < 0.5): fma(w, diff, self)
|
||||
# high (|w| >= 0.5): fma(-(1-w), diff, end)
|
||||
mask = weight.abs() >= 0.5
|
||||
neg_omw = -(1.0 - weight)
|
||||
w = torch.where(mask, neg_omw, weight)
|
||||
bases = [torch.where(mask, e, s) for s, e in zip(self, end)]
|
||||
w_list = [w] * len(diff)
|
||||
torch._foreach_addcmul_(bases, w_list, diff)
|
||||
for s, b in zip(self, bases):
|
||||
s.copy_(b)
|
||||
else:
|
||||
abs_weight = weight if weight >= 0 else -weight
|
||||
if abs_weight >= 0.5:
|
||||
# High formula: end + (-(1-w)) * diff → fma(-(1-w), diff, end)
|
||||
# Compute 1-w in target dtype to match CUDA rounding.
|
||||
d0 = self[0]
|
||||
neg_omw = -(1.0 - torch.tensor(weight, dtype=d0.dtype, device=d0.device))
|
||||
neg_omw_list = [neg_omw] * len(diff)
|
||||
for s, e in zip(self, end):
|
||||
s.copy_(e)
|
||||
torch._foreach_addcmul_(self, neg_omw_list, diff)
|
||||
else:
|
||||
# Low formula: start + w * diff → fma(w, diff, start)
|
||||
weights = [torch.full_like(d, weight) for d in diff]
|
||||
torch._foreach_addcmul_(self, weights, diff)
|
||||
return self
|
||||
|
||||
|
||||
def foreach_pow_scalar(
|
||||
scalar: Any, exps: Sequence[bool | complex | float | int]
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
return torch._foreach_pow([scalar for _ in exps], exps)
|
||||
|
||||
|
||||
def predicate(obj: object) -> bool:
|
||||
# This will cause the rest of dynamo to handle the if statement correctly, so we don't have to rewrite it here.
|
||||
# We can't just use bool() here since we can't trace into that in general.
|
||||
if obj:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def cmp_eq(a: object, b: object) -> bool:
|
||||
# Note that the commented `is` check should ideally be removed. This is a
|
||||
# CPython optimization that skips the __eq__ checks it the obj id's are
|
||||
# same. But, these lines adds many `is` nodes in the Fx graph for
|
||||
# SymNodeVariable. For now, we can just skip this check. This is STILL
|
||||
# correct because one of the __eq__ checks will pass later, just could be
|
||||
# slow in some corner cases.
|
||||
# if a is b:
|
||||
# return True
|
||||
if isinstance(a, type):
|
||||
# Default metaclass equality is identity-based. Preserve the reflected
|
||||
# operand fallback without tracing through type.__eq__.
|
||||
if type(a).__eq__ is type.__eq__:
|
||||
result = True if a is b else NotImplemented
|
||||
else:
|
||||
result = type(a).__eq__(a, b)
|
||||
else:
|
||||
result = a.__eq__(b)
|
||||
if result is NotImplemented:
|
||||
if isinstance(b, type):
|
||||
if type(b).__eq__ is type.__eq__:
|
||||
result = True if a is b else NotImplemented
|
||||
else:
|
||||
result = type(b).__eq__(b, a)
|
||||
else:
|
||||
result = b.__eq__(a)
|
||||
return result is not NotImplemented and result
|
||||
|
||||
|
||||
def cmp_ne(a: object, b: object) -> bool:
|
||||
if isinstance(a, type):
|
||||
if type(a).__ne__ is type.__ne__:
|
||||
result = False if a is b else NotImplemented
|
||||
else:
|
||||
result = type(a).__ne__(a, b)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
elif isinstance(type(a).__ne__, types.FunctionType):
|
||||
result = a.__ne__(b)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
# Fall through to try b.__ne__(a) or cmp_eq
|
||||
if isinstance(b, type):
|
||||
if type(b).__ne__ is type.__ne__:
|
||||
result = False if a is b else NotImplemented
|
||||
else:
|
||||
result = type(b).__ne__(b, a)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
elif isinstance(type(b).__ne__, types.FunctionType):
|
||||
result = b.__ne__(a)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
return not cmp_eq(a, b)
|
||||
|
||||
|
||||
def cmp_lt(a: Any, b: Any) -> bool:
|
||||
result = a.__lt__(b)
|
||||
if result is NotImplemented:
|
||||
raise TypeError(f"{type(a)} does not support the < operator")
|
||||
return result
|
||||
|
||||
|
||||
def cmp_le(a: Any, b: Any) -> bool:
|
||||
# Check if __le__ is overridden
|
||||
if isinstance(type(a).__le__, types.FunctionType):
|
||||
return a.__le__(b)
|
||||
return cmp_eq(a, b) or cmp_lt(a, b)
|
||||
|
||||
|
||||
def cmp_gt(a: Any, b: Any) -> bool:
|
||||
# Check if __gt__ is overridden
|
||||
if isinstance(type(a).__gt__, types.FunctionType):
|
||||
return a.__gt__(b)
|
||||
# a > b is equivalent to b < a
|
||||
return cmp_lt(b, a)
|
||||
|
||||
|
||||
def cmp_ge(a: Any, b: Any) -> bool:
|
||||
# Check if __ge__ is overridden
|
||||
if isinstance(type(a).__ge__, types.FunctionType):
|
||||
return a.__ge__(b)
|
||||
return cmp_eq(a, b) or cmp_gt(a, b)
|
||||
|
||||
|
||||
def group_tensors_by_device_and_dtype(
|
||||
tensorlistlist: list[list[torch.Tensor | None]], with_indices: bool = False
|
||||
) -> dict[tuple[torch.device, torch.dtype], tuple[list[list[Any]], list[int]]]:
|
||||
"""Pure Python implementation of torch._C._group_tensors_by_device_and_dtype.
|
||||
|
||||
Groups tensors by their device and dtype. This is useful before sending
|
||||
tensors off to a foreach implementation, which requires tensors to be on
|
||||
one device and dtype.
|
||||
|
||||
Args:
|
||||
tensorlistlist: A list of lists of tensors (tensors can be None).
|
||||
with_indices: If True, track original indices in the output.
|
||||
|
||||
Returns:
|
||||
A dict mapping (device, dtype) tuples to (grouped_tensorlistlist, indices).
|
||||
"""
|
||||
# Result dict: (device, dtype) -> (list of lists, indices)
|
||||
result: dict[
|
||||
tuple[torch.device, torch.dtype], tuple[list[list[Any]], list[int]]
|
||||
] = {}
|
||||
|
||||
if not tensorlistlist or not tensorlistlist[0]:
|
||||
return result
|
||||
|
||||
num_lists = len(tensorlistlist)
|
||||
num_tensors = len(tensorlistlist[0])
|
||||
|
||||
for idx in range(num_tensors):
|
||||
# Find the first non-None tensor at this index to get device and dtype
|
||||
first_tensor = None
|
||||
for tlist in tensorlistlist:
|
||||
if tlist is not None and idx < len(tlist) and tlist[idx] is not None:
|
||||
first_tensor = tlist[idx]
|
||||
break
|
||||
|
||||
if first_tensor is None:
|
||||
# All tensors at this index are None, skip
|
||||
continue
|
||||
|
||||
key = (first_tensor.device, first_tensor.dtype)
|
||||
|
||||
if key not in result:
|
||||
# Initialize empty lists for each tensorlist
|
||||
result[key] = ([[] for _ in range(num_lists)], [])
|
||||
|
||||
grouped_lists, indices = result[key]
|
||||
|
||||
# Add tensors from each list at this index
|
||||
for list_idx, tlist in enumerate(tensorlistlist):
|
||||
if tlist is not None and idx < len(tlist):
|
||||
grouped_lists[list_idx].append(tlist[idx])
|
||||
else:
|
||||
grouped_lists[list_idx].append(None)
|
||||
|
||||
if with_indices:
|
||||
indices.append(idx)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,33 @@
|
||||
"""
|
||||
Python polyfills for builtins
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, MutableMapping
|
||||
from typing import TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
try:
|
||||
import _collections # type: ignore[import-not-found]
|
||||
|
||||
@substitute_in_graph(_collections._count_elements)
|
||||
def _count_elements(
|
||||
mapping: MutableMapping[T, int],
|
||||
iterable: Iterable[T],
|
||||
) -> None:
|
||||
"Tally elements from the iterable."
|
||||
mapping_get = mapping.get
|
||||
for elem in iterable:
|
||||
mapping[elem] = mapping_get(elem, 0) + 1
|
||||
|
||||
__all__.append("_count_elements")
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Python polyfills for builtins
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import functools
|
||||
import operator
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
__all__ = [
|
||||
"all",
|
||||
"any",
|
||||
"cast",
|
||||
"enumerate",
|
||||
"sum",
|
||||
]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.all, can_constant_fold_through=True)
|
||||
def all(iterable: Iterable[object], /) -> bool:
|
||||
for elem in iterable:
|
||||
if not elem:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.any, can_constant_fold_through=True)
|
||||
def any(iterable: Iterable[object], /) -> bool:
|
||||
for elem in iterable:
|
||||
if elem:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.enumerate, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def enumerate(iterable: Iterable[_T], start: int = 0) -> Iterable[tuple[int, _T]]:
|
||||
if not isinstance(start, int):
|
||||
raise TypeError(
|
||||
f"{type(start).__name__!r} object cannot be interpreted as an integer"
|
||||
)
|
||||
|
||||
for x in iterable:
|
||||
yield start, x
|
||||
start += 1
|
||||
|
||||
|
||||
@substitute_in_graph(builtins.sum, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def sum(iterable: Iterable[_T], /, start: _T = 0) -> _T: # type: ignore[assignment]
|
||||
return functools.reduce(operator.add, iterable, start)
|
||||
|
||||
|
||||
class _CallableIterator:
|
||||
def __init__(self, fn, sentinel): # type: ignore[no-untyped-def]
|
||||
self.fn = fn
|
||||
self.sentinel = sentinel
|
||||
|
||||
def __iter__(self): # type: ignore[no-untyped-def]
|
||||
return self
|
||||
|
||||
def __next__(self): # type: ignore[no-untyped-def]
|
||||
# The iterator created in this case will call object with no arguments
|
||||
# for each call to its __next__() method;
|
||||
r = self.fn()
|
||||
|
||||
# If the value returned is equal to sentinel, StopIteration will be raised
|
||||
if r == self.sentinel:
|
||||
raise StopIteration
|
||||
|
||||
# otherwise the value will be returned.
|
||||
return r
|
||||
|
||||
|
||||
_sentinel_missing = object()
|
||||
|
||||
|
||||
# TODO(guilhermeleobas): use substitute_in_graph for iter()
|
||||
def iter_(fn_or_iterable, sentinel=_sentinel_missing, /): # type: ignore[no-untyped-def]
|
||||
# Without a second argument, object must be a collection object which supports
|
||||
# the iterable (__iter__) or the sequence protocol (__getitem__ with an integer
|
||||
# starting at 0)
|
||||
if sentinel is _sentinel_missing:
|
||||
iterable = fn_or_iterable
|
||||
if hasattr(iterable, "__iter__"):
|
||||
iterator = iterable.__iter__()
|
||||
if hasattr(iterator, "__next__"):
|
||||
return iterator
|
||||
else:
|
||||
raise TypeError(f"'{type(iterator)}' object is not iterable")
|
||||
if hasattr(iterable, "__getitem__"):
|
||||
# Needs to be a new function to avoid iter becoming a generator
|
||||
def sequence_protocol(iterable): # type: ignore[no-untyped-def]
|
||||
i = 0
|
||||
while True:
|
||||
try:
|
||||
yield iterable.__getitem__(i)
|
||||
i += 1
|
||||
except IndexError:
|
||||
break
|
||||
|
||||
return sequence_protocol(iterable)
|
||||
raise TypeError(f"'{type(iterable)}' object is not iterable")
|
||||
else:
|
||||
# If the second argument, sentinel, is given, then object must be a
|
||||
# callable object.
|
||||
fn = fn_or_iterable
|
||||
|
||||
if not isinstance(fn, Callable): # type: ignore[arg-type]
|
||||
raise TypeError("iter(v, w): v must be a callable")
|
||||
|
||||
return _CallableIterator(fn, sentinel)
|
||||
|
||||
|
||||
@substitute_in_graph(typing.cast, can_constant_fold_through=True)
|
||||
def cast(typ: type, val: _T) -> _T: # type: ignore[type-var]
|
||||
return val
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Python polyfills for copy
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = [
|
||||
"reduce_ex_user_defined_object",
|
||||
]
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@substitute_in_graph(object.__reduce_ex__, skip_signature_check=True) # type: ignore[arg-type]
|
||||
def reduce_ex_user_defined_object(obj: T, protocol: int, /) -> tuple: # type: ignore[type-arg]
|
||||
"""Traceable polyfill for object.__reduce_ex__ on user-defined objects.
|
||||
|
||||
Returns the same tuple that CPython's _common_reduce produces:
|
||||
(copyreg.__newobj__, (cls,), obj.__dict__, None, None).
|
||||
copy._reconstruct then calls cls.__new__(cls) and updates __dict__.
|
||||
"""
|
||||
import copyreg
|
||||
|
||||
cls = type(obj)
|
||||
return (
|
||||
copyreg.__newobj__, # pyrefly: ignore[missing-attribute]
|
||||
(cls,),
|
||||
obj.__dict__,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Python polyfills for functools
|
||||
"""
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = ["reduce"]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_U = TypeVar("_U")
|
||||
|
||||
|
||||
_initial_missing = object()
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/functools.html#functools.reduce
|
||||
@substitute_in_graph(functools.reduce)
|
||||
def reduce(
|
||||
function: Callable[[_U, _T], _U],
|
||||
iterable: Iterable[_T],
|
||||
initial: _U = _initial_missing, # type: ignore[assignment]
|
||||
/,
|
||||
) -> _U:
|
||||
it = iter(iterable)
|
||||
|
||||
value: _U
|
||||
if initial is _initial_missing:
|
||||
try:
|
||||
value = next(it) # type: ignore[assignment]
|
||||
except StopIteration:
|
||||
raise TypeError(
|
||||
"reduce() of empty iterable with no initial value",
|
||||
) from None
|
||||
else:
|
||||
value = initial
|
||||
|
||||
for element in it:
|
||||
value = function(value, element)
|
||||
|
||||
return value
|
||||
@@ -0,0 +1,41 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from torch._C import _fx_map_aggregate, _fx_map_arg
|
||||
from torch.fx.immutable_collections import immutable_dict, immutable_list
|
||||
from torch.fx.node import Node
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
@substitute_in_graph(_fx_map_arg, can_constant_fold_through=True)
|
||||
def map_arg(a: Any, fn: Callable[[Node], Any]) -> Any:
|
||||
return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
|
||||
|
||||
|
||||
@substitute_in_graph(_fx_map_aggregate, can_constant_fold_through=True)
|
||||
def map_aggregate(a: Any, fn: Callable[[Any], Any]) -> Any:
|
||||
result: Any
|
||||
if isinstance(a, tuple):
|
||||
it = (map_aggregate(elem, fn) for elem in a)
|
||||
# Support NamedTuple (if it has `_fields`) by repacking into original type.
|
||||
result = type(a)(*it) if hasattr(a, "_fields") else tuple(it)
|
||||
elif isinstance(a, list):
|
||||
result = immutable_list([map_aggregate(elem, fn) for elem in a])
|
||||
elif isinstance(a, dict):
|
||||
result = immutable_dict([(k, map_aggregate(v, fn)) for k, v in a.items()])
|
||||
elif isinstance(a, slice):
|
||||
result = slice(
|
||||
map_aggregate(a.start, fn),
|
||||
map_aggregate(a.stop, fn),
|
||||
map_aggregate(a.step, fn),
|
||||
)
|
||||
else:
|
||||
result = fn(a)
|
||||
return result
|
||||
|
||||
|
||||
__all__ = [
|
||||
"map_arg",
|
||||
"map_aggregate",
|
||||
]
|
||||
@@ -0,0 +1,119 @@
|
||||
"""
|
||||
Python polyfills for heapq
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
import importlib
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
# Partially copied from CPython test/support/import_helper.py
|
||||
# https://github.com/python/cpython/blob/bb8791c0b75b5970d109e5557bfcca8a578a02af/Lib/test/support/import_helper.py
|
||||
def _save_and_remove_modules(names: set[str]) -> dict[str, ModuleType]:
|
||||
orig_modules = {}
|
||||
prefixes = tuple(name + "." for name in names)
|
||||
for modname in list(sys.modules):
|
||||
if modname in names or modname.startswith(prefixes):
|
||||
orig_modules[modname] = sys.modules.pop(modname)
|
||||
return orig_modules
|
||||
|
||||
|
||||
def import_fresh_module(name: str, blocked: list[str]) -> ModuleType:
|
||||
# Keep track of modules saved for later restoration as well
|
||||
# as those which just need a blocking entry removed
|
||||
names = {name, *blocked}
|
||||
orig_modules = _save_and_remove_modules(names)
|
||||
for modname in blocked:
|
||||
sys.modules[modname] = None # type: ignore[assignment]
|
||||
|
||||
try:
|
||||
return importlib.import_module(name)
|
||||
finally:
|
||||
_save_and_remove_modules(names)
|
||||
sys.modules.update(orig_modules)
|
||||
|
||||
|
||||
# Import the pure Python heapq module, blocking the C extension
|
||||
py_heapq = import_fresh_module("heapq", blocked=["_heapq"])
|
||||
|
||||
|
||||
__all__ = [
|
||||
"_heapify_max",
|
||||
"_heappop_max",
|
||||
"_heapreplace_max",
|
||||
"heapify",
|
||||
"heappop",
|
||||
"heappush",
|
||||
"heappushpop",
|
||||
"heapreplace",
|
||||
"merge",
|
||||
"nlargest",
|
||||
"nsmallest",
|
||||
]
|
||||
|
||||
|
||||
@substitute_in_graph(heapq._heapify_max)
|
||||
def _heapify_max(heap: list[_T], /) -> None:
|
||||
return py_heapq._heapify_max(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq._heappop_max) # type: ignore[attr-defined]
|
||||
def _heappop_max(heap: list[_T]) -> _T:
|
||||
return py_heapq._heappop_max(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq._heapreplace_max) # type: ignore[attr-defined]
|
||||
def _heapreplace_max(heap: list[_T], item: _T) -> _T:
|
||||
return py_heapq._heapreplace_max(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heapify)
|
||||
def heapify(heap: list[_T], /) -> None:
|
||||
return py_heapq.heapify(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heappop)
|
||||
def heappop(heap: list[_T], /) -> _T:
|
||||
return py_heapq.heappop(heap)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heappush)
|
||||
def heappush(heap: list[_T], item: _T) -> None:
|
||||
return py_heapq.heappush(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heappushpop)
|
||||
def heappushpop(heap: list[_T], item: _T) -> _T:
|
||||
return py_heapq.heappushpop(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.heapreplace)
|
||||
def heapreplace(heap: list[_T], item: _T) -> _T:
|
||||
return py_heapq.heapreplace(heap, item)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.merge) # type: ignore[arg-type]
|
||||
def merge(*iterables, key=None, reverse=False): # type: ignore[no-untyped-def]
|
||||
return py_heapq.merge(*iterables, key=key, reverse=reverse)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.nlargest) # type: ignore[arg-type]
|
||||
def nlargest(n, iterable, key=None): # type: ignore[no-untyped-def]
|
||||
return py_heapq.nlargest(n, iterable, key=key)
|
||||
|
||||
|
||||
@substitute_in_graph(heapq.nsmallest) # type: ignore[arg-type]
|
||||
def nsmallest(n, iterable, key=None): # type: ignore[no-untyped-def]
|
||||
return py_heapq.nsmallest(n, iterable, key=key)
|
||||
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
Python polyfills for itertools
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import operator
|
||||
from collections.abc import Callable
|
||||
from typing import overload, TYPE_CHECKING, TypeAlias, TypeVar
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable, Iterator
|
||||
|
||||
|
||||
__all__ = [
|
||||
"accumulate",
|
||||
"chain",
|
||||
"chain_from_iterable",
|
||||
"compress",
|
||||
"cycle",
|
||||
"dropwhile",
|
||||
"filterfalse",
|
||||
"islice",
|
||||
"pairwise",
|
||||
"starmap",
|
||||
"takewhile",
|
||||
"tee",
|
||||
"zip_longest",
|
||||
]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_U = TypeVar("_U")
|
||||
_Predicate: TypeAlias = Callable[[_T], object]
|
||||
_T1 = TypeVar("_T1")
|
||||
_T2 = TypeVar("_T2")
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.chain
|
||||
@substitute_in_graph(itertools.chain, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def chain(*iterables: Iterable[_T]) -> Iterator[_T]:
|
||||
for iterable in iterables:
|
||||
yield from iterable
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.accumulate
|
||||
@substitute_in_graph(itertools.accumulate, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def accumulate(
|
||||
iterable: Iterable[_T],
|
||||
func: Callable[[_T, _T], _T] | None = None,
|
||||
*,
|
||||
initial: _T | None = None,
|
||||
) -> Iterator[_T]:
|
||||
# call iter outside of the generator to match cypthon behavior
|
||||
iterator = iter(iterable)
|
||||
if func is None:
|
||||
func = operator.add
|
||||
|
||||
def _accumulate(iterator: Iterator[_T]) -> Iterator[_T]:
|
||||
total = initial
|
||||
if total is None:
|
||||
try:
|
||||
total = next(iterator)
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
yield total
|
||||
for element in iterator:
|
||||
total = func(total, element)
|
||||
yield total
|
||||
|
||||
return _accumulate(iterator)
|
||||
|
||||
|
||||
@substitute_in_graph(itertools.chain.from_iterable) # type: ignore[arg-type]
|
||||
def chain_from_iterable(iterable: Iterable[Iterable[_T]], /) -> Iterator[_T]:
|
||||
# previous version of this code was:
|
||||
# return itertools.chain(*iterable)
|
||||
# If iterable is an infinite generator, this will lead to infinite recursion
|
||||
for it in iterable:
|
||||
yield from it
|
||||
|
||||
|
||||
chain.from_iterable = chain_from_iterable # type: ignore[attr-defined]
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.compress
|
||||
@substitute_in_graph(itertools.compress, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def compress(data: Iterable[_T], selectors: Iterable[_U], /) -> Iterator[_T]:
|
||||
return (datum for datum, selector in zip(data, selectors) if selector)
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.cycle
|
||||
@substitute_in_graph(itertools.cycle, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def cycle(iterable: Iterable[_T]) -> Iterator[_T]:
|
||||
iterator = iter(iterable)
|
||||
|
||||
def _cycle(iterator: Iterator[_T]) -> Iterator[_T]:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
saved = []
|
||||
for element in iterable:
|
||||
yield element
|
||||
saved.append(element)
|
||||
|
||||
while saved:
|
||||
for element in saved:
|
||||
yield element
|
||||
|
||||
return _cycle(iterator)
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.dropwhile
|
||||
@substitute_in_graph(itertools.dropwhile, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def dropwhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
|
||||
# dropwhile(lambda x: x < 5, [1, 4, 6, 3, 8]) -> 6 3 8
|
||||
if not callable(predicate):
|
||||
raise TypeError(f"'{type(predicate).__name__}' object is not callable")
|
||||
|
||||
iterator = iter(iterable)
|
||||
for x in iterator:
|
||||
if not predicate(x):
|
||||
yield x
|
||||
break
|
||||
|
||||
yield from iterator
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.takewhile
|
||||
@substitute_in_graph(itertools.takewhile, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def takewhile(predicate: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
|
||||
# takewhile(lambda x: x<5, [1,4,6,3,8]) → 1 4
|
||||
if not callable(predicate):
|
||||
raise TypeError(f"'{type(predicate).__name__}' object is not callable")
|
||||
|
||||
for x in iterable:
|
||||
if not predicate(x):
|
||||
break
|
||||
yield x
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[], _U],
|
||||
iterable: Iterable[tuple[()]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[_T], _U],
|
||||
iterable: Iterable[tuple[_T]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[_T, _T1], _U],
|
||||
iterable: Iterable[tuple[_T, _T1]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def starmap(
|
||||
function: Callable[[_T, _T1, _T2], _U],
|
||||
iterable: Iterable[tuple[_T, _T1, _T2]],
|
||||
/,
|
||||
) -> itertools.starmap[_U]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.starmap
|
||||
@substitute_in_graph(itertools.starmap, is_embedded_type=True) # type: ignore[arg-type]
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def starmap(function: Callable[..., _T], iterable: Iterable, /) -> Iterable[_T]:
|
||||
# starmap(pow, [(2,5), (3,2), (10,3)]) → 32 9 1000
|
||||
if not callable(function):
|
||||
raise TypeError(f"'{type(function).__name__}' object is not callable")
|
||||
|
||||
for args in iterable:
|
||||
yield function(*args)
|
||||
|
||||
|
||||
@substitute_in_graph(itertools.filterfalse, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def filterfalse(function: _Predicate[_T], iterable: Iterable[_T], /) -> Iterator[_T]:
|
||||
it = iter(iterable)
|
||||
if function is None:
|
||||
return filter(operator.not_, it)
|
||||
else:
|
||||
return filter(lambda x: not function(x), it)
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.islice
|
||||
@substitute_in_graph(itertools.islice, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def islice(iterable: Iterable[_T], /, *args: int | None) -> Iterator[_T]:
|
||||
s = slice(*args)
|
||||
start = 0 if s.start is None else s.start
|
||||
stop = s.stop
|
||||
step = 1 if s.step is None else s.step
|
||||
if start < 0 or (stop is not None and stop < 0) or step <= 0:
|
||||
raise ValueError(
|
||||
"Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.",
|
||||
)
|
||||
|
||||
if stop is None:
|
||||
# TODO: use indices = itertools.count() and merge implementation with the else branch
|
||||
# when we support infinite iterators
|
||||
next_i = start
|
||||
for i, element in enumerate(iterable):
|
||||
if i == next_i:
|
||||
yield element
|
||||
next_i += step
|
||||
else:
|
||||
indices = range(max(start, stop))
|
||||
next_i = start
|
||||
for i, element in zip(indices, iterable):
|
||||
if i == next_i:
|
||||
yield element
|
||||
next_i += step
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.pairwise
|
||||
@substitute_in_graph(itertools.pairwise, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def pairwise(iterable: Iterable[_T], /) -> Iterator[tuple[_T, _T]]:
|
||||
a = None
|
||||
first = True
|
||||
for b in iterable:
|
||||
if first:
|
||||
first = False
|
||||
else:
|
||||
yield a, b # type: ignore[misc]
|
||||
a = b
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.tee
|
||||
@substitute_in_graph(itertools.tee)
|
||||
def tee(iterable: Iterable[_T], n: int = 2, /) -> tuple[Iterator[_T], ...]:
|
||||
iterator = iter(iterable)
|
||||
shared_link = [None, None]
|
||||
|
||||
def _tee(link) -> Iterator[_T]: # type: ignore[no-untyped-def]
|
||||
try:
|
||||
while True:
|
||||
if link[1] is None:
|
||||
link[0] = next(iterator)
|
||||
link[1] = [None, None]
|
||||
value, link = link
|
||||
yield value
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
return tuple(_tee(shared_link) for _ in range(n))
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T1],
|
||||
/,
|
||||
*,
|
||||
fillvalue: _U = ...,
|
||||
) -> Iterator[tuple[_T1]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T1],
|
||||
iter2: Iterable[_T2],
|
||||
/,
|
||||
) -> Iterator[tuple[_T1 | None, _T2 | None]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T1],
|
||||
iter2: Iterable[_T2],
|
||||
/,
|
||||
*,
|
||||
fillvalue: _U = ...,
|
||||
) -> Iterator[tuple[_T1 | _U, _T2 | _U]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T],
|
||||
iter2: Iterable[_T],
|
||||
iter3: Iterable[_T],
|
||||
/,
|
||||
*iterables: Iterable[_T],
|
||||
) -> Iterator[tuple[_T | None, ...]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def zip_longest(
|
||||
iter1: Iterable[_T],
|
||||
iter2: Iterable[_T],
|
||||
iter3: Iterable[_T],
|
||||
/,
|
||||
*iterables: Iterable[_T],
|
||||
fillvalue: _U = ...,
|
||||
) -> Iterator[tuple[_T | _U, ...]]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/itertools.html#itertools.zip_longest
|
||||
@substitute_in_graph(itertools.zip_longest, is_embedded_type=True) # type: ignore[arg-type,misc]
|
||||
def zip_longest(
|
||||
*iterables: Iterable[_T],
|
||||
fillvalue: _U = None, # type: ignore[assignment]
|
||||
) -> Iterator[tuple[_T | _U, ...]]:
|
||||
# zip_longest('ABCD', 'xy', fillvalue='-') -> Ax By C- D-
|
||||
|
||||
iterators = list(map(iter, iterables))
|
||||
num_active = len(iterators)
|
||||
if not num_active:
|
||||
return
|
||||
|
||||
while True:
|
||||
values = []
|
||||
for i, iterator in enumerate(iterators):
|
||||
try:
|
||||
value = next(iterator)
|
||||
except StopIteration:
|
||||
num_active -= 1
|
||||
if not num_active:
|
||||
return
|
||||
iterators[i] = itertools.repeat(fillvalue) # type: ignore[arg-type]
|
||||
value = fillvalue # type: ignore[assignment]
|
||||
values.append(value)
|
||||
yield tuple(values)
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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
|
||||
|
||||
import torch.utils._pytree as python_pytree
|
||||
|
||||
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, ...] = (
|
||||
"_collections",
|
||||
"builtins",
|
||||
"copy",
|
||||
"functools",
|
||||
"itertools",
|
||||
"operator",
|
||||
"os",
|
||||
"struct",
|
||||
"sys",
|
||||
"fx",
|
||||
"tensor",
|
||||
"torch_c_nn",
|
||||
"traceback",
|
||||
)
|
||||
if python_pytree._cxx_pytree_dynamo_traceable:
|
||||
POLYFILLED_MODULE_NAMES += ("pytree",)
|
||||
|
||||
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))
|
||||
@@ -0,0 +1,119 @@
|
||||
"""
|
||||
Python polyfills for operator
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import Any, overload, TYPE_CHECKING, TypeVar
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
|
||||
# Most unary and binary operators are handled by BuiltinVariable (e.g., `pos`, `add`)
|
||||
__all__ = ["attrgetter", "itemgetter", "methodcaller", "countOf"]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_T1 = TypeVar("_T1")
|
||||
_T2 = TypeVar("_T2")
|
||||
_Ts = TypeVarTuple("_Ts")
|
||||
_U = TypeVar("_U")
|
||||
_U1 = TypeVar("_U1")
|
||||
_U2 = TypeVar("_U2")
|
||||
_Us = TypeVarTuple("_Us")
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def attrgetter(attr: str, /) -> Callable[[Any], _U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def attrgetter(
|
||||
attr1: str, attr2: str, /, *attrs: str
|
||||
) -> Callable[[Any], tuple[_U1, _U2, Unpack[_Us]]]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.attrgetter
|
||||
@substitute_in_graph(operator.attrgetter, is_embedded_type=True) # type: ignore[arg-type,misc]
|
||||
def attrgetter(*attrs: str) -> Callable[[Any], Any | tuple[Any, ...]]:
|
||||
if len(attrs) == 0:
|
||||
raise TypeError("attrgetter expected 1 argument, got 0")
|
||||
|
||||
if any(not isinstance(attr, str) for attr in attrs):
|
||||
raise TypeError("attribute name must be a string")
|
||||
|
||||
def resolve_attr(obj: Any, attr: str) -> Any:
|
||||
for name in attr.split("."):
|
||||
obj = getattr(obj, name)
|
||||
return obj
|
||||
|
||||
if len(attrs) == 1:
|
||||
attr = attrs[0]
|
||||
|
||||
def getter(obj: Any) -> Any:
|
||||
return resolve_attr(obj, attr)
|
||||
|
||||
else:
|
||||
|
||||
def getter(obj: Any) -> tuple[Any, ...]: # type: ignore[misc]
|
||||
return tuple(resolve_attr(obj, attr) for attr in attrs)
|
||||
|
||||
return getter
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def itemgetter(item: _T, /) -> Callable[[Any], _U]: ...
|
||||
|
||||
|
||||
@overload
|
||||
# pyrefly: ignore [inconsistent-overload]
|
||||
def itemgetter(
|
||||
item1: _T1, item2: _T2, /, *items: Unpack[_Ts]
|
||||
) -> Callable[[Any], tuple[_U1, _U2, Unpack[_Us]]]: ...
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.itemgetter
|
||||
@substitute_in_graph(operator.itemgetter, is_embedded_type=True) # type: ignore[arg-type,misc]
|
||||
def itemgetter(*items: Any) -> Callable[[Any], Any | tuple[Any, ...]]:
|
||||
if len(items) == 0:
|
||||
raise TypeError("itemgetter expected 1 argument, got 0")
|
||||
|
||||
if len(items) == 1:
|
||||
item = items[0]
|
||||
|
||||
def getter(obj: Any) -> Any:
|
||||
return obj[item]
|
||||
|
||||
else:
|
||||
|
||||
def getter(obj: Any) -> tuple[Any, ...]: # type: ignore[misc]
|
||||
return tuple(obj[item] for item in items)
|
||||
|
||||
return getter
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.methodcaller
|
||||
@substitute_in_graph(operator.methodcaller, is_embedded_type=True) # type: ignore[arg-type]
|
||||
def methodcaller(name: str, /, *args: Any, **kwargs: Any) -> Callable[[Any], Any]:
|
||||
if not isinstance(name, str):
|
||||
raise TypeError("method name must be a string")
|
||||
|
||||
def caller(obj: Any) -> Any:
|
||||
return getattr(obj, name)(*args, **kwargs)
|
||||
|
||||
return caller
|
||||
|
||||
|
||||
# Reference: https://docs.python.org/3/library/operator.html#operator.countOf
|
||||
@substitute_in_graph(operator.countOf, can_constant_fold_through=True) # type: ignore[arg-type,misc]
|
||||
def countOf(a: Iterable[_T], b: _T, /) -> int:
|
||||
return sum(it is b or it == b for it in a)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
Python polyfills for os
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import AnyStr
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = ["fspath"]
|
||||
|
||||
|
||||
# Copied from os.py in the standard library
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
@substitute_in_graph(os.fspath, can_constant_fold_through=True)
|
||||
def fspath(path: AnyStr | os.PathLike[AnyStr]) -> AnyStr:
|
||||
if isinstance(path, (str, bytes)):
|
||||
return path
|
||||
|
||||
path_type = type(path)
|
||||
try:
|
||||
path_repr = path_type.__fspath__(path) # type: ignore[arg-type]
|
||||
except AttributeError:
|
||||
if hasattr(path_type, "__fspath__"):
|
||||
raise
|
||||
raise TypeError(
|
||||
f"expected str, bytes or os.PathLike object, not {path_type.__name__}",
|
||||
) from None
|
||||
if isinstance(path_repr, (str, bytes)):
|
||||
return path_repr # type: ignore[return-value]
|
||||
raise TypeError(
|
||||
f"expected {path_type.__name__}.__fspath__() to return str or bytes, "
|
||||
f"not {type(path_repr).__name__}",
|
||||
)
|
||||
@@ -0,0 +1,759 @@
|
||||
# Owner(s): ["module: pytree"]
|
||||
|
||||
"""
|
||||
Python polyfills for torch.utils.pytree
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TYPE_CHECKING, TypeVar
|
||||
|
||||
import optree
|
||||
import optree._C
|
||||
import optree.utils
|
||||
from optree import (
|
||||
is_namedtuple,
|
||||
is_namedtuple_class,
|
||||
is_namedtuple_instance,
|
||||
is_structseq,
|
||||
is_structseq_class,
|
||||
is_structseq_instance,
|
||||
namedtuple_fields,
|
||||
structseq_fields,
|
||||
)
|
||||
|
||||
import torch.utils._cxx_pytree as cxx_pytree # noqa: F401
|
||||
import torch.utils._pytree as python_pytree
|
||||
from torch.utils._pytree import BUILTIN_TYPES, STANDARD_DICT_TYPES
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import builtins
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from typing_extensions import Self, TypeIs
|
||||
|
||||
from torch.utils._cxx_pytree import PyTree
|
||||
|
||||
|
||||
__all__ = [
|
||||
"is_namedtuple",
|
||||
"is_namedtuple_class",
|
||||
"is_namedtuple_instance",
|
||||
"is_structseq",
|
||||
"is_structseq_class",
|
||||
"is_structseq_instance",
|
||||
"namedtuple_fields",
|
||||
"structseq_fields",
|
||||
"treespec_leaf",
|
||||
"treespec_tuple",
|
||||
"treespec_dict",
|
||||
"tree_is_leaf",
|
||||
"tree_iter",
|
||||
"tree_leaves",
|
||||
"tree_flatten",
|
||||
"tree_flatten_with_path",
|
||||
"tree_structure",
|
||||
"tree_unflatten",
|
||||
]
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_KT = TypeVar("_KT")
|
||||
_VT = TypeVar("_VT")
|
||||
|
||||
|
||||
@substitute_in_graph(
|
||||
optree._C.is_dict_insertion_ordered,
|
||||
can_constant_fold_through=True,
|
||||
)
|
||||
def _(*args: Any, **kwargs: Any) -> bool:
|
||||
# In namespace 'torch', the dictionary is always traversed in insertion order.
|
||||
# This function returns True.
|
||||
raise ValueError(
|
||||
"Should not be called directly "
|
||||
"because the original function will be called in the constant fold path."
|
||||
)
|
||||
|
||||
|
||||
__name = ""
|
||||
for __name, __func in (
|
||||
("is_namedtuple", is_namedtuple),
|
||||
("is_namedtuple_class", is_namedtuple_class),
|
||||
("is_namedtuple_instance", is_namedtuple_instance),
|
||||
("is_structseq", is_structseq),
|
||||
("is_structseq_class", is_structseq_class),
|
||||
("is_structseq_instance", is_structseq_instance),
|
||||
("namedtuple_fields", namedtuple_fields),
|
||||
("structseq_fields", structseq_fields),
|
||||
):
|
||||
globals()[__name] = substitute_in_graph(
|
||||
__func, # type: ignore[arg-type]
|
||||
can_constant_fold_through=True,
|
||||
)(__func.__python_implementation__) # type: ignore[attr-defined]
|
||||
del __func
|
||||
del __name
|
||||
|
||||
|
||||
@substitute_in_graph(optree.tree_is_leaf, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def tree_is_leaf(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> bool:
|
||||
if (tree is None and none_is_leaf) or (is_leaf is not None and is_leaf(tree)):
|
||||
return True
|
||||
if optree.register_pytree_node.get(type(tree), namespace=namespace) is None:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@substitute_in_graph(optree.tree_iter, can_constant_fold_through=False) # type: ignore[arg-type]
|
||||
def tree_iter(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> Iterable[Any]:
|
||||
stack = [tree]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
if tree_is_leaf(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
):
|
||||
yield node
|
||||
continue
|
||||
|
||||
children, *_ = optree.tree_flatten_one_level(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
stack.extend(reversed(children))
|
||||
|
||||
|
||||
@substitute_in_graph(optree.tree_leaves, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def tree_leaves(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> list[Any]:
|
||||
return list(
|
||||
tree_iter(
|
||||
tree,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _Asterisk(str):
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(cls) -> Self:
|
||||
return super().__new__(cls, "*")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "*" # no quotes
|
||||
|
||||
|
||||
_asterisk = _Asterisk()
|
||||
del _Asterisk
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PyTreeSpec:
|
||||
"""Analog for :class:`optree.PyTreeSpec` in Python."""
|
||||
|
||||
_children: tuple[PyTreeSpec, ...]
|
||||
_type: builtins.type | None
|
||||
_metadata: Any
|
||||
_entries: tuple[Any, ...]
|
||||
_unflatten_func: Callable[[Any | None, Iterable[PyTree]], PyTree] | None
|
||||
none_is_leaf: bool
|
||||
namespace: str
|
||||
|
||||
num_nodes: int = field(init=False)
|
||||
num_leaves: int = field(init=False)
|
||||
num_children: int = field(init=False)
|
||||
|
||||
def __post_init__(self, /) -> None:
|
||||
if self._type is None:
|
||||
assert len(self._children) == 0
|
||||
assert self._metadata is None
|
||||
assert self._entries == ()
|
||||
assert self._unflatten_func is None
|
||||
num_nodes = 1
|
||||
num_leaves = 1
|
||||
num_children = 0
|
||||
else:
|
||||
assert callable(self._unflatten_func)
|
||||
num_nodes = 1
|
||||
num_leaves = 0
|
||||
for child in self._children:
|
||||
num_nodes += child.num_nodes
|
||||
num_leaves += child.num_leaves
|
||||
num_children = len(self._children)
|
||||
|
||||
object.__setattr__(self, "num_nodes", num_nodes)
|
||||
object.__setattr__(self, "num_leaves", num_leaves)
|
||||
object.__setattr__(self, "num_children", num_children)
|
||||
|
||||
def __repr__(self, /) -> str:
|
||||
def helper(treespec: PyTreeSpec) -> str:
|
||||
if treespec.is_leaf():
|
||||
assert treespec.type is None
|
||||
return _asterisk
|
||||
|
||||
assert treespec.type is not None
|
||||
assert callable(treespec._unflatten_func)
|
||||
children_representations = [
|
||||
helper(subspec) for subspec in treespec._children
|
||||
]
|
||||
if (
|
||||
treespec.type in BUILTIN_TYPES
|
||||
or (treespec.type is type(None) and not self.none_is_leaf)
|
||||
or optree.is_namedtuple_class(treespec.type)
|
||||
or optree.is_structseq_class(treespec.type)
|
||||
):
|
||||
return treespec._unflatten_func(
|
||||
treespec._metadata,
|
||||
children_representations,
|
||||
)
|
||||
return (
|
||||
f"CustomTreeNode({treespec.type.__name__}[{treespec._metadata!r}], "
|
||||
f"[{', '.join(children_representations)}])"
|
||||
)
|
||||
|
||||
inner = [
|
||||
str(helper(self)),
|
||||
*(["NoneIsLeaf"] if self.none_is_leaf else []),
|
||||
f"namespace={self.namespace!r}",
|
||||
]
|
||||
return f"PyTreeSpec({', '.join(inner)})"
|
||||
|
||||
def __len__(self, /) -> int:
|
||||
return self.num_leaves
|
||||
|
||||
@property
|
||||
def type(self, /) -> builtins.type | None:
|
||||
return self._type
|
||||
|
||||
def is_leaf(self, /) -> bool:
|
||||
return self.num_nodes == 1 and self.num_leaves == 1
|
||||
|
||||
def paths(self, /) -> list[tuple[Any, ...]]:
|
||||
def helper(treespec: PyTreeSpec, path_prefix: list[Any]) -> None:
|
||||
if treespec.is_leaf():
|
||||
paths.append(path_prefix)
|
||||
return
|
||||
|
||||
for entry, subspec in zip(
|
||||
treespec._entries,
|
||||
treespec._children,
|
||||
strict=True,
|
||||
):
|
||||
helper(subspec, path_prefix + [entry])
|
||||
|
||||
paths: list[list[Any]] = []
|
||||
helper(self, [])
|
||||
return [tuple(path) for path in paths]
|
||||
|
||||
def accessors(self, /) -> list[optree.PyTreeAccessor]:
|
||||
def helper(
|
||||
treespec: PyTreeSpec,
|
||||
entry_path_prefix: list[optree.PyTreeEntry],
|
||||
) -> None:
|
||||
if treespec.is_leaf():
|
||||
entry_paths.append(entry_path_prefix)
|
||||
return
|
||||
|
||||
node_type = treespec.type
|
||||
assert node_type is not None
|
||||
handler = optree.register_pytree_node.get(
|
||||
node_type, namespace=treespec.namespace
|
||||
)
|
||||
assert handler is not None
|
||||
kind: optree.PyTreeKind = handler.kind
|
||||
path_entry_type: type[optree.PyTreeEntry] = handler.path_entry_type
|
||||
|
||||
for entry, subspec in zip(
|
||||
treespec._entries,
|
||||
treespec._children,
|
||||
strict=True,
|
||||
):
|
||||
helper(
|
||||
subspec,
|
||||
entry_path_prefix + [path_entry_type(entry, node_type, kind)],
|
||||
)
|
||||
|
||||
entry_paths: list[list[optree.PyTreeEntry]] = []
|
||||
helper(self, [])
|
||||
return [optree.PyTreeAccessor(path) for path in entry_paths]
|
||||
|
||||
def children(self, /) -> list[PyTreeSpec]:
|
||||
return list(self._children)
|
||||
|
||||
def child(self, index: int, /) -> PyTreeSpec:
|
||||
return self._children[index]
|
||||
|
||||
def entries(self, /) -> list[Any]:
|
||||
return list(self._entries)
|
||||
|
||||
def entry(self, index: int, /) -> Any:
|
||||
return self._entries[index]
|
||||
|
||||
def flatten_up_to(self, tree: PyTree, /) -> list[PyTree]:
|
||||
def helper(
|
||||
treespec: PyTreeSpec,
|
||||
node: PyTree,
|
||||
subtrees: list[PyTree],
|
||||
) -> None:
|
||||
if treespec.is_leaf():
|
||||
subtrees.append(node)
|
||||
return
|
||||
|
||||
node_type = type(node)
|
||||
if treespec.type not in BUILTIN_TYPES:
|
||||
# Always require custom node types to match exactly
|
||||
if node_type != treespec.type:
|
||||
raise ValueError(
|
||||
f"Type mismatch; "
|
||||
f"expected {treespec.type!r}, but got {node_type!r}.",
|
||||
)
|
||||
|
||||
children, metadata, *_ = optree.tree_flatten_one_level(
|
||||
node,
|
||||
none_is_leaf=self.none_is_leaf,
|
||||
namespace=self.namespace,
|
||||
)
|
||||
if len(children) != treespec.num_children:
|
||||
raise ValueError(
|
||||
f"Node arity mismatch; "
|
||||
f"expected {treespec.num_children}, but got {len(children)}.",
|
||||
)
|
||||
if metadata != treespec._metadata:
|
||||
raise ValueError(
|
||||
f"Node context mismatch for custom node type {treespec.type!r}.",
|
||||
)
|
||||
else:
|
||||
# For builtin dictionary types, we allow some flexibility
|
||||
# Otherwise, we require exact matches
|
||||
both_standard_dict = (
|
||||
treespec.type in STANDARD_DICT_TYPES
|
||||
and node_type in STANDARD_DICT_TYPES
|
||||
)
|
||||
if not both_standard_dict and node_type != treespec.type:
|
||||
raise ValueError(
|
||||
f"Node type mismatch; "
|
||||
f"expected {treespec.type!r}, but got {node_type!r}.",
|
||||
)
|
||||
if len(node) != treespec.num_children:
|
||||
raise ValueError(
|
||||
f"Node arity mismatch; "
|
||||
f"expected {treespec.num_children}, but got {len(node)}.",
|
||||
)
|
||||
|
||||
if both_standard_dict:
|
||||
# dictionary types are compatible with each other
|
||||
expected_keys = treespec.entries()
|
||||
got_key_set = set(node)
|
||||
expected_key_set = set(expected_keys)
|
||||
if got_key_set != expected_key_set:
|
||||
missing_keys = expected_key_set.difference(got_key_set)
|
||||
extra_keys = got_key_set.difference(expected_key_set)
|
||||
message = ""
|
||||
if missing_keys:
|
||||
message += f"; missing key(s): {missing_keys}"
|
||||
if extra_keys:
|
||||
message += f"; extra key(s): {extra_keys}"
|
||||
raise ValueError(f"Node keys mismatch{message}.")
|
||||
children = [node[key] for key in expected_keys]
|
||||
else:
|
||||
# node_type is treespec.type
|
||||
children, metadata, *_ = optree.tree_flatten_one_level(
|
||||
node,
|
||||
none_is_leaf=self.none_is_leaf,
|
||||
namespace=self.namespace,
|
||||
)
|
||||
if (
|
||||
node_type is not deque # ignore mismatch of `maxlen` for deque
|
||||
) and metadata != treespec._metadata:
|
||||
raise ValueError(
|
||||
f"Node metadata mismatch for node type {treespec.type!r}; "
|
||||
f"expected {treespec._metadata!r}, but got {metadata!r}.", # namedtuple type mismatch
|
||||
)
|
||||
|
||||
for subtree, subspec in zip(children, treespec._children, strict=True):
|
||||
helper(subspec, subtree, subtrees)
|
||||
|
||||
subtrees: list[PyTree] = []
|
||||
helper(self, tree, subtrees)
|
||||
return subtrees
|
||||
|
||||
def unflatten(self, leaves: Iterable[Any], /) -> PyTree:
|
||||
if not isinstance(leaves, (list, tuple)):
|
||||
leaves = list(leaves)
|
||||
if len(leaves) != self.num_leaves:
|
||||
raise ValueError(
|
||||
f"treespec.unflatten(leaves): `leaves` has length {len(leaves)} "
|
||||
f"but the spec refers to a pytree that holds {self.num_leaves} "
|
||||
f"items ({self}).",
|
||||
)
|
||||
if self.is_leaf():
|
||||
return leaves[0]
|
||||
|
||||
# Recursively unflatten the children
|
||||
start = 0
|
||||
end = 0
|
||||
subtrees = []
|
||||
for subspec in self._children:
|
||||
end += subspec.num_leaves
|
||||
subtrees.append(subspec.unflatten(leaves[start:end]))
|
||||
start = end
|
||||
|
||||
assert callable(self._unflatten_func)
|
||||
return self._unflatten_func(self._metadata, subtrees)
|
||||
|
||||
|
||||
def _is_pytreespec_instance(obj: Any, /) -> TypeIs[PyTreeSpec | python_pytree.TreeSpec]:
|
||||
return isinstance(obj, (PyTreeSpec, python_pytree.TreeSpec))
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.treespec_leaf,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def treespec_leaf(
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "", # unused
|
||||
) -> PyTreeSpec:
|
||||
return PyTreeSpec(
|
||||
(),
|
||||
None,
|
||||
None,
|
||||
(),
|
||||
None,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace="",
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.treespec_tuple,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def treespec_tuple(
|
||||
iterable: Iterable[PyTreeSpec] = (),
|
||||
/,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> PyTreeSpec:
|
||||
children = tuple(iterable)
|
||||
if any(not _is_pytreespec_instance(child) for child in children):
|
||||
raise ValueError(f"Expected a tuple of PyTreeSpecs, got: {children!r}.")
|
||||
if any(child.none_is_leaf != none_is_leaf for child in children):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `none_is_leaf` value "
|
||||
f"as the parent; expected {none_is_leaf}, got: {children!r}.",
|
||||
)
|
||||
if any(child.namespace not in (namespace, "") for child in children):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `namespace` value "
|
||||
f"as the parent; expected {namespace!r}, got: {children!r}.",
|
||||
)
|
||||
handler = optree.register_pytree_node.get(tuple, namespace=namespace)
|
||||
assert handler is not None
|
||||
return PyTreeSpec(
|
||||
tuple(children),
|
||||
tuple,
|
||||
None,
|
||||
tuple(range(len(children))),
|
||||
handler.unflatten_func,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.treespec_dict,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def treespec_dict(
|
||||
mapping: Mapping[Any, PyTreeSpec] | Iterable[tuple[Any, PyTreeSpec]] = (),
|
||||
/,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
**kwargs: PyTreeSpec,
|
||||
) -> PyTreeSpec:
|
||||
dct = dict(mapping, **kwargs)
|
||||
if any(not _is_pytreespec_instance(child) for child in dct.values()):
|
||||
raise ValueError(f"Expected a dictionary of TreeSpecs, got: {dct!r}.")
|
||||
if any(child.none_is_leaf != none_is_leaf for child in dct.values()):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `none_is_leaf` value "
|
||||
f"as the parent; expected {none_is_leaf}, got: {dct!r}.",
|
||||
)
|
||||
if any(child.namespace not in (namespace, "") for child in dct.values()):
|
||||
raise ValueError(
|
||||
"All children PyTreeSpecs must have the same `namespace` value "
|
||||
f"as the parent; expected {namespace!r}, got: {dct!r}.",
|
||||
)
|
||||
|
||||
(
|
||||
children,
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func,
|
||||
) = optree.tree_flatten_one_level( # type: ignore[assignment,var-annotated]
|
||||
dct, # type: ignore[arg-type]
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
return PyTreeSpec(
|
||||
tuple(children), # type: ignore[arg-type]
|
||||
dict,
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func, # type: ignore[arg-type]
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_flatten,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_flatten(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[Any], PyTreeSpec]:
|
||||
def helper(node: PyTree, leaves: list[Any]) -> PyTreeSpec:
|
||||
if tree_is_leaf(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
):
|
||||
leaves.append(node)
|
||||
return PyTreeSpec(
|
||||
(),
|
||||
None,
|
||||
None,
|
||||
(),
|
||||
None,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
(
|
||||
children,
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func,
|
||||
) = optree.tree_flatten_one_level(
|
||||
node,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
# Recursively flatten the children
|
||||
subspecs = tuple(helper(child, leaves) for child in children)
|
||||
return PyTreeSpec(
|
||||
subspecs,
|
||||
type(node),
|
||||
metadata,
|
||||
entries,
|
||||
unflatten_func, # type: ignore[arg-type]
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
) # type: ignore[arg-type]
|
||||
|
||||
leaves: list[Any] = []
|
||||
treespec = helper(tree, leaves)
|
||||
return leaves, treespec
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree._C.flatten,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def _C_flatten(
|
||||
tree: PyTree,
|
||||
/,
|
||||
leaf_predicate: Callable[[PyTree], bool] | None = None,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[Any], PyTreeSpec]:
|
||||
return tree_flatten( # type: ignore[return-value]
|
||||
tree,
|
||||
is_leaf=leaf_predicate,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_flatten_with_path,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_flatten_with_path(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[tuple[Any, ...]], list[Any], PyTreeSpec]:
|
||||
leaves, treespec = tree_flatten(
|
||||
tree,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
return treespec.paths(), leaves, treespec # type: ignore[return-value]
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree._C.flatten_with_path,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def _C_flatten_with_path(
|
||||
tree: PyTree,
|
||||
/,
|
||||
leaf_predicate: Callable[[PyTree], bool] | None = None,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> tuple[list[tuple[Any, ...]], list[Any], PyTreeSpec]:
|
||||
return tree_flatten_with_path( # type: ignore[return-value]
|
||||
tree,
|
||||
is_leaf=leaf_predicate,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_structure,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_structure(
|
||||
tree: PyTree,
|
||||
/,
|
||||
is_leaf: Callable[[PyTree], bool] | None = None,
|
||||
*,
|
||||
none_is_leaf: bool = False,
|
||||
namespace: str = "",
|
||||
) -> PyTreeSpec:
|
||||
return tree_flatten( # type: ignore[return-value]
|
||||
tree,
|
||||
is_leaf=is_leaf,
|
||||
none_is_leaf=none_is_leaf,
|
||||
namespace=namespace,
|
||||
)[1]
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
optree.tree_unflatten,
|
||||
# We need to disable constant folding here because we want the function to reference the
|
||||
# PyTreeSpec class defined above, not the one in the C++ module.
|
||||
can_constant_fold_through=False,
|
||||
)
|
||||
def tree_unflatten(treespec: PyTreeSpec, leaves: Iterable[Any]) -> PyTree:
|
||||
if not _is_pytreespec_instance(treespec):
|
||||
raise TypeError(
|
||||
f"Expected `treespec` to be an instance of "
|
||||
f"PyTreeSpec but got item of type {type(treespec)}."
|
||||
)
|
||||
return treespec.unflatten(leaves)
|
||||
|
||||
|
||||
_none_registration = optree.register_pytree_node.get(type(None))
|
||||
assert _none_registration is not None
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
_none_registration.unflatten_func,
|
||||
can_constant_fold_through=True,
|
||||
skip_signature_check=True,
|
||||
)
|
||||
def none_unflatten(_: None, children: Iterable[_T], /) -> None:
|
||||
if len(list(children)) != 0:
|
||||
raise ValueError("Expected no children.")
|
||||
return None
|
||||
|
||||
|
||||
with optree.dict_insertion_ordered(False, namespace="torch"):
|
||||
_dict_registration = optree.register_pytree_node.get(dict)
|
||||
assert _dict_registration is not None
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
_dict_registration.flatten_func,
|
||||
can_constant_fold_through=True,
|
||||
skip_signature_check=True,
|
||||
)
|
||||
def dict_flatten(
|
||||
dct: dict[_KT, _VT], /
|
||||
) -> tuple[list[_VT], tuple[list[_KT], list[_KT]], tuple[_KT, ...]]:
|
||||
sorted_keys = optree.utils.total_order_sorted(dct)
|
||||
values = [dct[key] for key in sorted_keys]
|
||||
original_keys = list(dct)
|
||||
return values, (original_keys, sorted_keys), tuple(sorted_keys)
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
_dict_registration.unflatten_func,
|
||||
can_constant_fold_through=True,
|
||||
skip_signature_check=True,
|
||||
)
|
||||
def dict_unflatten(
|
||||
metadata: tuple[list[_KT], list[_KT]],
|
||||
values: Iterable[_VT],
|
||||
/,
|
||||
) -> dict[_KT, _VT]:
|
||||
original_keys, sorted_keys = metadata
|
||||
d = dict.fromkeys(original_keys)
|
||||
d.update(zip(sorted_keys, values, strict=True))
|
||||
return d # type: ignore[return-value]
|
||||
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
Python polyfills for struct
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from typing import Any
|
||||
from typing_extensions import Buffer
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = [
|
||||
"pack",
|
||||
"unpack",
|
||||
]
|
||||
|
||||
|
||||
@substitute_in_graph(struct.pack, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def pack(fmt: bytes | str, /, *v: Any) -> bytes:
|
||||
return struct.pack(fmt, *v)
|
||||
|
||||
|
||||
@substitute_in_graph(struct.unpack, can_constant_fold_through=True) # type: ignore[arg-type]
|
||||
def unpack(format: bytes | str, buffer: Buffer, /) -> tuple[Any, ...]:
|
||||
return struct.unpack(format, buffer)
|
||||
@@ -0,0 +1,39 @@
|
||||
"""
|
||||
Python polyfills for sys
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = [
|
||||
"intern",
|
||||
"getrecursionlimit",
|
||||
]
|
||||
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
@substitute_in_graph(sys.intern, can_constant_fold_through=True)
|
||||
def intern(string: str, /) -> str:
|
||||
return string
|
||||
|
||||
|
||||
@substitute_in_graph(sys.getrecursionlimit, can_constant_fold_through=True)
|
||||
def getrecursionlimit() -> int:
|
||||
return sys.getrecursionlimit()
|
||||
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
|
||||
@substitute_in_graph(sys.get_int_max_str_digits, can_constant_fold_through=True)
|
||||
def get_int_max_str_digits() -> int:
|
||||
return sys.get_int_max_str_digits()
|
||||
|
||||
@substitute_in_graph(sys.set_int_max_str_digits, can_constant_fold_through=True)
|
||||
def set_int_max_str_digits(maxdigits: int) -> None:
|
||||
sys.set_int_max_str_digits(maxdigits)
|
||||
|
||||
__all__ += ["get_int_max_str_digits", "set_int_max_str_digits"]
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
@substitute_in_graph( # type: ignore[arg-type]
|
||||
torch.Tensor._make_subclass
|
||||
)
|
||||
def make_subclass(
|
||||
cls: type[Any], data: torch.Tensor, requires_grad: bool = False, **kwargs: Any
|
||||
) -> Any:
|
||||
with torch._C.DisableTorchFunctionSubclass():
|
||||
# This is a rough approximation of `THPVariable_make_subclass`. It should
|
||||
# suffice for most of Dynamo tracing purposes.
|
||||
# https://github.com/pytorch/pytorch/blob/ccfde4dadfa3c342076a1ee387017f84dd4ad2f7/torch/csrc/autograd/python_variable.cpp#L597-L650
|
||||
assert len(kwargs) == 0, (
|
||||
"_make_subclass only supports requires_grad as keyword arg"
|
||||
)
|
||||
data = data.detach()
|
||||
|
||||
# Avoid unnecessary `requires_grad` mutation, which isn't supported in Dynamo.
|
||||
if data.requires_grad != requires_grad:
|
||||
data.requires_grad = requires_grad
|
||||
|
||||
# Dynamo can't yet handle upcasting to base tensor type via `as_subclass`.
|
||||
if cls is torch.Tensor:
|
||||
return torch.Tensor(data)
|
||||
|
||||
# Calling `as_subclass` because
|
||||
# 1. Dynamo knows how to handle it
|
||||
# 2. the C impls match at this point -- both `THPVariable_make_subclass` and
|
||||
# `THPVariable_as_subclass` calls `THPVariable_NewWithVar`.
|
||||
return data.as_subclass(cls)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"make_subclass",
|
||||
]
|
||||
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
Polyfills for torch._C._nn functions.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.overrides import _is_torch_function_mode_enabled, _pop_mode_temporarily
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
@substitute_in_graph(torch._C._nn._parse_to, skip_signature_check=True)
|
||||
def _parse_to_polyfill(
|
||||
*args: Any, **kwargs: Any
|
||||
) -> tuple[torch.device, torch.dtype, bool, torch.memory_format]:
|
||||
"""
|
||||
Polyfill for torch._C._nn._parse_to that parses arguments to nn.Module.to().
|
||||
|
||||
Signature mirrors torch._C._nn._parse_to which accepts:
|
||||
- to(device) - device as string or torch.device
|
||||
- to(dtype) - dtype as torch.dtype
|
||||
- to(tensor) - extracts device and dtype from tensor
|
||||
- to(device=..., dtype=..., non_blocking=..., memory_format=...)
|
||||
|
||||
Returns:
|
||||
tuple: (device, dtype, non_blocking, memory_format)
|
||||
"""
|
||||
# Check for __torch_function__ mode and dispatch using handle_torch_function pattern
|
||||
if _is_torch_function_mode_enabled():
|
||||
with _pop_mode_temporarily() as mode:
|
||||
result = mode.__torch_function__(
|
||||
torch._C._nn._parse_to, tuple(), args, kwargs or {}
|
||||
)
|
||||
if result is not NotImplemented:
|
||||
return result
|
||||
|
||||
# Default implementation
|
||||
device = None
|
||||
dtype = None
|
||||
non_blocking = False
|
||||
memory_format = None
|
||||
|
||||
# Handle positional arguments
|
||||
if len(args) == 1:
|
||||
arg = args[0]
|
||||
# Check if it's a tensor
|
||||
if isinstance(arg, torch.Tensor):
|
||||
device = arg.device
|
||||
dtype = arg.dtype
|
||||
# Check if it's a dtype
|
||||
elif isinstance(arg, torch.dtype):
|
||||
dtype = arg
|
||||
# Check if it's a device (string or torch.device)
|
||||
elif isinstance(arg, (str, torch.device)):
|
||||
device = torch.device(arg) if isinstance(arg, str) else arg
|
||||
else:
|
||||
raise TypeError(
|
||||
f"to() received an invalid combination of arguments. Got: {type(arg)}"
|
||||
)
|
||||
elif len(args) > 1:
|
||||
raise TypeError(
|
||||
f"to() received too many positional arguments. Got {len(args)}, expected at most 1"
|
||||
)
|
||||
|
||||
# Handle keyword arguments
|
||||
if "device" in kwargs:
|
||||
device_arg = kwargs["device"]
|
||||
if device_arg is not None:
|
||||
device = (
|
||||
torch.device(device_arg) if isinstance(device_arg, str) else device_arg
|
||||
)
|
||||
|
||||
if "dtype" in kwargs:
|
||||
dtype = kwargs["dtype"]
|
||||
|
||||
if "non_blocking" in kwargs:
|
||||
non_blocking = kwargs["non_blocking"]
|
||||
|
||||
if "memory_format" in kwargs:
|
||||
memory_format = kwargs["memory_format"]
|
||||
|
||||
# pyrefly: ignore[bad-return]
|
||||
return (device, dtype, non_blocking, memory_format)
|
||||
|
||||
|
||||
@substitute_in_graph(torch.__future__.get_swap_module_params_on_conversion)
|
||||
def get_swap_module_params_on_conversion_polyfill() -> bool:
|
||||
"""
|
||||
Polyfill for torch.__future__.get_swap_module_params_on_conversion.
|
||||
|
||||
Returns the actual value from the underlying global variable.
|
||||
"""
|
||||
# Access the module's global variable directly to avoid recursion
|
||||
import torch.__future__ as torch_future
|
||||
|
||||
return torch_future._swap_module_params_on_conversion
|
||||
|
||||
|
||||
@substitute_in_graph(torch._has_compatible_shallow_copy_type)
|
||||
def _has_compatible_shallow_copy_type_polyfill(
|
||||
input: torch.Tensor, from_: torch.Tensor
|
||||
) -> bool:
|
||||
"""
|
||||
Polyfill for torch._has_compatible_shallow_copy_type.
|
||||
|
||||
Checks if two tensors have compatible types for shallow copying.
|
||||
The C++ implementation checks if input's TensorImpl has compatible shallow copy type
|
||||
with from_'s key_set. We approximate this by checking if both tensors are the same type.
|
||||
"""
|
||||
# Check if both tensors are the same type (handles both regular tensors and subclasses)
|
||||
# This is more permissive than checking exact torch.Tensor type equality
|
||||
# but properly handles subclasses by allowing same-type shallow copies
|
||||
return type(input) is type(from_)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"_parse_to_polyfill",
|
||||
"get_swap_module_params_on_conversion_polyfill",
|
||||
"_has_compatible_shallow_copy_type_polyfill",
|
||||
]
|
||||
@@ -0,0 +1,39 @@
|
||||
"""
|
||||
Python polyfills for traceback
|
||||
"""
|
||||
|
||||
import traceback
|
||||
from traceback import StackSummary
|
||||
from types import TracebackType
|
||||
|
||||
from ..decorators import substitute_in_graph
|
||||
|
||||
|
||||
__all__ = ["extract_tb", "clear_frames"]
|
||||
|
||||
|
||||
@substitute_in_graph(traceback.extract_tb, can_constant_fold_through=True)
|
||||
def extract_tb(tb: TracebackType | None, limit: int | None = None) -> StackSummary:
|
||||
if tb is None:
|
||||
return traceback.StackSummary.from_list([])
|
||||
# pyrefly: ignore [implicit-any]
|
||||
frame_summary = []
|
||||
while tb is not None:
|
||||
if limit:
|
||||
if len(frame_summary) < limit:
|
||||
frame_summary.append(
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
tb.frame_summary
|
||||
)
|
||||
else:
|
||||
break
|
||||
else:
|
||||
frame_summary.append(tb.frame_summary) # pyrefly: ignore[missing-attribute]
|
||||
tb = tb.tb_next
|
||||
return traceback.StackSummary.from_list(frame_summary)
|
||||
|
||||
|
||||
@substitute_in_graph(traceback.clear_frames, can_constant_fold_through=True)
|
||||
def clear_frames(tb: TracebackType | None) -> None:
|
||||
# no-op
|
||||
return None
|
||||
@@ -0,0 +1,231 @@
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
from abc import abstractmethod
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
import torch
|
||||
from torch._dynamo.package import (
|
||||
_BackendId,
|
||||
_DynamoCacheEntry,
|
||||
DynamoCache,
|
||||
PrecompileCacheEntry,
|
||||
)
|
||||
|
||||
|
||||
"""
|
||||
Classes and implementations related to precompile
|
||||
"""
|
||||
|
||||
T = TypeVar("T")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BackendCacheArtifact(Generic[T]):
|
||||
"""
|
||||
Represents a single serializable backend artifact from a dynamo backend.
|
||||
Each BackendCacheArtifact has a key associated with it along with some
|
||||
serializable content.
|
||||
|
||||
Example implementation:
|
||||
|
||||
class MyPrecompileCacheArtifact(PrecompileCacheArtifact[MySerializableType]):
|
||||
my_field: int
|
||||
|
||||
def after_deserialization(self) -> MySerializableType:
|
||||
result = pickle.loads(self.content)
|
||||
# Do some extra work post deserialization
|
||||
result.my_post_deserialization_function(self.my_field)
|
||||
return result
|
||||
"""
|
||||
|
||||
key: str
|
||||
content: Any
|
||||
|
||||
@abstractmethod
|
||||
def after_deserialization(self) -> T:
|
||||
"""
|
||||
Code to be run after reading raw byte contents from disk.
|
||||
Generally converts self.content from raw bytes back into its original form.
|
||||
"""
|
||||
...
|
||||
|
||||
def edit_contents(self, edit_fn: Callable[..., Any]) -> None:
|
||||
"""
|
||||
Edit the contents of the artifact.
|
||||
"""
|
||||
self.content = edit_fn(self.content)
|
||||
|
||||
|
||||
class EagerCacheArtifact(BackendCacheArtifact[Any]):
|
||||
def after_deserialization(self) -> Any:
|
||||
return self.content
|
||||
|
||||
|
||||
class BypassDynamoCacheEntry(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class PrecompileContext:
|
||||
"""
|
||||
PrecompileContext is a special CacheArtifactManager for handling precompilation
|
||||
It uses the same interface as CacheArtifactManager, but handles deserialization differently: instead
|
||||
of placing each artifact into respective caches, it will stitch all the cache artifacts for a single key
|
||||
together and place it into a global Precompile Cache.
|
||||
|
||||
PrecompileContext has two main portions: dynamo_cache_entries and backend_cache_artifacts.
|
||||
When saving, PrecompileContext.serialize() will serialize all dynamo cache entries along with any PrecompileCacheArtifacts that
|
||||
are needed to save those dynamo cache entries.
|
||||
|
||||
The following artifact types are supported by PrecompileContext:
|
||||
- BundledAOTAutogradCacheArtifact
|
||||
|
||||
"""
|
||||
|
||||
# Protected by the compile_lock
|
||||
# _backend_artifacts_by_key organizes results by the key of each artifact.
|
||||
# Each object here must be serializable
|
||||
_backend_artifacts_by_key: dict[_BackendId, BackendCacheArtifact[Any]] = {}
|
||||
|
||||
# On call to `serialize()`, all cache artifacts in _dynamo_cache_entries are converted
|
||||
# into DynamoCacheArtifacts and added to _new_cache_artifacts for serialization
|
||||
_dynamo_cache_entries: dict[str, _DynamoCacheEntry] = {}
|
||||
|
||||
@classmethod
|
||||
def clear(cls) -> None:
|
||||
cls._backend_artifacts_by_key.clear()
|
||||
cls._dynamo_cache_entries.clear()
|
||||
|
||||
@classmethod
|
||||
def record_artifact(
|
||||
cls,
|
||||
artifact: BackendCacheArtifact[Any],
|
||||
) -> None:
|
||||
"""
|
||||
Records a backend artifact to be used with dynamo cache entries
|
||||
"""
|
||||
# Temporarily disable all dispatch modes (including FakeTensorMode) during
|
||||
# deepcopy to avoid issues with cloning fake tensors (e.g., device mesh
|
||||
# with meta tensors that fail when cloning due to device mismatches)
|
||||
from torch.utils._mode_utils import no_dispatch
|
||||
|
||||
with no_dispatch():
|
||||
cls._backend_artifacts_by_key[_BackendId(artifact.key)] = copy.deepcopy(
|
||||
artifact
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def record_dynamo_cache_entry(
|
||||
cls, cache_entry: _DynamoCacheEntry, key: str
|
||||
) -> None:
|
||||
cls._dynamo_cache_entries[key] = cache_entry
|
||||
|
||||
@classmethod
|
||||
def edit_artifact(cls, key: str, edit_fn: Callable[..., Any]) -> None:
|
||||
"""
|
||||
Edit the content of an existing artifact
|
||||
"""
|
||||
assert key in cls._backend_artifacts_by_key, f"Key {key} not found in artifacts"
|
||||
artifact = cls._backend_artifacts_by_key[_BackendId(key)]
|
||||
artifact.edit_contents(edit_fn)
|
||||
|
||||
@classmethod
|
||||
def serialize_artifact_by_key(cls, key: str) -> BackendCacheArtifact[Any] | None:
|
||||
"""
|
||||
Return the backend cache artifact with the associated key
|
||||
"""
|
||||
return cls._backend_artifacts_by_key.get(_BackendId(key), None)
|
||||
|
||||
@staticmethod
|
||||
def dump_debug_info(
|
||||
dynamo_entries: dict[str, _DynamoCacheEntry],
|
||||
backend_artifacts: dict[_BackendId, BackendCacheArtifact[Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Return a JSON serializable debug dump of all entries in the precompile context
|
||||
Called in serialize before serialization, and in populate_caches after deserialization
|
||||
"""
|
||||
# Print debug information
|
||||
debug_info: defaultdict[str, list[Any]] = defaultdict(list)
|
||||
for key, cache_entry in dynamo_entries.items():
|
||||
info = cache_entry.debug_info()
|
||||
info["key"] = key
|
||||
debug_info["dynamo"].append(info)
|
||||
|
||||
for artifact in backend_artifacts.values():
|
||||
debug_info["backends"].append(artifact.key)
|
||||
|
||||
return debug_info
|
||||
|
||||
@classmethod
|
||||
def save_to_dynamo_cache(cls) -> dict[str, Any]:
|
||||
precompile_cache_entries, debug_info = cls.create_cache_entries()
|
||||
for key, entry in precompile_cache_entries.items():
|
||||
DynamoCache.write(entry, key)
|
||||
return debug_info
|
||||
|
||||
@classmethod
|
||||
def create_cache_entries(
|
||||
cls,
|
||||
) -> tuple[dict[str, PrecompileCacheEntry], dict[str, Any]]:
|
||||
"""
|
||||
Grabs all the cache entries in the precompile context and
|
||||
stitches them together into full PrecompileCacheEntries.
|
||||
"""
|
||||
dynamo_entries = cls._dynamo_cache_entries
|
||||
backend_artifacts = cls._backend_artifacts_by_key
|
||||
|
||||
num_artifacts = len(dynamo_entries)
|
||||
|
||||
debug_info = PrecompileContext.dump_debug_info(
|
||||
dynamo_entries, backend_artifacts
|
||||
)
|
||||
debug_str = json.dumps(
|
||||
{
|
||||
"num_entries": num_artifacts,
|
||||
"artifacts": debug_info,
|
||||
},
|
||||
)
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "dynamo_cache_entries",
|
||||
"encoding": "json",
|
||||
},
|
||||
payload_fn=lambda: debug_str,
|
||||
expect_trace_id=False,
|
||||
)
|
||||
|
||||
precompile_cache_entries = {}
|
||||
|
||||
for key, cache_entry in dynamo_entries.items():
|
||||
try:
|
||||
result = PrecompileCacheEntry.from_cache_entry(
|
||||
cache_entry, backend_artifacts
|
||||
)
|
||||
if result is not None:
|
||||
precompile_cache_entries[key] = result
|
||||
except Exception as e:
|
||||
logger.warning("Failed to create cache entry %s", key, exc_info=True)
|
||||
|
||||
error = e
|
||||
data = json.dumps(
|
||||
{
|
||||
"key": key,
|
||||
"error": str(error),
|
||||
}
|
||||
)
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "dynamo_cache_exception",
|
||||
"encoding": "json",
|
||||
},
|
||||
payload_fn=lambda: data,
|
||||
)
|
||||
continue
|
||||
return precompile_cache_entries, debug_info
|
||||
@@ -0,0 +1,176 @@
|
||||
"""
|
||||
Dynamo profiling implementation.
|
||||
|
||||
This module provides profiling functionality for Dynamo, including:
|
||||
- ProfileMetrics: Class for collecting and aggregating performance metrics like
|
||||
execution time, operator counts, and fusion statistics
|
||||
- ProfileResult: Class for analyzing and reporting profiling results
|
||||
- Utilities for tracking missed/uncaptured operations
|
||||
- Functions for instrumenting FX graphs with profiling capabilities
|
||||
|
||||
The profiler helps measure and optimize the performance of Dynamo-compiled code
|
||||
by tracking both captured and total operations, timing, and graph statistics.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import os
|
||||
from typing import Any
|
||||
from typing_extensions import Self
|
||||
|
||||
import torch
|
||||
|
||||
from .utils import print_once
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ProfileMetrics:
|
||||
microseconds: float = 0.0
|
||||
operators: int = 0
|
||||
fusions: int = 0
|
||||
graphs: int = 0
|
||||
|
||||
def __iadd__(self, other: Self) -> Self:
|
||||
self.microseconds += other.microseconds
|
||||
self.operators += other.operators
|
||||
self.fusions += other.fusions
|
||||
return self
|
||||
|
||||
def __add__(self, other: ProfileMetrics) -> ProfileMetrics:
|
||||
assert isinstance(other, ProfileMetrics)
|
||||
return ProfileMetrics(
|
||||
self.microseconds + other.microseconds,
|
||||
self.operators + other.operators,
|
||||
self.fusions + other.fusions,
|
||||
)
|
||||
|
||||
def __truediv__(self, other: Any) -> ProfileMetrics:
|
||||
if isinstance(other, int):
|
||||
other = ProfileMetrics(other, other, other)
|
||||
return ProfileMetrics(
|
||||
self.microseconds / max(1, other.microseconds),
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
self.operators / max(1, other.operators),
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
self.fusions / max(1, other.fusions),
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.operators:4.0%} ops {self.microseconds:4.0%} time"
|
||||
|
||||
def tocsv(self) -> list[float]:
|
||||
return [self.operators, self.microseconds]
|
||||
|
||||
|
||||
class ProfileResult:
|
||||
def __init__(
|
||||
self, captured: ProfileMetrics, total: ProfileMetrics, unique_graphs: int
|
||||
) -> None:
|
||||
self.captured: ProfileMetrics = captured or ProfileMetrics()
|
||||
self.total: ProfileMetrics = total or ProfileMetrics()
|
||||
self.unique_graphs: int = unique_graphs
|
||||
|
||||
def __iadd__(self, other: Self) -> Self:
|
||||
self.captured += other.captured
|
||||
self.total += other.total
|
||||
self.unique_graphs += other.unique_graphs
|
||||
return self
|
||||
|
||||
def percent(self) -> ProfileMetrics:
|
||||
return self.captured / self.total
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
f"{self.unique_graphs:2} graphs {self.captured.graphs:2} graph calls "
|
||||
f"{self.captured.operators:4}/{self.total.operators:4} = "
|
||||
+ str(self.percent())
|
||||
)
|
||||
|
||||
def tocsv(self) -> list[Any]:
|
||||
return [
|
||||
self.unique_graphs,
|
||||
self.captured.graphs,
|
||||
self.captured.operators,
|
||||
self.total.operators,
|
||||
] + self.percent().tocsv()
|
||||
|
||||
|
||||
def should_print_missing() -> bool:
|
||||
return os.environ.get("TORCHDYNAMO_PRINT_MISSING") == "1"
|
||||
|
||||
|
||||
def print_missing(stack: list[str]) -> None:
|
||||
if any("/torch/autograd/profiler.py" in x for x in stack):
|
||||
return
|
||||
stack = [
|
||||
x for x in stack if ("<built-in" not in x and "site-packages/torch/" not in x)
|
||||
]
|
||||
print_once("MISSING", " >> ".join(stack[-3:]))
|
||||
|
||||
|
||||
class Profiler:
|
||||
unique_graphs: int = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.prof = torch.profiler.profile(
|
||||
activities=[torch.profiler.ProfilerActivity.CPU],
|
||||
with_stack=should_print_missing(),
|
||||
)
|
||||
|
||||
def results(self) -> ProfileResult:
|
||||
captured_regions = 0
|
||||
captured_ops = 0
|
||||
captured_microseconds = 0
|
||||
total_ops = 0
|
||||
total_microseconds = 0
|
||||
|
||||
last_op_end_time = -1
|
||||
captured_region_end_time = -1
|
||||
events = sorted(self.prof.events(), key=lambda x: x.time_range.start)
|
||||
for e in events:
|
||||
if e.name == "TORCHDYNAMO":
|
||||
captured_region_end_time = e.time_range.end
|
||||
captured_regions += 1
|
||||
# ignore `handle = torch.zeros(1)` in record_function.__init__()
|
||||
total_ops -= 1
|
||||
elif e.time_range.start >= last_op_end_time:
|
||||
last_op_end_time = e.time_range.end
|
||||
if e.time_range.end <= captured_region_end_time:
|
||||
captured_ops += 1
|
||||
captured_microseconds += e.time_range.elapsed_us()
|
||||
elif should_print_missing():
|
||||
print_missing(e.stack)
|
||||
total_ops += 1
|
||||
total_microseconds += e.time_range.elapsed_us()
|
||||
else:
|
||||
pass # ops recursively called from other ops (ignored)
|
||||
|
||||
unique_graphs = Profiler.unique_graphs
|
||||
Profiler.unique_graphs = 0
|
||||
# we counted one extra op that is part of the profiler setup code
|
||||
total_ops -= 1
|
||||
|
||||
return ProfileResult(
|
||||
captured=ProfileMetrics(
|
||||
microseconds=captured_microseconds,
|
||||
operators=captured_ops,
|
||||
fusions=captured_ops - captured_regions,
|
||||
graphs=captured_regions,
|
||||
),
|
||||
total=ProfileMetrics(
|
||||
microseconds=total_microseconds,
|
||||
operators=total_ops,
|
||||
fusions=total_ops - 1,
|
||||
),
|
||||
unique_graphs=unique_graphs,
|
||||
)
|
||||
|
||||
|
||||
def fx_insert_profiling(gm: torch.fx.GraphModule, example_inputs: list[Any]) -> Any:
|
||||
def _wrapped(*args: Any) -> Any:
|
||||
with torch.profiler.record_function("TORCHDYNAMO"):
|
||||
return gm.forward(*args)
|
||||
|
||||
Profiler.unique_graphs += 1
|
||||
return _wrapped
|
||||
@@ -0,0 +1,130 @@
|
||||
"""
|
||||
Python execution state recording and replay functionality.
|
||||
|
||||
This module provides mechanisms for capturing and replaying Python execution state:
|
||||
|
||||
- ModuleRecord: Tracks module access patterns and attribute usage
|
||||
- DummyModule: Lightweight module substitute for replay
|
||||
- ExecutionRecord: Manages execution context including globals, locals and builtins
|
||||
- ExecutionRecorder: Records variable states and module access during execution
|
||||
|
||||
The module enables serialization and reproduction of Python execution environments,
|
||||
particularly useful for debugging and testing frameworks that need to capture
|
||||
and recreate specific program states.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import field
|
||||
from io import BufferedReader, BufferedWriter
|
||||
from types import CellType, CodeType, ModuleType
|
||||
from typing import Any, IO
|
||||
from typing_extensions import Self
|
||||
|
||||
from torch.utils._import_utils import import_dill
|
||||
|
||||
|
||||
dill = import_dill()
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ModuleRecord:
|
||||
module: ModuleType
|
||||
accessed_attrs: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DummyModule:
|
||||
name: str
|
||||
is_torch: bool = False
|
||||
value: object = None
|
||||
|
||||
@property
|
||||
def __name__(self) -> str:
|
||||
return self.name
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ExecutionRecord:
|
||||
code: CodeType
|
||||
closure: tuple[CellType]
|
||||
globals: dict[str, Any] = field(default_factory=dict)
|
||||
locals: dict[str, Any] = field(default_factory=dict)
|
||||
builtins: dict[str, Any] = field(default_factory=dict)
|
||||
code_options: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def dump(self, f: IO[str] | BufferedWriter) -> None:
|
||||
assert dill is not None, "replay_record requires `pip install dill`"
|
||||
dill.dump(self, f)
|
||||
|
||||
@classmethod
|
||||
def load(cls, f: IO[bytes] | BufferedReader) -> Self:
|
||||
assert dill is not None, "replay_record requires `pip install dill`"
|
||||
return dill.load(f)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ExecutionRecorder:
|
||||
LOCAL_MOD_PREFIX = "___local_mod_"
|
||||
|
||||
code: CodeType
|
||||
closure: tuple[CellType]
|
||||
globals: dict[str, Any] = field(default_factory=dict)
|
||||
locals: dict[str, Any] = field(default_factory=dict)
|
||||
builtins: dict[str, Any] = field(default_factory=dict)
|
||||
code_options: dict[str, Any] = field(default_factory=dict)
|
||||
name_to_modrec: dict[str, ModuleRecord] = field(default_factory=dict)
|
||||
|
||||
def add_local_var(self, name: str, var: Any) -> None:
|
||||
if isinstance(var, ModuleType):
|
||||
self.locals[name] = self._add_mod(var)
|
||||
else:
|
||||
self.locals[name] = var
|
||||
|
||||
def add_global_var(self, name: str, var: Any) -> None:
|
||||
if isinstance(var, ModuleType):
|
||||
self.globals[name] = self._add_mod(var)
|
||||
else:
|
||||
self.globals[name] = var
|
||||
|
||||
def add_local_mod(self, name: str, mod: ModuleType) -> None:
|
||||
assert isinstance(mod, ModuleType)
|
||||
self.add_global_var(name, mod)
|
||||
|
||||
def record_module_access(self, mod: ModuleType, name: str, val: Any) -> None:
|
||||
if isinstance(val, ModuleType):
|
||||
self.name_to_modrec[mod.__name__].accessed_attrs[name] = self._add_mod(val)
|
||||
return
|
||||
|
||||
if mod.__name__ in self.name_to_modrec:
|
||||
self.name_to_modrec[mod.__name__].accessed_attrs[name] = val
|
||||
|
||||
def get_record(self) -> ExecutionRecord:
|
||||
return ExecutionRecord(
|
||||
self.code,
|
||||
self.closure,
|
||||
ExecutionRecorder._resolve_modules(self.globals),
|
||||
ExecutionRecorder._resolve_modules(self.locals),
|
||||
self.builtins.copy(),
|
||||
self.code_options.copy(),
|
||||
)
|
||||
|
||||
def _add_mod(self, mod: ModuleType) -> ModuleRecord:
|
||||
if mod.__name__ not in self.name_to_modrec:
|
||||
self.name_to_modrec[mod.__name__] = ModuleRecord(mod)
|
||||
|
||||
return self.name_to_modrec[mod.__name__]
|
||||
|
||||
@classmethod
|
||||
def _resolve_modules(cls, vars: dict[str, Any]) -> dict[str, Any]:
|
||||
def resolve_module(var: Any) -> Any:
|
||||
if not isinstance(var, ModuleRecord):
|
||||
return var
|
||||
|
||||
dummy_mod = DummyModule(var.module.__name__)
|
||||
for attr_name, attr_value in var.accessed_attrs.items():
|
||||
attr_value = resolve_module(attr_value)
|
||||
dummy_mod.__setattr__(attr_name, attr_value)
|
||||
|
||||
return dummy_mod
|
||||
|
||||
return {k: resolve_module(v) for k, v in vars.items()}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,637 @@
|
||||
"""
|
||||
Utilities for reproducing and debugging issues in Dynamo after graph capture.
|
||||
|
||||
This file provides tools and infrastructure for debugging problems that occur
|
||||
after Dynamo has captured the graph but before/during backend compilation.
|
||||
Key components include:
|
||||
|
||||
- Minification tools to reduce large graphs to minimal failing examples
|
||||
- Accuracy testing to validate compiled graph outputs match eager mode
|
||||
- Repro generation to create standalone reproduction scripts
|
||||
- Debug backends for capturing and analyzing failures
|
||||
- Utilities for saving/loading graph states and inputs
|
||||
|
||||
The tools here focus specifically on the post-graph-capture stage, making them
|
||||
useful for debugging backend compilation issues, AOTAutograd problems, and
|
||||
accuracy discrepancies between compiled and eager execution.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import textwrap
|
||||
from collections.abc import Callable, Sequence
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch._dynamo.debug_utils import (
|
||||
AccuracyError,
|
||||
backend_accuracy_fails,
|
||||
BUCK_CMD_PREFIX,
|
||||
BuckTargetWriter,
|
||||
extra_imports,
|
||||
generate_config_string,
|
||||
generate_env_vars_string,
|
||||
helper_for_dump_minify,
|
||||
InputReader,
|
||||
InputWriter,
|
||||
minifier_dir,
|
||||
NNModuleToString,
|
||||
NopInputReader,
|
||||
run_fwd_maybe_bwd,
|
||||
same_two_models,
|
||||
)
|
||||
from torch.fx.experimental.symbolic_shapes import fx_placeholder_targets
|
||||
from torch.hub import tqdm
|
||||
|
||||
from .. import config
|
||||
from ..backends.registry import CompilerFn, lookup_backend, register_debug_backend
|
||||
from ..debug_utils import clone_inputs_retaining_gradness
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
inductor_config = import_module("torch._inductor.config")
|
||||
use_buck = inductor_config.is_fbcode()
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# MAIN ENTRY POINT
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def _accuracy_fails(
|
||||
gm: torch.fx.GraphModule,
|
||||
example_inputs: Sequence[Any],
|
||||
compiler_fn: Callable[[torch.fx.GraphModule, list[Any]], torch.fx.GraphModule],
|
||||
) -> bool:
|
||||
return backend_accuracy_fails(
|
||||
gm,
|
||||
example_inputs,
|
||||
compiler_fn,
|
||||
only_fwd=config.repro_forward_only,
|
||||
ignore_non_fp=config.repro_ignore_non_fp,
|
||||
)
|
||||
|
||||
|
||||
class WrapBackendDebug:
|
||||
def __init__(
|
||||
self, unconfigured_compiler_fn: CompilerFn, compiler_name: str | None
|
||||
) -> None:
|
||||
functools.wraps(unconfigured_compiler_fn)(self)
|
||||
self._torchdynamo_orig_backend = unconfigured_compiler_fn
|
||||
self._compiler_name = compiler_name
|
||||
if hasattr(unconfigured_compiler_fn, "__name__"):
|
||||
self.__name__ = unconfigured_compiler_fn.__name__
|
||||
if hasattr(unconfigured_compiler_fn, "compiler_name"):
|
||||
self.__name__ = unconfigured_compiler_fn.compiler_name # type: ignore[attr-defined]
|
||||
if hasattr(unconfigured_compiler_fn, "get_compiler_config"):
|
||||
self.get_compiler_config = unconfigured_compiler_fn.get_compiler_config # type: ignore[attr-defined]
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[Any], **kwargs: Any
|
||||
) -> torch.fx.GraphModule:
|
||||
compiler_fn = functools.partial(self._torchdynamo_orig_backend, **kwargs)
|
||||
assert config.repro_after in ("dynamo", "aot", None)
|
||||
|
||||
if config.repro_after == "dynamo":
|
||||
|
||||
def add_paths(exc: Exception) -> None:
|
||||
exc.minifier_path = os.path.join(minifier_dir(), "minifier_launcher.py") # type: ignore[attr-defined]
|
||||
if use_buck:
|
||||
exc.buck_command = " ".join( # type: ignore[attr-defined]
|
||||
BUCK_CMD_PREFIX
|
||||
+ [BuckTargetWriter(exc.minifier_path).cmd_line_path] # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
if config.repro_level == 3:
|
||||
dump_to_minify_after_dynamo(gm, example_inputs, self._compiler_name)
|
||||
|
||||
# Check for either accuracy (level 4) or other type of failures.
|
||||
if config.repro_level == 4:
|
||||
# Check Accuracy
|
||||
compiled_gm = compiler_fn(copy.deepcopy(gm), example_inputs)
|
||||
if _accuracy_fails(gm, example_inputs, compiler_fn): # type: ignore[arg-type]
|
||||
log.warning(
|
||||
"Accuracy failed for the TorchDynamo produced graph. Creating script to minify the error."
|
||||
)
|
||||
dump_to_minify_after_dynamo(
|
||||
fx.GraphModule(gm, copy.deepcopy(gm.graph)),
|
||||
example_inputs,
|
||||
self._compiler_name,
|
||||
)
|
||||
exc = AccuracyError("Bad accuracy detected.")
|
||||
add_paths(exc)
|
||||
raise exc
|
||||
else:
|
||||
try:
|
||||
compiled_gm = compiler_fn(copy.deepcopy(gm), example_inputs)
|
||||
run_fwd_maybe_bwd(compiled_gm, example_inputs) # type: ignore[arg-type]
|
||||
except Exception as exc:
|
||||
log.warning(
|
||||
"Compiled Fx GraphModule failed. Creating script to minify the error."
|
||||
)
|
||||
if config.repro_level == 1:
|
||||
dump_state_fn = functools.partial(
|
||||
dump_backend_state, compiler_name=self._compiler_name
|
||||
)
|
||||
dump_state_fn(
|
||||
fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs
|
||||
)
|
||||
elif config.repro_level == 2:
|
||||
dump_to_minify_after_dynamo(
|
||||
fx.GraphModule(gm, copy.deepcopy(gm.graph)),
|
||||
example_inputs,
|
||||
self._compiler_name,
|
||||
)
|
||||
add_paths(exc)
|
||||
raise
|
||||
else:
|
||||
compiled_gm = compiler_fn(gm, example_inputs)
|
||||
|
||||
return compiled_gm # type: ignore[return-value]
|
||||
|
||||
|
||||
def wrap_backend_debug(
|
||||
unconfigured_compiler_fn: CompilerFn, compiler_name: str | None
|
||||
) -> WrapBackendDebug:
|
||||
"""
|
||||
A minifier decorator that wraps the TorchDynamo produced Fx graph modules.
|
||||
As opposed to wrap_compiler_debug, this wrapper intercepts at the
|
||||
TorchDynamo produced Fx Graph Module. This makes it backend-agnostic to some
|
||||
level, e.g., it is useful for minifying issues related to Aot Autograd
|
||||
tracing. If an error is found, we minify and save the minified repro in
|
||||
repro.tar.gz.
|
||||
"""
|
||||
return WrapBackendDebug(unconfigured_compiler_fn, compiler_name)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# REPRO DUMPERS
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def generate_dynamo_fx_repro_string(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str | None,
|
||||
check_accuracy: bool = False,
|
||||
*,
|
||||
stable_output: bool = False,
|
||||
save_dir: str | None = None,
|
||||
command: str = "run",
|
||||
) -> str:
|
||||
"""
|
||||
Generate a repro string for backend-agnostic minified version.
|
||||
"""
|
||||
|
||||
model_str = NNModuleToString.convert(gm)
|
||||
|
||||
# TODO: Figure out why torch.compile'd hash isn't work on this codepath
|
||||
writer = InputWriter(save_dir, stable_hash=True)
|
||||
for placeholder, arg in zip(fx_placeholder_targets(gm), args):
|
||||
if isinstance(arg, (int, torch.SymInt)):
|
||||
writer.symint(placeholder, arg)
|
||||
elif isinstance(arg, torch.Tensor):
|
||||
# TODO: improve these names with FQN
|
||||
writer.tensor(placeholder, arg)
|
||||
else:
|
||||
raise TypeError(f"arg is neither SymInt/int nor torch.Tensor, {arg}")
|
||||
load_args = "\n".join(writer.lines())
|
||||
|
||||
return textwrap.dedent(
|
||||
f"""
|
||||
{generate_env_vars_string(stable_output=stable_output)}
|
||||
from math import inf
|
||||
import torch
|
||||
from torch import tensor, device
|
||||
import torch.fx as fx
|
||||
import torch._dynamo
|
||||
from torch._dynamo.testing import rand_strided
|
||||
from torch._dynamo.debug_utils import run_fwd_maybe_bwd
|
||||
|
||||
{generate_config_string(stable_output=stable_output)}
|
||||
|
||||
{extra_imports}
|
||||
|
||||
{model_str}
|
||||
mod = Repro()
|
||||
|
||||
{load_args}
|
||||
|
||||
if __name__ == '__main__':
|
||||
from torch._dynamo.repro.after_dynamo import run_repro
|
||||
run_repro(mod, load_args, accuracy={check_accuracy!r}, command={command!r},
|
||||
save_dir={save_dir!r}, autocast={torch.is_autocast_enabled()!r}, backend={compiler_name!r})
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def dump_backend_repro_as_file(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str | None,
|
||||
check_accuracy: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Saves the repro to a repro.py file
|
||||
"""
|
||||
curdir = os.getcwd()
|
||||
subdir = os.path.join(os.getcwd(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
file_name = os.path.join(subdir, f"minified_{len(gm.graph.nodes)}_nodes.py")
|
||||
log.warning(
|
||||
"Writing checkpoint with %s nodes to %s", len(gm.graph.nodes), file_name
|
||||
)
|
||||
|
||||
with open(file_name, "w") as fd:
|
||||
fd.write(
|
||||
generate_dynamo_fx_repro_string(
|
||||
gm, args, compiler_name, check_accuracy, save_dir=subdir
|
||||
)
|
||||
)
|
||||
latest_repro = os.path.join(curdir, "repro.py")
|
||||
log.warning("Copying %s to %s for convenience", file_name, latest_repro)
|
||||
|
||||
if use_buck:
|
||||
BuckTargetWriter(latest_repro).write()
|
||||
|
||||
shutil.copyfile(file_name, latest_repro)
|
||||
|
||||
|
||||
def dump_backend_state(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str | None,
|
||||
check_accuracy: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Dumps the dynamo graph to repro the issue.
|
||||
1) It tries to convert Fx GraphModule to a string. If we can, it writes to a
|
||||
repro.py file.
|
||||
2) If we can't convert Fx GraphModule to a string, we use to_folder to save
|
||||
the module and save a tar file.
|
||||
"""
|
||||
assert NNModuleToString.can_convert_to_string(gm)
|
||||
return dump_backend_repro_as_file(gm, args, compiler_name, check_accuracy)
|
||||
# return dump_backend_repro_as_tarfile(gm, args, compiler_name)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# MINIFIER DUMPER
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def dump_to_minify_after_dynamo(
|
||||
gm: torch.fx.GraphModule, args: Sequence[Any], compiler_name: str | None
|
||||
) -> None:
|
||||
# TODO: factor this out
|
||||
subdir = os.path.join(minifier_dir(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
helper_for_dump_minify(
|
||||
generate_dynamo_fx_repro_string(
|
||||
gm,
|
||||
args,
|
||||
compiler_name,
|
||||
check_accuracy=config.repro_level == 4,
|
||||
save_dir=subdir,
|
||||
command="minify",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# MINIFIER BACKENDS
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
@register_debug_backend # type: ignore[arg-type]
|
||||
def dynamo_minifier_backend(
|
||||
gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: str | None
|
||||
) -> fx.GraphModule:
|
||||
from functorch.compile import minifier
|
||||
|
||||
compiler_fn = lookup_backend(compiler_name) # type: ignore[arg-type]
|
||||
|
||||
# TODO: It's inconsistent to pass SymInt inputs but REAL tensors.
|
||||
# We should pass ints and look at the GraphModule placeholders
|
||||
# to resolve them to SymInt (if necessary)
|
||||
example_inputs = [
|
||||
i.node.hint if isinstance(i, torch.SymInt) else i for i in example_inputs
|
||||
]
|
||||
|
||||
try:
|
||||
compiled_gm = compiler_fn(gm, example_inputs)
|
||||
run_fwd_maybe_bwd(compiled_gm, example_inputs) # type: ignore[arg-type]
|
||||
raise ValueError("No issue was detected")
|
||||
except Exception as exc:
|
||||
orig_failure = str(exc)
|
||||
log.warning(
|
||||
"Compiled Fx GraphModule failed. Creating script to minify the error."
|
||||
)
|
||||
dump_state_fn = functools.partial(
|
||||
dump_backend_state, compiler_name=compiler_name
|
||||
)
|
||||
dump_state_fn(fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs)
|
||||
fails_fn = functools.partial(
|
||||
backend_fails,
|
||||
compiler_fn=compiler_fn,
|
||||
orig_failure=orig_failure,
|
||||
)
|
||||
minifier(
|
||||
gm,
|
||||
example_inputs,
|
||||
module_fails=fails_fn,
|
||||
dump_state=dump_state_fn,
|
||||
)
|
||||
return gm
|
||||
|
||||
|
||||
@register_debug_backend # type: ignore[arg-type]
|
||||
def dynamo_accuracy_minifier_backend(
|
||||
gm: fx.GraphModule, example_inputs: Sequence[Any], compiler_name: str | None
|
||||
) -> fx.GraphModule:
|
||||
from functorch.compile import minifier
|
||||
|
||||
compiler_fn = lookup_backend(compiler_name) # type: ignore[arg-type]
|
||||
|
||||
# Set the eval mode to remove randomness.
|
||||
gm.eval()
|
||||
|
||||
# Check Accuracy
|
||||
if _accuracy_fails(gm, example_inputs, compiler_fn): # type: ignore[arg-type]
|
||||
log.warning("Accuracy failed for the TorchDynamo produced graph")
|
||||
dump_state_fn = functools.partial(
|
||||
dump_backend_state, compiler_name=compiler_name, check_accuracy=True
|
||||
)
|
||||
fails_fn = functools.partial(
|
||||
_accuracy_fails,
|
||||
compiler_fn=compiler_fn, # type: ignore[arg-type]
|
||||
)
|
||||
dump_state_fn(fx.GraphModule(gm, copy.deepcopy(gm.graph)), example_inputs)
|
||||
minifier(
|
||||
gm,
|
||||
example_inputs,
|
||||
module_fails=fails_fn,
|
||||
dump_state=dump_state_fn,
|
||||
)
|
||||
else:
|
||||
log.error("Input graph does not fail accuracy testing")
|
||||
return gm
|
||||
|
||||
|
||||
def backend_fails(
|
||||
gm: fx.GraphModule,
|
||||
example_inputs: Sequence[Any],
|
||||
compiler_fn: CompilerFn,
|
||||
orig_failure: Sequence[Any],
|
||||
) -> bool:
|
||||
"""
|
||||
Minifier uses this function to identify if the minified graph module fails
|
||||
with the same error.
|
||||
|
||||
One caveat is that minifier can potentially go into a wrong direction when
|
||||
the resulting graph module fails for a different reason. To avoid this, we
|
||||
save the string for the original exception and check similarity between new
|
||||
and old exception. They can be somewhat different in some cases, when the
|
||||
exception string depends on the failing node information. So, we have a
|
||||
loose similarity metric to guide the minifier path.
|
||||
"""
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
try:
|
||||
# Run the original gm to check eager validity
|
||||
run_fwd_maybe_bwd(gm, clone_inputs_retaining_gradness(example_inputs))
|
||||
compiled_gm = compiler_fn(gm, example_inputs) # type: ignore[arg-type]
|
||||
run_fwd_maybe_bwd(compiled_gm, clone_inputs_retaining_gradness(example_inputs)) # type: ignore[arg-type]
|
||||
except Exception as e:
|
||||
new_failure = str(e)
|
||||
if SequenceMatcher(None, orig_failure, new_failure).ratio() > 0.5:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# REPRO MAIN
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def run_load_args(options: Any, mod: torch.nn.Module, load_args: Any) -> list[Any]:
|
||||
if not hasattr(load_args, "_version"):
|
||||
log.warning(
|
||||
"load_args does not have a _version attribute, please file a bug to PyTorch "
|
||||
"and describe how you generate this repro script"
|
||||
)
|
||||
else:
|
||||
if load_args._version > 0:
|
||||
log.warning(
|
||||
"load_args is version %s, but this version of PyTorch only supports "
|
||||
"version 0. We will try to run it anyway but there may be an incompatibility; "
|
||||
"if so, try upgrading your version of PyTorch.",
|
||||
load_args._version,
|
||||
)
|
||||
|
||||
nop_reader = NopInputReader()
|
||||
load_args(nop_reader)
|
||||
|
||||
with tqdm(desc="Loading inputs", total=nop_reader.total) as pbar:
|
||||
input_reader = InputReader(save_dir=options.save_dir, pbar=pbar)
|
||||
load_args(input_reader)
|
||||
args = input_reader.args
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def repro_minify(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
|
||||
args = run_load_args(options, mod, load_args)
|
||||
|
||||
# Setup debug minifier compiler
|
||||
if not options.accuracy:
|
||||
compiler_fn = lookup_backend("dynamo_minifier_backend")
|
||||
else:
|
||||
compiler_fn = lookup_backend("dynamo_accuracy_minifier_backend")
|
||||
|
||||
if options.backend is None:
|
||||
raise RuntimeError(
|
||||
"Compiler name is None - this likely means that a custom compiler "
|
||||
"was called by torchdynamo. Please remove this error, import your "
|
||||
"custom compiler function, and replace the backend=None "
|
||||
"line in run_repro to backend=<my_imported_custom_function>"
|
||||
)
|
||||
|
||||
dynamo_minifier_backend = functools.partial(
|
||||
compiler_fn,
|
||||
compiler_name=options.backend, # type: ignore[call-arg]
|
||||
)
|
||||
opt_mod = torch._dynamo.optimize(dynamo_minifier_backend)(mod)
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=options.autocast):
|
||||
opt_mod(*args)
|
||||
|
||||
|
||||
def repro_run(options: Any, mod: torch.nn.Module, load_args: Any) -> None:
|
||||
opt_mod = torch._dynamo.optimize(options.backend)(mod)
|
||||
|
||||
if options.accuracy != "":
|
||||
mod.eval()
|
||||
opt_mod.eval() # type: ignore[union-attr]
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=options.autocast):
|
||||
# TODO: disable clone
|
||||
args = run_load_args(options, mod, load_args)
|
||||
assert same_two_models(mod, mod, args), "Eager itself failed" # type: ignore[arg-type]
|
||||
if not same_two_models(
|
||||
mod, # type: ignore[arg-type]
|
||||
opt_mod, # type: ignore[arg-type]
|
||||
args,
|
||||
only_fwd=config.repro_forward_only,
|
||||
ignore_non_fp=config.repro_ignore_non_fp,
|
||||
):
|
||||
raise AccuracyError("Dynamo failed")
|
||||
else:
|
||||
with torch.amp.autocast("cuda", enabled=options.autocast):
|
||||
args = run_load_args(options, mod, load_args)
|
||||
run_fwd_maybe_bwd(mod, args, only_fwd=options.only_fwd, disable_clone=True) # type: ignore[arg-type]
|
||||
del args
|
||||
|
||||
args = run_load_args(options, mod, load_args)
|
||||
run_fwd_maybe_bwd(
|
||||
opt_mod, # type: ignore[arg-type]
|
||||
args,
|
||||
only_fwd=options.only_fwd,
|
||||
disable_clone=True, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def run_repro(
|
||||
mod: torch.nn.Module,
|
||||
load_args: Any,
|
||||
*,
|
||||
command: str = "run",
|
||||
accuracy: bool | str = "",
|
||||
save_dir: str | None = None,
|
||||
autocast: bool = False,
|
||||
backend: str = "inductor",
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
for k in kwargs:
|
||||
log.warning(
|
||||
"Unrecognized kwarg %s; perhaps this repro was made on a newer version of PyTorch",
|
||||
k,
|
||||
)
|
||||
|
||||
if accuracy is True:
|
||||
accuracy = "accuracy"
|
||||
elif accuracy is False:
|
||||
accuracy = ""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description=f"""\
|
||||
An after_dynamo repro script, typically triggering a bug in Dynamo or
|
||||
AOTAutograd. When run with no arguments, this script defaults to running
|
||||
'{command}'. Extra flags may be available; to find out more, try '{command}
|
||||
--help'. There are also alternate subcommands available, see below.
|
||||
|
||||
default settings on this script:
|
||||
{accuracy=}
|
||||
{save_dir=}
|
||||
""",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
|
||||
def common_flags(parser: argparse.ArgumentParser) -> None:
|
||||
accuracy_group = parser.add_mutually_exclusive_group()
|
||||
accuracy_group.add_argument(
|
||||
"--no-accuracy",
|
||||
dest="accuracy",
|
||||
action="store_const",
|
||||
const="",
|
||||
default=accuracy,
|
||||
help="do not test accuracy, just run the module and see if it errors",
|
||||
)
|
||||
accuracy_group.add_argument(
|
||||
"--accuracy",
|
||||
action="store_const",
|
||||
const="accuracy",
|
||||
default=accuracy,
|
||||
help="test accuracy",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=save_dir,
|
||||
metavar="DIR",
|
||||
help="directory where saved inputs live",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save-dir",
|
||||
dest="save_dir",
|
||||
action="store_const",
|
||||
const=None,
|
||||
help="don't use any directory for saved inputs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-isolate",
|
||||
dest="isolate",
|
||||
action="store_false",
|
||||
default=False,
|
||||
help="no isolate (doesn't do anything for after_dynamo)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--autocast",
|
||||
default=autocast,
|
||||
action="store_true",
|
||||
help="use torch.cuda.amp.autocast",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-autocast",
|
||||
dest="autocast",
|
||||
action="store_false",
|
||||
help="don't use torch.cuda.amp.autocast",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
type=str,
|
||||
default=backend,
|
||||
metavar="BACKEND",
|
||||
help="torch.compile backend to use",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(
|
||||
dest="command", metavar="{run,minify}", required=True
|
||||
)
|
||||
|
||||
parser_run = subparsers.add_parser(
|
||||
"run",
|
||||
help="just run the repro",
|
||||
)
|
||||
common_flags(parser_run)
|
||||
parser_run.add_argument(
|
||||
"--only-fwd",
|
||||
action="store_true",
|
||||
help="don't run backwards compilation for testing",
|
||||
)
|
||||
|
||||
parser_minify = subparsers.add_parser(
|
||||
"minify", help="run the minifier on the repro"
|
||||
)
|
||||
common_flags(parser_minify)
|
||||
|
||||
args = None
|
||||
if len(sys.argv) <= 1:
|
||||
args = [command, *sys.argv[1:]]
|
||||
|
||||
options = parser.parse_args(args)
|
||||
COMMAND_FNS = {
|
||||
"minify": repro_minify,
|
||||
"run": repro_run,
|
||||
}
|
||||
COMMAND_FNS[options.command](options, mod, load_args)
|
||||
@@ -0,0 +1,662 @@
|
||||
"""
|
||||
Utilities for debugging and reproducing issues in Ahead of Time with Inductor (AOTI) compilation.
|
||||
|
||||
This file provides tools and utilities for:
|
||||
- Generating minimal reproducible test cases (minification)
|
||||
- Handling exported programs and graph modules
|
||||
- Creating debug repros for AOTI compilation issues
|
||||
- Supporting both accuracy testing and error reproduction
|
||||
- Managing configuration and environment for repro cases
|
||||
|
||||
The main components include:
|
||||
- Minification tools to reduce test cases while preserving errors
|
||||
- Repro generation utilities for exported programs
|
||||
- Error handling specific to AOTI compilation
|
||||
- Command-line interface for running and managing repros
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import textwrap
|
||||
from collections.abc import Sequence
|
||||
from importlib import import_module
|
||||
from typing import Any, IO
|
||||
|
||||
import torch
|
||||
from torch._dynamo.debug_utils import (
|
||||
_cuda_system_info_comment,
|
||||
BuckTargetWriter,
|
||||
extra_imports,
|
||||
generate_config_string,
|
||||
generate_env_vars_string,
|
||||
helper_for_dump_minify,
|
||||
InputReader,
|
||||
minifier_dir,
|
||||
NNModuleToString,
|
||||
NopInputReader,
|
||||
)
|
||||
from torch.export import ExportedProgram
|
||||
from torch.hub import tqdm
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
inductor_config = import_module("torch._inductor.config")
|
||||
use_buck = inductor_config.is_fbcode()
|
||||
|
||||
|
||||
class AOTIMinifierError(Exception):
|
||||
def __init__(self, original_exception: str | Exception) -> None:
|
||||
additional_message = "This error is caused by a bug in the AOTI minifier, please report a bug to PyTorch"
|
||||
full_message = f"{additional_message}: {str(original_exception)}"
|
||||
super().__init__(full_message)
|
||||
self.original_exception = original_exception
|
||||
|
||||
|
||||
def dump_to_minify(
|
||||
exported_program: ExportedProgram,
|
||||
compiler_name: str,
|
||||
command: str = "minify",
|
||||
options: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
If command is "minify":
|
||||
Dump exported_program to `debug_dir/minifier/minifier_launcher.py`, with minify command.
|
||||
If command is "run":
|
||||
Dump exported_program to `cwd/repro.py`, with run command.
|
||||
"""
|
||||
assert command in ["minify", "run"]
|
||||
|
||||
subdir = os.path.join(minifier_dir(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
|
||||
if command == "minify":
|
||||
out = io.StringIO()
|
||||
save_graph_repro_ep(
|
||||
out,
|
||||
compiler_name,
|
||||
exported_program=exported_program,
|
||||
save_dir=subdir,
|
||||
command="minify",
|
||||
config_patches=options,
|
||||
)
|
||||
return helper_for_dump_minify(out.getvalue())
|
||||
else:
|
||||
curdir = os.getcwd()
|
||||
file_name = os.path.join(curdir, "repro.py")
|
||||
try:
|
||||
with open(file_name, "w") as fd:
|
||||
save_graph_repro_ep(
|
||||
fd,
|
||||
compiler_name,
|
||||
exported_program=exported_program,
|
||||
config_patches=options,
|
||||
save_dir=subdir,
|
||||
command="run",
|
||||
module_in_comment=True,
|
||||
)
|
||||
log.warning("Writing repro file to %s", file_name)
|
||||
if use_buck:
|
||||
BuckTargetWriter(file_name).write()
|
||||
except OSError:
|
||||
log.warning("No write permissions for %s", file_name)
|
||||
|
||||
|
||||
def get_module_string(gm: torch.fx.GraphModule) -> str:
|
||||
def _convert_to_comment(s_: str) -> str:
|
||||
s = s_.split("\n")
|
||||
if len(s) == 1:
|
||||
return "# " + s_
|
||||
first = s.pop(0)
|
||||
for i in range(len(s)):
|
||||
line = s[i]
|
||||
if line.strip() != "":
|
||||
s[i] = "# " + line
|
||||
else:
|
||||
s[i] = ""
|
||||
s = "\n".join(s)
|
||||
s = first + "\n" + s
|
||||
return s
|
||||
|
||||
module_string = NNModuleToString.convert(gm)
|
||||
return _convert_to_comment(module_string)
|
||||
|
||||
|
||||
def save_graph_repro_ep(
|
||||
fd: IO[Any],
|
||||
compiler_name: str,
|
||||
*,
|
||||
exported_program: ExportedProgram | None = None,
|
||||
gm: torch.nn.Module | None = None,
|
||||
args: tuple[Any] | None = None,
|
||||
config_patches: dict[str, str] | None = None,
|
||||
stable_output: bool = False,
|
||||
save_dir: str | None = None,
|
||||
command: str = "run",
|
||||
accuracy: str | bool | None = None,
|
||||
check_str: str | None = None,
|
||||
module_in_comment: bool = False,
|
||||
strict: bool = False,
|
||||
) -> None:
|
||||
# Save graph for reproducing the error.
|
||||
# Either exported_program or gm will be saved, depending on which one is defined.
|
||||
# Only one of exported_program and gm should be defined.
|
||||
|
||||
if exported_program is None and gm is None:
|
||||
raise AOTIMinifierError("One of exported_program and gm must be defined")
|
||||
if exported_program is not None and gm is not None:
|
||||
raise AOTIMinifierError("Only one of exported_program and gm can be defined")
|
||||
if gm is not None and args is None:
|
||||
raise AOTIMinifierError("If gm is defined, args should also be defined")
|
||||
|
||||
if exported_program is None:
|
||||
assert gm is not None
|
||||
assert args is not None
|
||||
exported_program = torch.export.export(gm, args, strict=strict)
|
||||
elif gm is None:
|
||||
gm = exported_program.module(check_guards=False)
|
||||
|
||||
# save a graph preview using gm
|
||||
module_string = get_module_string(gm) # type: ignore[arg-type]
|
||||
fd.write(module_string)
|
||||
|
||||
# save a graph repro using exported_program
|
||||
fd.write(
|
||||
generate_compiler_repro_exported_program(
|
||||
exported_program,
|
||||
options=config_patches,
|
||||
stable_output=stable_output,
|
||||
save_dir=save_dir,
|
||||
)
|
||||
)
|
||||
if accuracy is None:
|
||||
accuracy = "_accuracy" in compiler_name
|
||||
fd.write("if __name__ == '__main__':\n")
|
||||
fd.write(" from torch._dynamo.repro.aoti import run_repro\n")
|
||||
fd.write(
|
||||
f" with torch.no_grad():\n"
|
||||
f" run_repro(exported_program, config_patches=config_patches, accuracy={accuracy!r}, command={command!r}, "
|
||||
f"save_dir={save_dir!r}, check_str={check_str!r})\n"
|
||||
)
|
||||
|
||||
|
||||
def dump_compiler_graph_state(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: Sequence[Any],
|
||||
compiler_name: str,
|
||||
*,
|
||||
config_patches: dict[str, str] | None = None,
|
||||
accuracy: str | bool | None = None,
|
||||
strict: bool = False,
|
||||
) -> None:
|
||||
subdir = os.path.join(minifier_dir(), "checkpoints")
|
||||
if not os.path.exists(subdir):
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
file_name = os.path.join(subdir, f"{len(gm.graph.nodes)}.py")
|
||||
log.warning(
|
||||
"Writing checkpoint with %s nodes to %s", len(gm.graph.nodes), file_name
|
||||
)
|
||||
with open(file_name, "w") as fd:
|
||||
save_graph_repro_ep(
|
||||
fd,
|
||||
compiler_name,
|
||||
gm=gm,
|
||||
args=tuple(args),
|
||||
config_patches=config_patches,
|
||||
save_dir=subdir,
|
||||
accuracy=accuracy,
|
||||
module_in_comment=True,
|
||||
strict=strict,
|
||||
)
|
||||
curdir = os.getcwd()
|
||||
repro_path = os.path.join(curdir, "repro.py")
|
||||
try:
|
||||
shutil.copyfile(file_name, repro_path)
|
||||
log.warning("Copying repro file for convenience to %s", repro_path)
|
||||
if use_buck:
|
||||
BuckTargetWriter(file_name).write()
|
||||
except OSError:
|
||||
log.warning("No write permissions for %s", repro_path)
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# DUMP REPROS
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
def generate_compiler_repro_exported_program(
|
||||
exported_program: ExportedProgram,
|
||||
*,
|
||||
options: dict[str, str] | None = None,
|
||||
stable_output: bool = False,
|
||||
save_dir: str | None = None,
|
||||
) -> str:
|
||||
model_str = textwrap.dedent(
|
||||
f"""
|
||||
{generate_env_vars_string(stable_output=stable_output)}
|
||||
import torch
|
||||
import torch._inductor.inductor_prims
|
||||
|
||||
{generate_config_string(stable_output=stable_output)}
|
||||
|
||||
isolate_fails_code_str = None
|
||||
|
||||
{extra_imports}
|
||||
|
||||
"""
|
||||
)
|
||||
if not stable_output:
|
||||
model_str += f"# torch version: {torch.version.__version__}\n"
|
||||
if hasattr(torch.version, "cuda"):
|
||||
model_str += f"# torch cuda version: {torch.version.cuda}\n"
|
||||
if hasattr(torch.version, "git_version"):
|
||||
model_str += f"# torch git version: {torch.version.git_version}\n\n\n"
|
||||
model_str += _cuda_system_info_comment()
|
||||
if save_dir:
|
||||
ep_path = os.path.join(save_dir, "exported_program.pt2")
|
||||
else:
|
||||
ep_path = "exported_program.pt2"
|
||||
torch.export.save(exported_program, ep_path)
|
||||
|
||||
model_str += f"exported_program = torch.export.load('{ep_path}')\n"
|
||||
model_str += "# print(exported_program.graph)\n"
|
||||
model_str += f"config_patches={options}\n"
|
||||
return model_str
|
||||
|
||||
|
||||
def repro_load_args(load_args: Any, save_dir: str | None) -> tuple[Any]:
|
||||
if not hasattr(load_args, "_version"):
|
||||
log.warning(
|
||||
"load_args does not have a _version attribute, please file a bug to PyTorch "
|
||||
"and describe how you generate this repro script"
|
||||
)
|
||||
else:
|
||||
if load_args._version > 0:
|
||||
log.warning(
|
||||
"load_args is version %s, but this version of PyTorch only supports "
|
||||
"version 0. We will try to run it anyway but there may be an incompatibility; "
|
||||
"if so, try upgrading your version of PyTorch.",
|
||||
load_args._version,
|
||||
)
|
||||
|
||||
nop_reader = NopInputReader()
|
||||
load_args(nop_reader)
|
||||
|
||||
with tqdm(desc="Loading inputs", total=nop_reader.total) as pbar:
|
||||
input_reader = InputReader(save_dir=save_dir, pbar=pbar)
|
||||
load_args(input_reader)
|
||||
args = input_reader.args
|
||||
|
||||
return tuple(args)
|
||||
|
||||
|
||||
def repro_common(
|
||||
options: Any, exported_program: ExportedProgram
|
||||
) -> tuple[torch.fx.GraphModule, Any, Any]:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
torch._inductor.config.generate_intermediate_hooks = True
|
||||
mod = exported_program.module(check_guards=False)
|
||||
args, kwargs = exported_program.example_inputs
|
||||
return mod, args, kwargs # type: ignore[return-value]
|
||||
|
||||
|
||||
def repro_get_args(
|
||||
options: Any,
|
||||
exported_program: ExportedProgram,
|
||||
config_patches: dict[str, Any] | None,
|
||||
) -> tuple[torch.fx.GraphModule, Any, Any]:
|
||||
mod, args, kwargs = repro_common(options, exported_program)
|
||||
return mod, args, kwargs
|
||||
|
||||
|
||||
def repro_run(
|
||||
options: Any,
|
||||
exported_program: ExportedProgram,
|
||||
config_patches: dict[str, Any] | None,
|
||||
) -> None:
|
||||
from torch._inductor import _aoti_compile_and_package_inner
|
||||
|
||||
gm, args, kwargs = repro_common(options, exported_program)
|
||||
|
||||
from torch.cuda import synchronize
|
||||
|
||||
_aoti_compile_and_package_inner(
|
||||
gm,
|
||||
args,
|
||||
kwargs,
|
||||
load_and_run=True,
|
||||
check_accuracy=options.accuracy,
|
||||
inductor_configs=config_patches,
|
||||
)
|
||||
|
||||
need_sync = False
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor) and arg.is_cuda:
|
||||
need_sync = True
|
||||
break
|
||||
|
||||
if need_sync:
|
||||
synchronize() # ensure segfaults are surfaced
|
||||
|
||||
|
||||
def export_for_aoti_minifier(
|
||||
gm: torch.nn.Module,
|
||||
tuple_inputs: tuple[Any],
|
||||
strict: bool = False,
|
||||
skip_export_error: bool = True,
|
||||
) -> torch.nn.Module | None:
|
||||
# Some graphs cannot be used for AOTI/export (illegal graphs), these should be
|
||||
# considered as graphs that don't fail in the minifier, so the minifier keeps searching.
|
||||
# In these case, we return None. Otherwise, we return the exported graph module.
|
||||
# This won't affect the minifier result because the minifier is only responsible for catching
|
||||
# errors in AOTI, not export.
|
||||
#
|
||||
# Please add to this list of illegal graphs if you change the implementation here.
|
||||
# - graph output is not allowed by export
|
||||
#
|
||||
# If skip_export_error=True, then the errors in export will not be raised, and the minifier
|
||||
# will keep exploring and ignore this graph.
|
||||
from torch._dynamo.exc import UserError, UserErrorType
|
||||
|
||||
try:
|
||||
ep = torch.export.export(gm, tuple_inputs, strict=strict)
|
||||
gm = ep.module(check_guards=False)
|
||||
return gm
|
||||
except Exception as e:
|
||||
if skip_export_error:
|
||||
return None
|
||||
if isinstance(e, UserError) and e.error_type == UserErrorType.INVALID_OUTPUT:
|
||||
# graph output is not allowed by export when strict=True
|
||||
return None
|
||||
if isinstance(e, RuntimeError):
|
||||
# graph output is not allowed by export when strict=False
|
||||
pattern = r"Found .* in output, which is not a known type\."
|
||||
if re.search(pattern, str(e)) is not None:
|
||||
return None
|
||||
raise AOTIMinifierError(e) from e
|
||||
# we should never reach here
|
||||
# pyrefly: ignore [unreachable]
|
||||
return None
|
||||
|
||||
|
||||
def repro_minify(
|
||||
options: Any,
|
||||
exported_program: ExportedProgram,
|
||||
config_patches: dict[str, Any] | None,
|
||||
) -> None:
|
||||
from functorch.compile import minifier
|
||||
from torch._inductor import _aoti_compile_and_package_inner
|
||||
from torch._inductor.compile_fx import _aoti_flatten_inputs
|
||||
|
||||
mod, args, kwargs = repro_common(options, exported_program)
|
||||
|
||||
# update serialized_in_spec and serialized_out_spec
|
||||
flat_example_inputs, inductor_configs = _aoti_flatten_inputs(
|
||||
mod, args, kwargs, options=config_patches
|
||||
)
|
||||
compiler_name = "aot_inductor"
|
||||
assert options.minifier_export_mode in ["dynamo", "python"]
|
||||
strict = options.minifier_export_mode == "dynamo"
|
||||
skip_export_error = options.skip_export_error
|
||||
|
||||
from torch.cuda import synchronize
|
||||
|
||||
need_sync = False
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor) and arg.is_cuda:
|
||||
need_sync = True
|
||||
break
|
||||
|
||||
def module_fails(
|
||||
gm: torch.fx.GraphModule,
|
||||
flat_example_inputs: list[Any],
|
||||
check_str: str | None = None,
|
||||
) -> bool:
|
||||
# Need to export first so the in_spec and out_spec are populated
|
||||
tuple_inputs = tuple(flat_example_inputs)
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
gm = export_for_aoti_minifier(
|
||||
gm, tuple_inputs, strict=strict, skip_export_error=skip_export_error
|
||||
)
|
||||
|
||||
# Some graphs cannot be used for AOTI/export (illegal graphs), these should be
|
||||
# considered as graphs that don't fail in the minifier, so the minifier keeps searching.
|
||||
if gm is None:
|
||||
return False
|
||||
|
||||
assert isinstance(gm, torch.fx.GraphModule)
|
||||
|
||||
try:
|
||||
_aoti_compile_and_package_inner(
|
||||
gm,
|
||||
tuple_inputs,
|
||||
load_and_run=True,
|
||||
check_accuracy=options.accuracy,
|
||||
inductor_configs=inductor_configs,
|
||||
)
|
||||
if need_sync:
|
||||
synchronize() # ensure segfaults are surfaced
|
||||
return False
|
||||
except Exception as e:
|
||||
if check_str is not None and check_str not in repr(e):
|
||||
return False
|
||||
return True
|
||||
|
||||
minifier(
|
||||
mod,
|
||||
flat_example_inputs,
|
||||
module_fails=functools.partial(module_fails, check_str=options.check_str),
|
||||
dump_state=functools.partial(
|
||||
dump_compiler_graph_state,
|
||||
compiler_name=compiler_name,
|
||||
config_patches=config_patches,
|
||||
accuracy=options.accuracy,
|
||||
strict=strict,
|
||||
),
|
||||
save_dir=options.save_dir,
|
||||
offload_to_disk=options.offload_to_disk,
|
||||
skip_offload=options.skip_saving_eager_intermediates,
|
||||
skip_sanity=options.skip_sanity,
|
||||
max_granularity=options.max_granularity,
|
||||
)
|
||||
|
||||
|
||||
def run_repro(
|
||||
exported_program: ExportedProgram,
|
||||
*,
|
||||
config_patches: dict[str, str] | None = None,
|
||||
command: str = "run",
|
||||
accuracy: bool | str = "",
|
||||
save_dir: str | None = None,
|
||||
tracing_mode: str | None = None,
|
||||
check_str: str | None = None,
|
||||
minifier_export_mode: str = "python",
|
||||
skip_export_error: bool = True,
|
||||
**more_kwargs: Any,
|
||||
) -> Any:
|
||||
for k in more_kwargs:
|
||||
log.warning(
|
||||
"Unrecognized kwarg %s; perhaps this repro was made on a newer version of PyTorch",
|
||||
k,
|
||||
)
|
||||
|
||||
if accuracy is True:
|
||||
accuracy = "accuracy"
|
||||
elif accuracy is False:
|
||||
accuracy = ""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description=f"""\
|
||||
An AOTI repro script, typically triggering a bug in PyTorch AOTInductor.
|
||||
When run with no arguments, this script defaults to running '{command}'.
|
||||
Extra flags may be available; to find out more, try '{command} --help'.
|
||||
There are also alternate subcommands available, see below.
|
||||
|
||||
default settings on this script:
|
||||
{accuracy=}
|
||||
{tracing_mode=}
|
||||
{save_dir=}
|
||||
{check_str=}
|
||||
""",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
|
||||
def common_flags(parser: argparse.ArgumentParser) -> None:
|
||||
accuracy_group = parser.add_mutually_exclusive_group()
|
||||
accuracy_group.add_argument(
|
||||
"--no-accuracy",
|
||||
dest="accuracy",
|
||||
action="store_const",
|
||||
const="",
|
||||
default=accuracy,
|
||||
help="do not test accuracy, just run the module and see if it errors",
|
||||
)
|
||||
accuracy_group.add_argument(
|
||||
"--accuracy",
|
||||
action="store_const",
|
||||
const="accuracy",
|
||||
default=accuracy,
|
||||
help="""\
|
||||
test if the RMSE between the compiled module and the fp64 reference is greater
|
||||
than eager and the fp64 reference. This is usually more reliable than the
|
||||
standard allclose test, as we expect numeric differences from compiling, often
|
||||
improving accuracy over eager. RMSE test allows for compiled module to
|
||||
diverge greatly from eager, as long as this divergence moves it closer to the
|
||||
'true' mathematical value of the network. Caveats: (1) double precision can
|
||||
still suffer from rounding error, so it is not a perfect reference (see for
|
||||
example 'Herbie: Automatically Improving Floating Point Accuracy') for
|
||||
approaches that detect the necessary working precision and compute it in
|
||||
arbitrary precision floating point; unfortunately, this is not practical for
|
||||
tensor computation; (2) if there are not enough samples in the output being
|
||||
compared, we may get unlucky and have an unlucky greater RMSE than eager; this
|
||||
could be overcome by applying a more rigorous statistical test at some
|
||||
p-value, which we leave for future work.
|
||||
""",
|
||||
)
|
||||
accuracy_group.add_argument(
|
||||
"--strict-accuracy",
|
||||
dest="accuracy",
|
||||
action="store_const",
|
||||
const="strict_accuracy",
|
||||
default=accuracy,
|
||||
help="""\
|
||||
by default, when doing accuracy minification we will reject reductions which
|
||||
change the divergence from a floating point divergence to a integral/boolean
|
||||
divergence. This is because some operations like ReLU involve temporarily
|
||||
sharp boundaries that smooth out again afterwards; without requiring
|
||||
divergence on floating point, the minifier will often fixate on divergent
|
||||
boolean tensor even though this is not the true source of the divergence.
|
||||
However, rejecting these reductions makes it more difficult for the minifier
|
||||
to make process. Using this option will let the minifier progress for ALL
|
||||
divergences--you just might not end up with a useful repro in the end.""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=save_dir,
|
||||
metavar="DIR",
|
||||
help="directory where saved inputs live",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-save-dir",
|
||||
dest="save_dir",
|
||||
action="store_const",
|
||||
const=None,
|
||||
help="don't use any directory for saved inputs",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(
|
||||
dest="command", metavar="{run,minify}", required=True
|
||||
)
|
||||
|
||||
parser_run = subparsers.add_parser(
|
||||
"run",
|
||||
help="just run the repro",
|
||||
)
|
||||
common_flags(parser_run)
|
||||
|
||||
parser_minify = subparsers.add_parser(
|
||||
"minify", help="run the minifier on the repro"
|
||||
)
|
||||
common_flags(parser_minify)
|
||||
parser_get_args = subparsers.add_parser("get_args", help="get the args")
|
||||
common_flags(parser_get_args)
|
||||
parser_minify.add_argument(
|
||||
"--skip-saving-eager-intermediates",
|
||||
action="store_true",
|
||||
help="skip saving eager intermediates on --minify",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--offload-to-disk",
|
||||
action="store_true",
|
||||
help="during minification, offload delta debugging intermediates to disk. Use if you're OOMing",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--skip-sanity",
|
||||
action="store_true",
|
||||
help="skip sanity check at beginning of minification on original graph",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--max-granularity",
|
||||
type=int,
|
||||
default=None,
|
||||
help="start at this granularity and work down; must be power of 2",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--check-str",
|
||||
type=str,
|
||||
default=check_str,
|
||||
help="require minified program to fail with error containing this string",
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--minifier-export-mode",
|
||||
type=str,
|
||||
default=minifier_export_mode,
|
||||
help=(
|
||||
"The export mode used in minifier, either dynamo or python."
|
||||
"`dynamo` corresponds to strict=True, and `python` corresponds to strict=False."
|
||||
),
|
||||
)
|
||||
parser_minify.add_argument(
|
||||
"--skip-export-error",
|
||||
type=bool,
|
||||
default=skip_export_error,
|
||||
help="Skip intermediate graphs that cannot be exported.",
|
||||
)
|
||||
|
||||
# Run the repro in the context of minification, inverting exit code meaning
|
||||
parser_minifier_query = subparsers.add_parser(
|
||||
"minifier-query",
|
||||
)
|
||||
common_flags(parser_minifier_query)
|
||||
parser_minifier_query.add_argument(
|
||||
"--check-str",
|
||||
type=str,
|
||||
default=check_str,
|
||||
help="require minified program to fail with error containing this string",
|
||||
)
|
||||
|
||||
args = None
|
||||
if len(sys.argv) <= 1:
|
||||
args = [command, *sys.argv[1:]]
|
||||
|
||||
options = parser.parse_args(args)
|
||||
COMMAND_FNS = {
|
||||
"minify": repro_minify,
|
||||
"run": repro_run,
|
||||
"get_args": repro_get_args,
|
||||
}
|
||||
return COMMAND_FNS[options.command](
|
||||
options, exported_program, config_patches=config_patches
|
||||
)
|
||||
@@ -0,0 +1,750 @@
|
||||
"""
|
||||
This module provides functionality for resuming Python execution at specific points in code,
|
||||
primarily used by PyTorch Dynamo for control flow handling and optimization. It implements
|
||||
bytecode transformation and execution state management to enable:
|
||||
|
||||
- Resuming execution at arbitrary points in Python bytecode
|
||||
- Managing context managers and their state across execution boundaries
|
||||
- Transforming and generating new code objects with preserved execution state
|
||||
- Supporting Python 3.11+ exception handling and block management
|
||||
- Restoring torch function mode stacks and other execution context
|
||||
|
||||
The module is critical for PyTorch Dynamo's ability to optimize code while preserving
|
||||
Python semantics and execution state.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import dataclasses
|
||||
import sys
|
||||
import types
|
||||
from collections.abc import Callable, Iterable
|
||||
from contextlib import AbstractContextManager
|
||||
from typing import Any, cast
|
||||
|
||||
from .bytecode_transformation import (
|
||||
add_push_null,
|
||||
bytecode_from_template,
|
||||
create_binary_subscr,
|
||||
create_call_function,
|
||||
create_call_function_ex,
|
||||
create_instruction,
|
||||
create_jump_absolute,
|
||||
create_load_const,
|
||||
Instruction,
|
||||
overwrite_instruction,
|
||||
transform_code_object,
|
||||
unique_id,
|
||||
)
|
||||
from .utils import ExactWeakKeyDictionary
|
||||
|
||||
|
||||
# taken from code.h in cpython
|
||||
CO_OPTIMIZED = 0x0001
|
||||
CO_NEWLOCALS = 0x0002
|
||||
CO_VARARGS = 0x0004
|
||||
CO_VARKEYWORDS = 0x0008
|
||||
CO_NESTED = 0x0010
|
||||
CO_GENERATOR = 0x0020
|
||||
CO_NOFREE = 0x0040
|
||||
CO_COROUTINE = 0x0080
|
||||
CO_ITERABLE_COROUTINE = 0x0100
|
||||
CO_ASYNC_GENERATOR = 0x0200
|
||||
|
||||
# trace_rules.py import this constant for consistency
|
||||
TORCH_DYNAMO_RESUME_IN_PREFIX = "torch_dynamo_resume_in"
|
||||
IS_TRACING_RESUME_PROLOGUE_VARNAME = "__is_tracing_resume_prologue"
|
||||
|
||||
|
||||
# If is_resume - this codegen is for a resume function
|
||||
def _initial_push_null(insts: list[Instruction]) -> None:
|
||||
if sys.version_info >= (3, 11):
|
||||
insts.append(create_instruction("PUSH_NULL"))
|
||||
if sys.version_info < (3, 13):
|
||||
insts.append(create_instruction("SWAP", arg=2))
|
||||
|
||||
|
||||
# Generates bytecode from template and splits the code where LOAD_FAST dummy is present.
|
||||
def _bytecode_from_template_with_split(
|
||||
template: Callable[..., Any],
|
||||
stack_index: int,
|
||||
varname_map: dict[str, Any] | None = None,
|
||||
) -> tuple[list[Instruction], list[Instruction]]:
|
||||
template_code = bytecode_from_template(template, varname_map=varname_map)
|
||||
template_code.append(create_instruction("POP_TOP"))
|
||||
|
||||
# adjust exception table entry depth
|
||||
for inst in template_code:
|
||||
if inst.exn_tab_entry:
|
||||
inst.exn_tab_entry.depth += stack_index
|
||||
|
||||
# search for LOAD_FAST dummy and replace it with 2 NOPs (we can break up the bytecode between them)
|
||||
dummy_idx, dummy_inst = next(
|
||||
(
|
||||
(i, inst)
|
||||
for i, inst in enumerate(template_code)
|
||||
if inst.opname in ("LOAD_FAST", "LOAD_FAST_BORROW")
|
||||
and inst.argval == "dummy"
|
||||
),
|
||||
(None, None),
|
||||
)
|
||||
assert dummy_idx is not None and dummy_inst is not None
|
||||
|
||||
# replace LOAD_FAST dummy with first NOP marking exception area
|
||||
overwrite_instruction(dummy_inst, [create_instruction("NOP")])
|
||||
|
||||
# POP_TOP follows LOAD_FAST dummy - replace with NOP marking end of exception area
|
||||
assert template_code[dummy_idx + 1].opname == "POP_TOP"
|
||||
overwrite_instruction(template_code[dummy_idx + 1], [create_instruction("NOP")])
|
||||
|
||||
return template_code[: dummy_idx + 1], template_code[dummy_idx + 1 :]
|
||||
|
||||
|
||||
def _try_except_tf_mode_template(dummy: Any, stack_var_name: Any) -> None:
|
||||
# NOTE: Make sure this name matches what is generated by symbolic_convert:import_source
|
||||
# on torch._dynamo.utils.
|
||||
# pyrefly: ignore [unknown-name]
|
||||
global __import_torch_dot__dynamo_dot_utils
|
||||
try:
|
||||
dummy
|
||||
except: # noqa: E722, B001
|
||||
__import_torch_dot__dynamo_dot_utils.set_torch_function_mode_stack( # type: ignore[name-defined]
|
||||
stack_var_name
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ReenterWith:
|
||||
stack_index: int
|
||||
target_values: tuple[Any, ...] | None = None
|
||||
|
||||
def try_except_torch_function_mode(
|
||||
self, code_options: dict[str, Any], cleanup: list[Instruction]
|
||||
) -> list[Instruction]:
|
||||
"""
|
||||
Codegen based off of:
|
||||
try:
|
||||
(rest)
|
||||
except:
|
||||
(restore previous tf mode stack)
|
||||
raise
|
||||
"""
|
||||
from .variables.torch_function import get_prev_stack_var_name
|
||||
|
||||
setup_try_except, epilogue = _bytecode_from_template_with_split(
|
||||
_try_except_tf_mode_template,
|
||||
self.stack_index,
|
||||
varname_map={"stack_var_name": get_prev_stack_var_name()},
|
||||
)
|
||||
cleanup[:] = epilogue + cleanup
|
||||
|
||||
return setup_try_except
|
||||
|
||||
# If we do not want to destroy the stack, we can do the same thing as a
|
||||
# `SETUP_WITH` block, only that we store the context manager in a local_symbol
|
||||
def try_finally(
|
||||
self, code_options: dict[str, Any], cleanup: list[Instruction]
|
||||
) -> list[Instruction]:
|
||||
"""
|
||||
Codegen based off of:
|
||||
load args
|
||||
enter context
|
||||
try:
|
||||
(rest)
|
||||
finally:
|
||||
exit context
|
||||
"""
|
||||
# NOTE: we assume that TOS is a context manager CLASS!
|
||||
# pyrefly: ignore [implicit-any]
|
||||
load_args = []
|
||||
if self.target_values:
|
||||
load_args = [create_load_const(val) for val in self.target_values]
|
||||
ctx_name = unique_id(f"___context_manager_{self.stack_index}")
|
||||
if ctx_name not in code_options["co_varnames"]:
|
||||
code_options["co_varnames"] += (ctx_name,)
|
||||
for name in ["__enter__", "__exit__"]:
|
||||
if name not in code_options["co_names"]:
|
||||
code_options["co_names"] += (name,)
|
||||
|
||||
create_ctx: list[Instruction] = []
|
||||
_initial_push_null(create_ctx)
|
||||
create_ctx.extend(
|
||||
[
|
||||
*load_args,
|
||||
*create_call_function(len(load_args), False),
|
||||
create_instruction("STORE_FAST", argval=ctx_name),
|
||||
]
|
||||
)
|
||||
|
||||
def _template(ctx: AbstractContextManager[Any], dummy: Any) -> None:
|
||||
ctx.__enter__()
|
||||
try:
|
||||
dummy
|
||||
finally:
|
||||
ctx.__exit__(None, None, None)
|
||||
|
||||
setup_try_finally, epilogue = _bytecode_from_template_with_split(
|
||||
_template, self.stack_index, varname_map={"ctx": ctx_name}
|
||||
)
|
||||
cleanup[:] = epilogue + cleanup
|
||||
return create_ctx + setup_try_finally
|
||||
|
||||
def __call__(
|
||||
self, code_options: dict[str, Any], cleanup: list[Instruction]
|
||||
) -> tuple[list[Instruction], Instruction | None]:
|
||||
"""
|
||||
Codegen based off of:
|
||||
with ctx(args):
|
||||
(rest)
|
||||
"""
|
||||
# NOTE: we assume that TOS is a context manager CLASS!
|
||||
# pyrefly: ignore [implicit-any]
|
||||
load_args = []
|
||||
if self.target_values:
|
||||
load_args = [create_load_const(val) for val in self.target_values]
|
||||
|
||||
create_ctx: list[Instruction] = []
|
||||
# Do not push NULL in Python 3.14+ since the NULL should be on the symbolic stack.
|
||||
if sys.version_info < (3, 14):
|
||||
_initial_push_null(create_ctx)
|
||||
create_ctx.extend(
|
||||
[
|
||||
*load_args,
|
||||
*create_call_function(len(load_args), False),
|
||||
]
|
||||
)
|
||||
|
||||
def _template(ctx: AbstractContextManager[Any], dummy: Any) -> None:
|
||||
with ctx:
|
||||
dummy
|
||||
|
||||
setup_with, epilogue = _bytecode_from_template_with_split(
|
||||
_template, self.stack_index
|
||||
)
|
||||
cleanup[:] = epilogue + cleanup
|
||||
|
||||
load_fast_ctx_inst = next(
|
||||
(
|
||||
inst
|
||||
for inst in setup_with
|
||||
if inst.opname in ("LOAD_FAST", "LOAD_FAST_BORROW")
|
||||
and inst.argval == "ctx"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert load_fast_ctx_inst is not None
|
||||
# ctx already loaded on stack before the template - no need to LOAD_FAST
|
||||
overwrite_instruction(load_fast_ctx_inst, [create_instruction("NOP")])
|
||||
|
||||
# 3.11+ only
|
||||
push_exc_info_gen = (
|
||||
inst for inst in epilogue if inst.opname == "PUSH_EXC_INFO"
|
||||
)
|
||||
push_exc_info_inst = next(push_exc_info_gen, None)
|
||||
# expect only 1 PUSH_EXC_INFO in epilogue
|
||||
assert next(push_exc_info_gen, None) is None
|
||||
|
||||
return create_ctx + setup_with, push_exc_info_inst
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ResumeFunctionMetadata:
|
||||
code: types.CodeType
|
||||
instructions: list[Instruction] = dataclasses.field(default_factory=list)
|
||||
# Python 3.11+ fields
|
||||
# NOTE: Python 3.11 removed blocks, but for our purposes, a "block" consists
|
||||
# of instructions of all exception table entries that have the same target.
|
||||
|
||||
# map from PUSH_EXC_INFO's in the prefix to original block target offset
|
||||
prefix_block_target_offset_remap: list[int] = dataclasses.field(
|
||||
default_factory=list
|
||||
)
|
||||
# per-offset map from new block target offsets to original block target offsets
|
||||
block_target_offset_remap: dict[tuple[int, int], dict[int, int]] = (
|
||||
dataclasses.field(default_factory=dict)
|
||||
)
|
||||
|
||||
|
||||
def _filter_iter(
|
||||
l1: Iterable[Any],
|
||||
l2: Iterable[Any],
|
||||
cond: Callable[[Any, Any], bool],
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Two-pointer conditional filter.
|
||||
e.g. _filter_iter(insts, sorted_offsets, lambda i, o: i.offset == o)
|
||||
returns the instructions with offsets in sorted_offsets
|
||||
"""
|
||||
it = iter(l2)
|
||||
res: list[Instruction] = []
|
||||
try:
|
||||
cur = next(it)
|
||||
for val in l1:
|
||||
if cond(val, cur):
|
||||
res.append(val)
|
||||
cur = next(it)
|
||||
except StopIteration:
|
||||
pass
|
||||
return res
|
||||
|
||||
|
||||
def _load_tuple_and_call(tup: tuple[Any, ...]) -> list[Instruction]:
|
||||
insts: list[Instruction] = []
|
||||
_initial_push_null(insts)
|
||||
insts.extend(create_load_const(val) for val in tup)
|
||||
insts.extend(create_call_function(len(tup), False))
|
||||
return insts
|
||||
|
||||
|
||||
class ContinueExecutionCache:
|
||||
cache = ExactWeakKeyDictionary()
|
||||
generated_code_metadata = ExactWeakKeyDictionary()
|
||||
|
||||
@classmethod
|
||||
def lookup(
|
||||
cls, code: types.CodeType, lineno: int, init_offset: int, *key: Any
|
||||
) -> types.CodeType:
|
||||
if code not in cls.cache:
|
||||
cls.cache[code] = {}
|
||||
key = tuple(key)
|
||||
if key not in cls.cache[code]:
|
||||
cls.cache[code][key] = cls.generate(code, lineno, init_offset, *key)
|
||||
return cls.cache[code][key]
|
||||
|
||||
@classmethod
|
||||
def generate(
|
||||
cls,
|
||||
code: types.CodeType,
|
||||
lineno: int,
|
||||
init_offset: int,
|
||||
resume_offset: int,
|
||||
setup_fn_target_offsets: tuple[int, ...], # only used in Python 3.11+
|
||||
nstack: int,
|
||||
argnames: tuple[str, ...],
|
||||
argnames_null: tuple[str, ...],
|
||||
setup_fns: tuple[ReenterWith, ...],
|
||||
handle_inactive_ctx: bool,
|
||||
stack_ctx_vars: tuple[tuple[int, tuple[Any, ...]], ...],
|
||||
argnames_ctx_vars: tuple[tuple[str, tuple[Any, ...]], ...],
|
||||
null_idxes: tuple[int, ...],
|
||||
# mainly used to ensure distinct code objects per stack trace,
|
||||
# which prevents excessive recompilation of inner frames
|
||||
nested_code_objs: tuple[types.CodeType],
|
||||
# Are we currently graph breaking on an instruction that doesn't push
|
||||
# its result to the stack? If so, and we are not the leaf resume, then we need to pop
|
||||
# the result of calling the next resume function.
|
||||
pop_nested_resume_result: bool,
|
||||
) -> types.CodeType:
|
||||
assert resume_offset is not None
|
||||
assert not (
|
||||
code.co_flags
|
||||
& (CO_GENERATOR | CO_COROUTINE | CO_ITERABLE_COROUTINE | CO_ASYNC_GENERATOR)
|
||||
)
|
||||
assert code.co_flags & CO_OPTIMIZED
|
||||
if code in ContinueExecutionCache.generated_code_metadata:
|
||||
return cls.generate_based_on_original_code_object(
|
||||
code,
|
||||
lineno,
|
||||
init_offset,
|
||||
resume_offset,
|
||||
setup_fn_target_offsets,
|
||||
nstack,
|
||||
argnames,
|
||||
argnames_null,
|
||||
setup_fns,
|
||||
handle_inactive_ctx,
|
||||
stack_ctx_vars,
|
||||
argnames_ctx_vars,
|
||||
null_idxes,
|
||||
nested_code_objs,
|
||||
pop_nested_resume_result,
|
||||
)
|
||||
|
||||
is_py311_plus = sys.version_info >= (3, 11)
|
||||
meta = ResumeFunctionMetadata(code)
|
||||
|
||||
def update(
|
||||
instructions: list[Instruction], code_options: dict[str, Any]
|
||||
) -> None:
|
||||
meta.instructions = copy.deepcopy(instructions)
|
||||
|
||||
args = ["__nested_resume_fns", "__nested_frame_values"]
|
||||
args += [f"___stack{i}" for i in range(nstack)]
|
||||
args.extend(v for v in argnames if v not in args)
|
||||
freevars = tuple(code_options["co_cellvars"] or []) + tuple(
|
||||
code_options["co_freevars"] or []
|
||||
)
|
||||
freevars = tuple(sorted(freevars))
|
||||
code_options["co_name"] = (
|
||||
f"{TORCH_DYNAMO_RESUME_IN_PREFIX}_{code_options['co_name']}_at_{lineno}"
|
||||
)
|
||||
if is_py311_plus:
|
||||
qualified_path = code_options["co_qualname"].rsplit(".", maxsplit=1)
|
||||
if len(qualified_path) == 1:
|
||||
code_options["co_qualname"] = code_options["co_name"]
|
||||
else:
|
||||
assert len(qualified_path) == 2
|
||||
module_name, co_name = qualified_path
|
||||
code_options["co_qualname"] = (
|
||||
f"{module_name}.{TORCH_DYNAMO_RESUME_IN_PREFIX}_{co_name}_at_{lineno}"
|
||||
)
|
||||
code_options["co_firstlineno"] = lineno
|
||||
code_options["co_cellvars"] = ()
|
||||
code_options["co_freevars"] = freevars
|
||||
code_options["co_argcount"] = len(args)
|
||||
code_options["co_posonlyargcount"] = 0
|
||||
code_options["co_kwonlyargcount"] = 0
|
||||
code_options["co_varnames"] = tuple(
|
||||
args
|
||||
+ [v for v in argnames_null if v not in args]
|
||||
+ [v for v in code_options["co_varnames"] if v not in args]
|
||||
+ [IS_TRACING_RESUME_PROLOGUE_VARNAME]
|
||||
)
|
||||
code_options["co_flags"] = code_options["co_flags"] & ~(
|
||||
CO_VARARGS | CO_VARKEYWORDS
|
||||
)
|
||||
target = next(i for i in instructions if i.offset == resume_offset)
|
||||
|
||||
prefix = []
|
||||
if is_py311_plus:
|
||||
if freevars:
|
||||
prefix.append(
|
||||
create_instruction("COPY_FREE_VARS", arg=len(freevars))
|
||||
)
|
||||
prefix.append(create_instruction("RESUME", arg=0))
|
||||
|
||||
# Set is_tracing_resume_prologue to prevent graph breaks.
|
||||
# This doesn't really do anything at runtime, but dynamo will trace this
|
||||
# and will know that we're in a resume function prologue.
|
||||
prefix.extend(
|
||||
[
|
||||
create_instruction("LOAD_CONST", argval=True),
|
||||
create_instruction(
|
||||
"STORE_FAST", argval=IS_TRACING_RESUME_PROLOGUE_VARNAME
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
cleanup: list[Instruction] = []
|
||||
hooks = {fn.stack_index: fn for fn in setup_fns}
|
||||
hook_target_offsets = {
|
||||
fn.stack_index: setup_fn_target_offsets[i]
|
||||
for i, fn in enumerate(setup_fns)
|
||||
}
|
||||
offset_to_inst = {inst.offset: inst for inst in instructions}
|
||||
# map old hook targets to new targets generated by the hook
|
||||
# pyrefly: ignore [implicit-any]
|
||||
old_hook_target_remap = {}
|
||||
stack_i = 0
|
||||
null_i = 0
|
||||
stack_ctx_vars_d = dict(stack_ctx_vars) # type: ignore[var-annotated,arg-type]
|
||||
for i in range(nstack + len(null_idxes)):
|
||||
if null_i < len(null_idxes) and null_idxes[null_i] == i:
|
||||
prefix.append(create_instruction("PUSH_NULL"))
|
||||
null_i += 1
|
||||
else:
|
||||
prefix.append(
|
||||
create_instruction("LOAD_FAST", argval=f"___stack{stack_i}")
|
||||
)
|
||||
if handle_inactive_ctx and stack_i in stack_ctx_vars_d:
|
||||
# NOTE: we assume that current stack var is a context manager CLASS!
|
||||
# Load args for context variable and construct it
|
||||
prefix.extend(_load_tuple_and_call(stack_ctx_vars_d[stack_i]))
|
||||
stack_i += 1
|
||||
|
||||
if i in hooks:
|
||||
hook = hooks.pop(i)
|
||||
hook_insts, exn_target = hook(code_options, cleanup)
|
||||
prefix.extend(hook_insts)
|
||||
if is_py311_plus:
|
||||
hook_target_offset = hook_target_offsets.pop(i)
|
||||
old_hook_target = offset_to_inst[hook_target_offset]
|
||||
meta.prefix_block_target_offset_remap.append(hook_target_offset)
|
||||
old_hook_target_remap[old_hook_target] = exn_target
|
||||
|
||||
if is_py311_plus:
|
||||
# reverse the mapping since targets of later/nested contexts are inserted
|
||||
# into the mapping later, but show up earlier in the prefix.
|
||||
meta.prefix_block_target_offset_remap = list(
|
||||
reversed(meta.prefix_block_target_offset_remap)
|
||||
)
|
||||
|
||||
assert not hooks
|
||||
|
||||
# NOTE: we assume that local var is a context manager CLASS!
|
||||
# initialize inactive context vars in argnames
|
||||
if handle_inactive_ctx:
|
||||
for name, vals in argnames_ctx_vars:
|
||||
prefix.append(create_instruction("LOAD_FAST", argval=name))
|
||||
prefix.extend(_load_tuple_and_call(vals))
|
||||
prefix.append(create_instruction("STORE_FAST", argval=name))
|
||||
|
||||
# 3.12+: store NULL into variables that were NULL
|
||||
if argnames_null:
|
||||
assert sys.version_info >= (3, 12)
|
||||
for v in argnames_null:
|
||||
assert v not in args
|
||||
prefix.extend(
|
||||
[
|
||||
create_instruction("PUSH_NULL"),
|
||||
create_instruction("STORE_FAST", argval=v),
|
||||
]
|
||||
)
|
||||
|
||||
# Call nested resume function
|
||||
if nested_code_objs:
|
||||
prefix.extend(
|
||||
[
|
||||
# set up __nested_resume_fns[-1] call
|
||||
*add_push_null(
|
||||
[
|
||||
create_instruction(
|
||||
"LOAD_FAST", argval="__nested_resume_fns"
|
||||
),
|
||||
create_instruction("LOAD_CONST", argval=-1),
|
||||
create_binary_subscr(),
|
||||
]
|
||||
),
|
||||
# del __nested_resume_fns[-1]
|
||||
create_instruction("LOAD_FAST", argval="__nested_resume_fns"),
|
||||
create_instruction("LOAD_CONST", argval=-1),
|
||||
create_instruction("DELETE_SUBSCR"),
|
||||
# load [__nested_resume_fns, __nested_frame_values]
|
||||
create_instruction("LOAD_FAST", argval="__nested_resume_fns"),
|
||||
create_instruction("LOAD_FAST", argval="__nested_frame_values"),
|
||||
create_instruction("BUILD_LIST", arg=2),
|
||||
# load __nested_frame_values[-1]
|
||||
create_instruction("LOAD_FAST", argval="__nested_frame_values"),
|
||||
create_instruction("LOAD_CONST", argval=-1),
|
||||
create_binary_subscr(),
|
||||
# create [
|
||||
# __nested_resume_fns,
|
||||
# __nested_frame_values,
|
||||
# *__nested_frame_values[-1],
|
||||
# ]
|
||||
create_instruction("LIST_EXTEND", arg=1),
|
||||
# del __nested_frame_values[-1]
|
||||
create_instruction("LOAD_FAST", argval="__nested_frame_values"),
|
||||
create_instruction("LOAD_CONST", argval=-1),
|
||||
create_instruction("DELETE_SUBSCR"),
|
||||
# delete __nested values
|
||||
create_instruction("DELETE_FAST", argval="__nested_resume_fns"),
|
||||
create_instruction(
|
||||
"DELETE_FAST", argval="__nested_frame_values"
|
||||
),
|
||||
# Set is_tracing_resume_prologue back to allow graph breaks
|
||||
# in the nested resume
|
||||
create_instruction("LOAD_CONST", argval=False),
|
||||
create_instruction(
|
||||
"STORE_FAST", argval=IS_TRACING_RESUME_PROLOGUE_VARNAME
|
||||
),
|
||||
# finish the call
|
||||
*create_call_function_ex(False, False),
|
||||
]
|
||||
)
|
||||
if pop_nested_resume_result:
|
||||
# pop the result of calling the nested resume function
|
||||
prefix.append(create_instruction("POP_TOP"))
|
||||
else:
|
||||
# Set is_tracing_resume_prologue back to allow graph breaks after the jump
|
||||
prefix.extend(
|
||||
[
|
||||
create_instruction("LOAD_CONST", argval=False),
|
||||
create_instruction(
|
||||
"STORE_FAST", argval=IS_TRACING_RESUME_PROLOGUE_VARNAME
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
prefix.append(create_jump_absolute(target))
|
||||
|
||||
# because the line number table monotonically increases from co_firstlineno
|
||||
# remove starts_line for any instructions before the graph break instruction
|
||||
# this will ensure the instructions after the break have the correct line numbers
|
||||
for inst in instructions:
|
||||
if inst.offset == target.offset:
|
||||
break
|
||||
inst.starts_line = None
|
||||
if sys.version_info >= (3, 11):
|
||||
inst.positions = None
|
||||
|
||||
if cleanup:
|
||||
prefix.extend(cleanup)
|
||||
prefix.extend(cls.unreachable_codes(code_options))
|
||||
|
||||
# remap original instructions' exception table entries
|
||||
if old_hook_target_remap:
|
||||
# pyrefly: ignore [unbound-name]
|
||||
assert is_py311_plus
|
||||
for inst in instructions:
|
||||
if (
|
||||
inst.exn_tab_entry
|
||||
and inst.exn_tab_entry.target in old_hook_target_remap
|
||||
):
|
||||
inst.exn_tab_entry.target = old_hook_target_remap[ # type: ignore[assignment]
|
||||
inst.exn_tab_entry.target
|
||||
]
|
||||
|
||||
# TODO(jansel): add dead code elimination here
|
||||
instructions[:] = prefix + instructions
|
||||
|
||||
new_code, _ = transform_code_object(code, update)
|
||||
ContinueExecutionCache.generated_code_metadata[new_code] = meta
|
||||
return new_code
|
||||
|
||||
@staticmethod
|
||||
def unreachable_codes(code_options: dict[str, Any]) -> list[Instruction]:
|
||||
"""Codegen a `raise None` to make analysis work for unreachable code"""
|
||||
return [
|
||||
create_load_const(None),
|
||||
create_instruction("RAISE_VARARGS", arg=1),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def generate_based_on_original_code_object(
|
||||
cls,
|
||||
code: types.CodeType,
|
||||
lineno: int,
|
||||
init_offset: int,
|
||||
resume_offset: int,
|
||||
setup_fn_target_offsets: tuple[int, ...],
|
||||
*args: Any,
|
||||
) -> types.CodeType:
|
||||
"""
|
||||
This handles the case of generating a resume into code generated
|
||||
to resume something else. We want to always generate starting
|
||||
from the original code object so that if control flow paths
|
||||
converge we only generated 1 resume function (rather than 2^n
|
||||
resume functions).
|
||||
"""
|
||||
|
||||
meta: ResumeFunctionMetadata = ContinueExecutionCache.generated_code_metadata[
|
||||
code
|
||||
]
|
||||
|
||||
def find_orig_offset(cur_offset: int) -> int:
|
||||
orig_offset = -1
|
||||
|
||||
def find_orig_offset_transform(
|
||||
instructions: list[Instruction], code_options: dict[str, Any]
|
||||
) -> None:
|
||||
nonlocal orig_offset
|
||||
(target,) = (i for i in instructions if i.offset == cur_offset)
|
||||
# match the functions starting at the last instruction as we have added a prefix
|
||||
new_target_tuple = tuple(
|
||||
i2
|
||||
for i1, i2 in zip(
|
||||
reversed(instructions), reversed(meta.instructions)
|
||||
)
|
||||
if i1 is target
|
||||
)
|
||||
|
||||
if not new_target_tuple:
|
||||
# Instruction with cur_offset in instructions was not found
|
||||
# in the original code - orig_offset left as -1.
|
||||
# Caller expected to handle this case.
|
||||
return
|
||||
|
||||
assert len(new_target_tuple) == 1
|
||||
new_target = new_target_tuple[0]
|
||||
|
||||
assert target.opcode == new_target.opcode
|
||||
assert new_target.offset is not None
|
||||
orig_offset = new_target.offset
|
||||
|
||||
transform_code_object(code, find_orig_offset_transform)
|
||||
return orig_offset
|
||||
|
||||
orig_init_offset = find_orig_offset(init_offset)
|
||||
# It is fine if the initial instruction is not found in the original code;
|
||||
# this means we graph broke in the prefix, which only happens with nested graph breaks.
|
||||
# We should not be running into ambiguous graph break issues here.
|
||||
orig_resume_offset = find_orig_offset(resume_offset)
|
||||
assert orig_resume_offset > -1, (
|
||||
"resume instruction not found in original code - this is a bug."
|
||||
)
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
# setup_fn_target_offsets currently contains the target offset of
|
||||
# each setup_fn, based on `code`. When we codegen the resume function
|
||||
# based on the original code object, `meta.code`, the offsets in
|
||||
# setup_fn_target_offsets must be based on `meta.code` instead.
|
||||
offset_key = (orig_init_offset, orig_resume_offset)
|
||||
# NOTE: we key by offset_key since the same resume function may graph
|
||||
# break in multiple places and we need different block_target_offset_remap's
|
||||
# for each graph break location. Keying by orig_resume_offset may not be enough
|
||||
# if 2 graph breaks on different initial offsets resume on the same instruction
|
||||
# (although this is rare and not tested anywhere).
|
||||
if offset_key not in meta.block_target_offset_remap:
|
||||
block_target_offset_remap = meta.block_target_offset_remap[
|
||||
offset_key
|
||||
# pyrefly: ignore [implicit-any]
|
||||
] = {}
|
||||
|
||||
def remap_block_offsets(
|
||||
instructions: list[Instruction], code_options: dict[str, Any]
|
||||
) -> None:
|
||||
# NOTE: each prefix block generates exactly one PUSH_EXC_INFO,
|
||||
# so we can tell which block a prefix PUSH_EXC_INFO belongs to,
|
||||
# by counting. Then we can use meta.prefix_block_target_offset_remap
|
||||
# to determine where in the original code the PUSH_EXC_INFO offset
|
||||
# replaced.
|
||||
prefix_blocks: list[Instruction] = []
|
||||
for inst in instructions:
|
||||
# NOTE meta.prefix_block_target_offset_remap is based off of how we codegen'd
|
||||
# context managers at the prefix/prologue of the resume function. It is the same for
|
||||
# every graph break in the same resume function, so we do not need to recompute
|
||||
# for each graph break (unlike for meta.block_target_offset_remap)
|
||||
if len(prefix_blocks) == len(
|
||||
meta.prefix_block_target_offset_remap
|
||||
):
|
||||
break
|
||||
if inst.opname == "PUSH_EXC_INFO":
|
||||
prefix_blocks.append(inst)
|
||||
|
||||
# remap block target offsets for blocks generated in the resume prefix
|
||||
for inst, o in zip(
|
||||
prefix_blocks, meta.prefix_block_target_offset_remap
|
||||
):
|
||||
block_target_offset_remap[cast(int, inst.offset)] = o
|
||||
|
||||
# current bytecode targets are after the prefix PUSH_EXC_INFO's
|
||||
cur_start_offset = (
|
||||
cast(int, prefix_blocks[-1].offset) if prefix_blocks else -1
|
||||
)
|
||||
# get the remaining block target offsets of the current bytecode
|
||||
cur_inst_offsets = sorted(
|
||||
n for n in setup_fn_target_offsets if n > cur_start_offset
|
||||
)
|
||||
targets = _filter_iter(
|
||||
instructions, cur_inst_offsets, lambda inst, o: inst.offset == o
|
||||
)
|
||||
# The original code and resume code should have matching suffixes.
|
||||
# Match the post-prefix block target offsets of the current resume code
|
||||
# and the original code.
|
||||
orig_targets = reversed(
|
||||
_filter_iter(
|
||||
zip(reversed(instructions), reversed(meta.instructions)),
|
||||
reversed(targets),
|
||||
lambda v1, v2: v1[0] is v2,
|
||||
)
|
||||
)
|
||||
for orig, cur in zip(orig_targets, targets):
|
||||
block_target_offset_remap[cur.offset] = orig[1].offset
|
||||
|
||||
transform_code_object(code, remap_block_offsets)
|
||||
|
||||
# if offset_key or offset is not in setup_fn_target_offsets, it is an error
|
||||
# that needs to be fixed
|
||||
setup_fn_target_offsets = tuple(
|
||||
meta.block_target_offset_remap[offset_key][n]
|
||||
for n in setup_fn_target_offsets
|
||||
)
|
||||
return ContinueExecutionCache.lookup(
|
||||
meta.code,
|
||||
lineno,
|
||||
orig_init_offset,
|
||||
orig_resume_offset,
|
||||
setup_fn_target_offsets,
|
||||
*args,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,70 @@
|
||||
"""This module implements tensor version operations for Dynamo tracing.
|
||||
|
||||
It provides primitives for handling tensor versioning during tracing, particularly in the
|
||||
context of functionalization where version operations are handled eagerly on fake tensors.
|
||||
|
||||
When we functionalize _tensor_version + _unsafe_set_version_counter, the ops disappear from
|
||||
the traced graph. We run them eagerly on the fake tensors used for tracing, in order to get
|
||||
past asserts that would fail in autograd.
|
||||
|
||||
Why is this ok?
|
||||
1) Versions on functional tensors do not make any sense since you cannot mutate a functional
|
||||
tensor.
|
||||
2) The whole point of version munging is to trick autograd into doing what we want, and after
|
||||
AotAutograd there is no longer any need for these ops.
|
||||
|
||||
Note this is similar to how no_grad is handled.
|
||||
"""
|
||||
|
||||
from contextlib import AbstractContextManager
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import SymInt
|
||||
from torch._prims import _make_prim, RETURN_TYPE
|
||||
from torch._subclasses import FakeTensorMode
|
||||
from torch._subclasses.functional_tensor import FunctionalTensorMode
|
||||
|
||||
|
||||
_tensor_version = _make_prim(
|
||||
schema="_tensor_version(Tensor self) -> SymInt",
|
||||
return_type=RETURN_TYPE.NEW,
|
||||
meta=torch.ops.aten._version.default,
|
||||
impl_aten=torch.ops.aten._version.default,
|
||||
doc="Tracable unbacked SymInt version of torch.Tensor._version",
|
||||
)
|
||||
|
||||
|
||||
@_tensor_version.py_impl(FakeTensorMode) # type: ignore[misc]
|
||||
def _tensor_version_fake(fake_mode: FakeTensorMode, self_tensor: Any) -> SymInt:
|
||||
"""
|
||||
The initial dynamo capture of _tensor_version + _unsafe_set_version_counter turns the
|
||||
`._version` into an unbacked SymInt so that we don't need to specialize on the `._version`
|
||||
of input tensors to the graph.
|
||||
"""
|
||||
assert fake_mode.shape_env is not None
|
||||
return fake_mode.shape_env.create_unbacked_symint()
|
||||
|
||||
|
||||
_unsafe_set_version_counter = _make_prim(
|
||||
schema="_unsafe_set_version_counter(Tensor[] tensors, SymInt[] versions) -> ()",
|
||||
return_type=RETURN_TYPE.NEW,
|
||||
meta=lambda self, version: None,
|
||||
impl_aten=torch._C._autograd._unsafe_set_version_counter,
|
||||
doc="Tracable+SymInt version of torch._C._autograd._unsafe_set_version_counter",
|
||||
)
|
||||
torch.fx.node.has_side_effect(_unsafe_set_version_counter)
|
||||
|
||||
|
||||
@_tensor_version.py_impl(FunctionalTensorMode) # type: ignore[misc]
|
||||
def _tensor_version_functional(mode: FunctionalTensorMode, self: Any) -> int:
|
||||
return self._version
|
||||
|
||||
|
||||
@_unsafe_set_version_counter.py_impl(FunctionalTensorMode) # type: ignore[misc]
|
||||
def _unsafe_set_version_counter_functional(
|
||||
ctx: AbstractContextManager[Any],
|
||||
tensors: tuple[torch.Tensor, ...],
|
||||
versions: tuple[int, ...],
|
||||
) -> None:
|
||||
torch._C._autograd._unsafe_set_version_counter(tensors, versions)
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Testing utilities for Dynamo, providing a specialized TestCase class and test running functionality.
|
||||
|
||||
This module extends PyTorch's testing framework with Dynamo-specific testing capabilities.
|
||||
It includes:
|
||||
- A custom TestCase class that handles Dynamo-specific setup/teardown
|
||||
- Test running utilities with dependency checking
|
||||
- Automatic reset of Dynamo state between tests
|
||||
- Proper handling of gradient mode state
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import unittest
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.testing
|
||||
from torch._dynamo import polyfills
|
||||
from torch._logging._internal import trace_log
|
||||
from torch.testing._internal.common_utils import ( # type: ignore[attr-defined]
|
||||
IS_WINDOWS,
|
||||
TEST_WITH_CROSSREF,
|
||||
TEST_WITH_TORCHDYNAMO,
|
||||
TestCase as TorchTestCase,
|
||||
)
|
||||
|
||||
from . import config, reset, utils
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def run_tests(needs: str | tuple[str, ...] = ()) -> None:
|
||||
from torch.testing._internal.common_utils import run_tests
|
||||
|
||||
if TEST_WITH_TORCHDYNAMO or TEST_WITH_CROSSREF:
|
||||
return # skip testing
|
||||
|
||||
if (
|
||||
not torch.xpu.is_available()
|
||||
and IS_WINDOWS
|
||||
and os.environ.get("TORCHINDUCTOR_WINDOWS_TESTS", "0") == "0"
|
||||
):
|
||||
return
|
||||
|
||||
if isinstance(needs, str):
|
||||
needs = (needs,)
|
||||
for need in needs:
|
||||
if need == "cuda":
|
||||
if not torch.cuda.is_available():
|
||||
return
|
||||
else:
|
||||
try:
|
||||
importlib.import_module(need)
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
run_tests()
|
||||
|
||||
|
||||
class TestCase(TorchTestCase):
|
||||
_exit_stack: contextlib.ExitStack
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
cls._exit_stack.close()
|
||||
super().tearDownClass()
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
super().setUpClass()
|
||||
cls._exit_stack = contextlib.ExitStack() # type: ignore[attr-defined]
|
||||
cls._exit_stack.enter_context( # type: ignore[attr-defined]
|
||||
config.patch(
|
||||
raise_on_ctx_manager_usage=True,
|
||||
suppress_errors=False,
|
||||
log_compilation_metrics=False,
|
||||
),
|
||||
)
|
||||
|
||||
def setUp(self) -> None:
|
||||
self._prior_is_grad_enabled = torch.is_grad_enabled()
|
||||
self._prior_nested_graph_breaks = config.nested_graph_breaks
|
||||
config.nested_graph_breaks = True
|
||||
super().setUp()
|
||||
reset()
|
||||
utils.counters.clear()
|
||||
self.handler = logging.NullHandler()
|
||||
trace_log.addHandler(self.handler)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
trace_log.removeHandler(self.handler)
|
||||
for k, v in utils.counters.items():
|
||||
log.debug("%s %s", k, v.most_common())
|
||||
reset()
|
||||
utils.counters.clear()
|
||||
torch._C._autograd._saved_tensors_hooks_enable()
|
||||
super().tearDown()
|
||||
if self._prior_is_grad_enabled is not torch.is_grad_enabled():
|
||||
log.warning("Running test changed grad mode")
|
||||
torch.set_grad_enabled(self._prior_is_grad_enabled)
|
||||
config.nested_graph_breaks = self._prior_nested_graph_breaks
|
||||
|
||||
def assertEqual(self, x: Any, y: Any, *args: Any, **kwargs: Any) -> None: # type: ignore[override]
|
||||
if (
|
||||
config.debug_disable_compile_counter
|
||||
and isinstance(x, utils.CompileCounterInt)
|
||||
or isinstance(y, utils.CompileCounterInt)
|
||||
):
|
||||
return
|
||||
return super().assertEqual(x, y, *args, **kwargs)
|
||||
|
||||
# assertExpectedInline might also need to be disabled for wrapped nested
|
||||
# graph break tests
|
||||
|
||||
|
||||
class CPythonTestCase(TestCase):
|
||||
"""
|
||||
Test class for CPython tests located in "test/dynamo/CPython/Py_version/*".
|
||||
|
||||
This class enables specific features that are disabled by default, such as
|
||||
tracing through unittest methods.
|
||||
"""
|
||||
|
||||
_stack: contextlib.ExitStack
|
||||
dynamo_strict_nopython = True
|
||||
|
||||
# Restore original unittest methods to simplify tracing CPython test cases.
|
||||
assertEqual = unittest.TestCase.assertEqual # type: ignore[assignment]
|
||||
assertNotEqual = unittest.TestCase.assertNotEqual # type: ignore[assignment]
|
||||
assertTrue = unittest.TestCase.assertTrue
|
||||
assertFalse = unittest.TestCase.assertFalse
|
||||
assertIs = unittest.TestCase.assertIs
|
||||
assertIsNot = unittest.TestCase.assertIsNot
|
||||
assertIsNone = unittest.TestCase.assertIsNone
|
||||
assertIsNotNone = unittest.TestCase.assertIsNotNone
|
||||
assertIn = unittest.TestCase.assertIn
|
||||
assertNotIn = unittest.TestCase.assertNotIn
|
||||
assertIsInstance = unittest.TestCase.assertIsInstance
|
||||
assertNotIsInstance = unittest.TestCase.assertNotIsInstance
|
||||
assertAlmostEqual = unittest.TestCase.assertAlmostEqual
|
||||
assertNotAlmostEqual = unittest.TestCase.assertNotAlmostEqual
|
||||
assertGreater = unittest.TestCase.assertGreater
|
||||
assertGreaterEqual = unittest.TestCase.assertGreaterEqual
|
||||
assertLess = unittest.TestCase.assertLess
|
||||
assertLessEqual = unittest.TestCase.assertLessEqual
|
||||
assertRegex = unittest.TestCase.assertRegex
|
||||
assertNotRegex = unittest.TestCase.assertNotRegex
|
||||
assertCountEqual = unittest.TestCase.assertCountEqual
|
||||
assertMultiLineEqual = polyfills.assert_multi_line_equal
|
||||
assertSequenceEqual = polyfills.assert_sequence_equal
|
||||
assertListEqual = unittest.TestCase.assertListEqual
|
||||
assertTupleEqual = unittest.TestCase.assertTupleEqual
|
||||
assertSetEqual = unittest.TestCase.assertSetEqual
|
||||
# pyrefly: ignore [bad-override]
|
||||
assertDictEqual = polyfills.assert_dict_equal
|
||||
# pyrefly: ignore [bad-override]
|
||||
assertRaises = unittest.TestCase.assertRaises
|
||||
# pyrefly: ignore [bad-override]
|
||||
assertRaisesRegex = unittest.TestCase.assertRaisesRegex
|
||||
assertWarns = unittest.TestCase.assertWarns
|
||||
assertWarnsRegex = unittest.TestCase.assertWarnsRegex
|
||||
assertLogs = unittest.TestCase.assertLogs
|
||||
fail = unittest.TestCase.fail
|
||||
failureException = unittest.TestCase.failureException
|
||||
|
||||
def compile_fn(
|
||||
self,
|
||||
fn: Callable[..., Any],
|
||||
backend: str | Callable[..., Any],
|
||||
nopython: bool,
|
||||
) -> Callable[..., Any]:
|
||||
# We want to compile only the test function, excluding any setup code
|
||||
# from unittest
|
||||
|
||||
method = getattr(self, self._testMethodName)
|
||||
method = torch._dynamo.optimize(backend, error_on_graph_break=nopython)(method)
|
||||
|
||||
setattr(self, self._testMethodName, method)
|
||||
return fn
|
||||
|
||||
def _dynamo_test_key(self) -> str:
|
||||
suffix = super()._dynamo_test_key()
|
||||
test_cls = self.__class__
|
||||
test_file = inspect.getfile(test_cls).split(os.sep)[-1].split(".")[0]
|
||||
py_ver = re.search(r"/([\d_]+)/", inspect.getfile(test_cls))
|
||||
if py_ver:
|
||||
py_ver = py_ver.group().strip(os.sep).replace("_", "") # type: ignore[assignment]
|
||||
else:
|
||||
return suffix
|
||||
return f"CPython{py_ver}-{test_file}-{suffix}"
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
cls._stack.close()
|
||||
super().tearDownClass()
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
# Skip test if python versions doesn't match
|
||||
prefix = os.path.join("dynamo", "cpython") + os.path.sep
|
||||
regex = re.escape(prefix) + r"\d_\d{2}"
|
||||
search_path = inspect.getfile(cls)
|
||||
m = re.search(regex, search_path)
|
||||
if m:
|
||||
test_py_ver = tuple(map(int, m.group().removeprefix(prefix).split("_")))
|
||||
py_ver = sys.version_info[:2]
|
||||
if py_ver != test_py_ver:
|
||||
expected = ".".join(map(str, test_py_ver))
|
||||
got = ".".join(map(str, py_ver))
|
||||
raise unittest.SkipTest(
|
||||
f"Test requires Python {expected} but got Python {got}"
|
||||
)
|
||||
else:
|
||||
raise unittest.SkipTest(
|
||||
f"Test requires a specific Python version but not found in path {inspect.getfile(cls)}"
|
||||
)
|
||||
|
||||
super().setUpClass()
|
||||
cls._stack = contextlib.ExitStack() # type: ignore[attr-defined]
|
||||
cls._stack.enter_context( # type: ignore[attr-defined]
|
||||
config.patch(
|
||||
enable_trace_unittest=True,
|
||||
),
|
||||
)
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def wrap_with_policy(self, method_name: str, policy: Callable) -> None:
|
||||
pass
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Functions used to test torch._dynamo.dont_skip_tracing.
|
||||
This file is located in torch/_dynamo so that it is skipped by trace rules.
|
||||
There is a special rule in trace_rules that doesn't skip this file when
|
||||
dont_skip_tracing is active.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def f1(x: torch.Tensor) -> torch.Tensor:
|
||||
return x + 1
|
||||
|
||||
|
||||
def f2(x: torch.Tensor) -> torch.Tensor:
|
||||
return x + 1
|
||||
|
||||
|
||||
def f3(x: torch.Tensor) -> torch.Tensor:
|
||||
return f2(x)
|
||||
|
||||
|
||||
def f4(x: torch.Tensor) -> torch.Tensor:
|
||||
x = f5(x, 1)
|
||||
x = torch._dynamo.dont_skip_tracing(f6)(x)
|
||||
x = f5(x, 8)
|
||||
return x
|
||||
|
||||
|
||||
def f5(x: torch.Tensor, n: int) -> torch.Tensor:
|
||||
if torch.compiler.is_compiling():
|
||||
return x + n
|
||||
return x
|
||||
|
||||
|
||||
def f6(x: torch.Tensor) -> torch.Tensor:
|
||||
x = f5(x, 2)
|
||||
torch._dynamo.graph_break()
|
||||
x = f5(x, 4)
|
||||
return x
|
||||
@@ -0,0 +1,323 @@
|
||||
"""Common utilities for testing Dynamo's minifier functionality.
|
||||
|
||||
This module provides the base infrastructure for running minification tests in Dynamo.
|
||||
It includes:
|
||||
- MinifierTestResult: A dataclass for storing and processing minifier test results
|
||||
- MinifierTestBase: A base test class with utilities for:
|
||||
- Running tests in isolated environments
|
||||
- Managing temporary directories and configurations
|
||||
- Executing minifier launcher scripts
|
||||
- Running and validating reproduction scripts
|
||||
- Supporting both compile-time and runtime error testing
|
||||
|
||||
The minifier helps reduce failing Dynamo compilations to minimal reproductions.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import traceback
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
import torch._dynamo
|
||||
import torch._dynamo.test_case
|
||||
from torch._dynamo.trace_rules import _as_posix_path
|
||||
from torch.utils._traceback import report_compile_source_on_error
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class MinifierTestResult:
|
||||
minifier_code: str
|
||||
repro_code: str
|
||||
|
||||
def _get_module(self, t: str) -> str:
|
||||
match = re.search(r"class Repro\(torch\.nn\.Module\):\s+([ ].*\n| *\n)+", t)
|
||||
assert match is not None, "failed to find module"
|
||||
r = match.group(0)
|
||||
r = re.sub(r"\s+$", "\n", r, flags=re.MULTILINE)
|
||||
r = re.sub(r"\n{3,}", "\n\n", r)
|
||||
return r.strip()
|
||||
|
||||
def get_exported_program_path(self) -> str | None:
|
||||
# Extract the exported program file path from AOTI minifier's repro.py
|
||||
# Regular expression pattern to match the file path
|
||||
pattern = r'torch\.export\.load\(\s*["\'](.*?)["\']\s*\)'
|
||||
# Search for the pattern in the text
|
||||
match = re.search(pattern, self.repro_code)
|
||||
# Extract and print the file path if a match is found
|
||||
if match:
|
||||
file_path = match.group(1)
|
||||
return file_path
|
||||
return None
|
||||
|
||||
def minifier_module(self) -> str:
|
||||
return self._get_module(self.minifier_code)
|
||||
|
||||
def repro_module(self) -> str:
|
||||
return self._get_module(self.repro_code)
|
||||
|
||||
|
||||
class MinifierTestBase(torch._dynamo.test_case.TestCase):
|
||||
DEBUG_DIR = tempfile.mkdtemp()
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
super().setUpClass()
|
||||
if not os.path.exists(cls.DEBUG_DIR):
|
||||
cls.DEBUG_DIR = tempfile.mkdtemp()
|
||||
cls._exit_stack.enter_context( # type: ignore[attr-defined]
|
||||
torch._dynamo.config.patch(debug_dir_root=cls.DEBUG_DIR)
|
||||
)
|
||||
# These configurations make new process startup slower. Disable them
|
||||
# for the minification tests to speed them up.
|
||||
cls._exit_stack.enter_context( # type: ignore[attr-defined]
|
||||
torch._inductor.config.patch(
|
||||
{
|
||||
# https://github.com/pytorch/pytorch/issues/100376
|
||||
"pattern_matcher": False,
|
||||
# multiprocess compilation takes a long time to warmup
|
||||
"compile_threads": 1,
|
||||
# https://github.com/pytorch/pytorch/issues/100378
|
||||
"cpp.vec_isa_ok": False,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
if os.getenv("PYTORCH_KEEP_TMPDIR", "0") != "1":
|
||||
shutil.rmtree(cls.DEBUG_DIR)
|
||||
else:
|
||||
print(f"test_minifier_common tmpdir kept at: {cls.DEBUG_DIR}")
|
||||
cls._exit_stack.close() # type: ignore[attr-defined]
|
||||
|
||||
def _gen_codegen_fn_patch_code(self, device: str, bug_type: str) -> str:
|
||||
assert bug_type in ("compile_error", "runtime_error", "accuracy")
|
||||
return f"""\
|
||||
{torch._dynamo.config.codegen_config()}
|
||||
{torch._inductor.config.codegen_config()}
|
||||
torch._inductor.config.{"cpp" if device == "cpu" else "triton"}.inject_relu_bug_TESTING_ONLY = {bug_type!r}
|
||||
"""
|
||||
|
||||
def _maybe_subprocess_run(
|
||||
self, args: Sequence[Any], *, isolate: bool, cwd: str | None = None
|
||||
) -> subprocess.CompletedProcess[bytes]:
|
||||
from torch._inductor.cpp_builder import normalize_path_separator
|
||||
|
||||
if not isolate:
|
||||
assert len(args) >= 2, args
|
||||
assert args[0] == "python3", args
|
||||
if args[1] == "-c":
|
||||
assert len(args) == 3, args
|
||||
code = args[2]
|
||||
args = ["-c"]
|
||||
else:
|
||||
assert len(args) >= 2, args
|
||||
with open(args[1]) as f:
|
||||
# Need normalize path of the code.
|
||||
code = normalize_path_separator(f.read())
|
||||
args = args[1:]
|
||||
|
||||
# WARNING: This is not a perfect simulation of running
|
||||
# the program out of tree. We only interpose on things we KNOW we
|
||||
# need to handle for tests. If you need more stuff, you will
|
||||
# need to augment this appropriately.
|
||||
|
||||
# NB: Can't use save_config because that will omit some fields,
|
||||
# but we must save and reset ALL fields
|
||||
dynamo_config = torch._dynamo.config.get_config_copy()
|
||||
inductor_config = torch._inductor.config.get_config_copy()
|
||||
try:
|
||||
stderr = io.StringIO()
|
||||
log_handler = logging.StreamHandler(stderr)
|
||||
log = logging.getLogger("torch._dynamo")
|
||||
log.addHandler(log_handler)
|
||||
try:
|
||||
prev_cwd = _as_posix_path(os.getcwd())
|
||||
if cwd is not None:
|
||||
cwd = _as_posix_path(cwd)
|
||||
os.chdir(cwd)
|
||||
with patch("sys.argv", args), report_compile_source_on_error():
|
||||
exec(code, {"__name__": "__main__", "__compile_source__": code})
|
||||
rc = 0
|
||||
except Exception:
|
||||
rc = 1
|
||||
traceback.print_exc(file=stderr)
|
||||
finally:
|
||||
log.removeHandler(log_handler)
|
||||
if cwd is not None:
|
||||
os.chdir(prev_cwd) # type: ignore[possibly-undefined]
|
||||
# Make sure we don't leave buggy compiled frames lying
|
||||
# around
|
||||
torch._dynamo.reset()
|
||||
finally:
|
||||
torch._dynamo.config.load_config(dynamo_config)
|
||||
torch._inductor.config.load_config(inductor_config)
|
||||
|
||||
# TODO: return a more appropriate data structure here
|
||||
return subprocess.CompletedProcess(
|
||||
args,
|
||||
rc,
|
||||
b"",
|
||||
stderr.getvalue().encode("utf-8"),
|
||||
)
|
||||
else:
|
||||
if cwd is not None:
|
||||
cwd = _as_posix_path(cwd)
|
||||
return subprocess.run(args, capture_output=True, cwd=cwd, check=False)
|
||||
|
||||
# Run `code` in a separate python process.
|
||||
# Returns the completed process state and the directory containing the
|
||||
# minifier launcher script, if `code` outputted it.
|
||||
def _run_test_code(
|
||||
self, code: str, *, isolate: bool
|
||||
) -> tuple[subprocess.CompletedProcess[bytes], str | Any]:
|
||||
proc = self._maybe_subprocess_run(
|
||||
["python3", "-c", code], isolate=isolate, cwd=self.DEBUG_DIR
|
||||
)
|
||||
|
||||
print("test stdout:", proc.stdout.decode("utf-8"))
|
||||
print("test stderr:", proc.stderr.decode("utf-8"))
|
||||
repro_dir_match = re.search(
|
||||
r"(\S+)minifier_launcher.py", proc.stderr.decode("utf-8")
|
||||
)
|
||||
if repro_dir_match is not None:
|
||||
return proc, repro_dir_match.group(1)
|
||||
return proc, None
|
||||
|
||||
# Runs the minifier launcher script in `repro_dir`
|
||||
def _run_minifier_launcher(
|
||||
self,
|
||||
repro_dir: str,
|
||||
isolate: bool,
|
||||
*,
|
||||
minifier_args: Sequence[Any] = (),
|
||||
repro_after: str | None = None,
|
||||
) -> tuple[subprocess.CompletedProcess[bytes], str]:
|
||||
self.assertIsNotNone(repro_dir)
|
||||
launch_file = _as_posix_path(os.path.join(repro_dir, "minifier_launcher.py"))
|
||||
with open(launch_file) as f:
|
||||
launch_code = f.read()
|
||||
|
||||
self.assertTrue(os.path.exists(launch_file))
|
||||
|
||||
args = ["python3", launch_file, "minify", *minifier_args]
|
||||
if not isolate and repro_after != "aot_inductor":
|
||||
# AOTI minifier doesn't have --no-isolate flag.
|
||||
# Everything in AOTI minifier is in no-isolate mode.
|
||||
args.append("--no-isolate")
|
||||
launch_proc = self._maybe_subprocess_run(args, isolate=isolate, cwd=repro_dir)
|
||||
print("minifier stdout:", launch_proc.stdout.decode("utf-8"))
|
||||
stderr = launch_proc.stderr.decode("utf-8")
|
||||
print("minifier stderr:", stderr)
|
||||
|
||||
self.assertNotIn("Input graph did not fail the tester", stderr)
|
||||
|
||||
return launch_proc, launch_code
|
||||
|
||||
# Runs the repro script in `repro_dir`
|
||||
def _run_repro(
|
||||
self, repro_dir: str, *, isolate: bool = True
|
||||
) -> tuple[subprocess.CompletedProcess[bytes], str]:
|
||||
self.assertIsNotNone(repro_dir)
|
||||
repro_file = _as_posix_path(os.path.join(repro_dir, "repro.py"))
|
||||
with open(repro_file) as f:
|
||||
repro_code = f.read()
|
||||
|
||||
self.assertTrue(os.path.exists(repro_file))
|
||||
|
||||
repro_proc = self._maybe_subprocess_run(
|
||||
["python3", repro_file], isolate=isolate, cwd=repro_dir
|
||||
)
|
||||
print("repro stdout:", repro_proc.stdout.decode("utf-8"))
|
||||
print("repro stderr:", repro_proc.stderr.decode("utf-8"))
|
||||
return repro_proc, repro_code
|
||||
|
||||
# Template for testing code.
|
||||
# `run_code` is the code to run for the test case.
|
||||
# `patch_code` is the code to be patched in every generated file; usually
|
||||
# just use this to turn on bugs via the config
|
||||
def _gen_test_code(self, run_code: str, repro_after: str, repro_level: int) -> str:
|
||||
repro_after_line = ""
|
||||
if repro_after == "aot_inductor":
|
||||
repro_after_line = (
|
||||
"torch._inductor.config.aot_inductor.dump_aoti_minifier = True"
|
||||
)
|
||||
elif repro_after:
|
||||
repro_after_line = f"""\
|
||||
torch._dynamo.config.repro_after = "{repro_after}"
|
||||
"""
|
||||
return f"""\
|
||||
import torch
|
||||
import torch._dynamo
|
||||
import torch._inductor
|
||||
{_as_posix_path(torch._dynamo.config.codegen_config())}
|
||||
{_as_posix_path(torch._inductor.config.codegen_config())}
|
||||
{repro_after_line}
|
||||
torch._dynamo.config.repro_level = {repro_level}
|
||||
torch._inductor.config.aot_inductor.repro_level = {repro_level}
|
||||
torch._dynamo.config.debug_dir_root = "{_as_posix_path(self.DEBUG_DIR)}"
|
||||
{run_code}
|
||||
"""
|
||||
|
||||
# Runs a full minifier test.
|
||||
# Minifier tests generally consist of 3 stages:
|
||||
# 1. Run the problematic code
|
||||
# 2. Run the generated minifier launcher script
|
||||
# 3. Run the generated repro script
|
||||
#
|
||||
# If possible, you should run the test with isolate=False; use
|
||||
# isolate=True only if the bug you're testing would otherwise
|
||||
# crash the process
|
||||
def _run_full_test(
|
||||
self,
|
||||
run_code: str,
|
||||
repro_after: str,
|
||||
expected_error: str | None,
|
||||
*,
|
||||
isolate: bool,
|
||||
minifier_args: Sequence[Any] = (),
|
||||
) -> MinifierTestResult | None:
|
||||
if isolate:
|
||||
repro_level = 3
|
||||
elif expected_error is None or expected_error == "AccuracyError":
|
||||
repro_level = 4
|
||||
else:
|
||||
repro_level = 2
|
||||
test_code = self._gen_test_code(run_code, repro_after, repro_level)
|
||||
print("running test", file=sys.stderr)
|
||||
test_proc, repro_dir = self._run_test_code(test_code, isolate=isolate)
|
||||
if expected_error is None:
|
||||
# Just check that there was no error
|
||||
self.assertEqual(test_proc.returncode, 0)
|
||||
|
||||
self.assertIsNone(repro_dir)
|
||||
return None
|
||||
# NB: Intentionally do not test return code; we only care about
|
||||
# actually generating the repro, we don't have to crash
|
||||
|
||||
self.assertIn(expected_error, test_proc.stderr.decode("utf-8"))
|
||||
|
||||
self.assertIsNotNone(repro_dir)
|
||||
print("running minifier", file=sys.stderr)
|
||||
_minifier_proc, minifier_code = self._run_minifier_launcher(
|
||||
repro_dir,
|
||||
isolate=isolate,
|
||||
minifier_args=minifier_args,
|
||||
repro_after=repro_after,
|
||||
)
|
||||
print("running repro", file=sys.stderr)
|
||||
repro_proc, repro_code = self._run_repro(repro_dir, isolate=isolate)
|
||||
|
||||
self.assertIn(expected_error, repro_proc.stderr.decode("utf-8"))
|
||||
self.assertNotEqual(repro_proc.returncode, 0)
|
||||
return MinifierTestResult(minifier_code=minifier_code, repro_code=repro_code)
|
||||
@@ -0,0 +1,612 @@
|
||||
"""Testing utilities and infrastructure for Dynamo.
|
||||
|
||||
This module provides a comprehensive set of testing utilities including:
|
||||
- Test result collection and validation
|
||||
- Graph manipulation and comparison tools
|
||||
- Test case management and execution helpers
|
||||
- Specialized test decorators for different Python versions and features
|
||||
- RNG state management
|
||||
- Compilation counting and monitoring
|
||||
- Debug utilities for bytecode transformation
|
||||
|
||||
The utilities in this module are used across Dynamo's test suite to ensure
|
||||
consistent testing patterns and proper test isolation.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import dis
|
||||
import functools
|
||||
import logging
|
||||
import os.path
|
||||
import random
|
||||
import re
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from collections.abc import Callable, Generator, Sequence
|
||||
from typing import Any, overload, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from torch import fx
|
||||
from torch._dynamo.backends.debugging import aot_eager
|
||||
from torch._dynamo.output_graph import OutputGraph
|
||||
|
||||
from . import config, eval_frame, optimize_assert, reset
|
||||
from .bytecode_transformation import (
|
||||
create_instruction,
|
||||
debug_checks,
|
||||
is_generator,
|
||||
transform_code_object,
|
||||
)
|
||||
from .guards import CheckFunctionManager, CompileId, GuardedCode
|
||||
from .types import ConvertFrameReturn, DynamoFrameType, wrap_guarded_code
|
||||
from .utils import CompileCounterInt, same
|
||||
|
||||
|
||||
np: types.ModuleType | None = None
|
||||
try:
|
||||
import numpy as np
|
||||
except ModuleNotFoundError:
|
||||
np = None
|
||||
|
||||
|
||||
unsupported = eval_frame.unsupported
|
||||
three = 3
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
def clone_me(x: torch.Tensor | None) -> torch.Tensor | None:
|
||||
if x is None:
|
||||
return None
|
||||
return x.detach().clone().requires_grad_(x.requires_grad)
|
||||
|
||||
|
||||
def remove_optimized_module_prefix(name: str) -> str:
|
||||
return re.sub(r"^_orig_mod[.]", "", name)
|
||||
|
||||
|
||||
def extract_graph_and_tracker(fn, *args, **kwargs): # type: ignore[no-untyped-def]
|
||||
from torch._dynamo.symbolic_convert import InstructionTranslator
|
||||
|
||||
gm = None
|
||||
region_tracker = None
|
||||
|
||||
def extract_graph_backend(_gm, *args, **kwargs): # type: ignore[no-untyped-def]
|
||||
nonlocal gm
|
||||
nonlocal region_tracker
|
||||
gm = _gm
|
||||
region_tracker = InstructionTranslator.current_tx().output.region_tracker
|
||||
return _gm
|
||||
|
||||
torch.compile(backend=extract_graph_backend, fullgraph=True)(fn)(*args, **kwargs)
|
||||
return gm.graph, region_tracker # type: ignore[union-attr]
|
||||
|
||||
|
||||
def extract_graph(fn, *args, **kwargs): # type: ignore[no-untyped-def]
|
||||
backend = AotEagerAndRecordGraphs()
|
||||
result = torch.compile(backend=backend)(fn)(*args, **kwargs)
|
||||
return result, backend.graphs, backend.fw_graphs, backend.bw_graphs
|
||||
|
||||
|
||||
def collect_results(
|
||||
model: torch.nn.Module, prediction: Any, loss: Any, example_inputs: Any
|
||||
) -> list[Any]:
|
||||
results = []
|
||||
results.append(prediction)
|
||||
results.append(loss)
|
||||
# if isinstance(loss, torch.Tensor) and loss.item() > 1:
|
||||
# log.warning(
|
||||
# f"High loss value alert - {loss:.2f}. Can result in unstable gradients."
|
||||
# )
|
||||
|
||||
grads = {}
|
||||
params = {}
|
||||
for name, param in model.named_parameters():
|
||||
if isinstance(model, eval_frame.OptimizedModule):
|
||||
name = remove_optimized_module_prefix(name)
|
||||
param_copy = param
|
||||
grad = param.grad
|
||||
# Treat None and zero grad as same
|
||||
if param.grad is None:
|
||||
grad = torch.zeros_like(param)
|
||||
grads[name + ".grad"] = grad
|
||||
params[name] = param_copy
|
||||
results.append(grads)
|
||||
results.append(params)
|
||||
buffers = {}
|
||||
for name, buffer in model.named_buffers():
|
||||
if isinstance(model, eval_frame.OptimizedModule):
|
||||
name = remove_optimized_module_prefix(name)
|
||||
buffers[name] = buffer
|
||||
results.append(buffers)
|
||||
for example in example_inputs:
|
||||
if isinstance(example, (tuple, list)):
|
||||
results.extend(inp.grad for inp in example if isinstance(inp, torch.Tensor))
|
||||
else:
|
||||
if isinstance(example, torch.Tensor):
|
||||
results.append(example.grad)
|
||||
return results
|
||||
|
||||
|
||||
def requires_bwd_pass(out: Any) -> bool:
|
||||
if isinstance(out, torch.Tensor):
|
||||
return out.requires_grad
|
||||
elif isinstance(out, (list, tuple)):
|
||||
return any(requires_bwd_pass(x) for x in out)
|
||||
elif out is None:
|
||||
return False
|
||||
elif isinstance(out, int):
|
||||
return False
|
||||
raise NotImplementedError("Don't know how to reduce", type(out))
|
||||
|
||||
|
||||
@overload
|
||||
def reduce_to_scalar_loss(out: torch.Tensor) -> torch.Tensor: ...
|
||||
|
||||
|
||||
@overload
|
||||
def reduce_to_scalar_loss(
|
||||
out: list[Any] | tuple[Any, ...] | dict[Any, Any],
|
||||
) -> float: ...
|
||||
|
||||
|
||||
def reduce_to_scalar_loss(out: Any) -> torch.Tensor | float:
|
||||
"""Reduce the output of a model to get scalar loss"""
|
||||
if isinstance(out, torch.Tensor):
|
||||
# Mean does not work on integer tensors
|
||||
return out.sum() / out.numel()
|
||||
elif isinstance(out, (list, tuple)):
|
||||
return sum(reduce_to_scalar_loss(x) for x in out) / len(out)
|
||||
elif type(out).__name__ in (
|
||||
"MaskedLMOutput",
|
||||
"Seq2SeqLMOutput",
|
||||
"CausalLMOutputWithCrossAttentions",
|
||||
):
|
||||
return reduce_to_scalar_loss(out.logits)
|
||||
elif type(out).__name__ == "SquashedNormal":
|
||||
return out.mean.sum()
|
||||
elif isinstance(out, dict):
|
||||
return sum(reduce_to_scalar_loss(value) for value in out.values()) / len(
|
||||
out.keys()
|
||||
)
|
||||
raise NotImplementedError("Don't know how to reduce", type(out))
|
||||
|
||||
|
||||
def debug_dir() -> str:
|
||||
path = os.path.join(os.path.dirname(__file__), "../debug")
|
||||
if not os.path.exists(path):
|
||||
os.mkdir(path)
|
||||
return path
|
||||
|
||||
|
||||
def debug_dump(name: str, code: types.CodeType, extra: str = "") -> None:
|
||||
with open(os.path.join(debug_dir(), name), "w") as fd:
|
||||
fd.write(
|
||||
f"{dis.Bytecode(code).info()}\n\n{dis.Bytecode(code).dis()}\n\n{extra}\n"
|
||||
)
|
||||
|
||||
|
||||
def debug_insert_nops(
|
||||
frame: DynamoFrameType, cache_size: int, hooks: Any, _: Any, *, skip: int = 0
|
||||
) -> ConvertFrameReturn:
|
||||
"""used to debug jump updates"""
|
||||
|
||||
def insert_nops(instructions: list[Any], code_options: Any) -> None:
|
||||
instructions.insert(0, create_instruction("NOP"))
|
||||
instructions.insert(0, create_instruction("NOP"))
|
||||
|
||||
metrics_context = torch._dynamo.utils.get_metrics_context()
|
||||
with torch._dynamo.utils.dynamo_timed("debug_insert_nops"), metrics_context:
|
||||
if is_generator(frame.f_code):
|
||||
return ConvertFrameReturn()
|
||||
|
||||
debug_checks(frame.f_code)
|
||||
code, _ = transform_code_object(frame.f_code, insert_nops)
|
||||
graph = OutputGraph(
|
||||
code_options={},
|
||||
compiler_fn=None,
|
||||
root_tx=None, # type: ignore[arg-type]
|
||||
export=False,
|
||||
export_constraints=[],
|
||||
frame_state={"_id": 0},
|
||||
# TODO: shouldn't this be f_locals/f_globals from frame?
|
||||
local_scope=locals(),
|
||||
global_scope=globals(),
|
||||
f_code=frame.f_code,
|
||||
torch_function_mode_stack=[],
|
||||
package=None,
|
||||
)
|
||||
|
||||
return wrap_guarded_code(
|
||||
GuardedCode(
|
||||
code,
|
||||
CheckFunctionManager(frame.f_code, graph).guard_manager, # type: ignore[arg-type]
|
||||
CompileId(frame_id=0, frame_compile_id=0),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class CompileCounter:
|
||||
def __init__(self) -> None:
|
||||
self.frame_count: int | CompileCounterInt = 0
|
||||
self.clear()
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
self.frame_count += 1
|
||||
for node in gm.graph.nodes:
|
||||
if "call" in node.op:
|
||||
self.op_count += 1
|
||||
return gm.forward
|
||||
|
||||
def clear(self) -> None:
|
||||
if config.debug_disable_compile_counter:
|
||||
self.frame_count = CompileCounterInt(0)
|
||||
else:
|
||||
self.frame_count = 0
|
||||
self.op_count = 0
|
||||
|
||||
|
||||
class CompileCounterWithBackend:
|
||||
def __init__(self, backend: str) -> None:
|
||||
self.frame_count: int | CompileCounterInt = 0
|
||||
self.backend = backend
|
||||
self.graphs: list[torch.fx.GraphModule] = []
|
||||
self.clear()
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
from .backends.registry import lookup_backend
|
||||
|
||||
self.frame_count += 1
|
||||
for node in gm.graph.nodes:
|
||||
if "call" in node.op:
|
||||
self.op_count += 1
|
||||
self.graphs.append(gm)
|
||||
return lookup_backend(self.backend)(gm, example_inputs)
|
||||
|
||||
def clear(self) -> None:
|
||||
if config.debug_disable_compile_counter:
|
||||
self.frame_count = CompileCounterInt(0)
|
||||
else:
|
||||
self.frame_count = 0
|
||||
self.op_count = 0
|
||||
self.graphs = []
|
||||
|
||||
|
||||
# Equivalent to backend="eager", but also records graphs that
|
||||
# we can assert on
|
||||
class EagerAndRecordGraphs:
|
||||
def __init__(self) -> None:
|
||||
self.graphs: list[torch.fx.GraphModule] = []
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
self.graphs.append(gm)
|
||||
return gm.forward
|
||||
|
||||
|
||||
class AotEagerAndRecordGraphs:
|
||||
def __init__(self) -> None:
|
||||
self.graphs: list[torch.fx.GraphModule] = []
|
||||
self.fw_graphs: list[torch.fx.GraphModule] = []
|
||||
self.bw_graphs: list[torch.fx.GraphModule] = []
|
||||
|
||||
def __call__(
|
||||
self, gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
self.graphs.append(gm)
|
||||
|
||||
def fw_compiler(
|
||||
gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
self.fw_graphs.append(gm)
|
||||
return gm.forward
|
||||
|
||||
def bw_compiler(
|
||||
gm: torch.fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
self.bw_graphs.append(gm)
|
||||
return gm.forward
|
||||
|
||||
return aot_eager(
|
||||
gm,
|
||||
example_inputs,
|
||||
fw_compiler=fw_compiler,
|
||||
bw_compiler=bw_compiler,
|
||||
)
|
||||
|
||||
|
||||
class InductorAndRecordGraphs:
|
||||
def __init__(self) -> None:
|
||||
self.graphs: list[torch.fx.GraphModule] = []
|
||||
self.inductor_graphs: list[torch.fx.GraphModule] = []
|
||||
|
||||
def __call__(self, gm, example_inputs): # type: ignore[no-untyped-def]
|
||||
import torch._inductor.compile_fx as compile_fx_mod
|
||||
|
||||
self.graphs.append(gm)
|
||||
|
||||
old_compile_fx_inner = compile_fx_mod._compile_fx_inner
|
||||
|
||||
def patched(*args, **kwargs): # type: ignore[no-untyped-def]
|
||||
self.inductor_graphs.append(args[0])
|
||||
return old_compile_fx_inner(*args, **kwargs)
|
||||
|
||||
with patch.object(compile_fx_mod, "_compile_fx_inner", new=patched):
|
||||
return compile_fx_mod.compile_fx(gm, example_inputs)
|
||||
|
||||
|
||||
def strip_comment(code: str) -> str:
|
||||
return re.sub(r"(?m)^ *#.*\n?", "", code)
|
||||
|
||||
|
||||
def remove_trailing_space(code: str) -> str:
|
||||
return "\n".join([line.rstrip() for line in code.split("\n")])
|
||||
|
||||
|
||||
def _squash_blank_lines(code: str) -> str:
|
||||
lines = code.split("\n")
|
||||
result: list[str] = []
|
||||
saw_blank = False
|
||||
for line in lines:
|
||||
if line.strip() == "":
|
||||
if saw_blank:
|
||||
continue
|
||||
saw_blank = True
|
||||
else:
|
||||
saw_blank = False
|
||||
result.append(line)
|
||||
return "\n".join(result)
|
||||
|
||||
|
||||
def normalize_gm(gm_str: str) -> str:
|
||||
# strip comments as comments have path to files which may differ from
|
||||
# system to system.
|
||||
stripped = strip_comment(gm_str)
|
||||
no_trailing = remove_trailing_space(stripped)
|
||||
return _squash_blank_lines(no_trailing)
|
||||
|
||||
|
||||
def empty_line_normalizer(code: str) -> str:
|
||||
"""
|
||||
Normalize code: remove empty lines.
|
||||
"""
|
||||
normal_code = re.sub(r"[\r\n]+", "\n", code)
|
||||
return normal_code
|
||||
|
||||
|
||||
def standard_test(
|
||||
self: Any,
|
||||
fn: Callable[..., Any],
|
||||
nargs: int,
|
||||
expected_ops: int | None = None,
|
||||
expected_ops_dynamic: int | None = None,
|
||||
expected_frame_count: int = 1,
|
||||
) -> None:
|
||||
if not config.assume_static_by_default and expected_ops_dynamic is not None:
|
||||
expected_ops = expected_ops_dynamic
|
||||
|
||||
actual = CompileCounter()
|
||||
|
||||
args1 = [torch.randn(10, 10) for _ in range(nargs)]
|
||||
args2 = [torch.randn(10, 10) for _ in range(nargs)]
|
||||
correct1 = fn(*args1)
|
||||
correct2 = fn(*args2)
|
||||
reset()
|
||||
opt_fn = optimize_assert(actual)(fn)
|
||||
val1a = opt_fn(*args1)
|
||||
val2a = opt_fn(*args2)
|
||||
val1b = opt_fn(*args1)
|
||||
val2b = opt_fn(*args2)
|
||||
reset()
|
||||
self.assertTrue(same(val1a, correct1))
|
||||
self.assertTrue(same(val1b, correct1))
|
||||
self.assertTrue(same(val2a, correct2))
|
||||
self.assertTrue(same(val2b, correct2))
|
||||
self.assertEqual(actual.frame_count, expected_frame_count)
|
||||
if expected_ops is not None:
|
||||
self.assertEqual(actual.op_count, expected_ops)
|
||||
|
||||
|
||||
def dummy_fx_compile(
|
||||
gm: fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> Callable[..., Any]:
|
||||
return gm.forward
|
||||
|
||||
|
||||
def format_speedup(
|
||||
speedup: float,
|
||||
pvalue: float,
|
||||
is_correct: bool = True,
|
||||
pvalue_threshold: float = 0.1,
|
||||
) -> str:
|
||||
if not is_correct:
|
||||
return "ERROR"
|
||||
if pvalue > pvalue_threshold:
|
||||
return f"{speedup:.3f}x SAME"
|
||||
return f"{speedup:.3f}x p={pvalue:.2f}"
|
||||
|
||||
|
||||
def rand_strided(
|
||||
size: Sequence[int],
|
||||
stride: Sequence[int],
|
||||
dtype: torch.dtype = torch.float32,
|
||||
device: str | torch.device = "cpu",
|
||||
extra_size: int = 0,
|
||||
) -> torch.Tensor:
|
||||
needed_size = extra_size
|
||||
if all(s > 0 for s in size):
|
||||
# only need to allocate if all sizes are non-zero
|
||||
needed_size += (
|
||||
sum((shape - 1) * stride for shape, stride in zip(size, stride)) + 1
|
||||
)
|
||||
if dtype.is_floating_point:
|
||||
if dtype == torch.float4_e2m1fn_x2:
|
||||
buffer = torch.randint(
|
||||
0, 256, (needed_size,), dtype=torch.uint8, device=device
|
||||
).view(torch.float4_e2m1fn_x2)
|
||||
elif dtype.itemsize == 1:
|
||||
"""
|
||||
normal distribution kernel is not implemented for fp8..
|
||||
Workaround that by creating a fp16 tensor and then cast.
|
||||
"""
|
||||
buffer = torch.randn(needed_size, dtype=torch.float16, device=device).to(
|
||||
dtype=dtype
|
||||
)
|
||||
else:
|
||||
buffer = torch.randn(needed_size, dtype=dtype, device=device)
|
||||
else:
|
||||
buffer = torch.zeros(size=[needed_size], dtype=dtype, device=device)
|
||||
return torch.as_strided(buffer, size, stride)
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def check_dynamic_shape_capture() -> bool:
|
||||
# This also mirrors config from `test/dynamo/test_dynamic_shapes.py:make_dynamic_cls`
|
||||
return not config.assume_static_by_default
|
||||
|
||||
|
||||
def _make_fn_with_patches(fn: Callable[_P, _T], *patches: Any) -> Callable[_P, _T]:
|
||||
@functools.wraps(fn)
|
||||
def _fn(*args: _P.args, **kwargs: _P.kwargs) -> _T:
|
||||
with contextlib.ExitStack() as stack:
|
||||
for module, attr, val in patches:
|
||||
stack.enter_context(patch.object(module, attr, val))
|
||||
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return _fn
|
||||
|
||||
|
||||
def make_test_cls_with_patches(
|
||||
cls: type,
|
||||
cls_prefix: str,
|
||||
fn_suffix: str,
|
||||
*patches: Any,
|
||||
xfail_prop: str | None = None,
|
||||
decorator: Callable[[Callable[..., Any]], Callable[..., Any]] = lambda x: x,
|
||||
) -> type:
|
||||
DummyTestClass = type(f"{cls_prefix}{cls.__name__}", cls.__bases__, {})
|
||||
DummyTestClass.__qualname__ = DummyTestClass.__name__
|
||||
|
||||
for name in dir(cls):
|
||||
if name.startswith("test_"):
|
||||
fn = getattr(cls, name)
|
||||
if not callable(fn):
|
||||
setattr(DummyTestClass, name, getattr(cls, name))
|
||||
continue
|
||||
new_name = f"{name}{fn_suffix}"
|
||||
new_fn = _make_fn_with_patches(fn, *patches)
|
||||
new_fn.__name__ = new_name
|
||||
if xfail_prop is not None and hasattr(fn, xfail_prop):
|
||||
new_fn = unittest.expectedFailure(new_fn)
|
||||
setattr(DummyTestClass, new_name, decorator(new_fn))
|
||||
# NB: Doesn't handle slots correctly, but whatever
|
||||
elif not hasattr(DummyTestClass, name):
|
||||
setattr(DummyTestClass, name, getattr(cls, name))
|
||||
|
||||
return DummyTestClass
|
||||
|
||||
|
||||
# test Python 3.11+ specific features
|
||||
def skipIfNotPy311(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
if sys.version_info >= (3, 11):
|
||||
return fn
|
||||
# pyrefly: ignore [bad-return, bad-argument-type]
|
||||
return unittest.skip(fn)
|
||||
|
||||
|
||||
def skipIfNotPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
if sys.version_info >= (3, 12):
|
||||
return fn
|
||||
return unittest.skip("Requires Python 3.12+")(fn)
|
||||
|
||||
|
||||
def skipIfOnlyNotPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
if sys.version_info >= (3, 13) or sys.version_info < (3, 12):
|
||||
return unittest.skip("Requires Python 3.12")(fn)
|
||||
return fn
|
||||
|
||||
|
||||
def xfailIfPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
if sys.version_info >= (3, 12):
|
||||
return unittest.expectedFailure(fn)
|
||||
return fn
|
||||
|
||||
|
||||
def skipIfPy312(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
if sys.version_info >= (3, 12):
|
||||
return unittest.skip("Not supported in Python 3.12+")(fn)
|
||||
return fn
|
||||
|
||||
|
||||
# Controls tests generated in test/inductor/test_torchinductor_dynamic_shapes.py
|
||||
# and test/dynamo/test_dynamic_shapes.py
|
||||
def expectedFailureDynamic(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
fn._expected_failure_dynamic = True # type: ignore[attr-defined]
|
||||
return fn
|
||||
|
||||
|
||||
# Controls tests generated in test/inductor/test_torchinductor_codegen_dynamic_shapes.py
|
||||
def expectedFailureCodegenDynamic(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
fn._expected_failure_codegen_dynamic = True # type: ignore[attr-defined]
|
||||
return fn
|
||||
|
||||
|
||||
# Controls test generated in test/inductor/test_cpp_wrapper.py
|
||||
def expectedFailureDynamicWrapper(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
fn._expected_failure_dynamic_wrapper = True # type: ignore[attr-defined]
|
||||
return fn
|
||||
|
||||
|
||||
def reset_rng_state(use_xla: bool = False) -> None:
|
||||
torch.manual_seed(1337)
|
||||
random.seed(1337)
|
||||
if np:
|
||||
np.random.seed(1337)
|
||||
if use_xla:
|
||||
import torch_xla.core.xla_model as xm
|
||||
|
||||
xm.set_rng_state(1337, str(xm.xla_device()))
|
||||
|
||||
|
||||
def _skipped_function_for_test_reconstruct(
|
||||
f: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
|
||||
) -> _T:
|
||||
return f(*args, **kwargs)
|
||||
|
||||
|
||||
_testing_invoke_subgraph_inductor_compile_captured_gms = None
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _testing_capture_invoke_subgraph_inductor_compile_gms() -> Generator[
|
||||
list[torch.fx.GraphModule]
|
||||
]:
|
||||
"""
|
||||
Context manager to capture graph modules compiled by invoke_subgraph_inductor_compile.
|
||||
|
||||
Usage:
|
||||
with _testing_capture_invoke_subgraph_inductor_compile_gms() as captured_gms:
|
||||
# code that triggers invoke_subgraph_inductor_compile
|
||||
pass
|
||||
# captured_gms will contain the list of captured graph modules
|
||||
"""
|
||||
global _testing_invoke_subgraph_inductor_compile_captured_gms
|
||||
# pyrefly: ignore [implicit-any]
|
||||
_testing_invoke_subgraph_inductor_compile_captured_gms = []
|
||||
try:
|
||||
yield _testing_invoke_subgraph_inductor_compile_captured_gms
|
||||
finally:
|
||||
_testing_invoke_subgraph_inductor_compile_captured_gms = None
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,140 @@
|
||||
"""This module contains the core type definitions and protocols used throughout Dynamo.
|
||||
|
||||
The types defined here fall into several categories:
|
||||
- Guard related types (GuardFn, GuardFail, GuardedCode): Used for tracking and managing guards that protect compiled code
|
||||
- Frame and cache types (FrameState, CacheEntry): Used for managing interpreter frame state and caching
|
||||
- Callback protocols (DynamoCallbackFn): Define the interface for frame evaluation callbacks
|
||||
- Hook protocols (DynamoGuardHook, ProfilerStartHook, ProfilerEndHook, BytecodeHook): Define various hook points for
|
||||
instrumentation and customization
|
||||
|
||||
These types provide the foundational interfaces that enable Dynamo's dynamic compilation and optimization system,
|
||||
ensuring type safety and clear contracts between different components of the system.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import types
|
||||
from collections.abc import Callable
|
||||
from typing import Any, NamedTuple, Protocol
|
||||
|
||||
# CacheEntry has a `guard_manager` field for the guard, and a `code` field for the code object.
|
||||
from torch._C._dynamo.eval_frame import (
|
||||
_CacheEntry as CacheEntry,
|
||||
_ExtraState as ExtraState,
|
||||
_FrameAction as FrameAction,
|
||||
_FrameExecStrategy as FrameExecStrategy,
|
||||
_PyInterpreterFrame as DynamoFrameType,
|
||||
)
|
||||
from torch._guards import CompileId, Guard
|
||||
|
||||
|
||||
# We use a dict to store additional data per frame.
|
||||
FrameState = dict[Any, Any]
|
||||
|
||||
|
||||
class GuardFail(NamedTuple):
|
||||
# A string repr of the piece of failed guard code we eval-ed
|
||||
reason: str
|
||||
# A code object where we failed a guard
|
||||
orig_code: types.CodeType
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class GuardFilterEntry:
|
||||
name: str
|
||||
has_value: bool
|
||||
value: object
|
||||
guard_type: str
|
||||
derived_guard_types: tuple[str, ...]
|
||||
is_global: bool
|
||||
orig_guard: Guard
|
||||
|
||||
|
||||
class GuardFn(Protocol):
|
||||
closure_vars: dict[str, object]
|
||||
args: list[str]
|
||||
code_parts: list[str]
|
||||
verbose_code_parts: list[str]
|
||||
global_scope: dict[str, object]
|
||||
guard_fail_fn: Callable[[GuardFail], None] | None
|
||||
cache_entry: CacheEntry | None
|
||||
extra_state: ExtraState | None
|
||||
|
||||
# maps locals of user function to bool
|
||||
def __call__(self, f_locals: dict[str, object]) -> bool: ...
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class GuardedCode:
|
||||
code: types.CodeType
|
||||
guard_manager: GuardFn
|
||||
compile_id: CompileId
|
||||
trace_annotation: str = "Unknown"
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ConvertFrameReturn:
|
||||
# default return is no compiled code (i.e. `return None`):
|
||||
# strategy is to skip non-recursively, for all future intercepted frames too
|
||||
|
||||
# eval frame execution strategy for this frame
|
||||
frame_exec_strategy: FrameExecStrategy = dataclasses.field(
|
||||
default_factory=lambda: FrameExecStrategy(FrameAction.SKIP, FrameAction.DEFAULT)
|
||||
)
|
||||
# also apply frame_exec strategy to future frames with same code
|
||||
apply_to_code: bool = True
|
||||
guarded_code: GuardedCode | None = None
|
||||
|
||||
|
||||
def wrap_guarded_code(guarded_code: GuardedCode) -> ConvertFrameReturn:
|
||||
return ConvertFrameReturn(
|
||||
frame_exec_strategy=FrameExecStrategy(FrameAction.DEFAULT, FrameAction.DEFAULT),
|
||||
guarded_code=guarded_code,
|
||||
)
|
||||
|
||||
|
||||
class DynamoCallbackFn(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
frame: DynamoFrameType,
|
||||
cache_entry: CacheEntry | None,
|
||||
frame_state: FrameState,
|
||||
) -> ConvertFrameReturn: ...
|
||||
|
||||
|
||||
DynamoCallback = DynamoCallbackFn | None | bool
|
||||
|
||||
|
||||
class DynamoGuardHook(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
guard_manager: GuardFn,
|
||||
code: types.CodeType,
|
||||
f_locals: dict[str, object],
|
||||
index: int,
|
||||
last: bool,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class DynamoGuardCompleteHook(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
cache_hit: bool,
|
||||
) -> bool: ...
|
||||
|
||||
|
||||
class ProfilerStartHook(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
name: str,
|
||||
# TODO(whc) how do I annotate a _RecordFunction here?
|
||||
) -> Any: ...
|
||||
|
||||
|
||||
class ProfilerEndHook(Protocol):
|
||||
def __call__(self, record: Any) -> None: ...
|
||||
|
||||
|
||||
class BytecodeHook(Protocol):
|
||||
def __call__(
|
||||
self, code: types.CodeType, new_code: types.CodeType
|
||||
) -> types.CodeType | None: ...
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,263 @@
|
||||
"""
|
||||
This package implements variable tracking and symbolic execution capabilities for Dynamo,
|
||||
which are essential for converting Python code into FX graphs. It provides a comprehensive
|
||||
set of variable types that handle different Python constructs during tracing.
|
||||
|
||||
Each variable type (like BuiltinVariable, TensorVariable, NNModuleVariable, etc.) is responsible
|
||||
for tracking and symbolically executing operations on specific Python objects. This enables
|
||||
Dynamo to:
|
||||
- Track the flow of values through Python code
|
||||
- Maintain correct semantics during graph conversion
|
||||
- Handle complex Python features like context managers, iterators, and custom objects
|
||||
- Support both eager and symbolic execution modes
|
||||
|
||||
The VariableTracker base class provides the foundation for all variable types, with each
|
||||
subclass implementing specific behavior for different Python constructs. This modular design
|
||||
allows Dynamo to accurately trace and optimize Python code while preserving its semantics.
|
||||
"""
|
||||
|
||||
from .base import VariableTracker
|
||||
from .builtin import (
|
||||
BaseBuiltinVariable,
|
||||
BuiltinVariable,
|
||||
DictBuiltinVariable,
|
||||
IterBuiltinVariable,
|
||||
ListBuiltinVariable,
|
||||
)
|
||||
from .constant import ConstantVariable
|
||||
from .ctx_manager import (
|
||||
CatchWarningsCtxManagerVariable,
|
||||
ContextWrappingVariable,
|
||||
CUDADeviceVariable,
|
||||
CudagraphOverrideVariable,
|
||||
DisabledSavedTensorsHooksVariable,
|
||||
DualLevelContextManager,
|
||||
DynamoConfigPatchVariable,
|
||||
ErrorOnGraphBreakVariable,
|
||||
FSDPParamGroupUseTrainingStateVariable,
|
||||
FxTracebackAnnotateVariable,
|
||||
GenericContextWrappingVariable,
|
||||
GradIncrementNestingCtxManagerVariable,
|
||||
GradInplaceRequiresGradCtxManagerVariable,
|
||||
GradModeVariable,
|
||||
InferenceModeVariable,
|
||||
JvpIncrementNestingCtxManagerVariable,
|
||||
SDPAKernelVariable,
|
||||
SetFwdGradEnabledContextManager,
|
||||
TemporarilyPopInterpreterStackCtxManagerVariable,
|
||||
VmapIncrementNestingCtxManagerVariable,
|
||||
WithEnterFunctionVariable,
|
||||
WithExitFunctionVariable,
|
||||
)
|
||||
from .dicts import (
|
||||
ConstDictVariable,
|
||||
DefaultDictVariable,
|
||||
DictItemsVariable,
|
||||
DunderDictVariable,
|
||||
MappingProxyVariable,
|
||||
NNModuleHooksDictVariable,
|
||||
)
|
||||
from .distributed import BackwardHookVariable, DistributedVariable
|
||||
from .functions import (
|
||||
BaseUserFunctionVariable,
|
||||
BuiltinMethodVariable,
|
||||
CollectionsNamedTupleFunction,
|
||||
CreateTMADescriptorExperimentalVariable,
|
||||
CreateTMADescriptorStableVariable,
|
||||
FunctionDecoratedByContextlibContextManagerVariable,
|
||||
FunctoolsPartialVariable,
|
||||
InspectSignatureVariable,
|
||||
LocalGeneratorFunctionVariable,
|
||||
LocalGeneratorObjectVariable,
|
||||
NestedUserFunctionVariable,
|
||||
PolyfilledFunctionVariable,
|
||||
PyTreeGetNodeTypeFunctionVariable,
|
||||
PyTreeTreeIsLeafFunctionVariable,
|
||||
SkipFunctionVariable,
|
||||
SparseTensorCreationSkipVariable,
|
||||
TMADescriptorExperimentalVariable,
|
||||
TMADescriptorStableVariable,
|
||||
TritonSetAllocatorVariable,
|
||||
UserFunctionVariable,
|
||||
UserMethodVariable,
|
||||
WrapperUserFunctionVariable,
|
||||
WrapperUserMethodVariable,
|
||||
)
|
||||
from .higher_order_ops import (
|
||||
FunctionalCallVariable,
|
||||
FunctorchHigherOrderVariable,
|
||||
ReparametrizeModuleCallVariable,
|
||||
TorchHigherOrderOperatorVariable,
|
||||
)
|
||||
from .iter import (
|
||||
CountIteratorVariable,
|
||||
FilterVariable,
|
||||
IteratorVariable,
|
||||
ItertoolsVariable,
|
||||
MapVariable,
|
||||
ObjectIteratorVariable,
|
||||
RepeatIteratorVariable,
|
||||
ZipVariable,
|
||||
)
|
||||
from .lazy import LazyConstantVariable, LazyVariableTracker
|
||||
from .lists import (
|
||||
BaseListVariable,
|
||||
ListIteratorVariable,
|
||||
ListVariable,
|
||||
RangeVariable,
|
||||
SliceVariable,
|
||||
TupleIteratorVariable,
|
||||
TupleVariable,
|
||||
)
|
||||
from .misc import (
|
||||
AutogradFunctionContextVariable,
|
||||
AutogradFunctionVariable,
|
||||
CellVariable,
|
||||
DeletedVariable,
|
||||
ExceptionVariable,
|
||||
GetAttrVariable,
|
||||
LambdaVariable,
|
||||
MethodWrapperVariable,
|
||||
NewGlobalVariable,
|
||||
NumpyVariable,
|
||||
ObjectVariable,
|
||||
PythonModuleVariable,
|
||||
RandomClassVariable,
|
||||
RandomVariable,
|
||||
StringFormatVariable,
|
||||
SuperVariable,
|
||||
TorchVersionVariable,
|
||||
TracebackVariable,
|
||||
TypingVariable,
|
||||
UnknownVariable,
|
||||
WeakRefVariable,
|
||||
)
|
||||
from .nn_module import (
|
||||
FSDPManagedNNModuleVariable,
|
||||
NNModuleVariable,
|
||||
UnspecializedBuiltinNNModuleVariable,
|
||||
UnspecializedNNModuleVariable,
|
||||
)
|
||||
from .optimizer import OptimizerVariable
|
||||
from .sdpa import SDPAParamsVariable
|
||||
from .sets import (
|
||||
DictKeySetVariable,
|
||||
FrozensetVariable,
|
||||
OrderedSetClassVariable,
|
||||
OrderedSetVariable,
|
||||
SetVariable,
|
||||
)
|
||||
from .streams import (
|
||||
CudaStreamVariable,
|
||||
EventVariable,
|
||||
StreamContextVariable,
|
||||
StreamVariable,
|
||||
)
|
||||
from .tensor import (
|
||||
DataPtrVariable,
|
||||
FakeItemVariable,
|
||||
NumpyNdarrayVariable,
|
||||
SymNodeVariable,
|
||||
TensorVariable,
|
||||
UnspecializedPythonVariable,
|
||||
UntypedStorageVariable,
|
||||
)
|
||||
from .torch import TorchCtxManagerClassVariable, TorchInGraphFunctionVariable
|
||||
from .user_defined import (
|
||||
FrozenDataClassVariable,
|
||||
InspectVariable,
|
||||
MutableMappingVariable,
|
||||
NamedTupleVariable,
|
||||
RemovableHandleVariable,
|
||||
StructSequenceVariable,
|
||||
UserDefinedClassVariable,
|
||||
UserDefinedConstantVariable,
|
||||
UserDefinedDictVariable,
|
||||
UserDefinedExceptionClassVariable,
|
||||
UserDefinedExceptionObjectVariable,
|
||||
UserDefinedListVariable,
|
||||
UserDefinedObjectVariable,
|
||||
UserDefinedSetVariable,
|
||||
UserDefinedTupleVariable,
|
||||
UserDefinedVariable,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AutogradFunctionContextVariable",
|
||||
"AutogradFunctionVariable",
|
||||
"BackwardHookVariable",
|
||||
"BaseBuiltinVariable",
|
||||
"BaseListVariable",
|
||||
"BuiltinVariable",
|
||||
"CatchWarningsCtxManagerVariable",
|
||||
"ConstantVariable",
|
||||
"ConstDictVariable",
|
||||
"DictBuiltinVariable",
|
||||
"ContextWrappingVariable",
|
||||
"CountIteratorVariable",
|
||||
"CreateTMADescriptorExperimentalVariable",
|
||||
"CreateTMADescriptorStableVariable",
|
||||
"CUDADeviceVariable",
|
||||
"CudagraphOverrideVariable",
|
||||
"DataPtrVariable",
|
||||
"DefaultDictVariable",
|
||||
"DeletedVariable",
|
||||
"DictKeySetVariable",
|
||||
"DynamoConfigPatchVariable",
|
||||
"FakeItemVariable",
|
||||
"GetAttrVariable",
|
||||
"GradModeVariable",
|
||||
"InspectSignatureVariable",
|
||||
"InspectVariable",
|
||||
"IterBuiltinVariable",
|
||||
"IteratorVariable",
|
||||
"ItertoolsVariable",
|
||||
"LambdaVariable",
|
||||
"LazyConstantVariable",
|
||||
"LazyVariableTracker",
|
||||
"ListBuiltinVariable",
|
||||
"ListIteratorVariable",
|
||||
"ListVariable",
|
||||
"NestedUserFunctionVariable",
|
||||
"CellVariable",
|
||||
"NewGlobalVariable",
|
||||
"NNModuleVariable",
|
||||
"NumpyNdarrayVariable",
|
||||
"NumpyVariable",
|
||||
"OptimizerVariable",
|
||||
"PolyfilledFunctionVariable",
|
||||
"PythonModuleVariable",
|
||||
"RangeVariable",
|
||||
"RemovableHandleVariable",
|
||||
"RepeatIteratorVariable",
|
||||
"SDPAParamsVariable",
|
||||
"ErrorOnGraphBreakVariable",
|
||||
"SkipFunctionVariable",
|
||||
"SliceVariable",
|
||||
"StringFormatVariable",
|
||||
"SuperVariable",
|
||||
"TemporarilyPopInterpreterStackCtxManagerVariable",
|
||||
"TensorVariable",
|
||||
"TMADescriptorExperimentalVariable",
|
||||
"TMADescriptorStableVariable",
|
||||
"TorchCtxManagerClassVariable",
|
||||
"TorchInGraphFunctionVariable",
|
||||
"TorchVersionVariable",
|
||||
"TupleVariable",
|
||||
"UnknownVariable",
|
||||
"UnspecializedNNModuleVariable",
|
||||
"UnspecializedPythonVariable",
|
||||
"UntypedStorageVariable",
|
||||
"UserDefinedClassVariable",
|
||||
"UserDefinedTupleVariable",
|
||||
"NamedTupleVariable",
|
||||
"StructSequenceVariable",
|
||||
"UserDefinedObjectVariable",
|
||||
"UserFunctionVariable",
|
||||
"UserMethodVariable",
|
||||
"VariableTracker",
|
||||
"WithEnterFunctionVariable",
|
||||
"WithExitFunctionVariable",
|
||||
"MappingProxyVariable",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,482 @@
|
||||
"""
|
||||
Constant variable tracking in Dynamo.
|
||||
|
||||
This module is fundamental to Dynamo's ability to track and propagate constant
|
||||
values during compilation, ensuring proper handling of Python literals and
|
||||
maintaining type safety through the compilation process.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import Any, Literal, overload, TYPE_CHECKING
|
||||
from typing_extensions import override
|
||||
|
||||
import torch
|
||||
from torch._dynamo.source import GetItemSource
|
||||
|
||||
from .. import variables
|
||||
from ..exc import raise_observed_exception, unimplemented
|
||||
from ..utils import common_constant_types, istype, np, raise_args_mismatch
|
||||
from .base import ValueMutationNew, VariableTracker
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from torch._dynamo.symbolic_convert import InstructionTranslator
|
||||
|
||||
from .functions import UserFunctionVariable
|
||||
|
||||
|
||||
class ConstantVariable(VariableTracker):
|
||||
"""
|
||||
Variable tracker for Python literals and basic immutable types, with automatic
|
||||
routing support for collection types (lists, tuples, sets, etc.).
|
||||
|
||||
The create() method intelligently constructs appropriate variable types for
|
||||
nested collections.
|
||||
"""
|
||||
|
||||
# PyLong_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/longobject.c#L6585
|
||||
# PyFloat_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/floatobject.c#L1880
|
||||
# PyBool_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/boolobject.c#L171
|
||||
# PyUnicode_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/unicodeobject.c#L14931
|
||||
# PyBytes_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/bytesobject.c#L3017
|
||||
# PyComplex_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/complexobject.c#L1099
|
||||
# _PyNone_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/object.c#L2022
|
||||
_cpython_type = (int, float, str, bytes, bool, type(None), complex, type(...))
|
||||
|
||||
@overload
|
||||
@staticmethod
|
||||
def create(value: None) -> ConstantVariable: ...
|
||||
|
||||
@overload
|
||||
@staticmethod
|
||||
def create(value: bool) -> ConstantVariable: ...
|
||||
|
||||
@overload
|
||||
@staticmethod
|
||||
def create(value: Any, **kwargs: Any) -> VariableTracker: ...
|
||||
|
||||
@staticmethod
|
||||
def create(value: Any, **kwargs: Any) -> VariableTracker:
|
||||
"""
|
||||
Create a `ConstantVariable` based on the given value, and supports
|
||||
automatic routing for collection types like `tuple` (in which case we'd
|
||||
create `ConstantVariable` for the leaf items).
|
||||
|
||||
NOTE: the caller must install the proper guards if needed; most often
|
||||
the guard will be `CONSTANT_MATCH`.
|
||||
"""
|
||||
# Return pre-allocated sentinels for None/True/False when there are
|
||||
# no extra kwargs (source, etc.) that would differentiate the instance.
|
||||
if not kwargs:
|
||||
match value:
|
||||
case None:
|
||||
return CONSTANT_VARIABLE_NONE
|
||||
case True:
|
||||
return CONSTANT_VARIABLE_TRUE
|
||||
case False:
|
||||
return CONSTANT_VARIABLE_FALSE
|
||||
|
||||
source = kwargs.get("source")
|
||||
|
||||
# Routing for supported collection literals.
|
||||
if isinstance(value, set):
|
||||
items = [ConstantVariable.create(x) for x in value]
|
||||
return variables.SetVariable(items, **kwargs) # type: ignore[arg-type]
|
||||
elif isinstance(value, frozenset):
|
||||
items = [ConstantVariable.create(x) for x in value]
|
||||
return variables.FrozensetVariable(items, **kwargs) # type: ignore[arg-type]
|
||||
elif isinstance(value, slice):
|
||||
slice_args = (value.start, value.stop, value.step)
|
||||
slice_args_vars = tuple(ConstantVariable.create(arg) for arg in slice_args)
|
||||
return variables.SliceVariable(slice_args_vars, **kwargs)
|
||||
elif isinstance(value, (list, tuple)):
|
||||
items = []
|
||||
for i, x in enumerate(value):
|
||||
item_source = GetItemSource(source, i) if source else None
|
||||
items.append(
|
||||
ConstantVariable.create(
|
||||
x,
|
||||
source=item_source,
|
||||
)
|
||||
)
|
||||
return variables.BaseListVariable.cls_for(type(value))(items, **kwargs)
|
||||
|
||||
return ConstantVariable(value, **kwargs)
|
||||
|
||||
def __init__(self, value: Any, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
assert ConstantVariable.is_base_literal(value), f"""
|
||||
Cannot construct `ConstantVariable` for value of type {type(value)}.
|
||||
|
||||
This failure likely due to PyTorch-internal use of `ConstantVariable` on
|
||||
non-literal python values, please try using `VariableTracker.build` instead. If
|
||||
you believe it's a necessary and legitimate use case (the value is immutable and
|
||||
can't easily be represented with another `VariableTracker` class), please add
|
||||
its type to `common_constant_types`.
|
||||
"""
|
||||
if np is not None and isinstance(value, np.number):
|
||||
self.value = value.item()
|
||||
else:
|
||||
self.value = value
|
||||
|
||||
def as_proxy(self) -> Any:
|
||||
return self.value
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"ConstantVariable({type(self.value).__name__}: {repr(self.value)})"
|
||||
|
||||
def as_python_constant(self) -> Any:
|
||||
return self.value
|
||||
|
||||
def is_python_constant(self) -> Literal[True]:
|
||||
return True
|
||||
|
||||
def is_symnode_like(self) -> bool:
|
||||
return isinstance(self.value, (int, bool))
|
||||
|
||||
def is_constant_match(self, *values: Any) -> bool:
|
||||
return self.value in values
|
||||
|
||||
def is_constant_none(self) -> bool:
|
||||
return self.value is None
|
||||
|
||||
@property
|
||||
def items(self) -> list[VariableTracker]:
|
||||
"""
|
||||
Need this when adding a BaseListVariable and a ConstantVariable together.
|
||||
Happens in detectron2.
|
||||
"""
|
||||
return self.unpack_var_sequence(tx=None)
|
||||
|
||||
def getitem_const(
|
||||
self, tx: InstructionTranslator, arg: VariableTracker
|
||||
) -> VariableTracker:
|
||||
return ConstantVariable.create(
|
||||
self.value[arg.as_python_constant()],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_base_literal(obj: object) -> bool:
|
||||
return type(obj) in common_constant_types
|
||||
|
||||
@staticmethod
|
||||
def is_literal(obj: object, cache: dict[int, object] | None = None) -> bool:
|
||||
if cache is None:
|
||||
cache = {}
|
||||
if id(obj) in cache:
|
||||
# no-op if there is a cyclical reference
|
||||
return True
|
||||
if type(obj) in (list, tuple, set, frozenset, torch.Size):
|
||||
cache[id(obj)] = obj
|
||||
return all(ConstantVariable.is_literal(x, cache) for x in obj) # type: ignore[attr-defined]
|
||||
return ConstantVariable.is_base_literal(obj)
|
||||
|
||||
def unpack_var_sequence(
|
||||
self, tx: InstructionTranslator | None
|
||||
) -> list[VariableTracker]:
|
||||
try:
|
||||
return [ConstantVariable.create(x) for x in self.as_python_constant()]
|
||||
except TypeError as e:
|
||||
raise NotImplementedError from e
|
||||
|
||||
def len_impl(self, tx: InstructionTranslator) -> VariableTracker:
|
||||
"""Generic len for any constant value (sequence or mapping)."""
|
||||
try:
|
||||
return ConstantVariable.create(len(self.value))
|
||||
except TypeError as e:
|
||||
raise_observed_exception(type(e), tx, args=list(e.args))
|
||||
|
||||
def sq_length(self, tx: InstructionTranslator) -> VariableTracker:
|
||||
"""Sequence length - delegates to len_impl for constants."""
|
||||
return self.len_impl(tx)
|
||||
|
||||
def mp_length(self, tx: InstructionTranslator) -> VariableTracker:
|
||||
"""Mapping length - delegates to len_impl for constants."""
|
||||
return self.len_impl(tx)
|
||||
|
||||
def const_getattr(self, tx: InstructionTranslator, name: str) -> VariableTracker:
|
||||
if not hasattr(self.value, name):
|
||||
raise_observed_exception(AttributeError, tx, args=[name])
|
||||
member = getattr(self.value, name)
|
||||
if callable(member):
|
||||
raise NotImplementedError
|
||||
return member
|
||||
|
||||
def call_method(
|
||||
self,
|
||||
tx: InstructionTranslator,
|
||||
name: str,
|
||||
args: list[VariableTracker],
|
||||
kwargs: dict[str, VariableTracker],
|
||||
) -> VariableTracker:
|
||||
from .tensor import SymNodeVariable
|
||||
|
||||
if name == "format" and istype(self.value, str):
|
||||
return variables.BuiltinVariable(str.format).call_function(
|
||||
tx, [self, *args], kwargs
|
||||
)
|
||||
elif name == "join" and istype(self.value, str):
|
||||
if kwargs or len(args) != 1:
|
||||
raise_args_mismatch(
|
||||
tx,
|
||||
name,
|
||||
"1 args and 0 kwargs",
|
||||
f"{len(args)} args and {len(kwargs)} kwargs",
|
||||
)
|
||||
arg_unpacked = args[0].force_unpack_var_sequence(tx)
|
||||
try:
|
||||
arg_const = [x.as_python_constant() for x in arg_unpacked]
|
||||
return ConstantVariable.create(self.value.join(arg_const))
|
||||
except NotImplementedError:
|
||||
return super().call_method(tx, name, args, kwargs)
|
||||
elif name == "__iter__" and istype(self.value, str):
|
||||
# this could be some generic iterator to avoid the circular import,
|
||||
# but ListIterator does what we want
|
||||
from .lists import ListIteratorVariable
|
||||
|
||||
return ListIteratorVariable(
|
||||
self.unpack_var_sequence(tx), mutation_type=ValueMutationNew()
|
||||
)
|
||||
|
||||
if any(isinstance(x, SymNodeVariable) for x in args):
|
||||
# Promote to SymNodeVariable for operations involving dynamic shapes.
|
||||
return variables.SymNodeVariable.create(
|
||||
tx, self.as_proxy(), self.value
|
||||
).call_method(tx, name, args, kwargs)
|
||||
|
||||
try:
|
||||
const_args = [a.as_python_constant() for a in args]
|
||||
const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
|
||||
except NotImplementedError:
|
||||
return super().call_method(tx, name, args, kwargs)
|
||||
|
||||
if isinstance(self.value, str) and name in str.__dict__:
|
||||
method = getattr(self.value, name)
|
||||
try:
|
||||
return ConstantVariable.create(method(*const_args, **const_kwargs))
|
||||
except Exception as e:
|
||||
raise_observed_exception(type(e), tx)
|
||||
elif isinstance(self.value, (float, int)) and hasattr(self.value, name):
|
||||
if not (args or kwargs):
|
||||
try:
|
||||
return ConstantVariable.create(getattr(self.value, name)())
|
||||
except (OverflowError, ValueError) as exc:
|
||||
raise_observed_exception(
|
||||
type(exc),
|
||||
tx,
|
||||
args=list(exc.args),
|
||||
)
|
||||
if (
|
||||
hasattr(operator, name)
|
||||
and len(args) == 1
|
||||
and args[0].is_python_constant()
|
||||
):
|
||||
add_target = const_args[0]
|
||||
op = getattr(operator, name)
|
||||
if isinstance(
|
||||
add_target, (torch.SymBool, torch.SymFloat, torch.SymInt)
|
||||
):
|
||||
# Addition between a non sym and sym makes a sym
|
||||
proxy = tx.output.create_proxy(
|
||||
"call_function", op, (self.value, add_target), {}
|
||||
)
|
||||
return SymNodeVariable.create(tx, proxy, add_target)
|
||||
else:
|
||||
try:
|
||||
return ConstantVariable.create(op(self.value, add_target))
|
||||
except Exception as e:
|
||||
raise_observed_exception(type(e), tx, args=list(e.args))
|
||||
elif isinstance(self.value, bytes) and name == "decode":
|
||||
method = getattr(self.value, name)
|
||||
return ConstantVariable.create(method(*const_args, **const_kwargs))
|
||||
elif type(self.value) is complex and name in complex.__dict__:
|
||||
method = getattr(self.value, name)
|
||||
try:
|
||||
return ConstantVariable.create(method(*const_args, **const_kwargs))
|
||||
except Exception as e:
|
||||
raise_observed_exception(type(e), tx)
|
||||
|
||||
if name == "__round__" and len(args) == 1 and args[0].is_python_constant():
|
||||
try:
|
||||
return ConstantVariable.create(
|
||||
round(self.value, args[0].as_python_constant())
|
||||
)
|
||||
except Exception as e:
|
||||
raise_observed_exception(type(e), tx, args=list(e.args))
|
||||
elif name == "__contains__" and len(args) == 1 and args[0].is_python_constant():
|
||||
assert not kwargs
|
||||
search = args[0].as_python_constant()
|
||||
try:
|
||||
result = search in self.value
|
||||
return ConstantVariable.create(result)
|
||||
except TypeError as e:
|
||||
raise_observed_exception(type(e), tx, args=list(e.args))
|
||||
return super().call_method(tx, name, args, kwargs)
|
||||
|
||||
def call_tree_map(
|
||||
self,
|
||||
tx: InstructionTranslator,
|
||||
tree_map_fn: UserFunctionVariable,
|
||||
map_fn: VariableTracker,
|
||||
rest: Sequence[VariableTracker],
|
||||
tree_map_kwargs: dict[str, VariableTracker],
|
||||
) -> VariableTracker:
|
||||
if self.value is None:
|
||||
none_is_leaf_var = tree_map_kwargs.get("none_is_leaf")
|
||||
if none_is_leaf_var is not None:
|
||||
try:
|
||||
none_is_leaf = bool(none_is_leaf_var.as_python_constant())
|
||||
except NotImplementedError:
|
||||
return self._tree_map_fallback(
|
||||
tx,
|
||||
tree_map_fn,
|
||||
map_fn,
|
||||
rest,
|
||||
tree_map_kwargs,
|
||||
)
|
||||
else:
|
||||
tree_map_module = getattr(
|
||||
getattr(tree_map_fn, "fn", None), "__module__", ""
|
||||
)
|
||||
# torch.utils._pytree and torch.utils._cxx_pytree treat None as a leaf
|
||||
# by default, while optree keeps it as an internal node unless
|
||||
# none_is_leaf=True is provided.
|
||||
none_is_leaf = not tree_map_module.startswith("optree")
|
||||
if none_is_leaf:
|
||||
return map_fn.call_function(tx, [self, *rest], {})
|
||||
else:
|
||||
for other in rest:
|
||||
if not other.is_constant_none():
|
||||
return self._tree_map_fallback(
|
||||
tx,
|
||||
tree_map_fn,
|
||||
map_fn,
|
||||
rest,
|
||||
tree_map_kwargs,
|
||||
)
|
||||
return self.clone()
|
||||
if isinstance(self.value, (int, float, bool, complex, str, bytes, torch.dtype)):
|
||||
return map_fn.call_function(tx, [self, *rest], {})
|
||||
return super().call_tree_map(
|
||||
tx,
|
||||
tree_map_fn,
|
||||
map_fn,
|
||||
rest,
|
||||
tree_map_kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
def call_obj_hasattr(
|
||||
self, tx: InstructionTranslator, name: str
|
||||
) -> ConstantVariable:
|
||||
result = hasattr(self.value, name)
|
||||
return variables.ConstantVariable.create(result)
|
||||
|
||||
def is_python_hashable(self) -> Literal[True]:
|
||||
return True
|
||||
|
||||
def get_python_hash(self) -> int:
|
||||
return hash(self.value)
|
||||
|
||||
def is_python_equal(self, other: object) -> bool:
|
||||
from .tensor import SymNodeVariable
|
||||
|
||||
if isinstance(other, SymNodeVariable):
|
||||
return self.as_python_constant() == other.evaluate_expr()
|
||||
return (
|
||||
isinstance(other, VariableTracker)
|
||||
and self.as_python_constant() == other.as_python_constant()
|
||||
)
|
||||
|
||||
def get_real_python_backed_value(self) -> object:
|
||||
return self.value
|
||||
|
||||
def nb_index_impl(
|
||||
self,
|
||||
tx: Any,
|
||||
) -> VariableTracker:
|
||||
# CPython: int and bool define nb_index (returns self for int,
|
||||
# int(self) for bool). All other constant types do not.
|
||||
if isinstance(self.value, (int, bool)):
|
||||
return ConstantVariable.create(operator.index(self.value))
|
||||
return super().nb_index_impl(tx)
|
||||
|
||||
def nb_int_impl(
|
||||
self,
|
||||
tx: Any,
|
||||
) -> VariableTracker:
|
||||
# CPython: int defines nb_int (long_long, returns copy).
|
||||
# bool inherits nb_int from int via slot inheritance.
|
||||
# float defines nb_int (truncates toward zero via PyLong_FromDouble).
|
||||
return ConstantVariable.create(int(self.value))
|
||||
|
||||
def nb_float_impl(
|
||||
self,
|
||||
tx: Any,
|
||||
) -> VariableTracker:
|
||||
# CPython: float defines nb_float (float_float, returns copy).
|
||||
# int defines nb_float (long_float, converts to float).
|
||||
# bool inherits nb_float from int via slot inheritance.
|
||||
return ConstantVariable.create(float(self.value))
|
||||
|
||||
|
||||
CONSTANT_VARIABLE_NONE = ConstantVariable(None)
|
||||
CONSTANT_VARIABLE_TRUE = ConstantVariable(True)
|
||||
CONSTANT_VARIABLE_FALSE = ConstantVariable(False)
|
||||
|
||||
|
||||
class FakeIdVariable(VariableTracker):
|
||||
"""A compile-time-only id value that can be used as a dict key but cannot
|
||||
be reconstructed across graph breaks.
|
||||
|
||||
When dynamo evaluates ``id(x)`` on a variable tracker that has no
|
||||
corresponding runtime object (e.g. a ``ConstDictVariable`` created during
|
||||
tracing), we mint a fake integer id. This variable holds that id and
|
||||
supports the minimal interface needed to participate as a dict key
|
||||
(hashing and equality). It intentionally blocks reconstruction so that a
|
||||
graph break does not silently bake a stale id into the resumed bytecode.
|
||||
"""
|
||||
|
||||
# PyLong_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/longobject.c#L6585
|
||||
_cpython_type = int
|
||||
|
||||
def __init__(self, value: int, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.value = value
|
||||
|
||||
def as_python_constant(self) -> int:
|
||||
return self.value
|
||||
|
||||
def is_python_constant(self) -> bool:
|
||||
return False
|
||||
|
||||
def python_type(self) -> type:
|
||||
return int
|
||||
|
||||
def is_python_hashable(self) -> bool:
|
||||
return True
|
||||
|
||||
def get_python_hash(self) -> int:
|
||||
return hash(self.value)
|
||||
|
||||
def is_python_equal(self, other: object) -> bool:
|
||||
if isinstance(other, (FakeIdVariable, ConstantVariable)):
|
||||
return self.value == other.as_python_constant()
|
||||
return False
|
||||
|
||||
def reconstruct(self, codegen: Any) -> None:
|
||||
unimplemented(
|
||||
gb_type="Reconstruction of FakeIdVariable",
|
||||
context=str(self.value),
|
||||
explanation=(
|
||||
"A fake id produced by id() on a compile-time container "
|
||||
"cannot be reconstructed across a graph break."
|
||||
),
|
||||
hints=[
|
||||
"Avoid using id() on containers in code that may graph-break.",
|
||||
],
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,258 @@
|
||||
"""
|
||||
Distributed computing variable tracking classes for PyTorch Dynamo.
|
||||
|
||||
This module implements variable tracking for distributed computing components:
|
||||
- Process Groups (for collective communication)
|
||||
- Device Meshes (for distributed tensor sharding)
|
||||
- Placement Types (for specifying distribution strategies)
|
||||
- Distributed Tensors and their operations
|
||||
- Backward hooks for distributed module operations
|
||||
|
||||
These classes are responsible for tracking distributed operations during graph
|
||||
compilation while maintaining proper guards and handling distributed-specific
|
||||
behaviors. They ensure correct handling of distributed components like process
|
||||
groups, device meshes, and placement strategies while preserving proper semantics
|
||||
for distributed tensor operations in the compiled code.
|
||||
|
||||
The implementation provides special handling for distributed package availability
|
||||
checks and proper tracking of distributed state and operations across processes.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from typing import Any, Literal, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.fx.experimental._backward_state import BackwardState
|
||||
|
||||
from .. import compiled_autograd
|
||||
from .._trace_wrapped_higher_order_op import trace_wrapped
|
||||
from ..exc import unimplemented
|
||||
from ..external_utils import call_module_hooks_from_backward_state
|
||||
from ..guards import GuardBuilder, install_guard
|
||||
from ..source import AttrSource
|
||||
from .base import VariableTracker
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._dynamo.symbolic_convert import InstructionTranslator
|
||||
|
||||
|
||||
class DistributedVariable(VariableTracker):
|
||||
"""
|
||||
The base distributed variable that encapsulates common methods
|
||||
for the distributed objects (i.e. ProcessGroup, DeviceMesh, etc.).
|
||||
Concrete distributed objects could inherit this class and add object
|
||||
specific logic.
|
||||
|
||||
i.e. It provides the check on the distributed package existence
|
||||
and hold the tracking value for the corresponding distributed object.
|
||||
"""
|
||||
|
||||
def __init__(self, value: Any, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
if not DistributedVariable.is_available():
|
||||
unimplemented(
|
||||
gb_type="torch.distributed package is not available!",
|
||||
context="",
|
||||
explanation="The PyTorch package doesn't include torch.distributed when building from source.",
|
||||
hints=[
|
||||
"Set USE_DISTRIBUTED=1 to enable it when building PyTorch from source."
|
||||
],
|
||||
)
|
||||
self.value = value
|
||||
|
||||
def python_type(self) -> type:
|
||||
return type(self.value)
|
||||
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
# check if the distributed package is available or not
|
||||
return torch.distributed.is_available()
|
||||
|
||||
def is_python_hashable(self) -> Literal[True]:
|
||||
return True
|
||||
|
||||
def get_python_hash(self) -> int:
|
||||
return hash(self.value)
|
||||
|
||||
def is_python_equal(self, other: object) -> bool:
|
||||
return (
|
||||
isinstance(other, VariableTracker)
|
||||
and self.as_python_constant() == other.as_python_constant()
|
||||
)
|
||||
|
||||
|
||||
def is_from_local(value: object) -> bool:
|
||||
if not DistributedVariable.is_available():
|
||||
return False
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
return inspect.isfunction(value) and value is DTensor.from_local
|
||||
|
||||
|
||||
def is_constant_pg_functions(value: object) -> bool:
|
||||
if not DistributedVariable.is_available():
|
||||
return False
|
||||
|
||||
from torch.distributed.distributed_c10d import (
|
||||
_get_group_size_by_name,
|
||||
_get_group_tag,
|
||||
_rank_not_in_group,
|
||||
_resolve_group_name_by_ranks_and_tag,
|
||||
get_process_group_ranks,
|
||||
)
|
||||
|
||||
constant_processgroup_functions = [
|
||||
_get_group_size_by_name,
|
||||
_get_group_tag,
|
||||
_rank_not_in_group,
|
||||
get_process_group_ranks,
|
||||
_resolve_group_name_by_ranks_and_tag,
|
||||
]
|
||||
|
||||
return inspect.isfunction(value) and value in constant_processgroup_functions
|
||||
|
||||
|
||||
class WorldMetaClassVariable(DistributedVariable):
|
||||
"""
|
||||
Tracks torch.distributed.GroupMember and torch.distributed.group, which are
|
||||
instances of the metaclass _WorldMeta.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def is_group_member_type(cls, value: object) -> bool:
|
||||
if not cls.is_available():
|
||||
return False
|
||||
|
||||
from torch.distributed.distributed_c10d import _WorldMeta
|
||||
|
||||
return type(value) is _WorldMeta
|
||||
|
||||
def python_type(self) -> type:
|
||||
return type(self.value)
|
||||
|
||||
def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
|
||||
if name == "WORLD":
|
||||
from .builder import SourcelessBuilder
|
||||
|
||||
assert self.source
|
||||
source = AttrSource(base=self.source, member="WORLD")
|
||||
install_guard(source.make_guard(GuardBuilder.ID_MATCH))
|
||||
return SourcelessBuilder.create(tx, self.value.WORLD)
|
||||
elif name == "NON_GROUP_MEMBER":
|
||||
assert self.source
|
||||
source = AttrSource(base=self.source, member="NON_GROUP_MEMBER")
|
||||
install_guard(source.make_guard(GuardBuilder.ID_MATCH))
|
||||
return VariableTracker.build(tx, self.value.NON_GROUP_MEMBER)
|
||||
return super().var_getattr(tx, name)
|
||||
|
||||
|
||||
class BackwardHookVariable(VariableTracker):
|
||||
"""
|
||||
Handles torch.utils.hooks.BackwardHook for module-level backward
|
||||
hooks.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
tx: "InstructionTranslator",
|
||||
module: VariableTracker,
|
||||
user_hooks: VariableTracker,
|
||||
user_pre_hooks: VariableTracker,
|
||||
) -> "BackwardHookVariable":
|
||||
if not compiled_autograd.compiled_autograd_enabled:
|
||||
unimplemented(
|
||||
gb_type="Module-level backwards hooks require compiled autograd.",
|
||||
context="",
|
||||
explanation="",
|
||||
hints=[
|
||||
"Enable compiled autograd by setting torch._dynamo.config.compiled_autograd = True."
|
||||
],
|
||||
)
|
||||
|
||||
def _in_graph_bw_hooks(
|
||||
bw_state: BackwardState,
|
||||
) -> torch.utils.hooks.BackwardHook:
|
||||
"""
|
||||
Rather than installing the user hooks in the graph (which
|
||||
don't survive AotAutograd), we install hooks that will call
|
||||
trace_wrapped in the backward pass that CompiledAutograd
|
||||
can turn into actual hook calls.
|
||||
"""
|
||||
return torch.utils.hooks.BackwardHook(
|
||||
None,
|
||||
(
|
||||
functools.partial(
|
||||
trace_wrapped,
|
||||
fn=call_module_hooks_from_backward_state,
|
||||
bw_state=bw_state,
|
||||
hooks_name=user_hooks_name,
|
||||
module_name=module_name,
|
||||
),
|
||||
),
|
||||
(
|
||||
functools.partial(
|
||||
trace_wrapped,
|
||||
fn=call_module_hooks_from_backward_state,
|
||||
bw_state=bw_state,
|
||||
hooks_name=user_pre_hooks_name,
|
||||
module_name=module_name,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
module_name, bw_state_proxy = tx.output.add_backward_state_hook(module, "mod")
|
||||
user_pre_hooks_name, _ = tx.output.add_backward_state_hook(user_pre_hooks)
|
||||
user_hooks_name, _ = tx.output.add_backward_state_hook(user_hooks)
|
||||
proxy = tx.output.create_proxy(
|
||||
"call_function",
|
||||
_in_graph_bw_hooks,
|
||||
(bw_state_proxy,),
|
||||
{},
|
||||
)
|
||||
proxy.node.meta["example_value"] = torch.utils.hooks.BackwardHook(None, (), ())
|
||||
return BackwardHookVariable(proxy, module, user_hooks, user_pre_hooks)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
proxy: torch.fx.Proxy,
|
||||
module: VariableTracker,
|
||||
user_hooks: VariableTracker,
|
||||
user_pre_hooks: VariableTracker,
|
||||
**options: Any,
|
||||
) -> None:
|
||||
super().__init__(**options)
|
||||
self.proxy = proxy
|
||||
self.module = module
|
||||
self.user_hooks = user_hooks
|
||||
self.user_pre_hooks = user_pre_hooks
|
||||
|
||||
def as_proxy(self) -> torch.fx.Proxy:
|
||||
return self.proxy
|
||||
|
||||
def call_method(
|
||||
self,
|
||||
tx: "InstructionTranslator",
|
||||
name: str,
|
||||
args: list[VariableTracker],
|
||||
kwargs: dict[str, VariableTracker],
|
||||
) -> VariableTracker:
|
||||
if name in ("setup_input_hook", "setup_output_hook"):
|
||||
return self._setup_hook(tx, name, *args, **kwargs)
|
||||
return super().call_method(tx, name, args, kwargs)
|
||||
|
||||
def _setup_hook(
|
||||
self, tx: "InstructionTranslator", hook_method_name: str, args: VariableTracker
|
||||
) -> VariableTracker:
|
||||
from .builder import wrap_fx_proxy
|
||||
|
||||
return wrap_fx_proxy(
|
||||
tx,
|
||||
tx.output.create_proxy(
|
||||
"call_method",
|
||||
hook_method_name,
|
||||
(self.as_proxy(), args.as_proxy()),
|
||||
{},
|
||||
),
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user