
    ɯwgn                     D   d dl Z d dlZd dlZd dlZd dlmZmZ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# ejH                  d       G d d             Z%ejH                   G d d             Z&d Z'd Z( G d d      Z)y)    N)AnycastDictListOptionalTuple   )
create_call_functioncreate_call_methodcreate_dup_topcreate_instructioncreate_jump_absolutecreate_load_methodInstructionInstructionExnTabEntrytransform_code_object	unique_id)ExactWeakKeyDictionary                @         i   torch_dynamo_resume_inc                     t         j                  dk\  rK| j                  t        d             t         j                  dk  r| j                  t        dd             y y y )N      	PUSH_NULL)r       SWAPr   arg)sysversion_infoappendr   )instss    c/home/mcse/projects/flask/flask-venv/lib/python3.12/site-packages/torch/_dynamo/resume_execution.py_initial_push_nullr,   '   sN    
7"'45g%LL+F:; & #    T)frozenc                   N    e Zd ZU eed<   dZeeedf      ed<   de	e
   fdZd Zy)ReenterWithstack_indexN.target_valuescleanupc                 J   g }| j                   r$| j                   D cg c]  }t        d|       }}t        d| j                         |d   vr|dxx   fz  cc<   dD ]  }||d   vs|dxx   |fz  cc<    t        t        j
                  dk  rdnd	      }t        d      }g }t        |       |j                  g |t        t        |      d
      t        d      t        d      t        d      t        d      t        d             t        j
                  dk  r|j                  t        d|             nHt        d      }	t        d      }
t        |	|
|| j                  dz   d
      |	_        |j                  |	       fd}t        j
                  dk  r*t        d      t        d      |g |       t        d      }nt        j
                  dk  r7t        d      g |       t        d|      | |       t        d      |}nyt        dd      }t        dd      }t        |||| j                  dz   d      |_        
g |       t        d|      | |       ||t        d      t        dd      |}||z   |d d  |S c c}w )!z
        Codegen based off of:
        load args
        enter context
        try:
            (rest)
        finally:
            exit context
        
LOAD_CONSTargval___context_manager_co_varnames)	__enter____exit__co_namesr   NOPPUSH_EXC_INFOF
STORE_FAST	LOAD_FASTr:   r   POP_TOPSETUP_FINALLYtargetr	   c                      t        d       t        d      t        dd       t               t               gt        d      t        d      S )Nr@   r6   r;   r5   r    rA   )r   r   r   r   )ctx_names   r+   create_resetz,ReenterWith.try_except.<locals>.create_resetr   sS    ";x@":."<=   $A& #9- r-   r    	   	POP_BLOCKBEGIN_FINALLYEND_FINALLYJUMP_FORWARDRERAISEr%   COPYr    r   T
POP_EXCEPTN)r2   r   r   r1   r'   r(   r,   extendr
   lenr   r   r)   r   exn_tab_entry)selfcode_optionsr3   	load_argsvalnameexcept_jump_targetcleanup_complete_jump_targetsetup_finallyexn_tab_beginexn_tab_endrG   epiloguefinally_exn_tab_endfinally_exn_tab_targetrF   s                  @r+   
try_exceptzReenterWith.try_except5   se    	  -- #<<I  243C3C2DEF<66'H;6'- 	4D<
33Z(TG3(	4 0%%/E_
 (:%'@$+-=) 	%c)ne< #<A #;x@	
 #;/ $A& #9-
	
 g%  "?;MN /u5M,U3K*@"  1$+M'   /		 f$";/"?3" 	
 #=1H '";/ #>:VW #	
  #9- -H #5YA"F%7A%F"/E"#&  1$0,  #>:VW #	
  $ ' #<0 #9!4 -H '
Ks   J c                    g }| j                   r$| j                   D cg c]  }t        d|       }}t        j                  dk  rrt        d      }t        d      }t        d      ||t        d      t        d      g|z   |d	d	 g |t        d
t	        |            t        d|      t        d      d	fS t        j                  dk  rt        d      }t        d      }t        d      }	t        d      t        dd	      t        d      t        d      t        d
d      t        d      t        d|	      |t        d|      t        d      |t        d      t        d      t        d      t        d      |	g|z   |d	d	 g |t        d
t	        |            t        d|      t        d      d	fS t        d      }t        d      }	d }
t        d      }t        d      }t        d      }t        dd      }t        dd      }t        |||| j                  dz   d      |_        t        |||| j                  dz   d      |_        t        |||| j                  dz   d      |_        | |
        |
        |
       gt        dd      t        d      t        d|	      |t        d      t        t        j                  dk  rd nd|      ||t        d      t        dd      |t        d      t        d      t        d      |	|z   |d	d	 g }t        |       |j                  g |t        t	        |      d      t        d!      |       ||fS c c}w )"zR
        Codegen based off of:
        with ctx(args):
            (rest)
        r5   r6   rH   WITH_CLEANUP_STARTrK   rJ   WITH_CLEANUP_FINISHrL   NCALL_FUNCTIONr%   
SETUP_WITHrC   rA   r   WITH_EXCEPT_STARTr=   DUP_TOPr    rM   POP_JUMP_IF_TRUErN   rP   c                      t        dd       S )Nr5   r6   r    r-   r+   create_load_nonez.ReenterWith.__call__.<locals>.create_load_none   s    ),tDDr-   r>   r   rO   r	   TFr       POP_JUMP_FORWARD_IF_TRUEBEFORE_WITH)r2   r   r'   r(   rR   r   r1   rS   r
   r,   rQ   )rT   rU   r3   rV   rW   with_cleanup_startbegin_finallywith_except_startpop_top_after_with_except_startrZ   rm   exn_tab_1_beginexn_tab_1_endexn_tab_1_targetexn_tab_2_endexn_tab_2_targetrets                    r+   __call__zReenterWith.__call__   sT    	  -- #<<I  f$!34H!I.?M";/""#89"=1 GAJ"?IG #<8JK #9-	
   ' 23F G.@.K++=e+D( #;/"<="9-"9-"?:"9-">:VW!"&/N #9-/"9-"9-"<0"9-,%& 'GAJ*"?IG #<8IJ #9-	
   /A.K++=e+D(E 1;O.u5M1/B.ya@M1&a@,B   1$-O) .D    1$.* =S//   1$=+9  " " "	
 &a/ #9- #>:VW ! ##67 #'''1 /+:	  !" !#$ #<0%& #9!4'( 0)* #<0+, #9--. #9-/0 -12 3GAJ6 &(Cs#JJ)#i.%@ '}5 $	 (((}s   M!)__name__
__module____qualname__int__annotations__r2   r   r   r   r   r   ra   r|   rl   r-   r+   r0   r0   .   s9    /3M8E#s(O,3r[0A rhG)r-   r0   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<   dZeeeef      ed<   y)ResumeFunctionMetadatacode)default_factoryinstructions prefix_block_target_offset_remapNblock_target_offset_remap)r}   r~   r   typesCodeTyper   dataclassesfieldlistr   r   r   r   r   r   r   r   rl   r-   r+   r   r   3  sf    
..&7k&7&7&ML${#M 3D+2C2C3$d3i  ;?xS#X7>r-   r   c                     t        |      }g }	 t        |      }| D ](  } |||      s|j                  |       t        |      }* 	 |S # t        $ r Y |S w xY w)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
    )iternextr)   StopIteration)l1l2conditrescurrW   s          r+   _filter_iterr   C  sl     
bBC2h 	CC~

32h	 J  Js   A
 A
 
	AAc                     g }t        |       | D ]  }|j                  t        d|               |j                  t	        t        |       d             |S )Nr5   r6   F)r,   r)   r   rQ   r
   rR   )tupr*   rW   s      r+   _load_tuple_and_callr   V  sQ    !Eu C'SABC	LL%c#h67Lr-   c                      e Zd Z e       Z e       Zed        Zedede	e   dede	e
   de	e
   de	e   de	ee	e   f   d	e	e
e	e   f   d
e	e   dej                  fd       Zedee   fd       Zedede	edf   fd       Zy)ContinueExecutionCachec                     || j                   vri | j                   |<   t        |      }|| j                   |   vr$ | j                  ||g| | j                   |   |<   | j                   |   |   S N)cachetuplegenerate)clsr   linenokeys       r+   lookupzContinueExecutionCache.lookupc  sk    syy  CIIdOCjciio%#/3<<f#Cs#CCIIdOC yys##r-   offsetsetup_fn_target_offsetsnstackargnamesargnames_null	setup_fnsstack_ctx_varsargnames_ctx_vars
null_idxesreturnc                    	
 J |j                   t        t        z  t        z  t        z  z  rJ |j                   t
        z  sJ |t        j                  v r j                  |	
      S t        j                  dk\  t        |      dt        t           dt        t        t         f   f
 	fd}t#        ||      }t        j                  |<   |S )Nr   r   rU   c                 l
   t        j                  |       "_        t        #      D cg c]  }d| 	 c}j	                  fdD               t        |d   xs g       t        |d   xs g       z   }t        t        |            }t         d|d    d! |d<    rV|d   j                  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        D cg c]	  }|vs| c}z   |d   D cg c]	  }|vs| c}z         |d<   |d   t        t        z   z  |d<   t        %fd| D              }g }	 rC|r%|	j                  t        dt        |                   |	j                  t        dd             g }
'D ci c]  }|j                  | }}t!        '      D ci c]  \  }}|j                  &|    }}}| D ci c]  }|j"                  | }}i }d}t%        (      }t        #      D ]  }|t        $      k  rD$|   ||z   k(  r9|	j                  t        d             |d
z  }|t        $      k  r$|   ||z   k(  r9|	j                  t        dd|              ||v rf|j'                  |      } |||
      \  }}|	j	                  |        r6|j'                  |      }||   }"j(                  j                  |       |||<   ||z   }||v s|	j	                  t+        ||                  r#t-        t/        "j(                              "_        |rJ D ]W  \  }}|	j                  t        d|             |	j	                  t+        |             |	j                  t        d|             Y rIt0        j2                  dk\  sJ D ]/  }|vsJ |	j	                  t        d      t        d|      g       1 |	j                  t5        |             | D ]?  }|j"                  |j"                  k(  r n$d |_        t0        j2                  dk\  s9d |_        A |
r1|	j	                  |
       |	j	                  j;                  |             |rY sJ | D ]P  }|j<                  s|j<                  j>                  |v s)||j<                  j>                     |j<                  _        R |	| z   | d d  y c c}w c c}w c c}w c c}w c c}}w c c}w )N___stackc              3   ,   K   | ]  }|vs|  y wr   rl   ).0vargss     r+   	<genexpr>zBContinueExecutionCache.generate.<locals>.update.<locals>.<genexpr>  s     =aq}=s   	co_cellvarsco_freevars_co_name_at_co_qualname.r	   )maxsplitr   co_firstlinenorl   co_argcountr   co_posonlyargcountco_kwonlyargcountr9   co_flagsc              3   B   K   | ]  }|j                   k(  s|  y wr   r   r   ir   s     r+   r   zBContinueExecutionCache.generate.<locals>.update.<locals>.<genexpr>  s     HQXX5G!H   COPY_FREE_VARSr%   RESUMEr"   r@   r6   r?   rn   r   ) copydeepcopyr   rangerQ   r   sortedTORCH_DYNAMO_RESUME_IN_PREFIXrsplitrR   
CO_VARARGSCO_VARKEYWORDSr   r)   r   r1   	enumerater   dictpopr   r   r   reversedr'   r(   r   starts_line	positionsunreachable_codesrS   rD   ))r   rU   r   freevarsqualified_pathmodule_namer   r   rD   prefixr3   fnhookshook_target_offsetsinstoffset_to_instold_hook_target_remapnull_idxes_istack_ctx_vars_dhook
hook_insts
exn_targethook_target_offsetold_hook_targetreal_irX   valsr   r   r   r   r   is_py311_plusr   metar   r   r   r   r   r   s)                              @r+   updatez/ContinueExecutionCache.generate.<locals>.update  s    $l ;D,1&M:qhqcN:DKK=8==\-8>B?%]+1rC H VH-.H 11<	3J2K4PVxX  !-m!<!C!CCRS!C!T~&!+2>y2IL/~.!333+9(K '-q)F(Gq	QUV\U]^ !% .4L)**,L'*2L'*-d)L'12L-.01L,-*/+=q}1=>*=9KQd]1KL+L'
 (4J'?^+C (L$ H\HHFFMM*+;XO 0qAB)+G2;<BR^^R'<E< 'y1#Ar  7 ::# # =IIDdkk4/INI$&!L#N36] R 3z?2"<0A4DDMM"4["AB A%L	 !3z?2"<0A4DD 0xPQs^TU: 99Q<D-1,-H*J
MM*-$-@-D-DQ-G**89K*L==DDEWXAK-o>\)-- MM"67G7O"PQ+R.  9=TBBC95 9 0 M
d0TJK24890dKLM ''7222& AD=(=MM.{;.|AF MM.v67
 % *;;&--/#' ##w.%)DN* g&c33LAB %$$}( D** ..559NN4I ..555**1 %|3LOM ;6 >K  =# Js/   T)	T
3T
	T!
T!
T&8T+T1)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   strr   r   )r   r   r   r   r   r   r   r   r   r   r   r   r   new_coder   r   s   ` ``````````  @@r+   r   zContinueExecutionCache.generatel  s    !!!MMl*-BBEWWY	
 
 }}|+++)AAA=='!  ((G3%d+I	4k!2 I	4$sCx. I	4 I	4 I	4V )v6CG66x@r-   c                 6    t        dd      t        dd      gS )zACodegen a `raise None` to make analysis work for unreachable coder5   Nr6   RAISE_VARARGSr	   r%   rk   )rU   s    r+   r   z(ContinueExecutionCache.unreachable_codes"  s"     |D9A6
 	
r-   .c                   	
 t         j                  |   	d
dt        t           dt        t
        t        f   f	
fd}t        ||       t        j                  dk\  r]	j                  s=i x	_
        dt        t           dt        t
        t        f   f	fd}t        ||       t        	fdD              t        j                  	j                  |
g| 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).
        Nr   rU   c                     fd| D        \  fdt        t        |       t        j                              D        \  }j                  |j                  k(  sJ |j                  y )Nc              3   B   K   | ]  }|j                   k(  s|  y wr   r   r   s     r+   r   ziContinueExecutionCache.generate_based_on_original_code_object.<locals>.find_new_offset.<locals>.<genexpr>?  s     GqAHH4FGr   c              3   0   K   | ]  \  }}|u r|  y wr   rl   )r   i1i2rD   s      r+   r   ziContinueExecutionCache.generate_based_on_original_code_object.<locals>.find_new_offset.<locals>.<genexpr>A  s$      B< s   )zipr   r   opcoder   )r   rU   
new_targetrD   r   
new_offsetr   s      @r+   find_new_offsetzVContinueExecutionCache.generate_based_on_original_code_object.<locals>.find_new_offset;  s_     HLGIV!(<"8(4CTCT:UVMZ
 ==J$5$5555#**Jr-   r   c                 f  
 g }| D ]F  }t        |      t        j                        k(  r n#|j                  dk(  s6|j                  |       H t	        |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                  <   # y )Nr>   c              3   .   K   | ]  }|kD  s	|  y wr   rl   )r   nold_start_offsets     r+   r   zmContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsets.<locals>.<genexpr>o  s      .aBR>R.s   
c                      | j                   |k(  S r   r   )r   os     r+   <lambda>zlContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsets.<locals>.<lambda>s  s    WXHX r-   c                     | d   |u S )Nr   rl   )v1v2s     r+   r  zlContinueExecutionCache.generate_based_on_original_code_object.<locals>.remap_block_offsets.<locals>.<lambda>x  s    r!u{ r-   r	   )rR   r   opnamer)   r  r   r   r   r   r   r   r   )r   rU   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_offsetsS  sI    8:M , 7}- AA2  ";;/9)0067 $'%t'L'L$ Na MN1$sDKK2HIN @MS-"3":":;RT % (. .#:. ($ +$&68XG #/H\2HT=N=N4OP2#K
 %(W$= NS@CA1#**=Nr-   c              3   <   K   | ]  }j                   |     y wr   )r   )r   r  r   s     r+   r   zPContinueExecutionCache.generate_based_on_original_code_object.<locals>.<genexpr>  s!      ,67..q1,s   )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'U(
 
	+{+	+;?S>	+ 	dO4w&
 11MOO)D,J(N"&{"3(NCGS>(NT &d,?@ ', ,;R, '# &,,IIvz+B
EI
 	
r-   N)r}   r~   r   r   r   r   classmethodr   r   r   r   r0   r   r   r   r   staticmethodr   r   r   r   rl   r-   r+   r   r   _  s-   "$E46$ $ s 	s
 "'ss s *s Szs %s c5:o.s !eCj1s #Js 
s sj 
4+< 
 
 Z
#&Z
AFsCxZ
 Z
r-   r   )*r   r   r'   r   typingr   r   r   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,   	dataclassr0   r   r   r   r   rl   r-   r+   <module>r%     s      
  9 9   * 
		   !9 < d#A) A) $A)H ? ? ?&f
 f
R	Hr-   