Cj8 4 dZddlZddlZddlZddlZddlZddlZddlZddlZddl m Z m Z m Z m Z mZddlmZmZmZmZddlmZddlmZmZmZmZddlmZerdd lmZejd Gd d Z ejd GddZ!ej"dkrdej!de!fdZ#n&ej"dkrdej!de!fdZ#n dej!de!fdZ#GddZ$ej"dkr de%de&fdZ'nej"dkr de%de&fdZ'nde%de&fdZ'de$ddde%de(dzd edzd!e!dzde!f d"Z)d!e!de!fd#Z*d!e!de&fd$Z+dd%ed&e&de!fd'Z,de-e!fd(Z.de!fd)Z/d*e(de-e!fd+Z0d,e!e-e!zde-e!fd-Z1d,e!e-e!zde-e!fd.Z2d/e(d0e&de-e!fd1Z3 dd3e&d0e&d4e&de-e!fd5Z4d/e(de-e!fd6Z5de%de!fd7Z6d!e!de!fd8Z7d*e(de-e!fd9Z8d:e%de(de(fd;Z9 dde&de-e!fd?Z:de(de-e!fd@Z;dAe(de-e!fdBZd*e(de!fdFZ?dGe(de@e-e(e e(e(gdfe e(gdfffdHZAd*e(de-e(fdIZBej"dkr(dGe(de@e-e(e ejCdze(gdfffdJZDejd GdKdLZEd*e(de-e(fdMZFdNe e(de(fdOZGdPe-eEddfdQZHdReIde-eEfdSZJdPe-eEdeIfdTZKdUe-e!dVe(de@eIeIffdWZLdXeMe(e!fdYe(de!dzfdZZNdUe e!ddfd[ZOePejQZRd\e!ddfd]ZSdUe-e!d^e(de!fd_ZTdUe-e!ddfd`ZUdaeIdUe-e!ddfdbZVdUe-e!de-eEfdcZWdPe-e ddeMe!e(fddfdeZXdUe-e!ddfdfZYdUe-e!ddfdgZZdUe-e!ddfdhZ[die!dje-e!de-e!fdkZ\dUe-e!de-e!fdlZ]dUe-e!ddfdmZ^dUe-e!ddfdnZ_dodpdqdrdsZ`dUe-e!ddfdtZadUe-e!ddfduZbdUe-e!ddfdvZcdwejddUe-e!ddfdxZedUe-e!de(fdyZfdze!de(fd{ZgdUee!ddfd|ZhdUee!ddfd}Zid~eIde%fdZjdwejdddfdZkePejlZmePejnZoePejpZqePejrZsdeMe%efd%ede(fdZt ddUe-e!deMe%efde defdzddfdZudUe-e!ddfdZvejwde-e%fdZx ddwejdde e-e!eMe%efgedfde&de@ejdedffdZydUe-e!de-e%deMe%efde@e-e!ejdffdZzdUee!deddfdZ{ddwejdde&de-e!fdZ|dUee!de-e!fdZ}ejw ddwejdde&dee!fdZ~ejZdde%de&de%fdZdwejdde&fdZ dde defde eefdzde&de&de-e!f dZdS)a This module provides utilities for analyzing, transforming and manipulating Python bytecode. It includes functionality for: - Converting between different bytecode formats and versions - Virtualizing jumps and managing jump targets - Handling exception tables and their entries - Managing instruction offsets and extended arguments - Providing a clean API for bytecode modification and transformation - Supporting Python version-specific bytecode features - Generating bytecode from template functions The module is designed to work across different Python versions (3.7+) and handles version-specific bytecode differences transparently. N)CallableIterableIteratorMappingSequence)AnycastOptional TYPE_CHECKING)config) get_indexofpropagate_line_numsremove_extra_line_numsstacksize_analysis)is_safe_constant)DynamoTracerOutputT)slotsc^eZdZUded<ded<ded<eed<eed<defdZd edefd Z d S) InstructionExnTabEntry Instructionstartendtargetdepthlastireturnc d|jd|jd|jd|jd|jd S)NzInstructionExnTabEntry(start=, end=z , target=z, depth=z, lasti=))rshort_inst_reprrrrrselfs SC:\Projects\DSCons\.venv\Lib\site-packages\torch\_dynamo\bytecode_transformation.py__repr__zInstructionExnTabEntry.__repr__1s 7DJ,F,F,H,H 7 78++-- 7 7k1133 7 7Z 7 7*. 7 7 7 oct|tsdS|j|juo;|j|juo-|j|juo|j|jko|j|jkSNF) isinstancerrrrrr)r#r's r$__eq__zInstructionExnTabEntry.__eq__9ss!344 5 J!' ! &AE! & qx' & ag% & ag%  r&N) __name__ __module__ __qualname____annotations__intboolstrr%objectr+r&r$rr)s    JJJ KKK #        4       r&rceZdZUdZeed<eed<edzed<eed<dZedzed<dZ edzed<d Z e ed <dZ e d ed <dZe ded <dZedzed<dZedzed<defdZdede fdZdefdZddZdS)rz$A mutable version of dis.InstructionopcodeopnameNargargvaloffset starts_lineFis_jump_targetz dis.Positions positionsr exn_tab_entryargreprrc t|SNidr"s r$__hash__zInstruction.__hash__Ws $xxr&othercBt|t|kSrArBr#rEs r$r+zInstruction.__eq__Zs$xx2e99$$r&c(d|jd|jdS)NzInstruction(opname=z , offset=r )r7r:r"s r$r!zInstruction.short_inst_repr]sIT[II4;IIIIr&c6|j|_|j|_dSrA)r;r=rGs r$copy_positionszInstruction.copy_positions`s ,r&)rErrN)r,r-r.__doc__r0r/r2rr:r;r<r1r=r rr>rr?rDr3r+r!rJr4r&r$rrEsA.. KKK KKK tOOO KKKFC$J"Kt""" ND   +/Ix(///&*FH] #***37M)D0777GS4Z#%F%t%%%%JJJJJ))))))r&r irc ~t|j|j|j|j|j|j|j|jSrA) rr6r7r8r9r: line_numberr<r=rOs r$convert_instructionrSg= H H E H H M  K   r&rM c ~t|j|j|j|j|j|j|j|jSrA) rr6r7r8r9r:r;r<r=rRs r$rSrSurTr&c tt|j|j|j|j|j|j|jdSrA)rr6r7r8r9r:r;r<rRs r$rSrSs; H H E H H M     r&ceZdZdefdZdS) _NotProvidedrcdS)NrZr4r"s r$r%z_NotProvided.__repr__s~r&N)r,r-r.r2r%r4r&r$rZrZs/#r&rZrM namec |dvS)N) LOAD_ATTR LOAD_GLOBALLOAD_SUPER_ATTRr4r^s r$inst_has_op_bitsrdsFFFr&c|dkS)Nrar4rcs r$rdrds }$$r&cdSr)r4rcs r$rdrdsur&)r8r9rr8r9rc0t|r|td|d}n(|du|tuz|duz}|dkrtd|$t|tstdt t j|||||S)a At most one of `arg`, `argval`, and `target` can be not None/_NotProvided. This is to prevent ambiguity, e.g. does create_instruction("LOAD_CONST", 5) mean load the constant at co_consts[5], or load the constant 5? If `arg` is not provided, it will be computed during assembly from `argval` or `target`. Bits in the args of instructions LOAD_GLOBAL, LOAD_ATTR (3.12+), and LOAD_SUPER_ATTR modify the behavior of the instruction. In this case, we allow both `arg` and `argval` to be set. The value of `arg` here is expected to be the value of the op bits and the true value of `arg` will be computed during assembly. If `arg` is not set, the bits are assumed to be 0. Nz*target cannot be specified for instructionrr z@only one of arg, argval, and target can be not None/_NotProvidedz#instruction arg must be int or None)r6r7r8r9r)rd RuntimeErrorrZr*r0rdisopmap)r^r8r9rcnts r$create_instructionrls0   KLL L ;C$6#=>&PTBTU 77R  z#s33@AAA ytVF   r&cLtjdkrdnd}t||S)NrU JUMP_FORWARD JUMP_ABSOLUTErsys version_inforl)rinsts r$create_jump_absoluterus+ -88>>oD d6 2 2 22r&c|jdvS)N)rnro)r7rps r$is_jump_absoluterws == ==r&valcheckedcj|r!t|std|td|S)z In general we should only create `LOAD_CONST` for immutable objects, but sometimes it's convenient _and safe_ for Dynamo create `LOAD_CONST` for mutable objects. In such cases, use `checked=False`. zunsafe constant LOAD_CONSTr9)rAssertionErrorrl)rxrys r$create_load_constr~sF ;$$ ; !9C!9!9:: : l3 7 7 77r&cNddlm}td|tdgS)a Create instructions that trigger the bytecode debugger to stop. Usage: codegen.extend_output(create_breakpoint()) When the bytecode debugger is active, execution will pause at this point. At runtime, these instructions have no effect on program state. r BREAKPOINT_MARKERr{r|POP_TOP)bytecode_debuggerrrlrs r$create_breakpointrs>544444 <0ABBB9%% r&cbtjdkrtddStdS)NrUCOPYr r8DUP_TOPrqr4r&r$create_dup_toprs1 7""!&a0000 i ( ((r&nc|dkrgStjdkrdt|ddDS|dkrtdgd|dz zgStd | gS) a Returns a "simple" sequence of instructions that rotates TOS to the n-th position in the stack. For Python < 3.11, returns a single ROT_* instruction. If no such instruction exists, an error is raised and the caller is expected to generate an equivalent sequence of instructions. For Python >= 3.11, any rotation can be expressed as a simple sequence of swaps. r rUc0g|]}td|S)SWAPr)rl.0rOs r$ z create_rot_n.. s&KKKa"6q111KKKr&ROT_)TWOTHREEFOURROT_Nr)rrrsrangerlrs r$ create_rot_nrs Avv  7""LK5Ar??KKKKAvv"6,D,D,DQU,K#KLLMM wA . . . //r& inst_or_instsc8t|tr|gn6t|tstdt ||dt dt ffd }dt ddffd }tjdkrEtdj r|ds |dntd gzntjd krmtdj r|ds |dnd j d kr|d s |d nhtd gznTtjd krDtd j r|d s |d ntd gzS)a Appends or prepends a PUSH_NULL instruction to `inst_or_insts`, depending on Python version. Used when you know that `inst_or_insts` generates a callable that will be called. NOTE: Assumes `inst_or_insts` is a single instruction or sequence of instructions that pushes exactly 1 object to the stack that is to be called. It is important that you include ALL instructions that construct the callable - not just the first instruction/a prefix. Will attempt to use the NULL push bit for instructions with such bits (LOAD_GLOBAL 3.11+, LOAD_ATTR 3.12+, LOAD_SUPER_ATTR). In this case, instructions WILL be modified. )expected inst_or_insts to be a list, got idxrcl|jtd|d|jdzdkSNexpected insts[z].arg to be not Noner r8r}rinstss r$inst_has_bit_setz'add_push_null..inst_has_bit_set.s? :> ! !L3!L!L!LMM MSz~!Q&&r&Nct|jtd|d|xjdzc_dSrrrs r$ set_inst_bitz#add_push_null..set_inst_bit3sA :> ! !L3!L!L!LMM M c !r&rLr PUSH_NULLr\rrarU) r*rlistr}typer0r1rrrsrdr7rl)rrrrs @r$ add_push_nullrs"---$1?-..  QD6F6Fr6J6J > L    / <<==EE W $ $ E"I, - - >6F6Fr6J6J > L     1X_ - -6F6Fq6I6I - LOOOO' 445=EE W $ $ E!HO , , >5E5Ea5H5H > LOOOO' 445=E Lr&ct|tr|g}n6t|tstdt ||}t jdkr|St jdkrdnd}||jdkrL||jtd|d ||jd zdkr||xjd zc_|St jdkr|td gz}ntd g|z}|S) zLike add_push_null, but the low bit of LOAD_ATTR/LOAD_SUPER_ATTR is not set, due to an expected CALL_FUNCTION_EX instruction. rrUrLrrraNrz$].arg to be not None for LOAD_GLOBALr r) r*rrr}rrrrsr7r8rl)rrrs r$add_push_null_call_function_exrTs5 ---$1?-..  QD ! K#KKK  :>A  " " #JNNa NNL 7""+K8899#K001E9 Lr&nargs push_nullctjdkrg}|r^|tdtjdkr|dzn|dz}|t |tjdkr$|td||td ||Std |gS) aS Creates a sequence of instructions that makes a function call. `push_null` is used in Python 3.11+ only. It is used in codegen when a function call is intended to be made with the NULL + fn convention, and we know that the NULL has not been pushed yet. We will push a NULL and rotate it to the correct position immediately before making the function call. `push_null` should be True if no NULL is pushed for the callable. Conversely, `push_null` should be False if a NULL was pushed for the callable. Prefer using `push_null=False` when possible since we will not need to rotate NULL to the right place, which is less efficient. Generally, you should codegen a function by using `add_push_null` then `create_call_function` with `push_null=False`. Example of when to set push_null False: insts = [ create_instruction("LOAD_GLOBAL", argval="torch"), create_instruction("LOAD_ATTR", argval="nn"), create_instruction("LOAD_ATTR", argval="functional"), create_instruction("LOAD_ATTR", argval="relu"), ] insts = add_push_null(insts) insts.append(create_instruction("LOAD_FAST", argval="x")) insts.extend(create_call_function(1, False)) Example of when to set push_null True: insts = [create_instruction("LOAD_FAST", x)] for should_wrap, wrapper_name in wrappers: if should_wrap: insts.extend([ create_instruction("LOAD_GLOBAL", argval="wrapper1"), create_instruction("SWAP", arg=2), *create_call_function(1, True), ) rUrrLr rr\PRECALLrCALL CALL_FUNCTION)rrrsappendrlextendr)rroutputrotss r$create_call_functionrxsR 7""  . MM,[99 : : : # 0G ; ;5199D MM,t,, - - -  g % % MM,YEBBB C C C (U;;;<<< E : : : ;;r&F has_kwargsignore_314_kwargs_pushc.tjdkrg}tjdkr(|s&|s$|tdd}|rx|tdtjdkrt |dznt |dz}|t ||tdt | |Stdt | gS) a( Assumes that in 3.14+, if has_kwargs=False, there is NOT a NULL on the TOS for the kwargs. This utility function will add a PUSH_NULL. If the caller has already pushed a NULL for the kwargs, then set ignore_314_kwargs_push=True so we don't push another NULL for the kwargs. rUrMrTrLrrMCALL_FUNCTION_EXr)rrrsrrlr0rr)rrrrrs r$create_call_function_exrs 7""   ' ' (* ( MM,[99 : : :J  . MM,[99 : : : #w..J!##__q(  MM,t,, - - - ();ZQQQRRR 1s: G G G HHr&ctjdkrtd|gStjdkr"td|td|gStd|gS)Nr\rrrUr CALL_METHODrq)rs r$create_call_methodrsv 7"""6u55566 7"" ye 4 4 4 v5 1 1 1   }% 8 8 8 99r&chtjdkrtdd|Std|S)Nr\r`r )r8r9 LOAD_METHODr|rqrcs r$create_load_methodrs8 7""!+1TBBBB mD 9 9 99r&cLtjdkrdnd}t||S)NrU BEFORE_WITH SETUP_WITHrprq)rr7s r$create_setup_withrs+!.'99]]|F fV 4 4 44r&c xtjdkrtd|gS|dkrgS|dkrtdgS|dkrtdtdgS td |dz td td d ttdtd tdtd d tdtd t dg t dtdtd|dz S)NrUrrr rROT_TWOrM ROT_THREE BUILD_LISTrr{rr| STORE_SUBSCRreverserrUNPACK_SEQUENCE)rrrsrlcreate_binary_subscrrrrs r$ create_swaprsN 7"""6q11122Avv a"9--.. a";//1CI1N1NOO& **9%%9%%  A    9%%  ,!a%888 r&r7cN|dvr*|dkrdnd}tjdkr|dz|zS|dz|zS|dvr|dkrdnd}|dz|zS|d krJtjd krd Stjdkrd |dzzStjd krd|dzzSd|dzzStd|)aReturn the depth of the callable from TOS for a call instruction. Depth is 1-indexed: TOS = 1, below-TOS = 2, etc. This encodes the per-version stack layouts for every call instruction variant. Note: on 3.11-3.12, CPython's CALL can have either NULL or self at the bottom depending on function vs method call. Dynamo's LOAD_METHOD normalizes method calls to NULL + bound_method, so the bottom is always NULL and the callable is always at depth arg + 1. Stack layouts (bottom-to-top within the call's frame): CALL (3.13+): callable NULL arg0 ... argN-1 CALL (3.11-3.12): NULL callable arg0 ... argN-1 CALL_KW (3.13+): callable NULL arg0 ... argN-1 kw_names CALL_FUNCTION (3.10): callable arg0 ... argN-1 CALL_FUNCTION_KW (3.10): callable arg0 ... argN-1 kw_names CALL_FUNCTION_EX: 3.14+: callable NULL args_tuple kwargs_or_null 3.13: callable NULL args_tuple [kwargs_dict] 3.11-3.12: NULL callable args_tuple [kwargs_dict] 3.10: callable args_tuple [kwargs_dict] )rCALL_KWrr rrLr)rCALL_FUNCTION_KWrrrrrMrUznot a call instruction: )rrrs ValueError)r7r8kw_extras r$get_call_callable_depthrs.$$$)++11  w & &7X% %7X% % 8 8 8"44411!Qw!! % % %  w & &1   ( (a=   ( (a= a= |rtdn t}t t |||gStjdkr2|rdnd}t |t |t|gS|rdnd}t |t |tdd t|gS) zQ BINARY_SLICE and STORE_SLICE (if `set` is True) for all Python versions rrr\ STORE_SLICE BINARY_SLICE BINARY_SUBSCR BUILD_SLICErr)rrrsrlrr~slice)rrr subscr_inst inst_names r$create_binary_slicerEs  7""27 S ~ . . .=Q=S=S  eE3// 0 0    W $ $%*>MM e $ $ c " " y ) )  ',@NN e $ $ c " " }! 4 4 4 y ) )   r&ctjdkrtd|gS|dkrtdgSgt|t t|dztdS)NrUrrr rr)rrrsrlrrrRs r$ create_copyrds 7"""6q11122Avv"9--..  Q   QU    Q  r&rcgttdtt|tjdkrdndzt ddtdS)Nr{r|rUrr Fr)rrlprintrrrrsr)rs r$create_print_on_stackrxsv  ),uEEE F F  U3#3w#>#>aaAF G G  a ' '  9%%  r&valuecgttdttd|tddtdS)Nr{r|r Fr)rrlrr)rs r$create_print_valuerse  ),uEEE F F <666  a ' '  9%%  r&cbtjdkrtdStddS)Nrr BINARY_OPrrqr4r&r$rrs1 '!!!/222 kr 2 2 22r&cntjdkr|dkrtdStd|S)Nrrr4 BUILD_TUPLEr)rrrsr~rlrs r$create_build_tuplers9 7""qAvv $$$ m 3 3 33r& first_linenoctjdddkr$tdtjddg|dddtdtddffd d td tddffd }d tddffd }||fS)z Used to create typing.CodeType.co_linetable See https://github.com/python/cpython/blob/main/Objects/lnotab_notes.txt This is the internal format of the line number table for Python 3.10 Nr)rM z+linetable_writer requires Python 3.10, got r byteno_delta lineno_deltarc2|dks|dkrtdt|d}tdt|d}|dkr|dkrtd||z}||z}||dzf|dk|dkdSdS)Nriz,byte_offset and line_offset cannot both be 0)maxminr}r)rr byte_offset line_offset linetables r$_updatez!linetable_writer.._updatesa<1#4#4a\3!7!788KdC c$:$:;;KaK1$4$4$%STTT K 'L K 'L   k;+=> ? ? ?a<1#4#4#4#4#4#4r& lineno_new byteno_newc<|z }|||z |dSrAr4)rrrrbytenolinenors r$updatez linetable_writer..updates7!F*  l+++!F* r& total_bytesc&|z dSrAr4)rrrrs r$rzlinetable_writer..ends  f$l33333r&)rrrsr}r0)rrrrrrrrs @@@@@r$linetable_writerr s" w&& P#:J2A2:N P P   I FL F@c@@@@@@@@3CD4444444444 fc !!r&c|dkrtd||dzg}|dz}|dkr3|dxxdzcc<||dz|dz}|dk3|S)z~ 6-bit chunk encoding of an unsigned integer See https://github.com/python/cpython/blob/3.11/Objects/locations.md rz+encode_varint requires non-negative n, got ?r@)r}r)rbs r$ encode_varintrs  1uuN1NNOOO RA!GA a%% "  R a a%% Hr&ctjdkrtdtjg|dtjdzdt ddffd }|fS)z Used to create typing.CodeType.co_linetable See https://github.com/python/cpython/blob/3.11/Objects/locations.md This is the internal format of the line number table for Python 3.11 rUz0linetable_311_writer requires Python 3.11+, got r=N inst_sizercrjnd}dtdtddffd }|d}n|z }||dkr||d|dz}|dk|||dS)Ndeltasizercd|cxkrdksntd|d}rZjSjLjEj>d|zdz jjz jdzjdzf}nd|zdz |dkr | dzdz}n|dz}t||D]$}t|%dS)Nrz"size must be in range (0, 8], got r4r )r}r end_lineno col_offsetend_col_offsetrrr)rr other_varintsrrr=s r$rz5linetable_311_writer..update.._updatesID A ()Td)T)TUUU24 >!(4!,8!,8!0<$$\D%81%<===",y/??!,q0!014%MM $$\D%81%<===199$f]a/EEaKE  u!5!5666&77A$$]1%5%5666677r&rr)rr0)r=rrrrrrs` r$rz$linetable_311_writer..updates-6@))DJ 7s 7# 7$ 7 7 7 7 7 7 7B! )F2 #a-- a(((Q a-- GL) , , , , ,r&)rrrsr}ri Positionsr0)rrrrs @@r$linetable_311_writerrs  g % % U3CSUU  - -cmd2- -s- -t- -- -- -- -- -- -- -^&  r&cBeZdZUeed<eed<eed<eed<eed<dS)ExceptionTableEntryrrrrrN)r,r-r.r0r/r1r4r&r$r r s= JJJ HHH KKK JJJ KKKKKr&r c*|dkrtd||dzg}|dz}|dkr#||dz|dz}|dk#|tt |dz D]}||xxdzcc<|S)zR Similar to `encode_varint`, but the 6-bit chunks are ordered in reverse. rz;encode_exception_table_varint requires non-negative n, got r r r r )r}rrrlen)rrrOs r$encode_exception_table_varintr#s 1uu M! M M    RA!GA a%% R a a%%IIKKK 3q66A:   !  Hr& bytes_iterczt|}|dz}|dzr!|dz}t|}||dzz}|dz!|S)z5 Inverse of `encode_exception_table_varint`. r r r )next)r$rrxs r$decode_exception_table_varintr'0s^ ZA b&C b&      q2v  b& Jr&tabctt|dz D]}||j||jkrA||j||dzjkr"||dzj||dzjksZt d|d|d||jd||jd|dzd||dzjd||dzjddS) z Verifies that a list of ExceptionTableEntries will make a well-formed jump table: entries are non-empty, sorted, and do not overlap. r z5exception table entries are not well-formed at index z: entry[z ]=(start=rz ), entry[r N)rr"rrr})r(rOs r$check_exception_tabler*=s2 3s88a<   FLCFJ & &A SQZ---AE  CAJN22 SSSSS%(V\SS9>BE2;??!CF&.1$C2;??!CF.{;;B!GEaLLE JJ*5#vueLL M M M N c""" sB BB=<B=cg}|D]}t|jdz}|dxxdzcc<|||j|jz dz}|t|dz|t|jdz|jdz|jz}|t|t|S)zd Inverse of parse_exception_table - encodes list of exception table entries into bytes. rrr )r#rrrrrrbytes)r(rentry first_entryr0r1s r$assemble_exception_tabler8es A443EK14DEE A&  U[(1, .v{;;<<< .u|q/@AABBBkQ%+ - .r223333 88Or& instructions firstlinenocNg}tjdkrt|\}}d}t|D]\}}|jdkr8d}|dz }dD]-} ||| zjdkr||| zj|_n.nt |dz|z}d}||j|d}|jpd} ||j | dzftt |dzdz D]} |dn}t|\}}} |D]O}|j ||j t||jpd} ||j | dzfP| t|t|t|fS) z)Do the opposite of dis.get_instructions()rUr EXTENDED_ARGr r rrMrr)rr)rrrsr enumerater7r=instruction_sizer8rr6rr r;r"r5) r9r:codelnotab update_linenonum_extrOrtrjr8_rs r$assemblerFwsD 7"" 4[ A A  .. $ $GAt{n,, 1 "A#AE*1^CC)5a!e)<)FD-T22a7'A  M$.) 4 4 4G(-aC KKcDj1 2 2 2+D11Q6:;; $ $ F#### $! $&&6k%B%B" s  3 3D+ d.D :::(-aC KKcDj1 2 2 2 2 CII ;;f %%r&offset_to_instr:cfdD]-}|||zjtjkr |||zcS.dS)zU Get the instruction located at a given offset, accounting for EXTENDED_ARGs )rrrr Nr6rir<)rGr:rs r$_get_instruction_by_offsetrJsM .. &1* % ,0@ @ @!&1*- - - - A 4r&cd|D}|D]B}|jtjvs|jtjvrt ||j|_CdS)z9Replace jump targets with pointers to make editing easierc,i|]}|j |j|SrA)r:rrts r$ z$virtualize_jumps..s*"dk6M T6M6M6Mr&N)r6rihasjabshasjrelrJr9r)r9 jump_targetsrts r$virtualize_jumpsrRsp&2LPP ;#+ % % )C)C4\4;OODKPPr& instructionctjdkrtdd|jvr!|jdd|_n9d|jvr!|jdd|_nt dt j|j|_|jtvrtd|jddS)NrUz+Cannot flip jump direction in Python < 3.11FORWARDBACKWARDz-Instruction is not a forward or backward jumpz expected z" to be a relative jump instruction) rrrsrhr7replaceAttributeErrorrirjr6 _REL_JUMPSr})rSs r$flip_jump_directionrZs '!!HIIIK&&&(/77 :NN  {) ) )(/77 INN LMMM;#56K++ N * N N N   ,+r&rc||}dD]2}||kr*|||z jtjkr |||z }2|S)zp i.e. get the first EXTENDED_ARG instruction (if any) when targeting instructions[idx] with a jump. r=rI)r9rrr:s r$_get_instruction_frontr\sW # F &==\#,7>#BRRR!#,/FF  Mr&c Tttjttj}|D]}|j|vr|jtjvr|jtd|jd|jj td|jd|j td|jd|jj |j kr8tj dkrtdd|jvrt|tj dkrd |jvrt|t|t|}|D]J}|j|vr=|jtd|jdt!|||j}|jtjvrRtj dkr3|j td t#|j d z |_ntd |j td |j td|jdt't#|j |j z t)|z |_|xjd zc_|j |_d|j |_LdS)zJFill in args for virtualized jump target after instructions may have movedNzjump instruction z has no targetzjump target for z has no offsetrUz*Got negative jump offset for Python < 3.11rUrVz"absolute jump target has no offsetrz+Python 3.11+ should not have absolute jumpsz"relative jump target has no offsetzrelative jump instruction zto )setrirOunionrPr6rr}r7r:rrrsrhrZupdate_offsetsrr\r0r8absr?r9r?)r9jumpsrtindexofrs r$devirtualize_jumpsrds    " "3s{#3#3 4 4E22 ;%  {#+--;&(GDKGGG;%-()WDK)W)W)WXXX;&(GDKGGG;% 33''11*+WXXX DK//+D111'722zT[7P7P+D111<   ,''G11 ;%  {"$%T%T%T%TUUU+L'$+:NOOF{ck))#g--},,-QRRR"6=1#455DHH&'TUUU=(()MNNN;&(PT[PPP  36Ft6L6LLMMQ -DK000DL911r& exn_tab_bytesc t|}d|D t  dt| dtt t ff fd }|\}}|D]w}|jtd|j d|j|j kr|\}}|j|j k|j|j krtj ||_ xdS#t$rYdSwxYw)zDReplace exception table entries with pointers to make editing easiercDi|]}tt|j|Sr4)r r0r:rMs r$rNz.virtualize_exception_table..s&LLLtd3 ,,dLLLr&rrct}tkr:|jkr)dz tkr|jk)dkrtddz }t t |jt |t |j|j|j }||fS)Nr rzAend_offset_idx must be positive in exception table virtualization) r&r"rr}rrJrrrr)r6 end_offset inst_entryend_offset_idx exn_tab_iterrGoffsetss r$stepz(virtualize_exception_table..step!s&&E W--'.2IUY2V2V!#W--'.2IUY2V2V""$W!!!34J/*>5;GG*>:FF*>5<HH  J*$ $r&N instruction z4 has no offset during exception table virtualization)r2sortedkeysr-tupler rr:r}r7rrcopyr>r.) rer9exn_tabrnr6rjrtrkrlrGrms @@@@r$virtualize_exception_tablerusk$M22GLL|LLLN^((**++GN==L'  %e/1GGH % % % % % % % % %2!DFFz ; ;D{"$+4;++++ ))$(DFF!z+ )){ek))%)Yz%:%:" ; ;      sB)C>> D  D c 8 i t|}|D]}|jrt|||jjj}|t dt t|jjjt|jjzdz }|t dt|||jj j}|t d||f}||jj |jj f}| vr* ||krt d|d |d|| |<t d }d g g d fd } |D]} r< dd|d kr$|  r dd|d k$ r dd |d cxkr|dcxkr ddksnt d|d dt dd } | |d kr6 t#| |d dz g dR|d  | r |  t%  S)zMCompute exception table in list format from instructions with exn_tab_entriesNz5exception table entry start instruction has no offsetrz3exception table entry end instruction has no offsetz6exception table entry target instruction has no offsetz.conflicting exception table entries for range z: z vs c$|d|d fSNrr r4ts r$z)compute_exception_table..{s1!u r&keyrrcrn}|dkrPtt|d|dg|R|ddzdSdSdS)zT Pop the key_stack and append an exception table entry if possible. r rrN)poprr r)r}exn_dictrt key_stacknextis r$rz$compute_exception_table..pops  #--//CA'CFE(:(:CFSXc]SSSA   # #r&rr zexception table key & is not properly nested within parent )rN)rr>r\rr:r}r r0rr?rrrrprqrrr r*)r9rcrtrrrr}rx keys_sortedrleftrrtrrs @@@@r$compute_exception_tablerIs>@H,''G# #   " *gd&8&>? }$K S$,0788"4#5#9::;  {$I,gd&8&?@ ~$L#,C4-3T5G5MNChC=C''(;;;"*3-;;58;; HSM.E.EFFFK E')I)+G # # # # # # # # # IbM!,s1v55 CEEE IbM!,s1v55  bM!$ALLLL#a&LLLLIbM!PPPPP$7S77%0_77 3   r&c t|}i|D]^}|jrU||jj||jjf}|vr#|j|krt d||j|<_fdt dD}t|||D]M}t||j||jdzD]!}tj |||_"NdS)z~ Copies exception table entries to all instructions in an entry's range. Supports nested exception table entries. z1conflicting exn_tab_entry for instruction at key c g|] }| Sr4r4)rr}entriess r$rz4propagate_inst_exn_table_entries..s* r&c$|d|d fSrxr4rys r$r{z2propagate_inst_exn_table_entries..sadQqTE]r&r|r N) rr>rrr}rprqrrrs)r9rcrtr}sorted_entriesr6rOrs @r$ propagate_inst_exn_table_entriesrsU ,''G=?G . .   .*01*./Cg~~%55(QCQQ -GCL &w||~~;R;R S S SN&ng>>> ==wu{+WUY-?!-CDD = =A,0Ie,<,r>rrrsr}rCaddrrr)r9rcexn_tab_entry_setrOrtr6s r$ check_inst_exn_tab_entries_validrs,''G\**4   '))$%VWWW$$%%):::$GAGG  ! !"T%7"8"8 9 9 9&E{'))$Oy''$M|7**$PEK(ACCCC1CCCCC$LALL%ek2LL6=ei6HLLL/r&c(d|D|dd<dS)Nc<g|]}|jtjk|Sr4rIrs r$rz'strip_extended_args..s'OOOQ!(c>N2N2Nq2N2N2Nr&r4)r9s r$strip_extended_argsrs!OO,OOOLOOOr&old_inst new_instsc|jr3|jj|ur%t|dkr|d|j_|ddD],}tj|j|_|j|_-|dj|_|dj|_|dj|_|dj|_|dj |_ |g|ddzS)Nr rr) r>rr"rsr=r6r7r8r9r)rrrts r$overwrite_instructionrs 3  " &( 2 2  NNQ  %.r]"!"" ,,!Yx'=>>!+l)HOl)HOQ<#HLl)HOl)HO : !"" %%r&ctjdkrtdddd}|D]9}|j|vr.||j|_tj|j|_:|S)zFLOAD_METHOD puts a NULL on the stack which causes issues, so remove itrUzrrsr=rirjr6r8r9)r9rrtrs r$remove_binary_store_slicerRs I** ;9 9 97**t{n/L/L244 01D1DWh1W1WXX ! 5d&8&<&D&D)4"&(, $2D(E(EK %$(NK !)M2DK'DKDHDK   [ ) ) )LOOOr&) LOAD_FASTr)LOAD_FAST_BORROWr) STORE_FASTr)rr)LOAD_FAST_LOAD_FAST!LOAD_FAST_BORROW_LOAD_FAST_BORROWSTORE_FAST_STORE_FASTSTORE_FAST_LOAD_FASTc2g}|D]}|jtvret|j\}}|j\}}t||t||g}|t ||u||||dd<dS)Nr|)r7 FUSED_INSTSr9rlrrr)r9rrtinst0inst1argval0argval1rs r$remove_fused_load_storerqsI # # ;+ % %&t{3LE5#{ GW#5999"5999M   24GG H H H H   T " " " "LOOOr&cg}|D]p}d|jvrPtddt|j|jg}|t ||[||q||dd<dS)NRETURNNOPGRAPH_BREAK_IF_LEAFr|)r7rlr9rrr)r9rrtrs r$$add_graph_break_if_leaf_instructionsrsI## t{ " ""51FGGG"4;t{CCCM   24GG H H H H   T " " " "LOOOr&cg}tj|D]\}}|jdkr|jdkr|jdrt t |D]}||j|ur |||_||jrr||jj |ur|||j_ ||jj |ur|||j_ ||jj|ur|||j_| || |d||dd<dS)Nrrrr) itertoolspairwiser7r9 startswithrr"rr>rrr)r9rrt next_instrOs r$'remove_graph_break_if_leaf_instructionsrsbI$-l;;##i K5  444 ++H5553|,,-- I I?)T11-6LO*?0I$A4:dBB>G Q5;#A48D@@rr7r9r8rrrsr}rl) r@r9 cell_and_freerrrtr first_var call_insts r$explicit_superrs%+0@0FBGM t  F|,,)) T d ;- ' 'DK7,B,B q)EyA~~!W,,1G1G$//(722 11$w..5.maybe_pop_nsLq  A &*+s/???   r&rirMr<rrirrr z.fix_extended_args removed instructions (added=z), expected non-negative) r0r6rir<r8rrlr"r})r9rrtaddedrs @r$fix_extended_argsrs$ "Fst  ;#* * *DHH X Q$(X-- KNNN MM,^RPPP Q Q Q MM,^RPPP Q Q Q MM,^QOOO P P P P X Q$(V++ KNNN MM,^RPPP Q Q Q MM,^QOOO P P P P X Q$(T// KNNN MM,^QOOO P P P d KK#l++ +E qyy $U $ $ $   LO Lr&rtc|ddl}tjdkr'd|jjjj|jdzzSdS)NrrUrr )torchrrrs_C_dynamo eval_framepy_opcode_cachesr6)rtrs r$r?r?sALLL 7""EH$/@MPQQRR 1r&c d}|D]A}|j|kr"td|jd|jd||t|z }BdS)Nrroz has offset z , expected )r:r}r7r?r9r:rts r$ check_offsetsrsv F)) ;&  Xt{XX XXPVXX  "4((( ))r&cFd}|D]}||_|t|z }dS)Nr)r:r?rs r$r`r` s< F)) "4((())r&argsctttt|}d|gt |zdt |d|dDgzD}dd|zS)NcLg|]!}dd|D"S) c3K|]}|dV dS)03Nr4)rxs r$ z)debug_bytes...s&((qA((((((r&)join)rr8s r$rzdebug_bytes..sC  ((C(((((r&c:g|]\}}t||kSr4)r0)rars r$rzdebug_bytes..s& < < F  & 1 1 11r&ct|dd\}}|j|jkr'tt|j|jdS)z.#str&T)safeN)transform_code_objectco_coder}r)r@doderEs r$ debug_checksr!sQ#D*;*;$GGGGD! |t|##[t|DDEEE$#r& code_optionsct|dD] \}}||ur|cS|dxx|fz cc<t|ddz S)N co_constsr )r>r")rrxrOvs r$get_const_indexr3sq,{3441 !88HHH #' |K( ) )A --r&varname_from_oparg.c   dtdD dtdtf fd }tjdkrY|t ddtd D d td d zD nt |st d i tjD]$} ||}| |<#t$rYnwxYw fdd D fdd d zD ttD] dtf fd } j dkr jturt dtjdkre jt d| jdzt#t jdzz _| j _܉ j dkr jturt dtjdkrf jt d| jdzt#t jdzz _| j _ j dkr jt d jturt d| jdzt#t jdzzdz _F j t$vr'tjdkrt d j dt' jt(s8t d j d t+ jt jdkr8t d j d!t jt) fd" jD}|d#d$z|dd%zz _ jt.vri|r]tjdkr/ j vr j _߉ j _ jt0vrZ|rN| j}tjd&kr j d'kr |dz}|dz}| _m jt2vr*|r j _ jt4vrX jKt7 j}|d#krt d( j| _dS))Nci|]\}}|| Sr4r4rrr^s r$rNzfix_vars..Es L L L93T3 L L Lr&co_namesr^rc < |}n#t$rtx}|<gd|Rd<tdtkr6tdtddtdYnwxYw|S)Nrzco_names length z" does not match names dict length )KeyErrorr"r})r^rrnamess r$get_name_indexz fix_vars..get_name_indexGs +CC    #E *C%+'Hj)A'H4'H'HL $< +,,E ::$.XsVVV)#tD#VVVr&rci|]\}}|| Sr4r4rs r$rNzfix_vars..Ys.   T #   r&rrz4varname_from_oparg must be callable for Python 3.11+c"i|] }|| Sr4r4rr^allnamess r$rNzfix_vars..isQQQTD(4.QQQr&c"i|] }|| Sr4r4rs r$rNzfix_vars..js/    (4.   r&c,jtuSrA)r9rZ)rOr9sr$should_compute_argz$fix_vars..should_compute_argps?)= =r&raz,LOAD_GLOBAL instruction must have argval setz9LOAD_GLOBAL instruction must have arg set in Python 3.11+r rr`z*LOAD_ATTR instruction must have argval setr\z7LOAD_ATTR instruction must have arg set in Python 3.12+rbz-LOAD_SUPER_ATTR instruction must have arg setz0LOAD_SUPER_ATTR instruction must have argval setrLzfused instruction z requires Python 3.13+z argval must be a tuple, got z argval must have length 2, got c3@K|]}|vr|n|VdSrAr4)rr^ freenamesvarnamess r$rzfix_vars..sM#'("2"2 $r&rr)rMr IMPORT_NAMEz"failed to find constant index for )r>r2r0rrrsr}callablercount IndexErrorrr"r1r7r9rZr8r rr*rrrr6 HAS_LOCALHAS_NAMEHAS_FREE HAS_CONSTr)r9rrr rr^r arg_tupler8rrrOr rs`` @@@@@r$fix_varsr ?sd M L ,z2J(K(K L L LE S S        '!!  ) !VWW WVVym9T/U/UVVV  &]+l=.II   *++ Y !WXX X?$$  C ))#..!$    QQQQ\-5PQQQ    $]3l=6QQ   3|$$ % %c*c* >D > > > > > > > ? !] 2 2A%55$%STTT7**?&.(S(6~l1o6L'M'MQR'Rl1o122Q6' Q##'5n\!_5K&L&L Q## !_ #{ 2 2A%55$%QRRR7**?&.(Q(6~l1o6L'M'MQR'Rl1o122Q6' Q##'5n\!_5K&L&L Q## !_ #'8 8 8A"*$%TUUUA%55$%WXXX  Q 6771<\!_011A57 O   !_ #{ 2 2'))$Wa)?WWWl1o4e<< $:a)?:: Q 677::<?)**a//$Ca)?CC%(a)?%@%@CC(O2I$-Q<1#419J"KLO   !_ #y 0 0!!## K$//$Q.h>>+4LO4J*KLO''*2<?3I*JLO' !_ #x / /!!## *$n\!_%;<<$//$Q.-??AIC4KC&) Q# !_ #x / /!!## H&/ Q0F&G Q# !_ #y 0 0A"*%lLO4JKK77(U\!_=SUU'* Q#Gc*c*sC,, C:9C:c|D]X}|jturH|jtvs*|jtvs|jt vs|jt vr|jdvrd|_YdS)N)rar`rb) r9rZr6rrrrr7r8)r9rts r$clear_instruction_argsr"st   K| + + y((;(**;(**;)++ #RRRDH  r&czdg}|d|gdtjdkr|d|d|dtjdkr|d|d d g|S) N co_argcountco_posonlyargcount) co_kwonlyargcount co_nlocals co_stacksizeco_flagsrrrr co_filenameco_namerU co_qualnameco_firstlineno co_linetableco_exceptiontablerr)rrrrrs)rqs r$ get_code_keysr0s ?DKK$%%%KK     7"" M"""KK !!!KK 7"" '(((KK    Kr&transformationsrrc\t}fd|D}t|d|dkr.tdt|dd|dt|}t ||||}t |||\}}||fS)Nc2i|]}|t|Sr4)getattr)rkr@s r$rNz)transform_code_object..s%666AAwtQ''666r&rr'zco_varnames length z does not match co_nlocals )r0r"r}cleaned_instructionsrclean_and_assemble_instructions) r@r1rrqrr9 tracer_outputrEbytecodes ` r$rr s ??D6666666L < &''< +EEE =#l=&A"B"B = = ,\ : = =   (d33L %%%#OL,??M1,lSSKAx ] ""r&rqc t|t|tdd<d}tjdkr!t jfd|D}|j}t||d}|r.-s#B#B#BLO#B#B#Br&)rTr-r.rr(r%z%code_options keys mismatch: expected z, got r/c g|] }| Sr4r4r<s r$rz3clean_and_assemble_instructions..Js)H)H)Ha,q/)H)H)Hr&)rrr"rrrstypesCodeType_varname_from_opargr r`rdr1rrrFrr^rqr}r8r)r9rqrrtmp_codedirtyr9rAs ` r$r7r7"s!,L999$\222!$\-%@!A!AL 7"">#B#B#B#BT#B#B#BC%9 \<r?rs r$rz'_clone_instructions..`sl   H H E H H M  K H O I r&) dictrrr>rrrrr)r9copiedremaprOr6s r$rIrI]sF" \6** + +EE$K   ?O #5 4ek"ei el#  AO Mr&c(ttttj|}t |t |tjdkr*t||j t|j |t|t||stjdkrt|tjdkrt!||tjdkrMt#|tjdkrt%|tjdkrt'|t(jrt-|tjdkrt/|t1||S)NrUr\rL)rrrSriget_instructionsrrrrrsrFrrur/rRrrrrrrr &debug_force_graph_break_on_leaf_returnrr`rdrJs r$rHrHsq/1Ed1K1KLLMML %%%, 7"" t~>>>"4#9<HHH\""" %%%  ?  g % % #L 1 1 1  g % % 4 . . .  w & &  - - -7**),7777**' 555  8 ? 0 > > > 7""|$$$<((( r& with_uuidc|dtt}|r-|dtjddz }|S)NrE-)r&_unique_id_counteruuiduuid4rW)r^rSrets r$ unique_idrZsT  . .D+,, . .C4 !4:<<!!))#s333 Jr&c"d}|j|zdkS)N r)r))r@ co_generators r$ is_generatorr^sL ML (A --r&fn varname_mapnoreturnnoprefixct|j}t||r/t|D]\}}|jdkr||dzd}n |D]-}d|_d|_|r|j|vr||j|_.|rtj dkrg}|D]x}|jdkrVtj d|_ d|_| || tdc| |y|}g}|D]"}|jdkr| |#t|dkr&|d|d ur|d nt|dkr|d |d ur_d |d _tj d |d _ d|d _t$|d _|d n"| td |D]S}t'|d } | j|_| j |_ | j|_| j|_| j|_T|S) aGenerates bytecode from a template function `fn` for use in dynamo bytecode generation. For example, we can generate Python-version-independent bytecode for looping through a dictionary and copying the values to a new dictionary. def template(d1, d2): for k, v in d1.items(): d2[k] = v or a try block: def template(): try: dummy1 except: dummy2 raise dummy3 Args: fn: a function template to generate bytecode from varname_map: a mapping of `fn`'s varnames to new names. This map will be applied to the generated bytecode's varnames. For example, local variables in `fn` can be replaced with new names that are generated by `OutputGraph.new_var`. noreturn: remove all RETURN_* bytecodes and replace them with a jump to the end of the bytecode. NOTE: any items pushed to the stack for return WILL remain on the stack! Append a POP_TOP if you don't want that item to be present. noprefix: remove prefix bytecodes (all bytecode before the first RESUME, inclusive). RESUMEr Nr\ RETURN_CONSTr{ RETURN_VALUErrr)r6__code__r"r>r7r;r=r9rrrsrirjr6rrlr"rr8rZrur) r_r`rarbrrOrtrreturns jump_insts r$bytecode_from_templaterjsN ! - -E5!!! ''  GAt{h&&a!egg'33   34;+55%dk2DK-/  w & &I + +;.00"%)L"9DK".DK$$T***$$%7%G%GHHHH$$T****E % %D{n,,t$$$ w<<1  uRy!8!8 IIbMMMM \\A  RyGBK''#(b  #&9U#3b  $b #/b   B /66777  / /1r;; '. '. $='. '. Lr&)T)FrA)NTT)rKrs dataclassesri functoolsrrrr>rWcollections.abcrrrrrtypingrr r r r bytecode_analysisrrrrutilsr output_graphr dataclassrrrsrSrZr2r1rdr0rlrurwr~rrrrrrrrrrrrrrrrrrrrrr rrrr r#r'r*r5r2r8rFrMrJrRr^rPrYrZr\rdrurrrrrrrrrrrrrr?rrr?rr`rrhaslocalrhasnamerhasfreerhasconstrrr r" lru_cacher0rr7rFr6rIrHrrVrZr^rjr4r&r$rys2     KKKKKKKKKKKKKK555555555555 $#####1000000T"""       #" 6T""")))))))#")>w  s  ;           s  ;        s  ;      wGsGtGGGGG   %s%t%%%%% st%!% ''' ' t' $J ' $  '  ''''V333333 >[>T>>>> 8 83 8 8 8 8 8 84 ,$) )))) 0C0D-0000.=k!22= +====@!k!22! +!!!!H4<4<4<k9J4<4<4<4C,>c,>c,>,>,>,>`7<  : $J /3  +    >34 ,(k):cd;&73k33334#4+4444 &"&" 49hSz4/0(C5$;2G GH&"&"&"&"R  S  T#Y     w>!>! tCy(CMD$8##>#DEE F>!>!>!>!BT"""#" S T#Y    & hsm     t$78T$%D1D,E,$':";$#&4 ,#&3#&5PUCV#&#&#&#&L k)* 47 4    P8K#8PTPPPPS     [  T     k):      >1T+%6>14>1>1>1>1B0 0 (,[(90  0 0 0 0 f[{#[ [[[[| $ % 04[#5E0F      .=4 3D=====<&4 3D&&&&&RPd;&7PDPPPP&&&*;&7& +&&&&. ${*; [@Q     d;&7 D    > D,= $    06)Q97   ${*;     &  tK7H  T      ${:K PT    6&&tK7H&T&&&&R"D$5"#""""J ; 3    ) 5)$)))))+!6)4)))) 2u 2 2 2 2 2Fu~F$FFFF C   3s{   3s{   C    .$sCx. .s .s . . . .59R*R*{#R*sCx.R*!c*T1R* R*R*R*R*jk):t" !tCy!!!!R ## .# k DcN+X6J-KK#  #  5>8$89 9: ####0(J{#(J+/9(JDHcN(J 4 en ,-(J(J(J(JV+8K+@+#+RV++++--u~-T-d;FW----"h{&;"[@Q""""J ', . $ k>%Y_&&CDS.u~.$....-1 ggcgc"T)gg g  + ggggggr&