o
    ߗid                     @  sd  U d dl mZ d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dl	Z	d dl
Z
d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlmZ d dlmZmZmZmZ d dlmZmZmZmZm Z m!Z!m"Z"m#Z#m$Z$ d dl%m&Z& d dlm'Z' d dl(Z(d dl)Z(d dl*m+Z+ d d	l,m-Z- d d
l.m/Z/ d dl0m1Z1m2Z2m3Z3m4Z4 d dl)m5Z5 d dl6m7Z7m8Z8m9Z9m:Z: d dl;m<Z< d dl=m>Z>m?Z? d dl@mAZB d dlCmDZD d dlEmFZF d dlGmHZHmIZI d dlJmKZKmLZL ddlMmNZNmOZOmPZP ddlQmRZRmSZS ddlTmUZUmVZVmWZWmXZXmYZY ddlZm[Z[m\Z\m]Z]m^Z^ ddl_m`Z`maZambZbmcZc ddlOmdZdmeZemfZfmgZgmhZhmiZimjZjmkZkmlZlmmZmmnZn ddlompZpmqZqmrZr ddlsmtZt ddlumvZv ddlwmxZx dd lymzZz dd!l{m|Z|m}Z}m~Z~mZ dd"lPmZ dd#lmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZmZ dd$lmZ d%ed&< zd dlZW n ey   dZY nw ejrdd'lmZ dd(lmZ dd)lmZmZmZ dd*lmZ e	eZe(jed+Ze(jed,Ze Ze#d-Ze&d.ZG d/d0 d0ZG d1d2 d2Ze Ze Zdad3ed4< eeB	ddd=d>ZddAdBZecjddFdGZ		HdddPdQZd ae ZdRedS< ddUdVZddWdXZG dYdZ dZZ	[	H	dddadbZd dclmZ d ddlmZ ejrnddelmZ e Zdfedg< ddkdlZ		dd dmdddZG dd dZdddZdddZdddZG dd dej̓ZG dd dZdddZdS )    )annotationsN)Path)CellTypeCodeTypeFunctionType
ModuleType)	AnyCallableDictListOptionalSetTupleTypeVarUnion)	ParamSpec)ReferenceType)GlobalStateGuard)get_compile_pg)TensorifyState)compile_contextCompileContext	CompileIdtracing)
structured)compile_time_strobelight_metajustknobs_check#maybe_upload_prof_stats_to_manifoldsignpost_event)_use_lazy_graph_module)ConstraintViolationErrorGuardOnDataDependentSymNode)_forward_from_src)_WaitCounter)DistributedDataParallel)_disable_current_modesis_in_torch_dispatch_mode)CapturedTracebackformat_traceback_short   )configexctrace_rules)remove_dead_coderemove_pointless_jumps) check_inst_exn_tab_entries_validInstructionis_generator propagate_inst_exn_table_entriestransform_code_object)CacheSizeRelevantForFramecompute_cache_sizeexceeds_cache_size_limitis_recompilation)always_optimize_code_objects
dynamo_tls	skip_codeTorchPatcher)augment_exc_messageBackendCompilerFailedFailOnRecompileLimitHitformat_error_msgInternalTorchDynamoErrorRecompileLimitExceededSkipCodeRecursiveExceptionTorchRuntimeErrorUncapturedHigherOrderOpErrorunimplementedUnsupported)CheckFunctionManager&get_and_maybe_log_recompilation_reasonGuardedCode)Hooks)put_code_state)ExecutionRecord)TORCH_DYNAMO_RESUME_IN_PREFIX)DistributedStateInstructionTranslator
LocalStateSpeculationLog)is_numpy)chromium_event_timedCleanupManagerCompileTimeInstructionCountercountersdynamo_timedformat_bytecodegen_record_file_nameget_metrics_contextincrement_frameis_namedtupleistype
LazyStringorig_code_mapreset_graph_break_dup_checkersetup_compile_debug	to_int_ustroubleshooting_urlwrite_record_to_file)#torch_function_mode_stack_state_mgrzOptional[ModuleType]np)
CompilerFn)WrapBackendDebug)BytecodeHook
CacheEntryDynamoFrameType)FrameStateSizeEntrybytecodegraph_breaks_T_Pc                   @  s   e Zd ZdS )TODO_UNKNOWNN)__name__
__module____qualname__ ru   ru   Y/home/ubuntu/transcripts/venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.pyrq      s    rq   c                   @  s4   e Zd ZdddZdddZdddZdddZdS )TrackerreturnNonec                 C  s   g | _ t | _d S N)seensetseen_idsselfru   ru   rv   __init__   s   zTracker.__init__
strong_objr   c                   sH   t |  jvr"t| fdd}j| j  d S d S )Nc                   s   j  S rz   )r}   remove)_idxr   ru   rv   <lambda>       zTracker.add.<locals>.<lambda>)idr}   weakrefrefr{   appendadd)r   r   objru   r   rv   r      s   
zTracker.additemboolc                 C  s   t || jv S rz   )r   r}   )r   r   ru   ru   rv   __contains__   s   zTracker.__contains__c                 C  s   | j   | j  d S rz   )r{   clearr}   r~   ru   ru   rv   r      s   
zTracker.clearN)rx   ry   )r   r   rx   ry   )r   r   rx   r   )rr   rs   rt   r   r   r   r   ru   ru   ru   rv   rw      s
    


rw   zOptional[GlobalStateGuard]initial_global_statesrcstrglobalsDict[str, Any]	co_fieldsOptional[Dict[str, str]]rx   r   c                 C  s   t | ||}t|j |S rz   )original_forward_from_srcr:   __code__)r   r   r   resultru   ru   rv   fx_forward_from_src_skip_result   s   
r   fnCallable[_P, _T]c                   s"   t  d
 fdd} |_|S )z
    Context manager to:
        1) Save/restore torch.is_grad_enabled() state
        2) Save/restore python random state
        3) Save/restore torch random state
        4) Monkey patch torch.fx.graph_module._forward_from_src
    args_P.argskwargs	_P.kwargsrx   ro   c                    s4  t  }t }tj  t }t }t }t	 }tj
 }d }	tj r/tj
 }	tj }
tjjj}ttjj_t }t }|tjj  |t zd | i |W |  tj dksjJ d|  tj| tjj| tj||d t | tj!| |	d urtj!|	 tj"|
 |tjj_|# sJ d|$  dW  d    S |  tj dksJ d|  tj| tjj| tj||d t | tj!| |	d urtj!|	 tj"|
 |tjj_|# sJ d|$  dw 1 sw   Y  d S )Nr   zQTorch function mode stack state changed while dynamo tracing, please report a bug)	warn_onlyzGlobal z7state changed while dynamo tracing, please report a bug)%r   torchis_grad_enabled_C_PreserveDispatchKeyGuardis_inference_mode_enabled$are_deterministic_algorithms_enabled-is_deterministic_algorithms_warn_only_enabledrandomgetstateget_rng_statecudais_available_get_cublas_allow_tf32fxgraph_moduler"   r   ra   
contextlib	ExitStackenter_context_symbolic_trace_maybe_revert_all_patchesre   close_len_torch_function_stack_set_grad_enabledautograd	grad_mode_enter_inference_modeuse_deterministic_algorithmssetstateset_rng_state_set_cublas_allow_tf32checkreason)r   r   guardsprior_grad_modeprior_inference_modeprior_deterministicprior_warn_onlypy_rng_statetorch_rng_statecuda_rng_state
allow_tf32prior_fwd_from_srccleanup
exit_stackr   ru   rv   _fn   s   











z"preserve_global_state.<locals>._fnNr   r   r   r   rx   ro   )	functoolswraps_torchdynamo_orig_callable)r   r   ru   r   rv   preserve_global_state   s   	0r   framerk   r   c                   s   | j tv rdS | j jD ].}|| jv r9| j| }t|tr)|jds&|tu r) dS t	r9t
jr9|t	u s6t|r9 dS qi d fdd | j D ]	} |rR dS qItd	| j j| j j| j j d
S )z+Check if the frame has torch.* related bitsTztorch.r   objectrx   r   c                   sZ  t  }|v r| S d|< t tjtjjfs&t tr.t tjjr.d|< | S t	j
rGtrGt tjs?t tjrGd|< | S t ttfr_tfdd D |< | S t tr{t  }tfdd|D |< | S t ttttdtfrd|< | S t rt drt fdd jD |< | S dS )	z)Recursively check if the obj has a tensorFTc                 3      | ]} |V  qd S rz   ru   .0v
has_tensorru   rv   	<genexpr>*      z:has_tensor_in_frame.<locals>.has_tensor.<locals>.<genexpr>c                 3  r   rz   ru   r   r   ru   rv   r   1  r   N_fieldsc                 3  s    | ]
} t |V  qd S rz   )getattrr   )r   r   ru   rv   r   7  s    )r   
isinstancer   TensornnModuler]   type
issubclassr*   trace_numpyrf   ndarraygenericlisttupleanydictvaluesr   intfloatr   r\   hasattrr   )r   obj_idr   r   r}   )r   rv   r     sF   


z'has_tensor_in_frame.<locals>.has_tensorz0skipping because no torch.* %s             %s %sFN)r   r   rx   r   )f_coder8   co_names	f_globalsr   r   rr   
startswithr   rf   r*   r   rR   f_localsr   logdebugco_nameco_filenameco_firstlineno)r   r   r   valueru   r   rv   has_tensor_in_frame  s4   




+r  Fe	Exceptioncoder   Optional[DynamoFrameType]exportry   c                 C  s:   d }t | drt| |}t|| j || _t| |d d S )Nexec_recordr  )r   rY   rd   r  record_filenamer<   )r  r  r   r  r	  ru   ru   rv   exception_handlerP  s   

r
  z/typing.Counter[Union[int, FrameStateSizeEntry]]FRAME_COMPILE_COUNTERfuncc                 C  s   t jrt| S | S rz   )r*   cprofilecprofile_wrapperr  ru   ru   rv   maybe_cprofilee  s   r  c                   s   t  d	 fdd}|S )
Nr   r   r   r   rx   ro   c               	     s  t  }|s
J dtdj dt|dd d}t }|  t		 }|j
g| R i |}t		 | }|  tdj|| t|}z|| W n tyb   td| Y nw td| |d	}	z'tjd
dddddt|gtjd}
tjdddt|	g|
jd td|	 W n" ty   td |tjjd |tjjd Y nw tt|  rt j!"d fdd |S )NzTrace id is Nonez/tmp/r   /z.profilez7### Cprofile for %s trace id [%s] took %.3f seconds ###zCannot write to %szRaw profile at %sz.svg	gprof2dotz-fpstatsz"--node-label=total-time-percentagez!--node-label=self-time-percentagez--node-label=total-time)stdoutdotz-Tsvgz-o)stdinz Generated SVG from profile at %szxFailed to generate SVG from profile -- dumping stats instead.Try installing gprof2dot and dot for a better visualization   linkc                     s
   d dS )Ncprofile_manifold_url)nameurlru   ru   manifold_linkru   rv   r     s   
 z;cprofile_wrapper.<locals>.profile_wrapper.<locals>.<lambda>)#r   current_trace_idr   rr   r   replacecProfileProfileenabletimeruncalldisabler   warningr  Stats
dump_statsPermissionError	exceptionwith_suffix
subprocessPopenPIPE
check_callr  FileNotFoundError
sort_statsSortKeyTIMEprint_stats
CUMULATIVEr   r   _loggingtrace_structured)r   r   trace_idprofile_pathprofstart_tsretvalprofile_latencypssvg_pathgprof2dot_processr  r  rv   profile_wrapperl  sr   

	
z)cprofile_wrapper.<locals>.profile_wrapperr   )r   r   )r  rA  ru   r  rv   r  k  s   <r  c                   @  s<   e Zd Z			d!d"ddZed#ddZddd$dd ZdS )%ConvertFrameAssertTFNcompiler_fnrg   	one_graphr   r  export_constraintsOptional[typing.Never]rx   ry   c                 C  s"   t   || _|| _|| _|| _d S rz   )r`   r   
_one_graph_export_export_constraints)r   rC  rD  r  rE  ru   ru   rv   r     s
   
zConvertFrameAssert.__init__*Callable[[CompilerFn], ConvertFrameAssert]c                       fddS )Nc                   s   t |  j j jS rz   )convert_frame_assertrG  rH  rI  backendr~   ru   rv   r     s    z8ConvertFrameAssert._clone_with_backend.<locals>.<lambda>ru   r~   ru   r~   rv   _clone_with_backend     z&ConvertFrameAssert._clone_with_backendr   skipr   rk   cache_entryOptional[CacheEntry]hooksrJ   frame_state*Dict[str, Union[int, FrameStateSizeEntry]]rR  r   Optional[GuardedCode]c                C  s  t   |j}t||}t| |tv rd S tjdr'tjd|j	kr'd S |j	dkr4|j
dr4d S |j	dkr;d S |j	dkrN|j
tjtjjrNd S |j	dkrZ|j
dkrZd S |j	dkrk|j
dkrkt|jskd S t|rstd	 t|syd S t ad
|vrt|d
< td7 a|d
 }t|tsJ t| }	t|  d7  < t||	}
tdd|j	|t|
|j
|j |j!|j"d |j	t#s|j	 d|j
 d|j  }t$j%&| t't(|
' t)|j|j*|j+|j|j,| j-| j.| j/| j0||||||
|d dW  d    S 1 sw   Y  d S )NTORCHDYNAMO_DEBUG_FUNCTIONz	<genexpr>)ztransformers/file_utils.pyztransformers/utils/generic.pyzdiffusers/utils/outputs.py__setattr__r   z<module><string>z<lambda>	generator_idr)   dynamoz_convert_frame_assert._compile)r   frame_id
compile_idr   r   
cache_sizeaccumulated_cache_size :)rV  r`  rR  )1r[   r   r5   input_codesr   output_codesosenvirongetr   r   endswithr   pathdirnamer   optim__file__r   
f_builtinsr1   rE   r  r   r   FRAME_COUNTERr   r   r  r   r   r   r   +num_cache_entries_with_same_id_matched_objsnum_cache_entriesrM   r9   traced_frame_infosr   r   r   _compiler   r   closurer   rG  rH  rI  )r   r   rS  rU  rV  rR  r  ra  r_  frame_compile_idr`  inforu   ru   rv   __call__  s   	






&zConvertFrameAssert.__call__TFN)
rC  rg   rD  r   r  r   rE  rF  rx   ry   )rx   rJ  r   rk   rS  rT  rU  rJ   rV  rW  rR  r   rx   rX  rr   rs   rt   r   propertyrO  rx  ru   ru   ru   rv   rB    s    rB  TrC  rg   rD  rE  rF  c                 C  s   t | |||S )z&Fully convert a frame into an FX graph)rB  )rC  rD  r  rE  ru   ru   rv   rL  7  s   rL  )OrderedDict)RemovableHandle)OutputGraphzDict[int, BytecodeHook]_bytecode_hookshookri   r~  c                 C  s   t t}| t|j< |S )zRegister hooks for bytecode generated by Dynamo. The hook can do some
    logging, as well as return a new code object to be used. Please refer
    to `BytecodeHook` for the hook signature.
    )r~  r  r   )r  handleru   ru   rv   register_bytecode_hookM  s   
r  rQ  Dict[str, object]localsbuiltinsru  Tuple[CellType]rU  rJ   rS  rT  ra  r4   rV  4Optional[Dict[str, Union[int, FrameStateSizeEntry]]]r`  r   rR  r   rX  c          -        s  ddl mm}mm} d
d d tj tdP	fd
d}dQ fdd}t	ddt
dQ	
fdd t }ttj tt tdddd | t t t  }r{t|t 	nd 	d t|r|rt|t|\}}|rdRdd }dSfd!d"}td#|tt||| t tjrt| d$tj rt!d%rt"| d&t#| d& t$d'j%j&j'd( d)(t)j*d( d*+  t,-t.tj/0d+fd,d- t12 }d }d }d }d }tj3j4j56  d }zKz||	|}t7  |W W r&i j8_9d.d/l4m:}  t;| }!|d u rld urlt<j=}"t<j>j=}#? }$t<j@jA}%t<jB}&d0d1 jCD }'d2d1 jDD }(tj3j4j5  nd }"d }#d }$d }%d }&ti }'ti }(t t12 | d3 
dTd6d7})|)tE }*i d8|!d9j%d:j&d;j'd<|jFd=|jGd>|"d?|#d@|$dA|%dB|&dC|dD|dE|dF|dG|'dH|(
|d utjHtjItjJtKL|*dtM
dI	}+|N|+ W  d    W  d    W  d    W  d    S  tOy[ }, zItP|,jQ}t;|,}tR|,|dJ tj/j0dKdLd- dMd- dN tST|,\}}tU|,tVtWtXtYtZt[|t\|f	rD t]tP|,jQ dOt;|, ^|,j_d d },~,ww rci j8_9d.d/l4m:}  t;| }!|d u rd urt<j=}"t<j>j=}#? }$t<j@jA}%t<jB}&d0d1 jCD }'d2d1 jDD }(tj3j4j5  nd }"d }#d }$d }%d }&ti }'ti }(t t12 | d3 
dTd6d7})|)tE }*i d8|!d9j%d:j&d;j'd<|jFd=|jGd>|"d?|#d@|$dA|%dB|&dC|dD|dE|dF|dG|'dH|(
|d utjHtjItjJtKL|*dtM
dI	}+|N|+ w 1 s+w   Y  W d    n	1 s;w   Y  W d    n1 sKw   Y  W d    d S W d    d S 1 sdw   Y  d S )UNr   )bisectBisectValidationExceptiontranslation_validation_enabledValidationExceptiong        instructionsList[Instruction]code_optionsr  rx   ry   c                   sn     t| 
	|dzbz0tjj      W d    n1 s3w   Y  W d    n1 sBw   Y  W n* tjyU   	    tj
tjtjfyb     tyr    rq jj  w W j  nj  w jd usJ jsJ j| d d < |j tjstd t|  t|  tt| | d d < d S )N)rV  speculation_logdistributed_statezJThe config.dead_code_elimination flag is deprecated, it's now always true.)restartrO   r   outputtracing_contextset_current_txrunr+   UnspecializeRestartAnalysisr   SpeculationRestartAnalysisTensorifyScalarRestartAnalysis	SkipFramer  	shape_envcall_cleanup_hooksoutput_instructionsupdater  r*   dead_code_eliminationwarningswarnr2   r/   r.   r-   )r  r  )r  r  ru  r  rC  r  r  rE  rV  r   r  rD  r  r  tf_mode_stacktracerr  ru   rv   	transform{  sj   
 

z_compile.<locals>.transformr  r   rD  r   rU  rJ   r  2Callable[[List[Instruction], Dict[str, Any]], Any]rX  c                   s   t  2}|tdddd |td  |tjj	  |t
   | |||W  d    S 1 s9w   Y  d S )Nz_compile.compile_innerentire_frame_compile!dynamo_cumulative_compile_time_us)
phase_namedynamo_compile_column_usz#pytorch.wait_counter.dynamo_compile)r   r   r   rW   r#   guardr   _dynamocallback_handlerinstall_callbacksrU   record)r  rD  rU  r  stack)_compile_innerru   rv   compile_inner  s    

 z_compile.<locals>.compile_innerr  )r  c                   sV  t    }}d#dd}|d| j| j| j|  d }t D ]}|t _zt	| |}W  np t
jyj }	 z2t|	t
js>t  tdtt|	j |	jpRtt|	 t   }|dkr`td W Y d }	~	qd }	~	w t
jy }	 z%t|	t
js|t  td|	| j| j| j |rtd W Y d }	~	 d S d }	~	ww d u sjd usJ d|d usJ |d| j| j| j| t D ]}
|
| |}|d ur|}q| t|< t | || d usJ d$dd}|d usJ || }||}d}|d| j!d |  d7 }|d|j!d |  7 }| j!d | |j!d | ks J |d}|d| j" d7 }|d|j" 7 }| j"|j"ks>J |d}|d| j# d7 }|d |j# 7 }| j#|j#ks\J |j$rg% rgd S j&d usoJ j't(j)|< t* |r~|j+nd }d urtnd!}d"| }t,||j-|}% s|j.d ur|.j& |S )%Nprefixr   r  filenameline_nor   r  r   rx   ry   c                 S  s*   t tjrt t| |||| d S d S rz   )bytecode_logisEnabledForloggingDEBUGr   rX   )r  r  r  r  r  ru   ru   rv   log_bytecode  s
   z6_compile.<locals>._compile_inner.<locals>.log_bytecodezORIGINAL BYTECODEzRestarting analysis due to %sd   z100+ RestartAnalysis() callsz.Skipping frame %s %s                     %s %sz%No graph captured with one_graph=Truez;compiler collective wasn't run before compilation completedzMODIFIED BYTECODEc                 S  s4   dd l }| j| j t| j|j@  t| j|j@  S )Nr   )inspectco_argcountco_kwonlyargcountr   co_flags
CO_VARARGSCO_VARKEYWORDS)r  r  ru   ru   rv   
count_args(  s   z4_compile.<locals>._compile_inner.<locals>.count_argszarg mismatch: zold code object has args z, znew code object has args zfree var mismatch: zold code object has free var znew code object has free var zcell var mismatch: zold code object has cell var znew code object has cell var UnknownzTorch-Compiled Region: )r  r   r  r   r  r   r  r   r  r   rx   ry   r  r   rx   r   )/r#  r   r   r   	itertoolscountr   ri  attemptr3   r+   RestartAnalysisr   r  r   r   r   rw  r^   r(   __traceback__r   restart_reasonr   r   rE   r  r   
all_statesr  r   r_   rf  co_varnamesco_freevarsco_cellvarsr  is_empty_graphr   cleanupsrT   instancerG   guard_fail_fnrI   guard_managerguard_export_fn)r  rD  rU  r  last_attempt_start_time
start_timer  out_coder  r  r  hook_outputr  total_argcount_oldtotal_argcount_newmsgcheck_fncompile_id_strannotation_strguarded_code)rS  r`  r  dynamo_time_before_restartr  restart_reasonsru   rv   r    s   	









z _compile.<locals>._compile_innerr^  T)reset_event_loglog_pt2_compile_eventr   c                 S  s   d| j  d| j d| j dS )N'z' (rd  ))r   r   r   )r  ru   ru   rv   format_func_info  s   z"_compile.<locals>.format_func_infoc                     s    sdS  d S )Nz$Unable to find recompilation reasonsru   ru   )recompile_reasonsru   rv   format_guard_failures  s   z'_compile.<locals>.format_guard_failuresztorch._dynamo hit config.%s (%s)
   function: %s
   last reason: %s
To log all recompilation reasons, use TORCH_LOGS="recompiles".
To diagnose recompilation issues, see %s.zG reached, because fail_on_cache_limit_hit = True this is a HARD failurez7pytorch/compiler:skip_code_recursive_on_cache_limit_hitz reachedzBtorchdynamo start compiling %s %s:%s, stack (elided %s frames):
%s    rQ  dynamo_startc                
     sL   dt tfddttjd d  j j	t
 jdg iS )Nr  c                   s   | d  kS )Nr  ru   )f)convert_frame_internru   rv   r     r   z,_compile.<locals>.<lambda>.<locals>.<lambda>   rQ  )liner  r  )r   r  	takewhiler   from_tracebackr'   extractsummaryr   r   intern_stringr   ru   )r  r  rR  ru   rv   r     s   


z_compile.<locals>.<lambda>r)   )
curr_framec                 S     h | ]}|j qS ru   rt   r   opru   ru   rv   	<setcomp>!      z_compile.<locals>.<setcomp>c                 S  r  ru   r  r  ru   ru   rv   r  "  s    g    eAdr   c                   s   h d  fdd|   D S )N>   verbosebase_dirrepro_afterrepro_levelTYPE_CHECKINGlog_file_namedebug_dir_rootrepro_tolerancerepro_forward_only_custom_ops_profile_save_config_ignorerepro_ignore_non_fpignore_logger_methodslog_compilation_metricssame_two_models_use_fp64traceable_tensor_subclassesreorderable_logging_functions)_autograd_backward_strict_mode_banned_ops+inject_BUILD_SET_unimplemented_TESTING_ONLYc                   s0   i | ]\}}| vr|t |trt|n|qS ru   )r   r|   r   )r   keyr   	blocklistru   rv   
<dictcomp>I  s
    z4_compile.<locals>.clean_for_json.<locals>.<dictcomp>)items)r  ru   r  rv   clean_for_json2  s   
z _compile.<locals>.clean_for_json	frame_keyr   r   r   ra  rb  guard_countshape_env_guard_countgraph_op_countgraph_node_countgraph_input_count	fail_typefail_reasonfail_user_frame_filenamefail_user_frame_linenonon_compliant_opscompliant_custom_ops)	r  dynamo_time_before_restart_shas_guarded_codeconfig_suppress_errors config_inline_inbuilt_nn_modulesspecialize_floatdynamo_config
is_forward%dynamo_compile_time_before_restart_usr  artifactc                   S  
   dddS )Ndynamo_errorstringr  encodingru   ru   ru   ru   rv   r        c                   S  s   t  S rz   	traceback
format_excru   ru   ru   rv   r     s    metadata_fn
payload_fnz: )r  r  r  r  rx   ry   )
r  r   rD  r   rU  rJ   r  r  rx   rX  )r  r   rx   r   )rx   r   )r  r   rx   r   )`torch.fx.experimental.validatorr  r  r  r  r   	overrides _get_current_function_mode_stackr   r   r  rZ   r   r*   use_lazy_graph_moduler   r   rS   r|   rQ   r   rN   rP   r7   rH   r6   r   r&  r   rc   fail_on_cache_limit_hitr>   &skip_code_recursive_on_cache_limit_hitr   rA   rE   r   r   r   r   joinr'   r  formatr   r  rn  r6  r7  r#  time_nsr  utilsReinplaceCountersr   rK   r  local_scoper  r   lenr   r  count_callsgraphnodesplaceholdersr   r!  get_config_copyrq  rr  suppress_errorsinline_inbuilt_nn_modulesr&  jsondumpsrb   update_outerr  r   rt   r
  r+   get_exc_messager   rF   rC   r=   AssertionErrorr    r!   rD   r@   with_tracebackr  )-r  r   r  r  ru  rC  rD  r  rE  rU  rS  ra  r   rV  r`  rR  r  r  r  r  metrics_context
compile_pgexceeded
limit_typer  r  start_time_nsr  r  r  r  r  r  r  r  r  r  r  r  r   r!  r  config_dictmetricsr  ru   )r  r  r  rS  ru  r  r`  rC  r  r  r  r  rE  rV  r   r  rD  r  r  r  rR  r  r  r  r  rv   rt  W  s  	0< 


.


	

  :}

-


	
  2T rt  c                   @  s2   e Zd ZdddZedd
dZ	ddddZdS )ConvertFramerC  rg   rU  rJ   rx   ry   c                 C  s   || _ t|dd| _|| _d S )NF)rD  )r   rL  _inner_convert_hooks)r   rC  rU  ru   ru   rv   r   s  s   
zConvertFrame.__init__*Callable[[WrapBackendDebug], ConvertFrame]c                   rK  )Nc                   s   t |  jS rz   )convert_framerZ  rM  r~   ru   rv   r   z  r   z2ConvertFrame._clone_with_backend.<locals>.<lambda>ru   r~   ru   r~   rv   rO  x  rP  z ConvertFrame._clone_with_backendr   r   rk   rS  rT  rV  rW  rR  r   ~Optional[Union[GuardedCode, torch._C._dynamo.eval_frame.SkipCodeRecursiveFlag, torch._C._dynamo.eval_frame.CacheLimitHitFlag]]c                   s  t d d  d7  < z| j|||||d d}t d d  d7  < |W S  ty } zt|tr1 t|t}t|trttj	rt
|drtt|j. |j}	dt|	}
d|
  tjjd	d
d  fddd tj dd W d    n1 s~w   Y  tjs|s t|dd }|j}t||||}|rtj|dd ntj|dd t|trtjj j!j"W  Y d }~S t|t#rtjj j!j$W  Y d }~S W Y d }~d S d }~ww )Nframestotalr)   rQ  okr`  r  z&Graph break: skip: from user code at:
r*  c                   S  r+  )Ndynamo_graph_break_reasonr-  r.  ru   ru   ru   ru   rv   r     r0  z'ConvertFrame.__call__.<locals>.<lambda>c                     s     dt   S )N
r1  ru   user_stack_traceru   rv   r     r  r4  T)exc_infor	  )%rV   rY  r  r   rD   rF   graph_break_logr  r  r  r   r   r   r`  
real_stackr=  r2  format_listr   r6  r7  r   r*   rI  r   r   r?   r   rw  r&  rB   r   r  
eval_frameskip_code_recursive_flagrA   cache_limit_hit_flag)r   r   rS  rU  rV  rR  r   r  	soft_fail
user_stackuser_stack_formattedr	  r  	error_msgru   rc  rv   rx  |  s^   








zConvertFrame.__call__N)rC  rg   rU  rJ   rx   ry   )rx   r[  )r   )r   rk   rS  rT  rU  rJ   rV  rW  rR  r   rx   r]  r{  ru   ru   ru   rv   rX  r  s    
	rX  c                 C  
   t | |S )zHTry to convert a frame into an FX graph, if error leave frame unmodified)rX  )rC  rU  ru   ru   rv   r\    s   
r\  r  c                 C  s   ddl m} tj}dt_t| d}t|}W d    n1 s!w   Y  tt	|j
 t
  |_
z%t|j|j
|j|j|j|ddd t tddd d i tddd W |t_d S |t_w )	Nr)   )eagerFrbr   *   i  )
rC  rD  r  rE  rU  ra  rS  r   rV  r`  )backends.debuggingrq  r*   replay_record_enabledopenrL   loadr   r  chainr   r  rt  r  r  r  ru  rJ   r4   r   )r  rq  original_replay_valin_filer  ru   ru   rv   replay  s4   r{  c                 C  s>   t jdk rdS t| D ]}|jdkr|jd   S qtd)N)      r   RESUMEr  z$RESUME instruction not found in code)sysversion_infodisget_instructionsopnameoffsetRuntimeError)r  instru   ru   rv   first_real_inst_idx  s   

r  c                   @  s   e Zd ZdddddZdS )ConvertFrameProtocolr   rQ  r   rk   rS  rT  rU  rJ   rV  rW  rR  r   rx   rX  c                C  s   d S rz   ru   )r   r   rS  rU  rV  rR  ru   ru   rv   rx    s   	zConvertFrameProtocol.__call__Nrz  )rr   rs   rt   rx  ru   ru   ru   rv   r    s    r  c                   @  s    e Zd ZdddZdddZdS )CatchErrorsWrappercallbackr  rU  rJ   rx   ry   c                 C  s   t ||  || _|| _d S rz   )r   r   r   rU  )r   r  rU  ru   ru   rv   r     s   
zCatchErrorsWrapper.__init__r   rk   rS  rT  rV  rW  rX  c              	   C  s  |d usJ t |j}tjdkr|jt|jk}n|jt|jk}|s5|s5tjs5t	ddrat
| jddsattjr_|r@d}nt |jrId}n
t	ddrQd}nd}td	|jj||jj d S |jjd
kro|jjdkrod S t dkrt }|rt0 ddlm} ||j| jjd}	t| jdsJ d| j|	j}
|
||| j|W  d    S 1 sw   Y  t. t  | j||| j|ddW  d    W  d    S 1 sw   Y  W d    d S 1 sw   Y  d S )N)r|     F)include_infra_modesrH  ztraced frame alreadyzin skipfileszSnon-infra torch dispatch mode present, this is not supported today in torch.compilezdynamo tracing is disabledz#skipping: %s (reason: %s, file: %s)r[  __new__ddp_optimizerr   )DDPOptimizer)bucket_bytes_capbackend_compile_fnrO  zJDDPOptimizer only supports callback fns that know how to clone themselves.r)   rQ  )r,   r   r   r  r  f_lastir  r*   r%  r&   r   r   r   r  r  r  r   r   r   _get_optimize_ddp_moder$   _get_active_ddp_modulecompile_lock"torch._dynamo.backends.distributedr  r  r   rO  
compile_fnrU  r%   )r   r   rS  rV  is_skipfilehas_started_executionskip_reason
ddp_moduler  r  hijacked_callbackru   ru   rv   rx     sv   



 RzCatchErrorsWrapper.__call__N)r  r  rU  rJ   rx   ry   )r   rk   rS  rT  rV  rW  rx   rX  )rr   rs   rt   r   rx  ru   ru   ru   rv   r    s    
r  r  c                 C  rp  rz   )r  )r  rU  ru   ru   rv   catch_errors_wrapperi  s   
r  rz   )r   r   r   r   r   r   rx   r   )r   r   rx   r   )r   rk   rx   r   )NF)
r  r  r  r   r   r  r  r   rx   ry   )r  r   rx   r   ry  )
rC  rg   rD  r   r  r   rE  rF  rx   rB  )r  ri   rx   r~  )NN)"r  r   r   r  r  r  r  r  ru  r  rC  rg   rD  r   r  r   rE  rF  rU  rJ   rS  rT  ra  r4   r   r  rV  r  r`  r   rR  r   rx   rX  )rC  rg   rU  rJ   rx   rX  )r  r   rx   ry   r  )r  r  rU  rJ   rx   r  )
__future__r   collectionsr   r   r  r   r  rK  r  rg  r  r   r,  r  	threadingr#  r2  typingr  r   pathlibr   typesr   r   r   r   r   r	   r
   r   r   r   r   r   r   typing_extensionsr   r   r   torch._loggingtorch._C._dynamo.guardsr   torch._dynamo.distributedr   torch._dynamo.symbolic_convertr   torch._guardsr   r   r   r   r   torch._utils_internalr   r   r   r   torch.fx._lazy_graph_moduler   %torch.fx.experimental.symbolic_shapesr    r!   torch.fx.graph_moduler"   r   torch.monitorr#   torch.nn.parallel.distributedr$   torch.utils._python_dispatchr%   r&   torch.utils._tracebackr'   r(   r  r*   r+   r,   bytecode_analysisr-   r.   bytecode_transformationr/   r0   r1   r2   r3   ra  r4   r5   r6   r7   ri  r8   r9   r:   r;   r<   r=   r>   r?   r@   rA   rB   rC   rD   rE   rF   r   rG   rH   rI   rU  rJ   pgorK   replay_recordrL   resume_executionrM   symbolic_convertrN   rO   rP   rQ   rR   r@  rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r]   r^   r_   r`   ra   rb   rc   rd   variables.torch_functionre   __annotations__numpyrf   ModuleNotFoundErrorr  backends.registryrg   repro.after_dynamorh   ri   rj   rk   variables.builderrl   	getLoggerrr   r   r6  getArtifactLoggerr  rf  RLockr  ro   rp   rq   rw   re  rf  r   r   r   r   #suppress_torch_distributed_warningsr  r
  rp  Counterr  r  r  rB  rL  r}  torch.utils.hooksr~  output_graphr  r  r  rt  rX  r\  r{  r  Protocolr  r  r  ru   ru   ru   rv   <module>   s   ,4P


>Q


A 

    
m

	O