mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
081c5b3adc
Summary: The primary problem we are setting out to solve here is fake tensor freshness. Before this PR, fake tensors after dynamo represented fake tensors *at the end* of trace, so subsequent retraces like aot_autograd would start off with fake tensors in the wrong (end result) state, rather than their expected fresh state. The solution here is to start a fresh fake mode, and re-fakify the tensors. The nuance comes from ensuring that symbols are uniformly created for the symbolic sizes and strides of the tensor. This PR is the result of *a lot* of back and forth with ezyang and eellison. Initially, the first pass at this was not super different from what we have in the PR - the broad strokes were the same: 1) We cache source->symbol in shape_env 2) We pass policy objects around, stored at dynamo fakificaiton time, and reused for later fakification 3) We create a new fake mode for backends (from https://github.com/pytorch/pytorch/pull/113605/files) This is ugly, and has some layering violations. We detoured our decision making through a few other alternatives. Immutable/mutable fake tensor mode was the most interesting alternative, https://github.com/pytorch/pytorch/pull/113653, and was struck down on concerns of complexity in fake mode combined with it not covering all edge cases. We also detoured on what to do about tensor memoization returning back potentially different tensors than requested, and if that was an anti pattern (it is) we want to hack in with the symbol cache (we don't). We went back to the drawing board here, but with a few concessions: 1) the cache for source->symbol must live outside of shape_env, for both lifecycle, and layering reasons 2) A good amount of work needs to be done to pipe policy around fake_mode and meta_utils correctly, to cover all the cases (ezyang did this) cc penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng imported-using-ghimport Test Plan: Imported from OSS Reviewed By: huydhn, Chillee Differential Revision: D51566250 Pulled By: voznesenskym Pull Request resolved: https://github.com/pytorch/pytorch/pull/114526 Approved by: https://github.com/Chillee, https://github.com/huydhn
731 lines
34 KiB
Python
731 lines
34 KiB
Python
import contextlib
|
|
import warnings
|
|
import weakref
|
|
from typing import ContextManager, List, Optional, Tuple, TYPE_CHECKING
|
|
|
|
import torch
|
|
from torch._C._functorch import (
|
|
_unwrap_functional_tensor,
|
|
_wrap_functional_tensor,
|
|
current_level,
|
|
peek_interpreter_stack,
|
|
TransformType,
|
|
)
|
|
from torch._guards import Source
|
|
|
|
from torch.multiprocessing.reductions import StorageWeakRef
|
|
from torch.utils._python_dispatch import (
|
|
is_traceable_wrapper_subclass,
|
|
transform_subclass,
|
|
)
|
|
from torch.utils.weak import WeakIdRef
|
|
|
|
if TYPE_CHECKING:
|
|
# Import the following modules during type checking to enable code intelligence features,
|
|
# Do not import unconditionally, as they import sympy and importing sympy is very slow
|
|
from torch.fx.experimental.symbolic_shapes import SymbolicContext
|
|
|
|
DimList = List
|
|
|
|
|
|
def safe_is_leaf(t):
|
|
try:
|
|
return t.is_leaf
|
|
except RuntimeError:
|
|
# inference mode can trigger this
|
|
return False
|
|
|
|
|
|
def safe_grad(t):
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", "The .grad attribute of a Tensor")
|
|
return t.grad
|
|
|
|
|
|
def assert_eq(a, b):
|
|
assert a == b, f"{a} != {b}"
|
|
|
|
|
|
def assert_metadata_eq(assert_eq, m1, m2, *, skip_symbolic=False):
|
|
def go(m1, m2):
|
|
assert_eq(m1.dtype, m2.dtype)
|
|
if not skip_symbolic:
|
|
assert_eq(m1.shape, m2.shape)
|
|
assert_eq(m1.requires_grad, m2.requires_grad)
|
|
assert_eq(m1.is_leaf, m2.is_leaf)
|
|
assert_eq(m1.grad_fn is None, m2.grad_fn is None)
|
|
assert_eq(m1.is_sparse, m2.is_sparse)
|
|
assert_eq(m1.is_inference(), m2.is_inference())
|
|
assert_eq(m1.is_conj(), m2.is_conj())
|
|
assert_eq(m1.is_neg(), m2.is_neg())
|
|
assert_eq(safe_grad(m1) is not None, safe_grad(m2) is not None)
|
|
if safe_grad(m1) is not None:
|
|
go(safe_grad(m1), safe_grad(m2))
|
|
if m1.is_sparse:
|
|
assert_eq(m1.dense_dim(), m2.dense_dim())
|
|
assert_eq(m1.sparse_dim(), m2.sparse_dim())
|
|
assert_eq(m1.is_coalesced(), m2.is_coalesced())
|
|
else:
|
|
if not skip_symbolic:
|
|
assert_eq(m1.stride(), m2.stride())
|
|
assert_eq(m1.storage_offset(), m2.storage_offset())
|
|
assert_eq(m1._is_view(), m2._is_view())
|
|
if m1._is_view():
|
|
go(m1._base, m2._base)
|
|
# TODO: test if is resizable (no direct query for this atm)
|
|
# TODO: audit AutogradMeta to see if it matches
|
|
# TODO: test forward AD
|
|
|
|
return go(m1, m2)
|
|
|
|
|
|
# This is a class for converting multiple tensors into meta tensors which
|
|
# share the same view/storage structure. The operation model is you allocate
|
|
# one of these, and then call it repeatedly on all the tensors you want to
|
|
# convert. It's important to use the same object for tensors you want to
|
|
# share storage because this is how we correlate shared storages to the same
|
|
# meta storages. This class will hold weak references to cached tenosrs
|
|
# and tensor storages.
|
|
class MetaConverter:
|
|
def __init__(self):
|
|
self.storage_memo = {}
|
|
self.tensor_memo: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
|
|
self.maybe_storages_to_delete = []
|
|
self.check_expired_frequency = 128
|
|
self.check_expired_count = 0
|
|
self.hit = 0
|
|
self.miss = 0
|
|
self.del_hook = None
|
|
self.arg_cnt = 0
|
|
|
|
def successful(self):
|
|
return self.hit > 0 and self.miss == 0
|
|
|
|
def check_for_expired_weak_storages(self):
|
|
new_li = []
|
|
stor_to_delete = []
|
|
for obj in self.maybe_storages_to_delete:
|
|
if not obj.expired():
|
|
new_li.append(obj)
|
|
else:
|
|
stor_to_delete.append(obj)
|
|
for obj in stor_to_delete:
|
|
self.storage_memo.pop(obj, None)
|
|
self.maybe_storages_to_delete = new_li
|
|
|
|
# if for some reason we have aquired many storages which have not expired
|
|
# even though a tensor with their storage has expired (aliasing or otherwise)
|
|
# check for expired storages less often so as to bound the amount of work we
|
|
# do checking for expired storages
|
|
self.check_expired_frequency = max(
|
|
self.check_expired_frequency, len(self.maybe_storages_to_delete)
|
|
)
|
|
|
|
def get_tensor_memo(self, t):
|
|
return self.tensor_memo.get(WeakIdRef(t), None)
|
|
|
|
def set_tensor_memo(self, t, v):
|
|
# hold a weak ref to self, otherwise it will be kept alive
|
|
# by the del_ten closure
|
|
self_weak_ref = weakref.ref(self)
|
|
if t.is_sparse or t.is_mkldnn:
|
|
weak_st = None
|
|
else:
|
|
weak_st = StorageWeakRef(t._typed_storage())
|
|
tensor_ref_key = WeakIdRef(t)
|
|
|
|
def del_ten():
|
|
# tensor outlives the converter
|
|
self_ref = self_weak_ref()
|
|
if self_ref is None:
|
|
return
|
|
# on shutdown, tensor_ref_key may not be in memo
|
|
self_ref.tensor_memo.pop(tensor_ref_key, None)
|
|
if weak_st and weak_st.expired():
|
|
self_ref.storage_memo.pop(weak_st, None)
|
|
elif weak_st is not None:
|
|
# [expired-storages]
|
|
# NB: even though the tensor has died,
|
|
# the deallocation of its storage can take longer,
|
|
# even when the storage has no other uses/views.
|
|
# In this case, the StorageWeakRef object will be kept alive
|
|
# longer than it needs to be, however the storage itself
|
|
# will be deallocated. We retain the possibly dead storages
|
|
# and periodically check if any of them are expired and
|
|
# can be freed.
|
|
self_ref.maybe_storages_to_delete.append(weak_st)
|
|
|
|
weakref.finalize(t, del_ten)
|
|
self.tensor_memo[tensor_ref_key] = v
|
|
|
|
# NB: doesn't actually return a storage, because meta storage is
|
|
# not supported
|
|
def meta_storage(self, s, callback):
|
|
# NB: TypedStorage is freshly allocated and cannot be used as hash
|
|
# key index.
|
|
|
|
# Use a Weak Ref to s in order to not leak memory
|
|
swr = StorageWeakRef(s)
|
|
if swr not in self.storage_memo:
|
|
self.storage_memo[swr] = callback(
|
|
lambda: torch.empty(s.size(), dtype=torch.uint8, device="meta")
|
|
).untyped_storage()
|
|
return self.storage_memo[swr]
|
|
|
|
# This function assumes that it's possible to do the conversion
|
|
# NB: name here is used in a conventional way by Dynamo; it corresponds
|
|
# precisely to the Source.name() of the tensor we're fakeifying and
|
|
# corresponds to a valid Python expression. When we construct sub-names
|
|
# as part of this process, we will maintain this invariant! (Even though
|
|
# other users of this may not need it this property to be upheld.)
|
|
def meta_tensor(
|
|
self,
|
|
t,
|
|
shape_env=None,
|
|
callback=lambda t: t(),
|
|
source: Optional[Source] = None,
|
|
symbolic_context: Optional["SymbolicContext"] = None,
|
|
):
|
|
from torch._subclasses.fake_tensor import FakeTensor
|
|
|
|
if source is None:
|
|
from torch._dynamo.source import ConstantSource
|
|
|
|
# TODO: make a dedicated UnknownSource for this?
|
|
source = ConstantSource(
|
|
f"__meta_utils_unknown_tensor{len(self.tensor_memo)}"
|
|
)
|
|
|
|
# This indicates you set no_dispatch() before calling into this
|
|
# function. This is an error: we may be creating fake tensors and
|
|
# will perform operations on them which need fake tensor mode to
|
|
# be active. You will segfault if you are in a no_dispatch() block.
|
|
assert not torch._C._dispatch_tls_local_exclude_set().has(
|
|
torch._C.DispatchKey.Python
|
|
)
|
|
arg_cnt = self.arg_cnt
|
|
self.arg_cnt += 1
|
|
|
|
# When we make as_strided calls, we end up generating a guard
|
|
# that the new as_strided tensor is in bounds for the old storage
|
|
# for the base (since as_strided calls can "bust" out of their
|
|
# bounding box.) This guard is unnecessary: if a user is able
|
|
# to provide us a tensor with the view base setup this way, we
|
|
# don't need to produce a guard, because the fact that they
|
|
# were able to produce the view base means its in bounds.
|
|
#
|
|
# Now, ordinarily, this guard would be harmless. However, the
|
|
# generated guard refers to variables bound on the base variable.
|
|
# At the moment, Dynamo doesn't actually guard on x._base, because
|
|
# according to Voz this results in a lot of spurious invalidations,
|
|
# and also if the user doesn't directly make use of _base, its
|
|
# pointless anyway (because programs should be parametric over
|
|
# whether or not the input tensor is a view or not--unless you're
|
|
# mutating the input, but that's a whole 'nother ballgame). So
|
|
# for expediency, we suppress these guards so we don't have to
|
|
# deal with this (yet, anyway.)
|
|
#
|
|
# NB: An old version of this code suppressed guards for ALL operations
|
|
# happening during meta conversion, not just as_strided calls.
|
|
# This is too aggressive: we do duck sizing and 0/1 simplification
|
|
# as we allocate variables, and we do need to register guards for
|
|
# these cases.
|
|
maybe_suppress = contextlib.nullcontext
|
|
if shape_env is not None:
|
|
maybe_suppress = shape_env.suppress_guards
|
|
|
|
def sym_sizes_strides_storage_offset(
|
|
t, src
|
|
) -> Tuple[Tuple[int, ...], Tuple[int, ...], int]:
|
|
if shape_env is not None:
|
|
if isinstance(t, FakeTensor) and t.fake_mode.shape_env is shape_env:
|
|
# Don't reallocate the sizes; the shape envs are the same,
|
|
# so reuse the old sizes/strides/etc
|
|
return (t.size(), t.stride(), t.storage_offset())
|
|
else:
|
|
return shape_env.create_symbolic_sizes_strides_storage_offset(
|
|
t,
|
|
src,
|
|
# Assume that the set of dims that are dynamic are the same between
|
|
# the wrapper tensor and any inner tensors.
|
|
# We can revisit this if this assumption does not hold
|
|
# for any important subclasses later.
|
|
symbolic_context=symbolic_context,
|
|
)
|
|
else:
|
|
assert symbolic_context is None
|
|
return (t.size(), t.stride(), t.storage_offset())
|
|
|
|
# see expired-storages
|
|
self.check_expired_count += 1
|
|
if self.check_expired_count >= self.check_expired_frequency:
|
|
self.check_for_expired_weak_storages()
|
|
self.check_expired_count = 0
|
|
|
|
if self.get_tensor_memo(t) is None:
|
|
with torch.inference_mode(t.is_inference()):
|
|
if t.is_sparse:
|
|
is_leaf = safe_is_leaf(t)
|
|
r = callback(
|
|
lambda: torch.ops.aten._sparse_coo_tensor_with_dims(
|
|
t.sparse_dim(),
|
|
t.dense_dim(),
|
|
t.shape,
|
|
dtype=t.dtype,
|
|
layout=torch.sparse_coo,
|
|
device="meta",
|
|
)
|
|
)
|
|
assert safe_is_leaf(r), "the callback you passed in doesn't detach"
|
|
# Note [is_coalesced is dispatched]
|
|
# Strangely enough, is_coalesced() is a dispatched operator,
|
|
# which means that it will get caught by fake tensor mode.
|
|
# Ordinarily this would error, but there's some logic in
|
|
# fake tensor ensure this doesn't happen.
|
|
r._coalesced_(t.is_coalesced())
|
|
if t.requires_grad:
|
|
r.requires_grad = True
|
|
if t.requires_grad and not is_leaf:
|
|
with torch.enable_grad():
|
|
r = r.clone()
|
|
r._coalesced_(t.is_coalesced())
|
|
elif t.is_mkldnn:
|
|
is_leaf = safe_is_leaf(t)
|
|
sizes, strides, _storage_offset = sym_sizes_strides_storage_offset(
|
|
t, source
|
|
)
|
|
r = callback(
|
|
lambda: torch.empty_strided(
|
|
sizes, strides, dtype=t.dtype, device="meta"
|
|
)
|
|
)
|
|
assert safe_is_leaf(r), "the callback you passed in doesn't detach"
|
|
if t.requires_grad:
|
|
r.requires_grad = True
|
|
if t.requires_grad and not is_leaf:
|
|
with torch.enable_grad():
|
|
r = r.clone()
|
|
elif t._is_view():
|
|
# Construct views in two steps: recursively meta-fy their
|
|
# base, and then create view(s) off that. NB: doing it
|
|
# directly from storage is WRONG because this won't cause
|
|
# version counters to get shared.
|
|
assert t._is_view()
|
|
|
|
from torch._dynamo.source import AttrSource
|
|
from torch.fx.experimental.symbolic_shapes import (
|
|
DimDynamic,
|
|
StatelessSymbolicContext,
|
|
)
|
|
|
|
if shape_env and not t.is_nested and not t._base.is_nested:
|
|
base_symbolic_context = StatelessSymbolicContext(
|
|
dynamic_sizes=[DimDynamic.STATIC] * t._base.dim(),
|
|
constraint_sizes=[None] * t._base.dim(),
|
|
)
|
|
else:
|
|
base_symbolic_context = None
|
|
base = self.meta_tensor(
|
|
t._base,
|
|
shape_env,
|
|
callback,
|
|
source=AttrSource(source, "_base"),
|
|
symbolic_context=base_symbolic_context,
|
|
)
|
|
|
|
def is_c_of_r(complex_dtype, real_dtype):
|
|
return (
|
|
utils.is_complex_dtype(complex_dtype)
|
|
and utils.corresponding_real_dtype(complex_dtype)
|
|
== real_dtype
|
|
)
|
|
|
|
# In some situations, MetaConverter may be called in a
|
|
# context where autograd is disabled. For the _is_view
|
|
# assert to pass, we have to setup the autograd view
|
|
# metadata anyway. Do this by reenabling the
|
|
# ADInplaceOrView key. This is kind of a hack.
|
|
old_exclude = torch._C._dispatch_tls_is_dispatch_key_excluded(
|
|
torch._C.DispatchKey.ADInplaceOrView
|
|
)
|
|
torch._C._dispatch_tls_set_dispatch_key_excluded(
|
|
torch._C.DispatchKey.ADInplaceOrView, False
|
|
)
|
|
try:
|
|
if base.dtype == t.dtype:
|
|
pass
|
|
elif is_c_of_r(base.dtype, t.dtype):
|
|
base = torch.view_as_real(base)
|
|
elif is_c_of_r(t.dtype, base.dtype):
|
|
base = torch.view_as_complex(base)
|
|
else:
|
|
# This is not guaranteed to succeed. If it fails, it
|
|
# means there is another dtype-converting view function
|
|
# that hasn't been handled here
|
|
base = base.view(t.dtype)
|
|
|
|
# This is very tricky. Naively, you might expect this
|
|
# to hold:
|
|
#
|
|
# if t.requires_grad and not safe_is_leaf(t)
|
|
# assert t._base.requires_grad
|
|
#
|
|
# But it's not true! As you can see in the following
|
|
# program:
|
|
#
|
|
# x = torch.zeros(4)
|
|
# y = x.view(1, 4)
|
|
# y.requires_grad = True
|
|
# z = y.view(1, 1, 4)
|
|
# assert z._base is x
|
|
#
|
|
# So we may have to do *two* views out of the base to
|
|
# recreate this situation.
|
|
def _view_from_base(base, t):
|
|
if t.is_nested:
|
|
# Nested tensors do not support as_strided, and
|
|
# hence,always have _view_func available.
|
|
#
|
|
# The unsafe version of _view_func omits
|
|
# checking whether the base passed in has the same
|
|
# metadata as the original base the view_func
|
|
# was originally executed with. (1) It is OK here,
|
|
# because we're calling it on the meta-ified base,
|
|
# so the metadata is guaranteed to be the same.
|
|
# (2) It is necessary because we don't actually
|
|
# want to guard on the base's metadata here.
|
|
return t._view_func_unsafe(base)
|
|
else:
|
|
(
|
|
sizes,
|
|
strides,
|
|
storage_offset,
|
|
) = sym_sizes_strides_storage_offset(t, source)
|
|
return base.as_strided(sizes, strides, storage_offset)
|
|
|
|
if safe_is_leaf(t):
|
|
# Leaf views that track view metadata are created by
|
|
# creating a view inside a no_grad block
|
|
with torch.no_grad(), maybe_suppress():
|
|
r = _view_from_base(base, t)
|
|
# As it's a leaf, we can directly assign requires_grad
|
|
r.requires_grad = t.requires_grad
|
|
else:
|
|
if t._base.requires_grad == t.requires_grad:
|
|
# Easy case, just run the view op
|
|
with torch.enable_grad(), maybe_suppress():
|
|
r = _view_from_base(base, t)
|
|
|
|
# NB: We don't actaully faithfully replicate
|
|
# autograd connectivity, but that doesn't matter
|
|
# today. See following for more info:
|
|
# https://gist.github.com/soulitzer/e03f015b314c3f5fcf80888c69390913
|
|
else:
|
|
# Obscure case. Create a leaf view and give it the
|
|
# correct requires_grad, then do the final view.
|
|
# NB: Can't have a non-leaf without requiring grad!
|
|
assert t.requires_grad
|
|
with torch.no_grad():
|
|
mid = base.view(base.shape)
|
|
mid.requires_grad = t.requires_grad
|
|
with torch.enable_grad(), maybe_suppress():
|
|
r = _view_from_base(mid, t)
|
|
# The CreationMeta influences whether or not inplace
|
|
# mutation is an error or not. So we need to make
|
|
# sure we properly propagate this as well.
|
|
torch._C._autograd._set_creation_meta(
|
|
r, torch._C._autograd._get_creation_meta(t)
|
|
)
|
|
finally:
|
|
torch._C._dispatch_tls_set_dispatch_key_excluded(
|
|
torch._C.DispatchKey.ADInplaceOrView, old_exclude
|
|
)
|
|
|
|
else:
|
|
is_leaf = safe_is_leaf(t)
|
|
if not t.is_nested:
|
|
# Nested tensor subclasses have special logic for
|
|
# creating symbolic size/strides/storage_offset
|
|
(
|
|
sizes,
|
|
strides,
|
|
storage_offset,
|
|
) = sym_sizes_strides_storage_offset(t, source)
|
|
|
|
def empty_create(inner_t, inner_src):
|
|
(
|
|
inner_sizes,
|
|
inner_strides,
|
|
inner_storage_offset,
|
|
) = sym_sizes_strides_storage_offset(inner_t, inner_src)
|
|
return torch.empty_strided(
|
|
inner_sizes,
|
|
inner_strides,
|
|
dtype=inner_t.dtype,
|
|
device="meta",
|
|
)
|
|
|
|
# If we have a subclass that desugars into dense tensors,
|
|
# perform our callback on each inner tensor.
|
|
if is_traceable_wrapper_subclass(t):
|
|
# Note: transform_subclass will use __tensor_unflatten__ to generate
|
|
# a fresh subclass wrapper, which is why sizes/strides are not passed in
|
|
# to the creation function here.
|
|
# We assume that if the inner tensors of the subclass are given symbolic sizes,
|
|
# their sizes will be used to construct the (symbolic) sizes of the wrapper tensor.
|
|
from torch._dynamo.source import AttrSource
|
|
|
|
if t.is_nested:
|
|
# Avoid circular import
|
|
from torch._dynamo.source import (
|
|
TensorProperty,
|
|
TensorPropertySource,
|
|
)
|
|
|
|
# For nested tensors, manually do transform_subclass
|
|
# so we can insert some special processing on ctx
|
|
attrs, ctx = t.__tensor_flatten__()
|
|
transformed_tensors_dict = {}
|
|
orig_shape_env = None
|
|
for attr in attrs:
|
|
inner_t = getattr(t, attr)
|
|
if orig_shape_env is None:
|
|
orig_shape_env = (
|
|
inner_t.fake_mode.shape_env
|
|
if isinstance(inner_t, FakeTensor)
|
|
else None
|
|
)
|
|
transformed_tensors_dict[attr] = callback(
|
|
lambda: empty_create(
|
|
inner_t, AttrSource(source, attr)
|
|
)
|
|
)
|
|
# We expect JaggedTensor to have a 'ragged_size' in
|
|
# its context
|
|
assert isinstance(ctx, dict)
|
|
assert "ragged_size" in ctx
|
|
assert isinstance(t._size[1], torch.SymInt)
|
|
if orig_shape_env is shape_env:
|
|
# It's already fake and the shape envs line up, reuse the old size
|
|
# Do not assert singleton_int; it may already
|
|
# be a variable
|
|
ctx["ragged_size"] = t._size[1]
|
|
else:
|
|
assert t._size[1].node.singleton_int() is not None
|
|
# Replace the eager ragged size with our freshly
|
|
# allocated jagged size that has a source
|
|
ctx["ragged_size"] = shape_env.create_symintnode(
|
|
shape_env.create_symbol(
|
|
t._size[1],
|
|
TensorPropertySource(
|
|
source, TensorProperty.SIZE, 1
|
|
),
|
|
),
|
|
hint=t._size[1],
|
|
)
|
|
r = type(t).__tensor_unflatten__(
|
|
transformed_tensors_dict, ctx
|
|
)
|
|
else:
|
|
r = transform_subclass(
|
|
t,
|
|
lambda attr, inner_t: callback(
|
|
lambda: empty_create(
|
|
inner_t,
|
|
AttrSource(source, attr),
|
|
)
|
|
),
|
|
)
|
|
else:
|
|
r = callback(
|
|
lambda: torch.empty_strided(
|
|
sizes,
|
|
strides,
|
|
dtype=t.dtype,
|
|
device="meta",
|
|
)
|
|
)
|
|
assert safe_is_leaf(r), "the callback you passed in doesn't detach"
|
|
if t.requires_grad:
|
|
r.requires_grad = t.requires_grad
|
|
if not is_leaf:
|
|
# Fake up some autograd history.
|
|
with torch.enable_grad():
|
|
# preserve_format is the default, but we want to
|
|
# emphasize how important it is to preserve
|
|
# format here
|
|
r = r.clone(memory_format=torch.preserve_format)
|
|
|
|
# Graph-Break for wrapped tensors
|
|
if torch._C._functorch.is_functorch_wrapped_tensor(t):
|
|
return NotImplemented
|
|
|
|
s = t.untyped_storage()
|
|
swr = StorageWeakRef(s)
|
|
if swr not in self.storage_memo and (
|
|
r.is_nested
|
|
or (
|
|
r.stride() == strides
|
|
and r.storage_offset() == storage_offset
|
|
)
|
|
):
|
|
# You're normal and happy, install the fresh storage into the memo
|
|
self.storage_memo[swr] = r.untyped_storage()
|
|
else:
|
|
# You're in crazy town; somehow you gave us a tensor
|
|
# that wasn't a view, but had nonzero storage offset,
|
|
# nontrivial strides (such that clone() couldn't
|
|
# preserve them), or already aliases with another
|
|
# tensor's storage. The most typical way to end
|
|
# up here is with set_. So use set_ to bludgeon this
|
|
# in.
|
|
r_s = self.meta_storage(s, callback=callback)
|
|
# NB: In principle, this should always work, but there
|
|
# is some subtle difference in the autograd metadata
|
|
# that means we will backprop the set_ call, even if
|
|
# r is declared as an input to grad.
|
|
# See https://github.com/pytorch/pytorch/issues/87956
|
|
# for the reproducer.
|
|
# NB: The in_kernel_invocation_manager here is necessary
|
|
# for fake tensor. If we run the set_ call with fake
|
|
# tensor on, r will improperly report that it is NOT a
|
|
# meta tensor but a cpu tensor, and then the set_ call
|
|
# will fail due to device mismatch. no_dispatch() is
|
|
# not enough, because the fake tensor will still claim
|
|
# to be a CPU tensor and you'll end up in the CPU
|
|
# kernel. Arguably this is a hack; a cleaner way to
|
|
# solve this is to have a FakeStorage concept which
|
|
# would report it's CPU device--no problem now! But
|
|
# this is difficult to do because we don't have storage
|
|
# subclasses. Relevant test is
|
|
# DynamicShapesFunctionTests::test_add_dynamic_shapes in
|
|
# test/dynamo/test_dynamic_shapes.py
|
|
maybe_fake_mgr: ContextManager[None] = contextlib.nullcontext()
|
|
from torch._subclasses.fake_tensor import (
|
|
in_kernel_invocation_manager,
|
|
maybe_get_fake_mode,
|
|
)
|
|
|
|
mb_fake_mode = maybe_get_fake_mode(r)
|
|
if mb_fake_mode is not None:
|
|
maybe_fake_mgr = in_kernel_invocation_manager(mb_fake_mode)
|
|
with maybe_fake_mgr, torch.no_grad():
|
|
r.set_(r_s, storage_offset, sizes, strides)
|
|
|
|
if safe_grad(t) is not None:
|
|
from torch._dynamo.source import AttrSource
|
|
|
|
r.grad = self.meta_tensor(
|
|
safe_grad(t),
|
|
shape_env,
|
|
callback,
|
|
source=AttrSource(source, "grad"),
|
|
symbolic_context=symbolic_context,
|
|
)
|
|
torch._C._set_conj(r, t.is_conj())
|
|
torch._C._set_neg(r, t.is_neg())
|
|
# This can be skipped if necessary for performance reasons
|
|
assert_metadata_eq(assert_eq, t, r, skip_symbolic=True)
|
|
self.set_tensor_memo(t, r)
|
|
|
|
return self.get_tensor_memo(t)
|
|
|
|
def __call__(
|
|
self,
|
|
t,
|
|
shape_env=None,
|
|
*,
|
|
callback=lambda t: t(),
|
|
source=None,
|
|
symbolic_context=None,
|
|
):
|
|
# TODO: zero tensors? We appear to have eliminated them by
|
|
# excluding complex for now
|
|
|
|
if isinstance(t, torch.Tensor) or is_traceable_wrapper_subclass(t):
|
|
if t.device.type != "xla" and any(
|
|
[
|
|
t.is_sparse_csr,
|
|
t.layout in [torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc],
|
|
t.is_quantized,
|
|
t._is_view() and t._base is not None and t._base.is_sparse,
|
|
torch._is_functional_tensor(t),
|
|
t.device.type in ("lazy"),
|
|
# We need a way to test if a tensor is batched but there
|
|
# is no official APi to do it
|
|
# torch._C._is_batched(t),
|
|
]
|
|
):
|
|
# TODO: sparse should support meta
|
|
# NB technically to('meta') does work but our logging
|
|
# instrumentation will see the meta conversions and the
|
|
# tests all break so we just exclude this. In any case
|
|
# the to conversion isn't really right anyhow.
|
|
|
|
if torch._is_functional_tensor(t) and t.device.type != "lazy":
|
|
if t._is_view():
|
|
raise RuntimeError(
|
|
"Cannot safely fakify a view because this process drops the view information right now."
|
|
)
|
|
|
|
st = peek_interpreter_stack()
|
|
assert (
|
|
st is None or st.key() == TransformType.Functionalize
|
|
), "Expect st to be either None or have Functionalize transform key."
|
|
if st is None:
|
|
# the case of AOTAutograd
|
|
torch._sync(t)
|
|
unwrap_t = torch._from_functional_tensor(t)
|
|
with torch._dispatch.python.suspend_functionalization():
|
|
fake_t = self.meta_tensor(
|
|
unwrap_t,
|
|
shape_env=shape_env,
|
|
callback=callback,
|
|
source=source,
|
|
symbolic_context=symbolic_context,
|
|
)
|
|
out = torch._to_functional_tensor(fake_t)
|
|
torch._mirror_autograd_meta_to(fake_t, out)
|
|
return out
|
|
else:
|
|
# torch.func.functionalize
|
|
reapply_views = torch._C._functionalization_reapply_views_tls()
|
|
unwrap_t = _unwrap_functional_tensor(t, reapply_views)
|
|
pop_st_ctx = (
|
|
torch._functorch.pyfunctorch.temporarily_pop_interpreter_stack()
|
|
)
|
|
with pop_st_ctx:
|
|
fake_t = self.meta_tensor(
|
|
unwrap_t,
|
|
shape_env=shape_env,
|
|
callback=callback,
|
|
source=source,
|
|
symbolic_context=symbolic_context,
|
|
)
|
|
return _wrap_functional_tensor(fake_t, current_level())
|
|
self.miss += 1
|
|
return NotImplemented
|
|
else:
|
|
self.hit += 1
|
|
r = self.meta_tensor(
|
|
t,
|
|
shape_env=shape_env,
|
|
callback=callback,
|
|
source=source,
|
|
symbolic_context=symbolic_context,
|
|
)
|
|
if type(t) is torch.nn.Parameter:
|
|
# NB: Cannot directly use Parameter constructor
|
|
# because that would force a detach, not desirable
|
|
r._is_param = True
|
|
return r
|
|
elif torch.overrides.is_tensor_like(t):
|
|
self.miss += 1
|
|
return NotImplemented
|
|
else:
|
|
# non-Tensor types don't count as hit or miss
|
|
return t
|
|
|
|
|
|
import torch._prims_common as utils
|