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
|
import warnings
from typing import Callable, Union
import torch
import torch.utils._pytree as pytree
from torch._ops import OpOverload
from torch._subclasses.fake_tensor import (
FakeTensorMode,
tree_flatten_only,
UnsupportedFakeTensorException,
)
from torch.utils._python_dispatch import TorchDispatchMode
from torch.utils._pytree import tree_flatten
aten = torch.ops.aten
def outputs_alias_inputs(outputs, inputs):
input_storages = {
inp.storage()._cdata
for inp in tree_flatten_only(torch.Tensor, inputs)
if torch._C._has_storage(inp)
}
return any(
torch._C._has_storage(out) and out.storage()._cdata in input_storages
for out in tree_flatten_only(torch.Tensor, outputs)
)
def outputs_are_inputs(outputs, inputs):
input_ids = {id(inp) for inp in tree_flatten_only(torch.Tensor, inputs)}
return any(id(out) in input_ids for out in tree_flatten_only(torch.Tensor, outputs))
def output_alias_each_other(outputs):
storages = set()
for out in tree_flatten_only(torch.Tensor, outputs):
if not torch._C._has_storage(out):
continue
stor = out.storage()._cdata
if stor in storages:
return True
storages.add(stor)
return False
class CrossRefFakeMode(TorchDispatchMode):
def __init__(
self,
ignore_op_fn: Union[Callable[[OpOverload], bool], None] = None,
*,
check_strides=True,
check_aliasing=True,
):
self.ignore_op_fn = (
ignore_op_fn if ignore_op_fn is not None else lambda fn: False
)
self.check_strides = check_strides
self.check_aliasing = check_aliasing
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
kwargs = kwargs or {}
fake_r = None
# empty_like excluded for now due to sparse complex
# aten._to_dense.default this one is getting called with csc
if (
func
not in (
aten.lift_fresh.default,
aten.lift_fresh_copy.default,
aten.set_.source_Storage_storage_offset,
)
and not self.ignore_op_fn(func)
and torch.Tag.dynamic_output_shape not in func.tags # type: ignore[attr-defined]
and torch.Tag.inplace_view not in func.tags # type: ignore[attr-defined]
and torch.Tag.data_dependent_output not in func.tags # type: ignore[attr-defined]
):
try:
with FakeTensorMode() as fake_mode:
fake_args, fake_kwargs = pytree.tree_map_only(
torch.Tensor, fake_mode.from_tensor, (args, kwargs)
)
with warnings.catch_warnings():
fake_r = func(*fake_args, **fake_kwargs)
except UnsupportedFakeTensorException:
pass
r = func(*args, **kwargs)
if fake_r is not None:
r_flat, _ = tree_flatten(r)
f_flat, _ = tree_flatten(fake_r)
assert len(r_flat) == len(
r_flat
), f"Mismatch {len(r_flat)} != {len(r_flat)} on {func}"
if self.check_aliasing:
r_aliasing = outputs_alias_inputs(r, (args, kwargs))
f_aliasing = outputs_alias_inputs(fake_r, (fake_args, fake_kwargs))
assert (
r_aliasing == f_aliasing
), f"Mismatch on {func}: {r_aliasing} != {f_aliasing}"
r_identity_eq = outputs_are_inputs(r, (args, kwargs))
f_identity_eq = outputs_are_inputs(fake_r, (fake_args, fake_kwargs))
assert (
r_identity_eq == f_identity_eq
), f"Mismatch on {func}: {r_identity_eq} != {f_identity_eq}"
r_output_alias_each_other = output_alias_each_other(r)
f_output_alias_each_other = output_alias_each_other(fake_r)
assert (
r_output_alias_each_other == f_output_alias_each_other
), f"Mismatch on {func}: {r_output_alias_each_other} != {f_output_alias_each_other}"
for r_out, fake_out in zip(tree_flatten(r)[0], tree_flatten(fake_r)[0]):
r_is_ten = isinstance(r_out, torch.Tensor)
assert r_is_ten == isinstance(
fake_out, torch.Tensor
), f"Mismatched number of tensor outputs on {func}"
if r_is_ten:
assert (
r_out.requires_grad == fake_out.requires_grad
), f"Mismatch on {func}"
if torch._C._has_storage(r_out):
r_offset = r_out.storage_offset()
f_offset = fake_out.storage_offset()
assert (
r_offset == f_offset
), f"Mismatch on {func}: {r_offset} != {f_offset}"
try:
torch._prims.utils.compare_tensor_meta(
r_out, fake_out, check_strides=self.check_strides
)
except Exception as e:
raise RuntimeError(f"Mismatch on {func}: {e}")
return r
|
