Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,107 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import copyreg
|
||||
import os.path as _osp
|
||||
import weakref
|
||||
|
||||
import torch
|
||||
from torch.utils import (
|
||||
backcompat as backcompat,
|
||||
collect_env as collect_env,
|
||||
data as data,
|
||||
deterministic as deterministic,
|
||||
hooks as hooks,
|
||||
)
|
||||
from torch.utils.backend_registration import (
|
||||
generate_methods_for_privateuse1_backend,
|
||||
rename_privateuse1_backend,
|
||||
)
|
||||
from torch.utils.cpp_backtrace import get_cpp_backtrace
|
||||
from torch.utils.throughput_benchmark import ThroughputBenchmark
|
||||
|
||||
|
||||
def set_module(obj, mod):
|
||||
"""
|
||||
Set the module attribute on a python object for a given object for nicer printing
|
||||
"""
|
||||
if not isinstance(mod, str):
|
||||
raise TypeError("The mod argument should be a string")
|
||||
obj.__module__ = mod
|
||||
|
||||
|
||||
cmake_prefix_path = _osp.join(_osp.dirname(_osp.dirname(__file__)), "share", "cmake")
|
||||
|
||||
|
||||
def swap_tensors(t1, t2):
|
||||
"""
|
||||
This function swaps the content of the two Tensor objects.
|
||||
At a high level, this will make t1 have the content of t2 while preserving
|
||||
its identity.
|
||||
|
||||
This will not work if t1 and t2 have different slots.
|
||||
"""
|
||||
# Ensure there are no weakrefs
|
||||
if weakref.getweakrefs(t1):
|
||||
raise RuntimeError("Cannot swap t1 because it has weakref associated with it")
|
||||
if weakref.getweakrefs(t2):
|
||||
raise RuntimeError("Cannot swap t2 because it has weakref associated with it")
|
||||
t1_slots = set(copyreg._slotnames(t1.__class__)) # type: ignore[attr-defined]
|
||||
t2_slots = set(copyreg._slotnames(t2.__class__)) # type: ignore[attr-defined]
|
||||
if t1_slots != t2_slots:
|
||||
raise RuntimeError("Cannot swap t1 and t2 if they have different slots")
|
||||
|
||||
def swap_attr(name):
|
||||
tmp = getattr(t1, name)
|
||||
setattr(t1, name, (getattr(t2, name)))
|
||||
setattr(t2, name, tmp)
|
||||
|
||||
def error_pre_hook(grad_outputs):
|
||||
raise RuntimeError(
|
||||
"Trying to execute AccumulateGrad node that was poisoned by swap_tensors "
|
||||
"this can happen when you try to run backward on a tensor that was swapped. "
|
||||
"For a module m with `torch.__future__.set_swap_module_params_on_conversion(True)` "
|
||||
"you should not change the device or dtype of the module (e.g. `m.cpu()` or `m.half()`) "
|
||||
"between running forward and backward. To resolve this, please only change the "
|
||||
"device/dtype before running forward (or after both forward and backward)."
|
||||
)
|
||||
|
||||
def check_use_count(t, name="t1"):
|
||||
use_count = t._use_count()
|
||||
error_str = (
|
||||
f"Expected use_count of {name} to be 1 or 2 with an AccumulateGrad node but got {use_count} "
|
||||
f"make sure you are not holding references to the tensor in other places."
|
||||
)
|
||||
if use_count > 1:
|
||||
if use_count == 2 and t.is_leaf:
|
||||
accum_grad_node = torch.autograd.graph.get_gradient_edge(t).node
|
||||
# Make sure that the accumulate_grad node was not lazy_init-ed by get_gradient_edge
|
||||
if t._use_count() == 2:
|
||||
accum_grad_node.register_prehook(error_pre_hook)
|
||||
else:
|
||||
raise RuntimeError(error_str)
|
||||
else:
|
||||
raise RuntimeError(error_str)
|
||||
|
||||
check_use_count(t1, "t1")
|
||||
check_use_count(t2, "t2")
|
||||
|
||||
# Swap the types
|
||||
# Note that this will fail if there are mismatched slots
|
||||
swap_attr("__class__")
|
||||
|
||||
# Swap the dynamic attributes
|
||||
swap_attr("__dict__")
|
||||
|
||||
# Swap the slots
|
||||
for slot in t1_slots:
|
||||
if hasattr(t1, slot) and hasattr(t2, slot):
|
||||
swap_attr(slot)
|
||||
elif hasattr(t1, slot):
|
||||
setattr(t2, slot, (getattr(t1, slot)))
|
||||
delattr(t1, slot)
|
||||
elif hasattr(t2, slot):
|
||||
setattr(t1, slot, (getattr(t2, slot)))
|
||||
delattr(t2, slot)
|
||||
|
||||
# Swap the at::Tensor they point to
|
||||
torch._C._swap_tensor_impl(t1, t2)
|
||||
@@ -0,0 +1,135 @@
|
||||
import base64
|
||||
import zlib
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
_ENCODING_VERSION: int = 1
|
||||
|
||||
__all__ = ["AppendingByteSerializer"]
|
||||
|
||||
|
||||
#######################################
|
||||
# Helper classes
|
||||
#######################################
|
||||
|
||||
CHECKSUM_DIGEST_SIZE = 4
|
||||
|
||||
|
||||
class BytesWriter:
|
||||
def __init__(self) -> None:
|
||||
# Reserve CHECKSUM_DIGEST_SIZE bytes for checksum
|
||||
self._data = bytearray(CHECKSUM_DIGEST_SIZE)
|
||||
|
||||
def write_uint64(self, i: int) -> None:
|
||||
self._data.extend(i.to_bytes(8, byteorder="big", signed=False))
|
||||
|
||||
def write_str(self, s: str) -> None:
|
||||
payload = base64.b64encode(s.encode("utf-8"))
|
||||
self.write_bytes(payload)
|
||||
|
||||
def write_bytes(self, b: bytes) -> None:
|
||||
self.write_uint64(len(b))
|
||||
self._data.extend(b)
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
digest = zlib.crc32(self._data[CHECKSUM_DIGEST_SIZE:]).to_bytes(
|
||||
4, byteorder="big", signed=False
|
||||
)
|
||||
if len(digest) != CHECKSUM_DIGEST_SIZE:
|
||||
raise AssertionError("Computed checksum digest has unexpected size")
|
||||
self._data[0:CHECKSUM_DIGEST_SIZE] = digest
|
||||
return bytes(self._data)
|
||||
|
||||
|
||||
class BytesReader:
|
||||
def __init__(self, data: bytes) -> None:
|
||||
# Check for data corruption
|
||||
if len(data) < CHECKSUM_DIGEST_SIZE:
|
||||
raise AssertionError("Input data is too short to contain checksum")
|
||||
digest = zlib.crc32(data[CHECKSUM_DIGEST_SIZE:]).to_bytes(
|
||||
4, byteorder="big", signed=False
|
||||
)
|
||||
if len(digest) != CHECKSUM_DIGEST_SIZE:
|
||||
raise AssertionError("Computed checksum digest has unexpected size")
|
||||
if data[0:CHECKSUM_DIGEST_SIZE] != digest:
|
||||
raise RuntimeError(
|
||||
"Bytes object is corrupted, checksum does not match. "
|
||||
f"Expected: {data[0:CHECKSUM_DIGEST_SIZE]!r}, Got: {digest!r}"
|
||||
)
|
||||
|
||||
self._data = data
|
||||
self._i = CHECKSUM_DIGEST_SIZE
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return len(self._data) == self._i
|
||||
|
||||
def read_uint64(self) -> int:
|
||||
result = int.from_bytes(
|
||||
self._data[self._i : self._i + 8], byteorder="big", signed=False
|
||||
)
|
||||
self._i += 8
|
||||
return result
|
||||
|
||||
def read_str(self) -> str:
|
||||
return base64.b64decode(self.read_bytes()).decode("utf-8")
|
||||
|
||||
def read_bytes(self) -> bytes:
|
||||
size = self.read_uint64()
|
||||
result = self._data[self._i : self._i + size]
|
||||
self._i += size
|
||||
return result
|
||||
|
||||
|
||||
#######################################
|
||||
# AppendingByteSerializer
|
||||
#######################################
|
||||
|
||||
|
||||
class AppendingByteSerializer(Generic[T]):
|
||||
"""
|
||||
Provides efficient serialization and deserialization of list of bytes
|
||||
Note that this does not provide any guarantees around byte order
|
||||
"""
|
||||
|
||||
_serialize_fn: Callable[[BytesWriter, T], None]
|
||||
_writer: BytesWriter
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
serialize_fn: Callable[[BytesWriter, T], None],
|
||||
) -> None:
|
||||
self._serialize_fn = serialize_fn
|
||||
self.clear()
|
||||
|
||||
def clear(self) -> None:
|
||||
self._writer = BytesWriter()
|
||||
# First 8-bytes are for version
|
||||
self._writer.write_uint64(_ENCODING_VERSION)
|
||||
|
||||
def append(self, data: T) -> None:
|
||||
self._serialize_fn(self._writer, data)
|
||||
|
||||
def extend(self, elems: Iterable[T]) -> None:
|
||||
for elem in elems:
|
||||
self.append(elem)
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
return self._writer.to_bytes()
|
||||
|
||||
@staticmethod
|
||||
def to_list(data: bytes, *, deserialize_fn: Callable[[BytesReader], T]) -> list[T]:
|
||||
reader = BytesReader(data)
|
||||
if reader.read_uint64() != _ENCODING_VERSION:
|
||||
raise AssertionError(
|
||||
f"Encoding version mismatch in AppendingByteSerializer.to_list, \
|
||||
got {reader.read_uint64()}"
|
||||
)
|
||||
|
||||
result: list[T] = []
|
||||
while not reader.is_finished():
|
||||
result.append(deserialize_fn(reader))
|
||||
return result
|
||||
@@ -0,0 +1,913 @@
|
||||
import contextlib
|
||||
import copy
|
||||
import hashlib
|
||||
import importlib
|
||||
import inspect
|
||||
import io
|
||||
import os
|
||||
import pickle
|
||||
import tokenize
|
||||
import unittest
|
||||
from collections.abc import Callable
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from types import FunctionType, ModuleType
|
||||
from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar
|
||||
from typing_extensions import deprecated
|
||||
|
||||
from torch._utils_internal import justknobs_check
|
||||
|
||||
|
||||
# Types saved/loaded in configs
|
||||
CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict)
|
||||
|
||||
# Immutable scalar types that don't need deepcopy when returned from configs.
|
||||
# Everything else is defensively copied to prevent accidental mutation.
|
||||
_IMMUTABLE_CONFIG_TYPES = (int, float, bool, type(None), str, tuple)
|
||||
|
||||
|
||||
# Duplicated, because mypy needs these types statically
|
||||
T = TypeVar("T", bound=int | float | bool | str | list | set | tuple | dict | None)
|
||||
|
||||
|
||||
_UNSET_SENTINEL = object()
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class _Config(Generic[T]):
|
||||
"""Represents a config with richer behaviour than just a default value.
|
||||
::
|
||||
i.e.
|
||||
foo = Config(justknob="//foo:bar", default=False)
|
||||
install_config_module(...)
|
||||
|
||||
This configs must be installed with install_config_module to be used
|
||||
|
||||
Precedence Order:
|
||||
alias: If set, the directly use the value of the alias.
|
||||
env_name_force: If set, this environment variable has precedence over
|
||||
everything after this.
|
||||
If multiple env variables are given, the precedence order is from
|
||||
left to right.
|
||||
user_override: If a user sets a value (i.e. foo.bar=True), that
|
||||
has precedence over everything after this. User overrides are thread-local.
|
||||
env_name_default: If set, this environment variable will override everything
|
||||
after this.
|
||||
If multiple env variables are given, the precedence order is from
|
||||
left to right.
|
||||
justknob: If this pytorch installation supports justknobs, that will
|
||||
override defaults, but will not override the user_override precedence.
|
||||
default: This value is the lowest precedence, and will be used if nothing is
|
||||
set.
|
||||
|
||||
Environment Variables:
|
||||
These are interpreted to be either "0" or "1" to represent true and false.
|
||||
|
||||
Arguments:
|
||||
justknob: the name of the feature / JK. In OSS this is unused.
|
||||
default: is the value to default this knob to in OSS.
|
||||
alias: The alias config to read instead.
|
||||
env_name_force: The environment variable, or list of, to read that is a FORCE
|
||||
environment variable. I.e. it overrides everything except for alias.
|
||||
env_name_default: The environment variable, or list of, to read that changes the
|
||||
default behaviour. I.e. user overrides take preference.
|
||||
"""
|
||||
|
||||
default: T | object
|
||||
justknob: str | None = None
|
||||
env_name_default: list[str] | None = None
|
||||
env_name_force: list[str] | None = None
|
||||
value_type: type | None = None
|
||||
alias: str | None = None
|
||||
# Deprecation support
|
||||
deprecated: bool = False
|
||||
deprecation_message: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.env_name_default = _Config.string_or_list_of_string_to_list(
|
||||
self.env_name_default
|
||||
)
|
||||
self.env_name_force = _Config.string_or_list_of_string_to_list(
|
||||
self.env_name_force
|
||||
)
|
||||
|
||||
if self.alias is not None:
|
||||
if (
|
||||
self.default is not _UNSET_SENTINEL
|
||||
or self.justknob is not None
|
||||
or self.env_name_default is not None
|
||||
or self.env_name_force is not None
|
||||
):
|
||||
raise AssertionError(
|
||||
"if alias is set, none of {default, justknob, \
|
||||
env_name_default and env_name_force} can be set"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def string_or_list_of_string_to_list(
|
||||
val: str | list[str] | None,
|
||||
) -> list[str] | None:
|
||||
if val is None:
|
||||
return None
|
||||
if isinstance(val, str):
|
||||
return [val]
|
||||
if not isinstance(val, list):
|
||||
raise AssertionError(f"val is not a list, got {type(val)}")
|
||||
return val
|
||||
|
||||
|
||||
# In runtime, we unbox the Config[T] to a T, but typechecker cannot see this,
|
||||
# so in order to allow for this dynamic behavior to work correctly with
|
||||
# typechecking we are going to lie to the typechecker that Config[T] returns
|
||||
# a T.
|
||||
if TYPE_CHECKING:
|
||||
|
||||
def Config(
|
||||
default: T | object = _UNSET_SENTINEL,
|
||||
justknob: str | None = None,
|
||||
env_name_default: str | list[str] | None = None,
|
||||
env_name_force: str | list[str] | None = None,
|
||||
value_type: type | None = None,
|
||||
alias: str | None = None,
|
||||
# Deprecation support
|
||||
deprecated: bool = False,
|
||||
deprecation_message: str | None = None,
|
||||
) -> T: ...
|
||||
|
||||
else:
|
||||
|
||||
def Config(
|
||||
default: T | object = _UNSET_SENTINEL,
|
||||
justknob: str | None = None,
|
||||
env_name_default: str | list[str] | None = None,
|
||||
env_name_force: str | list[str] | None = None,
|
||||
value_type: type | None = None,
|
||||
alias: str | None = None,
|
||||
# Deprecation support
|
||||
deprecated: bool = False,
|
||||
deprecation_message: str | None = None,
|
||||
) -> _Config[T]:
|
||||
return _Config(
|
||||
default=default,
|
||||
justknob=justknob,
|
||||
env_name_default=env_name_default,
|
||||
env_name_force=env_name_force,
|
||||
value_type=value_type,
|
||||
alias=alias,
|
||||
# Deprecation support
|
||||
deprecated=deprecated,
|
||||
deprecation_message=deprecation_message,
|
||||
)
|
||||
|
||||
|
||||
def _read_env_variable(name: str) -> bool | str | None:
|
||||
value = os.environ.get(name)
|
||||
if value == "1":
|
||||
return True
|
||||
if value == "0":
|
||||
return False
|
||||
return value
|
||||
|
||||
|
||||
def install_config_module(module: ModuleType) -> None:
|
||||
"""
|
||||
Converts a module-level config into a `ConfigModule()`.
|
||||
|
||||
See _config_typing.pyi for instructions on how to get the converted module to typecheck.
|
||||
"""
|
||||
|
||||
class ConfigModuleInstance(ConfigModule):
|
||||
# __annotations__ is written to by Sphinx autodoc
|
||||
_bypass_keys = set({"_is_dirty", "_hash_digest", "__annotations__"})
|
||||
|
||||
def visit(
|
||||
source: ModuleType | type,
|
||||
dest: ModuleType | SubConfigProxy,
|
||||
prefix: str,
|
||||
) -> None:
|
||||
"""Walk the module structure and move everything to module._config"""
|
||||
type_hints = inspect.get_annotations(source)
|
||||
for key, value in list(source.__dict__.items()):
|
||||
if (
|
||||
key.startswith("__")
|
||||
or isinstance(value, (ModuleType, FunctionType))
|
||||
or (
|
||||
hasattr(value, "__module__")
|
||||
and (
|
||||
value.__module__ == "typing"
|
||||
or value.__module__.startswith("collections.abc")
|
||||
)
|
||||
)
|
||||
# Handle from torch.utils._config_module import Config
|
||||
or (isinstance(value, type) and issubclass(value, _Config))
|
||||
):
|
||||
continue
|
||||
|
||||
name = f"{prefix}{key}"
|
||||
annotated_type = type_hints.get(key, None)
|
||||
if isinstance(value, CONFIG_TYPES):
|
||||
config[name] = _ConfigEntry(
|
||||
_Config(default=value, value_type=annotated_type), name
|
||||
)
|
||||
if dest is module:
|
||||
delattr(module, key)
|
||||
elif isinstance(value, _Config):
|
||||
if annotated_type is not None and value.value_type is None:
|
||||
value.value_type = annotated_type
|
||||
|
||||
config[name] = _ConfigEntry(value, name)
|
||||
|
||||
if dest is module:
|
||||
delattr(module, key)
|
||||
elif isinstance(value, type):
|
||||
if value.__module__ != module.__name__:
|
||||
raise AssertionError(
|
||||
f"subconfig class {value} must be defined in module {module.__name__}"
|
||||
)
|
||||
# a subconfig with `class Blah:` syntax
|
||||
proxy = SubConfigProxy(module, f"{name}.")
|
||||
visit(value, proxy, f"{name}.")
|
||||
if dest is module:
|
||||
setattr(dest, key, proxy)
|
||||
else:
|
||||
dest.__dict__[key] = proxy
|
||||
else:
|
||||
raise AssertionError(f"Unhandled config {key}={value} ({type(value)})")
|
||||
|
||||
config: dict[str, _ConfigEntry] = {}
|
||||
|
||||
compile_ignored_keys = get_assignments_with_compile_ignored_comments(module)
|
||||
|
||||
visit(module, module, "")
|
||||
module._config = config # type: ignore[attr-defined]
|
||||
module._compile_ignored_keys = compile_ignored_keys # type: ignore[attr-defined]
|
||||
module.__class__ = ConfigModuleInstance
|
||||
module._is_dirty = True # type: ignore[attr-defined]
|
||||
module._hash_digest = None # type: ignore[attr-defined]
|
||||
|
||||
|
||||
COMPILE_IGNORED_MARKER = "@compile_ignored"
|
||||
|
||||
|
||||
# Gets all the keys (i.e. assignments) with a @compile_ignored comment
|
||||
def get_assignments_with_compile_ignored_comments(module: ModuleType) -> set[str]:
|
||||
source_code = inspect.getsource(module)
|
||||
assignments = set()
|
||||
|
||||
# Tokenize the source code to retrieve comments
|
||||
tokens = tokenize.tokenize(io.BytesIO(source_code.encode("utf-8")).readline)
|
||||
current_comment = "", -1
|
||||
prev_name = ""
|
||||
|
||||
for token in tokens:
|
||||
if token.type == tokenize.COMMENT:
|
||||
prev_name = ""
|
||||
maybe_current = token.string.strip()
|
||||
if COMPILE_IGNORED_MARKER in maybe_current:
|
||||
if current_comment != ("", -1):
|
||||
raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
|
||||
current_comment = maybe_current, token.start[0]
|
||||
elif token.type == tokenize.NAME:
|
||||
# Only accept the first name token, to handle if you have
|
||||
# something like foo: Bar = ...
|
||||
if not prev_name:
|
||||
prev_name = token.string
|
||||
elif token.type == tokenize.OP and token.string == "=":
|
||||
# Check if the current assignment follows a comment
|
||||
# with COMPILE_IGNORED_MARKER
|
||||
if (
|
||||
COMPILE_IGNORED_MARKER in current_comment[0]
|
||||
and current_comment[1] == token.start[0] - 1
|
||||
):
|
||||
assignments.add(prev_name)
|
||||
current_comment = "", -1 # reset
|
||||
prev_name = ""
|
||||
if current_comment != ("", -1):
|
||||
raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
|
||||
return assignments
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ConfigEntry:
|
||||
# The default value specified in the configuration
|
||||
default: Any
|
||||
# The type of the configuration value
|
||||
value_type: type
|
||||
# The value specified by the user when they overrode the configuration
|
||||
# _UNSET_SENTINEL indicates the value is not set.
|
||||
user_override: ContextVar[object]
|
||||
# The justknob to check for this config
|
||||
justknob: str | None = None
|
||||
# environment variables are read at install time
|
||||
env_value_force: Any = _UNSET_SENTINEL
|
||||
env_value_default: Any = _UNSET_SENTINEL
|
||||
# Used to work arounds bad assumptions in unittest.mock.patch
|
||||
# The code to blame is
|
||||
# https://github.com/python/cpython/blob/94a7a4e22fb8f567090514785c69e65298acca42/Lib/unittest/mock.py#L1637
|
||||
# Essentially, mock.patch requires, that if __dict__ isn't accessible
|
||||
# (which it isn't), that after delattr is called on the object, the
|
||||
# object must throw when hasattr is called. Otherwise, it doesn't call
|
||||
# setattr again.
|
||||
# Technically we'll have an intermediate state of hiding the config while
|
||||
# mock.patch is unpatching itself, but it calls setattr after the delete
|
||||
# call so the final state is correct. It's just very unintuitive.
|
||||
# upstream bug - python/cpython#126886
|
||||
hide: bool = False
|
||||
alias: str | None = None
|
||||
# Deprecation support
|
||||
deprecated: bool = False
|
||||
deprecation_message: str | None = None
|
||||
_deprecation_warned: bool = False
|
||||
|
||||
def __init__(self, config: _Config, name: str) -> None:
|
||||
self.default = config.default
|
||||
self.value_type = (
|
||||
config.value_type if config.value_type is not None else type(self.default)
|
||||
)
|
||||
self.justknob = config.justknob
|
||||
self.alias = config.alias
|
||||
# Deprecation fields
|
||||
self.deprecated = config.deprecated
|
||||
self.deprecation_message = config.deprecation_message
|
||||
self._deprecation_warned = False
|
||||
|
||||
self.user_override = ContextVar(name, default=_UNSET_SENTINEL)
|
||||
if config.env_name_default is not None:
|
||||
for val in config.env_name_default:
|
||||
if (env_value := _read_env_variable(val)) is not None:
|
||||
self.env_value_default = env_value
|
||||
break
|
||||
if config.env_name_force is not None:
|
||||
for val in config.env_name_force:
|
||||
if (env_value := _read_env_variable(val)) is not None:
|
||||
self.env_value_force = env_value
|
||||
break
|
||||
|
||||
# Ensure justknobs and envvars are allowlisted types
|
||||
if self.justknob is not None and self.default is not None:
|
||||
if not isinstance(self.default, bool):
|
||||
raise AssertionError(
|
||||
f"justknobs only support booleans, {self.default} is not a boolean"
|
||||
)
|
||||
if self.value_type is not None and (
|
||||
config.env_name_default is not None or config.env_name_force is not None
|
||||
):
|
||||
if self.value_type not in (
|
||||
bool,
|
||||
str,
|
||||
Optional[bool], # noqa: UP045
|
||||
Optional[str], # noqa: UP045
|
||||
):
|
||||
raise AssertionError(
|
||||
f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither"
|
||||
)
|
||||
|
||||
|
||||
class ConfigModule(ModuleType):
|
||||
# NOTE: This should be kept in sync with _config_typing.pyi.
|
||||
|
||||
# The actual configuration settings. E.g., torch._dynamo.config.debug
|
||||
# would live as "debug" in the key, and torch._inductor.config.triton.cudagraphs
|
||||
# maps as "triton.cudagraphs". See discussion on the class for meaning of various sub items
|
||||
_config: dict[str, _ConfigEntry]
|
||||
_bypass_keys: set[str]
|
||||
_compile_ignored_keys: set[str]
|
||||
_is_dirty: bool
|
||||
_hash_digest: bytes | None
|
||||
|
||||
def __init__(self) -> None:
|
||||
raise NotImplementedError(
|
||||
f"use {__name__}.install_config_module(sys.modules[__name__])"
|
||||
)
|
||||
|
||||
def _warn_if_deprecated(self, name: str, config: _ConfigEntry) -> None:
|
||||
"""Issue deprecation warning for config if not already warned."""
|
||||
if config.deprecated and not config._deprecation_warned:
|
||||
import warnings
|
||||
|
||||
msg = f"{self.__name__}.{name} is deprecated"
|
||||
if config.deprecation_message:
|
||||
msg += f" and {config.deprecation_message}"
|
||||
msg += ". It will be removed in a future version of PyTorch."
|
||||
warnings.warn(msg, FutureWarning, stacklevel=3)
|
||||
config._deprecation_warned = True
|
||||
|
||||
def __setattr__(self, name: str, value: object) -> None:
|
||||
if name in self._bypass_keys:
|
||||
super().__setattr__(name, value)
|
||||
elif name not in self._config:
|
||||
raise AttributeError(f"{self.__name__}.{name} does not exist")
|
||||
else:
|
||||
# Issue deprecation warning on write (once per config)
|
||||
config = self._config[name]
|
||||
self._warn_if_deprecated(name, config)
|
||||
|
||||
if config.alias is not None:
|
||||
self._set_alias_val(config, value)
|
||||
else:
|
||||
config.user_override.set(value)
|
||||
self._is_dirty = True
|
||||
config.hide = False
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
try:
|
||||
config = self._config[name]
|
||||
|
||||
if config.hide:
|
||||
raise AttributeError(f"{self.__name__}.{name} does not exist")
|
||||
|
||||
# Issue deprecation warning on read (once per config)
|
||||
self._warn_if_deprecated(name, config)
|
||||
|
||||
alias_val = self._get_alias_val(config)
|
||||
if alias_val is not _UNSET_SENTINEL:
|
||||
return alias_val
|
||||
|
||||
if config.env_value_force is not _UNSET_SENTINEL:
|
||||
return config.env_value_force
|
||||
|
||||
user_override = config.user_override.get()
|
||||
if user_override is not _UNSET_SENTINEL:
|
||||
return user_override
|
||||
|
||||
if config.env_value_default is not _UNSET_SENTINEL:
|
||||
return config.env_value_default
|
||||
|
||||
if config.justknob is not None:
|
||||
# JK only supports bools and ints
|
||||
return justknobs_check(name=config.justknob, default=config.default)
|
||||
|
||||
# Reference types can still be modified, so copy them to
|
||||
# user_overrides to prevent accidental mutation of defaults.
|
||||
if not isinstance(config.default, _IMMUTABLE_CONFIG_TYPES):
|
||||
config.user_override.set(copy.deepcopy(config.default))
|
||||
return config.user_override.get()
|
||||
return config.default
|
||||
|
||||
except KeyError as e:
|
||||
# make hasattr() work properly
|
||||
raise AttributeError(f"{self.__name__}.{name} does not exist") from e
|
||||
|
||||
def __delattr__(self, name: str) -> None:
|
||||
self._is_dirty = True
|
||||
# must support delete because unittest.mock.patch deletes
|
||||
# then recreate things
|
||||
self._config[name].user_override.set(_UNSET_SENTINEL)
|
||||
self._config[name].hide = True
|
||||
|
||||
def _get_alias_module_and_name(
|
||||
self, entry: _ConfigEntry
|
||||
) -> tuple[ModuleType, str] | None:
|
||||
alias = entry.alias
|
||||
if alias is None:
|
||||
return None
|
||||
module_name, constant_name = alias.rsplit(".", 1)
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
except ImportError as e:
|
||||
raise AttributeError(f"config alias {alias} does not exist") from e
|
||||
return module, constant_name
|
||||
|
||||
def _get_alias_val(self, entry: _ConfigEntry) -> Any:
|
||||
data = self._get_alias_module_and_name(entry)
|
||||
if data is None:
|
||||
return _UNSET_SENTINEL
|
||||
module, constant_name = data
|
||||
constant_value = getattr(module, constant_name)
|
||||
return constant_value
|
||||
|
||||
def _set_alias_val(self, entry: _ConfigEntry, val: Any) -> None:
|
||||
data = self._get_alias_module_and_name(entry)
|
||||
if data is None:
|
||||
raise AssertionError(
|
||||
"alias data should not be None when setting alias value"
|
||||
)
|
||||
module, constant_name = data
|
||||
setattr(module, constant_name, val)
|
||||
|
||||
def _is_default(self, name: str) -> bool:
|
||||
"""
|
||||
Returns true if the config is at its default value.
|
||||
configs overridden by the env are not considered default.
|
||||
"""
|
||||
config_val = self._config[name]
|
||||
# The config is not overridden by the user, and the env_value_default
|
||||
# is different from the default value (meaning user has set the env to
|
||||
# change the default value).
|
||||
not_set_env_default = (
|
||||
config_val.env_value_default is _UNSET_SENTINEL
|
||||
or config_val.env_value_default == config_val.default
|
||||
)
|
||||
not_set_env_force = (
|
||||
config_val.env_value_force is _UNSET_SENTINEL
|
||||
or config_val.env_value_force == config_val.default
|
||||
)
|
||||
|
||||
unset = config_val.user_override.get() is _UNSET_SENTINEL
|
||||
# Handle reference types specially to avoid spammy warnings
|
||||
if not isinstance(config_val.default, _IMMUTABLE_CONFIG_TYPES):
|
||||
unset = unset or config_val.user_override.get() == config_val.default
|
||||
return unset and not_set_env_default and not_set_env_force
|
||||
|
||||
def _get_dict(
|
||||
self,
|
||||
ignored_keys: list[str] | None = None,
|
||||
ignored_prefixes: list[str] | None = None,
|
||||
skip_default: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Export a dictionary of current configuration keys and values.
|
||||
|
||||
This function is design to provide a single point which handles
|
||||
accessing config options and exporting them into a dictionary.
|
||||
This is used by a number of different user facing export methods
|
||||
which all have slightly different semantics re: how and what to
|
||||
skip.
|
||||
If a config is aliased, it skips this config.
|
||||
|
||||
Arguments:
|
||||
ignored_keys are keys that should not be exported.
|
||||
ignored_prefixes are prefixes that if a key matches should
|
||||
not be exported
|
||||
skip_default does two things. One if a key has not been modified
|
||||
it skips it.
|
||||
"""
|
||||
config: dict[str, Any] = {}
|
||||
for key, entry in self._config.items():
|
||||
if entry.alias is not None:
|
||||
continue
|
||||
if ignored_keys and key in ignored_keys:
|
||||
continue
|
||||
if ignored_prefixes:
|
||||
if any(key.startswith(prefix) for prefix in ignored_prefixes):
|
||||
continue
|
||||
if skip_default and self._is_default(key):
|
||||
continue
|
||||
|
||||
# Read value directly, bypassing __getattr__ overhead
|
||||
# (deprecation warnings, alias resolution).
|
||||
user_override = entry.user_override.get()
|
||||
if entry.env_value_force is not _UNSET_SENTINEL:
|
||||
val = entry.env_value_force
|
||||
elif user_override is not _UNSET_SENTINEL:
|
||||
val = user_override
|
||||
elif entry.env_value_default is not _UNSET_SENTINEL:
|
||||
val = entry.env_value_default
|
||||
elif entry.justknob is not None:
|
||||
val = justknobs_check(name=entry.justknob, default=entry.default)
|
||||
else:
|
||||
val = entry.default
|
||||
if not isinstance(val, _IMMUTABLE_CONFIG_TYPES):
|
||||
val = copy.deepcopy(val)
|
||||
config[key] = val
|
||||
|
||||
return config
|
||||
|
||||
def get_type(self, config_name: str) -> type:
|
||||
return self._config[config_name].value_type
|
||||
|
||||
def save_config(self) -> bytes:
|
||||
"""Convert config to a pickled blob"""
|
||||
ignored_keys = getattr(self, "_save_config_ignore", [])
|
||||
return pickle.dumps(
|
||||
self._get_dict(ignored_keys=ignored_keys),
|
||||
protocol=2,
|
||||
)
|
||||
|
||||
def save_config_portable(
|
||||
self, *, ignore_private_configs: bool = True
|
||||
) -> dict[str, Any]:
|
||||
"""Convert config to portable format"""
|
||||
prefixes = []
|
||||
if ignore_private_configs:
|
||||
prefixes.append("_")
|
||||
prefixes.extend(getattr(self, "_cache_config_ignore_prefix", []))
|
||||
config = self._get_dict(ignored_prefixes=prefixes)
|
||||
for key in getattr(self, "_cache_config_factory_keys", []):
|
||||
if key in config and config[key] is not None:
|
||||
instance = config[key]()
|
||||
if hasattr(instance, "uuid"):
|
||||
config[key] = instance.uuid()
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Config '{key}' is set to {config[key]} which does not "
|
||||
f"implement uuid(). Implement uuid() for cache key "
|
||||
f"participation."
|
||||
)
|
||||
return config
|
||||
|
||||
def codegen_config(self) -> str:
|
||||
"""Convert config to Python statements that replicate current config.
|
||||
This does NOT include config settings that are at default values.
|
||||
"""
|
||||
|
||||
# additional imports required
|
||||
imports = set()
|
||||
|
||||
def get_module_name(func: Callable, add_dot: bool) -> str:
|
||||
module_name = func.__module__
|
||||
if module_name == "builtins":
|
||||
module_name = ""
|
||||
if add_dot and module_name != "":
|
||||
module_name += "."
|
||||
return module_name
|
||||
|
||||
def add_import(func: Callable) -> None:
|
||||
module_name = get_module_name(func, False)
|
||||
if module_name:
|
||||
imports.add(module_name)
|
||||
|
||||
def list_of_callables_to_string(v: list | set) -> list[str]:
|
||||
return [f"{get_module_name(item, True)}{item.__name__}" for item in v]
|
||||
|
||||
def importable_callable(v: Any) -> bool:
|
||||
# functools.partial has no attributes below but is a callable
|
||||
return callable(v) and hasattr(v, "__module__") and hasattr(v, "__name__")
|
||||
|
||||
def get_config_line(mod, k, v) -> str: # type: ignore[no-untyped-def]
|
||||
"""
|
||||
Return a string version of the config line.
|
||||
Handle v when v is a callable, or a list/dict of callables. Add import statements for callables if necessary.
|
||||
We assume that the value of a single config won't be a mix of callables and non-callables.
|
||||
|
||||
Example output:
|
||||
import logging
|
||||
import _warnings
|
||||
torch._dynamo.config.reorderable_logging_functions = { _warnings.warn, logging.warn, print }
|
||||
"""
|
||||
if importable_callable(v):
|
||||
add_import(v)
|
||||
return f"{mod}.{k} = {get_module_name(v, True)}{v.__name__}"
|
||||
elif isinstance(v, (list, set)) and all(
|
||||
importable_callable(item) for item in v
|
||||
):
|
||||
for item in v:
|
||||
add_import(item)
|
||||
v_list = list_of_callables_to_string(v)
|
||||
if isinstance(v, list):
|
||||
return f"{mod}.{k} = {v_list}"
|
||||
else:
|
||||
return f"{mod}.{k} = {{ {', '.join(v_list)} }}"
|
||||
else:
|
||||
return f"{mod}.{k} = {v!r}"
|
||||
|
||||
lines = []
|
||||
mod = self.__name__
|
||||
for k, v in self._get_dict(
|
||||
ignored_keys=getattr(self, "_save_config_ignore", []), skip_default=True
|
||||
).items():
|
||||
lines.append(get_config_line(mod, k, v))
|
||||
for import_name in imports:
|
||||
lines.insert(0, f"import {import_name}")
|
||||
return "\n".join(lines)
|
||||
|
||||
def get_hash(self) -> bytes:
|
||||
"""Hashes the configs that are not compile_ignored"""
|
||||
if self._is_dirty or self._hash_digest is None:
|
||||
dict_to_hash = self._get_dict(ignored_keys=list(self._compile_ignored_keys))
|
||||
string_to_hash = repr(sorted(dict_to_hash.items()))
|
||||
self._hash_digest = hashlib.md5(
|
||||
string_to_hash.encode("utf-8"), usedforsecurity=False
|
||||
).digest()
|
||||
self._is_dirty = False
|
||||
return self._hash_digest
|
||||
|
||||
@deprecated(
|
||||
"`config.to_dict()` has been deprecated. It no longer changes the underlying config."
|
||||
" use `config.get_config_copy()` instead if you just want a copy of the config, or "
|
||||
"config.load_config if you need mutable access",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return self.get_config_copy()
|
||||
|
||||
@deprecated(
|
||||
"`config.shallow_copy_dict()` has been deprecated. It no longer changes the underlying config."
|
||||
" use `config.get_config_copy()` instead if you just want a copy of the config, or "
|
||||
"config.load_config if you need mutable access",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def shallow_copy_dict(self) -> dict[str, Any]:
|
||||
return self.get_config_copy()
|
||||
|
||||
def load_config(self, maybe_pickled_config: bytes | dict[str, Any]) -> None:
|
||||
"""Restore from a prior call to save_config() or shallow_copy_dict()"""
|
||||
if not isinstance(maybe_pickled_config, dict):
|
||||
config = pickle.loads(maybe_pickled_config)
|
||||
else:
|
||||
config = maybe_pickled_config
|
||||
for k, v in config.items():
|
||||
if k in self._config:
|
||||
setattr(self, k, v)
|
||||
else:
|
||||
from torch._dynamo.utils import warn_once
|
||||
|
||||
warn_once(f"key {k} with value {v} is not understood by this config")
|
||||
|
||||
def get_config_copy(self) -> dict[str, Any]:
|
||||
return self._get_dict()
|
||||
|
||||
def get_serializable_config_copy(self) -> dict[str, Any]:
|
||||
return self._get_dict(ignored_keys=getattr(self, "_save_config_ignore", []))
|
||||
|
||||
def patch(
|
||||
self,
|
||||
arg1: str | dict[str, Any] | None = None,
|
||||
arg2: Any = None,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> "ContextDecorator":
|
||||
"""
|
||||
Decorator and/or context manager to make temporary changes to a config. Note that patched settings are thread-local.
|
||||
|
||||
As a decorator:
|
||||
|
||||
@config.patch("name", val)
|
||||
@config.patch(name1=val1, name2=val2)
|
||||
@config.patch({"name1": val1, "name2", val2})
|
||||
def foo(...):
|
||||
...
|
||||
|
||||
As a context manager:
|
||||
|
||||
with config.patch("name", val):
|
||||
...
|
||||
"""
|
||||
changes: dict[str, Any]
|
||||
if arg1 is not None:
|
||||
if arg2 is not None:
|
||||
if not isinstance(arg1, str):
|
||||
raise AssertionError(
|
||||
"first argument must be a string when passing 2 positional args to patch"
|
||||
)
|
||||
# patch("key", True) syntax
|
||||
changes = {arg1: arg2}
|
||||
else:
|
||||
if not isinstance(arg1, dict):
|
||||
raise AssertionError(
|
||||
"first argument must be a dict when passing a single positional arg to patch"
|
||||
)
|
||||
# patch({"key": True}) syntax
|
||||
changes = arg1
|
||||
if kwargs:
|
||||
raise AssertionError(
|
||||
"cannot pass both positional and keyword arguments to patch"
|
||||
)
|
||||
else:
|
||||
# patch(key=True) syntax
|
||||
changes = kwargs
|
||||
if arg2 is not None:
|
||||
raise AssertionError(
|
||||
"second positional argument is only valid when first argument is a key string"
|
||||
)
|
||||
if not isinstance(changes, dict):
|
||||
raise AssertionError(f"expected `dict` got {type(changes)}")
|
||||
config = self
|
||||
|
||||
class ConfigPatch(ContextDecorator):
|
||||
def __init__(self) -> None:
|
||||
self.changes = changes
|
||||
self._prior: ContextVar[tuple[dict[str, Any], ...]] = ContextVar(
|
||||
f"{config.__name__}.ConfigPatch[{id(self)}]",
|
||||
default=(),
|
||||
)
|
||||
|
||||
def __enter__(self) -> None:
|
||||
prior: dict[str, Any] = {}
|
||||
for key in self.changes:
|
||||
# KeyError on invalid entry
|
||||
prior[key] = config.__getattr__(key)
|
||||
prior_stack = self._prior.get()
|
||||
self._prior.set((*prior_stack, prior))
|
||||
try:
|
||||
for k, v in self.changes.items():
|
||||
config.__setattr__(k, v)
|
||||
except Exception:
|
||||
self._prior.set(prior_stack)
|
||||
raise
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore[no-untyped-def]
|
||||
prior_stack = self._prior.get()
|
||||
if not prior_stack:
|
||||
raise AssertionError(
|
||||
"prior should not be empty when exiting ConfigPatch"
|
||||
)
|
||||
prior = prior_stack[-1]
|
||||
self._prior.set(prior_stack[:-1])
|
||||
for k, v in prior.items():
|
||||
config.__setattr__(k, v)
|
||||
|
||||
return ConfigPatch()
|
||||
|
||||
def _make_closure_patcher(self, **changes: dict[str, Any]) -> Any:
|
||||
"""
|
||||
A lower-overhead version of patch() for things on the critical path.
|
||||
|
||||
Usage:
|
||||
|
||||
# do this off the critical path
|
||||
change_fn = config.make_closure_patcher(foo=True)
|
||||
|
||||
...
|
||||
|
||||
revert = change_fn()
|
||||
try:
|
||||
...
|
||||
finally:
|
||||
revert()
|
||||
|
||||
"""
|
||||
config = self._config
|
||||
|
||||
def change() -> Callable[[], None]:
|
||||
prior = {k: config[k].user_override.get() for k in changes}
|
||||
for k, v in changes.items():
|
||||
self._config[k].user_override.set(v)
|
||||
|
||||
def revert() -> None:
|
||||
for k, v in prior.items():
|
||||
self._config[k].user_override.set(v)
|
||||
|
||||
return revert
|
||||
|
||||
return change
|
||||
|
||||
|
||||
class ContextDecorator(contextlib.ContextDecorator):
|
||||
"""
|
||||
Same as contextlib.ContextDecorator, but with support for
|
||||
`unittest.TestCase`
|
||||
"""
|
||||
|
||||
def __enter__(self) -> None:
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb) -> NoReturn: # type: ignore[no-untyped-def]
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
def __call__(self, func: Callable[[Any], Any]) -> Any:
|
||||
if isinstance(func, type) and issubclass(func, unittest.TestCase):
|
||||
|
||||
class _TestCase(func): # type: ignore[valid-type, misc]
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
self.__enter__()
|
||||
try:
|
||||
super().setUpClass()
|
||||
except Exception:
|
||||
self.__exit__(None, None, None)
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
try:
|
||||
super().tearDownClass()
|
||||
finally:
|
||||
self.__exit__(None, None, None)
|
||||
|
||||
_TestCase.__name__ = func.__name__
|
||||
_TestCase.__qualname__ = func.__qualname__
|
||||
_TestCase.__module__ = func.__module__
|
||||
|
||||
return _TestCase
|
||||
|
||||
return super().__call__(func)
|
||||
|
||||
|
||||
class SubConfigProxy:
|
||||
"""
|
||||
Shim to redirect to main config.
|
||||
`config.triton.cudagraphs` maps to _config["triton.cudagraphs"]
|
||||
"""
|
||||
|
||||
def __init__(self, config: object, prefix: str) -> None:
|
||||
# `super().__setattr__` to bypass custom `__setattr__`
|
||||
super().__setattr__("_config", config)
|
||||
super().__setattr__("_prefix", prefix)
|
||||
|
||||
def __setattr__(self, name: str, value: object) -> None:
|
||||
return self._config.__setattr__(self._prefix + name, value)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return self._config.__getattr__(self._prefix + name)
|
||||
|
||||
def __delattr__(self, name: str) -> None:
|
||||
return self._config.__delattr__(self._prefix + name)
|
||||
|
||||
|
||||
def get_tristate_env(name: str, default: Any = None) -> bool | None:
|
||||
value = os.environ.get(name)
|
||||
if value == "1":
|
||||
return True
|
||||
if value == "0":
|
||||
return False
|
||||
return default
|
||||
|
||||
|
||||
def inherit_fields_from(parent_cls):
|
||||
def wrapper(child_cls):
|
||||
for k, v in parent_cls.__dict__.items():
|
||||
# copy fields that are not private and not overridden
|
||||
if not k.startswith("_") and k not in child_cls.__dict__:
|
||||
setattr(child_cls, k, v)
|
||||
return child_cls
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,36 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
"""
|
||||
This was semi-automatically generated by running
|
||||
|
||||
stubgen torch.utils._config_module.py
|
||||
|
||||
And then manually extracting the methods of ConfigModule and converting them into top-level functions.
|
||||
|
||||
This file should be imported into any file that uses install_config_module like so:
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils._config_typing import * # noqa: F401, F403
|
||||
|
||||
from torch.utils._config_module import install_config_module
|
||||
|
||||
# adds patch, save_config, etc
|
||||
install_config_module(sys.modules[__name__])
|
||||
|
||||
Note that the import should happen before the call to install_config_module(), otherwise runtime errors may occur.
|
||||
"""
|
||||
|
||||
if not TYPE_CHECKING: # noqa: PYI002
|
||||
raise AssertionError("Do not use at runtime") # noqa: W291
|
||||
|
||||
def save_config() -> bytes: ...
|
||||
def save_config_portable(*, ignore_private_configs: bool = True) -> dict[str, Any]: ...
|
||||
def codegen_config() -> str: ...
|
||||
def get_hash() -> bytes: ...
|
||||
def to_dict() -> dict[str, Any]: ...
|
||||
def shallow_copy_dict() -> dict[str, Any]: ...
|
||||
def load_config(config: bytes | dict[str, Any]) -> None: ...
|
||||
def get_config_copy() -> dict[str, Any]: ...
|
||||
def get_serializable_config_copy() -> dict[str, Any]: ...
|
||||
def patch(arg1: str | dict[str, Any] | None = None, arg2: Any = None, **kwargs): ...
|
||||
@@ -0,0 +1,241 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# This module provides a FAST (on GPU) content addressable store for storages
|
||||
# (and tensors on top of them) with VERY WEAK portability guarantees (e.g.,
|
||||
# don't expect CPU/CUDA to address to the same hash, don't expect it to be
|
||||
# portable across devices) that is NOT cryptographically secure. In return,
|
||||
# we are able to hash 40G of tensor data on GPU in less than a second,
|
||||
# compared to running SHA-1 in CPU which would a minute or so. The primary
|
||||
# use case is for efficiently snapshotting intermediate tensor data for
|
||||
# offline debugging, but it's been put in this module in case you think of
|
||||
# another use case for it. The hash function could be replaced with a
|
||||
# straight reimplementation of SHA-1, which would give us much stronger
|
||||
# portability guarantees.
|
||||
#
|
||||
# WARNING: THERE IS NO BC/FC GUARANTEE FOR THIS FORMAT! If you need to format
|
||||
# shift the result, consider packing it into a single torch.save object
|
||||
# with traditional view sharing.
|
||||
#
|
||||
# Because of the weak portability guarantees, you can only write to the
|
||||
# content store from a single process; we don't provide any capability
|
||||
# of "reopening" a content store to add more things to it. But we don't
|
||||
# assume that you can keep all of the tensors you want to add to the store
|
||||
# in memory at once, because you probably can't! Nor do we assume that
|
||||
# you know a priori whether or not two storages can be deduplicated or not.
|
||||
#
|
||||
# Note: only storages are content-addressed; tensors are name addressed
|
||||
#
|
||||
# Note: our padding strategy means that [1, 0] and [1] int16 tensors would
|
||||
# map to the same (padded) storage. We think this will be immaterial for most
|
||||
# users.
|
||||
|
||||
import ctypes
|
||||
import functools
|
||||
import hashlib
|
||||
import os.path
|
||||
import struct
|
||||
from collections import defaultdict
|
||||
|
||||
import torch
|
||||
import torch._prims as prims
|
||||
import torch._utils
|
||||
import torch.nn.functional as F
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
|
||||
|
||||
def lazy_compile(**compile_kwargs):
|
||||
"""Lazily wrap a function with torch.compile on the first call
|
||||
|
||||
This avoids eagerly importing dynamo.
|
||||
"""
|
||||
|
||||
def decorate_fn(fn):
|
||||
@functools.wraps(fn)
|
||||
def compile_hook(*args, **kwargs):
|
||||
compiled_fn = torch.compile(fn, **compile_kwargs)
|
||||
globals()[fn.__name__] = functools.wraps(fn)(compiled_fn)
|
||||
return compiled_fn(*args, **kwargs)
|
||||
|
||||
return compile_hook
|
||||
|
||||
return decorate_fn
|
||||
|
||||
|
||||
# Use of torch.compile is mandatory for (1) good memory usage
|
||||
# and (2) xor_sum implementation. This is our first instance of
|
||||
# using PT2 to implement a kernel in PyTorch; if we get AOT capabilities
|
||||
# it would be good to apply it here.
|
||||
@lazy_compile(dynamic=True)
|
||||
def hash_storage_kernel(x):
|
||||
# The randint calls are carefully written to hit things we
|
||||
# have lowerings for in inductor. Lack of unsigned 32-bit integer
|
||||
# is a pain.
|
||||
a = torch.randint(
|
||||
-(2**31), 2**31, x.shape, device=x.device, dtype=torch.int32
|
||||
).abs()
|
||||
a = ((a % (2**31 - 1)) + 1).long()
|
||||
b = (
|
||||
torch.randint(-(2**31), 2**31, x.shape, device=x.device, dtype=torch.int32)
|
||||
.abs()
|
||||
.long()
|
||||
)
|
||||
# This is a standard shift-multiply universal hash family
|
||||
# plus xor sum hash, using Philox to generate random numbers.
|
||||
# Our Philox RNG is not deterministic across devices so
|
||||
# don't use this for stable hashing.
|
||||
#
|
||||
# This assumes fixed length so you're also obligated to bucket
|
||||
# by the length of tensor as well
|
||||
return prims.xor_sum((a * x + b).int(), [0])
|
||||
|
||||
|
||||
# Returns a hex digest of the data in the storage. Guaranteed to be
|
||||
# SHA-1 if stable_hash=True, otherwise it will consistent for a single
|
||||
# process run but not necessarily across processes.
|
||||
def hash_storage(storage: torch.UntypedStorage, *, stable_hash: bool = False) -> str:
|
||||
import torch._dynamo
|
||||
from torch._dynamo.utils import is_compile_supported
|
||||
|
||||
device_type = storage.device.type
|
||||
if stable_hash or not is_compile_supported(device_type):
|
||||
cpu_storage = storage.cpu()
|
||||
# TODO: make storage support buffer protocol so this isn't
|
||||
# necessary
|
||||
buf = (ctypes.c_byte * cpu_storage.nbytes()).from_address(
|
||||
cpu_storage.data_ptr()
|
||||
)
|
||||
sha1 = hashlib.sha1(usedforsecurity=False)
|
||||
sha1.update(buf)
|
||||
return sha1.hexdigest()
|
||||
|
||||
# TODO: factor this into a random utility
|
||||
if device_type == "cpu":
|
||||
generator = torch._C.default_generator
|
||||
elif device_type == "cuda":
|
||||
generator = torch.cuda.default_generators[storage.device.index]
|
||||
elif device_type == "mps":
|
||||
generator = torch.mps._get_default_mps_generator()
|
||||
elif device_type == "xpu":
|
||||
generator = torch.xpu.default_generators[storage.device.index]
|
||||
else:
|
||||
raise AssertionError(f"unhandled device type {device_type}")
|
||||
state = generator.get_state()
|
||||
try:
|
||||
generator.manual_seed(0)
|
||||
x = torch.empty(0, dtype=torch.uint8, device=storage.device).set_(storage) # type: ignore[call-overload]
|
||||
# The dtype-casting view cannot be compiled, and so the
|
||||
# padding/reshaping also needs to be done externally even
|
||||
# though it could be profitably fused
|
||||
pad = -x.numel() % 4
|
||||
if pad > 0:
|
||||
x = F.pad(x, (0, pad), "constant", 0)
|
||||
x = x.view(torch.int32)
|
||||
# We run the 32-bit hash five times with differing parameters to
|
||||
# reduce chance of collision
|
||||
ITER = 5
|
||||
cs = [hash_storage_kernel(x).item() for _ in range(ITER)]
|
||||
return struct.pack(">" + "i" * ITER, *cs).hex()
|
||||
finally:
|
||||
generator.set_state(state)
|
||||
|
||||
|
||||
class ContentStoreWriter:
|
||||
# Structure:
|
||||
# storages/
|
||||
# 00/
|
||||
# 0000..00
|
||||
# tensors/
|
||||
# name
|
||||
def __init__(self, loc: str, stable_hash: bool = False) -> None:
|
||||
self.loc: str = loc
|
||||
self.seen_storage_hashes: set[str] = set()
|
||||
self.stable_hash = stable_hash
|
||||
|
||||
# TODO: offer some sort of non-blocking API to speed things up
|
||||
def write_storage(self, storage: torch.UntypedStorage) -> str:
|
||||
h = hash_storage(storage, stable_hash=self.stable_hash)
|
||||
if h in self.seen_storage_hashes:
|
||||
return h
|
||||
# TODO: consider not using torch.save for this; we don't actually
|
||||
# need any metadata for the storage
|
||||
subfolder = os.path.join(self.loc, "storages")
|
||||
os.makedirs(subfolder, exist_ok=True)
|
||||
target = os.path.join(subfolder, h)
|
||||
if os.path.exists(target):
|
||||
return h
|
||||
torch.save(storage, target)
|
||||
self.seen_storage_hashes.add(h)
|
||||
return h
|
||||
|
||||
def compute_tensor_metadata(self, t: torch.Tensor, h=None):
|
||||
if h is None:
|
||||
h = hash_storage(t.untyped_storage(), stable_hash=self.stable_hash)
|
||||
return (
|
||||
t.dtype,
|
||||
h,
|
||||
t.storage_offset(),
|
||||
tuple(t.shape),
|
||||
t.stride(),
|
||||
torch._utils.get_tensor_metadata(t),
|
||||
)
|
||||
|
||||
def write_tensor(self, name: str, t: torch.Tensor) -> None:
|
||||
storage = t.untyped_storage()
|
||||
h = self.write_storage(storage)
|
||||
# TODO: Support more advanced snapshotting of requires_grad/grad/etc
|
||||
d, f = os.path.split(name)
|
||||
payload = self.compute_tensor_metadata(t, h=h)
|
||||
subfolder = os.path.join(self.loc, "tensors", d)
|
||||
os.makedirs(subfolder, exist_ok=True)
|
||||
torch.save(payload, os.path.join(subfolder, f))
|
||||
|
||||
|
||||
class ContentStoreReader:
|
||||
def __init__(self, loc: str, *, cache=True) -> None:
|
||||
self.loc = loc
|
||||
self.storage_cache: (
|
||||
dict[torch.device | None, dict[str, StorageWeakRef]] | None
|
||||
) = None
|
||||
if cache:
|
||||
self.storage_cache = defaultdict(dict)
|
||||
|
||||
def read_storage(self, h: str, *, device=None) -> torch.UntypedStorage:
|
||||
if device is not None:
|
||||
device = torch.device(device)
|
||||
ws = (
|
||||
self.storage_cache[device].get(h)
|
||||
if self.storage_cache is not None
|
||||
else None
|
||||
)
|
||||
s: torch.UntypedStorage | None
|
||||
if ws is not None:
|
||||
s = torch.UntypedStorage._new_with_weak_ptr(ws.cdata)
|
||||
if s is not None:
|
||||
return s
|
||||
s = torch.load(
|
||||
os.path.join(self.loc, "storages", h),
|
||||
weights_only=True,
|
||||
map_location=device,
|
||||
)._untyped_storage
|
||||
if s is None:
|
||||
raise AssertionError(
|
||||
f"expected storage for hash {h} in {os.path.join(self.loc, 'storages')}, got None"
|
||||
)
|
||||
if self.storage_cache is not None:
|
||||
self.storage_cache[device][h] = StorageWeakRef(s)
|
||||
return s
|
||||
|
||||
def read_tensor_metadata(self, name: str):
|
||||
fn = os.path.join(self.loc, "tensors", name)
|
||||
if not os.path.exists(fn):
|
||||
raise FileNotFoundError(fn)
|
||||
return torch.load(fn, weights_only=True)
|
||||
|
||||
def read_tensor(self, name: str, *, device=None) -> torch.Tensor:
|
||||
dtype, h, storage_offset, size, stride, metadata = self.read_tensor_metadata(
|
||||
name
|
||||
)
|
||||
storage = self.read_storage(h, device=device)
|
||||
t = torch.tensor([], dtype=dtype, device=storage.device)
|
||||
t.set_(storage, storage_offset, size, stride)
|
||||
torch._utils.set_tensor_metadata(t, metadata)
|
||||
return t
|
||||
@@ -0,0 +1,170 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Extra utilities for working with context managers that should have been
|
||||
# in the standard library but are not
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
import sys
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast, overload, TypeVar
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
# Used for annotating the decorator usage of _DecoratorContextManager (e.g.,
|
||||
# 'no_grad' and 'enable_grad').
|
||||
# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators
|
||||
FuncType = Callable[..., Any]
|
||||
F = TypeVar("F", bound=FuncType)
|
||||
|
||||
|
||||
def _wrap_generator(ctx_factory, func):
|
||||
"""
|
||||
Wrap each generator invocation with the context manager factory.
|
||||
|
||||
The input should be a function that returns a context manager,
|
||||
not a context manager itself, to handle one-shot context managers.
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
def generator_context(*args, **kwargs):
|
||||
gen = func(*args, **kwargs)
|
||||
|
||||
# Generators are suspended and unsuspended at `yield`, hence we
|
||||
# make sure the grad mode is properly set every time the execution
|
||||
# flow returns into the wrapped generator and restored when it
|
||||
# returns through our `yield` to our caller (see PR #49017).
|
||||
try:
|
||||
# Issuing `None` to a generator fires it up
|
||||
with ctx_factory():
|
||||
response = gen.send(None)
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Forward the response to our caller and get its next request
|
||||
request = yield response
|
||||
|
||||
except GeneratorExit:
|
||||
# Inform the still active generator about its imminent closure
|
||||
with ctx_factory():
|
||||
gen.close()
|
||||
raise
|
||||
|
||||
except BaseException: # noqa: B036
|
||||
# Propagate the exception thrown at us by the caller
|
||||
with ctx_factory():
|
||||
response = gen.throw(*sys.exc_info())
|
||||
|
||||
else:
|
||||
# Pass the last request to the generator and get its response
|
||||
with ctx_factory():
|
||||
response = gen.send(request)
|
||||
|
||||
# We let the exceptions raised above by the generator's `.throw` or
|
||||
# `.send` methods bubble up to our caller, except for StopIteration
|
||||
except StopIteration as e:
|
||||
# The generator informed us that it is done: take whatever its
|
||||
# returned value (if any) was and indicate that we're done too
|
||||
# by returning it (see docs for python's return-statement).
|
||||
return e.value
|
||||
|
||||
return generator_context
|
||||
|
||||
|
||||
def context_decorator(ctx, func):
|
||||
"""
|
||||
Like contextlib.ContextDecorator.
|
||||
|
||||
But with the following differences:
|
||||
1. Is done by wrapping, rather than inheritance, so it works with context
|
||||
managers that are implemented from C and thus cannot easily inherit from
|
||||
Python classes
|
||||
2. Wraps generators in the intuitive way (c.f. https://bugs.python.org/issue37743)
|
||||
3. Errors out if you try to wrap a class, because it is ambiguous whether
|
||||
or not you intended to wrap only the constructor
|
||||
|
||||
The input argument can either be a context manager (in which case it must
|
||||
be a multi-shot context manager that can be directly invoked multiple times)
|
||||
or a callable that produces a context manager.
|
||||
"""
|
||||
if callable(ctx) and hasattr(ctx, "__enter__"):
|
||||
raise AssertionError(
|
||||
f"Passed in {ctx} is both callable and also a valid context manager "
|
||||
"(has __enter__), making it ambiguous which interface to use. If you "
|
||||
"intended to pass a context manager factory, rewrite your call as "
|
||||
"context_decorator(lambda: ctx()); if you intended to pass a context "
|
||||
"manager directly, rewrite your call as context_decorator(lambda: ctx)"
|
||||
)
|
||||
|
||||
if not callable(ctx):
|
||||
|
||||
def ctx_factory():
|
||||
return ctx
|
||||
|
||||
else:
|
||||
ctx_factory = ctx
|
||||
|
||||
if inspect.isclass(func):
|
||||
raise RuntimeError(
|
||||
"Cannot decorate classes; it is ambiguous whether or not only the "
|
||||
"constructor or all methods should have the context manager applied; "
|
||||
"additionally, decorating a class at definition-site will prevent "
|
||||
"use of the identifier as a conventional type. "
|
||||
"To specify which methods to decorate, decorate each of them "
|
||||
"individually."
|
||||
)
|
||||
|
||||
if inspect.isgeneratorfunction(func):
|
||||
return _wrap_generator(ctx_factory, func)
|
||||
|
||||
@functools.wraps(func)
|
||||
def decorate_context(*args, **kwargs):
|
||||
# pyrefly: ignore [bad-context-manager]
|
||||
with ctx_factory():
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return decorate_context
|
||||
|
||||
|
||||
class _DecoratorContextManager:
|
||||
"""Allow a context manager to be used as a decorator."""
|
||||
|
||||
def __call__(self, orig_func: F) -> F:
|
||||
if inspect.isclass(orig_func):
|
||||
warnings.warn(
|
||||
"Decorating classes is deprecated and will be disabled in "
|
||||
"future versions. You should only decorate functions or methods. "
|
||||
"To preserve the current behavior of class decoration, you can "
|
||||
"directly decorate the `__init__` method and nothing else.",
|
||||
FutureWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
func = cast(F, lambda *args, **kwargs: orig_func(*args, **kwargs))
|
||||
else:
|
||||
func = orig_func
|
||||
|
||||
return cast(F, context_decorator(self.clone, func))
|
||||
|
||||
def __enter__(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def clone(self):
|
||||
# override this method if your children class takes __init__ parameters
|
||||
return self.__class__()
|
||||
|
||||
|
||||
class _NoParamDecoratorContextManager(_DecoratorContextManager):
|
||||
"""Allow a context manager to be used as a decorator without parentheses."""
|
||||
|
||||
@overload
|
||||
def __new__(cls, orig_func: F) -> F: ... # type: ignore[misc]
|
||||
@overload
|
||||
def __new__(cls, orig_func: None = None) -> Self: ...
|
||||
|
||||
def __new__(cls, orig_func: F | None = None) -> Self | F: # type: ignore[misc]
|
||||
if orig_func is None:
|
||||
return super().__new__(cls)
|
||||
return cls()(orig_func)
|
||||
@@ -0,0 +1,57 @@
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from re import match as _match
|
||||
|
||||
|
||||
def read_file(fname: Path | str) -> list[str]:
|
||||
with open(fname, encoding="utf-8") as f:
|
||||
return f.readlines()
|
||||
|
||||
|
||||
def _embed_headers(
|
||||
content: list[str], include_dirs: list[Path], processed_files: set[str]
|
||||
) -> str:
|
||||
for line_idx, cur_line in enumerate(content):
|
||||
# Eliminate warning: `#pragma once in main file`
|
||||
if cur_line.startswith("#pragma once"):
|
||||
content[line_idx] = ""
|
||||
continue
|
||||
m = _match('^\\s*#include\\s*[<"]([^>"]+)[>"]', cur_line)
|
||||
if m is None:
|
||||
continue
|
||||
for include_dir in include_dirs:
|
||||
path = include_dir / m[1]
|
||||
if not path.exists():
|
||||
continue
|
||||
if str(path) in processed_files:
|
||||
content[line_idx] = ""
|
||||
continue
|
||||
processed_files.add(str(path))
|
||||
content[line_idx] = _embed_headers(
|
||||
read_file(path), include_dirs, processed_files
|
||||
)
|
||||
break
|
||||
return "".join(content)
|
||||
|
||||
|
||||
def embed_headers(
|
||||
fname: str, include_dirs: Sequence[str] | Sequence[Path] | str | None = None
|
||||
) -> str:
|
||||
if include_dirs is None:
|
||||
base_dir = Path(__file__).parent.parent.parent
|
||||
include_dirs = [base_dir, base_dir / "aten" / "src"]
|
||||
elif isinstance(include_dirs, str):
|
||||
include_dirs = [Path(include_dirs)]
|
||||
else:
|
||||
include_dirs = [Path(x) for x in include_dirs]
|
||||
|
||||
return _embed_headers(read_file(fname), include_dirs, {fname})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
print(f"Usage:\n {sys.argv[0]} filename")
|
||||
sys.exit(1)
|
||||
print(embed_headers(sys.argv[1]))
|
||||
@@ -0,0 +1,63 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import collections
|
||||
|
||||
|
||||
Entry = collections.namedtuple("Entry", "version, hash")
|
||||
|
||||
|
||||
def update_hash(seed, value):
|
||||
# Good old boost::hash_combine
|
||||
# https://www.boost.org/doc/libs/1_35_0/doc/html/boost/hash_combine_id241013.html
|
||||
return seed ^ (hash(value) + 0x9E3779B9 + (seed << 6) + (seed >> 2))
|
||||
|
||||
|
||||
def hash_source_files(hash_value, source_files):
|
||||
for filename in source_files:
|
||||
with open(filename, "rb") as file:
|
||||
hash_value = update_hash(hash_value, file.read())
|
||||
return hash_value
|
||||
|
||||
|
||||
def hash_build_arguments(hash_value, build_arguments):
|
||||
for group in build_arguments:
|
||||
if group:
|
||||
for argument in group:
|
||||
hash_value = update_hash(hash_value, argument)
|
||||
return hash_value
|
||||
|
||||
|
||||
class ExtensionVersioner:
|
||||
def __init__(self) -> None:
|
||||
self.entries = {}
|
||||
|
||||
def get_version(self, name):
|
||||
entry = self.entries.get(name)
|
||||
return None if entry is None else entry.version
|
||||
|
||||
def bump_version_if_changed(
|
||||
self,
|
||||
name,
|
||||
source_files,
|
||||
build_arguments,
|
||||
build_directory,
|
||||
with_cuda,
|
||||
with_sycl,
|
||||
is_python_module,
|
||||
is_standalone,
|
||||
):
|
||||
hash_value = 0
|
||||
hash_value = hash_source_files(hash_value, source_files)
|
||||
hash_value = hash_build_arguments(hash_value, build_arguments)
|
||||
hash_value = update_hash(hash_value, build_directory)
|
||||
hash_value = update_hash(hash_value, with_cuda)
|
||||
hash_value = update_hash(hash_value, with_sycl)
|
||||
hash_value = update_hash(hash_value, is_python_module)
|
||||
hash_value = update_hash(hash_value, is_standalone)
|
||||
|
||||
entry = self.entries.get(name)
|
||||
if entry is None:
|
||||
self.entries[name] = entry = Entry(0, hash_value)
|
||||
elif hash_value != entry.hash:
|
||||
self.entries[name] = entry = Entry(entry.version + 1, hash_value)
|
||||
|
||||
return entry.version
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
DebugMode: a debugging TorchDispatchMode that intercepts and logs runtime calls.
|
||||
|
||||
See torch.utils._debug_mode._mode for the full implementation and docstring.
|
||||
"""
|
||||
|
||||
from torch.utils._debug_mode._calls import (
|
||||
_AnnotateCall,
|
||||
_DebugCall,
|
||||
_get_call_name,
|
||||
_OpCall,
|
||||
_OutputPlacementCall,
|
||||
_RedistributeCall,
|
||||
_TritonKernelCall,
|
||||
)
|
||||
from torch.utils._debug_mode._mode import (
|
||||
DebugInterpreter,
|
||||
DebugMode,
|
||||
get_active_debug_mode,
|
||||
)
|
||||
from torch.utils._debug_mode._utils import (
|
||||
_stringify_shape,
|
||||
hash_tensor_fn,
|
||||
norm_hash_fn,
|
||||
TensorIdTracker,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["DebugMode", "get_active_debug_mode"]
|
||||
@@ -0,0 +1,354 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
Call record classes for DebugMode: _DebugCall and its subclasses.
|
||||
"""
|
||||
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
# Import _utils module for mutable globals
|
||||
from torch.utils._debug_mode import _utils
|
||||
from torch.utils._debug_mode._utils import (
|
||||
_arg_to_str,
|
||||
_get_op_name,
|
||||
_get_stack_trace,
|
||||
_maybe_get_autograd_trace,
|
||||
REDISTRIBUTE_FUNC,
|
||||
TensorIdTracker,
|
||||
)
|
||||
from torch.utils._pytree import tree_all, tree_map
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._dynamo.device_interface import DeviceInterface
|
||||
|
||||
|
||||
class _DebugCall:
|
||||
"""Base class for tracking operator calls in DebugMode"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
call_depth: int,
|
||||
record: dict[str, Any] | None = None,
|
||||
log: dict[str, Any] | None = None,
|
||||
stack: bool = False,
|
||||
) -> None:
|
||||
self.call_depth = call_depth
|
||||
if stack:
|
||||
self.stack_trace = _get_stack_trace()
|
||||
self.fwd_stack_trace = _maybe_get_autograd_trace()
|
||||
|
||||
# results from dispatch hooks
|
||||
self.record = record
|
||||
self.log = log
|
||||
self.output_str: str | None = None
|
||||
|
||||
def stringify_args(
|
||||
self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
|
||||
) -> None:
|
||||
"""
|
||||
To reduce memory consumption, this method stringifies args/kwargs, stores the result, and deletes original args/kwargs.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Subclasses must implement stringify_args(), even if no-op"
|
||||
)
|
||||
|
||||
def stringify_output(
|
||||
self,
|
||||
output: Any,
|
||||
attributes: list[str],
|
||||
tensor_memo: TensorIdTracker | None = None,
|
||||
) -> None:
|
||||
"""Store stringified version of call output in self.output_str"""
|
||||
if tree_all(lambda x: x is None, output):
|
||||
return
|
||||
output_str = tree_map(lambda x: _arg_to_str(x, attributes, tensor_memo), output)
|
||||
self.output_str = f" -> {str(output_str)}"
|
||||
|
||||
def render(self, attributes: list[str]) -> str:
|
||||
raise NotImplementedError("Subclasses must implement string render()")
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.render([])
|
||||
|
||||
|
||||
class _OpCall(_DebugCall):
|
||||
"""Normal operator call"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
op,
|
||||
args: tuple,
|
||||
kwargs: dict,
|
||||
call_depth: int,
|
||||
stack: bool = False,
|
||||
) -> None:
|
||||
super().__init__(call_depth, stack=stack)
|
||||
self.op = op
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
self.args_str: str | None = None
|
||||
self.kwargs_str: str | None = None
|
||||
|
||||
def stringify_args(
|
||||
self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
|
||||
) -> None:
|
||||
self.args_str = ", ".join(
|
||||
_arg_to_str(arg, attributes, tensor_memo) for arg in self.args
|
||||
)
|
||||
if self.kwargs:
|
||||
self.kwargs_str = ", " + ", ".join(
|
||||
f"{k}={_arg_to_str(v, attributes, tensor_memo)}"
|
||||
for k, v in self.kwargs.items()
|
||||
)
|
||||
else:
|
||||
self.kwargs_str = ""
|
||||
del self.args
|
||||
del self.kwargs
|
||||
|
||||
def render(self, attributes: list[str]) -> str:
|
||||
if self.args_str is not None:
|
||||
args_str = self.args_str
|
||||
else:
|
||||
args_str = ", ".join(_arg_to_str(arg, attributes) for arg in self.args)
|
||||
|
||||
if self.kwargs_str is not None:
|
||||
kwargs_str = self.kwargs_str
|
||||
else:
|
||||
if self.kwargs:
|
||||
kwargs_str = ", " + ", ".join(
|
||||
f"{k}={_arg_to_str(v, attributes)}" for k, v in self.kwargs.items()
|
||||
)
|
||||
else:
|
||||
kwargs_str = ""
|
||||
|
||||
if isinstance(self.op, torch._ops.OpOverload):
|
||||
op_name = self.op.__qualname__
|
||||
elif hasattr(self.op, "__module__") and hasattr(self.op, "__name__"):
|
||||
op_name = f"{self.op.__module__}.{self.op.__name__}"
|
||||
else:
|
||||
op_name = str(self.op)
|
||||
|
||||
base_str = f"{op_name}({args_str}{kwargs_str})"
|
||||
|
||||
if self.output_str:
|
||||
base_str += self.output_str
|
||||
if self.log:
|
||||
base_str += f" # {self.log}"
|
||||
return base_str
|
||||
|
||||
def __iter__(self):
|
||||
# for BC; tuple(self) returns (op, args, kwargs, call_depth)
|
||||
if self.args_str is not None:
|
||||
yield from [self.op, self.args_str, self.kwargs_str, self.call_depth]
|
||||
else:
|
||||
yield from [self.op, self.args, self.kwargs, self.call_depth]
|
||||
|
||||
|
||||
class _RedistributeCall(_DebugCall):
|
||||
def __init__(
|
||||
self,
|
||||
arg,
|
||||
src_placement,
|
||||
dst_placement,
|
||||
transform_info_str,
|
||||
call_depth,
|
||||
stack=False,
|
||||
is_explicit=False,
|
||||
) -> None:
|
||||
super().__init__(call_depth, stack=stack)
|
||||
self.arg = arg
|
||||
self.src_placement = src_placement
|
||||
self.dst_placement = dst_placement
|
||||
self.transform_info_str = transform_info_str
|
||||
self.is_explicit = is_explicit
|
||||
self.is_outer_call = isinstance(arg, int)
|
||||
|
||||
self.arg_str: str | None = None
|
||||
|
||||
def stringify_args(
|
||||
self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
|
||||
) -> None:
|
||||
self.arg_str = f"{_arg_to_str(self.arg, attributes, tensor_memo)}"
|
||||
del self.arg
|
||||
|
||||
def render(self, attributes: list[str]) -> str:
|
||||
if self.arg_str is not None:
|
||||
arg_str = self.arg_str
|
||||
else:
|
||||
arg_str = f"{_arg_to_str(self.arg, attributes)}"
|
||||
|
||||
if self.transform_info_str is not None: # prioritize over src/dst placements
|
||||
placement_str = f"trace: {self.transform_info_str}"
|
||||
else:
|
||||
src_placement_str = _arg_to_str(self.src_placement, attributes)
|
||||
dst_placement_str = _arg_to_str(self.dst_placement, attributes)
|
||||
placement_str = f"{src_placement_str} -> {dst_placement_str}"
|
||||
|
||||
# DebugMode will add redistribute_input logs at 2 levels,
|
||||
# once per redistribute decision, and once per redistributed input.
|
||||
# We only annotate [implicit/explicit] logs on the former (outer-level call).
|
||||
if self.is_outer_call:
|
||||
annotation = " [implicit] "
|
||||
elif self.is_explicit:
|
||||
annotation = " [explicit] "
|
||||
else:
|
||||
annotation = ""
|
||||
|
||||
base_str = f"{REDISTRIBUTE_FUNC}{annotation}({arg_str}, {placement_str})"
|
||||
if self.output_str:
|
||||
base_str += self.output_str
|
||||
return base_str
|
||||
|
||||
def __iter__(self):
|
||||
# for BC; tuple(self) returns (op, placement info, kwargs, call_depth)
|
||||
if self.arg_str is not None:
|
||||
arg = self.arg_str
|
||||
else:
|
||||
arg = self.arg
|
||||
|
||||
yield REDISTRIBUTE_FUNC
|
||||
if self.transform_info_str:
|
||||
yield [arg, self.transform_info_str]
|
||||
else:
|
||||
yield [arg, self.src_placement, self.dst_placement]
|
||||
yield {}
|
||||
yield self.call_depth
|
||||
|
||||
|
||||
class _OutputPlacementCall(_DebugCall):
|
||||
"""Records output placement for a DTensor op."""
|
||||
|
||||
def __init__(self, placements_str: str, call_depth: int) -> None:
|
||||
super().__init__(call_depth)
|
||||
self.placements_str = placements_str
|
||||
|
||||
def stringify_args(
|
||||
self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
|
||||
) -> None:
|
||||
pass # Already stringified
|
||||
|
||||
def render(self, attributes: list[str]) -> str:
|
||||
return f"-> output: {self.placements_str}"
|
||||
|
||||
|
||||
class _TritonKernelCall(_DebugCall):
|
||||
"""Triton kernel call from Inductor"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_name: str,
|
||||
kwargs: dict[str, Any],
|
||||
call_depth: int,
|
||||
):
|
||||
super().__init__(call_depth)
|
||||
self.kernel_name = kernel_name
|
||||
self.kwargs = kwargs
|
||||
self.kwargs_str: str | None = None
|
||||
|
||||
self.pre_hashes: dict[str, Any] | None = None
|
||||
self.post_hashes: dict[str, Any] | None = None
|
||||
|
||||
def stringify_args(
|
||||
self, attributes: list[str], tensor_memo: TensorIdTracker | None = None
|
||||
) -> None:
|
||||
# Optionally hash kernel inputs before launch
|
||||
if hash_fn := _utils._TRITON_INPUT_HASH_FN:
|
||||
self.pre_hashes = {
|
||||
k: hash_fn(v)
|
||||
for k, v in self.kwargs.items()
|
||||
if isinstance(v, torch.Tensor)
|
||||
}
|
||||
|
||||
if self.kwargs:
|
||||
self.kwargs_str = ", ".join(
|
||||
f"{k}={_arg_to_str(v, attributes, tensor_memo)}"
|
||||
for k, v in self.kwargs.items()
|
||||
)
|
||||
else:
|
||||
self.kwargs_str = ""
|
||||
|
||||
def render(self, attributes: list[str]) -> str:
|
||||
base_str = f"[triton] {self.kernel_name}({self.kwargs_str})"
|
||||
if self.pre_hashes:
|
||||
pre_hashes_str = ", ".join(f"{k}: {v}" for k, v in self.pre_hashes.items())
|
||||
pre_hashes_str = (
|
||||
"\n "
|
||||
+ " " * self.call_depth
|
||||
+ f"# pre-kernel hashes: {{{pre_hashes_str}}}"
|
||||
)
|
||||
else:
|
||||
pre_hashes_str = ""
|
||||
if self.post_hashes:
|
||||
post_hashes_str = ", ".join(
|
||||
f"{k}: {v}" for k, v in self.post_hashes.items()
|
||||
)
|
||||
post_hashes_str = (
|
||||
"\n "
|
||||
+ " " * self.call_depth
|
||||
+ f"# post-kernel hashes: {{{post_hashes_str}}}"
|
||||
)
|
||||
else:
|
||||
post_hashes_str = ""
|
||||
return f"{base_str}{pre_hashes_str}{post_hashes_str}\n"
|
||||
|
||||
def finalize(self, device_interface: "DeviceInterface"):
|
||||
# synchronize -> hash/store kernel results
|
||||
device_interface.synchronize(device_interface.current_device())
|
||||
if _utils._RECORD_TRITON_OUTPUTS:
|
||||
self.record = {
|
||||
"output": {
|
||||
k: v.clone() if isinstance(v, torch.Tensor) else v
|
||||
for k, v in self.kwargs.items()
|
||||
}
|
||||
}
|
||||
if hash_fn := _utils._TRITON_OUTPUT_HASH_FN:
|
||||
self.post_hashes = {
|
||||
k: hash_fn(v)
|
||||
for k, v in self.kwargs.items()
|
||||
if isinstance(v, torch.Tensor)
|
||||
}
|
||||
|
||||
# don't store tensors
|
||||
del self.kwargs
|
||||
|
||||
def __iter__(self):
|
||||
yield from [self.kernel_name, (), self.kwargs_str, self.call_depth]
|
||||
|
||||
|
||||
class _AnnotateCall(_DebugCall):
|
||||
"""Custom annotation call"""
|
||||
|
||||
def __init__(
|
||||
self, tag: Any, header: str, call_depth: int, stack: bool = False
|
||||
) -> None:
|
||||
super().__init__(call_depth, stack=stack)
|
||||
self.tag = tag
|
||||
self.header = header
|
||||
|
||||
def render(self, attributes: list[str]) -> str:
|
||||
return f"[{self.header}] {self.tag}"
|
||||
|
||||
def __iter__(self):
|
||||
yield from [
|
||||
f"[{self.header}] {self.tag}",
|
||||
(),
|
||||
{},
|
||||
self.call_depth,
|
||||
]
|
||||
|
||||
|
||||
def _get_call_name(call: _DebugCall) -> str:
|
||||
"""String identifying _DebugCall (e.g. func, kernel, module name)"""
|
||||
if isinstance(call, _OpCall):
|
||||
return _get_op_name(call.op)
|
||||
elif isinstance(call, _TritonKernelCall):
|
||||
return call.kernel_name
|
||||
elif isinstance(call, _AnnotateCall):
|
||||
return f"[{call.header}] {call.tag}"
|
||||
elif isinstance(call, _RedistributeCall):
|
||||
return REDISTRIBUTE_FUNC
|
||||
else:
|
||||
return str(call)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,251 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
Utility functions for DebugMode: tensor formatting, hashing, stack traces, and hook runners.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import os
|
||||
import traceback
|
||||
import weakref
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch.fx.graph import _parse_stack_trace
|
||||
from torch.utils._dtype_abbrs import dtype_abbrs
|
||||
from torch.utils._pytree import tree_map
|
||||
from torch.utils._traceback import CapturedTraceback
|
||||
from torch.utils.weak import WeakIdRef
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils._debug_mode._calls import _DebugCall
|
||||
|
||||
|
||||
REDISTRIBUTE_FUNC = "redistribute_input"
|
||||
|
||||
# Tracks if we're in inductor benchmarking, and temporarily disables logging
|
||||
# (for ignoring autotuning kernel launches which don't affect the user-facing result)
|
||||
_IN_INDUCTOR_BENCHMARK: bool = False
|
||||
# For record_outputs, log_tensor_hashes hooks for triton kernels.
|
||||
# Stores kernel outputs in call.record["output"]
|
||||
_RECORD_TRITON_OUTPUTS: bool = False
|
||||
# Annotates kernel output hashes, and stores them in call.post_hashes
|
||||
_TRITON_OUTPUT_HASH_FN: Callable | None = None
|
||||
# Annotates kernel input hashes, and stores them in call.pre_hashes
|
||||
_TRITON_INPUT_HASH_FN: Callable | None = None
|
||||
|
||||
# registered dispatch call hooks
|
||||
_DISPATCH_RECORD_HOOKS: list[Callable] = []
|
||||
_DISPATCH_LOG_HOOKS: list[Callable] = []
|
||||
_DISPATCH_PRE_LOG_HOOKS: list[Callable] = []
|
||||
|
||||
|
||||
def _stringify_shape(shape) -> str:
|
||||
return f"[{', '.join([str(x) for x in shape])}]"
|
||||
|
||||
|
||||
def _stringify_device_mesh(mesh) -> str:
|
||||
return f"DM({', '.join([str(s) for s in mesh.shape])})"
|
||||
|
||||
|
||||
def _stringify_placement(placement) -> str:
|
||||
return f"[{', '.join([str(p) for p in placement])}]"
|
||||
|
||||
|
||||
def _stringify_attributes(tensor, attributes) -> str:
|
||||
pairs = {}
|
||||
for attr in attributes:
|
||||
if hasattr(tensor, attr):
|
||||
pairs[attr] = getattr(tensor, attr)
|
||||
if len(pairs) == 0:
|
||||
return ""
|
||||
return f"{{{', '.join([f'{k}={v}' for k, v in pairs.items()])}}}"
|
||||
|
||||
|
||||
def _stringify_dtensor_spec(spec) -> str:
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
|
||||
return DTensorSpec.format_shard_order_str(spec.placements, spec.shard_order)
|
||||
|
||||
|
||||
class TensorIdTracker:
|
||||
def __init__(self) -> None:
|
||||
self.tensor_memo: dict[WeakIdRef, int] = {}
|
||||
self.next_tensor_id = 0
|
||||
|
||||
def _id(self, tensor) -> int:
|
||||
with torch._C._DisablePythonDispatcher():
|
||||
o = WeakIdRef(tensor)
|
||||
|
||||
def del_memo() -> None:
|
||||
self.tensor_memo.pop(o, None)
|
||||
|
||||
weakref.finalize(tensor, del_memo)
|
||||
if o not in self.tensor_memo:
|
||||
self.tensor_memo[o] = self.next_tensor_id
|
||||
self.next_tensor_id += 1
|
||||
return self.tensor_memo[o]
|
||||
|
||||
|
||||
def _tensor_debug_string(tensor, attributes, tensor_memo=None) -> str:
|
||||
"""Convert tensor to debug string representation."""
|
||||
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
tensor_debug_str = f"{dtype_abbrs[tensor.dtype]}{_stringify_shape(tensor.shape)}{_stringify_attributes(tensor, attributes)}"
|
||||
id_str = f"${tensor_memo._id(tensor)}" if tensor_memo is not None else ""
|
||||
if isinstance(tensor, torch.distributed.tensor.DTensor):
|
||||
# omitted device mesh
|
||||
return f"dt{id_str}: {tensor_debug_str}| {_stringify_dtensor_spec(tensor._spec)}"
|
||||
elif isinstance(tensor, FakeTensor):
|
||||
return f"ft{id_str}: {tensor_debug_str}"
|
||||
else:
|
||||
return f"t{id_str}: {tensor_debug_str}"
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported tensor type: {type(tensor)}")
|
||||
|
||||
|
||||
def _arg_to_str(arg, attributes, tensor_memo=None) -> str:
|
||||
from torch.distributed.tensor._dtensor_spec import DTensorSpec
|
||||
|
||||
def to_str(x):
|
||||
if isinstance(x, torch.Tensor):
|
||||
return _tensor_debug_string(x, attributes, tensor_memo)
|
||||
elif isinstance(x, DTensorSpec):
|
||||
return _stringify_dtensor_spec(x)
|
||||
return x
|
||||
|
||||
arg = tree_map(to_str, arg)
|
||||
return str(arg)
|
||||
|
||||
|
||||
def norm_hash_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | float:
|
||||
"""
|
||||
from Observer. Computes a hash for a tensor by converting it to float (if needed), making it contiguous,
|
||||
replacing NaN/inf values with fixed numbers, and then computing the L1 norm in float64 or complex128.
|
||||
This is used to generate a deterministic summary value for tensor comparison.
|
||||
"""
|
||||
with torch._C._DisablePythonDispatcher():
|
||||
if not (t.is_floating_point() or t.is_complex()):
|
||||
t = t.float()
|
||||
t = t.contiguous()
|
||||
|
||||
if t.is_complex():
|
||||
t_float = t.to(dtype=torch.complex128)
|
||||
else:
|
||||
t_float = t.to(dtype=torch.float64)
|
||||
|
||||
out = t_float.norm(p=1)
|
||||
if use_scalar:
|
||||
return out.item()
|
||||
return out
|
||||
|
||||
|
||||
def _compute_rel_diff(hash1, hash2):
|
||||
# Relative difference: |hash1 - hash2| / max(|hash1|, |hash2|, eps)
|
||||
numerator = abs(hash1 - hash2)
|
||||
denominator = max(abs(hash1), abs(hash2), 1e-10)
|
||||
return numerator / denominator
|
||||
|
||||
|
||||
def hash_tensor_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | int:
|
||||
"""
|
||||
wrapper over torch.hash_tensor
|
||||
"""
|
||||
if isinstance(t, torch.distributed.tensor.DTensor):
|
||||
t = t.to_local()
|
||||
|
||||
if t.is_floating_point():
|
||||
t_clean = t.to(dtype=torch.float64)
|
||||
elif t.is_complex():
|
||||
t_clean = t.to(dtype=torch.complex128).view(torch.float64)
|
||||
else:
|
||||
t_clean = t.to(dtype=torch.int64)
|
||||
|
||||
if t.numel() > 0:
|
||||
out = torch.hash_tensor(t_clean)
|
||||
else:
|
||||
out = torch.zeros((), device=t_clean.device, dtype=torch.uint64)
|
||||
|
||||
if use_scalar:
|
||||
return out.item() # type: ignore[attribute]
|
||||
return out
|
||||
|
||||
|
||||
def _get_stack_trace() -> str:
|
||||
from torch.fx.experimental.symbolic_shapes import uninteresting_files
|
||||
|
||||
summary = CapturedTraceback.extract().summary()
|
||||
summary = summary[:-4] # filter out DebugMode frames
|
||||
summary = [
|
||||
frame for frame in summary if frame.filename not in uninteresting_files()
|
||||
]
|
||||
summary = traceback.StackSummary.from_list(summary)
|
||||
return "".join(summary.format())
|
||||
|
||||
|
||||
def _get_user_stack_trace(stack_trace_str: str) -> str | None:
|
||||
# Extract user code stack trace, filtering out torch internals.
|
||||
torch_dir = os.path.dirname(inspect.getfile(torch))
|
||||
filter_fn = lambda file, name, code: not file.startswith(torch_dir + os.path.sep) # noqa: E731
|
||||
trace = _parse_stack_trace(stack_trace_str, filter_fn=filter_fn)
|
||||
if trace:
|
||||
return f"File: {trace.file}:{trace.lineno} in {trace.name}, code: {trace.code}"
|
||||
return None
|
||||
|
||||
|
||||
def _maybe_get_autograd_trace() -> str | None:
|
||||
if torch._C._current_autograd_node() is not None:
|
||||
tb = torch._C._current_autograd_node().metadata.get("traceback_") # type: ignore[attr-defined]
|
||||
if tb:
|
||||
return "".join(tb)
|
||||
return None
|
||||
|
||||
|
||||
def _get_op_name(op) -> str:
|
||||
if isinstance(op, torch._ops.OpOverload):
|
||||
op_name = op.__qualname__
|
||||
elif hasattr(op, "__module__") and hasattr(op, "__name__"):
|
||||
op_name = f"{op.__module__}.{op.__name__}"
|
||||
else:
|
||||
op_name = str(op)
|
||||
return op_name
|
||||
|
||||
|
||||
def _run_hook(hook, *args):
|
||||
out = hook(*args)
|
||||
if out is not None and not isinstance(out, dict):
|
||||
raise AssertionError(f"hook must return None or dict, got {type(out).__name__}")
|
||||
return out
|
||||
|
||||
|
||||
def _run_dispatch_pre_log_hooks(call: "_DebugCall", func, types, args, kwargs) -> None:
|
||||
if _DISPATCH_PRE_LOG_HOOKS:
|
||||
for hook in _DISPATCH_PRE_LOG_HOOKS:
|
||||
hook_out = _run_hook(hook, func, types, args, kwargs, call)
|
||||
if hook_out is not None:
|
||||
# Store pre-hook results in call.log
|
||||
if call.log is None:
|
||||
call.log = {}
|
||||
call.log.update(hook_out)
|
||||
|
||||
|
||||
def _run_dispatch_hooks(call: "_DebugCall", func, types, args, kwargs, result) -> None:
|
||||
if _DISPATCH_RECORD_HOOKS:
|
||||
record = {}
|
||||
for hook in _DISPATCH_RECORD_HOOKS:
|
||||
hook_out = _run_hook(hook, func, types, args, kwargs, result)
|
||||
if hook_out is not None:
|
||||
record.update(hook_out)
|
||||
if record:
|
||||
call.record = record
|
||||
|
||||
if _DISPATCH_LOG_HOOKS:
|
||||
# Preserve existing log from pre-hooks (e.g., input_hash)
|
||||
if call.log is None:
|
||||
call.log = {}
|
||||
for hook in _DISPATCH_LOG_HOOKS:
|
||||
hook_out = _run_hook(hook, func, types, args, kwargs, result)
|
||||
if hook_out is not None:
|
||||
call.log.update(hook_out)
|
||||
@@ -0,0 +1,137 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
|
||||
import torch
|
||||
from torch._C import _len_torch_function_stack
|
||||
from torch.overrides import _pop_mode, _push_mode, TorchFunctionMode
|
||||
from torch.utils._contextlib import context_decorator
|
||||
|
||||
|
||||
CURRENT_DEVICE: torch.device | None = None
|
||||
|
||||
|
||||
@functools.lru_cache(1)
|
||||
def _device_constructors():
|
||||
return {
|
||||
# standard ones
|
||||
torch.empty,
|
||||
torch.empty_permuted,
|
||||
torch.empty_strided,
|
||||
torch.empty_quantized,
|
||||
torch.ones,
|
||||
torch.arange,
|
||||
torch.bartlett_window,
|
||||
torch.blackman_window,
|
||||
torch.eye,
|
||||
torch.fft.fftfreq,
|
||||
torch.fft.rfftfreq,
|
||||
torch.full,
|
||||
torch.hamming_window,
|
||||
torch.hann_window,
|
||||
torch.kaiser_window,
|
||||
torch.linspace,
|
||||
torch.logspace,
|
||||
torch.nested.nested_tensor,
|
||||
# This function doesn't actually take a device argument
|
||||
# torch.normal,
|
||||
torch.rand,
|
||||
torch.randn,
|
||||
torch.randint,
|
||||
torch.randperm,
|
||||
torch.range,
|
||||
torch.sparse_coo_tensor,
|
||||
torch.sparse_compressed_tensor,
|
||||
torch.sparse_csr_tensor,
|
||||
torch.sparse_csc_tensor,
|
||||
torch.sparse_bsr_tensor,
|
||||
torch.sparse_bsc_tensor,
|
||||
torch.tril_indices,
|
||||
torch.triu_indices,
|
||||
torch.zeros,
|
||||
torch.asarray,
|
||||
# weird ones
|
||||
torch.tensor,
|
||||
torch.as_tensor,
|
||||
torch.scalar_tensor,
|
||||
# *_like may contain device kwarg, but the user implicitly
|
||||
# expects a specific device even when kwarg unused.
|
||||
# torch.zeros_like,
|
||||
# torch.randint_like,
|
||||
# torch.randn_like,
|
||||
# torch.ones_like,
|
||||
# torch.full_like,
|
||||
# torch.empty_like,
|
||||
}
|
||||
|
||||
|
||||
# NB: This is directly called from C++ in torch/csrc/Device.cpp
|
||||
class DeviceContext(TorchFunctionMode):
|
||||
def __init__(self, device) -> None:
|
||||
self.device = torch.device(device)
|
||||
self.prev_mode: DeviceContext | None = None
|
||||
|
||||
def __enter__(self):
|
||||
global CURRENT_DEVICE
|
||||
self.old_device = CURRENT_DEVICE
|
||||
CURRENT_DEVICE = self.device
|
||||
# We need to put the device at the bottom of the stack
|
||||
# If we set default device within a function mode context
|
||||
# exiting that context mode will pop the device function mode off
|
||||
# of the stack incorrectly
|
||||
cur_stack = [_pop_mode() for _ in range(_len_torch_function_stack())]
|
||||
|
||||
_push_mode(self)
|
||||
|
||||
for mode in reversed(cur_stack):
|
||||
if isinstance(mode, DeviceContext):
|
||||
self.prev_mode = mode
|
||||
else:
|
||||
_push_mode(mode)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
global CURRENT_DEVICE
|
||||
CURRENT_DEVICE = self.old_device
|
||||
cur_stack = []
|
||||
# Invariant: there should only be one DeviceContext on the stack at any time
|
||||
# (At the bottom), pop all modes until we hit the bottom, assert it's a DeviceContext
|
||||
# or else someone else has popped it!
|
||||
for _ in range(_len_torch_function_stack() - 1):
|
||||
mode = _pop_mode()
|
||||
if isinstance(mode, DeviceContext):
|
||||
raise AssertionError(
|
||||
"Found nested DeviceContext on the mode stack where none expected"
|
||||
)
|
||||
cur_stack.append(mode)
|
||||
|
||||
if _len_torch_function_stack() > 0:
|
||||
mode = _pop_mode()
|
||||
if not isinstance(mode, DeviceContext):
|
||||
raise AssertionError(
|
||||
"Expected a DeviceContext at the bottom of the mode stack"
|
||||
)
|
||||
if self.prev_mode is not None:
|
||||
_push_mode(self.prev_mode)
|
||||
|
||||
for mode in reversed(cur_stack):
|
||||
_push_mode(mode)
|
||||
|
||||
def __torch_function__(self, func, types, args=(), kwargs=None):
|
||||
kwargs = kwargs or {}
|
||||
if func in _device_constructors() and kwargs.get("device") is None:
|
||||
kwargs["device"] = self.device
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
# NB: This is directly called from C++ in torch/csrc/Device.cpp
|
||||
def device_decorator(device, func):
|
||||
return context_decorator(lambda: device, func)
|
||||
|
||||
|
||||
def set_device(device):
|
||||
"""
|
||||
Set the default device inside of the wrapped function by decorating it with this function.
|
||||
|
||||
If you would like to use this as a context manager, use device as a
|
||||
context manager directly, e.g., ``with torch.device(device)``.
|
||||
"""
|
||||
return lambda func: device_decorator(torch.device(device), func)
|
||||
@@ -0,0 +1,30 @@
|
||||
import torch
|
||||
|
||||
|
||||
# Used for testing and logging
|
||||
dtype_abbrs = {
|
||||
torch.bfloat16: "bf16",
|
||||
torch.float64: "f64",
|
||||
torch.float32: "f32",
|
||||
torch.float16: "f16",
|
||||
torch.float8_e4m3fn: "f8e4m3fn",
|
||||
torch.float8_e5m2: "f8e5m2",
|
||||
torch.float8_e4m3fnuz: "f8e4m3fnuz",
|
||||
torch.float8_e5m2fnuz: "f8e5m2fnuz",
|
||||
torch.float8_e8m0fnu: "f8e8m0fnu",
|
||||
torch.float4_e2m1fn_x2: "f4e2m1fnx2",
|
||||
torch.complex32: "c32",
|
||||
torch.complex64: "c64",
|
||||
torch.complex128: "c128",
|
||||
torch.int8: "i8",
|
||||
torch.int16: "i16",
|
||||
torch.int32: "i32",
|
||||
torch.int64: "i64",
|
||||
torch.bool: "b8",
|
||||
torch.uint8: "u8",
|
||||
torch.uint16: "u16",
|
||||
torch.uint32: "u32",
|
||||
torch.uint64: "u64",
|
||||
torch.bits16: "b16",
|
||||
torch.bits1x8: "b1x8",
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
from collections.abc import Callable
|
||||
from typing import TypeVar
|
||||
|
||||
|
||||
F = TypeVar("F")
|
||||
|
||||
|
||||
# Allows one to expose an API in a private submodule publicly as per the definition
|
||||
# in PyTorch's public api policy.
|
||||
#
|
||||
# It is a temporary solution while we figure out if it should be the long-term solution
|
||||
# or if we should amend PyTorch's public api policy. The concern is that this approach
|
||||
# may not be very robust because it's not clear what __module__ is used for.
|
||||
# However, both numpy and jax overwrite the __module__ attribute of their APIs
|
||||
# without problem, so it seems fine.
|
||||
def exposed_in(module: str) -> Callable[[F], F]:
|
||||
def wrapper(fn: F) -> F:
|
||||
fn.__module__ = module
|
||||
return fn
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,41 @@
|
||||
from types import TracebackType
|
||||
from typing_extensions import Self
|
||||
|
||||
from filelock import FileLock as base_FileLock
|
||||
|
||||
from torch.monitor import _WaitCounter
|
||||
|
||||
|
||||
class FileLock(base_FileLock):
|
||||
"""
|
||||
This behaves like a normal file lock.
|
||||
|
||||
However, it adds waitcounters for acquiring and releasing the filelock
|
||||
as well as for the critical region within it.
|
||||
|
||||
pytorch.filelock.enter - While we're acquiring the filelock.
|
||||
pytorch.filelock.region - While we're holding the filelock and doing work.
|
||||
pytorch.filelock.exit - While we're releasing the filelock.
|
||||
"""
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
self.region_counter = _WaitCounter("pytorch.filelock.region").guard()
|
||||
with _WaitCounter("pytorch.filelock.enter").guard():
|
||||
result = super().__enter__()
|
||||
self.region_counter.__enter__()
|
||||
return result
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
self.region_counter.__exit__()
|
||||
with _WaitCounter("pytorch.filelock.exit").guard():
|
||||
# Returns nothing per
|
||||
# https://github.com/tox-dev/filelock/blob/57f488ff8fdc2193572efe102408fb63cfefe4e4/src/filelock/_api.py#L379
|
||||
super().__exit__(exc_type, exc_value, traceback)
|
||||
# Returns nothing per
|
||||
# https://github.com/pytorch/pytorch/blob/0f6bfc58a2cfb7a5c052bea618ab62becaf5c912/torch/csrc/monitor/python_init.cpp#L315
|
||||
return None
|
||||
@@ -0,0 +1,60 @@
|
||||
from typing import TypeAlias
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.autograd.grad_mode import no_grad
|
||||
|
||||
|
||||
def _get_foreach_kernels_supported_devices() -> list[str]:
|
||||
r"""Return the device type list that supports foreach kernels."""
|
||||
return ["cuda", "xpu", "mtia", torch._C._get_privateuse1_backend_name()]
|
||||
|
||||
|
||||
def _get_fused_kernels_supported_devices() -> list[str]:
|
||||
r"""Return the device type list that supports fused kernels in optimizer."""
|
||||
return [
|
||||
"mps",
|
||||
"cuda",
|
||||
"xpu",
|
||||
"hpu",
|
||||
"cpu",
|
||||
"mtia",
|
||||
torch._C._get_privateuse1_backend_name(),
|
||||
]
|
||||
|
||||
|
||||
TensorListList: TypeAlias = list[list[Tensor | None]]
|
||||
Indices: TypeAlias = list[int]
|
||||
_foreach_supported_types = [torch.Tensor]
|
||||
|
||||
|
||||
# This util function splits tensors into groups by device and dtype, which is useful before sending
|
||||
# tensors off to a foreach implementation, which requires tensors to be on one device and dtype.
|
||||
# If tensorlistlist contains more than one tensorlist, the following assumptions are made BUT NOT verified:
|
||||
# - tensorlists CAN be None
|
||||
# - all tensors in the first specified list cannot be None
|
||||
# - given an index i, all specified tensorlist[i]s match in dtype and device
|
||||
# with_indices (bool, optional): whether to track previous indices as the last list per dictionary entry.
|
||||
# It comes in handy if there are Nones or literals in the tensorlists that are getting scattered out.
|
||||
# Whereas mutating a tensor in the resulting split-up tensorlists WILL propagate changes back to the
|
||||
# original input tensorlists, changing up Nones/literals WILL NOT propagate, and manual propagation
|
||||
# may be necessary. Check out torch/optim/sgd.py for an example.
|
||||
@no_grad()
|
||||
def _group_tensors_by_device_and_dtype(
|
||||
tensorlistlist: TensorListList,
|
||||
with_indices: bool = False,
|
||||
) -> dict[tuple[torch.device, torch.dtype], tuple[TensorListList, Indices]]:
|
||||
return torch._C._group_tensors_by_device_and_dtype(tensorlistlist, with_indices)
|
||||
|
||||
|
||||
def _device_has_foreach_support(device: torch.device) -> bool:
|
||||
return (
|
||||
device.type in (_get_foreach_kernels_supported_devices() + ["cpu"])
|
||||
and not torch.jit.is_scripting()
|
||||
)
|
||||
|
||||
|
||||
def _has_foreach_support(tensors: list[Tensor], device: torch.device) -> bool:
|
||||
return _device_has_foreach_support(device) and all(
|
||||
t is None or type(t) in _foreach_supported_types for t in tensors
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from typing import Concatenate, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_T = TypeVar("_T")
|
||||
_C = TypeVar("_C")
|
||||
|
||||
# Sentinel used to indicate that cache lookup failed.
|
||||
_cache_sentinel = object()
|
||||
|
||||
|
||||
def cache_method(
|
||||
f: Callable[Concatenate[_C, _P], _T],
|
||||
) -> Callable[Concatenate[_C, _P], _T]:
|
||||
"""
|
||||
Like `@functools.cache` but for methods.
|
||||
|
||||
`@functools.cache` (and similarly `@functools.lru_cache`) shouldn't be used
|
||||
on methods because it caches `self`, keeping it alive
|
||||
forever. `@cache_method` ignores `self` so won't keep `self` alive (assuming
|
||||
no cycles with `self` in the parameters).
|
||||
|
||||
Footgun warning: This decorator completely ignores self's properties so only
|
||||
use it when you know that self is frozen or won't change in a meaningful
|
||||
way (such as the wrapped function being pure).
|
||||
"""
|
||||
cache_name = "_cache_method_" + f.__name__
|
||||
|
||||
@functools.wraps(f)
|
||||
def wrap(self: _C, *args: _P.args, **kwargs: _P.kwargs) -> _T:
|
||||
if kwargs:
|
||||
raise AssertionError("cache_method does not accept keyword arguments")
|
||||
if not (cache := getattr(self, cache_name, None)):
|
||||
cache = {}
|
||||
setattr(self, cache_name, cache)
|
||||
cached_value = cache.get(args, _cache_sentinel)
|
||||
if cached_value is not _cache_sentinel:
|
||||
return cached_value
|
||||
value = f(self, *args, **kwargs)
|
||||
cache[args] = value
|
||||
return value
|
||||
|
||||
return wrap
|
||||
@@ -0,0 +1,200 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def remove_triton_function_declaration(source_code: str) -> str:
|
||||
remove_head = re.sub(r"(\n.+\s\'\'\'\n)", "\n", source_code)
|
||||
remove_tail = re.sub(r"(\'\'\'\,.+)", "\n", remove_head)
|
||||
return remove_tail
|
||||
|
||||
|
||||
def remove_async_compile(source_code: str) -> str:
|
||||
remove_top_level = str.replace(source_code, "async_compile = AsyncCompile()", "")
|
||||
remove_compile = str.replace(remove_top_level, "async_compile.wait(globals())", "")
|
||||
remove_del = str.replace(remove_compile, "del async_compile", "")
|
||||
return remove_del
|
||||
|
||||
|
||||
def rename_kernels(source_code: str) -> str:
|
||||
pattern = r"(\w+)\s*=\s*async_compile\.triton\('triton_',\s"
|
||||
triton_kernel_decl = "def triton_"
|
||||
matches = [
|
||||
(match.end(), match.group(1))
|
||||
for match in re.finditer(pattern, source_code, re.DOTALL)
|
||||
]
|
||||
|
||||
# Starting from the last match to avoid issues with shifting indices after replacements
|
||||
for end_index, captured_string in reversed(matches):
|
||||
# Find the index of the next "B" after the current match
|
||||
index_of_B = source_code.find(triton_kernel_decl, end_index)
|
||||
if index_of_B != -1:
|
||||
# Replace the triton_kernel_decl with the captured string
|
||||
source_code = (
|
||||
source_code[:index_of_B]
|
||||
+ f"def {captured_string}"
|
||||
+ source_code[index_of_B + len(triton_kernel_decl) :]
|
||||
)
|
||||
else:
|
||||
# If triton_kernel_decl is not found after the current match, continue to the next
|
||||
continue
|
||||
|
||||
return source_code
|
||||
|
||||
|
||||
def merge_params(original_params: list[str], new_params: list[str]) -> list[str]:
|
||||
for idx in range(len(new_params)):
|
||||
if new_params[idx] == "T":
|
||||
new_params[idx] = original_params[idx]
|
||||
return new_params
|
||||
|
||||
|
||||
def add_launch_params(
|
||||
original: str, kernel_to_params: dict[str, tuple[str, str]]
|
||||
) -> str:
|
||||
# Regex to match the function call in the original string
|
||||
pattern = r"(\w+)\.run\((.*)\)"
|
||||
|
||||
def replace(match) -> str:
|
||||
# Extract parts from the regex match
|
||||
func_name = match.group(1)
|
||||
params = match.group(2)
|
||||
new_params, grid = kernel_to_params[func_name]
|
||||
new_params = merge_params(params.split(", "), new_params.split(", "))
|
||||
|
||||
# Format the new function call
|
||||
new_string = f"{func_name}[{grid}]({', '.join(new_params)})"
|
||||
return new_string
|
||||
|
||||
transformed = re.sub(pattern, replace, original)
|
||||
|
||||
remove_inductor_wrappers = re.sub(
|
||||
r"@triton_heuristics[^@]*@triton.jit",
|
||||
r"@triton.jit",
|
||||
transformed,
|
||||
flags=re.DOTALL,
|
||||
)
|
||||
|
||||
return remove_inductor_wrappers
|
||||
|
||||
|
||||
def process_file(
|
||||
input_filename: str, output_filename: str, auto_generate_params: bool = True
|
||||
) -> str:
|
||||
with open(input_filename) as file:
|
||||
source_code = file.read()
|
||||
|
||||
transformed_code = source_code
|
||||
if "def triton_(" in source_code:
|
||||
raise RuntimeError(
|
||||
"Need to run original Pytorch code generating kernels with TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1"
|
||||
)
|
||||
# transformed_code = rename_kernels(transformed_code)
|
||||
transformed_code = remove_triton_function_declaration(transformed_code)
|
||||
transformed_code = remove_async_compile(transformed_code)
|
||||
|
||||
launch_params_filename = f"{input_filename}.launch_params"
|
||||
|
||||
# Auto-generate launch_params if they don't exist and auto_generate_params is True
|
||||
if not os.path.exists(launch_params_filename) and auto_generate_params:
|
||||
print(f"Launch params file {launch_params_filename} not found. Generating...")
|
||||
try:
|
||||
# Set environment variable and run the input file
|
||||
env = os.environ.copy()
|
||||
env["TORCHINDUCTOR_DUMP_LAUNCH_PARAMS"] = "1"
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, input_filename],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=os.path.dirname(input_filename) or ".",
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f"Error running {input_filename}:")
|
||||
print(f"stdout: {result.stdout}")
|
||||
print(f"stderr: {result.stderr}")
|
||||
raise RuntimeError(
|
||||
f"Failed to generate launch params. Command failed with return code {result.returncode}"
|
||||
)
|
||||
|
||||
print(f"Successfully generated {launch_params_filename}")
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to generate launch params by running {input_filename}: {str(e)}"
|
||||
) from e
|
||||
|
||||
if not os.path.exists(launch_params_filename):
|
||||
raise RuntimeError(
|
||||
f"Missing {launch_params_filename}. Run `TORCHINDUCTOR_DUMP_LAUNCH_PARAMS=1 python {input_filename}` first."
|
||||
)
|
||||
|
||||
with open(launch_params_filename) as f:
|
||||
launch_params_meta = f.readlines()
|
||||
|
||||
split_params = [i.split("|") for i in launch_params_meta]
|
||||
kernel_args_grid = {a.strip(): (b.strip(), c.strip()) for a, b, c in split_params}
|
||||
transformed_code = add_launch_params(transformed_code, kernel_args_grid)
|
||||
|
||||
with open(output_filename, "w") as file:
|
||||
file.write(transformed_code)
|
||||
print(f"Successfully generated {output_filename}")
|
||||
return transformed_code
|
||||
|
||||
|
||||
def get_clean_triton(
|
||||
input_path: Path,
|
||||
output_path: Path = Path("triton_only_repro.py"),
|
||||
auto_generate_params: bool = True,
|
||||
):
|
||||
"""Run experiments and output results to file
|
||||
|
||||
Args:
|
||||
input_path (Path): Path to inductor generated output codede
|
||||
output_path (Path): Path to write out the new python file
|
||||
auto_generate_params (bool): Whether to automatically generate launch_params if missing
|
||||
"""
|
||||
return process_file(str(input_path), str(output_path), auto_generate_params)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
"""Sample usage:
|
||||
# Running sweep
|
||||
python _get_clean_triton.py output_code.py
|
||||
|
||||
# To disable auto-generation of launch params:
|
||||
python _get_clean_triton.py output_code.py --no-auto-generate
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Clean Inductor generated code to remove Inductor dependencies"
|
||||
)
|
||||
|
||||
# Add the arguments
|
||||
parser.add_argument(
|
||||
"input_path", type=Path, help="Path to inductor generated output code"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path",
|
||||
type=Path,
|
||||
default=Path("triton_only_repro.py"),
|
||||
help="Path to write out the clean triton output",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-auto-generate",
|
||||
action="store_true",
|
||||
help="Disable automatic generation of launch_params file",
|
||||
)
|
||||
|
||||
# Parse the arguments
|
||||
args = parser.parse_args()
|
||||
|
||||
# Call the function with parsed arguments
|
||||
result = get_clean_triton(
|
||||
args.input_path, args.output_path, not args.no_auto_generate
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
import functools
|
||||
|
||||
from torch.utils._triton import has_triton
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_helion_package() -> bool:
|
||||
try:
|
||||
import helion # type: ignore[import-untyped, import-not-found] # noqa: F401
|
||||
except ImportError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_helion() -> bool:
|
||||
return has_helion_package() and has_triton()
|
||||
@@ -0,0 +1,37 @@
|
||||
import functools
|
||||
import importlib.util
|
||||
from types import ModuleType
|
||||
|
||||
|
||||
def _check_module_exists(name: str) -> bool:
|
||||
r"""Returns if a top-level module with :attr:`name` exists *without**
|
||||
importing it. This is generally safer than try-catch block around a
|
||||
`import X`. It avoids third party libraries breaking assumptions of some of
|
||||
our tests, e.g., setting multiprocessing start method when imported
|
||||
(see librosa/#747, torchvision/#544).
|
||||
"""
|
||||
try:
|
||||
spec = importlib.util.find_spec(name)
|
||||
return spec is not None
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def dill_available() -> bool:
|
||||
return _check_module_exists("dill")
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def import_dill() -> ModuleType | None:
|
||||
if not dill_available():
|
||||
return None
|
||||
|
||||
import dill
|
||||
|
||||
# XXX: By default, dill writes the Pickler dispatch table to inject its
|
||||
# own logic there. This globally affects the behavior of the standard library
|
||||
# pickler for any user who transitively depends on this module!
|
||||
# Undo this extension to avoid altering the behavior of the pickler globally.
|
||||
dill.extend(use_dill=False)
|
||||
return dill
|
||||
@@ -0,0 +1,90 @@
|
||||
import inspect
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _signature_metadata(
|
||||
sig: inspect.Signature,
|
||||
) -> tuple[tuple[inspect.Parameter, ...], bool, int]:
|
||||
"""
|
||||
Returns tuple(sig.parameters.values()), if any has VAR_POSITIONAL or VAR_KEYWORD, and the max_positional
|
||||
"""
|
||||
params = tuple(sig.parameters.values())
|
||||
has_var_args = False
|
||||
max_positional = 0
|
||||
|
||||
for p in params:
|
||||
kind = p.kind
|
||||
if kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
|
||||
has_var_args = True
|
||||
if kind in (
|
||||
inspect.Parameter.POSITIONAL_ONLY,
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
):
|
||||
max_positional += 1
|
||||
|
||||
return params, has_var_args, max_positional
|
||||
|
||||
|
||||
def _fast_bind(
|
||||
sig: inspect.Signature, *args: Any, **kwargs: Any
|
||||
) -> inspect.BoundArguments:
|
||||
"""
|
||||
Fast path for inspect.Signature.bind() for signatures without
|
||||
VAR_POSITIONAL or VAR_KEYWORD parameters. Falls back to sig.bind()
|
||||
for signatures that contain *args or **kwargs.
|
||||
"""
|
||||
params, has_var_args, max_positional = _signature_metadata(sig)
|
||||
|
||||
# fallback for complex signatures
|
||||
if has_var_args:
|
||||
return sig.bind(*args, **kwargs)
|
||||
|
||||
len_args = len(args)
|
||||
|
||||
if len_args > max_positional:
|
||||
raise TypeError(
|
||||
f"Too many positional arguments: expected max {max_positional}, got {len_args}"
|
||||
)
|
||||
|
||||
arguments: dict[str, Any] = {}
|
||||
arg_i = 0
|
||||
|
||||
for p in params:
|
||||
name = p.name
|
||||
kind = p.kind
|
||||
|
||||
if kind is inspect.Parameter.POSITIONAL_ONLY:
|
||||
if name in kwargs:
|
||||
raise TypeError(
|
||||
f"Got some positional-only arguments passed as keyword arguments: '{name}'"
|
||||
)
|
||||
if arg_i < len_args:
|
||||
arguments[name] = args[arg_i]
|
||||
arg_i += 1
|
||||
elif p.default is inspect.Parameter.empty:
|
||||
raise TypeError(f"Missing required argument '{name}'")
|
||||
|
||||
elif kind is inspect.Parameter.POSITIONAL_OR_KEYWORD:
|
||||
if arg_i < len_args:
|
||||
if name in kwargs:
|
||||
raise TypeError(f"Multiple values for argument '{name}'")
|
||||
arguments[name] = args[arg_i]
|
||||
arg_i += 1
|
||||
elif name in kwargs:
|
||||
arguments[name] = kwargs[name]
|
||||
elif p.default is inspect.Parameter.empty:
|
||||
raise TypeError(f"Missing required argument '{name}'")
|
||||
|
||||
elif kind is inspect.Parameter.KEYWORD_ONLY:
|
||||
if name in kwargs:
|
||||
arguments[name] = kwargs[name]
|
||||
elif p.default is inspect.Parameter.empty:
|
||||
raise TypeError(f"Missing required argument '{name}'")
|
||||
|
||||
# disallow extra keyword arguments not in the signature
|
||||
# cause kwargs have been processed by sig.bind at the beginning
|
||||
for name in kwargs:
|
||||
if name not in sig.parameters:
|
||||
raise TypeError(f"Got an unexpected keyword argument '{name}'")
|
||||
|
||||
return inspect.BoundArguments(sig, arguments) # type: ignore[arg-type]
|
||||
@@ -0,0 +1,15 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import TypeVar
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
# returns if all are the same mode
|
||||
def all_same_mode(modes):
|
||||
return all(tuple(mode == modes[0] for mode in modes))
|
||||
|
||||
|
||||
no_dispatch = torch._C._DisableTorchDispatch
|
||||
@@ -0,0 +1,178 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import (
|
||||
Hashable,
|
||||
Iterable,
|
||||
Iterator,
|
||||
MutableSet,
|
||||
Reversible,
|
||||
Set as AbstractSet,
|
||||
)
|
||||
from typing import Any, cast, TypeVar
|
||||
|
||||
|
||||
T = TypeVar("T", bound=Hashable)
|
||||
T_co = TypeVar("T_co", bound=Hashable, covariant=True)
|
||||
|
||||
__all__ = ["OrderedSet"]
|
||||
|
||||
|
||||
class OrderedSet(MutableSet[T], Reversible[T]):
|
||||
"""
|
||||
Insertion ordered set, similar to OrderedDict.
|
||||
"""
|
||||
|
||||
__slots__ = ("_dict",)
|
||||
|
||||
def __init__(self, iterable: Iterable[T] | None = None) -> None:
|
||||
self._dict = dict.fromkeys(iterable, None) if iterable is not None else {}
|
||||
|
||||
@staticmethod
|
||||
def _from_dict(dict_inp: dict[T, None]) -> OrderedSet[T]:
|
||||
s: OrderedSet[T] = OrderedSet()
|
||||
s._dict = dict_inp
|
||||
return s
|
||||
|
||||
#
|
||||
# Required overridden abstract methods
|
||||
#
|
||||
def __contains__(self, elem: object) -> bool:
|
||||
return elem in self._dict
|
||||
|
||||
def __iter__(self) -> Iterator[T]:
|
||||
return iter(self._dict)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._dict)
|
||||
|
||||
def __reversed__(self) -> Iterator[T]:
|
||||
return reversed(self._dict)
|
||||
|
||||
def add(self, elem: T) -> None:
|
||||
self._dict[elem] = None
|
||||
|
||||
def discard(self, elem: T) -> None:
|
||||
self._dict.pop(elem, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
# overridden because MutableSet impl is slow
|
||||
self._dict.clear()
|
||||
|
||||
# Unimplemented set() methods in _collections_abc.MutableSet
|
||||
|
||||
@classmethod
|
||||
def _wrap_iter_in_set(cls, other: Any) -> Any:
|
||||
"""
|
||||
Wrap non-Set Iterables in OrderedSets
|
||||
|
||||
Some of the magic methods are more strict on input types than
|
||||
the public apis, so we need to wrap inputs in sets.
|
||||
"""
|
||||
|
||||
if not isinstance(other, AbstractSet) and isinstance(other, Iterable):
|
||||
return cls(other)
|
||||
else:
|
||||
return other
|
||||
|
||||
def pop(self) -> T:
|
||||
if not self:
|
||||
raise KeyError("pop from an empty set")
|
||||
return self._dict.popitem()[0]
|
||||
|
||||
def copy(self) -> OrderedSet[T]:
|
||||
return OrderedSet._from_dict(self._dict.copy())
|
||||
|
||||
def difference(self, *others: Iterable[T]) -> OrderedSet[T]:
|
||||
res = self.copy()
|
||||
res.difference_update(*others)
|
||||
return res
|
||||
|
||||
def difference_update(self, *others: Iterable[T]) -> None:
|
||||
for other in others:
|
||||
self -= other # type: ignore[arg-type]
|
||||
|
||||
def update(self, *others: Iterable[T]) -> None:
|
||||
for other in others:
|
||||
self |= other
|
||||
|
||||
def intersection(self, *others: Iterable[T]) -> OrderedSet[T]:
|
||||
res = self.copy()
|
||||
for other in others:
|
||||
if other is not self:
|
||||
res &= other # type: ignore[arg-type]
|
||||
return res
|
||||
|
||||
def intersection_update(self, *others: Iterable[T]) -> None:
|
||||
for other in others:
|
||||
self &= other # type: ignore[arg-type]
|
||||
|
||||
def issubset(self, other: Iterable[T]) -> bool:
|
||||
return self <= self._wrap_iter_in_set(other)
|
||||
|
||||
def issuperset(self, other: Iterable[T]) -> bool:
|
||||
return self >= self._wrap_iter_in_set(other)
|
||||
|
||||
def symmetric_difference(self, other: Iterable[T]) -> OrderedSet[T]:
|
||||
return self ^ other # type: ignore[operator]
|
||||
|
||||
def symmetric_difference_update(self, other: Iterable[T]) -> None:
|
||||
self ^= other # type: ignore[arg-type]
|
||||
|
||||
def union(self, *others: Iterable[T]) -> OrderedSet[T]:
|
||||
res = self.copy()
|
||||
for other in others:
|
||||
if other is self:
|
||||
continue
|
||||
res |= other
|
||||
return res
|
||||
|
||||
# Specify here for correct type inference, otherwise would
|
||||
# return AbstractSet[T]
|
||||
def __sub__(self, other: AbstractSet[T_co]) -> OrderedSet[T]:
|
||||
# following cpython set impl optimization
|
||||
if isinstance(other, OrderedSet) and (len(self) * 4) > len(other):
|
||||
out = self.copy()
|
||||
out -= other
|
||||
return out
|
||||
return cast(OrderedSet[T], super().__sub__(other))
|
||||
|
||||
def __ior__(self, other: Iterable[T]) -> OrderedSet[T]: # type: ignore[misc, override] # noqa: PYI034
|
||||
if isinstance(other, OrderedSet):
|
||||
self._dict.update(other._dict)
|
||||
return self
|
||||
return super().__ior__(other) # type: ignore[arg-type]
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if isinstance(other, OrderedSet):
|
||||
return self._dict == other._dict
|
||||
return super().__eq__(other)
|
||||
|
||||
def __ne__(self, other: object) -> bool:
|
||||
if isinstance(other, OrderedSet):
|
||||
return self._dict != other._dict
|
||||
return super().__ne__(other)
|
||||
|
||||
def __or__(self, other: AbstractSet[T_co]) -> OrderedSet[T]:
|
||||
return cast(OrderedSet[T], super().__or__(other))
|
||||
|
||||
def __and__(self, other: AbstractSet[T_co]) -> OrderedSet[T]:
|
||||
# MutableSet impl will iterate over other, iter over smaller of two sets
|
||||
if isinstance(other, OrderedSet) and len(self) < len(other):
|
||||
# pyrefly: ignore [unsupported-operation, bad-return]
|
||||
return other & self
|
||||
return cast(OrderedSet[T], super().__and__(other))
|
||||
|
||||
def __xor__(self, other: AbstractSet[T_co]) -> OrderedSet[T]:
|
||||
return cast(OrderedSet[T], super().__xor__(other))
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}({list(self)})"
|
||||
|
||||
def __getstate__(self) -> list[T]:
|
||||
return list(self._dict.keys())
|
||||
|
||||
def __setstate__(self, state: list[T]) -> None:
|
||||
self._dict = dict.fromkeys(state, None)
|
||||
|
||||
def __reduce__(self) -> tuple[type[OrderedSet[T]], tuple[list[T]]]:
|
||||
return (OrderedSet, (list(self),))
|
||||
@@ -0,0 +1,125 @@
|
||||
import functools
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_jax_package() -> bool:
|
||||
"""Check if JAX is installed."""
|
||||
try:
|
||||
import jax # noqa: F401 # type: ignore[import-not-found]
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_pallas_package() -> bool:
|
||||
"""Check if Pallas (JAX experimental) is available."""
|
||||
if not has_jax_package():
|
||||
return False
|
||||
try:
|
||||
from jax.experimental import ( # noqa: F401 # type: ignore[import-not-found]
|
||||
pallas as pl,
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_jax_version(fallback: tuple[int, int, int] = (0, 0, 0)) -> tuple[int, int, int]:
|
||||
"""Get JAX version as (major, minor, patch) tuple."""
|
||||
try:
|
||||
import jax # type: ignore[import-not-found]
|
||||
|
||||
version_parts = jax.__version__.split(".")
|
||||
major, minor, patch = (int(v) for v in version_parts[:3])
|
||||
return (major, minor, patch)
|
||||
except (ImportError, ValueError, AttributeError):
|
||||
return fallback
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_jax_cuda_backend() -> bool:
|
||||
"""Check if JAX has CUDA backend support with SM90+ (required by Mosaic GPU)."""
|
||||
if not has_jax_package():
|
||||
return False
|
||||
try:
|
||||
import jax # type: ignore[import-not-found]
|
||||
|
||||
# Check if CUDA backend is available
|
||||
devices = jax.devices("gpu")
|
||||
if len(devices) == 0:
|
||||
return False
|
||||
|
||||
# Mosaic GPU requires SM90+ (compute capability 9.0+)
|
||||
if torch.cuda.is_available():
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
if major < 9:
|
||||
return False
|
||||
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_jax_tpu_backend() -> bool:
|
||||
"""Check if JAX has TPU backend support."""
|
||||
if not has_jax_package():
|
||||
return False
|
||||
try:
|
||||
import jax # type: ignore[import-not-found]
|
||||
|
||||
# Check if TPU backend is available
|
||||
devices = jax.devices("tpu")
|
||||
return len(devices) > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_torch_tpu() -> bool:
|
||||
"""Check if torch_tpu is installed and available."""
|
||||
try:
|
||||
import torch_tpu.api # noqa: F401 # type: ignore[import]
|
||||
|
||||
# Verify hardware/runtime access
|
||||
torch_tpu.api.tpu_device()
|
||||
return True
|
||||
except (ImportError, RuntimeError):
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_cpu_pallas() -> bool:
|
||||
"""Checks for a full Pallas-on-CPU environment."""
|
||||
return has_pallas_package()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_cuda_pallas() -> bool:
|
||||
"""Checks for a full Pallas-on-CUDA environment."""
|
||||
return has_pallas_package() and torch.cuda.is_available() and has_jax_cuda_backend()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_tpu_pallas() -> bool:
|
||||
"""Checks for a full Pallas-on-TPU environment."""
|
||||
return has_pallas_package() and has_jax_tpu_backend() and has_torch_tpu()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_pallas() -> bool:
|
||||
"""
|
||||
Check if Pallas backend is fully available for use.
|
||||
|
||||
Requirements:
|
||||
- JAX package installed
|
||||
- Pallas (jax.experimental.pallas) available
|
||||
- A compatible backend (CUDA or TPU) is available in both PyTorch and JAX.
|
||||
"""
|
||||
return has_cpu_pallas() or has_cuda_pallas() or has_tpu_pallas()
|
||||
@@ -0,0 +1,965 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import warnings
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import cast, overload, Protocol, TYPE_CHECKING
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
import torch
|
||||
import torchgen
|
||||
import torchgen.model
|
||||
from torch._C import (
|
||||
_get_dispatch_stack_at,
|
||||
_len_torch_dispatch_stack,
|
||||
_pop_torch_dispatch_stack,
|
||||
_push_on_torch_dispatch_stack,
|
||||
DispatchKey,
|
||||
)
|
||||
from torch._C._dynamo.guards import set_is_in_mode_without_ignore_compile_internals
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# TODO: Limitations and things about enable_torch_dispatch_mode we should fix before exposing it:
|
||||
# - We need a better user-facing api for _DisableTorchDispatch that
|
||||
# is able to selectively disable __torch_dispatch__ of a particular class.
|
||||
# - It doesn't work with the tensor constructors (torch.tensor, torch.Tensor)
|
||||
# - Better name (see https://github.com/pytorch/pytorch/pull/63496#discussion_r694091694)
|
||||
|
||||
_is_in_torch_dispatch_mode = False
|
||||
_is_in_non_infra_torch_dispatch_mode = False
|
||||
# If inside any mode that has ignore_compile_internals() = False
|
||||
_is_in_any_mode_without_ignore_compile_internals = False
|
||||
|
||||
|
||||
def is_in_torch_dispatch_mode(include_infra_modes: bool = True) -> bool:
|
||||
return (
|
||||
_is_in_torch_dispatch_mode
|
||||
if include_infra_modes
|
||||
else _is_in_non_infra_torch_dispatch_mode
|
||||
)
|
||||
|
||||
|
||||
def is_in_any_mode_without_ignore_compile_internals() -> bool:
|
||||
return _is_in_any_mode_without_ignore_compile_internals
|
||||
|
||||
|
||||
def any_torch_dispatch_mode_on_stack() -> bool:
|
||||
stack_len = torch._C._len_torch_dispatch_stack()
|
||||
|
||||
for idx in range(stack_len):
|
||||
mode = _get_dispatch_stack_at(idx)
|
||||
|
||||
# Apply filters first
|
||||
if mode.is_infra_mode():
|
||||
continue
|
||||
|
||||
if mode.ignore_compile_internals():
|
||||
continue
|
||||
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class TorchDispatchMode:
|
||||
"""
|
||||
A ``TorchDispatchMode`` allows you to override the meaning of all
|
||||
``__torch_dispatch__`` overridable functions within a dynamic scope,
|
||||
without having to actually create a tensor subclass or manually
|
||||
monkey-patch functions in the PyTorch API. Some common situations
|
||||
where you should use a mode:
|
||||
|
||||
* You want to override the meaning of factory functions, or other
|
||||
functions that do not otherwise take a tensor as an argument
|
||||
(these cannot be overridden with tensor subclasses).
|
||||
|
||||
* You want to override the behavior of all functions without needing
|
||||
to wrap your inputs in tensor subclasses; e.g., if you are just
|
||||
interested in logging intermediate computations.
|
||||
|
||||
* You want to control the order of execution of various tensor
|
||||
subclasses explicitly, rather than implicitly via the return of
|
||||
``NotImplemented``.
|
||||
|
||||
Independent subclasses of :class:`TorchDispatchMode` are compositional:
|
||||
modes can be pushed onto a stack using ``with MyMode():``.
|
||||
When you call functions in the PyTorch API inside your
|
||||
``__torch_dispatch__`` implementation, by default, they will forward on to
|
||||
the next mode on the mode stack. If you want recursively call back into
|
||||
your current ``__torch_dispatch__`` implementation, either explicitly
|
||||
invoke ``self.__torch_dispatch__(...)``, or use the context manager
|
||||
``self`` to make PyTorch
|
||||
API self-referential (beware of infinite loops, in this case!)
|
||||
"""
|
||||
|
||||
# - When False, custom torch dispatch mode will error out explicitly when a hop
|
||||
# is called under the mode.
|
||||
# - When True, custom torch dispatch mode's __torch_dispatch__ will be triggered.
|
||||
# Mode authors can implement how the mode interacts with higher order operators.
|
||||
supports_higher_order_operators = False
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
if cls._should_skip_dynamo():
|
||||
if "__torch_dispatch__" in cls.__dict__:
|
||||
raw = cls.__dict__["__torch_dispatch__"]
|
||||
if not isinstance(raw, classmethod):
|
||||
cls.__torch_dispatch__ = torch._disable_dynamo(raw, recursive=True)
|
||||
|
||||
def __init__(self, _dispatch_key=None):
|
||||
if _dispatch_key is not None:
|
||||
if not isinstance(_dispatch_key, torch._C.DispatchKey):
|
||||
raise AssertionError("_dispatch_key must be a torch._C.DispatchKey")
|
||||
self.__dict__["_dispatch_key"] = _dispatch_key
|
||||
|
||||
self.old_dispatch_mode_flags: deque[bool] = deque()
|
||||
self.old_non_infra_dispatch_mode_flags: deque[bool] = deque()
|
||||
self.old_without_ignore_compile_internals_dispatch_mode_flags: deque[bool] = (
|
||||
deque()
|
||||
)
|
||||
|
||||
def _lazy_init_old_dispatch_mode_flags(self):
|
||||
if not hasattr(self, "old_dispatch_mode_flags"):
|
||||
self.old_dispatch_mode_flags: deque[bool] = deque() # type: ignore[no-redef]
|
||||
|
||||
if not hasattr(self, "old_non_infra_dispatch_mode_flags"):
|
||||
self.old_non_infra_dispatch_mode_flags: deque[bool] = deque() # type: ignore[no-redef]
|
||||
|
||||
if not hasattr(
|
||||
self, "old_without_ignore_compile_internals_dispatch_mode_flags"
|
||||
):
|
||||
self.old_without_ignore_compile_internals_dispatch_mode_flags: deque[ # type: ignore[no-redef]
|
||||
bool
|
||||
] = deque()
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
raise NotImplementedError
|
||||
|
||||
def __enter__(self):
|
||||
global _is_in_torch_dispatch_mode
|
||||
global _is_in_non_infra_torch_dispatch_mode
|
||||
global _is_in_any_mode_without_ignore_compile_internals
|
||||
|
||||
# Previously, there wasn't any state in this class' constructor
|
||||
# super calls were added to existing modes, but for any new modes
|
||||
# this will replicate the previous behavior of not strictly needing
|
||||
# to call super().__init__()
|
||||
self._lazy_init_old_dispatch_mode_flags()
|
||||
self.old_dispatch_mode_flags.append(_is_in_torch_dispatch_mode)
|
||||
_is_in_torch_dispatch_mode = True
|
||||
self.old_non_infra_dispatch_mode_flags.append(
|
||||
_is_in_non_infra_torch_dispatch_mode
|
||||
)
|
||||
_is_in_non_infra_torch_dispatch_mode = (
|
||||
_is_in_non_infra_torch_dispatch_mode or not self.is_infra_mode()
|
||||
)
|
||||
self.old_without_ignore_compile_internals_dispatch_mode_flags.append(
|
||||
_is_in_any_mode_without_ignore_compile_internals
|
||||
)
|
||||
_is_in_any_mode_without_ignore_compile_internals = (
|
||||
_is_in_any_mode_without_ignore_compile_internals
|
||||
or not self.ignore_compile_internals()
|
||||
)
|
||||
set_is_in_mode_without_ignore_compile_internals(
|
||||
_is_in_any_mode_without_ignore_compile_internals
|
||||
)
|
||||
_push_mode(self)
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
mb_dk_or_mode_key = self.__dict__.get("_dispatch_key", None)
|
||||
if mb_dk_or_mode_key is None:
|
||||
# Today, mode keys are not used at all in the per-dispatch-key-mode logic (for pre-dispatch)
|
||||
# We should probably revisit this.
|
||||
mb_dk_or_mode_key = self.__dict__.get("_mode_key", None)
|
||||
global _is_in_torch_dispatch_mode
|
||||
_is_in_torch_dispatch_mode = self.old_dispatch_mode_flags.pop()
|
||||
global _is_in_non_infra_torch_dispatch_mode
|
||||
_is_in_non_infra_torch_dispatch_mode = (
|
||||
self.old_non_infra_dispatch_mode_flags.pop()
|
||||
)
|
||||
global _is_in_any_mode_without_ignore_compile_internals
|
||||
_is_in_any_mode_without_ignore_compile_internals = (
|
||||
self.old_without_ignore_compile_internals_dispatch_mode_flags.pop()
|
||||
)
|
||||
set_is_in_mode_without_ignore_compile_internals(
|
||||
_is_in_any_mode_without_ignore_compile_internals
|
||||
)
|
||||
_pop_mode(mb_dk_or_mode_key)
|
||||
|
||||
@classmethod
|
||||
def push(cls, *args, **kwargs):
|
||||
warnings.warn(
|
||||
"`Mode.push()` is no longer necessary and can be replaced with just `with Mode()`",
|
||||
stacklevel=2,
|
||||
)
|
||||
instance = cls(*args, **kwargs)
|
||||
return instance
|
||||
|
||||
@classmethod
|
||||
def is_infra_mode(cls) -> bool:
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def _should_skip_dynamo(cls) -> bool:
|
||||
"""Skip Dynamo when the flag is set to True
|
||||
|
||||
This is temporary measure to rollout a feature
|
||||
that skips PT2 compilation inside __torch_dispatch__
|
||||
frames.
|
||||
|
||||
If this flag is off, we would expect following:
|
||||
|
||||
class YoloMode(TorchDispatchMode):
|
||||
@classmethod
|
||||
def _should_skip_dynamo(cls):
|
||||
return False
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
return torch.ops.aten.mul.Tensor(args[0], args[1])
|
||||
|
||||
x = torch.ones(5)
|
||||
with YoloMode():
|
||||
out = torch.compile(torch.add, backend=backend, fullgraph=True)(x, x)
|
||||
|
||||
# instead of recursively disabling, we are compiling into __torch_dispatch__
|
||||
assert len(backend.graphs) == 1
|
||||
"""
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def ignore_compile_internals(cls) -> bool:
|
||||
"""Ignore operators that are compiled via torch.compile.
|
||||
|
||||
If ``True``, then this TorchDispatchMode ignores operators that
|
||||
are optimized by :func:`torch.compile`. Mechanically, this involves
|
||||
turning off the TorchDispatchMode throughout the whole compilation process,
|
||||
and turning it back on for the runtime of the compiled artifact(s).
|
||||
For example,
|
||||
|
||||
@torch.compile
|
||||
def f(x):
|
||||
return x.sin().cos()
|
||||
|
||||
with LoggingMode():
|
||||
f(x)
|
||||
|
||||
The above example will not log anything if
|
||||
``LoggingMode.ignore_compile_internals()`` is True.
|
||||
torch.compile will fuse sin() and cos() into a single operation
|
||||
and this TorchDispatchMode will not be passed sin and cos.
|
||||
|
||||
If ``False`` (default), :func:`torch.compile` will respect
|
||||
the eager semantics of passing this TorchDispatchMode all
|
||||
operators that would have run during eager execution.
|
||||
The way this will usually happen is that :func:`torch.compile`
|
||||
will just fallback to eager-mode PyTorch.
|
||||
"""
|
||||
if cls.is_infra_mode():
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _get_current_dispatch_mode() -> TorchDispatchMode | None:
|
||||
"""
|
||||
Return the top user mode on the stack (the next one that would be
|
||||
executed) if there are any.
|
||||
"""
|
||||
stack_len = _len_torch_dispatch_stack()
|
||||
if stack_len > 0:
|
||||
return _get_dispatch_stack_at(stack_len - 1)
|
||||
return None
|
||||
|
||||
|
||||
def _detect_infra_mode(key):
|
||||
if key not in (
|
||||
torch._C._TorchDispatchModeKey.FUNCTIONAL,
|
||||
torch._C._TorchDispatchModeKey.PROXY,
|
||||
):
|
||||
raise AssertionError(
|
||||
f"key must be either FUNCTIONAL ({torch._C._TorchDispatchModeKey.FUNCTIONAL}) \
|
||||
or PROXY ({torch._C._TorchDispatchModeKey.PROXY}) _TorchDispatchModeKey, \
|
||||
got {key}"
|
||||
)
|
||||
from torch._ops import _get_dispatch_mode_pre_dispatch
|
||||
|
||||
pre_dispatch_mode = _get_dispatch_mode_pre_dispatch(key)
|
||||
post_dispatch_mode = torch._C._get_dispatch_mode(key)
|
||||
|
||||
if pre_dispatch_mode is not None and post_dispatch_mode is not None:
|
||||
raise AssertionError(
|
||||
"At most one of pre_dispatch_mode and post_dispatch_mode may be active"
|
||||
)
|
||||
|
||||
if pre_dispatch_mode is None:
|
||||
return post_dispatch_mode
|
||||
|
||||
return pre_dispatch_mode
|
||||
|
||||
|
||||
def _unset_infra_mode(key):
|
||||
from torch._ops import _get_dispatch_mode_pre_dispatch, unset_mode_pre_dispatch
|
||||
|
||||
pre_dispatch_mode = _get_dispatch_mode_pre_dispatch(key)
|
||||
post_dispatch_mode = torch._C._get_dispatch_mode(key)
|
||||
if pre_dispatch_mode and post_dispatch_mode:
|
||||
raise AssertionError(
|
||||
"Can't have active infra mode on both pre and post dispatch mode stack"
|
||||
)
|
||||
|
||||
if pre_dispatch_mode:
|
||||
mode = unset_mode_pre_dispatch(key)
|
||||
return mode
|
||||
if post_dispatch_mode:
|
||||
return torch._C._unset_dispatch_mode(key)
|
||||
|
||||
|
||||
def _disable_infra_mode(key):
|
||||
if key not in (
|
||||
torch._C._TorchDispatchModeKey.FUNCTIONAL,
|
||||
torch._C._TorchDispatchModeKey.PROXY,
|
||||
):
|
||||
raise AssertionError(
|
||||
"key must be either FUNCTIONAL or PROXY _TorchDispatchModeKey"
|
||||
)
|
||||
mode_unset = _unset_infra_mode(key)
|
||||
try:
|
||||
yield mode_unset
|
||||
finally:
|
||||
if mode_unset is not None:
|
||||
_push_mode(mode_unset)
|
||||
|
||||
|
||||
def _get_current_dispatch_mode_stack() -> list[TorchDispatchMode]:
|
||||
"""
|
||||
Returns the current stack of dispatch modes, with the most recent
|
||||
(i.e., the one that will be processed first) at the end of the
|
||||
list (standard stack convention).
|
||||
"""
|
||||
stack_len = _len_torch_dispatch_stack()
|
||||
return [_get_dispatch_stack_at(i) for i in range(stack_len)]
|
||||
|
||||
|
||||
def _push_mode(mode: TorchDispatchMode) -> None:
|
||||
k = mode._dispatch_key if hasattr(mode, "_dispatch_key") else None
|
||||
if k is not None and k != torch._C.DispatchKey.PreDispatch:
|
||||
raise AssertionError(
|
||||
"mode._dispatch_key must be None or DispatchKey.PreDispatch"
|
||||
)
|
||||
if k is None:
|
||||
_push_on_torch_dispatch_stack(mode)
|
||||
return
|
||||
|
||||
from torch._ops import _set_mode_pre_dispatch, get_cached_ops
|
||||
|
||||
# See Note [Not Caching Per-Dispatch-Key Mode Handlers]
|
||||
# Clear the cache of every op that has been used so far, for this particular key.
|
||||
ks = torch._C._functionality_to_backend_keys(k)
|
||||
for op in get_cached_ops():
|
||||
for key in ks:
|
||||
op._uncache_dispatch(key)
|
||||
_set_mode_pre_dispatch(mode)
|
||||
|
||||
|
||||
def _pop_mode(k: DispatchKey | torch._C._TorchDispatchModeKey | None = None):
|
||||
if k == torch._C.DispatchKey.PreDispatch: # type: ignore[attr-defined]
|
||||
from torch._ops import _pop_mode_from_pre_dispatch
|
||||
|
||||
return _pop_mode_from_pre_dispatch()
|
||||
|
||||
if k is None or isinstance(k, torch._C._TorchDispatchModeKey):
|
||||
return _pop_torch_dispatch_stack(k)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _pop_mode_temporarily(k: DispatchKey | None = None):
|
||||
old = _pop_mode(k)
|
||||
try:
|
||||
yield old
|
||||
finally:
|
||||
_push_mode(old)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _disable_current_modes():
|
||||
from torch._ops import (
|
||||
_len_torch_dispatch_stack_pre_dispatch,
|
||||
_pop_mode_from_pre_dispatch,
|
||||
)
|
||||
from torch._subclasses.functional_tensor import FunctionalTensorMode
|
||||
from torch._subclasses.schema_check_mode import SchemaCheckMode
|
||||
from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode
|
||||
|
||||
mode_len_pre_dispatch = _len_torch_dispatch_stack_pre_dispatch()
|
||||
old_pre_dispatch_modes = [
|
||||
_pop_mode_from_pre_dispatch() for _ in range(mode_len_pre_dispatch)
|
||||
]
|
||||
|
||||
has_proxy_mode_in_pre_dispatch = False
|
||||
has_functional_mode_in_pre_dispatch = False
|
||||
has_schema_check_mode_in_pre_dispatch = False
|
||||
|
||||
for i in old_pre_dispatch_modes:
|
||||
if isinstance(i, ProxyTorchDispatchMode):
|
||||
has_proxy_mode_in_pre_dispatch = True
|
||||
if isinstance(i, FunctionalTensorMode):
|
||||
has_functional_mode_in_pre_dispatch = True
|
||||
if isinstance(i, SchemaCheckMode):
|
||||
has_schema_check_mode_in_pre_dispatch = True
|
||||
|
||||
mode_len = _len_torch_dispatch_stack()
|
||||
old_modes = [_pop_mode() for _ in range(mode_len)]
|
||||
|
||||
for old in old_modes:
|
||||
if (
|
||||
isinstance(old, FunctionalTensorMode)
|
||||
and has_functional_mode_in_pre_dispatch
|
||||
):
|
||||
raise AssertionError(
|
||||
"Can't have FunctionalMode available both in PreDispatch and Python Key"
|
||||
)
|
||||
if isinstance(old, ProxyTorchDispatchMode) and has_proxy_mode_in_pre_dispatch:
|
||||
raise AssertionError(
|
||||
"Can't have ProxyTorchDispatchMode available both in PreDispatch and Python Key"
|
||||
)
|
||||
if isinstance(old, SchemaCheckMode) and has_schema_check_mode_in_pre_dispatch:
|
||||
raise AssertionError(
|
||||
"Can't have SchemaCheckMode available both in PreDispatch and Python Key"
|
||||
)
|
||||
|
||||
# Manually disable proxy and fake modes, if any are active
|
||||
try:
|
||||
yield old_pre_dispatch_modes + old_modes
|
||||
finally:
|
||||
for mode in reversed(old_modes):
|
||||
_push_mode(mode)
|
||||
for mode in reversed(old_pre_dispatch_modes):
|
||||
_push_mode(mode)
|
||||
|
||||
|
||||
class BaseTorchDispatchMode(TorchDispatchMode):
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
# Subtypes which have __tensor_flatten__ and __tensor_unflatten__.
|
||||
class TensorWithFlatten(Protocol):
|
||||
def __tensor_flatten__(self) -> tuple[Sequence[str], object]: ...
|
||||
|
||||
@staticmethod
|
||||
def __tensor_unflatten__(
|
||||
inner_tensors: int, flatten_spec: int, outer_size: int, outer_stride: int
|
||||
) -> torch.Tensor: ...
|
||||
|
||||
# It would be really nice to be able to say that the return of
|
||||
# is_traceable_wrapper_subclass() is Intersection[torch.Tensor,
|
||||
# TensorWithFlatten] - but that doesn't exist.
|
||||
|
||||
shape: torch._C.Size
|
||||
|
||||
@overload
|
||||
def stride(self, dim: None = None) -> tuple[int, ...]: ...
|
||||
|
||||
@overload
|
||||
def stride(self, dim: int) -> int: ...
|
||||
|
||||
@overload
|
||||
def size(self, dim: None = None) -> tuple[int, ...]: ...
|
||||
|
||||
@overload
|
||||
def size(self, dim: int) -> int: ...
|
||||
|
||||
def storage_offset(self) -> int: ...
|
||||
|
||||
def dim(self) -> int: ...
|
||||
|
||||
@overload
|
||||
def to(
|
||||
self,
|
||||
dtype: torch.types._dtype,
|
||||
non_blocking: bool = False,
|
||||
copy: bool = False,
|
||||
*,
|
||||
memory_format: torch.memory_format | None = None,
|
||||
) -> torch.Tensor: ...
|
||||
|
||||
@overload
|
||||
def to(
|
||||
self,
|
||||
device: torch._prims_common.DeviceLikeType | None = None,
|
||||
dtype: torch.types._dtype | None = None,
|
||||
non_blocking: bool = False,
|
||||
copy: bool = False,
|
||||
*,
|
||||
memory_format: torch.memory_format | None = None,
|
||||
) -> torch.Tensor: ...
|
||||
|
||||
@overload
|
||||
def to(
|
||||
self,
|
||||
other: torch.Tensor,
|
||||
non_blocking: bool = False,
|
||||
copy: bool = False,
|
||||
*,
|
||||
memory_format: torch.memory_format | None = None,
|
||||
) -> torch.Tensor: ...
|
||||
|
||||
|
||||
def is_traceable_wrapper_subclass(t: object) -> TypeIs[TensorWithFlatten]:
|
||||
"""
|
||||
Returns whether or not a tensor subclass that implements __torch_dispatch__
|
||||
is 'traceable' with torch.compile.
|
||||
In order for a tensor subclass to support TorchDispatchMode-style tracing in PT2,
|
||||
It must implement two magic methods: __tensor_flatten__ and __tensor_unflatten__.
|
||||
It is also expected to obey some restrictions around traceability and aliasing:
|
||||
* The subclass's __torch_dispatch__() implementation should desugar into pytorch
|
||||
dispatcher operations that can be traced into a graph.
|
||||
* The subclass should use return_and_correct_aliasing(). This is needed today to make
|
||||
sure that torch.compile does the right thing in a few cases around input mutation
|
||||
and output aliasing.
|
||||
|
||||
Expected magic method signatures:
|
||||
attrs, ctx = t.__tensor_flatten__()
|
||||
attrs: list of attribute name strings for inner tensors
|
||||
ctx: dict containing any other subclass-specific metadata needed for unflattening
|
||||
|
||||
t = MySubClass.__tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride)
|
||||
inner_tensors: dict mapping attribute name -> tensor for each inner tensor
|
||||
ctx: dict with subclass metadata in the form that __tensor_flatten__() produces
|
||||
outer_size: expected (possibly symbolic) size that the returned subclass
|
||||
instance should have. Note that this arg is useful for certain subclasses
|
||||
that require the shape info to be constructed. In most cases, this arg can be
|
||||
safely ignored.
|
||||
outer_stride: expected (possibly symbolic) stride that the returned subclass
|
||||
instance should have. Note that this arg is useful for certain subclasses
|
||||
that require the stride info to be constructed. In most cases, this arg can be
|
||||
safely ignored.
|
||||
"""
|
||||
is_subclass = isinstance(t, torch.Tensor) and type(t) is not torch.Tensor
|
||||
return (
|
||||
is_subclass
|
||||
and hasattr(t, "__tensor_flatten__")
|
||||
and hasattr(t, "__tensor_unflatten__")
|
||||
)
|
||||
|
||||
|
||||
def is_traceable_wrapper_subclass_type(t: type) -> TypeIs[type[TensorWithFlatten]]:
|
||||
"""Same as above, but takes a type argument instead of an instance."""
|
||||
return (
|
||||
issubclass(t, torch.Tensor)
|
||||
and t is not torch.Tensor
|
||||
and hasattr(t, "__tensor_flatten__")
|
||||
and hasattr(t, "__tensor_unflatten__")
|
||||
)
|
||||
|
||||
|
||||
def transform_subclass(t, callback, outer_size=None, outer_stride=None):
|
||||
"""
|
||||
Given a traceable, wrapper tensor subclass ``t`` that implements
|
||||
``__torch_dispatch__`` and holds some inner tensors,
|
||||
and a callback of type ``Callable[[str, torch.Tensor], torch.Tensor]``,
|
||||
`transform_subclass` will construct a fresh instance of the wrapper tensor subclass.
|
||||
It will do so by grabbing each inner tensor attribute from the wrapper,
|
||||
passing them into ``callback`` to get a transformed tensor,
|
||||
and putting each transformed tensor into the fresh tensor subclass instance.
|
||||
|
||||
Note: this function will not handle ensuring that the fresh subclass
|
||||
gets the same (autograd, and aliasing) metadata as the original tensor.
|
||||
This is generally handled in other subsystems like AOTAutograd.
|
||||
"""
|
||||
outer_size = outer_size if outer_size is not None else t.size()
|
||||
outer_stride = outer_stride if outer_stride is not None else t.stride()
|
||||
|
||||
attrs, ctx = t.__tensor_flatten__()
|
||||
transformed_tensors_dict = {}
|
||||
for attr in attrs:
|
||||
transformed_tensors_dict[attr] = callback(attr, getattr(t, attr))
|
||||
sub = type(t).__tensor_unflatten__(
|
||||
transformed_tensors_dict, ctx, outer_size, outer_stride
|
||||
)
|
||||
|
||||
# NB: Purposefully guard here to simplify the inner / outer symbols.
|
||||
# Using sym_eq() for symbolic comparison can result in an expression that's too
|
||||
# difficult to guard on, so we use == here.
|
||||
if sub.shape != outer_size:
|
||||
raise AssertionError(
|
||||
f"Expected return value from {type(t)}__tensor_unflatten__() to have "
|
||||
f"shape equal to {outer_size}, but got: {sub.shape}"
|
||||
)
|
||||
if sub.stride() != outer_stride:
|
||||
raise AssertionError(
|
||||
f"Expected return value from {type(t)}__tensor_unflatten__() to have "
|
||||
f"stride equal to {outer_stride}, but got: {sub.stride()}"
|
||||
)
|
||||
|
||||
return sub
|
||||
|
||||
|
||||
def _correct_storage_aliasing(func, schema_info, args, outs) -> None:
|
||||
"""
|
||||
Given: an OpOverload, a SchemaInfo (cached information from torchgen about schema),
|
||||
and the inputs/outputs to the OpOverload,
|
||||
this function checks to see if func is a view operator
|
||||
(by checking if any of the outputs in the op's schema
|
||||
are immutable aliases of inputs).
|
||||
If so, this function manually aliases the storage of the output tensor
|
||||
with its corresponding input tensor alias.
|
||||
It does this by unsafely overwriting the storage field of the output tensor
|
||||
to be the same storage as the input.
|
||||
"""
|
||||
if not isinstance(func, torch._ops.OpOverload):
|
||||
raise AssertionError(f"func must be an OpOverload, got {type(args)}")
|
||||
if not isinstance(args, tuple):
|
||||
raise AssertionError(f"args must be a tuple, got {type(args)}")
|
||||
if not isinstance(outs, (list, tuple)):
|
||||
raise AssertionError(f"outs must be a list or tuple, got {type(args)}")
|
||||
|
||||
def alias_non_inplace_storage(arg, ret) -> None:
|
||||
# This is hopefully a reasonable assert:
|
||||
# subclasses that rely on this API for output aliasing
|
||||
# should always return wrapper tensor subclasses for us to manually alias.
|
||||
# in theory if a subclass that needs this API wants to sometimes return
|
||||
# plain tensors, we could remove the assert and just not perform the aliasing,
|
||||
# but it seems safer to learn more about this case first.
|
||||
#
|
||||
# Performance note: This is all just to assert that the argument and result
|
||||
# types match, checking that is cheaper than is_traceable_wrapper_subclass_type,
|
||||
# and multiple returns are relatively unlikely, so just check up front!
|
||||
arg_type = type(arg)
|
||||
ret_type = type(ret)
|
||||
if arg_type is not ret_type and (
|
||||
is_traceable_wrapper_subclass_type(arg_type)
|
||||
or is_traceable_wrapper_subclass_type(ret_type)
|
||||
):
|
||||
ret_list = ret if isinstance(ret, list) else [ret]
|
||||
for r in ret_list:
|
||||
if type(arg) is not type(r):
|
||||
raise AssertionError(
|
||||
f"Called {str(func)} with input of type {type(arg)}\n"
|
||||
f"and output of type {type(ret)}. But expected types to match."
|
||||
)
|
||||
# Need to call a non-dispatcher helper, because we explicitly do **not**
|
||||
# want our subclass to intercept the set_() call.
|
||||
# instead, our subclass should directly have its storage swapped out.
|
||||
# we **explicitly** don't want to reset the sizes on ret, if the storage implies a size change.
|
||||
# Why?
|
||||
# The purpose of this API is *not* to change the size/strides of our output- we assume it's already correct.
|
||||
# We just want to "fix up" the storage aliasing, without modifying or output's metadata.
|
||||
# Example: out = inp.expand(inp.shape[0], inp.shape[0])
|
||||
# This requires swapping the storage of out to be the same as inp,
|
||||
# but we do *not* want it to change the sizes/strides that were compute for out.
|
||||
|
||||
if isinstance(ret, list):
|
||||
for r in ret:
|
||||
torch._functionalize_unsafe_set(r, arg)
|
||||
else:
|
||||
if not isinstance(ret, torch.Tensor):
|
||||
raise AssertionError(f"expected torch.Tensor, got {type(ret)}")
|
||||
torch._functionalize_unsafe_set(ret, arg)
|
||||
|
||||
for arg_idx, return_idx in schema_info.read_only_alias_match_indexes:
|
||||
alias_non_inplace_storage(args[arg_idx], outs[return_idx])
|
||||
|
||||
|
||||
def _get_write_alias(x) -> str | None:
|
||||
alias_set = x.alias_set
|
||||
if not alias_set or not x.is_write:
|
||||
return None
|
||||
# torchscript allows for complicated alias sets, but our dispatcher ops only really involve simple aliasing
|
||||
if len(alias_set) != 1:
|
||||
raise AssertionError("Expected alias_set to contain exactly one element")
|
||||
# timeit says next(iter(alias_set)) is faster than list(alias_set)[0] even for
|
||||
# set of size 1 on Python 3.13.
|
||||
return next(iter(alias_set))
|
||||
|
||||
|
||||
# This abstracts over the fact that in return_and_correct_aliasing,
|
||||
# we sometimes use torchgen schema parsing (for aten ops, since torchscript's schema parsing is sometimes buggy),
|
||||
# and sometimes use torchscript schema parsing (for custom ops, for which torchgen parsing is untested).
|
||||
@dataclass
|
||||
class AliasInfo:
|
||||
alias_set: set[str]
|
||||
is_write: bool
|
||||
name: str | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SchemaInfo:
|
||||
args: list[AliasInfo]
|
||||
outs: list[AliasInfo]
|
||||
|
||||
is_inplace_view_op: bool
|
||||
|
||||
# [_get_write_alias(x) for x in outs]. Guaranteed to contain no Nones; we coerce
|
||||
# all-Nones result to empty list instead, and we don't support
|
||||
# some-but-not-all-Nones.
|
||||
outs_write_aliases: list[str] | None
|
||||
|
||||
# List of (arg_idx, return_idx) where args[arg_idx].alias_set &
|
||||
# outs[out_idx].alias_set is not empty, and not args[arg_idx].is_write.
|
||||
read_only_alias_match_indexes: list[tuple[int, int]]
|
||||
|
||||
|
||||
# Given an OpOverload, returns schema information on it.
|
||||
# This is cached for efficiency, since it can involve running torchgen
|
||||
@functools.cache
|
||||
def get_alias_info(func) -> SchemaInfo:
|
||||
# For ATen ops: use torchgen (since torchscript parser doesn't handle alias annotations
|
||||
# properly for some ops that output tensorlists)
|
||||
if func.namespace == "aten":
|
||||
torchgen_schema_str = str(func._schema)
|
||||
if not torchgen_schema_str.startswith("aten::"):
|
||||
raise AssertionError(
|
||||
"Expected torchgen schema string to start with 'aten::'"
|
||||
)
|
||||
# remove the aten:: namespace, which is added by the torchscript parser,
|
||||
# and torchgen doesn't know how to handle
|
||||
torchgen_schema_str = torchgen_schema_str[6:]
|
||||
import re
|
||||
|
||||
# the torchscript parser ends up converting int[2]=1 into int[2]=[1, 1],
|
||||
# which torchgen chokes on.
|
||||
torchgen_schema_str = re.sub(r"=\[[0, ]+\]", "=0", torchgen_schema_str)
|
||||
torchgen_schema_str = re.sub(r"=\[[1, ]+\]", "=1", torchgen_schema_str)
|
||||
# for aten::rot90 / aten:fft_*
|
||||
torchgen_schema_str = re.sub(
|
||||
r"=\[(-?[0-9]+), (-?[0-9]+)\]", r"=[\1,\2]", torchgen_schema_str
|
||||
)
|
||||
torchgen_schema = torchgen.model.FunctionSchema.parse(torchgen_schema_str)
|
||||
arg_schemas = [
|
||||
AliasInfo(
|
||||
alias_set=(
|
||||
set() if a.annotation is None else set(a.annotation.alias_set)
|
||||
),
|
||||
is_write=a.annotation is not None and a.annotation.is_write,
|
||||
name=a.name,
|
||||
)
|
||||
for a in torchgen_schema.arguments.flat_all
|
||||
]
|
||||
out_schemas = [
|
||||
AliasInfo(
|
||||
alias_set=(
|
||||
set() if a.annotation is None else set(a.annotation.alias_set)
|
||||
),
|
||||
is_write=a.annotation is not None and a.annotation.is_write,
|
||||
name=a.name,
|
||||
)
|
||||
for a in torchgen_schema.returns
|
||||
]
|
||||
else:
|
||||
# For non-aten ops, torchgen is untested so we rely on torchscript schema parsing
|
||||
arg_schemas = [
|
||||
AliasInfo(
|
||||
alias_set=(
|
||||
set() if a.alias_info is None else set(a.alias_info.before_set)
|
||||
),
|
||||
is_write=a.alias_info is not None and a.alias_info.is_write,
|
||||
name=a.name,
|
||||
)
|
||||
for a in func._schema.arguments
|
||||
]
|
||||
out_schemas = [
|
||||
AliasInfo(
|
||||
alias_set=(
|
||||
set() if a.alias_info is None else set(a.alias_info.before_set)
|
||||
),
|
||||
is_write=a.alias_info is not None and a.alias_info.is_write,
|
||||
name=a.name,
|
||||
)
|
||||
for a in func._schema.returns
|
||||
]
|
||||
read_only_alias_match_indexes = []
|
||||
for arg_idx, schema_arg in enumerate(arg_schemas):
|
||||
for return_idx, schema_out in enumerate(out_schemas):
|
||||
is_read_only_alias_match = (
|
||||
schema_arg.alias_set & schema_out.alias_set
|
||||
) and not schema_arg.is_write
|
||||
if is_read_only_alias_match:
|
||||
read_only_alias_match_indexes.append((arg_idx, return_idx))
|
||||
|
||||
outs_write_aliases_list: list[str | None] = [
|
||||
_get_write_alias(r) for r in out_schemas
|
||||
]
|
||||
non_nones = sum(x is not None for x in outs_write_aliases_list)
|
||||
if non_nones == 0:
|
||||
outs_write_aliases: list[str] | None = None
|
||||
elif non_nones != len(outs_write_aliases_list):
|
||||
# simplifying assumption: we don't have **any** ops with return types like "-> (Tensor(a!), Tensor)"
|
||||
raise RuntimeError("Unsupported schema: " + str(func._schema))
|
||||
else:
|
||||
outs_write_aliases = cast(list[str], outs_write_aliases_list)
|
||||
|
||||
schema_info = SchemaInfo(
|
||||
args=arg_schemas,
|
||||
outs=out_schemas,
|
||||
# This check is surprisingly expensive because pybind11 enum_s are
|
||||
# inefficient. Just cache it.
|
||||
is_inplace_view_op=torch.Tag.inplace_view in func.tags,
|
||||
outs_write_aliases=outs_write_aliases,
|
||||
read_only_alias_match_indexes=read_only_alias_match_indexes,
|
||||
)
|
||||
return schema_info
|
||||
|
||||
|
||||
def autograd_would_have_decomposed(
|
||||
func: torch._ops.OpOverload, flat_args: Sequence[torch.Tensor | object]
|
||||
) -> bool:
|
||||
"""
|
||||
Suppose that an operator has CompositeImplicitAutograd decomp registered.
|
||||
Would autograd have used this decomposition? It will only use it if there
|
||||
isn't an explicit backend registration for the device as well. This function
|
||||
will tell if this would have occurred.
|
||||
|
||||
Why do we need to apply these decompositions later? When inference mode is
|
||||
on, the autograd key is bypassed entirely, so a lower level mode cannot rely
|
||||
on the decomposition have been applied. It's easy to accidentally never apply
|
||||
the decomposition, resulting in an operator showing up in a graph that
|
||||
is unexpected.
|
||||
|
||||
Why do we need to AVOID applying the decomposition when autograd wouldn't
|
||||
have decomposed? If autograd doesn't decompose, this means in eager mode
|
||||
we would have run the fused kernel. It must be possible to trace this
|
||||
fused kernel directly into the graph for fidelity with eager (NB: a user
|
||||
has the option of then further decomposing at proxy tensor mode via
|
||||
decomposition table, but we must preserve it to proxy mode to have the
|
||||
choice.)
|
||||
|
||||
Why does functionalization need to also perform the test here? This is
|
||||
because some CompositeImplicitAutograd decompositions are not functional.
|
||||
If we are eventually going to decompose, we need to do this while we can
|
||||
still turn functionalization back on, so those decompositions get functionalized.
|
||||
So an early decomposition in functionalization may still be necessary. Note that
|
||||
if proxy tensor decomposition process could turn functionalization back on, this
|
||||
wouldn't be necessary, and maybe that is a useful thing to do anyway because
|
||||
the decomposition table is user specified and a user could violate the functional
|
||||
decomp requirement with a bad decomp. If this happened, then you could always
|
||||
pass through functionalization.
|
||||
"""
|
||||
has_backend_registration = False
|
||||
for a in flat_args:
|
||||
if isinstance(a, torch.Tensor):
|
||||
backend_key = torch._C._parse_dispatch_key(
|
||||
torch._C._dispatch_key_for_device(a.device.type)
|
||||
)
|
||||
if backend_key is None:
|
||||
raise AssertionError(
|
||||
f"failed to parse dispatch key for device {a.device.type}"
|
||||
)
|
||||
# TODO: use func.has_kernel_for_dispatch_key(backend_key)
|
||||
# but this one checks py_impl and CompositeImplicitAutograd
|
||||
# incorrectly shows up as has backend reg here
|
||||
has_backend_registration = torch._C._dispatch_has_kernel_for_dispatch_key(
|
||||
func.name(), backend_key
|
||||
)
|
||||
|
||||
# in theory we should take all backend keys and take the highest priority one
|
||||
# to properly mimic the dispatcher,
|
||||
# this just grabs the first tensor and takes its device key
|
||||
break
|
||||
return not has_backend_registration
|
||||
|
||||
|
||||
def return_and_correct_aliasing(func, args, kwargs, out):
|
||||
"""
|
||||
This function should be used by wrapper tensor ``__torch_dispatch__`` subclasses
|
||||
that would like to work with torch.compile. It ensures that the subclass
|
||||
properly implements the aliasing behavior of every op,
|
||||
which is needed for correctness in AOTAutograd.
|
||||
This function will handle:
|
||||
|
||||
* When we see a view op, we will alias the storages of any
|
||||
input and output tensor subclasses
|
||||
|
||||
* When we see an inplace or out= op, we will directly
|
||||
return the corresponding input tensor, instead of returning
|
||||
a (potentially) fresh output tensor.
|
||||
"""
|
||||
|
||||
# Caching here because torchgen parsing is definitely not fast, and this function is called
|
||||
# once for every op in the graph during functionalization.
|
||||
schema_info = get_alias_info(func)
|
||||
|
||||
def get_arg_from_alias(output_alias, schema_info, args, kwargs):
|
||||
new_args, new_kwargs = torch.fx.operator_schemas.normalize_function( # type: ignore[misc]
|
||||
func, args=args, kwargs=kwargs
|
||||
)
|
||||
|
||||
arg_indices = [
|
||||
i for i, a in enumerate(schema_info.args) if output_alias in a.alias_set
|
||||
]
|
||||
# For any dispatcher op with an output alias, we expect it to map to exactly one alias in the schema's input arguments.
|
||||
if len(arg_indices) != 1:
|
||||
raise AssertionError(
|
||||
"Expected exactly one argument index for the given output alias"
|
||||
)
|
||||
idx = arg_indices[0]
|
||||
arg_info = schema_info.args[idx]
|
||||
if arg_info.name is not None and arg_info.name in new_kwargs:
|
||||
return new_kwargs[arg_info.name]
|
||||
return new_args[idx]
|
||||
|
||||
# Fix up the storages of any outs so that they point to the same storage as the input,
|
||||
# if func is a view op.
|
||||
_correct_storage_aliasing(
|
||||
func, schema_info, args, (out,) if not isinstance(out, tuple) else out
|
||||
)
|
||||
|
||||
# For inplace_view ops in particular, we'll try hard to make sure that the wrapper subclass's
|
||||
# metadata is set correctly.
|
||||
if schema_info.is_inplace_view_op:
|
||||
# no_dispatch() to make sure that we secretly change the metadata on the wrapper,
|
||||
# but don't end up dispatching the op anywhere else.
|
||||
mutated_args = [
|
||||
x
|
||||
for i, x in enumerate(args)
|
||||
if _get_write_alias(schema_info.args[i]) is not None
|
||||
]
|
||||
# Assumption: we have a very small number of inplace_view ops that follow a strict schema:
|
||||
# there is only a single argument that gets its metadata mutated.
|
||||
if len(mutated_args) != 1:
|
||||
raise AssertionError(
|
||||
"expected exactly one mutated arg for inplace_view ops"
|
||||
)
|
||||
# This check exists because we generally *do* want to update the metadata of any wrapper subclasses,
|
||||
# but FunctionalTensor is special: it overrides all size/stride calls to plumb to the inner tensor.
|
||||
# so we don't actually need to update the metadata (and attempting to do so causes errors)
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor
|
||||
|
||||
if not isinstance(mutated_args[0], FunctionalTensor):
|
||||
with torch.utils._mode_utils.no_dispatch():
|
||||
# See Note: [Fake Tensor Dispatch Keys]
|
||||
# we're borrowing the way it modifies dispatch key TLS.
|
||||
meta_in_tls = torch._C._meta_in_tls_dispatch_include()
|
||||
torch._C._set_meta_in_tls_dispatch_include(True)
|
||||
try:
|
||||
func(*args, **kwargs)
|
||||
finally:
|
||||
torch._C._set_meta_in_tls_dispatch_include(meta_in_tls)
|
||||
|
||||
# Next: we need to make sure to return inputs directly, if the output is a mutable alias (e.g. add_()).
|
||||
|
||||
schema_info_outs_write_aliases = schema_info.outs_write_aliases
|
||||
# simple case: none of our outputs have mutable aliases, so we can return the output as-is
|
||||
if schema_info_outs_write_aliases is None:
|
||||
return out
|
||||
|
||||
if len(schema_info_outs_write_aliases) == 1:
|
||||
return get_arg_from_alias(
|
||||
schema_info_outs_write_aliases[0], schema_info, args, kwargs
|
||||
)
|
||||
|
||||
# In the multi-return case, all aten ops return a tuple / list, so cast accordingly.
|
||||
outs_to_return = type(out)(
|
||||
[
|
||||
(get_arg_from_alias(write_alias, schema_info, args, kwargs))
|
||||
for write_alias in schema_info_outs_write_aliases
|
||||
]
|
||||
)
|
||||
return outs_to_return
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,151 @@
|
||||
import torch
|
||||
from torch._inductor.utils import get_device_tflops, get_gpu_dram_gbps
|
||||
from torch.fx.experimental.symbolic_shapes import (
|
||||
optimization_hint,
|
||||
statically_known_true,
|
||||
)
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from .flop_counter import flop_registry
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
_FLOAT_TYPES = OrderedSet(
|
||||
[
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
]
|
||||
)
|
||||
|
||||
# No fall-back kernel needed/exists for view ops
|
||||
_VIEW_OPS = OrderedSet(
|
||||
[
|
||||
aten.lift_fresh,
|
||||
aten.t,
|
||||
aten.transpose,
|
||||
aten.view,
|
||||
aten.detach,
|
||||
aten._unsafe_view,
|
||||
aten.split,
|
||||
aten.adjoint,
|
||||
aten.as_strided,
|
||||
aten.diagonal,
|
||||
aten.expand,
|
||||
aten.expand_as,
|
||||
aten.movedim,
|
||||
aten.permute,
|
||||
aten.select,
|
||||
aten.squeeze,
|
||||
aten.mT,
|
||||
aten.mH,
|
||||
aten.real,
|
||||
aten.imag,
|
||||
aten.view_as,
|
||||
aten.unflatten,
|
||||
aten.unfold,
|
||||
aten.unbind,
|
||||
aten.unsqueeze,
|
||||
aten.vsplit,
|
||||
aten.hsplit,
|
||||
aten.split_with_sizes,
|
||||
aten.swapaxes,
|
||||
aten.swapdims,
|
||||
aten.chunk,
|
||||
]
|
||||
)
|
||||
# We can ignore benchmarking tensor create ops
|
||||
_CREATE_OPS = OrderedSet(
|
||||
[
|
||||
aten.randint,
|
||||
aten.randn,
|
||||
aten.rand,
|
||||
aten.randn_like,
|
||||
aten.rand_like,
|
||||
aten.randint_like,
|
||||
aten.arange,
|
||||
aten.ones_like,
|
||||
aten.zeros_like,
|
||||
]
|
||||
)
|
||||
|
||||
_IGNORE_OPS = _VIEW_OPS | _CREATE_OPS
|
||||
|
||||
|
||||
def get_compute_time(func_packet, args, kwargs, out, out_dtypes) -> float: # type: ignore[no-untyped-def]
|
||||
"""
|
||||
Estimates the compute time of an aten operator.
|
||||
|
||||
Args:
|
||||
func_packet: The operator overload packet.
|
||||
args: The arguments to the operator.
|
||||
kwargs: The keyword arguments to the operator.
|
||||
out: The output of the operator.
|
||||
out_dtypes: The output data types.
|
||||
|
||||
Returns:
|
||||
float: The estimated compute time in nanoseconds.
|
||||
"""
|
||||
if func_packet in flop_registry:
|
||||
if len(out_dtypes) != 1:
|
||||
raise AssertionError(
|
||||
f"Only support single out dtype got {out_dtypes} for {func_packet}"
|
||||
)
|
||||
dtype = out_dtypes.pop()
|
||||
# This actually gives peta-FLOPs/s hence multiply by 1e15 to get the FLOPs/s
|
||||
peak_gpu_flops = get_device_tflops(dtype) * 1e15
|
||||
# We can expect to achieve 75% of theoretical peak flops
|
||||
factor = 0.75
|
||||
peak_empirical_flops = factor * peak_gpu_flops
|
||||
flop_count_func = flop_registry[func_packet]
|
||||
# We divide by a factor of 2 to get the MACs (multiply and accumulate)
|
||||
flop_count = flop_count_func(*args, **kwargs, out_val=out) / 2
|
||||
# We multiply by 1e9 to get the time in nano seconds
|
||||
compute_time = (flop_count / peak_empirical_flops) * 1e9
|
||||
return compute_time
|
||||
return 0.0
|
||||
|
||||
|
||||
def get_num_bytes(t: torch.Tensor) -> int:
|
||||
"""
|
||||
Calculates the memory consumption of a tensor.
|
||||
|
||||
Args:
|
||||
t (torch.Tensor): The input tensor.
|
||||
|
||||
Returns:
|
||||
int: The memory consumption of the tensor in bytes.
|
||||
"""
|
||||
real_numel = 1
|
||||
for size, stride in zip(t.shape, t.stride()):
|
||||
# For dims with stride=0 (expanded/broadcast), only 1 element accessed
|
||||
if not statically_known_true(stride == 0):
|
||||
real_numel *= optimization_hint(size, fallback=0)
|
||||
|
||||
return real_numel * t.element_size()
|
||||
|
||||
|
||||
def get_transfer_time(flat_args_kwargs, flat_outs) -> float: # type: ignore[no-untyped-def]
|
||||
"""
|
||||
Estimates the memory transfer time of input and output tensors.
|
||||
|
||||
Args:
|
||||
flat_args_kwargs (List[torch.Tensor]): The flat list of arguments and keyword arguments.
|
||||
flat_outs (List[torch.Tensor]): The flat list of outputs.
|
||||
|
||||
Returns:
|
||||
float: The estimated memory transfer time in nanoseconds.
|
||||
"""
|
||||
gpu_memory_bandwidth = get_gpu_dram_gbps()
|
||||
read_bytes = sum(
|
||||
get_num_bytes(t) for t in flat_args_kwargs if isinstance(t, torch.Tensor)
|
||||
)
|
||||
write_bytes = sum(
|
||||
get_num_bytes(t) for t in flat_outs if isinstance(t, torch.Tensor)
|
||||
)
|
||||
counted_bytes = read_bytes + write_bytes
|
||||
# The GPU memory bandwidth is in GB/s so the transfer time is in nanoseconds
|
||||
transfer_time = counted_bytes / gpu_memory_bandwidth
|
||||
return transfer_time
|
||||
@@ -0,0 +1,31 @@
|
||||
# NOTE! PLEASE KEEP THIS FILE *FREE* OF TORCH DEPS! IT SHOULD BE IMPORTABLE ANYWHERE.
|
||||
# IF YOU FEEL AN OVERWHELMING URGE TO ADD A TORCH DEP, MAKE A TRAMPOLINE FILE A LA torch._dynamo.utils
|
||||
# AND SCRUB AWAY TORCH NOTIONS THERE.
|
||||
import collections
|
||||
import functools
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable
|
||||
from typing import TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
simple_call_counter: OrderedDict[str, int] = collections.OrderedDict()
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
|
||||
def count_label(label: str) -> None:
|
||||
prev = simple_call_counter.setdefault(label, 0)
|
||||
simple_call_counter[label] = prev + 1
|
||||
|
||||
|
||||
def count(fn: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
if fn.__qualname__ not in simple_call_counter:
|
||||
simple_call_counter[fn.__qualname__] = 0
|
||||
simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
+313
@@ -0,0 +1,313 @@
|
||||
# mypy: disallow-untyped-defs
|
||||
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import time
|
||||
from collections.abc import Callable, Sequence
|
||||
from threading import Lock
|
||||
from typing import Any, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
logger = logging.getLogger("strobelight_function_profiler")
|
||||
|
||||
console_handler = logging.StreamHandler()
|
||||
formatter = logging.Formatter(
|
||||
"%(name)s, line %(lineno)d, %(asctime)s, %(levelname)s: %(message)s"
|
||||
)
|
||||
console_handler.setFormatter(formatter)
|
||||
|
||||
logger.addHandler(console_handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = False
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
|
||||
class StrobelightCLIProfilerError(Exception):
|
||||
"""
|
||||
Raised when an error happens during strobelight profiling
|
||||
"""
|
||||
|
||||
|
||||
def _pid_namespace_link(pid: int | None = None) -> str:
|
||||
"""Returns the link to the process's namespace, example: pid:[4026531836]"""
|
||||
PID_NAMESPACE_PATH = "/proc/{}/ns/pid"
|
||||
pid = pid or os.getpid()
|
||||
return os.readlink(PID_NAMESPACE_PATH.format(pid))
|
||||
|
||||
|
||||
def _pid_namespace(pid: int | None = None) -> int:
|
||||
"""Returns the process's namespace id"""
|
||||
pid = pid or os.getpid()
|
||||
link = _pid_namespace_link(pid)
|
||||
return int(link[link.find("[") + 1 : -1])
|
||||
|
||||
|
||||
def _command_to_string(command: Sequence[str]) -> str:
|
||||
return " ".join(command)
|
||||
|
||||
|
||||
class StrobelightCLIFunctionProfiler:
|
||||
"""
|
||||
Note: this is a meta only tool.
|
||||
|
||||
StrobelightCLIFunctionProfiler can be used to profile a python function and
|
||||
generate a strobelight link with the results. It works on meta servers but
|
||||
does not requires an fbcode target.
|
||||
When stop_at_error is false(default), error during profiling does not prevent
|
||||
the work function from running.
|
||||
|
||||
Check function_profiler_example.py for an example.
|
||||
"""
|
||||
|
||||
# This lock is used to make sure only one thread is running the profiler at any point.
|
||||
_lock = Lock()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
stop_at_error: bool = False,
|
||||
max_profile_duration_sec: int = 60 * 10,
|
||||
sample_each: float = 1e7, # sample each sample_each cycles.
|
||||
run_user_name: str = "pytorch-strobelight-ondemand",
|
||||
timeout_wait_for_running_sec: int = 60,
|
||||
timeout_wait_for_finished_sec: int = 60,
|
||||
recorded_env_variables: list[str] | None = None,
|
||||
sample_tags: list[str] | None = None,
|
||||
stack_max_len: int = 127,
|
||||
async_stack_max_len: int = 127,
|
||||
) -> None:
|
||||
self.stop_at_error = stop_at_error
|
||||
self.max_profile_duration_sec = max_profile_duration_sec
|
||||
self.sample_each = sample_each
|
||||
self.run_user_name = run_user_name
|
||||
self.timeout_wait_for_running_sec = timeout_wait_for_running_sec
|
||||
self.timeout_wait_for_finished_sec = timeout_wait_for_finished_sec
|
||||
# Results of the most recent run.
|
||||
# Tracks the strobelight run id of the most recent run
|
||||
self.current_run_id: int | None = None
|
||||
self.sample_tags = sample_tags
|
||||
|
||||
def _run_async(self) -> None:
|
||||
processId = os.getpid()
|
||||
namespace = _pid_namespace(processId)
|
||||
command = [
|
||||
"strobeclient",
|
||||
"run",
|
||||
"--profiler",
|
||||
"pyperf",
|
||||
"--event",
|
||||
"cycles",
|
||||
"--async",
|
||||
"--sample-interval",
|
||||
f"{int(self.sample_each)}",
|
||||
"--duration-ms",
|
||||
f"{int(self.max_profile_duration_sec * 1000)}",
|
||||
"--pid",
|
||||
f"{namespace}:{processId}",
|
||||
]
|
||||
|
||||
if self.sample_tags:
|
||||
command.append("--sample-tags")
|
||||
command.append(",".join(self.sample_tags))
|
||||
|
||||
logger.debug("running command: %s", _command_to_string(command))
|
||||
result = subprocess.run(command, capture_output=True)
|
||||
output = result.stderr.decode("utf-8")
|
||||
logger.debug("output:\n{%s}", output)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to start strobelight profiling, error in run_async:{output}"
|
||||
)
|
||||
|
||||
if match := re.search(r"INFO Run Id: (-?\d+)", output):
|
||||
self.current_run_id = int(match.group(1))
|
||||
return
|
||||
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to start strobelight profiling, unexpected result {output}"
|
||||
)
|
||||
|
||||
def _wait_for_running(self, counter: int = 0) -> None:
|
||||
if counter > 20:
|
||||
raise StrobelightCLIProfilerError(
|
||||
"wait_for_running called more than 20 times"
|
||||
)
|
||||
|
||||
command = ["strobeclient", "getRunStatus", "--run-id", f"{self.current_run_id}"]
|
||||
logger.debug("running command: %s", _command_to_string(command))
|
||||
result = subprocess.run(command, capture_output=True)
|
||||
output = result.stderr.decode("utf-8")
|
||||
logger.debug("output:\n{%s}", output)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to start strobelight profiling, error in wait_for_running:{output}"
|
||||
)
|
||||
|
||||
if match := re.search("Profile run status: (.*)", output):
|
||||
current_status = match.group(1)
|
||||
if current_status == "RUNNING":
|
||||
return
|
||||
elif current_status == "PREPARING":
|
||||
time.sleep(10)
|
||||
self._wait_for_running(counter + 1)
|
||||
return
|
||||
else:
|
||||
raise StrobelightCLIProfilerError(f"unexpected {current_status} phase")
|
||||
|
||||
raise StrobelightCLIProfilerError(f"unexpected output\n: {output} ")
|
||||
|
||||
def _stop_run(self) -> None:
|
||||
command = ["strobeclient", "stopRun", "--run-id", str(self.current_run_id)]
|
||||
logger.debug("running command: %s", _command_to_string(command))
|
||||
result = subprocess.run(command, capture_output=True)
|
||||
output = result.stderr.decode("utf-8")
|
||||
logger.debug("output:\n{%s}", output)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to stop strobelight profiling, return code is not 0 :{output}"
|
||||
)
|
||||
|
||||
if match := re.search("INFO ::1:(.*)", output):
|
||||
current_status = match.group(1)
|
||||
if current_status.__contains__("Success!"):
|
||||
return
|
||||
else:
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to stop strobelight profiling, got {current_status} result"
|
||||
)
|
||||
|
||||
raise StrobelightCLIProfilerError(f"unexpected output\n: {output} ")
|
||||
|
||||
def _get_results(self) -> None:
|
||||
command = ["strobeclient", "getRunStatus", "--run-id", str(self.current_run_id)]
|
||||
logger.debug("running command: %s", _command_to_string(command))
|
||||
result = subprocess.run(command, capture_output=True)
|
||||
output = result.stderr.decode("utf-8")
|
||||
logger.debug("output:\n{%s}", output)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to extract profiling results, return code is not 0 : {output}"
|
||||
)
|
||||
|
||||
if match := re.search("INFO ::1:(.*)", output):
|
||||
current_status = match.group(1)
|
||||
if current_status.__contains__("Profile run status: PROCESSING"):
|
||||
time.sleep(10)
|
||||
self._get_results()
|
||||
return
|
||||
elif not current_status.__contains__("Profile run finished with SUCCESS"):
|
||||
raise StrobelightCLIProfilerError(
|
||||
f"failed to extract profiling results, unexpected response {output}"
|
||||
)
|
||||
|
||||
for item in re.findall(
|
||||
r"(Total samples(.*)|GraphProfiler(.*)|Icicle view \(python stack\)(.*))",
|
||||
output,
|
||||
):
|
||||
logger.info(item[0])
|
||||
|
||||
def _stop_strobelight_no_throw(
|
||||
self,
|
||||
collect_results: bool,
|
||||
) -> None:
|
||||
try:
|
||||
# call stop run
|
||||
self._stop_run()
|
||||
logger.info("strobelight profiling stopped")
|
||||
|
||||
logger.debug("collection stopped")
|
||||
|
||||
if not collect_results:
|
||||
return
|
||||
|
||||
self._get_results()
|
||||
except Exception:
|
||||
logger.warning("error during stop_strobelight", exc_info=True)
|
||||
|
||||
# Return true if strobelight started and is running. Never throw.
|
||||
def _start_strobelight(self) -> bool:
|
||||
strobelight_started = False
|
||||
try:
|
||||
self._run_async()
|
||||
strobelight_started = True
|
||||
logger.info("strobelight run id is: %s", self.current_run_id)
|
||||
self._wait_for_running()
|
||||
logger.info("strobelight profiling running")
|
||||
return True
|
||||
|
||||
except Exception:
|
||||
logger.warning("error during start_strobelight:", exc_info=True)
|
||||
if strobelight_started:
|
||||
self._stop_strobelight_no_throw(collect_results=False)
|
||||
return False
|
||||
|
||||
def profile(
|
||||
self, work_function: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
|
||||
) -> _R | None:
|
||||
self.current_run_id = None
|
||||
|
||||
if locked := StrobelightCLIFunctionProfiler._lock.acquire(False):
|
||||
if not locked:
|
||||
if self.stop_at_error:
|
||||
raise StrobelightCLIProfilerError("concurrent runs not supported")
|
||||
|
||||
logger.warning("concurrent runs not supported")
|
||||
return work_function(*args, **kwargs)
|
||||
|
||||
started = self._start_strobelight()
|
||||
if not started:
|
||||
if self.stop_at_error:
|
||||
StrobelightCLIFunctionProfiler._lock.release()
|
||||
raise StrobelightCLIProfilerError(
|
||||
"failed to start strobelight profiling"
|
||||
)
|
||||
result = work_function(*args, **kwargs)
|
||||
StrobelightCLIFunctionProfiler._lock.release()
|
||||
return result
|
||||
|
||||
try:
|
||||
logger.debug("collection started")
|
||||
result = work_function(*args, **kwargs)
|
||||
self._stop_strobelight_no_throw(collect_results=True)
|
||||
StrobelightCLIFunctionProfiler._lock.release()
|
||||
return result
|
||||
except Exception as error:
|
||||
logger.warning("work function throw exception", exc_info=True)
|
||||
self._stop_strobelight_no_throw(collect_results=False)
|
||||
StrobelightCLIFunctionProfiler._lock.release()
|
||||
raise error
|
||||
return None
|
||||
|
||||
|
||||
# A function decorator that wraps profile, if no profiler is provided one with
|
||||
# default args is created. A function can be annotated as:
|
||||
# @strobelight()
|
||||
# @strobelight(profiler = StrobelightFunctionProfiler(stop_at_error=True,..))
|
||||
# @strobelight(stop_at_error=True,...)
|
||||
def strobelight(
|
||||
profiler: StrobelightCLIFunctionProfiler | None = None, **kwargs: Any
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R | None]]:
|
||||
if not profiler:
|
||||
profiler = StrobelightCLIFunctionProfiler(**kwargs)
|
||||
|
||||
def strobelight_inner(
|
||||
work_function: Callable[_P, _R],
|
||||
) -> Callable[_P, _R | None]:
|
||||
@functools.wraps(work_function)
|
||||
def wrapper_function(*args: _P.args, **kwargs: _P.kwargs) -> _R | None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return profiler.profile(work_function, *args, **kwargs)
|
||||
|
||||
return wrapper_function
|
||||
|
||||
return strobelight_inner
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,228 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This is a simple interpreter for Sympy expressions that dispatches to
|
||||
classes following the torch._inductor.virtualized calling convention.
|
||||
For directness, the interpreter takes the handler directly rather than
|
||||
consulting the TLS. It does not use most of the methods on the full
|
||||
handler; only those with corresponding Sympy expressions. To see an example
|
||||
of a full handler, see torch.utils._sympy.value_ranges.ValueRangeAnalysis.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
from sympy.logic.boolalg import Boolean as SympyBoolean, BooleanAtom
|
||||
|
||||
import torch
|
||||
|
||||
from .functions import (
|
||||
BitwiseFn_bitwise_and,
|
||||
BitwiseFn_bitwise_or,
|
||||
BitwiseFn_bitwise_xor,
|
||||
CeilToInt,
|
||||
CleanDiv,
|
||||
FloatPow,
|
||||
FloatTrueDiv,
|
||||
FloorDiv,
|
||||
FloorToInt,
|
||||
Identity,
|
||||
IntTrueDiv,
|
||||
IsNonOverlappingAndDenseIndicator,
|
||||
Max,
|
||||
Min,
|
||||
Mod,
|
||||
ModularIndexing,
|
||||
OpaqueUnaryFn_log2,
|
||||
PowByNatural,
|
||||
PythonMod,
|
||||
RoundDecimal,
|
||||
RoundToInt,
|
||||
ToFloat,
|
||||
TruncToFloat,
|
||||
TruncToInt,
|
||||
Where,
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# TODO: Dedupe this with SYMPY_INTERP
|
||||
|
||||
|
||||
@functools.cache
|
||||
def handlers():
|
||||
# TODO add CeilDiv (it doesn't appear in the index_expr)
|
||||
|
||||
# TODO default to some decompositions if the interpreter doesn't have them
|
||||
# like decomposing ModularIndexing or implementing Le(a,b) as Ge(b, a)
|
||||
|
||||
HANDLERS = {
|
||||
sympy.Or: "or_",
|
||||
sympy.And: "and_",
|
||||
sympy.Eq: "eq",
|
||||
sympy.Ne: "ne",
|
||||
sympy.Lt: "lt",
|
||||
sympy.Gt: "gt",
|
||||
sympy.Le: "le",
|
||||
sympy.Ge: "ge",
|
||||
sympy.Not: "not_",
|
||||
IntTrueDiv: "int_truediv",
|
||||
FloatTrueDiv: "truediv",
|
||||
FloorDiv: "floordiv",
|
||||
CleanDiv: "floordiv", # TODO: hmm?
|
||||
TruncToFloat: "trunc",
|
||||
Where: "where",
|
||||
sympy.Add: "add",
|
||||
sympy.Mul: "mul",
|
||||
FloatPow: "pow",
|
||||
PowByNatural: "pow_by_natural",
|
||||
# sympy simplifies x * x into Pow(x, 2), so we need to handle this.
|
||||
# Do NOT use builtin Pow for floats
|
||||
# TODO: There is a hazard here, if we have float * float it will
|
||||
# also get turned into Pow(float, 2) but we don't want this because
|
||||
# pow_by_natural is assumed to only be integers. Probably the fix is
|
||||
# to add a FloatMul to impede this optimization
|
||||
sympy.Pow: "pow_by_natural",
|
||||
Mod: "mod",
|
||||
PythonMod: "python_mod",
|
||||
# TODO: Inductor can generate these, but it's ill-specified which
|
||||
# semantics were intended here. Needs to be cleaned up along with
|
||||
# FloorDiv in a bigger cleanup
|
||||
sympy.Mod: "mod",
|
||||
sympy.Abs: "abs",
|
||||
sympy.log: "log",
|
||||
sympy.exp: "exp",
|
||||
sympy.Min: "minimum",
|
||||
sympy.Max: "maximum",
|
||||
Min: "minimum",
|
||||
Max: "maximum",
|
||||
ModularIndexing: "modular_indexing",
|
||||
sympy.functions.elementary.piecewise.ExprCondPair: "expr_cond_pair",
|
||||
sympy.Piecewise: "piecewise",
|
||||
Identity: "identity",
|
||||
IsNonOverlappingAndDenseIndicator: "is_non_overlapping_and_dense_indicator",
|
||||
RoundDecimal: "round_decimal",
|
||||
# TODO: do the rest of the opaque unary functions...
|
||||
OpaqueUnaryFn_log2: "log2",
|
||||
BitwiseFn_bitwise_and: "bitwise_and",
|
||||
BitwiseFn_bitwise_or: "bitwise_or",
|
||||
BitwiseFn_bitwise_xor: "bitwise_xor",
|
||||
}
|
||||
# TODO: This is kind of pointless, we shouldn't be generating sympy.sin
|
||||
# for these functions, they should be Opaque instead
|
||||
for name in ["cos", "sin", "tan", "sinh", "cosh", "tanh", "asin", "acos", "atan"]:
|
||||
HANDLERS[getattr(sympy, name)] = name
|
||||
|
||||
return HANDLERS
|
||||
|
||||
|
||||
ASSOCIATIVE_OPS = {"minimum", "maximum", "mul", "add", "and_", "or_"}
|
||||
|
||||
|
||||
def _run_sympy_handler(analysis, args, expr, index_dtype=torch.int64):
|
||||
# Special cases
|
||||
if isinstance(expr, sympy.Pow) and isinstance(
|
||||
expr.args[1], sympy.core.numbers.Half
|
||||
):
|
||||
return analysis.sqrt(args[0])
|
||||
if isinstance(expr, ToFloat):
|
||||
return analysis.to_dtype(args[0], torch.float64)
|
||||
|
||||
# These handlers are special because they take an extra dtype argument
|
||||
# specifying what they should convert to, and we need to appropriately set
|
||||
# this up when we convert from Sympy. A reasonable default when you
|
||||
# are translating is to conservatively do int64, and then narrow these
|
||||
# arguments later when you discover you can narrow the index range. But
|
||||
# if you already know that 32-bit indexing is OK, you can directly do the
|
||||
# sympy translation with index_dtype=torch.int32
|
||||
INDEX_DTYPE_HANDLERS = {
|
||||
TruncToInt: "trunc_to_int",
|
||||
sympy.floor: "floor_to_int",
|
||||
sympy.ceiling: "ceil_to_int",
|
||||
FloorToInt: "floor_to_int",
|
||||
CeilToInt: "ceil_to_int",
|
||||
RoundToInt: "round_to_int",
|
||||
}
|
||||
if (handler_name := INDEX_DTYPE_HANDLERS.get(expr.func)) is not None:
|
||||
return getattr(analysis, handler_name)(*args, index_dtype)
|
||||
|
||||
# Fastpath for n-ary integral addition
|
||||
if expr.func is sympy.Add and expr.is_integer and hasattr(analysis, "sym_sum"):
|
||||
r = analysis.sym_sum(args)
|
||||
log.debug("sym_sum(%s) -> %s", args, r)
|
||||
return r
|
||||
|
||||
if hasattr(expr.func, "_torch_handler_name"):
|
||||
handler_name = expr.func._torch_handler_name
|
||||
else:
|
||||
handler_name = handlers()[expr.func]
|
||||
handler = getattr(analysis, handler_name)
|
||||
try:
|
||||
if handler_name in ASSOCIATIVE_OPS:
|
||||
if len(args) <= 1:
|
||||
raise AssertionError("associative op needs >1 args")
|
||||
acc = handler(args[0], args[1])
|
||||
for i in range(2, len(args)):
|
||||
acc = handler(acc, args[i])
|
||||
log.debug("%s(%s) -> %s", handler_name, args, acc)
|
||||
return acc
|
||||
else:
|
||||
r = handler(*args)
|
||||
log.debug("%s(%s) -> %s", handler_name, args, r)
|
||||
return r
|
||||
except NotImplementedError:
|
||||
raise
|
||||
except Exception:
|
||||
log.warning("failed while executing %s(%s)", handler_name, args)
|
||||
raise
|
||||
|
||||
|
||||
_nil = object()
|
||||
|
||||
|
||||
def sympy_interp(
|
||||
analysis,
|
||||
env: dict[sympy.Symbol, Any],
|
||||
expr: sympy.Expr | SympyBoolean,
|
||||
*,
|
||||
index_dtype=torch.int64,
|
||||
missing_handler=None,
|
||||
):
|
||||
# Handle base cases
|
||||
dtype = None
|
||||
if isinstance(expr, BooleanAtom):
|
||||
dtype = torch.bool
|
||||
elif isinstance(expr, sympy.Integer):
|
||||
dtype = torch.int64
|
||||
elif isinstance(expr, sympy.Number):
|
||||
dtype = torch.double
|
||||
|
||||
if dtype is not None:
|
||||
return analysis.constant(expr, dtype)
|
||||
elif isinstance(expr, sympy.Symbol):
|
||||
if (r := env.get(expr, _nil)) is not _nil:
|
||||
return r
|
||||
elif missing_handler:
|
||||
return missing_handler(expr)
|
||||
else:
|
||||
raise KeyError(expr)
|
||||
|
||||
# Recursive case
|
||||
return _run_sympy_handler(
|
||||
analysis,
|
||||
[
|
||||
sympy_interp(
|
||||
analysis,
|
||||
env,
|
||||
arg,
|
||||
index_dtype=index_dtype,
|
||||
missing_handler=missing_handler,
|
||||
)
|
||||
for arg in expr.args
|
||||
],
|
||||
expr,
|
||||
index_dtype=index_dtype,
|
||||
)
|
||||
@@ -0,0 +1,417 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import mpmath.libmp as mlib # type: ignore[import-untyped]
|
||||
import sympy
|
||||
from sympy import Expr
|
||||
from sympy.core.decorators import _sympifyit
|
||||
from sympy.core.expr import AtomicExpr
|
||||
from sympy.core.numbers import Number
|
||||
from sympy.core.parameters import global_parameters
|
||||
from sympy.core.singleton import S, Singleton
|
||||
|
||||
|
||||
# pyrefly: ignore [invalid-inheritance]
|
||||
class IntInfinity(Number, metaclass=Singleton):
|
||||
r"""Positive integer infinite quantity.
|
||||
|
||||
Integer infinity is a value in an extended integers which
|
||||
is greater than all other integers. We distinguish it from
|
||||
sympy's existing notion of infinity in that it reports that
|
||||
it is_integer.
|
||||
|
||||
Infinity is a singleton, and can be accessed by ``S.IntInfinity``,
|
||||
or can be imported as ``int_oo``.
|
||||
"""
|
||||
|
||||
# NB: We can't actually mark this as infinite, as integer and infinite are
|
||||
# inconsistent assumptions in sympy. We also report that we are complex,
|
||||
# different from sympy.oo
|
||||
|
||||
is_integer = True
|
||||
is_commutative = True
|
||||
is_number = True
|
||||
is_extended_real = True
|
||||
is_comparable = True
|
||||
is_extended_positive = True
|
||||
is_prime = False
|
||||
|
||||
# Ensure we get dispatched to before plain numbers
|
||||
_op_priority = 100.0
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(cls):
|
||||
return AtomicExpr.__new__(cls)
|
||||
|
||||
def _sympystr(self, printer) -> str:
|
||||
return "int_oo"
|
||||
|
||||
def _eval_subs(self, old, new):
|
||||
if self == old:
|
||||
return new
|
||||
|
||||
# We could do these, not sure about it
|
||||
"""
|
||||
def _eval_evalf(self, prec=None):
|
||||
return Float('inf')
|
||||
|
||||
def evalf(self, prec=None, **options):
|
||||
return self._eval_evalf(prec)
|
||||
"""
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __add__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other in (S.Infinity, S.NegativeInfinity):
|
||||
return other
|
||||
if other in (S.NegativeIntInfinity, S.NaN):
|
||||
return S.NaN
|
||||
return self
|
||||
return Number.__add__(self, other)
|
||||
|
||||
__radd__ = __add__
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __sub__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other is S.Infinity:
|
||||
return S.NegativeInfinity
|
||||
if other is S.NegativeInfinity:
|
||||
return S.Infinity
|
||||
if other in (S.IntInfinity, S.NaN):
|
||||
return S.NaN
|
||||
return self
|
||||
return Number.__sub__(self, other)
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __rsub__(self, other):
|
||||
return (-self).__add__(other)
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __mul__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other.is_zero or other is S.NaN:
|
||||
return S.NaN
|
||||
if other.is_extended_positive:
|
||||
return self
|
||||
return S.NegativeIntInfinity
|
||||
return Number.__mul__(self, other)
|
||||
|
||||
__rmul__ = __mul__
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __truediv__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other in (
|
||||
S.Infinity,
|
||||
S.IntInfinity,
|
||||
S.NegativeInfinity,
|
||||
S.NegativeIntInfinity,
|
||||
S.NaN,
|
||||
):
|
||||
return S.NaN
|
||||
if other.is_extended_nonnegative:
|
||||
return S.Infinity # truediv produces float
|
||||
return S.NegativeInfinity # truediv produces float
|
||||
return Number.__truediv__(self, other)
|
||||
|
||||
def __abs__(self):
|
||||
return S.IntInfinity
|
||||
|
||||
def __neg__(self):
|
||||
return S.NegativeIntInfinity
|
||||
|
||||
def _eval_power(self, expt):
|
||||
if expt.is_extended_positive:
|
||||
return S.IntInfinity
|
||||
if expt.is_extended_negative:
|
||||
return S.Zero
|
||||
if expt is S.NaN:
|
||||
return S.NaN
|
||||
if expt is S.ComplexInfinity:
|
||||
return S.NaN
|
||||
if expt.is_extended_real is False and expt.is_number:
|
||||
from sympy.functions.elementary.complexes import re
|
||||
|
||||
expt_real = re(expt)
|
||||
if expt_real.is_positive:
|
||||
return S.ComplexInfinity
|
||||
if expt_real.is_negative:
|
||||
return S.Zero
|
||||
if expt_real.is_zero:
|
||||
return S.NaN
|
||||
|
||||
return self ** expt.evalf()
|
||||
|
||||
def _as_mpf_val(self, prec):
|
||||
return mlib.finf
|
||||
|
||||
def __hash__(self):
|
||||
return super().__hash__()
|
||||
|
||||
def __eq__(self, other):
|
||||
return other is S.IntInfinity
|
||||
|
||||
def __ne__(self, other):
|
||||
return other is not S.IntInfinity
|
||||
|
||||
def __gt__(self, other):
|
||||
if other is S.Infinity:
|
||||
return sympy.false # sympy.oo > int_oo
|
||||
elif other is S.IntInfinity:
|
||||
return sympy.false # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.true
|
||||
|
||||
def __ge__(self, other):
|
||||
if other is S.Infinity:
|
||||
return sympy.false # sympy.oo > int_oo
|
||||
elif other is S.IntInfinity:
|
||||
return sympy.true # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.true
|
||||
|
||||
def __lt__(self, other):
|
||||
if other is S.Infinity:
|
||||
return sympy.true # sympy.oo > int_oo
|
||||
elif other is S.IntInfinity:
|
||||
return sympy.false # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.false
|
||||
|
||||
def __le__(self, other):
|
||||
if other is S.Infinity:
|
||||
return sympy.true # sympy.oo > int_oo
|
||||
elif other is S.IntInfinity:
|
||||
return sympy.true # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.false
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __mod__(self, other):
|
||||
if not isinstance(other, Expr):
|
||||
return NotImplemented
|
||||
return S.NaN
|
||||
|
||||
__rmod__ = __mod__
|
||||
|
||||
def floor(self):
|
||||
return self
|
||||
|
||||
def ceiling(self):
|
||||
return self
|
||||
|
||||
|
||||
int_oo = S.IntInfinity
|
||||
|
||||
|
||||
def is_infinite(expr) -> bool:
|
||||
"""Check if an expression is any type of infinity (positive or negative).
|
||||
|
||||
This handles both sympy's built-in infinities (oo, -oo) and PyTorch's
|
||||
integer infinities (int_oo, -int_oo).
|
||||
|
||||
Note: We cannot rely on sympy's is_finite property because IntInfinity
|
||||
and NegativeIntInfinity have is_integer=True, which implies is_finite=True
|
||||
in sympy's assumption system.
|
||||
"""
|
||||
return expr in (
|
||||
S.Infinity,
|
||||
S.NegativeInfinity,
|
||||
S.IntInfinity,
|
||||
S.NegativeIntInfinity,
|
||||
)
|
||||
|
||||
|
||||
# pyrefly: ignore [invalid-inheritance]
|
||||
class NegativeIntInfinity(Number, metaclass=Singleton):
|
||||
"""Negative integer infinite quantity.
|
||||
|
||||
NegativeInfinity is a singleton, and can be accessed
|
||||
by ``S.NegativeInfinity``.
|
||||
|
||||
See Also
|
||||
========
|
||||
|
||||
IntInfinity
|
||||
"""
|
||||
|
||||
# Ensure we get dispatched to before plain numbers
|
||||
_op_priority = 100.0
|
||||
|
||||
is_integer = True
|
||||
is_extended_real = True
|
||||
is_commutative = True
|
||||
is_comparable = True
|
||||
is_extended_negative = True
|
||||
is_number = True
|
||||
is_prime = False
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(cls):
|
||||
return AtomicExpr.__new__(cls)
|
||||
|
||||
def _eval_subs(self, old, new):
|
||||
if self == old:
|
||||
return new
|
||||
|
||||
def _sympystr(self, printer) -> str:
|
||||
return "-int_oo"
|
||||
|
||||
"""
|
||||
def _eval_evalf(self, prec=None):
|
||||
return Float('-inf')
|
||||
|
||||
def evalf(self, prec=None, **options):
|
||||
return self._eval_evalf(prec)
|
||||
"""
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __add__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other is S.Infinity:
|
||||
return S.Infinity
|
||||
if other in (S.IntInfinity, S.NaN):
|
||||
return S.NaN
|
||||
return self
|
||||
return Number.__add__(self, other)
|
||||
|
||||
__radd__ = __add__
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __sub__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other is S.NegativeInfinity:
|
||||
return S.Infinity
|
||||
if other in (S.NegativeIntInfinity, S.NaN):
|
||||
return S.NaN
|
||||
return self
|
||||
return Number.__sub__(self, other)
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __rsub__(self, other):
|
||||
return (-self).__add__(other)
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __mul__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other.is_zero or other is S.NaN:
|
||||
return S.NaN
|
||||
if other.is_extended_positive:
|
||||
return self
|
||||
return S.IntInfinity
|
||||
return Number.__mul__(self, other)
|
||||
|
||||
__rmul__ = __mul__
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __truediv__(self, other):
|
||||
if isinstance(other, Number) and global_parameters.evaluate:
|
||||
if other in (
|
||||
S.Infinity,
|
||||
S.IntInfinity,
|
||||
S.NegativeInfinity,
|
||||
S.NegativeIntInfinity,
|
||||
S.NaN,
|
||||
):
|
||||
return S.NaN
|
||||
if other.is_extended_nonnegative:
|
||||
return self
|
||||
return S.Infinity # truediv returns float
|
||||
return Number.__truediv__(self, other)
|
||||
|
||||
def __abs__(self):
|
||||
return S.IntInfinity
|
||||
|
||||
def __neg__(self):
|
||||
return S.IntInfinity
|
||||
|
||||
def _eval_power(self, expt):
|
||||
if expt.is_number:
|
||||
if expt in (
|
||||
S.NaN,
|
||||
S.Infinity,
|
||||
S.NegativeInfinity,
|
||||
S.IntInfinity,
|
||||
S.NegativeIntInfinity,
|
||||
):
|
||||
return S.NaN
|
||||
|
||||
if isinstance(expt, sympy.Integer) and expt.is_extended_positive:
|
||||
if expt.is_odd:
|
||||
return S.NegativeIntInfinity
|
||||
else:
|
||||
return S.IntInfinity
|
||||
|
||||
inf_part = S.IntInfinity**expt
|
||||
s_part = S.NegativeOne**expt
|
||||
if inf_part == 0 and s_part.is_finite:
|
||||
return inf_part
|
||||
if (
|
||||
inf_part is S.ComplexInfinity
|
||||
and s_part.is_finite
|
||||
and not s_part.is_zero
|
||||
):
|
||||
return S.ComplexInfinity
|
||||
return s_part * inf_part
|
||||
|
||||
def _as_mpf_val(self, prec):
|
||||
return mlib.fninf
|
||||
|
||||
def __hash__(self):
|
||||
return super().__hash__()
|
||||
|
||||
def __eq__(self, other):
|
||||
return other is S.NegativeIntInfinity
|
||||
|
||||
def __ne__(self, other):
|
||||
return other is not S.NegativeIntInfinity
|
||||
|
||||
def __gt__(self, other):
|
||||
if other is S.NegativeInfinity:
|
||||
return sympy.true # -sympy.oo < -int_oo
|
||||
elif other is S.NegativeIntInfinity:
|
||||
return sympy.false # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.false
|
||||
|
||||
def __ge__(self, other):
|
||||
if other is S.NegativeInfinity:
|
||||
return sympy.true # -sympy.oo < -int_oo
|
||||
elif other is S.NegativeIntInfinity:
|
||||
return sympy.true # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.false
|
||||
|
||||
def __lt__(self, other):
|
||||
if other is S.NegativeInfinity:
|
||||
return sympy.false # -sympy.oo < -int_oo
|
||||
elif other is S.NegativeIntInfinity:
|
||||
return sympy.false # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.true
|
||||
|
||||
def __le__(self, other):
|
||||
if other is S.NegativeInfinity:
|
||||
return sympy.false # -sympy.oo < -int_oo
|
||||
elif other is S.NegativeIntInfinity:
|
||||
return sympy.true # consistency with sympy.oo
|
||||
else:
|
||||
return sympy.true
|
||||
|
||||
@_sympifyit("other", NotImplemented)
|
||||
def __mod__(self, other):
|
||||
if not isinstance(other, Expr):
|
||||
return NotImplemented
|
||||
return S.NaN
|
||||
|
||||
__rmod__ = __mod__
|
||||
|
||||
def floor(self):
|
||||
return self
|
||||
|
||||
def ceiling(self):
|
||||
return self
|
||||
|
||||
def as_powers_dict(self):
|
||||
return {S.NegativeOne: 1, S.IntInfinity: 1}
|
||||
@@ -0,0 +1,666 @@
|
||||
import sys
|
||||
|
||||
import sympy
|
||||
from sympy.printing.precedence import PRECEDENCE, precedence
|
||||
from sympy.printing.str import StrPrinter
|
||||
|
||||
|
||||
INDEX_TYPE = "int64_t"
|
||||
INDEX_TYPE_MAX = (1 << 63) - 1
|
||||
INDEX_TYPE_MIN = -1 << 63
|
||||
|
||||
|
||||
# This printer contains rules that are supposed to be generic for both C/C++ and
|
||||
# Python
|
||||
class ExprPrinter(StrPrinter):
|
||||
# override this so that _print_FloorDiv is used
|
||||
printmethod = "_torch_sympystr"
|
||||
|
||||
def _print_Mul(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, "*", precedence(expr))
|
||||
|
||||
def _print_Not(self, expr: sympy.Expr) -> str:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"not ({self._print(expr.args[0])})"
|
||||
|
||||
def _print_Add(self, expr: sympy.Expr, order: str | None = None) -> str:
|
||||
return self.stringify(expr.args, " + ", precedence(expr))
|
||||
|
||||
def _print_Relational(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, f" {expr.rel_op} ", precedence(expr))
|
||||
|
||||
def _print_BitwiseFn_bitwise_and(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " & ", PRECEDENCE["BitwiseAnd"])
|
||||
|
||||
def _print_BitwiseFn_bitwise_or(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " | ", PRECEDENCE["BitwiseOr"])
|
||||
|
||||
def _print_BitwiseFn_bitwise_xor(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " ^ ", PRECEDENCE["BitwiseXor"])
|
||||
|
||||
# NB: this is OK to put here, because Mod is only defined for positive
|
||||
# numbers, and so across C/Python its behavior is consistent
|
||||
def _print_Mod(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " % ", PRECEDENCE["Atom"] - 0.5)
|
||||
|
||||
def _print_FloatTrueDiv(self, expr: sympy.Expr) -> str:
|
||||
s = self.stringify(expr.args, " / ", PRECEDENCE["Atom"] - 0.5)
|
||||
return f"({s})"
|
||||
|
||||
def _print_CleanDiv(self, expr: sympy.Expr) -> str:
|
||||
return self._print_FloorDiv(expr)
|
||||
|
||||
def _print_Identity(self, expr: sympy.Expr) -> str:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return self._print(expr.args[0])
|
||||
|
||||
def _print_Float(self, expr: sympy.Expr) -> str:
|
||||
if expr._prec == 53:
|
||||
# IEEE-754 double precision have 53 bits. SymPy prints them with
|
||||
# 15 digits, but we need 17 for round-trip correctness
|
||||
return str(sympy.Float(expr, dps=17))
|
||||
else:
|
||||
# We don't use other precisions in pytorch
|
||||
return str(expr)
|
||||
|
||||
# This must be implemented because sympy will collect x * x into Pow(x, 2), without
|
||||
# any explicit intervention. We print it just like x * x, notably, we
|
||||
# never generate sympy.Pow with floats.
|
||||
#
|
||||
# NB: this pow by natural, you should never have used builtin sympy.pow
|
||||
# for FloatPow, and a symbolic exponent should be PowByNatural. These
|
||||
# means exp is guaranteed to be integer.
|
||||
def _print_Pow(self, expr: sympy.Expr) -> str:
|
||||
base, exp = expr.args
|
||||
if exp != int(exp):
|
||||
raise AssertionError(exp)
|
||||
exp = int(exp)
|
||||
if exp < 0:
|
||||
raise AssertionError(f"exponent must be non-negative, got {exp}")
|
||||
if exp > 0:
|
||||
return self.stringify([base] * exp, "*", PRECEDENCE["Mul"])
|
||||
return "1"
|
||||
|
||||
# Explicit NotImplemented functions are to prevent default sympy printing
|
||||
# behavior, which will just barf out ToFloat(...) to your IR. The error
|
||||
# message is better here because it tells you which printer class it needs
|
||||
# to go in.
|
||||
|
||||
def _print_ToFloat(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_ToFloat not implemented for {type(self)}")
|
||||
|
||||
def _print_Infinity(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_Infinity not implemented for {type(self)}")
|
||||
|
||||
def _print_NegativeInfinity(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(
|
||||
f"_print_NegativeInfinity not implemented for {type(self)}"
|
||||
)
|
||||
|
||||
def _print_NaN(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_NaN not implemented for {type(self)}")
|
||||
|
||||
def _print_FloorDiv(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_FloorDiv not implemented for {type(self)}")
|
||||
|
||||
def _print_PythonMod(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_PythonMod not implemented for {type(self)}")
|
||||
|
||||
def _print_IntTrueDiv(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_IntTrueDiv not implemented for {type(self)}")
|
||||
|
||||
def _print_PowByNatural(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(
|
||||
f"_print_PowByNatural not implemented for {type(self)}"
|
||||
)
|
||||
|
||||
def _print_FloatPow(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_FloatPow not implemented for {type(self)}")
|
||||
|
||||
def _print_TruncToInt(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_TruncToInt not implemented for {type(self)}")
|
||||
|
||||
def _print_RoundToInt(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(f"_print_RoundToInt not implemented for {type(self)}")
|
||||
|
||||
def _print_RoundDecimal(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(
|
||||
f"_print_RoundDecimal not implemented for {type(self)}"
|
||||
)
|
||||
|
||||
# NB: Some float operations are INTENTIONALLY not implemented for
|
||||
# printers. You can implement them as a quick unblock, but it is better
|
||||
# to ask yourself why we haven't done this computation in the Tensor
|
||||
# universe instead
|
||||
|
||||
def _print_TruncToFloat(self, expr: sympy.Expr) -> str:
|
||||
raise NotImplementedError(
|
||||
f"_print_TruncToFloat not implemented for {type(self)}"
|
||||
)
|
||||
|
||||
|
||||
class PythonPrinter(ExprPrinter):
|
||||
def _print_ToFloat(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("ToFloat expects exactly one argument")
|
||||
# NB: We use sym_float here because the printer is used for cache
|
||||
# serialization, and cache guards get evaluated with SymInt to
|
||||
# propagate guards to the parent ShapeEnv. However, this comes at a
|
||||
# runtime cost for guards involving float. If this is unacceptable
|
||||
# overhead, what you want to do is have two separate printers for
|
||||
# SymInt, one for when the inputs are guaranteed to be int, and
|
||||
# another for when they could be SymInt.
|
||||
#
|
||||
# NB: sym_min/sym_max also have this problem, but I chose not to fix
|
||||
# those.
|
||||
#
|
||||
# See https://github.com/pytorch/pytorch/issues/142507 for more
|
||||
# context.
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"torch.sym_float({self._print(expr.args[0])})"
|
||||
|
||||
def _print_And(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " and ", precedence(expr))
|
||||
|
||||
def _print_Or(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " or ", precedence(expr))
|
||||
|
||||
def _print_ModularIndexing(self, expr: sympy.Expr) -> str:
|
||||
x, div, mod = (
|
||||
self.parenthesize(arg, PRECEDENCE["Atom"] - 0.5) for arg in expr.args
|
||||
)
|
||||
if div != "1":
|
||||
x = f"({x} // {div})"
|
||||
return f"({x} % {mod})"
|
||||
|
||||
def _print_Infinity(self, expr: sympy.Expr) -> str:
|
||||
return "math.inf"
|
||||
|
||||
def _print_NegativeInfinity(self, expr: sympy.Expr) -> str:
|
||||
return "-math.inf"
|
||||
|
||||
def _print_NaN(self, expr: sympy.Expr) -> str:
|
||||
return "math.nan"
|
||||
|
||||
# WARNING: this is dangerous for Triton, which has C-style modulus
|
||||
def _print_PythonMod(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " % ", PRECEDENCE["Atom"] - 0.5)
|
||||
|
||||
# WARNING: this is dangerous for Triton, which has C-style modulus
|
||||
def _print_FloorDiv(self, expr: sympy.Expr) -> str:
|
||||
x, div = (self.parenthesize(arg, PRECEDENCE["Atom"] - 0.5) for arg in expr.args)
|
||||
return f"{x} // {div}"
|
||||
|
||||
# WARNING: this is dangerous for Triton, when lhs, rhs > 2**53, Python
|
||||
# does a special algorithm
|
||||
def _print_IntTrueDiv(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " / ", PRECEDENCE["Atom"] - 0.5)
|
||||
|
||||
def _helper_sqrt(self, expr: sympy.Expr) -> str:
|
||||
# NB: We use torch._sym_sqrt here instead of math.sqrt because the
|
||||
# guard expression may be evaluated with SymInt/SymFloat inputs (e.g.
|
||||
# during cache hit re-evaluation in evaluate_guards_expression).
|
||||
# math.sqrt on a SymFloat triggers evaluate_expr which forces
|
||||
# concretization/specialization of the symbol, creating spurious
|
||||
# guards that didn't exist in the original program.
|
||||
# torch._sym_sqrt properly propagates through the symbolic system
|
||||
# without forcing specialization.
|
||||
# See https://github.com/pytorch/pytorch/issues/152435
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"torch._sym_sqrt({self._print(expr)})"
|
||||
|
||||
def _print_OpaqueUnaryFn_sqrt(self, expr: sympy.Expr) -> str:
|
||||
return self._helper_sqrt(expr.args[0])
|
||||
|
||||
def _print_FloatPow(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " ** ", PRECEDENCE["Pow"])
|
||||
|
||||
# TODO: Not sure this works with Triton, even when base/exp are integral
|
||||
def _print_PowByNatural(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " ** ", PRECEDENCE["Pow"])
|
||||
|
||||
def _print_floor(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("floor expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.floor({self._print(expr.args[0])})"
|
||||
|
||||
def _print_FloorToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("FloorToInt expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.floor({self._print(expr.args[0])})"
|
||||
|
||||
def _print_TruncToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("TruncToInt expects exactly one argument")
|
||||
# This also could have been int(), they'll do the same thing for float
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.trunc({self._print(expr.args[0])})"
|
||||
|
||||
def _print_ceiling(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("ceiling expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.ceil({self._print(expr.args[0])})"
|
||||
|
||||
def _print_CeilToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("CeilToInt expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.ceil({self._print(expr.args[0])})"
|
||||
|
||||
def _print_Abs(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("Abs expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"abs({self._print(expr.args[0])})"
|
||||
|
||||
# NB: It's expected that we've made explicit any promotion in the sympy
|
||||
# expression, so it doesn't matter that Python max/min doesn't perform
|
||||
# promotion
|
||||
def _print_Max(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) < 2:
|
||||
raise AssertionError("Max expects at least two arguments")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"max({', '.join(map(self._print, expr.args))})"
|
||||
|
||||
def _print_Min(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) < 2:
|
||||
raise AssertionError("Min expects at least two arguments")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"min({', '.join(map(self._print, expr.args))})"
|
||||
|
||||
def _print_OpaqueUnaryFn_cos(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("cos expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.cos({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_cosh(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("cosh expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.cosh({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_acos(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("acos expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.acos({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_sin(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("sin expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.sin({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_sinh(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("sinh expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.sinh({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_asin(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("asin expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.asin({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_tan(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("tan expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.tan({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_tanh(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("tanh expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.tanh({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_atan(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("atan expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.atan({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_log2(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("log2 expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.log2({self._print(expr.args[0])})"
|
||||
|
||||
def _print_RoundToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("RoundToInt expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"round({self._print(expr.args[0])})"
|
||||
|
||||
def _print_RoundDecimal(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 2:
|
||||
raise AssertionError("RoundDecimal expects exactly two arguments")
|
||||
number, ndigits = expr.args
|
||||
if not isinstance(ndigits, sympy.Integer):
|
||||
raise TypeError("ndigits must be an instance of sympy.Integer")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"round({self._print(number)}, {ndigits})"
|
||||
|
||||
def _print_Piecewise(self, expr: sympy.Expr) -> str:
|
||||
# Convert Piecewise(expr_cond_pairs) to nested ternary expressions
|
||||
# Piecewise((e1, c1), (e2, c2), ..., (eN, cN))
|
||||
# becomes: e1 if c1 else (e2 if c2 else (... else eN))
|
||||
result: str | None = None
|
||||
for expr_i, cond_i in reversed(expr.args):
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
expr_str = self._print(expr_i)
|
||||
if cond_i == True: # noqa: E712
|
||||
# This is the default case
|
||||
result = expr_str
|
||||
else:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
cond_str = self._print(cond_i)
|
||||
if result is None:
|
||||
result = expr_str
|
||||
else:
|
||||
result = f"({expr_str} if {cond_str} else {result})"
|
||||
return result if result else "0"
|
||||
|
||||
|
||||
class CppPrinter(ExprPrinter):
|
||||
def _print_Integer(self, expr: sympy.Expr) -> str:
|
||||
suffix = "LL" if sys.platform in ["darwin", "win32"] else "L"
|
||||
i = int(expr)
|
||||
if i > INDEX_TYPE_MAX or i < INDEX_TYPE_MIN:
|
||||
raise OverflowError(f"{i} too big to convert to {INDEX_TYPE}")
|
||||
elif i == INDEX_TYPE_MIN:
|
||||
if i != (-1) << 63:
|
||||
raise AssertionError("unexpected minimum index type value")
|
||||
# Writing -9223372036854775808L makes the value overflow
|
||||
# as it is parsed as -(9223372036854775808L) by the C/C++ compiler
|
||||
return f"(-1{suffix} << 63)"
|
||||
return f"{i}{suffix}"
|
||||
|
||||
def _print_Where(self, expr: sympy.Expr) -> str:
|
||||
c, p, q = (
|
||||
self.parenthesize(arg, PRECEDENCE["Atom"] - 0.5) for arg in expr.args
|
||||
)
|
||||
return f"{c} ? {p} : {q}"
|
||||
|
||||
def _print_Or(self, expr: sympy.Expr) -> str:
|
||||
return self.stringify(expr.args, " || ", precedence(expr))
|
||||
|
||||
def _print_Piecewise(self, expr: sympy.Expr) -> str:
|
||||
# Convert Piecewise(expr_cond_pairs) to nested ternary operators
|
||||
# Piecewise((e1, c1), (e2, c2), ..., (eN, cN))
|
||||
# becomes: c1 ? e1 : (c2 ? e2 : (... : eN))
|
||||
result: str | None = None
|
||||
for expr_i, cond_i in reversed(expr.args):
|
||||
expr_str = self.parenthesize(expr_i, PRECEDENCE["Atom"] - 0.5)
|
||||
if cond_i == True: # noqa: E712
|
||||
# This is the default case
|
||||
result = expr_str
|
||||
else:
|
||||
cond_str = self.parenthesize(cond_i, PRECEDENCE["Atom"] - 0.5)
|
||||
if result is None:
|
||||
result = expr_str
|
||||
else:
|
||||
result = f"{cond_str} ? {expr_str} : {result}"
|
||||
return f"({result})" if result else "0"
|
||||
|
||||
def _print_ModularIndexing(self, expr: sympy.Expr) -> str:
|
||||
x, div, mod = expr.args
|
||||
x = self.doprint(x)
|
||||
if div != 1:
|
||||
div = self.doprint(div)
|
||||
if expr.is_integer:
|
||||
x = f"c10::div_floor_integer(static_cast<int64_t>({x}), static_cast<int64_t>({div}))"
|
||||
else:
|
||||
x = f"c10::div_floor_floating(static_cast<double>({x}), static_cast<double>({div}))"
|
||||
mod = self.doprint(mod)
|
||||
return f"(static_cast<{INDEX_TYPE}>({x}) % static_cast<{INDEX_TYPE}>({mod}))"
|
||||
|
||||
def _print_FloorDiv(self, expr: sympy.Expr) -> str:
|
||||
x, div = expr.args
|
||||
x = self.doprint(x)
|
||||
div = self.doprint(div)
|
||||
if expr.is_integer:
|
||||
return f"c10::div_floor_integer(static_cast<int64_t>({x}), static_cast<int64_t>({div}))"
|
||||
return f"c10::div_floor_floating(static_cast<double>({x}), static_cast<double>({div}))"
|
||||
|
||||
def _print_floor(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("floor expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
r = f"std::floor({self._print(expr.args[0])})"
|
||||
return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r
|
||||
|
||||
def _print_FloorToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("FloorToInt expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
r = f"std::floor({self._print(expr.args[0])})"
|
||||
return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r
|
||||
|
||||
def _print_TruncToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("TruncToInt expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
r = f"std::trunc({self._print(expr.args[0])})"
|
||||
return f"static_cast<{INDEX_TYPE}>({r})"
|
||||
|
||||
def _print_TruncToFloat(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("TruncToFloat expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::trunc({self._print(expr.args[0])})"
|
||||
|
||||
def _print_ToFloat(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("ToFloat expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"static_cast<double>({self._print(expr.args[0])})"
|
||||
|
||||
def _print_PythonMod(self, expr: sympy.Expr) -> str:
|
||||
x, div = expr.args
|
||||
x = self.doprint(x)
|
||||
div = self.doprint(div)
|
||||
return f"c10::div_mod({x}, {div})"
|
||||
|
||||
def _print_IntTrueDiv(self, expr: sympy.Expr) -> str:
|
||||
lhs, rhs = expr.args
|
||||
# TODO: This is only accurate up to 2**53
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"static_cast<double>({self._print(lhs)}) / static_cast<double>({self._print(rhs)})"
|
||||
|
||||
# TODO: PowByNatural: we need to implement our own int-int pow. Do NOT
|
||||
# use std::pow, that operates on floats
|
||||
def _print_PowByNatural(self, expr: sympy.Expr) -> str:
|
||||
# Implement the special-case of 2**x for now
|
||||
base, exp = expr.args
|
||||
if base == 2:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"(1 << ({self._print(exp)}))"
|
||||
raise NotImplementedError(
|
||||
f"_print_PowByNatural not implemented for {type(self)}"
|
||||
)
|
||||
|
||||
def _print_FloatPow(self, expr: sympy.Expr) -> str:
|
||||
base, exp = expr.args
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::pow({self._print(base)}, {self._print(exp)})"
|
||||
|
||||
def _print_Pow(self, expr: sympy.Expr) -> str:
|
||||
# Uses float constants to perform FP div
|
||||
base, exp = expr.args
|
||||
|
||||
if exp == 0.5 or exp == -0.5:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
base = self._print(base)
|
||||
return f"std::sqrt({base})" if exp == 0.5 else f"1.0/std::sqrt({base})"
|
||||
if exp.is_integer:
|
||||
exp = int(exp)
|
||||
if exp > 0:
|
||||
r = self.stringify([base] * exp, "*", PRECEDENCE["Mul"])
|
||||
elif exp < -1:
|
||||
r = (
|
||||
"1.0/("
|
||||
+ self.stringify([base] * abs(exp), "*", PRECEDENCE["Mul"])
|
||||
+ ")"
|
||||
)
|
||||
elif exp == -1:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
r = "1.0/" + self._print(base)
|
||||
else: # exp == 0
|
||||
r = "1.0"
|
||||
|
||||
return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r
|
||||
else:
|
||||
# TODO: float vs double
|
||||
return f"std::pow({base}, {float(exp)})"
|
||||
|
||||
def _print_Rational(self, expr: sympy.Expr) -> str:
|
||||
# Uses float constants to perform FP div
|
||||
if expr.q == 1:
|
||||
r = f"{expr.p}"
|
||||
else:
|
||||
r = f"{expr.p}.0/{expr.q}.0"
|
||||
return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r
|
||||
|
||||
def _print_ceiling(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("ceiling expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
r = f"std::ceil({self._print(expr.args[0])})"
|
||||
return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r
|
||||
|
||||
def _print_CeilToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("CeilToInt expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
r = f"std::ceil({self._print(expr.args[0])})"
|
||||
return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r
|
||||
|
||||
def _print_Min(self, expr: sympy.Expr) -> str:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
args = [self._print(a) for a in expr.args]
|
||||
if len(args) == 2:
|
||||
return f"std::min(static_cast<{INDEX_TYPE}>({args[0]}), static_cast<{INDEX_TYPE}>({args[1]}))"
|
||||
else:
|
||||
# Initializer list overload
|
||||
il = "{" + ", ".join(args) + "}"
|
||||
return f"std::min<{INDEX_TYPE}>({il})"
|
||||
|
||||
def _print_Max(self, expr: sympy.Expr) -> str:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
args = [self._print(a) for a in expr.args]
|
||||
if len(args) == 2:
|
||||
return f"std::max(static_cast<{INDEX_TYPE}>({args[0]}), static_cast<{INDEX_TYPE}>({args[1]}))"
|
||||
else:
|
||||
# Initializer list overload
|
||||
il = "{" + ", ".join(args) + "}"
|
||||
return f"std::max<{INDEX_TYPE}>({il})"
|
||||
|
||||
def _print_Abs(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("Abs expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::abs({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_cos(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("cos expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::cos({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_cosh(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("cosh expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::cosh({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_acos(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("acos expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::acos({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_sin(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("sin expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"math.sin({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_sinh(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("sinh expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::sinh({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_asin(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("asin expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::asin({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_tan(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("tan expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::tan({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_tanh(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("tanh expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::tanh({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_atan(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("atan expects exactly one argument")
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::atan({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_sqrt(self, expr: sympy.Expr) -> str:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::sqrt({self._print(expr.args[0])})"
|
||||
|
||||
def _print_OpaqueUnaryFn_log2(self, expr: sympy.Expr) -> str:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::log2({self._print(expr.args[0])})"
|
||||
|
||||
def _print_RoundToInt(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 1:
|
||||
raise AssertionError("RoundToInt expects exactly one argument")
|
||||
# TODO: dispatch to llrint depending on index type
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"std::lrint({self._print(expr.args[0])})"
|
||||
|
||||
def _print_RoundDecimal(self, expr: sympy.Expr) -> str:
|
||||
if len(expr.args) != 2:
|
||||
raise AssertionError("RoundDecimal expects exactly two arguments")
|
||||
number, ndigits = expr.args
|
||||
if number.is_integer:
|
||||
# ndigits < 0 should have been filtered by the sympy function
|
||||
if ndigits >= 0:
|
||||
raise AssertionError("ndigits must be negative for integer inputs")
|
||||
raise ValueError(
|
||||
f"For integer inputs, only non-negative ndigits are currently supported, but got {ndigits}."
|
||||
)
|
||||
number_str = self.parenthesize(number, PRECEDENCE["Mul"])
|
||||
return f"static_cast<double>(std::nearbyint(1e{ndigits} * {number_str}) * 1e{-ndigits})"
|
||||
|
||||
def _print_BooleanTrue(self, expr: sympy.Expr) -> str:
|
||||
return "true"
|
||||
|
||||
def _print_BooleanFalse(self, expr: sympy.Expr) -> str:
|
||||
return "false"
|
||||
|
||||
def _print_Infinity(self, expr: sympy.Expr) -> str:
|
||||
return "std::numeric_limits<double>::infinity()"
|
||||
|
||||
def _print_NegativeInfinity(self, expr: sympy.Expr) -> str:
|
||||
return f"-{self._print_Infinity(expr)}"
|
||||
|
||||
def _print_NaN(self, expr: sympy.Expr) -> str:
|
||||
return "std::numeric_limits<double>::quiet_NaN()"
|
||||
@@ -0,0 +1,615 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import math
|
||||
import operator
|
||||
from typing import NoReturn
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch.utils._sympy.functions import (
|
||||
_keep_float,
|
||||
BitwiseFn_bitwise_and,
|
||||
BitwiseFn_bitwise_or,
|
||||
BitwiseFn_bitwise_xor,
|
||||
FloatPow,
|
||||
FloatTrueDiv,
|
||||
FloorDiv,
|
||||
IntTrueDiv,
|
||||
Max,
|
||||
Min,
|
||||
Mod,
|
||||
OpaqueUnaryFn_exp,
|
||||
OpaqueUnaryFn_log,
|
||||
OpaqueUnaryFn_log2,
|
||||
OpaqueUnaryFn_sqrt,
|
||||
PowByNatural,
|
||||
RoundDecimal,
|
||||
RoundToInt,
|
||||
ToFloat,
|
||||
TruncToInt,
|
||||
)
|
||||
|
||||
|
||||
# The sympy interpretation of operators. It will also sometimes work with
|
||||
# plain int/float, but if you do certain operations you will get out a
|
||||
# sympy.Basic in the end. If you want the Python/FX traceable interpretation,
|
||||
# check PythonReferenceAnalysis.
|
||||
# NB: For magic methods this needs to use normal magic methods
|
||||
# so that test_magic_methods works
|
||||
class ReferenceAnalysis:
|
||||
@staticmethod
|
||||
def constant(c, dtype):
|
||||
return sympy.sympify(c)
|
||||
|
||||
@staticmethod
|
||||
def or_(a, b):
|
||||
return a | b
|
||||
|
||||
@staticmethod
|
||||
def and_(a, b):
|
||||
return a & b
|
||||
|
||||
@staticmethod
|
||||
def eq(a, b):
|
||||
if isinstance(a, sympy.Expr) or isinstance(b, sympy.Expr):
|
||||
return sympy.Eq(a, b)
|
||||
return a == b
|
||||
|
||||
@classmethod
|
||||
def ne(cls, a, b):
|
||||
return cls.not_(cls.eq(a, b))
|
||||
|
||||
@staticmethod
|
||||
def lt(a, b):
|
||||
return a < b
|
||||
|
||||
@staticmethod
|
||||
def gt(a, b):
|
||||
return a > b
|
||||
|
||||
@staticmethod
|
||||
def le(a, b):
|
||||
return a <= b
|
||||
|
||||
@staticmethod
|
||||
def ge(a, b):
|
||||
return a >= b
|
||||
|
||||
@staticmethod
|
||||
def not_(a):
|
||||
if isinstance(a, bool):
|
||||
raise AssertionError("not_ needs sympy expr")
|
||||
return ~a
|
||||
|
||||
@staticmethod
|
||||
def reciprocal(x):
|
||||
return FloatTrueDiv(1.0, x)
|
||||
|
||||
@staticmethod
|
||||
def square(x):
|
||||
return PowByNatural(x, 2)
|
||||
|
||||
@staticmethod
|
||||
def trunc_to_int(x, dtype):
|
||||
return TruncToInt(x)
|
||||
|
||||
@staticmethod
|
||||
def ceil_to_int(x, dtype):
|
||||
return sympy.ceiling(x)
|
||||
|
||||
@staticmethod
|
||||
def floor_to_int(x, dtype):
|
||||
return sympy.floor(x)
|
||||
|
||||
@staticmethod
|
||||
def floor(x):
|
||||
return _keep_float(sympy.floor)(x)
|
||||
|
||||
@staticmethod
|
||||
def ceil(x):
|
||||
return _keep_float(sympy.ceiling)(x)
|
||||
|
||||
@staticmethod
|
||||
def to_dtype(x, dtype):
|
||||
if dtype == torch.float64:
|
||||
return ToFloat(x)
|
||||
raise NotImplementedError(f"to_dtype {dtype} NYI")
|
||||
|
||||
@staticmethod
|
||||
def mod(x, y):
|
||||
return Mod(x, y)
|
||||
|
||||
@staticmethod
|
||||
def abs(x):
|
||||
return abs(x)
|
||||
|
||||
@staticmethod
|
||||
def neg(x):
|
||||
return -x
|
||||
|
||||
@staticmethod
|
||||
def truediv(a, b):
|
||||
return FloatTrueDiv(a, b)
|
||||
|
||||
@staticmethod
|
||||
def int_truediv(a, b):
|
||||
return IntTrueDiv(a, b)
|
||||
|
||||
@staticmethod
|
||||
def floordiv(a, b):
|
||||
return FloorDiv(a, b)
|
||||
|
||||
@staticmethod
|
||||
def truncdiv(a, b) -> NoReturn:
|
||||
raise NotImplementedError("TODO: truncdiv")
|
||||
|
||||
@staticmethod
|
||||
def add(a, b):
|
||||
return _keep_float(operator.add)(a, b)
|
||||
|
||||
@classmethod
|
||||
def sym_sum(cls, args):
|
||||
return sympy.Add(*args)
|
||||
|
||||
@staticmethod
|
||||
def mul(a, b):
|
||||
return _keep_float(operator.mul)(a, b)
|
||||
|
||||
@staticmethod
|
||||
def sub(a, b):
|
||||
return _keep_float(operator.sub)(a, b)
|
||||
|
||||
@staticmethod
|
||||
def exp(x):
|
||||
return OpaqueUnaryFn_exp(x)
|
||||
|
||||
@staticmethod
|
||||
def log(x):
|
||||
return OpaqueUnaryFn_log(x)
|
||||
|
||||
@staticmethod
|
||||
def log2(x):
|
||||
return OpaqueUnaryFn_log2(x)
|
||||
|
||||
@staticmethod
|
||||
def sqrt(x):
|
||||
return OpaqueUnaryFn_sqrt(x)
|
||||
|
||||
@staticmethod
|
||||
def pow(a, b):
|
||||
# pyrefly: ignore [bad-argument-count, bad-argument-type]
|
||||
return _keep_float(FloatPow)(a, b)
|
||||
|
||||
@staticmethod
|
||||
def pow_by_natural(a, b):
|
||||
return PowByNatural(a, b)
|
||||
|
||||
@staticmethod
|
||||
def minimum(a, b):
|
||||
return Min(a, b)
|
||||
|
||||
@staticmethod
|
||||
def maximum(a, b):
|
||||
return Max(a, b)
|
||||
|
||||
@staticmethod
|
||||
def round_to_int(a, dtype):
|
||||
return RoundToInt(a)
|
||||
|
||||
@staticmethod
|
||||
def round_decimal(a, b):
|
||||
return RoundDecimal(a, b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_and(a, b):
|
||||
return BitwiseFn_bitwise_and(a, b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_or(a, b):
|
||||
return BitwiseFn_bitwise_or(a, b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_xor(a, b):
|
||||
return BitwiseFn_bitwise_xor(a, b)
|
||||
|
||||
|
||||
# Unlike ReferenceAnalysis, does NOT sympyify, instead, works with plain
|
||||
# Python types and is FX traceable. Inheritance here is purely for code
|
||||
# sharing (TODO: considering splitting out a BaseReferenceAnalysis).
|
||||
class PythonReferenceAnalysis(ReferenceAnalysis):
|
||||
@staticmethod
|
||||
def constant(c, dtype):
|
||||
if dtype is torch.int64:
|
||||
return int(c)
|
||||
elif dtype is torch.double:
|
||||
return float(c)
|
||||
elif dtype is torch.bool:
|
||||
return bool(c)
|
||||
else:
|
||||
raise AssertionError(f"unrecognized dtype {dtype}")
|
||||
|
||||
@staticmethod
|
||||
def not_(a):
|
||||
return torch.sym_not(a)
|
||||
|
||||
@classmethod
|
||||
def sym_sum(cls, args):
|
||||
if len(args) == 0:
|
||||
return 0
|
||||
if len(args) == 1:
|
||||
return args[0]
|
||||
acc = cls.add(args[0], args[1])
|
||||
for i in range(2, len(args)):
|
||||
acc = cls.add(acc, args[i])
|
||||
return acc
|
||||
|
||||
@staticmethod
|
||||
def floordiv(a, b):
|
||||
return a // b
|
||||
|
||||
@staticmethod
|
||||
def mod(x, y):
|
||||
return x % y
|
||||
|
||||
@staticmethod
|
||||
def python_mod(x, y):
|
||||
return x % y
|
||||
|
||||
@staticmethod
|
||||
def truncdiv(a, b):
|
||||
return a / b
|
||||
|
||||
@staticmethod
|
||||
def to_dtype(x, dtype):
|
||||
if dtype == torch.float64:
|
||||
return torch.sym_float(x)
|
||||
raise NotImplementedError(f"to_dtype {dtype} NYI")
|
||||
|
||||
@staticmethod
|
||||
def exp(x) -> NoReturn:
|
||||
raise AssertionError("exp is not valid shape sympy expr")
|
||||
|
||||
@staticmethod
|
||||
def log(x) -> NoReturn:
|
||||
raise AssertionError("log is not valid shape sympy expr")
|
||||
|
||||
@staticmethod
|
||||
def log2(x):
|
||||
return torch._sym_log2(x) # type: ignore[attr-defined]
|
||||
|
||||
@staticmethod
|
||||
def sqrt(x):
|
||||
return torch._sym_sqrt(x) # type: ignore[attr-defined]
|
||||
|
||||
@staticmethod
|
||||
def minimum(a, b):
|
||||
return torch.sym_min(a, b)
|
||||
|
||||
@staticmethod
|
||||
def maximum(a, b):
|
||||
return torch.sym_max(a, b)
|
||||
|
||||
@staticmethod
|
||||
def floor_to_int(x, dtype):
|
||||
return math.floor(x)
|
||||
|
||||
@staticmethod
|
||||
def ceil_to_int(x, dtype):
|
||||
return math.ceil(x)
|
||||
|
||||
@staticmethod
|
||||
def floor(x):
|
||||
return float(math.floor(x))
|
||||
|
||||
@staticmethod
|
||||
def ceil(x):
|
||||
return float(math.ceil(x))
|
||||
|
||||
@staticmethod
|
||||
def truediv(a, b):
|
||||
return a / b
|
||||
|
||||
@staticmethod
|
||||
def pow(a, b):
|
||||
return a**b
|
||||
|
||||
@staticmethod
|
||||
def pow_by_natural(a, b):
|
||||
# Pray that safe_pow is not needed here lol. In particular, this
|
||||
# never participates in VR low/high ranges, so overflow should be
|
||||
# unlikely
|
||||
return a**b
|
||||
|
||||
@staticmethod
|
||||
def round_to_int(a, dtype):
|
||||
return round(a)
|
||||
|
||||
@staticmethod
|
||||
def round_decimal(a, b):
|
||||
return round(a, ndigits=b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_and(a, b):
|
||||
return a & b
|
||||
|
||||
@staticmethod
|
||||
def bitwise_or(a, b):
|
||||
return a | b
|
||||
|
||||
@staticmethod
|
||||
def bitwise_xor(a, b):
|
||||
return a ^ b
|
||||
|
||||
@staticmethod
|
||||
def expr_cond_pair(expr, cond):
|
||||
return (expr, cond)
|
||||
|
||||
@staticmethod
|
||||
def piecewise(*pairs):
|
||||
# Build nested sym_ite from right to left.
|
||||
# Piecewise((e1, c1), (e2, c2), ..., (en, True)) becomes
|
||||
# sym_ite(c1, e1, sym_ite(c2, e2, ... en))
|
||||
result = pairs[-1][0]
|
||||
for expr, cond in reversed(pairs[:-1]):
|
||||
result = torch.sym_ite(cond, expr, result)
|
||||
return result
|
||||
|
||||
|
||||
# Like PythonReferenceAnalysis, but some export-unfriendly choices of
|
||||
# operators to make things faster
|
||||
class OptimizedPythonReferenceAnalysis(PythonReferenceAnalysis):
|
||||
@staticmethod
|
||||
def sym_sum(args):
|
||||
return torch.sym_sum(args)
|
||||
|
||||
|
||||
def _to_dtype(x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
|
||||
return torch.ops.prims.convert_element_type.default(x, dtype)
|
||||
|
||||
|
||||
# Suppose we have some int/float arguments. This diagram commutes:
|
||||
#
|
||||
# int/float -- PythonReferenceAnalysis.op --> int/float
|
||||
# | |
|
||||
# | |
|
||||
# torch.tensor(..., dtype=torch.int64/torch.float64)
|
||||
# | |
|
||||
# V V
|
||||
# Tensor -- TensorReferenceAnalysis.op --> Tensor
|
||||
#
|
||||
# NB: int before and after must be representable in int64 (we will
|
||||
# insert guards accordingly.)
|
||||
#
|
||||
# This is guaranteed to be FX traceable with OpOverloads only.
|
||||
class TensorReferenceAnalysis:
|
||||
# NB: This is actually dead, because with Proxy tracing the factory
|
||||
# function isn't traced correctly. Here for completeness.
|
||||
@staticmethod
|
||||
def constant(c, dtype):
|
||||
d: int | float | bool
|
||||
if dtype is torch.int64:
|
||||
d = int(c)
|
||||
elif dtype is torch.double:
|
||||
d = float(c)
|
||||
elif dtype is torch.bool:
|
||||
d = bool(c)
|
||||
else:
|
||||
raise AssertionError(f"unrecognized dtype {dtype}")
|
||||
return torch.ops.aten.scalar_tensor.default(d, dtype=dtype)
|
||||
|
||||
@staticmethod
|
||||
def or_(a, b):
|
||||
return torch.ops.aten.logical_or.default(a, b)
|
||||
|
||||
@staticmethod
|
||||
def and_(a, b):
|
||||
return torch.ops.aten.logical_and.default(a, b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_and(a, b):
|
||||
return torch.ops.aten.bitwise_and(a, b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_or(a, b):
|
||||
return torch.ops.aten.bitwise_or(a, b)
|
||||
|
||||
@staticmethod
|
||||
def bitwise_xor(a, b):
|
||||
return torch.ops.aten.bitwise_xor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def eq(a, b):
|
||||
return torch.ops.aten.eq.Tensor(a, b)
|
||||
|
||||
@classmethod
|
||||
def ne(cls, a, b):
|
||||
return torch.ops.aten.ne.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def lt(a, b):
|
||||
return torch.ops.aten.lt.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def gt(a, b):
|
||||
return torch.ops.aten.gt.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def le(a, b):
|
||||
return torch.ops.aten.le.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def ge(a, b):
|
||||
return torch.ops.aten.ge.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def not_(a):
|
||||
return torch.ops.aten.logical_not.default(a)
|
||||
|
||||
@staticmethod
|
||||
def reciprocal(x):
|
||||
return torch.ops.aten.reciprocal.default(x)
|
||||
|
||||
@staticmethod
|
||||
def square(x):
|
||||
# TODO: maybe composite implicit autograd doesn't work here?
|
||||
return torch.ops.aten.square.default(x)
|
||||
|
||||
@staticmethod
|
||||
def trunc_to_int(x, dtype):
|
||||
return _to_dtype(torch.ops.aten.trunc.default(x), dtype)
|
||||
|
||||
@staticmethod
|
||||
def ceil_to_int(x, dtype):
|
||||
return _to_dtype(torch.ops.aten.ceil.default(x), dtype)
|
||||
|
||||
@staticmethod
|
||||
def floor_to_int(x, dtype):
|
||||
return _to_dtype(torch.ops.aten.floor.default(x), dtype)
|
||||
|
||||
@staticmethod
|
||||
def floor(x):
|
||||
return torch.ops.aten.floor.default(x)
|
||||
|
||||
@staticmethod
|
||||
def ceil(x):
|
||||
return torch.ops.aten.ceil.default(x)
|
||||
|
||||
@staticmethod
|
||||
def to_dtype(x, dtype):
|
||||
return _to_dtype(x, dtype)
|
||||
|
||||
@staticmethod
|
||||
def mod(x, y) -> NoReturn:
|
||||
# TODO: https://github.com/pytorch/pytorch/pull/133654
|
||||
raise NotImplementedError(
|
||||
"no C-style modulus operation available from frontend atm"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def abs(x):
|
||||
return torch.ops.aten.abs.default(x)
|
||||
|
||||
@staticmethod
|
||||
def neg(x):
|
||||
return torch.ops.aten.neg.default(x)
|
||||
|
||||
@staticmethod
|
||||
def truediv(a, b):
|
||||
return torch.ops.aten.true_divide.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def int_truediv(a, b):
|
||||
raise NotImplementedError(
|
||||
"Python int truediv difficult to implement in PyTorch atm"
|
||||
)
|
||||
|
||||
# TODO: This is wrong, CPython has a custom implementation of true
|
||||
# division that results in higher precision when the floats are
|
||||
# sufficiently large. Short term fix: add a guard here
|
||||
# pyrefly: ignore [unreachable]
|
||||
return torch.ops.aten.true_divide.default(
|
||||
_to_dtype(a, torch.float64), _to_dtype(b, torch.float64)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def floordiv(a, b):
|
||||
return torch.ops.aten.div.Tensor_mode(a, b, rounding_mode="floor")
|
||||
|
||||
@staticmethod
|
||||
def truncdiv(a, b) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"no C-style truncdiv operation available from frontend atm"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def add(a, b):
|
||||
return torch.ops.aten.add.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def mul(a, b):
|
||||
return torch.ops.aten.mul.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def sub(a, b):
|
||||
return torch.ops.aten.sub.Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def exp(x):
|
||||
return torch.ops.aten.exp.default(x)
|
||||
|
||||
@staticmethod
|
||||
def log(x):
|
||||
return torch.ops.aten.log.default(x)
|
||||
|
||||
@staticmethod
|
||||
def log2(x):
|
||||
return torch.ops.aten.log2.default(x)
|
||||
|
||||
@staticmethod
|
||||
def sqrt(x):
|
||||
return torch.ops.aten.sqrt.default(x)
|
||||
|
||||
@staticmethod
|
||||
def sin(x):
|
||||
return torch.ops.aten.sin.default(x)
|
||||
|
||||
@staticmethod
|
||||
def cos(x):
|
||||
return torch.ops.aten.cos.default(x)
|
||||
|
||||
@staticmethod
|
||||
def tanh(x):
|
||||
return torch.ops.aten.tanh.default(x)
|
||||
|
||||
@staticmethod
|
||||
def sinh(x):
|
||||
return torch.ops.aten.sinh.default(x)
|
||||
|
||||
@staticmethod
|
||||
def cosh(x):
|
||||
return torch.ops.aten.cosh.default(x)
|
||||
|
||||
@staticmethod
|
||||
def tan(x):
|
||||
return torch.ops.aten.tan.default(x)
|
||||
|
||||
@staticmethod
|
||||
def acos(x):
|
||||
return torch.ops.aten.acos.default(x)
|
||||
|
||||
@staticmethod
|
||||
def atan(x):
|
||||
return torch.ops.aten.atan.default(x)
|
||||
|
||||
@staticmethod
|
||||
def asin(x):
|
||||
return torch.ops.aten.asin.default(x)
|
||||
|
||||
@staticmethod
|
||||
def pow(a, b):
|
||||
return torch.ops.aten.pow.Tensor_Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def pow_by_natural(a, b):
|
||||
# NB: pow handles int x int fine
|
||||
return torch.ops.aten.pow.Tensor_Tensor(a, b)
|
||||
|
||||
@staticmethod
|
||||
def minimum(a, b):
|
||||
return torch.ops.aten.minimum.default(a, b)
|
||||
|
||||
@staticmethod
|
||||
def maximum(a, b):
|
||||
return torch.ops.aten.maximum.default(a, b)
|
||||
|
||||
@staticmethod
|
||||
def round_to_int(a, dtype):
|
||||
return torch.ops.aten.round.default(a)
|
||||
|
||||
@staticmethod
|
||||
def round_decimal(a, b) -> NoReturn:
|
||||
raise NotImplementedError(
|
||||
"round decimal doesn't support Tensor second argument atm"
|
||||
)
|
||||
|
||||
# return torch.ops.aten.round.decimals(a, b)
|
||||
@@ -0,0 +1,96 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import sympy
|
||||
from sympy.multipledispatch import dispatch
|
||||
|
||||
|
||||
__all__ = ["SingletonInt"]
|
||||
|
||||
|
||||
class SingletonInt(sympy.AtomicExpr):
|
||||
# This is probably not super important unless we are in multiple dispatch
|
||||
# situations with other more exotic Expr types.
|
||||
_op_priority = 99999
|
||||
|
||||
def __new__(cls, *args, coeff=None, **kwargs):
|
||||
instance = super().__new__(cls, *args, **kwargs)
|
||||
return instance
|
||||
|
||||
# The semantics of this class should match that of NestedIntSymNodeImpl in
|
||||
# c10/core/NestedIntSymNodeImpl.h
|
||||
def __init__(self, val, *, coeff=1) -> None:
|
||||
self._val = val
|
||||
self._coeff = coeff
|
||||
super().__init__()
|
||||
|
||||
# See NOTE [ Inequalities with nested int ]
|
||||
def _eval_Eq(self, other):
|
||||
if (
|
||||
isinstance(other, SingletonInt)
|
||||
and other._val == self._val
|
||||
and self._coeff == other._coeff
|
||||
):
|
||||
return sympy.true
|
||||
else:
|
||||
return sympy.false
|
||||
|
||||
# This is necessary so that calling expr.free_symbols on exprs that contain
|
||||
# this Singleton does not error
|
||||
@property
|
||||
def free_symbols(self):
|
||||
return set()
|
||||
|
||||
def __mul__(self, other):
|
||||
if isinstance(other, SingletonInt):
|
||||
raise ValueError(
|
||||
"SingletonInt cannot be multiplied by another SingletonInt"
|
||||
)
|
||||
return SingletonInt(self._val, coeff=self._coeff * other)
|
||||
|
||||
def __rmul__(self, other):
|
||||
if isinstance(other, SingletonInt):
|
||||
raise ValueError(
|
||||
"SingletonInt cannot be multiplied by another SingletonInt"
|
||||
)
|
||||
return SingletonInt(self._val, coeff=self._coeff * other)
|
||||
|
||||
# Make sure we promptly raise an error instead of falling back to building
|
||||
# an expression tree. There are probably more ops, how can we be exhaustive?
|
||||
def __add__(self, other):
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
def __sub__(self, other):
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
def __truediv__(self, other):
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
def __floordiv__(self, other):
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
def __mod__(self, other):
|
||||
raise NotImplementedError("NYI")
|
||||
|
||||
|
||||
# See NOTE [ Inequalities with nested int ]
|
||||
@dispatch(sympy.Integer, SingletonInt)
|
||||
def _eval_is_ge(a, b):
|
||||
if a < 2:
|
||||
return sympy.false
|
||||
raise ValueError("Symbolic SingletonInt: Relation is indeterminate")
|
||||
|
||||
|
||||
@dispatch(SingletonInt, sympy.Integer) # type: ignore[no-redef]
|
||||
def _eval_is_ge(a, b): # noqa: F811
|
||||
if b <= 2:
|
||||
return sympy.true
|
||||
raise ValueError("Symbolic SingletonInt: Relation is indeterminate")
|
||||
|
||||
|
||||
@dispatch(SingletonInt, SingletonInt) # type: ignore[no-redef]
|
||||
def _eval_is_ge(a, b): # noqa: F811
|
||||
if a._val == b._val:
|
||||
if a._coeff >= b._coeff:
|
||||
return sympy.true
|
||||
else:
|
||||
return sympy.false
|
||||
raise ValueError("Symbolic SingletonInt: Relation is indeterminate")
|
||||
@@ -0,0 +1,179 @@
|
||||
import logging
|
||||
|
||||
import sympy
|
||||
|
||||
from torch.utils._sympy.functions import FloorDiv
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_MIRROR_REL_OP: dict[type[sympy.Basic], type[sympy.Rel]] = {
|
||||
sympy.Eq: sympy.Eq,
|
||||
sympy.Ne: sympy.Ne,
|
||||
sympy.Ge: sympy.Le,
|
||||
sympy.Gt: sympy.Lt,
|
||||
sympy.Le: sympy.Ge,
|
||||
sympy.Lt: sympy.Gt,
|
||||
}
|
||||
|
||||
INEQUALITY_TYPES = (sympy.Gt, sympy.Ge, sympy.Lt, sympy.Le)
|
||||
|
||||
|
||||
def mirror_rel_op(type: type) -> type[sympy.Rel] | None:
|
||||
return _MIRROR_REL_OP.get(type)
|
||||
|
||||
|
||||
# Tries to simplify 'expr', so as to leave only 'thing' in the left-hand side.
|
||||
#
|
||||
# Returns a tuple of:
|
||||
# 1. The simplified expression
|
||||
# 2. The expression on the right-hand side
|
||||
#
|
||||
# Returns 'None' if it can't reach a state where the only thing in the left
|
||||
# hand side is 'thing'.
|
||||
#
|
||||
# 'trials': number of times 'try_solve' will try to isolate 'thing' to the
|
||||
# left-hand side.
|
||||
#
|
||||
# 'floordiv_inequality': flag to enable conversion of 'FloorDiv' into
|
||||
# inequalities.
|
||||
def try_solve(
|
||||
expr: sympy.Basic,
|
||||
thing: sympy.Basic,
|
||||
trials: int = 5,
|
||||
floordiv_inequality: bool = True,
|
||||
) -> tuple[sympy.Rel, sympy.Expr] | None:
|
||||
mirror = mirror_rel_op(type(expr))
|
||||
|
||||
# Ignore unsupported expressions:
|
||||
# - Those that are not relational operations
|
||||
# - Those that don't have a mirror (just avoiding unexpected classes)
|
||||
if not isinstance(expr, sympy.Rel) or mirror is None:
|
||||
log.debug("expression with unsupported type: %s", type(expr))
|
||||
return None
|
||||
|
||||
lhs_has_thing = expr.lhs.has(thing)
|
||||
rhs_has_thing = expr.rhs.has(thing)
|
||||
|
||||
# Give up when 'thing' appears on both sides of the relational expression.
|
||||
# That is because, as is, we assume the thing we are trying to isolate is
|
||||
# only on the right-hand side.
|
||||
if lhs_has_thing and rhs_has_thing:
|
||||
log.debug("thing (%s) found in both sides of expression: %s", thing, expr)
|
||||
return None
|
||||
|
||||
# Try considering both LHS and RHS by mirroring the original expression:
|
||||
# a < b ==> b > a
|
||||
expressions = []
|
||||
|
||||
# Add each version of 'expr' if 'thing' is in its left-hand side.
|
||||
if lhs_has_thing:
|
||||
expressions.append(expr)
|
||||
if rhs_has_thing:
|
||||
expressions.append(mirror(expr.rhs, expr.lhs))
|
||||
|
||||
for e in expressions:
|
||||
if e is None:
|
||||
continue
|
||||
|
||||
if not isinstance(e, sympy.Rel):
|
||||
raise AssertionError("expected sympy.Rel")
|
||||
|
||||
for _ in range(trials):
|
||||
trial = _try_isolate_lhs(e, thing, floordiv_inequality=floordiv_inequality)
|
||||
# Stop if there was no change in this trial.
|
||||
if trial == e:
|
||||
break
|
||||
e = trial # type: ignore[assignment]
|
||||
|
||||
# Return if we were able to isolate 'thing' on the left-hand side.
|
||||
if isinstance(e, sympy.Rel) and e.lhs == thing:
|
||||
log.debug("solved: %s ---> %s", expr, e)
|
||||
return e, e.rhs
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _try_isolate_lhs(
|
||||
e: sympy.Basic, thing: sympy.Basic, floordiv_inequality: bool
|
||||
) -> sympy.Basic:
|
||||
op = type(e)
|
||||
|
||||
if isinstance(e, sympy.Rel):
|
||||
# Move any constants in the left-hand side to the right-hand side.
|
||||
lhs_not_thing = (
|
||||
sum(a for a in e.lhs.args if not a.has(thing))
|
||||
if isinstance(e.lhs, sympy.Add)
|
||||
else 0
|
||||
)
|
||||
e = op(e.lhs - lhs_not_thing, e.rhs - lhs_not_thing) # type: ignore[attr-defined]
|
||||
|
||||
# Divide both sides by the factors that don't contain thing.
|
||||
if isinstance(e, sympy.Rel) and isinstance(e.lhs, sympy.Mul):
|
||||
lhs, rhs = e.args
|
||||
other = sympy.Mul(*[a for a in lhs.args if not a.has(thing)])
|
||||
|
||||
# If we can't tell whether 'other' is negative or positive, we do nothing.
|
||||
# That is because we don't know whether we have mirror the operation or not.
|
||||
# We also divide only when we know 'rhs' is not zero.
|
||||
if not (isinstance(e, INEQUALITY_TYPES) and other.is_negative is None) and not (
|
||||
not isinstance(e, INEQUALITY_TYPES) and rhs.is_zero
|
||||
):
|
||||
# Divide both sides by 'other'.
|
||||
lhs = lhs / other
|
||||
rhs = rhs / other
|
||||
|
||||
# If 'e' is an inequality and 'other' is negative, we have to
|
||||
# mirror the expression.
|
||||
if isinstance(e, INEQUALITY_TYPES) and other.is_negative:
|
||||
op = mirror_rel_op(op) # type: ignore[assignment]
|
||||
|
||||
if op is None:
|
||||
raise AssertionError("expected op to be not None")
|
||||
e = op(lhs, rhs)
|
||||
|
||||
################################################################################
|
||||
# left-hand side is FloorDiv
|
||||
################################################################################
|
||||
#
|
||||
# Given the expression: a // b op c
|
||||
# where 'op' is a relational operation, these rules only work if:
|
||||
# - b > 0
|
||||
# - c is an integer
|
||||
if (
|
||||
floordiv_inequality
|
||||
and isinstance(e, sympy.Rel)
|
||||
and isinstance(e.lhs, FloorDiv)
|
||||
and e.lhs.divisor.is_positive
|
||||
and e.rhs.is_integer
|
||||
):
|
||||
# a // b == expr
|
||||
# => a >= (b * expr) and a < (b * (expr + 1))
|
||||
if isinstance(e, sympy.Eq):
|
||||
numerator, denominator = e.lhs.args
|
||||
return sympy.And(
|
||||
sympy.Ge(numerator, (e.rhs * denominator)),
|
||||
sympy.Lt(numerator, ((e.rhs + 1) * denominator)),
|
||||
)
|
||||
# a // b != expr
|
||||
# => a < (b * expr) or a >= (b * (expr + 1))
|
||||
if isinstance(e, sympy.Ne):
|
||||
numerator, denominator = e.lhs.args
|
||||
return sympy.Or(
|
||||
sympy.Lt(numerator, (e.rhs * denominator)),
|
||||
sympy.Ge(numerator, ((e.rhs + 1) * denominator)),
|
||||
)
|
||||
# The transformations below only work if b is positive.
|
||||
# Note: we only have this information for constants.
|
||||
# a // b > expr => a >= b * (expr + 1)
|
||||
# a // b >= expr => a >= b * expr
|
||||
if isinstance(e, (sympy.Gt, sympy.Ge)):
|
||||
quotient = e.rhs if isinstance(e, sympy.Ge) else (e.rhs + 1)
|
||||
return sympy.Ge(e.lhs.args[0], (quotient * e.lhs.args[1]))
|
||||
# a // b < expr => a < b * expr
|
||||
# a // b <= expr => a < b * (expr + 1)
|
||||
if isinstance(e, (sympy.Lt, sympy.Le)):
|
||||
quotient = e.rhs if isinstance(e, sympy.Lt) else (e.rhs + 1)
|
||||
return sympy.Lt(e.lhs.args[0], (quotient * e.lhs.args[1]))
|
||||
|
||||
return e
|
||||
@@ -0,0 +1,101 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This file contains canonical definitions for our symbol naming conventions,
|
||||
across torch.fx.experimental.symbolic_shapes and torch._inductor. The
|
||||
intention is:
|
||||
|
||||
1. To make it easily greppable where all the sites we use a prefix are
|
||||
2. Make it possible to easily tell if we can introduce a new prefix without
|
||||
introducing a conflict
|
||||
|
||||
You can occasionally test if prefixes have been hardcoded by renaming prefixes
|
||||
in this file and seeing what breaks.
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable
|
||||
from enum import auto, Enum
|
||||
|
||||
import sympy
|
||||
|
||||
|
||||
class SymT(Enum):
|
||||
SIZE = auto()
|
||||
FLOAT = auto()
|
||||
UNBACKED_INT = auto()
|
||||
UNBACKED_FLOAT = auto()
|
||||
# Inductor: The intermediates in inner_fn tmp0, one generated per ops call.
|
||||
# If one of these shows up in an indexing expression, that means an
|
||||
# indirect load is happening.
|
||||
TMP = auto()
|
||||
# Inductor: Placeholder variable that is later replaced with TMP
|
||||
INDIRECT = auto()
|
||||
# Inductor: Some size expressions are replaced with a precomputed size ps0
|
||||
# which is computed host side, and then directly reused in the kernel, so
|
||||
# we don't repeatedly recompute it on device.
|
||||
PRECOMPUTED_SIZE = auto()
|
||||
# Inductor: An indexing variable i0 in loops IR which ranges over non-reduced
|
||||
# dim in the loop
|
||||
INDEX = auto()
|
||||
# Inductor: A reduction indexing (r0, r1) variables in loops IR which ranges over
|
||||
# reduced dim(s) in the loop
|
||||
R0_INDEX = auto()
|
||||
R1_INDEX = auto()
|
||||
# Inductor: In templated kernels torch._inductor.kernel, we have a hook to
|
||||
# store the final output and append epilogue fusions. To do this, we must
|
||||
# know what the indexes the outputs range over. NB: These will also
|
||||
# advertise as INDEX, this is... probably OK?
|
||||
TEMPLATE_INDEX = auto()
|
||||
# Inductor: iteration domain for blockIdx.x/blockIdx.y
|
||||
XBLOCK = auto()
|
||||
YBLOCK = auto()
|
||||
ZBLOCK = auto()
|
||||
# Inductor: this is used solely for dynamic_reshape_indexer
|
||||
VIEW = auto()
|
||||
# Alternate (non-modular) indexing used in halide kernels
|
||||
HALIDE = auto()
|
||||
|
||||
|
||||
# Invariant: there must not be a prefix which is a prefix of another string,
|
||||
# as this introduces ambiguity
|
||||
prefix_str = {
|
||||
SymT.SIZE: "s", # integer
|
||||
SymT.UNBACKED_INT: "u", # integer
|
||||
# Prefix z here is chosen to avoid false aliasing in symbol_is_type test
|
||||
# DO NOT add a "z" type. You also need to avoid conflicts on these
|
||||
# prefixes but this is somewhat easier to manage
|
||||
SymT.FLOAT: "zf",
|
||||
SymT.UNBACKED_FLOAT: "zuf",
|
||||
SymT.TMP: "tmp",
|
||||
SymT.PRECOMPUTED_SIZE: "ps",
|
||||
SymT.INDEX: "i",
|
||||
SymT.R0_INDEX: "r0_",
|
||||
SymT.R1_INDEX: "r1_",
|
||||
SymT.TEMPLATE_INDEX: "idx",
|
||||
SymT.XBLOCK: "x",
|
||||
SymT.YBLOCK: "y",
|
||||
SymT.ZBLOCK: "z",
|
||||
SymT.INDIRECT: "indirect", # false aliasing?
|
||||
SymT.VIEW: "view",
|
||||
SymT.HALIDE: "h",
|
||||
}
|
||||
|
||||
|
||||
def make_symbol(prefix: SymT, idx: int, **kwargs) -> sympy.Symbol:
|
||||
# TODO: maybe put the assumptions here directly
|
||||
return sympy.Symbol(f"{prefix_str[prefix]}{idx}", **kwargs)
|
||||
|
||||
|
||||
# This type is a little wider than it should be, because free_symbols says
|
||||
# that it contains Basic, rather than Symbol
|
||||
def symbol_is_type(sym: sympy.Basic, prefix: SymT | Iterable[SymT]) -> bool:
|
||||
if not isinstance(sym, sympy.Symbol):
|
||||
raise AssertionError("expected sympy.Symbol")
|
||||
name_str = sym.name.lower() # Match capitalized names like XBLOCK, RBLOCK
|
||||
if isinstance(prefix, SymT):
|
||||
return name_str.startswith(prefix_str[prefix])
|
||||
else:
|
||||
return name_str.startswith(tuple(prefix_str[p] for p in prefix))
|
||||
|
||||
|
||||
def free_symbol_is_type(e: sympy.Expr, prefix: SymT | Iterable[SymT]) -> bool:
|
||||
return any(symbol_is_type(v, prefix) for v in e.free_symbols)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,29 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
class Thunk(Generic[R]):
|
||||
"""
|
||||
A simple lazy evaluation implementation that lets you delay
|
||||
execution of a function. It properly handles releasing the
|
||||
function once it is forced.
|
||||
"""
|
||||
|
||||
f: Callable[[], R] | None
|
||||
r: R | None
|
||||
|
||||
__slots__ = ["f", "r"]
|
||||
|
||||
def __init__(self, f: Callable[[], R]) -> None:
|
||||
self.f = f
|
||||
self.r = None
|
||||
|
||||
def force(self) -> R:
|
||||
if self.f is None:
|
||||
return self.r # type: ignore[return-value]
|
||||
self.r = self.f()
|
||||
self.f = None
|
||||
return self.r
|
||||
@@ -0,0 +1,260 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
import inspect
|
||||
import os.path
|
||||
import tempfile
|
||||
import traceback
|
||||
from types import TracebackType
|
||||
|
||||
|
||||
# This file contains utilities for ensuring dynamically compile()'d
|
||||
# code fragments display their line numbers in backtraces.
|
||||
#
|
||||
# The constraints:
|
||||
#
|
||||
# - We don't have control over the user exception printer (in particular,
|
||||
# we cannot assume the linecache trick will work, c.f.
|
||||
# https://stackoverflow.com/q/50515651/23845 )
|
||||
#
|
||||
# - We don't want to create temporary files every time we compile()
|
||||
# some code; file creation should happen lazily only at exception
|
||||
# time. Arguably, you *should* be willing to write out your
|
||||
# generated Python code to file system, but in some situations
|
||||
# (esp. library code) it would violate user expectation to write
|
||||
# to the file system, so we try to avoid it. In particular, we'd
|
||||
# like to keep the files around, so users can open up the files
|
||||
# mentioned in the trace; if the file is invisible, we want to
|
||||
# avoid clogging up the filesystem.
|
||||
#
|
||||
# If this is not a constraint for you, there is a substantially simpler
|
||||
# way to implement the functionality in this PR: instead of using
|
||||
# eval/exec directly, just always write a Python file to filesystem
|
||||
# and compile that.
|
||||
#
|
||||
# - You have control over a context where the compiled code will get
|
||||
# executed, so that we can interpose while the stack is unwinding
|
||||
# (otherwise, we have no way to interpose on the exception printing
|
||||
# process.)
|
||||
#
|
||||
# There are two things you have to do to make use of the utilities here:
|
||||
#
|
||||
# - When you compile your source code, you must save its string source
|
||||
# in its f_globals under the magic name "__compile_source__"
|
||||
#
|
||||
# - Before running the compiled code, enter the
|
||||
# report_compile_source_on_error() context manager.
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def report_compile_source_on_error():
|
||||
try:
|
||||
yield
|
||||
except Exception as exc:
|
||||
tb = exc.__traceback__
|
||||
|
||||
# Walk the traceback, looking for frames that have
|
||||
# source attached
|
||||
stack = []
|
||||
while tb is not None:
|
||||
filename = tb.tb_frame.f_code.co_filename
|
||||
source = tb.tb_frame.f_globals.get("__compile_source__")
|
||||
|
||||
if filename == "<string>" and source is not None:
|
||||
# What black magic are we doing here? Intuitively, what
|
||||
# we would like to do is overwrite the co_filename on any
|
||||
# frames that were generated from exec/eval so that they
|
||||
# point to a temporary file that has the actual line
|
||||
# information, so Python's default error printer can print
|
||||
# useful line information on it.
|
||||
#
|
||||
# Writing out the temporary file is easy. But overwriting
|
||||
# co_filename is not! You can't modify the code object
|
||||
# associated with a frame. You can, however, reconstruct
|
||||
# a traceback with entirely new frames from scratch, so that's
|
||||
# what we do. But there's another problem, which is how to
|
||||
# make the frame?
|
||||
#
|
||||
# The black magic is we make a frankenstein frame and code
|
||||
# object which resembles the original frame/code enough so
|
||||
# that it will print properly under traceback and the default
|
||||
# error printer, but IT IS NOT THE ORIGINAL FRAME (you
|
||||
# couldn't, e.g., execute its code with different variables
|
||||
# and expect it to work.)
|
||||
|
||||
# Don't delete the temporary file so the user can inspect it
|
||||
# TODO: This creates a temporary file for every frame, but we
|
||||
# technically only need one per distinct __compile_source__
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", delete=False, suffix=".py"
|
||||
) as f:
|
||||
f.write(source)
|
||||
# Create a frame. Python doesn't let you construct
|
||||
# FrameType directly, so just make one with compile
|
||||
frame = tb.tb_frame
|
||||
code = compile("__inspect_currentframe()", f.name, "eval")
|
||||
code = code.replace(co_name=frame.f_code.co_name)
|
||||
# Python 3.11 only
|
||||
if hasattr(frame.f_code, "co_linetable"):
|
||||
# We can't copy ALL of the metadata over, because you
|
||||
# can cause Python to segfault this way. What exactly
|
||||
# do we need? We need enough information for
|
||||
# traceback to be able to print the exception
|
||||
# correctly. Code reading Lib/traceback.py reveals
|
||||
# that traceback calls code.co_positions() in order to
|
||||
# get the augmented line/col numbers. Objects/codeobject.c,
|
||||
# specifically _PyCode_InitAddressRange, reveals that
|
||||
# this iterator is initialized from co_linetable and
|
||||
# co_firstfileno. So copy these we must!
|
||||
code = code.replace( # type: ignore[call-arg]
|
||||
co_linetable=frame.f_code.co_linetable, # type: ignore[attr-defined]
|
||||
co_firstlineno=frame.f_code.co_firstlineno, # type: ignore[attr-defined]
|
||||
)
|
||||
fake_frame = eval(
|
||||
code,
|
||||
frame.f_globals,
|
||||
{**frame.f_locals, "__inspect_currentframe": inspect.currentframe},
|
||||
)
|
||||
fake_tb = TracebackType(None, fake_frame, tb.tb_lasti, tb.tb_lineno)
|
||||
stack.append(fake_tb)
|
||||
else:
|
||||
stack.append(tb)
|
||||
|
||||
tb = tb.tb_next
|
||||
|
||||
# Reconstruct the linked list
|
||||
tb_next = None
|
||||
for tb in reversed(stack):
|
||||
tb.tb_next = tb_next
|
||||
tb_next = tb
|
||||
|
||||
raise exc.with_traceback(tb_next) # noqa: B904
|
||||
|
||||
|
||||
def shorten_filename(fn, *, base=None):
|
||||
"""Shorten a source filepath, with the assumption that torch/ subdirectories don't need to be shown to user."""
|
||||
if base is None:
|
||||
base = os.path.dirname(os.path.dirname(__file__))
|
||||
# Truncate torch/foo.py to foo.py
|
||||
try:
|
||||
prefix = os.path.commonpath([fn, base])
|
||||
except ValueError:
|
||||
return fn
|
||||
else:
|
||||
return fn[len(prefix) + 1 :]
|
||||
|
||||
|
||||
def format_frame(frame, *, base=None, line=False) -> str:
|
||||
"""
|
||||
Format a FrameSummary in a short way, without printing full absolute path or code.
|
||||
|
||||
The idea is the result fits on a single line.
|
||||
"""
|
||||
extra_line = ""
|
||||
if line:
|
||||
extra_line = f"{frame.line} # "
|
||||
return f"{extra_line}{shorten_filename(frame.filename, base=base)}:{frame.lineno} in {frame.name}"
|
||||
|
||||
|
||||
def format_traceback_short(tb):
|
||||
"""Format a TracebackType in a short way, printing only the inner-most frame."""
|
||||
return format_frame(traceback.extract_tb(tb)[-1])
|
||||
|
||||
|
||||
class CapturedTraceback:
|
||||
__slots__ = ["tb", "skip"]
|
||||
|
||||
def __init__(self, tb, skip=0) -> None:
|
||||
self.tb = tb
|
||||
self.skip = skip
|
||||
|
||||
def cleanup(self) -> None:
|
||||
self.tb = None
|
||||
|
||||
def summary(self):
|
||||
import torch._C._profiler
|
||||
|
||||
if self.tb is None:
|
||||
# TODO: Maybe indicate that the traceback was elided?
|
||||
return traceback.StackSummary()
|
||||
|
||||
return _extract_symbolized_tb(
|
||||
torch._C._profiler.symbolize_tracebacks([self.tb])[0], self.skip
|
||||
)
|
||||
|
||||
def __getstate__(self):
|
||||
return (
|
||||
None,
|
||||
{
|
||||
"tb": None, # TB is not pickleable
|
||||
"skip": self.skip,
|
||||
},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def extract(*, script=False, cpp=False, skip=0):
|
||||
"""
|
||||
Like traceback.extract_stack(), but faster (approximately 20x faster); it
|
||||
is fast enough that you can unconditionally log stacks this way as part of
|
||||
normal execution. It returns a torch._C._profiler.CapturedTraceback
|
||||
object that must be formatted specially with format_captured_tb.
|
||||
|
||||
By default, this only reports Python backtraces (like extract_stack). You
|
||||
can set the script/cpp kwargs to also turn on TorchScript/C++ trace
|
||||
reporting.
|
||||
"""
|
||||
import torch._C._profiler
|
||||
|
||||
if script or cpp:
|
||||
if skip != 0:
|
||||
raise AssertionError("skip with script/cpp NYI")
|
||||
|
||||
return CapturedTraceback(
|
||||
torch._C._profiler.gather_traceback(python=True, script=script, cpp=cpp),
|
||||
# Elide extract() frame if we don't have script/cpp frames. If
|
||||
# we do have those frames, it doesn't work so force zero.
|
||||
0 if script or cpp else skip + 1,
|
||||
)
|
||||
|
||||
def format(self):
|
||||
"""
|
||||
Formats a single torch._C._profiler.CapturedTraceback into a list of
|
||||
strings equivalent to the output of traceback.format_list. Note that if
|
||||
pass it CapturedTraceback with C++ traces, it is better not to use this
|
||||
function and use the batch formatting API format_captured_tbs to amortize
|
||||
the cost of symbolization
|
||||
"""
|
||||
return traceback.format_list(self.summary())
|
||||
|
||||
@staticmethod
|
||||
def format_all(tbs):
|
||||
"""
|
||||
Bulk version of CapturedTraceback.format. Returns a list of list of strings.
|
||||
"""
|
||||
import torch._C._profiler
|
||||
|
||||
# Directly populate tracebacks that already have cached summaries
|
||||
rs: list[list[str] | None] = []
|
||||
delayed_idxs = []
|
||||
for i, tb in enumerate(tbs):
|
||||
if tb.tb is None:
|
||||
rs.append([])
|
||||
else:
|
||||
rs.append(None)
|
||||
delayed_idxs.append(i)
|
||||
|
||||
torch._C._profiler.symbolize_tracebacks([tbs[i].tb for i in delayed_idxs])
|
||||
for i in delayed_idxs:
|
||||
rs[i] = traceback.format_list(tbs[i].summary())
|
||||
|
||||
return rs
|
||||
|
||||
|
||||
def _extract_symbolized_tb(tb, skip):
|
||||
"""
|
||||
Given a symbolized traceback from symbolize_tracebacks, return a StackSummary object of
|
||||
pre-processed stack trace entries.
|
||||
"""
|
||||
stack = traceback.StackSummary()
|
||||
for f in reversed(tb[skip:]):
|
||||
stack.append(traceback.FrameSummary(f["filename"], f["line"], f["name"]))
|
||||
return stack
|
||||
@@ -0,0 +1,223 @@
|
||||
import functools
|
||||
import hashlib
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_triton_package() -> bool:
|
||||
try:
|
||||
import triton # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_triton_version(fallback: tuple[int, int] = (0, 0)) -> tuple[int, int]:
|
||||
try:
|
||||
import triton
|
||||
|
||||
major, minor = tuple(int(v) for v in triton.__version__.split(".")[:2])
|
||||
return (major, minor)
|
||||
except ImportError:
|
||||
return fallback
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _device_supports_tma() -> bool:
|
||||
import torch
|
||||
|
||||
return (
|
||||
torch.cuda.is_available()
|
||||
and torch.cuda.get_device_capability() >= (9, 0)
|
||||
and not torch.version.hip
|
||||
)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_triton_experimental_host_tma() -> bool:
|
||||
if has_triton_package():
|
||||
if _device_supports_tma():
|
||||
try:
|
||||
from triton.tools.experimental_descriptor import ( # noqa: F401
|
||||
create_1d_tma_descriptor,
|
||||
create_2d_tma_descriptor,
|
||||
)
|
||||
|
||||
try:
|
||||
from triton.tools.experimental_descriptor import enable_in_pytorch
|
||||
|
||||
return enable_in_pytorch()
|
||||
except ImportError:
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_triton_tensor_descriptor_host_tma() -> bool:
|
||||
if has_triton_package():
|
||||
if _device_supports_tma():
|
||||
try:
|
||||
from triton.tools.tensor_descriptor import ( # noqa: F401
|
||||
TensorDescriptor,
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_triton_tma() -> bool:
|
||||
return has_triton_tensor_descriptor_host_tma() or has_triton_experimental_host_tma()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_triton_tma_device() -> bool:
|
||||
if has_triton_package():
|
||||
import torch
|
||||
|
||||
if (
|
||||
torch.cuda.is_available()
|
||||
and torch.cuda.get_device_capability() >= (9, 0)
|
||||
and not torch.version.hip
|
||||
) or torch.xpu.is_available():
|
||||
# old API
|
||||
try:
|
||||
from triton.language.extra.cuda import ( # noqa: F401
|
||||
experimental_device_tensormap_create1d,
|
||||
experimental_device_tensormap_create2d,
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# new API
|
||||
try:
|
||||
from triton.language import make_tensor_descriptor # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_datacenter_blackwell_tma_device() -> bool:
|
||||
import torch
|
||||
|
||||
if (
|
||||
torch.cuda.is_available()
|
||||
and torch.cuda.get_device_capability() >= (10, 0)
|
||||
and torch.cuda.get_device_capability() < (11, 0)
|
||||
and not torch.version.hip
|
||||
):
|
||||
return has_triton_tma_device() and has_triton_tensor_descriptor_host_tma()
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@functools.lru_cache(None)
|
||||
def has_triton_stable_tma_api() -> bool:
|
||||
if has_triton_package():
|
||||
import torch
|
||||
|
||||
if (
|
||||
torch.cuda.is_available()
|
||||
and torch.cuda.get_device_capability() >= (9, 0)
|
||||
and not torch.version.hip
|
||||
) or torch.xpu.is_available():
|
||||
try:
|
||||
from triton.language import make_tensor_descriptor # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_triton() -> bool:
|
||||
if not has_triton_package():
|
||||
return False
|
||||
|
||||
from torch._inductor.config import triton_disable_device_detection
|
||||
|
||||
if triton_disable_device_detection:
|
||||
return False
|
||||
|
||||
from torch._dynamo.device_interface import get_interface_for_device
|
||||
|
||||
def cuda_extra_check(device_interface: Any) -> bool:
|
||||
return device_interface.Worker.get_device_properties().major >= 7
|
||||
|
||||
def cpu_extra_check(device_interface: Any) -> bool:
|
||||
import triton.backends
|
||||
|
||||
return "cpu" in triton.backends.backends
|
||||
|
||||
def _return_true(device_interface: Any) -> bool:
|
||||
return True
|
||||
|
||||
triton_supported_devices = {
|
||||
"cuda": cuda_extra_check,
|
||||
"xpu": _return_true,
|
||||
"cpu": cpu_extra_check,
|
||||
"mtia": _return_true,
|
||||
}
|
||||
|
||||
def is_device_compatible_with_triton() -> bool:
|
||||
for device, extra_check in triton_supported_devices.items():
|
||||
device_interface = get_interface_for_device(device)
|
||||
if device_interface.is_available() and extra_check(device_interface):
|
||||
return True
|
||||
return False
|
||||
|
||||
return is_device_compatible_with_triton()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def triton_backend() -> Any:
|
||||
from triton.compiler.compiler import make_backend
|
||||
from triton.runtime.driver import driver
|
||||
|
||||
target = driver.active.get_current_target()
|
||||
return make_backend(target)
|
||||
|
||||
|
||||
def _extern_libs_key(backend: Any) -> str:
|
||||
"""Return a cache key fragment for extern libs (e.g. libdevice.10.bc).
|
||||
|
||||
These files affect codegen but are not covered by triton_key() (Python
|
||||
sources only) or backend.hash() (ptxas version and arch only).
|
||||
"""
|
||||
opts = backend.parse_options({})
|
||||
extern_libs = getattr(opts, "extern_libs", None)
|
||||
if not extern_libs:
|
||||
return ""
|
||||
parts = []
|
||||
for name, path in sorted(extern_libs):
|
||||
if os.path.isfile(path):
|
||||
with open(path, "rb") as f:
|
||||
parts.append(f"{name}-{hashlib.sha256(f.read()).hexdigest()}")
|
||||
return "-".join(parts)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def triton_hash_with_backend() -> str:
|
||||
from torch._inductor.runtime.triton_compat import triton_key
|
||||
|
||||
backend = triton_backend()
|
||||
key = f"{triton_key()}-{backend.hash()}"
|
||||
|
||||
# Hash is upper case so that it can't contain any Python keywords.
|
||||
return hashlib.sha256(key.encode("utf-8")).hexdigest().upper()
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Miscellaneous utilities to aid with typing."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast, Concatenate, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
# Helper to turn Optional[T] into T when we know None either isn't
|
||||
# possible or should trigger an exception.
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def not_none(obj: T | None) -> T:
|
||||
if obj is None:
|
||||
raise TypeError("Invariant encountered: value was None when it should not be")
|
||||
return obj
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
_A1 = TypeVar("_A1")
|
||||
|
||||
|
||||
def copy_func_params(
|
||||
source_func: Callable[_P, Any],
|
||||
) -> Callable[[Callable[..., _R]], Callable[_P, _R]]:
|
||||
"""Cast the decorated function's call signature to the source_func's.
|
||||
|
||||
Usage:
|
||||
def upstream_func(a: int, b: float, *, double: bool = False) -> float: ...
|
||||
@copy_func_params(upstream_func)
|
||||
def enhanced(a: int, b: float, *args: Any, double: bool = False, **kwargs: Any) -> str: ...
|
||||
"""
|
||||
|
||||
def return_func(func: Callable[..., _R]) -> Callable[_P, _R]:
|
||||
return cast(Callable[_P, _R], func)
|
||||
|
||||
return return_func
|
||||
|
||||
|
||||
def copy_method_params(
|
||||
source_method: Callable[Concatenate[Any, _P], Any],
|
||||
) -> Callable[[Callable[..., _R]], Callable[Concatenate[_A1, _P], _R]]:
|
||||
"""Cast the decorated *method*'s call signature to the source_method's.
|
||||
Keeps the first argument type (e.g., self/cls).
|
||||
"""
|
||||
|
||||
def return_func(func: Callable[..., _R]) -> Callable[Concatenate[_A1, _P], _R]:
|
||||
return cast(Callable[Concatenate[_A1, _P], _R], func)
|
||||
|
||||
return return_func
|
||||
|
||||
|
||||
# stricter variants to preserve the origin callers Return Type too.
|
||||
# TODO: consider folding both these into the above variants with an optional
|
||||
# parameter to control whether to copy the return type or not.
|
||||
def copy_func_sig(
|
||||
source_func: Callable[_P, _R],
|
||||
) -> Callable[[Callable[..., _R]], Callable[_P, _R]]:
|
||||
"""Cast the decorated function's call signature and return type to the source_func's."""
|
||||
|
||||
def _return(func: Callable[..., _R]) -> Callable[_P, _R]:
|
||||
return cast(Callable[_P, _R], func)
|
||||
|
||||
return _return
|
||||
|
||||
|
||||
def copy_method_sig(
|
||||
source_method: Callable[Concatenate[_A1, _P], _R],
|
||||
) -> Callable[[Callable[..., _R]], Callable[Concatenate[_A1, _P], _R]]:
|
||||
"""Cast the decorated *method*'s call signature to the source_method and return type."""
|
||||
|
||||
def _return(func: Callable[..., _R]) -> Callable[Concatenate[_A1, _P], _R]:
|
||||
return cast(Callable[Concatenate[_A1, _P], _R], func)
|
||||
|
||||
return _return
|
||||
@@ -0,0 +1,86 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
from pathlib import Path
|
||||
from zipfile import ZipFile
|
||||
|
||||
|
||||
# Exclude some standard library modules to:
|
||||
# 1. Slim down the final zipped file size
|
||||
# 2. Remove functionality we don't want to support.
|
||||
DENY_LIST = [
|
||||
# Interface to unix databases
|
||||
"dbm",
|
||||
# ncurses bindings (terminal interfaces)
|
||||
"curses",
|
||||
# Tcl/Tk GUI
|
||||
"tkinter",
|
||||
"tkinter",
|
||||
# Tests for the standard library
|
||||
"test",
|
||||
"tests",
|
||||
"idle_test",
|
||||
"__phello__.foo.py",
|
||||
# importlib frozen modules. These are already baked into CPython.
|
||||
"_bootstrap.py",
|
||||
"_bootstrap_external.py",
|
||||
]
|
||||
|
||||
strip_file_dir = ""
|
||||
|
||||
|
||||
def remove_prefix(text, prefix):
|
||||
if text.startswith(prefix):
|
||||
return text[len(prefix) :]
|
||||
return text
|
||||
|
||||
|
||||
def write_to_zip(file_path, strip_file_path, zf, prepend_str="") -> None:
|
||||
stripped_file_path = prepend_str + remove_prefix(file_path, strip_file_dir + "/")
|
||||
path = Path(stripped_file_path)
|
||||
if path.name in DENY_LIST:
|
||||
return
|
||||
zf.write(file_path, stripped_file_path)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
global strip_file_dir
|
||||
parser = argparse.ArgumentParser(description="Zip py source")
|
||||
parser.add_argument("paths", nargs="*", help="Paths to zip.")
|
||||
parser.add_argument(
|
||||
"--install-dir", "--install_dir", help="Root directory for all output files"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--strip-dir",
|
||||
"--strip_dir",
|
||||
help="The absolute directory we want to remove from zip",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prepend-str",
|
||||
"--prepend_str",
|
||||
help="A string to prepend onto all paths of a file in the zip",
|
||||
default="",
|
||||
)
|
||||
parser.add_argument("--zip-name", "--zip_name", help="Output zip name")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
zip_file_name = args.install_dir + "/" + args.zip_name
|
||||
strip_file_dir = args.strip_dir
|
||||
prepend_str = args.prepend_str
|
||||
with ZipFile(zip_file_name, mode="w") as zf:
|
||||
for p in sorted(args.paths):
|
||||
if os.path.isdir(p):
|
||||
files = glob.glob(p + "/**/*.py", recursive=True)
|
||||
for file_path in sorted(files):
|
||||
# strip the absolute path
|
||||
write_to_zip(
|
||||
file_path, strip_file_dir + "/", zf, prepend_str=prepend_str
|
||||
)
|
||||
else:
|
||||
write_to_zip(p, strip_file_dir + "/", zf, prepend_str=prepend_str)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main() # pragma: no cover
|
||||
@@ -0,0 +1,27 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from torch._C import (
|
||||
_get_backcompat_broadcast_warn,
|
||||
_get_backcompat_keepdim_warn,
|
||||
_set_backcompat_broadcast_warn,
|
||||
_set_backcompat_keepdim_warn,
|
||||
)
|
||||
|
||||
|
||||
class Warning:
|
||||
def __init__(self, setter, getter) -> None:
|
||||
self.setter = setter
|
||||
self.getter = getter
|
||||
|
||||
def set_enabled(self, value) -> None:
|
||||
self.setter(value)
|
||||
|
||||
def get_enabled(self):
|
||||
return self.getter()
|
||||
|
||||
enabled = property(get_enabled, set_enabled)
|
||||
|
||||
|
||||
broadcast_warning = Warning(
|
||||
_set_backcompat_broadcast_warn, _get_backcompat_broadcast_warn
|
||||
)
|
||||
keepdim_warning = Warning(_set_backcompat_keepdim_warn, _get_backcompat_keepdim_warn)
|
||||
@@ -0,0 +1,520 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
from torch._C import _get_privateuse1_backend_name, _rename_privateuse1_backend
|
||||
from torch.overrides import handle_torch_function, has_torch_function_unary
|
||||
|
||||
|
||||
__all__ = [
|
||||
"rename_privateuse1_backend",
|
||||
"generate_methods_for_privateuse1_backend",
|
||||
]
|
||||
|
||||
# TODO: Should use `torch._C._get_privateuse1_backend_name()` to get
|
||||
# renamed-backend name for `privateuse1`, but the func will cause an
|
||||
# error with torch.jit.script, so we use the global variable named
|
||||
# `_privateuse1_backend_name`.
|
||||
_privateuse1_backend_name = "privateuseone"
|
||||
|
||||
|
||||
def rename_privateuse1_backend(backend_name: str) -> None:
|
||||
r"""
|
||||
Rename the privateuse1 backend device to make it more convenient to use as a device name within PyTorch APIs.
|
||||
|
||||
The steps are:
|
||||
|
||||
(1) (In C++) implement kernels for various torch operations, and register them
|
||||
to the PrivateUse1 dispatch key.
|
||||
(2) (In python) call torch.utils.rename_privateuse1_backend("foo")
|
||||
|
||||
You can now use "foo" as an ordinary device string in python.
|
||||
|
||||
Note: this API can only be called once per process. Attempting to change
|
||||
the external backend after it's already been set will result in an error.
|
||||
|
||||
Note(AMP): If you want to support AMP on your device, you can register a custom backend module.
|
||||
The backend must register a custom backend module with ``torch._register_device_module("foo", BackendModule)``.
|
||||
BackendModule needs to have the following API's:
|
||||
|
||||
(1) ``get_amp_supported_dtype() -> List[torch.dtype]``
|
||||
get the supported dtypes on your "foo" device in AMP, maybe the "foo" device supports one more dtype.
|
||||
|
||||
Note(random): If you want to support to set seed for your device, BackendModule needs to have the following API's:
|
||||
|
||||
(1) ``_is_in_bad_fork() -> bool``
|
||||
Return ``True`` if now it is in bad_fork, else return ``False``.
|
||||
|
||||
(2) ``manual_seed_all(seed int) -> None``
|
||||
Sets the seed for generating random numbers for your devices.
|
||||
|
||||
(3) ``device_count() -> int``
|
||||
Returns the number of "foo"s available.
|
||||
|
||||
(4) ``get_rng_state(device: Union[int, str, torch.device] = 'foo') -> Tensor``
|
||||
Returns a list of ByteTensor representing the random number states of all devices.
|
||||
|
||||
(5) ``set_rng_state(new_state: Tensor, device: Union[int, str, torch.device] = 'foo') -> None``
|
||||
Sets the random number generator state of the specified "foo" device.
|
||||
|
||||
And there are some common funcs:
|
||||
|
||||
(1) ``is_available() -> bool``
|
||||
Returns a bool indicating if "foo" is currently available.
|
||||
|
||||
(2) ``current_device() -> int``
|
||||
Returns the index of a currently selected device.
|
||||
|
||||
For more details, see https://pytorch.org/tutorials/advanced/extend_dispatcher.html#get-a-dispatch-key-for-your-backend
|
||||
For an existing example, see https://github.com/bdhirsh/pytorch_open_registration_example
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP("failing")
|
||||
>>> torch.utils.rename_privateuse1_backend("foo")
|
||||
# This will work, assuming that you've implemented the right C++ kernels
|
||||
# to implement torch.ones.
|
||||
>>> a = torch.ones(2, device="foo")
|
||||
|
||||
"""
|
||||
_rename_privateuse1_backend(backend_name)
|
||||
global _privateuse1_backend_name
|
||||
_privateuse1_backend_name = backend_name
|
||||
|
||||
|
||||
def _check_register_once(module, attr) -> None:
|
||||
if hasattr(module, attr):
|
||||
raise RuntimeError(
|
||||
f"The custom device module of {module} has already been registered with {attr}"
|
||||
)
|
||||
|
||||
|
||||
def _normalization_device(
|
||||
custom_backend_name: str, device: int | str | torch.device | None = None
|
||||
) -> int:
|
||||
def _get_current_device_index():
|
||||
_get_device_index = "current_device"
|
||||
if hasattr(torch, custom_backend_name) and hasattr(
|
||||
getattr(torch, custom_backend_name), _get_device_index
|
||||
):
|
||||
return getattr(getattr(torch, custom_backend_name), _get_device_index)()
|
||||
else:
|
||||
# The default device index is 0.
|
||||
return 0
|
||||
|
||||
if device is None:
|
||||
return _get_current_device_index()
|
||||
# if isinstance(device, str), this means that the parameter passed in is in the string format "foo:0"
|
||||
# convert str object to torch.device object, and then process it uniformly
|
||||
elif isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
|
||||
# variable device can only be torch.device type or int type
|
||||
if isinstance(device, torch.device):
|
||||
if device.type != custom_backend_name:
|
||||
raise RuntimeError(f"Invalid device, must be {custom_backend_name} device")
|
||||
elif device.index is None:
|
||||
device_idx = _get_current_device_index()
|
||||
else:
|
||||
device_idx = device.index
|
||||
# if isinstance(device, int), we can take the index number directly
|
||||
else:
|
||||
device_idx = device
|
||||
return device_idx
|
||||
|
||||
|
||||
def _generate_tensor_methods_for_privateuse1_backend(custom_backend_name: str) -> None:
|
||||
@property # type: ignore[misc]
|
||||
def wrap_tensor_backend(self: torch.Tensor) -> bool:
|
||||
if has_torch_function_unary(self):
|
||||
# TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185
|
||||
return handle_torch_function(wrap_tensor_backend.__get__, (self,), self) # type: ignore[attr-defined]
|
||||
return self.device.type == custom_backend_name
|
||||
|
||||
_check_register_once(torch.Tensor, f"is_{custom_backend_name}")
|
||||
wrap_tensor_backend.fget.__name__ = f"is_{custom_backend_name}" # type: ignore[attr-defined]
|
||||
setattr(torch.Tensor, f"is_{custom_backend_name}", wrap_tensor_backend)
|
||||
|
||||
def wrap_tensor_to(
|
||||
self: torch.Tensor,
|
||||
device: int | torch.device | None = None,
|
||||
non_blocking=False,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
r"""Perform Tensor device conversion. Call the to operator implementation.
|
||||
|
||||
.. note::
|
||||
If the ``self`` Tensor already
|
||||
has the correct :class:`torch.device`, then ``self`` is returned.
|
||||
Otherwise, the returned tensor is a copy of ``self`` with the desired :class:`torch.device`.
|
||||
|
||||
Args:
|
||||
device (int, optional): if specified, all parameters will be copied to that device
|
||||
non_blocking (bool): If ``True`` and the source is in pinned memory,
|
||||
the copy will be asynchronous with respect to the host. Otherwise,
|
||||
the argument has no effect.
|
||||
**kwargs (dict): For compatibility, may contain the key ``memory_format`` argument.
|
||||
"""
|
||||
if has_torch_function_unary(self):
|
||||
return handle_torch_function(
|
||||
wrap_tensor_to,
|
||||
(self,),
|
||||
self,
|
||||
device=device,
|
||||
non_blocking=False,
|
||||
**kwargs,
|
||||
)
|
||||
device_idx = _normalization_device(custom_backend_name, device)
|
||||
return self.to(
|
||||
device=torch.device(f"{custom_backend_name}:{device_idx}"),
|
||||
non_blocking=non_blocking,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
_check_register_once(torch.Tensor, custom_backend_name)
|
||||
wrap_tensor_to.__name__ = custom_backend_name
|
||||
setattr(torch.Tensor, custom_backend_name, wrap_tensor_to)
|
||||
|
||||
|
||||
def _generate_module_methods_for_privateuse1_backend(custom_backend_name: str) -> None:
|
||||
# Generate Module attributes and methods depends on Tensor methods,
|
||||
# so we need to check whether Tensor methods is already registered.
|
||||
if not hasattr(torch.Tensor, custom_backend_name):
|
||||
raise RuntimeError(
|
||||
f"Can not automatically generate {custom_backend_name}() method for torch.nn.Module."
|
||||
f"Because torch.Tensor doesn't has the method {custom_backend_name}()."
|
||||
f"For this error, you can try setting for_tensor=True."
|
||||
)
|
||||
|
||||
def wrap_module_to(
|
||||
# pyrefly: ignore [invalid-type-var]
|
||||
self: torch.nn.modules.module.T,
|
||||
device: int | torch.device | None = None,
|
||||
) -> torch.nn.modules.module.T: # pyrefly: ignore [invalid-type-var]
|
||||
r"""Move all model parameters and buffers to the custom device.
|
||||
|
||||
This also makes associated parameters and buffers different objects. So
|
||||
it should be called before constructing optimizer if the module will
|
||||
live on device while being optimized.
|
||||
|
||||
.. note::
|
||||
This method modifies the module in-place.
|
||||
|
||||
Args:
|
||||
device (int, optional): if specified, all parameters will be copied to that device
|
||||
"""
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return self._apply(lambda t: getattr(t, custom_backend_name)(device))
|
||||
|
||||
_check_register_once(torch.nn.Module, custom_backend_name)
|
||||
setattr(torch.nn.Module, custom_backend_name, wrap_module_to)
|
||||
|
||||
|
||||
def _generate_packed_sequence_methods_for_privateuse1_backend(
|
||||
custom_backend_name: str,
|
||||
) -> None:
|
||||
# Generate PackedSequence Module attributes and methods depends on Tensor methods,
|
||||
# so we need to check whether Tensor methods is already registered.
|
||||
if not hasattr(torch.Tensor, f"is_{custom_backend_name}") or not hasattr(
|
||||
torch.Tensor, custom_backend_name
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Can not automatically generate is_{custom_backend_name}() or "
|
||||
f"{custom_backend_name}() method for torch.nn.utils.rnn.PackedSequence."
|
||||
f"Because torch.Tensor doesn't has the method is_{custom_backend_name}()"
|
||||
f"or {custom_backend_name}()."
|
||||
f"For this error, you can try setting for_tensor=True."
|
||||
)
|
||||
|
||||
@property # type: ignore[misc]
|
||||
def wrap_tensor_backend(self: torch.nn.utils.rnn.PackedSequence) -> bool:
|
||||
return self.data.device.type == custom_backend_name
|
||||
|
||||
_check_register_once(torch.nn.utils.rnn.PackedSequence, f"is_{custom_backend_name}")
|
||||
setattr(
|
||||
torch.nn.utils.rnn.PackedSequence,
|
||||
f"is_{custom_backend_name}",
|
||||
wrap_tensor_backend,
|
||||
)
|
||||
|
||||
def wrap_module_to(
|
||||
self: torch.nn.utils.rnn.PackedSequence, *args, **kwargs
|
||||
) -> torch.nn.utils.rnn.PackedSequence:
|
||||
r"""Move all model parameters and buffers to the custom device.
|
||||
|
||||
This also makes associated parameters and buffers different objects. So
|
||||
it should be called before constructing optimizer if the module will
|
||||
live on device while being optimized.
|
||||
|
||||
.. note::
|
||||
This method modifies the module in-place.
|
||||
|
||||
Args:
|
||||
device (int, optional): if specified, all parameters will be copied to that device
|
||||
"""
|
||||
ex = torch.tensor((), dtype=self.data.dtype, device=self.data.device).to(
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
if ex.device.type == custom_backend_name:
|
||||
return self.to(*args, **kwargs)
|
||||
kwargs.update({"device": custom_backend_name})
|
||||
|
||||
return self.to(*args, **kwargs)
|
||||
|
||||
_check_register_once(torch.nn.utils.rnn.PackedSequence, custom_backend_name)
|
||||
setattr(torch.nn.utils.rnn.PackedSequence, custom_backend_name, wrap_module_to)
|
||||
|
||||
|
||||
def _generate_storage_methods_for_privateuse1_backend(
|
||||
custom_backend_name: str, unsupported_dtype: list[torch.dtype] | None = None
|
||||
) -> None:
|
||||
# Attribute is registered in the _StorageBase class
|
||||
# and UntypedStorage obtains through inheritance.
|
||||
@property # type: ignore[misc]
|
||||
def wrap_storage_backend(self: torch.storage._StorageBase) -> bool:
|
||||
r"""Return the internal :class:`torch.UntypedStorage`."""
|
||||
return self.device.type == custom_backend_name
|
||||
|
||||
_check_register_once(torch.storage._StorageBase, f"is_{custom_backend_name}")
|
||||
setattr(
|
||||
torch.storage._StorageBase, f"is_{custom_backend_name}", wrap_storage_backend
|
||||
)
|
||||
|
||||
def wrap_storage_to(self, device=None, non_blocking=False):
|
||||
r"""Return a copy of this object in custom device memory.
|
||||
|
||||
If this object is already in device memory and on the correct device, then
|
||||
no copy is performed and the original object is returned.
|
||||
|
||||
Args:
|
||||
device (int): The destination device id. Defaults to the current device.
|
||||
non_blocking (bool): If ``True`` and the source is in pinned memory,
|
||||
the copy will be asynchronous with respect to the host. Otherwise,
|
||||
the argument has no effect.
|
||||
"""
|
||||
# There should be a judgment related to storage device and a judgment related to storage type,
|
||||
# but it depends on the extended function, so this part is temporarily omitted in the automatic generation.
|
||||
device_idx = _normalization_device(custom_backend_name, device)
|
||||
|
||||
if getattr(self, f"is_{custom_backend_name}"):
|
||||
# storage has already on expected device.
|
||||
if self.get_device() == device_idx:
|
||||
return self
|
||||
# For sparse storage, custom need to extend the implementation by themselves.
|
||||
if self.is_sparse:
|
||||
raise RuntimeError(
|
||||
f"Can not support a sparse storage move to {custom_backend_name} backend"
|
||||
)
|
||||
# create untyped_storage and copy data
|
||||
untyped_storage = torch.UntypedStorage(
|
||||
self.size(), device=torch.device(f"{custom_backend_name}:{device_idx}")
|
||||
)
|
||||
untyped_storage.copy_(self, non_blocking)
|
||||
return untyped_storage
|
||||
|
||||
_check_register_once(torch.storage._StorageBase, custom_backend_name)
|
||||
setattr(torch.storage._StorageBase, custom_backend_name, wrap_storage_to)
|
||||
|
||||
# Register the corresponding attribute for the TypedStorage class.
|
||||
# When the TypedStorage class is removed, the registration is also removed.
|
||||
|
||||
@property # type: ignore[misc]
|
||||
def wrap_typed_storage_backend(self: torch.storage.TypedStorage) -> bool:
|
||||
torch.storage._warn_typed_storage_removal()
|
||||
return self._untyped_storage.device.type == custom_backend_name
|
||||
|
||||
_check_register_once(torch.TypedStorage, f"is_{custom_backend_name}")
|
||||
setattr(
|
||||
torch.storage.TypedStorage,
|
||||
f"is_{custom_backend_name}",
|
||||
wrap_typed_storage_backend,
|
||||
)
|
||||
|
||||
def wrap_typed_storage_to(
|
||||
self: torch.storage.TypedStorage, device=None, non_blocking=False, **kwargs
|
||||
) -> torch.storage.TypedStorage:
|
||||
torch.storage._warn_typed_storage_removal()
|
||||
if unsupported_dtype and self.dtype in unsupported_dtype:
|
||||
raise RuntimeError(
|
||||
f"Cannot create {custom_backend_name} storage "
|
||||
f"as {self.dtype} dtype is not supported by this backend"
|
||||
)
|
||||
custom_backend_storage: torch.UntypedStorage = getattr(
|
||||
self._untyped_storage, custom_backend_name
|
||||
)(device, non_blocking, **kwargs)
|
||||
return self._new_wrapped_storage(custom_backend_storage)
|
||||
|
||||
_check_register_once(torch.TypedStorage, custom_backend_name)
|
||||
setattr(torch.TypedStorage, custom_backend_name, wrap_typed_storage_to)
|
||||
|
||||
|
||||
def generate_methods_for_privateuse1_backend(
|
||||
for_tensor: bool = True,
|
||||
for_module: bool = True,
|
||||
for_packed_sequence: bool = True,
|
||||
for_storage: bool = False,
|
||||
unsupported_dtype: list[torch.dtype] | None = None,
|
||||
) -> None:
|
||||
r"""
|
||||
Automatically generate attributes and methods for the custom backend after rename privateuse1 backend.
|
||||
|
||||
In the default scenario, storage-related methods will not be generated automatically.
|
||||
|
||||
When you implement kernels for various torch operations, and register them to the PrivateUse1 dispatch key.
|
||||
And call the function torch.rename_privateuse1_backend("foo") to rename your backend name.
|
||||
At this point, you can easily register specific methods and attributes by calling this function.
|
||||
Just like torch.Tensor.foo(), torch.Tensor.is_foo, torch.Storage.foo(), torch.Storage.is_foo.
|
||||
|
||||
Note: We recommend you use generic functions (check devices are equal or to(device=)).
|
||||
We provide these methods for convenience only and they will be "monkey patched" onto the objects
|
||||
and so will not be properly typed. For Storage methods generate, if you need to support sparse data storage,
|
||||
you need to extend the implementation yourself.
|
||||
|
||||
Args:
|
||||
for_tensor (bool): whether register related methods for torch.Tensor class.
|
||||
for_module (bool): whether register related methods for torch.nn.Module class.
|
||||
for_storage (bool): whether register related methods for torch.Storage class.
|
||||
unsupported_dtype (List[torch.dtype]): takes effect only when the storage method needs to be generated,
|
||||
indicating that the storage does not support the torch.dtype type.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP("failing")
|
||||
>>> torch.utils.rename_privateuse1_backend("foo")
|
||||
>>> torch.utils.generate_methods_for_privateuse1_backend()
|
||||
# Then automatically generate backend-related attributes and methods.
|
||||
>>> a = torch.tensor(2).foo()
|
||||
>>> a.is_foo
|
||||
>>> hasattr(torch.nn.Module, 'foo')
|
||||
"""
|
||||
custom_backend_name = _get_privateuse1_backend_name()
|
||||
|
||||
if for_tensor:
|
||||
_generate_tensor_methods_for_privateuse1_backend(custom_backend_name)
|
||||
|
||||
if for_module:
|
||||
_generate_module_methods_for_privateuse1_backend(custom_backend_name)
|
||||
|
||||
if for_storage:
|
||||
_generate_storage_methods_for_privateuse1_backend(
|
||||
custom_backend_name, unsupported_dtype
|
||||
)
|
||||
|
||||
if for_packed_sequence:
|
||||
_generate_packed_sequence_methods_for_privateuse1_backend(custom_backend_name)
|
||||
|
||||
|
||||
def _get_custom_mod_func(func_name: str):
|
||||
r"""
|
||||
Return the func named `func_name` defined in custom device module. If not defined,
|
||||
return `None`. And the func is registered with `torch.utils.rename_privateuse1_backend('foo')`
|
||||
and `torch._register_device_module('foo', BackendModule)`.
|
||||
If the custom device module or the func is not defined, it will give warning or error message.
|
||||
Args:
|
||||
func_name (str): return the callable func named func_name defined in custom device module.
|
||||
Example::
|
||||
class DummyfooModule:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return True
|
||||
@staticmethod
|
||||
def func_name(*args, **kwargs):
|
||||
....
|
||||
torch.utils.rename_privateuse1_backend("foo")
|
||||
torch._register_device_module("foo", DummyfooModule)
|
||||
foo_is_available_func = torch.utils.backend_registration._get_custom_mod_func("is_available")
|
||||
if foo_is_available_func:
|
||||
foo_is_available = foo_is_available_func()
|
||||
func_ = torch.utils.backend_registration._get_custom_mod_func("func_name")
|
||||
if func_:
|
||||
result = func_(*args, **kwargs)
|
||||
Attention: This function is not meant to be used directly by users, which is why
|
||||
it is marked as private. It is a convenience function for backend implementers to
|
||||
more easily call the hooks into their backend extensions.
|
||||
"""
|
||||
if not isinstance(func_name, str):
|
||||
raise AssertionError(f"func_name must be `str`, but got `{type(func_name)}`.")
|
||||
backend_name = _get_privateuse1_backend_name()
|
||||
custom_device_mod = getattr(torch, backend_name, None)
|
||||
function = getattr(custom_device_mod, func_name, None)
|
||||
if custom_device_mod is None or function is None:
|
||||
message = f"Try to call torch.{backend_name}.{func_name}. The backend must register a custom backend "
|
||||
message += f"module with `torch._register_device_module('{backend_name}', BackendModule)`. And "
|
||||
message += f"BackendModule needs to have the following API's:\n `{func_name}(*args, **kwargs)`. \n"
|
||||
raise RuntimeError(message)
|
||||
return function
|
||||
|
||||
|
||||
class _DummyBackendModule:
|
||||
def is_initialized(self) -> bool:
|
||||
return True
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return True
|
||||
|
||||
def current_device(self) -> int:
|
||||
return 0
|
||||
|
||||
def _is_in_bad_fork(self) -> bool:
|
||||
return False
|
||||
|
||||
def manual_seed_all(self, seed: int) -> None:
|
||||
pass
|
||||
|
||||
def device_count(self) -> int:
|
||||
return 1
|
||||
|
||||
|
||||
class _DummyPrivateUse1Hook(torch._C._acc.PrivateUse1Hooks):
|
||||
def is_available(self) -> bool:
|
||||
return True
|
||||
|
||||
def has_primary_context(self, dev_id) -> bool:
|
||||
return True
|
||||
|
||||
def is_built(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _DummyDeviceGuard(torch._C._acc.DeviceGuard):
|
||||
def type_(self):
|
||||
return torch._C._autograd.DeviceType.PrivateUse1
|
||||
|
||||
|
||||
def _setup_privateuseone_for_python_backend(
|
||||
rename=None, backend_module=None, hook=None, device_guard=None
|
||||
) -> None:
|
||||
"""This function will prepare the PrivateUse1 dispatch key to be used as a python backend.
|
||||
|
||||
WARNING: this API is experimental and might change without notice.
|
||||
|
||||
Formally, this registers things that Pytorch expects a registered backend
|
||||
in C++ to have: including device guards, hooks, and backend modules and what not.
|
||||
|
||||
after this call, one can use `torch.library` to write Ops for this dispatch key
|
||||
and expect it to behave like a backend registered in C++.
|
||||
|
||||
See the unit test at test/test_privateuseone_python_backend.py for more details.
|
||||
|
||||
Args:
|
||||
rename: str | None, if passed in, we will rename privateuseone backend to
|
||||
the name given.
|
||||
backend_module: object | None, if passed in None, we will use DummyBackendModule
|
||||
hook: object | None, if passed in None, we will use DummyPrivateUse1Hook
|
||||
device_guard: object | None, if passed in None, we will use DummyDeviceGuard
|
||||
"""
|
||||
# NOTE: the ordering of which these functions are called is important.
|
||||
if rename is not None:
|
||||
torch.utils.rename_privateuse1_backend(rename)
|
||||
else:
|
||||
rename = "privateuseone"
|
||||
torch.utils.generate_methods_for_privateuse1_backend()
|
||||
if backend_module is None:
|
||||
backend_module = _DummyBackendModule()
|
||||
if hook is None:
|
||||
hook = _DummyPrivateUse1Hook()
|
||||
if device_guard is None:
|
||||
device_guard = _DummyDeviceGuard()
|
||||
torch._register_device_module(rename, backend_module)
|
||||
torch._C._acc.register_python_privateuseone_hook(hook)
|
||||
torch._C._acc.register_python_privateuseone_device_guard(device_guard)
|
||||
@@ -0,0 +1,6 @@
|
||||
from torch.utils.benchmark.utils.common import * # noqa: F403
|
||||
from torch.utils.benchmark.utils.timer import * # noqa: F403
|
||||
from torch.utils.benchmark.utils.compare import * # noqa: F403
|
||||
from torch.utils.benchmark.utils.fuzzer import * # noqa: F403
|
||||
from torch.utils.benchmark.utils.valgrind_wrapper.timer_interface import * # noqa: F403
|
||||
from torch.utils.benchmark.utils.sparse_fuzzer import * # noqa: F403
|
||||
@@ -0,0 +1,99 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Example of Timer and Compare APIs:
|
||||
|
||||
$ python -m examples.compare
|
||||
"""
|
||||
|
||||
import pickle
|
||||
import sys
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
import torch.utils.benchmark as benchmark_utils
|
||||
|
||||
|
||||
class FauxTorch:
|
||||
"""Emulate different versions of pytorch.
|
||||
|
||||
In normal circumstances this would be done with multiple processes
|
||||
writing serialized measurements, but this simplifies that model to
|
||||
make the example clearer.
|
||||
"""
|
||||
def __init__(self, real_torch, extra_ns_per_element) -> None:
|
||||
self._real_torch = real_torch
|
||||
self._extra_ns_per_element = extra_ns_per_element
|
||||
|
||||
def extra_overhead(self, result):
|
||||
# time.sleep has a ~65 us overhead, so only fake a
|
||||
# per-element overhead if numel is large enough.
|
||||
numel = int(result.numel())
|
||||
if numel > 5000:
|
||||
time.sleep(numel * self._extra_ns_per_element * 1e-9)
|
||||
return result
|
||||
|
||||
def add(self, *args, **kwargs):
|
||||
return self.extra_overhead(self._real_torch.add(*args, **kwargs))
|
||||
|
||||
def mul(self, *args, **kwargs):
|
||||
return self.extra_overhead(self._real_torch.mul(*args, **kwargs))
|
||||
|
||||
def cat(self, *args, **kwargs):
|
||||
return self.extra_overhead(self._real_torch.cat(*args, **kwargs))
|
||||
|
||||
def matmul(self, *args, **kwargs):
|
||||
return self.extra_overhead(self._real_torch.matmul(*args, **kwargs))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
tasks = [
|
||||
("add", "add", "torch.add(x, y)"),
|
||||
("add", "add (extra +0)", "torch.add(x, y + zero)"),
|
||||
]
|
||||
|
||||
serialized_results = []
|
||||
repeats = 2
|
||||
timers = [
|
||||
benchmark_utils.Timer(
|
||||
stmt=stmt,
|
||||
globals={
|
||||
"torch": torch if branch == "master" else FauxTorch(torch, overhead_ns),
|
||||
"x": torch.ones((size, 4)),
|
||||
"y": torch.ones((1, 4)),
|
||||
"zero": torch.zeros(()),
|
||||
},
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description=f"size: {size}",
|
||||
env=branch,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
for branch, overhead_ns in [("master", None), ("my_branch", 1), ("severe_regression", 5)]
|
||||
for label, sub_label, stmt in tasks
|
||||
for size in [1, 10, 100, 1000, 10000, 50000]
|
||||
for num_threads in [1, 4]
|
||||
]
|
||||
|
||||
for i, timer in enumerate(timers * repeats):
|
||||
serialized_results.append(pickle.dumps(
|
||||
timer.blocked_autorange(min_run_time=0.05)
|
||||
))
|
||||
print(f"\r{i + 1} / {len(timers) * repeats}", end="")
|
||||
sys.stdout.flush()
|
||||
print()
|
||||
|
||||
comparison = benchmark_utils.Compare([
|
||||
pickle.loads(i) for i in serialized_results
|
||||
])
|
||||
|
||||
print("== Unformatted " + "=" * 80 + "\n" + "/" * 95 + "\n")
|
||||
comparison.print()
|
||||
|
||||
print("== Formatted " + "=" * 80 + "\n" + "/" * 93 + "\n")
|
||||
comparison.trim_significant_figures()
|
||||
comparison.colorize()
|
||||
comparison.print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,86 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Example of the Timer and Fuzzer APIs:
|
||||
|
||||
$ python -m examples.fuzzer
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import torch.utils.benchmark as benchmark_utils
|
||||
|
||||
|
||||
def main() -> None:
|
||||
add_fuzzer = benchmark_utils.Fuzzer(
|
||||
parameters=[
|
||||
[
|
||||
benchmark_utils.FuzzedParameter(
|
||||
name=f"k{i}",
|
||||
minval=16,
|
||||
maxval=16 * 1024,
|
||||
distribution="loguniform",
|
||||
) for i in range(3)
|
||||
],
|
||||
benchmark_utils.FuzzedParameter(
|
||||
name="d",
|
||||
distribution={2: 0.6, 3: 0.4},
|
||||
),
|
||||
],
|
||||
tensors=[
|
||||
[
|
||||
benchmark_utils.FuzzedTensor(
|
||||
name=name,
|
||||
size=("k0", "k1", "k2"),
|
||||
dim_parameter="d",
|
||||
probability_contiguous=0.75,
|
||||
min_elements=64 * 1024,
|
||||
max_elements=128 * 1024,
|
||||
) for name in ("x", "y")
|
||||
],
|
||||
],
|
||||
seed=0,
|
||||
)
|
||||
|
||||
n = 250
|
||||
measurements = []
|
||||
for i, (tensors, tensor_properties, _) in enumerate(add_fuzzer.take(n=n)):
|
||||
x, x_order = tensors["x"], str(tensor_properties["x"]["order"])
|
||||
y, y_order = tensors["y"], str(tensor_properties["y"]["order"])
|
||||
shape = ", ".join(tuple(f'{i:>4}' for i in x.shape))
|
||||
|
||||
description = "".join([
|
||||
f"{x.numel():>7} | {shape:<16} | ",
|
||||
f"{'contiguous' if x.is_contiguous() else x_order:<12} | ",
|
||||
f"{'contiguous' if y.is_contiguous() else y_order:<12} | ",
|
||||
])
|
||||
|
||||
timer = benchmark_utils.Timer(
|
||||
stmt="x + y",
|
||||
globals=tensors,
|
||||
description=description,
|
||||
)
|
||||
|
||||
measurements.append(timer.blocked_autorange(min_run_time=0.1))
|
||||
measurements[-1].metadata = {"numel": x.numel()}
|
||||
print(f"\r{i + 1} / {n}", end="")
|
||||
sys.stdout.flush()
|
||||
print()
|
||||
|
||||
# More string munging to make pretty output.
|
||||
print(f"Average attempts per valid config: {1. / (1. - add_fuzzer.rejection_rate):.1f}")
|
||||
|
||||
def time_fn(m):
|
||||
return m.median / m.metadata["numel"]
|
||||
measurements.sort(key=time_fn)
|
||||
|
||||
template = f"{{:>6}}{' ' * 19}Size Shape{' ' * 13}X order Y order\n{'-' * 80}"
|
||||
print(template.format("Best:"))
|
||||
for m in measurements[:15]:
|
||||
print(f"{time_fn(m) * 1e9:>4.1f} ns / element {m.description}")
|
||||
|
||||
print("\n" + template.format("Worst:"))
|
||||
for m in measurements[-15:]:
|
||||
print(f"{time_fn(m) * 1e9:>4.1f} ns / element {m.description}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Example use of Timer and op fuzzers to measure kernel performance.
|
||||
|
||||
$ python -m examples.op_benchmark
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from torch.utils.benchmark import Timer
|
||||
from torch.utils.benchmark.op_fuzzers.binary import BinaryOpFuzzer
|
||||
from torch.utils.benchmark.op_fuzzers.unary import UnaryOpFuzzer
|
||||
import operator
|
||||
|
||||
|
||||
_MEASURE_TIME = 1.0
|
||||
|
||||
|
||||
def assert_dicts_equal(dict_0, dict_1) -> None:
|
||||
"""Builtin dict comparison will not compare numpy arrays.
|
||||
e.g.
|
||||
x = {"a": np.ones((2, 1))}
|
||||
x == x # Raises ValueError
|
||||
"""
|
||||
if set(dict_0.keys()) != set(dict_0.keys()):
|
||||
raise AssertionError("dicts must have the same keys")
|
||||
if all(np.all(v != dict_1[k]) for k, v in dict_0.items() if k != "dtype"):
|
||||
raise AssertionError("dict values differ for keys other than 'dtype'")
|
||||
|
||||
|
||||
def run(n, stmt, fuzzer_cls) -> None:
|
||||
float_iter = fuzzer_cls(seed=0, dtype=torch.float32).take(n)
|
||||
int_iter = fuzzer_cls(seed=0, dtype=torch.int32).take(n)
|
||||
raw_results = []
|
||||
for i, (float_values, int_values) in enumerate(zip(float_iter, int_iter, strict=True)):
|
||||
float_tensors, float_tensor_params, float_params = float_values
|
||||
int_tensors, int_tensor_params, int_params = int_values
|
||||
|
||||
# This benchmark assumes that the two fuzzers generate identically
|
||||
# sized and strided Tensors, since the same seed is used.
|
||||
assert_dicts_equal(float_params, int_params)
|
||||
assert_dicts_equal(float_tensor_params["x"], int_tensor_params["x"])
|
||||
|
||||
float_measurement, int_measurement = (
|
||||
Timer(
|
||||
stmt,
|
||||
globals=tensors,
|
||||
).blocked_autorange(min_run_time=_MEASURE_TIME)
|
||||
for tensors in (float_tensors, int_tensors)
|
||||
)
|
||||
|
||||
descriptions = []
|
||||
for name in float_tensors:
|
||||
shape_str = "(" + ", ".join([
|
||||
f"2 ** {int(np.log2(i))}"
|
||||
if 2 ** int(np.log2(i)) == i and i > 1
|
||||
else str(i)
|
||||
for i in float_tensors[name].shape
|
||||
]) + ")"
|
||||
order = float_tensor_params[name]["order"]
|
||||
order_str = ("" if all(order == np.arange(len(order))) else str(tuple(order)))
|
||||
steps = float_tensor_params[name]["steps"]
|
||||
steps_str = str(steps) if sum(steps) > len(steps) else ""
|
||||
descriptions.append((name, shape_str, order_str, steps_str))
|
||||
raw_results.append((float_measurement, int_measurement, descriptions))
|
||||
|
||||
print(f"\r{i + 1} / {n}", end="")
|
||||
print()
|
||||
|
||||
parsed_results, name_len, shape_len, order_len, steps_len = [], 0, 0, 0, 0
|
||||
for float_measurement, int_measurement, descriptions in raw_results:
|
||||
t_float = float_measurement.median * 1e6
|
||||
t_int = int_measurement.median * 1e6
|
||||
rel_diff = abs(t_float - t_int) / (t_float + t_int) * 2
|
||||
parsed_results.append((t_float, t_int, rel_diff, descriptions))
|
||||
for name, shape, order, steps in descriptions:
|
||||
name_len = max(name_len, len(name))
|
||||
shape_len = max(shape_len, len(shape))
|
||||
order_len = max(order_len, len(order))
|
||||
steps_len = max(steps_len, len(steps))
|
||||
|
||||
parsed_results.sort(key=operator.itemgetter(2))
|
||||
|
||||
print(f"stmt: {stmt}")
|
||||
print(f" diff faster{'':>17}{' ' * name_len} ", end="")
|
||||
print(f"{'shape'.ljust(shape_len)}{'':>16}{'order'.ljust(order_len)}", end="")
|
||||
print(f" steps\n{'-' * 100}")
|
||||
for results, spacer in [(parsed_results[:10], "..."), (parsed_results[-10:], "")]:
|
||||
for t_float, t_int, rel_diff, descriptions in results:
|
||||
time_str = [f"{rel_diff * 100:>4.1f}% {'int' if t_int < t_float else 'float':<20}"]
|
||||
time_str.extend(["".ljust(len(time_str[0])) for _ in descriptions[:-1]])
|
||||
for t_str, (name, shape, order, steps) in zip(time_str, descriptions, strict=True):
|
||||
name = f"{name}:".ljust(name_len + 1)
|
||||
shape = shape.ljust(shape_len + 10)
|
||||
order = order.ljust(order_len)
|
||||
print(f"{t_str} {name} {shape}| {order} | {steps}")
|
||||
print(spacer)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
run(n=100, stmt="torch.median(x, dim=0)", fuzzer_cls=UnaryOpFuzzer)
|
||||
run(n=100, stmt="torch.square(x)", fuzzer_cls=UnaryOpFuzzer)
|
||||
run(n=100, stmt="x + y", fuzzer_cls=BinaryOpFuzzer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
"""Trivial use of Timer API:
|
||||
|
||||
$ python -m examples.simple_timeit
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
import torch.utils.benchmark as benchmark_utils
|
||||
|
||||
|
||||
def main() -> None:
|
||||
timer = benchmark_utils.Timer(
|
||||
stmt="x + y",
|
||||
globals={"x": torch.ones((4, 8)), "y": torch.ones((1, 8))},
|
||||
label="Broadcasting add (4x8)",
|
||||
)
|
||||
|
||||
for i in range(3):
|
||||
print(f"Run: {i}\n{'-' * 40}")
|
||||
print(f"timeit:\n{timer.timeit(10000)}\n")
|
||||
print(f"autorange:\n{timer.blocked_autorange()}\n\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Microbenchmarks for the torch.fft module"""
|
||||
from argparse import ArgumentParser
|
||||
from collections import namedtuple
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
import torch.fft
|
||||
from torch.utils import benchmark
|
||||
from torch.utils.benchmark.op_fuzzers.spectral import SpectralOpFuzzer
|
||||
|
||||
|
||||
def _dim_options(ndim):
|
||||
if ndim == 1:
|
||||
return [None]
|
||||
elif ndim == 2:
|
||||
return [0, 1, None]
|
||||
elif ndim == 3:
|
||||
return [0, 1, 2, (0, 1), (0, 2), None]
|
||||
raise ValueError(f"Expected ndim in range 1-3, got {ndim}")
|
||||
|
||||
|
||||
def run_benchmark(name: str, function: object, dtype: torch.dtype, seed: int, device: str, samples: int,
|
||||
probability_regular: float):
|
||||
cuda = device == 'cuda'
|
||||
spectral_fuzzer = SpectralOpFuzzer(seed=seed, dtype=dtype, cuda=cuda,
|
||||
probability_regular=probability_regular)
|
||||
results = []
|
||||
for tensors, tensor_params, params in spectral_fuzzer.take(samples):
|
||||
shape = [params['k0'], params['k1'], params['k2']][:params['ndim']]
|
||||
str_shape = ' x '.join([f"{s:<4}" for s in shape])
|
||||
sub_label = f"{str_shape} {'' if tensor_params['x']['is_contiguous'] else '(discontiguous)'}"
|
||||
for dim in _dim_options(params['ndim']):
|
||||
for nthreads in (1, 4, 16) if not cuda else (1,):
|
||||
measurement = benchmark.Timer(
|
||||
stmt='func(x, dim=dim)',
|
||||
globals={'func': function, 'x': tensors['x'], 'dim': dim},
|
||||
label=f"{name}_{device}",
|
||||
sub_label=sub_label,
|
||||
description=f"dim={dim}",
|
||||
num_threads=nthreads,
|
||||
).blocked_autorange(min_run_time=1)
|
||||
measurement.metadata = {
|
||||
'name': name,
|
||||
'device': device,
|
||||
'dim': dim,
|
||||
'shape': shape,
|
||||
}
|
||||
measurement.metadata.update(tensor_params['x'])
|
||||
results.append(measurement)
|
||||
return results
|
||||
|
||||
|
||||
Benchmark = namedtuple('Benchmark', ['name', 'function', 'dtype'])
|
||||
BENCHMARKS = [
|
||||
Benchmark('fft_real', torch.fft.fftn, torch.float32),
|
||||
Benchmark('fft_complex', torch.fft.fftn, torch.complex64),
|
||||
Benchmark('ifft', torch.fft.ifftn, torch.complex64),
|
||||
Benchmark('rfft', torch.fft.rfftn, torch.float32),
|
||||
Benchmark('irfft', torch.fft.irfftn, torch.complex64),
|
||||
]
|
||||
BENCHMARK_MAP = {b.name: b for b in BENCHMARKS}
|
||||
BENCHMARK_NAMES = [b.name for b in BENCHMARKS]
|
||||
DEVICE_NAMES = ['cpu', 'cuda']
|
||||
|
||||
def _output_csv(file, results) -> None:
|
||||
file.write('benchmark,device,num_threads,numel,shape,contiguous,dim,mean (us),median (us),iqr (us)\n')
|
||||
for measurement in results:
|
||||
metadata = measurement.metadata
|
||||
device, dim, shape, name, numel, contiguous = (
|
||||
metadata['device'], metadata['dim'], metadata['shape'],
|
||||
metadata['name'], metadata['numel'], metadata['is_contiguous'])
|
||||
|
||||
if isinstance(dim, Iterable):
|
||||
dim_str = '-'.join(str(d) for d in dim)
|
||||
else:
|
||||
dim_str = str(dim)
|
||||
shape_str = 'x'.join(str(s) for s in shape)
|
||||
|
||||
print(name, device, measurement.task_spec.num_threads, numel, shape_str, contiguous, dim_str, # type: ignore[possibly-undefined]
|
||||
measurement.mean * 1e6, measurement.median * 1e6, measurement.iqr * 1e6,
|
||||
sep=',', file=file)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--device', type=str, choices=DEVICE_NAMES, nargs='+', default=DEVICE_NAMES)
|
||||
parser.add_argument('--bench', type=str, choices=BENCHMARK_NAMES, nargs='+', default=BENCHMARK_NAMES)
|
||||
parser.add_argument('--seed', type=int, default=0)
|
||||
parser.add_argument('--samples', type=int, default=10)
|
||||
parser.add_argument('--probability-regular', '--probability_regular', type=float, default=1.0)
|
||||
parser.add_argument('-o', '--output', type=str)
|
||||
args = parser.parse_args()
|
||||
|
||||
num_benchmarks = len(args.device) * len(args.bench)
|
||||
i = 0
|
||||
results = []
|
||||
for device in args.device:
|
||||
for bench in (BENCHMARK_MAP[b] for b in args.bench):
|
||||
results += run_benchmark(
|
||||
name=bench.name, function=bench.function, dtype=bench.dtype,
|
||||
seed=args.seed, device=device, samples=args.samples,
|
||||
probability_regular=args.probability_regular)
|
||||
i += 1
|
||||
print(f'Completed {bench.name} benchmark on {device} ({i} of {num_benchmarks})')
|
||||
|
||||
if args.output is not None:
|
||||
with open(args.output, 'w') as f:
|
||||
_output_csv(f, results)
|
||||
|
||||
compare = benchmark.Compare(results)
|
||||
compare.trim_significant_figures()
|
||||
compare.colorize()
|
||||
compare.print()
|
||||
@@ -0,0 +1,107 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedTensor
|
||||
|
||||
|
||||
_MIN_DIM_SIZE = 16
|
||||
_MAX_DIM_SIZE = 16 * 1024 ** 2
|
||||
_POW_TWO_SIZES = tuple(2 ** i for i in range(
|
||||
int(np.log2(_MIN_DIM_SIZE)),
|
||||
int(np.log2(_MAX_DIM_SIZE)) + 1,
|
||||
))
|
||||
|
||||
|
||||
class BinaryOpFuzzer(Fuzzer):
|
||||
def __init__(self, seed, dtype=torch.float32, cuda=False) -> None:
|
||||
super().__init__(
|
||||
parameters=[
|
||||
# Dimensionality of x and y. (e.g. 1D, 2D, or 3D.)
|
||||
FuzzedParameter("dim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
|
||||
|
||||
# Shapes for `x` and `y`.
|
||||
# It is important to test all shapes, however
|
||||
# powers of two are especially important and therefore
|
||||
# warrant special attention. This is done by generating
|
||||
# both a value drawn from all integers between the min and
|
||||
# max allowed values, and another from only the powers of two
|
||||
# (both distributions are loguniform) and then randomly
|
||||
# selecting between the two.
|
||||
# Moreover, `y` will occasionally have singleton
|
||||
# dimensions in order to test broadcasting.
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_any_{i}",
|
||||
minval=_MIN_DIM_SIZE,
|
||||
maxval=_MAX_DIM_SIZE,
|
||||
distribution="loguniform",
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_pow2_{i}",
|
||||
distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES}
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k_any_{i}"): 0.8,
|
||||
ParameterAlias(f"k_pow2_{i}"): 0.2,
|
||||
},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"y_k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k{i}"): 0.8,
|
||||
1: 0.2,
|
||||
},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
|
||||
# Steps for `x` and `y`. (Benchmarks strided memory access.)
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"{name}_step_{i}",
|
||||
distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04},
|
||||
)
|
||||
for i in range(3)
|
||||
for name in ("x", "y")
|
||||
],
|
||||
|
||||
# Repeatable entropy for downstream applications.
|
||||
FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"),
|
||||
],
|
||||
tensors=[
|
||||
FuzzedTensor(
|
||||
name="x",
|
||||
size=("k0", "k1", "k2"),
|
||||
steps=("x_step_0", "x_step_1", "x_step_2"),
|
||||
probability_contiguous=0.75,
|
||||
min_elements=4 * 1024,
|
||||
max_elements=32 * 1024 ** 2,
|
||||
max_allocation_bytes=2 * 1024**3, # 2 GB
|
||||
dim_parameter="dim",
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
FuzzedTensor(
|
||||
name="y",
|
||||
size=("y_k0", "y_k1", "y_k2"),
|
||||
steps=("x_step_0", "x_step_1", "x_step_2"),
|
||||
probability_contiguous=0.75,
|
||||
max_allocation_bytes=2 * 1024**3, # 2 GB
|
||||
dim_parameter="dim",
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
],
|
||||
seed=seed,
|
||||
)
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedSparseTensor
|
||||
|
||||
|
||||
_MIN_DIM_SIZE = 16
|
||||
_MAX_DIM_SIZE = 16 * 1024 ** 2
|
||||
_POW_TWO_SIZES = tuple(2 ** i for i in range(
|
||||
int(np.log2(_MIN_DIM_SIZE)),
|
||||
int(np.log2(_MAX_DIM_SIZE)) + 1,
|
||||
))
|
||||
|
||||
|
||||
class BinaryOpSparseFuzzer(Fuzzer):
|
||||
def __init__(self, seed, dtype=torch.float32, cuda=False) -> None:
|
||||
super().__init__(
|
||||
parameters=[
|
||||
# Dimensionality of x and y. (e.g. 1D, 2D, or 3D.)
|
||||
FuzzedParameter("dim_parameter", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
|
||||
FuzzedParameter(
|
||||
name="sparse_dim",
|
||||
distribution={1: 0.4, 2: 0.4, 3: 0.2},
|
||||
strict=True
|
||||
),
|
||||
# Shapes for `x` and `y`.
|
||||
# It is important to test all shapes, however
|
||||
# powers of two are especially important and therefore
|
||||
# warrant special attention. This is done by generating
|
||||
# both a value drawn from all integers between the min and
|
||||
# max allowed values, and another from only the powers of two
|
||||
# (both distributions are loguniform) and then randomly
|
||||
# selecting between the two.
|
||||
# Moreover, `y` will occasionally have singleton
|
||||
# dimensions in order to test broadcasting.
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_any_{i}",
|
||||
minval=_MIN_DIM_SIZE,
|
||||
maxval=_MAX_DIM_SIZE,
|
||||
distribution="loguniform",
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_pow2_{i}",
|
||||
distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES}
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k_any_{i}"): 0.8,
|
||||
ParameterAlias(f"k_pow2_{i}"): 0.2,
|
||||
},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"y_k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k{i}"): 1.0},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
FuzzedParameter(
|
||||
name="density",
|
||||
distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3},
|
||||
),
|
||||
FuzzedParameter(
|
||||
name="coalesced",
|
||||
distribution={True: 0.5, False: 0.5},
|
||||
),
|
||||
# Repeatable entropy for downstream applications.
|
||||
FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"),
|
||||
],
|
||||
tensors=[
|
||||
FuzzedSparseTensor(
|
||||
name="x",
|
||||
size=("k0", "k1", "k2"),
|
||||
dim_parameter="dim_parameter",
|
||||
sparse_dim="sparse_dim",
|
||||
density="density",
|
||||
coalesced="coalesced",
|
||||
min_elements=4 * 1024,
|
||||
max_elements=32 * 1024 ** 2,
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
FuzzedSparseTensor(
|
||||
name="y",
|
||||
size=("y_k0", "y_k1", "y_k2"),
|
||||
dim_parameter="dim_parameter",
|
||||
sparse_dim="sparse_dim",
|
||||
density="density",
|
||||
coalesced="coalesced",
|
||||
min_elements=4 * 1024,
|
||||
max_elements=32 * 1024 ** 2,
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
],
|
||||
seed=seed,
|
||||
)
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.types import _dtype
|
||||
|
||||
from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedSparseTensor
|
||||
|
||||
__all__ = ["UnaryOpSparseFuzzer"]
|
||||
|
||||
_MIN_DIM_SIZE = 16
|
||||
_MAX_DIM_SIZE = 16 * 1024 ** 2
|
||||
_POW_TWO_SIZES = tuple(2 ** i for i in range(
|
||||
int(np.log2(_MIN_DIM_SIZE)),
|
||||
int(np.log2(_MAX_DIM_SIZE)) + 1,
|
||||
))
|
||||
|
||||
class UnaryOpSparseFuzzer(Fuzzer):
|
||||
def __init__(self, seed: int | None, dtype: _dtype | None = None, cuda: bool = False) -> None:
|
||||
if dtype is None:
|
||||
dtype = getattr(torch, 'float32', None)
|
||||
super().__init__(
|
||||
parameters=[
|
||||
# Sparse dim parameter of x. (e.g. 1D, 2D, or 3D.)
|
||||
FuzzedParameter("dim_parameter", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
|
||||
FuzzedParameter(
|
||||
name="sparse_dim",
|
||||
distribution={1: 0.4, 2: 0.4, 3: 0.2},
|
||||
strict=True
|
||||
),
|
||||
# Shapes for `x`.
|
||||
# It is important to test all shapes, however
|
||||
# powers of two are especially important and therefore
|
||||
# warrant special attention. This is done by generating
|
||||
# both a value drawn from all integers between the min and
|
||||
# max allowed values, and another from only the powers of two
|
||||
# (both distributions are loguniform) and then randomly
|
||||
# selecting between the two.
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_any_{i}",
|
||||
minval=_MIN_DIM_SIZE,
|
||||
maxval=_MAX_DIM_SIZE,
|
||||
distribution="loguniform",
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_pow2_{i}",
|
||||
distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES}
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k_any_{i}"): 0.8,
|
||||
ParameterAlias(f"k_pow2_{i}"): 0.2,
|
||||
},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
FuzzedParameter(
|
||||
name="density",
|
||||
distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3},
|
||||
),
|
||||
FuzzedParameter(
|
||||
name="coalesced",
|
||||
distribution={True: 0.5, False: 0.5},
|
||||
),
|
||||
FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"),
|
||||
],
|
||||
tensors=[
|
||||
FuzzedSparseTensor(
|
||||
name="x",
|
||||
size=("k0", "k1", "k2"),
|
||||
dim_parameter="dim_parameter",
|
||||
sparse_dim="sparse_dim",
|
||||
min_elements=4 * 1024,
|
||||
max_elements=32 * 1024 ** 2,
|
||||
density="density",
|
||||
coalesced="coalesced",
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
],
|
||||
seed=seed,
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch.utils import benchmark
|
||||
from torch.utils.benchmark import FuzzedParameter, FuzzedTensor, ParameterAlias
|
||||
|
||||
|
||||
__all__ = ['SpectralOpFuzzer']
|
||||
|
||||
MIN_DIM_SIZE = 16
|
||||
MAX_DIM_SIZE = 16 * 1024
|
||||
|
||||
def power_range(upper_bound, base):
|
||||
return (base ** i for i in range(int(math.log(upper_bound, base)) + 1))
|
||||
|
||||
# List of regular numbers from MIN_DIM_SIZE to MAX_DIM_SIZE
|
||||
# These numbers factorize into multiples of prime factors 2, 3, and 5 only
|
||||
# and are usually the fastest in FFT implementations.
|
||||
REGULAR_SIZES = []
|
||||
for i in power_range(MAX_DIM_SIZE, 2):
|
||||
for j in power_range(MAX_DIM_SIZE // i, 3):
|
||||
ij = i * j
|
||||
for k in power_range(MAX_DIM_SIZE // ij, 5):
|
||||
ijk = ij * k
|
||||
if ijk > MIN_DIM_SIZE:
|
||||
REGULAR_SIZES.append(ijk)
|
||||
REGULAR_SIZES.sort()
|
||||
|
||||
class SpectralOpFuzzer(benchmark.Fuzzer):
|
||||
def __init__(self, *, seed: int, dtype=torch.float64,
|
||||
cuda: bool = False, probability_regular: float = 1.0) -> None:
|
||||
super().__init__(
|
||||
parameters=[
|
||||
# Dimensionality of x. (e.g. 1D, 2D, or 3D.)
|
||||
FuzzedParameter("ndim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
|
||||
|
||||
# Shapes for `x`.
|
||||
# It is important to test all shapes, however
|
||||
# regular sizes are especially important to the FFT and therefore
|
||||
# warrant special attention. This is done by generating
|
||||
# both a value drawn from all integers between the min and
|
||||
# max allowed values, and another from only the regular numbers
|
||||
# (both distributions are loguniform) and then randomly
|
||||
# selecting between the two.
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_any_{i}",
|
||||
minval=MIN_DIM_SIZE,
|
||||
maxval=MAX_DIM_SIZE,
|
||||
distribution="loguniform",
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_regular_{i}",
|
||||
distribution={size: 1. / len(REGULAR_SIZES) for size in REGULAR_SIZES}
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k_regular_{i}"): probability_regular,
|
||||
ParameterAlias(f"k_any_{i}"): 1 - probability_regular,
|
||||
},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
|
||||
# Steps for `x`. (Benchmarks strided memory access.)
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"step_{i}",
|
||||
distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04},
|
||||
) for i in range(3)
|
||||
],
|
||||
],
|
||||
tensors=[
|
||||
FuzzedTensor(
|
||||
name="x",
|
||||
size=("k0", "k1", "k2"),
|
||||
steps=("step_0", "step_1", "step_2"),
|
||||
probability_contiguous=0.75,
|
||||
min_elements=4 * 1024,
|
||||
max_elements=32 * 1024 ** 2,
|
||||
max_allocation_bytes=2 * 1024**3, # 2 GB
|
||||
dim_parameter="ndim",
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
],
|
||||
seed=seed,
|
||||
)
|
||||
@@ -0,0 +1,82 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedTensor
|
||||
|
||||
|
||||
_MIN_DIM_SIZE = 16
|
||||
_MAX_DIM_SIZE = 16 * 1024 ** 2
|
||||
_POW_TWO_SIZES = tuple(2 ** i for i in range(
|
||||
int(np.log2(_MIN_DIM_SIZE)),
|
||||
int(np.log2(_MAX_DIM_SIZE)) + 1,
|
||||
))
|
||||
|
||||
|
||||
class UnaryOpFuzzer(Fuzzer):
|
||||
def __init__(self, seed, dtype=torch.float32, cuda=False) -> None:
|
||||
super().__init__(
|
||||
parameters=[
|
||||
# Dimensionality of x. (e.g. 1D, 2D, or 3D.)
|
||||
FuzzedParameter("dim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
|
||||
|
||||
# Shapes for `x`.
|
||||
# It is important to test all shapes, however
|
||||
# powers of two are especially important and therefore
|
||||
# warrant special attention. This is done by generating
|
||||
# both a value drawn from all integers between the min and
|
||||
# max allowed values, and another from only the powers of two
|
||||
# (both distributions are loguniform) and then randomly
|
||||
# selecting between the two.
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_any_{i}",
|
||||
minval=_MIN_DIM_SIZE,
|
||||
maxval=_MAX_DIM_SIZE,
|
||||
distribution="loguniform",
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k_pow2_{i}",
|
||||
distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES}
|
||||
) for i in range(3)
|
||||
],
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"k{i}",
|
||||
distribution={
|
||||
ParameterAlias(f"k_any_{i}"): 0.8,
|
||||
ParameterAlias(f"k_pow2_{i}"): 0.2,
|
||||
},
|
||||
strict=True,
|
||||
) for i in range(3)
|
||||
],
|
||||
|
||||
# Steps for `x`. (Benchmarks strided memory access.)
|
||||
[
|
||||
FuzzedParameter(
|
||||
name=f"x_step_{i}",
|
||||
distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04},
|
||||
) for i in range(3)
|
||||
],
|
||||
|
||||
# Repeatable entropy for downstream applications.
|
||||
FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"),
|
||||
],
|
||||
tensors=[
|
||||
FuzzedTensor(
|
||||
name="x",
|
||||
size=("k0", "k1", "k2"),
|
||||
steps=("x_step_0", "x_step_1", "x_step_2"),
|
||||
probability_contiguous=0.75,
|
||||
min_elements=4 * 1024,
|
||||
max_elements=32 * 1024 ** 2,
|
||||
max_allocation_bytes=2 * 1024**3, # 2 GB
|
||||
dim_parameter="dim",
|
||||
dtype=dtype,
|
||||
cuda=cuda,
|
||||
),
|
||||
],
|
||||
seed=seed,
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
from typing import Any
|
||||
from collections.abc import Callable
|
||||
from typing_extensions import Protocol, runtime_checkable
|
||||
|
||||
|
||||
class TimerClass(Protocol):
|
||||
"""This is the portion of the `timeit.Timer` API used by benchmark utils."""
|
||||
def __init__(
|
||||
self,
|
||||
stmt: str,
|
||||
setup: str,
|
||||
timer: Callable[[], float],
|
||||
globals: dict[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
...
|
||||
|
||||
def timeit(self, number: int) -> float:
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class TimeitModuleType(Protocol):
|
||||
"""Modules generated from `timeit_template.cpp`."""
|
||||
def timeit(self, number: int) -> float:
|
||||
...
|
||||
|
||||
|
||||
class CallgrindModuleType(Protocol):
|
||||
"""Replicates the valgrind endpoints in `torch._C`.
|
||||
|
||||
These bindings are used to collect Callgrind profiles on earlier versions
|
||||
of PyTorch and will eventually be removed.
|
||||
"""
|
||||
__file__: str
|
||||
__name__: str
|
||||
|
||||
def _valgrind_supported_platform(self) -> bool:
|
||||
...
|
||||
|
||||
def _valgrind_toggle(self) -> None:
|
||||
...
|
||||
@@ -0,0 +1,359 @@
|
||||
"""Base shared classes and utilities."""
|
||||
|
||||
import collections
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import textwrap
|
||||
import time
|
||||
from typing import cast, Any
|
||||
from collections.abc import Iterable, Iterator
|
||||
import uuid
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = ["TaskSpec", "Measurement", "select_unit", "unit_to_english", "trim_sigfig", "ordered_unique", "set_torch_threads"]
|
||||
|
||||
|
||||
_MAX_SIGNIFICANT_FIGURES = 4
|
||||
_MIN_CONFIDENCE_INTERVAL = 25e-9 # 25 ns
|
||||
|
||||
# Measurement will include a warning if the distribution is suspect. All
|
||||
# runs are expected to have some variation; these parameters set the
|
||||
# thresholds.
|
||||
_IQR_WARN_THRESHOLD = 0.1
|
||||
_IQR_GROSS_WARN_THRESHOLD = 0.25
|
||||
|
||||
|
||||
@dataclasses.dataclass(init=True, repr=False, eq=True, frozen=True)
|
||||
class TaskSpec:
|
||||
"""Container for information used to define a Timer. (except globals)"""
|
||||
stmt: str
|
||||
setup: str
|
||||
global_setup: str = ""
|
||||
label: str | None = None
|
||||
sub_label: str | None = None
|
||||
description: str | None = None
|
||||
env: str | None = None
|
||||
num_threads: int = 1
|
||||
|
||||
@property
|
||||
def title(self) -> str:
|
||||
"""Best effort attempt at a string label for the measurement."""
|
||||
if self.label is not None:
|
||||
return self.label + (f": {self.sub_label}" if self.sub_label else "")
|
||||
elif "\n" not in self.stmt:
|
||||
return self.stmt + (f": {self.sub_label}" if self.sub_label else "")
|
||||
return (
|
||||
f"stmt:{f' ({self.sub_label})' if self.sub_label else ''}\n"
|
||||
f"{textwrap.indent(self.stmt, ' ')}"
|
||||
)
|
||||
|
||||
def setup_str(self) -> str:
|
||||
return (
|
||||
"" if (self.setup == "pass" or not self.setup)
|
||||
else f"setup:\n{textwrap.indent(self.setup, ' ')}" if "\n" in self.setup
|
||||
else f"setup: {self.setup}"
|
||||
)
|
||||
|
||||
def summarize(self) -> str:
|
||||
"""Build TaskSpec portion of repr string for other containers."""
|
||||
sections = [
|
||||
self.title,
|
||||
self.description or "",
|
||||
self.setup_str(),
|
||||
]
|
||||
return "\n".join([f"{i}\n" if "\n" in i else i for i in sections if i])
|
||||
|
||||
_TASKSPEC_FIELDS = tuple(i.name for i in dataclasses.fields(TaskSpec))
|
||||
|
||||
|
||||
@dataclasses.dataclass(init=True, repr=False)
|
||||
class Measurement:
|
||||
"""The result of a Timer measurement.
|
||||
|
||||
This class stores one or more measurements of a given statement. It is
|
||||
serializable and provides several convenience methods
|
||||
(including a detailed __repr__) for downstream consumers.
|
||||
"""
|
||||
number_per_run: int
|
||||
raw_times: list[float]
|
||||
task_spec: TaskSpec
|
||||
metadata: dict[Any, Any] | None = None # Reserved for user payloads.
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._sorted_times: tuple[float, ...] = ()
|
||||
self._warnings: tuple[str, ...] = ()
|
||||
self._median: float = -1.0
|
||||
self._mean: float = -1.0
|
||||
self._p25: float = -1.0
|
||||
self._p75: float = -1.0
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Forward TaskSpec fields for convenience.
|
||||
if name in _TASKSPEC_FIELDS:
|
||||
return getattr(self.task_spec, name)
|
||||
return super().__getattribute__(name)
|
||||
|
||||
# =========================================================================
|
||||
# == Convenience methods for statistics ===================================
|
||||
# =========================================================================
|
||||
#
|
||||
# These methods use raw time divided by number_per_run; this is an
|
||||
# extrapolation and hides the fact that different number_per_run will
|
||||
# result in different amortization of overheads, however if Timer has
|
||||
# selected an appropriate number_per_run then this is a non-issue, and
|
||||
# forcing users to handle that division would result in a poor experience.
|
||||
@property
|
||||
def times(self) -> list[float]:
|
||||
return [t / self.number_per_run for t in self.raw_times]
|
||||
|
||||
@property
|
||||
def median(self) -> float:
|
||||
self._lazy_init()
|
||||
return self._median
|
||||
|
||||
@property
|
||||
def mean(self) -> float:
|
||||
self._lazy_init()
|
||||
return self._mean
|
||||
|
||||
@property
|
||||
def iqr(self) -> float:
|
||||
self._lazy_init()
|
||||
return self._p75 - self._p25
|
||||
|
||||
@property
|
||||
def significant_figures(self) -> int:
|
||||
"""Approximate significant figure estimate.
|
||||
|
||||
This property is intended to give a convenient way to estimate the
|
||||
precision of a measurement. It only uses the interquartile region to
|
||||
estimate statistics to try to mitigate skew from the tails, and
|
||||
uses a static z value of 1.645 since it is not expected to be used
|
||||
for small values of `n`, so z can approximate `t`.
|
||||
|
||||
The significant figure estimation used in conjunction with the
|
||||
`trim_sigfig` method to provide a more human interpretable data
|
||||
summary. __repr__ does not use this method; it simply displays raw
|
||||
values. Significant figure estimation is intended for `Compare`.
|
||||
"""
|
||||
self._lazy_init()
|
||||
n_total = len(self._sorted_times)
|
||||
lower_bound = int(n_total // 4)
|
||||
upper_bound = int(torch.tensor(3 * n_total / 4).ceil())
|
||||
interquartile_points: tuple[float, ...] = self._sorted_times[lower_bound:upper_bound]
|
||||
std = torch.tensor(interquartile_points).std(unbiased=False).item()
|
||||
sqrt_n = torch.tensor(len(interquartile_points)).sqrt().item()
|
||||
|
||||
# Rough estimates. These are by no means statistically rigorous.
|
||||
confidence_interval = max(1.645 * std / sqrt_n, _MIN_CONFIDENCE_INTERVAL)
|
||||
relative_ci = torch.tensor(self._median / confidence_interval).log10().item()
|
||||
num_significant_figures = int(torch.tensor(relative_ci).floor())
|
||||
return min(max(num_significant_figures, 1), _MAX_SIGNIFICANT_FIGURES)
|
||||
|
||||
@property
|
||||
def has_warnings(self) -> bool:
|
||||
self._lazy_init()
|
||||
return bool(self._warnings)
|
||||
|
||||
def _lazy_init(self) -> None:
|
||||
if self.raw_times and not self._sorted_times:
|
||||
self._sorted_times = tuple(sorted(self.times))
|
||||
_sorted_times = torch.tensor(self._sorted_times, dtype=torch.float64)
|
||||
self._median = _sorted_times.quantile(.5).item()
|
||||
self._mean = _sorted_times.mean().item()
|
||||
self._p25 = _sorted_times.quantile(.25).item()
|
||||
self._p75 = _sorted_times.quantile(.75).item()
|
||||
|
||||
def add_warning(msg: str) -> None:
|
||||
rel_iqr = self.iqr / self.median * 100
|
||||
self._warnings += (
|
||||
f" WARNING: Interquartile range is {rel_iqr:.1f}% "
|
||||
f"of the median measurement.\n {msg}",
|
||||
)
|
||||
|
||||
if not self.meets_confidence(_IQR_GROSS_WARN_THRESHOLD):
|
||||
add_warning("This suggests significant environmental influence.")
|
||||
elif not self.meets_confidence(_IQR_WARN_THRESHOLD):
|
||||
add_warning("This could indicate system fluctuation.")
|
||||
|
||||
|
||||
def meets_confidence(self, threshold: float = _IQR_WARN_THRESHOLD) -> bool:
|
||||
return self.iqr / self.median < threshold
|
||||
|
||||
@property
|
||||
def title(self) -> str:
|
||||
return self.task_spec.title
|
||||
|
||||
@property
|
||||
def env(self) -> str:
|
||||
return (
|
||||
"Unspecified env" if self.taskspec.env is None
|
||||
else cast(str, self.taskspec.env)
|
||||
)
|
||||
|
||||
@property
|
||||
def as_row_name(self) -> str:
|
||||
return self.sub_label or self.stmt or "[Unknown]"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""
|
||||
Example repr:
|
||||
<utils.common.Measurement object at 0x7f395b6ac110>
|
||||
Broadcasting add (4x8)
|
||||
Median: 5.73 us
|
||||
IQR: 2.25 us (4.01 to 6.26)
|
||||
372 measurements, 100 runs per measurement, 1 thread
|
||||
WARNING: Interquartile range is 39.4% of the median measurement.
|
||||
This suggests significant environmental influence.
|
||||
"""
|
||||
self._lazy_init()
|
||||
skip_line, newline = "MEASUREMENT_REPR_SKIP_LINE", "\n"
|
||||
n = len(self._sorted_times)
|
||||
time_unit, time_scale = select_unit(self._median)
|
||||
iqr_filter = '' if n >= 4 else skip_line
|
||||
|
||||
repr_str = f"""
|
||||
{super().__repr__()}
|
||||
{self.task_spec.summarize()}
|
||||
{'Median: ' if n > 1 else ''}{self._median / time_scale:.2f} {time_unit}
|
||||
{iqr_filter}IQR: {self.iqr / time_scale:.2f} {time_unit} ({self._p25 / time_scale:.2f} to {self._p75 / time_scale:.2f})
|
||||
{n} measurement{'s' if n > 1 else ''}, {self.number_per_run} runs {'per measurement,' if n > 1 else ','} {self.num_threads} thread{'s' if self.num_threads > 1 else ''}
|
||||
{newline.join(self._warnings)}""".strip() # noqa: B950
|
||||
|
||||
return "\n".join(l for l in repr_str.splitlines(keepends=False) if skip_line not in l)
|
||||
|
||||
@staticmethod
|
||||
def merge(measurements: Iterable["Measurement"]) -> list["Measurement"]:
|
||||
"""Convenience method for merging replicates.
|
||||
|
||||
Merge will extrapolate times to `number_per_run=1` and will not
|
||||
transfer any metadata. (Since it might differ between replicates)
|
||||
"""
|
||||
grouped_measurements: collections.defaultdict[TaskSpec, list[Measurement]] = collections.defaultdict(list)
|
||||
for m in measurements:
|
||||
grouped_measurements[m.task_spec].append(m)
|
||||
|
||||
def merge_group(task_spec: TaskSpec, group: list["Measurement"]) -> "Measurement":
|
||||
times: list[float] = []
|
||||
for m in group:
|
||||
# Different measurements could have different `number_per_run`,
|
||||
# so we call `.times` which normalizes the results.
|
||||
times.extend(m.times)
|
||||
|
||||
return Measurement(
|
||||
number_per_run=1,
|
||||
raw_times=times,
|
||||
task_spec=task_spec,
|
||||
metadata=None,
|
||||
)
|
||||
|
||||
return [merge_group(t, g) for t, g in grouped_measurements.items()]
|
||||
|
||||
|
||||
def select_unit(t: float) -> tuple[str, float]:
|
||||
"""Determine how to scale times for O(1) magnitude.
|
||||
|
||||
This utility is used to format numbers for human consumption.
|
||||
"""
|
||||
time_unit = {-3: "ns", -2: "us", -1: "ms"}.get(int(torch.tensor(t).log10().item() // 3), "s")
|
||||
time_scale = {"ns": 1e-9, "us": 1e-6, "ms": 1e-3, "s": 1}[time_unit]
|
||||
return time_unit, time_scale
|
||||
|
||||
|
||||
def unit_to_english(u: str) -> str:
|
||||
return {
|
||||
"ns": "nanosecond",
|
||||
"us": "microsecond",
|
||||
"ms": "millisecond",
|
||||
"s": "second",
|
||||
}[u]
|
||||
|
||||
|
||||
def trim_sigfig(x: float, n: int) -> float:
|
||||
"""Trim `x` to `n` significant figures. (e.g. 3.14159, 2 -> 3.10000)"""
|
||||
if n != int(n):
|
||||
raise AssertionError("Number of significant figures must be an integer")
|
||||
magnitude = int(torch.tensor(x).abs().log10().ceil().item())
|
||||
scale = 10 ** (magnitude - n)
|
||||
return float(torch.tensor(x / scale).round() * scale)
|
||||
|
||||
|
||||
def ordered_unique(elements: Iterable[Any]) -> list[Any]:
|
||||
return list(collections.OrderedDict(dict.fromkeys(elements)).keys())
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def set_torch_threads(n: int) -> Iterator[None]:
|
||||
prior_num_threads = torch.get_num_threads()
|
||||
try:
|
||||
torch.set_num_threads(n)
|
||||
yield
|
||||
finally:
|
||||
torch.set_num_threads(prior_num_threads)
|
||||
|
||||
|
||||
def _make_temp_dir(prefix: str | None = None, gc_dev_shm: bool = False) -> str:
|
||||
"""Create a temporary directory. The caller is responsible for cleanup.
|
||||
|
||||
This function is conceptually similar to `tempfile.mkdtemp`, but with
|
||||
the key additional feature that it will use shared memory if the
|
||||
`BENCHMARK_USE_DEV_SHM` environment variable is set. This is an
|
||||
implementation detail, but an important one for cases where many Callgrind
|
||||
measurements are collected at once. (Such as when collecting
|
||||
microbenchmarks.)
|
||||
|
||||
This is an internal utility, and is exported solely so that microbenchmarks
|
||||
can reuse the util.
|
||||
"""
|
||||
use_dev_shm: bool = (os.getenv("BENCHMARK_USE_DEV_SHM") or "").lower() in ("1", "true")
|
||||
if use_dev_shm:
|
||||
root = "/dev/shm/pytorch_benchmark_utils"
|
||||
if os.name != "posix":
|
||||
raise AssertionError(f"tmpfs (/dev/shm) is POSIX only, current platform is {os.name}")
|
||||
if not os.path.exists("/dev/shm"):
|
||||
raise AssertionError("This system does not appear to support tmpfs (/dev/shm).")
|
||||
os.makedirs(root, exist_ok=True)
|
||||
|
||||
# Because we're working in shared memory, it is more important than
|
||||
# usual to clean up ALL intermediate files. However we don't want every
|
||||
# worker to walk over all outstanding directories, so instead we only
|
||||
# check when we are sure that it won't lead to contention.
|
||||
if gc_dev_shm:
|
||||
for i in os.listdir(root):
|
||||
owner_file = os.path.join(root, i, "owner.pid")
|
||||
if not os.path.exists(owner_file):
|
||||
continue
|
||||
|
||||
with open(owner_file) as f:
|
||||
owner_pid = int(f.read())
|
||||
|
||||
if owner_pid == os.getpid():
|
||||
continue
|
||||
|
||||
try:
|
||||
# https://stackoverflow.com/questions/568271/how-to-check-if-there-exists-a-process-with-a-given-pid-in-python
|
||||
os.kill(owner_pid, 0)
|
||||
|
||||
except OSError:
|
||||
print(f"Detected that {os.path.join(root, i)} was orphaned in shared memory. Cleaning up.")
|
||||
shutil.rmtree(os.path.join(root, i))
|
||||
|
||||
else:
|
||||
root = tempfile.gettempdir()
|
||||
|
||||
# We include the time so names sort by creation time, and add a UUID
|
||||
# to ensure we don't collide.
|
||||
name = f"{prefix or tempfile.gettempprefix()}__{int(time.time())}__{uuid.uuid4()}"
|
||||
path = os.path.join(root, name)
|
||||
os.makedirs(path, exist_ok=False)
|
||||
|
||||
if use_dev_shm:
|
||||
with open(os.path.join(path, "owner.pid"), "w") as f:
|
||||
f.write(str(os.getpid()))
|
||||
|
||||
return path
|
||||
@@ -0,0 +1,345 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Display class to aggregate and print the results of many measurements."""
|
||||
import collections
|
||||
import enum
|
||||
import itertools as it
|
||||
|
||||
from torch.utils.benchmark.utils import common
|
||||
from torch import tensor as _tensor
|
||||
import operator
|
||||
|
||||
__all__ = ["Colorize", "Compare"]
|
||||
|
||||
BEST = "\033[92m"
|
||||
GOOD = "\033[34m"
|
||||
BAD = "\033[2m\033[91m"
|
||||
VERY_BAD = "\033[31m"
|
||||
BOLD = "\033[1m"
|
||||
TERMINATE = "\033[0m"
|
||||
|
||||
|
||||
class Colorize(enum.Enum):
|
||||
NONE = "none"
|
||||
COLUMNWISE = "columnwise"
|
||||
ROWWISE = "rowwise"
|
||||
|
||||
|
||||
# Classes to separate internal bookkeeping from what is rendered.
|
||||
class _Column:
|
||||
def __init__(
|
||||
self,
|
||||
grouped_results: list[tuple[common.Measurement | None, ...]],
|
||||
time_scale: float,
|
||||
time_unit: str,
|
||||
trim_significant_figures: bool,
|
||||
highlight_warnings: bool,
|
||||
) -> None:
|
||||
self._grouped_results = grouped_results
|
||||
self._flat_results = [*it.chain.from_iterable(grouped_results)]
|
||||
self._time_scale = time_scale
|
||||
self._time_unit = time_unit
|
||||
self._trim_significant_figures = trim_significant_figures
|
||||
self._highlight_warnings = (
|
||||
highlight_warnings
|
||||
and any(r.has_warnings for r in self._flat_results if r)
|
||||
)
|
||||
leading_digits = [
|
||||
int(_tensor(r.median / self._time_scale).log10().ceil()) if r else None
|
||||
for r in self._flat_results
|
||||
]
|
||||
unit_digits = max(d for d in leading_digits if d is not None)
|
||||
decimal_digits = min(
|
||||
max(m.significant_figures - digits, 0)
|
||||
for digits, m in zip(leading_digits, self._flat_results, strict=True)
|
||||
if (m is not None) and (digits is not None)
|
||||
) if self._trim_significant_figures else 1
|
||||
length = unit_digits + decimal_digits + (1 if decimal_digits else 0)
|
||||
self._template = f"{{:>{length}.{decimal_digits}f}}{{:>{7 if self._highlight_warnings else 0}}}"
|
||||
|
||||
def get_results_for(self, group):
|
||||
return self._grouped_results[group]
|
||||
|
||||
def num_to_str(self, value: float | None, estimated_sigfigs: int, spread: float | None):
|
||||
if value is None:
|
||||
return " " * len(self.num_to_str(1, estimated_sigfigs, None))
|
||||
|
||||
if self._trim_significant_figures:
|
||||
value = common.trim_sigfig(value, estimated_sigfigs)
|
||||
|
||||
return self._template.format(
|
||||
value,
|
||||
f" (! {spread * 100:.0f}%)" if self._highlight_warnings and spread is not None else "")
|
||||
|
||||
|
||||
def optional_min(seq):
|
||||
l = list(seq)
|
||||
return None if len(l) == 0 else min(l)
|
||||
|
||||
|
||||
class _Row:
|
||||
def __init__(self, results, row_group, render_env, env_str_len,
|
||||
row_name_str_len, time_scale, colorize, num_threads=None) -> None:
|
||||
super().__init__()
|
||||
self._results = results
|
||||
self._row_group = row_group
|
||||
self._render_env = render_env
|
||||
self._env_str_len = env_str_len
|
||||
self._row_name_str_len = row_name_str_len
|
||||
self._time_scale = time_scale
|
||||
self._colorize = colorize
|
||||
self._columns: tuple[_Column, ...] = ()
|
||||
self._num_threads = num_threads
|
||||
|
||||
def register_columns(self, columns: tuple[_Column, ...]) -> None:
|
||||
self._columns = columns
|
||||
|
||||
def as_column_strings(self):
|
||||
concrete_results = [r for r in self._results if r is not None]
|
||||
env = f"({concrete_results[0].env})" if self._render_env else ""
|
||||
env = env.ljust(self._env_str_len + 4)
|
||||
output = [" " + env + concrete_results[0].as_row_name]
|
||||
for m, col in zip(self._results, self._columns or (), strict=False):
|
||||
if m is None:
|
||||
output.append(col.num_to_str(None, 1, None))
|
||||
else:
|
||||
output.append(col.num_to_str(
|
||||
m.median / self._time_scale,
|
||||
m.significant_figures,
|
||||
m.iqr / m.median if m.has_warnings else None
|
||||
))
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def color_segment(segment, value, best_value):
|
||||
if value <= best_value * 1.01 or value <= best_value + 100e-9:
|
||||
return BEST + BOLD + segment + TERMINATE * 2
|
||||
if value <= best_value * 1.1:
|
||||
return GOOD + BOLD + segment + TERMINATE * 2
|
||||
if value >= best_value * 5:
|
||||
return VERY_BAD + BOLD + segment + TERMINATE * 2
|
||||
if value >= best_value * 2:
|
||||
return BAD + segment + TERMINATE * 2
|
||||
|
||||
return segment
|
||||
|
||||
def row_separator(self, overall_width):
|
||||
return (
|
||||
[f"{self._num_threads} threads: ".ljust(overall_width, "-")]
|
||||
if self._num_threads is not None else []
|
||||
)
|
||||
|
||||
def finalize_column_strings(self, column_strings, col_widths):
|
||||
best_values = [-1 for _ in column_strings]
|
||||
if self._colorize == Colorize.ROWWISE:
|
||||
row_min = min(r.median for r in self._results if r is not None)
|
||||
best_values = [row_min for _ in column_strings]
|
||||
elif self._colorize == Colorize.COLUMNWISE:
|
||||
best_values = [
|
||||
optional_min(r.median for r in column.get_results_for(self._row_group) if r is not None)
|
||||
for column in (self._columns or ())
|
||||
]
|
||||
|
||||
row_contents = [column_strings[0].ljust(col_widths[0])]
|
||||
for col_str, width, result, best_value in zip(column_strings[1:], col_widths[1:], self._results, best_values, strict=False):
|
||||
col_str = col_str.center(width)
|
||||
if self._colorize != Colorize.NONE and result is not None and best_value is not None:
|
||||
col_str = self.color_segment(col_str, result.median, best_value)
|
||||
row_contents.append(col_str)
|
||||
return row_contents
|
||||
|
||||
|
||||
class Table:
|
||||
def __init__(
|
||||
self,
|
||||
results: list[common.Measurement],
|
||||
colorize: Colorize,
|
||||
trim_significant_figures: bool,
|
||||
highlight_warnings: bool
|
||||
) -> None:
|
||||
if len({r.label for r in results}) != 1:
|
||||
raise AssertionError("All results must share the same label")
|
||||
|
||||
self.results = results
|
||||
self._colorize = colorize
|
||||
self._trim_significant_figures = trim_significant_figures
|
||||
self._highlight_warnings = highlight_warnings
|
||||
self.label = results[0].label
|
||||
self.time_unit, self.time_scale = common.select_unit(
|
||||
min(r.median for r in results)
|
||||
)
|
||||
|
||||
self.row_keys = common.ordered_unique([self.row_fn(i) for i in results])
|
||||
self.row_keys.sort(key=operator.itemgetter(slice(2))) # preserve stmt order
|
||||
self.column_keys = common.ordered_unique([self.col_fn(i) for i in results])
|
||||
self.rows, self.columns = self.populate_rows_and_columns()
|
||||
|
||||
@staticmethod
|
||||
def row_fn(m: common.Measurement) -> tuple[int, str | None, str]:
|
||||
return m.num_threads, m.env, m.as_row_name
|
||||
|
||||
@staticmethod
|
||||
def col_fn(m: common.Measurement) -> str | None:
|
||||
return m.description
|
||||
|
||||
def populate_rows_and_columns(self) -> tuple[tuple[_Row, ...], tuple[_Column, ...]]:
|
||||
rows: list[_Row] = []
|
||||
columns: list[_Column] = []
|
||||
ordered_results: list[list[common.Measurement | None]] = [
|
||||
[None for _ in self.column_keys]
|
||||
for _ in self.row_keys
|
||||
]
|
||||
row_position = {key: i for i, key in enumerate(self.row_keys)}
|
||||
col_position = {key: i for i, key in enumerate(self.column_keys)}
|
||||
for r in self.results:
|
||||
i = row_position[self.row_fn(r)]
|
||||
j = col_position[self.col_fn(r)]
|
||||
ordered_results[i][j] = r
|
||||
|
||||
unique_envs = {r.env for r in self.results}
|
||||
render_env = len(unique_envs) > 1
|
||||
env_str_len = max(len(i) for i in unique_envs) if render_env else 0
|
||||
|
||||
row_name_str_len = max(len(r.as_row_name) for r in self.results)
|
||||
|
||||
prior_num_threads = -1
|
||||
prior_env = ""
|
||||
row_group = -1
|
||||
rows_by_group: list[list[list[common.Measurement | None]]] = []
|
||||
for (num_threads, env, _), row in zip(self.row_keys, ordered_results, strict=True):
|
||||
thread_transition = (num_threads != prior_num_threads)
|
||||
if thread_transition:
|
||||
prior_num_threads = num_threads
|
||||
prior_env = ""
|
||||
row_group += 1
|
||||
rows_by_group.append([])
|
||||
rows.append(
|
||||
_Row(
|
||||
results=row,
|
||||
row_group=row_group,
|
||||
render_env=(render_env and env != prior_env),
|
||||
env_str_len=env_str_len,
|
||||
row_name_str_len=row_name_str_len,
|
||||
time_scale=self.time_scale,
|
||||
colorize=self._colorize,
|
||||
num_threads=num_threads if thread_transition else None,
|
||||
)
|
||||
)
|
||||
rows_by_group[-1].append(row)
|
||||
prior_env = env
|
||||
|
||||
for i in range(len(self.column_keys)):
|
||||
grouped_results = [tuple(row[i] for row in g) for g in rows_by_group]
|
||||
column = _Column(
|
||||
grouped_results=grouped_results,
|
||||
time_scale=self.time_scale,
|
||||
time_unit=self.time_unit,
|
||||
trim_significant_figures=self._trim_significant_figures,
|
||||
highlight_warnings=self._highlight_warnings,)
|
||||
columns.append(column)
|
||||
|
||||
rows_tuple, columns_tuple = tuple(rows), tuple(columns)
|
||||
for ri in rows_tuple:
|
||||
ri.register_columns(columns_tuple)
|
||||
return rows_tuple, columns_tuple
|
||||
|
||||
def render(self) -> str:
|
||||
string_rows = [[""] + self.column_keys]
|
||||
string_rows.extend(r.as_column_strings() for r in self.rows)
|
||||
num_cols = max(len(i) for i in string_rows)
|
||||
for sr in string_rows:
|
||||
sr.extend(["" for _ in range(num_cols - len(sr))])
|
||||
|
||||
col_widths = [max(len(j) for j in i) for i in zip(*string_rows, strict=True)]
|
||||
finalized_columns = [" | ".join(i.center(w) for i, w in zip(string_rows[0], col_widths, strict=True))]
|
||||
overall_width = len(finalized_columns[0])
|
||||
for string_row, row in zip(string_rows[1:], self.rows, strict=True):
|
||||
finalized_columns.extend(row.row_separator(overall_width))
|
||||
finalized_columns.append(" | ".join(row.finalize_column_strings(string_row, col_widths)))
|
||||
|
||||
newline = "\n"
|
||||
has_warnings = self._highlight_warnings and any(ri.has_warnings for ri in self.results)
|
||||
return f"""
|
||||
[{(' ' + (self.label or '') + ' ').center(overall_width - 2, '-')}]
|
||||
{newline.join(finalized_columns)}
|
||||
|
||||
Times are in {common.unit_to_english(self.time_unit)}s ({self.time_unit}).
|
||||
{'(! XX%) Measurement has high variance, where XX is the IQR / median * 100.' + newline if has_warnings else ""}"""[1:]
|
||||
|
||||
|
||||
class Compare:
|
||||
"""Helper class for displaying the results of many measurements in a
|
||||
formatted table.
|
||||
|
||||
The table format is based on the information fields provided in
|
||||
:class:`torch.utils.benchmark.Timer` (`description`, `label`, `sub_label`,
|
||||
`num_threads`, etc).
|
||||
|
||||
The table can be directly printed using :meth:`print` or casted as a `str`.
|
||||
|
||||
For a full tutorial on how to use this class, see:
|
||||
https://pytorch.org/tutorials/recipes/recipes/benchmark.html
|
||||
|
||||
Args:
|
||||
results: List of Measurement to display.
|
||||
"""
|
||||
def __init__(self, results: list[common.Measurement]) -> None:
|
||||
self._results: list[common.Measurement] = []
|
||||
self.extend_results(results)
|
||||
self._trim_significant_figures = False
|
||||
self._colorize = Colorize.NONE
|
||||
self._highlight_warnings = False
|
||||
|
||||
def __str__(self) -> str:
|
||||
return "\n".join(self._render())
|
||||
|
||||
def extend_results(self, results) -> None:
|
||||
"""Append results to already stored ones.
|
||||
|
||||
All added results must be instances of ``Measurement``.
|
||||
"""
|
||||
for r in results:
|
||||
if not isinstance(r, common.Measurement):
|
||||
raise ValueError(
|
||||
"Expected an instance of `Measurement`, " f"got {type(r)} instead."
|
||||
)
|
||||
self._results.extend(results)
|
||||
|
||||
def trim_significant_figures(self) -> None:
|
||||
"""Enables trimming of significant figures when building the formatted table."""
|
||||
self._trim_significant_figures = True
|
||||
|
||||
def colorize(self, rowwise=False) -> None:
|
||||
"""Colorize formatted table.
|
||||
|
||||
Colorize columnwise by default.
|
||||
"""
|
||||
self._colorize = Colorize.ROWWISE if rowwise else Colorize.COLUMNWISE
|
||||
|
||||
def highlight_warnings(self) -> None:
|
||||
"""Enables warning highlighting when building formatted table."""
|
||||
self._highlight_warnings = True
|
||||
|
||||
def print(self) -> None:
|
||||
"""Print formatted table"""
|
||||
print(str(self))
|
||||
|
||||
def _render(self):
|
||||
results = common.Measurement.merge(self._results)
|
||||
grouped_results = self._group_by_label(results)
|
||||
output = [self._layout(group) for group in grouped_results.values()]
|
||||
return output
|
||||
|
||||
def _group_by_label(self, results: list[common.Measurement]):
|
||||
grouped_results: collections.defaultdict[str, list[common.Measurement]] = collections.defaultdict(list)
|
||||
for r in results:
|
||||
grouped_results[r.label].append(r)
|
||||
return grouped_results
|
||||
|
||||
def _layout(self, results: list[common.Measurement]):
|
||||
table = Table(
|
||||
results,
|
||||
self._colorize,
|
||||
self._trim_significant_figures,
|
||||
self._highlight_warnings
|
||||
)
|
||||
return table.render()
|
||||
@@ -0,0 +1,195 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any, cast
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch._dynamo
|
||||
from torch._dynamo.testing import CompileCounterWithBackend
|
||||
from torch.utils.benchmark import Timer
|
||||
|
||||
|
||||
__all__ = ["bench_all", "benchmark_compile"]
|
||||
|
||||
|
||||
_warned_tensor_cores = False
|
||||
_default_float_32_precision = torch.get_float32_matmul_precision()
|
||||
|
||||
try:
|
||||
|
||||
from tabulate import tabulate
|
||||
|
||||
HAS_TABULATE = True
|
||||
except ModuleNotFoundError:
|
||||
HAS_TABULATE = False
|
||||
tabulate = None # type: ignore[assignment]
|
||||
print("tabulate is not installed, please pip install tabulate to use this utility")
|
||||
|
||||
if HAS_TABULATE:
|
||||
def _enable_tensor_cores() -> None:
|
||||
global _warned_tensor_cores
|
||||
|
||||
if torch.cuda.is_available():
|
||||
if torch.backends.cuda.matmul.allow_tf32 is False and torch.cuda.get_device_capability() >= (8, 0):
|
||||
torch.set_float32_matmul_precision("high")
|
||||
if not _warned_tensor_cores:
|
||||
print("Your GPU supports tensor cores")
|
||||
print("we will enable it automatically by setting `torch.set_float32_matmul_precision('high')`")
|
||||
_warned_tensor_cores = True
|
||||
|
||||
def _disable_tensor_cores() -> None:
|
||||
torch.set_float32_matmul_precision(_default_float_32_precision)
|
||||
|
||||
def bench_loop(
|
||||
model: torch.nn.Module | Callable,
|
||||
sample_input: torch.Tensor | Any,
|
||||
num_iters: int = 5,
|
||||
optimizer: torch.optim.Optimizer | None = None,
|
||||
loss_fn: Callable | None = None,
|
||||
):
|
||||
# Define the statement and setup for the benchmark
|
||||
if optimizer and loss_fn:
|
||||
# Training mode
|
||||
stmt = """
|
||||
output = model(sample_input)
|
||||
loss = loss_fn(output) if loss_fn else output.sum()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
"""
|
||||
else:
|
||||
# Inference mode
|
||||
stmt = "model(sample_input)"
|
||||
|
||||
# Create the Timer object
|
||||
timer = Timer(
|
||||
stmt=stmt,
|
||||
globals={"model": model, "sample_input": sample_input, "optimizer": optimizer, "loss_fn": loss_fn},
|
||||
)
|
||||
|
||||
|
||||
result = timer.timeit(number=num_iters)
|
||||
|
||||
# Get the average time per iteration in milliseconds
|
||||
avg_time = result.mean * 1000
|
||||
return round(avg_time, 2)
|
||||
|
||||
def benchmark_compile(
|
||||
model: torch.nn.Module | Callable,
|
||||
sample_input: torch.Tensor | Any,
|
||||
num_iters: int = 5,
|
||||
backend: str | None = None,
|
||||
mode: str | None = "default",
|
||||
optimizer: torch.optim.Optimizer | None = None,
|
||||
loss_fn : torch.nn.Module | Callable | None = None,
|
||||
):
|
||||
"""
|
||||
Use this utility to benchmark torch.compile
|
||||
"""
|
||||
if backend:
|
||||
try:
|
||||
torch._dynamo.reset()
|
||||
compile_counter_with_backend = CompileCounterWithBackend(backend)
|
||||
opt_model = torch.compile(model, backend=compile_counter_with_backend, mode=mode)
|
||||
|
||||
# Compilation only happens after the first inference
|
||||
compilation_time = bench_loop(opt_model, sample_input, 1, optimizer, loss_fn)
|
||||
|
||||
running_time = bench_loop(opt_model, sample_input, num_iters, optimizer, loss_fn)
|
||||
|
||||
if compile_counter_with_backend.frame_count == 0:
|
||||
raise RuntimeError("No compilation occurred during benchmarking.")
|
||||
|
||||
if compile_counter_with_backend.frame_count > 1:
|
||||
raise RuntimeError("Recompilation occurred during benchmarking.")
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
print(f"Failed to compile {backend} with mode {mode}")
|
||||
return None, None
|
||||
else:
|
||||
opt_model = model
|
||||
compilation_time = None
|
||||
running_time = bench_loop(opt_model, sample_input, num_iters, optimizer, loss_fn)
|
||||
|
||||
compilation_time = round(compilation_time, 2) if compilation_time else None
|
||||
running_time = round(running_time, 2) if running_time else None
|
||||
|
||||
|
||||
return compilation_time, running_time
|
||||
|
||||
|
||||
def bench_all(
|
||||
model : torch.nn.Module | Callable,
|
||||
sample_input: torch.Tensor | Any,
|
||||
num_iters : int = 5,
|
||||
optimizer: torch.optim.Optimizer | None = None,
|
||||
loss_fn : torch.nn.Module | Callable | None = None,
|
||||
):
|
||||
"""
|
||||
This is a simple utility that can be used to benchmark torch.compile
|
||||
In particular it ensures that your GPU is setup to use tensor cores if it supports its
|
||||
It also tries out all the main backends and prints a table of results so you can easily compare them all
|
||||
Many of the backendds have their own optional dependencies so please pip install them separately
|
||||
|
||||
You will get one table for inference and another for training
|
||||
If you'd like to leverage this utility for training make sure to pass in a torch.optim.Optimizer
|
||||
|
||||
The important warnings are
|
||||
Your GPU supports tensor cores
|
||||
we will enable it automatically by setting `torch.set_float32_matmul_precision('high')`
|
||||
|
||||
If a compilation fails for any reason including the dependency not being included
|
||||
then we will print Failed to compile {backend} with mode {mode}
|
||||
"""
|
||||
field_names = ["Train/Inference", "Backend", "Mode", "Compilation Time", "Average Running Time"]
|
||||
table = []
|
||||
|
||||
|
||||
eager_time = None
|
||||
torch._dynamo.reset()
|
||||
_, eager_time = benchmark_compile(model, sample_input, num_iters, None, None, optimizer)
|
||||
table.append(
|
||||
[("Training" if optimizer else "Inference"), "Eager", "-", "-", f"{eager_time} ms"]
|
||||
)
|
||||
|
||||
for backend in torch._dynamo.list_backends():
|
||||
|
||||
if backend == "inductor":
|
||||
mode_options = cast(list[str | None], list(torch._inductor.list_mode_options().keys())) + [None]
|
||||
for mode in mode_options:
|
||||
if mode == "default":
|
||||
continue
|
||||
torch._dynamo.reset()
|
||||
try:
|
||||
if torch.cuda.is_available():
|
||||
_enable_tensor_cores()
|
||||
compilation_time, running_time = benchmark_compile(
|
||||
model, sample_input, num_iters, backend, mode, optimizer, loss_fn)
|
||||
finally:
|
||||
if torch.cuda.is_available():
|
||||
_disable_tensor_cores()
|
||||
table.append([
|
||||
("Training" if optimizer else "Inference"),
|
||||
# pyrefly: ignore [redundant-condition]
|
||||
backend if backend else "-",
|
||||
mode if mode is not None else "-",
|
||||
f"{compilation_time} ms " if compilation_time else "-",
|
||||
f"{running_time} ms " if running_time else "-",
|
||||
])
|
||||
|
||||
else:
|
||||
torch._dynamo.reset()
|
||||
compilation_time, running_time = benchmark_compile(
|
||||
model, sample_input, num_iters, backend, None, optimizer, loss_fn)
|
||||
|
||||
if running_time is not None:
|
||||
table.append([
|
||||
("Training" if optimizer else "Inference"),
|
||||
backend, "-",
|
||||
f"{compilation_time} ms " or "-",
|
||||
f"{running_time} ms ",
|
||||
])
|
||||
|
||||
|
||||
# pyrefly: ignore [not-callable]
|
||||
return tabulate(table, headers=field_names, tablefmt="github")
|
||||
@@ -0,0 +1,175 @@
|
||||
"""JIT C++ strings into executables."""
|
||||
import atexit
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import textwrap
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.utils.benchmark.utils._stubs import CallgrindModuleType, TimeitModuleType
|
||||
from torch.utils.benchmark.utils.common import _make_temp_dir
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
|
||||
LOCK = threading.Lock()
|
||||
SOURCE_ROOT = os.path.split(os.path.abspath(__file__))[0]
|
||||
|
||||
# We calculate uuid once at import time so that separate processes will have
|
||||
# separate build roots, but threads will share the same build root.
|
||||
# `cpp_extension` uses build root as part of the cache key, so per-invocation
|
||||
# uuid's (e.g. different build root per _compile_template call) would lead to
|
||||
# a 0% cache hit rate and spurious recompilation. Consider the following:
|
||||
# ```
|
||||
# setup = "auto x = torch::ones({1024, 1024});"
|
||||
# stmt = "torch::mm(x, x);"
|
||||
# for num_threads in [1, 2, 4, 8]:
|
||||
# print(Timer(stmt, setup, num_threads=num_threads, language="c++").blocked_autorange())
|
||||
# ````
|
||||
# `setup` and `stmt` do not change, so we can reuse the executable from the
|
||||
# first pass through the loop.
|
||||
_BUILD_ROOT: str | None = None
|
||||
|
||||
def _get_build_root() -> str:
|
||||
global _BUILD_ROOT
|
||||
if _BUILD_ROOT is None:
|
||||
_BUILD_ROOT = _make_temp_dir(prefix="benchmark_utils_jit_build")
|
||||
# pyrefly: ignore [missing-argument]
|
||||
atexit.register(shutil.rmtree, _BUILD_ROOT)
|
||||
return _BUILD_ROOT
|
||||
|
||||
|
||||
# BACK_TESTING_NOTE:
|
||||
# There are two workflows where this code could be used. One is the obvious
|
||||
# case where someone simply builds or installs PyTorch and uses Timer.
|
||||
# The other is that the entire `torch/utils/benchmark` folder from a CURRENT
|
||||
# PyTorch checkout is copy-pasted into a much OLDER version of the PyTorch
|
||||
# source code. This is what we refer to here as "back testing". The rationale
|
||||
# is that we might want to use current tooling to study some aspect of an
|
||||
# earlier version of PyTorch. (e.g. a regression.)
|
||||
#
|
||||
# The problem is that Timer relies on several aspects of core PyTorch, namely
|
||||
# some binding functions for Valgrind symbols in `torch._C` and the
|
||||
# `torch.__config__._cxx_flags()` method. If we were to naively copy code
|
||||
# around this wouldn't work as the symbols of interest aren't present in
|
||||
# earlier versions of PyTorch. In order to work around this, we must add back
|
||||
# testing shims. These shims will never activate during normal use, but will
|
||||
# allow Timer to function outside of the "correct" version of PyTorch by
|
||||
# emulating functionality that was added later.
|
||||
#
|
||||
# These shims are temporary, and as Timer becomes more integrated with
|
||||
# PyTorch the cost and complexity of such shims will increase. Once back
|
||||
# testing is no longer required (which is to say we have done enough historic
|
||||
# analysis and the shims no longer justify their maintenance and code
|
||||
# complexity costs) back testing paths will be removed.
|
||||
|
||||
CXX_FLAGS: list[str] | None
|
||||
if hasattr(torch.__config__, "_cxx_flags"):
|
||||
try:
|
||||
CXX_FLAGS = torch.__config__._cxx_flags().strip().split()
|
||||
if CXX_FLAGS is not None and "-g" not in CXX_FLAGS:
|
||||
CXX_FLAGS.append("-g")
|
||||
# remove "-W" flags to allow build benchmarks
|
||||
# with a relaxed constraint of compiler versions
|
||||
if CXX_FLAGS is not None:
|
||||
CXX_FLAGS = list(filter(lambda x: not x.startswith("-W"), CXX_FLAGS))
|
||||
|
||||
except RuntimeError:
|
||||
# We are in FBCode.
|
||||
CXX_FLAGS = None
|
||||
else:
|
||||
# FIXME: Remove when back testing is no longer required.
|
||||
CXX_FLAGS = ["-O2", "-fPIC", "-g"]
|
||||
|
||||
EXTRA_INCLUDE_PATHS: list[str] = [os.path.join(SOURCE_ROOT, "valgrind_wrapper")]
|
||||
CONDA_PREFIX = os.getenv("CONDA_PREFIX")
|
||||
if CONDA_PREFIX is not None:
|
||||
# Load will automatically search /usr/include, but not conda include.
|
||||
EXTRA_INCLUDE_PATHS.append(os.path.join(CONDA_PREFIX, "include"))
|
||||
|
||||
|
||||
COMPAT_CALLGRIND_BINDINGS: CallgrindModuleType | None = None
|
||||
def get_compat_bindings() -> CallgrindModuleType:
|
||||
with LOCK:
|
||||
global COMPAT_CALLGRIND_BINDINGS
|
||||
if COMPAT_CALLGRIND_BINDINGS is None:
|
||||
COMPAT_CALLGRIND_BINDINGS = cpp_extension.load(
|
||||
name="callgrind_bindings",
|
||||
sources=[os.path.join(
|
||||
SOURCE_ROOT,
|
||||
"valgrind_wrapper",
|
||||
"compat_bindings.cpp"
|
||||
)],
|
||||
extra_cflags=CXX_FLAGS,
|
||||
extra_include_paths=EXTRA_INCLUDE_PATHS,
|
||||
)
|
||||
return COMPAT_CALLGRIND_BINDINGS
|
||||
|
||||
|
||||
def _compile_template(
|
||||
*,
|
||||
stmt: str,
|
||||
setup: str,
|
||||
global_setup: str,
|
||||
src: str,
|
||||
is_standalone: bool
|
||||
) -> Any:
|
||||
for before, after, indentation in (
|
||||
("// GLOBAL_SETUP_TEMPLATE_LOCATION", global_setup, 0),
|
||||
("// SETUP_TEMPLATE_LOCATION", setup, 4),
|
||||
("// STMT_TEMPLATE_LOCATION", stmt, 8)
|
||||
):
|
||||
# C++ doesn't care about indentation so this code isn't load
|
||||
# bearing the way it is with Python, but this makes the source
|
||||
# look nicer if a human has to look at it.
|
||||
src = re.sub(
|
||||
before,
|
||||
textwrap.indent(after, " " * indentation)[indentation:],
|
||||
src
|
||||
)
|
||||
|
||||
# We want to isolate different Timers. However `cpp_extension` will
|
||||
# cache builds which will significantly reduce the cost of repeated
|
||||
# invocations.
|
||||
with LOCK:
|
||||
name = f"timer_cpp_{abs(hash(src))}"
|
||||
build_dir = os.path.join(_get_build_root(), name)
|
||||
os.makedirs(build_dir, exist_ok=True)
|
||||
|
||||
src_path = os.path.join(build_dir, "timer_src.cpp")
|
||||
with open(src_path, "w") as f:
|
||||
f.write(src)
|
||||
|
||||
# `cpp_extension` has its own locking scheme, so we don't need our lock.
|
||||
return cpp_extension.load(
|
||||
name=name,
|
||||
sources=[src_path],
|
||||
build_directory=build_dir,
|
||||
extra_cflags=CXX_FLAGS,
|
||||
extra_include_paths=EXTRA_INCLUDE_PATHS,
|
||||
is_python_module=not is_standalone,
|
||||
is_standalone=is_standalone,
|
||||
)
|
||||
|
||||
|
||||
def compile_timeit_template(*, stmt: str, setup: str, global_setup: str) -> TimeitModuleType:
|
||||
template_path: str = os.path.join(SOURCE_ROOT, "timeit_template.cpp")
|
||||
with open(template_path) as f:
|
||||
src: str = f.read()
|
||||
|
||||
module = _compile_template(stmt=stmt, setup=setup, global_setup=global_setup, src=src, is_standalone=False)
|
||||
if not isinstance(module, TimeitModuleType):
|
||||
raise AssertionError("compiled module is not a TimeitModuleType")
|
||||
return module
|
||||
|
||||
|
||||
def compile_callgrind_template(*, stmt: str, setup: str, global_setup: str) -> str:
|
||||
template_path: str = os.path.join(SOURCE_ROOT, "valgrind_wrapper", "timer_callgrind_template.cpp")
|
||||
with open(template_path) as f:
|
||||
src: str = f.read()
|
||||
|
||||
target = _compile_template(stmt=stmt, setup=setup, global_setup=global_setup, src=src, is_standalone=True)
|
||||
if not isinstance(target, str):
|
||||
raise AssertionError("compiled target path is not a string")
|
||||
return target
|
||||
@@ -0,0 +1,469 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import itertools as it
|
||||
from typing import Any
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Fuzzer",
|
||||
"FuzzedParameter", "ParameterAlias",
|
||||
"FuzzedTensor",
|
||||
]
|
||||
|
||||
|
||||
_DISTRIBUTIONS = (
|
||||
"loguniform",
|
||||
"uniform",
|
||||
)
|
||||
|
||||
|
||||
class FuzzedParameter:
|
||||
"""Specification for a parameter to be generated during fuzzing."""
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
minval: int | float | None = None,
|
||||
maxval: int | float | None = None,
|
||||
distribution: str | dict[Any, float] | None = None,
|
||||
strict: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
name:
|
||||
A string name with which to identify the parameter.
|
||||
FuzzedTensors can reference this string in their
|
||||
specifications.
|
||||
minval:
|
||||
The lower bound for the generated value. See the description
|
||||
of `distribution` for type behavior.
|
||||
maxval:
|
||||
The upper bound for the generated value. Type behavior is
|
||||
identical to `minval`.
|
||||
distribution:
|
||||
Specifies the distribution from which this parameter should
|
||||
be drawn. There are three possibilities:
|
||||
- "loguniform"
|
||||
Samples between `minval` and `maxval` (inclusive) such
|
||||
that the probabilities are uniform in log space. As a
|
||||
concrete example, if minval=1 and maxval=100, a sample
|
||||
is as likely to fall in [1, 10) as it is [10, 100].
|
||||
- "uniform"
|
||||
Samples are chosen with uniform probability between
|
||||
`minval` and `maxval` (inclusive). If either `minval`
|
||||
or `maxval` is a float then the distribution is the
|
||||
continuous uniform distribution; otherwise samples
|
||||
are constrained to the integers.
|
||||
- dict:
|
||||
If a dict is passed, the keys are taken to be choices
|
||||
for the variables and the values are interpreted as
|
||||
probabilities. (And must sum to one.)
|
||||
If a dict is passed, `minval` and `maxval` must not be set.
|
||||
Otherwise, they must be set.
|
||||
strict:
|
||||
If a parameter is strict, it will not be included in the
|
||||
iterative resampling process which Fuzzer uses to find a
|
||||
valid parameter configuration. This allows an author to
|
||||
prevent skew from resampling for a given parameter (for
|
||||
instance, a low size limit could inadvertently bias towards
|
||||
Tensors with fewer dimensions) at the cost of more iterations
|
||||
when generating parameters.
|
||||
"""
|
||||
self._name = name
|
||||
self._minval = minval
|
||||
self._maxval = maxval
|
||||
self._distribution = self._check_distribution(distribution)
|
||||
self.strict = strict
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
def sample(self, state):
|
||||
if self._distribution == "loguniform":
|
||||
return self._loguniform(state)
|
||||
|
||||
if self._distribution == "uniform":
|
||||
return self._uniform(state)
|
||||
|
||||
if isinstance(self._distribution, dict):
|
||||
return self._custom_distribution(state)
|
||||
|
||||
def _check_distribution(self, distribution):
|
||||
if not isinstance(distribution, dict):
|
||||
if distribution not in _DISTRIBUTIONS:
|
||||
raise AssertionError(f"Unknown distribution: {distribution}")
|
||||
else:
|
||||
if any(i < 0 for i in distribution.values()):
|
||||
raise AssertionError("Probabilities cannot be negative")
|
||||
if not abs(sum(distribution.values()) - 1) > 1e-5:
|
||||
raise AssertionError("Distribution is not normalized")
|
||||
if self._minval is not None:
|
||||
raise AssertionError("When passing a custom distribution, 'minval' must be None")
|
||||
if self._maxval is not None:
|
||||
raise AssertionError("When passing a custom distribution, 'maxval' must be None")
|
||||
|
||||
return distribution
|
||||
|
||||
def _loguniform(self, state):
|
||||
import numpy as np
|
||||
output = int(2 ** state.uniform(
|
||||
low=np.log2(self._minval) if self._minval is not None else None,
|
||||
high=np.log2(self._maxval) if self._maxval is not None else None,
|
||||
))
|
||||
if self._minval is not None and output < self._minval:
|
||||
return self._minval
|
||||
if self._maxval is not None and output > self._maxval:
|
||||
return self._maxval
|
||||
return output
|
||||
|
||||
def _uniform(self, state):
|
||||
if isinstance(self._minval, int) and isinstance(self._maxval, int):
|
||||
return int(state.randint(low=self._minval, high=self._maxval + 1))
|
||||
return state.uniform(low=self._minval, high=self._maxval)
|
||||
|
||||
def _custom_distribution(self, state):
|
||||
import numpy as np
|
||||
# If we directly pass the keys to `choice`, numpy will convert
|
||||
# them to numpy dtypes.
|
||||
index = state.choice(
|
||||
np.arange(len(self._distribution)),
|
||||
p=tuple(self._distribution.values()))
|
||||
return list(self._distribution.keys())[index]
|
||||
|
||||
|
||||
class ParameterAlias:
|
||||
"""Indicates that a parameter should alias the value of another parameter.
|
||||
|
||||
When used in conjunction with a custom distribution, this allows fuzzed
|
||||
tensors to represent a broader range of behaviors. For example, the
|
||||
following sometimes produces Tensors which broadcast:
|
||||
|
||||
Fuzzer(
|
||||
parameters=[
|
||||
FuzzedParameter("x_len", 4, 1024, distribution="uniform"),
|
||||
|
||||
# `y` will either be size one, or match the size of `x`.
|
||||
FuzzedParameter("y_len", distribution={
|
||||
0.5: 1,
|
||||
0.5: ParameterAlias("x_len")
|
||||
}),
|
||||
],
|
||||
tensors=[
|
||||
FuzzedTensor("x", size=("x_len",)),
|
||||
FuzzedTensor("y", size=("y_len",)),
|
||||
],
|
||||
)
|
||||
|
||||
Chains of alias' are allowed, but may not contain cycles.
|
||||
"""
|
||||
def __init__(self, alias_to) -> None:
|
||||
self.alias_to = alias_to
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"ParameterAlias[alias_to: {self.alias_to}]"
|
||||
|
||||
|
||||
def dtype_size(dtype):
|
||||
if dtype == torch.bool:
|
||||
return 1
|
||||
if dtype.is_floating_point or dtype.is_complex:
|
||||
return int(torch.finfo(dtype).bits / 8)
|
||||
return int(torch.iinfo(dtype).bits / 8)
|
||||
|
||||
|
||||
def prod(values, base=1):
|
||||
"""np.prod can overflow, so for sizes the product should be done in Python.
|
||||
|
||||
Even though np.prod type promotes to int64, it can still overflow in which
|
||||
case the negative value will pass the size check and OOM when attempting to
|
||||
actually allocate the Tensor.
|
||||
"""
|
||||
return functools.reduce(lambda x, y: int(x) * int(y), values, base)
|
||||
|
||||
|
||||
class FuzzedTensor:
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
size: tuple[str | int, ...],
|
||||
steps: tuple[str | int, ...] | None = None,
|
||||
probability_contiguous: float = 0.5,
|
||||
min_elements: int | None = None,
|
||||
max_elements: int | None = None,
|
||||
max_allocation_bytes: int | None = None,
|
||||
dim_parameter: str | None = None,
|
||||
roll_parameter: str | None = None,
|
||||
dtype=torch.float32,
|
||||
cuda=False,
|
||||
tensor_constructor: Callable | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
name:
|
||||
A string identifier for the generated Tensor.
|
||||
size:
|
||||
A tuple of integers or strings specifying the size of the generated
|
||||
Tensor. String values will replaced with a concrete int during the
|
||||
generation process, while ints are simply passed as literals.
|
||||
steps:
|
||||
An optional tuple with the same length as `size`. This indicates
|
||||
that a larger Tensor should be allocated, and then sliced to
|
||||
produce the generated Tensor. For instance, if size is (4, 8)
|
||||
and steps is (1, 4), then a tensor `t` of size (4, 32) will be
|
||||
created and then `t[:, ::4]` will be used. (Allowing one to test
|
||||
Tensors with strided memory.)
|
||||
probability_contiguous:
|
||||
A number between zero and one representing the chance that the
|
||||
generated Tensor has a contiguous memory layout. This is achieved by
|
||||
randomly permuting the shape of a Tensor, calling `.contiguous()`,
|
||||
and then permuting back. This is applied before `steps`, which can
|
||||
also cause a Tensor to be non-contiguous.
|
||||
min_elements:
|
||||
The minimum number of parameters that this Tensor must have for a
|
||||
set of parameters to be valid. (Otherwise they are resampled.)
|
||||
max_elements:
|
||||
Like `min_elements`, but setting an upper bound.
|
||||
max_allocation_bytes:
|
||||
Like `max_elements`, but for the size of Tensor that must be
|
||||
allocated prior to slicing for `steps` (if applicable). For
|
||||
example, a FloatTensor with size (1024, 1024) and steps (4, 4)
|
||||
would have 1M elements, but would require a 64 MB allocation.
|
||||
dim_parameter:
|
||||
The length of `size` and `steps` will be truncated to this value.
|
||||
This allows Tensors of varying dimensions to be generated by the
|
||||
Fuzzer.
|
||||
dtype:
|
||||
The PyTorch dtype of the generated Tensor.
|
||||
cuda:
|
||||
Whether to place the Tensor on a GPU.
|
||||
tensor_constructor:
|
||||
Callable which will be used instead of the default Tensor
|
||||
construction method. This allows the author to enforce properties
|
||||
of the Tensor (e.g. it can only have certain values). The dtype and
|
||||
concrete shape of the Tensor to be created will be passed, and
|
||||
concrete values of all parameters will be passed as kwargs. Note
|
||||
that transformations to the result (permuting, slicing) will be
|
||||
performed by the Fuzzer; the tensor_constructor is only responsible
|
||||
for creating an appropriately sized Tensor.
|
||||
"""
|
||||
self._name = name
|
||||
self._size = size
|
||||
self._steps = steps
|
||||
self._probability_contiguous = probability_contiguous
|
||||
self._min_elements = min_elements
|
||||
self._max_elements = max_elements
|
||||
self._max_allocation_bytes = max_allocation_bytes
|
||||
self._dim_parameter = dim_parameter
|
||||
self._dtype = dtype
|
||||
self._cuda = cuda
|
||||
self._tensor_constructor = tensor_constructor
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
@staticmethod
|
||||
def default_tensor_constructor(size, dtype, **kwargs):
|
||||
if dtype.is_floating_point or dtype.is_complex:
|
||||
return torch.rand(size=size, dtype=dtype, device="cpu")
|
||||
else:
|
||||
return torch.randint(1, 127, size=size, dtype=dtype, device="cpu")
|
||||
|
||||
def _make_tensor(self, params, state):
|
||||
import numpy as np
|
||||
size, steps, allocation_size = self._get_size_and_steps(params)
|
||||
constructor = (
|
||||
self._tensor_constructor or
|
||||
self.default_tensor_constructor
|
||||
)
|
||||
|
||||
raw_tensor = constructor(size=allocation_size, dtype=self._dtype, **params)
|
||||
if self._cuda:
|
||||
raw_tensor = raw_tensor.cuda()
|
||||
|
||||
# Randomly permute the Tensor and call `.contiguous()` to force re-ordering
|
||||
# of the memory, and then permute it back to the original shape.
|
||||
dim = len(size)
|
||||
order = np.arange(dim)
|
||||
if state.rand() > self._probability_contiguous:
|
||||
while dim > 1 and np.all(order == np.arange(dim)):
|
||||
order = state.permutation(raw_tensor.dim())
|
||||
|
||||
raw_tensor = raw_tensor.permute(tuple(order)).contiguous()
|
||||
raw_tensor = raw_tensor.permute(tuple(np.argsort(order)))
|
||||
|
||||
slices = [slice(0, size * step, step) for size, step in zip(size, steps, strict=True)]
|
||||
tensor = raw_tensor[tuple(slices)]
|
||||
|
||||
properties = {
|
||||
"numel": int(tensor.numel()),
|
||||
"order": order,
|
||||
"steps": steps,
|
||||
"is_contiguous": tensor.is_contiguous(),
|
||||
"dtype": str(self._dtype),
|
||||
}
|
||||
|
||||
return tensor, properties
|
||||
|
||||
def _get_size_and_steps(self, params):
|
||||
dim = (
|
||||
params[self._dim_parameter]
|
||||
if self._dim_parameter is not None
|
||||
else len(self._size)
|
||||
)
|
||||
|
||||
def resolve(values, dim):
|
||||
"""Resolve values into concrete integers."""
|
||||
values = tuple(params.get(i, i) for i in values)
|
||||
if len(values) > dim:
|
||||
values = values[:dim]
|
||||
if len(values) < dim:
|
||||
values = values + tuple(1 for _ in range(dim - len(values)))
|
||||
return values
|
||||
|
||||
size = resolve(self._size, dim)
|
||||
steps = resolve(self._steps or (), dim)
|
||||
allocation_size = tuple(size_i * step_i for size_i, step_i in zip(size, steps, strict=True))
|
||||
return size, steps, allocation_size
|
||||
|
||||
def satisfies_constraints(self, params) -> bool:
|
||||
size, _, allocation_size = self._get_size_and_steps(params)
|
||||
# Product is computed in Python to avoid integer overflow.
|
||||
num_elements = prod(size)
|
||||
if num_elements < 0:
|
||||
raise AssertionError("Computed number of elements is negative")
|
||||
|
||||
allocation_bytes = prod(allocation_size, base=dtype_size(self._dtype))
|
||||
|
||||
def nullable_greater(left, right):
|
||||
if left is None or right is None:
|
||||
return False
|
||||
return left > right
|
||||
|
||||
return not any((
|
||||
nullable_greater(num_elements, self._max_elements),
|
||||
nullable_greater(self._min_elements, num_elements),
|
||||
nullable_greater(allocation_bytes, self._max_allocation_bytes),
|
||||
))
|
||||
|
||||
|
||||
class Fuzzer:
|
||||
def __init__(
|
||||
self,
|
||||
parameters: list[FuzzedParameter | list[FuzzedParameter]],
|
||||
tensors: list[FuzzedTensor | list[FuzzedTensor]],
|
||||
constraints: list[Callable] | None = None,
|
||||
seed: int | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
parameters:
|
||||
List of FuzzedParameters which provide specifications
|
||||
for generated parameters. Iterable elements will be
|
||||
unpacked, though arbitrary nested structures will not.
|
||||
tensors:
|
||||
List of FuzzedTensors which define the Tensors which
|
||||
will be created each step based on the parameters for
|
||||
that step. Iterable elements will be unpacked, though
|
||||
arbitrary nested structures will not.
|
||||
constraints:
|
||||
List of callables. They will be called with params
|
||||
as kwargs, and if any of them return False the current
|
||||
set of parameters will be rejected.
|
||||
seed:
|
||||
Seed for the RandomState used by the Fuzzer. This will
|
||||
also be used to set the PyTorch random seed so that random
|
||||
ops will create reproducible Tensors.
|
||||
"""
|
||||
import numpy as np
|
||||
if seed is None:
|
||||
seed = int(np.random.RandomState().randint(0, 2 ** 32 - 1, dtype=np.int64))
|
||||
self._seed = seed
|
||||
self._parameters = Fuzzer._unpack(parameters, FuzzedParameter)
|
||||
self._tensors = Fuzzer._unpack(tensors, FuzzedTensor)
|
||||
self._constraints = constraints or ()
|
||||
|
||||
p_names = {p.name for p in self._parameters}
|
||||
t_names = {t.name for t in self._tensors}
|
||||
name_overlap = p_names.intersection(t_names)
|
||||
if name_overlap:
|
||||
raise ValueError(f"Duplicate names in parameters and tensors: {name_overlap}")
|
||||
|
||||
self._rejections = 0
|
||||
self._total_generated = 0
|
||||
|
||||
@staticmethod
|
||||
def _unpack(values, cls):
|
||||
return tuple(it.chain.from_iterable(
|
||||
[[i] if isinstance(i, cls) else i for i in values]
|
||||
))
|
||||
|
||||
def take(self, n):
|
||||
import numpy as np
|
||||
state = np.random.RandomState(self._seed)
|
||||
torch.manual_seed(state.randint(low=0, high=2 ** 63, dtype=np.int64))
|
||||
for _ in range(n):
|
||||
params = self._generate(state)
|
||||
tensors = {}
|
||||
tensor_properties = {}
|
||||
for t in self._tensors:
|
||||
tensor, properties = t._make_tensor(params, state)
|
||||
tensors[t.name] = tensor
|
||||
tensor_properties[t.name] = properties
|
||||
yield tensors, tensor_properties, params
|
||||
|
||||
@property
|
||||
def rejection_rate(self):
|
||||
if not self._total_generated:
|
||||
return 0.
|
||||
return self._rejections / self._total_generated
|
||||
|
||||
def _generate(self, state):
|
||||
strict_params: dict[str, float | int | ParameterAlias] = {}
|
||||
for _ in range(1000):
|
||||
candidate_params: dict[str, float | int | ParameterAlias] = {}
|
||||
for p in self._parameters:
|
||||
if p.strict:
|
||||
if p.name in strict_params:
|
||||
candidate_params[p.name] = strict_params[p.name]
|
||||
else:
|
||||
candidate_params[p.name] = p.sample(state)
|
||||
strict_params[p.name] = candidate_params[p.name]
|
||||
else:
|
||||
candidate_params[p.name] = p.sample(state)
|
||||
|
||||
candidate_params = self._resolve_aliases(candidate_params)
|
||||
|
||||
self._total_generated += 1
|
||||
if not all(f(candidate_params) for f in self._constraints):
|
||||
self._rejections += 1
|
||||
continue
|
||||
|
||||
if not all(t.satisfies_constraints(candidate_params) for t in self._tensors):
|
||||
self._rejections += 1
|
||||
continue
|
||||
|
||||
return candidate_params
|
||||
raise ValueError("Failed to generate a set of valid parameters.")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_aliases(params):
|
||||
params = dict(params)
|
||||
alias_count = sum(isinstance(v, ParameterAlias) for v in params.values())
|
||||
|
||||
keys = list(params.keys())
|
||||
while alias_count:
|
||||
for k in keys:
|
||||
v = params[k]
|
||||
if isinstance(v, ParameterAlias):
|
||||
params[k] = params[v.alias_to]
|
||||
alias_count_new = sum(isinstance(v, ParameterAlias) for v in params.values())
|
||||
if alias_count == alias_count_new:
|
||||
raise ValueError(f"ParameterAlias cycle detected\n{params}")
|
||||
|
||||
alias_count = alias_count_new
|
||||
|
||||
return params
|
||||
+122
@@ -0,0 +1,122 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from numbers import Number
|
||||
import torch
|
||||
from torch.utils.benchmark import FuzzedTensor
|
||||
import math
|
||||
|
||||
class FuzzedSparseTensor(FuzzedTensor):
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
size: tuple[str | int, ...],
|
||||
min_elements: int | None = None,
|
||||
max_elements: int | None = None,
|
||||
dim_parameter: str | None = None,
|
||||
sparse_dim: str | None = None,
|
||||
nnz: str | None = None,
|
||||
density: str | None = None,
|
||||
coalesced: str | None = None,
|
||||
dtype=torch.float32,
|
||||
cuda=False
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
name:
|
||||
A string identifier for the generated Tensor.
|
||||
size:
|
||||
A tuple of integers or strings specifying the size of the generated
|
||||
Tensor. String values will replaced with a concrete int during the
|
||||
generation process, while ints are simply passed as literals.
|
||||
min_elements:
|
||||
The minimum number of parameters that this Tensor must have for a
|
||||
set of parameters to be valid. (Otherwise they are resampled.)
|
||||
max_elements:
|
||||
Like `min_elements`, but setting an upper bound.
|
||||
dim_parameter:
|
||||
The length of `size` will be truncated to this value.
|
||||
This allows Tensors of varying dimensions to be generated by the
|
||||
Fuzzer.
|
||||
sparse_dim:
|
||||
The number of sparse dimensions in a sparse tensor.
|
||||
density:
|
||||
This value allows tensors of varying sparsities to be generated by the Fuzzer.
|
||||
coalesced:
|
||||
The sparse tensor format permits uncoalesced sparse tensors,
|
||||
where there may be duplicate coordinates in the indices.
|
||||
dtype:
|
||||
The PyTorch dtype of the generated Tensor.
|
||||
cuda:
|
||||
Whether to place the Tensor on a GPU.
|
||||
"""
|
||||
super().__init__(name=name, size=size, min_elements=min_elements,
|
||||
max_elements=max_elements, dim_parameter=dim_parameter, dtype=dtype, cuda=cuda)
|
||||
self._density = density
|
||||
self._coalesced = coalesced
|
||||
self._sparse_dim = sparse_dim
|
||||
|
||||
@staticmethod
|
||||
def sparse_tensor_constructor(size, dtype, sparse_dim, nnz, is_coalesced):
|
||||
"""sparse_tensor_constructor creates a sparse tensor with coo format.
|
||||
|
||||
Note that when `is_coalesced` is False, the number of elements is doubled but the number of indices
|
||||
represents the same amount of number of non zeros `nnz`, i.e, this is virtually the same tensor
|
||||
with the same sparsity pattern. Moreover, most of the sparse operation will use coalesce() method
|
||||
and what we want here is to get a sparse tensor with the same `nnz` even if this is coalesced or not.
|
||||
|
||||
In the other hand when `is_coalesced` is True the number of elements is reduced in the coalescing process
|
||||
by an unclear amount however the probability to generate duplicates indices are low for most of the cases.
|
||||
This decision was taken on purpose to maintain the construction cost as low as possible.
|
||||
"""
|
||||
if isinstance(size, Number):
|
||||
size = [size] * sparse_dim
|
||||
if all(size[d] <= 0 for d in range(sparse_dim)) and nnz != 0:
|
||||
raise AssertionError('invalid arguments')
|
||||
v_size = [nnz] + list(size[sparse_dim:])
|
||||
if dtype.is_floating_point:
|
||||
v = torch.rand(size=v_size, dtype=dtype, device="cpu")
|
||||
else:
|
||||
v = torch.randint(1, 127, size=v_size, dtype=dtype, device="cpu")
|
||||
|
||||
i = torch.rand(sparse_dim, nnz, device="cpu")
|
||||
i.mul_(torch.tensor(size[:sparse_dim]).unsqueeze(1).to(i))
|
||||
i = i.to(torch.long)
|
||||
|
||||
if not is_coalesced:
|
||||
v = torch.cat([v, torch.randn_like(v)], 0)
|
||||
i = torch.cat([i, i], 1)
|
||||
|
||||
x = torch.sparse_coo_tensor(i, v, torch.Size(size))
|
||||
if is_coalesced:
|
||||
x = x.coalesce()
|
||||
return x
|
||||
|
||||
def _make_tensor(self, params, state):
|
||||
size, _, _ = self._get_size_and_steps(params)
|
||||
density = params['density']
|
||||
nnz = math.ceil(sum(size) * density)
|
||||
if nnz > sum(size):
|
||||
raise AssertionError('nnz cannot exceed total number of elements')
|
||||
|
||||
is_coalesced = params['coalesced']
|
||||
sparse_dim = params['sparse_dim'] if self._sparse_dim else len(size)
|
||||
sparse_dim = min(sparse_dim, len(size))
|
||||
tensor = self.sparse_tensor_constructor(size, self._dtype, sparse_dim, nnz, is_coalesced)
|
||||
|
||||
if self._cuda:
|
||||
tensor = tensor.cuda()
|
||||
sparse_dim = tensor.sparse_dim()
|
||||
dense_dim = tensor.dense_dim()
|
||||
is_hybrid = len(size[sparse_dim:]) > 0
|
||||
|
||||
properties = {
|
||||
"numel": int(tensor.numel()),
|
||||
"shape": tensor.size(),
|
||||
"is_coalesced": tensor.is_coalesced(),
|
||||
"density": density,
|
||||
"sparsity": 1.0 - density,
|
||||
"sparse_dim": sparse_dim,
|
||||
"dense_dim": dense_dim,
|
||||
"is_hybrid": is_hybrid,
|
||||
"dtype": str(self._dtype),
|
||||
}
|
||||
return tensor, properties
|
||||
+43
@@ -0,0 +1,43 @@
|
||||
/* C++ template for Timer.timeit
|
||||
|
||||
This template will be consumed by `cpp_jit.py`, and will replace:
|
||||
`GLOBAL_SETUP_TEMPLATE_LOCATION`,
|
||||
`SETUP_TEMPLATE_LOCATION`
|
||||
and
|
||||
`STMT_TEMPLATE_LOCATION`
|
||||
sections with user provided statements.
|
||||
*/
|
||||
#include <chrono>
|
||||
|
||||
#include <c10/util/irange.h>
|
||||
#include <torch/csrc/utils/pybind.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
// Global setup. (e.g. #includes)
|
||||
// GLOBAL_SETUP_TEMPLATE_LOCATION
|
||||
|
||||
double timeit(int n) {
|
||||
pybind11::gil_scoped_release no_gil;
|
||||
|
||||
// Setup
|
||||
// SETUP_TEMPLATE_LOCATION
|
||||
|
||||
{
|
||||
// Warmup
|
||||
// STMT_TEMPLATE_LOCATION
|
||||
}
|
||||
|
||||
// Main loop
|
||||
auto start_time = std::chrono::high_resolution_clock::now();
|
||||
for (const auto loop_idx : c10::irange(n)) {
|
||||
(void)loop_idx;
|
||||
// STMT_TEMPLATE_LOCATION
|
||||
}
|
||||
auto end_time = std::chrono::high_resolution_clock::now();
|
||||
return std::chrono::duration<double>(end_time - start_time).count();
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("timeit", &timeit);
|
||||
}
|
||||
@@ -0,0 +1,533 @@
|
||||
"""Timer class based on the timeit.Timer class, but torch aware."""
|
||||
import enum
|
||||
import timeit
|
||||
import textwrap
|
||||
from typing import overload, Any, NoReturn
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch.utils.benchmark.utils import common, cpp_jit
|
||||
from torch.utils.benchmark.utils._stubs import TimerClass, TimeitModuleType
|
||||
from torch.utils.benchmark.utils.valgrind_wrapper import timer_interface as valgrind_timer_interface
|
||||
|
||||
|
||||
__all__ = ["Timer", "timer", "Language"]
|
||||
|
||||
|
||||
if torch.accelerator.is_available():
|
||||
def timer() -> float:
|
||||
torch.accelerator.synchronize()
|
||||
return timeit.default_timer()
|
||||
else:
|
||||
timer = timeit.default_timer
|
||||
|
||||
|
||||
class Language(enum.Enum):
|
||||
PYTHON = 0
|
||||
CPP = 1
|
||||
|
||||
|
||||
class CPPTimer:
|
||||
def __init__(
|
||||
self,
|
||||
stmt: str,
|
||||
setup: str,
|
||||
global_setup: str,
|
||||
timer: Callable[[], float],
|
||||
globals: dict[str, Any],
|
||||
) -> None:
|
||||
if timer is not timeit.default_timer:
|
||||
raise NotImplementedError(
|
||||
"PyTorch was built with accelerators and an accelerator is present; however "
|
||||
"Timer does not yet support accelerator measurements. If your "
|
||||
"code is CPU only, pass `timer=timeit.default_timer` to the "
|
||||
"Timer's constructor to indicate this. (Note that this will "
|
||||
"produce incorrect results if an accelerator is in fact used, as "
|
||||
"Timer will not synchronize the accelerator.)"
|
||||
)
|
||||
|
||||
if globals:
|
||||
raise ValueError("C++ timing does not support globals.")
|
||||
|
||||
self._stmt: str = textwrap.dedent(stmt)
|
||||
self._setup: str = textwrap.dedent(setup)
|
||||
self._global_setup: str = textwrap.dedent(global_setup)
|
||||
self._timeit_module: TimeitModuleType | None = None
|
||||
|
||||
def timeit(self, number: int) -> float:
|
||||
if self._timeit_module is None:
|
||||
self._timeit_module = cpp_jit.compile_timeit_template(
|
||||
stmt=self._stmt,
|
||||
setup=self._setup,
|
||||
global_setup=self._global_setup,
|
||||
)
|
||||
|
||||
return self._timeit_module.timeit(number)
|
||||
|
||||
|
||||
class Timer:
|
||||
"""Helper class for measuring execution time of PyTorch statements.
|
||||
|
||||
For a full tutorial on how to use this class, see:
|
||||
https://pytorch.org/tutorials/recipes/recipes/benchmark.html
|
||||
|
||||
The PyTorch Timer is based on `timeit.Timer` (and in fact uses
|
||||
`timeit.Timer` internally), but with several key differences:
|
||||
|
||||
1) Runtime aware:
|
||||
Timer will perform warmups (important as some elements of PyTorch are
|
||||
lazily initialized), set threadpool size so that comparisons are
|
||||
apples-to-apples, and synchronize asynchronous accelerator functions when
|
||||
necessary.
|
||||
|
||||
2) Focus on replicates:
|
||||
When measuring code, and particularly complex kernels / models,
|
||||
run-to-run variation is a significant confounding factor. It is
|
||||
expected that all measurements should include replicates to quantify
|
||||
noise and allow median computation, which is more robust than mean.
|
||||
To that effect, this class deviates from the `timeit` API by
|
||||
conceptually merging `timeit.Timer.repeat` and `timeit.Timer.autorange`.
|
||||
(Exact algorithms are discussed in method docstrings.) The `timeit`
|
||||
method is replicated for cases where an adaptive strategy is not
|
||||
desired.
|
||||
|
||||
3) Optional metadata:
|
||||
When defining a Timer, one can optionally specify `label`, `sub_label`,
|
||||
`description`, and `env`. (Defined later) These fields are included in
|
||||
the representation of result object and by the `Compare` class to group
|
||||
and display results for comparison.
|
||||
|
||||
4) Instruction counts
|
||||
In addition to wall times, Timer can run a statement under Callgrind
|
||||
and report instructions executed.
|
||||
|
||||
Directly analogous to `timeit.Timer` constructor arguments:
|
||||
|
||||
`stmt`, `setup`, `timer`, `globals`
|
||||
|
||||
PyTorch Timer specific constructor arguments:
|
||||
|
||||
`label`, `sub_label`, `description`, `env`, `num_threads`
|
||||
|
||||
Args:
|
||||
stmt: Code snippet to be run in a loop and timed.
|
||||
|
||||
setup: Optional setup code. Used to define variables used in `stmt`
|
||||
|
||||
global_setup: (C++ only)
|
||||
Code which is placed at the top level of the file for things like
|
||||
`#include` statements.
|
||||
|
||||
timer:
|
||||
Callable which returns the current time. If PyTorch was built
|
||||
without accelerators or there is no accelerator present, this defaults to
|
||||
`timeit.default_timer`; otherwise it will synchronize accelerators before
|
||||
measuring the time.
|
||||
|
||||
globals:
|
||||
A dict which defines the global variables when `stmt` is being
|
||||
executed. This is the other method for providing variables which
|
||||
`stmt` needs.
|
||||
|
||||
label:
|
||||
String which summarizes `stmt`. For instance, if `stmt` is
|
||||
"torch.nn.functional.relu(torch.add(x, 1, out=out))"
|
||||
one might set label to "ReLU(x + 1)" to improve readability.
|
||||
|
||||
sub_label:
|
||||
Provide supplemental information to disambiguate measurements
|
||||
with identical stmt or label. For instance, in our example
|
||||
above sub_label might be "float" or "int", so that it is easy
|
||||
to differentiate:
|
||||
"ReLU(x + 1): (float)"
|
||||
|
||||
"ReLU(x + 1): (int)"
|
||||
when printing Measurements or summarizing using `Compare`.
|
||||
|
||||
description:
|
||||
String to distinguish measurements with identical label and
|
||||
sub_label. The principal use of `description` is to signal to
|
||||
`Compare` the columns of data. For instance one might set it
|
||||
based on the input size to create a table of the form: ::
|
||||
|
||||
| n=1 | n=4 | ...
|
||||
------------- ...
|
||||
ReLU(x + 1): (float) | ... | ... | ...
|
||||
ReLU(x + 1): (int) | ... | ... | ...
|
||||
|
||||
|
||||
using `Compare`. It is also included when printing a Measurement.
|
||||
|
||||
env:
|
||||
This tag indicates that otherwise identical tasks were run in
|
||||
different environments, and are therefore not equivalent, for
|
||||
instance when A/B testing a change to a kernel. `Compare` will
|
||||
treat Measurements with different `env` specification as distinct
|
||||
when merging replicate runs.
|
||||
|
||||
num_threads:
|
||||
The size of the PyTorch threadpool when executing `stmt`. Single
|
||||
threaded performance is important as both a key inference workload
|
||||
and a good indicator of intrinsic algorithmic efficiency, so the
|
||||
default is set to one. This is in contrast to the default PyTorch
|
||||
threadpool size which tries to utilize all cores.
|
||||
"""
|
||||
|
||||
_timer_cls: type[TimerClass] = timeit.Timer
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stmt: str = "pass",
|
||||
setup: str = "pass",
|
||||
global_setup: str = "",
|
||||
timer: Callable[[], float] = timer,
|
||||
globals: dict[str, Any] | None = None,
|
||||
label: str | None = None,
|
||||
sub_label: str | None = None,
|
||||
description: str | None = None,
|
||||
env: str | None = None,
|
||||
num_threads: int = 1,
|
||||
language: Language | str = Language.PYTHON,
|
||||
) -> None:
|
||||
if not isinstance(stmt, str):
|
||||
raise ValueError("Currently only a `str` stmt is supported.")
|
||||
|
||||
# We copy `globals` to prevent mutations from leaking.
|
||||
# (For instance, `eval` adds the `__builtins__` key)
|
||||
self._globals = dict(globals or {})
|
||||
|
||||
timer_kwargs = {}
|
||||
if language in (Language.PYTHON, "py", "python"):
|
||||
# Include `torch` if not specified as a convenience feature.
|
||||
self._globals.setdefault("torch", torch)
|
||||
self._language: Language = Language.PYTHON
|
||||
if global_setup:
|
||||
raise ValueError(
|
||||
f"global_setup is C++ only, got `{global_setup}`. Most "
|
||||
"likely this code can simply be moved to `setup`."
|
||||
)
|
||||
|
||||
elif language in (Language.CPP, "cpp", "c++"):
|
||||
if self._timer_cls is not timeit.Timer:
|
||||
raise AssertionError("_timer_cls has already been swapped.")
|
||||
self._timer_cls = CPPTimer
|
||||
setup = ("" if setup == "pass" else setup)
|
||||
self._language = Language.CPP
|
||||
timer_kwargs["global_setup"] = global_setup
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid language `{language}`.")
|
||||
|
||||
# Convenience adjustment so that multi-line code snippets defined in
|
||||
# functions do not IndentationError (Python) or look odd (C++). The
|
||||
# leading newline removal is for the initial newline that appears when
|
||||
# defining block strings. For instance:
|
||||
# textwrap.dedent("""
|
||||
# print("This is a stmt")
|
||||
# """)
|
||||
# produces '\nprint("This is a stmt")\n'.
|
||||
#
|
||||
# Stripping this down to 'print("This is a stmt")' doesn't change
|
||||
# what gets executed, but it makes __repr__'s nicer.
|
||||
stmt = textwrap.dedent(stmt)
|
||||
stmt = (stmt[1:] if stmt and stmt[0] == "\n" else stmt).rstrip()
|
||||
setup = textwrap.dedent(setup)
|
||||
setup = (setup[1:] if setup and setup[0] == "\n" else setup).rstrip()
|
||||
|
||||
|
||||
self._timer = self._timer_cls(
|
||||
stmt=stmt,
|
||||
setup=setup,
|
||||
timer=timer,
|
||||
globals=valgrind_timer_interface.CopyIfCallgrind.unwrap_all(self._globals),
|
||||
**timer_kwargs,
|
||||
)
|
||||
self._task_spec = common.TaskSpec(
|
||||
stmt=stmt,
|
||||
setup=setup,
|
||||
global_setup=global_setup,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description=description,
|
||||
env=env,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
|
||||
def _timeit(self, number: int) -> float:
|
||||
# Even calling a timer in C++ takes ~50 ns, so no real operation should
|
||||
# take less than 1 ns. (And this prevents divide by zero errors.)
|
||||
return max(self._timer.timeit(number), 1e-9)
|
||||
|
||||
def timeit(self, number: int = 1000000) -> common.Measurement:
|
||||
"""Mirrors the semantics of timeit.Timer.timeit().
|
||||
|
||||
Execute the main statement (`stmt`) `number` times.
|
||||
https://docs.python.org/3/library/timeit.html#timeit.Timer.timeit
|
||||
"""
|
||||
with common.set_torch_threads(self._task_spec.num_threads):
|
||||
# Warmup
|
||||
self._timeit(number=max(int(number // 100), 2))
|
||||
|
||||
return common.Measurement(
|
||||
number_per_run=number,
|
||||
raw_times=[self._timeit(number=number)],
|
||||
task_spec=self._task_spec
|
||||
)
|
||||
|
||||
def repeat(self, repeat: int = -1, number: int = -1) -> None:
|
||||
raise NotImplementedError("See `Timer.blocked_autorange.`")
|
||||
|
||||
def autorange(self, callback: Callable[[int, float], NoReturn] | None = None) -> None:
|
||||
raise NotImplementedError("See `Timer.blocked_autorange.`")
|
||||
|
||||
def _threaded_measurement_loop(
|
||||
self,
|
||||
number: int,
|
||||
time_hook: Callable[[], float],
|
||||
stop_hook: Callable[[list[float]], bool],
|
||||
min_run_time: float,
|
||||
max_run_time: float | None = None,
|
||||
callback: Callable[[int, float], NoReturn] | None = None
|
||||
) -> list[float]:
|
||||
total_time = 0.0
|
||||
can_stop = False
|
||||
times: list[float] = []
|
||||
with common.set_torch_threads(self._task_spec.num_threads):
|
||||
while (total_time < min_run_time) or (not can_stop):
|
||||
time_spent = time_hook()
|
||||
times.append(time_spent)
|
||||
total_time += time_spent
|
||||
if callback:
|
||||
callback(number, time_spent)
|
||||
can_stop = stop_hook(times)
|
||||
if max_run_time and total_time > max_run_time:
|
||||
break
|
||||
return times
|
||||
|
||||
def _estimate_block_size(self, min_run_time: float) -> int:
|
||||
with common.set_torch_threads(self._task_spec.num_threads):
|
||||
# Estimate the block size needed for measurement to be negligible
|
||||
# compared to the inner loop. This also serves as a warmup.
|
||||
overhead = torch.tensor([self._timeit(0) for _ in range(5)]).median().item()
|
||||
number = 1
|
||||
while True:
|
||||
time_taken = self._timeit(number)
|
||||
relative_overhead = overhead / time_taken
|
||||
if relative_overhead <= 1e-4 and time_taken >= min_run_time / 1000:
|
||||
break
|
||||
if time_taken > min_run_time:
|
||||
break
|
||||
# Avoid overflow in C++ pybind11 interface
|
||||
if number * 10 > 2147483647:
|
||||
break
|
||||
number *= 10
|
||||
return number
|
||||
|
||||
def blocked_autorange(
|
||||
self,
|
||||
callback: Callable[[int, float], NoReturn] | None = None,
|
||||
min_run_time: float = 0.2,
|
||||
) -> common.Measurement:
|
||||
"""Measure many replicates while keeping timer overhead to a minimum.
|
||||
|
||||
At a high level, blocked_autorange executes the following pseudo-code::
|
||||
|
||||
`setup`
|
||||
|
||||
total_time = 0
|
||||
while total_time < min_run_time
|
||||
start = timer()
|
||||
for _ in range(block_size):
|
||||
`stmt`
|
||||
total_time += (timer() - start)
|
||||
|
||||
Note the variable `block_size` in the inner loop. The choice of block
|
||||
size is important to measurement quality, and must balance two
|
||||
competing objectives:
|
||||
|
||||
1) A small block size results in more replicates and generally
|
||||
better statistics.
|
||||
|
||||
2) A large block size better amortizes the cost of `timer`
|
||||
invocation, and results in a less biased measurement. This is
|
||||
important because accelerator synchronization time is non-trivial
|
||||
(order single to low double digit microseconds) and would
|
||||
otherwise bias the measurement.
|
||||
|
||||
blocked_autorange sets block_size by running a warmup period,
|
||||
increasing block size until timer overhead is less than 0.1% of
|
||||
the overall computation. This value is then used for the main
|
||||
measurement loop.
|
||||
|
||||
Returns:
|
||||
A `Measurement` object that contains measured runtimes and
|
||||
repetition counts, and can be used to compute statistics.
|
||||
(mean, median, etc.)
|
||||
"""
|
||||
number = self._estimate_block_size(min_run_time)
|
||||
|
||||
def time_hook() -> float:
|
||||
return self._timeit(number)
|
||||
|
||||
def stop_hook(times: list[float]) -> bool:
|
||||
return True
|
||||
|
||||
times = self._threaded_measurement_loop(
|
||||
number, time_hook, stop_hook,
|
||||
min_run_time=min_run_time,
|
||||
callback=callback)
|
||||
|
||||
return common.Measurement(
|
||||
number_per_run=number,
|
||||
raw_times=times,
|
||||
task_spec=self._task_spec
|
||||
)
|
||||
|
||||
def adaptive_autorange(
|
||||
self,
|
||||
threshold: float = 0.1,
|
||||
*,
|
||||
min_run_time: float = 0.01,
|
||||
max_run_time: float = 10.0,
|
||||
callback: Callable[[int, float], NoReturn] | None = None,
|
||||
) -> common.Measurement:
|
||||
"""Similar to `blocked_autorange` but also checks for variablility in measurements
|
||||
and repeats until iqr/median is smaller than `threshold` or `max_run_time` is reached.
|
||||
|
||||
|
||||
At a high level, adaptive_autorange executes the following pseudo-code::
|
||||
|
||||
`setup`
|
||||
|
||||
times = []
|
||||
while times.sum < max_run_time
|
||||
start = timer()
|
||||
for _ in range(block_size):
|
||||
`stmt`
|
||||
times.append(timer() - start)
|
||||
|
||||
enough_data = len(times)>3 and times.sum > min_run_time
|
||||
small_iqr=times.iqr/times.mean<threshold
|
||||
|
||||
if enough_data and small_iqr:
|
||||
break
|
||||
|
||||
Args:
|
||||
threshold: value of iqr/median threshold for stopping
|
||||
|
||||
min_run_time: total runtime needed before checking `threshold`
|
||||
|
||||
max_run_time: total runtime for all measurements regardless of `threshold`
|
||||
|
||||
Returns:
|
||||
A `Measurement` object that contains measured runtimes and
|
||||
repetition counts, and can be used to compute statistics.
|
||||
(mean, median, etc.)
|
||||
"""
|
||||
number = self._estimate_block_size(min_run_time=0.05)
|
||||
|
||||
def time_hook() -> float:
|
||||
return self._timeit(number)
|
||||
|
||||
def stop_hook(times: list[float]) -> bool:
|
||||
if len(times) > 3:
|
||||
return common.Measurement(
|
||||
number_per_run=number,
|
||||
raw_times=times,
|
||||
task_spec=self._task_spec
|
||||
).meets_confidence(threshold=threshold)
|
||||
return False
|
||||
times = self._threaded_measurement_loop(
|
||||
number, time_hook, stop_hook, min_run_time, max_run_time, callback=callback)
|
||||
|
||||
return common.Measurement(
|
||||
number_per_run=number,
|
||||
raw_times=times,
|
||||
task_spec=self._task_spec
|
||||
)
|
||||
|
||||
@overload
|
||||
def collect_callgrind(
|
||||
self,
|
||||
number: int,
|
||||
*,
|
||||
repeats: None,
|
||||
collect_baseline: bool,
|
||||
retain_out_file: bool,
|
||||
) -> valgrind_timer_interface.CallgrindStats:
|
||||
...
|
||||
|
||||
@overload
|
||||
def collect_callgrind(
|
||||
self,
|
||||
number: int,
|
||||
*,
|
||||
repeats: int,
|
||||
collect_baseline: bool,
|
||||
retain_out_file: bool,
|
||||
) -> tuple[valgrind_timer_interface.CallgrindStats, ...]:
|
||||
...
|
||||
|
||||
def collect_callgrind(
|
||||
self,
|
||||
number: int = 100,
|
||||
*,
|
||||
repeats: int | None = None,
|
||||
collect_baseline: bool = True,
|
||||
retain_out_file: bool = False,
|
||||
) -> Any:
|
||||
"""Collect instruction counts using Callgrind.
|
||||
|
||||
Unlike wall times, instruction counts are deterministic
|
||||
(modulo non-determinism in the program itself and small amounts of
|
||||
jitter from the Python interpreter.) This makes them ideal for detailed
|
||||
performance analysis. This method runs `stmt` in a separate process
|
||||
so that Valgrind can instrument the program. Performance is severely
|
||||
degraded due to the instrumentation, however this is ameliorated by
|
||||
the fact that a small number of iterations is generally sufficient to
|
||||
obtain good measurements.
|
||||
|
||||
In order to use this method `valgrind`, `callgrind_control`, and
|
||||
`callgrind_annotate` must be installed.
|
||||
|
||||
Because there is a process boundary between the caller (this process)
|
||||
and the `stmt` execution, `globals` cannot contain arbitrary in-memory
|
||||
data structures. (Unlike timing methods) Instead, globals are
|
||||
restricted to builtins, `nn.Modules`'s, and TorchScripted functions/modules
|
||||
to reduce the surprise factor from serialization and subsequent
|
||||
deserialization. The `GlobalsBridge` class provides more detail on this
|
||||
subject. Take particular care with nn.Modules: they rely on pickle and
|
||||
you may need to add an import to `setup` for them to transfer properly.
|
||||
|
||||
By default, a profile for an empty statement will be collected and
|
||||
cached to indicate how many instructions are from the Python loop which
|
||||
drives `stmt`.
|
||||
|
||||
Returns:
|
||||
A `CallgrindStats` object which provides instruction counts and
|
||||
some basic facilities for analyzing and manipulating results.
|
||||
"""
|
||||
if not isinstance(self._task_spec.stmt, str):
|
||||
raise ValueError("`collect_callgrind` currently only supports string `stmt`")
|
||||
|
||||
if repeats is not None and repeats < 1:
|
||||
raise ValueError("If specified, `repeats` must be >= 1")
|
||||
|
||||
# Check that the statement is valid. It doesn't guarantee success, but it's much
|
||||
# simpler and quicker to raise an exception for a faulty `stmt` or `setup` in
|
||||
# the parent process rather than the valgrind subprocess.
|
||||
self._timeit(1)
|
||||
is_python = (self._language == Language.PYTHON)
|
||||
if not is_python and self._globals:
|
||||
raise AssertionError("_timer globals are only supported for Python timers")
|
||||
result = valgrind_timer_interface.wrapper_singleton().collect_callgrind(
|
||||
task_spec=self._task_spec,
|
||||
globals=self._globals,
|
||||
number=number,
|
||||
repeats=repeats or 1,
|
||||
collect_baseline=collect_baseline and is_python,
|
||||
is_python=is_python,
|
||||
retain_out_file=retain_out_file,
|
||||
)
|
||||
|
||||
return (result[0] if repeats is None else result)
|
||||
+129
@@ -0,0 +1,129 @@
|
||||
|
||||
/*
|
||||
----------------------------------------------------------------
|
||||
|
||||
Notice that the following BSD-style license applies to this one
|
||||
file (callgrind.h) only. The rest of Valgrind is licensed under the
|
||||
terms of the GNU General Public License, version 2, unless
|
||||
otherwise indicated. See the COPYING file in the source
|
||||
distribution for details.
|
||||
|
||||
----------------------------------------------------------------
|
||||
|
||||
This file is part of callgrind, a valgrind tool for cache simulation
|
||||
and call tree tracing.
|
||||
|
||||
Copyright (C) 2003-2017 Josef Weidendorfer. All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions
|
||||
are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
2. The origin of this software must not be misrepresented; you must
|
||||
not claim that you wrote the original software. If you use this
|
||||
software in a product, an acknowledgment in the product
|
||||
documentation would be appreciated but is not required.
|
||||
|
||||
3. Altered source versions must be plainly marked as such, and must
|
||||
not be misrepresented as being the original software.
|
||||
|
||||
4. The name of the author may not be used to endorse or promote
|
||||
products derived from this software without specific prior written
|
||||
permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS
|
||||
OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
||||
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
|
||||
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
|
||||
GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
|
||||
WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
|
||||
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
|
||||
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
----------------------------------------------------------------
|
||||
|
||||
Notice that the above BSD-style license applies to this one file
|
||||
(callgrind.h) only. The entire rest of Valgrind is licensed under
|
||||
the terms of the GNU General Public License, version 2. See the
|
||||
COPYING file in the source distribution for details.
|
||||
|
||||
----------------------------------------------------------------
|
||||
*/
|
||||
|
||||
#ifndef __CALLGRIND_H
|
||||
#define __CALLGRIND_H
|
||||
|
||||
#include "valgrind.h"
|
||||
|
||||
/* !! ABIWARNING !! ABIWARNING !! ABIWARNING !! ABIWARNING !!
|
||||
This enum comprises an ABI exported by Valgrind to programs
|
||||
which use client requests. DO NOT CHANGE THE ORDER OF THESE
|
||||
ENTRIES, NOR DELETE ANY -- add new ones at the end.
|
||||
|
||||
The identification ('C','T') for Callgrind has historical
|
||||
reasons: it was called "Calltree" before. Besides, ('C','G') would
|
||||
clash with cachegrind.
|
||||
*/
|
||||
|
||||
typedef
|
||||
enum {
|
||||
VG_USERREQ__DUMP_STATS = VG_USERREQ_TOOL_BASE('C','T'),
|
||||
VG_USERREQ__ZERO_STATS,
|
||||
VG_USERREQ__TOGGLE_COLLECT,
|
||||
VG_USERREQ__DUMP_STATS_AT,
|
||||
VG_USERREQ__START_INSTRUMENTATION,
|
||||
VG_USERREQ__STOP_INSTRUMENTATION
|
||||
} Vg_CallgrindClientRequest;
|
||||
|
||||
/* Dump current state of cost centers, and zero them afterwards */
|
||||
#define CALLGRIND_DUMP_STATS \
|
||||
VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__DUMP_STATS, \
|
||||
0, 0, 0, 0, 0)
|
||||
|
||||
/* Dump current state of cost centers, and zero them afterwards.
|
||||
The argument is appended to a string stating the reason which triggered
|
||||
the dump. This string is written as a description field into the
|
||||
profile data dump. */
|
||||
#define CALLGRIND_DUMP_STATS_AT(pos_str) \
|
||||
VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__DUMP_STATS_AT, \
|
||||
pos_str, 0, 0, 0, 0)
|
||||
|
||||
/* Zero cost centers */
|
||||
#define CALLGRIND_ZERO_STATS \
|
||||
VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__ZERO_STATS, \
|
||||
0, 0, 0, 0, 0)
|
||||
|
||||
/* Toggles collection state.
|
||||
The collection state specifies whether the happening of events
|
||||
should be noted or if they are to be ignored. Events are noted
|
||||
by increment of counters in a cost center */
|
||||
#define CALLGRIND_TOGGLE_COLLECT \
|
||||
VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__TOGGLE_COLLECT, \
|
||||
0, 0, 0, 0, 0)
|
||||
|
||||
/* Start full callgrind instrumentation if not already switched on.
|
||||
When cache simulation is done, it will flush the simulated cache;
|
||||
this will lead to an artificial cache warmup phase afterwards with
|
||||
cache misses which would not have happened in reality. */
|
||||
#define CALLGRIND_START_INSTRUMENTATION \
|
||||
VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__START_INSTRUMENTATION, \
|
||||
0, 0, 0, 0, 0)
|
||||
|
||||
/* Stop full callgrind instrumentation if not already switched off.
|
||||
This flushes Valgrinds translation cache, and does no additional
|
||||
instrumentation afterwards, which effectivly will run at the same
|
||||
speed as the "none" tool (ie. at minimal slowdown).
|
||||
Use this to bypass Callgrind aggregation for uninteresting code parts.
|
||||
To start Callgrind in this mode to ignore the setup phase, use
|
||||
the option "--instr-atstart=no". */
|
||||
#define CALLGRIND_STOP_INSTRUMENTATION \
|
||||
VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__STOP_INSTRUMENTATION, \
|
||||
0, 0, 0, 0, 0)
|
||||
|
||||
#endif /* __CALLGRIND_H */
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
/* Used to collect profiles of old versions of PyTorch. */
|
||||
#include <callgrind.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
|
||||
bool _valgrind_supported_platform() {
|
||||
#if defined(NVALGRIND)
|
||||
return false;
|
||||
#else
|
||||
return true;
|
||||
#endif
|
||||
}
|
||||
|
||||
void _valgrind_toggle() {
|
||||
#if defined(NVALGRIND)
|
||||
TORCH_CHECK(false, "Valgrind is not supported.");
|
||||
#else
|
||||
CALLGRIND_TOGGLE_COLLECT;
|
||||
#endif
|
||||
}
|
||||
|
||||
void _valgrind_toggle_and_dump_stats() {
|
||||
#if defined(NVALGRIND)
|
||||
TORCH_CHECK(false, "Valgrind is not supported.");
|
||||
#else
|
||||
// NB: See note in Module.cpp
|
||||
CALLGRIND_TOGGLE_COLLECT;
|
||||
CALLGRIND_DUMP_STATS;
|
||||
#endif
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(callgrind_bindings, m) {
|
||||
m.def("_valgrind_supported_platform", &_valgrind_supported_platform);
|
||||
m.def("_valgrind_toggle", &_valgrind_toggle);
|
||||
m.def("_valgrind_toggle_and_dump_stats", &_valgrind_dump_stats);
|
||||
}
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
/* C++ template for Timer.collect_callgrind
|
||||
|
||||
This template will be consumed by `cpp_jit.py`, and will replace:
|
||||
`GLOBAL_SETUP_TEMPLATE_LOCATION`,
|
||||
`SETUP_TEMPLATE_LOCATION`
|
||||
and
|
||||
`STMT_TEMPLATE_LOCATION`
|
||||
sections with user provided statements.
|
||||
*/
|
||||
|
||||
#include <c10/util/irange.h>
|
||||
#include <callgrind.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
// Global setup. (e.g. #includes)
|
||||
// GLOBAL_SETUP_TEMPLATE_LOCATION
|
||||
|
||||
#if defined(NVALGRIND)
|
||||
static_assert(false);
|
||||
#endif
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
// This file should only be called inside of `Timer`, so we can adopt a
|
||||
// very simple and rigid argument parsing scheme.
|
||||
TORCH_CHECK(argc == 9);
|
||||
TORCH_CHECK(std::string(argv[1]) == "--number");
|
||||
auto number = std::stoi(argv[2]);
|
||||
|
||||
TORCH_CHECK(
|
||||
std::string(argv[3]) == "--number-warmup" ||
|
||||
std::string(argv[3]) == "--number_warmup");
|
||||
auto number_warmup = std::stoi(argv[4]);
|
||||
|
||||
TORCH_CHECK(std::string(argv[5]) == "--repeats");
|
||||
auto repeats = std::stoi(argv[6]);
|
||||
|
||||
TORCH_CHECK(
|
||||
std::string(argv[7]) == "--number-threads" ||
|
||||
std::string(argv[7]) == "--number_threads");
|
||||
auto number_threads = std::stoi(argv[8]);
|
||||
torch::set_num_threads(number_threads);
|
||||
|
||||
// Setup
|
||||
// SETUP_TEMPLATE_LOCATION
|
||||
|
||||
// Warmup
|
||||
for (const auto i : c10::irange(number_warmup)) {
|
||||
(void)i;
|
||||
// STMT_TEMPLATE_LOCATION
|
||||
}
|
||||
|
||||
// Main loop
|
||||
for (const auto repeat : c10::irange(repeats)) {
|
||||
(void)repeat;
|
||||
CALLGRIND_TOGGLE_COLLECT;
|
||||
|
||||
for (const auto i : c10::irange(number)) {
|
||||
(void)i;
|
||||
// STMT_TEMPLATE_LOCATION
|
||||
}
|
||||
|
||||
// NB: See note in Module.cpp
|
||||
CALLGRIND_TOGGLE_COLLECT;
|
||||
CALLGRIND_DUMP_STATS;
|
||||
}
|
||||
}
|
||||
+919
@@ -0,0 +1,919 @@
|
||||
"""Intermediate layer between `Timer` and `valgrind`."""
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import enum
|
||||
import dataclasses
|
||||
import itertools as it
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
from typing import (
|
||||
cast, Any, NamedTuple,
|
||||
TYPE_CHECKING)
|
||||
|
||||
import torch
|
||||
from torch.utils.benchmark.utils import common, cpp_jit
|
||||
import operator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterator
|
||||
|
||||
from torch.utils.benchmark.utils._stubs import CallgrindModuleType
|
||||
|
||||
|
||||
__all__ = ["FunctionCount", "FunctionCounts", "CallgrindStats", "CopyIfCallgrind"]
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
CompletedProcessType = subprocess.CompletedProcess[str]
|
||||
else:
|
||||
CompletedProcessType = subprocess.CompletedProcess
|
||||
|
||||
|
||||
class FunctionCount(NamedTuple):
|
||||
# TODO(#105471): Rename the count field
|
||||
count: int # type: ignore[assignment]
|
||||
function: str
|
||||
|
||||
|
||||
@dataclasses.dataclass(repr=False, eq=False, frozen=True)
|
||||
class FunctionCounts:
|
||||
"""Container for manipulating Callgrind results.
|
||||
|
||||
It supports:
|
||||
1) Addition and subtraction to combine or diff results.
|
||||
2) Tuple-like indexing.
|
||||
3) A `denoise` function which strips CPython calls which are known to
|
||||
be non-deterministic and quite noisy.
|
||||
4) Two higher order methods (`filter` and `transform`) for custom
|
||||
manipulation.
|
||||
"""
|
||||
_data: tuple[FunctionCount, ...]
|
||||
inclusive: bool
|
||||
truncate_rows: bool = True
|
||||
|
||||
# For normal use, torch._tensor_str.PRINT_OPTS.linewidth determines
|
||||
# the print settings. This is simply to allow hermetic unit tests.
|
||||
_linewidth: int | None = None
|
||||
|
||||
def __iter__(self) -> Iterator[FunctionCount]:
|
||||
yield from self._data
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._data)
|
||||
|
||||
def __getitem__(self, item: Any) -> FunctionCount | FunctionCounts:
|
||||
data: FunctionCount | tuple[FunctionCount, ...] = self._data[item]
|
||||
return (
|
||||
FunctionCounts(cast(tuple[FunctionCount, ...], data), self.inclusive, truncate_rows=False)
|
||||
if isinstance(data, tuple) else data
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
count_len = 0
|
||||
for c, _ in self:
|
||||
# Account for sign in string length.
|
||||
count_len = max(count_len, len(str(c)) + int(c < 0))
|
||||
|
||||
lines = []
|
||||
linewidth = self._linewidth or torch._tensor_str.PRINT_OPTS.linewidth
|
||||
fn_str_len = max(linewidth - count_len - 4, 40)
|
||||
for c, fn in self:
|
||||
if len(fn) > fn_str_len:
|
||||
left_len = int((fn_str_len - 5) // 2)
|
||||
fn = fn[:left_len] + " ... " + fn[-(fn_str_len - left_len - 5):]
|
||||
lines.append(f" {c:>{count_len}} {fn}")
|
||||
|
||||
if self.truncate_rows and len(lines) > 18:
|
||||
lines = lines[:9] + ["...".rjust(count_len + 2)] + lines[-9:]
|
||||
|
||||
if not self.inclusive:
|
||||
lines.extend(["", f"Total: {self.sum()}"])
|
||||
|
||||
return "\n".join([super().__repr__()] + lines)
|
||||
|
||||
def __add__(
|
||||
self,
|
||||
other: FunctionCounts,
|
||||
) -> FunctionCounts:
|
||||
return self._merge(other, lambda c: c)
|
||||
|
||||
def __sub__(
|
||||
self,
|
||||
other: FunctionCounts,
|
||||
) -> FunctionCounts:
|
||||
return self._merge(other, operator.neg)
|
||||
|
||||
def __mul__(self, other: int | float) -> FunctionCounts:
|
||||
return self._from_dict({
|
||||
fn: int(c * other) for c, fn in self._data
|
||||
}, self.inclusive)
|
||||
|
||||
def transform(self, map_fn: Callable[[str], str]) -> FunctionCounts:
|
||||
"""Apply `map_fn` to all of the function names.
|
||||
|
||||
This can be used to regularize function names (e.g. stripping irrelevant
|
||||
parts of the file path), coalesce entries by mapping multiple functions
|
||||
to the same name (in which case the counts are added together), etc.
|
||||
"""
|
||||
counts: collections.defaultdict[str, int] = collections.defaultdict(int)
|
||||
for c, fn in self._data:
|
||||
counts[map_fn(fn)] += c
|
||||
|
||||
return self._from_dict(counts, self.inclusive)
|
||||
|
||||
def filter(self, filter_fn: Callable[[str], bool]) -> FunctionCounts:
|
||||
"""Keep only the elements where `filter_fn` applied to function name returns True."""
|
||||
return FunctionCounts(tuple(i for i in self if filter_fn(i.function)), self.inclusive)
|
||||
|
||||
def sum(self) -> int:
|
||||
return sum(c for c, _ in self)
|
||||
|
||||
def denoise(self) -> FunctionCounts:
|
||||
"""Remove known noisy instructions.
|
||||
|
||||
Several instructions in the CPython interpreter are rather noisy. These
|
||||
instructions involve unicode to dictionary lookups which Python uses to
|
||||
map variable names. FunctionCounts is generally a content agnostic
|
||||
container, however this is sufficiently important for obtaining
|
||||
reliable results to warrant an exception."""
|
||||
return self.filter(lambda fn: "dictobject.c:lookdict_unicode" not in fn)
|
||||
|
||||
def _merge(
|
||||
self,
|
||||
second: FunctionCounts,
|
||||
merge_fn: Callable[[int], int]
|
||||
) -> FunctionCounts:
|
||||
if self.inclusive != second.inclusive:
|
||||
raise AssertionError("Cannot merge inclusive and exclusive counts.")
|
||||
counts: collections.defaultdict[str, int] = collections.defaultdict(int)
|
||||
for c, fn in self:
|
||||
counts[fn] += c
|
||||
|
||||
for c, fn in second:
|
||||
counts[fn] += merge_fn(c)
|
||||
|
||||
return self._from_dict(counts, self.inclusive)
|
||||
|
||||
@staticmethod
|
||||
def _from_dict(counts: dict[str, int], inclusive: bool) -> FunctionCounts:
|
||||
flat_counts = (FunctionCount(c, fn) for fn, c in counts.items() if c)
|
||||
return FunctionCounts(tuple(sorted(flat_counts, reverse=True)), inclusive)
|
||||
|
||||
|
||||
@dataclasses.dataclass(repr=False, eq=False, frozen=True)
|
||||
class CallgrindStats:
|
||||
"""Top level container for Callgrind results collected by Timer.
|
||||
|
||||
Manipulation is generally done using the FunctionCounts class, which is
|
||||
obtained by calling `CallgrindStats.stats(...)`. Several convenience
|
||||
methods are provided as well; the most significant is
|
||||
`CallgrindStats.as_standardized()`.
|
||||
"""
|
||||
task_spec: common.TaskSpec
|
||||
number_per_run: int
|
||||
built_with_debug_symbols: bool
|
||||
baseline_inclusive_stats: FunctionCounts
|
||||
baseline_exclusive_stats: FunctionCounts
|
||||
stmt_inclusive_stats: FunctionCounts
|
||||
stmt_exclusive_stats: FunctionCounts
|
||||
stmt_callgrind_out: str | None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
base_stats = self.baseline_exclusive_stats
|
||||
output = f"""
|
||||
{super().__repr__()}
|
||||
{self.task_spec.summarize()}
|
||||
{'':>25}All{'':>10}Noisy symbols removed
|
||||
Instructions: {self.counts(denoise=False):>12}{'':>15}{self.counts(denoise=True):>12}
|
||||
Baseline: {base_stats.sum():>12}{'':>15}{base_stats.denoise().sum():>12}
|
||||
{self.number_per_run} runs per measurement, {self.task_spec.num_threads} thread{'s' if self.task_spec.num_threads > 1 else ''}
|
||||
""".strip()
|
||||
if not self.built_with_debug_symbols:
|
||||
output += textwrap.dedent("""
|
||||
Warning: PyTorch was not built with debug symbols.
|
||||
Source information may be limited. Rebuild with
|
||||
REL_WITH_DEB_INFO=1 for more detailed results.""")
|
||||
return output
|
||||
|
||||
def stats(self, inclusive: bool = False) -> FunctionCounts:
|
||||
"""Returns detailed function counts.
|
||||
|
||||
Conceptually, the FunctionCounts returned can be thought of as a tuple
|
||||
of (count, path_and_function_name) tuples.
|
||||
|
||||
`inclusive` matches the semantics of callgrind. If True, the counts
|
||||
include instructions executed by children. `inclusive=True` is useful
|
||||
for identifying hot spots in code; `inclusive=False` is useful for
|
||||
reducing noise when diffing counts from two different runs. (See
|
||||
CallgrindStats.delta(...) for more details)
|
||||
"""
|
||||
return self.stmt_inclusive_stats if inclusive else self.stmt_exclusive_stats
|
||||
|
||||
def counts(self, *, denoise: bool = False) -> int:
|
||||
"""Returns the total number of instructions executed.
|
||||
|
||||
See `FunctionCounts.denoise()` for an explanation of the `denoise` arg.
|
||||
"""
|
||||
stats = self.stmt_exclusive_stats
|
||||
return (stats.denoise() if denoise else stats).sum()
|
||||
|
||||
# FIXME: Once 3.7 is the minimum version, type annotate `other` per PEP 563
|
||||
def delta(
|
||||
self,
|
||||
other: CallgrindStats,
|
||||
inclusive: bool = False,
|
||||
) -> FunctionCounts:
|
||||
"""Diff two sets of counts.
|
||||
|
||||
One common reason to collect instruction counts is to determine the
|
||||
the effect that a particular change will have on the number of instructions
|
||||
needed to perform some unit of work. If a change increases that number, the
|
||||
next logical question is "why". This generally involves looking at what part
|
||||
if the code increased in instruction count. This function automates that
|
||||
process so that one can easily diff counts on both an inclusive and
|
||||
exclusive basis.
|
||||
"""
|
||||
return self.stats(inclusive=inclusive) - other.stats(inclusive=inclusive)
|
||||
|
||||
def as_standardized(self) -> CallgrindStats:
|
||||
"""Strip library names and some prefixes from function strings.
|
||||
|
||||
When comparing two different sets of instruction counts, on stumbling
|
||||
block can be path prefixes. Callgrind includes the full filepath
|
||||
when reporting a function (as it should). However, this can cause
|
||||
issues when diffing profiles. If a key component such as Python
|
||||
or PyTorch was built in separate locations in the two profiles, which
|
||||
can result in something resembling::
|
||||
|
||||
23234231 /tmp/first_build_dir/thing.c:foo(...)
|
||||
9823794 /tmp/first_build_dir/thing.c:bar(...)
|
||||
...
|
||||
53453 .../aten/src/Aten/...:function_that_actually_changed(...)
|
||||
...
|
||||
-9823794 /tmp/second_build_dir/thing.c:bar(...)
|
||||
-23234231 /tmp/second_build_dir/thing.c:foo(...)
|
||||
|
||||
Stripping prefixes can ameliorate this issue by regularizing the
|
||||
strings and causing better cancellation of equivalent call sites
|
||||
when diffing.
|
||||
"""
|
||||
def strip(stats: FunctionCounts) -> FunctionCounts:
|
||||
transforms = (
|
||||
# PyTorch may have been built in different locations.
|
||||
(r"^.+build/\.\./", "build/../"),
|
||||
(r"^.+/" + re.escape("build/aten/"), "build/aten/"),
|
||||
|
||||
# "Python" and "Objects" come from CPython.
|
||||
(r"^.+/" + re.escape("Python/"), "Python/"),
|
||||
(r"^.+/" + re.escape("Objects/"), "Objects/"),
|
||||
|
||||
# Strip library name. e.g. `libtorch.so`
|
||||
(r"\s\[.+\]$", ""),
|
||||
)
|
||||
|
||||
for before, after in transforms:
|
||||
stats = stats.transform(lambda fn: re.sub(before, after, fn))
|
||||
|
||||
return stats
|
||||
|
||||
return CallgrindStats(
|
||||
task_spec=self.task_spec,
|
||||
number_per_run=self.number_per_run,
|
||||
built_with_debug_symbols=self.built_with_debug_symbols,
|
||||
baseline_inclusive_stats=strip(self.baseline_inclusive_stats),
|
||||
baseline_exclusive_stats=strip(self.baseline_exclusive_stats),
|
||||
stmt_inclusive_stats=strip(self.stmt_inclusive_stats),
|
||||
stmt_exclusive_stats=strip(self.stmt_exclusive_stats),
|
||||
|
||||
# `as_standardized` will change symbol names, so the contents will
|
||||
# no longer map directly to `callgrind.out`
|
||||
stmt_callgrind_out=None,
|
||||
)
|
||||
|
||||
|
||||
class Serialization(enum.Enum):
|
||||
PICKLE = 0
|
||||
TORCH = 1
|
||||
TORCH_JIT = 2
|
||||
|
||||
|
||||
_GLOBALS_ALLOWED_TYPES: dict[Serialization, tuple[Any, ...]] = {
|
||||
Serialization.PICKLE: (str, bytes, bool, int, float, complex),
|
||||
Serialization.TORCH_JIT: (torch.jit.ScriptFunction, torch.jit.ScriptModule),
|
||||
Serialization.TORCH: (torch.nn.Module,),
|
||||
}
|
||||
|
||||
|
||||
class CopyIfCallgrind:
|
||||
"""Signal that a global may be replaced with a deserialized copy.
|
||||
|
||||
See `GlobalsBridge` for why this matters.
|
||||
"""
|
||||
def __init__(self, value: Any, *, setup: str | None = None) -> None:
|
||||
for method, supported_types in _GLOBALS_ALLOWED_TYPES.items():
|
||||
if any(isinstance(value, t) for t in supported_types):
|
||||
self._value: Any = value
|
||||
self._setup: str | None = setup
|
||||
self._serialization: Serialization = method
|
||||
break
|
||||
else:
|
||||
supported_str = "\n".join([
|
||||
getattr(t, "__name__", repr(t))
|
||||
for t in it.chain(_GLOBALS_ALLOWED_TYPES.values())])
|
||||
|
||||
raise ValueError(
|
||||
f"Unsupported type: {type(value)}\n"
|
||||
f"`collect_callgrind` restricts globals to the following types:\n"
|
||||
f"{textwrap.indent(supported_str, ' ')}"
|
||||
)
|
||||
|
||||
@property
|
||||
def value(self) -> Any:
|
||||
return self._value
|
||||
|
||||
@property
|
||||
def setup(self) -> str | None:
|
||||
return self._setup
|
||||
|
||||
@property
|
||||
def serialization(self) -> Serialization:
|
||||
return self._serialization
|
||||
|
||||
@staticmethod
|
||||
def unwrap_all(globals: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
k: (v.value if isinstance(v, CopyIfCallgrind) else v)
|
||||
for k, v in globals.items()
|
||||
}
|
||||
|
||||
|
||||
class GlobalsBridge:
|
||||
"""Handle the transfer of (certain) globals when collecting Callgrind statistics.
|
||||
|
||||
Key takeaway: Any globals passed must be wrapped in `CopyIfCallgrind` to
|
||||
work with `Timer.collect_callgrind`.
|
||||
|
||||
Consider the following code snippet:
|
||||
```
|
||||
import pickle
|
||||
import timeit
|
||||
|
||||
class Counter:
|
||||
value = 0
|
||||
|
||||
def __call__(self):
|
||||
self.value += 1
|
||||
|
||||
counter = Counter()
|
||||
timeit.Timer("counter()", globals={"counter": counter}).timeit(10)
|
||||
print(counter.value) # 10
|
||||
|
||||
timeit.Timer(
|
||||
"counter()",
|
||||
globals={"counter": pickle.loads(pickle.dumps(counter))}
|
||||
).timeit(20)
|
||||
print(counter.value) # Still 10
|
||||
```
|
||||
|
||||
In the first case, `stmt` is executed using the objects in `globals`;
|
||||
however, the addition of serialization and deserialization changes the
|
||||
semantics and may meaningfully change behavior.
|
||||
|
||||
This is a practical consideration when collecting Callgrind statistics.
|
||||
Unlike `exec` based execution (which `timeit` uses under the hood) which
|
||||
can share in-memory data structures with the caller, Callgrind collection
|
||||
requires an entirely new process in order to run under Valgrind. This means
|
||||
that any data structures used for statement execution will have to be
|
||||
serialized and deserialized in the subprocess.
|
||||
|
||||
In order to avoid surprising semantics from (user invisible) process
|
||||
boundaries, what can be passed through `globals` is severely restricted
|
||||
for `Timer.collect_callgrind`. It is expected that most setup should be
|
||||
achievable (albeit perhaps less ergonomically) by passing a `setup`
|
||||
string.
|
||||
|
||||
There are, however, exceptions. One such class are TorchScripted functions.
|
||||
Because they require a concrete file with source code it is not possible
|
||||
to define them using a `setup` string. Another group are torch.nn.Modules,
|
||||
whose construction can be complex and prohibitively cumbersome to coerce
|
||||
into a `setup` string. Finally, most builtin types are sufficiently well
|
||||
behaved and sufficiently common to warrant allowing as well. (e.g.
|
||||
`globals={"n": 1}` is very convenient.)
|
||||
|
||||
Fortunately, all have well defined serialization semantics. This class
|
||||
is responsible for enabling the Valgrind subprocess to use elements in
|
||||
`globals` so long as they are an allowed type.
|
||||
|
||||
Caveats:
|
||||
The user is required to acknowledge this serialization by wrapping
|
||||
elements in `globals` with `CopyIfCallgrind`.
|
||||
|
||||
While ScriptFunction and ScriptModule are expected to save and load
|
||||
quite robustly, it is up to the user to ensure that an nn.Module can
|
||||
un-pickle successfully.
|
||||
|
||||
`torch.Tensor` and `np.ndarray` are deliberately excluded. The
|
||||
serialization/deserialization process perturbs the representation of a
|
||||
tensor in ways that could result in incorrect measurements. For example,
|
||||
if a tensor lives in pinned CPU memory, this fact would not be preserved
|
||||
by a dump, and that will in turn change the performance of certain CUDA
|
||||
operations.
|
||||
"""
|
||||
|
||||
def __init__(self, globals: dict[str, Any], data_dir: str) -> None:
|
||||
self._globals: dict[str, CopyIfCallgrind] = {}
|
||||
self._data_dir = data_dir
|
||||
if not os.path.exists(data_dir):
|
||||
os.mkdir(data_dir)
|
||||
|
||||
if globals.get("torch", torch) is not torch:
|
||||
raise ValueError("`collect_callgrind` does not support mocking out `torch`.")
|
||||
|
||||
for name, value in globals.items():
|
||||
if name in ("torch", "__builtins__"):
|
||||
# Torch will be imported by the collection script, and
|
||||
# __builtins__ is added by Timer.
|
||||
continue
|
||||
|
||||
if not isinstance(value, CopyIfCallgrind):
|
||||
raise ValueError(
|
||||
"`collect_callgrind` requires that globals be wrapped in "
|
||||
"`CopyIfCallgrind` so that serialization is explicit."
|
||||
)
|
||||
|
||||
self._globals[name] = value
|
||||
|
||||
def construct(self) -> str:
|
||||
load_lines = []
|
||||
for name, wrapped_value in self._globals.items():
|
||||
if wrapped_value.setup is not None:
|
||||
load_lines.append(textwrap.dedent(wrapped_value.setup))
|
||||
|
||||
if wrapped_value.serialization == Serialization.PICKLE:
|
||||
path = os.path.join(self._data_dir, f"{name}.pkl")
|
||||
load_lines.append(
|
||||
f"with open({repr(path)}, 'rb') as f:\n {name} = pickle.load(f)")
|
||||
with open(path, "wb") as f:
|
||||
pickle.dump(wrapped_value.value, f)
|
||||
|
||||
elif wrapped_value.serialization == Serialization.TORCH:
|
||||
path = os.path.join(self._data_dir, f"{name}.pt")
|
||||
# TODO: Figure out if we can use torch.serialization.add_safe_globals here
|
||||
# Using weights_only=False after the change in
|
||||
# https://dev-discuss.pytorch.org/t/bc-breaking-change-torch-load-is-being-flipped-to-use-weights-only-true-by-default-in-the-nightlies-after-137602/2573
|
||||
load_lines.append(f"{name} = torch.load({repr(path)}, weights_only=False)")
|
||||
torch.save(wrapped_value.value, path)
|
||||
|
||||
elif wrapped_value.serialization == Serialization.TORCH_JIT:
|
||||
path = os.path.join(self._data_dir, f"{name}.pt")
|
||||
load_lines.append(f"{name} = torch.jit.load({repr(path)})")
|
||||
with open(path, "wb") as f:
|
||||
torch.jit.save(wrapped_value.value, f) # type: ignore[no-untyped-call]
|
||||
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unknown serialization method: {wrapped_value.serialization}")
|
||||
|
||||
return "\n".join(load_lines)
|
||||
|
||||
|
||||
class _ValgrindWrapper:
|
||||
def __init__(self) -> None:
|
||||
self._bindings_module: CallgrindModuleType | None = None
|
||||
valgrind_symbols = (
|
||||
"_valgrind_supported_platform",
|
||||
"_valgrind_toggle",
|
||||
"_valgrind_toggle_and_dump_stats",
|
||||
)
|
||||
if all(hasattr(torch._C, symbol) for symbol in valgrind_symbols):
|
||||
self._supported_platform: bool = torch._C._valgrind_supported_platform()
|
||||
|
||||
else:
|
||||
print("Callgrind bindings are not present in `torch._C`. JIT-ing bindings.")
|
||||
self._bindings_module = cpp_jit.get_compat_bindings()
|
||||
if not all(hasattr(self._bindings_module, symbol) for symbol in valgrind_symbols):
|
||||
raise AssertionError("JIT-compiled callgrind bindings are missing required symbols")
|
||||
self._supported_platform = self._bindings_module._valgrind_supported_platform()
|
||||
|
||||
self._commands_available: dict[str, bool] = {}
|
||||
if self._supported_platform:
|
||||
# Only bother checking on supported platforms.
|
||||
for cmd in ("valgrind", "callgrind_control", "callgrind_annotate"):
|
||||
self._commands_available[cmd] = not subprocess.run(
|
||||
["which", cmd],
|
||||
capture_output=True,
|
||||
check=False,
|
||||
).returncode
|
||||
|
||||
self._build_type: str | None = None
|
||||
build_search = re.search("BUILD_TYPE=(.+),", torch.__config__.show()) # type: ignore[no-untyped-call]
|
||||
if build_search is not None:
|
||||
self._build_type = build_search.groups()[0].split(",")[0]
|
||||
|
||||
def _validate(self) -> None:
|
||||
if not self._supported_platform:
|
||||
raise OSError("Valgrind is not supported on this platform.")
|
||||
|
||||
missing_cmds = [cmd for cmd, available in self._commands_available.items() if not available]
|
||||
if missing_cmds:
|
||||
raise OSError("Missing: " + ", ".join(missing_cmds))
|
||||
|
||||
def collect_callgrind(
|
||||
self,
|
||||
task_spec: common.TaskSpec,
|
||||
globals: dict[str, Any],
|
||||
*,
|
||||
number: int,
|
||||
repeats: int,
|
||||
collect_baseline: bool,
|
||||
is_python: bool,
|
||||
retain_out_file: bool,
|
||||
) -> tuple[CallgrindStats, ...]:
|
||||
"""Collect stats, and attach a reference run which can be used to filter interpreter overhead."""
|
||||
self._validate()
|
||||
if not is_python and collect_baseline:
|
||||
raise AssertionError("collect_baseline is only supported for Python timers")
|
||||
|
||||
*task_stats, baseline_stats = self._invoke(
|
||||
task_spec=task_spec,
|
||||
globals=globals,
|
||||
number=number,
|
||||
repeats=repeats,
|
||||
collect_baseline=collect_baseline,
|
||||
is_python=is_python,
|
||||
retain_out_file=retain_out_file,
|
||||
)
|
||||
if len(task_stats) != repeats:
|
||||
raise AssertionError("Unexpected number of task stats returned from _invoke")
|
||||
|
||||
return tuple(
|
||||
CallgrindStats(
|
||||
task_spec=task_spec,
|
||||
number_per_run=number,
|
||||
built_with_debug_symbols=self._build_type == "RelWithDebInfo",
|
||||
baseline_inclusive_stats=baseline_stats[0],
|
||||
baseline_exclusive_stats=baseline_stats[1],
|
||||
stmt_inclusive_stats=stmt_inclusive_stats,
|
||||
stmt_exclusive_stats=stmt_exclusive_stats,
|
||||
stmt_callgrind_out=out_contents,
|
||||
)
|
||||
for stmt_inclusive_stats, stmt_exclusive_stats, out_contents in task_stats
|
||||
)
|
||||
|
||||
def _invoke(
|
||||
self,
|
||||
*,
|
||||
task_spec: common.TaskSpec,
|
||||
globals: dict[str, Any],
|
||||
number: int,
|
||||
repeats: int,
|
||||
collect_baseline: bool,
|
||||
is_python: bool,
|
||||
retain_out_file: bool,
|
||||
) -> tuple[tuple[FunctionCounts, FunctionCounts, str | None], ...]:
|
||||
"""Core invocation method for Callgrind collection.
|
||||
|
||||
Valgrind operates by effectively replacing the CPU with an emulated
|
||||
version which allows it to instrument any code at the cost of severe
|
||||
performance degradation. This has the practical effect that in order
|
||||
to collect Callgrind statistics, a new process has to be created
|
||||
running under `valgrind`. The steps for this process are:
|
||||
|
||||
1) Create a scratch directory.
|
||||
2) Codegen a run script. (_ValgrindWrapper._construct_script)
|
||||
Inside the run script:
|
||||
* Validate that Python and torch match the parent process
|
||||
* Validate that it is indeed running under valgrind
|
||||
* Execute `setup` and warm up `stmt`
|
||||
* Begin collecting stats
|
||||
* Run the `stmt` loop
|
||||
* Stop collecting stats
|
||||
3) Parse the run results.
|
||||
4) Cleanup the scratch directory.
|
||||
"""
|
||||
working_dir = common._make_temp_dir(prefix="callgrind")
|
||||
data_dir = os.path.join(working_dir, "data")
|
||||
script_file = os.path.join(working_dir, "timer_callgrind.py")
|
||||
callgrind_out = os.path.join(working_dir, "callgrind.out")
|
||||
error_log = os.path.join(working_dir, "error.txt")
|
||||
stat_log = os.path.join(working_dir, "callgrind_stat.txt")
|
||||
stdout_stderr_log = os.path.join(working_dir, "stdout_stderr.log")
|
||||
|
||||
def run(args: list[str], **kwargs: Any) -> tuple[CompletedProcessType, str]:
|
||||
# https://thraxil.org/users/anders/posts/2008/03/13/Subprocess-Hanging-PIPE-is-your-enemy/
|
||||
with open(stdout_stderr_log, "wb") as f_stdout_stderr:
|
||||
invocation = subprocess.run(
|
||||
args,
|
||||
stdout=f_stdout_stderr,
|
||||
stderr=subprocess.STDOUT,
|
||||
**kwargs,
|
||||
)
|
||||
with open(stdout_stderr_log) as f:
|
||||
return invocation, f.read()
|
||||
|
||||
try:
|
||||
if is_python:
|
||||
if self._bindings_module is not None:
|
||||
shutil.copy(
|
||||
self._bindings_module.__file__,
|
||||
os.path.join(working_dir, os.path.split(self._bindings_module.__file__)[1])
|
||||
)
|
||||
|
||||
script_file = os.path.join(working_dir, "timer_callgrind.py")
|
||||
with open(script_file, "w") as f:
|
||||
f.write(self._construct_script(
|
||||
task_spec,
|
||||
globals=GlobalsBridge(globals, data_dir),
|
||||
number=number,
|
||||
repeats=repeats,
|
||||
collect_baseline=collect_baseline,
|
||||
error_log=error_log,
|
||||
stat_log=stat_log,
|
||||
bindings=self._bindings_module))
|
||||
|
||||
run_loop_cmd = ["python", script_file]
|
||||
else:
|
||||
if collect_baseline:
|
||||
raise AssertionError("collect_baseline must be False for non-Python timers")
|
||||
run_loop_exec = cpp_jit.compile_callgrind_template(
|
||||
stmt=task_spec.stmt,
|
||||
setup=task_spec.setup,
|
||||
global_setup=task_spec.global_setup,
|
||||
)
|
||||
run_loop_cmd = [
|
||||
run_loop_exec,
|
||||
"--number", str(number),
|
||||
"--number-warmup", str(min(number, 10)),
|
||||
"--repeats", str(repeats),
|
||||
"--number-threads", str(task_spec.num_threads),
|
||||
]
|
||||
|
||||
valgrind_invocation, valgrind_invocation_output = run([
|
||||
"valgrind",
|
||||
"--tool=callgrind",
|
||||
f"--callgrind-out-file={callgrind_out}",
|
||||
"--dump-line=yes",
|
||||
"--dump-instr=yes",
|
||||
"--instr-atstart=yes",
|
||||
"--collect-atstart=no",
|
||||
] + run_loop_cmd)
|
||||
|
||||
if valgrind_invocation.returncode:
|
||||
error_report = ""
|
||||
if os.path.exists(error_log):
|
||||
with open(error_log) as f:
|
||||
error_report = f.read()
|
||||
if not error_report:
|
||||
error_report = "Unknown error.\n" + valgrind_invocation_output
|
||||
|
||||
raise OSError(f"Failed to collect callgrind profile:\n{error_report}")
|
||||
|
||||
def parse_output(fpath: str, inclusive: bool) -> FunctionCounts:
|
||||
_annotate_invocation, annotate_invocation_output = run([
|
||||
"callgrind_annotate",
|
||||
f"--inclusive={'yes' if inclusive else 'no'}",
|
||||
"--threshold=100",
|
||||
"--show-percs=no",
|
||||
fpath
|
||||
], check=True)
|
||||
|
||||
total_pattern = re.compile(r"^([0-9,]+)\s+PROGRAM TOTALS")
|
||||
begin_pattern = re.compile(r"Ir\s+file:function")
|
||||
function_pattern = re.compile(r"^\s*([0-9,]+)\s+(.+:.+)$")
|
||||
|
||||
class ScanState(enum.Enum):
|
||||
SCANNING_FOR_TOTAL = 0
|
||||
SCANNING_FOR_START = 1
|
||||
PARSING = 2
|
||||
|
||||
scan_state = ScanState.SCANNING_FOR_TOTAL
|
||||
fn_counts = []
|
||||
for l in annotate_invocation_output.splitlines(keepends=False):
|
||||
if scan_state == ScanState.SCANNING_FOR_TOTAL:
|
||||
total_match = total_pattern.match(l)
|
||||
if total_match:
|
||||
program_totals = int(total_match.groups()[0].replace(",", ""))
|
||||
scan_state = ScanState.SCANNING_FOR_START
|
||||
|
||||
elif scan_state == ScanState.SCANNING_FOR_START:
|
||||
if begin_pattern.match(l):
|
||||
scan_state = ScanState.PARSING
|
||||
|
||||
else:
|
||||
if scan_state != ScanState.PARSING:
|
||||
raise AssertionError("Failed to enter PARSING state while parsing callgrind_annotate output")
|
||||
fn_match = function_pattern.match(l)
|
||||
if fn_match:
|
||||
ir_str, file_function = fn_match.groups()
|
||||
ir = int(ir_str.replace(",", ""))
|
||||
if ir == program_totals: # type: ignore[possibly-undefined]
|
||||
# Callgrind includes some top level red herring symbols when
|
||||
# a program dumps multiple profiles.
|
||||
continue
|
||||
fn_counts.append(FunctionCount(ir, file_function))
|
||||
|
||||
elif re.match(r"-+", l):
|
||||
# Ignore heading separator lines.
|
||||
continue
|
||||
|
||||
else:
|
||||
break
|
||||
|
||||
if scan_state != ScanState.PARSING:
|
||||
raise AssertionError(f"Failed to parse {fpath}")
|
||||
return FunctionCounts(tuple(sorted(fn_counts, reverse=True)), inclusive=inclusive)
|
||||
|
||||
def read_results(i: int) -> tuple[FunctionCounts, FunctionCounts, str | None]:
|
||||
if i == repeats and not collect_baseline:
|
||||
# Null baseline.
|
||||
return (
|
||||
FunctionCounts((), inclusive=True),
|
||||
FunctionCounts((), inclusive=False),
|
||||
None,
|
||||
)
|
||||
|
||||
fpath = f"{callgrind_out}.{i + 1}" # Callgrind one-indexes files.
|
||||
callgrind_out_contents: str | None = None
|
||||
if retain_out_file:
|
||||
with open(fpath) as f:
|
||||
callgrind_out_contents = f.read()
|
||||
|
||||
return (
|
||||
parse_output(fpath, inclusive=True),
|
||||
parse_output(fpath, inclusive=False),
|
||||
callgrind_out_contents
|
||||
)
|
||||
|
||||
return tuple(read_results(i) for i in range(repeats + 1))
|
||||
finally:
|
||||
shutil.rmtree(working_dir)
|
||||
|
||||
@staticmethod
|
||||
def _construct_script(
|
||||
task_spec: common.TaskSpec,
|
||||
globals: GlobalsBridge,
|
||||
*,
|
||||
number: int,
|
||||
repeats: int,
|
||||
collect_baseline: bool,
|
||||
error_log: str,
|
||||
stat_log: str,
|
||||
bindings: CallgrindModuleType | None,
|
||||
) -> str:
|
||||
def block_stmt(stmt: str, indent: int = 0) -> str:
|
||||
"""Partially unroll benchmark loop.
|
||||
|
||||
The naive template looks something like:
|
||||
"for _ in range({number}): {stmt}"
|
||||
|
||||
However a loop in Python is surprisingly expensive, and significantly
|
||||
increases the number of background Python instructions. So instead we
|
||||
partially unroll the loops, with a block size of 100 chosen to keep
|
||||
the instruction overhead from `range` low while also not ballooning
|
||||
the size of the generated file.
|
||||
"""
|
||||
block_size = 100
|
||||
loop_count = number // block_size
|
||||
if loop_count == 1:
|
||||
# There is no point in having `for _ in range(1): ...` rather
|
||||
# than just `...`, and this lets us save shave a few background
|
||||
# instructions.
|
||||
loop_count = 0
|
||||
remainder = number - block_size * loop_count
|
||||
blocked_stmt = ""
|
||||
|
||||
if loop_count:
|
||||
unrolled_stmts = textwrap.indent("\n".join([stmt] * block_size), " " * 4)
|
||||
blocked_stmt += f"for _ in range({loop_count}):\n{unrolled_stmts}\n"
|
||||
|
||||
if remainder:
|
||||
blocked_stmt += "\n".join([stmt] * remainder)
|
||||
|
||||
return textwrap.indent(blocked_stmt, " " * indent)
|
||||
|
||||
pass_baseline = (
|
||||
"callgrind_bindings._valgrind_toggle()\n"
|
||||
f"{block_stmt('pass')}\n"
|
||||
"callgrind_bindings._valgrind_toggle_and_dump_stats()"
|
||||
)
|
||||
|
||||
return textwrap.dedent(r"""
|
||||
import gc
|
||||
import os
|
||||
import pickle
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
# Mitigate https://github.com/pytorch/pytorch/issues/37377
|
||||
# which can sometimes cause the subprocess call to fail.
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
torch.set_num_threads({num_threads})
|
||||
|
||||
{bindings_import}
|
||||
|
||||
PID = os.getpid()
|
||||
|
||||
def log_failure(msg):
|
||||
with open({error_log_repr}, "wt") as f:
|
||||
f.write(msg)
|
||||
sys.exit(1)
|
||||
|
||||
def check_result(completed_process):
|
||||
if completed_process.returncode:
|
||||
log_failure(f"Command failed: {{' '.join(completed_process.args)}}")
|
||||
return completed_process
|
||||
|
||||
# =============================================================================
|
||||
# == Check that subprocess matches parent =====================================
|
||||
# =============================================================================
|
||||
if os.path.realpath(sys.executable) != "{parent_interpreter}":
|
||||
log_failure(
|
||||
"Interpreter mismatch:\n"
|
||||
f" {{os.path.realpath(sys.executable)}}\n vs.\n {parent_interpreter}"
|
||||
)
|
||||
|
||||
if torch.__file__ != "{torch_file}":
|
||||
log_failure(
|
||||
"PyTorch does not match expected file:\n"
|
||||
f" {{torch.__file__}}\n vs.\n {torch_file}"
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# == User specified setup =====================================================
|
||||
# =============================================================================
|
||||
# Load serialized globals
|
||||
{load_globals}
|
||||
|
||||
# User setup str
|
||||
{setup}
|
||||
|
||||
for _ in range({warmup_number}):
|
||||
{indented_stmt}
|
||||
|
||||
# =============================================================================
|
||||
# == Callgrind management =====================================================
|
||||
# =============================================================================
|
||||
with open("{stat_log}", "wb") as stat_file:
|
||||
# If many instances of callgrind are running at once, the output of
|
||||
# `callgrind_control` may exceed 16kb which would cause `subprocess.PIPE`
|
||||
# to deadlock. So instead we use a file.
|
||||
callgrind_stat = check_result(subprocess.run(
|
||||
["callgrind_control", "--stat"],
|
||||
stdout=stat_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
))
|
||||
|
||||
with open("{stat_log}", "rt") as stat_file:
|
||||
stat_lines = stat_file.read().splitlines()
|
||||
|
||||
if f"PID {{PID}}: python {{__file__}}" not in stat_lines:
|
||||
log_failure("Process does not appear to be running callgrind.")
|
||||
|
||||
gc.collect()
|
||||
time.sleep(0.01)
|
||||
|
||||
# =============================================================================
|
||||
# == User code block ==========================================================
|
||||
# =============================================================================
|
||||
for _ in range({repeats}):
|
||||
callgrind_bindings._valgrind_toggle()
|
||||
{blocked_stmt}
|
||||
callgrind_bindings._valgrind_toggle_and_dump_stats()
|
||||
gc.collect()
|
||||
|
||||
{baseline}
|
||||
""").strip().format(
|
||||
indented_stmt=textwrap.indent(task_spec.stmt, " " * 4),
|
||||
blocked_stmt=block_stmt(task_spec.stmt, indent=4),
|
||||
baseline=(pass_baseline if collect_baseline else ""),
|
||||
number=number,
|
||||
repeats=repeats,
|
||||
load_globals=globals.construct(),
|
||||
setup=task_spec.setup,
|
||||
warmup_number=min(number, 10),
|
||||
num_threads=task_spec.num_threads,
|
||||
error_log_repr=repr(error_log),
|
||||
stat_log=stat_log,
|
||||
parent_interpreter=os.path.realpath(sys.executable),
|
||||
torch_file=torch.__file__,
|
||||
bindings_import=(
|
||||
"import torch._C as callgrind_bindings" if bindings is None
|
||||
else f"import {bindings.__name__} as callgrind_bindings"),
|
||||
)
|
||||
|
||||
|
||||
CALLGRIND_SINGLETON: _ValgrindWrapper | None = None
|
||||
def wrapper_singleton() -> _ValgrindWrapper:
|
||||
global CALLGRIND_SINGLETON
|
||||
if CALLGRIND_SINGLETON is None:
|
||||
CALLGRIND_SINGLETON = _ValgrindWrapper()
|
||||
return CALLGRIND_SINGLETON
|
||||
+7157
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,472 @@
|
||||
#!/usr/bin/env python3
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any, TypeVar, NamedTuple
|
||||
from collections.abc import Callable, Sequence
|
||||
import textwrap
|
||||
import torch
|
||||
from torch._C import TupleType, ListType
|
||||
from torch.jit._recursive import wrap_cpp_module
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
MAX_RAW_TENSOR_SIZE = 16
|
||||
|
||||
class InflatableArg(NamedTuple):
|
||||
"""Helper type for bundled inputs.
|
||||
|
||||
'value' is the compressed/deflated input that is stored in the model. Value
|
||||
must be of the same type as the argument to the function that it is a deflated
|
||||
input for.
|
||||
|
||||
'fmt' is a formattable code string that is executed to inflate the compressed data into
|
||||
the appropriate input. It can use 'value' as an input to the format str. It must result
|
||||
in a value of the same type as 'value'.
|
||||
|
||||
'fmt_fn' is a formattable function code string that is executed to inflate the compressed
|
||||
data into the appropriate input. It must result in a value of the same type as 'value'.
|
||||
The function name should be the formattable part of the string.
|
||||
|
||||
Note: Only top level InflatableArgs can be inflated. i.e. you cannot place
|
||||
an inflatable arg inside of some other structure. You should instead create
|
||||
an inflatable arg such that the fmt code string returns the full structure
|
||||
of your input.
|
||||
"""
|
||||
|
||||
value: Any
|
||||
fmt: str = "{}"
|
||||
fmt_fn: str = ""
|
||||
|
||||
|
||||
def bundle_inputs(
|
||||
model: torch.jit.ScriptModule,
|
||||
inputs: Sequence[tuple[Any, ...]] | dict[Callable, Sequence[tuple[Any, ...]] | None] | None,
|
||||
info: list[str] | dict[Callable, list[str]] | None = None,
|
||||
*,
|
||||
_receive_inflate_expr: list[str] | None = None,
|
||||
) -> torch.jit.ScriptModule:
|
||||
"""Create and return a copy of the specified model with inputs attached.
|
||||
|
||||
The original model is not mutated or changed in any way.
|
||||
|
||||
Models with bundled inputs can be invoked in a uniform manner by
|
||||
benchmarking and code coverage tools.
|
||||
|
||||
If inputs is passed in as a list then the inputs will be bundled for 'forward'.
|
||||
If inputs is instead passed in as a map then all the methods specified in the map
|
||||
will have their corresponding inputs bundled. Info should match watchever type is
|
||||
chosen for the inputs.
|
||||
|
||||
The returned model will support the following methods:
|
||||
|
||||
`get_all_bundled_inputs_for_<function_name>() -> List[Tuple[Any, ...]]`
|
||||
Returns a list of tuples suitable for passing to the model like
|
||||
`for inp in model.get_all_bundled_inputs_for_foo(): model.foo(*inp)`
|
||||
|
||||
`get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]`
|
||||
Returns a dictionary mapping function names to a metadata dictionary.
|
||||
This nested dictionary maps preset strings like:
|
||||
'get_inputs_function_name' -> the name of a function attribute in this model that can be
|
||||
run to get back a list of inputs corresponding to that function.
|
||||
'info' -> the user provided extra information about the bundled inputs
|
||||
|
||||
If forward has bundled inputs then these following functions will also be defined on the returned module:
|
||||
|
||||
`get_all_bundled_inputs() -> List[Tuple[Any, ...]]`
|
||||
Returns a list of tuples suitable for passing to the model like
|
||||
`for inp in model.get_all_bundled_inputs(): model(*inp)`
|
||||
|
||||
`get_num_bundled_inputs() -> int`
|
||||
Equivalent to `len(model.get_all_bundled_inputs())`,
|
||||
but slightly easier to call from C++.
|
||||
|
||||
Inputs can be specified in one of two ways:
|
||||
|
||||
- The model can define `_generate_bundled_inputs_for_<function_name>`.
|
||||
If the user chooses this method inputs[<function>] should map to None
|
||||
|
||||
- The `inputs` argument to this function can be a dictionary mapping functions to a
|
||||
list of inputs, of the same form that will be returned by get_all_bundled_inputs_for_<function_name>.
|
||||
Alternatively if only bundling inputs for forward the map can be omitted and a singular list of inputs
|
||||
can be provided instead.
|
||||
|
||||
The type of the inputs is List[Tuple[Any, ...]]. The outer list corresponds with a
|
||||
list of inputs, the inner tuple is the list of args that together make up one input.
|
||||
For inputs of functions that take one arg, this will be a tuple of length one. The Any, ...
|
||||
is the actual data that makes up the args, e.g. a tensor.
|
||||
|
||||
Info is an optional parameter that maps functions to a list of strings providing extra information about that
|
||||
function's bundled inputs. Alternatively if only bundling inputs for forward the map can be omitted and
|
||||
a singular list of information can be provided instead. This could be descriptions, expected outputs, etc.
|
||||
- Ex: info={model.forward : ['man eating icecream', 'an airplane', 'a dog']}
|
||||
|
||||
This function will attempt to optimize arguments so that (e.g.)
|
||||
arguments like `torch.zeros(1000)` will be represented compactly.
|
||||
Only top-level arguments will be optimized.
|
||||
Tensors in lists or tuples will not.
|
||||
"""
|
||||
if not isinstance(model, torch.jit.ScriptModule):
|
||||
raise Exception("Only ScriptModule is supported.") # noqa: TRY002
|
||||
|
||||
ignored_methods, ignored_attrs = _get_bundled_inputs_attributes_and_methods(model)
|
||||
clone = torch._C._hack_do_not_use_clone_module_with_class( # type: ignore[attr-defined]
|
||||
model._c,
|
||||
ignored_methods,
|
||||
ignored_attrs,
|
||||
)
|
||||
|
||||
# The above cloning function returns a torch._C.scriptmodule and we need a torch.jit.scriptmodule.
|
||||
# Fortunately there is a function in _recursive that does exactly that conversion.
|
||||
cloned_module = wrap_cpp_module(clone)
|
||||
if isinstance(inputs, dict):
|
||||
if not isinstance(info, dict) and info is not None:
|
||||
raise AssertionError("If inputs is a dict, info must be a dict or None")
|
||||
augment_many_model_functions_with_bundled_inputs(cloned_module, inputs, _receive_inflate_expr, info)
|
||||
else:
|
||||
if not isinstance(info, list) and info is not None:
|
||||
raise AssertionError("If inputs is a list, info must be a list or None")
|
||||
augment_model_with_bundled_inputs(cloned_module, inputs, _receive_inflate_expr, info)
|
||||
return cloned_module
|
||||
|
||||
def augment_model_with_bundled_inputs(
|
||||
model: torch.jit.ScriptModule,
|
||||
inputs: Sequence[tuple[Any, ...]] | None = None,
|
||||
_receive_inflate_expr: list[str] | None = None, # For debugging.
|
||||
info: list[str] | None = None, # Optional argument to provide info about forward or its inputs
|
||||
skip_size_check=False,
|
||||
) -> None:
|
||||
"""Add bundled sample inputs to a model for the forward function.
|
||||
|
||||
Models with bundled inputs can be invoked in a uniform manner by
|
||||
benchmarking and code coverage tools.
|
||||
|
||||
Augmented models will support the following methods:
|
||||
|
||||
`get_all_bundled_inputs() -> List[Tuple[Any, ...]]`
|
||||
Returns a list of tuples suitable for passing to the model like
|
||||
`for inp in model.get_all_bundled_inputs(): model(*inp)`
|
||||
|
||||
`get_num_bundled_inputs() -> int`
|
||||
Equivalent to `len(model.get_all_bundled_inputs())`,
|
||||
but slightly easier to call from C++.
|
||||
|
||||
`get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]`
|
||||
Returns a dictionary mapping function names to a metadata dictionary.
|
||||
This nested dictionary maps preset strings like:
|
||||
'get_inputs_function_name' -> the name of a function attribute in this model that can be
|
||||
run to get back a list of inputs corresponding to that function.
|
||||
'info' -> the user provided extra information about the bundled inputs
|
||||
|
||||
Inputs can be specified in one of two ways:
|
||||
|
||||
- The model can define `_generate_bundled_inputs_for_forward`.
|
||||
If the user chooses this method inputs should be None
|
||||
|
||||
- `inputs` is a list of inputs of form List[Tuple[Any, ...]]. A list of tuples where the elements
|
||||
of each tuple are the args that make up one input.
|
||||
"""
|
||||
if not isinstance(model, torch.jit.ScriptModule):
|
||||
raise Exception("Only ScriptModule is supported.") # noqa: TRY002
|
||||
|
||||
forward: Callable = model.forward
|
||||
|
||||
# Sometimes forward won't have a name attached so just in case
|
||||
if not hasattr(forward, "__name__"):
|
||||
forward.__name__ = 'forward'
|
||||
augment_many_model_functions_with_bundled_inputs(
|
||||
model,
|
||||
inputs={forward : inputs},
|
||||
_receive_inflate_expr=_receive_inflate_expr,
|
||||
info={forward : info} if info else None,
|
||||
skip_size_check=skip_size_check,
|
||||
)
|
||||
|
||||
|
||||
def augment_many_model_functions_with_bundled_inputs(
|
||||
model: torch.jit.ScriptModule,
|
||||
inputs: dict[Callable, Sequence[tuple[Any, ...]] | None],
|
||||
_receive_inflate_expr: list[str] | None = None, # For debugging.
|
||||
info: dict[Callable, list[str]] | None = None, # Optional argument to provide info about the function or its inputs
|
||||
skip_size_check=False,
|
||||
) -> None:
|
||||
"""Add bundled sample inputs to a model for an arbitrary list of public functions.
|
||||
|
||||
Models with bundled inputs can be invoked in a uniform manner by
|
||||
benchmarking and code coverage tools.
|
||||
|
||||
Augmented models will support the following methods:
|
||||
|
||||
`get_all_bundled_inputs_for_<function_name>() -> List[Tuple[Any, ...]]`
|
||||
Returns a list of tuples suitable for passing to the model like
|
||||
`for inp in model.get_all_bundled_inputs_for_foo(): model.foo(*inp)`
|
||||
|
||||
`get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]`
|
||||
Returns a dictionary mapping function names to a metadata dictionary.
|
||||
This nested dictionary maps preset strings like:
|
||||
'get_inputs_function_name' -> the name of a function attribute in this model that can be
|
||||
run to get back a list of inputs corresponding to that function.
|
||||
'info' -> the user provided extra information about the bundled inputs
|
||||
|
||||
If forward has bundled inputs then these following functions are also defined:
|
||||
|
||||
`get_all_bundled_inputs() -> List[Tuple[Any, ...]]`
|
||||
Returns a list of tuples suitable for passing to the model like
|
||||
`for inp in model.get_all_bundled_inputs(): model(*inp)`
|
||||
|
||||
`get_num_bundled_inputs() -> int`
|
||||
Equivalent to `len(model.get_all_bundled_inputs())`,
|
||||
but slightly easier to call from C++.
|
||||
|
||||
Inputs can be specified in one of two ways:
|
||||
|
||||
- The model can define `_generate_bundled_inputs_for_<function_name>`.
|
||||
If the user chooses this method inputs[<function>] should map to None
|
||||
|
||||
- The `inputs` argument to this function can be a dictionary mapping functions to a
|
||||
list of inputs, of the same form that will be returned by get_all_bundled_inputs_for_<function_name>.
|
||||
The type of the inputs is List[Tuple[Any, ...]]. The outer list corresponds with a
|
||||
list of inputs, the inner tuple is the list of args that together make up one input.
|
||||
For inputs of functions that take one arg, this will be a tuple of length one. The Any, ...
|
||||
is the actual data that makes up the args, e.g. a tensor.
|
||||
|
||||
Info is an optional parameter that maps functions to a list of strings providing extra information about that
|
||||
function's bundled inputs. This could be descriptions, expected outputs, etc.
|
||||
- Ex: info={model.forward : ['man eating icecream', 'an airplane', 'a dog']}
|
||||
|
||||
This function will attempt to optimize arguments so that (e.g.)
|
||||
arguments like `torch.zeros(1000)` will be represented compactly.
|
||||
Only top-level arguments will be optimized.
|
||||
Tensors in lists or tuples will not.
|
||||
"""
|
||||
if not isinstance(model, torch.jit.ScriptModule):
|
||||
raise Exception("Only ScriptModule is supported.") # noqa: TRY002
|
||||
|
||||
if not inputs:
|
||||
raise Exception("Please provide inputs for at least 1 function") # noqa: TRY002
|
||||
|
||||
if hasattr(model, "get_all_bundled_inputs") or hasattr(model, "get_bundled_inputs_functions_and_info"):
|
||||
raise Exception( # noqa: TRY002
|
||||
"Models can only be augmented with bundled inputs once. "
|
||||
"This Model seems to have already been augmented with "
|
||||
"bundled inputs. Please start afresh with one that "
|
||||
"doesn't have bundled inputs.",
|
||||
)
|
||||
|
||||
get_bundled_inputs_functions_and_info_template = ""
|
||||
|
||||
for function, input_list in inputs.items():
|
||||
if hasattr(function, "__name__"):
|
||||
function_name = function.__name__
|
||||
else:
|
||||
if hasattr(function, "name"):
|
||||
function_name = function.name # type: ignore[attr-defined]
|
||||
else:
|
||||
raise Exception( # noqa: TRY002
|
||||
'At least one of your functions has no attribute name please ensure all have one. m.foo.name = "foo"')
|
||||
|
||||
|
||||
if input_list is not None and not isinstance(input_list, Sequence):
|
||||
raise TypeError(f"Error inputs for function {function_name} is not a Sequence")
|
||||
|
||||
function_arg_types = [arg.type for arg in function.schema.arguments[1:]] # type: ignore[attr-defined]
|
||||
deflated_inputs_type: ListType = ListType(TupleType(function_arg_types))
|
||||
model._c._register_attribute(f"_bundled_inputs_deflated_{function_name}", deflated_inputs_type, [])
|
||||
|
||||
if hasattr(model, "_generate_bundled_inputs_for_" + function_name):
|
||||
if input_list is not None:
|
||||
raise Exception( # noqa: TRY002
|
||||
f"inputs[{function_name}] is not None, but _generate_bundled_inputs_for_{function_name} is already defined"
|
||||
)
|
||||
# Model author already defined _generate_bundled_inputs_for_<function_name>.
|
||||
elif input_list is None or len(input_list) == 0:
|
||||
raise Exception( # noqa: TRY002
|
||||
f"inputs for {function_name} must be specified if "
|
||||
f"_generate_bundled_inputs_for_{function_name} is not already defined"
|
||||
)
|
||||
else:
|
||||
# Iterate over the inputs and args in each input.
|
||||
# Accumulate `deflated_inputs` as (possibly) compressed values
|
||||
# and `parts` to be joined into the expression that unpacks them.
|
||||
deflated_inputs = []
|
||||
parts = []
|
||||
for inp_idx, args in enumerate(input_list):
|
||||
if not isinstance(args, tuple) and not isinstance(args, list): # type: ignore[arg-type]
|
||||
raise TypeError(
|
||||
f"Error bundled input for function {function_name} idx: {inp_idx} is not a Tuple or a List"
|
||||
)
|
||||
deflated_args = []
|
||||
parts.append("(")
|
||||
for arg_idx, arg in enumerate(args):
|
||||
inflate_helper_fn_name = _get_inflate_helper_fn_name(arg_idx, inp_idx, function_name)
|
||||
deflated, inflater, helper_definition = _inflate_expr(
|
||||
arg,
|
||||
f"deflated[{inp_idx}][{arg_idx}]",
|
||||
inflate_helper_fn_name,
|
||||
skip_size_check=skip_size_check,
|
||||
)
|
||||
deflated_args.append(deflated)
|
||||
parts.append(f" {inflater},")
|
||||
if helper_definition:
|
||||
model.define(textwrap.dedent(helper_definition))
|
||||
deflated_inputs.append(tuple(deflated_args))
|
||||
parts.append("),")
|
||||
parts.append("")
|
||||
expr = "\n".join(parts)
|
||||
|
||||
# Back-channel return this expr for debugging.
|
||||
if _receive_inflate_expr is not None:
|
||||
_receive_inflate_expr.append(expr)
|
||||
setattr(model, f"_bundled_inputs_deflated_{function_name}", deflated_inputs)
|
||||
definition = textwrap.dedent("""
|
||||
def _generate_bundled_inputs_for_{name}(self):
|
||||
deflated = self._bundled_inputs_deflated_{name}
|
||||
return [
|
||||
{expr}
|
||||
]
|
||||
""").format(expr=expr, name=function_name)
|
||||
model.define(definition)
|
||||
|
||||
# Define get_all_bundled_inputs_for_<function_name> that caches the generated inputs.
|
||||
model.define(textwrap.dedent("""
|
||||
def get_all_bundled_inputs_for_{name}(self):
|
||||
all_inputs = self._generate_bundled_inputs_for_{name}()
|
||||
assert all_inputs is not None
|
||||
return all_inputs
|
||||
""").format(name=function_name))
|
||||
|
||||
# Add to the high level helper methods
|
||||
inputs_info = repr(info[function]) if info and function in info else '[]'
|
||||
get_bundled_inputs_functions_and_info_template += f"""
|
||||
temp_dict : Dict[str,List[str]] = {{}}
|
||||
info: List[str] = {inputs_info}
|
||||
|
||||
temp_dict['info'] = info
|
||||
temp_dict['get_inputs_function_name'] = ['get_all_bundled_inputs_for_{function_name}']
|
||||
all_inputs['{function_name}'] = temp_dict
|
||||
"""
|
||||
|
||||
# To ensure backwards compatibility and a streamlined api for forward these wrappers are provided
|
||||
if function_name == 'forward':
|
||||
model.define(textwrap.dedent("""
|
||||
def get_all_bundled_inputs(self):
|
||||
return self.get_all_bundled_inputs_for_forward()
|
||||
"""))
|
||||
model.define(textwrap.dedent("""
|
||||
def get_num_bundled_inputs(self):
|
||||
return len(self.get_all_bundled_inputs_for_forward())
|
||||
"""))
|
||||
|
||||
# Define some high level helper methods that act on all bundled inputs
|
||||
model.define(textwrap.dedent(f"""
|
||||
def get_bundled_inputs_functions_and_info(self):
|
||||
all_inputs : Dict[str, Dict[str,List[str]]] = {{}}
|
||||
{get_bundled_inputs_functions_and_info_template}
|
||||
return all_inputs
|
||||
"""))
|
||||
|
||||
def _inflate_expr(
|
||||
arg: T, ref: str, inflate_helper_fn_name: str, skip_size_check: bool = False
|
||||
) -> tuple[T | torch.Tensor, str, str | None]:
|
||||
# Allow custom inflation expressions any object.
|
||||
# For example, calling custom image-decoding ops.
|
||||
# Or just use "{}" as the format string to ignore size limits.
|
||||
if isinstance(arg, InflatableArg):
|
||||
if arg.fmt_fn:
|
||||
if arg.fmt not in ["{}", ""]:
|
||||
raise Exception( # noqa: TRY002
|
||||
f"Bundled input argument at position '{ref}' has "
|
||||
f"both arg.fmt_fn => \n{arg.fmt_fn} "
|
||||
f"\n and arg.fmt => {arg.fmt}. "
|
||||
"Please choose `arg.fmt` if the deflater is straightforward or "
|
||||
"`arg.fmt_fn` if you need a function."
|
||||
)
|
||||
|
||||
helper_definition = arg.fmt_fn.format(inflate_helper_fn_name)
|
||||
expr = f"self.{inflate_helper_fn_name}({ref})"
|
||||
|
||||
return arg.value, expr, helper_definition
|
||||
else:
|
||||
return arg.value, arg.fmt.format(ref), None
|
||||
|
||||
if isinstance(arg, torch.Tensor):
|
||||
# Small-storage tensors can just be saved directly.
|
||||
if arg._typed_storage().size() <= MAX_RAW_TENSOR_SIZE or skip_size_check:
|
||||
return arg, ref, None
|
||||
# Small contiguous tensors can be cloned to have small storage.
|
||||
# TODO: Should we do this even for non-contiguous tensors?
|
||||
if arg.is_contiguous() and arg.numel() <= MAX_RAW_TENSOR_SIZE:
|
||||
return arg.clone(), ref, None
|
||||
# Example inputs commonly come from torch.zeros, torch.ones, or torch.full.
|
||||
# These can be represented compactly.
|
||||
for fmt in [torch.contiguous_format, torch.channels_last]:
|
||||
if arg.is_contiguous(memory_format=fmt) and (arg == arg.flatten()[0]).all().item():
|
||||
return (arg.flatten()[0].clone().expand(*arg.size()),
|
||||
f"{ref}.contiguous(memory_format={fmt})", None)
|
||||
# Prevent big tensors from being bundled by default.
|
||||
# TODO: Provide more useful diagnostics.
|
||||
raise Exception( # noqa: TRY002
|
||||
f"Bundled input argument at position '{ref}' is "
|
||||
f"a tensor with storage size {arg._typed_storage().size()}. "
|
||||
f"You probably don't want to bundle this as an input. "
|
||||
)
|
||||
else:
|
||||
return arg, ref, None
|
||||
|
||||
def _get_bundled_inputs_attributes_and_methods(script_module: torch.jit.ScriptModule) -> tuple[list[str], list[str]]:
|
||||
methods: list[str] = []
|
||||
attributes: list[str] = []
|
||||
|
||||
# Has bundled inputs for forward
|
||||
if hasattr(script_module, 'get_all_bundled_inputs'):
|
||||
methods.append('get_all_bundled_inputs')
|
||||
methods.append('get_num_bundled_inputs')
|
||||
methods.append('run_on_bundled_input')
|
||||
|
||||
if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
|
||||
methods.append('get_bundled_inputs_functions_and_info')
|
||||
all_info = script_module.get_bundled_inputs_functions_and_info()
|
||||
for function_name in all_info:
|
||||
methods.append("get_all_bundled_inputs_for_" + function_name)
|
||||
methods.append("_generate_bundled_inputs_for_" + function_name)
|
||||
attributes.append("_bundled_inputs_deflated_" + function_name)
|
||||
|
||||
bundled_inputs_fn = getattr(
|
||||
script_module,
|
||||
f"get_all_bundled_inputs_for_{function_name}"
|
||||
)
|
||||
num_bundled_inputs: int = len(bundled_inputs_fn())
|
||||
|
||||
# Check inflate helper functions for each function, argument and bundled input
|
||||
func = getattr(script_module, function_name)
|
||||
for arg_idx in range(len(func.schema.arguments) - 1):
|
||||
for input_idx in range(num_bundled_inputs):
|
||||
helper_fn_name = _get_inflate_helper_fn_name(
|
||||
arg_idx=arg_idx,
|
||||
input_idx=input_idx,
|
||||
function_name=function_name
|
||||
)
|
||||
# if the arg has an InflatableArg with fmt_fn, add the helper function name
|
||||
if hasattr(script_module, helper_fn_name):
|
||||
methods.append(helper_fn_name)
|
||||
|
||||
return (methods, attributes)
|
||||
|
||||
|
||||
def _get_inflate_helper_fn_name(
|
||||
arg_idx: int,
|
||||
input_idx: int,
|
||||
function_name: str,
|
||||
) -> str:
|
||||
return f"_inflate_helper_for_{function_name}_input_{input_idx}_arg_{arg_idx}"
|
||||
|
||||
|
||||
|
||||
def bundle_randn(*size, dtype=None):
|
||||
"""Generate a tensor that will be inflated with torch.randn."""
|
||||
stub = torch.zeros(1, dtype=dtype).expand(*size)
|
||||
return InflatableArg(value=stub, fmt="torch.randn_like({})")
|
||||
|
||||
|
||||
def bundle_large_tensor(t):
|
||||
"""Wrap a tensor to allow bundling regardless of size."""
|
||||
return InflatableArg(value=t, fmt="{}")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,944 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
# Unlike the rest of the PyTorch this file must be python2 compliant.
|
||||
# This script outputs relevant system environment info
|
||||
# Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
|
||||
import datetime
|
||||
import json
|
||||
import locale
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from collections import namedtuple
|
||||
from typing import cast as _cast, Dict as _Dict
|
||||
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
TORCH_AVAILABLE = True
|
||||
except (ImportError, NameError, AttributeError, OSError):
|
||||
TORCH_AVAILABLE = False
|
||||
|
||||
# System Environment Information
|
||||
SystemEnv = namedtuple(
|
||||
"SystemEnv",
|
||||
[
|
||||
"torch_version",
|
||||
"is_debug_build",
|
||||
"cuda_compiled_version",
|
||||
"gcc_version",
|
||||
"clang_version",
|
||||
"cmake_version",
|
||||
"os",
|
||||
"libc_version",
|
||||
"python_version",
|
||||
"python_platform",
|
||||
"is_cuda_available",
|
||||
"cuda_runtime_version",
|
||||
"cuda_module_loading",
|
||||
"nvidia_driver_version",
|
||||
"nvidia_gpu_models",
|
||||
"cudnn_version",
|
||||
"is_xpu_available",
|
||||
"pip_version", # 'pip' or 'pip3'
|
||||
"pip_packages",
|
||||
"conda_packages",
|
||||
"hip_compiled_version",
|
||||
"hip_runtime_version",
|
||||
"miopen_runtime_version",
|
||||
"caching_allocator_config",
|
||||
"is_xnnpack_available",
|
||||
"cpu_info",
|
||||
],
|
||||
)
|
||||
|
||||
COMMON_PATTERNS = [
|
||||
"torch",
|
||||
"numpy",
|
||||
"triton",
|
||||
"optree",
|
||||
]
|
||||
|
||||
NVIDIA_PATTERNS = [
|
||||
"cuda-cudart",
|
||||
"cuda-cupti",
|
||||
"cuda-libraries",
|
||||
"cuda-opencl",
|
||||
"cuda-nvrtc",
|
||||
"cuda-runtime",
|
||||
"cublas",
|
||||
"cudnn",
|
||||
"cufft",
|
||||
"curand",
|
||||
"cusolver",
|
||||
"cusparse",
|
||||
"nccl",
|
||||
"nvjitlink",
|
||||
"nvtx",
|
||||
]
|
||||
|
||||
ONEAPI_PATTERNS = [
|
||||
"dpcpp-cpp-rt",
|
||||
"intel-cmplr-lib-rt",
|
||||
"intel-cmplr-lib-ur",
|
||||
"intel-cmplr-lic-rt",
|
||||
"intel-opencl-rt",
|
||||
"intel-sycl-rt",
|
||||
"mkl",
|
||||
"onemkl-sycl-blas",
|
||||
"onemkl-sycl-dft",
|
||||
"onemkl-sycl-lapack",
|
||||
"onemkl-sycl-rng",
|
||||
"onemkl-sycl-sparse",
|
||||
"intel-openmp",
|
||||
"tbb",
|
||||
"impi-rt",
|
||||
"impi-devel",
|
||||
"oneccl",
|
||||
"oneccl-devel",
|
||||
"intel-pti",
|
||||
"umf",
|
||||
"tcmlib",
|
||||
]
|
||||
|
||||
CONDA_PATTERNS = [
|
||||
"cudatoolkit",
|
||||
"soumith",
|
||||
"mkl",
|
||||
"magma",
|
||||
]
|
||||
|
||||
PIP_PATTERNS = [
|
||||
"mypy",
|
||||
"flake8",
|
||||
"onnx",
|
||||
]
|
||||
|
||||
|
||||
def run(command):
|
||||
"""Return (return-code, stdout, stderr)."""
|
||||
shell = type(command) is str
|
||||
p = subprocess.Popen(
|
||||
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell
|
||||
)
|
||||
raw_output, raw_err = p.communicate()
|
||||
rc = p.returncode
|
||||
if get_platform() == "win32":
|
||||
enc = "oem"
|
||||
else:
|
||||
enc = locale.getpreferredencoding()
|
||||
output = raw_output.decode(enc)
|
||||
err = raw_err.decode(enc)
|
||||
return rc, output.strip(), err.strip()
|
||||
|
||||
|
||||
def run_and_read_all(run_lambda, command):
|
||||
"""Run command using run_lambda; reads and returns entire output if rc is 0."""
|
||||
rc, out, _ = run_lambda(command)
|
||||
if rc != 0:
|
||||
return None
|
||||
return out
|
||||
|
||||
|
||||
def run_and_parse_first_match(run_lambda, command, regex):
|
||||
"""Run command using run_lambda, returns the first regex match if it exists."""
|
||||
rc, out, _ = run_lambda(command)
|
||||
if rc != 0:
|
||||
return None
|
||||
match = re.search(regex, out)
|
||||
if match is None:
|
||||
return None
|
||||
return match.group(1)
|
||||
|
||||
|
||||
def run_and_return_first_line(run_lambda, command):
|
||||
"""Run command using run_lambda and returns first line if output is not empty."""
|
||||
rc, out, _ = run_lambda(command)
|
||||
if rc != 0:
|
||||
return None
|
||||
return out.split("\n")[0]
|
||||
|
||||
|
||||
def get_conda_packages(run_lambda, patterns=None):
|
||||
if patterns is None:
|
||||
patterns = CONDA_PATTERNS + COMMON_PATTERNS + NVIDIA_PATTERNS + ONEAPI_PATTERNS
|
||||
conda = os.environ.get("CONDA_EXE", "conda")
|
||||
out = run_and_read_all(run_lambda, "{} list".format(conda))
|
||||
if out is None:
|
||||
return out
|
||||
|
||||
return "\n".join(
|
||||
line
|
||||
for line in out.splitlines()
|
||||
if not line.startswith("#") and any(name in line for name in patterns)
|
||||
)
|
||||
|
||||
|
||||
def get_gcc_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, "gcc --version", r"gcc (.*)")
|
||||
|
||||
|
||||
def get_clang_version(run_lambda):
|
||||
return run_and_parse_first_match(
|
||||
run_lambda, "clang --version", r"clang version (.*)"
|
||||
)
|
||||
|
||||
|
||||
def get_cmake_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, "cmake --version", r"cmake (.*)")
|
||||
|
||||
|
||||
def get_nvidia_driver_version(run_lambda):
|
||||
if get_platform() == "darwin":
|
||||
cmd = "kextstat | grep -i cuda"
|
||||
return run_and_parse_first_match(
|
||||
run_lambda, cmd, r"com[.]nvidia[.]CUDA [(](.*?)[)]"
|
||||
)
|
||||
smi = get_nvidia_smi()
|
||||
return run_and_parse_first_match(run_lambda, smi, r"Driver Version: (.*?) ")
|
||||
|
||||
|
||||
def get_gpu_info(run_lambda):
|
||||
if get_platform() == "darwin" or (
|
||||
TORCH_AVAILABLE
|
||||
and hasattr(torch.version, "hip")
|
||||
and torch.version.hip is not None
|
||||
):
|
||||
if TORCH_AVAILABLE and torch.cuda.is_available():
|
||||
if torch.version.hip is not None:
|
||||
prop = torch.cuda.get_device_properties(0)
|
||||
if hasattr(prop, "gcnArchName"):
|
||||
gcnArch = " ({})".format(prop.gcnArchName)
|
||||
else:
|
||||
gcnArch = "NoGCNArchNameOnOldPyTorch"
|
||||
else:
|
||||
gcnArch = ""
|
||||
return torch.cuda.get_device_name(None) + gcnArch
|
||||
return None
|
||||
smi = get_nvidia_smi()
|
||||
uuid_regex = re.compile(r" \(UUID: .+?\)")
|
||||
rc, out, _ = run_lambda(smi + " -L")
|
||||
if rc != 0:
|
||||
return None
|
||||
# Anonymize GPUs by removing their UUID
|
||||
return re.sub(uuid_regex, "", out)
|
||||
|
||||
|
||||
def get_running_cuda_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, "nvcc --version", r"release .+ V(.*)")
|
||||
|
||||
|
||||
def get_cudnn_version(run_lambda):
|
||||
"""Return a list of libcudnn.so; it's hard to tell which one is being used."""
|
||||
if get_platform() == "win32":
|
||||
system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
|
||||
cuda_path = os.environ.get("CUDA_PATH", "%CUDA_PATH%")
|
||||
where_cmd = os.path.join(system_root, "System32", "where")
|
||||
cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path)
|
||||
elif get_platform() == "darwin":
|
||||
# CUDA libraries and drivers can be found in /usr/local/cuda/. See
|
||||
# https://docs.nvidia.com/cuda/archive/9.0/cuda-installation-guide-mac-os-x/index.html#installation
|
||||
# https://docs.nvidia.com/deeplearning/cudnn/installation/latest/
|
||||
# Use CUDNN_LIBRARY when cudnn library is installed elsewhere.
|
||||
cudnn_cmd = "ls /usr/local/cuda/lib/libcudnn*"
|
||||
else:
|
||||
cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev'
|
||||
rc, out, _ = run_lambda(cudnn_cmd)
|
||||
# find will return 1 if there are permission errors or if not found
|
||||
if len(out) == 0 or (rc != 1 and rc != 0):
|
||||
l = os.environ.get("CUDNN_LIBRARY")
|
||||
if l is not None and os.path.isfile(l):
|
||||
return os.path.realpath(l)
|
||||
return None
|
||||
files_set = set()
|
||||
for fn in out.split("\n"):
|
||||
fn = os.path.realpath(fn) # eliminate symbolic links
|
||||
if os.path.isfile(fn):
|
||||
files_set.add(fn)
|
||||
if not files_set:
|
||||
return None
|
||||
# Alphabetize the result because the order is non-deterministic otherwise
|
||||
files = sorted(files_set)
|
||||
if len(files) == 1:
|
||||
return files[0]
|
||||
result = "\n".join(files)
|
||||
return "Probably one of the following:\n{}".format(result)
|
||||
|
||||
|
||||
def get_nvidia_smi():
|
||||
# Note: nvidia-smi is currently available only on Windows and Linux
|
||||
smi = "nvidia-smi"
|
||||
if get_platform() == "win32":
|
||||
system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
|
||||
program_files_root = os.environ.get("PROGRAMFILES", "C:\\Program Files")
|
||||
legacy_path = os.path.join(
|
||||
program_files_root, "NVIDIA Corporation", "NVSMI", smi
|
||||
)
|
||||
new_path = os.path.join(system_root, "System32", smi)
|
||||
smis = [new_path, legacy_path]
|
||||
for candidate_smi in smis:
|
||||
if os.path.exists(candidate_smi):
|
||||
smi = '"{}"'.format(candidate_smi)
|
||||
break
|
||||
return smi
|
||||
|
||||
|
||||
def _detect_linux_pkg_manager():
|
||||
if get_platform() != "linux":
|
||||
return "N/A"
|
||||
for mgr_name in ["dpkg", "dnf", "yum", "zypper"]:
|
||||
rc, _, _ = run(f"which {mgr_name}")
|
||||
if rc == 0:
|
||||
return mgr_name
|
||||
return "N/A"
|
||||
|
||||
|
||||
def get_linux_pkg_version(run_lambda, pkg_name):
|
||||
pkg_mgr = _detect_linux_pkg_manager()
|
||||
if pkg_mgr == "N/A":
|
||||
return "N/A"
|
||||
|
||||
grep_version = {
|
||||
"dpkg": {
|
||||
"field_index": 2,
|
||||
"command": "dpkg -l | grep {}",
|
||||
},
|
||||
"dnf": {
|
||||
"field_index": 1,
|
||||
"command": "dnf list | grep {}",
|
||||
},
|
||||
"yum": {
|
||||
"field_index": 1,
|
||||
"command": "yum list | grep {}",
|
||||
},
|
||||
"zypper": {
|
||||
"field_index": 2,
|
||||
"command": "zypper info {} | grep Version",
|
||||
},
|
||||
}
|
||||
|
||||
field_index: int = int(_cast(int, grep_version[pkg_mgr]["field_index"]))
|
||||
cmd: str = str(grep_version[pkg_mgr]["command"])
|
||||
cmd = cmd.format(pkg_name)
|
||||
ret = run_and_read_all(run_lambda, cmd)
|
||||
if ret is None or ret == "":
|
||||
return "N/A"
|
||||
lst = re.sub(" +", " ", ret).split(" ")
|
||||
if len(lst) <= field_index:
|
||||
return "N/A"
|
||||
return lst[field_index]
|
||||
|
||||
|
||||
def get_intel_gpu_driver_version(run_lambda):
|
||||
lst = []
|
||||
platform = get_platform()
|
||||
if platform == "linux":
|
||||
pkgs = { # type: ignore[var-annotated]
|
||||
"dpkg": {
|
||||
"intel-opencl-icd",
|
||||
"libze1",
|
||||
"level-zero",
|
||||
},
|
||||
"dnf": {
|
||||
"intel-opencl",
|
||||
"level-zero",
|
||||
},
|
||||
"yum": {
|
||||
"intel-opencl",
|
||||
"level-zero",
|
||||
},
|
||||
"zypper": {
|
||||
"intel-opencl",
|
||||
"level-zero",
|
||||
},
|
||||
}.get(_detect_linux_pkg_manager(), {})
|
||||
for pkg in pkgs:
|
||||
ver = get_linux_pkg_version(run_lambda, pkg)
|
||||
if ver != "N/A":
|
||||
lst.append(f"* {pkg}:\t{ver}")
|
||||
if platform in ["win32", "cygwin"]:
|
||||
txt = run_and_read_all(
|
||||
run_lambda,
|
||||
'powershell.exe "gwmi -Class Win32_PnpSignedDriver | where{$_.DeviceClass -eq \\"DISPLAY\\"\
|
||||
-and $_.Manufacturer -match \\"Intel\\"} | Select-Object -Property DeviceName,DriverVersion,DriverDate\
|
||||
| ConvertTo-Json"',
|
||||
)
|
||||
try:
|
||||
obj = json.loads(txt)
|
||||
if type(obj) is list:
|
||||
for o in obj:
|
||||
lst.append(
|
||||
f'* {o["DeviceName"]}: {o["DriverVersion"]} ({o["DriverDate"]})'
|
||||
)
|
||||
else:
|
||||
lst.append(f'* {obj["DriverVersion"]} ({obj["DriverDate"]})')
|
||||
except ValueError as e:
|
||||
lst.append(txt)
|
||||
lst.append(str(e))
|
||||
return "\n".join(lst)
|
||||
|
||||
|
||||
def get_intel_gpu_onboard(run_lambda):
|
||||
lst: list[str] = []
|
||||
platform = get_platform()
|
||||
if platform == "linux":
|
||||
txt = run_and_read_all(run_lambda, "xpu-smi discovery -j")
|
||||
if txt:
|
||||
try:
|
||||
obj = json.loads(txt)
|
||||
device_list = obj.get("device_list", [])
|
||||
if isinstance(device_list, list) and device_list:
|
||||
lst.extend(f'* {device["device_name"]}' for device in device_list)
|
||||
else:
|
||||
lst.append("N/A")
|
||||
except (ValueError, TypeError) as e:
|
||||
lst.append(txt)
|
||||
lst.append(str(e))
|
||||
else:
|
||||
lst.append("N/A")
|
||||
if platform in ["win32", "cygwin"]:
|
||||
txt = run_and_read_all(
|
||||
run_lambda,
|
||||
'powershell.exe "gwmi -Class Win32_PnpSignedDriver | where{$_.DeviceClass -eq \\"DISPLAY\\"\
|
||||
-and $_.Manufacturer -match \\"Intel\\"} | Select-Object -Property DeviceName | ConvertTo-Json"',
|
||||
)
|
||||
if txt:
|
||||
try:
|
||||
obj = json.loads(txt)
|
||||
if isinstance(obj, list) and obj:
|
||||
lst.extend(f'* {device["DeviceName"]}' for device in obj)
|
||||
else:
|
||||
lst.append(f'* {obj.get("DeviceName", "N/A")}')
|
||||
except ValueError as e:
|
||||
lst.append(txt)
|
||||
lst.append(str(e))
|
||||
else:
|
||||
lst.append("N/A")
|
||||
return "\n".join(lst)
|
||||
|
||||
|
||||
def get_intel_gpu_detected(run_lambda):
|
||||
if not TORCH_AVAILABLE or not hasattr(torch, "xpu"):
|
||||
return "N/A"
|
||||
|
||||
device_count = torch.xpu.device_count()
|
||||
if device_count == 0:
|
||||
return "N/A"
|
||||
|
||||
devices = [
|
||||
f"* [{i}] {torch.xpu.get_device_properties(i)}" for i in range(device_count)
|
||||
]
|
||||
return "\n".join(devices)
|
||||
|
||||
|
||||
# example outputs of CPU infos
|
||||
# * linux
|
||||
# Architecture: x86_64
|
||||
# CPU op-mode(s): 32-bit, 64-bit
|
||||
# Address sizes: 46 bits physical, 48 bits virtual
|
||||
# Byte Order: Little Endian
|
||||
# CPU(s): 128
|
||||
# On-line CPU(s) list: 0-127
|
||||
# Vendor ID: GenuineIntel
|
||||
# Model name: Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
|
||||
# CPU family: 6
|
||||
# Model: 106
|
||||
# Thread(s) per core: 2
|
||||
# Core(s) per socket: 32
|
||||
# Socket(s): 2
|
||||
# Stepping: 6
|
||||
# BogoMIPS: 5799.78
|
||||
# Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr
|
||||
# sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl
|
||||
# xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16
|
||||
# pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand
|
||||
# hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced
|
||||
# fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap
|
||||
# avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1
|
||||
# xsaves wbnoinvd ida arat avx512vbmi pku ospke avx512_vbmi2 gfni vaes vpclmulqdq
|
||||
# avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid md_clear flush_l1d arch_capabilities
|
||||
# Virtualization features:
|
||||
# Hypervisor vendor: KVM
|
||||
# Virtualization type: full
|
||||
# Caches (sum of all):
|
||||
# L1d: 3 MiB (64 instances)
|
||||
# L1i: 2 MiB (64 instances)
|
||||
# L2: 80 MiB (64 instances)
|
||||
# L3: 108 MiB (2 instances)
|
||||
# NUMA:
|
||||
# NUMA node(s): 2
|
||||
# NUMA node0 CPU(s): 0-31,64-95
|
||||
# NUMA node1 CPU(s): 32-63,96-127
|
||||
# Vulnerabilities:
|
||||
# Itlb multihit: Not affected
|
||||
# L1tf: Not affected
|
||||
# Mds: Not affected
|
||||
# Meltdown: Not affected
|
||||
# Mmio stale data: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
|
||||
# Retbleed: Not affected
|
||||
# Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
|
||||
# Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
|
||||
# Spectre v2: Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
|
||||
# Srbds: Not affected
|
||||
# Tsx async abort: Not affected
|
||||
# * win32
|
||||
# Architecture=9
|
||||
# CurrentClockSpeed=2900
|
||||
# DeviceID=CPU0
|
||||
# Family=179
|
||||
# L2CacheSize=40960
|
||||
# L2CacheSpeed=
|
||||
# Manufacturer=GenuineIntel
|
||||
# MaxClockSpeed=2900
|
||||
# Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
|
||||
# ProcessorType=3
|
||||
# Revision=27142
|
||||
#
|
||||
# Architecture=9
|
||||
# CurrentClockSpeed=2900
|
||||
# DeviceID=CPU1
|
||||
# Family=179
|
||||
# L2CacheSize=40960
|
||||
# L2CacheSpeed=
|
||||
# Manufacturer=GenuineIntel
|
||||
# MaxClockSpeed=2900
|
||||
# Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
|
||||
# ProcessorType=3
|
||||
# Revision=27142
|
||||
|
||||
|
||||
def get_cpu_info(run_lambda):
|
||||
rc, out, err = 0, "", ""
|
||||
if get_platform() == "linux":
|
||||
rc, out, err = run_lambda("lscpu")
|
||||
elif get_platform() == "win32":
|
||||
rc, out, err = run_lambda(
|
||||
'powershell.exe "gwmi -Class Win32_Processor | Select-Object -Property Name,Manufacturer,Family,\
|
||||
Architecture,ProcessorType,DeviceID,CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision\
|
||||
| ConvertTo-Json"'
|
||||
)
|
||||
if rc == 0:
|
||||
lst = []
|
||||
try:
|
||||
obj = json.loads(out)
|
||||
if type(obj) is list:
|
||||
for o in obj:
|
||||
lst.append("----------------------")
|
||||
lst.extend([f"{k}: {v}" for (k, v) in o.items()])
|
||||
else:
|
||||
lst.extend([f"{k}: {v}" for (k, v) in obj.items()])
|
||||
except ValueError as e:
|
||||
lst.append(out)
|
||||
lst.append(str(e))
|
||||
out = "\n".join(lst)
|
||||
elif get_platform() == "darwin":
|
||||
rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
|
||||
cpu_info = "None"
|
||||
if rc == 0:
|
||||
cpu_info = out
|
||||
else:
|
||||
cpu_info = err
|
||||
return cpu_info
|
||||
|
||||
|
||||
def get_platform():
|
||||
if sys.platform.startswith("linux"):
|
||||
return "linux"
|
||||
elif sys.platform.startswith("win32"):
|
||||
return "win32"
|
||||
elif sys.platform.startswith("cygwin"):
|
||||
return "cygwin"
|
||||
elif sys.platform.startswith("darwin"):
|
||||
return "darwin"
|
||||
else:
|
||||
return sys.platform
|
||||
|
||||
|
||||
def get_mac_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, "sw_vers -productVersion", r"(.*)")
|
||||
|
||||
|
||||
def get_windows_version(run_lambda):
|
||||
ret = run_and_read_all(
|
||||
run_lambda,
|
||||
'powershell.exe "gwmi -Class Win32_OperatingSystem | Select-Object -Property Caption,\
|
||||
OSArchitecture,Version | ConvertTo-Json"',
|
||||
)
|
||||
try:
|
||||
obj = json.loads(ret)
|
||||
ret = f'{obj["Caption"]} ({obj["Version"]} {obj["OSArchitecture"]})'
|
||||
except ValueError as e:
|
||||
ret += f"\n{str(e)}"
|
||||
return ret
|
||||
|
||||
|
||||
def get_lsb_version(run_lambda):
|
||||
return run_and_parse_first_match(
|
||||
run_lambda, "lsb_release -a", r"Description:\t(.*)"
|
||||
)
|
||||
|
||||
|
||||
def check_release_file(run_lambda):
|
||||
return run_and_parse_first_match(
|
||||
run_lambda, "cat /etc/*-release", r'PRETTY_NAME="(.*)"'
|
||||
)
|
||||
|
||||
|
||||
def get_os(run_lambda):
|
||||
from platform import machine
|
||||
|
||||
platform = get_platform()
|
||||
|
||||
if platform in ["win32", "cygwin"]:
|
||||
return get_windows_version(run_lambda)
|
||||
|
||||
if platform == "darwin":
|
||||
version = get_mac_version(run_lambda)
|
||||
if version is None:
|
||||
return None
|
||||
return "macOS {} ({})".format(version, machine())
|
||||
|
||||
if platform == "linux":
|
||||
# Ubuntu/Debian based
|
||||
desc = get_lsb_version(run_lambda)
|
||||
if desc is not None:
|
||||
return "{} ({})".format(desc, machine())
|
||||
|
||||
# Try reading /etc/*-release
|
||||
desc = check_release_file(run_lambda)
|
||||
if desc is not None:
|
||||
return "{} ({})".format(desc, machine())
|
||||
|
||||
return "{} ({})".format(platform, machine())
|
||||
|
||||
# Unknown platform
|
||||
return platform
|
||||
|
||||
|
||||
def get_python_platform():
|
||||
import platform
|
||||
|
||||
return platform.platform()
|
||||
|
||||
|
||||
def get_libc_version():
|
||||
import platform
|
||||
|
||||
if get_platform() != "linux":
|
||||
return "N/A"
|
||||
return "-".join(platform.libc_ver())
|
||||
|
||||
|
||||
def get_pip_packages(run_lambda, patterns=None):
|
||||
"""Return `pip list` output. Note: will also find conda-installed pytorch and numpy packages."""
|
||||
if patterns is None:
|
||||
patterns = PIP_PATTERNS + COMMON_PATTERNS + NVIDIA_PATTERNS + ONEAPI_PATTERNS
|
||||
|
||||
pip_version = "pip3" if sys.version_info.major == 3 else "pip"
|
||||
|
||||
os.environ["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
|
||||
# People generally have pip as `pip` or `pip3`
|
||||
# But here it is invoked as `python -mpip`
|
||||
out = run_and_read_all(
|
||||
run_lambda, [sys.executable, "-mpip", "list", "--format=freeze"]
|
||||
)
|
||||
if out is None:
|
||||
return pip_version, out
|
||||
|
||||
filtered_out = "\n".join(
|
||||
line for line in out.splitlines() if any(name in line for name in patterns)
|
||||
)
|
||||
|
||||
return pip_version, filtered_out
|
||||
|
||||
|
||||
def get_cachingallocator_config() -> _Dict[str, str]:
|
||||
"""Return the caching allocator configuration from environment variables.
|
||||
"""
|
||||
# pyrefly: ignore [bad-return]
|
||||
return {
|
||||
var: os.environ.get(var)
|
||||
for var in (
|
||||
"PYTORCH_CUDA_ALLOC_CONF",
|
||||
"PYTORCH_HIP_ALLOC_CONF",
|
||||
"PYTORCH_ALLOC_CONF",
|
||||
)
|
||||
if os.environ.get(var)
|
||||
}
|
||||
|
||||
|
||||
def get_cuda_module_loading_config():
|
||||
if TORCH_AVAILABLE and torch.cuda.is_available():
|
||||
torch.cuda.init()
|
||||
config = os.environ.get("CUDA_MODULE_LOADING", "")
|
||||
return config
|
||||
else:
|
||||
return "N/A"
|
||||
|
||||
|
||||
def is_xnnpack_available():
|
||||
if TORCH_AVAILABLE:
|
||||
import torch.backends.xnnpack
|
||||
|
||||
return str(torch.backends.xnnpack.enabled) # type: ignore[attr-defined]
|
||||
else:
|
||||
return "N/A"
|
||||
|
||||
|
||||
def get_env_info():
|
||||
"""
|
||||
Collects environment information to aid in debugging.
|
||||
|
||||
The returned environment information contains details on torch version, is debug build
|
||||
or not, cuda compiled version, gcc version, clang version, cmake version, operating
|
||||
system, libc version, python version, python platform, CUDA availability, CUDA
|
||||
runtime version, CUDA module loading config, GPU model and configuration, Nvidia
|
||||
driver version, cuDNN version, pip version and versions of relevant pip and
|
||||
conda packages, HIP runtime version, MIOpen runtime version,
|
||||
Caching allocator config, XNNPACK availability and CPU information.
|
||||
|
||||
Returns:
|
||||
SystemEnv (namedtuple): A tuple containing various environment details
|
||||
and system information.
|
||||
"""
|
||||
run_lambda = run
|
||||
pip_version, pip_list_output = get_pip_packages(run_lambda)
|
||||
|
||||
if TORCH_AVAILABLE:
|
||||
version_str = torch.__version__
|
||||
debug_mode_str = str(torch.version.debug)
|
||||
cuda_available_str = str(torch.cuda.is_available())
|
||||
cuda_version_str = torch.version.cuda
|
||||
xpu_available_str = str(torch.xpu.is_available())
|
||||
if torch.xpu.is_available():
|
||||
xpu_available_str = (
|
||||
f"{xpu_available_str}\n"
|
||||
+ f"XPU used to build PyTorch: {torch.version.xpu}\n"
|
||||
+ f"Intel GPU driver version:\n{get_intel_gpu_driver_version(run_lambda)}\n"
|
||||
+ f"Intel GPU models onboard:\n{get_intel_gpu_onboard(run_lambda)}\n"
|
||||
+ f"Intel GPU models detected:\n{get_intel_gpu_detected(run_lambda)}"
|
||||
)
|
||||
if (
|
||||
not hasattr(torch.version, "hip") or torch.version.hip is None
|
||||
): # cuda version
|
||||
hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A"
|
||||
else: # HIP version
|
||||
|
||||
def get_version_or_na(cfg, prefix):
|
||||
_lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s]
|
||||
return _lst[0] if _lst else "N/A"
|
||||
|
||||
cfg = torch._C._show_config().split("\n")
|
||||
hip_runtime_version = get_version_or_na(cfg, "HIP Runtime")
|
||||
miopen_runtime_version = get_version_or_na(cfg, "MIOpen")
|
||||
cuda_version_str = "N/A"
|
||||
hip_compiled_version = torch.version.hip
|
||||
else:
|
||||
version_str = debug_mode_str = cuda_available_str = cuda_version_str = xpu_available_str = "N/A" # type: ignore[assignment]
|
||||
hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A"
|
||||
|
||||
sys_version = sys.version.replace("\n", " ")
|
||||
|
||||
conda_packages = get_conda_packages(run_lambda)
|
||||
|
||||
return SystemEnv(
|
||||
torch_version=version_str,
|
||||
is_debug_build=debug_mode_str,
|
||||
python_version="{} ({}-bit runtime)".format(
|
||||
sys_version, sys.maxsize.bit_length() + 1
|
||||
),
|
||||
python_platform=get_python_platform(),
|
||||
is_cuda_available=cuda_available_str,
|
||||
cuda_compiled_version=cuda_version_str,
|
||||
cuda_runtime_version=get_running_cuda_version(run_lambda),
|
||||
cuda_module_loading=get_cuda_module_loading_config(),
|
||||
nvidia_gpu_models=get_gpu_info(run_lambda),
|
||||
nvidia_driver_version=get_nvidia_driver_version(run_lambda),
|
||||
cudnn_version=get_cudnn_version(run_lambda),
|
||||
is_xpu_available=xpu_available_str,
|
||||
hip_compiled_version=hip_compiled_version,
|
||||
hip_runtime_version=hip_runtime_version,
|
||||
miopen_runtime_version=miopen_runtime_version,
|
||||
pip_version=pip_version,
|
||||
pip_packages=pip_list_output,
|
||||
conda_packages=conda_packages,
|
||||
os=get_os(run_lambda),
|
||||
libc_version=get_libc_version(),
|
||||
gcc_version=get_gcc_version(run_lambda),
|
||||
clang_version=get_clang_version(run_lambda),
|
||||
cmake_version=get_cmake_version(run_lambda),
|
||||
caching_allocator_config=get_cachingallocator_config(),
|
||||
is_xnnpack_available=is_xnnpack_available(),
|
||||
cpu_info=get_cpu_info(run_lambda),
|
||||
)
|
||||
|
||||
|
||||
env_info_fmt = """
|
||||
PyTorch version: {torch_version}
|
||||
Is debug build: {is_debug_build}
|
||||
CUDA used to build PyTorch: {cuda_compiled_version}
|
||||
ROCM used to build PyTorch: {hip_compiled_version}
|
||||
|
||||
OS: {os}
|
||||
GCC version: {gcc_version}
|
||||
Clang version: {clang_version}
|
||||
CMake version: {cmake_version}
|
||||
Libc version: {libc_version}
|
||||
|
||||
Python version: {python_version}
|
||||
Python platform: {python_platform}
|
||||
Is CUDA available: {is_cuda_available}
|
||||
CUDA runtime version: {cuda_runtime_version}
|
||||
CUDA_MODULE_LOADING set to: {cuda_module_loading}
|
||||
GPU models and configuration: {nvidia_gpu_models}
|
||||
Nvidia driver version: {nvidia_driver_version}
|
||||
cuDNN version: {cudnn_version}
|
||||
Is XPU available: {is_xpu_available}
|
||||
HIP runtime version: {hip_runtime_version}
|
||||
MIOpen runtime version: {miopen_runtime_version}
|
||||
Is XNNPACK available: {is_xnnpack_available}
|
||||
Caching allocator config: {caching_allocator_config}
|
||||
|
||||
CPU:
|
||||
{cpu_info}
|
||||
|
||||
Versions of relevant libraries:
|
||||
{pip_packages}
|
||||
{conda_packages}
|
||||
""".strip()
|
||||
|
||||
|
||||
def pretty_str(envinfo):
|
||||
def replace_nones(dct, replacement="Could not collect"):
|
||||
for key in dct:
|
||||
if dct[key] is not None:
|
||||
continue
|
||||
dct[key] = replacement
|
||||
return dct
|
||||
|
||||
def replace_bools(dct, true="Yes", false="No"):
|
||||
for key in dct:
|
||||
if dct[key] is True:
|
||||
dct[key] = true
|
||||
elif dct[key] is False:
|
||||
dct[key] = false
|
||||
return dct
|
||||
|
||||
def prepend(text, tag="[prepend]"):
|
||||
lines = text.split("\n")
|
||||
updated_lines = [tag + line for line in lines]
|
||||
return "\n".join(updated_lines)
|
||||
|
||||
def replace_if_empty(text, replacement="No relevant packages"):
|
||||
if text is not None and len(text) == 0:
|
||||
return replacement
|
||||
return text
|
||||
|
||||
def maybe_start_on_next_line(string):
|
||||
# If `string` is multiline, prepend a \n to it.
|
||||
if string is not None and len(string.split("\n")) > 1:
|
||||
return "\n{}\n".format(string)
|
||||
return string
|
||||
|
||||
mutable_dict = envinfo._asdict()
|
||||
|
||||
# If nvidia_gpu_models is multiline, start on the next line
|
||||
mutable_dict["nvidia_gpu_models"] = maybe_start_on_next_line(
|
||||
envinfo.nvidia_gpu_models
|
||||
)
|
||||
|
||||
# If the machine doesn't have CUDA, report some fields as 'No CUDA'
|
||||
dynamic_cuda_fields = [
|
||||
"cuda_runtime_version",
|
||||
"nvidia_gpu_models",
|
||||
"nvidia_driver_version",
|
||||
]
|
||||
all_cuda_fields = dynamic_cuda_fields + ["cudnn_version"]
|
||||
all_dynamic_cuda_fields_missing = all(
|
||||
mutable_dict[field] is None for field in dynamic_cuda_fields
|
||||
)
|
||||
if (
|
||||
TORCH_AVAILABLE
|
||||
and not torch.cuda.is_available()
|
||||
and all_dynamic_cuda_fields_missing
|
||||
):
|
||||
for field in all_cuda_fields:
|
||||
mutable_dict[field] = "No CUDA"
|
||||
if envinfo.cuda_compiled_version is None:
|
||||
mutable_dict["cuda_compiled_version"] = "None"
|
||||
|
||||
# Replace True with Yes, False with No
|
||||
mutable_dict = replace_bools(mutable_dict)
|
||||
|
||||
# Replace all None objects with 'Could not collect'
|
||||
mutable_dict = replace_nones(mutable_dict)
|
||||
|
||||
# If either of these are '', replace with 'No relevant packages'
|
||||
mutable_dict["pip_packages"] = replace_if_empty(mutable_dict["pip_packages"])
|
||||
mutable_dict["conda_packages"] = replace_if_empty(mutable_dict["conda_packages"])
|
||||
|
||||
# Tag conda and pip packages with a prefix
|
||||
# If they were previously None, they'll show up as ie '[conda] Could not collect'
|
||||
if mutable_dict["pip_packages"]:
|
||||
mutable_dict["pip_packages"] = prepend(
|
||||
mutable_dict["pip_packages"], "[{}] ".format(envinfo.pip_version)
|
||||
)
|
||||
if mutable_dict["conda_packages"]:
|
||||
mutable_dict["conda_packages"] = prepend(
|
||||
mutable_dict["conda_packages"], "[conda] "
|
||||
)
|
||||
mutable_dict["cpu_info"] = envinfo.cpu_info
|
||||
mutable_dict["caching_allocator_config"] = envinfo.caching_allocator_config
|
||||
if not envinfo.caching_allocator_config:
|
||||
mutable_dict["caching_allocator_config"] = "N/A"
|
||||
return env_info_fmt.format(**mutable_dict)
|
||||
|
||||
|
||||
def get_pretty_env_info():
|
||||
"""
|
||||
Returns a pretty string of environment information.
|
||||
|
||||
This function retrieves environment information by calling the `get_env_info` function
|
||||
and then formats the information into a human-readable string. The retrieved environment
|
||||
information is listed in the document of `get_env_info`.
|
||||
This function is used in `python collect_env.py` that should be executed when reporting a bug.
|
||||
|
||||
Returns:
|
||||
str: A pretty string of the environment information.
|
||||
"""
|
||||
return pretty_str(get_env_info())
|
||||
|
||||
|
||||
def main() -> None:
|
||||
print("Collecting environment information...")
|
||||
output = get_pretty_env_info()
|
||||
print(output)
|
||||
|
||||
if (
|
||||
TORCH_AVAILABLE
|
||||
and hasattr(torch, "utils")
|
||||
and hasattr(torch.utils, "_crash_handler")
|
||||
):
|
||||
minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
|
||||
if sys.platform == "linux" and os.path.exists(minidump_dir):
|
||||
dumps = [
|
||||
os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir)
|
||||
]
|
||||
latest = max(dumps, key=os.path.getctime)
|
||||
ctime = os.path.getctime(latest)
|
||||
creation_time = datetime.datetime.fromtimestamp(ctime).strftime(
|
||||
"%Y-%m-%d %H:%M:%S"
|
||||
)
|
||||
msg = (
|
||||
"\n*** Detected a minidump at {} created on {}, ".format(
|
||||
latest, creation_time
|
||||
)
|
||||
+ "if this is related to your bug please include it when you file a report ***"
|
||||
)
|
||||
print(msg, file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,12 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from torch._C import _get_cpp_backtrace
|
||||
|
||||
def get_cpp_backtrace(frames_to_skip=0, maximum_number_of_frames=64) -> str:
|
||||
r"""
|
||||
Return a string containing the C++ stack trace of the current thread.
|
||||
|
||||
Args:
|
||||
frames_to_skip (int): the number of frames to skip from the top of the stack
|
||||
maximum_number_of_frames (int): the maximum number of frames to return
|
||||
"""
|
||||
return _get_cpp_backtrace(frames_to_skip, maximum_number_of_frames)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,78 @@
|
||||
from torch.utils.data.dataloader import (
|
||||
_DatasetKind,
|
||||
DataLoader,
|
||||
default_collate,
|
||||
default_convert,
|
||||
get_worker_info,
|
||||
)
|
||||
from torch.utils.data.datapipes._decorator import (
|
||||
argument_validation,
|
||||
functional_datapipe,
|
||||
guaranteed_datapipes_determinism,
|
||||
non_deterministic,
|
||||
runtime_validation,
|
||||
runtime_validation_disabled,
|
||||
)
|
||||
from torch.utils.data.datapipes.datapipe import (
|
||||
DataChunk,
|
||||
DFIterDataPipe,
|
||||
IterDataPipe,
|
||||
MapDataPipe,
|
||||
)
|
||||
from torch.utils.data.dataset import (
|
||||
ChainDataset,
|
||||
ConcatDataset,
|
||||
Dataset,
|
||||
IterableDataset,
|
||||
random_split,
|
||||
StackDataset,
|
||||
Subset,
|
||||
TensorDataset,
|
||||
)
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from torch.utils.data.sampler import (
|
||||
BatchSampler,
|
||||
RandomSampler,
|
||||
Sampler,
|
||||
SequentialSampler,
|
||||
SubsetRandomSampler,
|
||||
WeightedRandomSampler,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchSampler",
|
||||
"ChainDataset",
|
||||
"ConcatDataset",
|
||||
"DFIterDataPipe",
|
||||
"DataChunk",
|
||||
"DataLoader",
|
||||
"Dataset",
|
||||
"DistributedSampler",
|
||||
"IterDataPipe",
|
||||
"IterableDataset",
|
||||
"MapDataPipe",
|
||||
"RandomSampler",
|
||||
"Sampler",
|
||||
"SequentialSampler",
|
||||
"StackDataset",
|
||||
"Subset",
|
||||
"SubsetRandomSampler",
|
||||
"TensorDataset",
|
||||
"WeightedRandomSampler",
|
||||
"_DatasetKind",
|
||||
"argument_validation",
|
||||
"default_collate",
|
||||
"default_convert",
|
||||
"functional_datapipe",
|
||||
"get_worker_info",
|
||||
"guaranteed_datapipes_determinism",
|
||||
"non_deterministic",
|
||||
"random_split",
|
||||
"runtime_validation",
|
||||
"runtime_validation_disabled",
|
||||
]
|
||||
|
||||
# Please keep this list sorted
|
||||
if __all__ != sorted(__all__):
|
||||
raise AssertionError("__all__ is not sorted")
|
||||
@@ -0,0 +1,53 @@
|
||||
r"""Utility classes & functions for data loading. Code in this folder is mostly used by ../dataloder.py.
|
||||
|
||||
A lot of multiprocessing is used in data loading, which only supports running
|
||||
functions defined in global environment (py2 can't serialize static methods).
|
||||
Therefore, for code tidiness we put these functions into different files in this
|
||||
folder.
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import sys
|
||||
|
||||
# old private location of the ExceptionWrapper that some users rely on:
|
||||
from torch._utils import ExceptionWrapper
|
||||
|
||||
|
||||
IS_WINDOWS = sys.platform == "win32"
|
||||
|
||||
|
||||
MP_STATUS_CHECK_INTERVAL = 5.0
|
||||
r"""Interval (in seconds) to check status of processes to avoid hanging in
|
||||
multiprocessing data loading. This is mainly used in getting data from
|
||||
another process, in which case we need to periodically check whether the
|
||||
sender is alive to prevent hanging."""
|
||||
|
||||
|
||||
python_exit_status = False
|
||||
r"""Whether Python is shutting down. This flag is guaranteed to be set before
|
||||
the Python core library resources are freed, but Python may already be exiting
|
||||
for some time when this is set.
|
||||
|
||||
Hook to set this flag is `_set_python_exit_flag`, and is inspired by a similar
|
||||
hook in Python 3.7 multiprocessing library:
|
||||
https://github.com/python/cpython/blob/d4d60134b29290049e28df54f23493de4f1824b6/Lib/multiprocessing/util.py#L277-L327
|
||||
"""
|
||||
|
||||
|
||||
try:
|
||||
import numpy
|
||||
|
||||
HAS_NUMPY = True
|
||||
except ModuleNotFoundError:
|
||||
HAS_NUMPY = False
|
||||
|
||||
|
||||
def _set_python_exit_flag() -> None:
|
||||
global python_exit_status
|
||||
python_exit_status = True
|
||||
|
||||
|
||||
atexit.register(_set_python_exit_flag)
|
||||
|
||||
|
||||
from . import collate, fetch, pin_memory, signal_handling, worker
|
||||
@@ -0,0 +1,401 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Contains definitions of the methods used by the _BaseDataLoaderIter workers.
|
||||
|
||||
These methods are used to collate samples fetched from dataset into Tensor(s).
|
||||
These **needs** to be in global scope since Py2 doesn't support serializing
|
||||
static methods.
|
||||
|
||||
`default_collate` and `default_convert` are exposed to users via 'dataloader.py'.
|
||||
"""
|
||||
|
||||
import collections
|
||||
import contextlib
|
||||
import copy
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
np_str_obj_array_pattern = re.compile(r"[SaUO]")
|
||||
|
||||
|
||||
def default_convert(data):
|
||||
r"""
|
||||
Convert each NumPy array element into a :class:`torch.Tensor`.
|
||||
|
||||
If the input is a `Sequence`, `Collection`, or `Mapping`, it tries to convert each element inside to a :class:`torch.Tensor`.
|
||||
If the input is not an NumPy array, it is left unchanged.
|
||||
This is used as the default function for collation when both `batch_sampler` and `batch_size`
|
||||
are NOT defined in :class:`~torch.utils.data.DataLoader`.
|
||||
|
||||
The general input type to output type mapping is similar to that
|
||||
of :func:`~torch.utils.data.default_collate`. See the description there for more details.
|
||||
|
||||
Args:
|
||||
data: a single data point to be converted
|
||||
|
||||
Examples:
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # Example with `int`
|
||||
>>> default_convert(0)
|
||||
0
|
||||
>>> # Example with NumPy array
|
||||
>>> default_convert(np.array([0, 1]))
|
||||
tensor([0, 1])
|
||||
>>> # Example with NamedTuple
|
||||
>>> Point = namedtuple("Point", ["x", "y"])
|
||||
>>> default_convert(Point(0, 0))
|
||||
Point(x=0, y=0)
|
||||
>>> default_convert(Point(np.array(0), np.array(0)))
|
||||
Point(x=tensor(0), y=tensor(0))
|
||||
>>> # Example with List
|
||||
>>> default_convert([np.array([0, 1]), np.array([2, 3])])
|
||||
[tensor([0, 1]), tensor([2, 3])]
|
||||
"""
|
||||
elem_type = type(data)
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data
|
||||
elif (
|
||||
elem_type.__module__ == "numpy"
|
||||
and elem_type.__name__ != "str_"
|
||||
and elem_type.__name__ != "string_"
|
||||
):
|
||||
# array of string classes and object
|
||||
if (
|
||||
elem_type.__name__ == "ndarray"
|
||||
and np_str_obj_array_pattern.search(data.dtype.str) is not None
|
||||
):
|
||||
return data
|
||||
return torch.as_tensor(data)
|
||||
elif isinstance(data, collections.abc.Mapping):
|
||||
try:
|
||||
if isinstance(data, collections.abc.MutableMapping):
|
||||
# The mapping type may have extra properties, so we can't just
|
||||
# use `type(data)(...)` to create the new mapping.
|
||||
# Create a clone and update it if the mapping type is mutable.
|
||||
clone = copy.copy(data)
|
||||
clone.update({key: default_convert(data[key]) for key in data})
|
||||
return clone
|
||||
else:
|
||||
return elem_type({key: default_convert(data[key]) for key in data})
|
||||
except TypeError:
|
||||
# The mapping type may not support `copy()` / `update(mapping)`
|
||||
# or `__init__(iterable)`.
|
||||
return {key: default_convert(data[key]) for key in data}
|
||||
elif isinstance(data, tuple) and hasattr(data, "_fields"): # namedtuple
|
||||
return elem_type(*(default_convert(d) for d in data))
|
||||
elif isinstance(data, tuple):
|
||||
return [default_convert(d) for d in data] # Backwards compatibility.
|
||||
elif isinstance(data, collections.abc.Sequence) and not isinstance(
|
||||
data, (str, bytes)
|
||||
):
|
||||
try:
|
||||
if isinstance(data, collections.abc.MutableSequence):
|
||||
# The sequence type may have extra properties, so we can't just
|
||||
# use `type(data)(...)` to create the new sequence.
|
||||
# Create a clone and update it if the sequence type is mutable.
|
||||
clone = copy.copy(data) # type: ignore[arg-type]
|
||||
for i, d in enumerate(data):
|
||||
clone[i] = default_convert(d)
|
||||
return clone
|
||||
else:
|
||||
return elem_type([default_convert(d) for d in data])
|
||||
except TypeError:
|
||||
# The sequence type may not support `copy()` / `__setitem__(index, item)`
|
||||
# or `__init__(iterable)` (e.g., `range`).
|
||||
return [default_convert(d) for d in data]
|
||||
else:
|
||||
return data
|
||||
|
||||
|
||||
default_collate_err_msg_format = (
|
||||
"default_collate: batch must contain tensors, numpy arrays, numbers, "
|
||||
"dicts or lists; found {}"
|
||||
)
|
||||
|
||||
|
||||
def collate(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
r"""
|
||||
General collate function that handles collection type of element within each batch.
|
||||
|
||||
The function also opens function registry to deal with specific element types. `default_collate_fn_map`
|
||||
provides default collate functions for tensors, numpy arrays, numbers and strings.
|
||||
|
||||
Args:
|
||||
batch: a single batch to be collated
|
||||
collate_fn_map: Optional dictionary mapping from element type to the corresponding collate function.
|
||||
If the element type isn't present in this dictionary,
|
||||
this function will go through each key of the dictionary in the insertion order to
|
||||
invoke the corresponding collate function if the element type is a subclass of the key.
|
||||
|
||||
Examples:
|
||||
>>> def collate_tensor_fn(batch, *, collate_fn_map):
|
||||
... # Extend this function to handle batch of tensors
|
||||
... return torch.stack(batch, 0)
|
||||
>>> def custom_collate(batch):
|
||||
... collate_map = {torch.Tensor: collate_tensor_fn}
|
||||
... return collate(batch, collate_fn_map=collate_map)
|
||||
>>> # Extend `default_collate` by in-place modifying `default_collate_fn_map`
|
||||
>>> default_collate_fn_map.update({torch.Tensor: collate_tensor_fn})
|
||||
|
||||
Note:
|
||||
Each collate function requires a positional argument for batch and a keyword argument
|
||||
for the dictionary of collate functions as `collate_fn_map`.
|
||||
"""
|
||||
elem = batch[0]
|
||||
elem_type = type(elem)
|
||||
|
||||
if collate_fn_map is not None:
|
||||
if elem_type in collate_fn_map:
|
||||
return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)
|
||||
|
||||
for collate_type in collate_fn_map:
|
||||
if isinstance(elem, collate_type):
|
||||
return collate_fn_map[collate_type](
|
||||
batch, collate_fn_map=collate_fn_map
|
||||
)
|
||||
|
||||
if isinstance(elem, collections.abc.Mapping):
|
||||
try:
|
||||
if isinstance(elem, collections.abc.MutableMapping):
|
||||
# The mapping type may have extra properties, so we can't just
|
||||
# use `type(data)(...)` to create the new mapping.
|
||||
# Create a clone and update it if the mapping type is mutable.
|
||||
clone = copy.copy(elem)
|
||||
clone.update(
|
||||
{
|
||||
key: collate(
|
||||
[d[key] for d in batch], collate_fn_map=collate_fn_map
|
||||
)
|
||||
for key in elem
|
||||
}
|
||||
)
|
||||
return clone
|
||||
else:
|
||||
return elem_type(
|
||||
{
|
||||
key: collate(
|
||||
[d[key] for d in batch], collate_fn_map=collate_fn_map
|
||||
)
|
||||
for key in elem
|
||||
}
|
||||
)
|
||||
except TypeError:
|
||||
# The mapping type may not support `copy()` / `update(mapping)`
|
||||
# or `__init__(iterable)`.
|
||||
return {
|
||||
key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map)
|
||||
for key in elem
|
||||
}
|
||||
elif isinstance(elem, tuple) and hasattr(elem, "_fields"): # namedtuple
|
||||
return elem_type(
|
||||
*(
|
||||
collate(samples, collate_fn_map=collate_fn_map)
|
||||
for samples in zip(*batch, strict=False)
|
||||
)
|
||||
)
|
||||
elif isinstance(elem, collections.abc.Sequence):
|
||||
# check to make sure that the elements in batch have consistent size
|
||||
it = iter(batch)
|
||||
elem_size = len(next(it))
|
||||
|
||||
if not all(len(elem) == elem_size for elem in it):
|
||||
raise RuntimeError("each element in list of batch should be of equal size")
|
||||
transposed = list(
|
||||
zip(*batch, strict=False)
|
||||
) # It may be accessed twice, so we use a list.
|
||||
|
||||
if isinstance(elem, tuple):
|
||||
return [
|
||||
collate(samples, collate_fn_map=collate_fn_map)
|
||||
for samples in transposed
|
||||
] # Backwards compatibility.
|
||||
else:
|
||||
try:
|
||||
if isinstance(elem, collections.abc.MutableSequence):
|
||||
# The sequence type may have extra properties, so we can't just
|
||||
# use `type(data)(...)` to create the new sequence.
|
||||
# Create a clone and update it if the sequence type is mutable.
|
||||
clone = copy.copy(elem) # type: ignore[arg-type]
|
||||
for i, samples in enumerate(transposed):
|
||||
clone[i] = collate(samples, collate_fn_map=collate_fn_map)
|
||||
return clone
|
||||
else:
|
||||
return elem_type(
|
||||
[
|
||||
collate(samples, collate_fn_map=collate_fn_map)
|
||||
for samples in transposed
|
||||
]
|
||||
)
|
||||
except TypeError:
|
||||
# The sequence type may not support `copy()` / `__setitem__(index, item)`
|
||||
# or `__init__(iterable)` (e.g., `range`).
|
||||
return [
|
||||
collate(samples, collate_fn_map=collate_fn_map)
|
||||
for samples in transposed
|
||||
]
|
||||
|
||||
raise TypeError(default_collate_err_msg_format.format(elem_type))
|
||||
|
||||
|
||||
def collate_tensor_fn(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
elem = batch[0]
|
||||
out = None
|
||||
if elem.is_nested:
|
||||
raise RuntimeError(
|
||||
"Batches of nested tensors are not currently supported by the default collate_fn; "
|
||||
"please provide a custom collate_fn to handle them appropriately."
|
||||
)
|
||||
if elem.layout in {
|
||||
torch.sparse_coo,
|
||||
torch.sparse_csr,
|
||||
torch.sparse_bsr,
|
||||
torch.sparse_csc,
|
||||
torch.sparse_bsc,
|
||||
}:
|
||||
raise RuntimeError(
|
||||
"Batches of sparse tensors are not currently supported by the default collate_fn; "
|
||||
"please provide a custom collate_fn to handle them appropriately."
|
||||
)
|
||||
if torch.utils.data.get_worker_info() is not None:
|
||||
# If we're in a background process, concatenate directly into a
|
||||
# shared memory tensor to avoid an extra copy
|
||||
numel = sum(x.numel() for x in batch)
|
||||
storage = elem._typed_storage()._new_shared(numel, device=elem.device)
|
||||
out = elem.new(storage).resize_(len(batch), *list(elem.size()))
|
||||
return torch.stack(batch, 0, out=out)
|
||||
|
||||
|
||||
def collate_numpy_array_fn(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
elem = batch[0]
|
||||
# array of string classes and object
|
||||
if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
|
||||
raise TypeError(default_collate_err_msg_format.format(elem.dtype))
|
||||
|
||||
return collate([torch.as_tensor(b) for b in batch], collate_fn_map=collate_fn_map)
|
||||
|
||||
|
||||
def collate_numpy_scalar_fn(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
return torch.as_tensor(batch)
|
||||
|
||||
|
||||
def collate_float_fn(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
return torch.tensor(batch, dtype=torch.float64)
|
||||
|
||||
|
||||
def collate_int_fn(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
return torch.tensor(batch)
|
||||
|
||||
|
||||
def collate_str_fn(
|
||||
batch,
|
||||
*,
|
||||
collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None,
|
||||
):
|
||||
return batch
|
||||
|
||||
|
||||
default_collate_fn_map: dict[type | tuple[type, ...], Callable] = {
|
||||
torch.Tensor: collate_tensor_fn
|
||||
}
|
||||
with contextlib.suppress(ImportError):
|
||||
import numpy as np
|
||||
|
||||
# For both ndarray and memmap (subclass of ndarray)
|
||||
default_collate_fn_map[np.ndarray] = collate_numpy_array_fn
|
||||
# See scalars hierarchy: https://numpy.org/doc/stable/reference/arrays.scalars.html
|
||||
# Skip string scalars
|
||||
default_collate_fn_map[(np.bool_, np.number, np.object_)] = collate_numpy_scalar_fn
|
||||
default_collate_fn_map[float] = collate_float_fn
|
||||
default_collate_fn_map[int] = collate_int_fn
|
||||
default_collate_fn_map[str] = collate_str_fn
|
||||
default_collate_fn_map[bytes] = collate_str_fn
|
||||
|
||||
|
||||
def default_collate(batch):
|
||||
r"""
|
||||
Take in a batch of data and put the elements within the batch into a tensor with an additional outer dimension - batch size.
|
||||
|
||||
The exact output type can be a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a
|
||||
Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type.
|
||||
This is used as the default function for collation when
|
||||
`batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`.
|
||||
|
||||
Here is the general input type (based on the type of the element within the batch) to output type mapping:
|
||||
|
||||
* :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size)
|
||||
* NumPy Arrays -> :class:`torch.Tensor`
|
||||
* `float` -> :class:`torch.Tensor`
|
||||
* `int` -> :class:`torch.Tensor`
|
||||
* `str` -> `str` (unchanged)
|
||||
* `bytes` -> `bytes` (unchanged)
|
||||
* `Mapping[K, V_i]` -> `Mapping[K, default_collate([V_1, V_2, ...])]`
|
||||
* `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[default_collate([V1_1, V1_2, ...]),
|
||||
default_collate([V2_1, V2_2, ...]), ...]`
|
||||
* `Sequence[V1_i, V2_i, ...]` -> `Sequence[default_collate([V1_1, V1_2, ...]),
|
||||
default_collate([V2_1, V2_2, ...]), ...]`
|
||||
|
||||
Args:
|
||||
batch: a single batch to be collated
|
||||
|
||||
Examples:
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # Example with a batch of `int`s:
|
||||
>>> default_collate([0, 1, 2, 3])
|
||||
tensor([0, 1, 2, 3])
|
||||
>>> # Example with a batch of `str`s:
|
||||
>>> default_collate(["a", "b", "c"])
|
||||
['a', 'b', 'c']
|
||||
>>> # Example with `Map` inside the batch:
|
||||
>>> default_collate([{"A": 0, "B": 1}, {"A": 100, "B": 100}])
|
||||
{'A': tensor([ 0, 100]), 'B': tensor([ 1, 100])}
|
||||
>>> # Example with `NamedTuple` inside the batch:
|
||||
>>> Point = namedtuple("Point", ["x", "y"])
|
||||
>>> default_collate([Point(0, 0), Point(1, 1)])
|
||||
Point(x=tensor([0, 1]), y=tensor([0, 1]))
|
||||
>>> # Example with `Tuple` inside the batch:
|
||||
>>> default_collate([(0, 1), (2, 3)])
|
||||
[tensor([0, 2]), tensor([1, 3])]
|
||||
>>> # Example with `List` inside the batch:
|
||||
>>> default_collate([[0, 1], [2, 3]])
|
||||
[tensor([0, 2]), tensor([1, 3])]
|
||||
>>> # Two options to extend `default_collate` to handle specific type
|
||||
>>> # Option 1: Write custom collate function and invoke `default_collate`
|
||||
>>> def custom_collate(batch):
|
||||
... elem = batch[0]
|
||||
... if isinstance(elem, CustomType): # Some custom condition
|
||||
... return ...
|
||||
... else: # Fall back to `default_collate`
|
||||
... return default_collate(batch)
|
||||
>>> # Option 2: In-place modify `default_collate_fn_map`
|
||||
>>> def collate_customtype_fn(batch, *, collate_fn_map=None):
|
||||
... return ...
|
||||
>>> default_collate_fn_map.update(CustomType, collate_customtype_fn)
|
||||
>>> default_collate(batch) # Handle `CustomType` automatically
|
||||
"""
|
||||
return collate(batch, collate_fn_map=default_collate_fn_map)
|
||||
@@ -0,0 +1,57 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Contains definitions of the methods used by the _BaseDataLoaderIter to fetch data from an iterable-style or map-style dataset.
|
||||
|
||||
This logic is shared in both single- and multi-processing data loading.
|
||||
"""
|
||||
|
||||
from typing import NoReturn
|
||||
|
||||
|
||||
class _BaseDatasetFetcher:
|
||||
def __init__(self, dataset, auto_collation, collate_fn, drop_last) -> None:
|
||||
self.dataset = dataset
|
||||
self.auto_collation = auto_collation
|
||||
self.collate_fn = collate_fn
|
||||
self.drop_last = drop_last
|
||||
|
||||
def fetch(self, possibly_batched_index) -> NoReturn:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class _IterableDatasetFetcher(_BaseDatasetFetcher):
|
||||
def __init__(self, dataset, auto_collation, collate_fn, drop_last) -> None:
|
||||
super().__init__(dataset, auto_collation, collate_fn, drop_last)
|
||||
self.dataset_iter = iter(dataset)
|
||||
self.ended = False
|
||||
|
||||
def fetch(self, possibly_batched_index):
|
||||
if self.ended:
|
||||
raise StopIteration
|
||||
|
||||
if self.auto_collation:
|
||||
data = []
|
||||
for _ in possibly_batched_index:
|
||||
try:
|
||||
data.append(next(self.dataset_iter))
|
||||
except StopIteration:
|
||||
self.ended = True
|
||||
break
|
||||
if len(data) == 0 or (
|
||||
self.drop_last and len(data) < len(possibly_batched_index)
|
||||
):
|
||||
raise StopIteration
|
||||
else:
|
||||
data = next(self.dataset_iter)
|
||||
return self.collate_fn(data)
|
||||
|
||||
|
||||
class _MapDatasetFetcher(_BaseDatasetFetcher):
|
||||
def fetch(self, possibly_batched_index):
|
||||
if self.auto_collation:
|
||||
if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
|
||||
data = self.dataset.__getitems__(possibly_batched_index)
|
||||
else:
|
||||
data = [self.dataset[idx] for idx in possibly_batched_index]
|
||||
else:
|
||||
data = self.dataset[possibly_batched_index]
|
||||
return self.collate_fn(data)
|
||||
@@ -0,0 +1,108 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Contains definitions of the methods used by the _BaseDataLoaderIter to put fetched tensors into pinned memory.
|
||||
|
||||
These **needs** to be in global scope since Py2 doesn't support serializing
|
||||
static methods.
|
||||
"""
|
||||
|
||||
import collections
|
||||
import copy
|
||||
import queue
|
||||
|
||||
import torch
|
||||
from torch._utils import ExceptionWrapper
|
||||
|
||||
from . import MP_STATUS_CHECK_INTERVAL
|
||||
|
||||
|
||||
def _pin_memory_loop(in_queue, out_queue, device_id, done_event, device) -> None:
|
||||
# This setting is thread local, and prevents the copy in pin_memory from
|
||||
# consuming all CPU cores.
|
||||
torch.set_num_threads(1)
|
||||
|
||||
torch.multiprocessing._set_thread_name("pt_data_pin")
|
||||
torch.accelerator.set_device_index(device_id)
|
||||
|
||||
def do_one_step() -> None:
|
||||
try:
|
||||
r = in_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
|
||||
except queue.Empty:
|
||||
return
|
||||
idx, data = r
|
||||
if not done_event.is_set() and not isinstance(data, ExceptionWrapper):
|
||||
try:
|
||||
data = pin_memory(data, device)
|
||||
except Exception:
|
||||
data = ExceptionWrapper(
|
||||
where=f"in pin memory thread for device {device_id}"
|
||||
)
|
||||
r = (idx, data)
|
||||
while not done_event.is_set():
|
||||
try:
|
||||
out_queue.put(r, timeout=MP_STATUS_CHECK_INTERVAL)
|
||||
break
|
||||
except queue.Full:
|
||||
continue
|
||||
|
||||
# See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the
|
||||
# logic of this function.
|
||||
while not done_event.is_set():
|
||||
# Make sure that we don't preserve any object from one iteration
|
||||
# to the next
|
||||
do_one_step()
|
||||
|
||||
|
||||
def pin_memory(data, device=None):
|
||||
if isinstance(data, torch.Tensor):
|
||||
return data.pin_memory()
|
||||
|
||||
if hasattr(data, "pin_memory"):
|
||||
return data.pin_memory()
|
||||
|
||||
if isinstance(data, (str, bytes)):
|
||||
return data
|
||||
|
||||
if isinstance(data, collections.abc.Mapping):
|
||||
try:
|
||||
if isinstance(data, collections.abc.MutableMapping):
|
||||
# The mapping type may have extra properties, so we can't just
|
||||
# use `type(data)(...)` to create the new mapping.
|
||||
# Create a clone and update it if the mapping type is mutable.
|
||||
clone = copy.copy(data)
|
||||
clone.update(
|
||||
{k: pin_memory(sample, device) for k, sample in data.items()}
|
||||
)
|
||||
return clone
|
||||
else:
|
||||
# pyrefly: ignore [bad-instantiation]
|
||||
return type(data)(
|
||||
# pyrefly: ignore [bad-argument-count]
|
||||
{k: pin_memory(sample, device) for k, sample in data.items()}
|
||||
) # type: ignore[call-arg]
|
||||
except TypeError:
|
||||
# The mapping type may not support `copy()` / `update(mapping)`
|
||||
# or `__init__(iterable)`.
|
||||
return {k: pin_memory(sample, device) for k, sample in data.items()}
|
||||
|
||||
if isinstance(data, tuple):
|
||||
if hasattr(data, "_fields"): # namedtuple
|
||||
return type(data)(*(pin_memory(sample, device) for sample in data))
|
||||
return type(data)(pin_memory(sample, device) for sample in data)
|
||||
|
||||
if isinstance(data, collections.abc.Sequence):
|
||||
try:
|
||||
if isinstance(data, collections.abc.MutableSequence):
|
||||
# The sequence type may have extra properties, so we can't just
|
||||
# use `type(data)(...)` to create the new sequence.
|
||||
# Create a clone and update it if the sequence type is mutable.
|
||||
clone = copy.copy(data) # type: ignore[arg-type]
|
||||
for i, item in enumerate(data):
|
||||
clone[i] = pin_memory(item, device)
|
||||
return clone
|
||||
return type(data)([pin_memory(sample, device) for sample in data]) # type: ignore[call-arg]
|
||||
except TypeError:
|
||||
# The sequence type may not support `copy()` / `__setitem__(index, item)`
|
||||
# or `__init__(iterable)` (e.g., `range`).
|
||||
return [pin_memory(sample, device) for sample in data]
|
||||
|
||||
return data
|
||||
@@ -0,0 +1,80 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Signal handling for multiprocessing data loading.
|
||||
|
||||
NOTE [ Signal handling in multiprocessing data loading ]
|
||||
|
||||
In cases like DataLoader, if a worker process dies due to bus error/segfault
|
||||
or just hang, the main process will hang waiting for data. This is difficult
|
||||
to avoid on PyTorch side as it can be caused by limited shm, or other
|
||||
libraries users call in the workers. In this file and `DataLoader.cpp`, we make
|
||||
our best effort to provide some error message to users when such unfortunate
|
||||
events happen.
|
||||
|
||||
When a _BaseDataLoaderIter starts worker processes, their pids are registered in a
|
||||
defined in `DataLoader.cpp`: id(_BaseDataLoaderIter) => Collection[ Worker pids ]
|
||||
via `_set_worker_pids`.
|
||||
|
||||
When an error happens in a worker process, the main process received a SIGCHLD,
|
||||
and Python will eventually call the handler registered below
|
||||
(in `_set_SIGCHLD_handler`). In the handler, the `_error_if_any_worker_fails`
|
||||
call checks all registered worker pids and raise proper error message to
|
||||
prevent main process from hanging waiting for data from worker.
|
||||
|
||||
Additionally, at the beginning of each worker's `_utils.worker._worker_loop`,
|
||||
`_set_worker_signal_handlers` is called to register critical signal handlers
|
||||
(e.g., for SIGSEGV, SIGBUS, SIGFPE, SIGTERM) in C, which just prints an error
|
||||
message to stderr before triggering the default handler. So a message will also
|
||||
be printed from the worker process when it is killed by such signals.
|
||||
|
||||
See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for the reasoning of
|
||||
this signal handling design and other mechanism we implement to make our
|
||||
multiprocessing data loading robust to errors.
|
||||
"""
|
||||
|
||||
import signal
|
||||
import threading
|
||||
|
||||
# Some of the following imported functions are not used in this file, but are to
|
||||
# be used `_utils.signal_handling.XXXXX`.
|
||||
from torch._C import ( # noqa: F401
|
||||
_error_if_any_worker_fails,
|
||||
_remove_worker_pids,
|
||||
_set_worker_pids,
|
||||
_set_worker_signal_handlers,
|
||||
)
|
||||
|
||||
from . import IS_WINDOWS
|
||||
|
||||
|
||||
_SIGCHLD_handler_set = False
|
||||
r"""Whether SIGCHLD handler is set for DataLoader worker failures. Only one
|
||||
handler needs to be set for all DataLoaders in a process."""
|
||||
|
||||
|
||||
def _set_SIGCHLD_handler() -> None:
|
||||
# Windows doesn't support SIGCHLD handler
|
||||
if IS_WINDOWS:
|
||||
return
|
||||
# can't set signal in child threads
|
||||
if not isinstance(threading.current_thread(), threading._MainThread): # type: ignore[attr-defined]
|
||||
return
|
||||
global _SIGCHLD_handler_set
|
||||
if _SIGCHLD_handler_set:
|
||||
return
|
||||
previous_handler = signal.getsignal(signal.SIGCHLD)
|
||||
if not callable(previous_handler):
|
||||
# This doesn't catch default handler, but SIGCHLD default handler is a
|
||||
# no-op.
|
||||
previous_handler = None
|
||||
|
||||
def handler(signum, frame) -> None:
|
||||
# This following call uses `waitid` with WNOHANG from C side. Therefore,
|
||||
# Python can still get and update the process status successfully.
|
||||
_error_if_any_worker_fails()
|
||||
if previous_handler is not None:
|
||||
if not callable(previous_handler):
|
||||
raise AssertionError("previous_handler is not callable")
|
||||
previous_handler(signum, frame)
|
||||
|
||||
signal.signal(signal.SIGCHLD, handler)
|
||||
_SIGCHLD_handler_set = True
|
||||
@@ -0,0 +1,399 @@
|
||||
# mypy: allow-untyped-defs
|
||||
r"""Contains definitions of the methods used by the _BaseDataLoaderIter workers.
|
||||
|
||||
These **needs** to be in global scope since Py2 doesn't support serializing
|
||||
static methods.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import queue
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch._utils import ExceptionWrapper
|
||||
|
||||
from . import HAS_NUMPY, IS_WINDOWS, MP_STATUS_CHECK_INTERVAL, signal_handling
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
if IS_WINDOWS:
|
||||
import ctypes
|
||||
from ctypes.wintypes import BOOL, DWORD, HANDLE
|
||||
|
||||
# On Windows, the parent ID of the worker process remains unchanged when the manager process
|
||||
# is gone, and the only way to check it through OS is to let the worker have a process handle
|
||||
# of the manager and ask if the process status has changed.
|
||||
class ManagerWatchdog:
|
||||
def __init__(self) -> None:
|
||||
self.manager_pid = os.getppid()
|
||||
|
||||
# mypy cannot detect this code is windows only
|
||||
self.kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) # type: ignore[attr-defined]
|
||||
self.kernel32.OpenProcess.argtypes = (DWORD, BOOL, DWORD)
|
||||
self.kernel32.OpenProcess.restype = HANDLE
|
||||
self.kernel32.WaitForSingleObject.argtypes = (HANDLE, DWORD)
|
||||
self.kernel32.WaitForSingleObject.restype = DWORD
|
||||
|
||||
# Value obtained from https://msdn.microsoft.com/en-us/library/ms684880.aspx
|
||||
SYNCHRONIZE = 0x00100000
|
||||
self.manager_handle = self.kernel32.OpenProcess(
|
||||
SYNCHRONIZE, 0, self.manager_pid
|
||||
)
|
||||
|
||||
if not self.manager_handle:
|
||||
raise ctypes.WinError(ctypes.get_last_error()) # type: ignore[attr-defined]
|
||||
|
||||
self.manager_dead = False
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
if not self.manager_dead:
|
||||
# Value obtained from https://msdn.microsoft.com/en-us/library/windows/desktop/ms687032.aspx
|
||||
self.manager_dead = (
|
||||
self.kernel32.WaitForSingleObject(self.manager_handle, 0) == 0
|
||||
)
|
||||
return not self.manager_dead
|
||||
|
||||
else:
|
||||
|
||||
class ManagerWatchdog: # type: ignore[no-redef]
|
||||
def __init__(self) -> None:
|
||||
self.manager_pid = os.getppid()
|
||||
self.manager_dead = False
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
if not self.manager_dead:
|
||||
self.manager_dead = os.getppid() != self.manager_pid
|
||||
return not self.manager_dead
|
||||
|
||||
|
||||
_worker_info: WorkerInfo | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WorkerInfo:
|
||||
"""Information about the current DataLoader worker process or thread.
|
||||
|
||||
Attributes:
|
||||
id: The current worker id (0 to num_workers - 1)
|
||||
num_workers: Total number of workers
|
||||
seed: Random seed set for this worker
|
||||
dataset: Copy of the dataset object in this worker
|
||||
rng: Optional RNG state container. Defaults to None.
|
||||
worker_method: Optional worker method ("multiprocessing" or "thread"). Defaults to "multiprocessing".
|
||||
"""
|
||||
|
||||
id: int
|
||||
num_workers: int
|
||||
seed: int
|
||||
dataset: Dataset
|
||||
rng: _RNG | None = None
|
||||
worker_method: str | None = "multiprocessing"
|
||||
|
||||
|
||||
def get_worker_info() -> WorkerInfo | None:
|
||||
r"""Returns the information about the current
|
||||
:class:`~torch.utils.data.DataLoader` iterator worker process.
|
||||
|
||||
When called in a worker, this returns an object guaranteed to have the
|
||||
following attributes:
|
||||
|
||||
* :attr:`id`: the current worker id.
|
||||
* :attr:`num_workers`: the total number of workers.
|
||||
* :attr:`seed`: the random seed set for the current worker. This value is
|
||||
determined by main process RNG and the worker id. See
|
||||
:class:`~torch.utils.data.DataLoader`'s documentation for more details.
|
||||
* :attr:`dataset`: the copy of the dataset object in **this** process. Note
|
||||
that this will be a different object in a different process than the one
|
||||
in the main process.
|
||||
|
||||
When called in the main process, this returns ``None``.
|
||||
|
||||
.. note::
|
||||
When used in a :attr:`worker_init_fn` passed over to
|
||||
:class:`~torch.utils.data.DataLoader`, this method can be useful to
|
||||
set up each worker process differently, for instance, using ``worker_id``
|
||||
to configure the ``dataset`` object to only read a specific fraction of a
|
||||
sharded dataset, or use ``seed`` to seed other libraries used in dataset
|
||||
code.
|
||||
"""
|
||||
return _worker_info
|
||||
|
||||
|
||||
r"""Dummy class used to signal the end of an IterableDataset"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _IterableDatasetStopIteration:
|
||||
worker_id: int
|
||||
|
||||
|
||||
r"""Dummy class used to resume the fetching when worker reuse is enabled"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ResumeIteration:
|
||||
seed: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RNG:
|
||||
"""Container for thread-local random number generator state.
|
||||
|
||||
Used by thread workers to maintain separate RNG state per worker thread
|
||||
to avoid race conditions.
|
||||
|
||||
Attributes:
|
||||
random_generator: Python random.Random generator for this thread
|
||||
torch_generator: PyTorch Generator for this thread
|
||||
numpy_generator: NumPy Generator for this thread (None if numpy not available)
|
||||
"""
|
||||
|
||||
random_generator: random.Random
|
||||
torch_generator: torch.Generator
|
||||
numpy_generator: object | None = None
|
||||
|
||||
|
||||
# The function `_generate_state` is adapted from `numpy.random.SeedSequence`
|
||||
# from https://github.com/numpy/numpy/blob/main/numpy/random/bit_generator.pyx
|
||||
# It's MIT licensed, here is the copyright:
|
||||
|
||||
# Copyright (c) 2015 Melissa E. O'Neill
|
||||
# Copyright (c) 2019 NumPy Developers
|
||||
#
|
||||
# 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.
|
||||
|
||||
|
||||
# This function generates an array of int32 as the seed for
|
||||
# `numpy.random`, in order to prevent state collision due to same
|
||||
# seed and algorithm for `numpy.random` and `random` modules.
|
||||
# TODO: Implement `SeedSequence` like object for `torch.random`
|
||||
def _generate_state(base_seed, worker_id):
|
||||
INIT_A = 0x43B0D7E5
|
||||
MULT_A = 0x931E8875
|
||||
INIT_B = 0x8B51F9DD
|
||||
MULT_B = 0x58F38DED
|
||||
MIX_MULT_L = 0xCA01F9DD
|
||||
MIX_MULT_R = 0x4973F715
|
||||
XSHIFT = 4 * 8 // 2
|
||||
MASK32 = 0xFFFFFFFF
|
||||
|
||||
entropy = [worker_id, base_seed & MASK32, base_seed >> 32, 0]
|
||||
pool = [0] * 4
|
||||
|
||||
hash_const_A = INIT_A
|
||||
|
||||
def hash(value):
|
||||
nonlocal hash_const_A
|
||||
value = (value ^ hash_const_A) & MASK32
|
||||
hash_const_A = (hash_const_A * MULT_A) & MASK32
|
||||
value = (value * hash_const_A) & MASK32
|
||||
value = (value ^ (value >> XSHIFT)) & MASK32
|
||||
return value
|
||||
|
||||
def mix(x, y):
|
||||
result_x = (MIX_MULT_L * x) & MASK32
|
||||
result_y = (MIX_MULT_R * y) & MASK32
|
||||
result = (result_x - result_y) & MASK32
|
||||
result = (result ^ (result >> XSHIFT)) & MASK32
|
||||
return result
|
||||
|
||||
# Add in the entropy to the pool.
|
||||
for i in range(len(pool)):
|
||||
pool[i] = hash(entropy[i])
|
||||
|
||||
# Mix all bits together so late bits can affect earlier bits.
|
||||
for i_src in range(len(pool)):
|
||||
for i_dst in range(len(pool)):
|
||||
if i_src != i_dst:
|
||||
pool[i_dst] = mix(pool[i_dst], hash(pool[i_src]))
|
||||
|
||||
hash_const_B = INIT_B
|
||||
state = []
|
||||
for i_dst in range(4):
|
||||
data_val = pool[i_dst]
|
||||
data_val = (data_val ^ hash_const_B) & MASK32
|
||||
hash_const_B = (hash_const_B * MULT_B) & MASK32
|
||||
data_val = (data_val * hash_const_B) & MASK32
|
||||
data_val = (data_val ^ (data_val >> XSHIFT)) & MASK32
|
||||
state.append(data_val)
|
||||
return state
|
||||
|
||||
|
||||
def _worker_loop(
|
||||
dataset_kind,
|
||||
dataset,
|
||||
index_queue,
|
||||
data_queue,
|
||||
done_event,
|
||||
auto_collation,
|
||||
collate_fn,
|
||||
drop_last,
|
||||
base_seed,
|
||||
init_fn,
|
||||
worker_id,
|
||||
num_workers,
|
||||
persistent_workers,
|
||||
shared_seed,
|
||||
) -> None:
|
||||
# See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the
|
||||
# logic of this function.
|
||||
|
||||
try:
|
||||
# Initialize C side signal handlers for SIGBUS and SIGSEGV. Python signal
|
||||
# module's handlers are executed after Python returns from C low-level
|
||||
# handlers, likely when the same fatal signal had already happened
|
||||
# again.
|
||||
# https://docs.python.org/3/library/signal.html#execution-of-python-signal-handlers
|
||||
signal_handling._set_worker_signal_handlers()
|
||||
|
||||
torch.multiprocessing._set_thread_name("pt_data_worker")
|
||||
|
||||
torch.set_num_threads(1)
|
||||
seed = base_seed + worker_id
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if HAS_NUMPY:
|
||||
np_seed = _generate_state(base_seed, worker_id)
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(np_seed)
|
||||
|
||||
from torch.utils.data import IterDataPipe
|
||||
from torch.utils.data.graph_settings import apply_random_seed
|
||||
|
||||
shared_rng = torch.Generator()
|
||||
if isinstance(dataset, IterDataPipe):
|
||||
if shared_seed is None:
|
||||
raise AssertionError(
|
||||
"shared_seed must be provided for IterDataPipe workers"
|
||||
)
|
||||
shared_rng.manual_seed(shared_seed)
|
||||
dataset = apply_random_seed(dataset, shared_rng)
|
||||
|
||||
global _worker_info
|
||||
_worker_info = WorkerInfo(
|
||||
id=worker_id, num_workers=num_workers, seed=seed, dataset=dataset
|
||||
)
|
||||
|
||||
from torch.utils.data import _DatasetKind
|
||||
|
||||
init_exception = None
|
||||
|
||||
try:
|
||||
if init_fn is not None:
|
||||
init_fn(worker_id)
|
||||
|
||||
fetcher = _DatasetKind.create_fetcher(
|
||||
dataset_kind, dataset, auto_collation, collate_fn, drop_last
|
||||
)
|
||||
except Exception:
|
||||
init_exception = ExceptionWrapper(
|
||||
where=f"in DataLoader worker process {worker_id}"
|
||||
)
|
||||
|
||||
# When using Iterable mode, some worker can exit earlier than others due
|
||||
# to the IterableDataset behaving differently for different workers.
|
||||
# When such things happen, an `_IterableDatasetStopIteration` object is
|
||||
# sent over to the main process with the ID of this worker, so that the
|
||||
# main process won't send more tasks to this worker, and will send
|
||||
# `None` to this worker to properly exit it.
|
||||
#
|
||||
# Note that we cannot set `done_event` from a worker as it is shared
|
||||
# among all processes. Instead, we set the `iteration_end` flag to
|
||||
# signify that the iterator is exhausted. When either `done_event` or
|
||||
# `iteration_end` is set, we skip all processing step and just wait for
|
||||
# `None`.
|
||||
iteration_end = False
|
||||
|
||||
watchdog = ManagerWatchdog()
|
||||
|
||||
while watchdog.is_alive():
|
||||
try:
|
||||
r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
|
||||
except queue.Empty:
|
||||
continue
|
||||
if isinstance(r, _ResumeIteration):
|
||||
# Acknowledge the main process
|
||||
data_queue.put((r, None))
|
||||
iteration_end = False
|
||||
|
||||
if isinstance(dataset, IterDataPipe):
|
||||
if r.seed is None:
|
||||
raise AssertionError(
|
||||
"resume iteration seed is None for IterDataPipe"
|
||||
)
|
||||
shared_rng.manual_seed(r.seed)
|
||||
dataset = apply_random_seed(dataset, shared_rng)
|
||||
|
||||
# Recreate the fetcher for worker-reuse policy
|
||||
fetcher = _DatasetKind.create_fetcher(
|
||||
dataset_kind, dataset, auto_collation, collate_fn, drop_last
|
||||
)
|
||||
continue
|
||||
elif r is None:
|
||||
# Received the final signal
|
||||
if not done_event.is_set() and not iteration_end:
|
||||
raise AssertionError(
|
||||
"Received final signal but neither done_event nor iteration_end is set"
|
||||
)
|
||||
break
|
||||
elif done_event.is_set() or iteration_end:
|
||||
# `done_event` is set. But I haven't received the final signal
|
||||
# (None) yet. I will keep continuing until get it, and skip the
|
||||
# processing steps.
|
||||
continue
|
||||
idx, index = r
|
||||
data: _IterableDatasetStopIteration | ExceptionWrapper
|
||||
if init_exception is not None:
|
||||
data = init_exception
|
||||
init_exception = None
|
||||
else:
|
||||
try:
|
||||
data = fetcher.fetch(index) # type: ignore[possibly-undefined]
|
||||
except Exception as e:
|
||||
if (
|
||||
isinstance(e, StopIteration)
|
||||
and dataset_kind == _DatasetKind.Iterable
|
||||
):
|
||||
data = _IterableDatasetStopIteration(worker_id)
|
||||
# Set `iteration_end`
|
||||
# (1) to save future `next(...)` calls, and
|
||||
# (2) to avoid sending multiple `_IterableDatasetStopIteration`s.
|
||||
iteration_end = True
|
||||
else:
|
||||
# It is important that we don't store exc_info in a variable.
|
||||
# `ExceptionWrapper` does the correct thing.
|
||||
# See NOTE [ Python Traceback Reference Cycle Problem ]
|
||||
data = ExceptionWrapper(
|
||||
where=f"in DataLoader worker process {worker_id}"
|
||||
)
|
||||
data_queue.put((idx, data))
|
||||
del data, idx, index, r # save memory
|
||||
except KeyboardInterrupt:
|
||||
# Main process will raise KeyboardInterrupt anyways.
|
||||
pass
|
||||
if done_event.is_set():
|
||||
data_queue.cancel_join_thread()
|
||||
data_queue.close()
|
||||
@@ -0,0 +1,11 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing_extensions import deprecated as _deprecated
|
||||
|
||||
|
||||
@_deprecated(
|
||||
"Usage of `backward_compatibility.worker_init_fn` is deprecated "
|
||||
"as `DataLoader` automatically applies sharding in every worker",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def worker_init_fn(worker_id) -> None:
|
||||
pass
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
from torch.utils.data.datapipes import dataframe as dataframe, iter as iter, map as map
|
||||
@@ -0,0 +1,213 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import Any, get_type_hints
|
||||
|
||||
from torch.utils.data.datapipes._typing import _DataPipeMeta
|
||||
from torch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe
|
||||
|
||||
|
||||
######################################################
|
||||
# Functional API
|
||||
######################################################
|
||||
class functional_datapipe:
|
||||
name: str
|
||||
|
||||
def __init__(self, name: str, enable_df_api_tracing=False) -> None:
|
||||
"""
|
||||
Define a functional datapipe.
|
||||
|
||||
Args:
|
||||
enable_df_api_tracing - if set, any returned DataPipe would accept
|
||||
DataFrames API in tracing mode.
|
||||
"""
|
||||
self.name = name
|
||||
self.enable_df_api_tracing = enable_df_api_tracing
|
||||
|
||||
def __call__(self, cls):
|
||||
if issubclass(cls, IterDataPipe):
|
||||
if isinstance(cls, type): # type: ignore[arg-type]
|
||||
if not isinstance(cls, _DataPipeMeta):
|
||||
raise TypeError(
|
||||
"`functional_datapipe` can only decorate IterDataPipe"
|
||||
)
|
||||
# with non_deterministic decorator
|
||||
else:
|
||||
if not isinstance(cls, non_deterministic) and not (
|
||||
hasattr(cls, "__self__")
|
||||
and isinstance(cls.__self__, non_deterministic)
|
||||
):
|
||||
raise TypeError(
|
||||
"`functional_datapipe` can only decorate IterDataPipe"
|
||||
)
|
||||
IterDataPipe.register_datapipe_as_function(
|
||||
self.name, cls, enable_df_api_tracing=self.enable_df_api_tracing
|
||||
)
|
||||
elif issubclass(cls, MapDataPipe):
|
||||
MapDataPipe.register_datapipe_as_function(self.name, cls)
|
||||
|
||||
return cls
|
||||
|
||||
|
||||
######################################################
|
||||
# Determinism
|
||||
######################################################
|
||||
_determinism: bool = False
|
||||
|
||||
|
||||
class guaranteed_datapipes_determinism:
|
||||
prev: bool
|
||||
|
||||
def __init__(self) -> None:
|
||||
global _determinism
|
||||
self.prev = _determinism
|
||||
_determinism = True
|
||||
|
||||
def __enter__(self) -> None:
|
||||
pass
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
||||
global _determinism
|
||||
_determinism = self.prev
|
||||
|
||||
|
||||
class non_deterministic:
|
||||
cls: type[IterDataPipe] | None = None
|
||||
# TODO: Lambda for picking
|
||||
deterministic_fn: Callable[..., bool]
|
||||
|
||||
def __init__(self, arg: type[IterDataPipe] | Callable[..., bool]) -> None:
|
||||
# 1. Decorator doesn't have any argument
|
||||
if isinstance(arg, type): # type: ignore[arg-type]
|
||||
if not issubclass(arg, IterDataPipe): # type: ignore[arg-type]
|
||||
raise TypeError(
|
||||
"Only `IterDataPipe` can be decorated with `non_deterministic`"
|
||||
f", but {arg.__name__} is found"
|
||||
)
|
||||
self.cls = arg # type: ignore[assignment]
|
||||
# 2. Decorator has an argument of a function
|
||||
# This class should behave differently given different inputs. Use this
|
||||
# function to verify the determinism for each instance.
|
||||
# When the function returns True, the instance is non-deterministic. Otherwise,
|
||||
# the instance is a deterministic DataPipe.
|
||||
elif isinstance(arg, Callable): # type:ignore[arg-type]
|
||||
self.deterministic_fn = arg
|
||||
else:
|
||||
raise TypeError(f"{arg} can not be decorated by non_deterministic")
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
global _determinism
|
||||
# Decorate IterDataPipe
|
||||
if self.cls is not None:
|
||||
if _determinism:
|
||||
raise TypeError(
|
||||
f"{self.cls.__name__} is non-deterministic, but you set 'guaranteed_datapipes_determinism'. "
|
||||
"You can turn off determinism for this DataPipe if that is acceptable "
|
||||
"for your application"
|
||||
)
|
||||
return self.cls(*args, **kwargs) # type: ignore[call-arg]
|
||||
|
||||
# Decorate with a functional argument
|
||||
if not (
|
||||
isinstance(args[0], type) and issubclass(args[0], IterDataPipe) # type: ignore[arg-type]
|
||||
):
|
||||
raise TypeError(
|
||||
f"Only `IterDataPipe` can be decorated, but {args[0].__name__} is found"
|
||||
)
|
||||
self.cls = args[0]
|
||||
return self.deterministic_wrapper_fn
|
||||
|
||||
def deterministic_wrapper_fn(self, *args, **kwargs) -> IterDataPipe:
|
||||
res = self.deterministic_fn(*args, **kwargs)
|
||||
if not isinstance(res, bool):
|
||||
raise TypeError(
|
||||
"deterministic_fn of `non_deterministic` decorator is required "
|
||||
f"to return a boolean value, but {type(res)} is found"
|
||||
)
|
||||
global _determinism
|
||||
if _determinism and res:
|
||||
raise TypeError(
|
||||
f"{self.cls.__name__} is non-deterministic with the inputs, but you set " # type: ignore[union-attr]
|
||||
"'guaranteed_datapipes_determinism'. You can turn off determinism "
|
||||
"for this DataPipe if that is acceptable for your application"
|
||||
)
|
||||
return self.cls(*args, **kwargs) # type: ignore[call-arg, misc]
|
||||
|
||||
|
||||
######################################################
|
||||
# Type validation
|
||||
######################################################
|
||||
# Validate each argument of DataPipe with hint as a subtype of the hint.
|
||||
def argument_validation(f):
|
||||
signature = inspect.signature(f)
|
||||
hints = get_type_hints(f)
|
||||
|
||||
@wraps(f)
|
||||
def wrapper(*args, **kwargs):
|
||||
bound = signature.bind(*args, **kwargs)
|
||||
for argument_name, value in bound.arguments.items():
|
||||
if argument_name in hints and isinstance(
|
||||
hints[argument_name], _DataPipeMeta
|
||||
):
|
||||
hint = hints[argument_name]
|
||||
if not isinstance(value, IterDataPipe):
|
||||
raise TypeError(
|
||||
f"Expected argument '{argument_name}' as a IterDataPipe, but found {type(value)}"
|
||||
)
|
||||
if not value.type.issubtype(hint.type):
|
||||
raise TypeError(
|
||||
f"Expected type of argument '{argument_name}' as a subtype of "
|
||||
f"hint {hint.type}, but found {value.type}"
|
||||
)
|
||||
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# Default value is True
|
||||
_runtime_validation_enabled: bool = True
|
||||
|
||||
|
||||
class runtime_validation_disabled:
|
||||
prev: bool
|
||||
|
||||
def __init__(self) -> None:
|
||||
global _runtime_validation_enabled
|
||||
self.prev = _runtime_validation_enabled
|
||||
_runtime_validation_enabled = False
|
||||
|
||||
def __enter__(self) -> None:
|
||||
pass
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
||||
global _runtime_validation_enabled
|
||||
_runtime_validation_enabled = self.prev
|
||||
|
||||
|
||||
# Runtime checking
|
||||
# Validate output data is subtype of return hint
|
||||
def runtime_validation(f):
|
||||
# TODO:
|
||||
# Can be extended to validate '__getitem__' and nonblocking
|
||||
if f.__name__ != "__iter__":
|
||||
raise TypeError(
|
||||
f"Can not decorate function {f.__name__} with 'runtime_validation'"
|
||||
)
|
||||
|
||||
@wraps(f)
|
||||
def wrapper(self):
|
||||
global _runtime_validation_enabled
|
||||
if not _runtime_validation_enabled:
|
||||
yield from f(self)
|
||||
else:
|
||||
it = f(self)
|
||||
for d in it:
|
||||
if not self.type.issubtype_of_instance(d):
|
||||
raise RuntimeError(
|
||||
f"Expected an instance as subtype of {self.type}, but found {d}({type(d)})"
|
||||
)
|
||||
yield d
|
||||
|
||||
return wrapper
|
||||
+279
@@ -0,0 +1,279 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import inspect
|
||||
from enum import Enum
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class _SnapshotState(Enum):
|
||||
r"""
|
||||
These are the snapshotting-related states that IterDataPipes can be in.
|
||||
|
||||
`NotStarted` - allows you to restore a snapshot and create an iterator with reset
|
||||
`Restored` - cannot restore again, allows you to create an iterator without resetting the DataPipe
|
||||
`Iterating` - can restore, will reset if you create a new iterator
|
||||
"""
|
||||
|
||||
NotStarted = 0
|
||||
Restored = 1
|
||||
Iterating = 2
|
||||
|
||||
|
||||
def _simplify_obj_name(obj) -> str:
|
||||
"""Simplify the display strings of objects for the purpose of rendering within DataPipe error messages."""
|
||||
if inspect.isfunction(obj):
|
||||
return obj.__name__
|
||||
else:
|
||||
return repr(obj)
|
||||
|
||||
|
||||
def _strip_datapipe_from_name(name: str) -> str:
|
||||
return name.replace("IterDataPipe", "").replace("MapDataPipe", "")
|
||||
|
||||
|
||||
def _generate_input_args_string(obj):
|
||||
"""Generate a string for the input arguments of an object."""
|
||||
signature = inspect.signature(obj.__class__)
|
||||
input_param_names = set(signature.parameters.keys())
|
||||
result = []
|
||||
for name, value in inspect.getmembers(obj):
|
||||
if name in input_param_names:
|
||||
result.append((name, _simplify_obj_name(value)))
|
||||
return ", ".join([f"{name}={value}" for name, value in result])
|
||||
|
||||
|
||||
def _generate_iterdatapipe_msg(datapipe, simplify_dp_name: bool = False):
|
||||
output_string = (
|
||||
f"{datapipe.__class__.__name__}({_generate_input_args_string(datapipe)})"
|
||||
)
|
||||
if simplify_dp_name:
|
||||
output_string = _strip_datapipe_from_name(output_string)
|
||||
return output_string
|
||||
|
||||
|
||||
def _gen_invalid_iterdatapipe_msg(datapipe) -> str:
|
||||
return (
|
||||
"This iterator has been invalidated because another iterator has been created "
|
||||
f"from the same IterDataPipe: {_generate_iterdatapipe_msg(datapipe)}\n"
|
||||
"This may be caused multiple references to the same IterDataPipe. We recommend "
|
||||
"using `.fork()` if that is necessary."
|
||||
)
|
||||
|
||||
|
||||
_feedback_msg = (
|
||||
"\nFor feedback regarding this single iterator per IterDataPipe constraint, feel free "
|
||||
"to comment on this issue: https://github.com/pytorch/data/issues/45."
|
||||
)
|
||||
|
||||
|
||||
def _check_iterator_valid(datapipe, iterator_id, next_method_exists=False) -> None:
|
||||
r"""
|
||||
Given an instance of a DataPipe and an iterator ID, check if the IDs match, and if not, raises an exception.
|
||||
|
||||
In the case of ChildDataPipe, the ID gets compared to the one stored in `main_datapipe` as well.
|
||||
"""
|
||||
if next_method_exists:
|
||||
# This is the case where `IterDataPipe` has both `__iter__` and `__next__`.
|
||||
# The `_valid_iterator_id` should either be never set (`None`), or set by at most one
|
||||
# iterator (`0`). Otherwise, it means there are multiple iterators.
|
||||
if datapipe._valid_iterator_id is not None and datapipe._valid_iterator_id != 0:
|
||||
extra_msg = "\nNote that this exception is raised inside your IterDataPipe's a `__next__` method"
|
||||
raise RuntimeError(
|
||||
_gen_invalid_iterdatapipe_msg(datapipe) + extra_msg + _feedback_msg
|
||||
)
|
||||
elif (
|
||||
hasattr(datapipe, "_is_child_datapipe") and datapipe._is_child_datapipe is True
|
||||
):
|
||||
if hasattr(datapipe, "_check_valid_iterator_id"):
|
||||
if not datapipe._check_valid_iterator_id(iterator_id):
|
||||
raise RuntimeError(
|
||||
"This iterator has been invalidated, because a new iterator has been created "
|
||||
f"from one of the ChildDataPipes of "
|
||||
f"{_generate_iterdatapipe_msg(datapipe.main_datapipe)}."
|
||||
+ _feedback_msg
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"ChildDataPipe must have method `_check_valid_iterator_id`."
|
||||
)
|
||||
elif datapipe._valid_iterator_id != iterator_id:
|
||||
raise RuntimeError(_gen_invalid_iterdatapipe_msg(datapipe) + _feedback_msg)
|
||||
|
||||
|
||||
def _set_datapipe_valid_iterator_id(datapipe):
|
||||
"""Given a DataPipe, updates its valid iterator ID and reset the DataPipe."""
|
||||
if hasattr(datapipe, "_is_child_datapipe") and datapipe._is_child_datapipe is True:
|
||||
if hasattr(datapipe, "_set_main_datapipe_valid_iterator_id"):
|
||||
datapipe._set_main_datapipe_valid_iterator_id() # reset() is called within this method when appropriate
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"ChildDataPipe must have method `_set_main_datapipe_valid_iterator_id`."
|
||||
)
|
||||
else:
|
||||
if datapipe._valid_iterator_id is None:
|
||||
datapipe._valid_iterator_id = 0
|
||||
else:
|
||||
datapipe._valid_iterator_id += 1
|
||||
datapipe.reset()
|
||||
return datapipe._valid_iterator_id
|
||||
|
||||
|
||||
def hook_iterator(namespace) -> None:
|
||||
r"""
|
||||
Define a hook that is applied to all `__iter__` of metaclass `_DataPipeMeta`.
|
||||
|
||||
This is done for the purpose of profiling and checking if an iterator is still valid.
|
||||
"""
|
||||
|
||||
def profiler_record_fn_context(datapipe):
|
||||
if not hasattr(datapipe, "_profile_name"):
|
||||
datapipe._profile_name = _generate_iterdatapipe_msg(
|
||||
datapipe, simplify_dp_name=True
|
||||
)
|
||||
return torch.autograd.profiler.record_function(datapipe._profile_name)
|
||||
|
||||
class IteratorDecorator:
|
||||
r"""
|
||||
Wrap the iterator and modifying its `__next__` method.
|
||||
|
||||
This decorator is applied to DataPipes of which `__iter__` method is NOT a generator function.
|
||||
Those `__iter__` method commonly returns `self` but not necessarily.
|
||||
"""
|
||||
|
||||
def __init__(self, iterator, datapipe, iterator_id, has_next_method) -> None:
|
||||
self.iterator = iterator
|
||||
self.datapipe = datapipe
|
||||
self.iterator_id = iterator_id
|
||||
self._profiler_enabled = torch.autograd._profiler_enabled()
|
||||
# Check if `__iter__` returns `self` and `DataPipe` has `__next__`
|
||||
self.self_and_has_next_method = (
|
||||
self.iterator is self.datapipe and has_next_method
|
||||
)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def _get_next(self):
|
||||
"""Return next with logic related to iterator validity, profiler, and incrementation of samples yielded."""
|
||||
_check_iterator_valid(self.datapipe, self.iterator_id)
|
||||
result = next(self.iterator)
|
||||
if not self.self_and_has_next_method:
|
||||
self.datapipe._number_of_samples_yielded += 1
|
||||
return result
|
||||
|
||||
def __next__(self):
|
||||
# TODO: Add try-except to in-place reduce traceback from the Exception
|
||||
# See: https://github.com/pytorch/data/issues/284
|
||||
if self._profiler_enabled:
|
||||
with profiler_record_fn_context(self.datapipe):
|
||||
return self._get_next()
|
||||
else: # Decided against using `contextlib.nullcontext` for performance reasons
|
||||
return self._get_next()
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.iterator, name)
|
||||
|
||||
func = namespace["__iter__"]
|
||||
|
||||
# ``__iter__`` of IterDataPipe is a generator function
|
||||
if inspect.isgeneratorfunction(func):
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrap_generator(*args, **kwargs):
|
||||
gen = func(*args, **kwargs)
|
||||
datapipe = args[0]
|
||||
if datapipe._fast_forward_iterator:
|
||||
it = datapipe._fast_forward_iterator
|
||||
datapipe._fast_forward_iterator = None
|
||||
datapipe._snapshot_state = _SnapshotState.Iterating
|
||||
while True:
|
||||
try:
|
||||
yield next(it)
|
||||
except StopIteration:
|
||||
return
|
||||
iterator_id = _set_datapipe_valid_iterator_id(
|
||||
datapipe
|
||||
) # This ID is tied to each created iterator
|
||||
_profiler_enabled = torch.autograd._profiler_enabled()
|
||||
try:
|
||||
if _profiler_enabled:
|
||||
with profiler_record_fn_context(datapipe):
|
||||
response = gen.send(None)
|
||||
else:
|
||||
response = gen.send(None)
|
||||
|
||||
while True:
|
||||
datapipe._number_of_samples_yielded += 1
|
||||
request = yield response
|
||||
# Pass through here every time `__next__` is called
|
||||
if _profiler_enabled:
|
||||
with profiler_record_fn_context(datapipe):
|
||||
_check_iterator_valid(datapipe, iterator_id)
|
||||
response = gen.send(request)
|
||||
else: # Decided against using `contextlib.nullcontext` for performance reasons
|
||||
_check_iterator_valid(datapipe, iterator_id)
|
||||
response = gen.send(request)
|
||||
except StopIteration:
|
||||
return
|
||||
except Exception as e:
|
||||
# TODO: Simplify the traceback message to skip over `response = gen.send(None)`
|
||||
# Part of https://github.com/pytorch/data/issues/284
|
||||
datapipe = args[0]
|
||||
msg = "thrown by __iter__ of"
|
||||
single_iterator_msg = "single iterator per IterDataPipe constraint"
|
||||
if hasattr(e.args, "__len__"):
|
||||
full_msg = f"{msg} {datapipe.__class__.__name__}({_generate_input_args_string(datapipe)})"
|
||||
if len(e.args) == 0 or not isinstance(
|
||||
e.args[0], str
|
||||
): # If an exception message doesn't exist
|
||||
e.args = (f"\nThis exception is {full_msg}",)
|
||||
elif msg not in e.args[0] and single_iterator_msg not in e.args[0]:
|
||||
e.args = (
|
||||
e.args[0] + f"\nThis exception is {full_msg}",
|
||||
) + e.args[1:]
|
||||
raise
|
||||
|
||||
namespace["__iter__"] = wrap_generator
|
||||
else: # ``__iter__`` of IterDataPipe is NOT a generator function
|
||||
# IterDataPipe is an iterator with both ``__iter__`` and ``__next__``
|
||||
# And ``__iter__`` may or may not return `self`
|
||||
if "__next__" in namespace: # If `__next__` exists, put a wrapper around it
|
||||
next_func = namespace["__next__"]
|
||||
|
||||
@functools.wraps(next_func)
|
||||
def wrap_next(*args, **kwargs):
|
||||
datapipe = args[0]
|
||||
if torch.autograd._profiler_enabled():
|
||||
with profiler_record_fn_context(datapipe):
|
||||
result = next_func(*args, **kwargs)
|
||||
else:
|
||||
result = next_func(*args, **kwargs)
|
||||
datapipe._number_of_samples_yielded += 1
|
||||
return result
|
||||
|
||||
namespace["__next__"] = wrap_next
|
||||
|
||||
# Note that if the `__next__` and `__iter__` do something completely unrelated. It may cause issue but
|
||||
# the user will be violating the iterator protocol. Potential issue:
|
||||
# 1. Valid iterator ID may not update or checked properly
|
||||
# 2. The number of samples yielded will be miscounted
|
||||
|
||||
# Regardless if `__next__` exists or not, `__iter__` needs a wrapper to track the number of valid iterators
|
||||
@functools.wraps(func)
|
||||
def wrap_iter(*args, **kwargs):
|
||||
iter_ret = func(*args, **kwargs)
|
||||
datapipe = args[0]
|
||||
datapipe._snapshot_state = _SnapshotState.Iterating
|
||||
if datapipe._fast_forward_iterator:
|
||||
iter_ret = datapipe._fast_forward_iterator
|
||||
datapipe._fast_forward_iterator = None
|
||||
return iter_ret
|
||||
iterator_id = _set_datapipe_valid_iterator_id(
|
||||
datapipe
|
||||
) # This ID is tied to each created iterator
|
||||
return IteratorDecorator(
|
||||
iter_ret, datapipe, iterator_id, "__next__" in namespace
|
||||
)
|
||||
|
||||
namespace["__iter__"] = wrap_iter
|
||||
@@ -0,0 +1,483 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Taking reference from official Python typing
|
||||
# https://github.com/python/cpython/blob/master/Lib/typing.py
|
||||
|
||||
import collections
|
||||
import functools
|
||||
import numbers
|
||||
import sys
|
||||
|
||||
# Please check [Note: TypeMeta and TypeAlias]
|
||||
# In case of metaclass conflict due to ABCMeta or _ProtocolMeta
|
||||
# For Python 3.9, only Protocol in typing uses metaclass
|
||||
from abc import ABCMeta
|
||||
from collections.abc import Iterator
|
||||
|
||||
# TODO: Use TypeAlias when Python 3.6 is deprecated
|
||||
from typing import (
|
||||
_eval_type, # pyrefly: ignore [missing-module-attribute]
|
||||
_GenericAlias, # pyrefly: ignore [missing-module-attribute]
|
||||
_tp_cache, # pyrefly: ignore [missing-module-attribute]
|
||||
_type_check, # pyrefly: ignore [missing-module-attribute]
|
||||
_type_repr,
|
||||
Any,
|
||||
ForwardRef,
|
||||
Generic,
|
||||
get_type_hints,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
from torch.utils.data.datapipes._hook_iterator import _SnapshotState, hook_iterator
|
||||
|
||||
|
||||
class GenericMeta(ABCMeta): # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
|
||||
class Integer(numbers.Integral):
|
||||
pass
|
||||
|
||||
|
||||
class Boolean(numbers.Integral):
|
||||
pass
|
||||
|
||||
|
||||
# Python 'type' object is not subscriptable
|
||||
# Tuple[int, List, dict] -> valid
|
||||
# tuple[int, list, dict] -> invalid
|
||||
# Map Python 'type' to abstract base class
|
||||
TYPE2ABC = {
|
||||
bool: Boolean,
|
||||
int: Integer,
|
||||
float: numbers.Real,
|
||||
complex: numbers.Complex,
|
||||
dict: dict,
|
||||
list: list,
|
||||
set: set,
|
||||
tuple: tuple,
|
||||
None: type(None),
|
||||
}
|
||||
|
||||
|
||||
def issubtype(left, right, recursive=True):
|
||||
r"""
|
||||
Check if the left-side type is a subtype of the right-side type.
|
||||
|
||||
If any of type is a composite type like `Union` and `TypeVar` with
|
||||
bounds, it would be expanded into a list of types and check all
|
||||
of left-side types are subtypes of either one from right-side types.
|
||||
"""
|
||||
left = TYPE2ABC.get(left, left)
|
||||
right = TYPE2ABC.get(right, right)
|
||||
|
||||
if right is Any or left == right:
|
||||
return True
|
||||
|
||||
if isinstance(right, _GenericAlias):
|
||||
if getattr(right, "__origin__", None) is Generic:
|
||||
return True
|
||||
|
||||
if right is type(None):
|
||||
return False
|
||||
|
||||
# Right-side type
|
||||
constraints = _decompose_type(right)
|
||||
|
||||
if len(constraints) == 0 or Any in constraints:
|
||||
return True
|
||||
|
||||
if left is Any:
|
||||
return False
|
||||
|
||||
# Left-side type
|
||||
variants = _decompose_type(left)
|
||||
|
||||
# all() will return True for empty variants
|
||||
if len(variants) == 0:
|
||||
return False
|
||||
|
||||
return all(
|
||||
_issubtype_with_constraints(variant, constraints, recursive)
|
||||
for variant in variants
|
||||
)
|
||||
|
||||
|
||||
def _decompose_type(t, to_list=True):
|
||||
if isinstance(t, TypeVar):
|
||||
if t.__bound__ is not None:
|
||||
ts = [t.__bound__]
|
||||
else:
|
||||
# For T_co, __constraints__ is ()
|
||||
ts = list(t.__constraints__)
|
||||
elif hasattr(t, "__origin__") and t.__origin__ == Union:
|
||||
ts = t.__args__
|
||||
else:
|
||||
if not to_list:
|
||||
return None
|
||||
ts = [t]
|
||||
# Ignored: Generator has incompatible item type "object"; expected "Type[Any]"
|
||||
ts = [TYPE2ABC.get(_t, _t) for _t in ts] # type: ignore[misc]
|
||||
return ts
|
||||
|
||||
|
||||
def _issubtype_with_constraints(variant, constraints, recursive=True):
|
||||
r"""
|
||||
Check if the variant is a subtype of either one from constraints.
|
||||
|
||||
For composite types like `Union` and `TypeVar` with bounds, they
|
||||
would be expanded for testing.
|
||||
"""
|
||||
if variant in constraints:
|
||||
return True
|
||||
|
||||
# [Note: Subtype for Union and TypeVar]
|
||||
# Python typing is able to flatten Union[Union[...]] or Union[TypeVar].
|
||||
# But it couldn't flatten the following scenarios:
|
||||
# - Union[int, TypeVar[Union[...]]]
|
||||
# - TypeVar[TypeVar[...]]
|
||||
# So, variant and each constraint may be a TypeVar or a Union.
|
||||
# In these cases, all of inner types from the variant are required to be
|
||||
# extracted and verified as a subtype of any constraint. And, all of
|
||||
# inner types from any constraint being a TypeVar or a Union are
|
||||
# also required to be extracted and verified if the variant belongs to
|
||||
# any of them.
|
||||
|
||||
# Variant
|
||||
vs = _decompose_type(variant, to_list=False)
|
||||
|
||||
# Variant is TypeVar or Union
|
||||
if vs is not None:
|
||||
return all(_issubtype_with_constraints(v, constraints, recursive) for v in vs)
|
||||
|
||||
# Variant is not TypeVar or Union
|
||||
if hasattr(variant, "__origin__") and variant.__origin__ is not None:
|
||||
v_origin = variant.__origin__
|
||||
# In Python-3.9 typing library untyped generics do not have args
|
||||
v_args = getattr(variant, "__args__", None)
|
||||
else:
|
||||
v_origin = variant
|
||||
v_args = None
|
||||
|
||||
# Constraints
|
||||
for constraint in constraints:
|
||||
cs = _decompose_type(constraint, to_list=False)
|
||||
|
||||
# Constraint is TypeVar or Union
|
||||
if cs is not None:
|
||||
if _issubtype_with_constraints(variant, cs, recursive):
|
||||
return True
|
||||
# Constraint is not TypeVar or Union
|
||||
else:
|
||||
# __origin__ can be None for plain list, tuple, ... in Python 3.6
|
||||
if hasattr(constraint, "__origin__") and constraint.__origin__ is not None:
|
||||
c_origin = constraint.__origin__
|
||||
if v_origin == c_origin:
|
||||
if not recursive:
|
||||
return True
|
||||
# In Python-3.9 typing library untyped generics do not have args
|
||||
c_args = getattr(constraint, "__args__", None)
|
||||
if c_args is None or len(c_args) == 0:
|
||||
return True
|
||||
if (
|
||||
v_args is not None
|
||||
and len(v_args) == len(c_args)
|
||||
and all(
|
||||
issubtype(v_arg, c_arg)
|
||||
for v_arg, c_arg in zip(v_args, c_args, strict=True)
|
||||
)
|
||||
):
|
||||
return True
|
||||
# Tuple[int] -> Tuple
|
||||
else:
|
||||
if v_origin == constraint:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def issubinstance(data, data_type):
|
||||
if not issubtype(type(data), data_type, recursive=False):
|
||||
return False
|
||||
|
||||
# In Python-3.9 typing library __args__ attribute is not defined for untyped generics
|
||||
dt_args = getattr(data_type, "__args__", None)
|
||||
if isinstance(data, tuple):
|
||||
if dt_args is None or len(dt_args) == 0:
|
||||
return True
|
||||
if len(dt_args) != len(data):
|
||||
return False
|
||||
return all(issubinstance(d, t) for d, t in zip(data, dt_args, strict=True))
|
||||
elif isinstance(data, (list, set)):
|
||||
if dt_args is None or len(dt_args) == 0:
|
||||
return True
|
||||
t = dt_args[0]
|
||||
return all(issubinstance(d, t) for d in data)
|
||||
elif isinstance(data, dict):
|
||||
if dt_args is None or len(dt_args) == 0:
|
||||
return True
|
||||
kt, vt = dt_args
|
||||
return all(
|
||||
issubinstance(k, kt) and issubinstance(v, vt) for k, v in data.items()
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
# [Note: TypeMeta and TypeAlias]
|
||||
# In order to keep compatibility for Python 3.6, use Meta for the typing.
|
||||
# TODO: When PyTorch drops the support for Python 3.6, it can be converted
|
||||
# into the Alias system and using `__class_getitem__` for DataPipe. The
|
||||
# typing system will gain benefit of performance and resolving metaclass
|
||||
# conflicts as elaborated in https://www.python.org/dev/peps/pep-0560/
|
||||
|
||||
|
||||
class _DataPipeType:
|
||||
r"""Save type annotation in `param`."""
|
||||
|
||||
def __init__(self, param) -> None:
|
||||
self.param = param
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return _type_repr(self.param)
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, _DataPipeType):
|
||||
return self.param == other.param
|
||||
return NotImplemented
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.param)
|
||||
|
||||
def issubtype(self, other):
|
||||
if isinstance(other.param, _GenericAlias):
|
||||
if getattr(other.param, "__origin__", None) is Generic:
|
||||
return True
|
||||
if isinstance(other, _DataPipeType):
|
||||
return issubtype(self.param, other.param)
|
||||
if isinstance(other, type):
|
||||
return issubtype(self.param, other)
|
||||
raise TypeError(f"Expected '_DataPipeType' or 'type', but found {type(other)}")
|
||||
|
||||
def issubtype_of_instance(self, other):
|
||||
return issubinstance(other, self.param)
|
||||
|
||||
|
||||
# Default type for DataPipe without annotation
|
||||
_T_co = TypeVar("_T_co", covariant=True)
|
||||
# pyrefly: ignore [invalid-annotation]
|
||||
_DEFAULT_TYPE = _DataPipeType(Generic[_T_co])
|
||||
|
||||
|
||||
class _DataPipeMeta(GenericMeta):
|
||||
r"""
|
||||
Metaclass for `DataPipe`.
|
||||
|
||||
Add `type` attribute and `__init_subclass__` based on the type, and validate the return hint of `__iter__`.
|
||||
|
||||
Note that there is subclass `_IterDataPipeMeta` specifically for `IterDataPipe`.
|
||||
"""
|
||||
|
||||
type: _DataPipeType
|
||||
|
||||
def __new__(cls, name, bases, namespace, **kwargs):
|
||||
return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload]
|
||||
|
||||
# TODO: the statements below are not reachable by design as there is a bug and typing is low priority for now.
|
||||
cls.__origin__ = None
|
||||
if "type" in namespace:
|
||||
return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload]
|
||||
|
||||
namespace["__type_class__"] = False
|
||||
# For plain derived class without annotation
|
||||
for base in bases:
|
||||
if isinstance(base, _DataPipeMeta):
|
||||
return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload]
|
||||
|
||||
namespace.update(
|
||||
{"type": _DEFAULT_TYPE, "__init_subclass__": _dp_init_subclass}
|
||||
)
|
||||
return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload]
|
||||
|
||||
def __init__(self, name, bases, namespace, **kwargs) -> None:
|
||||
super().__init__(name, bases, namespace, **kwargs) # type: ignore[call-overload]
|
||||
|
||||
# TODO: Fix isinstance bug
|
||||
@_tp_cache
|
||||
def _getitem_(self, params):
|
||||
if params is None:
|
||||
raise TypeError(f"{self.__name__}[t]: t can not be None")
|
||||
if isinstance(params, str):
|
||||
params = ForwardRef(params)
|
||||
if not isinstance(params, tuple):
|
||||
params = (params,)
|
||||
|
||||
msg = f"{self.__name__}[t]: t must be a type"
|
||||
params = tuple(_type_check(p, msg) for p in params)
|
||||
|
||||
if isinstance(self.type.param, _GenericAlias):
|
||||
orig = getattr(self.type.param, "__origin__", None)
|
||||
if isinstance(orig, type) and orig is not Generic:
|
||||
p = self.type.param[params] # type: ignore[index]
|
||||
t = _DataPipeType(p)
|
||||
l = len(str(self.type)) + 2
|
||||
name = self.__name__[:-l]
|
||||
name = name + "[" + str(t) + "]"
|
||||
bases = (self,) + self.__bases__
|
||||
return self.__class__(
|
||||
name,
|
||||
bases,
|
||||
{
|
||||
"__init_subclass__": _dp_init_subclass,
|
||||
"type": t,
|
||||
"__type_class__": True,
|
||||
},
|
||||
)
|
||||
|
||||
if len(params) > 1:
|
||||
raise TypeError(
|
||||
f"Too many parameters for {self} actual {len(params)}, expected 1"
|
||||
)
|
||||
|
||||
t = _DataPipeType(params[0])
|
||||
|
||||
if not t.issubtype(self.type):
|
||||
raise TypeError(
|
||||
f"Can not subclass a DataPipe[{t}] from DataPipe[{self.type}]"
|
||||
)
|
||||
|
||||
# Types are equal, fast path for inheritance
|
||||
if self.type == t:
|
||||
return self
|
||||
|
||||
name = self.__name__ + "[" + str(t) + "]"
|
||||
bases = (self,) + self.__bases__
|
||||
|
||||
return self.__class__(
|
||||
name,
|
||||
bases,
|
||||
{"__init_subclass__": _dp_init_subclass, "__type_class__": True, "type": t},
|
||||
)
|
||||
|
||||
# TODO: Fix isinstance bug
|
||||
def _eq_(self, other):
|
||||
if not isinstance(other, _DataPipeMeta):
|
||||
return NotImplemented
|
||||
if self.__origin__ is None or other.__origin__ is None: # type: ignore[has-type]
|
||||
return self is other
|
||||
return (
|
||||
self.__origin__ == other.__origin__ # type: ignore[has-type]
|
||||
and self.type == other.type
|
||||
)
|
||||
|
||||
# TODO: Fix isinstance bug
|
||||
def _hash_(self):
|
||||
return hash((self.__name__, self.type))
|
||||
|
||||
|
||||
class _IterDataPipeMeta(_DataPipeMeta):
|
||||
r"""
|
||||
Metaclass for `IterDataPipe` and inherits from `_DataPipeMeta`.
|
||||
|
||||
Add various functions for behaviors specific to `IterDataPipe`.
|
||||
"""
|
||||
|
||||
def __new__(cls, name, bases, namespace, **kwargs):
|
||||
if "reset" in namespace:
|
||||
reset_func = namespace["reset"]
|
||||
|
||||
@functools.wraps(reset_func)
|
||||
def conditional_reset(*args, **kwargs) -> None:
|
||||
r"""
|
||||
Only execute DataPipe's `reset()` method if `_SnapshotState` is `Iterating` or `NotStarted`.
|
||||
|
||||
This allows recently restored DataPipe to preserve its restored state during the initial `__iter__` call.
|
||||
"""
|
||||
datapipe = args[0]
|
||||
if datapipe._snapshot_state in (
|
||||
_SnapshotState.Iterating,
|
||||
_SnapshotState.NotStarted,
|
||||
):
|
||||
# Reset `NotStarted` is necessary because the `source_datapipe` of a DataPipe might have
|
||||
# already begun iterating.
|
||||
datapipe._number_of_samples_yielded = 0
|
||||
datapipe._fast_forward_iterator = None
|
||||
reset_func(*args, **kwargs)
|
||||
datapipe._snapshot_state = _SnapshotState.Iterating
|
||||
|
||||
namespace["reset"] = conditional_reset
|
||||
|
||||
if "__iter__" in namespace:
|
||||
hook_iterator(namespace)
|
||||
return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload]
|
||||
|
||||
|
||||
def _dp_init_subclass(sub_cls, *args, **kwargs) -> None:
|
||||
# Add function for datapipe instance to reinforce the type
|
||||
sub_cls.reinforce_type = reinforce_type
|
||||
|
||||
# TODO:
|
||||
# - add global switch for type checking at compile-time
|
||||
|
||||
# Ignore internal type class
|
||||
if getattr(sub_cls, "__type_class__", False):
|
||||
return
|
||||
|
||||
# Check if the string type is valid
|
||||
if isinstance(sub_cls.type.param, ForwardRef):
|
||||
base_globals = sys.modules[sub_cls.__module__].__dict__
|
||||
try:
|
||||
param = _eval_type(sub_cls.type.param, base_globals, locals())
|
||||
sub_cls.type.param = param
|
||||
except TypeError as e:
|
||||
raise TypeError(
|
||||
f"{sub_cls.type.param.__forward_arg__} is not supported by Python typing"
|
||||
) from e
|
||||
|
||||
if "__iter__" in sub_cls.__dict__:
|
||||
iter_fn = sub_cls.__dict__["__iter__"]
|
||||
hints = get_type_hints(iter_fn)
|
||||
if "return" in hints:
|
||||
return_hint = hints["return"]
|
||||
# Plain Return Hint for Python 3.6
|
||||
if return_hint == Iterator:
|
||||
return
|
||||
if not (
|
||||
hasattr(return_hint, "__origin__")
|
||||
and (
|
||||
return_hint.__origin__ == Iterator
|
||||
or return_hint.__origin__ == collections.abc.Iterator
|
||||
)
|
||||
):
|
||||
raise TypeError(
|
||||
"Expected 'Iterator' as the return annotation for `__iter__` of {}"
|
||||
", but found {}".format(
|
||||
sub_cls.__name__, _type_repr(hints["return"])
|
||||
)
|
||||
)
|
||||
data_type = return_hint.__args__[0]
|
||||
if not issubtype(data_type, sub_cls.type.param):
|
||||
raise TypeError(
|
||||
f"Expected return type of '__iter__' as a subtype of {sub_cls.type},"
|
||||
f" but found {_type_repr(data_type)} for {sub_cls.__name__}"
|
||||
)
|
||||
|
||||
|
||||
def reinforce_type(self, expected_type):
|
||||
r"""
|
||||
Reinforce the type for DataPipe instance.
|
||||
|
||||
And the 'expected_type' is required to be a subtype of the original type
|
||||
hint to restrict the type requirement of DataPipe instance.
|
||||
"""
|
||||
if isinstance(expected_type, tuple):
|
||||
expected_type = tuple[expected_type] # type: ignore[valid-type]
|
||||
_type_check(expected_type, msg="'expected_type' must be a type")
|
||||
|
||||
if not issubtype(expected_type, self.type.param):
|
||||
raise TypeError(
|
||||
f"Expected 'expected_type' as subtype of {self.type}, but found {_type_repr(expected_type)}"
|
||||
)
|
||||
|
||||
self.type = _DataPipeType(expected_type)
|
||||
return self
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
from torch.utils.data.datapipes.dataframe.dataframes import (
|
||||
CaptureDataFrame,
|
||||
DFIterDataPipe,
|
||||
)
|
||||
from torch.utils.data.datapipes.dataframe.datapipes import DataFramesAsTuplesPipe
|
||||
|
||||
|
||||
__all__ = ["CaptureDataFrame", "DFIterDataPipe", "DataFramesAsTuplesPipe"]
|
||||
|
||||
# Please keep this list sorted
|
||||
if __all__ != sorted(__all__):
|
||||
raise AssertionError("__all__ is not sorted")
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any
|
||||
|
||||
|
||||
_pandas: Any = None
|
||||
_WITH_PANDAS: bool | None = None
|
||||
|
||||
|
||||
def _try_import_pandas() -> bool:
|
||||
try:
|
||||
import pandas # type: ignore[import]
|
||||
|
||||
global _pandas
|
||||
_pandas = pandas
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
# pandas used only for prototyping, will be shortly replaced with TorchArrow
|
||||
def _with_pandas() -> bool:
|
||||
global _WITH_PANDAS
|
||||
if _WITH_PANDAS is None:
|
||||
_WITH_PANDAS = _try_import_pandas()
|
||||
return _WITH_PANDAS
|
||||
|
||||
|
||||
class PandasWrapper:
|
||||
@classmethod
|
||||
def create_dataframe(cls, data, columns):
|
||||
if not _with_pandas():
|
||||
raise RuntimeError("DataFrames prototype requires pandas to function")
|
||||
return _pandas.DataFrame(data, columns=columns) # type: ignore[union-attr]
|
||||
|
||||
@classmethod
|
||||
def is_dataframe(cls, data):
|
||||
if not _with_pandas():
|
||||
return False
|
||||
return isinstance(data, _pandas.core.frame.DataFrame) # type: ignore[union-attr]
|
||||
|
||||
@classmethod
|
||||
def is_column(cls, data):
|
||||
if not _with_pandas():
|
||||
return False
|
||||
return isinstance(data, _pandas.core.series.Series) # type: ignore[union-attr]
|
||||
|
||||
@classmethod
|
||||
def iterate(cls, data):
|
||||
if not _with_pandas():
|
||||
raise RuntimeError("DataFrames prototype requires pandas to function")
|
||||
yield from data.itertuples(index=False)
|
||||
|
||||
@classmethod
|
||||
def concat(cls, buffer):
|
||||
if not _with_pandas():
|
||||
raise RuntimeError("DataFrames prototype requires pandas to function")
|
||||
return _pandas.concat(buffer) # type: ignore[union-attr]
|
||||
|
||||
@classmethod
|
||||
def get_item(cls, data, idx):
|
||||
if not _with_pandas():
|
||||
raise RuntimeError("DataFrames prototype requires pandas to function")
|
||||
return data[idx : idx + 1]
|
||||
|
||||
@classmethod
|
||||
def get_len(cls, df):
|
||||
if not _with_pandas():
|
||||
raise RuntimeError("DataFrames prototype requires pandas to function")
|
||||
return len(df.index)
|
||||
|
||||
@classmethod
|
||||
def get_columns(cls, df):
|
||||
if not _with_pandas():
|
||||
raise RuntimeError("DataFrames prototype requires pandas to function")
|
||||
return list(df.columns.values.tolist())
|
||||
|
||||
|
||||
# When you build own implementation just override it with dataframe_wrapper.set_df_wrapper(new_wrapper_class)
|
||||
default_wrapper = PandasWrapper
|
||||
|
||||
|
||||
def get_df_wrapper():
|
||||
return default_wrapper
|
||||
|
||||
|
||||
def set_df_wrapper(wrapper) -> None:
|
||||
global default_wrapper
|
||||
default_wrapper = wrapper
|
||||
|
||||
|
||||
def create_dataframe(data, columns=None):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.create_dataframe(data, columns)
|
||||
|
||||
|
||||
def is_dataframe(data):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.is_dataframe(data)
|
||||
|
||||
|
||||
def get_columns(data):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.get_columns(data)
|
||||
|
||||
|
||||
def is_column(data):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.is_column(data)
|
||||
|
||||
|
||||
def concat(buffer):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.concat(buffer)
|
||||
|
||||
|
||||
def iterate(data):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.iterate(data)
|
||||
|
||||
|
||||
def get_item(data, idx):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.get_item(data, idx)
|
||||
|
||||
|
||||
def get_len(df):
|
||||
wrapper = get_df_wrapper()
|
||||
return wrapper.get_len(df)
|
||||
+468
@@ -0,0 +1,468 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from torch.utils.data.datapipes._decorator import functional_datapipe
|
||||
from torch.utils.data.datapipes.dataframe.structures import DataChunkDF
|
||||
from torch.utils.data.datapipes.datapipe import DFIterDataPipe, IterDataPipe
|
||||
|
||||
|
||||
# TODO(VitalyFedyunin): Add error when two different traces get combined
|
||||
|
||||
__all__ = [
|
||||
"Capture",
|
||||
"CaptureA",
|
||||
"CaptureAdd",
|
||||
"CaptureCall",
|
||||
"CaptureControl",
|
||||
"CaptureDataFrame",
|
||||
"CaptureDataFrameWithDataPipeOps",
|
||||
"CaptureF",
|
||||
"CaptureGetAttr",
|
||||
"CaptureGetItem",
|
||||
"CaptureInitial",
|
||||
"CaptureLikeMock",
|
||||
"CaptureMul",
|
||||
"CaptureSetItem",
|
||||
"CaptureSub",
|
||||
"CaptureVariable",
|
||||
"CaptureVariableAssign",
|
||||
"DataFrameTracer",
|
||||
"DataFrameTracedOps",
|
||||
"disable_capture",
|
||||
"get_val",
|
||||
]
|
||||
|
||||
|
||||
def disable_capture() -> None:
|
||||
CaptureControl.disabled = True
|
||||
|
||||
|
||||
class CaptureControl:
|
||||
disabled = False
|
||||
|
||||
|
||||
class DataFrameTracedOps(DFIterDataPipe):
|
||||
def __init__(self, source_datapipe, output_var) -> None:
|
||||
super().__init__()
|
||||
self.source_datapipe = source_datapipe
|
||||
self.output_var = output_var
|
||||
|
||||
def __iter__(self):
|
||||
for item in self.source_datapipe:
|
||||
yield self.output_var.apply_ops(item)
|
||||
|
||||
|
||||
# TODO(VitalyFedyunin): Extract this list from the DFIterDataPipe registered functions
|
||||
DATAPIPES_OPS = [
|
||||
"_dataframes_as_tuples",
|
||||
"groupby",
|
||||
"_dataframes_filter",
|
||||
"map",
|
||||
"to_datapipe",
|
||||
"shuffle",
|
||||
"concat",
|
||||
"batch",
|
||||
"_dataframes_per_row",
|
||||
"_dataframes_concat",
|
||||
"_dataframes_shuffle",
|
||||
]
|
||||
|
||||
UNIMPLEMENTED_ATTR = ["__deepcopy__", "__setstate__", "is_shardable", "apply_sharding"]
|
||||
|
||||
|
||||
class Capture:
|
||||
# TODO: All operations are shared across entire InitialCapture, need to figure out what if we join two captures
|
||||
|
||||
def __init__(self, schema_df=None) -> None:
|
||||
self.ctx = {"operations": [], "variables": [], "schema_df": schema_df}
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self._ops_str()
|
||||
|
||||
def _ops_str(self):
|
||||
res = ""
|
||||
for op in self.ctx["operations"]:
|
||||
if len(res) > 0:
|
||||
res += "\n"
|
||||
res += str(op)
|
||||
return res
|
||||
|
||||
def __getstate__(self):
|
||||
# TODO(VitalyFedyunin): Currently can't pickle (why?)
|
||||
self.ctx["schema_df"] = None
|
||||
for var in self.ctx["variables"]:
|
||||
var.calculated_value = None
|
||||
state = {}
|
||||
for item in self.__dict__:
|
||||
state[item] = getattr(self, item)
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
for k, v in state.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
def __getattr__(self, attrname):
|
||||
if attrname == "kwarg" or attrname == "kwargs":
|
||||
raise RuntimeError("no kwargs!")
|
||||
if attrname == "__deepcopy__":
|
||||
raise AttributeError
|
||||
result = CaptureGetAttr(self, attrname, ctx=self.ctx)
|
||||
return result
|
||||
|
||||
def __getitem__(self, key):
|
||||
return CaptureGetItem(self, key, ctx=self.ctx)
|
||||
|
||||
def __setitem__(self, key, value) -> None:
|
||||
self.ctx["operations"].append(CaptureSetItem(self, key, value, ctx=self.ctx))
|
||||
|
||||
def __add__(self, add_val):
|
||||
res = CaptureAdd(self, add_val, ctx=self.ctx)
|
||||
var = CaptureVariable(res, ctx=self.ctx)
|
||||
self.ctx["operations"].append(
|
||||
CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)
|
||||
)
|
||||
return var
|
||||
|
||||
def __sub__(self, add_val):
|
||||
res = CaptureSub(self, add_val, ctx=self.ctx)
|
||||
var = CaptureVariable(res, ctx=self.ctx)
|
||||
self.ctx["operations"].append(
|
||||
CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)
|
||||
)
|
||||
return var
|
||||
|
||||
def __mul__(self, add_val):
|
||||
res = CaptureMul(self, add_val, ctx=self.ctx)
|
||||
var = CaptureVariable(res, ctx=self.ctx)
|
||||
t = CaptureVariableAssign(variable=var, value=res, ctx=self.ctx)
|
||||
self.ctx["operations"].append(t)
|
||||
return var
|
||||
|
||||
def _is_context_empty(self):
|
||||
return len(self.ctx["operations"]) == 0 and len(self.ctx["variables"]) == 0
|
||||
|
||||
def apply_ops_2(self, dataframe) -> None:
|
||||
# TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer)
|
||||
self.ctx["variables"][0].calculated_value = dataframe
|
||||
for op in self.ctx["operations"]:
|
||||
op.execute()
|
||||
|
||||
@property
|
||||
def columns(self):
|
||||
self.apply_ops_2(self.ctx["schema_df"])
|
||||
value = self.execute()
|
||||
return value.columns
|
||||
|
||||
# TODO(VitalyFedyunin): Add tests
|
||||
# TODO(VitalyFedyunin): Need to join context if one of them are empty because we used capture
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
# TODO: Check if args or kwargs have more than one different context
|
||||
if self._is_context_empty():
|
||||
# TODO: Allow CaptureA to take context from mock
|
||||
for arg in args:
|
||||
if isinstance(arg, Capture) and not arg._is_context_empty():
|
||||
self.ctx = arg.ctx
|
||||
break
|
||||
if self._is_context_empty():
|
||||
for k, v in kwargs.items():
|
||||
if isinstance(k, Capture) and not k._is_context_empty():
|
||||
self.ctx = k.ctx
|
||||
break
|
||||
if isinstance(v, Capture) and not v._is_context_empty():
|
||||
self.ctx = v.ctx
|
||||
break
|
||||
|
||||
res = CaptureCall(self, ctx=self.ctx, args=args, kwargs=kwargs)
|
||||
var = CaptureVariable(None, ctx=self.ctx)
|
||||
t = CaptureVariableAssign(ctx=self.ctx, variable=var, value=res)
|
||||
self.ctx["operations"].append(t)
|
||||
return var
|
||||
|
||||
|
||||
class CaptureF(Capture):
|
||||
def __init__(self, ctx=None, **kwargs) -> None:
|
||||
super().__init__()
|
||||
if ctx is None:
|
||||
self.ctx = {"operations": [], "variables": []}
|
||||
else:
|
||||
self.ctx = ctx
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
class CaptureA(CaptureF):
|
||||
def __str__(self) -> str:
|
||||
return f"{self.kwargs['name']}"
|
||||
|
||||
def execute(self):
|
||||
value = self.kwargs["real_attribute"]
|
||||
return value
|
||||
|
||||
|
||||
class CaptureLikeMock:
|
||||
def __init__(self, name) -> None:
|
||||
import unittest.mock as mock
|
||||
|
||||
# TODO(VitalyFedyunin): Do not use private function here, copy own implementation instead.
|
||||
get_target, attribute = mock._get_target(name) # type: ignore[attr-defined]
|
||||
self.get_target = get_target
|
||||
self.attribute = attribute
|
||||
self.name = name
|
||||
|
||||
def __enter__(self):
|
||||
self.save = getattr(self.get_target(), self.attribute)
|
||||
capt = CaptureA(name=self.name, real_attribute=self.save)
|
||||
setattr(self.get_target(), self.attribute, capt)
|
||||
|
||||
def __exit__(self, *exc_info):
|
||||
setattr(self.get_target(), self.attribute, self.save)
|
||||
|
||||
|
||||
class CaptureCall(Capture):
|
||||
def __init__(self, callable, ctx=None, **kwargs) -> None:
|
||||
super().__init__()
|
||||
if ctx is None:
|
||||
self.ctx = {"operations": [], "variables": []}
|
||||
else:
|
||||
self.ctx = ctx
|
||||
self.kwargs = kwargs
|
||||
self.callable = callable
|
||||
|
||||
def __str__(self) -> str:
|
||||
return "{callable}({args},{kwargs})".format(
|
||||
callable=self.callable, **self.kwargs
|
||||
)
|
||||
|
||||
def execute(self):
|
||||
# TODO: VitalyFedyunin execute kwargs and maybe nested structures
|
||||
executed_args = []
|
||||
for arg in self.kwargs["args"]:
|
||||
if isinstance(arg, Capture):
|
||||
executed_args.append(arg.execute())
|
||||
else:
|
||||
executed_args.append(arg)
|
||||
left = get_val(self.callable)
|
||||
return left(*executed_args, **self.kwargs["kwargs"])
|
||||
|
||||
|
||||
class CaptureVariableAssign(CaptureF):
|
||||
def __str__(self) -> str:
|
||||
variable = self.kwargs["variable"]
|
||||
value = self.kwargs["value"]
|
||||
return f"{variable} = {value}"
|
||||
|
||||
def execute(self) -> None:
|
||||
self.kwargs["variable"].calculated_value = self.kwargs["value"].execute()
|
||||
|
||||
|
||||
class CaptureVariable(Capture):
|
||||
# TODO(VitalyFedyunin): This should be atomic and thread safe
|
||||
names_idx = 0
|
||||
|
||||
def __init__(self, value, ctx) -> None:
|
||||
super().__init__()
|
||||
if CaptureControl.disabled:
|
||||
raise RuntimeError("Attempting to create capture variable with capture off")
|
||||
self.ctx = ctx
|
||||
self.value = value
|
||||
self.name = f"var_{CaptureVariable.names_idx}"
|
||||
CaptureVariable.names_idx += 1
|
||||
self.ctx["variables"].append(self)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.name
|
||||
|
||||
def execute(self):
|
||||
return self.calculated_value
|
||||
|
||||
def apply_ops(self, dataframe):
|
||||
# TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer)
|
||||
self.ctx["variables"][0].calculated_value = dataframe
|
||||
for op in self.ctx["operations"]:
|
||||
op.execute()
|
||||
return self.calculated_value
|
||||
|
||||
|
||||
class CaptureGetItem(Capture):
|
||||
def __init__(self, left, key, ctx) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self.left = left
|
||||
self.key = key
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.left}[{get_val(self.key)}]"
|
||||
|
||||
def execute(self):
|
||||
left = self.left.execute()
|
||||
return left[self.key]
|
||||
|
||||
|
||||
class CaptureSetItem(Capture):
|
||||
def __init__(self, left, key, value, ctx) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self.left = left
|
||||
self.key = key
|
||||
self.value = value
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.left}[{get_val(self.key)}] = {self.value}"
|
||||
|
||||
def execute(self) -> None:
|
||||
left = self.left.execute()
|
||||
value = self.value.execute()
|
||||
left[self.key] = value
|
||||
|
||||
|
||||
class CaptureAdd(Capture):
|
||||
def __init__(self, left, right, ctx) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self.left = left
|
||||
self.right = right
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.left} + {self.right}"
|
||||
|
||||
def execute(self):
|
||||
return get_val(self.left) + get_val(self.right)
|
||||
|
||||
|
||||
class CaptureMul(Capture):
|
||||
def __init__(self, left, right, ctx) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self.left = left
|
||||
self.right = right
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.left} * {self.right}"
|
||||
|
||||
def execute(self):
|
||||
return get_val(self.left) * get_val(self.right)
|
||||
|
||||
|
||||
class CaptureSub(Capture):
|
||||
def __init__(self, left, right, ctx) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self.left = left
|
||||
self.right = right
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.left} - {self.right}"
|
||||
|
||||
def execute(self):
|
||||
return get_val(self.left) - get_val(self.right)
|
||||
|
||||
|
||||
class CaptureGetAttr(Capture):
|
||||
def __init__(self, src, name, ctx) -> None:
|
||||
super().__init__()
|
||||
self.ctx = ctx
|
||||
self.src = src
|
||||
self.name = name
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.src}.{self.name}"
|
||||
|
||||
def execute(self):
|
||||
val = get_val(self.src)
|
||||
return getattr(val, self.name)
|
||||
|
||||
|
||||
def get_val(capture):
|
||||
if isinstance(capture, Capture):
|
||||
return capture.execute()
|
||||
elif isinstance(capture, str):
|
||||
return f'"{capture}"'
|
||||
else:
|
||||
return capture
|
||||
|
||||
|
||||
class CaptureInitial(CaptureVariable):
|
||||
def __init__(self, schema_df=None) -> None:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
new_ctx: dict[str, list[Any]] = {
|
||||
"operations": [],
|
||||
"variables": [],
|
||||
"schema_df": schema_df,
|
||||
}
|
||||
super().__init__(None, new_ctx)
|
||||
self.name = f"input_{self.name}"
|
||||
|
||||
|
||||
class CaptureDataFrame(CaptureInitial):
|
||||
pass
|
||||
|
||||
|
||||
class CaptureDataFrameWithDataPipeOps(CaptureDataFrame):
|
||||
def as_datapipe(self):
|
||||
return DataFrameTracedOps(self.ctx["variables"][0].source_datapipe, self)
|
||||
|
||||
def raw_iterator(self):
|
||||
return self.as_datapipe().__iter__()
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._dataframes_as_tuples())
|
||||
|
||||
def batch(self, batch_size=10, drop_last: bool = False, wrapper_class=DataChunkDF):
|
||||
dp = self._dataframes_per_row()._dataframes_concat(batch_size)
|
||||
dp = dp.as_datapipe().batch(1, drop_last=drop_last, wrapper_class=wrapper_class)
|
||||
dp._dp_contains_dataframe = True
|
||||
return dp
|
||||
|
||||
def groupby(
|
||||
self,
|
||||
group_key_fn,
|
||||
*,
|
||||
buffer_size=10000,
|
||||
group_size=None,
|
||||
guaranteed_group_size=None,
|
||||
drop_remaining=False,
|
||||
):
|
||||
dp = self._dataframes_per_row()
|
||||
dp = dp.as_datapipe().groupby(
|
||||
group_key_fn,
|
||||
buffer_size=buffer_size,
|
||||
group_size=group_size,
|
||||
guaranteed_group_size=guaranteed_group_size,
|
||||
drop_remaining=drop_remaining,
|
||||
)
|
||||
return dp
|
||||
|
||||
def shuffle(self, *args, **kwargs):
|
||||
return self._dataframes_shuffle(*args, **kwargs)
|
||||
|
||||
def filter(self, *args, **kwargs):
|
||||
return self._dataframes_filter(*args, **kwargs)
|
||||
|
||||
def collate(self, *args, **kwargs) -> NoReturn:
|
||||
raise RuntimeError("Can't collate unbatched DataFrames stream")
|
||||
|
||||
def __getattr__(self, attrname): # ?
|
||||
if attrname in UNIMPLEMENTED_ATTR:
|
||||
raise AttributeError("Attempting to get ", attrname)
|
||||
if attrname in DATAPIPES_OPS:
|
||||
return (self.as_datapipe()).__getattr__(attrname)
|
||||
return super().__getattr__(attrname)
|
||||
|
||||
|
||||
@functional_datapipe("trace_as_dataframe")
|
||||
class DataFrameTracer(CaptureDataFrameWithDataPipeOps, IterDataPipe): # type: ignore[misc]
|
||||
source_datapipe: Any | None = None
|
||||
|
||||
# TODO(VitalyFedyunin): Must implement all special functions of datapipes
|
||||
|
||||
def set_shuffle_settings(self, *args, **kwargs) -> None:
|
||||
pass
|
||||
|
||||
def is_shardable(self) -> bool:
|
||||
return False
|
||||
|
||||
def __init__(self, source_datapipe, schema_df=None) -> None:
|
||||
self.source_datapipe = source_datapipe
|
||||
if schema_df is None:
|
||||
schema_df = next(iter(self.source_datapipe))
|
||||
super().__init__(schema_df=schema_df)
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
from torch.utils.data.datapipes._decorator import functional_datapipe
|
||||
from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper
|
||||
from torch.utils.data.datapipes.datapipe import DFIterDataPipe, IterDataPipe
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConcatDataFramesPipe",
|
||||
"DataFramesAsTuplesPipe",
|
||||
"ExampleAggregateAsDataFrames",
|
||||
"FilterDataFramesPipe",
|
||||
"PerRowDataFramesPipe",
|
||||
"ShuffleDataFramesPipe",
|
||||
]
|
||||
|
||||
|
||||
@functional_datapipe("_dataframes_as_tuples")
|
||||
class DataFramesAsTuplesPipe(IterDataPipe):
|
||||
def __init__(self, source_datapipe) -> None:
|
||||
super().__init__()
|
||||
self.source_datapipe = source_datapipe
|
||||
|
||||
def __iter__(self):
|
||||
for df in self.source_datapipe:
|
||||
# for record in df.to_records(index=False):
|
||||
yield from df_wrapper.iterate(df)
|
||||
|
||||
|
||||
@functional_datapipe("_dataframes_per_row", enable_df_api_tracing=True)
|
||||
class PerRowDataFramesPipe(DFIterDataPipe):
|
||||
def __init__(self, source_datapipe) -> None:
|
||||
self.source_datapipe = source_datapipe
|
||||
|
||||
def __iter__(self):
|
||||
for df in self.source_datapipe:
|
||||
# TODO(VitalyFedyunin): Replacing with TorchArrow only API, as we are dropping pandas as followup
|
||||
for i in range(len(df)):
|
||||
yield df[i : i + 1]
|
||||
|
||||
|
||||
@functional_datapipe("_dataframes_concat", enable_df_api_tracing=True)
|
||||
class ConcatDataFramesPipe(DFIterDataPipe):
|
||||
def __init__(self, source_datapipe, batch=3) -> None:
|
||||
self.source_datapipe = source_datapipe
|
||||
self.n_batch = batch
|
||||
|
||||
def __iter__(self):
|
||||
buffer = []
|
||||
for df in self.source_datapipe:
|
||||
buffer.append(df)
|
||||
if len(buffer) == self.n_batch:
|
||||
yield df_wrapper.concat(buffer)
|
||||
buffer = []
|
||||
if buffer:
|
||||
yield df_wrapper.concat(buffer)
|
||||
|
||||
|
||||
@functional_datapipe("_dataframes_shuffle", enable_df_api_tracing=True)
|
||||
class ShuffleDataFramesPipe(DFIterDataPipe):
|
||||
def __init__(self, source_datapipe) -> None:
|
||||
self.source_datapipe = source_datapipe
|
||||
|
||||
def __iter__(self):
|
||||
size = None
|
||||
all_buffer: list[Any] = []
|
||||
for df in self.source_datapipe:
|
||||
if size is None:
|
||||
size = df_wrapper.get_len(df)
|
||||
all_buffer.extend(
|
||||
df_wrapper.get_item(df, i) for i in range(df_wrapper.get_len(df))
|
||||
)
|
||||
random.shuffle(all_buffer)
|
||||
buffer = []
|
||||
for df in all_buffer:
|
||||
buffer.append(df)
|
||||
if len(buffer) == size:
|
||||
yield df_wrapper.concat(buffer)
|
||||
buffer = []
|
||||
if buffer:
|
||||
yield df_wrapper.concat(buffer)
|
||||
|
||||
|
||||
@functional_datapipe("_dataframes_filter", enable_df_api_tracing=True)
|
||||
class FilterDataFramesPipe(DFIterDataPipe):
|
||||
def __init__(self, source_datapipe, filter_fn) -> None:
|
||||
self.source_datapipe = source_datapipe
|
||||
self.filter_fn = filter_fn
|
||||
|
||||
def __iter__(self):
|
||||
size = None
|
||||
all_buffer = []
|
||||
filter_res = []
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
for df in self.source_datapipe:
|
||||
if size is None:
|
||||
size = len(df.index)
|
||||
for i in range(len(df.index)):
|
||||
all_buffer.append(df[i : i + 1])
|
||||
filter_res.append(self.filter_fn(df.iloc[i]))
|
||||
|
||||
buffer = []
|
||||
for df, res in zip(all_buffer, filter_res, strict=True):
|
||||
if res:
|
||||
buffer.append(df)
|
||||
if len(buffer) == size:
|
||||
yield df_wrapper.concat(buffer)
|
||||
buffer = []
|
||||
if buffer:
|
||||
yield df_wrapper.concat(buffer)
|
||||
|
||||
|
||||
@functional_datapipe("_to_dataframes_pipe", enable_df_api_tracing=True)
|
||||
class ExampleAggregateAsDataFrames(DFIterDataPipe):
|
||||
def __init__(self, source_datapipe, dataframe_size=10, columns=None) -> None:
|
||||
self.source_datapipe = source_datapipe
|
||||
self.columns = columns
|
||||
self.dataframe_size = dataframe_size
|
||||
|
||||
def _as_list(self, item):
|
||||
try:
|
||||
return list(item)
|
||||
except (
|
||||
Exception
|
||||
): # TODO(VitalyFedyunin): Replace with better iterable exception
|
||||
return [item]
|
||||
|
||||
def __iter__(self):
|
||||
aggregate = []
|
||||
for item in self.source_datapipe:
|
||||
aggregate.append(self._as_list(item))
|
||||
if len(aggregate) == self.dataframe_size:
|
||||
yield df_wrapper.create_dataframe(aggregate, columns=self.columns)
|
||||
aggregate = []
|
||||
if len(aggregate) > 0:
|
||||
yield df_wrapper.create_dataframe(aggregate, columns=self.columns)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user