262 lines
9.9 KiB
Python
262 lines
9.9 KiB
Python
# mypy: ignore-errors
|
|
|
|
"""
|
|
Constant and enum variable tracking in Dynamo.
|
|
|
|
This module is fundamental to Dynamo's ability to track and propagate constant
|
|
values during compilation, ensuring proper handling of Python literals and
|
|
maintaining type safety through the compilation process.
|
|
"""
|
|
|
|
import operator
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
from torch._dynamo.source import AttrSource, GetItemSource
|
|
|
|
from .. import graph_break_hints, variables
|
|
from ..exc import raise_observed_exception, unimplemented_v2
|
|
from ..utils import cmp_name_to_op_mapping, common_constant_types, istype, np
|
|
from .base import VariableTracker
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from torch._dynamo.symbolic_convert import InstructionTranslator
|
|
|
|
|
|
class ConstantVariable(VariableTracker):
|
|
"""
|
|
Variable tracker for Python literals and basic immutable types, with automatic
|
|
routing support for collection types (lists, tuples, sets, etc.).
|
|
|
|
The create() method intelligently constructs appropriate variable types for
|
|
nested collections.
|
|
"""
|
|
|
|
@staticmethod
|
|
def create(value, **kwargs) -> VariableTracker:
|
|
"""
|
|
Create a `ConstantVariable` based on the given value, and supports
|
|
automatic routing for collection types like `tuple` (in which case we'd
|
|
create `ConstantVariable` for the leaf items).
|
|
|
|
NOTE: the caller must install the proper guards if needed; most often
|
|
the guard will be `CONSTANT_MATCH`.
|
|
"""
|
|
source = kwargs.get("source", None)
|
|
|
|
# Routing for supported collection literals.
|
|
if isinstance(value, set):
|
|
items = [ConstantVariable.create(x) for x in value]
|
|
return variables.SetVariable(items, **kwargs)
|
|
elif isinstance(value, frozenset):
|
|
items = [ConstantVariable.create(x) for x in value]
|
|
return variables.FrozensetVariable(items, **kwargs)
|
|
elif isinstance(value, (list, tuple)):
|
|
items = []
|
|
for i, x in enumerate(value):
|
|
item_source = GetItemSource(source, i) if source else None
|
|
items.append(
|
|
ConstantVariable.create(
|
|
x,
|
|
source=item_source,
|
|
)
|
|
)
|
|
return variables.BaseListVariable.cls_for(type(value))(items, **kwargs)
|
|
|
|
return ConstantVariable(value, **kwargs)
|
|
|
|
def __init__(self, value, **kwargs) -> None:
|
|
super().__init__(**kwargs)
|
|
assert ConstantVariable.is_base_literal(value), f"""
|
|
Cannot construct `ConstantVariable` for value of type {type(value)}.
|
|
|
|
This failure likely due to PyTorch-internal use of `ConstantVariable` on
|
|
non-literal python values, please try using `VariableTracker.build` instead. If
|
|
you believe it's a necessary and legitimate use case (the value is immutable and
|
|
can't easily be represented with another `VariableTracker` class), please add
|
|
its type to `common_constant_types`.
|
|
"""
|
|
if np is not None and isinstance(value, np.number):
|
|
self.value = value.item()
|
|
else:
|
|
self.value = value
|
|
|
|
def as_proxy(self):
|
|
return self.value
|
|
|
|
def __repr__(self) -> str:
|
|
return f"ConstantVariable({type(self.value).__name__}: {repr(self.value)})"
|
|
|
|
def as_python_constant(self):
|
|
return self.value
|
|
|
|
def is_python_constant(self):
|
|
return True
|
|
|
|
@property
|
|
def items(self):
|
|
"""
|
|
Need this when adding a BaseListVariable and a ConstantVariable together.
|
|
Happens in detectron2.
|
|
"""
|
|
return self.unpack_var_sequence(tx=None)
|
|
|
|
def getitem_const(self, tx: "InstructionTranslator", arg: VariableTracker):
|
|
return ConstantVariable.create(
|
|
self.value[arg.as_python_constant()],
|
|
)
|
|
|
|
@staticmethod
|
|
def is_base_literal(obj):
|
|
return type(obj) in common_constant_types
|
|
|
|
@staticmethod
|
|
def is_literal(obj):
|
|
if type(obj) in (list, tuple, set, frozenset, torch.Size):
|
|
return all(ConstantVariable.is_literal(x) for x in obj)
|
|
return ConstantVariable.is_base_literal(obj)
|
|
|
|
def unpack_var_sequence(self, tx):
|
|
try:
|
|
return [ConstantVariable.create(x) for x in self.as_python_constant()]
|
|
except TypeError as e:
|
|
raise NotImplementedError from e
|
|
|
|
def const_getattr(self, tx: "InstructionTranslator", name):
|
|
if not hasattr(self.value, name):
|
|
raise NotImplementedError
|
|
member = getattr(self.value, name)
|
|
if callable(member):
|
|
raise NotImplementedError
|
|
return member
|
|
|
|
def call_method(
|
|
self,
|
|
tx,
|
|
name,
|
|
args: "list[VariableTracker]",
|
|
kwargs: "dict[str, VariableTracker]",
|
|
) -> "VariableTracker":
|
|
from .tensor import SymNodeVariable
|
|
|
|
if name == "format" and istype(self.value, str):
|
|
return variables.BuiltinVariable(str.format).call_function(
|
|
tx, [self, *args], kwargs
|
|
)
|
|
elif name == "join" and istype(self.value, str):
|
|
assert len(args) == 1 and len(kwargs) == 0
|
|
arg_unpacked = args[0].force_unpack_var_sequence(tx)
|
|
try:
|
|
arg_const = [x.as_python_constant() for x in arg_unpacked]
|
|
return ConstantVariable.create(self.value.join(arg_const))
|
|
except NotImplementedError:
|
|
return super().call_method(tx, name, args, kwargs)
|
|
|
|
if any(isinstance(x, SymNodeVariable) for x in args):
|
|
# Promote to SymNodeVariable for operations involving dynamic shapes.
|
|
return variables.SymNodeVariable(self.as_proxy(), self.value).call_method(
|
|
tx, name, args, kwargs
|
|
)
|
|
|
|
try:
|
|
const_args = [a.as_python_constant() for a in args]
|
|
const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
|
|
except NotImplementedError:
|
|
return super().call_method(tx, name, args, kwargs)
|
|
|
|
if isinstance(self.value, str) and name in str.__dict__.keys():
|
|
method = getattr(self.value, name)
|
|
try:
|
|
return ConstantVariable.create(method(*const_args, **const_kwargs))
|
|
except Exception as e:
|
|
raise_observed_exception(type(e), tx)
|
|
elif isinstance(self.value, (float, int)):
|
|
if not (args or kwargs):
|
|
return ConstantVariable.create(getattr(self.value, name)())
|
|
if (
|
|
hasattr(operator, name)
|
|
and len(args) == 1
|
|
and args[0].is_python_constant()
|
|
):
|
|
add_target = const_args[0]
|
|
op = getattr(operator, name)
|
|
if isinstance(
|
|
add_target, (torch.SymBool, torch.SymFloat, torch.SymInt)
|
|
):
|
|
# Addition between a non sym and sym makes a sym
|
|
proxy = tx.output.create_proxy(
|
|
"call_function", op, (self.value, add_target), {}
|
|
)
|
|
return SymNodeVariable.create(tx, proxy, add_target)
|
|
else:
|
|
return ConstantVariable.create(op(self.value, add_target))
|
|
elif isinstance(self.value, bytes) and name == "decode":
|
|
method = getattr(self.value, name)
|
|
return ConstantVariable.create(method(*const_args, **const_kwargs))
|
|
|
|
if name == "__len__" and not (args or kwargs):
|
|
return ConstantVariable.create(len(self.value))
|
|
elif name == "__round__" and len(args) == 1 and args[0].is_python_constant():
|
|
return ConstantVariable.create(
|
|
round(self.value, args[0].as_python_constant())
|
|
)
|
|
elif name == "__contains__" and len(args) == 1 and args[0].is_python_constant():
|
|
assert not kwargs
|
|
search = args[0].as_python_constant()
|
|
result = search in self.value
|
|
return ConstantVariable.create(result)
|
|
return super().call_method(tx, name, args, kwargs)
|
|
|
|
def call_obj_hasattr(
|
|
self, tx: "InstructionTranslator", name: str
|
|
) -> "VariableTracker":
|
|
result = hasattr(self.value, name)
|
|
return variables.ConstantVariable.create(result)
|
|
|
|
|
|
class EnumVariable(VariableTracker):
|
|
"""VariableTracker for enum.Enum and enum.IntEnum instances
|
|
|
|
Provides specialized handling for Python enum types, supporting
|
|
both standard Enum and IntEnum with proper value tracking and comparison.
|
|
"""
|
|
|
|
def __init__(self, value, **kwargs) -> None:
|
|
super().__init__(**kwargs)
|
|
self.value = value
|
|
|
|
@classmethod
|
|
def create(cls, cls_type, value_vt, options):
|
|
if isinstance(value_vt, variables.ConstantVariable):
|
|
for member in list(cls_type):
|
|
if member.value == value_vt.as_python_constant():
|
|
return cls(member, **options)
|
|
unimplemented_v2(
|
|
gb_type="Failed to construct Enum variable",
|
|
context=f"value: {value_vt}, allowed enum values: {list(cls_type)}",
|
|
explanation="Attempted to construct an Enum value that is non-constant (e.g. int, string) "
|
|
"or is not an acceptable value for the Enum. "
|
|
f"Acceptable values for Enum `{cls_type}`: {list(cls_type)}.",
|
|
hints=[*graph_break_hints.USER_ERROR, *graph_break_hints.SUPPORTABLE],
|
|
)
|
|
|
|
def as_proxy(self):
|
|
if isinstance(self.value, int):
|
|
return int(self.value) # convert IntEnum to a normal int
|
|
return self.value
|
|
|
|
def __repr__(self) -> str:
|
|
return f"EnumVariable({type(self.value)})"
|
|
|
|
def as_python_constant(self):
|
|
return self.value
|
|
|
|
def var_getattr(self, tx: "InstructionTranslator", name):
|
|
if not hasattr(self.value, name):
|
|
raise NotImplementedError
|
|
if name in cmp_name_to_op_mapping:
|
|
return variables.GetAttrVariable(self, name)
|
|
member = getattr(self.value, name)
|
|
source = self.source and AttrSource(self.source, name)
|
|
return VariableTracker.build(tx, member, source=source)
|