builder.py 158 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546354735483549355035513552355335543555355635573558355935603561356235633564356535663567356835693570357135723573357435753576357735783579358035813582358335843585358635873588358935903591359235933594359535963597359835993600360136023603360436053606360736083609361036113612361336143615361636173618361936203621362236233624362536263627362836293630363136323633363436353636363736383639364036413642364336443645364636473648364936503651365236533654365536563657365836593660366136623663366436653666366736683669367036713672367336743675367636773678367936803681368236833684368536863687368836893690369136923693369436953696369736983699370037013702370337043705370637073708370937103711371237133714371537163717371837193720372137223723372437253726372737283729373037313732373337343735373637373738373937403741374237433744374537463747374837493750375137523753375437553756375737583759376037613762376337643765
  1. # mypy: ignore-errors
  2. """
  3. This module contains classes and utilities for building variable trackers in Dynamo.
  4. Variable trackers are used to convert Python values into symbolic representations
  5. that can be traced and transformed during graph capture.
  6. The key classes are:
  7. - VariableBuilder: Handles source-tracked objects that need guards and proper
  8. reconstruction in the output graph. Used for inputs, module attributes, etc.
  9. - SourcelessBuilder: Handles ephemeral objects created during tracing that don't
  10. need source tracking or guards. Used for temporary lists, intermediate values, etc.
  11. Variable trackers enable Dynamo to track the flow of values through the program,
  12. maintain guards for dynamic properties, and reconstruct values in the output graph.
  13. The builders in this module handle converting Python values into appropriate
  14. VariableTracker instances based on their type and usage context.
  15. """
  16. import abc
  17. import collections
  18. import contextlib
  19. import copy
  20. import dataclasses
  21. import enum
  22. import functools
  23. import inspect
  24. import itertools
  25. import logging
  26. import math
  27. import operator
  28. import random
  29. import re
  30. import sys
  31. import traceback
  32. import types
  33. import weakref
  34. from collections.abc import MutableMapping
  35. from typing import Any, Callable, NamedTuple, Optional, TYPE_CHECKING, Union
  36. import sympy
  37. import torch
  38. from torch import SymInt
  39. from torch._dispatch.python import enable_python_dispatcher
  40. from torch._dynamo.utils import (
  41. get_metrics_context,
  42. is_int_specialization_case,
  43. is_torch_sym,
  44. set_feature_use,
  45. )
  46. from torch._guards import TracingContext
  47. from torch._higher_order_ops.flat_apply import flat_apply
  48. from torch._higher_order_ops.torchbind import call_torchbind
  49. from torch._ops import HigherOrderOperator
  50. from torch._subclasses.fake_tensor import FakeTensor, is_fake, maybe_get_fake_mode
  51. from torch._subclasses.meta_utils import is_sparse_any, safe_grad
  52. from torch._utils_internal import justknobs_check
  53. from torch.fx.experimental._backward_state import BackwardState
  54. from torch.fx.experimental._dynamism import normalize_source_name
  55. from torch.fx.experimental.symbolic_shapes import (
  56. _constrain_range_for_size,
  57. _nested_int_aware_sort,
  58. DimDynamic,
  59. RelaxedUnspecConstraint,
  60. StatefulSymbolicContext,
  61. SubclassSymbolicContext,
  62. SymbolicContext,
  63. SymIntSymbolicContext,
  64. TrackedFake,
  65. )
  66. from torch.fx.immutable_collections import immutable_dict, immutable_list
  67. from torch.nn.utils._expanded_weights import ExpandedWeight
  68. from torch.utils._python_dispatch import (
  69. is_traceable_wrapper_subclass,
  70. is_traceable_wrapper_subclass_type,
  71. )
  72. from torch.utils._sympy.value_ranges import ValueRanges
  73. from torch.utils.weak import TensorWeakRef
  74. from .. import config, graph_break_hints, mutation_guard, replay_record, trace_rules
  75. from ..device_interface import get_registered_device_interfaces
  76. from ..exc import InternalTorchDynamoError, raise_observed_exception, unimplemented_v2
  77. from ..guards import GuardBuilder, install_guard, make_dupe_guard
  78. from ..pgo import (
  79. auto_dynamic,
  80. auto_unset,
  81. FrameStateSizeEntry,
  82. InferStride,
  83. process_automatic_dynamic,
  84. )
  85. from ..side_effects import SideEffects
  86. from ..source import (
  87. AttrProxySource,
  88. AttrSource,
  89. CallMethodItemSource,
  90. ChainedSource,
  91. ConstDictKeySource,
  92. ConvertIntSource,
  93. DictGetItemSource,
  94. DictSubclassGetItemSource,
  95. FloatTensorSource,
  96. GetItemSource,
  97. GradSource,
  98. is_constant_source,
  99. is_from_closure_source,
  100. is_from_global_source,
  101. is_from_nonlocal_source,
  102. is_from_optimizer_source,
  103. is_from_unspecialized_nn_module_source,
  104. ListGetItemSource,
  105. LocalSource,
  106. NonSerializableSetGetItemSource,
  107. NumpyTensorSource,
  108. OptimizerSource,
  109. RandomValueSource,
  110. Source,
  111. SubclassAttrListSource,
  112. TupleIteratorGetItemSource,
  113. UnspecializedBuiltinNNModuleSource,
  114. UnspecializedNNModuleSource,
  115. )
  116. from ..utils import (
  117. _extract_tensor_dict,
  118. build_checkpoint_variable,
  119. build_invoke_subgraph_variable,
  120. clone_input,
  121. common_constant_types,
  122. dict_keys,
  123. get_fake_value,
  124. get_items_from_dict,
  125. get_locals_to_steal,
  126. get_static_address_type,
  127. is_frozen_dataclass,
  128. is_function,
  129. is_function_or_wrapper,
  130. is_invoke_subgraph,
  131. is_lru_cache_wrapped_function,
  132. is_namedtuple,
  133. is_parameter_freezing,
  134. is_typing,
  135. is_utils_checkpoint,
  136. is_wrapper_or_member_descriptor,
  137. istype,
  138. namedtuple_fields,
  139. odict_values,
  140. proxy_args_kwargs,
  141. range_iterator,
  142. set_example_value,
  143. tensor_always_has_static_shape,
  144. tuple_iterator,
  145. tuple_iterator_getitem,
  146. tuple_iterator_len,
  147. unwrap_with_attr_name_if_wrapper,
  148. wrap_fake_exception,
  149. )
  150. from .base import (
  151. AttributeMutationNew,
  152. typestr,
  153. ValueMutationExisting,
  154. ValueMutationNew,
  155. VariableTracker,
  156. VariableTrackerMeta,
  157. )
  158. from .builtin import BuiltinVariable
  159. from .constant import ConstantVariable, EnumVariable
  160. from .ctx_manager import (
  161. AutocastModeVariable,
  162. DynamoConfigPatchVariable,
  163. ErrorOnGraphBreakVariable,
  164. EventVariable,
  165. NullContextVariable,
  166. PreserveVersionContextVariable,
  167. StreamContextVariable,
  168. StreamVariable,
  169. )
  170. from .dicts import (
  171. ConstDictVariable,
  172. DefaultDictVariable,
  173. DictKeySetVariable,
  174. FrozensetVariable,
  175. MappingProxyVariable,
  176. SetVariable,
  177. )
  178. from .distributed import (
  179. DeviceMeshVariable,
  180. PlacementClassVariable,
  181. PlacementVariable,
  182. ProcessGroupVariable,
  183. WorldMetaClassVariable,
  184. )
  185. from .functions import (
  186. BuiltinMethodVariable,
  187. CollectionsNamedTupleFunction,
  188. CollectiveFunctionRewriteVariable,
  189. CreateTMADescriptorExperimentalVariable,
  190. CreateTMADescriptorStableVariable,
  191. FunctoolsPartialVariable,
  192. FunctoolsWrapsVariable,
  193. SysFunctionVariable,
  194. TracebackVariable,
  195. TritonKernelVariable,
  196. UserFunctionVariable,
  197. UserMethodVariable,
  198. WrapperUserFunctionVariable,
  199. )
  200. from .higher_order_ops import TorchHigherOrderOperatorVariable
  201. from .iter import ItertoolsVariable
  202. from .lazy import LazyVariableTracker
  203. from .lists import (
  204. BaseListVariable,
  205. ListIteratorVariable,
  206. ListVariable,
  207. NamedTupleVariable,
  208. RangeVariable,
  209. SizeVariable,
  210. SliceVariable,
  211. TupleIteratorVariable,
  212. TupleVariable,
  213. )
  214. from .misc import (
  215. AutogradEngineVariable,
  216. AutogradFunctionContextVariable,
  217. AutogradFunctionVariable,
  218. ComptimeVariable,
  219. DebuggingVariable,
  220. DelayGraphBreakVariable,
  221. GetAttrVariable,
  222. GetSetDescriptorVariable,
  223. LambdaVariable,
  224. LoggingLoggerVariable,
  225. MethodWrapperVariable,
  226. NumpyDTypeVariable,
  227. NumpyTypeInfoVariable,
  228. NumpyVariable,
  229. PythonModuleVariable,
  230. RandomClassVariable,
  231. RandomVariable,
  232. RegexPatternVariable,
  233. SavedTensorBox,
  234. TorchVersionVariable,
  235. TypingVariable,
  236. WeakRefVariable,
  237. )
  238. from .nn_module import (
  239. FSDPManagedNNModuleVariable,
  240. UnspecializedBuiltinNNModuleVariable,
  241. UnspecializedNNModuleVariable,
  242. )
  243. from .optimizer import OptimizerVariable
  244. from .script_object import TorchScriptObjectVariable
  245. from .sdpa import SDPAParamsVariable
  246. from .tensor import (
  247. NumpyNdarrayVariable,
  248. supported_const_comparison_op_values,
  249. SymNodeVariable,
  250. TensorSubclassVariable,
  251. TensorVariable,
  252. UnspecializedPythonVariable,
  253. )
  254. from .torch import (
  255. DispatchKeySetVariable,
  256. FuncTorchInterpreterVariable,
  257. TorchCtxManagerClassVariable,
  258. TorchInGraphFunctionVariable,
  259. )
  260. from .torch_function import (
  261. TensorWithTFOverrideVariable,
  262. torch_function_mode_stack_state_mgr,
  263. TorchFunctionModeVariable,
  264. )
  265. from .user_defined import (
  266. FrozenDataClassVariable,
  267. IntWrapperVariable,
  268. KeyedJaggedTensorVariable,
  269. MutableMappingVariable,
  270. SourcelessGraphModuleVariable,
  271. UserDefinedClassVariable,
  272. UserDefinedDictVariable,
  273. UserDefinedExceptionClassVariable,
  274. UserDefinedListVariable,
  275. UserDefinedObjectVariable,
  276. UserDefinedSetVariable,
  277. UserDefinedTupleVariable,
  278. )
  279. try:
  280. import numpy as np
  281. except ModuleNotFoundError:
  282. np = None
  283. if TYPE_CHECKING:
  284. from torch._dynamo.codegen import PyCodegen
  285. from torch._dynamo.symbolic_convert import InstructionTranslator
  286. log = logging.getLogger(__name__)
  287. static_inputs_log = torch._logging.getArtifactLogger(
  288. __name__, "cudagraph_static_inputs"
  289. )
  290. DimList = list
  291. def safe_has_grad(t):
  292. with torch._logging.hide_warnings(torch._logging._internal.safe_grad_filter):
  293. return hasattr(t, "grad")
  294. class _missing:
  295. pass
  296. @dataclasses.dataclass
  297. class GraphArg:
  298. source: Source
  299. # TODO: storing a SymInt here but not a FakeTensor is a pretty strange
  300. # thing to do. Probably should have example (which stores an int) and
  301. # fake_example
  302. _example: Union[TensorWeakRef, torch.SymInt]
  303. # When True, this indicates that this GraphArg is a Python quantity (e.g.,
  304. # a float or int) which we pass to the FX graph as a Tensor. This
  305. # controls how we codegen calls into the Dynamo graph: we will call
  306. # torch.as_tensor on the quantity before passing it in.
  307. #
  308. # Note that we typically do not pass dynamic integers as tensors, because
  309. # they will most frequently just be used for size computation. But this
  310. # is a policy decision that we can change our mind on; in particular, when
  311. # an int comes from a random number generator (e.g., random.randint), we
  312. # DO pass it as a tensor.
  313. #
  314. # It's also worth noting that our current tracing rules for
  315. # pass_arg_as_tensor as subtly broken: we just pun the variable as a
  316. # 0d scalar Tensor and pray that the semantics are the same. Which they
  317. # often are, but not necessarily. ezyang(May 2024) plans to fix this
  318. # soon.
  319. pass_arg_as_tensor: bool
  320. fake_tensor: Optional[torch._subclasses.fake_tensor.FakeTensor]
  321. # UnspecializedPythonVariable often masquerades as a tensor.
  322. # We MUST NOT generate shape guard code
  323. # that actually tries to access tensor properties on these values.
  324. # is_tensor lets us tell if this graph arg actually is a tensor
  325. # or not.
  326. is_tensor: bool = True
  327. # Sometimes, the Tensor we pass to example is freshly allocated (smh).
  328. # Then we cannot only keep a weak reference to it. This lets you
  329. # stash a strong reference too.
  330. example_strong_ref: Optional[torch.Tensor] = None
  331. @property
  332. def example(self):
  333. if isinstance(self._example, TensorWeakRef):
  334. r = self._example()
  335. assert r is not None
  336. return r
  337. else:
  338. return self._example
  339. def __post_init__(self):
  340. if isinstance(self._example, torch.Tensor):
  341. self._example = TensorWeakRef(self._example)
  342. assert is_fake(self.fake_tensor)
  343. def reconstruct(self, codegen: "PyCodegen"):
  344. codegen(self.source)
  345. def erase(self):
  346. self._example = None
  347. self.example_strong_ref = None
  348. def __eq__(self, other):
  349. return self.source.name() == other.source.name()
  350. class BackwardStateGraphArg(GraphArg):
  351. def __init__(self) -> None:
  352. super().__init__(
  353. source=None,
  354. _example=BackwardState(),
  355. pass_arg_as_tensor=False,
  356. fake_tensor=None,
  357. is_tensor=False,
  358. )
  359. def reconstruct(self, codegen: "PyCodegen"):
  360. assert codegen.tx.output.backward_state_var
  361. codegen.add_push_null(
  362. lambda: codegen.load_import_from(BackwardState.__module__, "BackwardState")
  363. )
  364. codegen.call_function(0, False)
  365. codegen.dup_top()
  366. codegen.store(codegen.tx.output.backward_state_var)
  367. # All class-based iterators in itertools
  368. # NOTE: use id() because some objects are not hashable, it will raise error during lookup
  369. ITERTOOLS_TYPE_IDS: frozenset[int] = frozenset(
  370. id(member)
  371. for name, member in vars(itertools).items()
  372. if not name.startswith("_") and inspect.isclass(member)
  373. )
  374. # Will be updated later in substitute_in_graph in torch/_dynamo/polyfills/itertools.py
  375. ITERTOOLS_POLYFILLED_TYPE_IDS: set[int] = set()
  376. # Capture fn pointer at import time
  377. # This is to guard against trying to mark the iterated tensors
  378. # as static in case user overrides fn ptr
  379. og_module_named_buffers_fn_ptr = torch.nn.Module.named_buffers
  380. og_module_named_parameters_fn_ptr = torch.nn.Module.named_parameters
  381. class VariableBuilder:
  382. """Wrap a python value in a VariableTracker() instance"""
  383. def __init__(
  384. self,
  385. tx,
  386. source: Source,
  387. ) -> None:
  388. assert source is not None, (
  389. "Consider SourcelessBuilder for ephemeral objects, usually objects created locally."
  390. )
  391. assert TracingContext.try_get() is not None, "Expected active TracingContext"
  392. super().__init__()
  393. self.tx = tx
  394. self.source = source
  395. self.name = source.name()
  396. def __call__(self, value):
  397. if value in self.tx.output.side_effects:
  398. side_effect_result = self.tx.output.side_effects[value]
  399. dup_guard = make_dupe_guard(self.source, side_effect_result.source)
  400. if dup_guard:
  401. self.install_guards(dup_guard)
  402. return side_effect_result
  403. cached_vt = self.tx.output.variable_tracker_cache.lookup(value, self.source)
  404. if cached_vt:
  405. return cached_vt
  406. vt = self._wrap(value)
  407. if vt.source is None:
  408. vt.source = self.source
  409. def _is_deduplicable_sym_variable(value, vt):
  410. # Constants like 0, 1, 2, etc. can be unspecialized as SymNodeVariables sometimes, but we
  411. # should NOT track them. If we use a single SymNodeVariable instance to track them
  412. # across multiple uses, then guards created for one usage will incorrectly apply to
  413. # all other usages of that constant, leading to unnecessary recompilations.
  414. return is_torch_sym(value) and isinstance(vt, SymNodeVariable)
  415. if (
  416. (
  417. self._can_lift_attrs_to_inputs(vt)
  418. or _is_deduplicable_sym_variable(value, vt)
  419. )
  420. and value not in self.tx.output.side_effects
  421. and not is_wrapper_or_member_descriptor(value)
  422. ):
  423. vt = self.tx.output.side_effects.track_object_existing(value, vt)
  424. self.tx.output.variable_tracker_cache.add(value, self.source, vt)
  425. return vt
  426. def _can_lift_attrs_to_inputs(self, vt):
  427. return type(vt) in {
  428. TensorVariable,
  429. TensorWithTFOverrideVariable,
  430. UserDefinedObjectVariable,
  431. NumpyNdarrayVariable,
  432. }
  433. def get_source(self):
  434. return self.source
  435. def install_guards(self, *guards):
  436. source = self.get_source()
  437. try:
  438. tmp = [source.make_guard(guard) for guard in guards]
  439. except NotImplementedError:
  440. return None
  441. install_guard(*tmp, skip=1)
  442. return {}
  443. @classmethod
  444. def _type_dispatch(cls):
  445. return cls._type_dispatch_impl(config.trace_numpy)
  446. @classmethod
  447. @functools.cache
  448. def _type_dispatch_impl(cls, trace_numpy):
  449. # NB: Careful not to close over self to avoid ref cycle from lru_cache
  450. entries = [
  451. (
  452. (
  453. torch.Tensor,
  454. torch.nn.Parameter,
  455. torch._subclasses.FakeTensor,
  456. torch._subclasses.functional_tensor.FunctionalTensor,
  457. ),
  458. cls.wrap_tensor,
  459. ),
  460. (
  461. (tuple, list, odict_values, collections.deque, torch.Size),
  462. cls.wrap_listlike,
  463. ),
  464. (tuple_iterator, cls.wrap_tuple_iterator),
  465. (range_iterator, cls.wrap_range_iterator),
  466. ((slice, range), cls.wrap_slice_range),
  467. (tuple(common_constant_types), cls.wrap_literal),
  468. (re.Pattern, cls.wrap_regex_pattern),
  469. (weakref.ReferenceType, cls.wrap_weakref),
  470. (torch.utils.hooks.RemovableHandle, cls.wrap_removable_handle),
  471. (torch.jit.ScriptFunction, cls.wrap_jit_function),
  472. (types.MappingProxyType, cls.wrap_mapping_proxy),
  473. ]
  474. if trace_numpy and np:
  475. entries.append((np.ndarray, cls.wrap_numpy_ndarray))
  476. result = {}
  477. for ts, fn in entries:
  478. for t in ts if isinstance(ts, tuple) else (ts,):
  479. assert t not in result
  480. result[t] = fn
  481. return result
  482. def wrap_regex_pattern(self, value: re.Pattern):
  483. # TODO(jansel): something like a REPR_MATCH might be more robust here
  484. self.install_guards(GuardBuilder.ID_MATCH)
  485. return RegexPatternVariable(value)
  486. def wrap_weakref(self, value: weakref.ReferenceType):
  487. self.install_guards(GuardBuilder.TYPE_MATCH)
  488. return WeakRefVariable.build(self.tx, value, source=self.source)
  489. def wrap_removable_handle(self, value):
  490. # This means that the removable handle was created in some other frame.
  491. # Our current infra requires the hook to be registered and removed in
  492. # the same frame. So graph break.
  493. # Related test - PYTORCH_TEST_WITH_DYNAMO=1 python test/test_autograd.py -k TestAutograd.test_hooks
  494. unimplemented_v2(
  495. gb_type="Attempted to represent unregistered RemovableHandle",
  496. context="",
  497. explanation="Dynamo attempted to build a representation of a torch.utils.hooks.RemovableHandle, "
  498. "which is not supported. This happens because the RemovableHandle was created in another frame.",
  499. hints=[],
  500. )
  501. def wrap_jit_function(self, value):
  502. self.install_guards(GuardBuilder.TYPE_MATCH)
  503. return WrapperUserFunctionVariable(
  504. value, "_torchdynamo_inline", source=self.source
  505. )
  506. def wrap_mapping_proxy(self, value):
  507. self.install_guards(GuardBuilder.TYPE_MATCH)
  508. # This might be suboptimal compared to dict guards. But mappingproxy is
  509. # not very common, so its ok to guard on all keys.
  510. self.install_guards(GuardBuilder.MAPPING_KEYS_CHECK)
  511. all_const = all(ConstantVariable.is_literal(k) for k in value.keys())
  512. if not all_const:
  513. unimplemented_v2(
  514. gb_type="non-const keys in mappingproxy",
  515. context=f"non-const keys: {[k for k in value.keys() if not ConstantVariable.is_literal(k)]}",
  516. explanation="Dynamo expects mappingproxy keys to be constants.",
  517. hints=[
  518. "Ensure your mappingproxy keys are constants (e.g. int, float, strings)",
  519. ],
  520. )
  521. def build_key_value(k, v):
  522. key = ConstantVariable.create(k)
  523. source_key = k
  524. source_value = GetItemSource(self.get_source(), source_key)
  525. res_value = LazyVariableTracker.create(v, source_value)
  526. return key, res_value
  527. items = dict(build_key_value(k, v) for k, v in value.items())
  528. # Create a dict_vt to be used in the mapping proxy variable
  529. dict_vt = ConstDictVariable(items, source=None)
  530. result = MappingProxyVariable(dict_vt, source=self.source)
  531. return self.tx.output.side_effects.track_mutable(value, result)
  532. @classmethod
  533. @functools.cache
  534. def _id_dispatch(
  535. cls,
  536. ) -> dict[int, Callable[["VariableBuilder", Any], VariableTracker]]:
  537. from ..comptime import comptime
  538. entries = [
  539. (comptime, lambda self, value: ComptimeVariable()),
  540. (
  541. dataclasses.fields,
  542. lambda self, value: LambdaVariable(
  543. _dataclasses_fields_lambda,
  544. source=self.source,
  545. **self.install_guards(GuardBuilder.FUNCTION_MATCH),
  546. ),
  547. ),
  548. (torch.__version__, lambda self, value: TorchVersionVariable()),
  549. ]
  550. result = {}
  551. for ts, fn in entries:
  552. for t in ts if isinstance(ts, (tuple, list)) else (ts,):
  553. assert t not in result
  554. result[id(t)] = fn
  555. return result
  556. def _wrap(self, value):
  557. # import here to avoid circular dependencies
  558. from torch.utils._triton import (
  559. has_triton,
  560. has_triton_experimental_host_tma,
  561. has_triton_tensor_descriptor_host_tma,
  562. )
  563. from ..decorators import (
  564. DynamoConfigPatchProxy,
  565. ErrorOnGraphBreakDecoratorContextManager,
  566. )
  567. if has_triton():
  568. from triton.runtime.autotuner import Autotuner
  569. from triton.runtime.jit import JITFunction
  570. else:
  571. class JITFunction:
  572. pass
  573. class Autotuner:
  574. pass
  575. # default implementations, in case we don't have triton (or the wrong triton version)
  576. def create_1d_tma_descriptor():
  577. pass
  578. def create_2d_tma_descriptor():
  579. pass
  580. class TensorDescriptor:
  581. @staticmethod
  582. def from_tensor():
  583. pass
  584. if has_triton_experimental_host_tma():
  585. from triton.tools.experimental_descriptor import ( # noqa: F811
  586. create_1d_tma_descriptor,
  587. create_2d_tma_descriptor,
  588. )
  589. if has_triton_tensor_descriptor_host_tma():
  590. from triton.tools.tensor_descriptor import TensorDescriptor # noqa: F811
  591. # Handle exact type() match
  592. type_dispatch = self._type_dispatch().get(type(value))
  593. if type_dispatch is not None:
  594. return type_dispatch(self, value)
  595. # Handle exact id() match
  596. id_dispatch = self._id_dispatch().get(id(value))
  597. if id_dispatch is not None:
  598. return id_dispatch(self, value)
  599. # Everything else (NB: order matters!)
  600. if (
  601. isinstance(value, torch.Tensor)
  602. and type(value)
  603. not in (
  604. # These torch-native subclasses have overly restrictive
  605. # `__torch_function__` which prevents Dynamo from reading their
  606. # tensor attributes like `is_nested` or calling methods like
  607. # `_is_view`.
  608. torch.nn.parameter.UninitializedBuffer,
  609. torch.nn.parameter.UninitializedParameter,
  610. ExpandedWeight,
  611. )
  612. and type(value) not in config.nontraceable_tensor_subclasses
  613. ):
  614. if (
  615. type(value).__torch_dispatch__ is torch.Tensor.__torch_dispatch__
  616. or is_traceable_wrapper_subclass(value)
  617. ):
  618. return self.wrap_tensor(value)
  619. if is_namedtuple(value):
  620. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  621. output = [
  622. LazyVariableTracker.create(
  623. getattr(value, name),
  624. source=AttrSource(self.source, name),
  625. )
  626. for name in namedtuple_fields(type(value))
  627. ]
  628. result = NamedTupleVariable(
  629. output, tuple_cls=type(value), source=self.source
  630. )
  631. return result
  632. elif istype(value, (dict, collections.defaultdict, collections.OrderedDict)):
  633. self.install_guards(GuardBuilder.TYPE_MATCH)
  634. all_const = all(ConstantVariable.is_literal(k) for k in value.keys())
  635. # For all_const, we don't have to guard on anything yet. We guard on
  636. # keys lazily by adding a dict_getitem entry for each accessed key.
  637. # For cases where we need to guard on all keys, we lazily put guards
  638. # during the dict call_method (check dicts.py)
  639. if not all_const:
  640. # Guard on the key order
  641. # This is not ideal, i.e., there is no need to guard on the key
  642. # order. But we guard on the key order because of the complexity
  643. #
  644. # 1) For non-constant objects, we can't save the key in the
  645. # guard context because it can be memory heavy. We can add
  646. # weakrefs but this complicates the accesses.
  647. #
  648. # 2) For non-constant objects, we also have to guard on the keys
  649. # (like TENSOR_MATCH on tensor). We might also have guards on
  650. # the attributes of the keys (like tensor.grad). To make this
  651. # work in tree structure is complicated.
  652. #
  653. # So, instead we guard on the key order. While guarding on key
  654. # order, we just save the indices and use it to access keys and
  655. # values. Indices are cheap to save.
  656. self.tx.output.guard_on_key_order.add(self.source)
  657. # We need all the keys to be hashable. We do this within the
  658. # _HashableTracker class in dicts.py
  659. def build_key_value(i, k, v):
  660. base = self.get_source()
  661. if all_const:
  662. key = ConstantVariable.create(k)
  663. source_key = k
  664. else:
  665. source_key = ConstDictKeySource(base, i)
  666. key = LazyVariableTracker.create(k, source_key)
  667. source_value = DictGetItemSource(base, source_key)
  668. res_value = LazyVariableTracker.create(v, source_value)
  669. return key, res_value
  670. # Ensure that we call dict.keys and not value.keys (which can call
  671. # overridden keys method). In the C++ guards, we relied on
  672. # PyDict_Next to traverse the dictionary, which uses the internal
  673. # data structure and does not call the overridden keys method.
  674. result = dict(
  675. build_key_value(i, k, v)
  676. for i, (k, v) in enumerate(get_items_from_dict(value))
  677. )
  678. if istype(value, collections.defaultdict):
  679. factory_source = AttrSource(self.source, "default_factory")
  680. result = DefaultDictVariable(
  681. result,
  682. type(value),
  683. default_factory=VariableBuilder(self.tx, factory_source)(
  684. value.default_factory
  685. ),
  686. source=self.source,
  687. )
  688. else:
  689. result = ConstDictVariable(
  690. result, user_cls=type(value), source=self.source
  691. )
  692. return self.tx.output.side_effects.track_mutable(value, result)
  693. elif isinstance(value, torch.nn.Module):
  694. return self.wrap_module(value)
  695. elif ConstantVariable.is_literal(value): # non-atomic literals
  696. return self.wrap_literal(value)
  697. elif isinstance(value, torch.overrides.TorchFunctionMode):
  698. var = TorchFunctionModeVariable(value, source=self.source)
  699. self.tx.output.side_effects.track_object_existing(value, var)
  700. return var
  701. elif istype(value, set):
  702. if any(isinstance(x, torch.Tensor) for x in value):
  703. unimplemented_v2(
  704. gb_type="Attempted to wrap a set with tensors",
  705. context="Python set containing torch.Tensor elements",
  706. explanation=(
  707. "Dynamo cannot trace sets of tensors. To get a stable ordering, "
  708. "Dynamo needs to convert the set into a list and the order might not be "
  709. "stable if the set contains tensors."
  710. ),
  711. hints=[
  712. "Use a dictionary where the keys are tensors.",
  713. *graph_break_hints.SUPPORTABLE,
  714. ],
  715. )
  716. self.install_guards(GuardBuilder.TYPE_MATCH)
  717. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  718. # The list gives a ordering for the set items. The ordering is based
  719. # on the Python hash and it is not related to object ordering inside
  720. # the set object. The order being incorrect at runtime will lead to
  721. # a recompilation.
  722. L = list(value)
  723. items = [
  724. LazyVariableTracker.create(
  725. v, source=NonSerializableSetGetItemSource(self.source, i)
  726. )
  727. for i, v in enumerate(L)
  728. ]
  729. result = SetVariable(items, source=self.source)
  730. return self.tx.output.side_effects.track_object_existing(value, result)
  731. elif istype(value, frozenset) and all(
  732. (
  733. # For DBR quantization, we could get a frozenset of torch funcs.
  734. (type(x) is types.BuiltinMethodType and x.__module__ == "torch")
  735. or
  736. # Another commonly used frozenset of types.
  737. x in torch.utils._pytree.BUILTIN_TYPES
  738. )
  739. for x in value
  740. ):
  741. # For the limited cases of frozenset here, we know the items won't
  742. # change across runs, so we can safely create sourceless VTs for
  743. # them and only guard on the frozenset id.
  744. # TODO support source for sets and remove the special logics here.
  745. items = [SourcelessBuilder.create(self.tx, v) for v in value]
  746. self.install_guards(GuardBuilder.ID_MATCH)
  747. return FrozensetVariable(items, source=self.source)
  748. elif isinstance(
  749. value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
  750. ):
  751. self.install_guards(GuardBuilder.ID_MATCH)
  752. return EnumVariable(value=value, source=self.source)
  753. elif DebuggingVariable.is_reorderable_logging_function(value):
  754. # Put this above builtin_callable so that print() can be handled
  755. # along with other builtin debugging functions
  756. self.install_guards(GuardBuilder.BUILTIN_MATCH)
  757. return DebuggingVariable(value, source=self.source)
  758. elif isinstance(value, logging.Logger):
  759. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  760. return LoggingLoggerVariable(value, source=self.source)
  761. elif is_utils_checkpoint(value):
  762. return build_checkpoint_variable(source=self.source)
  763. elif is_invoke_subgraph(value):
  764. return build_invoke_subgraph_variable(source=self.source)
  765. elif isinstance(value, functools.partial):
  766. func_src = AttrSource(self.get_source(), "func")
  767. func_obj = VariableBuilder(self.tx, func_src)(value.func)
  768. args = []
  769. args_source = AttrSource(self.get_source(), "args")
  770. for i, arg in enumerate(value.args):
  771. args.append(
  772. VariableBuilder(self.tx, GetItemSource(args_source, i))(arg)
  773. )
  774. keywords = {}
  775. keywords_source = AttrSource(self.get_source(), "keywords")
  776. for k, v in value.keywords.items():
  777. if not ConstantVariable.is_literal(k):
  778. unimplemented_v2(
  779. gb_type="functools.partial() with non-literal keyword",
  780. context=f"non-literal keyword: {k}",
  781. explanation="functools.partial() expects literal/string keywords",
  782. hints=[*graph_break_hints.USER_ERROR],
  783. )
  784. keywords[k] = VariableBuilder(
  785. self.tx, DictGetItemSource(keywords_source, k)
  786. )(v)
  787. install_guard(
  788. self.get_source().make_guard(GuardBuilder.TYPE_MATCH),
  789. keywords_source.make_guard(GuardBuilder.DICT_KEYS_MATCH),
  790. args_source.make_guard(GuardBuilder.SEQUENCE_LENGTH),
  791. )
  792. return FunctoolsPartialVariable(func_obj, args, keywords)
  793. elif is_typing(value):
  794. # typing.List, typing.Mapping, etc.
  795. self.install_guards(GuardBuilder.ID_MATCH)
  796. return TypingVariable(
  797. value,
  798. source=self.source,
  799. )
  800. elif np is not None and isinstance(value, np.generic):
  801. # numpy array scalars: convert to 0D arrays
  802. return self.wrap_numpy_ndarray(np.asarray(value))
  803. elif trace_rules.is_numpy(value):
  804. assert np
  805. self.install_guards(
  806. GuardBuilder.FUNCTION_MATCH
  807. if callable(value)
  808. else GuardBuilder.TYPE_MATCH
  809. )
  810. return NumpyVariable(value, source=self.source)
  811. elif trace_rules.is_numpy_dtype(value):
  812. self.install_guards(GuardBuilder.ID_MATCH)
  813. return NumpyDTypeVariable(value, source=self.source)
  814. elif trace_rules.is_numpy_type_info(value):
  815. if isinstance(value, np.iinfo):
  816. self.install_guards(GuardBuilder.TYPE_MATCH)
  817. dt_source = AttrSource(self.source, "dtype")
  818. install_guard(dt_source.make_guard(GuardBuilder.ID_MATCH))
  819. else:
  820. self.install_guards(GuardBuilder.ID_MATCH)
  821. return NumpyTypeInfoVariable(value, source=self.source)
  822. # NB: These can't be put in type_dispatch, they have to run later
  823. elif CollectiveFunctionRewriteVariable.can_rewrite(value):
  824. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  825. return CollectiveFunctionRewriteVariable.create(
  826. self.tx,
  827. value,
  828. source=self.source,
  829. )
  830. elif istype(value, torch.autograd.function.FunctionMeta):
  831. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  832. return AutogradFunctionVariable(
  833. value,
  834. source=self.source,
  835. )
  836. elif isinstance(value, torch.autograd.function.FunctionCtx):
  837. actual_saved_tensors = None
  838. try:
  839. actual_saved_tensors = value.saved_tensors
  840. except RuntimeError:
  841. pass
  842. saved_tensors = []
  843. guards = [self.source.make_guard(GuardBuilder.TYPE_MATCH)]
  844. if isinstance(actual_saved_tensors, tuple):
  845. saved_tensors_source = AttrSource(self.source, "saved_tensors")
  846. guards.append(
  847. saved_tensors_source.make_guard(GuardBuilder.SEQUENCE_LENGTH)
  848. )
  849. for i, v in enumerate(actual_saved_tensors):
  850. saved_tensors.append(
  851. VariableBuilder(
  852. self.tx, GetItemSource(saved_tensors_source, i)
  853. )(v)
  854. )
  855. install_guard(*guards)
  856. return self.tx.output.side_effects.track_object_existing(
  857. value,
  858. AutogradFunctionContextVariable(
  859. value,
  860. source=self.source,
  861. saved_tensors=SavedTensorBox(saved_tensors),
  862. ),
  863. )
  864. elif (
  865. isinstance(value, types.MethodType)
  866. and istype(
  867. getattr(value, "__self__", None), torch.autograd.function.FunctionMeta
  868. )
  869. and getattr(value, "__name__", "") == "apply"
  870. and value == getattr(value.__self__, "apply", None)
  871. ):
  872. # handle aliased autograd function `apply` calls
  873. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  874. return GetAttrVariable(
  875. AutogradFunctionVariable(
  876. value.__self__, source=AttrSource(self.source, member="__self__")
  877. ),
  878. "apply",
  879. )
  880. elif isinstance(value, torch._C._ImperativeEngine):
  881. self.install_guards(GuardBuilder.ID_MATCH)
  882. return AutogradEngineVariable(value, source=self.source)
  883. elif (
  884. value
  885. is torch._dynamo.external_utils.FakeCompiledAutogradEngine._exec_final_callbacks_stub
  886. ):
  887. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  888. return LambdaVariable(
  889. lambda: UserFunctionVariable(
  890. torch._dynamo.external_utils.FakeCompiledAutogradEngine.exec_final_callbacks,
  891. ).call_function(
  892. self.tx,
  893. (self.tx.output.side_effects.get_ca_final_callbacks_var(),),
  894. {},
  895. )
  896. )
  897. elif isinstance(value, DynamoConfigPatchProxy):
  898. return DynamoConfigPatchVariable(value.changes)
  899. elif isinstance(value, ErrorOnGraphBreakDecoratorContextManager):
  900. return ErrorOnGraphBreakVariable(value.error_on_graph_break)
  901. elif callable(value) and trace_rules.lookup_callable(value) is not None:
  902. if trace_rules.is_callable_allowed(value):
  903. self.tx.output.has_user_defined_allowed_in_graph = True
  904. return trace_rules.lookup_callable(value).create_with_source(
  905. value, source=self.source
  906. )
  907. elif np and isinstance(value, np.number):
  908. return self.wrap_unspecialized_primitive(value)
  909. elif isinstance(value, HigherOrderOperator):
  910. if value is torch._higher_order_ops.invoke_subgraph:
  911. unimplemented_v2(
  912. gb_type="Attempted to wrap torch._higher_order_ops.invoke_subgraph",
  913. context="",
  914. explanation="Directly using invoke_subgraph is not supported. Use nested_compile_region",
  915. hints=[],
  916. )
  917. self.install_guards(GuardBuilder.TYPE_MATCH, GuardBuilder.NAME_MATCH)
  918. return TorchHigherOrderOperatorVariable.make(value, source=self.source)
  919. elif isinstance(value, torch.cuda.StreamContext):
  920. self.install_guards(GuardBuilder.ID_MATCH)
  921. stream_source = AttrSource(self.source, "stream")
  922. stream_var = VariableBuilder(self.tx, stream_source)(value.stream)
  923. return StreamContextVariable.create(self.tx, stream_var)
  924. elif isinstance(value, torch.Stream):
  925. self.install_guards(GuardBuilder.ID_MATCH)
  926. stream_proxy = self.tx.output.create_proxy(
  927. "call_function",
  928. type(value),
  929. (),
  930. {
  931. "stream_id": value.stream_id,
  932. "device_index": value.device_index,
  933. "device_type": value.device_type,
  934. },
  935. )
  936. set_example_value(stream_proxy.node, value)
  937. return StreamVariable(
  938. stream_proxy,
  939. value,
  940. value.device,
  941. source=self.source,
  942. )
  943. elif isinstance(value, (torch._C._SDPAParams)):
  944. self.install_guards(GuardBuilder.TYPE_MATCH)
  945. return SDPAParamsVariable.create(self.tx, value, self.source)
  946. elif isinstance(value, torch._functorch.pyfunctorch.FuncTorchInterpreter):
  947. self.install_guards(GuardBuilder.ID_MATCH)
  948. return FuncTorchInterpreterVariable(value)
  949. elif isinstance(value, torch.Event):
  950. self.install_guards(GuardBuilder.ID_MATCH)
  951. torch._dynamo.utils.store_user_object_weakref(value)
  952. event_proxy = self.tx.output.create_proxy(
  953. "call_function",
  954. torch._dynamo.utils.get_user_object_from_id,
  955. (id(value),),
  956. {},
  957. )
  958. set_example_value(event_proxy.node, value)
  959. return EventVariable(
  960. event_proxy,
  961. value,
  962. source=self.source,
  963. )
  964. elif (
  965. istype(value, contextlib.nullcontext)
  966. and inspect.getattr_static(value, "enter_result", None) is None
  967. ):
  968. self.install_guards(GuardBuilder.TYPE_MATCH)
  969. return NullContextVariable(source=self.source)
  970. elif KeyedJaggedTensorVariable.is_matching_object(value):
  971. self.install_guards(GuardBuilder.TYPE_MATCH)
  972. result = KeyedJaggedTensorVariable(value, source=self.source)
  973. # TODO: this doing it manually is bad
  974. return self.tx.output.side_effects.track_object_existing(value, result)
  975. elif isinstance(value, torch.optim.Optimizer):
  976. self.install_guards(GuardBuilder.ID_MATCH)
  977. self.source = OptimizerSource(self.source)
  978. return OptimizerVariable(value, source=self.source)
  979. elif isinstance(value, torch.DispatchKeySet):
  980. self.install_guards(GuardBuilder.DISPATCH_KEY_SET_MATCH)
  981. return DispatchKeySetVariable(value)
  982. elif WorldMetaClassVariable.is_group_member_type(value):
  983. return WorldMetaClassVariable(value, source=self.source)
  984. elif ProcessGroupVariable.is_process_group(value):
  985. self.install_guards(GuardBuilder.ID_MATCH)
  986. return ProcessGroupVariable(value, source=self.source)
  987. elif DeviceMeshVariable.is_device_mesh(value):
  988. # TODO: see if we need to add custom guard instead of a simple ID_MATCH
  989. self.install_guards(GuardBuilder.EQUALS_MATCH)
  990. return DeviceMeshVariable(value, source=self.source)
  991. elif PlacementClassVariable.is_placement_type(value):
  992. # TODO: see if we need to add custom guard instead of a simple ID_MATCH
  993. self.install_guards(GuardBuilder.ID_MATCH)
  994. return PlacementClassVariable(value, source=self.source)
  995. elif PlacementVariable.is_placement(value):
  996. # TODO: see if we need to add custom guard instead of a simple ID_MATCH
  997. self.install_guards(GuardBuilder.EQUALS_MATCH)
  998. return PlacementVariable(
  999. value,
  1000. source=self.source,
  1001. )
  1002. elif (
  1003. id(value) in ITERTOOLS_TYPE_IDS
  1004. and id(value) not in ITERTOOLS_POLYFILLED_TYPE_IDS
  1005. ):
  1006. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  1007. return ItertoolsVariable(value, source=self.source)
  1008. elif is_torch_sym(value):
  1009. # Note: this doesn't handle nested symints.
  1010. # For SymBool input, we reuse the infra for SymInt by simulating SymBool with a SymInt in dynamo.
  1011. # Concretely,
  1012. # 1. We create a SymInt in dynamo's shape_env, whose source is constructed as ConvertIntSource(self.source).
  1013. # so that guards on the SymInts can be effectively applied on the original SymBool in user program.
  1014. # 2. We create a SymBool based on the SymInt in dynamo's ShapeEnv. Because the original user program
  1015. # depends on the value being a SymBool. This allows dynamo to interpret the user's program correctly.
  1016. source = (
  1017. self.source
  1018. if isinstance(value, torch.SymInt)
  1019. else ConvertIntSource(self.source)
  1020. )
  1021. if value.node.has_hint():
  1022. new_symint = (
  1023. self.tx.output.shape_env.create_unspecified_symint_and_symbol(
  1024. int(value.node.hint),
  1025. source,
  1026. dynamic_dim=DimDynamic.DYNAMIC,
  1027. )
  1028. )
  1029. else:
  1030. if isinstance(value, torch.SymBool):
  1031. # We need to create an unbacked symint to replace the unbacked symbool.
  1032. new_symint = self.tx.output.shape_env.create_unbacked_symint()
  1033. else:
  1034. # TODO (yidi): we need to figure out a way to propagate the guards
  1035. # we accumulated when tracing the subggraph to outer shape_env. For normal symints,
  1036. # this is automatically done by evaluating the guards once but this
  1037. # will cause data-dependent error when we evaluate the outer unbacked symints.
  1038. # The test case that triggers this graph break is test_cond_unbacked_symint_closure
  1039. unimplemented_v2(
  1040. gb_type="Attempted to wrap unbacked SymInt",
  1041. context="",
  1042. explanation="Unbacked SymInt input is not supported yet.",
  1043. hints=[*graph_break_hints.SUPPORTABLE],
  1044. )
  1045. sym_node_proxy = self.tx.output.root_tracer.create_graph_input(
  1046. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1047. type(new_symint),
  1048. new_symint,
  1049. source=source,
  1050. )
  1051. sym_node_proxy.node.meta["grapharg"] = GraphArg(
  1052. source,
  1053. new_symint,
  1054. False,
  1055. None,
  1056. is_tensor=False,
  1057. example_strong_ref=new_symint,
  1058. )
  1059. # We bind the new_symint to graph input.
  1060. sym_expr = new_symint.node.expr
  1061. assert isinstance(sym_expr, sympy.Symbol), (
  1062. f"{sym_expr} is not a basic Symbol."
  1063. )
  1064. self.tx.output.tracked_fakes.append(TrackedFake(new_symint, source, None))
  1065. tracing_symint = (
  1066. new_symint if isinstance(value, torch.SymInt) else new_symint == 1
  1067. ) # cast it back to symbool for tracing
  1068. return SymNodeVariable(sym_node_proxy, tracing_symint)
  1069. elif isinstance(value, (JITFunction, Autotuner)):
  1070. self.install_guards(GuardBuilder.ID_MATCH)
  1071. return TritonKernelVariable(
  1072. value,
  1073. None, # No kernel idx provided
  1074. None, # No grid provided
  1075. source=self.source,
  1076. )
  1077. elif value is create_1d_tma_descriptor:
  1078. return CreateTMADescriptorExperimentalVariable(rank=1)
  1079. elif value is create_2d_tma_descriptor:
  1080. return CreateTMADescriptorExperimentalVariable(rank=2)
  1081. elif value is TensorDescriptor.from_tensor:
  1082. return CreateTMADescriptorStableVariable()
  1083. elif isinstance(value, torch.amp.autocast_mode.autocast):
  1084. self.install_guards(GuardBuilder.ID_MATCH)
  1085. return AutocastModeVariable(
  1086. target_values=[
  1087. value.device,
  1088. value.fast_dtype,
  1089. value._enabled,
  1090. value._cache_enabled,
  1091. ],
  1092. source=self.source,
  1093. )
  1094. elif TorchCtxManagerClassVariable.is_matching_cls(value):
  1095. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  1096. return TorchCtxManagerClassVariable(value, source=self.source)
  1097. elif inspect.getattr_static(value, "__script_if_tracing_wrapper", False):
  1098. self.install_guards(GuardBuilder.TYPE_MATCH)
  1099. return WrapperUserFunctionVariable(
  1100. value, "__original_fn", source=self.source
  1101. )
  1102. elif is_lru_cache_wrapped_function(value):
  1103. self.install_guards(GuardBuilder.TYPE_MATCH)
  1104. return WrapperUserFunctionVariable(value, "__wrapped__", source=self.source)
  1105. elif value is traceback.clear_frames:
  1106. return TracebackVariable(source=self.source)
  1107. elif value is sys.exc_info or (
  1108. sys.version_info >= (3, 11) and value is sys.exception
  1109. ):
  1110. return SysFunctionVariable(value, source=self.source)
  1111. elif is_function_or_wrapper(value) and inspect.getattr_static(
  1112. value, "_torchdynamo_inline", False
  1113. ):
  1114. self.install_guards(GuardBuilder.TYPE_MATCH)
  1115. return WrapperUserFunctionVariable(
  1116. value, "_torchdynamo_inline", source=self.source
  1117. )
  1118. elif value is functools.wraps:
  1119. self.install_guards(GuardBuilder.ID_MATCH)
  1120. return FunctoolsWrapsVariable(value, source=self.source)
  1121. elif value is collections.namedtuple:
  1122. self.install_guards(GuardBuilder.ID_MATCH)
  1123. return CollectionsNamedTupleFunction(value, source=self.source)
  1124. elif isinstance(
  1125. value, types.BuiltinMethodType
  1126. ) and BuiltinMethodVariable.is_supported_builtin_method(value):
  1127. self.install_guards(GuardBuilder.ID_MATCH)
  1128. return BuiltinMethodVariable(value, source=self.source)
  1129. elif is_function(value) and value in (float.fromhex, float.hex):
  1130. self.install_guards(GuardBuilder.ID_MATCH)
  1131. return GetAttrVariable(
  1132. BuiltinVariable(float, source=self.source),
  1133. value.__name__,
  1134. )
  1135. elif is_function_or_wrapper(value):
  1136. value, attr_name = unwrap_with_attr_name_if_wrapper(value)
  1137. # For these wrappers, Dynamo points to the wrapped function,
  1138. # so source needs to be updated as well.
  1139. if attr_name is not None:
  1140. self.source = AttrSource(self.source, attr_name)
  1141. return trace_rules.lookup(value).create_with_source(
  1142. value, source=self.source
  1143. )
  1144. elif value is random.Random:
  1145. self.install_guards(GuardBuilder.ID_MATCH)
  1146. return RandomClassVariable(source=self.source)
  1147. elif istype(value, random.Random) and RandomVariable.is_supported_random_obj(
  1148. value
  1149. ):
  1150. self.install_guards(GuardBuilder.TYPE_MATCH)
  1151. result = RandomVariable(value, source=self.source)
  1152. self.tx.output.side_effects.track_mutable(value, result)
  1153. return result
  1154. # Don't use istype, since some python modules are not subclasses of types.ModuleType directly.
  1155. # E.g, type(torch.ops) -> <class 'torch._ops._Ops'>,
  1156. # type(torch.backends.cudnn) -> <class 'torch.backends.cudnn.CudnnModule'>
  1157. elif isinstance(value, (types.ModuleType, replay_record.DummyModule)):
  1158. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  1159. result = PythonModuleVariable(
  1160. value,
  1161. source=self.source,
  1162. )
  1163. self.tx.output.side_effects.track_object_existing(value, result)
  1164. return result
  1165. elif isinstance(value, types.MethodType) and isinstance(
  1166. value.__self__, (torch.nn.Module, torch.utils._pytree.TreeSpec)
  1167. ):
  1168. # don't let MethodTypes fall through to UserDefinedObject,
  1169. # which doesn't support 'CALL_FUNCTION'
  1170. # TODO(whc): Why do we limit this to methods on NNModules?
  1171. # I don't have a good reason for this, but it preserves the existing behavior
  1172. # for MBartForConditionalGeneration, which generates many graph breaks and OOMs otherwise.
  1173. # I suspect we probably want to relax this check and dig deeper there.
  1174. # In order to construct a MethodVariable in Dynamo, we start with an actual method obj from python,
  1175. # but need to separately wrap its underlying `__func__` and its `self` argument. We wrap `self` here
  1176. # and then `__func__` gets wrapped inside UserMethodVariable.
  1177. self_obj = VariableBuilder(
  1178. self.tx, source=AttrSource(self.source, "__self__")
  1179. )(value.__self__)
  1180. assert self_obj and isinstance(self_obj, VariableTracker), (
  1181. "Failed to produce a valid self obj"
  1182. )
  1183. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  1184. return UserMethodVariable(
  1185. value.__func__,
  1186. self_obj,
  1187. source=self.source,
  1188. )
  1189. elif isinstance(value, types.GetSetDescriptorType):
  1190. # GetSet descriptors are C functions attached to an attribute lookup
  1191. # using PyGetSetDef. Python, on attribute lookup, can decide to
  1192. # create a new object on the fly, and therefore the `id` of the
  1193. # descriptors is not guaranteed to be same for different attribute
  1194. # accesses. Since these are unlikely to change during the program
  1195. # execution, we can skip guarding on them.
  1196. return GetSetDescriptorVariable(value)
  1197. elif isinstance(value, types.MethodWrapperType):
  1198. # Method-wrappers are written in C, and they are not guaranteed to
  1199. # return the same object on attribute lookup. Therefore, we cannot
  1200. # insert a FUNCTION_MATCH guard here. method-wrappers are very
  1201. # unlikely to change, so its ok to skip the guard here.
  1202. return MethodWrapperVariable(value)
  1203. elif issubclass(type(value), type) and issubclass(value, BaseException):
  1204. # match user defined exceptions
  1205. self.install_guards(GuardBuilder.ID_MATCH)
  1206. return UserDefinedExceptionClassVariable(value)
  1207. elif issubclass(type(value), type):
  1208. if value in (
  1209. torch.utils.hooks.BackwardHook,
  1210. torch.nn.Parameter,
  1211. torch.nn.Buffer,
  1212. ):
  1213. # TODO(jansel): combine this case with the one above
  1214. return trace_rules.lookup(value).create_with_source(
  1215. value, source=self.source
  1216. )
  1217. if value is torch.autograd._unsafe_preserve_version_counter:
  1218. self.install_guards(GuardBuilder.FUNCTION_MATCH)
  1219. return PreserveVersionContextVariable.constructor(self.tx)
  1220. if (
  1221. # `value` must be a strict subclass of `torch.Tensor`
  1222. issubclass(value, torch.Tensor)
  1223. and value is not torch.Tensor
  1224. # `TensorSubclassVariable` is not for subclass that overrides
  1225. # `torch_dispatch`.
  1226. and value.__torch_dispatch__ is torch.Tensor.__torch_dispatch__
  1227. # `TensorSubclassVariable` would lead to construction of
  1228. # `TensorWithTFOverrideVariable`, but we don't want that for
  1229. # traceable wrapper subclasses (we wrap those subclass instances
  1230. # into `TensorVariable`).
  1231. and not is_traceable_wrapper_subclass_type(value)
  1232. ):
  1233. return TensorSubclassVariable(value, source=self.source)
  1234. if not is_from_closure_source(self.source):
  1235. # For closure source, the variable comes from LOAD_SUPER_ATTR,
  1236. # which calls self.__class__. This is internal Cpython
  1237. # implementation, and it is rare for the user to modify
  1238. # self.__class__ manually.
  1239. # For other cases, this is a userdefined class, so install an
  1240. # ID_MATCH even if its a global variable.
  1241. self.install_guards(GuardBuilder.ID_MATCH)
  1242. return UserDefinedClassVariable(
  1243. value,
  1244. source=self.source,
  1245. )
  1246. elif TorchScriptObjectVariable.is_matching_cls(type(value)):
  1247. from ..source import (
  1248. FlattenScriptObjectSource,
  1249. ScriptObjectQualifiedNameSource,
  1250. )
  1251. if torch._library.fake_class_registry.tracing_with_real(value):
  1252. proxy = self.tx.output.root_tracer.create_graph_input(
  1253. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1254. type(value),
  1255. value,
  1256. source=self.source,
  1257. )
  1258. # setting is_unspecialized=False to not insert a as_tensor call in reconstruct by default
  1259. # setting example to be real value because these example values will be used
  1260. # as example_inputs for user compiler.
  1261. proxy.node.meta["grapharg"] = GraphArg(
  1262. self.source, value, False, None, False, value
  1263. )
  1264. return TorchScriptObjectVariable.create(
  1265. proxy,
  1266. value,
  1267. source=self.source,
  1268. )
  1269. # This exists to allow a smoother transition.
  1270. # The implications are:
  1271. # The script objects won't be tracked as proxies.
  1272. # Methods on these objects won't show up in the graph.
  1273. # The original script object might be mutated.
  1274. if not hasattr(value, "__obj_flatten__"):
  1275. return self.wrap_user_defined(value)
  1276. # Install the guards on the fully qualified name of the script object
  1277. LazyVariableTracker.realize_all(
  1278. VariableBuilder(self.tx, ScriptObjectQualifiedNameSource(self.source))(
  1279. value._type().qualified_name() # type: ignore[attr-defined]
  1280. )
  1281. )
  1282. # Install the guards on the content of the script object by setting the source
  1283. # to be FlattenScriptObjectSource, which calls __obj_flatten__() to get the contents.
  1284. LazyVariableTracker.realize_all(
  1285. VariableBuilder(self.tx, FlattenScriptObjectSource(self.source))(
  1286. value.__obj_flatten__()
  1287. )
  1288. )
  1289. fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
  1290. self.tx.output.fake_mode, value
  1291. )
  1292. proxy = self.tx.output.root_tracer.create_graph_input(
  1293. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1294. type(value),
  1295. fake_script_obj,
  1296. source=self.source,
  1297. )
  1298. # setting is_unspecialized=False to not insert a as_tensor call in reconstruct by default
  1299. # setting example to be real value because these example values will be used
  1300. # as example_inputs for user compiler.
  1301. proxy.node.meta["grapharg"] = GraphArg(
  1302. self.source, value, False, None, False, fake_script_obj
  1303. )
  1304. return TorchScriptObjectVariable.create(
  1305. proxy,
  1306. fake_script_obj,
  1307. source=self.source,
  1308. )
  1309. elif (
  1310. isinstance(value, (dict, collections.OrderedDict))
  1311. and type(value).__new__ is dict.__new__
  1312. ):
  1313. # Construct a dict_vt that will reside inside the UserDefinedDictVariable
  1314. self.install_guards(GuardBuilder.TYPE_MATCH)
  1315. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1316. # Guard on the key order
  1317. self.tx.output.guard_on_key_order.add(self.source)
  1318. # We need all the keys to be hashable. We do this within the
  1319. # _HashableTracker class in dicts.py
  1320. def build_key_value(i, k, v):
  1321. base = self.get_source()
  1322. source_key = ConstDictKeySource(base, i)
  1323. key = LazyVariableTracker.create(k, source_key)
  1324. source_value = DictSubclassGetItemSource(base, source_key)
  1325. res_value = LazyVariableTracker.create(v, source_value)
  1326. return key, res_value
  1327. # Ensure that we call dict.keys and not value.keys (which can call
  1328. # overridden keys method). In the C++ guards, we relied on
  1329. # PyDict_Next to traverse the dictionary, which uses the internal
  1330. # data structure and does not call the overridden keys method.
  1331. result = dict(
  1332. build_key_value(i, k, v)
  1333. for i, (k, v) in enumerate(get_items_from_dict(value))
  1334. )
  1335. dict_vt = ConstDictVariable(
  1336. result,
  1337. user_cls=(
  1338. collections.OrderedDict
  1339. if isinstance(value, collections.OrderedDict)
  1340. else dict
  1341. ),
  1342. mutation_type=ValueMutationExisting(),
  1343. source=self.source,
  1344. )
  1345. # Force this to reconstruct on mutation to keep the reconstruction
  1346. # bytecode simple
  1347. dict_vt.should_reconstruct_all = True
  1348. result = UserDefinedDictVariable(value, dict_vt=dict_vt, source=self.source)
  1349. return self.tx.output.side_effects.track_object_existing(value, result)
  1350. elif isinstance(value, tuple):
  1351. self.install_guards(GuardBuilder.TYPE_MATCH)
  1352. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1353. # NB - Be careful in not triggering user code. Guards also work on
  1354. # the underlying tuple data structure.
  1355. output = [
  1356. LazyVariableTracker.create(
  1357. tuple.__getitem__(value, i),
  1358. source=GetItemSource(self.get_source(), i),
  1359. )
  1360. for i in range(tuple.__len__(value))
  1361. ]
  1362. tuple_vt = TupleVariable(
  1363. output, source=self.source, mutation_type=ValueMutationExisting()
  1364. )
  1365. result = UserDefinedTupleVariable(
  1366. value, tuple_vt=tuple_vt, source=self.source
  1367. )
  1368. return self.tx.output.side_effects.track_object_existing(value, result)
  1369. elif isinstance(value, list):
  1370. self.install_guards(GuardBuilder.TYPE_MATCH)
  1371. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1372. # NB - Be careful in not triggering user code. Guards also work on
  1373. # the underlying list data structure.
  1374. output = [
  1375. LazyVariableTracker.create(
  1376. list.__getitem__(value, i),
  1377. source=ListGetItemSource(self.get_source(), i),
  1378. )
  1379. for i in range(list.__len__(value))
  1380. ]
  1381. list_vt = ListVariable(
  1382. output, source=self.source, mutation_type=ValueMutationExisting()
  1383. )
  1384. result = UserDefinedListVariable(value, list_vt=list_vt, source=self.source)
  1385. return self.tx.output.side_effects.track_object_existing(value, result)
  1386. elif isinstance(value, (set, frozenset)):
  1387. self.install_guards(GuardBuilder.TYPE_MATCH)
  1388. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1389. L = list(dict.fromkeys(value))
  1390. output = [
  1391. LazyVariableTracker.create(
  1392. list.__getitem__(L, i),
  1393. source=NonSerializableSetGetItemSource(self.get_source(), i),
  1394. )
  1395. for i in range(list.__len__(L))
  1396. ]
  1397. set_vt_cls = SetVariable if isinstance(value, set) else FrozensetVariable
  1398. set_vt = set_vt_cls(
  1399. output, source=self.source, mutation_type=ValueMutationExisting()
  1400. )
  1401. result = UserDefinedSetVariable(value, set_vt=set_vt, source=self.source)
  1402. return self.tx.output.side_effects.track_object_existing(value, result)
  1403. elif issubclass(type(value), MutableMapping):
  1404. self.install_guards(GuardBuilder.TYPE_MATCH)
  1405. result = MutableMappingVariable(value, source=self.source)
  1406. return self.tx.output.side_effects.track_object_existing(value, result)
  1407. elif is_frozen_dataclass(value):
  1408. self.install_guards(GuardBuilder.TYPE_MATCH)
  1409. result = FrozenDataClassVariable.create(self.tx, value, source=self.source)
  1410. return self.tx.output.side_effects.track_object_existing(value, result)
  1411. elif isinstance(value, dict_keys):
  1412. if all(ConstantVariable.is_literal(k) for k in value):
  1413. # If the dict_keys object is passed from outside the compile region, it must either be passed along with
  1414. # the corresponding dict object or treated as a set (when only the keys are passed into the compiled region).
  1415. # - If it is passed along with the dict, the dict object itself is already guarded.
  1416. # - If only the dict_keys object is passed, we add EQUALS_MATCH and SEQUENCE_LENGTH guards
  1417. # to ensure it remains unchanged across multiple runs.
  1418. items = [SourcelessBuilder.create(self.tx, v) for v in value]
  1419. install_guard(
  1420. self.get_source().make_guard(GuardBuilder.SEQUENCE_LENGTH),
  1421. self.get_source().make_guard(GuardBuilder.EQUALS_MATCH),
  1422. )
  1423. return DictKeySetVariable(items, source=self.source)
  1424. else:
  1425. unimplemented_v2(
  1426. gb_type="non-const keys in dict_keys",
  1427. context=f"non-const keys: {[k for k in value if not ConstantVariable.is_literal(k)]}",
  1428. explanation="Dynamo expects dict_keys keys to be constants.",
  1429. hints=[
  1430. "Ensure your dict_keys keys are constants (e.g. int, float, strings)",
  1431. ],
  1432. )
  1433. elif IntWrapperVariable.is_matching_object(value):
  1434. from torch.export.dynamic_shapes import _DimHintType
  1435. if value.dynamism is None or value.dynamism.type == _DimHintType.STATIC:
  1436. return self.wrap_symint(value.val)
  1437. elif value.dynamism.type == _DimHintType.DYNAMIC:
  1438. log.debug(
  1439. "%s marked %s via IntWrapper",
  1440. self.source.name(),
  1441. DimDynamic.DYNAMIC,
  1442. )
  1443. return self.wrap_symint(
  1444. value.val,
  1445. dynamism=DimDynamic.DYNAMIC,
  1446. context=SymIntSymbolicContext(
  1447. constraint=RelaxedUnspecConstraint(warn_only=False)
  1448. ),
  1449. )
  1450. elif value.dynamism.type == _DimHintType.AUTO:
  1451. log.debug(
  1452. "%s marked %s via IntWrapper",
  1453. self.source.name(),
  1454. DimDynamic.DYNAMIC,
  1455. )
  1456. return self.wrap_symint(value.val, dynamism=DimDynamic.DYNAMIC)
  1457. else:
  1458. raise RuntimeError(f"Undefined dynamism {value.dynamism}")
  1459. else:
  1460. return self.wrap_user_defined(value)
  1461. def wrap_user_defined(self, value: Any):
  1462. self.install_guards(GuardBuilder.TYPE_MATCH)
  1463. result = UserDefinedObjectVariable(value, source=self.source)
  1464. if not SideEffects.cls_supports_mutation_side_effects(type(value)):
  1465. # don't allow STORE_ATTR mutation with custom __setattr__
  1466. return result
  1467. return self.tx.output.side_effects.track_object_existing(value, result)
  1468. def wrap_listlike(self, value: Union[tuple, list, odict_values, NamedTuple]):
  1469. for item in value:
  1470. if item is value:
  1471. unimplemented_v2(
  1472. gb_type="list elements are pointing to the list itself",
  1473. context="",
  1474. explanation="Dynamo does not support lists whose items reference to itself",
  1475. hints=["Avoid using self referential list"],
  1476. )
  1477. if config.specialize_int and type(value) is torch.Size:
  1478. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  1479. return ConstantVariable.create(value=value)
  1480. # One can index a tensor with a list/tuple. Therefore, we need to
  1481. # have a stricter match.
  1482. self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
  1483. # Tuples are immutable objects, so we should mark its items static. This
  1484. # avoids wrapping of tuple items as symints. This helps for nn module
  1485. # attributes like conv2d strides, dilations.
  1486. if (
  1487. istype(value, tuple)
  1488. and all(ConstantVariable.is_literal(item) for item in value)
  1489. and self.source.guard_source().is_unspecialized_nn_module()
  1490. ):
  1491. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  1492. return TupleVariable([ConstantVariable.create(item) for item in value])
  1493. output = [
  1494. LazyVariableTracker.create(
  1495. item,
  1496. source=GetItemSource(self.get_source(), i),
  1497. )
  1498. for i, item in enumerate(value)
  1499. ]
  1500. maybe_gm = self.tx.output.local_scope.get("self")
  1501. if isinstance(
  1502. self.source, LocalSource
  1503. ) and self.source.local_name in get_locals_to_steal(maybe_gm):
  1504. # The input tensor list to dynamo from compiled autograd may contain activations
  1505. # which are freed as they are used in inductor. Dynamo's default behavior is to
  1506. # lift all tensors to the graph inputs, but this will cause dynamo to hold an
  1507. # extra reference to the activation tensors and increase peak memory usage.
  1508. # To allow freeing ASAP, we keep the list as graph argument to the dynamo output
  1509. # graph, and unpack it locally.
  1510. # e.g. instead of `def forward(self, L_inputs_0_, L_inputs_1_, ...):`, we have
  1511. # `def forward(self, L_inputs_):`
  1512. source = self.source
  1513. assert isinstance(value, list)
  1514. tensor_list_proxy = self.tx.output.root_tracer.create_graph_input(
  1515. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1516. type(value),
  1517. value,
  1518. source=source,
  1519. )
  1520. tensor_list_proxy.node.meta["steal_arg"] = True
  1521. list_variable = wrap_fx_proxy_cls(
  1522. target_cls=TensorVariable,
  1523. tx=self.tx,
  1524. proxy=tensor_list_proxy,
  1525. example_value=value,
  1526. subclass_type=None,
  1527. source=source,
  1528. )
  1529. # Apply relevant logic from `VariableTracker.build(value[i])`
  1530. # (except for the `create_graph_input` stuff).
  1531. guards = []
  1532. for i, tensor_variable in enumerate(list_variable.items):
  1533. source_i = GetItemSource(base=source, index=i, index_is_slice=False)
  1534. # access unpacked tensor from this list instead of from a lifted arg
  1535. self.tx.output.input_source_to_var[source_i] = tensor_variable
  1536. tensor_variable.proxy.node.meta["tensor_dict"] = _extract_tensor_dict(
  1537. value[i]
  1538. )
  1539. guard = functools.partial(
  1540. GuardBuilder.TENSOR_MATCH, value=TensorWeakRef(value[i])
  1541. )
  1542. guards.append(source_i.make_guard(guard))
  1543. install_guard(*guards, skip=1)
  1544. grapharg = GraphArg(
  1545. source,
  1546. value,
  1547. pass_arg_as_tensor=False,
  1548. fake_tensor=None,
  1549. is_tensor=False,
  1550. )
  1551. tensor_list_proxy.node.meta["grapharg"] = grapharg
  1552. # The following is very important for maintaining the "python object
  1553. # <==> variable tracker" 1-to-1 mapping, which is mainly handled via
  1554. # `side_effects`. Note that constructing `tensor_variable` above
  1555. # already adds it to graph arg, but we never registered it with
  1556. # `side_effects`. The preemptive `realize` calls here basically
  1557. # does that registration (at the end of `self.__call__`).
  1558. #
  1559. # A slightly cleaner alternative is to register the
  1560. # `tensor_variable`s above with `side_effects` directly, and just
  1561. # return the `list_variable`, but that breaks some tensor-subclass
  1562. # related tests like `test_inputs_aliasing_bytecode_stack_restore`,
  1563. # because `tensor_variable` is constructed via
  1564. # `handle_traced_output`, which doesn't really expect/handle tensor
  1565. # subclass.
  1566. #
  1567. # Eventually, we expect to fix remove all of these by having Dynamo
  1568. # auto-boxing inputs to the compiled graph, see
  1569. # https://github.com/pytorch/pytorch/issues/153701.
  1570. for vt in output:
  1571. vt.realize()
  1572. result = BaseListVariable.cls_for_instance(value)(output, source=self.source)
  1573. if istype(value, (list, collections.deque)):
  1574. return self.tx.output.side_effects.track_mutable(value, result)
  1575. return result
  1576. def wrap_tuple_iterator(self, value: tuple_iterator):
  1577. self.install_guards(GuardBuilder.TUPLE_ITERATOR_LEN)
  1578. output = [
  1579. VariableBuilder(self.tx, TupleIteratorGetItemSource(self.get_source(), i))(
  1580. tuple_iterator_getitem(value, i)
  1581. )
  1582. for i in range(tuple_iterator_len(value))
  1583. ]
  1584. result = TupleIteratorVariable(output, source=self.source)
  1585. return self.tx.output.side_effects.track_mutable(value, result)
  1586. def wrap_range_iterator(self, value: range_iterator):
  1587. self.install_guards(GuardBuilder.RANGE_ITERATOR_MATCH)
  1588. # Get all the values from the range iterator; no need to install guards
  1589. # on items since `RANGE_ITERATOR_MATCH` guarantees the same items.
  1590. items = [ConstantVariable.create(v) for v in copy.deepcopy(value)]
  1591. result = ListIteratorVariable(items, source=self.source)
  1592. return self.tx.output.side_effects.track_mutable(value, result)
  1593. def wrap_slice_range(self, value: Union[slice, range]):
  1594. items = [
  1595. VariableBuilder(self.tx, AttrSource(self.get_source(), k))(
  1596. getattr(value, k)
  1597. )
  1598. for k in ("start", "stop", "step")
  1599. ]
  1600. self.install_guards(GuardBuilder.TYPE_MATCH)
  1601. if isinstance(value, slice):
  1602. return SliceVariable(items, source=self.source)
  1603. else:
  1604. return RangeVariable(items, source=self.source)
  1605. def mark_static_input(self, value: torch.Tensor, guard: bool):
  1606. from ..decorators import mark_static_address
  1607. static_inputs_log.debug(
  1608. "Marking static input %s, id: %s)", self.source.name(), id(value)
  1609. )
  1610. mark_static_address(value, guard=guard)
  1611. # Check if we've seen this tensor before and update graph metadata if needed
  1612. # As long as this runs before AOT this is sound
  1613. if value in self.tx.output.side_effects:
  1614. var = self.tx.output.side_effects[value]
  1615. var.proxy.node.meta["tensor_dict"]["_dynamo_static_input_type"] = (
  1616. value._dynamo_static_input_type
  1617. )
  1618. def wrap_module(self, value: torch.nn.Module):
  1619. from ..eval_frame import OptimizedModule
  1620. if len(value.__dict__) == 0:
  1621. unimplemented_v2(
  1622. gb_type="Uninitialized nn.Module",
  1623. context=typestr(value),
  1624. explanation=f"Attempted to trace an uninitialized nn.Module of type {typestr(value)}.",
  1625. hints=[
  1626. *graph_break_hints.USER_ERROR,
  1627. "Ensure your nn.Module instance has called `super().__init__()`.",
  1628. ],
  1629. )
  1630. if istype(value, OptimizedModule):
  1631. # Check if the optimized module was disabled
  1632. if inspect.getattr_static(value.forward, "_torchdynamo_disable", False):
  1633. # This bytecode is mostly of kind LOAD_ATTR or LOAD_METHOD. If
  1634. # we graph break here, Dynamo does not know how to create
  1635. # continuation functions for such bytecodes. So, we delay the
  1636. # graph break to CALL_FUNCTION.
  1637. msg = inspect.getattr_static(
  1638. value.forward, "_torchdynamo_disable_msg", None
  1639. )
  1640. return DelayGraphBreakVariable(
  1641. source=self.source,
  1642. msg=f"Optimized `nn.Module` is wrapped with `torch.compiler.disable` (reason: {msg})",
  1643. )
  1644. self.install_guards(GuardBuilder.TYPE_MATCH)
  1645. self.source = AttrSource(self.source, "_orig_mod")
  1646. return self.wrap_module(value._orig_mod)
  1647. if (
  1648. isinstance(value, (torch.nn.RNN, torch.nn.GRU, torch.nn.LSTM))
  1649. and not config.allow_rnn
  1650. ):
  1651. unimplemented_v2(
  1652. gb_type="Attempted to wrap RNN, GRU, or LSTM",
  1653. context=str(value),
  1654. explanation="Dynamo does not support RNN, GRU, or LSTM.",
  1655. hints=[*graph_break_hints.SUPPORTABLE],
  1656. )
  1657. if getattr(value, "_is_fsdp_managed_module", False):
  1658. # See note [Dynamo treats FSDP wrapped modules as UnspecializedNNModule]
  1659. # in fully_sharded_data_parallel.py for more information
  1660. # we can't do this assert inside FSDP constructor,
  1661. # since we don't know yet whether dynamo will be used
  1662. if not getattr(value, "_fsdp_use_orig_params", False):
  1663. unimplemented_v2(
  1664. gb_type="FSDP with use_orig_params=False",
  1665. context="",
  1666. explanation="Dynamo only supports FSDP with use_orig_params=True",
  1667. hints=[],
  1668. )
  1669. # Note on FSDP guarding
  1670. # Eager FSDP already assumes (requires, but without enforcement)
  1671. # that users don't mutate their model parameters/structure after
  1672. # FSDP wrapping, because FSDP wouldn't notice or update its
  1673. # FlatParams.
  1674. #
  1675. # Therefore, torch.compile can skip guarding on params or submodule
  1676. # structure of fsdp_managed modules, by using FSDPNNModuleSource as
  1677. # the guard source. This behavior is gated on
  1678. # config.skip_fsdp_guards.
  1679. self.install_guards(GuardBuilder.TYPE_MATCH)
  1680. result = FSDPManagedNNModuleVariable(value, source=self.get_source())
  1681. if not SideEffects.cls_supports_mutation_side_effects(type(value)):
  1682. # don't allow STORE_ATTR mutation with custom __setattr__
  1683. return result
  1684. return self.tx.output.side_effects.track_object_existing(value, result)
  1685. elif mutation_guard.is_dynamic_nn_module(value, self.tx.export):
  1686. # created dynamically, don't specialize on it
  1687. # Note [Tracing a torch.compiled function]
  1688. # when make_fx tracing a compiled function, we need
  1689. if isinstance(value, torch.fx.experimental.proxy_tensor._AttrProxy):
  1690. value = value.get_base()
  1691. self.source = AttrProxySource(self.source)
  1692. if torch._dynamo.config.inline_inbuilt_nn_modules:
  1693. freezing = is_parameter_freezing()
  1694. # Guard against the case where user may overwrite named parameters
  1695. # / named buffers
  1696. # NOTE: This is not likely to happen but worth guarding to avoid
  1697. # exception
  1698. if (
  1699. callable(value.named_parameters)
  1700. and value.named_parameters.__func__
  1701. is og_module_named_parameters_fn_ptr
  1702. ):
  1703. try: # catch TypeErrors in named_parameters() from unserializable nn modules
  1704. for _, p in value.named_parameters():
  1705. self.mark_static_input(p, guard=freezing)
  1706. except TypeError as e:
  1707. raise_observed_exception(type(e), self.tx, args=list(e.args))
  1708. if (
  1709. callable(value.named_buffers)
  1710. and value.named_buffers.__func__ is og_module_named_buffers_fn_ptr
  1711. ):
  1712. try: # catch TypeErrors in named_parameters() from unserializable nn modules
  1713. for _, b in value.named_buffers():
  1714. self.mark_static_input(b, guard=freezing)
  1715. except TypeError as e:
  1716. raise_observed_exception(type(e), self.tx, args=list(e.args))
  1717. if freezing:
  1718. # we need to add the module to tracing context
  1719. # in order to allow its params to get invalidated
  1720. # this will get cleaned up once compile ends
  1721. self.tx.output.nn_modules[self.name] = value
  1722. if (
  1723. value.__module__.startswith(("torch.nn.modules", "torch.ao."))
  1724. and not value.__module__.startswith("torch.nn.modules.container")
  1725. ) or getattr(value.__class__, "_dynamo_marked_static", False):
  1726. new_source = self.source
  1727. if config.inline_inbuilt_nn_modules and (
  1728. not self.tx.output.export or config.install_free_tensors
  1729. ):
  1730. # Export corner case - look at test_repros.py test_inlining_cornercase
  1731. new_source = UnspecializedBuiltinNNModuleSource(self.source)
  1732. result = UnspecializedBuiltinNNModuleVariable(value, source=new_source)
  1733. install_guard(new_source.make_guard(GuardBuilder.TYPE_MATCH))
  1734. else:
  1735. new_source = self.source
  1736. if config.inline_inbuilt_nn_modules and (
  1737. not self.tx.output.export or config.install_free_tensors
  1738. ):
  1739. # Export corner case - look at test_repros.py test_inlining_cornercase
  1740. new_source = UnspecializedNNModuleSource(self.source)
  1741. result = UnspecializedNNModuleVariable(value, source=new_source)
  1742. install_guard(new_source.make_guard(GuardBuilder.TYPE_MATCH))
  1743. if not SideEffects.cls_supports_mutation_side_effects(type(value)):
  1744. # don't allow STORE_ATTR mutation with custom __setattr__
  1745. return result
  1746. return self.tx.output.side_effects.track_object_existing(value, result)
  1747. elif issubclass(
  1748. value.__class__, torch.nn.parallel.distributed.DistributedDataParallel
  1749. ):
  1750. self.install_guards(GuardBuilder.TYPE_MATCH)
  1751. return UnspecializedNNModuleVariable(value, source=self.get_source())
  1752. else:
  1753. return self.tx.output.register_attr_or_module(
  1754. value,
  1755. self.name,
  1756. source=self.get_source(),
  1757. # Guards are added inside register_attr_or_module
  1758. )
  1759. def wrap_literal(self, value):
  1760. if type(value) is int:
  1761. # allowlist has higher precedence over specialization control.
  1762. if is_dynamic_source(self.source.name()):
  1763. log.debug("%s marked dynamic via source whitelist", self.source.name())
  1764. return self.wrap_symint(value, dynamism=DimDynamic.DYNAMIC)
  1765. if is_unbacked_source(self.source.name()):
  1766. log.debug("%s marked unbacked via source whitelist", self.source.name())
  1767. return self.wrap_symint(value, dynamism=DimDynamic.SIZE_LIKE_UNBACKED)
  1768. if not config.specialize_int:
  1769. # unspecializing int by default, but still
  1770. # specialize for the following conditions
  1771. if is_int_specialization_case(value, self.source):
  1772. recompile_hint = None
  1773. if (
  1774. self.source.guard_source().is_unspecialized_builtin_nn_module()
  1775. or self.source.guard_source().is_unspecialized_nn_module()
  1776. ):
  1777. # This means that it is an integer from a NN module.
  1778. # Dynamo considers nn module int attributes to be static
  1779. # (a good heuristic). But a user might want to mark the
  1780. # int attribute to be a symint, so track this integer
  1781. # for recompilation later.
  1782. recompile_hint = (
  1783. "torch.compile considers integer attributes of the nn.Module to be static. "
  1784. "If you are observing recompilation, you might want to make this integer dynamic "
  1785. "using torch._dynamo.config.allow_unspec_int_on_nn_module = True, or convert this "
  1786. "integer into a tensor."
  1787. )
  1788. process_automatic_dynamic(
  1789. self.tx,
  1790. self.source.name(),
  1791. FrameStateSizeEntry.make_scalar(value),
  1792. is_unspecialized_nn_module=self.source.guard_source().is_unspecialized_nn_module(),
  1793. )
  1794. self.install_guards(
  1795. functools.partial(
  1796. GuardBuilder.EQUALS_MATCH, recompile_hint=recompile_hint
  1797. )
  1798. )
  1799. return ConstantVariable.create(value=value, source=self.source)
  1800. return self.wrap_symint(value)
  1801. elif not config.specialize_float and type(value) is float:
  1802. return self.wrap_symfloat(value)
  1803. else:
  1804. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  1805. result = ConstantVariable.create(value=value, source=self.source)
  1806. if isinstance(value, (list, set)):
  1807. return self.tx.output.side_effects.track_mutable(value, result)
  1808. return result
  1809. def assert_not_wrapped_by_this_graph(self, value: torch.Tensor):
  1810. if is_fake(value) and maybe_get_fake_mode(value) is self.tx.fake_mode:
  1811. raise InternalTorchDynamoError(
  1812. "Cannot wrap a Tensor that has already been",
  1813. "wrapped by this instance of Dynamo",
  1814. )
  1815. def wrap_tensor(self, value: torch.Tensor):
  1816. source = self.get_source()
  1817. # We cannot already be tracking the tensor, which implies
  1818. # it would have already been wrapped
  1819. assert value not in self.tx.output.side_effects
  1820. is_static_input = get_static_address_type(value) is not None
  1821. if (
  1822. config.inline_inbuilt_nn_modules
  1823. and not is_static_input
  1824. and (
  1825. isinstance(value, torch.nn.Parameter)
  1826. # mark tensor attributes of nn modules static. This is done to keep inline_inbuilt_nn_modules behavior
  1827. # compatible with previous behavior.
  1828. or (source and source.guard_source().is_unspecialized_nn_module())
  1829. )
  1830. ):
  1831. self.mark_static_input(value, guard=is_parameter_freezing())
  1832. is_static_input = True
  1833. # Install any tensors which are "free" variables; that is:
  1834. # 1. Globals
  1835. # 2. NonLocals
  1836. # 3. tensors that are attributes of nn module
  1837. should_install_free_tensor = config.install_free_tensors and (
  1838. is_from_global_source(source)
  1839. or is_from_nonlocal_source(source)
  1840. or is_from_unspecialized_nn_module_source(source)
  1841. )
  1842. make_graph_attribute = is_static_input and (
  1843. not config.inline_inbuilt_nn_modules
  1844. or is_parameter_freezing()
  1845. or torch._dynamo.config.prepare_freezing
  1846. )
  1847. if should_install_free_tensor or (
  1848. (source.guard_source().is_specialized_nn_module() or make_graph_attribute)
  1849. and not source.guard_source().is_fsdp_module()
  1850. ):
  1851. self.assert_not_wrapped_by_this_graph(value)
  1852. return self.tx.output.register_attr_or_module(
  1853. value, self.name, source=source
  1854. )
  1855. if get_static_address_type(value) == "guarded":
  1856. # If it's a guarded tensor, we can install the parameter directly
  1857. # into the Fx graph instead of lifting it as an input. Lifting
  1858. # offers no benefit, such as regional compilation, since we still
  1859. # guard on the tensor's ID. Moreover, installing it in the Fx graph
  1860. # eliminates the pre-graph bytecode required to extract the tensor
  1861. # from locals/globals, reducing overhead. This can lead to
  1862. # significant cost savings, especially for optimizers handling many
  1863. # tensors.
  1864. self.install_guards(GuardBuilder.ID_MATCH)
  1865. self.assert_not_wrapped_by_this_graph(value)
  1866. return self.tx.output.register_attr_or_module(
  1867. value, self.name, source=source
  1868. )
  1869. if is_constant_source(source):
  1870. self.assert_not_wrapped_by_this_graph(value)
  1871. return self.tx.output.register_attr_or_module(
  1872. value,
  1873. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1874. source=source,
  1875. # Guards are added inside register_attr_or_module
  1876. )
  1877. # NB: this just says we accessed a tensor from the same source again
  1878. # (e.g., a tensor lives in a global foo, and we LOAD_GLOBAL it twice).
  1879. # This is distinct from two distinct sources mapping to the same
  1880. # Tensor (per id())! No guard is necessary here. See below for the
  1881. # other case.
  1882. is_duplicate_tensor = source in self.tx.output.input_source_to_var
  1883. if is_duplicate_tensor:
  1884. return self.tx.output.input_source_to_var[source]
  1885. options = {}
  1886. subclass_type = infer_subclass_type(value)
  1887. if subclass_type is not None:
  1888. self.install_guards(GuardBuilder.TYPE_MATCH)
  1889. if get_static_address_type(value) == "guarded":
  1890. self.install_guards(GuardBuilder.ID_MATCH)
  1891. # By this point, we should have deduplicated all tensors
  1892. self.assert_not_wrapped_by_this_graph(value)
  1893. if (
  1894. isinstance(value, torch.Tensor)
  1895. and value.is_nested
  1896. and not isinstance(value, torch.nested._internal.nested_tensor.NestedTensor)
  1897. ):
  1898. unimplemented_v2(
  1899. gb_type="Attempted to wrap strided NestedTensor",
  1900. context="",
  1901. explanation="torch.compile does not support strided NestedTensor",
  1902. hints=[],
  1903. )
  1904. # TODO(pearu,sparse-team) - Add the corresponding SPARSE_TENSOR_MATCH guards
  1905. if (
  1906. isinstance(value, torch.Tensor)
  1907. and is_sparse_any(value)
  1908. and (not self.tx.export or not config.capture_sparse_compute)
  1909. ):
  1910. # A hot fix for sparse tensors + torch.compile. Support for
  1911. # export + sparsity is being added but we need to create
  1912. # SPARSE_TENSOR_GUARDS for guards to work properly.
  1913. unimplemented_v2(
  1914. gb_type="Attempted to wrap sparse Tensor",
  1915. context="",
  1916. explanation="torch.compile does not support sparse Tensors",
  1917. hints=[*graph_break_hints.SUPPORTABLE],
  1918. )
  1919. if (
  1920. safe_has_grad(value)
  1921. and safe_grad(value) is not None
  1922. and value.dtype != safe_grad(value).dtype
  1923. ):
  1924. unimplemented_v2(
  1925. gb_type="dtype mismatch between tensor and its gradient",
  1926. context=f"tensor dtype: {value.dtype}; grad dtype: {safe_grad(value).dtype}",
  1927. explanation="Inconsistent dtype between tensor and its gradient. "
  1928. "This can happen in FSDP and crashes meta tensor creation.",
  1929. hints=[*graph_break_hints.SUPPORTABLE],
  1930. )
  1931. # tx.output has multiple tracers if we're introspecting HigherOrderOperator.
  1932. # When we've discovered an untracked tensor, then we actually need
  1933. # to get Dynamo to track the tensor (which is what this function does)
  1934. # and put it as a graph input on the root tracer. Later on,
  1935. # if the input is actually used in the body of the HigherOrderOperator,
  1936. # then the relevant SubgraphTracer will lift it to being an input of
  1937. # the subgraph.
  1938. # See NOTE [HigherOrderOperator tracing design] for more details.
  1939. example_value = wrap_to_fake_tensor_and_record(
  1940. value, tx=self.tx, is_tensor=True, source=source
  1941. )
  1942. tensor_proxy = self.tx.output.root_tracer.create_graph_input(
  1943. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  1944. type(value),
  1945. example_value,
  1946. source=source,
  1947. )
  1948. cache_real_value_when_export(self.tx, tensor_proxy, value)
  1949. tensor_variable = wrap_fx_proxy(
  1950. tx=self.tx,
  1951. proxy=tensor_proxy,
  1952. example_value=example_value,
  1953. subclass_type=subclass_type,
  1954. source=source,
  1955. **options,
  1956. )
  1957. if value._is_view():
  1958. # If value is a view, add its base tensor to the tracked fakes list.
  1959. # This is so we are able to access the correct source for its symbolic
  1960. # shape values, in case we need them.
  1961. wrap_to_fake_tensor_and_record(
  1962. value._base,
  1963. tx=self.tx,
  1964. source=AttrSource(source, "_base"),
  1965. is_tensor=True,
  1966. )
  1967. guard_type = GuardBuilder.TENSOR_MATCH
  1968. if isinstance(source, GradSource) and is_from_optimizer_source(source):
  1969. guard_type = GuardBuilder.NOT_NONE_MATCH
  1970. self.install_guards(
  1971. functools.partial(
  1972. guard_type,
  1973. value=(
  1974. value
  1975. if isinstance(source, NumpyTensorSource)
  1976. else TensorWeakRef(value)
  1977. ),
  1978. )
  1979. )
  1980. # We install TYPE_MATCH guards for traceable wrapper subclass object,
  1981. # and recursively install corresponding guard for each inner attribute.
  1982. if is_traceable_wrapper_subclass(value):
  1983. self.install_guards(GuardBuilder.TENSOR_SUBCLASS_METADATA_MATCH)
  1984. self.install_guards(GuardBuilder.TYPE_MATCH)
  1985. install_guard(
  1986. SubclassAttrListSource(source).make_guard(GuardBuilder.EQUALS_MATCH)
  1987. )
  1988. attrs, _ = value.__tensor_flatten__()
  1989. for attr in attrs:
  1990. inner_value = getattr(value, attr)
  1991. inner_source = AttrSource(self.source, attr)
  1992. LazyVariableTracker.realize_all(
  1993. VariableBuilder(self.tx, inner_source)(inner_value)
  1994. )
  1995. self.tx.output.input_source_to_var[source] = tensor_variable
  1996. assert "tensor_dict" not in tensor_proxy.node.meta
  1997. tensor_proxy.node.meta["tensor_dict"] = _extract_tensor_dict(value)
  1998. # Note: this information is conveyed via subclass_type now
  1999. fake_tensor_value = tensor_variable.proxy.node.meta["example_value"]
  2000. if maybe_get_fake_mode(fake_tensor_value) is not self.tx.fake_mode:
  2001. raise InternalTorchDynamoError("Wrapped Tensor must be this graph's fake")
  2002. grapharg = GraphArg(source, value, False, fake_tensor_value)
  2003. tensor_proxy.node.meta["grapharg"] = grapharg
  2004. return tensor_variable
  2005. def wrap_numpy_ndarray(self, value):
  2006. assert np is not None
  2007. assert isinstance(value, np.ndarray)
  2008. source = NumpyTensorSource(self.get_source())
  2009. from torch._numpy import _util
  2010. readonly = not value.flags.writeable
  2011. if readonly:
  2012. try:
  2013. value.flags.writeable = True
  2014. except ValueError:
  2015. # One can not easily make nditer elements writable,
  2016. # but warning is not the end of the world
  2017. assert isinstance(value.base, np.nditer)
  2018. with torch_function_mode_stack_state_mgr.temp_restore_stack():
  2019. try:
  2020. tensor_value = _util._try_convert_to_tensor(value)
  2021. if readonly:
  2022. from torch._prims_common import clone_preserve_strides
  2023. tensor_value = clone_preserve_strides(tensor_value)
  2024. except NotImplementedError as e:
  2025. # failed to convert to tensor, graph break
  2026. unimplemented_v2(
  2027. gb_type="failed to convert numpy.ndarray to Tensor",
  2028. context=str(value),
  2029. explanation="Exception encountered when attempting to convert numpy.ndarray to Tensor",
  2030. hints=[],
  2031. from_exc=e,
  2032. )
  2033. # We do this because we want the full behavior of guarding the numpy ndarray as if it were
  2034. # a tensor. It's a little annoying to make a VT to throw out, but there's so many side effects here
  2035. # that there's not another great way to do this atm.
  2036. # This creates the right graphargs, as well as registration for guards in tensor names and shape env.
  2037. LazyVariableTracker.realize_all(VariableBuilder(self.tx, source)(tensor_value))
  2038. example_value = wrap_to_fake_tensor_and_record(
  2039. tensor_value,
  2040. tx=self.tx,
  2041. is_tensor=False,
  2042. source=source,
  2043. )
  2044. proxy = self.tx.output.root_tracer.create_graph_input(
  2045. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2046. type(tensor_value),
  2047. example_value,
  2048. source=source,
  2049. )
  2050. cache_real_value_when_export(self.tx, proxy, tensor_value)
  2051. options = {"source": source}
  2052. numpy_ndarray_variable = wrap_fx_proxy_cls(
  2053. target_cls=NumpyNdarrayVariable,
  2054. tx=self.tx,
  2055. proxy=proxy,
  2056. example_value=example_value,
  2057. **options,
  2058. )
  2059. self.tx.output.input_source_to_var[source] = numpy_ndarray_variable
  2060. example_value = numpy_ndarray_variable.proxy.node.meta["example_value"]
  2061. # pass_arg_as_tensor should be true because we are wrapping a np.ndarray as argument input, and it needs to be
  2062. # converted to a tensor.
  2063. grapharg = GraphArg(
  2064. source,
  2065. tensor_value,
  2066. pass_arg_as_tensor=True,
  2067. fake_tensor=example_value,
  2068. is_tensor=True,
  2069. example_strong_ref=tensor_value,
  2070. )
  2071. proxy.node.meta["grapharg"] = grapharg
  2072. # TODO - Why do we need to set the source of the np ndarray vt back to
  2073. # original source. Many tests fails.
  2074. numpy_ndarray_variable.source = self.source
  2075. return numpy_ndarray_variable
  2076. def wrap_symint(
  2077. self,
  2078. value,
  2079. dynamism: Optional[DimDynamic] = None,
  2080. context: Optional[SymIntSymbolicContext] = None,
  2081. ):
  2082. assert type(value) is int
  2083. if self.name in self.tx.output.unspec_variable_map:
  2084. return self.tx.output.unspec_variable_map[self.name]
  2085. shape_env = self.tx.output.shape_env
  2086. if TracingContext.get().force_unspec_int_unbacked_size_like:
  2087. wrapped_value = shape_env.create_unbacked_symint()
  2088. _constrain_range_for_size(wrapped_value)
  2089. self.tx.output.tracked_fakes.append(
  2090. TrackedFake(wrapped_value, self.source, None)
  2091. )
  2092. # NB: We do not do float. For motivation, see
  2093. # https://docs.google.com/document/d/1INSCdYu1PxXcr43HrD82OudeEuS-qxQe1yZmLg2wy6A/edit
  2094. # but the general idea is that we generate kernels that can
  2095. # take unspecialized floats and use them in sizevar computation
  2096. elif not is_constant_source(self.get_source()):
  2097. if dynamism is None and torch._dynamo.config.specialize_int:
  2098. # If specialize_int is False, also return
  2099. # a constant (but this should have been handled
  2100. # in the caller, TBH). But if `dynamism` is set, then actually
  2101. # turn it into a symint
  2102. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2103. return ConstantVariable.create(value=value, source=self.source)
  2104. name = self.source.name()
  2105. frame_state_entry = process_automatic_dynamic(
  2106. self.tx,
  2107. name,
  2108. FrameStateSizeEntry.make_scalar(value),
  2109. is_unspecialized_nn_module=self.source.guard_source().is_unspecialized_nn_module(),
  2110. )
  2111. # TODO: This should be dynamic, as we in general do not
  2112. # know if bare integers are actually going to be sizevars
  2113. # and it is inappropriate to eagerly duck size them with
  2114. # real sizevars
  2115. normalized_source_name = normalize_source_name(self.source.name())
  2116. base_source = self.source
  2117. if isinstance(base_source, ChainedSource):
  2118. base_source = base_source.get_base()
  2119. if dynamism is not None:
  2120. dynamic_dim = dynamism
  2121. elif (
  2122. config.automatic_dynamic_shapes
  2123. and frame_state_entry.scalar is auto_dynamic
  2124. ):
  2125. set_feature_use("dynamo.automatic_dynamic_shapes", True)
  2126. dynamic_dim = get_automatic_dynamic_shapes_mark_as()
  2127. elif (
  2128. isinstance(base_source, LocalSource)
  2129. and base_source.dynamism is not None
  2130. and dict(base_source.dynamism).get(normalized_source_name, {0: False})[
  2131. 0
  2132. ]
  2133. ) or not config.assume_static_by_default:
  2134. dynamic_dim = DimDynamic.DYNAMIC
  2135. else: # assume_static_by_default
  2136. # TODO: dynamic_dim = DimDynamic.STATIC should work but
  2137. # for some reason it doesn't
  2138. if frame_state_entry.scalar is auto_dynamic:
  2139. set_feature_use("dynamo.automatic_dynamic_shapes", False)
  2140. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2141. return ConstantVariable.create(value=value)
  2142. wrapped_value = shape_env.create_unspecified_symint_and_symbol(
  2143. value,
  2144. source=self.source,
  2145. dynamic_dim=dynamic_dim,
  2146. )
  2147. self.tx.output.tracked_fakes.append(
  2148. TrackedFake(wrapped_value, self.source, context)
  2149. )
  2150. else:
  2151. assert is_constant_source(self.get_source())
  2152. # TODO: Do I actually need guard for constant source?
  2153. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2154. return ConstantVariable.create(value=value, source=self.source)
  2155. assert not isinstance(self.get_source(), RandomValueSource)
  2156. install_guard(self.get_source().make_guard(GuardBuilder.TYPE_MATCH))
  2157. options = {"source": self.get_source()}
  2158. proxy = self.tx.output.root_tracer.create_graph_input(
  2159. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2160. type(wrapped_value),
  2161. wrapped_value,
  2162. source=self.get_source(),
  2163. )
  2164. sym_expr = wrapped_value.node.expr
  2165. assert isinstance(sym_expr, sympy.Symbol), f"{sym_expr} is not a basic Symbol."
  2166. self.tx.output.root_tracer.bound_symbols[sym_expr] = proxy
  2167. unspec_var = SymNodeVariable(proxy, wrapped_value, **options)
  2168. self.tx.output.unspec_variable_map[self.name] = unspec_var
  2169. if not is_constant_source(self.get_source()):
  2170. proxy.node.meta["grapharg"] = GraphArg(
  2171. self.get_source(),
  2172. wrapped_value,
  2173. pass_arg_as_tensor=False,
  2174. fake_tensor=None,
  2175. is_tensor=False,
  2176. example_strong_ref=wrapped_value,
  2177. )
  2178. return unspec_var
  2179. def wrap_symfloat(self, value):
  2180. # SymFloat wrapping is special. We first wrap it in the same way we
  2181. # do an unspecialized primitive, and then we item() it into a
  2182. # SymFloat. Removal of the item() call is left to a later FX pass,
  2183. # mostly because that pass is more easily done after we have lowered
  2184. # to ATen ops. (Dynamo doesn't do decomposition right now).
  2185. if self.name in self.tx.output.unspec_variable_map:
  2186. return self.tx.output.unspec_variable_map[self.name]
  2187. frame_state_entry = process_automatic_dynamic(
  2188. self.tx,
  2189. self.source.name(),
  2190. FrameStateSizeEntry.make_scalar(value),
  2191. is_unspecialized_nn_module=self.source.guard_source().is_unspecialized_nn_module(),
  2192. )
  2193. # NB: we specialize on nan input, because our guard modeling in
  2194. # ShapeEnv cannot deal with nan
  2195. if (
  2196. torch._dynamo.config.specialize_float
  2197. or is_constant_source(self.get_source())
  2198. or math.isnan(value)
  2199. or math.isinf(value)
  2200. # We don't support cudagraphs for now. Without this cudagraphs
  2201. # break because they expect all cuda inputs but our tensorified
  2202. # float will be a f64[] cpu tensor. Fixes the following test
  2203. # when specialize_float=False
  2204. # python test/inductor/test_compiled_optimizers.py CompiledOptimizerTests.test_rmsprop_weight_decay_maximize_capturable_cuda # noqa: B950
  2205. or torch._inductor.config.triton.cudagraphs
  2206. or justknobs_check("pytorch/compiler:unspecialize_float_killswitch", False)
  2207. or (
  2208. config.assume_static_by_default
  2209. and frame_state_entry.scalar is not auto_dynamic
  2210. )
  2211. ):
  2212. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2213. return ConstantVariable.create(value=value, source=self.source)
  2214. # NB: At the point we've gotten here, we don't assume static by
  2215. # default. Since we have a guard mechanism, there isn't really any
  2216. # downside to trying to be dynamic for float all the time. Unlike
  2217. # ints, this won't make codegen perf worse. Modest cost to compile
  2218. # time.
  2219. wrapped_value = torch.tensor(value, dtype=torch.float64)
  2220. # We don't support specializing floats for grad checking tensors
  2221. # See https://github.com/pytorch/pytorch/pull/140828 for more
  2222. # context.
  2223. if torch._C._functorch.is_gradtrackingtensor(wrapped_value):
  2224. self.install_guards(GuardBuilder.CONSTANT_MATCH)
  2225. return ConstantVariable.create(value=value, source=self.source)
  2226. # TODO: Switch RandomValueSource over to use this, this is more
  2227. # accurate
  2228. assert not isinstance(self.get_source(), RandomValueSource)
  2229. install_guard(self.get_source().make_guard(GuardBuilder.TYPE_MATCH))
  2230. # The FloatTensorSource here is just for pedantic correctness: if you
  2231. # guard against an UnspecializedPythonVariable, you need to guard
  2232. # against the tensor-ified version of the local, otherwise it's not a
  2233. # Tensor. However, we never let the UnspecializedPythonVariable escape
  2234. # here, so there should never actually be any guards against this
  2235. # source.
  2236. source = FloatTensorSource(self.get_source())
  2237. options = {"source": source, "raw_value": value}
  2238. # TODO: Maybe the tensor-ification should be built into the source,
  2239. # rather than by special pattern match
  2240. example_value = wrap_to_fake_tensor_and_record(
  2241. wrapped_value, tx=self.tx, is_tensor=False, source=source
  2242. )
  2243. proxy = self.tx.output.root_tracer.create_graph_input(
  2244. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2245. type(wrapped_value),
  2246. example_value,
  2247. source=source,
  2248. )
  2249. cache_real_value_when_export(self.tx, proxy, wrapped_value)
  2250. unspec_var = wrap_fx_proxy_cls(
  2251. UnspecializedPythonVariable,
  2252. tx=self.tx,
  2253. proxy=proxy,
  2254. example_value=example_value,
  2255. **options,
  2256. )
  2257. assert isinstance(unspec_var, UnspecializedPythonVariable)
  2258. self.tx.output.unspec_variable_map[self.name] = unspec_var
  2259. if self.tx.export and not isinstance(self.get_source(), LocalSource):
  2260. raise AssertionError(
  2261. f"Dynamo attempts to add additional input during export: value={wrapped_value}, source={self.get_source()}"
  2262. )
  2263. fake_tensor_value = None
  2264. example_value = unspec_var.proxy.node.meta["example_value"]
  2265. assert is_fake(example_value)
  2266. fake_tensor_value = example_value
  2267. assert fake_tensor_value.fake_mode is self.tx.fake_mode, (
  2268. f"fake mode ({fake_tensor_value.fake_mode}) from fake tensor metadata doesn't match mode"
  2269. "({self.tx.fake_mode}) from InstructionTranslator"
  2270. )
  2271. # There's something a bit incoherent about pass_arg_as_tensor,
  2272. # specifically regarding sources.
  2273. #
  2274. # Specifically, suppose we have "x: float" local argument. We
  2275. # eventually end up with an UnspecializedPythonVariable denoting
  2276. # torch.as_tensor(x)... but it's source is still L['x'] (which if you
  2277. # accessed it directly is a float!) So you gotta be careful when
  2278. # setting up your guards, because it's still going to be a float at
  2279. # this point, the conversion happens only precisely at the point we're
  2280. # actually calling the FX graph. This happens to be what we want for
  2281. # shape guard generation, but it's kind of unintuitive.
  2282. proxy.node.meta["grapharg"] = GraphArg(
  2283. self.get_source(),
  2284. wrapped_value,
  2285. pass_arg_as_tensor=True,
  2286. fake_tensor=fake_tensor_value,
  2287. is_tensor=False,
  2288. example_strong_ref=wrapped_value,
  2289. )
  2290. # Directly do item to bypass capture_scalar_outputs
  2291. r = wrap_fx_proxy(
  2292. self.tx,
  2293. self.tx.output.create_proxy(
  2294. "call_method",
  2295. "item",
  2296. *proxy_args_kwargs([unspec_var], {}),
  2297. ),
  2298. )
  2299. self.tx.output.tracked_fakes.append(TrackedFake(r.sym_num, self.source, None))
  2300. get_metrics_context().set("tensorify_float_attempt", True, overwrite=True)
  2301. return r
  2302. def wrap_unspecialized_primitive(self, value):
  2303. if self.name in self.tx.output.unspec_variable_map:
  2304. return self.tx.output.unspec_variable_map[self.name]
  2305. wrapped_value = torch.tensor(value)
  2306. if not isinstance(self.get_source(), RandomValueSource):
  2307. install_guard(self.get_source().make_guard(GuardBuilder.TYPE_MATCH))
  2308. options = {"source": self.get_source()}
  2309. options.update({"raw_value": value})
  2310. example_value = wrap_to_fake_tensor_and_record(
  2311. wrapped_value, tx=self.tx, is_tensor=False, source=self.get_source()
  2312. )
  2313. proxy = self.tx.output.root_tracer.create_graph_input(
  2314. re.sub(r"[^a-zA-Z0-9]+", "_", self.name),
  2315. type(wrapped_value),
  2316. example_value,
  2317. source=self.get_source(),
  2318. )
  2319. cache_real_value_when_export(self.tx, proxy, wrapped_value)
  2320. unspec_var = wrap_fx_proxy_cls(
  2321. UnspecializedPythonVariable,
  2322. tx=self.tx,
  2323. proxy=proxy,
  2324. example_value=example_value,
  2325. **options,
  2326. )
  2327. self.tx.output.unspec_variable_map[self.name] = unspec_var
  2328. if not is_constant_source(self.get_source()):
  2329. if self.tx.export and not isinstance(self.get_source(), LocalSource):
  2330. raise AssertionError(
  2331. f"Dynamo attempts to add additional input during export: value={wrapped_value}, source={self.get_source()}"
  2332. )
  2333. fake_tensor_value = None
  2334. if isinstance(unspec_var, ConstantVariable):
  2335. # TODO: when can this happen?
  2336. example_value = unspec_var.value
  2337. else:
  2338. example_value = unspec_var.proxy.node.meta["example_value"]
  2339. assert is_fake(example_value)
  2340. fake_tensor_value = example_value
  2341. assert fake_tensor_value.fake_mode is self.tx.fake_mode, (
  2342. f"fake mode ({fake_tensor_value.fake_mode}) from fake tensor metadata doesn't match mode"
  2343. "({self.tx.fake_mode}) from InstructionTranslator"
  2344. )
  2345. proxy.node.meta["grapharg"] = GraphArg(
  2346. self.get_source(),
  2347. wrapped_value,
  2348. pass_arg_as_tensor=True,
  2349. fake_tensor=fake_tensor_value,
  2350. is_tensor=False,
  2351. example_strong_ref=wrapped_value,
  2352. )
  2353. return unspec_var
  2354. def _dataclasses_fields_lambda(obj):
  2355. if isinstance(obj, UserDefinedObjectVariable):
  2356. value = obj.value
  2357. else:
  2358. unimplemented_v2(
  2359. gb_type="dataclass fields failure",
  2360. context=f"obj: {obj}; variable type: {type(obj)}",
  2361. explanation=f"Dataclass fields handling fails for {obj}. Expected it to be a user-defined object.",
  2362. hints=[],
  2363. )
  2364. items = []
  2365. for field in dataclasses.fields(value):
  2366. source = None
  2367. if obj.source:
  2368. base_src = AttrSource(obj.source, "__dataclass_fields__")
  2369. source = DictGetItemSource(base_src, field.name)
  2370. items.append(UserDefinedObjectVariable(field, source=source))
  2371. return TupleVariable(items)
  2372. def _clone_input(value, fake_mode):
  2373. if isinstance(value, torch.Tensor):
  2374. # tensor subclasses will not be converted to FakeTensors and need to be cloned
  2375. if not (
  2376. isinstance(value, FakeTensor)
  2377. or (
  2378. # Is functional tensor fakeified by this instance of Dynamo
  2379. torch._is_functional_tensor(value)
  2380. and maybe_get_fake_mode(value) is fake_mode
  2381. )
  2382. or value.is_nested
  2383. ):
  2384. # NB: ensure strides are preserved
  2385. value = clone_input(value)
  2386. return value
  2387. def wrap_fx_proxy(
  2388. tx, proxy, example_value=None, subclass_type=None, **options
  2389. ) -> VariableTracker:
  2390. kwargs = {
  2391. "tx": tx,
  2392. "proxy": proxy,
  2393. "example_value": example_value,
  2394. "subclass_type": subclass_type,
  2395. **options,
  2396. }
  2397. if subclass_type is None:
  2398. return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
  2399. else:
  2400. result = wrap_fx_proxy_cls(target_cls=TensorWithTFOverrideVariable, **kwargs)
  2401. result.install_global(tx)
  2402. return result
  2403. def cache_real_value_when_export(tx, proxy, example_value):
  2404. if tx.export:
  2405. # The legacy behavior for real value cache with subclasses was
  2406. # to perform a clone WITHOUT preserving the subclass. It's
  2407. # not entirely clear this is what you actually want though.
  2408. with torch._C.DisableTorchFunctionSubclass():
  2409. proxy.tracer.real_value_cache[proxy.node] = _clone_input(
  2410. example_value, tx.fake_mode
  2411. )
  2412. # Note: Unfortunate split due to some gross classes existing that subclass TensorVariable
  2413. # Should be compositional instead
  2414. #
  2415. # This is a horribly complicated function that does too many things, to
  2416. # explain what it does, let's first talk about the classic usage wrap_fx_proxy
  2417. # for a TensorVariable. There are two primary modes of use:
  2418. #
  2419. # 1. Wrapping a pre-existing Tensor. In this case, example_value is set
  2420. # to the pre-existing Tensor. (Note that this example_value will NOT
  2421. # be the final example_value we put into node.meta['example_value'],
  2422. # instead it is converted into a fake tensor using
  2423. # wrap_to_fake_tensor_and_record and registered as a graph input.)
  2424. #
  2425. # 2. "Wrapping" the result of some Tensor operation Dynamo traced over. In
  2426. # this case, example_value is None (and we are going to figure it out
  2427. # ourselves using FakeTensors, via get_fake_value, which will run
  2428. # the operation represented by the (singular!) FX node referenced by
  2429. # the passed in proxy.)
  2430. #
  2431. # The expectation is you end up with a Tensor output, and everything is
  2432. # straightforwardly traced into the graph.
  2433. #
  2434. # In all cases, the returned `TensorVariable` subclass will have an `example_value`
  2435. # and that `example_value` must be a `FakeTensor` produced by the currently running
  2436. # instance of Dynamo.
  2437. #
  2438. # Upon closer inspection, you may notice that there are a slurry of non-Tensor
  2439. # output cases in handle_traced_output. What gives? Well, we sometimes trace operations into the
  2440. # graph that don't involve tensors.
  2441. #
  2442. # * Some operators return tuples; we need to recursively handle their
  2443. # contents
  2444. #
  2445. # * Some operators have side effects that will affect subsequent AOTAutograd
  2446. # tracing but don't otherwise return anything.
  2447. #
  2448. # * Some operators return symbolic ints/floats/bools which can go in the
  2449. # graph and be traced (but only if they're actually symbolic! If they're
  2450. # static you don't want to put them in the graph, which means you
  2451. # shouldn't call this function.)
  2452. #
  2453. # The common theme is that you only use this function WHEN YOU ARE TRACING
  2454. # SOMETHING INTO THE GRAPH. This is sort of obvious, because you can't call
  2455. # this function without a proxy.
  2456. def wrap_fx_proxy_cls(
  2457. target_cls, tx, proxy, example_value=None, subclass_type=None, **options
  2458. ):
  2459. if example_value is None:
  2460. return _wrap_fx_proxy(
  2461. target_cls, tx, proxy, example_value, subclass_type, **options
  2462. )
  2463. elif isinstance(example_value, torch.Tensor):
  2464. return _wrap_fx_preexisting_tensor(
  2465. target_cls, tx, proxy, example_value, subclass_type, **options
  2466. )
  2467. else:
  2468. # This will skip tracing an op and recursively reinvoke wrap_fx_proxy_cls on supported
  2469. # data structures. In essence this just handles tracing some other value which may
  2470. # contain Fake Tensors or is otherwise proxyable.
  2471. return handle_traced_output(
  2472. example_value, tx, proxy, options, subclass_type, target_cls
  2473. )
  2474. # This is 1 above (wrapping a preexisting tensor)
  2475. def _wrap_fx_preexisting_tensor(
  2476. target_cls, tx, proxy, tensor, subclass_type=None, **options
  2477. ):
  2478. from ..symbolic_convert import InstructionTranslatorBase
  2479. assert isinstance(tensor, torch.Tensor), (
  2480. f"_wrap_fx_preexisting_tensor expected tensor, got {type(tensor)}"
  2481. )
  2482. assert isinstance(tx, InstructionTranslatorBase)
  2483. if "guards" in options and options["guards"] is not None:
  2484. tx.output.guards.update(options["guards"])
  2485. # Placeholders always carry example_value in node.meta.
  2486. # non-placeholders always have no example_value in node.meta
  2487. if proxy.node.op == "placeholder":
  2488. assert "example_value" in proxy.node.meta, (
  2489. f"placeholder {proxy} doesn't have 'example_value' in node.meta"
  2490. )
  2491. else:
  2492. assert "example_value" not in proxy.node.meta, (
  2493. f"{proxy.node.meta['example_value']}"
  2494. )
  2495. # See NOTE: [Deferring tensor pack/unpack hooks until runtime]
  2496. with torch._dynamo.utils._disable_saved_tensors_hooks_during_tracing():
  2497. # Handle recursive calls here
  2498. if maybe_get_fake_mode(tensor) is tx.fake_mode:
  2499. pass
  2500. else:
  2501. cache_real_value_when_export(tx, proxy, tensor)
  2502. if tx.export:
  2503. # The legacy behavior for real value cache with subclasses was
  2504. # to perform a clone WITHOUT preserving the subclass. It's
  2505. # not entirely clear this is what you actually want though.
  2506. with torch._C.DisableTorchFunctionSubclass():
  2507. proxy.tracer.real_value_cache[proxy.node] = _clone_input(
  2508. tensor, tx.fake_mode
  2509. )
  2510. # NB: If we're ignoring subclass, then the expectation is you will
  2511. # take the returned TensorVariable and wrap it into a more
  2512. # accurate TensorVariable that is able to track subclass-ness;
  2513. # otherwise this is wrong!
  2514. kwargs = {
  2515. "is_tensor": target_cls
  2516. in (TensorVariable, TensorWithTFOverrideVariable),
  2517. }
  2518. assert "source" in options and options["source"] is not None
  2519. kwargs["source"] = options["source"]
  2520. tensor = wrap_to_fake_tensor_and_record(tensor, tx=tx, **kwargs)
  2521. if tensor.device.type != "meta" and (
  2522. maybe_get_fake_mode(tensor) is not tx.fake_mode
  2523. ):
  2524. raise InternalTorchDynamoError(
  2525. "`tensor` needs to be a `FakeTensor`"
  2526. f"wrapped by this instance of Dynamo. Found: {tensor}"
  2527. )
  2528. return construct_tensor_variable(
  2529. target_cls, tx, proxy, tensor, subclass_type, options
  2530. )
  2531. # This is 2 in the above comment (wrapping the output of a traced op)
  2532. def _wrap_fx_proxy(
  2533. target_cls, tx, proxy, example_value=None, subclass_type=None, **options
  2534. ):
  2535. from ..symbolic_convert import InstructionTranslatorBase
  2536. assert isinstance(tx, InstructionTranslatorBase)
  2537. if "guards" in options and options["guards"] is not None:
  2538. tx.output.guards.update(options["guards"])
  2539. assert "example_value" not in proxy.node.meta, f"{proxy.node.meta['example_value']}"
  2540. # See NOTE: [Deferring tensor pack/unpack hooks until runtime]
  2541. with torch._dynamo.utils._disable_saved_tensors_hooks_during_tracing():
  2542. # with preserve_rng_state():
  2543. # only allow_non_graph_fake in this instance because we handle the non-fake
  2544. # cases properly below.
  2545. example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
  2546. return handle_traced_output(
  2547. example_value, tx, proxy, options, subclass_type, target_cls
  2548. )
  2549. # This handles wrapping of the output of an op traced into the graph
  2550. def handle_traced_output(example_value, tx, proxy, options, subclass_type, target_cls):
  2551. import torch._functorch.vmap
  2552. import torch._subclasses.fake_tensor
  2553. import torch._utils
  2554. if isinstance(example_value, torch.Tensor):
  2555. var = construct_tensor_variable(
  2556. target_cls, tx, proxy, example_value, subclass_type, options
  2557. )
  2558. # NOTE: [Side effect tracking for newly constructed tensor]
  2559. # For newly constructed objects that have mutable attributes, we usually
  2560. # construct their VariableTracker via `track_object_new`, but since
  2561. # tensor variable construction is a bit different, we handle them
  2562. # specially here. This ensures that codegen will actually generate the
  2563. # attribute mutations on this tensor.
  2564. #
  2565. # NOTE we pass a dummy object as the `item` argument to avoid
  2566. # constructing a dummy _tensor_ object. The object isn't used for
  2567. # newly constructed VTs anyways.
  2568. tx.output.side_effects._track_obj(
  2569. proxy, var, mutation_type_cls=AttributeMutationNew
  2570. )
  2571. return var
  2572. elif (
  2573. hasattr(proxy.node.target, "__name__")
  2574. and proxy.node.target.__name__ == "set_state"
  2575. and isinstance(proxy.node.target.__self__, torch._C.Generator)
  2576. or proxy.node.target == torch.random.set_rng_state
  2577. ):
  2578. return TorchInGraphFunctionVariable(proxy.node.target)
  2579. elif (
  2580. proxy.node.target == torch._C._DisableFuncTorch
  2581. or proxy.node.target == torch.cuda._is_in_bad_fork
  2582. ):
  2583. return UserDefinedObjectVariable(example_value)
  2584. elif istype(example_value, torch.Size) and all(
  2585. isinstance(x, int) for x in example_value
  2586. ):
  2587. sizes = [ConstantVariable.create(x) for x in example_value]
  2588. return SizeVariable(sizes, **options)
  2589. elif isinstance(example_value, (tuple, list)):
  2590. set_example_value(proxy.node, example_value)
  2591. unpacked = []
  2592. for i, val in enumerate(example_value):
  2593. if val is None:
  2594. # nn.MultiheadAttention() can return None, see issue #175
  2595. unpacked.append(
  2596. ConstantVariable.create(None, **options),
  2597. )
  2598. else:
  2599. proxy_i = proxy.tracer.create_proxy(
  2600. kind="call_function",
  2601. target=operator.getitem,
  2602. args=(proxy, i),
  2603. kwargs={},
  2604. )
  2605. if "source" in options:
  2606. # This path should only trigger for list stealing, so it's
  2607. # safe to use `GetItemSource`.
  2608. assert isinstance(example_value, list)
  2609. source = options["source"]
  2610. options_i = options.copy()
  2611. options_i["source"] = GetItemSource(
  2612. base=source, index=i, index_is_slice=False
  2613. )
  2614. else:
  2615. # use the same options object as parent
  2616. options_i = options
  2617. # WARNING: this assumes the same target_cls as this tuple/list call
  2618. unpacked.append(
  2619. wrap_fx_proxy_cls(
  2620. target_cls=target_cls,
  2621. tx=tx,
  2622. proxy=proxy_i,
  2623. example_value=val,
  2624. **options_i,
  2625. )
  2626. )
  2627. if isinstance(example_value, torch.Size):
  2628. # NB: Keep the old proxy around. See SizeVariable for an
  2629. # explanation why
  2630. return SizeVariable(unpacked, proxy, **options)
  2631. elif istype(example_value, tuple):
  2632. return TupleVariable(unpacked, **options)
  2633. elif istype(example_value, (list, immutable_list)):
  2634. return ListVariable(unpacked, **options)
  2635. else:
  2636. assert (
  2637. example_value.__class__.__module__ == "torch.return_types"
  2638. or hasattr(example_value, "_fields")
  2639. ), (
  2640. f"expected {example_value.__class__.__module__} == torch.return_types or named tuple but got {type(example_value)}"
  2641. )
  2642. return NamedTupleVariable(unpacked, example_value.__class__, **options)
  2643. elif example_value is None or proxy.node.target is torch.manual_seed:
  2644. return ConstantVariable.create(None, **options)
  2645. elif isinstance(example_value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
  2646. tx.output.current_tracer.track_produced_symints(example_value, proxy)
  2647. set_example_value(proxy.node, example_value)
  2648. return SymNodeVariable(proxy, example_value, **options)
  2649. elif (
  2650. inspect.isclass(proxy.node.target)
  2651. and issubclass(proxy.node.target, torch.Stream)
  2652. ) or proxy.node.target in [
  2653. device_interface.current_stream
  2654. for _, device_interface in get_registered_device_interfaces()
  2655. ]:
  2656. set_example_value(proxy.node, example_value)
  2657. return StreamVariable(proxy, example_value, example_value.device, **options)
  2658. elif (
  2659. inspect.isclass(proxy.node.target)
  2660. and issubclass(proxy.node.target, torch.Event)
  2661. ) or proxy.node.target in [
  2662. device_interface.Event
  2663. for _, device_interface in get_registered_device_interfaces()
  2664. ]:
  2665. set_example_value(proxy.node, example_value)
  2666. return EventVariable(proxy, example_value, **options)
  2667. elif proxy.node.target == "query" and proxy.node.op == "call_method":
  2668. set_example_value(proxy.node, example_value)
  2669. return ConstantVariable(example_value, **options)
  2670. elif (
  2671. example_value is not None
  2672. and isinstance(example_value, torch.Event)
  2673. and proxy.node.target == "record_event"
  2674. and proxy.node.op == "call_method"
  2675. ):
  2676. set_example_value(proxy.node, example_value)
  2677. return EventVariable(proxy, example_value, **options)
  2678. elif isinstance(example_value, int) and (
  2679. proxy.node.target
  2680. in [
  2681. torch.sym_int,
  2682. getattr,
  2683. operator.getitem,
  2684. torch._utils._element_size,
  2685. torch.seed,
  2686. operator.mod,
  2687. torch._functorch.vmap._validate_and_get_batch_size,
  2688. torch._functorch.predispatch._vmap_increment_nesting,
  2689. torch._functorch.predispatch._vmap_decrement_nesting,
  2690. # some mac builds are missing torch.distributed.get_rank()
  2691. getattr(torch.distributed, "get_rank", _missing),
  2692. getattr(torch.distributed, "get_world_size", _missing),
  2693. # This always wants to be in the graph, even if the constraint
  2694. # results in a constant int
  2695. torch._constrain_as_size,
  2696. ]
  2697. or (
  2698. # TODO: this is a little sus, because we didn't check what the self is
  2699. proxy.node.op == "call_method" and proxy.node.target in ["bit_length"]
  2700. )
  2701. ):
  2702. set_example_value(proxy.node, example_value)
  2703. return ConstantVariable.create(example_value, **options)
  2704. elif isinstance(example_value, torch.backends.cuda.SDPAParams):
  2705. from .sdpa import SDPAParamsVariable
  2706. set_example_value(proxy.node, example_value)
  2707. return SDPAParamsVariable(proxy, **options)
  2708. elif isinstance(example_value, bool) and (
  2709. proxy.node.target
  2710. in [
  2711. torch._C._are_functorch_transforms_active,
  2712. torch._C._functorch.is_batchedtensor,
  2713. torch.backends.cuda.is_flash_attention_available,
  2714. torch.backends.cuda.can_use_flash_attention,
  2715. torch.backends.cuda.can_use_efficient_attention,
  2716. "is_integer",
  2717. ]
  2718. + list(supported_const_comparison_op_values.keys())
  2719. ):
  2720. set_example_value(proxy.node, example_value)
  2721. return ConstantVariable.create(example_value, **options)
  2722. elif isinstance(example_value, (int, float, bool)) and (
  2723. proxy.node.target is call_torchbind
  2724. or proxy.node.target is flat_apply
  2725. or (proxy.node.op == "call_method" and proxy.node.target == "item")
  2726. ):
  2727. set_example_value(proxy.node, example_value)
  2728. return ConstantVariable.create(example_value, **options)
  2729. elif isinstance(example_value, float) or proxy.node.target in ["hex", "__round__"]:
  2730. set_example_value(proxy.node, example_value)
  2731. return ConstantVariable.create(example_value, **options)
  2732. else:
  2733. unimplemented_v2(
  2734. gb_type="torch.* op returned non-Tensor",
  2735. context=f"example_value type: {typestr(example_value)}; op: {proxy.node.op}; target: {proxy.node.target}",
  2736. explanation="torch.* ops that return a non-Tensor cannot be traced into the Dynamo FX graph output",
  2737. hints=[],
  2738. )
  2739. def infer_subclass_type(value):
  2740. if type(value) in (
  2741. torch.Tensor,
  2742. torch.nn.Parameter,
  2743. torch._subclasses.fake_tensor.FakeTensor,
  2744. torch._subclasses.functional_tensor.FunctionalTensor,
  2745. ) or is_traceable_wrapper_subclass(value):
  2746. # Ordinarily, we would fakeify a tensor so that it can get dynamic
  2747. # shapes and be computed on without triggering actual operations.
  2748. # However, how can we fakeify a tensor subclass? Ordinary
  2749. # inheritance (nor multiple inheritance) won't work work.
  2750. #
  2751. # Instead, our plan is to *manually simulate* the tensor subclass
  2752. # inheriting from a fake tensor with dynamo. This means our
  2753. # data representation for a tensor subclass will be a fake tensor
  2754. # + tensor subclass type + any extra data the subclass may have
  2755. # been storing on the tensor. Because all Python accesses are
  2756. # mediated through TensorWithTFOverrideVariable, we can ensure
  2757. # that we dispatch differently, e.g., according to
  2758. # __torch_function__
  2759. #
  2760. # To simplify things for now, the __dict__ tracking bits haven't
  2761. # been implemented yet, but they can be added into this design at
  2762. # a later point in time.
  2763. return None
  2764. else:
  2765. return type(value)
  2766. def get_specialized_props(target_cls, tx, example_value, subclass_type):
  2767. specialized_props = target_cls.specialize(example_value)
  2768. # TODO: not sure about this fake mode test
  2769. if (
  2770. isinstance(example_value, torch._subclasses.fake_tensor.FakeTensor)
  2771. and example_value.fake_mode is tx.fake_mode
  2772. ):
  2773. if subclass_type:
  2774. tensor_type = subclass_type
  2775. elif isinstance(example_value, torch.nn.Parameter):
  2776. tensor_type = torch.nn.Parameter
  2777. elif isinstance(example_value, torch.nn.Buffer):
  2778. tensor_type = torch.nn.Buffer
  2779. else:
  2780. tensor_type = torch.Tensor
  2781. specialized_props["class_type"] = tensor_type
  2782. return specialized_props
  2783. def construct_tensor_variable(
  2784. target_cls, tx, proxy, example_value, subclass_type, options
  2785. ):
  2786. """
  2787. Actually construct a tensor variable after all the pre-processing from
  2788. wrapping a pre-existing or newly created tensor value.
  2789. """
  2790. # NB: In most (all?) cases, this does not actually do a clone.
  2791. # (WARNING: this means that if we mutate metadata on the fake
  2792. # tensor, the stored example value will update too!)
  2793. example_value = _clone_input(example_value, tx.fake_mode)
  2794. set_example_value(proxy.node, example_value)
  2795. # We bind the unbacked symints in sizes/trdies of tensor lazily.
  2796. # So that subgraphs can access the unbacked symbol's proxy in parent graph
  2797. # when lifting unbacked symbols of input tensors to subgraph inputs.
  2798. # We do it lazily because the tensor may not be used in subgraphs.
  2799. if proxy.node.op != "placeholder":
  2800. tx.output.current_tracer.track_produced_symints(example_value, proxy)
  2801. options.update(get_specialized_props(target_cls, tx, example_value, subclass_type))
  2802. return target_cls(proxy, **options)
  2803. def get_automatic_dynamic_shapes_mark_as():
  2804. if config.automatic_dynamic_shapes_mark_as == "dynamic":
  2805. return DimDynamic.DYNAMIC
  2806. elif config.automatic_dynamic_shapes_mark_as == "unbacked":
  2807. return DimDynamic.SIZE_LIKE_UNBACKED
  2808. elif config.automatic_dynamic_shapes_mark_as == "oblivious":
  2809. return DimDynamic.OBLIVIOUS_SIZE
  2810. else:
  2811. raise ValueError(
  2812. f"invalid automatic_dynamic_shapes_mark_as = {config.automatic_dynamic_shapes_mark_as}"
  2813. )
  2814. _DYNAMIC_SOURCES: Optional[set[str]] = None
  2815. _DYNAMIC_SOURCES_CONFIG_HASH: Optional[int] = None
  2816. def get_dynamic_sources() -> set[str]:
  2817. global _DYNAMIC_SOURCES, _DYNAMIC_SOURCES_CONFIG_HASH
  2818. current_hash = hash(torch.compiler.config.dynamic_sources)
  2819. # If we have already calculated the sources and the config hasn't changed, return cached result
  2820. if _DYNAMIC_SOURCES is not None and _DYNAMIC_SOURCES_CONFIG_HASH == current_hash:
  2821. return _DYNAMIC_SOURCES
  2822. # Config has changed or first time, (re)calculate the sources
  2823. _DYNAMIC_SOURCES = {
  2824. s
  2825. for s in torch.compiler.config.dynamic_sources.replace(" ", "").split(",")
  2826. if s
  2827. }
  2828. _DYNAMIC_SOURCES_CONFIG_HASH = current_hash
  2829. return _DYNAMIC_SOURCES
  2830. def is_dynamic_source(source_name: str) -> bool:
  2831. dynamic_sources = get_dynamic_sources()
  2832. for pattern in dynamic_sources:
  2833. if pattern == source_name or re.match(pattern, source_name):
  2834. log.debug(
  2835. "%s was marked dynamic due to dynamic source allowlist pattern: %s",
  2836. source_name,
  2837. pattern,
  2838. )
  2839. return True
  2840. return False
  2841. def record_automatic_dynamic(
  2842. tx: "InstructionTranslator", name: str, e: torch.Tensor
  2843. ) -> FrameStateSizeEntry:
  2844. # This mimics stride inference algorithm in _create_symbolic_sizes_strides_storage_offset
  2845. ex_size = e.size()
  2846. if not is_sparse_any(e):
  2847. ex_stride = e.stride()
  2848. dim = e.dim()
  2849. stride = [None] * dim
  2850. pending = [(ex_stride[i], -i) for i in range(dim)]
  2851. pending.sort(key=_nested_int_aware_sort)
  2852. candidates = {}
  2853. for i_stride, neg_i in pending:
  2854. i = -neg_i
  2855. stride[i] = candidates.get(i_stride, i_stride)
  2856. candidates.setdefault(i_stride * ex_size[i], InferStride(i))
  2857. else:
  2858. stride = []
  2859. return process_automatic_dynamic(
  2860. tx, name, FrameStateSizeEntry.make_tensor(tuple(ex_size), tuple(stride))
  2861. )
  2862. _UNBACKED_SOURCES: Optional[set[str]] = None
  2863. _UNBACKED_SOURCES_CONFIG_HASH: Optional[int] = None
  2864. def get_unbacked_sources() -> set[str]:
  2865. global _UNBACKED_SOURCES, _UNBACKED_SOURCES_CONFIG_HASH
  2866. current_hash = hash(torch.compiler.config.unbacked_sources)
  2867. # If we have already calculated the sources and the config hasn't changed, return cached result
  2868. if _UNBACKED_SOURCES is not None and _UNBACKED_SOURCES_CONFIG_HASH == current_hash:
  2869. return _UNBACKED_SOURCES
  2870. # Config has changed or first time, (re)calculate the sources
  2871. _UNBACKED_SOURCES = {
  2872. s
  2873. for s in torch.compiler.config.unbacked_sources.replace(" ", "").split(",")
  2874. if s
  2875. }
  2876. _UNBACKED_SOURCES_CONFIG_HASH = current_hash
  2877. return _UNBACKED_SOURCES
  2878. def is_unbacked_source(source_name: str) -> bool:
  2879. unbacked_sources = get_unbacked_sources()
  2880. for pattern in unbacked_sources:
  2881. if pattern == source_name or re.match(pattern, source_name):
  2882. log.debug(
  2883. "%s was marked unbacked due to unbacked source allowlist pattern: %s",
  2884. source_name,
  2885. pattern,
  2886. )
  2887. return True
  2888. return False
  2889. # Performs automatic dynamic dim determination.
  2890. # Returns a SymbolicContext
  2891. def _automatic_dynamic(
  2892. e, tx, source, static_shapes, outer_only=False
  2893. ) -> SymbolicContext:
  2894. # strided NT not supported
  2895. if e.is_nested and not isinstance(
  2896. e, torch.nested._internal.nested_tensor.NestedTensor
  2897. ):
  2898. unimplemented_v2(
  2899. gb_type="Encountered strided NestedTensor in automatic dynamic dim determination",
  2900. context="",
  2901. explanation="torch.compile does not support strided NestedTensor",
  2902. hints=[],
  2903. )
  2904. name = source.name()
  2905. prior_policy = tx.output.tracing_context.tensor_to_context.get(e, None)
  2906. shape_env_to_source_to_symbol_cache = (
  2907. prior_policy.shape_env_to_source_to_symbol_cache if prior_policy else None
  2908. )
  2909. # Get base context if the tensor is a view
  2910. view_base_context: Optional[SymbolicContext] = None
  2911. if e._is_view():
  2912. base_source = AttrSource(source, "_base")
  2913. view_base_context = _automatic_dynamic(e._base, tx, base_source, static_shapes)
  2914. if is_traceable_wrapper_subclass(e) and not outer_only:
  2915. # Get symbolic context for outer tensor
  2916. outer_context = _automatic_dynamic(
  2917. e, tx, source, static_shapes, outer_only=True
  2918. )
  2919. # Get symbolic contexts for inner tensors
  2920. inner_contexts = {} # mapping from attr -> symbolic context
  2921. attrs, _ = type(e).__tensor_flatten__(e)
  2922. for attr in attrs:
  2923. inner_tensor = getattr(e, attr)
  2924. inner_source = AttrSource(source, attr)
  2925. inner_contexts[attr] = _automatic_dynamic(
  2926. inner_tensor, tx, inner_source, static_shapes
  2927. )
  2928. return SubclassSymbolicContext(
  2929. dynamic_sizes=outer_context.dynamic_sizes,
  2930. dynamic_strides=outer_context.dynamic_strides,
  2931. constraint_sizes=outer_context.constraint_sizes,
  2932. constraint_strides=outer_context.constraint_strides,
  2933. view_base_context=view_base_context,
  2934. tensor_source=outer_context.tensor_source,
  2935. shape_env_to_source_to_symbol_cache=outer_context.shape_env_to_source_to_symbol_cache,
  2936. inner_contexts=inner_contexts,
  2937. )
  2938. if static_shapes and not is_dynamic_source(name):
  2939. return StatefulSymbolicContext(
  2940. dynamic_sizes=[DimDynamic.STATIC] * e.dim(),
  2941. dynamic_strides=[DimDynamic.INFER_STRIDE] * e.dim(),
  2942. constraint_sizes=[None] * e.dim(),
  2943. constraint_strides=[None] * e.dim(),
  2944. view_base_context=view_base_context,
  2945. tensor_source=source,
  2946. shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
  2947. )
  2948. # We preserve the dynamism of inputs. For example, when users call
  2949. # make_fx(torch.cond, tracing_mode="symbolic")(*args), inputs have SymInt sizes.
  2950. from torch.fx.experimental.symbolic_shapes import is_nested_int
  2951. if any(isinstance(s, SymInt) and not is_nested_int(s) for s in e.size()):
  2952. return StatefulSymbolicContext(
  2953. dynamic_sizes=[
  2954. DimDynamic.DYNAMIC if isinstance(s, SymInt) else DimDynamic.STATIC
  2955. for s in e.size()
  2956. ],
  2957. dynamic_strides=[DimDynamic.INFER_STRIDE] * e.dim(),
  2958. constraint_sizes=[None] * e.dim(),
  2959. constraint_strides=[None] * e.dim(),
  2960. view_base_context=view_base_context,
  2961. tensor_source=source,
  2962. shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
  2963. )
  2964. # Prep for automatic dynamic
  2965. frame_state_entry = record_automatic_dynamic(tx, name, e)
  2966. # TODO: index export_constraints ahead of time so we don't have to
  2967. # do a linear scan every time here
  2968. t_id = id(e)
  2969. dim2constraint = {}
  2970. def update_dim2constraint(dim, constraint_range, name):
  2971. if dim in dim2constraint:
  2972. from torch.fx.experimental.symbolic_shapes import StrictMinMaxConstraint
  2973. old_constraint_range, old_name = dim2constraint[dim]
  2974. new_constraint_range = StrictMinMaxConstraint(
  2975. vr=constraint_range.vr & old_constraint_range.vr,
  2976. warn_only=False,
  2977. )
  2978. # It is possible for (non-None) old_name and name to be different
  2979. # but this will only happen the corresponding Dims can be derived equal.
  2980. new_name = old_name or name
  2981. dim2constraint[dim] = new_constraint_range, new_name
  2982. else:
  2983. dim2constraint[dim] = constraint_range, name
  2984. from torch.export.dynamic_shapes import _RelaxedConstraint
  2985. if tx.output.export_constraints:
  2986. for constraint in tx.output.export_constraints:
  2987. if isinstance(constraint, _RelaxedConstraint):
  2988. continue
  2989. if constraint.t_id == t_id:
  2990. update_dim2constraint(
  2991. constraint.dim, constraint.constraint_range, constraint.name
  2992. )
  2993. dynamic_sizes = []
  2994. dynamic_strides = []
  2995. constraint_sizes = []
  2996. constraint_strides = []
  2997. specialize_on = []
  2998. for i in range(e.dim()):
  2999. # NB: mark dynamic has precedence over static
  3000. marked_strict_unbacked = i in getattr(
  3001. e, "_dynamo_strict_unbacked_indices", set()
  3002. )
  3003. marked_unbacked = i in getattr(e, "_dynamo_unbacked_indices", set())
  3004. marked_dynamic = i in getattr(e, "_dynamo_dynamic_indices", set())
  3005. marked_weak_dynamic = i in getattr(e, "_dynamo_weak_dynamic_indices", set())
  3006. marked_static = i in getattr(e, "_dynamo_static_indices", set())
  3007. specialize_on.append(getattr(e, "_specialize_on", {}).get(i, []))
  3008. # Reflect the user directive in the frame_state
  3009. # For dynamic, apply None always
  3010. normalized_source_name = normalize_source_name(source.name())
  3011. base_source = source
  3012. if isinstance(base_source, ChainedSource):
  3013. base_source = base_source.get_base()
  3014. if marked_dynamic or (
  3015. isinstance(base_source, LocalSource)
  3016. and base_source.dynamism is not None
  3017. and dict(base_source.dynamism).get(normalized_source_name, {i: False})[i]
  3018. ):
  3019. # TODO: This can be batched
  3020. # TODO: Doing this here is kind of sus, maybe better to set this
  3021. # up when we initially created the FrameStateSizeEntry to bong
  3022. # into the mutable state
  3023. log.debug("automatic dynamic %s marked dynamic", name)
  3024. mark_size = [auto_unset] * e.dim()
  3025. mark_size[i] = auto_dynamic
  3026. frame_state_entry |= FrameStateSizeEntry.make_size(size=mark_size)
  3027. # NB: both static and dynamic have precedence over
  3028. automatic_dynamic_size = (
  3029. config.automatic_dynamic_shapes and frame_state_entry.is_size_dynamic(i)
  3030. )
  3031. # NB: previously, if size was dynamic, we wouldn't make its stride
  3032. # dynamic. But now, because of InferStride concept, we will properly
  3033. # not make stride dynamic even if it's wobbling
  3034. automatic_dynamic_stride = (
  3035. config.automatic_dynamic_shapes and frame_state_entry.is_stride_dynamic(i)
  3036. )
  3037. if is_dynamic_source(name):
  3038. log.debug("%s marked dynamic via source whitelist", name)
  3039. automatic_dynamic_size = True
  3040. if is_unbacked_source(name):
  3041. log.debug("%s marked unbacked via source whitelist", name)
  3042. automatic_dynamic_size = True
  3043. automatic_dynamic = automatic_dynamic_size or automatic_dynamic_stride
  3044. # We will process constraints first, as they will imply that we
  3045. # have a dynamic dimension
  3046. # Precedence: export constraints > eager constraints
  3047. constraint = dim2constraint.get(i)
  3048. if constraint is None:
  3049. constraint_size = None
  3050. constraint_stride = None
  3051. if marked_dynamic and not config.allow_ignore_mark_dynamic:
  3052. # constraint_stride is deliberaly kept None because no easy way to provide value ranges for mark dynamic
  3053. constraint_stride = None
  3054. if hasattr(e, "_dynamo_dynamic_range"):
  3055. dim_range = [
  3056. dr for dr in e._dynamo_dynamic_range if dr.dim == i
  3057. ].pop()
  3058. if dim_range.min is None and dim_range.max is None:
  3059. constraint_size = RelaxedUnspecConstraint(warn_only=False)
  3060. else:
  3061. from torch.fx.experimental.symbolic_shapes import (
  3062. StrictMinMaxConstraint,
  3063. )
  3064. constraint_size = StrictMinMaxConstraint(
  3065. vr=ValueRanges(lower=dim_range.min, upper=dim_range.max),
  3066. warn_only=False,
  3067. )
  3068. else:
  3069. constraint_size = RelaxedUnspecConstraint(warn_only=False)
  3070. elif marked_strict_unbacked:
  3071. constraint_size = RelaxedUnspecConstraint(warn_only=False)
  3072. elif not marked_static and automatic_dynamic:
  3073. set_feature_use("dynamo.automatic_dynamic_shapes", True)
  3074. if automatic_dynamic_size:
  3075. constraint_size = RelaxedUnspecConstraint(warn_only=True)
  3076. if automatic_dynamic_stride:
  3077. constraint_stride = RelaxedUnspecConstraint(warn_only=True)
  3078. else:
  3079. if not marked_static and not config.automatic_dynamic_shapes:
  3080. set_feature_use("dynamo.automatic_dynamic_shapes", False)
  3081. constraint_size = None
  3082. constraint_stride = None
  3083. else:
  3084. constraint_size, name_ = constraint
  3085. constraint_stride = None
  3086. dim_name = f"{name}.size()[{i}]"
  3087. tx.output.shape_env.source_name_to_debug_name[dim_name] = name_
  3088. constraint_sizes.append(constraint_size)
  3089. constraint_strides.append(constraint_stride)
  3090. if marked_unbacked or is_unbacked_source(name):
  3091. dynamic_size = DimDynamic.SIZE_LIKE_UNBACKED
  3092. elif (
  3093. constraint_size is not None
  3094. or marked_dynamic
  3095. or marked_weak_dynamic
  3096. or is_nested_int(e.size()[i])
  3097. ):
  3098. # NB: We could assert static_shapes is False here, but it
  3099. # seems better to allow the user to override symbolic_context in this
  3100. # case
  3101. if automatic_dynamic:
  3102. dynamic_size = get_automatic_dynamic_shapes_mark_as()
  3103. else:
  3104. dynamic_size = DimDynamic.DYNAMIC
  3105. elif static_shapes or config.assume_static_by_default or marked_static:
  3106. dynamic_size = DimDynamic.STATIC
  3107. else:
  3108. # TODO: When does this show up?
  3109. dynamic_size = DimDynamic.DUCK
  3110. if constraint_stride is not None:
  3111. dynamic_stride = DimDynamic.DYNAMIC
  3112. else:
  3113. dynamic_stride = DimDynamic.INFER_STRIDE
  3114. dynamic_sizes.append(dynamic_size)
  3115. dynamic_strides.append(dynamic_stride)
  3116. return StatefulSymbolicContext(
  3117. dynamic_sizes=dynamic_sizes,
  3118. dynamic_strides=dynamic_strides,
  3119. constraint_sizes=constraint_sizes,
  3120. constraint_strides=constraint_strides,
  3121. specialize_on=specialize_on,
  3122. view_base_context=view_base_context,
  3123. tensor_source=source,
  3124. shape_env_to_source_to_symbol_cache=shape_env_to_source_to_symbol_cache,
  3125. )
  3126. # See note [Tensor Fakification and Symbol Caching]
  3127. def wrap_to_fake_tensor_and_record(
  3128. e, tx, *, source: Optional[Source], is_tensor: bool, parent_context=None
  3129. ):
  3130. if (
  3131. type(e) in (torch.Tensor, torch.nn.Parameter, FakeTensor)
  3132. or isinstance(e, torch.Tensor)
  3133. or is_traceable_wrapper_subclass(e)
  3134. ):
  3135. assert source is not None
  3136. static_shapes, _reason = tensor_always_has_static_shape(
  3137. e,
  3138. is_tensor,
  3139. tensor_source=source,
  3140. )
  3141. if not parent_context:
  3142. symbolic_context = _automatic_dynamic(e, tx, source, static_shapes)
  3143. else:
  3144. # Parent contexts are passed in when we are recursively creating
  3145. # fake tensors for subclasses. A better design would be not to create a
  3146. # parent/child relationship, but to recursively call _automatic_dynamic
  3147. # as we recursively call wrap_to_fake_tensor_and_record. This runs
  3148. # into bugs around how meta_utils knows and works to create fake tensors
  3149. # with tensor subclasses. Ideally, dynamo would drive both the recursive
  3150. # wrap_to_fake_tensor_and_record and _automatic_dynamic policy creation.
  3151. assert isinstance(source, AttrSource)
  3152. inner_context_name = source.member
  3153. symbolic_context = parent_context.inner_contexts[inner_context_name]
  3154. log.debug(
  3155. "wrap_to_fake %s %s %s %s",
  3156. source.name(),
  3157. tuple(e.shape),
  3158. symbolic_context,
  3159. type(e),
  3160. )
  3161. # Note [enable_python_dispatcher in dynamo]
  3162. # Dynamo disables itself when it runs fake tensor prop, which means that tensor subclasses
  3163. # have no way to know (purely based off of global state) if they are currently being run under compile or not.
  3164. # we use enable_python_dispatcher mainly to tweak the DispatchKeyState so that subclass authors
  3165. # can check it to know if they are running in an eager context or not
  3166. with enable_python_dispatcher():
  3167. fake_e = wrap_fake_exception(
  3168. lambda: tx.fake_mode.from_tensor(
  3169. e,
  3170. source=source,
  3171. symbolic_context=symbolic_context,
  3172. )
  3173. )
  3174. if (
  3175. source is not None
  3176. and isinstance(fake_e, FakeTensor)
  3177. and (sym_val := fake_e.item_memo) is not None
  3178. ):
  3179. tx.output.tracked_fakes.append(
  3180. TrackedFake(sym_val, CallMethodItemSource(source), symbolic_context)
  3181. )
  3182. if is_traceable_wrapper_subclass(fake_e):
  3183. attrs, _ = fake_e.__tensor_flatten__()
  3184. for attr in attrs:
  3185. fake_inner = getattr(fake_e, attr)
  3186. inner = getattr(e, attr)
  3187. inner_source = AttrSource(source, attr)
  3188. wrap_to_fake_tensor_and_record(
  3189. inner,
  3190. tx,
  3191. source=inner_source,
  3192. is_tensor=isinstance(fake_inner, torch.Tensor),
  3193. parent_context=symbolic_context,
  3194. )
  3195. tx.output.tracing_context.tensor_to_context[e] = symbolic_context
  3196. if is_sparse_any(fake_e):
  3197. # TODO: for TensorGuards, this eventually may need more
  3198. # fields for the size/stride of any other constituents
  3199. values = fake_e._values() if fake_e.is_sparse else fake_e.values()
  3200. tx.output.input_source_to_sizes_strides[source] = {
  3201. "size": fake_e.size(),
  3202. # TODO: revise this, but for now this stride instead of ()
  3203. # avoids SegFault with PYTORCH_TEST_WITH_DYNAMO=1
  3204. "stride": (1,) * fake_e.ndim,
  3205. "values_size": values.size(),
  3206. "values_stride": values.stride(),
  3207. }
  3208. else:
  3209. tx.output.input_source_to_sizes_strides[source] = {
  3210. "size": fake_e.size(),
  3211. "stride": fake_e.stride(),
  3212. }
  3213. if (
  3214. is_tensor
  3215. and not (static_shapes and source.is_specialized_nn_module())
  3216. and not is_constant_source(source)
  3217. ):
  3218. tx.output.tracked_fakes.append(
  3219. TrackedFake(fake_e, source, symbolic_context)
  3220. )
  3221. tx.output.tracked_fakes_id_to_source[id(e)].append(source)
  3222. return fake_e
  3223. else:
  3224. return e
  3225. class SourcelessBuilder:
  3226. """
  3227. Like builder, but stateless and does not require a source. Useful for simple type->VT objects, or objects
  3228. that are being created/evaporated during inlining (ex: consider a locally made list of tensors we then iterate over
  3229. .), such a list should not show up as an artifact from inputs, nor in reconstruction, nor in the graph. However,
  3230. there may be reasons to represent it as a ListVariable internally.
  3231. NOTE - Objects produced here are born UNGUARDED due to the nature of sources!
  3232. NOTE - This class is very new! It will have some rough edges, but it was created to stem the bleeding of giant
  3233. if/else type->VariableTracker trees that were cropping up all over dynamo.
  3234. """
  3235. def __init__(self) -> None:
  3236. raise AssertionError("Use SourcelessBuilder.create()")
  3237. @staticmethod
  3238. def create(tx: "InstructionTranslator", value) -> VariableTracker:
  3239. value_type = type(value)
  3240. fast_handler = SourcelessBuilder._type_handlers.get(value_type)
  3241. if fast_handler:
  3242. return fast_handler(tx, value)
  3243. if isinstance(value, VariableTracker):
  3244. # This is always valid to call, and useful for recursive calls.
  3245. return value
  3246. elif isinstance(value, dataclasses._HAS_DEFAULT_FACTORY_CLASS):
  3247. return UserDefinedObjectVariable(value)
  3248. elif ConstantVariable.is_literal(value):
  3249. return ConstantVariable.create(value)
  3250. elif callable(value) and trace_rules.lookup_callable(value) is not None:
  3251. if trace_rules.is_callable_allowed(value):
  3252. tx.output.has_user_defined_allowed_in_graph = True
  3253. return trace_rules.lookup_callable(value)(value)
  3254. elif callable(value) and UserDefinedClassVariable.is_supported_new_method(
  3255. value
  3256. ):
  3257. # NamedTuple._make uses an alias of tuple.__new__
  3258. obj = trace_rules.lookup_callable(value.__self__)(value.__self__)
  3259. return GetAttrVariable(obj, "__new__")
  3260. elif is_function_or_wrapper(value):
  3261. return trace_rules.lookup(value)(value)
  3262. elif isinstance(
  3263. value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
  3264. ):
  3265. return EnumVariable(value)
  3266. elif isinstance(value, (type, abc.ABCMeta)):
  3267. return UserDefinedClassVariable(value)
  3268. elif isinstance(value, types.MethodWrapperType):
  3269. return MethodWrapperVariable(value)
  3270. elif (
  3271. isinstance(value, types.MethodType)
  3272. # We only want to support sourceless class objects here
  3273. # An instance variable is not allowed and it should have source
  3274. and isinstance(value.__self__, (type, abc.ABCMeta))
  3275. ):
  3276. # value is a classmethod
  3277. assert getattr(value.__self__, value.__func__.__name__) == value
  3278. cls_obj_vt = SourcelessBuilder.create(tx, value.__self__)
  3279. try:
  3280. return cls_obj_vt.var_getattr(tx, value.__func__.__name__)
  3281. except NotImplementedError:
  3282. pass # failthrough to unimplemented branch
  3283. elif isinstance(value, torch.fx.graph_module.GraphModule):
  3284. return SourcelessGraphModuleVariable(value)
  3285. elif isinstance(
  3286. value, (torch.utils._pytree.TreeSpec, torch.utils._pytree.LeafSpec)
  3287. ):
  3288. return UserDefinedObjectVariable(value)
  3289. elif PlacementVariable.is_placement(value):
  3290. return PlacementVariable(value)
  3291. elif DeviceMeshVariable.is_device_mesh(value):
  3292. return DeviceMeshVariable(value)
  3293. elif value is functools.wraps:
  3294. return FunctoolsWrapsVariable(value)
  3295. elif isinstance(value, re.Pattern):
  3296. return RegexPatternVariable(value)
  3297. elif isinstance(value, torch._dynamo.variables.lazy.LazySymNodeFormatString):
  3298. return ConstantVariable.create(str(value))
  3299. elif isinstance(value, type(torch._higher_order_ops.flex_attention_backward)):
  3300. return torch._dynamo.variables.higher_order_ops.FlexAttentionBackwardHighOrderVariable(
  3301. value
  3302. )
  3303. elif isinstance(value, types.GenericAlias):
  3304. return TypingVariable(value)
  3305. elif is_namedtuple(value):
  3306. output = [
  3307. SourcelessBuilder.create(tx, getattr(value, name))
  3308. for name in namedtuple_fields(type(value))
  3309. ]
  3310. return NamedTupleVariable(output, tuple_cls=type(value))
  3311. elif (
  3312. isinstance(value, torch.SymInt)
  3313. and value.node.expr in tx.output.bound_symbols
  3314. ):
  3315. proxy = tx.output.bound_symbols[value.node.expr]
  3316. return SymNodeVariable.create(tx, proxy)
  3317. unimplemented_v2(
  3318. gb_type="Unexpected type in sourceless builder",
  3319. context=f"{value_type.__module__}.{value_type.__qualname__}",
  3320. explanation=f"SourcelessBuilder.create does not know how to wrap {value_type}",
  3321. hints=[*graph_break_hints.DYNAMO_BUG],
  3322. )
  3323. @staticmethod
  3324. def wrap_constant_literal(value):
  3325. assert ConstantVariable.is_literal(value)
  3326. return ConstantVariable.create(value=value)
  3327. @staticmethod
  3328. def make_type_handlers():
  3329. create = SourcelessBuilder.create
  3330. handlers = {}
  3331. for t in common_constant_types:
  3332. handlers[t] = lambda tx, value: ConstantVariable(value)
  3333. handlers[set] = lambda tx, value: SetVariable(
  3334. [create(tx, x) for x in value], mutation_type=ValueMutationNew()
  3335. )
  3336. handlers[dict] = lambda tx, value: ConstDictVariable(
  3337. {create(tx, k): create(tx, v) for k, v in value.items()},
  3338. type(value),
  3339. mutation_type=ValueMutationNew(),
  3340. )
  3341. handlers[list] = lambda tx, value: ListVariable(
  3342. [create(tx, x) for x in value], mutation_type=ValueMutationNew()
  3343. )
  3344. handlers[tuple] = lambda tx, value: TupleVariable(
  3345. [create(tx, x) for x in value]
  3346. )
  3347. handlers[torch.Size] = lambda tx, value: SizeVariable(
  3348. [create(tx, x) for x in value]
  3349. )
  3350. handlers[collections.OrderedDict] = handlers[dict]
  3351. handlers[immutable_dict] = handlers[dict]
  3352. handlers[immutable_list] = handlers[list]
  3353. handlers[random.Random] = lambda tx, value: RandomClassVariable()
  3354. handlers[types.ModuleType] = lambda tx, value: PythonModuleVariable(value)
  3355. handlers[torch.DispatchKeySet] = lambda tx, value: DispatchKeySetVariable(
  3356. value, mutation_type=ValueMutationNew()
  3357. )
  3358. handlers[torch._functorch.pyfunctorch.FuncTorchInterpreter] = (
  3359. lambda tx, value: FuncTorchInterpreterVariable(
  3360. value, mutation_type=ValueMutationNew()
  3361. )
  3362. )
  3363. handlers[torch.distributions.constraints._Real] = (
  3364. lambda tx, value: UserDefinedObjectVariable(
  3365. value, mutation_type=ValueMutationNew()
  3366. )
  3367. )
  3368. handlers[torch.distributions.constraints._Interval] = (
  3369. lambda tx, value: UserDefinedObjectVariable(
  3370. value, mutation_type=ValueMutationNew()
  3371. )
  3372. )
  3373. handlers[torch.distributions.constraints.Constraint] = (
  3374. lambda tx, value: UserDefinedObjectVariable(
  3375. value, mutation_type=ValueMutationNew()
  3376. )
  3377. )
  3378. def passthrough(tx: "InstructionTranslator", value):
  3379. return value
  3380. for cls in VariableTrackerMeta.all_subclasses:
  3381. handlers[cls] = passthrough
  3382. return handlers
  3383. SourcelessBuilder._type_handlers = SourcelessBuilder.make_type_handlers()
  3384. class SourcelessUserDefinedObjectBuilder:
  3385. """
  3386. SourceLessBuilder does not return a UserDefinedObjectVariable, but in some
  3387. cases it might be ok to return UserDefinedObjects. In such case, use this
  3388. builder.
  3389. """
  3390. def __init__(self) -> None:
  3391. raise AssertionError("Use SourcelessUserDefinedObjectBuilder.create()")
  3392. @staticmethod
  3393. def create(tx: "InstructionTranslator", value) -> VariableTracker:
  3394. value_type = type(value)
  3395. if issubclass(value_type, MutableMapping):
  3396. return MutableMappingVariable(value, mutation_type=ValueMutationNew())
  3397. elif isinstance(value, torch.nn.Module):
  3398. return UnspecializedNNModuleVariable(
  3399. value, mutation_type=ValueMutationNew()
  3400. )
  3401. else:
  3402. return UserDefinedObjectVariable(value, mutation_type=ValueMutationNew())