Spaces:
Running
Running
File size: 14,040 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 |
# mypy: ignore-errors
import collections
from enum import Enum
from typing import Any, Callable, Dict, List
from .. import variables
from ..current_scope_id import current_scope_id
from ..exc import unimplemented
from ..source import AttrSource, Source
from ..utils import identity, istype
class MutableLocalSource(Enum):
"""
If the VariableTracker.mutable_local represents a Variable that:
- already existed that Dynamo began tracking while introspection (Existing)
- is a new variable that is created during Dynamo introspection (Local)
"""
Existing = 0
Local = 1
class ParentsTracker:
"""
This is a perf optimization to limit the number of objects we need to visit in tx.replace_all.
This must be a seperate object so that it is not cloned in apply.
"""
def __init__(self):
# logically this is a set, but we use a dict to ensure deterministic ordering
self.parents: Dict[ParentsTracker, bool] = dict()
def add(self, parent):
self.parents[parent] = True
def recursive_parents(self):
rv = dict(self.parents)
worklist = list(self.parents)
while worklist:
for parent in worklist.pop().parents:
if parent not in rv:
assert isinstance(parent, ParentsTracker)
rv[parent] = True
worklist.append(parent)
return rv.keys()
class MutableLocalBase:
"""
Base class for Variable.mutable_local
"""
def __init__(self, typ: MutableLocalSource):
# In HigherOrderOperator tracing, we need to distinguish
# between MutableLocals inside the HigherOrderOperator and
# ones outside it. For example, it is not safe to mutate
# `a` in the following example because it was constructed
# in a different scope.
#
# def f(x):
# a = 1
# def g(x):
# nonlocal a
# a = 2
# return x
# return wrap(g, x) + a
#
# We use self.scope to distinguish this.
# scope == 0: The object was an existing variable
# scope == 1: The object was created while Dynamo
# was introspecting a function
# (and no HigherOrderOps were involved)
# scope >= 2: The object was created through
# Dynamo introspection of a HigherOrderOp.
# The exact number corresponds to the level
# of nested HigherOrderOps.
if typ is MutableLocalSource.Existing:
self.scope = 0
elif typ is MutableLocalSource.Local:
self.scope = current_scope_id()
else:
unimplemented(f"Unsupported MutableLocalSource: {typ}")
class MutableLocal(MutableLocalBase):
"""
Marker used to indicate this (list, iter, etc) was constructed in
local scope and can be mutated safely in analysis without leaking
state.
"""
def __init__(self):
super().__init__(MutableLocalSource.Local)
def __hash__(self):
return id(self)
def __eq__(self, other):
return self is other
def _is_top_level_scope(scope_id):
return scope_id == 1
def is_side_effect_safe(m: MutableLocalBase):
scope_id = current_scope_id()
# In the top-level scope (if no HigherOrderOperators are involved),
# we are allowed to modify variables created in this scope as well
# as existing variables.
if _is_top_level_scope(scope_id):
return True
# Otherwise, only allow local mutation of variables created in the current scope
return m.scope == scope_id
class VariableTrackerMeta(type):
def __call__(cls, *args, **kwargs):
"""Call __post_init__"""
obj = type.__call__(cls, *args, **kwargs)
obj.__post_init__(*args, **kwargs)
return obj
def __instancecheck__(cls, instance) -> bool:
"""Make isinstance work with LazyVariableTracker"""
if type.__instancecheck__(
variables.LazyVariableTracker, instance
) and cls not in (
VariableTracker,
variables.LazyVariableTracker,
):
instance = instance.realize()
return type.__instancecheck__(cls, instance)
class VariableTracker(metaclass=VariableTrackerMeta):
"""
Base class for tracked locals and stack values
VariableTracker instances are immutable and should be copied in
order to change them.
"""
# fields to leave unmodified in apply()
_nonvar_fields = {
"value",
"guards",
"source",
"mutable_local",
"parents_tracker",
"user_code_variable_name",
}
def clone(self, **kwargs):
"""Shallow copy with some (optional) changes"""
args = dict(self.__dict__)
args.update(kwargs)
return self.__class__(**args)
@classmethod
def copy(cls, value):
"""Deeper (but not full) copy, leaving FX and user objects alone"""
return cls.apply(identity, value)
@classmethod
def apply(
cls,
fn: Callable[["VariableTracker"], "VariableTracker"],
value,
cache=None,
skip_fn=lambda _: False, # Whether we should skip applying to this var
):
"""
Walk this object and call fn on all the VariableTracker
instances
"""
if cache is None:
cache = dict()
idx = id(value)
if idx in cache:
return cache[idx][0]
if isinstance(value, VariableTracker):
if not skip_fn(value):
def update_object_dict(v):
changed = False
rv = v.__dict__
for key in rv.keys():
if key not in v._nonvar_fields:
prior = rv[key]
rv[key] = cls.apply(fn, prior, cache, skip_fn)
changed = changed or prior is not rv[key]
return v
value = value.unwrap()
was_realized = value.is_realized()
result = fn(update_object_dict(value))
if not was_realized and value.is_realized():
# running fn() resulted in value getting realized,
# which means we missed updating the contents of result
result = update_object_dict(result.unwrap())
else:
result = fn(value)
if result is not None:
result = result.unwrap()
elif istype(value, list):
result = [cls.apply(fn, v, cache, skip_fn) for v in value]
elif istype(value, tuple):
result = tuple(cls.apply(fn, v, cache, skip_fn) for v in value)
elif istype(value, (dict, collections.OrderedDict)):
result = {
k: cls.apply(fn, v, cache, skip_fn) for k, v in list(value.items())
}
else:
result = value
# save `value` to keep it alive and ensure id() isn't reused
cache[idx] = (result, value)
return result
def __repr__(self):
return f"{self.__class__.__name__}()"
def python_type(self):
"""
Abstract method to be implemented by subclasses of VariableTracker.
This method should return the type represented by the instance of the subclass.
The purpose is to provide a standardized way to retrieve the Python type information
of the variable being tracked.
Returns:
type: The Python type (such as int, str, list, etc.) of the variable tracked by
the subclass. If the type cannot be determined or is not relevant,
leaving it undefined or invoking super() is always sound.
Note:
This is an abstract method and may be overridden in subclasses.
Example:
class SetVariable(VariableTracker):
def python_type(self):
return set
Raises:
NotImplementedError: If the method is not implemented in a subclass.
"""
raise NotImplementedError(f"{self} has no type")
def as_python_constant(self):
"""For constants"""
raise NotImplementedError(f"{self} is not a constant")
def guard_as_python_constant(self):
"""Similar to as_python_constant(), but add ID_MATCH guards to try to force things to become constants"""
try:
return self.as_python_constant()
except NotImplementedError as e:
unimplemented(str(e))
def is_python_constant(self):
try:
self.as_python_constant()
return True
except NotImplementedError:
return False
def make_guard(self, fn):
if self.source:
return self.source.make_guard(fn)
raise NotImplementedError()
def const_getattr(self, tx, name: str) -> Any:
"""getattr(self, name) returning a python constant"""
raise NotImplementedError()
def var_getattr(self, tx, name: str) -> "VariableTracker":
"""getattr(self, name) returning a new variable"""
value = self.const_getattr(tx, name)
if not variables.ConstantVariable.is_literal(value):
raise NotImplementedError()
source = None
if self.source:
source = AttrSource(self.source, name)
return variables.ConstantVariable.create(value, source=source)
def is_proxy(self):
try:
self.as_proxy()
return True
except NotImplementedError:
return False
def as_proxy(self):
raise NotImplementedError(str(self))
def maybe_fx_node(self):
try:
proxy = self.as_proxy()
import torch.fx
if isinstance(proxy, torch.fx.Proxy):
return proxy.node
return None
except NotImplementedError:
return None
def reconstruct(self, codegen):
raise NotImplementedError()
def can_reconstruct(self, tx):
"""If it is possible to reconstruct the Python object this
VariableTracker represents."""
assert tx is tx.output.root_tx, "Only root tx can reconstruct"
try:
from ..codegen import PyCodegen
cg = PyCodegen(tx)
self.reconstruct(cg)
return True
except NotImplementedError:
return False
def unpack_var_sequence(self, tx) -> List["VariableTracker"]:
raise NotImplementedError()
def has_unpack_var_sequence(self, tx) -> bool:
try:
self.unpack_var_sequence(tx)
return True
except NotImplementedError:
return False
def inspect_parameter_names(self) -> List[str]:
unimplemented(f"inspect_parameter_names: {self}")
def call_hasattr(self, tx, name: str) -> "VariableTracker":
unimplemented(f"hasattr {self.__class__.__name__} {name}")
def call_function(
self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
) -> "VariableTracker":
unimplemented(f"call_function {self} {args} {kwargs}")
def call_method(
self,
tx,
name,
args: "List[VariableTracker]",
kwargs: "Dict[str, VariableTracker]",
) -> "VariableTracker":
if name == "__len__" and self.has_unpack_var_sequence(tx):
assert not (args or kwargs)
return variables.ConstantVariable.create(len(self.unpack_var_sequence(tx)))
elif (
name == "__getattr__"
and len(args) == 1
and args[0].is_python_constant()
and not kwargs
):
return self.var_getattr(tx, args[0].as_python_constant())
raise unimplemented(f"call_method {self} {name} {args} {kwargs}")
def rename(self, tx, name):
return self
def realize(self) -> "VariableTracker":
"""Used by LazyVariableTracker to build the real VariableTracker"""
return self
def recursive_realize(self):
"""Realize all objects under this"""
return VariableTracker.apply(lambda x: x.realize(), self)
def unwrap(self) -> "VariableTracker":
"""Used by LazyVariableTracker to return the real VariableTracker if it already exists"""
return self
def is_realized(self):
"""Used by LazyVariableTracker to indicate an unrealized node"""
return True
def __init__(
self,
*,
source: Source = None,
mutable_local: MutableLocal = None,
parents_tracker: ParentsTracker = None,
):
super().__init__()
self.source = source
self.mutable_local = mutable_local
self.parents_tracker = parents_tracker
def __post_init__(self, *args, **kwargs):
if self.parents_tracker is None:
self.parents_tracker = ParentsTracker()
# visit children 1 level deep and ensure parent is set properly
VariableTracker.apply(
lambda node: node.parents_tracker.add(self.parents_tracker),
[v for k, v in self.__dict__.items() if k not in self._nonvar_fields],
skip_fn=lambda _: True,
)
def typestr(*objs):
if len(objs) == 1:
(obj,) = objs
if isinstance(obj, VariableTracker):
return str(obj)
else:
return type(obj).__name__
else:
return " ".join(map(typestr, objs))
|