
    %Vjii                     x   d 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 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mZ ddlmZ dZd	Zd
ZdZdZdZ dZ!dZ"dZ#dZ$dZ%dZ&de'e         ddfdZ(	 d-dede
f         de)dee*e+e
f                  de,e'e         e'e         f         fdZ-de
de
ddfdZ. ej/        d            G d! d"                      Z0ej/         G d# d$                      Z1d%ee
         d&ee
         d'ee
e
ge2f         de'e
         fd(Z3d)e,e
df         de'e         fd*Z4 G d+ d,          Z5dS ).a  
This module provides functionality for resuming Python execution at specific points in code,
primarily used by PyTorch Dynamo for control flow handling and optimization. It implements
bytecode transformation and execution state management to enable:

- Resuming execution at arbitrary points in Python bytecode
- Managing context managers and their state across execution boundaries
- Transforming and generating new code objects with preserved execution state
- Supporting Python 3.11+ exception handling and block management
- Restoring torch function mode stacks and other execution context

The module is critical for PyTorch Dynamo's ability to optimize code while preserving
Python semantics and execution state.
    N)Iterable)AbstractContextManager)AnyCallablecastOptional   )
add_push_nullbytecode_from_templatecreate_call_functioncreate_instructioncreate_jump_absolutecreate_load_constInstructionoverwrite_instructiontransform_code_object	unique_id)ExactWeakKeyDictionary                @         i   torch_dynamo_resume_in__is_tracing_resume_prologueinstsreturnc                     t           j        dk    rX|                     t          d                     t           j        dk     r(|                     t          dd                     d S d S d S )N      	PUSH_NULL)r#      SWAPr   arg)sysversion_infoappendr   )r   s    X/root/voice-cloning/.venv/lib/python3.11/site-packages/torch/_dynamo/resume_execution.py_initial_push_nullr.   8   sm    
7""'44555g%%LL+F:::;;;;; #"%%    template.stack_indexvarname_mapc                    t          | |          }|                    t          d                     |D ]}|j        r|j        xj        |z  c_        t          d t          |          D             d          \  }}||J t          |t          d          g           ||dz            j        dk    sJ t          ||dz            t          d          g           |d |dz            ||dz   d          fS )Nr2   POP_TOPc              3   P   K   | ]!\  }}|j         d k    |j        dk    ||fV  "dS )	LOAD_FASTdummyNopnameargval).0iinsts      r-   	<genexpr>z5_bytecode_from_template_with_split.<locals>.<genexpr>O   sO       	
 	
4{k))dkW.D.D I.D.D.D.D	
 	
r/   )NNNOPr	   )	r   r,   r   exn_tab_entrydepthnext	enumerater   r:   )r0   r1   r2   template_coder>   	dummy_idx
dummy_insts          r-   "_bytecode_from_template_with_splitrH   @   s@   
 +8MMMM+I66777  4 4 	4$$3$$ !	
 	
$]33	
 	
 	

 	 Iz  Z%;%;; *'9%'@'@&ABBB Q'.);;;;-	A69KE9R9R8STTT9q=)=Q+IIIr/   r8   stack_var_namec                 L    	 |  d S #  t                               |            xY wN)$__import_torch_dot__dynamo_dot_utilsset_torch_function_mode_stack)r8   rI   s     r-   _try_except_tf_mode_templaterN   b   s9    ,JJ	
 	
 	
 	s    #T)frozenc            	          e Zd ZU eed<   dZeeedf                  ed<   de	e
ef         dee         dee         fdZde	e
ef         dee         dee         fd	Zde	e
ef         dee         deee         ee         f         fd
ZdS )ReenterWithr1   N.target_valuescode_optionscleanupr    c                 x    ddl m} t          t          | j        d |            i          \  }}||z   |dd<   |S )z
        Codegen based off of:
        try:
            (rest)
        except:
            (restore previous tf mode stack)
            raise
        r	   )get_prev_stack_var_namerI   r4   N)variables.torch_functionrV   rH   rN   r1   )selfrS   rT   rV   setup_try_exceptepilogues         r-   try_except_torch_function_modez*ReenterWith.try_except_torch_function_modet   sf     	FEEEEE%G()+B+B+D+DE&
 &
 &
"(
 '
r/   c                    g }| j         rd | j         D             }t          d| j                   }||d         vr|dxx         |fz  cc<   dD ]}||d         vr|dxx         |fz  cc<   g }t          |           |                    g |t          t          |          d          t          d|                     d	t          t                   d
t          ddfd}t          || j        d	|i          \  }}	|	|z   |dd<   ||z   S )z
        Codegen based off of:
        load args
        enter context
        try:
            (rest)
        finally:
            exit context
        c                 ,    g | ]}t          |          S  r   r<   vals     r-   
<listcomp>z+ReenterWith.try_finally.<locals>.<listcomp>   !    NNNC*3//NNNr/   ___context_manager_co_varnames	__enter____exit__co_namesF
STORE_FASTr;   ctxr8   r    Nc                     |                                   	 | |                     d d d            d S # |                     d d d            w xY wrK   rf   rl   r8   s     r-   	_templatez*ReenterWith.try_finally.<locals>._template   sQ    MMOOO/T4.....T4....s	   1 A
r4   )rR   r   r1   r.   extendr   lenr   r   r   rH   )
rX   rS   rT   	load_argsctx_namename
create_ctxro   setup_try_finallyrZ   s
             r-   try_finallyzReenterWith.try_finally   s    	 	ONN4;MNNNIE43CEEFF<666'''H;6'''- 	4 	4D<
333Z(((TG3((((*
:&&&%c)nne<< #<AAA	
 	
 	
	/1#6 	/s 	/t 	/ 	/ 	/ 	/ 'It'eX5F'
 '
 '
#8 '
---r/   c                 $   g }| j         rd | j         D             }g }t          |           |                    g |t          t	          |          d                     dt
          t                   dt          ddfd}t          || j                  \  }}||z   |dd<   t          d |D             d          }|J t          |t          d	          g           d
 |D             }	t          |	d          }
t          |	d          J ||z   |
fS )zR
        Codegen based off of:
        with ctx(args):
            (rest)
        c                 ,    g | ]}t          |          S r^   r_   r`   s     r-   rb   z(ReenterWith.__call__.<locals>.<listcomp>   rc   r/   Frl   r8   r    Nc                 @    | 5  | d d d            d S # 1 swxY w Y   d S rK   r^   rn   s     r-   ro   z'ReenterWith.__call__.<locals>._template   sq                       s   c              3   F   K   | ]}|j         d k    |j        dk    |V  dS )r7   rl   Nr9   r<   r>   s     r-   r?   z'ReenterWith.__call__.<locals>.<genexpr>   sG        ;+--$+2F2F 2F2F2F2F r/   r@   c              3   0   K   | ]}|j         d k    |V  dS )PUSH_EXC_INFON)r:   r|   s     r-   r?   z'ReenterWith.__call__.<locals>.<genexpr>   s7       
 
)G)GD)G)G)G)G
 
r/   )rR   r.   rp   r   rq   r   r   rH   r1   rC   r   r   )rX   rS   rT   rr   ru   ro   
setup_withrZ   load_fast_ctx_instpush_exc_info_genpush_exc_info_insts              r-   __call__zReenterWith.__call__   s    	 	ONN4;MNNNI(*
:&&&%c)nne<<	
 	
 	
	1#6 	s 	t 	 	 	 	  Bt' 
  

H '
! &  
 
 
 "---03Ee3L3L2MNNN
 
%
 
 
 ""3T::%t,,444J&(:::r/   )__name__
__module____qualname__int__annotations__rR   r   tupler   dictstrlistr   r[   rw   r   r^   r/   r-   rQ   rQ   o   s	        /3M8E#s(O,333  cN 59+5F 	k	       0,. cN,.59+5F,.	k	,. ,. ,. ,.\3; cN3;59+5F3;	tK (;"77	83; 3; 3; 3; 3; 3;r/   rQ   c                       e Zd ZU ej        ed<    ej        e          Z	ee
         ed<    ej        e          Zee         ed<    ej        e          Zeeeeef         f         ed<   dS )ResumeFunctionMetadatacode)default_factoryinstructions prefix_block_target_offset_remapblock_target_offset_remapN)r   r   r   typesCodeTyper   dataclassesfieldr   r   r   r   r   r   r   r^   r/   r-   r   r      s         
.&7k&7&M&M&ML${#MMM 3D+2C3 3 3$d3i    <M;;L< < <tCc3h$78     r/   r   l1l2condc                     t          |          }g }	 t          |          }| D ]2} |||          r$|                    |           t          |          }3n# t          $ r Y nw xY w|S )z
    Two-pointer conditional filter.
    e.g. _filter_iter(insts, sorted_offsets, lambda i, o: i.offset == o)
    returns the instructions with offsets in sorted_offsets
    )iterrC   r,   StopIteration)r   r   r   itrescurra   s          r-   _filter_iterr     s     
bBC2hh 	 	CtC~~ 

32hh	    Js   AA 
A%$A%tupc                     g }t          |           |                    d | D                        |                    t          t          |           d                     |S )Nc              3   4   K   | ]}t          |          V  d S rK   r_   r`   s     r-   r?   z'_load_tuple_and_call.<locals>.<genexpr>  s+      77C"3''777777r/   F)r.   rp   r   rq   )r   r   s     r-   _load_tuple_and_callr     s^    !Eu	LL773777777	LL%c#hh66777Lr/   c                   (   e Zd Z e            Z e            Zedej        de	de
dej        fd            Zedej        de	de	dee	df         d	e	d
eedf         deedf         deedf         deee	ee
df         f         df         deeeee
df         f         df         dee	df         deej                 dej        fd            Zedeee
f         dee         fd            Zedej        de	de	dee	df         de
dej        fd            ZdS )ContinueExecutionCacher   linenokeyr    c                     || j         vr
i | j         |<   t          |          }|| j         |         vr | j        ||g|R  | j         |         |<   | j         |         |         S rK   )cacher   generate)clsr   r   r   s       r-   lookupzContinueExecutionCache.lookup%  sm    sy   CIdOCjjcio%%#/3<f#Cs#C#C#CCIdOC ys##r/   offsetsetup_fn_target_offsets.nstackargnamesargnames_null	setup_fnsstack_ctx_varsargnames_ctx_vars
null_idxesnested_code_objsc                    	
 J |j         t          t          z  t          z  t          z  z  rJ |j         t
          z  sJ |t          j        v r                      |	
          S t          j
        dk    t          |          dt          t                   dt          t          t           f         dd f
 	fd}t#          ||          \  }}t          j        |<   |S )Nr"   r   rS   r    c           
         t          j        |            _        ddgd t          "          D             z                      fdD                        t          |d         pg           t          |d         pg           z   }t          t          |                    }t           d|d          d	 |d<   rk|d
                             dd          }t          |          dk    r|d         |d
<   n/t          |          dk    sJ |\  }}| dt           d| d	 |d
<   |d<   d|d<   ||d<   t                    |d<   d|d<   d|d<   t          fdD             z   fd|d         D             z   t          gz             |d<   |d         t          t          z   z  |d<   t          $fd| D                       }g }rW|r1|                    t          dt          |                               |                    t          dd                     |                    t          dd          t          d t                    g           g }d! &D             }	%fd"t!          &          D             }
d# | D             }i }d}t#          '          }t          "          D ]<}|t          #          k     rX#|         ||z   k    rI|                    t          d$                     |dz  }|t          #          k     r#|         ||z   k    I|                    t          d%d&|                      ||	v rw|	                    |          } |||          \  }}|                    |           r<|
                    |          }||         } j                            |           |||<   ||v r(|                    t)          ||                              >r&t+          t-           j                             _        |	rJ D ]o\  }}|                    t          d%|                     |                    t)          |                     |                    t          d |                     prPt.          j        d'k    sJ D ];}|vsJ |                    t          d$          t          d |          g           <!r|                    g t3          t          d%d          t          dd(          t          d)          g          t          d%d          t          dd(          t          d*          t          d%d          t          d%d          t          d+d          t          d%d          t          dd(          t          d)          t          d,d          t          d%d          t          dd(          t          d*          t          d-d          t          d-d          t          dd.          t          d t                    t          d/d                     n:|                    t          dd.          t          d t                    g           |                    t5          |                     | D ]2}|j        |j        k    r nd |_        t.          j        d0k    rd |_        3|r=|                    |           |                                        |                     |r:sJ | D ]3}|j        r*|j        j         |v r||j        j                  |j        _         4|| z   | d d <   d S )1N__nested_resume_fns__nested_frame_valuesc                     g | ]}d | S )___stackr^   )r<   r=   s     r-   rb   zCContinueExecutionCache.generate.<locals>.update.<locals>.<listcomp>_  s    ;;;^^^;;;r/   c              3   $   K   | ]
}|v|V  d S rK   r^   r<   vargss     r-   r?   zBContinueExecutionCache.generate.<locals>.update.<locals>.<genexpr>`  s'      ==aq}}}}}}==r/   co_cellvarsco_freevars_co_name_at_co_qualname.r	   )maxsplitr   co_firstlinenor^   co_argcountr   co_posonlyargcountco_kwonlyargcountc                     g | ]}|v|	S r^   r^   r   s     r-   rb   zCContinueExecutionCache.generate.<locals>.update.<locals>.<listcomp>z  s    ===q}}1}}}r/   c                     g | ]}|v|	S r^   r^   r   s     r-   rb   zCContinueExecutionCache.generate.<locals>.update.<locals>.<listcomp>{  s    KKKQd]]1]]]r/   re   co_flagsc              3   2   K   | ]}|j         k    |V  d S rK   r   r<   r=   r   s     r-   r?   zBContinueExecutionCache.generate.<locals>.update.<locals>.<genexpr>  s/      HHQX5G5G!5G5G5G5GHHr/   COPY_FREE_VARSr(   RESUME
LOAD_CONSTTrk   rj   c                     i | ]
}|j         |S r^   r1   )r<   fns     r-   
<dictcomp>zCContinueExecutionCache.generate.<locals>.update.<locals>.<dictcomp>  s    <<<BR^R<<<r/   c                 2    i | ]\  }}|j         |         S r^   r   )r<   r=   r   r   s      r-   r   zCContinueExecutionCache.generate.<locals>.update.<locals>.<dictcomp>  s7     # # #Ar  7 :# # #r/   c                     i | ]
}|j         |S r^   r   r|   s     r-   r   zCContinueExecutionCache.generate.<locals>.update.<locals>.<dictcomp>  s    IIIDdk4IIIr/   r%   r7   r   )r#      BINARY_SUBSCRDELETE_SUBSCR
BUILD_LISTLIST_EXTENDDELETE_FASTFCALL_FUNCTION_EXr"   )!copydeepcopyr   rangerp   r   sortedTORCH_DYNAMO_RESUME_IN_PREFIXrsplitrq   "IS_TRACING_RESUME_PROLOGUE_VARNAME
CO_VARARGSCO_VARKEYWORDSrC   r,   r   rD   r   popr   r   r   reversedr*   r+   r
   r   r   starts_line	positionsunreachable_codesrA   target)(r   rS   freevarsqualified_pathmodule_namer   r   prefixrT   hookshook_target_offsetsoffset_to_instold_hook_target_remapnull_idxes_istack_ctx_vars_dr=   hook
hook_insts
exn_targethook_target_offsetold_hook_targetrt   valsr   r>   r   r   r   r   r   is_py311_plusr   metar   r   r   r   r   r   r   s(                            @r-   updatez/ContinueExecutionCache.generate.<locals>.updateY  s	    !%l ; ;D)+BCD;;U6]];;;;DKK====8======\-8>B??%]+1rC C H VH--..H0XX<	3JXXPVXX #  	!-m!<!C!CCRS!C!T!T~&&!++2>y2IL//~..!3333+9(K&^^)F^^^^V\^^ !/ .4L)**,L'*2L'*-d))L'12L-.01L,-*/====m===>KKKKl=9KKKL 667+ +L' (4J'?^+C (L$ HHHH\HHHHHFF C MM*+;XOOO   0qAAABBB
 MM&|DAAA&$-O     *,G<<)<<<E# # # #&y11# # # JILIIIN$&!L#N336]] M M 3z??22"<0A4DDDMM"4["A"ABBB A%L	 !3z??22"<0A4DDD 0^PQ^^TTTUUU:: 99Q<<D-1T,-H-H*J
MM*---$ L-@-D-DQ-G-G**89K*L=DDEWXXXAK-o>((( MM"67G7J"K"KLLL  9=TBCC9 95 9 0 M M
d0TJJJKKK24889990dKKKLLLL  	'72222&  AD====MM.{;;.|AFFF      ;.& 2$/8M!" !" !" !3< K K K 2? C C . +;?TUUU. +<CCC. +?;;.  +;?TUUU!." +;?VWWW#.$ +<Q???%.( +;?VWWW).* +<CCC+., +?;;-.8 +=a@@@9.< +;?VWWW=.> +<CCC?.@ +?;;A.D +=AVWWWE.F +)2I  G.P +<FFFQ.R +(1S  S.Z ++=1EEE[.0 0 0 0f *<FFF*(1S     MM.v66777
 % * *;&-//E#' #w..%)DN Cg&&&c33LAABBB % 	$$$}(  D* .59NNN4I .55*1
 %|3LOOOr/   )r   CO_GENERATORCO_COROUTINECO_ITERABLE_COROUTINECO_ASYNC_GENERATORCO_OPTIMIZEDr   generated_code_metadata&generate_based_on_original_code_objectr*   r+   r   r   r   r   r   r   r   )r   r   r   r   r   r   r   r   r   r   r   r   r   r  new_coder   r  r  s   ` ```````````   @@r-   r   zContinueExecutionCache.generate.  s   $ !!!Ml*-BBEWWY	
 	
 
 }|++++)AAA=='!    (G3%d++V	4{+V	4;?S>V	4V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4 V	4p ,D&99!CG6x@r/   rS   c                 B    t          d          t          dd          gS )zACodegen a `raise None` to make analysis work for unreachable codeNRAISE_VARARGSr	   r(   )r   r   )rS   s    r-   r   z(ContinueExecutionCache.unreachable_codes5  s)     d##A666
 	
r/   r   c                   	
 t           j        |         	d
dt          t                   dt          t
          t          f         ddf	
fd}t          ||           
dk    sJ t          j	        dk    rv
	j
        vrQi x	j
        
<   dt          t                   dt          t
          t          f         ddf	fd	}t          ||           t          	
fd
D                       t          j        	j        |
g|R  S )a>  
        This handles the case of generating a resume into code generated
        to resume something else.  We want to always generate starting
        from the original code object so that if control flow paths
        converge we only generated 1 resume function (rather than 2^n
        resume functions).
        r   r   rS   r    Nc                     fd| D             \  fdt          t          |           t          j                            D             \  }j        |j        k    sJ |j        J |j        d S )Nc              3   2   K   | ]}|j         k    |V  d S rK   r   r   s     r-   r?   ziContinueExecutionCache.generate_based_on_original_code_object.<locals>.find_new_offset.<locals>.<genexpr>W  s/      GGqAH4F4F4F4F4F4FGGr/   c              3   *   K   | ]\  }}|u 	|V  d S rK   r^   )r<   i1i2r   s      r-   r?   ziContinueExecutionCache.generate_based_on_original_code_object.<locals>.find_new_offset.<locals>.<genexpr>Y  s9        B<< <<< r/   )zipr   r   opcoder   )r   rS   
new_targetr   r  
new_offsetr   s      @r-   find_new_offsetzVContinueExecutionCache.generate_based_on_original_code_object.<locals>.find_new_offsetS  s     HGGGLGGGIV   !(<"8"8(4CT:U:UVV  MZ
 =J$55555$000#*JJJr/   r   r"   c                   
 g }| D ]I}t          |          t          j                  k    r n!|j        dk    r|                    |           Jt	          |j                  D ]"\  }}|t          t          |j                  <   #|r t          t          |d         j                  nd
t          
fdD                       }t          | |d           }t          t	          t          |           t          j                            |d           }t	          ||          D ]\  }}	|d         j        |	j        <   d S )Nr~   r   c              3   (   K   | ]}|k    |V  d S rK   r^   )r<   nold_start_offsets     r-   r?   zmContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsets.<locals>.<genexpr>  s7       . .aBR>R>R>R>R>R>R. .r/   c                     | j         |k    S rK   r   )r>   os     r-   <lambda>zlContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsets.<locals>.<lambda>  s    WXHX r/   c                     | d         |u S )Nr   r^   )v1v2s     r-   r#  zlContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsets.<locals>.<lambda>  s    r!u{ r/   r	   )rq   r   r:   r,   r  r   r   r   r   r   r   r   )r   rS   prefix_blocksr>   r"  old_inst_offsetstargetsnew_targetsnewoldr   r   r  r   s             @r-   remap_block_offsetszZContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsetso  s    8:M , 7 7}-- A2 2   "E;/99)00666 $'%t'L$ $ N Na MN1$sDK2H2HII @MTS-"3":;;;RT % (. . . . .#:. . . ( ($ +$&68X8X G #/H\22HT=N4O4OPP22# #K
 %(W$=$= N NS@CA1#*==N Nr/   c              3   >   K   | ]}j                  |         V  d S rK   )r   )r<   r  r  r  s     r-   r?   zPContinueExecutionCache.generate_based_on_original_code_object.<locals>.<genexpr>  sD       , , .z:1=, , , , , ,r/   )r   r  r   r   r   r   r   r   r*   r+   r   r   r   r   )r   r   r   r   r   r   r  r-  r   r  r  s      ``   @@@r-   r  z=ContinueExecutionCache.generate_based_on_original_code_object=  s   " (>'U(
 
	+{+	+;?S>	+	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	dO444Qw&&
 !??? )D,J-(N"&{"3(NCGS>(N(N (N (N (N (N (N (N (NT &d,?@@@ ', , , , , ,0, , , ' '# &,Ivz+B
EI
 
 
 	
r/   N)r   r   r   r   r   r  classmethodr   r   r   r   r   r   r   rQ   r   staticmethodr   r   r   r   r  r^   r/   r-   r   r   !  s;       ""$$E4466$%. $# $S $U^ $ $ $ [$ DnD D 	D
 "'sCxD D S/D S#XD c)*D eCsCx$893>?D !sE#s(O';!<c!ABD #s(OD  /D  
!D D D [DL 
S#X 
4;L 
 
 
 \
 d
nd
 d
 	d

 "'sCxd
 d
 
d
 d
 d
 [d
 d
 d
r/   r   rK   )6__doc__r   r   r*   r   collections.abcr   
contextlibr   typingr   r   r   r   bytecode_transformationr
   r   r   r   r   r   r   r   r   r   utilsr   r  CO_NEWLOCALSr   r   	CO_NESTEDr  	CO_NOFREEr	  r
  r  r   r   r   r.   r   r   r   r   rH   rN   	dataclassrQ   r   boolr   r   r   r^   r/   r-   <module>r<     sY         



  $ $ $ $ $ $ - - - - - - 0 0 0 0 0 0 0 0 0 0 0 0                        * ) ) ) ) ) 
		   !9 %C "<d;/ <D < < < < -1J JsCx JJ $sCx.)J 4d;//0	J J J JD
 
S 
T 
 
 
 
 d###}; }; }; }; }; }; }; $#};@        " C:t#
$ 
#Y	   .eCHo ${2C    A
 A
 A
 A
 A
 A
 A
 A
 A
 A
r/   