pytorch/torch/_dynamo/side_effects.py
Michael Voznesensky 02f6a8126e Support a simple subset of functions as backward hooks on intermediate tensors (#109537)
The main thrust of the initial effort here was to capture `register_hook` calls on tensors in compile regions. The first part of this was done in https://github.com/pytorch/pytorch/pull/108903 wherein we added support for register_hook input tensors.

The distinction between input and intermediary is due to implementation differences.

There are 2 kinds of hooks:

1) Hooks on objects with sources (inputs, params)
2) Hooks on objects w/o sources (intermediaries, and outputs).

Note: As outputs can be made simple by how dynamo handles residuals, they could actually be handled as if they were inputs, but, for the sake of this PR, we will refer to hooks as either hooks on inputs (sourced), or hooks on intermediaries (not sourced).

**The plan:**

For tensors w/ a source: (The PR above)
We record registered hooks, store them as a global, and associate them with the tensor in residuals. This means that when dynamo goes to create the frame, where we produce bytecode to stitch together our PT2 modified bytecode with the original eager code, we call register_hook. This registration of hooks in residuals is sound because (a) it happens right after a Pt2 frame region ends and (b) we know that the tensor is alive in f_locals, f_globals, or a module in the users invoking frame. This means we can soundly know it will be around to invoke register_hook on. As long as we guard on the identity of the lifted function, this is sound to do.

For tensors w/o a source: (This PR)

Ostensibly, the most correct and complete solution would be to smuggle hooks into a runtime wrapper in aot_autograd, where all the items the hooks close over are lifted to inputs as necessary and passed alongside the user provided function. This is necessary so that we can properly trace out and capture all the mutations within the user defined hook at backwards time.

This is too complicated - so, we limited the scope of this initial PR to a simple subset of hooks:

- Hooks must have a source (be known to us already, not a lambda or intermediary defined function)
- We must be tracing under compiled autograd

**The flow**:

We use the HOP added in https://github.com/pytorch/pytorch/pull/109690/files, referred to as the HOP below.

1) We intercept register_hook calls and wrap the user defined fn in the HOP
2) We write a `_register_hook_trampoline` to the graph that is a local no-arg function that is invoked as a call_function in the dynamo graph
3) aot_autograd inlines through it during its trace, and sees the HOP
4) the HOP preserves itself in the graph - it does not get traced into
5) During backwards, compiled_autograd installs the HOP under a hook call
6) When compiled_autograd enters compilation over its generated graph, dynamo traces the contents of the hook

Pull Request resolved: https://github.com/pytorch/pytorch/pull/109537
Approved by: https://github.com/ezyang
2023-10-11 01:35:37 +00:00

536 lines
22 KiB
Python

import collections
import inspect
from typing import Any, Dict, List, Optional
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: 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: Source, cls_source: 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[AttributeMutation, 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 collections.OrderedDict()
self.store_attr_mutations = store_attr_mutations or collections.OrderedDict()
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=collections.OrderedDict(self.id_to_variable),
store_attr_mutations=collections.OrderedDict(
(k, collections.OrderedDict(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 = collections.OrderedDict(
(k, VariableTracker.apply(fn, v, cache, skip_fn))
for k, v in self.id_to_variable.items()
)
self.store_attr_mutations = collections.OrderedDict(
(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] = collections.OrderedDict()
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,
source: Source,
item: Any,
variable: VariableTracker,
mutable_cls=MutableSideEffects,
):
"""Start tracking a new variable for mutation"""
variable = variable.clone(mutable_local=mutable_cls(source), source=source)
self.id_to_variable[id(item)] = variable
self.keepalive.append(item)
return variable
track_list = _track_obj
track_dict = _track_obj
def track_object_existing(
self,
source: Source,
item: Any,
variable: VariableTracker,
):
return self._track_obj(
source, 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: 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 = collections.OrderedDict(
(k, v) for k, v in self.id_to_variable.items() if is_live(v)
)
self.store_attr_mutations = collections.OrderedDict(
(k, v) for k, v in self.store_attr_mutations.items() if is_live(k)
)
def mutation(self, oldvar, newvar):
self.check_allowed_side_effect(oldvar)
return newvar.clone(
mutable_local=MutableSideEffects(oldvar.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):
idx = len(self.tensor_hooks.keys())
self.tensor_hooks[idx] = (tensor, hook, handle)
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,
) 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("register_hook")])
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)])
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()