Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
BIN
Binary file not shown.
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,15 @@
|
||||
from torch import Tensor
|
||||
from torch.types import _dtype, _int, Device
|
||||
|
||||
# Defined in torch/csrc/acc/Module.cpp
|
||||
class PrivateUse1Hooks:
|
||||
def has_primary_context(self, device_index: _int) -> bool: ...
|
||||
def is_built(self) -> bool: ...
|
||||
def is_available(self) -> bool: ...
|
||||
|
||||
class DeviceGuard:
|
||||
def type_(self) -> Device: ...
|
||||
|
||||
def register_python_privateuseone_device_guard(guard: DeviceGuard) -> bool: ...
|
||||
def register_python_privateuseone_hook(hook: PrivateUse1Hooks) -> bool: ...
|
||||
def create_empty_tensor(shape: tuple[_int, ...], dtype: _dtype) -> Tensor: ...
|
||||
@@ -0,0 +1,164 @@
|
||||
from ctypes import c_void_p
|
||||
from typing import overload, Protocol
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
# Defined in torch/csrc/inductor/aoti_runner/pybind.cpp
|
||||
|
||||
# Tensor to AtenTensorHandle
|
||||
def unsafe_alloc_void_ptrs_from_tensors(tensors: list[Tensor]) -> list[c_void_p]: ...
|
||||
def unsafe_alloc_void_ptr_from_tensor(tensor: Tensor) -> c_void_p: ...
|
||||
|
||||
# AtenTensorHandle to Tensor
|
||||
def alloc_tensors_by_stealing_from_void_ptrs(
|
||||
handles: list[c_void_p],
|
||||
) -> list[Tensor]: ...
|
||||
def alloc_tensor_by_stealing_from_void_ptr(
|
||||
handle: c_void_p,
|
||||
) -> Tensor: ...
|
||||
|
||||
class AOTIModelContainerRunner(Protocol):
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerCpu:
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerCuda:
|
||||
@overload
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str
|
||||
) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str, cubin_dir: str
|
||||
) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerXpu:
|
||||
@overload
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str
|
||||
) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str, kernel_bin_dir: str
|
||||
) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerMps:
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
# Defined in torch/csrc/inductor/aoti_package/pybind.cpp
|
||||
class AOTIModelPackageLoader:
|
||||
def __init__(
|
||||
self,
|
||||
model_package_path: str,
|
||||
model_name: str,
|
||||
run_single_threaded: bool,
|
||||
num_runners: int,
|
||||
device_index: int,
|
||||
) -> None: ...
|
||||
def get_metadata(self) -> dict[str, str]: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def boxed_run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_fqns(self) -> list[str]: ...
|
||||
def load_constants(
|
||||
self,
|
||||
constants_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
check_full_update: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
@@ -0,0 +1,163 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._C._profiler import (
|
||||
_ProfilerEvent,
|
||||
ActiveProfilerType,
|
||||
ProfilerActivity,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
# Defined in torch/csrc/autograd/init.cpp
|
||||
|
||||
class DeviceType(Enum):
|
||||
CPU = ...
|
||||
CUDA = ...
|
||||
XPU = ...
|
||||
MKLDNN = ...
|
||||
OPENGL = ...
|
||||
OPENCL = ...
|
||||
IDEEP = ...
|
||||
HIP = ...
|
||||
FPGA = ...
|
||||
MAIA = ...
|
||||
XLA = ...
|
||||
MTIA = ...
|
||||
MPS = ...
|
||||
HPU = ...
|
||||
Meta = ...
|
||||
Vulkan = ...
|
||||
Metal = ...
|
||||
PrivateUse1 = ...
|
||||
|
||||
class ProfilerEvent:
|
||||
def cpu_elapsed_us(self, other: ProfilerEvent) -> float: ...
|
||||
def cpu_memory_usage(self) -> int: ...
|
||||
def cuda_elapsed_us(self, other: ProfilerEvent) -> float: ...
|
||||
def privateuse1_elapsed_us(self, other: ProfilerEvent) -> float: ...
|
||||
def cuda_memory_usage(self) -> int: ...
|
||||
def device(self) -> int: ...
|
||||
def handle(self) -> int: ...
|
||||
def has_cuda(self) -> bool: ...
|
||||
def is_remote(self) -> bool: ...
|
||||
def kind(self) -> int: ...
|
||||
def name(self) -> str: ...
|
||||
def node_id(self) -> int: ...
|
||||
def sequence_nr(self) -> int: ...
|
||||
def shapes(self) -> list[list[int]]: ...
|
||||
def thread_id(self) -> int: ...
|
||||
def flops(self) -> float: ...
|
||||
def is_async(self) -> bool: ...
|
||||
|
||||
class _KinetoEvent:
|
||||
def name(self) -> str: ...
|
||||
def overload_name(self) -> str: ...
|
||||
def device_index(self) -> int: ...
|
||||
def device_resource_id(self) -> int: ...
|
||||
def start_ns(self) -> int: ...
|
||||
def end_ns(self) -> int: ...
|
||||
def duration_ns(self) -> int: ...
|
||||
def is_async(self) -> bool: ...
|
||||
def linked_correlation_id(self) -> int: ...
|
||||
def external_id(self) -> int: ...
|
||||
def shapes(self) -> list[list[int]]: ...
|
||||
def dtypes(self) -> list[str]: ...
|
||||
def concrete_inputs(self) -> list[Any]: ...
|
||||
def kwinputs(self) -> dict[str, Any]: ...
|
||||
def device_type(self) -> DeviceType: ...
|
||||
def start_thread_id(self) -> int: ...
|
||||
def end_thread_id(self) -> int: ...
|
||||
def correlation_id(self) -> int: ...
|
||||
def fwd_thread_id(self) -> int: ...
|
||||
def stack(self) -> list[str]: ...
|
||||
def scope(self) -> int: ...
|
||||
def sequence_nr(self) -> int: ...
|
||||
def flops(self) -> int: ...
|
||||
def cuda_elapsed_us(self) -> int: ...
|
||||
def privateuse1_elapsed_us(self) -> int: ...
|
||||
def is_user_annotation(self) -> bool: ...
|
||||
def is_python_function(self) -> bool: ...
|
||||
def is_hidden_event(self) -> bool: ...
|
||||
def metadata_json(self) -> str: ...
|
||||
def activity_type(self) -> str: ...
|
||||
def extra_meta(self) -> dict[str, str]: ...
|
||||
def flow_id(self) -> int: ...
|
||||
def flow_type(self) -> int: ...
|
||||
def flow_start(self) -> bool: ...
|
||||
def structured_input_shapes(self) -> list[list[int] | list[list[int]]]: ...
|
||||
def structured_input_strides(self) -> list[list[int] | list[list[int]]]: ...
|
||||
def python_id(self) -> int: ...
|
||||
def python_parent_id(self) -> int: ...
|
||||
def python_module_id(self) -> int: ...
|
||||
|
||||
class _ProfilerResult:
|
||||
def events(self) -> list[_KinetoEvent]: ...
|
||||
def legacy_events(self) -> list[list[ProfilerEvent]]: ...
|
||||
def save(self, path: str) -> None: ...
|
||||
def experimental_event_tree(self) -> list[_ProfilerEvent]: ...
|
||||
def trace_start_ns(self) -> int: ...
|
||||
|
||||
class SavedTensor:
|
||||
def unpack(self) -> torch.Tensor: ...
|
||||
|
||||
def _make_saved_tensor(
|
||||
tensor: torch.Tensor,
|
||||
is_output: bool,
|
||||
is_inplace_on_view: bool = False,
|
||||
) -> SavedTensor: ...
|
||||
def _enable_profiler(
|
||||
config: ProfilerConfig,
|
||||
activities: set[ProfilerActivity],
|
||||
) -> None: ...
|
||||
def _prepare_profiler(
|
||||
config: ProfilerConfig,
|
||||
activities: set[ProfilerActivity],
|
||||
activity_filter: dict[ProfilerActivity, set[str]] = ...,
|
||||
) -> None: ...
|
||||
def _toggle_collection_dynamic(
|
||||
enable: bool,
|
||||
activities: set[ProfilerActivity],
|
||||
) -> None: ...
|
||||
def _disable_profiler() -> _ProfilerResult: ...
|
||||
def _profiler_enabled() -> bool: ...
|
||||
def _add_metadata_json(key: str, value: str) -> None: ...
|
||||
def _kineto_step() -> None: ...
|
||||
def _get_current_graph_task_keep_graph() -> bool: ...
|
||||
def _get_sequence_nr() -> int: ...
|
||||
def kineto_available() -> bool: ...
|
||||
def _record_function_with_args_enter(name: str, *args) -> torch.Tensor: ...
|
||||
def _record_function_with_args_exit(handle: torch.Tensor) -> None: ...
|
||||
def _supported_activities() -> set[ProfilerActivity]: ...
|
||||
def _enable_record_function(enable: bool) -> None: ...
|
||||
def _set_empty_test_observer(is_global: bool, sampling_prob: float) -> None: ...
|
||||
def _push_saved_tensors_default_hooks(
|
||||
pack_hook: Callable[[torch.Tensor], Any],
|
||||
unpack_hook: Callable[[Any], torch.Tensor],
|
||||
) -> None: ...
|
||||
def _pop_saved_tensors_default_hooks() -> None: ...
|
||||
def _top_saved_tensors_default_hooks(
|
||||
ignore_is_tracing: bool,
|
||||
) -> tuple[Callable[[torch.Tensor], Any], Callable[[Any], torch.Tensor]]: ...
|
||||
def _unsafe_set_version_counter(
|
||||
t: tuple[torch.Tensor, ...], prev_version: tuple[int, ...]
|
||||
) -> None: ...
|
||||
def _enable_profiler_legacy(config: ProfilerConfig) -> None: ...
|
||||
def _disable_profiler_legacy() -> list[list[ProfilerEvent]]: ...
|
||||
def _profiler_type() -> ActiveProfilerType: ...
|
||||
def _saved_tensors_hooks_enable() -> None: ...
|
||||
def _saved_tensors_hooks_disable(message: str, fail_if_non_empty=True) -> None: ...
|
||||
def _saved_tensors_hooks_get_disabled_error_message() -> str | None: ...
|
||||
def _saved_tensors_hooks_set_tracing(is_tracing: bool) -> bool: ...
|
||||
|
||||
class CreationMeta(Enum):
|
||||
DEFAULT = ...
|
||||
IN_CUSTOM_FUNCTION = ...
|
||||
MULTI_OUTPUT_NODE = ...
|
||||
NO_GRAD_MODE = ...
|
||||
INFERENCE_MODE = ...
|
||||
|
||||
def _set_creation_meta(t: torch.Tensor, creation_meta: CreationMeta) -> None: ...
|
||||
def _get_creation_meta(t: torch.Tensor) -> CreationMeta: ...
|
||||
@@ -0,0 +1,8 @@
|
||||
from typing import Any
|
||||
|
||||
from torch.types import _bool
|
||||
|
||||
# Defined in torch/csrc/cpu/Module.cpp
|
||||
|
||||
def _init_amx() -> _bool: ...
|
||||
def _get_cpu_capability() -> dict[str, Any]: ...
|
||||
@@ -0,0 +1,14 @@
|
||||
from enum import IntEnum
|
||||
|
||||
# Defined in torch/csrc/cuda/shared/cudnn.cpp
|
||||
is_cuda: bool
|
||||
|
||||
def getRuntimeVersion() -> tuple[int, int, int]: ...
|
||||
def getCompileVersion() -> tuple[int, int, int]: ...
|
||||
def getVersionInt() -> int: ...
|
||||
|
||||
class RNNMode(IntEnum):
|
||||
rnn_relu = ...
|
||||
rnn_tanh = ...
|
||||
lstm = ...
|
||||
gru = ...
|
||||
@@ -0,0 +1 @@
|
||||
def getVersionInt() -> int: ...
|
||||
@@ -0,0 +1,21 @@
|
||||
# This module is defined in torch/csrc/distributed/python_placement.cpp
|
||||
|
||||
class Placement:
|
||||
def is_partial(self, reduce_op: str | None = None) -> bool: ...
|
||||
def is_replicate(self) -> bool: ...
|
||||
def is_shard(self, dim: int | None = None) -> bool: ...
|
||||
|
||||
class Shard(Placement):
|
||||
dim: int
|
||||
def __init__(self, dim: int): ...
|
||||
|
||||
class StridedShard(Placement):
|
||||
dim: int
|
||||
split_factor: int
|
||||
def __init__(self, dim: int, *, split_factor: int): ...
|
||||
|
||||
class Replicate(Placement): ...
|
||||
|
||||
class Partial(Placement):
|
||||
reduce_op: str
|
||||
def __init__(self, reduce_op: str | None = None): ...
|
||||
@@ -0,0 +1,26 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
# This module is defined in torch/csrc/distributed/autograd/init.cpp
|
||||
|
||||
class DistAutogradContext:
|
||||
def _context_id(self) -> int: ...
|
||||
def _recv_functions(self) -> dict[int, Any]: ...
|
||||
def _send_functions(self) -> dict[int, Any]: ...
|
||||
def _known_worker_ids(self) -> set[int]: ...
|
||||
|
||||
def _new_context() -> DistAutogradContext: ...
|
||||
def _release_context(context_id: int) -> None: ...
|
||||
def _get_max_id() -> int: ...
|
||||
def _is_valid_context(worker_id: int) -> bool: ...
|
||||
def _retrieve_context(context_id: int) -> DistAutogradContext: ...
|
||||
def _current_context() -> DistAutogradContext: ...
|
||||
def _init(worker_id: int) -> None: ...
|
||||
def _get_debug_info() -> dict[str, str]: ...
|
||||
def backward(
|
||||
context_id: int,
|
||||
roots: list[torch.Tensor],
|
||||
retain_graph: bool = False,
|
||||
) -> None: ...
|
||||
def get_gradients(context_id: int) -> dict[torch.Tensor, torch.Tensor]: ...
|
||||
@@ -0,0 +1,900 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# mypy: disable-error-code="type-arg"
|
||||
from collections.abc import Callable
|
||||
from datetime import timedelta
|
||||
from enum import Enum
|
||||
from typing import Any, overload
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch._C import ScriptObject
|
||||
from torch._C._autograd import DeviceType
|
||||
from torch.distributed.distributed_c10d import GroupName
|
||||
from torch.futures import Future
|
||||
|
||||
# This module is defined in torch/csrc/distributed/c10d/init.cpp
|
||||
|
||||
_DEFAULT_FIRST_BUCKET_BYTES: int
|
||||
_DEFAULT_NO_TIMEOUT: timedelta
|
||||
_DEFAULT_PG_TIMEOUT: timedelta
|
||||
_DEFAULT_PG_NCCL_TIMEOUT: timedelta
|
||||
|
||||
class BuiltinCommHookType(Enum):
|
||||
ALLREDUCE = ...
|
||||
FP16_COMPRESS = ...
|
||||
|
||||
def _register_comm_hook(reducer: Reducer, state: Any, comm_hook: Any): ...
|
||||
def _register_builtin_comm_hook(
|
||||
reducer: Reducer,
|
||||
comm_hook_type: BuiltinCommHookType,
|
||||
): ...
|
||||
def _set_global_rank(rank: int) -> None: ...
|
||||
def _hash_tensors(tensors: list[Tensor]) -> int: ...
|
||||
|
||||
class GradBucket:
|
||||
def index(self) -> int: ...
|
||||
def buffer(self) -> Tensor: ...
|
||||
def gradients(self) -> list[Tensor]: ...
|
||||
def is_last(self) -> bool: ...
|
||||
def set_buffer(self, tensor: Tensor) -> None: ...
|
||||
def parameters(self) -> list[Tensor]: ...
|
||||
|
||||
class Reducer:
|
||||
def __init__(
|
||||
self,
|
||||
params: list[Tensor],
|
||||
bucket_indices: list[list[int]],
|
||||
per_bucket_size_limits: list[int],
|
||||
process_group: ProcessGroup,
|
||||
expect_sparse_gradients: list[bool] = ...,
|
||||
bucket_bytes_cap: int = ..., # kDefaultBucketBytesCap in reducer.hpp
|
||||
find_unused_parameters: bool = ...,
|
||||
gradient_as_bucket_view: bool = ...,
|
||||
param_to_name_mapping: dict[int, str] = ...,
|
||||
first_bucket_types_cap: int = ..., # kDefaultFirstBucketBytes in reducer.hpp
|
||||
skip_all_reduce_unused_params: bool = ...,
|
||||
use_python_reducer: bool = ...,
|
||||
bucket_bytes_cap_list: list[int] = ...,
|
||||
batched_grad_copy: bool = ...,
|
||||
) -> None: ...
|
||||
def prepare_for_forward(self) -> None: ...
|
||||
def prepare_for_backward(self, output: list[Tensor]) -> None: ...
|
||||
def get_backward_stats(self) -> list[int]: ...
|
||||
def _install_post_backward_futures(self, futures: list[Future]) -> None: ...
|
||||
def _rebuild_buckets(self) -> bool: ...
|
||||
def _get_zeros_like_grad_buckets(self) -> list[GradBucket]: ...
|
||||
def _push_all_rebuilt_params(self) -> None: ...
|
||||
def _set_forward_pass_work_handle(
|
||||
self,
|
||||
work: Work,
|
||||
use_static_world_size: bool,
|
||||
): ...
|
||||
def _get_local_used_map(self) -> Tensor: ...
|
||||
def _set_ddp_runtime_logging_sample_rate(self, sample_rate: int) -> None: ...
|
||||
def _set_static_graph(self) -> None: ...
|
||||
def _run_comm_hook(self, bucket: GradBucket) -> Future: ...
|
||||
def set_logger(self, logger: Logger) -> None: ...
|
||||
def _remove_autograd_hooks(self) -> None: ...
|
||||
def _check_reducer_finalized(self) -> None: ...
|
||||
def _set_sparse_metadata(self, global_unique_ids: dict[str, Tensor]) -> None: ...
|
||||
def _reset_state(self) -> None: ...
|
||||
def _update_process_group(self, new_process_group: ProcessGroup) -> None: ...
|
||||
|
||||
class DDPLoggingData:
|
||||
strs_map: dict[str, str]
|
||||
ints_map: dict[str, int]
|
||||
|
||||
class Logger:
|
||||
def __init__(self, reducer: Reducer) -> None: ...
|
||||
def set_construction_data_and_log(
|
||||
self,
|
||||
module_name: str,
|
||||
device_ids: list[int],
|
||||
output_device: int,
|
||||
broadcast_buffers: bool,
|
||||
has_sync_bn: bool,
|
||||
static_graph: bool,
|
||||
) -> None: ...
|
||||
def set_runtime_stats_and_log(self) -> None: ...
|
||||
def set_error_and_log(self, error: str) -> None: ...
|
||||
def _get_ddp_logging_data(self) -> DDPLoggingData: ...
|
||||
def _set_comm_hook_name(self, comm_hook: str) -> None: ...
|
||||
def _set_uneven_input_join(self) -> None: ...
|
||||
def _set_static_graph(self) -> None: ...
|
||||
|
||||
class _WorkerServer:
|
||||
port: int
|
||||
|
||||
def __init__(self, host_or_file: str, port: int = ...) -> None: ...
|
||||
def shutdown(self) -> None: ...
|
||||
|
||||
class DebugLevel(Enum):
|
||||
OFF = ...
|
||||
INFO = ...
|
||||
DETAIL = ...
|
||||
|
||||
def get_debug_level() -> DebugLevel: ...
|
||||
def set_debug_level(level: DebugLevel) -> None: ...
|
||||
def set_debug_level_from_env() -> None: ...
|
||||
|
||||
class ReduceOp:
|
||||
# pyrefly: ignore # unknown-name
|
||||
def __init__(self, op: RedOpType) -> None: ...
|
||||
|
||||
# pyrefly: ignore # unknown-name
|
||||
SUM: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
AVG: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
PRODUCT: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
MIN: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
MAX: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
BAND: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
BOR: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
BXOR: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
PREMUL_SUM: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
UNUSED: RedOpType = ...
|
||||
|
||||
# mypy error being ignored:
|
||||
# Detected enum "torch._C._distributed_c10d.ReduceOp.RedOpType" in a type
|
||||
# stub with zero members. There is a chance this is due to a recent change
|
||||
# in the semantics of enum membership. If so, use `member = value` to mark
|
||||
# an enum member, instead of `member: type`
|
||||
class RedOpType(Enum): ... # type: ignore[misc]
|
||||
|
||||
class BroadcastOptions:
|
||||
rootRank: int
|
||||
rootTensor: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class AllreduceOptions:
|
||||
reduceOp: ReduceOp
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
sparseIndices: Tensor | None
|
||||
|
||||
class AllreduceCoalescedOptions(AllreduceOptions): ...
|
||||
|
||||
class ReduceOptions:
|
||||
reduceOp: ReduceOp
|
||||
rootRank: int
|
||||
rootTensor: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class AllgatherOptions:
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class GatherOptions:
|
||||
rootRank: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class ScatterOptions:
|
||||
rootRank: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class ReduceScatterOptions:
|
||||
reduceOp: ReduceOp
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class BarrierOptions:
|
||||
device_ids: list[int]
|
||||
device: torch.device
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class AllToAllOptions:
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class Store:
|
||||
def set(self, key: str, value: str) -> None: ...
|
||||
def get(self, key: str) -> bytes: ...
|
||||
def add(self, key: str, value: int) -> int: ...
|
||||
def check(self, keys: list[str]) -> bool: ...
|
||||
def compare_set(
|
||||
self,
|
||||
key: str,
|
||||
expected_value: str,
|
||||
desired_value: str,
|
||||
) -> bytes: ...
|
||||
def delete_key(self, key: str) -> bool: ...
|
||||
def multi_get(self, keys: list[str]) -> list[bytes]: ...
|
||||
def num_keys(self) -> int: ...
|
||||
def set_timeout(self, timeout: timedelta) -> None: ...
|
||||
@overload
|
||||
def wait(self, keys: list[str]) -> None: ...
|
||||
@overload
|
||||
def wait(self, keys: list[str], timeout: timedelta) -> None: ...
|
||||
def queue_pop(self, key: str, block: bool = True) -> bytes: ...
|
||||
def queue_push(self, key: str, value: bytes | str) -> None: ...
|
||||
def queue_len(self, key: str) -> int: ...
|
||||
def list_keys(self) -> list[str]: ...
|
||||
|
||||
class FileStore(Store):
|
||||
def __init__(self, path: str, numWorkers: int = ...) -> None: ...
|
||||
|
||||
class HashStore(Store):
|
||||
def __init__(self) -> None: ...
|
||||
|
||||
class TCPStore(Store):
|
||||
def __init__(
|
||||
self,
|
||||
host_name: str,
|
||||
port: int,
|
||||
world_size: int | None = ...,
|
||||
is_master: bool = ...,
|
||||
timeout: timedelta = ...,
|
||||
wait_for_workers: bool = ...,
|
||||
multi_tenant: bool = ...,
|
||||
master_listen_fd: int | None = ...,
|
||||
use_libuv: bool | None = ...,
|
||||
) -> None: ...
|
||||
@property
|
||||
def host(self) -> str: ...
|
||||
@property
|
||||
def port(self) -> int: ...
|
||||
|
||||
class PrefixStore(Store):
|
||||
def __init__(self, prefix: str, store: Store) -> None: ...
|
||||
@property
|
||||
def underlying_store(self) -> Store: ...
|
||||
|
||||
class _ControlCollectives:
|
||||
def barrier(self, key: str, timeout: timedelta, blocking: bool) -> None: ...
|
||||
def broadcast_send(self, key: str, data: str, timeout: timedelta) -> None: ...
|
||||
def broadcast_recv(self, key: str, timeout: timedelta) -> str: ...
|
||||
def gather_send(self, key: str, data: str, timeout: timedelta) -> None: ...
|
||||
def gather_recv(self, key: str, timeout: timedelta) -> str: ...
|
||||
def scatter_send(self, key: str, data: str, timeout: timedelta) -> None: ...
|
||||
def scatter_recv(self, key: str, timeout: timedelta) -> str: ...
|
||||
def all_gather(self, key: str, data: str, timeout: timedelta) -> str: ...
|
||||
def all_sum(self, key: str, data: int, timeout: timedelta) -> int: ...
|
||||
|
||||
class _StoreCollectives(_ControlCollectives):
|
||||
def __init__(self, store: Store, rank: int, world_size: int) -> None: ...
|
||||
|
||||
class _DistributedBackendOptions:
|
||||
def __init__(self) -> None: ...
|
||||
@property
|
||||
def store(self) -> Store: ...
|
||||
@store.setter
|
||||
def store(self, store: Store) -> None: ...
|
||||
@property
|
||||
def group_rank(self) -> int: ...
|
||||
@group_rank.setter
|
||||
def group_rank(self, rank: int) -> None: ...
|
||||
@property
|
||||
def group_size(self) -> int: ...
|
||||
@group_size.setter
|
||||
def group_size(self, size: int) -> None: ...
|
||||
@property
|
||||
def timeout(self) -> timedelta: ...
|
||||
@timeout.setter
|
||||
def timeout(self, timeout: timedelta) -> None: ...
|
||||
@property
|
||||
def group_id(self) -> str: ...
|
||||
@group_id.setter
|
||||
def group_id(self, group_id: str) -> None: ...
|
||||
@property
|
||||
def global_ranks_in_group(self) -> list[int]: ...
|
||||
@global_ranks_in_group.setter
|
||||
def global_ranks_in_group(self, ranks: list[int]) -> None: ...
|
||||
|
||||
class Work:
|
||||
def is_completed(self) -> bool: ...
|
||||
def is_success(self) -> bool: ...
|
||||
def exception(self) -> Any: ...
|
||||
def wait(self, timeout: timedelta = ...) -> bool: ...
|
||||
def block_current_stream(self) -> None: ...
|
||||
def get_future(self) -> Future: ...
|
||||
def source_rank(self) -> int: ...
|
||||
def _source_rank(self) -> int: ...
|
||||
def result(self) -> list[Tensor]: ...
|
||||
def synchronize(self) -> None: ...
|
||||
def boxed(self) -> ScriptObject: ...
|
||||
@staticmethod
|
||||
def unbox(obj: ScriptObject) -> Work: ...
|
||||
|
||||
class Backend:
|
||||
class Options:
|
||||
def __init__(self, backend: str, timeout: timedelta = ...) -> None: ...
|
||||
@property
|
||||
def backend(self) -> str: ...
|
||||
@property
|
||||
def _timeout(self) -> timedelta: ...
|
||||
@_timeout.setter
|
||||
def _timeout(self, val: timedelta) -> None: ...
|
||||
global_ranks_in_group: list[int]
|
||||
group_name: GroupName
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
size: int,
|
||||
) -> None: ...
|
||||
@property
|
||||
def supports_splitting(self) -> bool: ...
|
||||
@property
|
||||
def supports_coalescing(self) -> bool: ...
|
||||
@property
|
||||
def supports_time_estimate(self) -> bool: ...
|
||||
def set_timeout(self, timeout: timedelta) -> None: ...
|
||||
@property
|
||||
def options(self) -> Options: ...
|
||||
def rank(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def name(self) -> str: ...
|
||||
def abort(self) -> None: ...
|
||||
def shutdown(self) -> None: ...
|
||||
def eager_connect_single_device(self, device: torch.device | None) -> None: ...
|
||||
def _set_sequence_number_for_group(self) -> None: ...
|
||||
def _set_default_timeout(self, timeout: timedelta) -> None: ...
|
||||
def get_error(self) -> ErrorType: ...
|
||||
def supports_tensor_alloc(self, device: torch.device) -> bool: ...
|
||||
def allocate_tensor(
|
||||
self,
|
||||
size: int,
|
||||
*,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> Tensor: ...
|
||||
@property
|
||||
def mem_allocator(self) -> Any: ...
|
||||
|
||||
class ProcessGroup:
|
||||
class BackendType(Enum):
|
||||
UNDEFINED = ...
|
||||
GLOO = ...
|
||||
NCCL = ...
|
||||
UCC = ...
|
||||
MPI = ...
|
||||
XCCL = ...
|
||||
CUSTOM = ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
) -> None: ...
|
||||
def rank(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def get_group_store(self) -> Store: ...
|
||||
def split_group(
|
||||
self,
|
||||
new_ranks: list[int],
|
||||
timeout: timedelta | None = None,
|
||||
opts: Backend.Options | None = None,
|
||||
group_name: GroupName | None = None,
|
||||
group_desc: str | None = None,
|
||||
) -> ProcessGroup | None: ...
|
||||
def merge_remote_group(
|
||||
self,
|
||||
store: Store,
|
||||
size: int,
|
||||
timeout: timedelta,
|
||||
group_name: GroupName | None = None,
|
||||
group_desc: str | None = None,
|
||||
) -> ProcessGroup: ...
|
||||
def abort(self) -> None: ...
|
||||
def set_timeout(self, timeout: timedelta) -> None: ...
|
||||
def shutdown(self) -> None: ...
|
||||
@overload
|
||||
def broadcast(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def broadcast(
|
||||
self,
|
||||
tensor: Tensor,
|
||||
root: int,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allreduce(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts: AllreduceOptions = ...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allreduce(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allreduce(
|
||||
self,
|
||||
tensor: Tensor,
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def allreduce_coalesced(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
def reduce_scatter_tensor_coalesced(
|
||||
self,
|
||||
outputTensors: list[Tensor],
|
||||
inputTensors: list[Tensor],
|
||||
opts: ReduceScatterOptions | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce(
|
||||
self,
|
||||
tensor: Tensor,
|
||||
root: int,
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allgather(
|
||||
self,
|
||||
output_tensors: list[list[Tensor]],
|
||||
input_tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allgather(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensor: Tensor,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def _allgather_base(
|
||||
self,
|
||||
output: Tensor,
|
||||
input: Tensor,
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
def allgather_coalesced(
|
||||
self,
|
||||
output_lists: list[list[Tensor]],
|
||||
input_list: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
def allgather_into_tensor_coalesced(
|
||||
self,
|
||||
output_lists: list[Tensor],
|
||||
input_list: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def gather(
|
||||
self,
|
||||
output_tensors: list[list[Tensor]],
|
||||
input_tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def gather(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensor: Tensor,
|
||||
root: int,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def scatter(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensors: list[list[Tensor]],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def scatter(
|
||||
self,
|
||||
output_tensor: Tensor,
|
||||
input_tensors: list[Tensor],
|
||||
root: int,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce_scatter(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensors: list[list[Tensor]],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce_scatter(
|
||||
self,
|
||||
output_tensors: Tensor,
|
||||
input_tensor: list[Tensor],
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def _reduce_scatter_base(
|
||||
self,
|
||||
outputTensor: Tensor,
|
||||
inputTensor: Tensor,
|
||||
opts: ReduceScatterOptions | None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall_base(
|
||||
self,
|
||||
output_tensor: Tensor,
|
||||
input_tensor: Tensor,
|
||||
output_split_sizes: list[int],
|
||||
input_split_sizes: list[int],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall_base(
|
||||
self,
|
||||
output: Tensor,
|
||||
input: Tensor,
|
||||
output_split_sizes: list[int],
|
||||
input_split_sizes: list[int],
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall(
|
||||
self,
|
||||
output_tensor: list[Tensor],
|
||||
input_tensor: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall(
|
||||
self,
|
||||
output: list[Tensor],
|
||||
input: list[Tensor],
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def send(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
dstRank: int,
|
||||
tag: int,
|
||||
) -> Work: ...
|
||||
def recv(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
srcRank: int,
|
||||
tag: int,
|
||||
) -> Work: ...
|
||||
def recv_anysource(self, tensors: list[Tensor], tag: int) -> Work: ...
|
||||
@overload
|
||||
def barrier(self, opts=...) -> Work: ...
|
||||
@overload
|
||||
def barrier(self, timeout: timedelta | None = None) -> Work: ...
|
||||
def boxed(self) -> ScriptObject: ...
|
||||
@staticmethod
|
||||
def unbox(obj: ScriptObject) -> ProcessGroup: ...
|
||||
def _start_coalescing(self, device: torch.device) -> None: ...
|
||||
def _end_coalescing(self, device: torch.device) -> Work: ...
|
||||
def _get_backend_name(self) -> str: ...
|
||||
def _backend_id(self, backend_type: BackendType) -> int: ...
|
||||
@property
|
||||
def _device_types(self) -> list[torch.device]: ...
|
||||
def _get_backend(self, device: torch.device) -> Backend: ...
|
||||
def _set_default_backend(self, backend_type: BackendType) -> None: ...
|
||||
def _register_backend(
|
||||
self,
|
||||
device: torch.device,
|
||||
backend_type: BackendType,
|
||||
backend: Backend | None,
|
||||
) -> None: ...
|
||||
def _set_group_name(self, name: GroupName) -> None: ...
|
||||
def _set_group_desc(self, desc: str) -> None: ...
|
||||
def name(self) -> str: ...
|
||||
def _has_hooks(self) -> bool: ...
|
||||
def _wait_for_pending_works(self) -> None: ...
|
||||
def _set_sequence_number_for_group(self) -> None: ...
|
||||
@property
|
||||
def bound_device_id(self) -> torch.device | None: ...
|
||||
@bound_device_id.setter
|
||||
def bound_device_id(self, device: torch.device | None) -> None: ...
|
||||
@property
|
||||
def group_name(self) -> GroupName: ...
|
||||
@property
|
||||
def group_desc(self) -> str: ...
|
||||
|
||||
class FakeProcessGroup(Backend):
|
||||
@staticmethod
|
||||
def _create_internal(rank: int, world_size: int) -> FakeProcessGroup: ...
|
||||
|
||||
class FakeWork(Work):
|
||||
seq_id: int
|
||||
def __init__(self) -> None: ...
|
||||
def wait(self, timeout: timedelta = ...) -> bool: ...
|
||||
def getFuture(self) -> Future: ...
|
||||
|
||||
class PythonCallbackWork(Work):
|
||||
def __init__(self, callback: Callable[[timedelta], bool]) -> None: ...
|
||||
def wait(self, timeout: timedelta = ...) -> bool: ...
|
||||
def get_future(self) -> Future: ...
|
||||
|
||||
class ProcessGroupGloo(Backend):
|
||||
class Device: ...
|
||||
|
||||
class Options(Backend.Options):
|
||||
devices: list[ProcessGroupGloo.Device]
|
||||
threads: int
|
||||
|
||||
def __init__(self): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
timeout: timedelta,
|
||||
) -> None: ...
|
||||
@staticmethod
|
||||
def create_device(hostname="", interface="", lazy_init=None) -> Device: ...
|
||||
@staticmethod
|
||||
def create_default_device(lazy_init=None) -> Device: ...
|
||||
def _set_default_timeout(self, timeout) -> None: ...
|
||||
@property
|
||||
def options(self) -> Options: ... # type: ignore[override]
|
||||
|
||||
class _ProcessGroupWrapper(Backend):
|
||||
def __init__(self, pg: Backend, gloo_pg: ProcessGroupGloo) -> None: ...
|
||||
wrapped_pg: Backend
|
||||
@property
|
||||
def options(self) -> Backend.Options: ...
|
||||
def get_error(self) -> ErrorType: ...
|
||||
|
||||
class ErrorType(Enum):
|
||||
SUCCESS = ...
|
||||
TIMEOUT = ...
|
||||
COMM_ERROR = ...
|
||||
REMOTE_ERROR = ...
|
||||
|
||||
class ProcessGroupNCCL(Backend):
|
||||
class NCCLConfig:
|
||||
blocking: int
|
||||
cga_cluster_size: int
|
||||
min_ctas: int
|
||||
max_ctas: int
|
||||
def unsafe_get_ptr(self) -> int: ...
|
||||
|
||||
class Options(Backend.Options):
|
||||
config: ProcessGroupNCCL.NCCLConfig
|
||||
is_high_priority_stream: bool
|
||||
split_from: ProcessGroupNCCL
|
||||
split_color: int
|
||||
|
||||
def __init__(self, is_high_priority_stream: bool = False): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
options: Options,
|
||||
) -> None: ...
|
||||
def _group_start(self) -> None: ...
|
||||
def _group_end(self) -> None: ...
|
||||
def _start_time_estimate(self) -> None: ...
|
||||
def _end_time_estimate(self) -> float: ...
|
||||
def _set_default_timeout(self, timeout) -> None: ...
|
||||
def perform_nocolor_split(self, device: torch.device) -> None: ...
|
||||
def register_mem_pool(self, pool: torch.cuda.MemPool) -> None: ...
|
||||
def deregister_mem_pool(self, pool: torch.cuda.MemPool) -> None: ...
|
||||
def comm_split_count(self) -> int: ...
|
||||
def _add_ephemeral_timeout(self, timeout: timedelta) -> None: ...
|
||||
def abort(self) -> None: ...
|
||||
def _is_initialized(self) -> bool: ...
|
||||
@property
|
||||
def uid(self) -> int: ...
|
||||
@property
|
||||
def options(self) -> Options: ... # type: ignore[override]
|
||||
@staticmethod
|
||||
def get_build_nccl_version(self) -> tuple[int, int, int]: ...
|
||||
@staticmethod
|
||||
def get_runtime_nccl_version(self) -> tuple[int, int, int]: ...
|
||||
|
||||
class ProcessGroupUCC(Backend):
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
timeout: timedelta,
|
||||
) -> None: ...
|
||||
|
||||
class ProcessGroupMPI(Backend):
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
size: int,
|
||||
pgComm: int,
|
||||
) -> None: ...
|
||||
@staticmethod
|
||||
def create(ranks: list[int]) -> ProcessGroupMPI: ...
|
||||
|
||||
def _compute_bucket_assignment_by_size(
|
||||
tensors: list[Tensor],
|
||||
bucket_size_limits: list[int],
|
||||
expect_sparse_gradient: list[bool] = ...,
|
||||
tensor_indices: list[int] = ...,
|
||||
) -> tuple[list[list[int]], list[int]]: ...
|
||||
def _broadcast_coalesced(
|
||||
process_group: ProcessGroup,
|
||||
tensors: list[Tensor],
|
||||
buffer_size: int,
|
||||
src: int,
|
||||
): ...
|
||||
def _test_python_store(store: Store): ...
|
||||
def _verify_params_across_processes(
|
||||
process_group: ProcessGroup,
|
||||
params: list[Tensor],
|
||||
logger: Logger | None,
|
||||
): ...
|
||||
def _make_nccl_premul_sum(factor: float | list[Tensor]) -> ReduceOp: ...
|
||||
def _register_process_group(
|
||||
group_name: GroupName,
|
||||
process_group: ProcessGroup,
|
||||
) -> None: ...
|
||||
def _resolve_process_group(group_name: GroupName) -> ProcessGroup: ...
|
||||
def _register_work(tensor: torch.Tensor, work: Work) -> ProcessGroup: ...
|
||||
def _get_work_registry_size() -> int: ...
|
||||
def _set_allow_inflight_collective_as_graph_input(
|
||||
value: bool,
|
||||
) -> None: ...
|
||||
def _allow_inflight_collective_as_graph_input() -> bool: ...
|
||||
def _unregister_all_process_groups() -> None: ...
|
||||
def _unregister_process_group(group_name: GroupName) -> None: ...
|
||||
|
||||
# Initializes the device state in CUmodule so that it's able to perform NVSHMEM
|
||||
# operations. CUmodule is a pointer to a CUDA module, carried by a int64 in
|
||||
# Python. At C++ interface, it is converted to a uintptr_t.
|
||||
def _nvshmemx_cumodule_init(module: int) -> None: ...
|
||||
|
||||
# Check if NVSHMEM is available on current system.
|
||||
def _is_nvshmem_available() -> bool: ...
|
||||
|
||||
class _SymmetricMemory:
|
||||
@staticmethod
|
||||
def set_group_info(
|
||||
group_name: str,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
store: Store,
|
||||
) -> None: ...
|
||||
@staticmethod
|
||||
def empty_strided_p2p(
|
||||
size: torch.types._size,
|
||||
stride: torch.types._size,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
group_name: str | None = None,
|
||||
alloc_id: int | None = None,
|
||||
) -> torch.Tensor: ...
|
||||
@staticmethod
|
||||
def has_multicast_support(
|
||||
device_type: DeviceType,
|
||||
device_idx: int,
|
||||
) -> bool: ...
|
||||
# Set Symmetric Memory allocation backend.
|
||||
@staticmethod
|
||||
def set_backend(name: str) -> None: ...
|
||||
@staticmethod
|
||||
def get_backend(device: torch.device) -> str | None: ...
|
||||
@staticmethod
|
||||
def is_symm_mem_tensor(tensor: torch.Tensor) -> bool: ...
|
||||
@staticmethod
|
||||
def get_mempool_allocator(device: torch.device) -> Any: ...
|
||||
signal_pad_size: int
|
||||
@property
|
||||
def rank(self) -> int: ...
|
||||
@property
|
||||
def world_size(self) -> int: ...
|
||||
@staticmethod
|
||||
def rendezvous(
|
||||
tensor: torch.Tensor, group_name: str | None = None
|
||||
) -> _SymmetricMemory: ...
|
||||
def get_buffer(
|
||||
self,
|
||||
rank: int,
|
||||
sizes: torch.types._size,
|
||||
dtype: torch.dtype,
|
||||
storage_offset: int | None = 0,
|
||||
) -> torch.Tensor: ...
|
||||
def get_signal_pad(
|
||||
self,
|
||||
rank: int,
|
||||
sizes: torch.types._size = [],
|
||||
dtype: torch.dtype | None = None,
|
||||
storage_offset: int | None = 0,
|
||||
) -> torch.Tensor: ...
|
||||
def barrier(self, channel: int = 0, timeout_ms: int = 0) -> None: ...
|
||||
def put_signal(
|
||||
self,
|
||||
dst_rank: int,
|
||||
channel: int = 0,
|
||||
timeout_ms: int = 0,
|
||||
) -> None: ...
|
||||
def wait_signal(
|
||||
self,
|
||||
src_rank: int,
|
||||
channel: int = 0,
|
||||
timeout_ms: int = 0,
|
||||
) -> None: ...
|
||||
def get_remote_tensor(
|
||||
self,
|
||||
peer: int,
|
||||
sizes: torch.types._size,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor: ...
|
||||
@staticmethod
|
||||
def memset32(
|
||||
tensor: torch.Tensor, offset: int, val: int, count: int = 1
|
||||
) -> torch.Tensor: ...
|
||||
@staticmethod
|
||||
def stream_write_value32(
|
||||
tensor: torch.Tensor, offset: int, val: int
|
||||
) -> torch.Tensor: ...
|
||||
@property
|
||||
def buffer_ptrs(self) -> list[int]: ...
|
||||
@property
|
||||
def buffer_ptrs_dev(self) -> int: ...
|
||||
@property
|
||||
def signal_pad_ptrs(self) -> list[int]: ...
|
||||
@property
|
||||
def signal_pad_ptrs_dev(self) -> int: ...
|
||||
@property
|
||||
def multicast_ptr(self) -> int: ...
|
||||
@property
|
||||
def buffer_size(self) -> int: ...
|
||||
@property
|
||||
def device(self) -> torch.device: ...
|
||||
|
||||
class ProcessGroupXCCL(Backend):
|
||||
class Options(Backend.Options):
|
||||
is_high_priority_stream: bool
|
||||
|
||||
def __init__(self, is_high_priority_stream: bool = False): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
options: Options,
|
||||
) -> None: ...
|
||||
@property
|
||||
def options(self) -> Options: ... # type: ignore[override]
|
||||
|
||||
def _set_process_group(pg: ProcessGroup) -> None: ...
|
||||
def _current_process_group() -> ProcessGroup: ...
|
||||
|
||||
class _Request:
|
||||
def body(self) -> bytes: ...
|
||||
def get_param(self, str) -> str: ...
|
||||
|
||||
class _Response:
|
||||
def set_content(self, content: str | bytes, content_type: str) -> None: ...
|
||||
def set_status(self, status: int) -> None: ...
|
||||
|
||||
def _register_handler(
|
||||
name: str, handler: Callable[[_Request, _Response], None]
|
||||
) -> None: ...
|
||||
def _set_comm_profiling_name(name: str) -> None: ...
|
||||
def _get_comm_profiling_name() -> str: ...
|
||||
@@ -0,0 +1,188 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# mypy: disable-error-code="type-arg"
|
||||
from datetime import timedelta
|
||||
from typing import Any, Generic, overload, TypeVar
|
||||
|
||||
import torch
|
||||
from torch._C import Future
|
||||
from torch._C._autograd import ProfilerEvent
|
||||
from torch._C._distributed_c10d import Store
|
||||
from torch._C._profiler import ProfilerConfig
|
||||
|
||||
# This module is defined in torch/csrc/distributed/rpc/init.cpp
|
||||
|
||||
_DEFAULT_INIT_METHOD: str
|
||||
_DEFAULT_NUM_WORKER_THREADS: int
|
||||
_UNSET_RPC_TIMEOUT: float
|
||||
_DEFAULT_RPC_TIMEOUT_SEC: float
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
class RpcBackendOptions:
|
||||
rpc_timeout: float
|
||||
init_method: str
|
||||
def __init__(
|
||||
self,
|
||||
rpc_timeout: float = ...,
|
||||
init_method: str = ...,
|
||||
) -> None: ...
|
||||
|
||||
class WorkerInfo:
|
||||
def __init__(self, name: str, worker_id: int) -> None: ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
@property
|
||||
def id(self) -> int: ...
|
||||
def __eq__(self, other: object) -> bool: ...
|
||||
|
||||
class RpcAgent:
|
||||
def join(self, shutdown: bool = False, timeout: float = 0): ...
|
||||
def sync(self): ...
|
||||
def shutdown(self): ...
|
||||
@overload
|
||||
def get_worker_info(self) -> WorkerInfo: ...
|
||||
@overload
|
||||
def get_worker_info(self, workerName: str) -> WorkerInfo: ...
|
||||
def get_worker_infos(self) -> list[WorkerInfo]: ...
|
||||
def _get_device_map(self, dst: WorkerInfo) -> dict[torch.device, torch.device]: ...
|
||||
def get_debug_info(self) -> dict[str, str]: ...
|
||||
def get_metrics(self) -> dict[str, str]: ...
|
||||
|
||||
class PyRRef(Generic[_T]):
|
||||
def __init__(self, value: _T, type_hint: Any = None) -> None: ...
|
||||
def is_owner(self) -> bool: ...
|
||||
def confirmed_by_owner(self) -> bool: ...
|
||||
def owner(self) -> WorkerInfo: ...
|
||||
def owner_name(self) -> str: ...
|
||||
def to_here(self, timeout: float = ...) -> _T: ...
|
||||
def local_value(self) -> Any: ...
|
||||
def rpc_sync(self, timeout: float = ...) -> Any: ...
|
||||
def rpc_async(self, timeout: float = ...) -> Any: ...
|
||||
def remote(self, timeout: float = ...) -> Any: ...
|
||||
def _serialize(self) -> tuple: ...
|
||||
@staticmethod
|
||||
def _deserialize(tp: tuple) -> PyRRef: ...
|
||||
def _get_type(self) -> type[_T]: ...
|
||||
def _get_future(self) -> Future[_T]: ...
|
||||
def _get_profiling_future(self) -> Future[_T]: ...
|
||||
def _set_profiling_future(self, profilingFuture: Future[_T]): ...
|
||||
|
||||
class _TensorPipeRpcBackendOptionsBase(RpcBackendOptions):
|
||||
num_worker_threads: int
|
||||
device_maps: dict[str, dict[torch.device, torch.device]]
|
||||
devices: list[torch.device]
|
||||
def __init__(
|
||||
self,
|
||||
num_worker_threads: int,
|
||||
_transports: list | None,
|
||||
_channels: list | None,
|
||||
rpc_timeout: float = ...,
|
||||
init_method: str = ...,
|
||||
device_maps: dict[str, dict[torch.device, torch.device]] = {}, # noqa: B006
|
||||
devices: list[torch.device] = [], # noqa: B006
|
||||
) -> None: ...
|
||||
def _set_device_map(
|
||||
self,
|
||||
to: str,
|
||||
device_map: dict[torch.device, torch.device],
|
||||
): ...
|
||||
|
||||
class TensorPipeAgent(RpcAgent):
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
name: str,
|
||||
worker_id: int,
|
||||
world_size: int | None,
|
||||
opts: _TensorPipeRpcBackendOptionsBase,
|
||||
reverse_device_maps: dict[str, dict[torch.device, torch.device]],
|
||||
devices: list[torch.device],
|
||||
) -> None: ...
|
||||
def join(self, shutdown: bool = False, timeout: float = 0): ...
|
||||
def shutdown(self): ...
|
||||
@overload
|
||||
def get_worker_info(self) -> WorkerInfo: ...
|
||||
@overload
|
||||
def get_worker_info(self, workerName: str) -> WorkerInfo: ...
|
||||
@overload
|
||||
def get_worker_info(self, id: int) -> WorkerInfo: ...
|
||||
def get_worker_infos(self) -> list[WorkerInfo]: ...
|
||||
def _get_device_map(self, dst: WorkerInfo) -> dict[torch.device, torch.device]: ...
|
||||
def _update_group_membership(
|
||||
self,
|
||||
worker_info: WorkerInfo,
|
||||
my_devices: list[torch.device],
|
||||
reverse_device_map: dict[str, dict[torch.device, torch.device]],
|
||||
is_join: bool,
|
||||
): ...
|
||||
def _get_backend_options(self) -> _TensorPipeRpcBackendOptionsBase: ...
|
||||
@property
|
||||
def is_static_group(self) -> bool: ...
|
||||
@property
|
||||
def store(self) -> Store: ...
|
||||
|
||||
def _is_current_rpc_agent_set() -> bool: ...
|
||||
def _get_current_rpc_agent() -> RpcAgent: ...
|
||||
def _set_and_start_rpc_agent(agent: RpcAgent): ...
|
||||
def _reset_current_rpc_agent(): ...
|
||||
def _delete_all_user_and_unforked_owner_rrefs(timeout: timedelta = ...): ...
|
||||
def _destroy_rref_context(ignoreRRefLeak: bool): ...
|
||||
def _rref_context_get_debug_info() -> dict[str, str]: ...
|
||||
def _cleanup_python_rpc_handler(): ...
|
||||
def _invoke_rpc_builtin(
|
||||
dst: WorkerInfo,
|
||||
opName: str,
|
||||
rpcTimeoutSeconds: float,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
): ...
|
||||
def _invoke_rpc_python_udf(
|
||||
dst: WorkerInfo,
|
||||
pickledPythonUDF: str,
|
||||
tensors: list[torch.Tensor],
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
): ...
|
||||
def _invoke_rpc_torchscript(
|
||||
dstWorkerName: str,
|
||||
qualifiedNameStr: str,
|
||||
argsTuple: tuple,
|
||||
kwargsDict: dict,
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
): ...
|
||||
def _invoke_remote_builtin(
|
||||
dst: WorkerInfo,
|
||||
opName: str,
|
||||
rpcTimeoutSeconds: float,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
): ...
|
||||
def _invoke_remote_python_udf(
|
||||
dst: WorkerInfo,
|
||||
pickledPythonUDF: str,
|
||||
tensors: list[torch.Tensor],
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
): ...
|
||||
def _invoke_remote_torchscript(
|
||||
dstWorkerName: WorkerInfo,
|
||||
qualifiedNameStr: str,
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
): ...
|
||||
def get_rpc_timeout() -> float: ...
|
||||
def enable_gil_profiling(flag: bool): ...
|
||||
def _set_rpc_timeout(rpcTimeoutSeconds: float): ...
|
||||
|
||||
class RemoteProfilerManager:
|
||||
@staticmethod
|
||||
def set_current_profiling_key(key: str): ...
|
||||
|
||||
def _enable_server_process_global_profiler(new_config: ProfilerConfig): ...
|
||||
def _disable_server_process_global_profiler() -> list[list[list[ProfilerEvent]]]: ...
|
||||
def _set_profiler_node_id(default_node_id: int): ...
|
||||
def _enable_jit_rref_pickle(): ...
|
||||
def _disable_jit_rref_pickle(): ...
|
||||
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
from torch._C._distributed_c10d import Store
|
||||
from torch._C._distributed_rpc import _TensorPipeRpcBackendOptionsBase, TensorPipeAgent
|
||||
|
||||
# This module is defined in torch/csrc/distributed/rpc/testing/init.cpp
|
||||
|
||||
class FaultyTensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):
|
||||
def __init__(
|
||||
self,
|
||||
num_worker_threads: int,
|
||||
rpc_timeout: float,
|
||||
init_method: str,
|
||||
messages_to_fail: list[str],
|
||||
messages_to_delay: dict[str, float],
|
||||
num_fail_sends: int,
|
||||
) -> None: ...
|
||||
num_send_recv_threads: int
|
||||
messages_to_fail: list[str]
|
||||
messages_to_delay: dict[str, float]
|
||||
num_fail_sends: int
|
||||
|
||||
class FaultyTensorPipeAgent(TensorPipeAgent):
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
name: str,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
options: FaultyTensorPipeRpcBackendOptions,
|
||||
reverse_device_maps: dict[str, dict[torch.device, torch.device]],
|
||||
devices: list[torch.device],
|
||||
) -> None: ...
|
||||
@@ -0,0 +1,69 @@
|
||||
from . import compiled_autograd, eval_frame, guards # noqa: F401
|
||||
|
||||
def strip_function_call(name: str) -> str: ...
|
||||
def is_valid_var_name(name: str) -> bool | int: ...
|
||||
def get_type_slots(obj: type | object) -> tuple[int, int, int, int]: ...
|
||||
def has_slot(slots: int, slot_bit: int) -> bool: ...
|
||||
|
||||
class PySequenceSlots:
|
||||
SQ_LENGTH: int
|
||||
SQ_CONCAT: int
|
||||
SQ_REPEAT: int
|
||||
SQ_ITEM: int
|
||||
SQ_CONTAINS: int
|
||||
SQ_ASS_ITEM: int
|
||||
SQ_INPLACE_CONCAT: int
|
||||
SQ_INPLACE_REPEAT: int
|
||||
|
||||
class PyMappingSlots:
|
||||
MP_LENGTH: int
|
||||
MP_SUBSCRIPT: int
|
||||
MP_ASS_SUBSCRIPT: int
|
||||
|
||||
class PyNumberSlots:
|
||||
NB_ADD: int
|
||||
NB_SUBTRACT: int
|
||||
NB_MULTIPLY: int
|
||||
NB_REMAINDER: int
|
||||
NB_POWER: int
|
||||
NB_NEGATIVE: int
|
||||
NB_POSITIVE: int
|
||||
NB_ABSOLUTE: int
|
||||
NB_BOOL: int
|
||||
NB_INVERT: int
|
||||
NB_LSHIFT: int
|
||||
NB_RSHIFT: int
|
||||
NB_AND: int
|
||||
NB_XOR: int
|
||||
NB_OR: int
|
||||
NB_INT: int
|
||||
NB_FLOAT: int
|
||||
NB_INPLACE_ADD: int
|
||||
NB_INPLACE_SUBTRACT: int
|
||||
NB_INPLACE_MULTIPLY: int
|
||||
NB_INPLACE_REMAINDER: int
|
||||
NB_INPLACE_POWER: int
|
||||
NB_INPLACE_LSHIFT: int
|
||||
NB_INPLACE_RSHIFT: int
|
||||
NB_INPLACE_AND: int
|
||||
NB_INPLACE_XOR: int
|
||||
NB_INPLACE_OR: int
|
||||
NB_FLOOR_DIVIDE: int
|
||||
NB_TRUE_DIVIDE: int
|
||||
NB_INPLACE_FLOOR_DIVIDE: int
|
||||
NB_INPLACE_TRUE_DIVIDE: int
|
||||
NB_INDEX: int
|
||||
NB_MATRIX_MULTIPLY: int
|
||||
NB_INPLACE_MATRIX_MULTIPLY: int
|
||||
|
||||
class PyTypeSlots:
|
||||
TP_HASH: int
|
||||
TP_ITER: int
|
||||
TP_ITERNEXT: int
|
||||
TP_CALL: int
|
||||
TP_REPR: int
|
||||
TP_RICHCOMPARE: int
|
||||
TP_GETATTRO: int
|
||||
TP_SETATTRO: int
|
||||
TP_DESCR_GET: int
|
||||
TP_DESCR_SET: int
|
||||
@@ -0,0 +1,13 @@
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch import Tensor
|
||||
from torch._dynamo.compiled_autograd import AutogradCompilerInstance
|
||||
|
||||
def set_autograd_compiler(
|
||||
autograd_compiler: Callable[[], AutogradCompilerInstance] | None,
|
||||
dynamic: bool,
|
||||
) -> tuple[Callable[[], AutogradCompilerInstance] | None, bool]: ...
|
||||
def clear_cache() -> None: ...
|
||||
def is_cache_empty() -> bool: ...
|
||||
def set_verbose_logger(fn: Callable[[str], None] | None) -> bool: ...
|
||||
def call_cpp_tensor_pre_hooks(idx: int, grad: Tensor) -> Tensor: ...
|
||||
@@ -0,0 +1,99 @@
|
||||
import enum
|
||||
import types
|
||||
from collections.abc import Callable
|
||||
from typing import overload
|
||||
|
||||
from torch._dynamo.guards import GuardManagerWrapper
|
||||
from torch._dynamo.types import DynamoCallback, DynamoGuardCompleteHook, DynamoGuardHook
|
||||
from torch._guards import CompileId
|
||||
|
||||
def set_eval_frame(callback: DynamoCallback) -> DynamoCallback: ...
|
||||
def set_skip_guard_eval_unsafe(value: bool) -> bool: ...
|
||||
def get_eval_frame_callback() -> DynamoCallback: ...
|
||||
def reset_code(code: types.CodeType) -> None: ...
|
||||
def unsupported(obj1: object, obj2: object) -> object: ...
|
||||
def set_code_exec_strategy(
|
||||
code: types.CodeType, strategy: _FrameExecStrategy
|
||||
) -> None: ...
|
||||
def set_guard_error_hook(hook: DynamoGuardHook) -> None: ...
|
||||
def set_guard_complete_hook(
|
||||
hook: DynamoGuardCompleteHook | None,
|
||||
) -> DynamoGuardCompleteHook | None: ...
|
||||
def raise_sigtrap() -> None: ...
|
||||
def set_c_recursion_limit(limit: int) -> None: ...
|
||||
def get_c_recursion_limit() -> int: ...
|
||||
|
||||
class _CacheEntry:
|
||||
def check_fn(self, *args: object, **kwargs: object) -> bool: ...
|
||||
def update_diff_guard_root_manager(self) -> None: ...
|
||||
code: types.CodeType
|
||||
compile_id: CompileId
|
||||
# If we run into circular issues, just use object
|
||||
guard_manager: GuardManagerWrapper
|
||||
backend: Callable
|
||||
next: _CacheEntry | None
|
||||
|
||||
class _PrecompileEntry:
|
||||
guard_manager: GuardManagerWrapper
|
||||
|
||||
class _ExtraState:
|
||||
def invalidate(
|
||||
self, cache_entry: _CacheEntry, guard_manager: GuardManagerWrapper
|
||||
) -> None: ...
|
||||
|
||||
class _FrameAction(enum.IntEnum):
|
||||
DEFAULT = 0
|
||||
SKIP = 1
|
||||
RUN_ONLY = 2
|
||||
|
||||
class _FrameExecStrategy:
|
||||
cur_action: _FrameAction
|
||||
recursive_action: _FrameAction
|
||||
|
||||
@overload
|
||||
def __init__(self) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, cur_action: _FrameAction, recursive_action: _FrameAction
|
||||
) -> None: ...
|
||||
|
||||
# This is an object that encapsulates the Python FrameType, and exposes
|
||||
# properties Dynamo cares about for a frame.
|
||||
class _PyInterpreterFrame:
|
||||
f_code: types.CodeType
|
||||
f_locals: dict[str, object]
|
||||
f_globals: dict[str, object]
|
||||
f_builtins: dict[str, object]
|
||||
f_lasti: int
|
||||
f_lineno: int
|
||||
f_back: types.FrameType
|
||||
# A tuple containing cell objects captured by this frame.
|
||||
closure: tuple[types.CellType]
|
||||
|
||||
def _debug_get_cache_entry_list(code: types.CodeType) -> list[_CacheEntry]: ...
|
||||
def _get_frame_value_stack_with_depth(
|
||||
frame: types.FrameType, depth: int
|
||||
) -> list[object]: ...
|
||||
def set_bytecode_debugger_callback(
|
||||
callback: Callable[[types.CodeType], None] | None,
|
||||
) -> None: ...
|
||||
def get_bytecode_debugger_callback() -> Callable[[types.CodeType], None] | None: ...
|
||||
def register_breakpoint_code(code: types.CodeType) -> None: ...
|
||||
|
||||
# Sentinel for NULL stack values returned by _get_frame_value_stack_at_depth
|
||||
class NullStackValue: ...
|
||||
|
||||
NULL_STACK_VALUE: NullStackValue
|
||||
|
||||
py_opcode_caches: list[int]
|
||||
|
||||
def code_framelocals_names(code: types.CodeType) -> tuple[str, ...]: ...
|
||||
def _load_precompile_entry(
|
||||
code: types.CodeType,
|
||||
guard_manager: GuardManagerWrapper,
|
||||
dynamo_code: types.CodeType,
|
||||
) -> None: ...
|
||||
def _reset_precompile_entries(code: types.CodeType) -> None: ...
|
||||
def _debug_get_precompile_entries(code: types.CodeType) -> list[_PrecompileEntry]: ...
|
||||
def set_fullgraph_compiled_frame_count(value: int) -> int: ...
|
||||
def set_fullgraph_error_on_nested_compile(value: bool) -> bool: ...
|
||||
@@ -0,0 +1,500 @@
|
||||
import enum
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TypeAlias
|
||||
|
||||
import torch
|
||||
|
||||
# TODO: We should move the `GuardManagerType`
|
||||
# defined in `guards.py` here and update other
|
||||
# imports
|
||||
GuardManagerType: TypeAlias = enum.Enum
|
||||
|
||||
class GlobalStateGuard:
|
||||
def check(self) -> bool: ...
|
||||
def reason(self) -> str: ...
|
||||
|
||||
class LeafGuard:
|
||||
def verbose_code_parts(self) -> list[str]: ...
|
||||
|
||||
class RelationalGuard: ...
|
||||
|
||||
class GuardDebugInfo:
|
||||
verbose_code_parts: list[str]
|
||||
result: bool
|
||||
num_guards_executed: int
|
||||
user_stack: traceback.StackSummary | None
|
||||
|
||||
class GuardManager:
|
||||
def check(self, value: Any) -> bool: ...
|
||||
def check_verbose(self, value: Any) -> GuardDebugInfo: ...
|
||||
|
||||
# Accessors
|
||||
def globals_dict_manager(
|
||||
self,
|
||||
f_globals: dict[str, Any],
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def framelocals_manager(
|
||||
self,
|
||||
key: tuple[str, int],
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def dict_getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def grad_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def generic_getattr_manager(
|
||||
self,
|
||||
attr: str,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_generic_dict_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def list_getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tuple_getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def set_getitem_manager(
|
||||
self,
|
||||
index: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def func_defaults_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def func_kwdefaults_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tuple_iterator_getitem_manager(
|
||||
self,
|
||||
index: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def weakref_call_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def call_function_no_args_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def global_weakref_manager(
|
||||
self,
|
||||
global_name: str,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def type_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def getattr_manager(
|
||||
self,
|
||||
attr: str,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tensor_property_size_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tensor_property_shape_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tensor_property_storage_offset_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def indexed_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def lambda_manager(
|
||||
self,
|
||||
python_lambda: Callable[..., Any],
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_root(self) -> RootGuardManager: ...
|
||||
def get_source(self) -> str: ...
|
||||
def fail_count(self) -> int: ...
|
||||
def get_child_managers(self) -> list[GuardManager]: ...
|
||||
def repr(self) -> str: ...
|
||||
def type_of_guarded_value(self) -> str: ...
|
||||
def get_leaf_guards(self) -> list[LeafGuard]: ...
|
||||
def get_accessors(self) -> list[GuardManager]: ...
|
||||
def is_guarded_value_immutable(self) -> bool: ...
|
||||
def is_tag_safe(self) -> bool: ...
|
||||
def is_tag_safe_root(self) -> bool: ...
|
||||
def has_no_accessors(self) -> bool: ...
|
||||
def has_object_aliasing_guard(self) -> bool: ...
|
||||
def get_type_of_guarded_value(self) -> type: ...
|
||||
def type_dict_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def type_mro_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def code_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def closure_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
# Leaf guards
|
||||
def add_lambda_guard(
|
||||
self,
|
||||
user_lambda: Callable[..., Any],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_id_match_guard(
|
||||
self,
|
||||
id_val: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_equals_match_guard(
|
||||
self,
|
||||
equals_val: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_global_state_guard(
|
||||
self,
|
||||
initial_state: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_torch_function_mode_stack_guard(
|
||||
self,
|
||||
initial_stack: list[Any],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_mapping_keys_guard(
|
||||
self,
|
||||
value: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dict_length_check_guard(
|
||||
self,
|
||||
value: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_length_check_guard(
|
||||
self,
|
||||
value: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_true_match_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_false_match_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_none_match_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_not_none_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dispatch_key_set_guard(
|
||||
self,
|
||||
dispatch_key: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_tensor_match_guard(
|
||||
self,
|
||||
value: Any,
|
||||
sizes: list[int],
|
||||
strides: list[int],
|
||||
tensor_name: str,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
ptype: Any,
|
||||
dispatch_keys: Any,
|
||||
) -> None: ...
|
||||
def add_dynamic_indices_guard(
|
||||
self,
|
||||
value: set[Any],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_no_hasattr_guard(
|
||||
self,
|
||||
attr_name: str,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dict_contains_guard(
|
||||
self,
|
||||
contains: bool,
|
||||
key: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_type_match_guard(
|
||||
self,
|
||||
value: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dict_version_guard(
|
||||
self,
|
||||
value: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_set_contains_guard(
|
||||
self,
|
||||
contains: bool,
|
||||
item: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dual_level_match_guard(
|
||||
self,
|
||||
level: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_float_is_nan_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_complex_is_nan_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_tuple_iterator_length_guard(
|
||||
self,
|
||||
length: int,
|
||||
type_id: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_range_iterator_match_guard(
|
||||
self,
|
||||
start: int,
|
||||
stop: int,
|
||||
step: int,
|
||||
type_id: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_default_device_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def mark_tag_safe(self) -> None: ...
|
||||
def mark_tag_safe_root(self) -> None: ...
|
||||
|
||||
class RootGuardManager(GuardManager):
|
||||
def get_epilogue_lambda_guards(self) -> list[LeafGuard]: ...
|
||||
def add_epilogue_lambda_guard(
|
||||
self,
|
||||
guard: LeafGuard,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def clone_manager(
|
||||
self, clone_filter_fn: Callable[[GuardManager], bool]
|
||||
) -> RootGuardManager: ...
|
||||
def attach_compile_id(self, compile_id: str) -> None: ...
|
||||
|
||||
class DictGuardManager(GuardManager):
|
||||
def get_key_manager(
|
||||
self,
|
||||
index: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_value_manager(
|
||||
self,
|
||||
index: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_key_value_managers(
|
||||
self,
|
||||
) -> dict[int, tuple[GuardManager, GuardManager]]: ...
|
||||
|
||||
# Guard accessor stubs
|
||||
class GuardAccessor: ...
|
||||
class DictGetItemGuardAccessor(GuardAccessor): ...
|
||||
class GetGenericDictGuardAccessor(GuardAccessor): ...
|
||||
class TypeDictGuardAccessor(GuardAccessor): ...
|
||||
class TypeMROGuardAccessor(GuardAccessor): ...
|
||||
class ClosureGuardAccessor(GuardAccessor): ...
|
||||
class TupleGetItemGuardAccessor(GuardAccessor): ...
|
||||
class TypeGuardAccessor(GuardAccessor): ...
|
||||
class CodeGuardAccessor(GuardAccessor): ...
|
||||
class FuncDefaultsGuardAccessor(GuardAccessor): ...
|
||||
class FuncKwDefaultsGuardAccessor(GuardAccessor): ...
|
||||
|
||||
class GetAttrGuardAccessor(GuardAccessor):
|
||||
def get_attr_name(self) -> str: ...
|
||||
|
||||
def install_object_aliasing_guard(
|
||||
x: GuardManager,
|
||||
y: GuardManager,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def install_no_tensor_aliasing_guard(
|
||||
guard_managers: list[GuardManager],
|
||||
tensor_names: list[str],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def install_storage_overlapping_guard(
|
||||
overlapping_guard_managers: list[GuardManager],
|
||||
non_overlapping_guard_managers: list[GuardManager],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def install_symbolic_shape_guard(
|
||||
guard_managers: list[GuardManager],
|
||||
nargs_int: int,
|
||||
nargs_float: int,
|
||||
py_addr: int,
|
||||
py_addr_keep_alive: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def profile_guard_manager(
|
||||
guard_manager: GuardManager,
|
||||
f_locals: dict[str, Any],
|
||||
n_iters: int,
|
||||
) -> float: ...
|
||||
|
||||
class TensorGuards:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dynamic_dims_sizes: list[torch.SymInt | None] | None = None,
|
||||
dynamic_dims_strides: list[torch.SymInt | None] | None = None,
|
||||
) -> None: ...
|
||||
def check(self, *args: Any) -> bool: ...
|
||||
def check_verbose(
|
||||
self, *args: Any, tensor_check_names: list[str] | None = None
|
||||
) -> bool | str: ...
|
||||
|
||||
def assert_size_stride(
|
||||
item: torch.Tensor,
|
||||
size: torch.types._size,
|
||||
stride: torch.types._size,
|
||||
op_name: str | None = None,
|
||||
) -> None: ...
|
||||
def assert_alignment(
|
||||
item: torch.Tensor,
|
||||
alignment: int,
|
||||
op_name: str | None = None,
|
||||
) -> None: ...
|
||||
def copy_misaligned(item: torch.Tensor) -> torch.Tensor: ...
|
||||
def check_obj_id(obj: object, expected: int) -> bool: ...
|
||||
def check_type_id(obj: object, expected: int) -> bool: ...
|
||||
def dict_version(d: dict[Any, Any]) -> int: ...
|
||||
def compute_overlapping_tensors(
|
||||
tensors: list[torch.Tensor], symbolic: bool = True
|
||||
) -> set[int]: ...
|
||||
def set_is_in_mode_without_ignore_compile_internals(value: bool) -> None: ...
|
||||
@@ -0,0 +1,9 @@
|
||||
# Defined in torch/csrc/export/pybind.cpp
|
||||
class CppExportedProgram: ...
|
||||
|
||||
def deserialize_exported_program(
|
||||
serialized_program: str,
|
||||
) -> CppExportedProgram: ...
|
||||
def serialize_exported_program(
|
||||
cpp_exported_program: CppExportedProgram,
|
||||
) -> str: ...
|
||||
@@ -0,0 +1,25 @@
|
||||
# Defined in torch/csrc/export/pt2_archive_constants.h
|
||||
|
||||
ARCHIVE_ROOT_NAME: str = ...
|
||||
ARCHIVE_FORMAT_PATH: str = ...
|
||||
ARCHIVE_FORMAT_VALUE: str = ...
|
||||
ARCHIVE_VERSION_PATH: str = ...
|
||||
ARCHIVE_VERSION_VALUE: str = ...
|
||||
MODELS_DIR: str = ...
|
||||
MODELS_FILENAME_FORMAT: str = ...
|
||||
AOTINDUCTOR_DIR: str = ...
|
||||
MTIA_DIR: str = ...
|
||||
WEIGHTS_DIR: str = ...
|
||||
WEIGHTS_CONFIG_FILENAME_FORMAT: str = ...
|
||||
WEIGHT_FILENAME_PREFIX: str = ...
|
||||
CONSTANTS_DIR: str = ...
|
||||
CONSTANTS_CONFIG_FILENAME_FORMAT: str = ...
|
||||
TENSOR_CONSTANT_FILENAME_PREFIX: str = ...
|
||||
CUSTOM_OBJ_FILENAME_PREFIX: str = ...
|
||||
SAMPLE_INPUTS_DIR: str = ...
|
||||
SAMPLE_INPUTS_FILENAME_FORMAT: str = ...
|
||||
EXECUTORCH_DIR: str = ...
|
||||
EXTRA_DIR: str = ...
|
||||
MODULE_INFO_PATH: str = ...
|
||||
XL_MODEL_WEIGHTS_DIR: str = ...
|
||||
XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH: str = ...
|
||||
@@ -0,0 +1,16 @@
|
||||
from torch import Tensor
|
||||
from torch.types import _bool
|
||||
|
||||
# Defined in torch/csrc/functionalization/Module.cpp
|
||||
|
||||
class ViewMeta:
|
||||
has_symbolic_inputs: _bool
|
||||
|
||||
# Returns the list of ViewMeta instances of the given functional tensor.
|
||||
#
|
||||
# Although we do have python bindings for their types, we won't
|
||||
# expose them here, since they should not be used by users.
|
||||
def get_view_meta_sequence(tensor: Tensor) -> list[ViewMeta]: ...
|
||||
|
||||
# Applies the ViewMeta sequence on top of the given base.
|
||||
def apply_view_meta_sequence(base: Tensor, sequence: list[ViewMeta]) -> Tensor: ...
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import AnyStr, overload
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
class UndefinedGrad:
|
||||
def __init__(self) -> None: ...
|
||||
def __call__(self, *inputs: Tensor) -> list[Tensor]: ...
|
||||
|
||||
class DelayedError:
|
||||
def __init__(self, msg: AnyStr, num_inputs: int) -> None: ...
|
||||
|
||||
# __call__ should really be a higher-kinded type:
|
||||
# def __call__(self, arg: Tensor) -> Tensor: ...
|
||||
# def __call__(self, *args: Tensor * num_inputs) -> Tuple[Tensor * num_inputs]: ...
|
||||
|
||||
@overload
|
||||
def __call__(self, i0: Tensor) -> Tensor: ...
|
||||
@overload
|
||||
def __call__(self, *args: Tensor) -> tuple[Tensor, ...]: ...
|
||||
@@ -0,0 +1,100 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from enum import Enum
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
# Defined in torch/csrc/functorch/init.cpp
|
||||
|
||||
def _assert_wrapped_functional(
|
||||
input_tensor: Tensor, wrapped_tensor: Tensor
|
||||
) -> None: ...
|
||||
def _func_decrement_nesting() -> int: ...
|
||||
def _func_increment_nesting(reapply_views: bool) -> int: ...
|
||||
def _propagate_functional_input_mutation(
|
||||
input_tensor: Tensor, wrapped_tensor: Tensor
|
||||
) -> None: ...
|
||||
def set_inplace_requires_grad_allowed(allowed: bool) -> None: ...
|
||||
def get_inplace_requires_grad_allowed() -> bool: ...
|
||||
def _set_dynamic_layer_keys_included(included: bool) -> None: ...
|
||||
def get_unwrapped(tensor: Tensor) -> Tensor: ...
|
||||
def is_batchedtensor(tensor: Tensor) -> bool: ...
|
||||
def is_functionaltensor(tensor: Tensor) -> bool: ...
|
||||
def is_functorch_wrapped_tensor(tensor: Tensor) -> bool: ...
|
||||
def is_gradtrackingtensor(tensor: Tensor) -> bool: ...
|
||||
def is_legacy_batchedtensor(tensor: Tensor) -> bool: ...
|
||||
def maybe_get_bdim(tensor: Tensor) -> int: ...
|
||||
def maybe_get_level(tensor: Tensor) -> int: ...
|
||||
def maybe_current_level() -> int | None: ...
|
||||
def unwrap_if_dead(tensor: Tensor) -> Tensor: ...
|
||||
def _unwrap_for_grad(tensor: Tensor, level: int) -> Tensor: ...
|
||||
def _wrap_for_grad(tensor: Tensor, level: int) -> Tensor: ...
|
||||
def _unwrap_batched(tensor: Tensor, level: int) -> tuple[Tensor, int | None]: ...
|
||||
def current_level() -> int: ...
|
||||
def count_jvp_interpreters() -> int: ...
|
||||
def _add_batch_dim(tensor: Tensor, bdim: int, level: int) -> Tensor: ...
|
||||
def _remove_batch_dim(
|
||||
tensor: Tensor, level: int, batch_size: int, out_dim: int
|
||||
) -> Tensor: ...
|
||||
def _maybe_unsafe_set_level(tensor: Tensor, level: int) -> None: ...
|
||||
def set_single_level_autograd_function_allowed(allowed: bool) -> None: ...
|
||||
def get_single_level_autograd_function_allowed() -> bool: ...
|
||||
def _unwrap_functional_tensor(tensor: Tensor, reapply_views: bool) -> Tensor: ...
|
||||
def _wrap_functional_tensor(tensor: Tensor, level: int) -> Tensor: ...
|
||||
def _vmap_increment_nesting(batch_size: int, randomness: str) -> int: ...
|
||||
def _vmap_decrement_nesting() -> int: ...
|
||||
def _grad_increment_nesting() -> int: ...
|
||||
def _grad_decrement_nesting() -> int: ...
|
||||
def _jvp_increment_nesting() -> int: ...
|
||||
def _jvp_decrement_nesting() -> int: ...
|
||||
|
||||
# Defined in aten/src/ATen/functorch/Interpreter.h
|
||||
class TransformType(Enum):
|
||||
Torch = ...
|
||||
Vmap = ...
|
||||
Grad = ...
|
||||
Jvp = ...
|
||||
Functionalize = ...
|
||||
|
||||
class RandomnessType(Enum):
|
||||
Error = ...
|
||||
Same = ...
|
||||
Different = ...
|
||||
|
||||
class CInterpreter:
|
||||
def key(self) -> TransformType: ...
|
||||
def level(self) -> int: ...
|
||||
def serialize(self) -> bytes: ...
|
||||
@staticmethod
|
||||
def deserialize(bytes) -> CInterpreter: ...
|
||||
|
||||
class CGradInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def lift(self, Tensor) -> Tensor: ...
|
||||
def prevGradMode(self) -> bool: ...
|
||||
|
||||
class CJvpInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def lift(self, Tensor) -> Tensor: ...
|
||||
def prevFwdGradMode(self) -> bool: ...
|
||||
|
||||
class CFunctionalizeInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def key(self) -> TransformType: ...
|
||||
def level(self) -> int: ...
|
||||
def functionalizeAddBackViews(self) -> bool: ...
|
||||
|
||||
class CVmapInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def key(self) -> TransformType: ...
|
||||
def level(self) -> int: ...
|
||||
def batchSize(self) -> int: ...
|
||||
def randomness(self) -> RandomnessType: ...
|
||||
|
||||
class DynamicLayer: ...
|
||||
|
||||
def get_dynamic_layer_stack_depth() -> int: ...
|
||||
def get_interpreter_stack() -> list[CInterpreter]: ...
|
||||
def peek_interpreter_stack() -> CInterpreter: ...
|
||||
def pop_dynamic_layer_stack() -> DynamicLayer: ...
|
||||
def pop_dynamic_layer_stack_and_undo_to_depth(int) -> None: ...
|
||||
def push_dynamic_layer_stack(dl: DynamicLayer) -> int: ...
|
||||
@@ -0,0 +1,4 @@
|
||||
# Defined in torch/csrc/instruction_counter/Module.cpp
|
||||
|
||||
def start() -> int: ...
|
||||
def end(id: int) -> int: ...
|
||||
@@ -0,0 +1,5 @@
|
||||
# Defined in torch/csrc/itt.cpp
|
||||
def is_available() -> None: ...
|
||||
def rangePush(message: str) -> None: ...
|
||||
def rangePop() -> None: ...
|
||||
def mark(message: str) -> None: ...
|
||||
@@ -0,0 +1,200 @@
|
||||
from typing import Any
|
||||
|
||||
# Defined in torch/csrc/jit/python/python_tree_views.cpp
|
||||
|
||||
class SourceRange:
|
||||
def highlight(self) -> str: ...
|
||||
@property
|
||||
def start(self) -> int: ...
|
||||
@property
|
||||
def end(self) -> int: ...
|
||||
|
||||
class SourceRangeFactory:
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
filename: Any,
|
||||
file_lineno: int,
|
||||
leading_whitespace_chars: int,
|
||||
) -> None: ...
|
||||
def make_range(self, line: int, start_col: int, end_col: int) -> SourceRange: ...
|
||||
def make_raw_range(self, start: int, end: int) -> SourceRange: ...
|
||||
@property
|
||||
def source(self) -> str: ...
|
||||
|
||||
class TreeView:
|
||||
def range(self) -> SourceRange: ...
|
||||
def dump(self) -> None: ...
|
||||
|
||||
class Ident(TreeView):
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None: ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
|
||||
class Param(TreeView):
|
||||
def __init__(self, type: Any | None, name: Ident, kwarg_only: bool) -> None: ...
|
||||
|
||||
class Attribute(TreeView):
|
||||
def __init__(self, name: Ident, value: Any) -> None: ...
|
||||
|
||||
# Literals
|
||||
def TrueLiteral(range: SourceRange) -> Any: ...
|
||||
def FalseLiteral(range: SourceRange) -> Any: ...
|
||||
def NoneLiteral(range: SourceRange) -> Any: ...
|
||||
|
||||
# Tree nodes
|
||||
class Stmt(TreeView):
|
||||
def __init__(self, thing: TreeView) -> None: ...
|
||||
|
||||
class Expr(TreeView): ...
|
||||
|
||||
class Def(TreeView):
|
||||
def __init__(self, name: Ident, decl: Any, body: list[Stmt]) -> None: ...
|
||||
def decl(self) -> Any: ...
|
||||
def name(self) -> Ident: ...
|
||||
|
||||
class Property(TreeView):
|
||||
def __init__(
|
||||
self, r: SourceRange, name: Ident, getter: Def, setter: Def | None
|
||||
) -> None: ...
|
||||
def name(self) -> Ident: ...
|
||||
def getter_name(self) -> str: ...
|
||||
def setter_name(self) -> Ident | None: ...
|
||||
|
||||
class ClassDef(TreeView):
|
||||
def __init__(
|
||||
self, name: Ident, body: list[Stmt], props: list[Property], assigns: list[Any]
|
||||
) -> None: ...
|
||||
|
||||
class Decl(TreeView):
|
||||
def __init__(
|
||||
self, r: SourceRange, params: list[Param], return_type: Expr | None
|
||||
) -> None: ...
|
||||
|
||||
class Delete(Stmt):
|
||||
def __init__(self, range: SourceRange, targets: list[Expr]) -> None: ...
|
||||
|
||||
class WithItem(Expr):
|
||||
def __init__(self, range: SourceRange, target: Expr, var: Any | None) -> None: ...
|
||||
|
||||
class Assign(Stmt):
|
||||
def __init__(
|
||||
self, lhs: list[Expr], rhs: Expr, type: Expr | None = None
|
||||
) -> None: ...
|
||||
|
||||
class AugAssign(Stmt):
|
||||
def __init__(self, lhs: Expr, kind_str: str, rhs: Expr) -> None: ...
|
||||
|
||||
class Return(Stmt):
|
||||
def __init__(self, range: SourceRange, value: Expr | None) -> None: ...
|
||||
|
||||
class Raise(Stmt):
|
||||
def __init__(self, range: SourceRange, expr: Expr) -> None: ...
|
||||
|
||||
class Assert(Stmt):
|
||||
def __init__(self, range: SourceRange, test: Expr, msg: Expr | None) -> None: ...
|
||||
|
||||
class Pass(Stmt):
|
||||
def __init__(self, range: SourceRange) -> None: ...
|
||||
|
||||
class Break(Stmt): ...
|
||||
class Continue(Stmt): ...
|
||||
|
||||
class Dots(Expr, TreeView):
|
||||
def __init__(self, range: SourceRange) -> None: ...
|
||||
|
||||
class If(Stmt):
|
||||
def __init__(
|
||||
self,
|
||||
range: SourceRange,
|
||||
cond: Expr,
|
||||
true_branch: list[Stmt],
|
||||
false_branch: list[Stmt],
|
||||
) -> None: ...
|
||||
|
||||
class While(Stmt):
|
||||
def __init__(self, range: SourceRange, cond: Expr, body: list[Stmt]) -> None: ...
|
||||
|
||||
class With(Stmt):
|
||||
def __init__(
|
||||
self, range: SourceRange, targets: list[WithItem], body: list[Stmt]
|
||||
) -> None: ...
|
||||
|
||||
class For(Stmt):
|
||||
def __init__(
|
||||
self,
|
||||
range: SourceRange,
|
||||
targets: list[Expr],
|
||||
itrs: list[Expr],
|
||||
body: list[Stmt],
|
||||
) -> None: ...
|
||||
|
||||
class ExprStmt(Stmt):
|
||||
def __init__(self, expr: Expr) -> None: ...
|
||||
|
||||
class Var(Expr):
|
||||
def __init__(self, name: Ident) -> None: ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
|
||||
class BinOp(Expr):
|
||||
def __init__(self, kind: str, lhs: Expr, rhs: Expr) -> None: ...
|
||||
|
||||
class UnaryOp(Expr):
|
||||
def __init__(self, range: SourceRange, kind: str, expr: Expr) -> None: ...
|
||||
|
||||
class Const(Expr):
|
||||
def __init__(self, range: SourceRange, value: str) -> None: ...
|
||||
|
||||
class StringLiteral(Expr):
|
||||
def __init__(self, range: SourceRange, value: str) -> None: ...
|
||||
|
||||
class Apply(Expr):
|
||||
def __init__(
|
||||
self, expr: Expr, args: list[Expr], kwargs: list[Attribute]
|
||||
) -> None: ...
|
||||
|
||||
class Select(Expr):
|
||||
def __init__(self, expr: Expr, field: Ident) -> None: ...
|
||||
|
||||
class TernaryIf(Expr):
|
||||
def __init__(self, cond: Expr, true_expr: Expr, false_expr: Expr) -> None: ...
|
||||
|
||||
class ListComp(Expr):
|
||||
def __init__(
|
||||
self, range: SourceRange, elt: Expr, target: Expr, iter: Expr
|
||||
) -> None: ...
|
||||
|
||||
class DictComp(Expr):
|
||||
def __init__(
|
||||
self, range: SourceRange, key: Expr, value: Expr, target: Expr, iter: Expr
|
||||
) -> None: ...
|
||||
|
||||
class ListLiteral(Expr):
|
||||
def __init__(self, range: SourceRange, args: list[Expr]) -> None: ...
|
||||
|
||||
class TupleLiteral(Expr):
|
||||
def __init__(self, range: SourceRange, args: list[Expr]) -> None: ...
|
||||
|
||||
class DictLiteral(Expr):
|
||||
def __init__(
|
||||
self, range: SourceRange, keys: list[Expr], values: list[Expr]
|
||||
) -> None: ...
|
||||
|
||||
class Subscript(Expr):
|
||||
def __init__(self, base: Expr, subscript_exprs: list[Expr]) -> None: ...
|
||||
|
||||
class SliceExpr(Expr):
|
||||
def __init__(
|
||||
self,
|
||||
range: SourceRange,
|
||||
lower: Expr | None,
|
||||
upper: Expr | None,
|
||||
step: Expr | None,
|
||||
) -> None: ...
|
||||
|
||||
class Starred(Expr):
|
||||
def __init__(self, range: SourceRange, expr: Expr) -> None: ...
|
||||
|
||||
class EmptyTypeAnnotation(TreeView):
|
||||
def __init__(self, range: SourceRange) -> None: ...
|
||||
@@ -0,0 +1,26 @@
|
||||
from torch import Tensor
|
||||
|
||||
# defined in torch/csrc/lazy/python/init.cpp
|
||||
def _mark_step(device: str, devices: list[str], wait: bool) -> None: ...
|
||||
def _wait_device_ops(devices: list[str]) -> None: ...
|
||||
def _reset_metrics() -> None: ...
|
||||
def _counter_names() -> list[str]: ...
|
||||
def _counter_value(name: str) -> int: ...
|
||||
def _metrics_report() -> str: ...
|
||||
def _get_graph_hash(tensors: list[Tensor]) -> str: ...
|
||||
def _sync_multi(
|
||||
tensors: list[Tensor],
|
||||
devices: list[str],
|
||||
wait: bool = True,
|
||||
sync_ltc_data: bool = True,
|
||||
) -> None: ...
|
||||
def _get_tensor_id(tensor: Tensor) -> int: ...
|
||||
def _get_tensors_text(tensors: list[Tensor]) -> str: ...
|
||||
def _get_tensors_dot(tensors: list[Tensor]) -> str: ...
|
||||
def _get_tensors_backend(tensors: list[Tensor]) -> str: ...
|
||||
def _get_force_fallback() -> str: ...
|
||||
def _set_force_fallback(newval: str) -> None: ...
|
||||
def _clear_ir_cache() -> None: ...
|
||||
def _dump_ir_cache(filename: str) -> None: ...
|
||||
def _set_reuse_ir(val: bool) -> None: ...
|
||||
def _get_default_device_type() -> str: ...
|
||||
@@ -0,0 +1,12 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# defined in torch/csrc/lazy/python/init.cpp
|
||||
|
||||
from typing import Any
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
def _init(): ...
|
||||
def _get_tensors_ts_device_data_node(
|
||||
tensors: list[Tensor],
|
||||
) -> tuple[list[int], list[Any]]: ...
|
||||
def _run_cached_graph(hash_str: str, graph_inputs: list[Any]) -> list[Tensor]: ...
|
||||
@@ -0,0 +1,58 @@
|
||||
# Defined in torch/csrc/monitor/python_init.cpp
|
||||
|
||||
import datetime
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from types import TracebackType
|
||||
|
||||
class Aggregation(Enum):
|
||||
VALUE = ...
|
||||
MEAN = ...
|
||||
COUNT = ...
|
||||
SUM = ...
|
||||
MAX = ...
|
||||
MIN = ...
|
||||
|
||||
class Stat:
|
||||
name: str
|
||||
count: int
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
aggregations: list[Aggregation],
|
||||
window_size: int,
|
||||
max_samples: int = -1,
|
||||
) -> None: ...
|
||||
def add(self, v: float) -> None: ...
|
||||
def get(self) -> dict[Aggregation, float]: ...
|
||||
|
||||
class Event:
|
||||
name: str
|
||||
timestamp: datetime.datetime
|
||||
data: dict[str, int | float | bool | str]
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
timestamp: datetime.datetime,
|
||||
data: dict[str, int | float | bool | str],
|
||||
) -> None: ...
|
||||
|
||||
def log_event(e: Event) -> None: ...
|
||||
|
||||
class EventHandlerHandle: ...
|
||||
|
||||
def register_event_handler(handler: Callable[[Event], None]) -> EventHandlerHandle: ...
|
||||
def unregister_event_handler(handle: EventHandlerHandle) -> None: ...
|
||||
|
||||
class _WaitCounterTracker:
|
||||
def __enter__(self) -> None: ...
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None = None,
|
||||
exc_value: BaseException | None = None,
|
||||
traceback: TracebackType | None = None,
|
||||
) -> None: ...
|
||||
|
||||
class _WaitCounter:
|
||||
def __init__(self, key: str) -> None: ...
|
||||
def guard(self) -> _WaitCounterTracker: ...
|
||||
@@ -0,0 +1,347 @@
|
||||
# @generated by tools/pyi/gen_pyi.py from torch/_C/_nn.pyi.in
|
||||
# mypy: disable-error-code="type-arg"
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Literal, overload
|
||||
|
||||
from torch import memory_format, Tensor
|
||||
from torch.types import _bool, _device, _dtype, _int, _size
|
||||
|
||||
# Defined in tools/autograd/templates/python_nn_functions.cpp
|
||||
|
||||
def adaptive_avg_pool2d(input: Tensor, output_size: _int | _size) -> Tensor: ...
|
||||
def adaptive_avg_pool3d(input: Tensor, output_size: _int | _size) -> Tensor: ...
|
||||
def adaptive_max_pool2d(
|
||||
input: Tensor,
|
||||
output_size: _int | _size,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def adaptive_max_pool3d(
|
||||
input: Tensor,
|
||||
output_size: _int | _size,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def avg_pool2d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
ceil_mode: bool = False,
|
||||
count_include_pad: bool = True,
|
||||
divisor_override: int | None = None,
|
||||
) -> Tensor: ...
|
||||
def avg_pool3d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
ceil_mode: bool = False,
|
||||
count_include_pad: bool = True,
|
||||
divisor_override: int | None = None,
|
||||
) -> Tensor: ...
|
||||
def binary_cross_entropy(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def col2im(
|
||||
input: Tensor,
|
||||
output_size: _int | _size,
|
||||
kernel_size: _int | _size,
|
||||
dilation: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
) -> Tensor: ...
|
||||
def cross_entropy_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
ignore_index: int = -100,
|
||||
label_smoothing: float = 0.0,
|
||||
) -> Tensor: ...
|
||||
def elu(
|
||||
input: Tensor,
|
||||
alpha: float = 1.0,
|
||||
scale: float = 1.0,
|
||||
input_scale: float = 1.0,
|
||||
) -> Tensor: ...
|
||||
def elu_(input: Tensor, alpha: float = ...) -> Tensor: ...
|
||||
def fractional_max_pool2d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
output_size: _int | _size,
|
||||
_random_samples: Tensor,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def fractional_max_pool3d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
output_size: _int | _size,
|
||||
_random_samples: Tensor,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def gelu(input: Tensor, approximate: str = ...) -> Tensor: ...
|
||||
def glu(input: Tensor, dim: int = -1) -> Tensor: ...
|
||||
def hardsigmoid(input: Tensor, *, out: Tensor | None = None) -> Tensor: ...
|
||||
def hardsigmoid_(input: Tensor) -> Tensor: ...
|
||||
def hardswish(input: Tensor) -> Tensor: ...
|
||||
def hardswish_(input: Tensor) -> Tensor: ...
|
||||
def hardtanh(
|
||||
input: Tensor,
|
||||
min_val: float = ...,
|
||||
max_val: float = ...,
|
||||
*,
|
||||
out: Tensor | None = None,
|
||||
) -> Tensor: ...
|
||||
def hardtanh_(
|
||||
input: Tensor,
|
||||
min_val: float = ...,
|
||||
max_val: float = ...,
|
||||
) -> Tensor: ...
|
||||
def huber_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
delta: float = 1.0,
|
||||
) -> Tensor: ...
|
||||
def im2col(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
dilation: _int | _size,
|
||||
padding: _int | _size,
|
||||
stride: _int | _size,
|
||||
) -> Tensor: ...
|
||||
def l1_loss(input: Tensor, target: Tensor, reduction: str = ...) -> Tensor: ...
|
||||
def leaky_relu(
|
||||
input: Tensor,
|
||||
negative_slope: float = ...,
|
||||
*,
|
||||
out: Tensor | None = None,
|
||||
) -> Tensor: ...
|
||||
def leaky_relu_(input: Tensor, negative_slope: float = ...) -> Tensor: ...
|
||||
def linear(
|
||||
input: Tensor,
|
||||
weight: Tensor,
|
||||
bias: Tensor | None = None,
|
||||
) -> Tensor: ...
|
||||
def log_sigmoid(input: Tensor) -> Tensor: ...
|
||||
def max_pool2d_with_indices(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
dilation: _int | _size = 1,
|
||||
ceil_mode: bool = False,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def max_pool3d_with_indices(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
dilation: _int | _size = 1,
|
||||
ceil_mode: bool = False,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def max_unpool2d(
|
||||
input: Tensor,
|
||||
indices: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
) -> Tensor: ...
|
||||
def max_unpool3d(
|
||||
input: Tensor,
|
||||
indices: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
stride: _int | _size,
|
||||
padding: _int | _size,
|
||||
) -> Tensor: ...
|
||||
def mish(input: Tensor) -> Tensor: ...
|
||||
def mish_(input: Tensor) -> Tensor: ...
|
||||
def mse_loss(input: Tensor, target: Tensor, reduction: str = ...) -> Tensor: ...
|
||||
def multi_margin_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
p: float = 1.0,
|
||||
margin: float = 1.0,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def multilabel_margin_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def nll_loss_nd(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
ignore_index: int = -100,
|
||||
) -> Tensor: ...
|
||||
def one_hot(tensor: Tensor, num_classes: int = ...) -> Tensor: ...
|
||||
def pad(
|
||||
input: Tensor,
|
||||
pad: Sequence[int],
|
||||
mode: str = ...,
|
||||
value: float | None = None,
|
||||
) -> Tensor: ...
|
||||
def relu6(input: Tensor) -> Tensor: ...
|
||||
def relu6_(input: Tensor) -> Tensor: ...
|
||||
def scaled_dot_product_attention(
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
attn_mask: Tensor | None = None,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
enable_gqa: bool = False,
|
||||
) -> Tensor: ...
|
||||
def silu(input: Tensor) -> Tensor: ...
|
||||
def silu_(input: Tensor) -> Tensor: ...
|
||||
def smooth_l1_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
beta: float = 1.0,
|
||||
) -> Tensor: ...
|
||||
def soft_margin_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def softplus(
|
||||
input: Tensor,
|
||||
beta: float = ...,
|
||||
threshold: float = ...,
|
||||
) -> Tensor: ...
|
||||
def softshrink(input: Tensor, lambd: float = ...) -> Tensor: ...
|
||||
|
||||
# Defined in aten/src/ATen/native/mkldnn/Linear.cpp
|
||||
def mkldnn_linear(input: Tensor, weight: Tensor, bias: Tensor | None) -> Tensor: ...
|
||||
|
||||
# Defined at aten/src/ATen/native/mkldnn/MKLDNNConversions.cpp
|
||||
def mkldnn_reorder_conv2d_weight(
|
||||
self: Tensor,
|
||||
padding: list,
|
||||
stride: list,
|
||||
dilatation: list,
|
||||
groups: int,
|
||||
) -> Tensor: ...
|
||||
def mkldnn_reorder_conv3d_weight(
|
||||
self: Tensor,
|
||||
padding: list,
|
||||
stride: list,
|
||||
dilatation: list,
|
||||
groups: int,
|
||||
) -> Tensor: ...
|
||||
|
||||
# Defined in aten/src/ATen/native/mkldnn/Prelu.cpp
|
||||
def mkldnn_prelu(input: Tensor, weight: Tensor) -> Tensor: ...
|
||||
|
||||
# Defined at tools/autograd/templates/python_nn_functions.cpp
|
||||
@overload
|
||||
def _parse_to(
|
||||
device: _device,
|
||||
dtype: _dtype,
|
||||
non_blocking: _bool,
|
||||
copy: _bool,
|
||||
*,
|
||||
memory_format: memory_format,
|
||||
) -> tuple[_device, _dtype, _bool, memory_format]: ...
|
||||
@overload
|
||||
def _parse_to(
|
||||
dtype: _dtype,
|
||||
non_blocking: _bool,
|
||||
copy: _bool,
|
||||
*,
|
||||
memory_format: memory_format,
|
||||
) -> tuple[_device, _dtype, _bool, memory_format]: ...
|
||||
@overload
|
||||
def _parse_to(
|
||||
tensor: Tensor,
|
||||
non_blocking: _bool,
|
||||
copy: _bool,
|
||||
*,
|
||||
memory_format: memory_format,
|
||||
) -> tuple[_device, _dtype, _bool, memory_format]: ...
|
||||
|
||||
# Defined in aten/src/ATen/native/PackedSequence.cpp
|
||||
def pad_sequence(
|
||||
sequences: list[Tensor] | tuple[Tensor, ...],
|
||||
batch_first: bool = False,
|
||||
padding_value: float = 0.0,
|
||||
padding_side: Literal["left", "right"] = "right",
|
||||
) -> Tensor: ...
|
||||
|
||||
# Upsample functions used by torch.nn.functional.interpolate
|
||||
def upsample_nearest1d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_nearest2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_nearest3d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_nearest_exact1d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_nearest_exact2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_nearest_exact3d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_linear1d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_bilinear2d_aa(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_bilinear2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_trilinear3d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_bicubic2d_aa(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_lanczos2d_aa(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_bicubic2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def flatten_dense_tensors(tensors: list[Tensor]) -> Tensor: ...
|
||||
def unflatten_dense_tensors(flat: Tensor, tensors: list[Tensor]) -> list[Tensor]: ...
|
||||
@@ -0,0 +1,9 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Defined in torch/csrc/cuda/shared/nvtx.cpp
|
||||
def rangePushA(message: str) -> int: ...
|
||||
def rangePop() -> int: ...
|
||||
def rangeStartA(message: str) -> int: ...
|
||||
def rangeEnd(int) -> None: ...
|
||||
def markA(message: str) -> None: ...
|
||||
def deviceRangeStart(message: str, stream: int) -> object: ...
|
||||
def deviceRangeEnd(range_handle: object, stream: int) -> None: ...
|
||||
@@ -0,0 +1,39 @@
|
||||
# Defined in torch/csrc/onnx/init.cpp
|
||||
|
||||
from enum import Enum
|
||||
|
||||
PRODUCER_VERSION: str
|
||||
|
||||
class TensorProtoDataType(Enum):
|
||||
UNDEFINED = ...
|
||||
FLOAT = ...
|
||||
UINT8 = ...
|
||||
INT8 = ...
|
||||
UINT16 = ...
|
||||
INT16 = ...
|
||||
INT32 = ...
|
||||
INT64 = ...
|
||||
STRING = ...
|
||||
BOOL = ...
|
||||
FLOAT16 = ...
|
||||
DOUBLE = ...
|
||||
UINT32 = ...
|
||||
UINT64 = ...
|
||||
COMPLEX64 = ...
|
||||
COMPLEX128 = ...
|
||||
BFLOAT16 = ...
|
||||
FLOAT8E5M2 = ...
|
||||
FLOAT8E4M3FN = ...
|
||||
FLOAT8E5M2FNUZ = ...
|
||||
FLOAT8E4M3FNUZ = ...
|
||||
|
||||
class OperatorExportTypes(Enum):
|
||||
ONNX = ...
|
||||
ONNX_ATEN = ...
|
||||
ONNX_ATEN_FALLBACK = ...
|
||||
ONNX_FALLTHROUGH = ...
|
||||
|
||||
class TrainingMode(Enum):
|
||||
EVAL = ...
|
||||
PRESERVE = ...
|
||||
TRAINING = ...
|
||||
@@ -0,0 +1,248 @@
|
||||
from enum import Enum
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from torch._C import device, dtype, layout
|
||||
|
||||
# defined in torch/csrc/profiler/python/init.cpp
|
||||
|
||||
class RecordScope(Enum):
|
||||
FUNCTION = ...
|
||||
BACKWARD_FUNCTION = ...
|
||||
TORCHSCRIPT_FUNCTION = ...
|
||||
KERNEL_FUNCTION_DTYPE = ...
|
||||
CUSTOM_CLASS = ...
|
||||
BUILD_FEATURE = ...
|
||||
LITE_INTERPRETER = ...
|
||||
USER_SCOPE = ...
|
||||
STATIC_RUNTIME_OP = ...
|
||||
STATIC_RUNTIME_MODEL = ...
|
||||
|
||||
class ProfilerState(Enum):
|
||||
Disabled = ...
|
||||
CPU = ...
|
||||
CUDA = ...
|
||||
NVTX = ...
|
||||
ITT = ...
|
||||
PRIVATEUSE1 = ...
|
||||
KINETO = ...
|
||||
KINETO_GPU_FALLBACK = ...
|
||||
KINETO_PRIVATEUSE1_FALLBACK = ...
|
||||
KINETO_PRIVATEUSE1 = ...
|
||||
|
||||
class ActiveProfilerType(Enum):
|
||||
NONE = ...
|
||||
LEGACY = ...
|
||||
KINETO = ...
|
||||
NVTX = ...
|
||||
ITT = ...
|
||||
PRIVATEUSE1 = ...
|
||||
|
||||
class ProfilerActivity(Enum):
|
||||
CPU = ...
|
||||
CUDA = ...
|
||||
XPU = ...
|
||||
MTIA = ...
|
||||
HPU = ...
|
||||
PrivateUse1 = ...
|
||||
|
||||
class _EventType(Enum):
|
||||
TorchOp = ...
|
||||
Backend = ...
|
||||
Allocation = ...
|
||||
OutOfMemory = ...
|
||||
PyCall = ...
|
||||
PyCCall = ...
|
||||
Kineto = ...
|
||||
|
||||
class _ExperimentalConfig:
|
||||
def __init__(
|
||||
self,
|
||||
profiler_metrics: list[str] = ...,
|
||||
profiler_measure_per_kernel: bool = ...,
|
||||
verbose: bool = ...,
|
||||
performance_events: list[str] = ...,
|
||||
enable_cuda_sync_events: bool = ...,
|
||||
profile_all_threads: bool = ...,
|
||||
) -> None: ...
|
||||
|
||||
class ProfilerConfig:
|
||||
def __init__(
|
||||
self,
|
||||
state: ProfilerState,
|
||||
report_input_shapes: bool,
|
||||
profile_memory: bool,
|
||||
with_stack: bool,
|
||||
with_flops: bool,
|
||||
with_modules: bool,
|
||||
experimental_config: _ExperimentalConfig,
|
||||
trace_id: str | None = None,
|
||||
) -> None: ...
|
||||
|
||||
class _ProfilerEvent:
|
||||
start_tid: int
|
||||
start_time_ns: int
|
||||
children: list[_ProfilerEvent]
|
||||
|
||||
# TODO(robieta): remove in favor of `self.typed`
|
||||
extra_fields: (
|
||||
_ExtraFields_TorchOp
|
||||
| _ExtraFields_Backend
|
||||
| _ExtraFields_Allocation
|
||||
| _ExtraFields_OutOfMemory
|
||||
| _ExtraFields_PyCall
|
||||
| _ExtraFields_PyCCall
|
||||
| _ExtraFields_Kineto
|
||||
)
|
||||
|
||||
@property
|
||||
def typed(
|
||||
self,
|
||||
) -> (
|
||||
tuple[Literal[_EventType.TorchOp], _ExtraFields_TorchOp]
|
||||
| tuple[Literal[_EventType.Backend], _ExtraFields_Backend]
|
||||
| tuple[Literal[_EventType.Allocation], _ExtraFields_Allocation]
|
||||
| tuple[Literal[_EventType.OutOfMemory], _ExtraFields_OutOfMemory]
|
||||
| tuple[Literal[_EventType.PyCall], _ExtraFields_PyCall]
|
||||
| tuple[Literal[_EventType.PyCCall], _ExtraFields_PyCCall]
|
||||
| tuple[Literal[_EventType.Kineto], _ExtraFields_Kineto]
|
||||
): ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
@property
|
||||
def tag(self) -> _EventType: ...
|
||||
@property
|
||||
def id(self) -> int: ...
|
||||
@property
|
||||
def parent(self) -> _ProfilerEvent | None: ...
|
||||
@property
|
||||
def correlation_id(self) -> int: ...
|
||||
@property
|
||||
def end_time_ns(self) -> int: ...
|
||||
@property
|
||||
def duration_time_ns(self) -> int: ...
|
||||
|
||||
class _TensorMetadata:
|
||||
impl_ptr: int | None
|
||||
storage_data_ptr: int | None
|
||||
id: int | None
|
||||
|
||||
@property
|
||||
def allocation_id(self) -> int | None: ...
|
||||
@property
|
||||
def layout(self) -> layout: ...
|
||||
@property
|
||||
def device(self) -> device: ...
|
||||
@property
|
||||
def dtype(self) -> dtype: ...
|
||||
@property
|
||||
def sizes(self) -> list[int]: ...
|
||||
@property
|
||||
def strides(self) -> list[int]: ...
|
||||
|
||||
Scalar: TypeAlias = int | float | bool | complex
|
||||
Input: TypeAlias = _TensorMetadata | list[_TensorMetadata] | Scalar | None
|
||||
|
||||
class _ExtraFields_TorchOp:
|
||||
name: str
|
||||
sequence_number: int
|
||||
allow_tf32_cublas: bool
|
||||
|
||||
@property
|
||||
def inputs(self) -> list[Input]: ...
|
||||
@property
|
||||
def scope(self) -> RecordScope: ...
|
||||
|
||||
class _ExtraFields_Backend: ...
|
||||
|
||||
class _ExtraFields_Allocation:
|
||||
ptr: int
|
||||
id: int | None
|
||||
alloc_size: int
|
||||
total_allocated: int
|
||||
total_reserved: int
|
||||
|
||||
@property
|
||||
def allocation_id(self) -> int | None: ...
|
||||
@property
|
||||
def device(self) -> device: ...
|
||||
|
||||
class _ExtraFields_OutOfMemory: ...
|
||||
|
||||
class _PyFrameState:
|
||||
line_number: int
|
||||
function_name: str
|
||||
|
||||
@property
|
||||
def file_name(self) -> str: ...
|
||||
|
||||
class _NNModuleInfo:
|
||||
@property
|
||||
def self_ptr(self) -> int: ...
|
||||
@property
|
||||
def cls_ptr(self) -> int: ...
|
||||
@property
|
||||
def cls_name(self) -> str: ...
|
||||
@property
|
||||
def parameters(
|
||||
self,
|
||||
) -> list[tuple[str, _TensorMetadata, _TensorMetadata | None]]: ...
|
||||
|
||||
class _OptimizerInfo:
|
||||
@property
|
||||
def parameters(
|
||||
self,
|
||||
) -> list[
|
||||
tuple[
|
||||
# Parameter
|
||||
_TensorMetadata,
|
||||
#
|
||||
# Gradient (if present during optimizer.step())
|
||||
_TensorMetadata | None,
|
||||
#
|
||||
# Optimizer state for Parameter as (name, tensor) pairs
|
||||
list[tuple[str, _TensorMetadata]],
|
||||
]
|
||||
]: ...
|
||||
|
||||
class _ExtraFields_PyCCall:
|
||||
@property
|
||||
def caller(self) -> _PyFrameState: ...
|
||||
|
||||
class _ExtraFields_PyCall:
|
||||
@property
|
||||
def callsite(self) -> _PyFrameState: ...
|
||||
@property
|
||||
def caller(self) -> _PyFrameState: ...
|
||||
@property
|
||||
def module(self) -> _NNModuleInfo | None: ...
|
||||
@property
|
||||
def optimizer(self) -> _OptimizerInfo | None: ...
|
||||
|
||||
class _ExtraFields_Kineto: ...
|
||||
|
||||
def _add_execution_trace_observer(output_file_path: str) -> bool: ...
|
||||
def _remove_execution_trace_observer() -> None: ...
|
||||
def _enable_execution_trace_observer() -> None: ...
|
||||
def _disable_execution_trace_observer() -> None: ...
|
||||
def _set_record_concrete_inputs_enabled_val(val: bool) -> None: ...
|
||||
def _set_fwd_bwd_enabled_val(val: bool) -> None: ...
|
||||
def _set_cuda_sync_enabled_val(val: bool) -> None: ...
|
||||
|
||||
class CapturedTraceback: ...
|
||||
|
||||
def gather_traceback(python: bool, script: bool, cpp: bool) -> CapturedTraceback: ...
|
||||
|
||||
# The Dict has name, filename, line
|
||||
def symbolize_tracebacks(
|
||||
to_symbolize: list[CapturedTraceback],
|
||||
) -> list[list[dict[str, str]]]: ...
|
||||
|
||||
class _RecordFunctionFast:
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
input_values: list | tuple | None = None,
|
||||
keyword_values: dict | None = None,
|
||||
) -> None: ...
|
||||
def __enter__(self) -> None: ...
|
||||
def __exit__(self, *exc_info: object) -> None: ...
|
||||
@@ -0,0 +1,3 @@
|
||||
# Defined in torch/csrc/utils/verbose.cpp
|
||||
def mkl_set_verbose(enable: int) -> int: ...
|
||||
def mkldnn_set_verbose(level: int) -> int: ...
|
||||
@@ -0,0 +1,11 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from torch._C import LiteScriptModule, ScriptModule
|
||||
|
||||
def _load_mobile_module_from_file(filename: str): ...
|
||||
def _load_mobile_module_from_bytes(bytes_: bytes): ...
|
||||
def _load_jit_module_from_file(filename: str): ...
|
||||
def _load_jit_module_from_bytes(bytes_: bytes): ...
|
||||
def _save_mobile_module(m: LiteScriptModule, filename: str): ...
|
||||
def _save_jit_module(m: ScriptModule, filename: str): ...
|
||||
def _save_mobile_module_to_bytes(m: LiteScriptModule) -> bytes: ...
|
||||
def _save_jit_module_to_bytes(m: ScriptModule) -> bytes: ...
|
||||
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
This makes the functions in torch._C._VariableFunctions available as
|
||||
torch._VF.<funcname>
|
||||
without mypy being able to find them.
|
||||
|
||||
A subset of those functions are mapped to ATen functions in
|
||||
torch/jit/_builtins.py
|
||||
|
||||
See https://github.com/pytorch/pytorch/issues/21478 for the reason for
|
||||
introducing torch._VF
|
||||
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class VFModule(types.ModuleType):
|
||||
vf: types.ModuleType
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__(name)
|
||||
self.vf = torch._C._VariableFunctions
|
||||
|
||||
def __getattr__(self, name: str) -> object:
|
||||
return getattr(self.vf, name)
|
||||
|
||||
|
||||
sys.modules[__name__] = VFModule(__name__)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
import torch
|
||||
|
||||
|
||||
def show() -> str:
|
||||
"""
|
||||
Return a human-readable string with descriptions of the
|
||||
configuration of PyTorch.
|
||||
"""
|
||||
return torch._C._show_config()
|
||||
|
||||
|
||||
# TODO: In principle, we could provide more structured version/config
|
||||
# information here. For now only CXX_FLAGS is exposed, as Timer
|
||||
# uses them.
|
||||
def _cxx_flags() -> str:
|
||||
"""Returns the CXX_FLAGS used when building PyTorch."""
|
||||
return torch._C._cxx_flags()
|
||||
|
||||
|
||||
def parallel_info() -> str:
|
||||
r"""Returns detailed string with parallelization settings"""
|
||||
return torch._C._parallel_info()
|
||||
@@ -0,0 +1,75 @@
|
||||
_overwrite_module_params_on_conversion: bool = False
|
||||
_swap_module_params_on_conversion: bool = False
|
||||
|
||||
|
||||
def set_overwrite_module_params_on_conversion(value: bool) -> None:
|
||||
"""
|
||||
Sets whether to assign new tensors to the parameters instead of changing the
|
||||
existing parameters in-place when converting an ``nn.Module``.
|
||||
|
||||
When enabled, the following methods will assign new parameters to the module:
|
||||
|
||||
#. ``module.{device}()`` (e.g. :meth:`nn.Module.cuda()`) for moving a module between devices
|
||||
#. ``module.{dtype}()`` (e.g. :meth:`nn.Module.float()`) for converting a module to a different dtype
|
||||
#. :meth:`nn.Module.to`
|
||||
#. :meth:`nn.Module.to_empty`
|
||||
|
||||
Args:
|
||||
value (bool): Whether to assign new tensors or not.
|
||||
|
||||
"""
|
||||
global _overwrite_module_params_on_conversion
|
||||
_overwrite_module_params_on_conversion = value
|
||||
|
||||
|
||||
def get_overwrite_module_params_on_conversion() -> bool:
|
||||
"""
|
||||
Returns whether to assign new tensors to the parameters instead of changing the
|
||||
existing parameters in-place when converting an :class:`torch.nn.Module`. Defaults to ``False``.
|
||||
|
||||
See :func:`~torch.__future__.set_overwrite_module_params_on_conversion` for more information.
|
||||
"""
|
||||
return _overwrite_module_params_on_conversion
|
||||
|
||||
|
||||
def set_swap_module_params_on_conversion(value: bool) -> None:
|
||||
"""
|
||||
Sets whether to use :func:`~torch.utils.swap_tensors` instead of setting ``.data`` to
|
||||
change the existing parameters in-place when converting an ``nn.Module`` and instead
|
||||
of ``param.copy_(state_dict[key])`` when loading a state dict into an ``nn.Module``.
|
||||
|
||||
.. note::
|
||||
This function takes precedence over :func:`~torch.__future__.get_overwrite_module_params_on_conversion`
|
||||
|
||||
When enabled, the following methods will swap the existing parameters in-place:
|
||||
|
||||
#. ``module.{device}()`` (e.g. :meth:`nn.Module.cuda()`) for moving a module between devices
|
||||
#. ``module.{dtype}()`` (e.g. :meth:`nn.Module.float()`) for converting a module to a different dtype
|
||||
#. :meth:`nn.Module.to`
|
||||
#. :meth:`nn.Module.to_empty`
|
||||
#. :meth:`nn.Module.load_state_dict`
|
||||
|
||||
The semantics for :meth:`~nn.Module.load_state_dict` when this is set are as follows:
|
||||
|
||||
#. For each parameter/buffer, its corresponding ``state_dict['key']`` is transformed via
|
||||
:meth:`~torch.Tensor.module_load` (i.e. ``res = param.module_load(state_dict['key'])``)
|
||||
#. If necessary, ``res`` will be wrapped in an :class:`~nn.Parameter`
|
||||
#. The parameter/buffer in the module will be swapped via :func:`~torch.utils.swap_tensors`
|
||||
with ``res``
|
||||
|
||||
Args:
|
||||
value (bool): Whether to use :func:`~torch.utils.swap_tensors` or not.
|
||||
|
||||
"""
|
||||
global _swap_module_params_on_conversion
|
||||
_swap_module_params_on_conversion = value
|
||||
|
||||
|
||||
def get_swap_module_params_on_conversion() -> bool:
|
||||
"""
|
||||
Returns whether to use :func:`~torch.utils.swap_tensors` instead of setting .data to
|
||||
change the existing parameters in-place when converting an ``nn.Module``. Defaults to ``False``.
|
||||
|
||||
See :func:`~torch.__future__.set_swap_module_params_on_conversion` for more information.
|
||||
"""
|
||||
return _swap_module_params_on_conversion
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,666 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) 2005-2010 ActiveState Software Inc.
|
||||
# Copyright (c) 2013 Eddy Petrișor
|
||||
|
||||
# flake8: noqa
|
||||
|
||||
"""
|
||||
This file is directly from
|
||||
https://github.com/ActiveState/appdirs/blob/3fe6a83776843a46f20c2e5587afcffe05e03b39/appdirs.py
|
||||
|
||||
The license of https://github.com/ActiveState/appdirs copied below:
|
||||
|
||||
|
||||
# This is the MIT license
|
||||
|
||||
Copyright (c) 2010 ActiveState Software Inc.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a
|
||||
copy of this software and associated documentation files (the
|
||||
"Software"), to deal in the Software without restriction, including
|
||||
without limitation the rights to use, copy, modify, merge, publish,
|
||||
distribute, sublicense, and/or sell copies of the Software, and to
|
||||
permit persons to whom the Software is furnished to do so, subject to
|
||||
the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included
|
||||
in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
|
||||
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
|
||||
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
||||
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
|
||||
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
|
||||
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
"""
|
||||
|
||||
"""Utilities for determining application-specific dirs.
|
||||
|
||||
See <https://github.com/ActiveState/appdirs> for details and usage.
|
||||
"""
|
||||
# Dev Notes:
|
||||
# - Windows "Known Folders": https://learn.microsoft.com/en-us/windows/win32/shell/csidl
|
||||
# - macOS File System Programming Guide: https://developer.apple.com/library/archive/documentation/FileManagement/Conceptual/FileSystemProgrammingGuide/Introduction/Introduction.html
|
||||
# - XDG spec for Un*x: https://standards.freedesktop.org/basedir-spec/basedir-spec-latest.html
|
||||
|
||||
__version__ = "1.4.4"
|
||||
__version_info__ = tuple(int(segment) for segment in __version__.split("."))
|
||||
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
unicode = str
|
||||
|
||||
if sys.platform.startswith("java"):
|
||||
import platform
|
||||
|
||||
os_name = platform.java_ver()[3][0]
|
||||
if os_name.startswith("Windows"): # "Windows XP", "Windows 7", etc.
|
||||
system = "win32"
|
||||
elif os_name.startswith("Mac"): # "Mac OS X", etc.
|
||||
system = "darwin"
|
||||
else: # "Linux", "SunOS", "FreeBSD", etc.
|
||||
# Setting this to "linux2" is not ideal, but only Windows or Mac
|
||||
# are actually checked for and the rest of the module expects
|
||||
# *sys.platform* style strings.
|
||||
system = "linux2"
|
||||
else:
|
||||
system = sys.platform
|
||||
|
||||
|
||||
def user_data_dir(appname=None, appauthor=None, version=None, roaming=False):
|
||||
r"""Return full path to the user-specific data dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"roaming" (boolean, default False) can be set True to use the Windows
|
||||
roaming appdata directory. That means that for users on a Windows
|
||||
network setup for roaming profiles, this user data will be
|
||||
sync'd on login. See
|
||||
<http://technet.microsoft.com/en-us/library/cc766489(WS.10).aspx>
|
||||
for a discussion of issues.
|
||||
|
||||
Typical user data directories are:
|
||||
Mac OS X: ~/Library/Application Support/<AppName>
|
||||
Unix: ~/.local/share/<AppName> # or in $XDG_DATA_HOME, if defined
|
||||
Win XP (not roaming): C:\Documents and Settings\<username>\Application Data\<AppAuthor>\<AppName>
|
||||
Win XP (roaming): C:\Documents and Settings\<username>\Local Settings\Application Data\<AppAuthor>\<AppName>
|
||||
Win 7 (not roaming): C:\Users\<username>\AppData\Local\<AppAuthor>\<AppName>
|
||||
Win 7 (roaming): C:\Users\<username>\AppData\Roaming\<AppAuthor>\<AppName>
|
||||
|
||||
For Unix, we follow the XDG spec and support $XDG_DATA_HOME.
|
||||
That means, by default "~/.local/share/<AppName>".
|
||||
"""
|
||||
if system == "win32":
|
||||
if appauthor is None:
|
||||
appauthor = appname
|
||||
const = roaming and "CSIDL_APPDATA" or "CSIDL_LOCAL_APPDATA"
|
||||
path = os.path.normpath(_get_win_folder(const))
|
||||
if appname:
|
||||
if appauthor is not False:
|
||||
path = os.path.join(path, appauthor, appname)
|
||||
else:
|
||||
path = os.path.join(path, appname)
|
||||
elif system == "darwin":
|
||||
path = os.path.expanduser("~/Library/Application Support/")
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
else:
|
||||
path = os.getenv("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
return path
|
||||
|
||||
|
||||
def site_data_dir(appname=None, appauthor=None, version=None, multipath=False):
|
||||
r"""Return full path to the user-shared data dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"multipath" is an optional parameter only applicable to *nix
|
||||
which indicates that the entire list of data dirs should be
|
||||
returned. By default, the first item from XDG_DATA_DIRS is
|
||||
returned, or '/usr/local/share/<AppName>',
|
||||
if XDG_DATA_DIRS is not set
|
||||
|
||||
Typical site data directories are:
|
||||
Mac OS X: /Library/Application Support/<AppName>
|
||||
Unix: /usr/local/share/<AppName> or /usr/share/<AppName>
|
||||
Win XP: C:\Documents and Settings\All Users\Application Data\<AppAuthor>\<AppName>
|
||||
Vista: (Fail! "C:\ProgramData" is a hidden *system* directory on Vista.)
|
||||
Win 7: C:\ProgramData\<AppAuthor>\<AppName> # Hidden, but writeable on Win 7.
|
||||
|
||||
For Unix, this is using the $XDG_DATA_DIRS[0] default.
|
||||
|
||||
WARNING: Do not use this on Windows. See the Vista-Fail note above for why.
|
||||
"""
|
||||
if system == "win32":
|
||||
if appauthor is None:
|
||||
appauthor = appname
|
||||
path = os.path.normpath(_get_win_folder("CSIDL_COMMON_APPDATA"))
|
||||
if appname:
|
||||
if appauthor is not False:
|
||||
path = os.path.join(path, appauthor, appname)
|
||||
else:
|
||||
path = os.path.join(path, appname)
|
||||
elif system == "darwin":
|
||||
path = os.path.expanduser("/Library/Application Support")
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
else:
|
||||
# XDG default for $XDG_DATA_DIRS
|
||||
# only first, if multipath is False
|
||||
path = os.getenv(
|
||||
"XDG_DATA_DIRS", os.pathsep.join(["/usr/local/share", "/usr/share"])
|
||||
)
|
||||
pathlist = [
|
||||
os.path.expanduser(x.rstrip(os.sep)) for x in path.split(os.pathsep)
|
||||
]
|
||||
if appname:
|
||||
if version:
|
||||
appname = os.path.join(appname, version)
|
||||
pathlist = [os.sep.join([x, appname]) for x in pathlist]
|
||||
|
||||
if multipath:
|
||||
path = os.pathsep.join(pathlist)
|
||||
else:
|
||||
path = pathlist[0]
|
||||
return path
|
||||
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
return path
|
||||
|
||||
|
||||
def user_config_dir(appname=None, appauthor=None, version=None, roaming=False):
|
||||
r"""Return full path to the user-specific config dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"roaming" (boolean, default False) can be set True to use the Windows
|
||||
roaming appdata directory. That means that for users on a Windows
|
||||
network setup for roaming profiles, this user data will be
|
||||
sync'd on login. See
|
||||
<http://technet.microsoft.com/en-us/library/cc766489(WS.10).aspx>
|
||||
for a discussion of issues.
|
||||
|
||||
Typical user config directories are:
|
||||
Mac OS X: ~/Library/Preferences/<AppName>
|
||||
Unix: ~/.config/<AppName> # or in $XDG_CONFIG_HOME, if defined
|
||||
Win *: same as user_data_dir
|
||||
|
||||
For Unix, we follow the XDG spec and support $XDG_CONFIG_HOME.
|
||||
That means, by default "~/.config/<AppName>".
|
||||
"""
|
||||
if system == "win32":
|
||||
path = user_data_dir(appname, appauthor, None, roaming)
|
||||
elif system == "darwin":
|
||||
path = os.path.expanduser("~/Library/Preferences/")
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
else:
|
||||
path = os.getenv("XDG_CONFIG_HOME", os.path.expanduser("~/.config"))
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
return path
|
||||
|
||||
|
||||
def site_config_dir(appname=None, appauthor=None, version=None, multipath=False):
|
||||
r"""Return full path to the user-shared data dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"multipath" is an optional parameter only applicable to *nix
|
||||
which indicates that the entire list of config dirs should be
|
||||
returned. By default, the first item from XDG_CONFIG_DIRS is
|
||||
returned, or '/etc/xdg/<AppName>', if XDG_CONFIG_DIRS is not set
|
||||
|
||||
Typical site config directories are:
|
||||
Mac OS X: same as site_data_dir
|
||||
Unix: /etc/xdg/<AppName> or $XDG_CONFIG_DIRS[i]/<AppName> for each value in
|
||||
$XDG_CONFIG_DIRS
|
||||
Win *: same as site_data_dir
|
||||
Vista: (Fail! "C:\ProgramData" is a hidden *system* directory on Vista.)
|
||||
|
||||
For Unix, this is using the $XDG_CONFIG_DIRS[0] default, if multipath=False
|
||||
|
||||
WARNING: Do not use this on Windows. See the Vista-Fail note above for why.
|
||||
"""
|
||||
if system == "win32":
|
||||
path = site_data_dir(appname, appauthor)
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
elif system == "darwin":
|
||||
path = os.path.expanduser("/Library/Preferences")
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
else:
|
||||
# XDG default for $XDG_CONFIG_DIRS
|
||||
# only first, if multipath is False
|
||||
path = os.getenv("XDG_CONFIG_DIRS", "/etc/xdg")
|
||||
pathlist = [
|
||||
os.path.expanduser(x.rstrip(os.sep)) for x in path.split(os.pathsep)
|
||||
]
|
||||
if appname:
|
||||
if version:
|
||||
appname = os.path.join(appname, version)
|
||||
pathlist = [os.sep.join([x, appname]) for x in pathlist]
|
||||
|
||||
if multipath:
|
||||
path = os.pathsep.join(pathlist)
|
||||
else:
|
||||
path = pathlist[0]
|
||||
return path
|
||||
|
||||
|
||||
def user_cache_dir(appname=None, appauthor=None, version=None, opinion=True):
|
||||
r"""Return full path to the user-specific cache dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"opinion" (boolean) can be False to disable the appending of
|
||||
"Cache" to the base app data dir for Windows. See
|
||||
discussion below.
|
||||
|
||||
Typical user cache directories are:
|
||||
Mac OS X: ~/Library/Caches/<AppName>
|
||||
Unix: ~/.cache/<AppName> (XDG default)
|
||||
Win XP: C:\Documents and Settings\<username>\Local Settings\Application Data\<AppAuthor>\<AppName>\Cache
|
||||
Vista: C:\Users\<username>\AppData\Local\<AppAuthor>\<AppName>\Cache
|
||||
|
||||
On Windows the only suggestion in the MSDN docs is that local settings go in
|
||||
the `CSIDL_LOCAL_APPDATA` directory. This is identical to the non-roaming
|
||||
app data dir (the default returned by `user_data_dir` above). Apps typically
|
||||
put cache data somewhere *under* the given dir here. Some examples:
|
||||
...\Mozilla\Firefox\Profiles\<ProfileName>\Cache
|
||||
...\Acme\SuperApp\Cache\1.0
|
||||
OPINION: This function appends "Cache" to the `CSIDL_LOCAL_APPDATA` value.
|
||||
This can be disabled with the `opinion=False` option.
|
||||
"""
|
||||
if system == "win32":
|
||||
if appauthor is None:
|
||||
appauthor = appname
|
||||
path = os.path.normpath(_get_win_folder("CSIDL_LOCAL_APPDATA"))
|
||||
if appname:
|
||||
if appauthor is not False:
|
||||
path = os.path.join(path, appauthor, appname)
|
||||
else:
|
||||
path = os.path.join(path, appname)
|
||||
if opinion:
|
||||
path = os.path.join(path, "Cache")
|
||||
elif system == "darwin":
|
||||
path = os.path.expanduser("~/Library/Caches")
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
else:
|
||||
path = os.getenv("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
return path
|
||||
|
||||
|
||||
def user_state_dir(appname=None, appauthor=None, version=None, roaming=False):
|
||||
r"""Return full path to the user-specific state dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"roaming" (boolean, default False) can be set True to use the Windows
|
||||
roaming appdata directory. That means that for users on a Windows
|
||||
network setup for roaming profiles, this user data will be
|
||||
sync'd on login. See
|
||||
<http://technet.microsoft.com/en-us/library/cc766489(WS.10).aspx>
|
||||
for a discussion of issues.
|
||||
|
||||
Typical user state directories are:
|
||||
Mac OS X: same as user_data_dir
|
||||
Unix: ~/.local/state/<AppName> # or in $XDG_STATE_HOME, if defined
|
||||
Win *: same as user_data_dir
|
||||
|
||||
For Unix, we follow this Debian proposal <https://wiki.debian.org/XDGBaseDirectorySpecification#state>
|
||||
to extend the XDG spec and support $XDG_STATE_HOME.
|
||||
|
||||
That means, by default "~/.local/state/<AppName>".
|
||||
"""
|
||||
if system in ["win32", "darwin"]:
|
||||
path = user_data_dir(appname, appauthor, None, roaming)
|
||||
else:
|
||||
path = os.getenv("XDG_STATE_HOME", os.path.expanduser("~/.local/state"))
|
||||
if appname:
|
||||
path = os.path.join(path, appname)
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
return path
|
||||
|
||||
|
||||
def user_log_dir(appname=None, appauthor=None, version=None, opinion=True):
|
||||
r"""Return full path to the user-specific log dir for this application.
|
||||
|
||||
"appname" is the name of application.
|
||||
If None, just the system directory is returned.
|
||||
"appauthor" (only used on Windows) is the name of the
|
||||
appauthor or distributing body for this application. Typically
|
||||
it is the owning company name. This falls back to appname. You may
|
||||
pass False to disable it.
|
||||
"version" is an optional version path element to append to the
|
||||
path. You might want to use this if you want multiple versions
|
||||
of your app to be able to run independently. If used, this
|
||||
would typically be "<major>.<minor>".
|
||||
Only applied when appname is present.
|
||||
"opinion" (boolean) can be False to disable the appending of
|
||||
"Logs" to the base app data dir for Windows, and "log" to the
|
||||
base cache dir for Unix. See discussion below.
|
||||
|
||||
Typical user log directories are:
|
||||
Mac OS X: ~/Library/Logs/<AppName>
|
||||
Unix: ~/.cache/<AppName>/log # or under $XDG_CACHE_HOME if defined
|
||||
Win XP: C:\Documents and Settings\<username>\Local Settings\Application Data\<AppAuthor>\<AppName>\Logs
|
||||
Vista: C:\Users\<username>\AppData\Local\<AppAuthor>\<AppName>\Logs
|
||||
|
||||
On Windows the only suggestion in the MSDN docs is that local settings
|
||||
go in the `CSIDL_LOCAL_APPDATA` directory. (Note: I'm interested in
|
||||
examples of what some windows apps use for a logs dir.)
|
||||
|
||||
OPINION: This function appends "Logs" to the `CSIDL_LOCAL_APPDATA`
|
||||
value for Windows and appends "log" to the user cache dir for Unix.
|
||||
This can be disabled with the `opinion=False` option.
|
||||
"""
|
||||
if system == "darwin":
|
||||
path = os.path.join(os.path.expanduser("~/Library/Logs"), appname)
|
||||
elif system == "win32":
|
||||
path = user_data_dir(appname, appauthor, version)
|
||||
version = False
|
||||
if opinion:
|
||||
path = os.path.join(path, "Logs")
|
||||
else:
|
||||
path = user_cache_dir(appname, appauthor, version)
|
||||
version = False
|
||||
if opinion:
|
||||
path = os.path.join(path, "log")
|
||||
if appname and version:
|
||||
path = os.path.join(path, version)
|
||||
return path
|
||||
|
||||
|
||||
class AppDirs:
|
||||
"""Convenience wrapper for getting application dirs."""
|
||||
|
||||
def __init__(
|
||||
self, appname=None, appauthor=None, version=None, roaming=False, multipath=False
|
||||
):
|
||||
self.appname = appname
|
||||
self.appauthor = appauthor
|
||||
self.version = version
|
||||
self.roaming = roaming
|
||||
self.multipath = multipath
|
||||
|
||||
@property
|
||||
def user_data_dir(self):
|
||||
return user_data_dir(
|
||||
self.appname, self.appauthor, version=self.version, roaming=self.roaming
|
||||
)
|
||||
|
||||
@property
|
||||
def site_data_dir(self):
|
||||
return site_data_dir(
|
||||
self.appname, self.appauthor, version=self.version, multipath=self.multipath
|
||||
)
|
||||
|
||||
@property
|
||||
def user_config_dir(self):
|
||||
return user_config_dir(
|
||||
self.appname, self.appauthor, version=self.version, roaming=self.roaming
|
||||
)
|
||||
|
||||
@property
|
||||
def site_config_dir(self):
|
||||
return site_config_dir(
|
||||
self.appname, self.appauthor, version=self.version, multipath=self.multipath
|
||||
)
|
||||
|
||||
@property
|
||||
def user_cache_dir(self):
|
||||
return user_cache_dir(self.appname, self.appauthor, version=self.version)
|
||||
|
||||
@property
|
||||
def user_state_dir(self):
|
||||
return user_state_dir(self.appname, self.appauthor, version=self.version)
|
||||
|
||||
@property
|
||||
def user_log_dir(self):
|
||||
return user_log_dir(self.appname, self.appauthor, version=self.version)
|
||||
|
||||
|
||||
# ---- internal support stuff
|
||||
|
||||
|
||||
def _get_win_folder_from_registry(csidl_name):
|
||||
"""This is a fallback technique at best. I'm not sure if using the
|
||||
registry for this guarantees us the correct answer for all CSIDL_*
|
||||
names.
|
||||
"""
|
||||
import winreg as _winreg
|
||||
|
||||
shell_folder_name = {
|
||||
"CSIDL_APPDATA": "AppData",
|
||||
"CSIDL_COMMON_APPDATA": "Common AppData",
|
||||
"CSIDL_LOCAL_APPDATA": "Local AppData",
|
||||
}[csidl_name]
|
||||
|
||||
key = _winreg.OpenKey(
|
||||
_winreg.HKEY_CURRENT_USER,
|
||||
r"Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders",
|
||||
)
|
||||
dir, _type = _winreg.QueryValueEx(key, shell_folder_name)
|
||||
return dir
|
||||
|
||||
|
||||
def _get_win_folder_with_pywin32(csidl_name):
|
||||
from win32com.shell import shell, shellcon
|
||||
|
||||
dir = shell.SHGetFolderPath(0, getattr(shellcon, csidl_name), 0, 0)
|
||||
# Try to make this a unicode path because SHGetFolderPath does
|
||||
# not return unicode strings when there is unicode data in the
|
||||
# path.
|
||||
try:
|
||||
dir = unicode(dir)
|
||||
|
||||
# Downgrade to short path name if have highbit chars. See
|
||||
# <http://bugs.activestate.com/show_bug.cgi?id=85099>.
|
||||
has_high_char = False
|
||||
for c in dir:
|
||||
if ord(c) > 255:
|
||||
has_high_char = True
|
||||
break
|
||||
if has_high_char:
|
||||
try:
|
||||
import win32api
|
||||
|
||||
dir = win32api.GetShortPathName(dir)
|
||||
except ImportError:
|
||||
pass
|
||||
except UnicodeError:
|
||||
pass
|
||||
return dir
|
||||
|
||||
|
||||
def _get_win_folder_with_ctypes(csidl_name):
|
||||
import ctypes
|
||||
|
||||
csidl_const = {
|
||||
"CSIDL_APPDATA": 26,
|
||||
"CSIDL_COMMON_APPDATA": 35,
|
||||
"CSIDL_LOCAL_APPDATA": 28,
|
||||
}[csidl_name]
|
||||
|
||||
buf = ctypes.create_unicode_buffer(1024)
|
||||
ctypes.windll.shell32.SHGetFolderPathW(None, csidl_const, None, 0, buf)
|
||||
|
||||
# Downgrade to short path name if have highbit chars. See
|
||||
# <http://bugs.activestate.com/show_bug.cgi?id=85099>.
|
||||
has_high_char = False
|
||||
for c in buf:
|
||||
if ord(c) > 255:
|
||||
has_high_char = True
|
||||
break
|
||||
if has_high_char:
|
||||
buf2 = ctypes.create_unicode_buffer(1024)
|
||||
if ctypes.windll.kernel32.GetShortPathNameW(buf.value, buf2, 1024):
|
||||
buf = buf2
|
||||
|
||||
return buf.value
|
||||
|
||||
|
||||
def _get_win_folder_with_jna(csidl_name):
|
||||
import array
|
||||
|
||||
from com.sun import jna
|
||||
from com.sun.jna.platform import win32
|
||||
|
||||
buf_size = win32.WinDef.MAX_PATH * 2
|
||||
buf = array.zeros("c", buf_size)
|
||||
shell = win32.Shell32.INSTANCE
|
||||
shell.SHGetFolderPath(
|
||||
None,
|
||||
getattr(win32.ShlObj, csidl_name),
|
||||
None,
|
||||
win32.ShlObj.SHGFP_TYPE_CURRENT,
|
||||
buf,
|
||||
)
|
||||
dir = jna.Native.toString(buf.tostring()).rstrip("\0")
|
||||
|
||||
# Downgrade to short path name if have highbit chars. See
|
||||
# <http://bugs.activestate.com/show_bug.cgi?id=85099>.
|
||||
has_high_char = False
|
||||
for c in dir:
|
||||
if ord(c) > 255:
|
||||
has_high_char = True
|
||||
break
|
||||
if has_high_char:
|
||||
buf = array.zeros("c", buf_size)
|
||||
kernel = win32.Kernel32.INSTANCE
|
||||
if kernel.GetShortPathName(dir, buf, buf_size):
|
||||
dir = jna.Native.toString(buf.tostring()).rstrip("\0")
|
||||
|
||||
return dir
|
||||
|
||||
|
||||
if system == "win32":
|
||||
try:
|
||||
import win32com.shell
|
||||
|
||||
_get_win_folder = _get_win_folder_with_pywin32
|
||||
except ImportError:
|
||||
try:
|
||||
from ctypes import windll
|
||||
|
||||
_get_win_folder = _get_win_folder_with_ctypes
|
||||
except ImportError:
|
||||
try:
|
||||
import com.sun.jna
|
||||
|
||||
_get_win_folder = _get_win_folder_with_jna
|
||||
except ImportError:
|
||||
_get_win_folder = _get_win_folder_from_registry
|
||||
|
||||
|
||||
# ---- self test code
|
||||
|
||||
if __name__ == "__main__":
|
||||
appname = "MyApp"
|
||||
appauthor = "MyCompany"
|
||||
|
||||
props = (
|
||||
"user_data_dir",
|
||||
"user_config_dir",
|
||||
"user_cache_dir",
|
||||
"user_state_dir",
|
||||
"user_log_dir",
|
||||
"site_data_dir",
|
||||
"site_config_dir",
|
||||
)
|
||||
|
||||
print(f"-- app dirs {__version__} --")
|
||||
|
||||
print("-- app dirs (with optional 'version')")
|
||||
dirs = AppDirs(appname, appauthor, version="1.0")
|
||||
for prop in props:
|
||||
print(f"{prop}: {getattr(dirs, prop)}")
|
||||
|
||||
print("\n-- app dirs (without optional 'version')")
|
||||
dirs = AppDirs(appname, appauthor)
|
||||
for prop in props:
|
||||
print(f"{prop}: {getattr(dirs, prop)}")
|
||||
|
||||
print("\n-- app dirs (without optional 'appauthor')")
|
||||
dirs = AppDirs(appname)
|
||||
for prop in props:
|
||||
print(f"{prop}: {getattr(dirs, prop)}")
|
||||
|
||||
print("\n-- app dirs (with disabled 'appauthor')")
|
||||
dirs = AppDirs(appname, appauthor=False)
|
||||
for prop in props:
|
||||
print(f"{prop}: {getattr(dirs, prop)}")
|
||||
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
import torch
|
||||
|
||||
# pyrefly: ignore [bad-dunder-all]
|
||||
__all__ = ['Await']
|
||||
|
||||
W = TypeVar("W")
|
||||
|
||||
class _PyAwaitMeta(type(torch._C._Await), type(Generic)): # type: ignore[misc, no-redef]
|
||||
pass
|
||||
|
||||
class _Await(torch._C._Await, Generic[W], metaclass=_PyAwaitMeta):
|
||||
r"""
|
||||
Wrapper around a ``torch._C.Await`` which encapsulates delayed execution
|
||||
of a callable. All manipulations happen with functions ``torch.jit._awaitable``,
|
||||
``torch.jit._awaitable_wait``, ``torch.jit._awaitable_nowait``.
|
||||
|
||||
Torch scriptable manipulations:
|
||||
``torch.jit._awaitable(func, *args)``
|
||||
Creates ``Await[W]`` object, where W is return type of func.
|
||||
|
||||
Returns:
|
||||
``torch.jit._awaitable_wait(Await[W])``
|
||||
Returns the result of the function, specified at ``_awaitable``, with specified arguments.
|
||||
|
||||
Returns:
|
||||
The result of type ``W`` of the function call. The result is owned by ``Await[W]``
|
||||
and returned on all following ``_awaitable_wait`` calls.
|
||||
|
||||
|
||||
``torch.jit._awaitable_nowait(W)``
|
||||
Returns:
|
||||
Trivial ``Await[W]`` with specified result.
|
||||
|
||||
|
||||
Only in eager mode:
|
||||
``fn() -> Callable[Tuple[Any], W]``
|
||||
Returns:
|
||||
Specified at ``_awaitable`` python function ``func``.
|
||||
|
||||
``args() -> Tuple[Any]``
|
||||
Returns:
|
||||
Specified at ``_awaitable`` python args.
|
||||
|
||||
``is_nowait() -> _bool``
|
||||
Returns:
|
||||
``True`` if this object was created via ``_awaitable_nowait`` call (trivial `Await[W]`).
|
||||
|
||||
In eager mode ``Await[W]`` can be used as ``W`` i.e. attributes of W can be called on ``Await[W]``,
|
||||
``_awaitable_wait()`` call will be transparently added.
|
||||
"""
|
||||
@@ -0,0 +1,56 @@
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import torch._C
|
||||
|
||||
|
||||
class _ClassNamespace(types.ModuleType):
|
||||
def __init__(self, name: str) -> None:
|
||||
super().__init__("torch.classes" + name)
|
||||
self.name = name
|
||||
|
||||
def __getattr__(self, attr: str) -> Any:
|
||||
proxy = torch._C._get_custom_class_python_wrapper(self.name, attr)
|
||||
if proxy is None:
|
||||
raise RuntimeError(f"Class {self.name}.{attr} not registered!")
|
||||
return proxy
|
||||
|
||||
|
||||
class _Classes(types.ModuleType):
|
||||
__file__ = "_classes.py"
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__("torch.classes")
|
||||
|
||||
def __getattr__(self, name: str) -> _ClassNamespace:
|
||||
namespace = _ClassNamespace(name)
|
||||
setattr(self, name, namespace)
|
||||
return namespace
|
||||
|
||||
@property
|
||||
def loaded_libraries(self) -> Any:
|
||||
return torch.ops.loaded_libraries
|
||||
|
||||
def load_library(self, path: str) -> None:
|
||||
"""
|
||||
Loads a shared library from the given path into the current process.
|
||||
|
||||
The library being loaded may run global initialization code to register
|
||||
custom classes with the PyTorch JIT runtime. This allows dynamically
|
||||
loading custom classes. For this, you should compile your class
|
||||
and the static registration code into a shared library object, and then
|
||||
call ``torch.classes.load_library('path/to/libcustom.so')`` to load the
|
||||
shared object.
|
||||
|
||||
After the library is loaded, it is added to the
|
||||
``torch.classes.loaded_libraries`` attribute, a set that may be inspected
|
||||
for the paths of all libraries loaded using this function.
|
||||
|
||||
Args:
|
||||
path (str): A path to a shared library to load.
|
||||
"""
|
||||
torch.ops.load_library(path)
|
||||
|
||||
|
||||
# The classes "namespace"
|
||||
classes = _Classes()
|
||||
@@ -0,0 +1,60 @@
|
||||
"""
|
||||
APIs related to torch.compile which lazily import torch._dynamo to avoid
|
||||
circular dependencies.
|
||||
"""
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from typing import overload, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
@overload
|
||||
def _disable_dynamo(
|
||||
fn: Callable[_P, _T], recursive: bool = True
|
||||
) -> Callable[_P, _T]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def _disable_dynamo(
|
||||
fn: None = None, recursive: bool = True
|
||||
) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: ...
|
||||
|
||||
|
||||
def _disable_dynamo(
|
||||
fn: Callable[_P, _T] | None = None, recursive: bool = True
|
||||
) -> Callable[_P, _T] | Callable[[Callable[_P, _T]], Callable[_P, _T]]:
|
||||
"""
|
||||
This API should be only used inside torch, external users should still use
|
||||
torch._dynamo.disable. The main goal of this API is to avoid circular
|
||||
imports issues that is common while using _dynamo.disable inside torch
|
||||
itself.
|
||||
|
||||
This API avoids it by lazily importing torch._dynamo from the import time to
|
||||
the invocation of the decorated function.
|
||||
"""
|
||||
if fn is not None:
|
||||
|
||||
@functools.wraps(fn)
|
||||
def inner(*args: _P.args, **kwargs: _P.kwargs) -> _T:
|
||||
# cache this on the first invocation to avoid adding too much overhead.
|
||||
disable_fn = getattr(fn, "__dynamo_disable", None)
|
||||
if disable_fn is None:
|
||||
import torch._dynamo
|
||||
|
||||
# We can safely turn off functools.wraps here because the inner
|
||||
# already wraps fn in the outer scope.
|
||||
disable_fn = torch._dynamo.disable(fn, recursive, wrapping=False)
|
||||
fn.__dynamo_disable = disable_fn # type: ignore[attr-defined]
|
||||
|
||||
return disable_fn(*args, **kwargs)
|
||||
|
||||
return inner
|
||||
else:
|
||||
# decorator usage like @_disable_dynamo(recursive=False). The resulting
|
||||
# object expects the original decorated function as the arg.
|
||||
return functools.partial(_disable_dynamo, recursive=recursive)
|
||||
@@ -0,0 +1,314 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
from collections import namedtuple
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
|
||||
|
||||
# NOTE [CustomOp autograd kernel indirection]
|
||||
# We register `inner` as the autograd kernel for this custom_op.
|
||||
# `inner` either calls the autograd formula registered by the user,
|
||||
# or goes into an `autograd_not_implemented` kernel.
|
||||
#
|
||||
# The reason why this indirection exists is
|
||||
# so that we can swap out the autograd kernel (the PyTorch dispatcher
|
||||
# doesn't actually allow us to do this). By default, we want
|
||||
# the `autograd_not_implemented` behavior, but then the user may come
|
||||
# and register something that is actually a backward formula
|
||||
def autograd_kernel_indirection(custom_op):
|
||||
autograd_fallback = autograd_not_implemented(custom_op)
|
||||
|
||||
def inner(*args, **kwargs):
|
||||
if custom_op._has_impl("autograd"):
|
||||
kernel = custom_op._get_impl("autograd").func
|
||||
return kernel(*args, **kwargs)
|
||||
# As explained in NOTE ["backward", "save_for_backward", and "autograd"],
|
||||
# after the user gives us "backward" and "save_for_backward", we generate
|
||||
# the "autograd" impl. If the user only provided one, then we tell
|
||||
# the user they've done something wrong.
|
||||
if custom_op._has_impl("save_for_backward") or custom_op._has_impl("backward"):
|
||||
missing = (
|
||||
"save_for_backward" if custom_op._has_impl("backward") else "backward"
|
||||
)
|
||||
found = "save_for_backward" if missing == "backward" else "backward"
|
||||
loc = custom_op._get_impl(found).location
|
||||
raise RuntimeError(
|
||||
f"We found a '{found}' registration for {custom_op} at "
|
||||
f"{loc} but were unable to find a '{missing}' registration. "
|
||||
f"To use the CustomOp API to register a backward formula, "
|
||||
f"please provide us both a backward function and a "
|
||||
f"'save for backward' function via `impl_backward` and "
|
||||
f"`impl_save_for_backward` respectively."
|
||||
)
|
||||
return autograd_fallback(*args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
# TODO(#101191): Use the actual C++ autograd not implemented fallback,
|
||||
# or change the default autograd fallback to the autograd not implemented fallback.
|
||||
def autograd_not_implemented(custom_op):
|
||||
def kernel(*args, **kwargs):
|
||||
if torch.is_grad_enabled() and pytree.tree_any(
|
||||
lambda x: isinstance(x, torch.Tensor) and x.requires_grad, (args, kwargs)
|
||||
):
|
||||
raise RuntimeError("Autograd has not been implemented for operator")
|
||||
with torch._C._AutoDispatchBelowAutograd():
|
||||
return custom_op(*args, **kwargs)
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
def mark_non_differentiable(ctx, output, output_differentiability):
|
||||
# Output types are restricted to be:
|
||||
# - Tensor
|
||||
# - Tensor[]
|
||||
# - int, bool, Scalar, float
|
||||
# See _check_can_register_backward
|
||||
if output_differentiability is not None:
|
||||
if not isinstance(output, tuple):
|
||||
tuple_output = (output,)
|
||||
else:
|
||||
tuple_output = output # type: ignore[assignment]
|
||||
if len(output_differentiability) != len(tuple_output):
|
||||
raise AssertionError(
|
||||
f"output_differentiability length {len(output_differentiability)} "
|
||||
f"!= output length {len(tuple_output)}"
|
||||
)
|
||||
non_differentiable_tensors = []
|
||||
for idx, (differentiable, out) in enumerate(
|
||||
zip(output_differentiability, tuple_output)
|
||||
):
|
||||
if isinstance(out, torch.Tensor):
|
||||
if not differentiable:
|
||||
non_differentiable_tensors.append(out)
|
||||
continue
|
||||
if isinstance(out, list):
|
||||
if not differentiable:
|
||||
non_differentiable_tensors.extend(out)
|
||||
continue
|
||||
if differentiable:
|
||||
raise RuntimeError(
|
||||
f"With output_differentiability={output_differentiability}. "
|
||||
f"At idx {idx}, we received an object of type {type(out)} that "
|
||||
f"is not a Tensor, so it cannot have be marked as differentiable in "
|
||||
f"output_differentiability."
|
||||
)
|
||||
if non_differentiable_tensors:
|
||||
ctx.mark_non_differentiable(*non_differentiable_tensors)
|
||||
|
||||
|
||||
def construct_autograd_kernel(
|
||||
schema,
|
||||
output_differentiability,
|
||||
custom_op,
|
||||
op_overload,
|
||||
save_for_backward_fn,
|
||||
backward_fn,
|
||||
):
|
||||
def apply(*args):
|
||||
flat_args, spec = pytree.tree_flatten(args)
|
||||
out_spec = None
|
||||
|
||||
def forward(ctx, *flat_args):
|
||||
ctx.set_materialize_grads(True)
|
||||
args = pytree.tree_unflatten(list(flat_args), spec)
|
||||
with torch._C._AutoDispatchBelowAutograd():
|
||||
output = op_overload(*args)
|
||||
|
||||
# We use the info about args to give better error messages in backward
|
||||
args_info = namedtuple_args(schema, pytree.tree_map(type, args))
|
||||
|
||||
save_for_backward_fn_inputs = namedtuple_args(schema, args)
|
||||
to_save = save_for_backward_fn(save_for_backward_fn_inputs, output)
|
||||
|
||||
save_pytree_for_backward(ctx, (to_save, args_info))
|
||||
mark_non_differentiable(ctx, output, output_differentiability)
|
||||
|
||||
nonlocal out_spec
|
||||
flat_output, out_spec = pytree.tree_flatten(output)
|
||||
return tuple(flat_output)
|
||||
|
||||
def backward(ctx, *flat_grad_output):
|
||||
if out_spec is None:
|
||||
raise AssertionError("out_spec is unexpectedly None")
|
||||
grads = pytree.tree_unflatten(list(flat_grad_output), out_spec)
|
||||
saved, args_info = unpack_saved(ctx)
|
||||
# There is nothing on the ctx object for now, it is just there so
|
||||
# that we can add additional things in the future.
|
||||
inner_ctx = object()
|
||||
if not isinstance(grads, tuple):
|
||||
grads = (grads,)
|
||||
grad_inputs_dict = backward_fn(inner_ctx, saved, *grads)
|
||||
|
||||
# Massage the grad_inputs_dict to a form acceptable by
|
||||
# autograd.Function.
|
||||
validate_grad_inputs_dict(grad_inputs_dict, custom_op, args_info)
|
||||
return grad_inputs_dict_to_flat_tuple(grad_inputs_dict, args_info)
|
||||
|
||||
generated_cls = gen_autograd_function(
|
||||
custom_op._opname + "_customop", forward, backward
|
||||
)
|
||||
|
||||
flat_output = generated_cls.apply(*flat_args)
|
||||
if out_spec is None:
|
||||
raise AssertionError("out_spec is unexpectedly None")
|
||||
return pytree.tree_unflatten(list(flat_output), out_spec)
|
||||
|
||||
return apply
|
||||
|
||||
|
||||
def gen_autograd_function(name, forward, backward):
|
||||
generated_cls = type(
|
||||
name,
|
||||
(torch.autograd.Function,),
|
||||
{
|
||||
"forward": staticmethod(forward),
|
||||
"backward": staticmethod(backward),
|
||||
},
|
||||
)
|
||||
return generated_cls
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def namedtuple_args_cls(schema):
|
||||
attribs = [arg.name for arg in schema.arguments.flat_all]
|
||||
name = str(schema.name) + "_args"
|
||||
# mypy doesn't support dynamic namedtuple name
|
||||
tuple_cls = namedtuple(name, attribs) # type: ignore[misc]
|
||||
return tuple_cls
|
||||
|
||||
|
||||
def namedtuple_args(schema, args):
|
||||
if not isinstance(args, tuple):
|
||||
raise AssertionError(f"expected tuple, got {type(args)}")
|
||||
tuple_cls = namedtuple_args_cls(schema)
|
||||
return tuple_cls(*args)
|
||||
|
||||
|
||||
def validate_grad_inputs_dict(grad_inputs_dict, forward_op, args_info):
|
||||
def error(what):
|
||||
backward = forward_op._get_impl("backward")
|
||||
raise RuntimeError(
|
||||
f"In the backward function defined for {forward_op} at "
|
||||
f"{backward.location} using the CustomOp API, {what}"
|
||||
)
|
||||
|
||||
if not isinstance(grad_inputs_dict, dict):
|
||||
error(
|
||||
f"expected the output of the backward function to be a dict but "
|
||||
f"got {type(grad_inputs_dict)}"
|
||||
)
|
||||
|
||||
expected_keys = {
|
||||
arg.name
|
||||
for arg in forward_op._schema.arguments.flat_all
|
||||
if arg.type.is_tensor_like()
|
||||
}
|
||||
actual_keys = grad_inputs_dict.keys()
|
||||
if expected_keys != actual_keys:
|
||||
error(
|
||||
f"expected the returned grad_input dict to have keys "
|
||||
f"{expected_keys} but got {actual_keys}. The backward "
|
||||
f"function must return a gradient (can be None) for each arg "
|
||||
f"to the CustomOp that may be a Tensor or Sequence[Tensor]. "
|
||||
f"Args declared to be non-Tensor-like types should not appear "
|
||||
f"in the grad_input dict"
|
||||
)
|
||||
|
||||
for name, grad in grad_inputs_dict.items():
|
||||
arg_info = getattr(args_info, name)
|
||||
|
||||
if isinstance(arg_info, list):
|
||||
if not isinstance(grad, (tuple, list)):
|
||||
error(
|
||||
f"for input '{name}' expected the grad_input dict to "
|
||||
f"hold a list of gradients but got object of type "
|
||||
f"{type(grad)}."
|
||||
)
|
||||
if len(grad) != len(arg_info):
|
||||
error(
|
||||
f"for input '{name}' expected the grad_input dict to "
|
||||
f"hold a list of {len(arg_info)} gradients but got "
|
||||
f"{len(grad)}"
|
||||
)
|
||||
for idx, (g, info) in enumerate(zip(grad, arg_info)):
|
||||
if g is None:
|
||||
continue
|
||||
if not isinstance(g, torch.Tensor):
|
||||
error(
|
||||
f"for input '{name}' expected the grad_input dict to "
|
||||
f"hold a list of None or Tensor gradients but got "
|
||||
f"object of {type(g)} at index {idx}"
|
||||
)
|
||||
if not issubclass(info, torch.Tensor):
|
||||
error(
|
||||
f"for input '{name}', got a Tensor as the gradient "
|
||||
f"for the {idx}-th value but expected None because "
|
||||
f"the {idx}-th value was not a Tensor (it was "
|
||||
f"type {arg_info}"
|
||||
)
|
||||
continue
|
||||
|
||||
if grad is None:
|
||||
continue
|
||||
if not isinstance(grad, torch.Tensor):
|
||||
error(
|
||||
f"got object of type {type(grad)} as the gradient for input "
|
||||
f"'{name}', "
|
||||
f"but expected the gradient to be either None or a Tensor"
|
||||
)
|
||||
if not issubclass(arg_info, torch.Tensor):
|
||||
error(
|
||||
f"got a Tensor as the gradient for input '{name}' but "
|
||||
f"expected None as the gradient because input '{name}' "
|
||||
f"was not a Tensor (it was type {arg_info})."
|
||||
)
|
||||
|
||||
|
||||
def grad_inputs_dict_to_flat_tuple(grad_inputs_dict, args_info):
|
||||
result = []
|
||||
for name, arg_info in args_info._asdict().items():
|
||||
if name not in grad_inputs_dict:
|
||||
result.append(pytree.tree_map(lambda x: None, arg_info))
|
||||
continue
|
||||
result.append(grad_inputs_dict[name])
|
||||
return tuple(pytree.tree_leaves(result))
|
||||
|
||||
|
||||
# Saves "stuff" (a pytree) onto the ctx object. Use unpack_saved to unpack it.
|
||||
# autograd.Function prefers that users use ctx.save_for_backward to
|
||||
# save Tensors (to avoid reference cycles) and for non-Tensors to go onto the
|
||||
# ctx object.
|
||||
def save_pytree_for_backward(ctx, stuff):
|
||||
flat_stuff, spec = pytree.tree_flatten(stuff)
|
||||
num_elts = len(flat_stuff)
|
||||
tensor_idxs = [
|
||||
idx for idx, thing in enumerate(flat_stuff) if isinstance(thing, torch.Tensor)
|
||||
]
|
||||
non_tensor_idxs = [
|
||||
idx
|
||||
for idx, thing in enumerate(flat_stuff)
|
||||
if not isinstance(thing, torch.Tensor)
|
||||
]
|
||||
tensors = [thing for thing in flat_stuff if isinstance(thing, torch.Tensor)]
|
||||
non_tensors = [thing for thing in flat_stuff if not isinstance(thing, torch.Tensor)]
|
||||
|
||||
ctx.spec = spec
|
||||
ctx.num_elts = num_elts
|
||||
ctx.save_for_backward(*tensors)
|
||||
ctx.tensor_idxs = tensor_idxs
|
||||
ctx.saved_non_tensors = non_tensors
|
||||
ctx.non_tensor_idxs = non_tensor_idxs
|
||||
|
||||
|
||||
# Inverse operation to save_pytree_for_backward
|
||||
def unpack_saved(ctx):
|
||||
flat_stuff = [None] * ctx.num_elts
|
||||
for tensor, idx in zip(ctx.saved_tensors, ctx.tensor_idxs):
|
||||
flat_stuff[idx] = tensor
|
||||
for non_tensor, idx in zip(ctx.saved_non_tensors, ctx.non_tensor_idxs):
|
||||
flat_stuff[idx] = non_tensor
|
||||
stuff = pytree.tree_unflatten(flat_stuff, ctx.spec)
|
||||
return stuff
|
||||
@@ -0,0 +1,719 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import dataclasses
|
||||
import functools
|
||||
import inspect
|
||||
import sys
|
||||
import typing
|
||||
import warnings
|
||||
import weakref
|
||||
|
||||
import torch
|
||||
import torch._C as _C
|
||||
import torch._library.infer_schema
|
||||
import torch.library as library
|
||||
from torch._library.infer_schema import infer_schema
|
||||
from torch.library import get_ctx
|
||||
from torchgen.model import (
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
OperatorName,
|
||||
SchemaKind,
|
||||
)
|
||||
|
||||
from .autograd import autograd_kernel_indirection, construct_autograd_kernel
|
||||
|
||||
|
||||
"""
|
||||
torch._custom_op is deprecated. We shipped a production-ready version of it into torch.library.
|
||||
Please use those APIs instead.
|
||||
"""
|
||||
|
||||
__all__ = ["custom_op", "CustomOp", "get_ctx"]
|
||||
|
||||
|
||||
SUPPORTED_DEVICE_TYPE_TO_KEY = {
|
||||
"cpu": "CPU",
|
||||
"cuda": "CUDA",
|
||||
}
|
||||
|
||||
# We will not let users register CustomOps with anything that could look like
|
||||
# PyTorch internals to avoid confusion.
|
||||
RESERVED_NS = {
|
||||
"prim",
|
||||
"prims",
|
||||
"aten",
|
||||
"at",
|
||||
"torch",
|
||||
"pytorch",
|
||||
}
|
||||
|
||||
|
||||
def warn_deprecated():
|
||||
warnings.warn(
|
||||
"torch._custom_op is deprecated and will be removed in PyTorch 2.6, please "
|
||||
"use the equivalent torch.library API instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
|
||||
def custom_op(qualname: str, manual_schema: str | None = None) -> typing.Callable:
|
||||
r"""
|
||||
This API is deprecated, please use torch.library.custom_op instead
|
||||
"""
|
||||
warn_deprecated()
|
||||
|
||||
def inner(func):
|
||||
if not inspect.isfunction(func):
|
||||
raise ValueError(
|
||||
f"custom_op(...)(func): Expected `func` to be a Python "
|
||||
f"function, got: {type(func)}"
|
||||
)
|
||||
|
||||
ns, name = parse_qualname(qualname)
|
||||
validate_namespace(ns)
|
||||
if func.__name__ != name:
|
||||
raise ValueError(
|
||||
f"custom_op(qualname='{qualname}', ...)(func): expected `func` "
|
||||
f"to have name '{name}' but got '{func.__name__}'. "
|
||||
f"Please either change the name of `func` or the qualname that "
|
||||
f"is passed to `custom_op`"
|
||||
)
|
||||
|
||||
schema = (
|
||||
infer_schema(func, mutates_args=())
|
||||
if manual_schema is None
|
||||
else manual_schema
|
||||
)
|
||||
schema_str = f"{name}{schema}"
|
||||
function_schema = FunctionSchema.parse(schema_str)
|
||||
validate_schema(function_schema)
|
||||
if manual_schema is not None:
|
||||
validate_function_matches_schema(function_schema, func)
|
||||
|
||||
lib = library.Library(ns, "FRAGMENT")
|
||||
lib.define(schema_str)
|
||||
ophandle = find_ophandle_or_throw(ns, function_schema.name)
|
||||
result = CustomOp(
|
||||
lib, ns, function_schema, name, ophandle, _private_access=True
|
||||
)
|
||||
|
||||
result.__name__ = func.__name__ # pyrefly: ignore [bad-assignment]
|
||||
result.__module__ = func.__module__
|
||||
result.__doc__ = func.__doc__
|
||||
|
||||
library.impl(lib, result._opname, "Autograd")(
|
||||
autograd_kernel_indirection(weakref.proxy(result))
|
||||
)
|
||||
|
||||
torch._C._dispatch_set_report_error_callback(
|
||||
ophandle, functools.partial(report_error_callback, weakref.proxy(result))
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
# Global dictionary holding references to all CustomOp objects
|
||||
# Yes, it keeps all CustomOps alive (see NOTE [CustomOp lifetime])
|
||||
# Used to query the CustomOp associated with a specific C++ dispatcher operator.
|
||||
# An example usage is FakeTensor: FakeTensor checks if a specific operator
|
||||
# has an implementation registered via the CustomOp API.
|
||||
# Indexed by qualname (e.g. aten::foo)
|
||||
global_registry: dict[str, "CustomOp"] = {}
|
||||
|
||||
|
||||
class CustomOp:
|
||||
r"""
|
||||
This API is deprecated, please use torch.library.custom_op instead
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, lib, cpp_ns, schema, operator_name, ophandle, *, _private_access=False
|
||||
):
|
||||
super().__init__()
|
||||
warn_deprecated()
|
||||
if not _private_access:
|
||||
raise RuntimeError(
|
||||
"The CustomOp constructor is private and we do not guarantee "
|
||||
"BC for it. Please use custom_op(...) to create a CustomOp object"
|
||||
)
|
||||
name = f"{cpp_ns}::{operator_name}"
|
||||
self._schema = schema
|
||||
self._cpp_ns = cpp_ns
|
||||
self._lib: library.Library = lib
|
||||
self._ophandle: _C._DispatchOperatorHandle = ophandle
|
||||
# Has the name of the op, e.g. "foo". We cache here for convenience.
|
||||
self._opname: str = operator_name
|
||||
# this is _opname but with namespace. e.g. "custom::foo"
|
||||
self._qualname: str = name
|
||||
self.__name__ = None # mypy requires this
|
||||
# NB: Some of these impls are registered as kernels to DispatchKeys.
|
||||
# Modifying the _impls dict directly won't do anything in that case.
|
||||
self._impls: dict[str, FuncAndLocation | None] = {}
|
||||
# See NOTE [CustomOp autograd kernel indirection]
|
||||
self._registered_autograd_kernel_indirection = False
|
||||
|
||||
global_registry[self._qualname] = self
|
||||
|
||||
def _register_autograd_kernel_indirection(self):
|
||||
if self._registered_autograd_kernel_indirection:
|
||||
raise AssertionError("autograd kernel indirection already registered")
|
||||
self._lib.impl(
|
||||
self._opname, autograd_kernel_indirection(weakref.proxy(self)), "Autograd"
|
||||
)
|
||||
self._registered_autograd_kernel_indirection = True
|
||||
|
||||
# Records the impl and the source location in self._impls
|
||||
# Note that this doesn't cause torch.library to use the impl, that
|
||||
# needs to be done in a separate self._lib.impl call.
|
||||
def _register_impl(self, kind, func, stacklevel=2):
|
||||
if self._has_impl(kind):
|
||||
func_and_location = self._impls[kind]
|
||||
if func_and_location is None:
|
||||
raise AssertionError("func_and_location is unexpectedly None")
|
||||
location = func_and_location.location
|
||||
raise RuntimeError(
|
||||
f"Attempting to register a {kind} impl for operator {self._qualname} "
|
||||
f"that already has a {kind} impl registered from Python at "
|
||||
f"{location}. This is not supported."
|
||||
)
|
||||
frame = inspect.getframeinfo(sys._getframe(stacklevel))
|
||||
location = f"{frame.filename}:{frame.lineno}"
|
||||
self._impls[kind] = FuncAndLocation(func, location)
|
||||
|
||||
def _get_impl(self, kind):
|
||||
return self._impls[kind]
|
||||
|
||||
def _has_impl(self, kind):
|
||||
return kind in self._impls
|
||||
|
||||
def _destroy(self):
|
||||
# NOTE: [CustomOp lifetime]
|
||||
# A CustomOp, once created, lives forever. The mechanism is that the
|
||||
# global registry holds a reference to it. However, to make testing
|
||||
# easier, we want to be able to destroy CustomOp objects.
|
||||
# CustomOp._destroy does the job, though it leaves the CustomOp
|
||||
# in a garbage state.
|
||||
del self._lib
|
||||
|
||||
opnamespace = getattr(torch.ops, self._cpp_ns)
|
||||
if hasattr(opnamespace, self._opname):
|
||||
delattr(opnamespace, self._opname)
|
||||
|
||||
del global_registry[self._qualname]
|
||||
|
||||
def __repr__(self):
|
||||
return f'<CustomOp(op="{self._qualname}")>'
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
# Bypass torch.ops.* and directly do OperatorHandle::callBoxed.
|
||||
# Using torch.ops.* is a bit of a pain (it can be slow and it has lifetime
|
||||
# issues from caching operators that make testing CustomOp difficult).
|
||||
result = _C._dispatch_call_boxed(self._ophandle, *args, **kwargs)
|
||||
return result
|
||||
|
||||
def impl(
|
||||
self,
|
||||
device_types: str | typing.Iterable[str],
|
||||
_stacklevel=2,
|
||||
) -> typing.Callable:
|
||||
r"""
|
||||
This API is deprecated, please use torch.library.custom_op instead
|
||||
"""
|
||||
if isinstance(device_types, str):
|
||||
device_types = [device_types]
|
||||
for device_type in device_types:
|
||||
validate_device_type(device_type)
|
||||
|
||||
def inner(f):
|
||||
for device_type in set(device_types):
|
||||
self._check_doesnt_have_library_impl(device_type)
|
||||
self._register_impl(device_type, f, stacklevel=_stacklevel)
|
||||
dispatch_key = SUPPORTED_DEVICE_TYPE_TO_KEY[device_type]
|
||||
library.impl(self._lib, self._opname, dispatch_key)(f)
|
||||
return f
|
||||
|
||||
return inner
|
||||
|
||||
def _check_doesnt_have_library_impl(self, device_type):
|
||||
if self._has_impl(device_type):
|
||||
return
|
||||
key = SUPPORTED_DEVICE_TYPE_TO_KEY[device_type]
|
||||
if _C._dispatch_has_computed_kernel_for_dispatch_key(self._qualname, key):
|
||||
raise RuntimeError(
|
||||
f"impl(..., device_types={device_type}): the operator {self._qualname} "
|
||||
f"already has an implementation for this device type via a "
|
||||
f"pre-existing torch.library or TORCH_LIBRARY registration."
|
||||
)
|
||||
|
||||
def impl_factory(self) -> typing.Callable:
|
||||
r"""Register an implementation for a factory function."""
|
||||
|
||||
def inner(f):
|
||||
self._register_impl("factory", f)
|
||||
library.impl(self._lib, self._opname, "BackendSelect")(f)
|
||||
return f
|
||||
|
||||
return inner
|
||||
|
||||
def impl_abstract(self, _stacklevel=2) -> typing.Callable:
|
||||
r"""
|
||||
This API is deprecated, please use torch.library.custom_op instead
|
||||
"""
|
||||
|
||||
def inner(f):
|
||||
self._check_doesnt_have_library_meta_impl()
|
||||
self._register_impl("abstract", f, stacklevel=_stacklevel)
|
||||
location = self._get_impl("abstract").location
|
||||
|
||||
qualname = self._qualname
|
||||
|
||||
# Handle DispatchKey.Meta registration
|
||||
@functools.wraps(f)
|
||||
def f_with_ctx(*args, **kwargs):
|
||||
def error_on_ctx():
|
||||
raise RuntimeError(
|
||||
f"Attempted to call get_ctx() for the meta implementation "
|
||||
f"for {qualname}."
|
||||
f"You have presumably called get_ctx() because the operator "
|
||||
f"has a data-dependent output shape; if so, there is no "
|
||||
f"such meta implementation and this error is the correct "
|
||||
f"behavior. Otherwise, please remove the call to get_ctx() "
|
||||
f"in the implementation registered with impl_abstract "
|
||||
f"at {location}"
|
||||
)
|
||||
|
||||
with torch._library.fake_impl.set_ctx_getter(error_on_ctx):
|
||||
return f(*args, **kwargs)
|
||||
|
||||
self._lib.impl(self._opname, f_with_ctx, "Meta")
|
||||
return f
|
||||
|
||||
return inner
|
||||
|
||||
def _check_can_register_backward(self):
|
||||
def error(detail):
|
||||
raise RuntimeError(
|
||||
f"Cannot use torch._custom_ops APIs to register backward "
|
||||
f"formula for {detail}. Got operator "
|
||||
f"{self._qualname} with schema: {schema}"
|
||||
)
|
||||
|
||||
schema = self._schema
|
||||
if schema.kind() != SchemaKind.functional:
|
||||
error("non-functional operator")
|
||||
|
||||
rets = schema.returns
|
||||
if not schema.returns:
|
||||
error("operator with no returns")
|
||||
|
||||
if len(rets) <= 0:
|
||||
raise AssertionError(f"expected at least one return, got {len(rets)}")
|
||||
is_non_mutating_view = any(
|
||||
r.annotation is not None and not r.annotation.is_write for r in rets
|
||||
)
|
||||
if is_non_mutating_view:
|
||||
error("operator that returns views")
|
||||
|
||||
# We make assumptions about the schema's return types.
|
||||
allowed_return_types = {
|
||||
BaseType(BaseTy.int): "int",
|
||||
BaseType(BaseTy.SymInt): "SymInt",
|
||||
BaseType(BaseTy.bool): "bool",
|
||||
BaseType(BaseTy.float): "float",
|
||||
BaseType(BaseTy.Tensor): "Tensor",
|
||||
ListType(BaseType(BaseTy.Tensor), None): "List[Tensor]",
|
||||
}
|
||||
for ret in schema.returns:
|
||||
if ret.type in allowed_return_types:
|
||||
continue
|
||||
error(
|
||||
f"operator with return not in {list(allowed_return_types.values())} (got {ret.type})"
|
||||
)
|
||||
|
||||
def _check_doesnt_have_library_autograd_impl(self):
|
||||
if self._registered_autograd_kernel_indirection:
|
||||
return
|
||||
|
||||
if _C._dispatch_has_kernel_for_dispatch_key(
|
||||
self._qualname, "CompositeImplicitAutograd"
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"impl_backward/impl_save_for_backward: the operator {self._qualname} "
|
||||
f"already has an implementation for this device type via a "
|
||||
f"pre-existing registration to DispatchKey::CompositeImplicitAutograd."
|
||||
f"CompositeImplicitAutograd operators do not need an autograd formula; "
|
||||
f"instead, the operator will decompose into its constituents and those "
|
||||
f"can have autograd formulas defined on them."
|
||||
)
|
||||
|
||||
# We can improve this by adding "all Autograd<BACKEND> keys", but
|
||||
# realistically people will just be using this API for CPU/CUDA for now.
|
||||
for key in ["Autograd", "AutogradCPU", "AutogradCUDA"]:
|
||||
if _C._dispatch_has_kernel_for_dispatch_key(self._qualname, key):
|
||||
raise RuntimeError(
|
||||
f"impl_backward/impl_save_for_backward: "
|
||||
f"the operator {self._qualname} already has an Autograd kernel "
|
||||
f"registered to DispatchKey::{key} vi a pre-existing "
|
||||
f"torch.library or TORCH_LIBRARY registration. Please either "
|
||||
f"remove those registrations or don't use the torch._custom_ops APIs"
|
||||
)
|
||||
|
||||
def _check_doesnt_have_library_meta_impl(self):
|
||||
if self._has_impl("abstract"):
|
||||
return
|
||||
|
||||
# If the user's operator is CompositeExplicitAutograd,
|
||||
# allow them to impl_abstract. This is being pragmatic
|
||||
# (existing custom ops may have CompositeExplicitAutograd
|
||||
# registration that don't work with Meta kernels, so this
|
||||
# gives them an escape hatch).
|
||||
if _C._dispatch_has_kernel_for_dispatch_key(
|
||||
self._qualname, "CompositeExplicitAutograd"
|
||||
) and not _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "Meta"):
|
||||
return
|
||||
|
||||
# Otherwise, if the user's already has a Meta kernel or their
|
||||
# op is CompositeImplicitAutograd or some other alias dispatch key,
|
||||
# raise.
|
||||
|
||||
# Special case for CompositeImplicitAutograd
|
||||
if _C._dispatch_has_kernel_for_dispatch_key(
|
||||
self._qualname, "CompositeImplicitAutograd"
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"impl_abstract(...): the operator {self._qualname} "
|
||||
f"already has an implementation for this device type via a "
|
||||
f"pre-existing registration to DispatchKey::CompositeImplicitAutograd."
|
||||
f"CompositeImplicitAutograd operators do not need an abstract impl; "
|
||||
f"instead, the operator will decompose into its constituents and those "
|
||||
f"can have abstract impls defined on them."
|
||||
)
|
||||
|
||||
if _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "Meta"):
|
||||
raise RuntimeError(
|
||||
f"impl_abstract(...): the operator {self._qualname} "
|
||||
f"already has an DispatchKey::Meta implementation via a "
|
||||
f"pre-existing torch.library or TORCH_LIBRARY registration. "
|
||||
f"Please either remove that registration or don't call impl_abstract."
|
||||
)
|
||||
|
||||
# NOTE ["backward", "save_for_backward", and "autograd"]
|
||||
# As a part of the explicit autograd API, a user must provide us
|
||||
# a "save_for_backward" function and a "backward" function.
|
||||
# When both of these have been provided, then we automatically
|
||||
# construct the "autograd" kernel.
|
||||
def _register_autograd_kernel(self):
|
||||
if not self._has_impl("backward"):
|
||||
raise AssertionError("backward impl must be registered first")
|
||||
if not self._has_impl("save_for_backward"):
|
||||
raise AssertionError("save_for_backward impl must be registered first")
|
||||
kernel = construct_autograd_kernel(
|
||||
self._schema,
|
||||
self._output_differentiability,
|
||||
self,
|
||||
get_op(self._qualname),
|
||||
self._get_impl("save_for_backward").func,
|
||||
self._get_impl("backward").func,
|
||||
)
|
||||
self._register_impl("autograd", kernel)
|
||||
|
||||
def impl_save_for_backward(self, _stacklevel=2):
|
||||
r"""Register a function that tells us what to save for backward.
|
||||
|
||||
Please see impl_backward for more details.
|
||||
"""
|
||||
|
||||
def inner(f):
|
||||
self._check_can_register_backward()
|
||||
self._check_doesnt_have_library_autograd_impl()
|
||||
if not self._registered_autograd_kernel_indirection:
|
||||
self._register_autograd_kernel_indirection()
|
||||
self._register_impl("save_for_backward", f, stacklevel=_stacklevel)
|
||||
if self._has_impl("backward"):
|
||||
self._register_autograd_kernel()
|
||||
|
||||
return inner
|
||||
|
||||
def impl_backward(self, output_differentiability=None, _stacklevel=2):
|
||||
r"""
|
||||
This API is deprecated, please use torch.library.custom_op instead
|
||||
"""
|
||||
if output_differentiability is not None:
|
||||
|
||||
def yell():
|
||||
raise RuntimeError(
|
||||
f"impl_backward(output_differentiability): expected "
|
||||
f"output_differentiability to be a list of bools with "
|
||||
f"length equal to the number of outputs of this CustomOp "
|
||||
f"got: {output_differentiability}"
|
||||
)
|
||||
|
||||
if not isinstance(output_differentiability, list):
|
||||
yell()
|
||||
for diff in output_differentiability:
|
||||
if not isinstance(diff, bool):
|
||||
yell()
|
||||
if len(self._schema.returns) != len(output_differentiability):
|
||||
yell()
|
||||
|
||||
def inner(f):
|
||||
self._check_can_register_backward()
|
||||
self._check_doesnt_have_library_autograd_impl()
|
||||
if not self._registered_autograd_kernel_indirection:
|
||||
self._register_autograd_kernel_indirection()
|
||||
self._register_impl("backward", f, stacklevel=_stacklevel)
|
||||
self._output_differentiability = output_differentiability
|
||||
if self._has_impl("save_for_backward"):
|
||||
self._register_autograd_kernel()
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class FuncAndLocation:
|
||||
func: typing.Callable
|
||||
location: str
|
||||
|
||||
|
||||
def find_ophandle_or_throw(cpp_ns: str, operator_name: OperatorName):
|
||||
overload_name = (
|
||||
"" if operator_name.overload_name is None else operator_name.overload_name
|
||||
)
|
||||
return _C._dispatch_find_schema_or_throw(
|
||||
f"{cpp_ns}::{str(operator_name.name)}", overload_name
|
||||
)
|
||||
|
||||
|
||||
def validate_namespace(ns: str) -> None:
|
||||
if "." in ns:
|
||||
raise ValueError(
|
||||
f'custom_op(..., ns="{ns}"): expected ns to not contain any . (and be a '
|
||||
f"valid variable name)"
|
||||
)
|
||||
if ns in RESERVED_NS:
|
||||
raise ValueError(
|
||||
f"custom_op(..., ns='{ns}'): '{ns}' is a reserved namespace, "
|
||||
f"please choose something else. "
|
||||
)
|
||||
|
||||
|
||||
def validate_schema(schema: FunctionSchema) -> None:
|
||||
if not torch._library.utils.is_functional_schema(schema):
|
||||
raise ValueError(
|
||||
f"custom_op only supports functional operators "
|
||||
f"(ops that do not mutate any inputs, do not return "
|
||||
f"views of the inputs, and has at least one return). "
|
||||
f"Got the following non-functional schema: {schema}"
|
||||
)
|
||||
|
||||
# For simplicity: don't allow self arguments
|
||||
if schema.arguments.self_arg is not None:
|
||||
raise ValueError(
|
||||
f"custom_op does not support arguments named 'self'. Please "
|
||||
f"rename your argument. Got: {schema}"
|
||||
)
|
||||
|
||||
|
||||
def parse_qualname(qualname: str) -> tuple[str, str]:
|
||||
names = qualname.split("::", 1)
|
||||
if len(names) != 2:
|
||||
raise ValueError(
|
||||
f"Expected there to be a namespace in {qualname}, i.e. The "
|
||||
f"operator name should look something like ns::foo"
|
||||
)
|
||||
if "." in names[1]:
|
||||
raise ValueError(
|
||||
f"The torch.custom_ops APIs do not handle overloads, "
|
||||
f"i.e. operator names with '.' in them. "
|
||||
f"Please name your operator something like ns::foo. "
|
||||
f"Got: {qualname}"
|
||||
)
|
||||
return names[0], names[1]
|
||||
|
||||
|
||||
def validate_device_type(device_type: str) -> None:
|
||||
if device_type not in SUPPORTED_DEVICE_TYPE_TO_KEY:
|
||||
raise ValueError(
|
||||
f"CustomOp.impl(device_types=[{device_type}, ...]): we only support device_type "
|
||||
f"in {SUPPORTED_DEVICE_TYPE_TO_KEY.keys()}."
|
||||
)
|
||||
|
||||
|
||||
def supported_param(param: inspect.Parameter) -> bool:
|
||||
return param.kind in (
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
inspect.Parameter.KEYWORD_ONLY,
|
||||
)
|
||||
|
||||
|
||||
def validate_function_matches_schema(
|
||||
schema: FunctionSchema, func: typing.Callable
|
||||
) -> None:
|
||||
sig = inspect.signature(func)
|
||||
|
||||
if not all(supported_param(p) for _, p in sig.parameters.items()):
|
||||
raise ValueError(
|
||||
f"custom_op(..., manual_schema)(func): positional-only args, "
|
||||
f"varargs, and kwargs are not supported. Please rewrite `func` "
|
||||
f"to not have them. Got `func` with signature: {sig}"
|
||||
)
|
||||
|
||||
if (
|
||||
any(
|
||||
p.annotation is not inspect.Parameter.empty
|
||||
for _, p in sig.parameters.items()
|
||||
)
|
||||
or sig.return_annotation is not inspect.Signature.empty
|
||||
):
|
||||
raise ValueError(
|
||||
f"custom_op(..., manual_schema)(func): When passing in a manual "
|
||||
f"schema, we expect `func` to have no type annotations to avoid "
|
||||
f"ambiguity. Got `func` with signature: {sig}"
|
||||
)
|
||||
|
||||
positional = [
|
||||
(name, param)
|
||||
for name, param in sig.parameters.items()
|
||||
if param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
|
||||
]
|
||||
kwargonly = [
|
||||
(name, param)
|
||||
for name, param in sig.parameters.items()
|
||||
if param.kind == inspect.Parameter.KEYWORD_ONLY
|
||||
]
|
||||
|
||||
def error():
|
||||
raise ValueError(
|
||||
f"custom_op(..., manual_schema)(func): When passing in a manual "
|
||||
f"schema, we expect `func`'s signature to match `manual_schema` "
|
||||
f"(aside from type annotations). "
|
||||
f"func's signature: {sig}, manual_schema: {schema}"
|
||||
)
|
||||
|
||||
def error_default_args():
|
||||
raise ValueError(
|
||||
f"custom_op(..., manual_schema)(func): "
|
||||
f"neither func nor manual_schema should have default "
|
||||
f"arguments. Got "
|
||||
f"func's signature: {sig}, manual_schema: {schema}"
|
||||
)
|
||||
|
||||
def compare(sig_args, schema_args):
|
||||
if len(sig_args) != len(schema_args):
|
||||
error()
|
||||
for (name, param), arg in zip(sig_args, schema_args):
|
||||
if name != arg.name:
|
||||
error()
|
||||
if param.default is not inspect.Parameter.empty or arg.default is not None:
|
||||
error_default_args()
|
||||
|
||||
compare(positional, schema.arguments.flat_positional)
|
||||
compare(kwargonly, schema.arguments.flat_kwarg_only)
|
||||
|
||||
|
||||
def report_error_callback(custom_op: typing.Any, key: str) -> None:
|
||||
if key == "Undefined":
|
||||
raise NotImplementedError(
|
||||
f"{custom_op}: There were no Tensor inputs to this operator "
|
||||
f"(e.g. you passed an empty list of Tensors). If your operator is a "
|
||||
f"factory function (that is, it takes no Tensors and constructs "
|
||||
f"a new one), then please use CustomOp.impl_factory to register "
|
||||
f"an implementation for it"
|
||||
)
|
||||
if key == "Meta":
|
||||
raise NotImplementedError(
|
||||
f"{custom_op}: when running with device='Meta' tensors: there is no "
|
||||
f"abstract impl registered for this CustomOp. Please register one via "
|
||||
f"CustomOp.impl_abstract to get this CustomOp to work with Meta tensors"
|
||||
)
|
||||
if key in ("CPU", "CUDA"):
|
||||
device = key.lower()
|
||||
raise NotImplementedError(
|
||||
f"{custom_op}: when running with device='{device}' tensors: there is no "
|
||||
f"{device} impl registered for this CustomOp. Please register one via "
|
||||
f"CustomOp.impl(device_type='{device}')"
|
||||
)
|
||||
raise NotImplementedError(
|
||||
f"{custom_op}: No implementation for dispatch key {key}. It is likely "
|
||||
f"that we have not added this functionality yet, please either open an "
|
||||
f"issue or if you're feeling adventurous, use the low-level "
|
||||
f"torch.library API"
|
||||
)
|
||||
|
||||
|
||||
def custom_op_from_existing(op):
|
||||
ns = op.namespace
|
||||
lib = torch.library.Library(ns, "FRAGMENT")
|
||||
name = op.name().split("::")[-1]
|
||||
schema_str = str(op._schema)
|
||||
# CustomOp expects the schema string without the namespace
|
||||
schema_str = schema_str.rsplit("::", maxsplit=1)[-1]
|
||||
schema = FunctionSchema.parse(schema_str)
|
||||
return CustomOp(lib, ns, schema, name, op, _private_access=True)
|
||||
|
||||
|
||||
def get_op(qualname):
|
||||
def error_not_found():
|
||||
raise ValueError(
|
||||
f"Could not find the operator {qualname}. Please make sure you have "
|
||||
f"already registered the operator and (if registered from C++) "
|
||||
f"loaded it via torch.ops.load_library."
|
||||
)
|
||||
|
||||
ns, name = parse_qualname(qualname)
|
||||
if not hasattr(torch.ops, ns):
|
||||
error_not_found()
|
||||
opnamespace = getattr(torch.ops, ns)
|
||||
if not hasattr(opnamespace, name):
|
||||
error_not_found()
|
||||
packet = getattr(opnamespace, name)
|
||||
if not hasattr(packet, "default"):
|
||||
error_not_found()
|
||||
return packet.default
|
||||
|
||||
|
||||
def _find_custom_op(qualname, also_check_torch_library=False):
|
||||
if qualname in global_registry:
|
||||
return global_registry[qualname]
|
||||
if not also_check_torch_library:
|
||||
raise RuntimeError(
|
||||
f'Could not find custom op "{qualname}". Did you register it via '
|
||||
f"the torch._custom_ops API?"
|
||||
)
|
||||
overload = get_op(qualname)
|
||||
result = custom_op_from_existing(overload)
|
||||
return result
|
||||
|
||||
|
||||
def get_abstract_impl(qualname):
|
||||
if qualname not in torch._custom_op.impl.global_registry:
|
||||
return None
|
||||
custom_op = torch._custom_op.impl.global_registry[qualname]
|
||||
if custom_op is None:
|
||||
return None
|
||||
if not custom_op._has_impl("abstract"):
|
||||
return None
|
||||
return custom_op._get_impl("abstract").func
|
||||
|
||||
|
||||
def _custom_op_with_schema(qualname, schema, needs_fixed_stride_order=True):
|
||||
ns, name = qualname.split("::")
|
||||
schema_str = f"{name}{schema}"
|
||||
function_schema = FunctionSchema.parse(schema_str)
|
||||
validate_schema(function_schema)
|
||||
tags = [torch._C.Tag.needs_fixed_stride_order] if needs_fixed_stride_order else []
|
||||
lib = library.Library(ns, "FRAGMENT")
|
||||
lib.define(schema_str, tags=tags)
|
||||
ophandle = find_ophandle_or_throw(ns, function_schema.name)
|
||||
result = CustomOp(lib, ns, function_schema, name, ophandle, _private_access=True)
|
||||
result._register_autograd_kernel_indirection()
|
||||
|
||||
torch._C._dispatch_set_report_error_callback(
|
||||
ophandle, functools.partial(report_error_callback, weakref.proxy(result))
|
||||
)
|
||||
return get_op(qualname)
|
||||
@@ -0,0 +1,326 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import inspect
|
||||
|
||||
from torch._custom_op.impl import (
|
||||
_custom_op_with_schema,
|
||||
_find_custom_op,
|
||||
infer_schema,
|
||||
parse_qualname,
|
||||
validate_namespace,
|
||||
)
|
||||
from torch.library import get_ctx
|
||||
|
||||
|
||||
__all__ = [
|
||||
"custom_op",
|
||||
"impl",
|
||||
"impl_abstract",
|
||||
"get_ctx",
|
||||
"impl_save_for_backward",
|
||||
"impl_backward",
|
||||
]
|
||||
|
||||
|
||||
def custom_op(qualname, func_or_schema=None):
|
||||
r"""Register a new custom operator
|
||||
|
||||
In PyTorch, defining an op (short for "operator") is a two step-process:
|
||||
- we need to define the op (by providing an operator name and schema)
|
||||
- we need to implement behavior for how the operator interacts with
|
||||
various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc.
|
||||
|
||||
This entrypoint defines the custom operator (the first step)
|
||||
you must then perform the second step by calling various
|
||||
``impl_*`` APIs.
|
||||
|
||||
This API may be used as a decorator (see examples).
|
||||
|
||||
For a detailed guide on custom ops, please see
|
||||
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
|
||||
|
||||
Arguments:
|
||||
qualname (str): Should be a string that looks like
|
||||
"namespace::operator_name". Operators in PyTorch need a namespace to
|
||||
avoid name collisions; a given operator may only be created once.
|
||||
If you are writing a Python library, we recommend the namespace to
|
||||
be the name of your top-level module.
|
||||
func_or_schema (Union[Callable, str]): Each PyTorch operator needs a
|
||||
schema that tells PyTorch the types of the inputs/outputs.
|
||||
If this is a Callable, we will automatically infer the schema from
|
||||
the type annotations on the function (see examples). Otherwise,
|
||||
if you don't want to use type annotations, you may provide us the
|
||||
schema string.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> import torch
|
||||
>>> import numpy as np
|
||||
>>> from torch import Tensor
|
||||
>>>
|
||||
>>> # Step 1: define the custom op.
|
||||
>>> # We need to provide the API a "prototype function"
|
||||
>>> # (a function that returns NotImplementedError), from which
|
||||
>>> # we will infer the types of the inputs and outputs.
|
||||
>>> @torch._custom_ops.custom_op("mylibrary::numpy_sin")
|
||||
>>> def numpy_sin(x: Tensor) -> Tensor:
|
||||
>>> raise NotImplementedError
|
||||
>>>
|
||||
>>> # The custom op is now accessible via the torch.ops module:
|
||||
>>> torch.ops.mylibrary.numpy_sin
|
||||
>>>
|
||||
>>> # Step 2: Register an implementation for various PyTorch subsystems
|
||||
>>>
|
||||
>>> # Register an implementation for CPU tensors
|
||||
>>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cpu")
|
||||
>>> def numpy_sin_impl_cpu(x):
|
||||
>>> return torch.from_numpy(np.sin(x.numpy()))
|
||||
>>>
|
||||
>>> # Register an implementation for CUDA tensors
|
||||
>>> @torch._custom_ops.impl("mylibrary::numpy_sin", device_types="cuda")
|
||||
>>> def numpy_sin_impl_cuda(x):
|
||||
>>> return torch.from_numpy(np.sin(x.cpu().numpy())).to(x.device)
|
||||
>>>
|
||||
>>> x = torch.randn(3)
|
||||
>>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cpu
|
||||
>>>
|
||||
>>> x_cuda = x.cuda()
|
||||
>>> torch.ops.mylibrary.numpy_sin(x) # calls numpy_sin_impl_cuda
|
||||
|
||||
"""
|
||||
ns, name = parse_qualname(qualname)
|
||||
validate_namespace(ns)
|
||||
|
||||
def inner(func):
|
||||
if not inspect.isfunction(func):
|
||||
raise ValueError(
|
||||
f"custom_op(...)(func): Expected `func` to be a Python "
|
||||
f"function, got: {type(func)}"
|
||||
)
|
||||
|
||||
if func.__name__ != name:
|
||||
raise ValueError(
|
||||
f"custom_op(qualname='{qualname}', ...)(func): expected `func` "
|
||||
f"to have name '{name}' but got '{func.__name__}'. "
|
||||
f"Please either change the name of `func` or the qualname that "
|
||||
f"is passed to `custom_op`"
|
||||
)
|
||||
|
||||
schema = infer_schema(func, mutates_args=())
|
||||
_custom_op_with_schema(qualname, schema)
|
||||
return func
|
||||
|
||||
if func_or_schema is None:
|
||||
return inner
|
||||
if isinstance(func_or_schema, str):
|
||||
_custom_op_with_schema(qualname, func_or_schema)
|
||||
else:
|
||||
return inner(func_or_schema)
|
||||
|
||||
|
||||
def impl(qualname, *, device_types=("cpu", "cuda"), func=None):
|
||||
r"""Register an implementation for a device type for this custom op.
|
||||
|
||||
If the op is passed multiple Tensor inputs with different device
|
||||
types, it will dispatch to the registered implementation for the highest
|
||||
priority device type among those present.
|
||||
The supported device types, in order of priority, are {'cuda', 'cpu'}.
|
||||
|
||||
This API may be used as a decorator (see examples).
|
||||
|
||||
For a detailed guide on custom ops, please see
|
||||
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
|
||||
|
||||
Arguments:
|
||||
device_types (str or Iterable[str]): the device type(s) to register the function for.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> import torch
|
||||
>>> import numpy as np
|
||||
>>> from torch import Tensor
|
||||
>>>
|
||||
>>> # Step 1: define the custom op.
|
||||
>>> # We need to provide the API a "prototype function"
|
||||
>>> # (a function that returns NotImplementedError), from which
|
||||
>>> # we will infer the types of the inputs and outputs.
|
||||
>>> @torch._custom_ops.custom_op("mylibrary::numpy_cos")
|
||||
>>> def numpy_cos(x: Tensor) -> Tensor:
|
||||
>>> raise NotImplementedError
|
||||
>>>
|
||||
>>> # The custom op is now accessible via the torch.ops module:
|
||||
>>> torch.ops.mylibrary.numpy_cos
|
||||
>>>
|
||||
>>> # Step 2: Register an implementation for various PyTorch subsystems
|
||||
>>>
|
||||
>>> # Register an implementation for CPU tensors
|
||||
>>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cpu")
|
||||
>>> def numpy_cos_impl_cpu(x):
|
||||
>>> return torch.from_numpy(np.cos(x.numpy()))
|
||||
>>>
|
||||
>>> # Register an implementation for CUDA tensors
|
||||
>>> @torch._custom_ops.impl("mylibrary::numpy_cos", device_types="cuda")
|
||||
>>> def numpy_cos_impl_cuda(x):
|
||||
>>> return torch.from_numpy(np.cos(x.cpu().numpy())).to(x.device)
|
||||
>>>
|
||||
>>> x = torch.randn(3)
|
||||
>>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cpu
|
||||
>>>
|
||||
>>> x_cuda = x.cuda()
|
||||
>>> torch.ops.mylibrary.numpy_cos(x) # calls numpy_cos_impl_cuda
|
||||
|
||||
"""
|
||||
|
||||
def inner(func):
|
||||
custom_op = _find_custom_op(qualname, also_check_torch_library=True)
|
||||
custom_op.impl(device_types, _stacklevel=3)(func)
|
||||
return func
|
||||
|
||||
if func is None:
|
||||
return inner
|
||||
return inner(func)
|
||||
|
||||
|
||||
def impl_abstract(qualname, *, func=None):
|
||||
r"""Register an abstract implementation for this operator.
|
||||
|
||||
An "abstract implementation" specifies the behavior of this operator on
|
||||
Tensors that carry no data. Given some input Tensors with certain properties
|
||||
(sizes/strides/storage_offset/device), it specifies what the properties of
|
||||
the output Tensors are.
|
||||
|
||||
The abstract implementation has the same signature as the operator.
|
||||
It is run for both FakeTensors and meta tensors. To write an abstract
|
||||
implementation, assume that all Tensor inputs to the operator are
|
||||
regular CPU/CUDA/Meta tensors, but they do not have storage, and
|
||||
you are trying to return regular CPU/CUDA/Meta tensor(s) as output.
|
||||
The abstract implementation must consist of only PyTorch operations
|
||||
(and may not directly access the storage or data of any input or
|
||||
intermediate Tensors).
|
||||
|
||||
This API may be used as a decorator (see examples).
|
||||
|
||||
For a detailed guide on custom ops, please see
|
||||
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
|
||||
|
||||
Examples::
|
||||
>>> import numpy as np
|
||||
>>> from torch import Tensor
|
||||
>>>
|
||||
>>> # Example 1: an operator without data-dependent output shape
|
||||
>>> @torch._custom_ops.custom_op("mylibrary::custom_linear")
|
||||
>>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
|
||||
>>> raise NotImplementedError
|
||||
>>>
|
||||
>>> @torch._custom_ops.impl_abstract("mylibrary::custom_linear")
|
||||
>>> def custom_linear_abstract(x, weight):
|
||||
>>> assert x.dim() == 2
|
||||
>>> assert weight.dim() == 2
|
||||
>>> assert bias.dim() == 1
|
||||
>>> assert x.shape[1] == weight.shape[1]
|
||||
>>> assert weight.shape[0] == bias.shape[0]
|
||||
>>> assert x.device == weight.device
|
||||
>>>
|
||||
>>> return (x @ weight.t()) + bias
|
||||
>>>
|
||||
>>> # Example 2: an operator with data-dependent output shape
|
||||
>>> @torch._custom_ops.custom_op('mylibrary::custom_nonzero')
|
||||
>>> def custom_nonzero(x: Tensor) -> Tensor:
|
||||
>>> ...
|
||||
>>>
|
||||
>>> @torch._custom_ops.impl_abstract("mylibrary::custom_nonzero")
|
||||
>>> def custom_nonzero_abstract(x):
|
||||
>>> # Number of nonzero-elements is data-dependent.
|
||||
>>> # Since we cannot peek at the data in an abstract impl,
|
||||
>>> # we use the ctx object to construct a new symint that
|
||||
>>> # represents the data-dependent size.
|
||||
>>> ctx = torch._custom_ops.get_ctx()
|
||||
>>> nnz = ctx.create_unbacked_symint()
|
||||
>>> shape = [x.dim(), nnz]
|
||||
>>> result = x.new_empty(shape, dtype=torch.long)
|
||||
>>> return result
|
||||
>>>
|
||||
>>> @torch._custom_ops.impl("mylibrary::custom_nonzero")
|
||||
>>> def custom_nonzero_impl(x):
|
||||
>>> x_np = to_numpy(x)
|
||||
>>> res = np.stack(np.nonzero(x_np), axis=1)
|
||||
>>> # unbacked symbolic ints in PyTorch must be >= 2, so we
|
||||
>>> # constrain the range to at least 2
|
||||
>>> if res.shape[0] <= 1:
|
||||
>>> raise RuntimeError("not supported")
|
||||
>>> return torch.tensor(res, device=x.device)
|
||||
|
||||
"""
|
||||
import torch.library
|
||||
|
||||
return torch.library.register_fake(qualname, func, _stacklevel=2)
|
||||
|
||||
|
||||
def impl_save_for_backward(qualname, *, func=None):
|
||||
r"""Register a function that tells us what to save for backward.
|
||||
|
||||
Please see :func:`impl_backward` for more details.
|
||||
"""
|
||||
|
||||
def inner(func):
|
||||
custom_op = _find_custom_op(qualname, also_check_torch_library=True)
|
||||
custom_op.impl_save_for_backward(_stacklevel=3)(func)
|
||||
return func
|
||||
|
||||
if func is None:
|
||||
return inner
|
||||
return inner(func)
|
||||
|
||||
|
||||
def impl_backward(qualname, output_differentiability=None, *, func=None):
|
||||
r"""Registers a backward formula for an operator.
|
||||
|
||||
In order for an operator to work with autograd, you need to register
|
||||
a backward formula. There are two pieces to this:
|
||||
1. You must give us a function to specify what to save for backward.
|
||||
Call this the "save for backward" function.
|
||||
2. You must give us a function that computes gradients. Call this the
|
||||
"backward" function.
|
||||
|
||||
Use `impl_save_for_backward` to define a "save for backward" function
|
||||
that specifies what gets saved for backward. The function should accept
|
||||
two arguments ``(inputs, output)`` and return the quantities to be saved
|
||||
for backward.
|
||||
|
||||
During runtime, when you call the operator in a forwards pass, PyTorch
|
||||
will invoke the "save for backward" function with the inputs and output
|
||||
of the operator.
|
||||
|
||||
Use `impl_backward` to define the "backward" function. The backward
|
||||
function must accept ``(ctx, saved, *grads)``:
|
||||
- ``ctx`` is a context object where we may provide information
|
||||
- ``saved`` is exactly what gets returned from the "save for backward"
|
||||
function
|
||||
- ``grads`` is one or more gradients. The number of gradients matches
|
||||
the number of outputs of the operator.
|
||||
|
||||
The backward function must return a dict that maps the name of
|
||||
an input to the operator to its corresponding gradient. All inputs that
|
||||
were declared to be Tensors in the operator definition must be accounted
|
||||
for in the dict. The gradient may be a Tensor or None.
|
||||
|
||||
For a detailed guide on custom ops, please see
|
||||
https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
|
||||
|
||||
"""
|
||||
|
||||
def inner(func):
|
||||
custom_op = _find_custom_op(qualname, also_check_torch_library=True)
|
||||
custom_op.impl_backward(output_differentiability, _stacklevel=3)(func)
|
||||
return func
|
||||
|
||||
if func is None:
|
||||
return inner
|
||||
return inner(func)
|
||||
|
||||
|
||||
def _destroy(qualname):
|
||||
"""De-registers a custom op. For testing purposes only"""
|
||||
custom_op = _find_custom_op(qualname)
|
||||
custom_op._destroy()
|
||||
@@ -0,0 +1,560 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Sequence
|
||||
from functools import lru_cache, partial, wraps
|
||||
from itertools import chain
|
||||
from typing import Optional, TYPE_CHECKING, TypeVar, Union
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.export.decomp_utils import CustomDecompTable
|
||||
|
||||
import torch
|
||||
import torch.library
|
||||
from torch._ops import HigherOrderOperator, OperatorBase, OpOverload, OpOverloadPacket
|
||||
from torch._prims_common import CustomOutParamAnnotation
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor
|
||||
from torch.utils import _pytree as pytree
|
||||
|
||||
|
||||
__all__ = [
|
||||
"decomposition_table",
|
||||
"pre_autograd_decomposition_table",
|
||||
"meta_table",
|
||||
"register_decomposition",
|
||||
"get_decompositions",
|
||||
"core_aten_decompositions",
|
||||
"_should_decompose_because_unsafe_op",
|
||||
]
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
# TODO: relax key type here; torch registrations should be possible to; but
|
||||
# right now this type is accurate
|
||||
global_decomposition_table: dict[str, dict[torch._ops.OperatorBase, Callable]] = (
|
||||
defaultdict(dict)
|
||||
)
|
||||
|
||||
decomposition_table = global_decomposition_table["post_autograd"]
|
||||
pre_autograd_decomposition_table = global_decomposition_table["pre_autograd"]
|
||||
meta_table = global_decomposition_table["meta"]
|
||||
|
||||
|
||||
def _should_decompose_because_unsafe_op(op: torch._ops.OperatorBase) -> bool:
|
||||
"""
|
||||
Returns True if the op must always decompose in export/compile tracing system
|
||||
|
||||
In export, we always decompose certain CIA ops that are tagged with
|
||||
maybe_aliasing_or_mutating because we statically need to know if the op is
|
||||
mutating or not. But these CIA ops could have different behaviour in runtime.
|
||||
|
||||
native_batch_norm is a prim op which has a wrong schema and it needs to be replaced
|
||||
with correct schema. But until then, we will force decompose it via this tag.
|
||||
"""
|
||||
if not isinstance(op, torch._ops.OpOverload):
|
||||
return False
|
||||
if torch.Tag.maybe_aliasing_or_mutating in op.tags:
|
||||
return True
|
||||
return op is torch.ops.aten.native_batch_norm.default
|
||||
|
||||
|
||||
def _add_op_to_registry(registry, op, fn):
|
||||
"""
|
||||
This is an internal API for adding an op to the decomposition table.
|
||||
|
||||
If op is OpOverload, it will be added to the registry directly.
|
||||
If op is OpOverloadPacket, all the valid op_overloads in the packet will be added to the registry.
|
||||
"""
|
||||
overloads: list[torch._ops.OperatorBase] = []
|
||||
if isinstance(op, HigherOrderOperator):
|
||||
# There's no concept of overloads for HigherOrderOperator
|
||||
registry[op] = fn
|
||||
return
|
||||
elif isinstance(op, OpOverload):
|
||||
overloads.append(op)
|
||||
else:
|
||||
if not isinstance(op, OpOverloadPacket):
|
||||
raise AssertionError(f"expected OpOverloadPacket, got {type(op)}")
|
||||
for ol in op.overloads():
|
||||
overloads.append(getattr(op, ol))
|
||||
|
||||
for op_overload in overloads:
|
||||
if op_overload in registry:
|
||||
raise RuntimeError(f"duplicate registrations for {op_overload}")
|
||||
# TorchScript dumps a bunch of extra nonsense overloads
|
||||
# which don't have corresponding dispatcher entries, we need
|
||||
# to filter those out, e.g aten.add.float_int
|
||||
if torch._C._dispatch_has_kernel(op_overload.name()):
|
||||
registry[op_overload] = fn
|
||||
|
||||
|
||||
def _convert_out_params(f):
|
||||
out_annotation = f.__annotations__.get("out")
|
||||
|
||||
# If there are no out params, do not wrap the function.
|
||||
if not out_annotation:
|
||||
return f
|
||||
|
||||
# Hack to detect when out is a Tuple. There seems to be no pretty way of doing this
|
||||
if getattr(out_annotation, "__origin__", None) is tuple:
|
||||
sig = inspect.signature(f)
|
||||
out_names = sig.return_annotation._fields
|
||||
# If out is a tuple, we need to register a function that unpacks all the out
|
||||
# elements as this is what native_functions.yaml expects
|
||||
|
||||
@wraps(f)
|
||||
def _fn(*args, **kwargs):
|
||||
out_kwargs = tuple(kwargs.pop(o, None) for o in out_names)
|
||||
# Either all of the out kwargs are set or none of them
|
||||
is_none = out_kwargs[0] is None
|
||||
if not all((o is None) == is_none for o in out_kwargs):
|
||||
raise AssertionError(
|
||||
f"all out kwargs must be set or none of them, got {out_kwargs}"
|
||||
)
|
||||
return f(*args, **kwargs, out=None if is_none else out_kwargs)
|
||||
|
||||
out_params = [
|
||||
inspect.Parameter(
|
||||
o,
|
||||
kind=inspect.Parameter.KEYWORD_ONLY,
|
||||
default=None,
|
||||
annotation=t,
|
||||
)
|
||||
for o, t in zip(out_names, out_annotation.__args__)
|
||||
]
|
||||
# Drop the out parameter and concatenate the new kwargs in the signature
|
||||
params = chain((v for k, v in sig.parameters.items() if k != "out"), out_params)
|
||||
_fn.__signature__ = inspect.Signature( # type: ignore[attr-defined]
|
||||
parameters=params, # type: ignore[arg-type]
|
||||
return_annotation=sig.return_annotation,
|
||||
)
|
||||
# Drop the out parameter and concatenate the new kwargs in the annotations
|
||||
_fn.__annotations__ = {k: v for k, v in f.__annotations__.items() if k != "out"}
|
||||
for o in out_params:
|
||||
_fn.__annotations__[o.name] = o.annotation
|
||||
|
||||
# Propagate that this function is wrapped by `out_wrapper`
|
||||
_fn._torch_decompositions_out_wrapper = f._torch_decompositions_out_wrapper # type: ignore[attr-defined]
|
||||
|
||||
return _fn
|
||||
|
||||
# Alternatively, there may be a single tensor out parameter with a name
|
||||
# other than "out". This will need special treatment and is indicated by an
|
||||
# annotation, which we will remove here so it is not exposed after wrapping.
|
||||
custom_out_param_name = f.__annotations__.pop(CustomOutParamAnnotation, None)
|
||||
if custom_out_param_name:
|
||||
|
||||
@wraps(f)
|
||||
def _fn(*args, **kwargs):
|
||||
out_kwarg = kwargs.pop(custom_out_param_name, None)
|
||||
return f(*args, **kwargs, out=out_kwarg)
|
||||
|
||||
out_param = inspect.Parameter(
|
||||
custom_out_param_name,
|
||||
kind=inspect.Parameter.KEYWORD_ONLY,
|
||||
default=None,
|
||||
annotation=out_annotation,
|
||||
)
|
||||
|
||||
# Drop the out parameter and concatenate the new kwarg in the signature
|
||||
sig = inspect.signature(f)
|
||||
params = chain(
|
||||
(v for k, v in sig.parameters.items() if k != "out"), (out_param,)
|
||||
)
|
||||
_fn.__signature__ = inspect.Signature( # type: ignore[attr-defined]
|
||||
parameters=params, # type: ignore[arg-type]
|
||||
return_annotation=sig.return_annotation,
|
||||
)
|
||||
|
||||
# Drop the out parameter and concatenate the new kwargs in the annotations
|
||||
_fn.__annotations__ = {k: v for k, v in f.__annotations__.items() if k != "out"}
|
||||
_fn.__annotations__[out_param.name] = out_param.annotation
|
||||
|
||||
return _fn
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def register_decomposition(
|
||||
aten_op, registry=None, *, type="post_autograd", unsafe=False
|
||||
) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
|
||||
"""
|
||||
A decorator to register a function as a decomposition to the Python
|
||||
decomposition table. Use it like this::
|
||||
|
||||
@register_decomposition(torch.ops.aten.clamp_min)
|
||||
def clamp_min(x):
|
||||
return torch.clamp(self, min=min)
|
||||
|
||||
If you are writing a new decomposition, consider contributing it
|
||||
directly to PyTorch in torch._decomp.decompositions.
|
||||
|
||||
This API is experimental; we are almost certainly going to extend
|
||||
the API when we make decompositions eligible for use in transforms (e.g.,
|
||||
autograd) and not just backend tracing, where we then need to know if a
|
||||
decomposition can be used to simulate a transform.
|
||||
|
||||
By default, we also will register it to the Meta key of dispatcher,
|
||||
and replace the c++ Meta implementation if there is already one.
|
||||
|
||||
unsafe kwarg is for reuse of this function for registering non-function
|
||||
things
|
||||
"""
|
||||
|
||||
if type not in {"post_autograd", "pre_autograd", "meta"}:
|
||||
raise AssertionError(
|
||||
f"type must be one of post_autograd, pre_autograd, or meta, got {type}"
|
||||
)
|
||||
|
||||
def decomposition_decorator(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
orig_fn = fn
|
||||
if not unsafe:
|
||||
fn = _convert_out_params(fn)
|
||||
|
||||
nonlocal registry
|
||||
if registry is None:
|
||||
registry = global_decomposition_table[type]
|
||||
|
||||
def register(op):
|
||||
_add_op_to_registry(registry, op, fn)
|
||||
|
||||
# To handle allowing multiple aten_ops at once
|
||||
pytree.tree_map_(register, aten_op)
|
||||
return orig_fn
|
||||
|
||||
return decomposition_decorator
|
||||
|
||||
|
||||
def get_decompositions(
|
||||
aten_ops: Sequence[torch._ops.OperatorBase | OpOverloadPacket],
|
||||
type: str = "post_autograd",
|
||||
) -> dict[torch._ops.OperatorBase, Callable]:
|
||||
"""
|
||||
Retrieve a dictionary of decompositions corresponding to the list of
|
||||
operator overloads and overload packets passed as input. Overload
|
||||
packets will include all decomposed overloads in the packet. If there is
|
||||
no decomposition for a requested operator, it is silently ignored.
|
||||
|
||||
This API is experimental; we are almost certainly going to give an alternate,
|
||||
more recommended formulation, where a user provides the set of operators
|
||||
they know how to implement, and we provide decompositions for everything
|
||||
not in this set.
|
||||
"""
|
||||
if type not in {"post_autograd", "pre_autograd", "meta"}:
|
||||
raise AssertionError(
|
||||
f"type must be one of post_autograd, pre_autograd, or meta, got {type}"
|
||||
)
|
||||
|
||||
registry = global_decomposition_table[type]
|
||||
packets_to_overloads = defaultdict(list)
|
||||
|
||||
for opo in registry:
|
||||
if isinstance(opo, (OpOverload, OpOverloadPacket)):
|
||||
packets_to_overloads[opo.overloadpacket].append(opo)
|
||||
decompositions: dict[torch._ops.OperatorBase, Callable] = {}
|
||||
for op in aten_ops:
|
||||
if isinstance(op, OpOverloadPacket) and op in packets_to_overloads:
|
||||
for op_overload in packets_to_overloads[op]:
|
||||
decompositions[op_overload] = registry[op_overload]
|
||||
elif isinstance(op, (torch._ops.OperatorBase)) and op in registry:
|
||||
decompositions[op] = registry[op]
|
||||
return decompositions
|
||||
|
||||
|
||||
def remove_decompositions(
|
||||
decompositions: dict[torch._ops.OperatorBase, Callable],
|
||||
aten_ops: Sequence[OpOverload | OpOverloadPacket],
|
||||
) -> None:
|
||||
"""
|
||||
Given a dictionary of decompositions obtained from get_decompositions(), removes
|
||||
operators associated with a list of operator overloads and overload packets passed
|
||||
as input. If the decomposition dictionary does not contain a decomposition that is
|
||||
specified to be removed, it is silently ignored.
|
||||
"""
|
||||
for op in aten_ops:
|
||||
if isinstance(op, OpOverloadPacket):
|
||||
for overload_name in op.overloads():
|
||||
opo = getattr(op, overload_name)
|
||||
decompositions.pop(opo, None)
|
||||
elif isinstance(op, OpOverload):
|
||||
decompositions.pop(op, None)
|
||||
|
||||
|
||||
# populate the table
|
||||
import torch._decomp.decompositions
|
||||
import torch._refs
|
||||
|
||||
|
||||
def core_aten_decompositions() -> "CustomDecompTable":
|
||||
from torch.export.exported_program import default_decompositions
|
||||
|
||||
return default_decompositions()
|
||||
|
||||
|
||||
# See NOTE [Core ATen Ops]
|
||||
#
|
||||
# list was copied from torch/_inductor/decomposition.py
|
||||
# excluding decompositions that results in prim ops
|
||||
# Resulting opset of decomposition is core aten ops
|
||||
def _core_aten_decompositions_post_autograd() -> dict[
|
||||
torch._ops.OperatorBase, Callable
|
||||
]:
|
||||
aten = torch.ops.aten
|
||||
return get_decompositions(
|
||||
[
|
||||
aten.addcdiv,
|
||||
aten.addcdiv_,
|
||||
aten.addcmul,
|
||||
aten.addcmul_,
|
||||
aten.addr,
|
||||
aten.affine_grid_generator,
|
||||
aten.alias_copy,
|
||||
aten.all,
|
||||
aten.aminmax,
|
||||
aten.arange.default,
|
||||
aten.arange.start,
|
||||
aten.avg_pool2d_backward,
|
||||
aten.baddbmm,
|
||||
aten.binary_cross_entropy,
|
||||
aten.binary_cross_entropy_backward,
|
||||
aten.binary_cross_entropy_with_logits,
|
||||
aten.block_diag,
|
||||
aten.bernoulli.p,
|
||||
aten.bernoulli.default,
|
||||
aten.celu,
|
||||
aten.celu_,
|
||||
aten.channel_shuffle,
|
||||
aten.clamp_max,
|
||||
aten.clamp_min,
|
||||
aten.col2im,
|
||||
aten.count_nonzero,
|
||||
aten.linalg_cross,
|
||||
aten.cudnn_batch_norm,
|
||||
aten.cudnn_batch_norm_backward,
|
||||
aten.miopen_batch_norm_backward,
|
||||
aten.deg2rad,
|
||||
aten.deg2rad_,
|
||||
aten.detach,
|
||||
aten.diag_embed,
|
||||
aten.diagonal_backward,
|
||||
aten.diagonal_copy,
|
||||
aten.dot,
|
||||
aten.vdot,
|
||||
aten.elu_,
|
||||
aten.elu_backward,
|
||||
aten._embedding_bag,
|
||||
aten.embedding_dense_backward,
|
||||
aten.empty_like,
|
||||
aten._euclidean_dist.default,
|
||||
aten.expand_as,
|
||||
aten.expand_copy,
|
||||
aten.eye,
|
||||
aten.fill,
|
||||
aten.fill_,
|
||||
aten.floor_divide,
|
||||
aten.frac,
|
||||
aten.frac_,
|
||||
aten._fused_moving_avg_obs_fq_helper,
|
||||
aten.gelu_,
|
||||
aten.gelu_backward,
|
||||
aten.glu,
|
||||
aten.glu_backward,
|
||||
aten.hardshrink,
|
||||
aten.hardsigmoid,
|
||||
aten.hardsigmoid_,
|
||||
aten.hardsigmoid_backward,
|
||||
aten.hardswish,
|
||||
aten.hardswish_,
|
||||
aten.hardswish_backward,
|
||||
aten.hardtanh_,
|
||||
aten.hardtanh_backward,
|
||||
aten.hann_window,
|
||||
aten.heaviside,
|
||||
aten.heaviside_,
|
||||
aten.huber_loss,
|
||||
aten.huber_loss_backward,
|
||||
aten.im2col,
|
||||
aten.index_add.out,
|
||||
aten.index_add.default,
|
||||
aten.index_add_,
|
||||
aten.index_copy.out,
|
||||
aten.index_copy.default,
|
||||
aten.index_copy_,
|
||||
aten.index_fill.int_Scalar,
|
||||
aten.index_fill.int_Tensor,
|
||||
aten.index_fill.int_Scalar_out,
|
||||
aten.index_fill.int_Tensor_out,
|
||||
aten.index_fill_,
|
||||
aten.isin,
|
||||
aten.isneginf,
|
||||
aten.isposinf,
|
||||
aten.l1_loss,
|
||||
aten._lazy_clone,
|
||||
aten._test_parallel_materialize,
|
||||
aten.leaky_relu_,
|
||||
aten.leaky_relu_backward,
|
||||
aten.lerp,
|
||||
aten.lerp_,
|
||||
aten.linspace,
|
||||
aten.logaddexp,
|
||||
aten.logaddexp2,
|
||||
aten.logit,
|
||||
aten.logit_,
|
||||
aten.logit_backward,
|
||||
aten.log_sigmoid_backward,
|
||||
aten.log_sigmoid_forward,
|
||||
aten._log_softmax_backward_data,
|
||||
aten.logspace,
|
||||
aten.logsumexp.default,
|
||||
aten.masked_fill,
|
||||
aten.masked_fill_,
|
||||
aten.max_unpool2d,
|
||||
aten.max_unpool3d,
|
||||
aten.mish,
|
||||
aten.mish_,
|
||||
aten.mish_backward,
|
||||
aten.mse_loss,
|
||||
aten.mse_loss_backward,
|
||||
aten.multi_margin_loss,
|
||||
aten.multilabel_margin_loss_forward,
|
||||
aten.mv,
|
||||
aten.mvlgamma,
|
||||
aten.mvlgamma_,
|
||||
aten.nansum,
|
||||
aten.nan_to_num,
|
||||
aten.nan_to_num_,
|
||||
aten.narrow,
|
||||
aten.native_batch_norm_backward,
|
||||
aten.native_dropout_backward,
|
||||
aten.native_group_norm_backward,
|
||||
aten.native_layer_norm_backward,
|
||||
aten._fused_rms_norm,
|
||||
aten._fused_rms_norm_backward,
|
||||
aten.new_empty,
|
||||
aten.new_full,
|
||||
aten.new_ones,
|
||||
aten.new_zeros,
|
||||
aten.nll_loss2d_forward,
|
||||
aten.nll_loss2d_backward,
|
||||
aten.nll_loss_backward,
|
||||
aten.nll_loss_forward,
|
||||
aten.norm.ScalarOpt_dtype,
|
||||
aten.norm.Scalar,
|
||||
aten.norm.ScalarOpt_dim_dtype,
|
||||
aten.norm.ScalarOpt_dim,
|
||||
aten.norm.dtype_out,
|
||||
aten.norm.out,
|
||||
aten.norm.names_dtype_out,
|
||||
aten.norm.names_out,
|
||||
aten.norm.ScalarOpt_dtype_out,
|
||||
aten.norm.Scalar_out,
|
||||
aten.ones,
|
||||
aten.ones_like,
|
||||
aten.pixel_shuffle,
|
||||
aten.pixel_unshuffle,
|
||||
aten._prelu_kernel,
|
||||
aten._prelu_kernel_backward,
|
||||
aten._reshape_alias,
|
||||
aten.rad2deg,
|
||||
aten.rad2deg_,
|
||||
aten.reflection_pad1d,
|
||||
aten.reflection_pad1d_backward,
|
||||
aten.reflection_pad2d,
|
||||
aten.reflection_pad2d_backward,
|
||||
aten.reflection_pad3d,
|
||||
aten.reflection_pad3d_backward,
|
||||
aten.replication_pad1d,
|
||||
aten.replication_pad2d,
|
||||
aten.replication_pad3d,
|
||||
aten.renorm,
|
||||
aten.renorm_,
|
||||
aten.replication_pad2d,
|
||||
aten.resize_as,
|
||||
aten.roll,
|
||||
aten.rot90,
|
||||
aten.rrelu_with_noise,
|
||||
aten.rrelu_with_noise_,
|
||||
aten.rsub,
|
||||
aten._safe_softmax,
|
||||
aten._scaled_dot_product_flash_attention_for_cpu.default,
|
||||
aten.select_backward,
|
||||
aten.select_scatter,
|
||||
aten.sgn,
|
||||
aten.sgn_,
|
||||
aten.sigmoid_backward,
|
||||
aten.silu,
|
||||
aten.silu_,
|
||||
aten.silu_backward.grad_input,
|
||||
aten.silu_backward,
|
||||
aten.sinc,
|
||||
aten.sinc_,
|
||||
aten.slice_backward,
|
||||
aten.smooth_l1_loss,
|
||||
aten.smooth_l1_loss_backward,
|
||||
aten.soft_margin_loss,
|
||||
aten.soft_margin_loss_backward,
|
||||
aten._softmax_backward_data,
|
||||
aten.softplus,
|
||||
aten.softplus_backward,
|
||||
aten.softshrink,
|
||||
aten.special_entr,
|
||||
aten.special_log_ndtr,
|
||||
aten.special_xlog1py,
|
||||
aten.split.Tensor,
|
||||
aten.split_with_sizes_copy,
|
||||
aten.squeeze_copy,
|
||||
aten.squeeze.default,
|
||||
aten.squeeze.dim,
|
||||
aten.std.correction,
|
||||
aten.std.out,
|
||||
aten.std.correction_out,
|
||||
aten.std.names_out,
|
||||
aten.std.correction_names_out,
|
||||
aten.std_mean.correction,
|
||||
aten.std_mean.correction_out,
|
||||
aten.stack,
|
||||
aten.sum.default,
|
||||
aten.sum.out,
|
||||
aten.t,
|
||||
aten.t_copy,
|
||||
aten.take,
|
||||
aten.tanh_backward,
|
||||
aten.threshold,
|
||||
aten.threshold_,
|
||||
aten.threshold_backward,
|
||||
aten.trace,
|
||||
aten.transpose.int,
|
||||
aten.transpose_copy,
|
||||
aten.tril,
|
||||
aten.tril_,
|
||||
aten.triu,
|
||||
aten.triu_,
|
||||
aten.unbind,
|
||||
aten.unfold_backward,
|
||||
aten.unfold_copy,
|
||||
aten._unsafe_index,
|
||||
aten._unsafe_index_put,
|
||||
aten._unsafe_masked_index,
|
||||
aten._unsafe_masked_index_put_accumulate,
|
||||
aten.unsafe_split.Tensor,
|
||||
aten.unsafe_split_with_sizes,
|
||||
aten.unsqueeze_copy,
|
||||
aten._unsafe_view,
|
||||
aten.upsample_linear1d,
|
||||
aten.upsample_bilinear2d.out,
|
||||
aten.upsample_trilinear3d.out,
|
||||
aten.upsample_nearest2d_backward,
|
||||
aten.view_as_complex,
|
||||
aten.xlogy,
|
||||
aten.xlogy_,
|
||||
aten.zero,
|
||||
aten.zero_,
|
||||
aten.zeros,
|
||||
aten.zeros_like,
|
||||
aten._chunk_cat,
|
||||
aten._weight_norm_interface,
|
||||
]
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,344 @@
|
||||
# mypy: allow-untyped-decorators
|
||||
# mypy: allow-untyped-defs
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch._decomp
|
||||
from torch import Tensor
|
||||
from torch._prims_common.wrappers import _maybe_remove_out_wrapper
|
||||
|
||||
|
||||
decomposition_table = torch._decomp.decomposition_table
|
||||
decomposition_table_for_jvp: dict[torch._ops.OperatorBase, Callable] = {}
|
||||
register_decomposition = torch._decomp.register_decomposition
|
||||
aten = torch.ops.aten
|
||||
|
||||
# NOTE: [forward-mode AD decompositions mechanism]
|
||||
#
|
||||
# The mechanism is in VariableType,
|
||||
# IF any inputs have forward grad
|
||||
# AND there is no forward AD formula implemented
|
||||
# AND the functions are actually differentiable
|
||||
# run the decomposition
|
||||
# See run_jit_decomposition_with_args_for_jvp
|
||||
# We currently use python decompositions that we torchscript.
|
||||
#
|
||||
# Note that we would be building the backward graph at the decomposed level
|
||||
# too, but that is OK, because we would've errored out otherwise anyway.
|
||||
#
|
||||
# TODO: The mechanism we are using to register decompositions doesn't
|
||||
# seem to be exclusively used for jvp. So open question here is whether
|
||||
# torch/csrc/jit/runtime/decomposition_registry.cpp is being used for other things.
|
||||
# If that is the case, we may go down the decomposition path unexpectedly
|
||||
# (and possibly produce an unintelligible error) vs erroring out earlier and
|
||||
# printing that the forward AD formula is not implemented.
|
||||
#
|
||||
# The solution to this may be to have an explicitly white list control when
|
||||
# to enable the decomposition.
|
||||
|
||||
|
||||
def maybe_register_decomposition(op):
|
||||
def decorator(f):
|
||||
try:
|
||||
return register_decomposition(op)(f)
|
||||
except Exception:
|
||||
return f
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# Functions where we need a special decomposition for jvp but there's another version that
|
||||
# should be used more generally (ex. for jvp we need to recompute the mean and variance for
|
||||
# the backwards of a normalization function. Without jvp, it should use the saved value)
|
||||
decomposition_table_for_jvp = {}
|
||||
|
||||
|
||||
def register_decomposition_for_jvp(fn):
|
||||
return register_decomposition(fn, registry=decomposition_table_for_jvp)
|
||||
|
||||
|
||||
def _register_jit_decomposition_for_jvp(decomp, use_python=False):
|
||||
if decomp in decomposition_table_for_jvp:
|
||||
decomposition_table_used = decomposition_table_for_jvp
|
||||
elif decomp in decomposition_table:
|
||||
decomposition_table_used = decomposition_table
|
||||
else:
|
||||
raise RuntimeError(f"could not find decomposition for {decomp}")
|
||||
decomp_fn = decomposition_table_used[decomp]
|
||||
|
||||
# `out_wrapper` extends a decompositions signature with
|
||||
# an `out` parameter. However jit will use the unwrapped function's
|
||||
# signature instead so we need to unwrap here to prevent an error
|
||||
decomp_fn = _maybe_remove_out_wrapper(decomp_fn)
|
||||
|
||||
if use_python:
|
||||
decomp_fn = torch.jit.ignore(decomp_fn)
|
||||
sig = inspect.signature(decomp_fn)
|
||||
|
||||
# Create a string wrapping the function from the signature
|
||||
# example output:
|
||||
# def wrapped_decomp(x: torch.Tensor, y: int, z: int):
|
||||
# return decomp_fn(x, y, z)
|
||||
# Thanks copilot!
|
||||
def get_function_def(sig):
|
||||
param_def = [f"{param_str}" for param_str in sig.parameters.values()]
|
||||
param_use = [f"{param_str}" for param_str in sig.parameters]
|
||||
|
||||
return f"def wrapped_decomp({', '.join(param_def)}):\n return decomp_fn({', '.join(param_use)})\n"
|
||||
|
||||
f_str = get_function_def(sig)
|
||||
graph = torch.jit.CompilationUnit(f_str).wrapped_decomp.graph
|
||||
else:
|
||||
graph = torch.jit.script(decomp_fn).graph
|
||||
torch.jit._register_decomposition(decomp, graph)
|
||||
|
||||
|
||||
# The only decompositions here are temporary or hacks for the purposes of jvp
|
||||
|
||||
|
||||
# TODO: do these also belong here?
|
||||
@maybe_register_decomposition(aten.trace.default)
|
||||
def trace(self: Tensor) -> Tensor:
|
||||
return torch.sum(torch.diag(self))
|
||||
|
||||
|
||||
@maybe_register_decomposition(aten.log_sigmoid_forward.default)
|
||||
def log_sigmoid_forward(self: Tensor) -> tuple[Tensor, Tensor]:
|
||||
min = torch.minimum(self.new_zeros(()), self)
|
||||
z = torch.exp(-torch.abs(self))
|
||||
if self.is_cuda or self.is_xpu:
|
||||
buffer = self.new_zeros((0,))
|
||||
else:
|
||||
buffer = z
|
||||
return min - torch.log1p(z), buffer
|
||||
|
||||
|
||||
def recompute_mean_var(
|
||||
input: Tensor, rstd: Tensor, inner_dim_indices: list[int], keepdim: bool
|
||||
):
|
||||
# for most norm decompositions, it will be the same as the core version except for here.
|
||||
# We recompute the mean and variance so that they track gradients through input
|
||||
|
||||
mean = torch.mean(input, dim=inner_dim_indices, keepdim=keepdim)
|
||||
var = torch.var(input, dim=inner_dim_indices, unbiased=False, keepdim=keepdim)
|
||||
eps = torch.pow(1 / rstd, 2) - var # this makes me so sad inside
|
||||
eps = eps.detach()
|
||||
rstd = 1 / torch.sqrt(var + eps)
|
||||
return mean, rstd
|
||||
|
||||
|
||||
@register_decomposition_for_jvp(aten.native_layer_norm_backward)
|
||||
def native_layer_norm_backward(
|
||||
grad_out: Tensor,
|
||||
input: Tensor,
|
||||
normalized_shape: list[int],
|
||||
mean: Tensor,
|
||||
rstd: Tensor,
|
||||
weight: Tensor | None,
|
||||
bias: Tensor | None,
|
||||
output_mask: list[bool],
|
||||
) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
|
||||
input_shape = input.shape
|
||||
input_ndim = input.dim()
|
||||
|
||||
axis = input_ndim - len(normalized_shape)
|
||||
inner_dims = input_shape[axis:]
|
||||
outer_dims = input_shape[:axis]
|
||||
inner_dim_indices = list(range(axis, input_ndim))
|
||||
outer_dim_indices = list(range(axis))
|
||||
|
||||
N = 1
|
||||
for i in inner_dims:
|
||||
N *= i
|
||||
M = 1
|
||||
for i in outer_dims:
|
||||
M *= i
|
||||
if M <= 0 or N <= 0:
|
||||
return (
|
||||
input.new_zeros(input_shape),
|
||||
input.new_zeros(input_shape[axis:]),
|
||||
input.new_zeros(input_shape[axis:]),
|
||||
)
|
||||
|
||||
mean_, rstd_ = recompute_mean_var(input, rstd, inner_dim_indices, keepdim=True)
|
||||
|
||||
x_hat = (input - mean_) * rstd_
|
||||
if weight is not None:
|
||||
grad_x_hat = grad_out * weight
|
||||
else:
|
||||
grad_x_hat = grad_out
|
||||
a = grad_x_hat * N
|
||||
b = torch.sum(grad_x_hat, inner_dim_indices, True)
|
||||
c1 = torch.mul(grad_x_hat, x_hat)
|
||||
c2 = torch.sum(c1, inner_dim_indices, True)
|
||||
c3 = torch.mul(x_hat, c2)
|
||||
inner = a - b - c3
|
||||
|
||||
if output_mask[0]:
|
||||
d_input: Tensor | None = (rstd_ / N) * inner
|
||||
else:
|
||||
d_input = torch.zeros_like(input) # should be None but doesn't work with vjp
|
||||
|
||||
if output_mask[1] and weight is not None:
|
||||
if len(outer_dim_indices) > 0:
|
||||
d_weight: Tensor | None = torch.sum(
|
||||
grad_out * x_hat, outer_dim_indices, False
|
||||
)
|
||||
else:
|
||||
d_weight = grad_out * x_hat
|
||||
elif weight is not None:
|
||||
d_weight = torch.zeros_like(weight) # should be None but doesn't work with vjp
|
||||
else:
|
||||
d_weight = torch.zeros(()) # should be None but doesn't work with vjp
|
||||
|
||||
if output_mask[2] and bias is not None:
|
||||
if len(outer_dim_indices) > 0:
|
||||
d_bias: Tensor | None = torch.sum(grad_out, outer_dim_indices, False)
|
||||
else:
|
||||
d_bias = grad_out.clone()
|
||||
elif bias is not None:
|
||||
d_bias = torch.zeros_like(bias) # should be None but doesn't work with vjp
|
||||
else:
|
||||
d_bias = torch.zeros(()) # should be None but doesn't work with vjp
|
||||
|
||||
return (d_input, d_weight, d_bias)
|
||||
|
||||
|
||||
def prod(x: list[int]):
|
||||
r = 1
|
||||
for i in x:
|
||||
r *= i
|
||||
return r
|
||||
|
||||
|
||||
@register_decomposition_for_jvp(aten.native_batch_norm_backward)
|
||||
def native_batch_norm_backward(
|
||||
grad_out: Tensor,
|
||||
input: Tensor,
|
||||
weight: Tensor | None,
|
||||
running_mean: Tensor | None,
|
||||
running_var: Tensor | None,
|
||||
save_mean: Tensor | None,
|
||||
save_invstd: Tensor | None,
|
||||
train: bool,
|
||||
eps: float,
|
||||
output_mask: list[bool],
|
||||
) -> tuple[Tensor, Tensor | None, Tensor | None]:
|
||||
input_shape = input.shape
|
||||
input_rank = input.dim()
|
||||
if input_rank < 2:
|
||||
raise AssertionError(f"rank of the input must be at least 2, got {input_rank}")
|
||||
|
||||
axis = 1
|
||||
num_features = prod(input_shape) / input_shape[axis] # type: ignore[arg-type]
|
||||
mean = save_mean
|
||||
invstd = save_invstd
|
||||
if train:
|
||||
if save_mean is None or save_invstd is None:
|
||||
raise AssertionError(
|
||||
"when train=True, save_mean and save_invstd are required"
|
||||
)
|
||||
|
||||
reduciton_dims = [0] + list(range(2, input.dim()))
|
||||
if invstd is None:
|
||||
raise AssertionError("invstd must not be None for typing")
|
||||
mean, invstd = recompute_mean_var(input, invstd, reduciton_dims, keepdim=False)
|
||||
else:
|
||||
if running_mean is None or running_var is None:
|
||||
raise AssertionError(
|
||||
"running_mean and running_var must not be None when train=False"
|
||||
)
|
||||
mean = running_mean
|
||||
invstd = torch.rsqrt(running_var + eps)
|
||||
|
||||
if invstd is None or mean is None:
|
||||
raise AssertionError(
|
||||
f"invstd and mean must not be None, got invstd={invstd}, mean={mean}"
|
||||
)
|
||||
|
||||
broadcast_mask = [1] * input_rank
|
||||
broadcast_mask[axis] = input_shape[axis]
|
||||
|
||||
reduction_axes: list[int] = []
|
||||
for i in range(input_rank):
|
||||
if i != axis:
|
||||
reduction_axes.append(i)
|
||||
|
||||
mean = torch.reshape(mean, broadcast_mask)
|
||||
norm = 1.0 / num_features
|
||||
grad_output_sum = torch.sum(grad_out, reduction_axes)
|
||||
dot_p = torch.sum(grad_out * (input - mean), reduction_axes)
|
||||
|
||||
grad_mean = torch.reshape(grad_output_sum * norm, broadcast_mask)
|
||||
proj_scale = torch.reshape(torch.mul(dot_p * norm, invstd * invstd), broadcast_mask)
|
||||
|
||||
if weight is None:
|
||||
grad_scale = torch.reshape(invstd, broadcast_mask) * 1.0
|
||||
else:
|
||||
grad_scale = torch.reshape(invstd * weight, broadcast_mask)
|
||||
|
||||
if train:
|
||||
proj = (input - mean) * proj_scale
|
||||
grad_input = ((grad_out - proj) - grad_mean) * grad_scale
|
||||
else:
|
||||
grad_input = grad_out * grad_scale
|
||||
|
||||
if output_mask[1]:
|
||||
grad_weight = dot_p * invstd
|
||||
elif weight is not None:
|
||||
grad_weight = torch.zeros_like(
|
||||
weight
|
||||
) # should be None but doesn't work with vjp
|
||||
else:
|
||||
grad_weight = torch.zeros(()) # should be None but doesn't work with vjp
|
||||
|
||||
if output_mask[2]:
|
||||
grad_bias = grad_output_sum
|
||||
else:
|
||||
grad_bias = torch.zeros_like(
|
||||
grad_output_sum
|
||||
) # should be None but doesn't work with vjp
|
||||
|
||||
return (grad_input, grad_weight, grad_bias)
|
||||
|
||||
|
||||
@register_decomposition_for_jvp(aten.batch_norm_backward)
|
||||
def batch_norm_backward(
|
||||
grad_out: Tensor,
|
||||
input: Tensor,
|
||||
weight: Tensor,
|
||||
running_mean: Tensor | None,
|
||||
running_var: Tensor | None,
|
||||
save_mean: Tensor | None,
|
||||
save_var: Tensor | None,
|
||||
update: bool,
|
||||
eps: float,
|
||||
output_mask: list[bool],
|
||||
reserve: Tensor,
|
||||
) -> tuple[Tensor, Tensor | None, Tensor | None]:
|
||||
return native_batch_norm_backward(
|
||||
grad_out,
|
||||
input,
|
||||
weight,
|
||||
running_mean,
|
||||
running_var,
|
||||
save_mean,
|
||||
save_var,
|
||||
update,
|
||||
eps,
|
||||
output_mask,
|
||||
)
|
||||
|
||||
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.trace.default, use_python=True)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.nll_loss_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.nll_loss2d_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten._log_softmax_backward_data.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten._softmax_backward_data.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.log_sigmoid_forward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.native_layer_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.native_batch_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.cudnn_batch_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.batch_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.miopen_batch_norm_backward.default)
|
||||
@@ -0,0 +1,272 @@
|
||||
# mypy: allow-untyped-decorators
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch._decomp as decomp
|
||||
from torch._decomp import get_decompositions
|
||||
from torch._ops import OpOverload
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
rng_decompositions: dict[str, dict[OpOverload, Callable]] = defaultdict(dict)
|
||||
|
||||
|
||||
def register_rng_decomposition(aten_op):
|
||||
return decomp.register_decomposition(aten_op, rng_decompositions)
|
||||
|
||||
|
||||
def throw_on_non_cuda(device):
|
||||
raise RuntimeError(
|
||||
f"You are trying to functionalize a {device.type} RNG operator but {device.type} does not "
|
||||
f"use Philox/counter-based RNG. Therefore, functionalizing a {device.type} RNG operator is "
|
||||
"not supported. We are discussing the possibility of a Philox-based RNG implementation for CPU."
|
||||
)
|
||||
|
||||
|
||||
# TODO - We have to register many more distributions here, and also higher level
|
||||
# ops like dropout which have fused implementation and can hide the rand inside.
|
||||
@register_rng_decomposition(aten.rand)
|
||||
def rand(shape, dtype=None, layout=torch.strided, device=None, pin_memory=False):
|
||||
if device and device.type != "cuda":
|
||||
throw_on_non_cuda(device)
|
||||
seed, offset = PhiloxStateTracker.get_state_as_tuple()
|
||||
dtype = dtype or torch.float32
|
||||
out, offset_jump = torch.ops.rngprims.philox_rand(
|
||||
shape, seed, offset, None, device, dtype
|
||||
)
|
||||
PhiloxStateTracker.advance_offset(offset_jump)
|
||||
return out
|
||||
|
||||
|
||||
@register_rng_decomposition(aten.rand_like)
|
||||
def rand_like(
|
||||
x: torch.Tensor,
|
||||
dtype=None,
|
||||
layout=None,
|
||||
device=None,
|
||||
pin_memory=False,
|
||||
memory_format=torch.preserve_format,
|
||||
):
|
||||
device = device or x.device
|
||||
if device.type != "cuda":
|
||||
throw_on_non_cuda(device)
|
||||
dtype = dtype or x.dtype
|
||||
seed, offset = PhiloxStateTracker.get_state_as_tuple()
|
||||
out, offset_jump = torch.ops.rngprims.philox_rand(
|
||||
x.shape, seed, offset, None, device, dtype
|
||||
)
|
||||
PhiloxStateTracker.advance_offset(offset_jump)
|
||||
return out
|
||||
|
||||
|
||||
class PhiloxState:
|
||||
"""
|
||||
Represents a PhiloxRngState - (seed, offset) where offset = base_offset +
|
||||
relative_offset. seed and base_offset basically point to the rng state just
|
||||
before tracing starts. relative offset tracks the totally consumed offset at
|
||||
trace time.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.seed = torch.tensor(())
|
||||
self.base_offset = torch.tensor(())
|
||||
self.relative_offset = 0
|
||||
self.offset_advanced_at_least_once = False
|
||||
|
||||
def validate_state(self):
|
||||
if self.seed.numel() == 0 or self.base_offset.numel() == 0:
|
||||
raise AssertionError(
|
||||
f"seed and base_offset must not be empty, got "
|
||||
f"seed.numel()={self.seed.numel()}, base_offset.numel()={self.base_offset.numel()}"
|
||||
)
|
||||
|
||||
def advance_offset(self, consumed_offset):
|
||||
self.offset_advanced_at_least_once = True
|
||||
self.relative_offset = self.relative_offset + consumed_offset
|
||||
|
||||
def set_state(self, seed, base_offset, relative_offset=0):
|
||||
self.seed = seed
|
||||
self.base_offset = base_offset
|
||||
self.relative_offset = relative_offset
|
||||
|
||||
def get_state_as_tuple(self):
|
||||
self.validate_state()
|
||||
return (self.seed, self.base_offset + self.relative_offset)
|
||||
|
||||
def get_state_as_tensor(self):
|
||||
# Only needed because we override get_rng_state.
|
||||
self.validate_state()
|
||||
return torch.stack([self.seed, self.base_offset + self.relative_offset])
|
||||
|
||||
def set_state_from_tensor(self, state):
|
||||
# Only needed because we override set_rng_state.
|
||||
self.seed, self.base_offset = torch.unbind(state)
|
||||
self.relative_offset = 0
|
||||
|
||||
|
||||
class PhiloxStateTracker:
|
||||
"""
|
||||
Singleton class to track the philox rng state during AOT Autograd tracing.
|
||||
For each aot tracing instance, AOT Autograd resets this tracker and keeps
|
||||
track of both forward and backward offsets. At runtime, we only care about
|
||||
the total consumed forward and backward offsets. For dynamic shapes, these
|
||||
offsets are a function of input shapes. Therefore, the AOT generated graphs
|
||||
have additional outputs that compute total consumed forward and backward
|
||||
offsets.
|
||||
"""
|
||||
|
||||
running_state: PhiloxState
|
||||
fwd_state: PhiloxState
|
||||
bwd_state: PhiloxState
|
||||
|
||||
def __enter__(self):
|
||||
PhiloxStateTracker.reset()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_cal, exc_tb):
|
||||
PhiloxStateTracker.reset()
|
||||
|
||||
@classmethod
|
||||
def reset(cls):
|
||||
cls.running_state = PhiloxState()
|
||||
cls.fwd_state = PhiloxState()
|
||||
cls.bwd_state = PhiloxState()
|
||||
|
||||
@classmethod
|
||||
def mark_beginning_of_forward(cls):
|
||||
# Tells the tracker to use fwd_state as the running state
|
||||
cls.running_state = cls.fwd_state
|
||||
|
||||
@classmethod
|
||||
def mark_beginning_of_backward(cls):
|
||||
# Tells the tracker to use bwd_state as the running state
|
||||
cls.running_state = cls.bwd_state
|
||||
|
||||
@classmethod
|
||||
def record_state(cls, seed, offset, mode):
|
||||
# Records the seed and offset tensors. These tensors are used to invoke
|
||||
# the philox_rand functional primitives.
|
||||
if mode == "forward":
|
||||
cls.fwd_state.set_state(seed, offset)
|
||||
cls.mark_beginning_of_forward()
|
||||
else:
|
||||
if mode != "backward":
|
||||
raise AssertionError(f"mode must be 'backward', got {mode}")
|
||||
cls.bwd_state.set_state(seed, offset)
|
||||
|
||||
@classmethod
|
||||
def get_state_as_tensor(cls):
|
||||
# The only reason this exists is because we override get_rng_state and
|
||||
# set_rng_state during tracing. get_rng_state expects a tensor output,
|
||||
# so return (seed, offset) tuple upset other parts of the program like
|
||||
# ctx.saved_tensors.
|
||||
|
||||
# A bad consequence is that if user saves and restores rng state, we
|
||||
# have little bit of ugliness in the generated code, where we first
|
||||
# concat the (seed, offset) to create a tensor for get_rng_state, and
|
||||
# then split it back to get (seed, offset) tuple in set_rng_state.
|
||||
|
||||
# TODO: Investigate if there is be a better way to wrap the tuple in a
|
||||
# false Tensor object, and then desugar it later on.
|
||||
return cls.running_state.get_state_as_tensor()
|
||||
|
||||
@classmethod
|
||||
def get_state_as_tuple(cls):
|
||||
return cls.running_state.get_state_as_tuple()
|
||||
|
||||
@classmethod
|
||||
def set_state_from_tensor(cls, x):
|
||||
# This is only needed because we override set_rng_state. Look at the
|
||||
# comment in get_state_from_tensor method.
|
||||
cls.running_state.set_state_from_tensor(x)
|
||||
|
||||
@classmethod
|
||||
def advance_offset(cls, consumed_offset):
|
||||
cls.running_state.advance_offset(consumed_offset)
|
||||
|
||||
@classmethod
|
||||
def get_current_relative_offset(cls):
|
||||
return cls.running_state.relative_offset
|
||||
|
||||
@staticmethod
|
||||
def multiple_of_4(offset):
|
||||
# torch cuda rng state offset must be a multiple of 4. For inductor, as
|
||||
# we sum up all the numel, the result might not be a multiple of 4. This
|
||||
# method achieves that.
|
||||
return (offset + 3) // 4 * 4
|
||||
|
||||
@classmethod
|
||||
def get_updated_fwd_offset(cls):
|
||||
# Short circuit if no rand ops were observed
|
||||
if not cls.fwd_state.offset_advanced_at_least_once:
|
||||
return cls.fwd_state.base_offset
|
||||
return cls.multiple_of_4(
|
||||
cls.fwd_state.base_offset + cls.fwd_state.relative_offset
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_updated_bwd_offset(cls):
|
||||
# Short circuit if no rand ops were observed
|
||||
if not cls.bwd_state.offset_advanced_at_least_once:
|
||||
return cls.bwd_state.base_offset
|
||||
return cls.multiple_of_4(
|
||||
cls.bwd_state.base_offset + cls.bwd_state.relative_offset
|
||||
)
|
||||
|
||||
|
||||
# Adding more decompositions which eventually use rand_like inside decomps.
|
||||
# Adding these in rng_decompositions ensures the functionalization of rand_like
|
||||
# ops used in these decomps. The list is copied from inductor codebase, which
|
||||
# uses it for similar purpose.
|
||||
#
|
||||
# Caution - These decomps do not have same accuracy as that of eager. However,
|
||||
# we can't just disable them with a config flag like fallback_random, because
|
||||
# for functionalization of rng ops, we have to decompose these ops.
|
||||
extra_random_decomps = get_decompositions(
|
||||
[
|
||||
aten.cauchy,
|
||||
aten.cauchy_,
|
||||
aten.exponential,
|
||||
aten.exponential_,
|
||||
aten.geometric,
|
||||
aten.geometric_,
|
||||
aten.native_dropout,
|
||||
aten.normal,
|
||||
aten.normal_,
|
||||
aten.normal_functional,
|
||||
aten.log_normal,
|
||||
aten.log_normal_,
|
||||
aten.rrelu_with_noise,
|
||||
aten.rrelu_with_noise_,
|
||||
aten.uniform_,
|
||||
]
|
||||
)
|
||||
register_extra_random_decomp = functools.partial(
|
||||
decomp.register_decomposition, registry=extra_random_decomps
|
||||
)
|
||||
|
||||
|
||||
@register_extra_random_decomp([aten.bernoulli_])
|
||||
def bernoulli_(self, p=0.5):
|
||||
if self.device == torch.device("cpu"):
|
||||
return NotImplemented
|
||||
return self.copy_(torch.rand_like(self, dtype=torch.float32) < p)
|
||||
|
||||
|
||||
@register_extra_random_decomp([aten.bernoulli.p])
|
||||
def bernoulli_p(self, p=0.5, *, generator=None):
|
||||
if self.device == torch.device("cpu"):
|
||||
return NotImplemented
|
||||
if generator is not None:
|
||||
raise AssertionError(f"generator must be None, got {generator}")
|
||||
return torch.rand_like(self, dtype=torch.float32) < p
|
||||
|
||||
|
||||
rng_decompositions.update(extra_random_decomps) # type: ignore[arg-type]
|
||||
@@ -0,0 +1,204 @@
|
||||
import itertools
|
||||
import unittest.mock
|
||||
from collections.abc import Callable, Generator, Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import torch
|
||||
import torch._C
|
||||
import torch._ops
|
||||
import torch.utils._python_dispatch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._C import DispatchKey
|
||||
|
||||
|
||||
__all__ = ["enable_python_dispatcher", "no_python_dispatcher", "enable_pre_dispatch"]
|
||||
|
||||
no_python_dispatcher = torch._C._DisablePythonDispatcher
|
||||
enable_python_dispatcher = torch._C._EnablePythonDispatcher
|
||||
enable_pre_dispatch = torch._C._EnablePreDispatch
|
||||
|
||||
CROSSREF_FUNCTIONALIZE = False
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_T = TypeVar("_T")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
|
||||
def all_py_loaded_overloads() -> Iterator[torch._ops.OpOverload]:
|
||||
"""
|
||||
Warning: the set of overloads this will report is very subtle. It is precisely
|
||||
the set of torch.ops functions that have actually been accessed from Python
|
||||
(e.g., we actually called torch.ops.aten.blah at some point. This is DIFFERENT
|
||||
from the set of registered operators, which will in general be a larger set,
|
||||
as this would include all operators which we ran C++ static initializers or
|
||||
Python operator registration on. This does not eagerly populate the list on
|
||||
torch.ops.aten; this list is lazy!
|
||||
|
||||
In other words, this is good for traversing over everything that has an
|
||||
OpOverload object allocated in Python. We use it for cache invalidation, but
|
||||
don't rely on this list being complete.
|
||||
|
||||
Note that even if we did report all C++ registered overloads, this isn't guaranteed
|
||||
to be complete either, as a subsequent lazy load of a library which triggers more
|
||||
registrations could add more things to the set.
|
||||
"""
|
||||
for ns in torch.ops:
|
||||
packets = getattr(torch.ops, ns)
|
||||
for op_name in packets:
|
||||
packet = getattr(packets, op_name)
|
||||
for overload in packet:
|
||||
yield getattr(packet, overload)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def suspend_functionalization() -> Generator[None, None, None]:
|
||||
f_tls = torch._C._dispatch_tls_is_dispatch_key_included(
|
||||
torch._C.DispatchKey.Functionalize
|
||||
)
|
||||
f_rv = torch._C._functionalization_reapply_views_tls()
|
||||
if f_tls:
|
||||
torch._disable_functionalization()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if f_tls:
|
||||
torch._enable_functionalization(reapply_views=f_rv)
|
||||
|
||||
|
||||
def check_tensor_metadata_matches(
|
||||
nv: torch.Tensor, rv: torch.Tensor, desc: Callable[[], str]
|
||||
) -> None:
|
||||
if not callable(desc):
|
||||
raise AssertionError(f"desc must be callable, got {type(desc)}")
|
||||
if nv.size() != rv.size():
|
||||
raise AssertionError(f"{desc()}: sizes {nv.size()} != {rv.size()}")
|
||||
if nv.dtype != rv.dtype:
|
||||
raise AssertionError(f"{desc()}: dtype {nv.dtype} != {rv.dtype}")
|
||||
same_strides, idx = torch._prims_common.check_significant_strides(
|
||||
nv, rv, only_cuda=False
|
||||
)
|
||||
if not same_strides:
|
||||
raise AssertionError(
|
||||
f"{desc()}: strides {nv.stride()} != {rv.stride()} (mismatch at index {idx})"
|
||||
)
|
||||
|
||||
|
||||
def check_metadata_matches(n: object, r: object, desc: Callable[[], str]) -> None:
|
||||
if not callable(desc):
|
||||
raise AssertionError(f"desc must be callable, got {type(desc)}")
|
||||
n_vals, _n_spec = pytree.tree_flatten(n)
|
||||
r_vals, _r_spec = pytree.tree_flatten(r)
|
||||
# TODO: test the specs match; empirically sometimes we have a tuple
|
||||
# on one side and a list on the other
|
||||
if len(n_vals) != len(r_vals):
|
||||
raise AssertionError(f"{len(n_vals)} != {len(r_vals)}")
|
||||
for i, nv, rv in zip(range(len(n_vals)), n_vals, r_vals):
|
||||
if not isinstance(rv, torch.Tensor):
|
||||
continue
|
||||
check_tensor_metadata_matches(nv, rv, lambda: f"{desc()} output {i}")
|
||||
|
||||
|
||||
class Lit:
|
||||
def __init__(self, s: str) -> None:
|
||||
self.s = s
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.s
|
||||
|
||||
|
||||
def _fmt(a: object) -> object:
|
||||
if isinstance(a, torch.Tensor):
|
||||
return Lit(
|
||||
f"torch.empty_strided({tuple(a.size())}, {a.stride()}, dtype={a.dtype})"
|
||||
)
|
||||
else:
|
||||
return a
|
||||
|
||||
|
||||
def make_crossref_functionalize(
|
||||
op: torch._ops.OpOverload[_P, _T], final_key: DispatchKey
|
||||
) -> Callable[_P, _T] | DispatchKey:
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
# This case is pretty weird, suppress it for now
|
||||
if op is torch.ops.aten.lift_fresh.default:
|
||||
return final_key
|
||||
|
||||
def handler(*args: _P.args, **kwargs: _P.kwargs) -> _T:
|
||||
fake_mode = FakeTensorMode()
|
||||
|
||||
def fakeify_defun(t: _R) -> _R | torch._subclasses.fake_tensor.FakeTensor:
|
||||
if isinstance(t, torch.Tensor):
|
||||
if torch._is_functional_tensor(t):
|
||||
r = torch._from_functional_tensor(t)
|
||||
# NB: This assumes that the inner tensor sizes/strides match
|
||||
# the outer tensor sizes/strides. This doesn't necessarily have to
|
||||
# be the case, see discussion at
|
||||
# https://github.com/pytorch/pytorch/pull/87610/files/401ddeda1d769bedc88a12de332c7357b60e51a4#r1007264456
|
||||
if t.size() != r.size():
|
||||
raise AssertionError(f"size mismatch: {t.size()} != {r.size()}")
|
||||
if t.stride() != r.stride():
|
||||
raise AssertionError(
|
||||
f"stride mismatch: {t.stride()} != {r.stride()}"
|
||||
)
|
||||
else:
|
||||
r = t
|
||||
# TODO: suppress guards
|
||||
return fake_mode.from_tensor(r)
|
||||
return t
|
||||
|
||||
def maybe_detach(t: _R) -> _R | torch.Tensor:
|
||||
if isinstance(t, torch.Tensor):
|
||||
return t.detach()
|
||||
else:
|
||||
return t
|
||||
|
||||
# TODO: This probably does the wrong thing if you're running other
|
||||
# substantive modes with the normal op outside here
|
||||
with (
|
||||
torch.utils._python_dispatch._disable_current_modes(),
|
||||
suspend_functionalization(),
|
||||
):
|
||||
f_args, f_kwargs = pytree.tree_map(fakeify_defun, (args, kwargs))
|
||||
orig_f_args, orig_f_kwargs = pytree.tree_map(
|
||||
maybe_detach, (f_args, f_kwargs)
|
||||
)
|
||||
with fake_mode:
|
||||
f_r = op(*f_args, **f_kwargs) # pyrefly: ignore [invalid-param-spec]
|
||||
r = op._op_dk(final_key, *args, **kwargs)
|
||||
|
||||
def desc() -> str:
|
||||
fmt_args = ", ".join(
|
||||
itertools.chain(
|
||||
(repr(pytree.tree_map(_fmt, a)) for a in orig_f_args),
|
||||
(
|
||||
f"{k}={pytree.tree_map(_fmt, v)}"
|
||||
for k, v in orig_f_kwargs.items()
|
||||
),
|
||||
)
|
||||
)
|
||||
return f"{op}({fmt_args})"
|
||||
|
||||
check_metadata_matches(f_r, r, desc)
|
||||
return r
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
# NB: enabling this is slow, don't do it in a hot loop. This is purely
|
||||
# for debugging purposes.
|
||||
@contextmanager
|
||||
def enable_crossref_functionalize() -> Generator[None, None, None]:
|
||||
for op in all_py_loaded_overloads():
|
||||
op._uncache_dispatch(torch._C.DispatchKey.Functionalize)
|
||||
try:
|
||||
with (
|
||||
enable_python_dispatcher(),
|
||||
unittest.mock.patch("torch._dispatch.python.CROSSREF_FUNCTIONALIZE", True),
|
||||
):
|
||||
yield
|
||||
finally:
|
||||
for op in all_py_loaded_overloads():
|
||||
op._uncache_dispatch(torch._C.DispatchKey.Functionalize)
|
||||
@@ -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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user