mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
fbeca60b1f
1. Removes calls to `replace_all` and `clone` and makes VTs mutable. 2. Properly handles Tuple Iterator mutation. Previously TupleIterator variables would only be properly reconstructed if they were advanced at least once in a frame. On calls to `next`, the source information would be lost (due to constructing a new iterator without using builder), which would ensure that during codegen the variable would be reconstructed from scratch. Now that VTs are mutated, the source is never lost, so we need to properly track mutation and handle it by replaying calls to `next` at the end of the modified bytecode. 3. Added test for checking iadd side effects, this was missing in our unit test coverage. 4. Fixed two incorrect sources, DelayGraphBreakVariable, and UserMethodVariable both relied on setting the source to AttrSource(parent, name) at the callsite of `var_getattr`. 5. Fixed a bug in inplace adding for lists, it would set the resulting VariableTracker's source to `None` which would utilize a different reconstruct path in codegen. Now this is handled explicitly by reconstructing vars when allow_cache=`False`, so that during side effect replay, the mutated var is correctly updated. In subsequent PRs: * Refactoring side effect tracking to be significantly simpler (I think we only need an `is_modified` flag) * Refactor `next_variables` iterator to match the signature of `next` * Remove all references to `options` in the code * Refactor VTs representing mutable collections to implement their own mutation update handling * Remove clone and/or make it specific to lists for creating slices * Add mutation tracking/replay for sets * Add mutation tracking/replay for iter.py * Removing setting source in builder (it's set at the top level after a var is returned) Pull Request resolved: https://github.com/pytorch/pytorch/pull/113725 Approved by: https://github.com/jansel
536 lines
22 KiB
Python
536 lines
22 KiB
Python
import inspect
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
import torch.nn
|
|
|
|
from . import utils, variables
|
|
from .bytecode_transformation import (
|
|
create_call_function,
|
|
create_call_method,
|
|
create_instruction,
|
|
)
|
|
from .codegen import PyCodegen
|
|
from .exc import unimplemented
|
|
from .source import LocalSource, Source
|
|
from .utils import nn_module_new, object_new
|
|
from .variables.base import (
|
|
is_side_effect_safe,
|
|
MutableLocalBase,
|
|
MutableLocalSource,
|
|
VariableTracker,
|
|
)
|
|
|
|
|
|
class MutableSideEffects(MutableLocalBase):
|
|
"""
|
|
VariableTracker.mutable_local marker to indicate a list passed as
|
|
an input that if we mutate we need to re-apply those mutations after
|
|
the graph runs.
|
|
"""
|
|
|
|
def __init__(self, source: Source, is_modified: bool = False):
|
|
super().__init__(MutableLocalSource.Existing)
|
|
self.source = source
|
|
self.is_modified = is_modified
|
|
|
|
|
|
class AttributeMutation(MutableLocalBase):
|
|
"""
|
|
VariableTracker.mutable_local marker to track changes to attributes
|
|
"""
|
|
|
|
def __init__(self, typ: MutableLocalSource, source: Optional[Source]):
|
|
super().__init__(typ)
|
|
self.source = source
|
|
|
|
|
|
class AttributeMutationExisting(AttributeMutation):
|
|
def __init__(self, source: Source):
|
|
super().__init__(MutableLocalSource.Existing, source)
|
|
self.source = source
|
|
|
|
|
|
class AttributeMutationNew(AttributeMutation):
|
|
def __init__(self, source: Optional[Source], cls_source: Optional[Source]):
|
|
super().__init__(MutableLocalSource.Local, source)
|
|
self.cls_source = cls_source
|
|
|
|
|
|
class SideEffects:
|
|
"""
|
|
Track side effects (list mutation, setattr, etc) that need to be
|
|
applied after an FX graph is run.
|
|
"""
|
|
|
|
id_to_variable: Dict[int, VariableTracker]
|
|
store_attr_mutations: Dict[MutableLocalBase, Dict[str, VariableTracker]]
|
|
keepalive: List[Any]
|
|
|
|
def __init__(
|
|
self,
|
|
id_to_variable=None,
|
|
store_attr_mutations=None,
|
|
keepalive=None,
|
|
save_for_backward=None,
|
|
tensor_hooks=None,
|
|
):
|
|
super().__init__()
|
|
self.id_to_variable = id_to_variable or {}
|
|
self.store_attr_mutations = store_attr_mutations or {}
|
|
self.keepalive = keepalive or []
|
|
self.save_for_backward = save_for_backward or []
|
|
self.tensor_hooks = tensor_hooks or {}
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
assert isinstance(other, SideEffects)
|
|
# NB: do NOT test keepalive
|
|
return (
|
|
self.id_to_variable == other.id_to_variable
|
|
and self.store_attr_mutations == other.store_attr_mutations
|
|
and self.save_for_backward == other.save_for_backward
|
|
and self.tensor_hooks == other.tensor_hooks
|
|
)
|
|
|
|
def diff(self, other: "SideEffects") -> Optional[str]:
|
|
if self.id_to_variable != other.id_to_variable:
|
|
sk_itv = self.id_to_variable.keys()
|
|
ok_itv = other.id_to_variable.keys()
|
|
if sk_itv != ok_itv:
|
|
return f"id_to_variable keys: {sk_itv} != {ok_itv}"
|
|
# Feel free to augment this with more fancy diffing logic
|
|
# if needed for debugging
|
|
return "id_to_variable: unknown diff"
|
|
elif self.store_attr_mutations != other.store_attr_mutations:
|
|
sk_sam = self.store_attr_mutations.keys()
|
|
ok_sam = other.store_attr_mutations.keys()
|
|
if sk_sam != ok_sam:
|
|
return f"store_attr_mutations keys: {sk_sam} != {ok_sam}"
|
|
return "store_attr_mutations: unknown diff"
|
|
elif self.save_for_backward != other.save_for_backward:
|
|
return "save_for_backward"
|
|
else:
|
|
return None
|
|
|
|
def clone(self):
|
|
"""Create a shallow copy"""
|
|
return self.__class__(
|
|
id_to_variable=dict(self.id_to_variable),
|
|
store_attr_mutations={
|
|
k: dict(v) for k, v in self.store_attr_mutations.items()
|
|
},
|
|
keepalive=list(self.keepalive),
|
|
save_for_backward=self.save_for_backward,
|
|
tensor_hooks=self.tensor_hooks,
|
|
)
|
|
|
|
def apply(self, fn, cache=None, skip_fn=lambda _: False):
|
|
if cache is None:
|
|
cache = dict()
|
|
|
|
self.id_to_variable = {
|
|
k: VariableTracker.apply(fn, v, cache, skip_fn)
|
|
for k, v in self.id_to_variable.items()
|
|
}
|
|
self.store_attr_mutations = {
|
|
k: VariableTracker.apply(fn, v, cache, skip_fn)
|
|
for k, v in self.store_attr_mutations.items()
|
|
}
|
|
self.save_for_backward = VariableTracker.apply(
|
|
fn, self.save_for_backward, cache, skip_fn
|
|
)
|
|
self.tensor_hooks = VariableTracker.apply(fn, self.tensor_hooks, cache, skip_fn)
|
|
|
|
def __contains__(self, item):
|
|
return id(item) in self.id_to_variable
|
|
|
|
def __getitem__(self, item):
|
|
return self.id_to_variable[id(item)]
|
|
|
|
def check_allowed_side_effect(self, item):
|
|
from torch._dynamo.variables.misc import AutogradFunctionContextVariable
|
|
|
|
# People do things like self.dim = dim inside autograd.Function.
|
|
# These are benign.
|
|
if isinstance(item, AutogradFunctionContextVariable):
|
|
return True
|
|
if not is_side_effect_safe(item.mutable_local):
|
|
unimplemented(
|
|
"HigherOrderOperator: Mutating a variable not in the current scope (SideEffects)"
|
|
)
|
|
|
|
def store_attr(self, item: VariableTracker, name: str, value: VariableTracker):
|
|
assert self.is_attribute_mutation(item)
|
|
self.check_allowed_side_effect(item)
|
|
if item.mutable_local not in self.store_attr_mutations:
|
|
self.store_attr_mutations[item.mutable_local] = {}
|
|
self.store_attr_mutations[item.mutable_local][name] = value
|
|
|
|
def load_attr(self, item, name, deleted_ok=False):
|
|
assert self.is_attribute_mutation(item)
|
|
result = self.store_attr_mutations[item.mutable_local][name]
|
|
if not deleted_ok and isinstance(result, variables.DeletedVariable):
|
|
unimplemented("read deleted attribute")
|
|
return result
|
|
|
|
def store_cell(self, cellvar, value):
|
|
assert isinstance(cellvar, variables.NewCellVariable)
|
|
assert isinstance(value, variables.VariableTracker)
|
|
self.store_attr(cellvar, "cell_contents", value)
|
|
|
|
def load_cell(self, cellvar):
|
|
assert isinstance(cellvar, variables.NewCellVariable)
|
|
return self.load_attr(cellvar, "cell_contents")
|
|
|
|
def load_global(self, gvar: VariableTracker, name: str):
|
|
assert isinstance(gvar, variables.VariableTracker)
|
|
return self.load_attr(gvar, name)
|
|
|
|
def store_global(self, gvar: VariableTracker, name: str, value: VariableTracker):
|
|
assert isinstance(gvar, variables.VariableTracker)
|
|
assert isinstance(value, variables.VariableTracker)
|
|
self.store_attr(gvar, name, value)
|
|
|
|
@staticmethod
|
|
def cls_supports_mutation_side_effects(cls):
|
|
return inspect.getattr_static(cls, "__setattr__", None) in (
|
|
object.__setattr__,
|
|
torch.nn.Module.__setattr__,
|
|
)
|
|
|
|
def is_attribute_mutation(self, item):
|
|
return isinstance(item.mutable_local, AttributeMutation)
|
|
|
|
def has_pending_mutation(self, item):
|
|
return self.is_attribute_mutation(item) and bool(
|
|
self.store_attr_mutations.get(item.mutable_local)
|
|
)
|
|
|
|
def is_modified(self, item):
|
|
if isinstance(item.mutable_local, AttributeMutationNew):
|
|
return True
|
|
if self.is_attribute_mutation(item):
|
|
return item.mutable_local in self.store_attr_mutations
|
|
return item.mutable_local.is_modified
|
|
|
|
def _track_obj(
|
|
self,
|
|
item: Any,
|
|
variable: VariableTracker,
|
|
mutable_cls=MutableSideEffects,
|
|
):
|
|
"""Start tracking a new variable for mutation"""
|
|
assert variable.source is not None
|
|
variable.mutable_local = mutable_cls(variable.source)
|
|
self.id_to_variable[id(item)] = variable
|
|
self.keepalive.append(item)
|
|
return variable
|
|
|
|
track_mutable = _track_obj
|
|
|
|
def track_object_existing(
|
|
self,
|
|
item: Any,
|
|
variable: VariableTracker,
|
|
):
|
|
return self._track_obj(item, variable, mutable_cls=AttributeMutationExisting)
|
|
|
|
def track_object_new(
|
|
self,
|
|
cls_source: Source,
|
|
user_cls: Any,
|
|
variable_cls: Any,
|
|
options,
|
|
):
|
|
if user_cls is torch.autograd.function.FunctionCtx:
|
|
obj = torch.autograd.Function()
|
|
elif issubclass(user_cls, torch.nn.Module):
|
|
obj = nn_module_new(user_cls)
|
|
else:
|
|
obj = object_new(user_cls)
|
|
variable = variable_cls(
|
|
obj,
|
|
mutable_local=AttributeMutationNew(None, cls_source),
|
|
**options,
|
|
)
|
|
self.id_to_variable[id(obj)] = variable
|
|
self.keepalive.append(obj)
|
|
return variable
|
|
|
|
def track_cell_new(
|
|
self,
|
|
):
|
|
obj = object()
|
|
variable = variables.NewCellVariable(
|
|
mutable_local=AttributeMutationNew(None, None),
|
|
)
|
|
self.id_to_variable[id(obj)] = variable
|
|
self.keepalive.append(obj)
|
|
return variable
|
|
|
|
def track_cell_existing(self, source: Source, item: Any):
|
|
variable = variables.NewCellVariable(
|
|
mutable_local=AttributeMutationExisting(source),
|
|
)
|
|
self.id_to_variable[id(item)] = variable
|
|
self.keepalive.append(item)
|
|
return variable
|
|
|
|
def track_global_existing(self, source: Source, item: Any):
|
|
variable = variables.NewGlobalVariable(
|
|
mutable_local=AttributeMutationExisting(source),
|
|
)
|
|
self.id_to_variable[id(item)] = variable
|
|
self.keepalive.append(item)
|
|
return variable
|
|
|
|
def track_save_for_backward(self, ctx, args):
|
|
assert isinstance(ctx, variables.AutogradFunctionContextVariable)
|
|
self.save_for_backward.append((ctx, args))
|
|
|
|
def prune_dead_object_new(self, tx):
|
|
live_new_objects = set()
|
|
skip_obj = None
|
|
|
|
def visit(var: VariableTracker):
|
|
if (
|
|
isinstance(var.mutable_local, AttributeMutationNew)
|
|
and var.mutable_local is not skip_obj
|
|
):
|
|
live_new_objects.add(var.mutable_local)
|
|
return var
|
|
|
|
def is_live(var: Union[MutableLocalBase, VariableTracker]):
|
|
if isinstance(var, AttributeMutationNew):
|
|
return var in live_new_objects
|
|
if isinstance(var, VariableTracker):
|
|
return is_live(var.mutable_local)
|
|
return True
|
|
|
|
VariableTracker.apply(visit, (tx.stack, tx.symbolic_locals))
|
|
for var in self.id_to_variable.values():
|
|
if not isinstance(var.mutable_local, AttributeMutationNew):
|
|
VariableTracker.apply(visit, var)
|
|
|
|
for skip_obj, setattrs in self.store_attr_mutations.items():
|
|
VariableTracker.apply(visit, setattrs)
|
|
|
|
self.id_to_variable = {
|
|
k: v for k, v in self.id_to_variable.items() if is_live(v)
|
|
}
|
|
self.store_attr_mutations = {
|
|
k: v for k, v in self.store_attr_mutations.items() if is_live(k)
|
|
}
|
|
|
|
def mutation(self, var):
|
|
self.check_allowed_side_effect(var)
|
|
if isinstance(var.mutable_local, MutableSideEffects):
|
|
var.mutable_local = MutableSideEffects(var.mutable_local.source, True)
|
|
|
|
def _get_modified_vars(self):
|
|
return [var for var in self.id_to_variable.values() if self.is_modified(var)]
|
|
|
|
def codegen_save_tempvars(self, cg: PyCodegen):
|
|
for var in self._get_modified_vars():
|
|
if isinstance(
|
|
var.mutable_local, (AttributeMutationExisting, AttributeMutationNew)
|
|
) and isinstance(var, variables.NewCellVariable):
|
|
cg.load_import_from(utils.__name__, "make_cell")
|
|
cg.extend_output(create_call_function(0, True))
|
|
cg.add_cache(var)
|
|
if isinstance(var.mutable_local, AttributeMutationNew):
|
|
var.mutable_local.source = LocalSource(cg.tempvars[var])
|
|
elif isinstance(var.mutable_local, AttributeMutationNew):
|
|
if isinstance(var, variables.AutogradFunctionContextVariable):
|
|
unimplemented("AutogradFunctionContextVariable escaped")
|
|
if "__call_nn_module_init" in self.store_attr_mutations.get(
|
|
var.mutable_local, {}
|
|
):
|
|
assert isinstance(var, variables.UnspecializedNNModuleVariable)
|
|
cg.load_import_from(utils.__name__, "nn_module_new")
|
|
else:
|
|
cg.load_import_from(utils.__name__, "object_new")
|
|
cg(var.mutable_local.cls_source)
|
|
cg.extend_output(create_call_function(1, True))
|
|
cg.add_cache(var)
|
|
var.mutable_local.source = LocalSource(cg.tempvars[var])
|
|
elif var in cg.tempvars:
|
|
assert cg.tempvars.get(var) is None
|
|
# subsequent usage should point to the original variable
|
|
cg(var.mutable_local.source)
|
|
cg.add_cache(var)
|
|
|
|
for ctx, args in self.save_for_backward:
|
|
cg(ctx.source)
|
|
cg.extend_output(
|
|
[create_instruction("LOAD_METHOD", argval="save_for_backward")]
|
|
)
|
|
for arg in args:
|
|
cg(arg)
|
|
cg.extend_output(
|
|
[
|
|
*create_call_method(len(args)),
|
|
create_instruction("POP_TOP"),
|
|
]
|
|
)
|
|
|
|
def register_hook(self, tensor, hook, handle, name):
|
|
idx = len(self.tensor_hooks.keys())
|
|
self.tensor_hooks[idx] = (tensor, hook, handle, name)
|
|
assert not handle.idx
|
|
handle.idx = idx
|
|
|
|
def remove_hook(self, idx):
|
|
del self.tensor_hooks[idx]
|
|
|
|
def codegen_hooks(self, cg):
|
|
for (
|
|
tensor,
|
|
hook,
|
|
handle,
|
|
name,
|
|
) in self.tensor_hooks.values():
|
|
# Note: [On tensor.register_hook]
|
|
#
|
|
# register_hook on a tensor, AKA backward hooks, have slightly nuanced differences in how they are implemented
|
|
# when it comes to hooks on objects with sources (inputs, params) vs objects without sources (intermediaries).
|
|
#
|
|
# For tensors with a source, we bypass direct inclusion of register_hook calls in the graph.
|
|
# Instead, these are tracked and stashed as a global variable, enabling their association with tensors in
|
|
# the residuals. During dynamo's frame creation, these hooks are invoked seamlessly on known reconstructible/fetch-able
|
|
# tensors. Because a source indicates knowledge of this object outside the torch compile region, and
|
|
# because we are running residuals firmly before .backward() can be run, it is sound to invoke
|
|
# `register_hook` on a known tensor.
|
|
#
|
|
# For tensors without a source, we support a limited subset of hooks. Global functions only, and
|
|
# compiled_autograd must be enabled or we will graph break.
|
|
#
|
|
# Handling the Handle: When a user retains the register_hook result in a handle, we intercept the
|
|
# STORE_FAST operation to record the user-designated local variable name. This ensures the reconstructed
|
|
# bytecode retains this name. If no handle is defined, we simply pop the generated value to keep the
|
|
# stack intact.
|
|
#
|
|
# Dynamo Tensor Hooks Workflow:
|
|
# - Functions passed to register_hook are lifted globally.
|
|
# - For tensors with sources:
|
|
# - In the "side_effects" phase of codegen, we iterate over tensors with hooks to:
|
|
# - Generate the tensor.
|
|
# - Issue a register_hook call on the tensor, linking to the globally stored function.
|
|
# - Incorporate a handle if one was established in the eager phase.
|
|
# - For tensors without sources:
|
|
# - We don't generate any instructions for registering a hook.
|
|
# - Handles from intermediary hooks are NYI.
|
|
# - We produce a call function that utilizes the trace_wrapped higher order op, closing over it.
|
|
# - We then manually insert the call function above into the graph.
|
|
# - The handle's exact user-specified name, "user_code_variable_name", is discerned and associated during STORE_FAST.
|
|
assert tensor.source, "Hooks on non input tensors NYI - should not get here"
|
|
cg(tensor)
|
|
cg.extend_output([cg.create_load_attr(name)])
|
|
cg(hook)
|
|
cg.extend_output(create_call_function(1, True))
|
|
# Let's go over how handles work.
|
|
#
|
|
# A handle is created from invoking `register_hook` on a tensor. A handle can be referenced at any
|
|
# time after that, or never. In dynamo, we track and associate a name with a handle (user_code_variable_name) to
|
|
# determine if a handle is accessed. If a handle has no user_code_variable_name, we just pop the produced value
|
|
# off the top of the stack, discarding the handle.
|
|
#
|
|
# If a handle is seen, we store it under that name. This is extremely important, because, the handle
|
|
# can be generated at any time after this point, and can be generated multiple times! If we were to defer
|
|
# actual codegen of the handle object until we saw a codegen call to it - then we would end up generating multiple
|
|
# register_hook calls, which is incorrect. This turns the codegen reconstruct(handle) call for the handle into
|
|
# essentially a lookup.
|
|
if (
|
|
hasattr(handle, "user_code_variable_name")
|
|
and handle.user_code_variable_name
|
|
):
|
|
# register_hook stored with variable name assigned to the handle
|
|
cg.extend_output([cg.create_store(handle.user_code_variable_name)])
|
|
else:
|
|
# register_hook stored w/o a variable name assigned to the handle
|
|
cg.extend_output([create_instruction("POP_TOP")])
|
|
|
|
def codegen_update_mutated(self, cg: PyCodegen):
|
|
suffixes = []
|
|
for var in self._get_modified_vars():
|
|
if isinstance(var, variables.ListVariable):
|
|
# old[:] = new
|
|
cg(var, allow_cache=False)
|
|
cg(var.mutable_local.source)
|
|
cg.extend_output(
|
|
[
|
|
cg.create_load_const(None),
|
|
cg.create_load_const(None),
|
|
create_instruction("BUILD_SLICE", arg=2),
|
|
]
|
|
)
|
|
suffixes.append([create_instruction("STORE_SUBSCR")])
|
|
elif isinstance(var, variables.ConstDictVariable):
|
|
cg.tx.output.update_co_names("clear")
|
|
cg.tx.output.update_co_names("update")
|
|
|
|
cg(var.mutable_local.source)
|
|
cg.extend_output([create_instruction("LOAD_METHOD", argval="update")])
|
|
cg(var, allow_cache=False)
|
|
|
|
cg(var.mutable_local.source)
|
|
cg.extend_output([create_instruction("LOAD_METHOD", argval="clear")])
|
|
|
|
suffixes.append(
|
|
[
|
|
*create_call_method(0), # clear
|
|
create_instruction("POP_TOP"),
|
|
*create_call_method(1), # update
|
|
create_instruction("POP_TOP"),
|
|
]
|
|
)
|
|
elif self.is_attribute_mutation(var):
|
|
for name, value in self.store_attr_mutations.get(
|
|
var.mutable_local, {}
|
|
).items():
|
|
if isinstance(var, variables.NewGlobalVariable):
|
|
cg.tx.output.update_co_names(name)
|
|
cg(value)
|
|
suffixes.append(
|
|
[create_instruction("STORE_GLOBAL", argval=name)]
|
|
)
|
|
elif name == "__call_nn_module_init":
|
|
pass # handled in codegen_save_tempvars
|
|
elif isinstance(value, variables.DeletedVariable):
|
|
if isinstance(
|
|
var.mutable_local, AttributeMutationExisting
|
|
) and hasattr(getattr(var, "value", None), name):
|
|
cg.tx.output.update_co_names(name)
|
|
cg(var.mutable_local.source)
|
|
suffixes.append(
|
|
[create_instruction("DELETE_ATTR", argval=name)]
|
|
)
|
|
else:
|
|
cg.tx.output.update_co_names(name)
|
|
cg(value)
|
|
cg(var.mutable_local.source)
|
|
suffixes.append([create_instruction("STORE_ATTR", argval=name)])
|
|
elif isinstance(var, variables.TupleIteratorVariable):
|
|
for _ in range(var.index):
|
|
cg.load_import_from(utils.__name__, "iter_next")
|
|
cg(var.mutable_local.source)
|
|
cg.extend_output(create_call_function(1, True))
|
|
cg.append_output(create_instruction("POP_TOP"))
|
|
else:
|
|
raise AssertionError(type(var))
|
|
|
|
# do all the actual mutations at the very end to handle dependencies
|
|
for suffix in reversed(suffixes):
|
|
cg.extend_output(suffix)
|
|
|
|
def is_empty(self):
|
|
return not (
|
|
any(map(self.is_modified, self.id_to_variable.values()))
|
|
or self.tensor_hooks
|
|
or self.save_for_backward
|
|
or self.tensor_hooks
|
|
)
|
|
|
|
def clear(self):
|
|
self.keepalive.clear()
|
|
self.id_to_variable.clear()
|