-
Notifications
You must be signed in to change notification settings - Fork 21.3k
/
proxy_tensor.py
821 lines (696 loc) · 33.4 KB
/
proxy_tensor.py
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
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import contextlib
import functools
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import torch
import torch.utils._pytree as pytree
from torch.fx import Tracer, GraphModule
from torch._subclasses.fake_tensor import FakeTensorMode
from torch._dispatch.python import enable_python_dispatcher
import torch.fx as fx
from torch.fx.passes.shape_prop import _extract_tensor_metadata
from contextlib import contextmanager, nullcontext
import inspect
from dataclasses import dataclass
import weakref
import operator
from torch.utils._stats import count
from torch.utils._python_dispatch import (
TorchDispatchMode,
_pop_mode_temporarily,
_get_current_dispatch_mode,
)
from torch._subclasses import FakeTensor
from .symbolic_shapes import ShapeEnv, SymDispatchMode, SymNode
from torch.fx import Proxy
import torch.fx.traceback as fx_traceback
from torch import SymInt, SymFloat, SymBool
from torch.utils.weak import WeakTensorKeyDictionary
__all__ = ["PythonKeyTracer", "dispatch_trace", "make_fx", "DecompositionInterpreter", "py_sym_types", "get_innermost_proxy_mode"]
aten = torch.ops.aten
prim = torch.ops.prim
CURRENT_DECOMPOSITION_TABLE: Dict[torch._ops.OpOverload, Callable] = {}
CONSTANT_NUMEL_LIMIT = 1
# We currently convert all SymInt to proxies before we use them.
# This could plausibly be handled at the Dynamo level.
pytree._register_pytree_node(torch.Size, lambda x: (list(x), None), lambda xs, _: tuple(xs))
def fake_signature(fn, nargs):
"""FX gets confused by varargs, de-confuse it"""
argnames = ",".join(f"arg{i}" for i in range(nargs))
return eval(f"lambda {argnames}: fn({argnames})", {"fn": fn})
@contextmanager
def decompose(decomposition_table):
global CURRENT_DECOMPOSITION_TABLE
old_decomposition_table = CURRENT_DECOMPOSITION_TABLE
CURRENT_DECOMPOSITION_TABLE = decomposition_table
try:
yield CURRENT_DECOMPOSITION_TABLE
finally:
CURRENT_DECOMPOSITION_TABLE = old_decomposition_table
# ensure we cannot collide with other properties
proxy_slot = object()
no_default = object()
py_sym_types = (SymInt, SymFloat, SymBool)
def is_sym_node(node):
assert hasattr(node, 'meta'), "All nodes traced with proxy_tensor should have meta"
return "val" in node.meta and isinstance(node.meta['val'], py_sym_types)
def set_proxy_slot(obj, tracer, proxy):
if isinstance(obj, torch.Tensor):
# We DO want to clobber proxies whenever we run an inplace operation
# on a tensor, and it affects the metadata on the proxy.
tracer.tensor_tracker[obj] = proxy
else:
# NB: Never clobber pre-existing proxy. Although the proxies
# are in principle equivalent, when we do graph partitioning
# we need there not to be spurious dependencies on tangent inputs.
# This works because primals get their SymInts set first, and
# THEN later we allocate tangent inputs. Make sure if a SymInt
# is derivable from a primal that we use that.
assert isinstance(obj, SymNode), type(obj)
if obj not in tracer.symnode_tracker:
tracer.symnode_tracker[obj] = proxy
def has_proxy_slot(obj, tracer):
assert isinstance(obj, (torch.Tensor, SymNode)), type(obj)
return get_proxy_slot(obj, tracer, False, lambda _: True)
# the default argument is what to return if the slot is not set.
# the transform argument is handy if you need to extract a subfield from
# the successfully looked up result (but NOT the default.)
def get_proxy_slot(obj, tracer, default=no_default, transform=lambda x: x):
if isinstance(obj, torch.Tensor):
tracker = tracer.tensor_tracker
else:
assert isinstance(obj, SymNode), type(obj)
tracker = tracer.symnode_tracker
if obj not in tracker:
if default is no_default:
raise RuntimeError(f"{obj} is not tracked with proxy for {tracer}")
return default
return transform(tracker[obj])
def snapshot_fake(val):
return val.detach()
def unwrap_proxy(proxy_mode, e):
if isinstance(e, torch.Tensor):
return get_proxy_slot(e, proxy_mode.tracer, e, lambda e: e.proxy)
elif isinstance(e, (torch.SymInt, torch.SymFloat, torch.SymBool)):
return get_proxy_slot(e.node, proxy_mode.tracer, e, lambda e: e())
else:
return e
# What invariants do we have for the 'val' set on the FX node? It has accurate
# metadata... but only for metadata that exists "below" all other subsystems
# (most notably autograd, but also vmap, functorch transforms, etc). This means
# you can get the dtype, shape, stride, storage, but you CANNOT get requires_grad,
# grad_fn, _base (_base actually may be set due to recursive call to
# ADInplaceOrView, but you shouldn't rely on it.)
def set_meta(proxy, val):
if isinstance(val, FakeTensor):
proxy.node.meta['val'] = snapshot_fake(val)
proxy.node.meta['tensor_meta'] = _extract_tensor_metadata(val)
elif isinstance(val, py_sym_types):
proxy.node.meta['val'] = val
elif isinstance(val, (list, tuple)):
if any(isinstance(x, FakeTensor) for x in val):
proxy.node.meta['val'] = [snapshot_fake(x) if isinstance(x, FakeTensor) else None for x in val]
elif isinstance(val, torch.Tensor):
if not val.is_sparse:
proxy.node.meta['tensor_meta'] = _extract_tensor_metadata(val)
# NB: Kinda hacky, but we should try to get val as the metadata
# everywhere
# TODO: This doesn't properly track storages. A more robust
# approach would be to maintain a per-trace FakeTensorMode and
# from_real_tensor to create fake values (don't forget to
# snapshot_fake)
fake_tensor_mode = FakeTensorMode(allow_fallback_kernels=True)
with fake_tensor_mode:
proxy.node.meta['val'] = torch.empty_strided(val.shape, val.stride(), device=val.device, dtype=val.dtype)
return proxy
def thunkify(f, *args, **kwargs):
"""
Delays computation of f until it's called again
Also caches the result
"""
return functools.lru_cache(1)(functools.partial(f, *args, **kwargs))
def track_tensor(tensor, proxy, *, constant, tracer):
def try_set_proxy_slot(outer_s, proxy_callable, *args):
assert callable(proxy_callable)
if isinstance(outer_s, SymInt):
inner_s = outer_s.node
set_proxy_slot(inner_s, tracer, thunkify(proxy_callable, outer_s, *args))
# The basic idea is that we need to associate each tensor/SymInt
# with a Proxy. How do we setup this association? We just store
# the proxy on the proxy slot of the object, keyed on the tracer
# (so that if we have multiple tracers at the same time, they
# don't clobber each other.)
for i, s in enumerate(tensor.shape):
try_set_proxy_slot(s, lambda x, i: set_meta(torch.ops.aten.sym_size(proxy, i), x), i)
for i, s in enumerate(tensor.stride()):
try_set_proxy_slot(s, lambda x, i: set_meta(torch.ops.aten.sym_stride(proxy, i), x), i)
try_set_proxy_slot(tensor.numel(), lambda x: set_meta(torch.ops.aten.sym_numel(proxy), x))
try_set_proxy_slot(tensor.storage_offset(), lambda x: set_meta(torch.ops.aten.sym_storage_offset(proxy), x))
set_proxy_slot(tensor, tracer, _ProxyTensor(proxy, constant))
def track_tensor_tree(inner_res, proxy_res, *, constant, tracer):
def wrap_with_proxy(e, proxy, constant):
if isinstance(e, torch.Tensor):
track_tensor(e, proxy, tracer=tracer, constant=constant)
set_meta(proxy, e)
elif isinstance(e, py_sym_types):
# NB: eagerly set meta here, so that the numbering is in order
set_meta(proxy, e)
set_proxy_slot(e.node, tracer, lambda: proxy)
elif isinstance(e, list):
# example use case: allreduce_ returns ([tensor], work)
for idx, ee in enumerate(e):
wrap_with_proxy(ee, proxy[idx], get_constant(idx))
def get_constant(idx):
if constant is None:
return None
else:
return constant[idx]
# Unfortunately, tree_map cannot directly be used here. As the resulting
# object may be a proxy that represents a tuple, we may need to
# explicitly unwrap the proxy by simulating the flattening operations.
if isinstance(inner_res, (tuple, list)):
if isinstance(proxy_res, fx.Proxy):
set_meta(proxy_res, inner_res)
for idx, e in enumerate(inner_res):
wrap_with_proxy(e, proxy_res[idx], get_constant(idx))
elif isinstance(inner_res, py_sym_types + (torch.Tensor,)):
wrap_with_proxy(inner_res, proxy_res, constant)
return inner_res
def maybe_disable_fake_tensor_mode():
# TODO: figure out if this API generally makes sense and bake it into the
# library
mb_fake_mode = _get_current_dispatch_mode()
if isinstance(mb_fake_mode, FakeTensorMode):
return _pop_mode_temporarily()
else:
return nullcontext()
@dataclass
class _ProxyTensor:
proxy: Proxy
constant: Optional[torch.Tensor]
def fetch_sym_proxy(tracer):
def inner(e):
n = e.node
if n.constant is not None:
return n.constant
else:
# NB: we REQUIRE all symints to be tracked
return get_proxy_slot(n, tracer)()
return inner
def fetch_tensor_proxy(tracer):
return lambda t: get_proxy_slot(t, tracer, t)
HANDLED_TYPES = (torch.Tensor, torch.nn.Parameter)
@contextlib.contextmanager
def inside_mode(proxy_mode):
old = proxy_mode.is_inside_mode
proxy_mode.is_inside_mode = True
try:
yield
finally:
proxy_mode.is_inside_mode = old
def proxy_call(proxy_mode, func, pre_autograd, args, kwargs):
def can_handle_tensor(x):
return type(x) in HANDLED_TYPES or has_proxy_slot(x, proxy_mode.tracer)
# If there are any tensor subclasses, we need to handle those tensor subclasses first
# TODO: we could use types to test this
if not pytree.tree_all_only(torch.Tensor, can_handle_tensor, (args, kwargs)):
return NotImplemented
if func in CURRENT_DECOMPOSITION_TABLE:
with proxy_mode:
r = CURRENT_DECOMPOSITION_TABLE[func](*args, **kwargs)
if r is not NotImplemented:
return r
# For pre-autograd tracing, we do not want to run CompositeImplicit decomps.
if not pre_autograd:
with proxy_mode:
r = func.decompose(*args, **kwargs)
if r is not NotImplemented:
return r
tracer = proxy_mode.tracer
f_args, f_kwargs = pytree.tree_map_only(torch.Tensor, fetch_tensor_proxy(tracer), (args, kwargs))
# If there are SymInts, we also should not consider this constant.
# However, fake tensor handling of SymInts is sufficiently broken that
# I couldn't write a test for this case
all_constant = (
pytree.tree_all_only(_ProxyTensor, lambda t: t.constant is not None, (f_args, f_kwargs))
# TODO: maybe constant SymInts should also be allowed? Not sure if
# this can happen
and pytree.tree_all_only((SymInt, SymFloat, SymBool), lambda _: False, (args, kwargs))
)
if torch.Tag.data_dependent_output in func.tags: # type: ignore[attr-defined]
# Check if all of the Tensor inputs are constants
if all_constant:
const_args, const_kwargs = pytree.tree_map_only(
_ProxyTensor, lambda t: t.constant, (f_args, f_kwargs)
)
with maybe_disable_fake_tensor_mode():
return func(*const_args, **const_kwargs)
# If any of the Tensor inputs are "real" (not FakeTensor), we may
# incorrectly burn in constants by allowing this access. Raise
# an error in this case
if pytree.tree_all_only(torch.Tensor, lambda t: not isinstance(t, FakeTensor), (args, kwargs)):
raise RuntimeError(
f"It appears that you're trying to get value out of a tracing tensor with {func} - erroring out! "
"It's likely that this is caused by data-dependent control flow or similar. "
"It may be possible to trace this with dynamic shapes; try setting tracing_mode='symbolic' "
"in your make_fx call."
)
proxy_args, proxy_kwargs = pytree.tree_map_only(
(SymInt, SymFloat, SymBool),
fetch_sym_proxy(proxy_mode.tracer),
pytree.tree_map_only(_ProxyTensor, lambda e: e.proxy, (f_args, f_kwargs))
)
# When we trace through a torch.tensor invocation, you never actually
# see a torch.ops.aten.tensor call. Instead, the way this function is
# implemented internally is that we allocate a plain tensor (this is
# *guaranteed* to be a plain tensor, we disable all modes when doing
# so), and then call at::lift_fresh on it (to give modes a chance to do
# their stuff). Furthermore, the tensor argument to lift_fresh is guaranteed
# to be freshly allocated, so we want lift_fresh to be a no-op (directly
# returning the input argument).
#
# Here is the basic problem: when we trace this sequence of executions
# into an FX graph, what happens to this call sequence? Traditionally,
# tensor constants get interned as buffers on the FX GraphModule. But
# this is dangerous. Consider:
#
# x = torch.tensor(1)
# x.add_(2)
#
# Naively, this traces into:
#
# t = self._tensor_constant0 # initialized to torch.tensor(1)
# x = torch.ops.aten.lift_fresh(t)
# x.add_(2)
#
# If lift_fresh returns t directly, the subsequent add_ call will
# modify the tensor constant. Really, the problem is we've violated
# the invariant the the argument to lift is fresh. So what we should
# preserve the invariant by replacing lift_fresh with lift_fresh_copy:
#
# t = self._tensor_constant0 # initialized to torch.tensor(1)
# x = torch.ops.aten.lift_fresh_copy(t)
# x.add_(2)
#
# This is what the overload modification does.
if func is torch.ops.aten.lift_fresh.default:
func = torch.ops.aten.lift_fresh_copy.default
# See Note [Per-Dispatch-Key Modes Must Be Reentrant]
# If our mode is on multiple mode stacks (e.g. the Autograd and Python mode stacks)
# then we only want it to trace out proxies the first time that we hit an op.
if proxy_mode.is_inside_mode:
return func(*args, **kwargs)
proxy_out = proxy_mode.tracer.create_proxy('call_function', func, proxy_args, proxy_kwargs,
name=proxy_mode.tracer.graph._target_to_str(func.overloadpacket.__name__))
# This makes DCE marginally less likely to DCE inplace operations.
# It is not strictly necessary
# Kind of a hacky way to test if an op is in-place or not
if func.overloadpacket.__name__[-1] == "_" and func.overloadpacket.__name__[0] != "_":
if isinstance(args[0], List):
# e.g., c10d::allreduce_ returns a list of tensors as the first element
# in the output.
for i, a in enumerate(args[0]):
a.proxy = proxy_out[0][i]
else:
args[0].proxy = proxy_out
with inside_mode(proxy_mode):
out = func(*args, **kwargs)
# In some circumstances, we will be tracing in a situation where a tensor
# is *statically* known to be a constant (currently, this only happens if
# you run torch.tensor; deterministic factory functions like torch.arange
# don't get this treatment). When the tensor in question is small, it's
# helpful to due constant propagation in case we call item() (in which
# case we can return the constant value that is known, rather than give
# an error.) The logic here tests if constant propagation is possible
# (because all of the inputs are constant). If so, we disable fake tensor
# mode (if it is on) and do true compute on the constant.
#
# It's worth highlighting that we're making a policy decision here.
# There is a potential that the tensor is actually quite large, and we
# don't actually want to run the compute. The tensor being quite large
# is one of the reasons why factory functions don't get this treatment
# (since they can be quite large; if a parameter is initialized to a
# constant value it will be!) Similarly, there is also a potential
# to run an operator that blows up the size of a small tensor; we don't
# protect against this case, but we could force, e.g., only single
# element constant computation by testing the numel of the result before
# propagating const-ness. Similarly, we don't require the constant to
# live on CPU, but we could.
any_constant = pytree.tree_any_only(_ProxyTensor, lambda t: t.constant is not None, (f_args, f_kwargs))
constant = None
# If this is a lift, the input tensor is guaranteed to be a
# constant, so we keep a copy of the original argument along so
# we can query it if we're asked to item() it at some later point
if func is torch.ops.aten.lift_fresh_copy.default and out.numel() <= CONSTANT_NUMEL_LIMIT:
with maybe_disable_fake_tensor_mode():
constant = args[0].clone()
elif (
torch.Tag.nondeterministic_seeded not in func.tags # type: ignore[attr-defined]
and all_constant
and any_constant
and pytree.tree_all_only(torch.Tensor, lambda t: t.numel() <= CONSTANT_NUMEL_LIMIT, out)
):
# NB: do NOT include factories as constants
with maybe_disable_fake_tensor_mode():
const_args, const_kwargs = pytree.tree_map_only(
_ProxyTensor, lambda t: t.constant, (f_args, f_kwargs)
)
constant = func(*const_args, **const_kwargs)
else:
constant = None
track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer)
return out
class PythonKeyTracer(Tracer):
def __init__(self):
super().__init__(autowrap_modules=())
self.tensor_tracker = WeakTensorKeyDictionary()
self.symnode_tracker = weakref.WeakKeyDictionary() # type: ignore[var-annotated]
# In general, we don't want to make modules leaves. In principle, users of
# this tracer might want to override this in order to turn a couple specific
# modules into leaves in the traced graph.
def call_module(
self, m: torch.nn.Module, forward: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any]
) -> Any:
return forward(*args, **kwargs)
# We don't want to turn getattr calls into proxies. So we just return the actual value.
def getattr(self, attr, attr_val, parameter_proxy_cache):
return attr_val
def create_arg(self, a: Any):
if isinstance(a, torch.nn.Parameter):
for n, p in self.root.named_parameters():
if a is p:
return self.create_node('get_attr', n, (), {})
qualname: Optional[str] = None
if not qualname:
i = 0
while True:
qualname = f'_param_constant{i}'
if not hasattr(self.root, qualname):
break
i += 1
setattr(self.root, qualname, a)
return self.create_node('get_attr', qualname, (), {})
elif isinstance(a, (SymInt, SymFloat, SymBool)):
assert a.node.constant is not None
return a.node.constant
return super().create_arg(a)
def dispatch_trace(
root: Union[torch.nn.Module, Callable],
tracer: Tracer,
concrete_args: Optional[Tuple[Any, ...]] = None,
) -> GraphModule:
graph = tracer.trace(root, concrete_args)
name = root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__
return GraphModule(tracer.root, graph, name)
def wrap_key(f, tensors, tracer, pre_autograd: bool):
flat_tensors, tensors_spec = pytree.tree_flatten(tensors)
dk = torch._C.DispatchKey.AutogradFunctionality if pre_autograd else None
@functools.wraps(f)
def wrapped(*proxies):
flat_proxies, proxies_spec = pytree.tree_flatten(proxies)
assert len(flat_proxies) == len(flat_tensors)
with _pop_mode_temporarily(dk) as m:
assert isinstance(m, ProxyTorchDispatchMode)
track_tensor_tree(flat_tensors, flat_proxies, constant=None, tracer=tracer)
out = f(*tensors)
out = pytree.tree_map_only(
torch.Tensor,
lambda t: get_proxy_slot(t, tracer, t, lambda x: x.proxy),
out
)
out = pytree.tree_map_only(
(SymInt, SymFloat, SymBool),
lambda t: get_proxy_slot(t.node, tracer)(),
out
)
return out
return wrapped
ORIGINAL_ATEN = None
@contextmanager
def set_original_aten_op(func):
global ORIGINAL_ATEN
if ORIGINAL_ATEN is None and fx_traceback.has_preserved_node_meta():
ORIGINAL_ATEN = func
fx_traceback.current_meta['original_aten'] = func
try:
yield
finally:
ORIGINAL_ATEN = None
fx_traceback.current_meta['original_aten'] = None
else:
yield
class ProxyTorchDispatchMode(TorchDispatchMode):
def __init__(self, tracer, tracing_mode, pre_autograd=False):
dk = torch._C.DispatchKey.AutogradFunctionality if pre_autograd else None
super().__init__(dk)
self.tracer = tracer
self.tracing_mode = tracing_mode
self.enable_tracing = True
self.pre_autograd = pre_autograd
self.is_inside_mode = False
self.sym_mode = ProxySymDispatchMode(tracer)
self.trace_state = {}
self._managers = []
@count
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
with self.sym_mode.enable(False), set_original_aten_op(func):
return self.inner_torch_dispatch(func, types, args, kwargs)
def __enter__(self):
# sym mode first, then us...
m = self.sym_mode.enable(True)
self._managers.append(m)
m.__enter__()
return super().__enter__()
def __exit__(self, exc_type, exc_value, traceback):
m = self._managers.pop()
# ...exit us first, then sym mode
b = super().__exit__(exc_type, exc_value, traceback)
if not b:
return m.__exit__(exc_type, exc_value, traceback)
else:
return m.__exit__(None, None, None)
def inner_torch_dispatch(self, func, types, args=(), kwargs=None):
if not self.enable_tracing:
return func(*args, **kwargs)
if func in [prim.device.default]:
return func(*args, **kwargs)
return proxy_call(self, func, self.pre_autograd, args, kwargs)
class ProxySymDispatchMode(SymDispatchMode):
def __init__(self, tracer):
super().__init__()
self.tracer = tracer
# When false, we don't trace operations. If you do this, you MUST
# call track_tensor/track_tensor_tree on all results of the operation
# to ensure we can adequately track the results
self.enable_tracing = True
@contextmanager
def enable(self, b):
old = self.enable_tracing
self.enable_tracing = b
try:
yield
finally:
self.enable_tracing = old
def _compute_proxy(self, func, args, out: Union[SymInt, SymFloat, SymBool]):
n_args = tuple(
get_proxy_slot(a.node, self.tracer)().node if isinstance(a, py_sym_types) else a
for a in args
)
# func doesn't have a __torch_function__ that Proxy can interpose, so
# we gotta do it manually
n_out = self.tracer.create_node("call_function", func, n_args, {})
p_out = fx.Proxy(n_out, self.tracer)
set_meta(p_out, out)
return p_out
def __sym_dispatch__(self, func, types, args, kwargs):
if not self.enable_tracing:
return func(*args, **kwargs)
# Peephole optimize multiply by one
# NB: be careful not to trigger guards here!
if func == operator.mul:
if isinstance(args[1], int) and args[1] == 1:
return args[0]
elif isinstance(args[0], int) and args[0] == 1:
return args[1]
# For speed, we assume there are no nested data structures
# (otherwise we could use tree_map)
# We also assume there are no keyword arguments.
assert not kwargs
out = func(*args, **kwargs)
# If func returned a constant, we don't need to trace; we have
# determined that the result is constant (no matter if the inputs
# were symbolic) and it is no longer necessary to trace the
# computation. This could occur if func triggered some guards.
if isinstance(out, py_sym_types):
# Delays tracing out the proxies on this op until we actually need it
p_out_thunk = thunkify(self._compute_proxy, func=func, args=args, out=out)
set_proxy_slot(out.node, self.tracer, p_out_thunk)
return out
# TODO: I'm not sure what the point of this class is; you can just
# make_fx through a regular Interpreter
class DecompositionInterpreter(torch.fx.Interpreter):
def __init__(self, module: torch.fx.GraphModule, new_graph: torch.fx.Graph, decomposition_table=None, **kwargs):
super().__init__(module, **kwargs)
self.new_graph = new_graph
self.tracer = torch.fx.proxy.GraphAppendingTracer(self.new_graph)
# Blegh
self.tracer.tensor_tracker = WeakTensorKeyDictionary() # type: ignore[attr-defined]
self.tracer.symnode_tracker = weakref.WeakKeyDictionary() # type: ignore[attr-defined]
self.decomposition_table = decomposition_table
if self.decomposition_table is None:
self.decomposition_table = {}
self.mode = ProxyTorchDispatchMode(self.tracer, tracing_mode="real")
def placeholder(self, target, args, kwargs):
out = super().placeholder(target, args, kwargs)
proxy = torch.fx.Proxy(self.new_graph.placeholder(target), self.tracer)
track_tensor_tree(out, proxy, constant=None, tracer=self.tracer)
# TODO handle case where the first character of target is '*'
return out
def get_attr(self, target, args, kwargs):
out = super().get_attr(target, args, kwargs)
proxy = torch.fx.Proxy(self.new_graph.get_attr(target), self.tracer)
track_tensor_tree(out, proxy, constant=None, tracer=self.tracer)
return out
# call_function, call_method, call_module get traced automatically by the outer mode.
def output(self, target, args, kwargs):
out = super().output(target, args, kwargs)
def unwrap(e):
return get_proxy_slot(e, self.tracer, e, lambda x: x.proxy.node)
self.new_graph.output(pytree.tree_map(unwrap, out))
return out
def run(self, *args, **kwargs):
# Should enter the mode at least once for being able to restore it later
# See: https://github.com/pytorch/pytorch/pull/82549#discussion_r934782025
with decompose(self.decomposition_table), self.mode:
return super().run(*args, **kwargs)
def wrapper_and_args_for_make_fx(func, args, kwargs):
# make_fx doesn't support kwargs, so we need to do this flattening
# and then unflatten the args before calling func
flat_args, spec = pytree.tree_flatten((args, kwargs))
def wrapped(flat_args):
fn_args, fn_kwargs = pytree.tree_unflatten(flat_args, spec)
return func(*fn_args, **fn_kwargs)
return wrapped, flat_args
@contextmanager
def disable_autocast_cache():
old_value = torch.is_autocast_cache_enabled()
torch.set_autocast_cache_enabled(False)
try:
yield
finally:
torch.set_autocast_cache_enabled(old_value)
def make_fx(f, decomposition_table=None, tracing_mode="real", _allow_non_fake_inputs=False, *, pre_autograd=False):
assert tracing_mode in ["real", "fake", "symbolic"]
if decomposition_table is None:
decomposition_table = {}
@functools.wraps(f)
def wrapped(*args):
phs = pytree.tree_map(lambda _: fx.PH, args) # type: ignore[attr-defined]
fx_tracer = PythonKeyTracer()
fake_tensor_mode: Any = nullcontext()
if tracing_mode == "real":
fake_tensor_mode = nullcontext()
elif tracing_mode == "fake":
import torch._dynamo
fake_tensor_mode = torch._dynamo.utils.detect_fake_mode(args)
if fake_tensor_mode is None:
fake_tensor_mode = FakeTensorMode(
allow_fallback_kernels=True,
allow_non_fake_inputs=_allow_non_fake_inputs)
elif tracing_mode == "symbolic":
import torch._dynamo
fake_tensor_mode = torch._dynamo.utils.detect_fake_mode(args)
if fake_tensor_mode is None:
shape_env = ShapeEnv()
fake_tensor_mode = FakeTensorMode(
allow_fallback_kernels=False,
allow_non_fake_inputs=_allow_non_fake_inputs,
shape_env=shape_env)
else:
shape_env = fake_tensor_mode.shape_env
assert shape_env is not None, "shape_env should be set if tracing with 'symbolic'"
else:
raise AssertionError(f"Unexpected tracing type: {tracing_mode}")
python_dispatcher_mode: Any = nullcontext()
# pre-autograd tracing uses per-dispatch-key modes,
# which requires the python dispatcher
if tracing_mode == "symbolic" or pre_autograd:
python_dispatcher_mode = enable_python_dispatcher()
proxy_mode = ProxyTorchDispatchMode(fx_tracer, tracing_mode, pre_autograd=pre_autograd)
arg_count = 0
def wrap_fake(x):
nonlocal arg_count
if isinstance(x, torch.Tensor):
# TODO: it would be nice to line these up with the names
# FX will choose for the placeholders, but we don't
# actually know what the names will be at this point yet
# NB: the Source here is actually meaningless
from torch._dynamo.source import ConstantSource
source = ConstantSource(f"input{arg_count}")
arg_count += 1
return fake_tensor_mode.from_tensor(x, source=source) # type: ignore[attr-defined]
return x
sym_mode = proxy_mode.sym_mode
wrap_fn_map = {
"real": lambda x: x,
"fake": wrap_fake,
"symbolic": wrap_fake,
}
args = pytree.tree_map(wrap_fn_map[tracing_mode], args)
if not hasattr(inspect.unwrap(f), '__code__') or inspect.unwrap(f).__code__.co_flags & inspect.CO_VARARGS:
# FX doesn't support varargs, so we gotta fake up a wrapper
# TODO: Would be nice to fix this at the source...
func = fake_signature(f, len(phs))
else:
func = f
# We disable the autocast cache as the autocast cache causes type conversions on parameters to
# check a cache, which introduces untracked tensors into the graph
#
# We also disable tracing by any other tensor proxy-based tracers except the current. The
# purpose of `make_fx` is to produce graphmodules as a side effect; its internal execution is
# thus irrelevant to any external functional trace.
with decompose(decomposition_table), fake_tensor_mode, python_dispatcher_mode, \
sym_mode, proxy_mode, disable_autocast_cache(), disable_proxy_modes_tracing(enable_current=True):
t = dispatch_trace(wrap_key(func, args, fx_tracer, pre_autograd), tracer=fx_tracer, concrete_args=tuple(phs))
# TODO: kind of a bad way to do it, should maybe figure out a better way
if tracing_mode == "symbolic":
t.shape_env = shape_env # type: ignore[assignment]
return t
return wrapped
def get_torch_dispatch_modes():
return torch.utils._python_dispatch._get_current_dispatch_mode_stack()
def get_innermost_proxy_mode():
for m in reversed(torch.utils._python_dispatch._get_current_dispatch_mode_stack()):
if isinstance(m, ProxyTorchDispatchMode):
return m
return None
@contextlib.contextmanager
def disable_proxy_modes_tracing(enable_current=False):
modes = get_torch_dispatch_modes()
proxy_tensor_modes = [m for m in modes if isinstance(m, ProxyTorchDispatchMode)]
if enable_current:
proxy_tensor_modes = proxy_tensor_modes[:-1]
olds = [(m.enable_tracing, m.sym_mode.enable_tracing) for m in proxy_tensor_modes]
for proxy_mode in proxy_tensor_modes:
proxy_mode.enable_tracing = False
proxy_mode.sym_mode.enable_tracing = False
try:
yield
finally:
for proxy_mode, (old, old_sym) in zip(proxy_tensor_modes, olds):
proxy_mode.enable_tracing = old
proxy_mode.sym_mode.enable_tracing = old_sym
def get_isolated_graphmodule(func, args, kwargs, tracing_mode="real"):
"""A helper function used to get the GraphModule for the given func.
It's expected to be used in the ProxyTensor tracing context.
It detaches the args and kwargs from the current tracer so that the trace of
the current graph module can be created without any side-effects.
"""
wrapped, all_args = wrapper_and_args_for_make_fx(func, args, kwargs)
with disable_proxy_modes_tracing():
gm = make_fx(wrapped, tracing_mode=tracing_mode)(all_args)
return gm