+
    –°³jJ ã            	      ó  € ^ RI Ht ^ RIt^ RIt^ RIt^ RIt^ RIt^ RIt^ RIt^ RI	t	^ RI
t
^ RIt^ RIt^ RIt^ RIHt ^ RIHt ^ RIHtHt ^ RIHtHtHtHt ^ RIt^ RIHt ^ RIt^ RIt^ RIHu H t! ^ RIH"t# ^ R	I$H%t%H&t&H't' ^ R
I(H)t) ^ RI*H+t+ ^ RI,H-t- ^ RI.H/t/H0t0 ^ RI1H2t2 ^ RI3H4t4H5t5H6t6 ^ RI7H8t8H9t9H:t:H;t;H<t< ^ RI=H>t> ^ RI?H@t@ ^ RIAHBtB ^ RICHDtDHEtEHFtF ^ RIGHHtH ^ RIIHJtJHKtK ^RILHMtMHNtNHOtPHQtQ ^RIRHStS ^RIQHTtTHUtU ^RIVHWtW ^RIXHYtYHZtZ ^RI[H\t\H]t]H^t^ ^RI_H`t`HataHbtbHctc ^RIHdtdHeteHftfHgtgHhthHitiHjtjHktkHltlHmtmHntnHotoHptpHqtq ^RIrHsts ^R ItHutuHvtvHwtwHxtxHytyHztz ^R!I{H|t| ^R"I}H~t~ ^R#IH€t€H�t�H‚t‚ ]'       d'   ^ R$IHƒtƒH„t„H…t… ^ RI†t†^R%I‡Hˆtˆ ^R&IQH‰t‰ ^R'IŠH‹t‹ ^R(IŒH�t� ]EP                  ! ]�4      t�]x! 4       EP"                  t’R) R* lt“R+ R, lt”]•]EP,                  ]PD                  ]—]˜]™]•]™]™3,          R,          3,          tš]•]EP,                  ]PD                  ]—R-]—3,          t›]QEP8                  ]y,          t�]R..R3,          tž]EP>                   ! R/ R04      4       t R1 R2 lt¡R3 R4 lt¢R5 R6 lt£R7 t¤]¥]—]™3,          t¦]•]™]P6                  ,          R83,          ]]¦.]•]™R83,          3,          ,          t§R�R9 R: llt¨R‘R; R< llt©R= R> ltª]EP>                   ! R? R@4      4       t« ! RA RB4      t¬ ! RC RD]4      t­ ! RE R.4      t®]EP>                   ! RF RG]®4      4       t¯]EP>                   ! RH RI]®4      4       t°]EP>                   ! RJ RK]®4      4       t±]EP>                   ! RL RM]®4      4       t²]EP>                   ! RN RO]®4      4       t³]`t´RP RQ ltµ]EP>                   ! RR RS]®4      4       t¶ ! RT RU]®4      t·]EP>                   ! RV RW]®4      4       t¸]EP>                   ! RX RY]®4      4       t¹]EP>                   ! RZ R[]®4      4       tº]EP>                   ! R\ R]]®4      4       t»]EP>                   ! R^ R_]®4      4       t¼]EP>                   ! R` Ra]®4      4       t½]EP>                   ! Rb Rc]®4      4       t¾]EP>                   ! Rd Re]¶4      4       t¿]EP>                   ! Rf Rg]·4      4       tÀ]EP>                   ! Rh Ri]®4      4       tÁ]EP>                   ! Rj Rk]®4      4       tÂ]EP>                   ! Rl Rm]®4      4       tÃ ! Rn Ro]®4      tÄ ! Rp Rq4      tÅ]EP>                   ! Rr Rs]¾4      4       tÆ]EP>                   ! Rt Ru]¾4      4       tÇ]EP>                   ! Rv Rw]¾4      4       tÈ]EP>                   ! Rx Ry]¾4      4       tÉ ! Rz R{]¾4      tÊ]EP>                   ! R| R}]®4      4       tË]EP>                   ! R~ R]®4      4       tÌ]EP>                   ! R€ R�]®4      4       tÍ]EP>                   ! R‚ Rƒ]®4      4       tÎ]EP>                   ! R„ R…]®4      4       tÏ]EP>                   ! R† R‡]®4      4       tÐ]EP>                   ! Rˆ R‰]®4      4       tÑ]EP>                   ! RŠ R‹]®4      4       tÒ]—tÓ]¾]l,          tÔ ! RŒ R�]v4      tÕ ! RŽ R�]Õ4      tÖR# )’é    )ÚannotationsN)ÚCallable)ÚEnum)ÚchainÚcount)ÚAnyÚLiteralÚProtocolÚTYPE_CHECKING)ÚExpr©Údtype)ÚcountersÚdynamo_timedÚget_debug_dir)ÚDebugPrinterManager)ÚMultiKernelState)Ú	cache_dir)Úget_opaque_obj_reprÚis_opaque_constant_type)Útrace_structured)Ú_coor_check_current_acceleratorÚ_coor_current_acceleratorÚ_coor_enabled)ÚCallMethodKeyÚConvertIntKeyÚDivideByKeyÚresolve_unbacked_bindingsÚSymTypes)Ú_get_qualified_name)ÚStorageWeakRef)Ú
OrderedSet)ÚCleanDivÚMaxÚMin)ÚSingletonInt)Úsymbol_is_typeÚSymT)Úasync_compileÚconfigÚdebugÚir)Úoutput_code_log)ÚIRNodeÚReinterpretView)Útriton_heuristics)ÚDevicePropertiesÚ
TritonMeta)ÚDEFAULT_STREAMÚDEFAULT_STREAM_IDXÚSTREAM_NAME_TEMPLATE)ÚCOOR_DEVICE_IDX_VARÚcoor_device_strÚget_raw_stream_nameÚget_stream_name)Úcache_on_selfÚDeferredLineBaseÚDelayReplaceLineÚget_benchmark_nameÚget_dtype_sizeÚIndentedBufferÚ#is_codegen_graph_partition_subgraphÚis_using_cudagraph_partitionÚLineContextÚmake_codegen_bufferÚsympy_productÚ	sympy_strÚ
sympy_subsÚtriton_version_uses_attrs_dict)ÚV)ÚArgNameÚCodeGenÚDeferredLineÚPythonPrinterÚWorkspaceArgÚWorkspaceZeroMode)Úcexpr)ÚCUSTOM_EXTERN_KERNEL_CODEGEN)Ú	config_ofÚshould_unwrap_unspec_argÚsignature_to_meta)ÚIterableÚIteratorÚSequence)ÚGraphLowering)ÚExternKernel)ÚBaseSchedulerNode)ÚFxConverterc               ó    € V ^8„  d   QhRRRR/# )é   Úexprú
sympy.ExprÚreturn© )Úformats   "Ú`/home/luca/screen-checking/.venv/lib/python3.14/site-packages/torch/_inductor/codegen/wrapper.pyÚ__annotate__rc   o   s   € ÷ ,ñ ,°:ð ,À*ñ ,ó    c                ó®   € \         P                  ! \         P                  ! V 4      4      p \         P                  ! V 4      w  rV^8X  d   V# \	        W4      # ©é   )ÚsympyÚtogetherÚsympifyÚfractionr#   )r]   Ú	numeratorÚdenominators   &  rb   Ú(_rewrite_symbol_solution_for_int_codegenrn   o   sC   € Ü�>Š>œ%Ÿ-š-¨Ó-Ó.€DÜ"Ÿ^š^¨DÓ1Ñ€IØ�aÔØÐÜ�IÓ+Ð+rd   c               ó    € V ^8„  d   QhRRRR/# )r\   Úobjr   r_   r`   )ra   s   "rb   rc   rc   w   s   € ÷ ñ ˜Cð  Cñ rd   c                ó¢  a € \        S \        4      '       d8   S P                  4        UUu/ uF  w  r\        V4      \        V4      bK  	  upp# \        S \        4      '       d   S  Uu. uF  p\        V4      NK  	  up# \        S \
        4      '       dB   \        S R4      '       d0   \        \        S 4      R4      ! V 3R lS P                   4       4      # \        S \
        4      '       d-   \
        ;QJ d    . R S  4       F  NK  	  5# ! R S  4       4      # \        S \        4      '       d   \        S P                  4      # S # u uppi u upi )zWConvert Enum values to their underlying value for valid Python repr in code generation.Ú_fieldsÚ_makec              3  óN   <"  € T F  p\        \        SV4      4      x € K  	  R # 5i©N)Ú_sanitize_for_reprÚgetattr)Ú.0Úfieldrp   s   & €rb   Ú	<genexpr>Ú%_sanitize_for_repr.<locals>.<genexpr>   s$   øé € ð +
ÙAL¸Ôœw s¨EÓ2×3Ð3Ãùs   ƒ"%c              3  ó8   "  € T F  p\        V4      x € K  	  R # 5iru   )rv   )rx   Úvs   & rb   rz   r{   ƒ   s   é € Ð8±C¨qÔ'¨×*Ð*³Cùó   ‚)Ú
isinstanceÚdictÚitemsrv   ÚlistÚtupleÚhasattrrw   Útyperr   r   Úvalue)rp   Úkr}   s   f  rb   rv   rv   w   s  ø€ ä�#”t×ÒØILÏÉÌÔUÉÁÀÔ" 1Ó%Ô'9¸!Ó'<Ò<ÉÒUÐUÜ�#”t×ÒÙ/2Ó3©s¨!Ô" 1Ö%©sÑ3Ð3ä�#”u×Ò¤'¨#¨y×"9Ò"9Ü”t˜C“y 'Ô*ô +
ØADÇÂó+
ó 
ð 	
ô �#”u×ÒßŒuÑ8±CÓ8�uÐ8ˆuÑ8±CÓ8Ó8Ð8Ü�#”t×ÒÜ! #§)¡)Ó,Ð,Ø€Jùó Vùâ3s   « EÁ)Ezir.CommBufferTypeÚWrapperLinec                  óJ   € ] tR t^�t$ R]R&   R]R&   R]R&   R]R&   R]R	&   R
tR# )ÚBenchmarkStorageGroupÚstrÚbuffer_nameÚintÚnbytesr   Údevicer   z&dict[str, tuple[list[int], list[int]]]Úinputsr`   N)Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__Ú__annotations__Ú__static_attributes__r`   rd   rb   rŠ   rŠ   �   s   ‡ àÓØƒKØƒKØƒJØ2×2rd   rŠ   c               ó    € V ^8„  d   QhRRRR/# )r\   ÚnodeÚ
BufferLiker_   ÚReuseKeyr`   )ra   s   "rb   rc   rc   ˜   s   € ÷ ñ ˜:ð ¨(ñ rd   c                ó  € \         P                  P                  V 4      pV P                  4       \         P                  P                  9  p\         P                  P
                  P                  V P                  4       4      p\         P                  P
                  P                  V P                  4       4      pV P                  4       V P                  4       \        \         P                  P                  P                  V4      4      VVV3# ru   )rH   ÚgraphÚget_allocation_storage_sizeÚget_nameÚunaligned_buffersÚ	schedulerÚget_buf_streamÚget_buf_mempoolÚget_device_or_errorÚ	get_dtyperE   ÚsizevarsÚsimplify)r˜   Ústorage_sizeÚ	alignmentÚstreamÚmempools   &    rb   Úbuffer_reuse_keyr«   ˜   s¸   € Ü—7‘7×6Ñ6°tÓ<€LØ—‘“¤q§w¡w×'@Ñ'@Ñ@€IÜ�W‰W×Ñ×-Ñ-¨d¯m©m«oÓ>€FÜ�g‰g×Ñ×/Ñ/°·±³Ó@€Gà× Ñ Ó"Ø�‰Óô 	”!—'‘'×"Ñ"×+Ñ+¨LÓ9Ó:ØØØð
ð 
rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r˜   r™   r_   ÚCommBufferReuseKeyr`   )ra   s   "rb   rc   rc   ª   s   € ÷ ñ  
ð Ð/Añ rd   c                óœ  € \         P                  P                  V 4      pV P                  4       p\	        V\
        P                  4      '       g   \        R \        V4       24      hV P                  4       V P                  4       \        \         P                  P                  P                  V4      4      VP                  VP                  3# )úexpected CommBufferLayout, got )rH   rœ   r�   Úget_output_specr   r,   ÚCommBufferLayoutÚAssertionErrorr…   r£   r¤   rE   r¥   r¦   Úcomm_buffer_typeÚ
group_name)r˜   r§   Úlayouts   &  rb   Úcomm_buffer_reuse_keyr¶   ª   s�   € ä—7‘7×6Ñ6°tÓ<€LØ×!Ñ!Ó#€FÜ�fœb×1Ñ1×2Ò2ÜÐ>¼tÀF»|¸nÐMÓNÐNà× Ñ Ó"Ø�‰ÓÜ”!—'‘'×"Ñ"×+Ñ+¨LÓ9Ó:Ø×ÑØ×Ñðð rd   c               ó    € V ^8„  d   QhRRRR/# )r\   Ú	input_bufr™   Ú
output_bufr`   )ra   s   "rb   rc   rc   ¹   s   € ÷ ñ  Zð ¸Zñ rd   c                ó´  € V P                  4       VP                  4       8w  d   R # V P                  4       VP                  4       8w  d   R # \        P                  P                  P                  \        P                  P                  V 4      4      p\        P                  P                  P                  \        P                  P                  V4      4      p\        V4      \        V4      8X  gg   \        P                  P                  P                  VRV,          4      '       d2   \        P                  P                  P                  W24      '       d   R# R # )Fgffffffî?T)
r£   r¤   rH   rœ   r¥   r¦   r�   rE   Ústatically_known_geqÚstatically_known_leq)r¸   r¹   Ú
input_sizeÚoutput_sizes   &&  rb   Úcan_match_buffer_sizer¿   ¹   sñ   € ð ×$Ñ$Ó&¨*×*HÑ*HÓ*JÔJÙà×ÑÓ 
× 4Ñ 4Ó 6Ô6Ùä—‘×!Ñ!×*Ñ*Ü	�‰×+Ñ+¨IÓ6ó€Jô —'‘'×"Ñ"×+Ñ+Ü	�‰×+Ñ+¨JÓ7ó€Kô 	�*Ó¤¨;Ó!7Ô7ô 	
�‰×Ñ×-Ñ-¨k¸4À*Õ;L×MÒMÜ�G‰G×Ñ×1Ñ1°+×JÒJáárd   c                óÞ  € \        V \        P                  4      '       d@   V P                  4       pVP                  VP
                  VP                  VP                  R3# . pT p\        V\        P                  \        P                  \        P                  34      '       d6   VP                  4       pVf   R# VP                  V4       VP                  pKt  \        V\        P                  4      '       g   R# V F0  pVP                  VP                  4       P                  8w  g   K.  Ru # 	  VP                  4       pVP                  VP
                  VP                  VP                  R3# )zí
Collapse a chain of ReinterpretView <- StorageBox
<- ReinterpretView <- StorageBox.... <- buffer wrappers if every layer
has the same offset as the innermost (base) buffer.

Returns:
    (size, stride, offset, dtype, collapsible: bool)
T)NNNNF)r   r,   ÚBufferÚ
get_layoutÚsizeÚstrideÚoffsetr   Ú	TensorBoxÚ
StorageBoxr/   ÚappendÚdata)rÉ   ÚlayÚlayoutsÚcurÚbase_lays   &    rb   Úcodegen_reinterpret_view_helperrÎ   Ø   s  € ô �$œŸ	™	×"Ò"Ø�o‰oÓˆØ�x‰x˜Ÿ™ S§Z¡Z°·±¸DÐ@Ð@à€GØ
€CÜ
�Sœ2Ÿ<™<¬¯©¼×8JÑ8JÐK×
LÒ
LØ�n‰nÓˆØŠ;Ø0Ð0Ø�‰�sÔØ�h‰hŠä�cœ2Ÿ9™9×%Ò%Ø,Ð,ó ˆØ�:‰:˜Ÿ™Ó)×0Ñ0Ö0Ø0Ò0ñ ð �~‰~Ó€HØ�=‰=˜(Ÿ/™/¨8¯?©?¸H¿N¹NÈDÐPÐPrd   .c               ó0   € V ^8„  d   QhRRRRRRRRR	R
RR/# )r\   Únamer‹   Úconfigszlist[triton.Config]Úgridszlist[TritonGrid]ÚwrapperúPythonWrapperCodegen | NoneÚoriginal_fxnode_nameú
str | Noner_   ztuple[str, str]r`   )ra   s   "rb   rc   rc     sL   € ÷ o&ñ o&Ø
ðo&à ðo&ð ðo&ð )ð	o&ð
 %ðo&ð ño&rd   c                ó.  a aaa€ \        4       oR  R loRR VV3R lllpRR V VV3R lllpRS  2pV! RV R	24       S'       d;   \        P                  P                  '       d   SP                  P                  4       M\        P                  ! 4       pSP                  4       ;_uu_ 4        T;_uu_ 4        \        P                  P                  '       dj   V'       db   \        P                  P                  '       dB   V\        P                  P                  9   d#   \        P                  P                  V,          p	MR.\        V4      ,          p	\        V4      ^8X  d+   V! V^ ,          V	^ ,          4      w  r«V! R
V
 2R
V 24       EM?\        V4      ^8:  d   \        R\        V4       24      h\        V4      \        V4      8w  d$   \        R\        V4       R\        V4       24      h\        4       p\        \        W!V	4      R RR7       Fµ  w  r­p. pVP                   '       dE   VP                    F4  pVR9  g   K  VP#                  RV RVP                   V,           24       K6  	  V'       d   RP%                  V4      pMRpV! W«4      w  r«RV RV
 2pVV9   d   K•  VP'                  V4       V! VRV RV 24       K·  	  RRR4       RRR4       VSP)                  4       3#   + '       g   i     L*; i  + '       g   i     L5; i)c               ó    € V ^8„  d   QhRRRR/# )r\   Úitemúint | sympy.Exprr_   r^   r`   )ra   s   "rb   rc   Ú6user_defined_kernel_grid_fn_code.<locals>.__annotate__
  s   € ÷ Mñ MÐ%5ð M¸*ñ Mrd   c                ór   € \        V \        P                  4      '       d   V # \        P                  ! V 4      # ru   )r   rh   r   ÚInteger)rÙ   s   &rb   Ú_convert_to_sympy_exprÚ@user_defined_kernel_grid_fn_code.<locals>._convert_to_sympy_expr
  s&   € Ü! $¬¯
©
×3Ò3ˆtÐL¼¿ºÀtÓ9LÐLrd   Nc               ó    € V ^8„  d   QhRRRR/# )r\   ÚgridÚ
TritonGridÚexample_gridzTritonGrid | Noner`   )ra   s   "rb   rc   rÛ     s   € ÷ 
ñ 
Øð
à'ñ
rd   c                ó   <€ Se   \        V 4      '       d   W 3# \        ;QJ d    . V3R lV  4       F  NK  	  5M! V3R lV  4       4      pV'       g   TpSP                  V4      \        P                  P
                  '       dH   SP                  \        ;QJ d    . V3R lV 4       F  NK  	  54      3# ! V3R lV 4       4      4      3# R3# )zÿ
This function returns a tuple of two values: the first one is for the real grid
which is used in the generated code; the second one is an example grid with
concrete values which is used in the autotune block to run the generated
kernels at compile time.
Nc              3  ó4   <"  € T F  pS! V4      x € K  	  R # 5iru   r`   )rx   ÚgrÞ   s   & €rb   rz   ÚKuser_defined_kernel_grid_fn_code.<locals>.determine_grid.<locals>.<genexpr>  s   øé € ÐC¹d¸Ñ1°!×4Ð4»dùs   ƒc              3  óZ   <"  € T F   pSP                  V\        V4      4      x € K"  	  R # 5iru   ©Úgenerate_example_arg_valuer…   )rx   ræ   rÓ   s   & €rb   rz   rç   "  s,   øé € ð á!-˜Að  ×:Ñ:¸1¼dÀ1»g×FÐFÛ!-ùó   ƒ(+)Úcallablerƒ   Úcodegen_python_shape_tupler*   ÚtritonÚautotune_at_compile_time)rá   rã   Ú
sympy_gridrÞ   rÓ   s   && €€rb   Údetermine_gridÚ8user_defined_kernel_grid_fn_code.<locals>.determine_grid  sÃ   ø€ ð Š?œh tŸnšnà�:Ðç”UÔC¹dÓC—U‘UÔC¹dÓCÓCˆ
ßØ%ˆLà×.Ñ.¨zÓ:ô —=‘=×9×9Ð9ð ×2Ñ2ß”Eô á!-ó—Eóð
ð 	
ð ô á!-óó óð
ð 	
ð ð
ð 	
rd   c               ó    € V ^8„  d   QhRRRR/# )r\   Úliner‹   rã   rÖ   r`   )ra   s   "rb   rc   rÛ   ,  s   € ÷ Jñ J˜ð J¨:ñ Jrd   c                óò   <€ SP                  V 4       S'       d]   \        P                  P                  '       d;   SSP                  9  d(   SP
                  P                  T;'       g    T 4       R # R # R # R # ru   )Ú	writeliner*   rî   rï   Úkernel_autotune_namesÚkernel_autotune_calls)rô   rã   rÐ   ÚoutputrÓ   s   &&€€€rb   rö   Ú3user_defined_kernel_grid_fn_code.<locals>.writeline,  s]   ø€ Ø×Ñ˜ÔçÜ—‘×6×6Ð6Ø˜G×9Ñ9Ô9à×)Ñ)×3Ñ3°L×4HÐ4HÀDÖIñ :ñ 7ñ rd   Úgrid_wrapper_for_údef z(meta):úreturn zexpected more than 1 grid, got z)expected len(grids) == len(configs), got ú != c                ó:   € \        V ^,          P                  4      # rf   ©ÚlenÚkwargs©Úxs   &rb   Ú<lambda>Ú2user_defined_kernel_grid_fn_code.<locals>.<lambda>X  s   € œc ! A¥$§+¡+Ô.rd   T©ÚkeyÚreversezmeta['z'] == ú and ÚTruezif z	: return ru   )Úmatrix_instr_nonkdimÚwaves_per_euÚkpack)r?   r*   rî   rï   rø   ÚindentÚ
contextlibÚnullcontextrH   rœ   Úautotuning_gridsr  r²   r"   ÚsortedÚzipr  rÈ   ÚjoinÚaddÚgetvalue)rÐ   rÑ   rÒ   rÓ   rÕ   rñ   rö   Úfn_nameÚkernel_autotune_calls_indentÚexample_gridsrá   rã   ÚseenÚcÚ
guardslistÚkwargÚguardsÚ	statementrÞ   rù   s   f&&f&             @@rb   Ú user_defined_kernel_grid_fn_coder!    s“  û€ ô Ó€FõM÷
ó 
÷>Jô Jð " $ Ð(€GÙ��W�I˜WÐ%Ô&÷ ”v—}‘}×=×=Ð=ð 	×%Ñ%×,Ñ,Ô.ä×#Ò#Ó%ð !ð
 
�‰��×6Ò6ä�M‰M×2×2Ð2ß$Ü—‘×(×(Ð(Ø$¬¯©×(@Ñ(@Ô@äŸG™G×4Ñ4Ð5IÕJ‰Mà!˜F¤S¨£ZÕ/ˆMÜˆu‹:˜Œ?Ù!/°°aµ¸-ÈÕ:JÓ!KÑˆDÙ˜ ˜vÐ&¨'°,°Ð(@ÖAä�5‹z˜QŒÜ$Ð'FÄsÈ5ÃzÀlÐ%SÓTÐTÜ�5‹zœS ›\Ô)Ü$Ø?ÄÀEÃ
¸|È4ÔPSÐT[ÓP\È~Ð^óð ô %/£LˆDô *0Ü�E MÓ2Ù.Ø÷*Ñ%�˜ð
  �
Ø—8—8�8à!"§¤˜Ø ð )ö ð
 '×-Ñ-°°u°g¸VÀAÇHÁHÈUÅOÐCTÐ.UÖVñ "*÷ Ø$Ÿ\™\¨*Ó5‘Fà#�FÙ%3°DÓ%GÑ"�Ø! & ¨°4°&Ð9�	Ø Ô$ÙØ—‘˜Ô#Ù˜) s¨6¨(°)¸L¸>Ð%JÖKñ1*÷5 7�ðh �F—O‘OÓ%Ð%Ð%÷i 7×6ú��ús>   Â.LÂ9(K1	Ã"K1	ÄEK1	É3K1	É<AK1	ËLË1LË<LÌL	c               ó    € V ^8„  d   QhRRRR/# )r\   Úepilogue_fusionú$tuple[ir.ComputedBuffer, str] | Noner_   r‹   r`   )ra   s   "rb   rc   rc   s  s    € ÷ q&ñ q&ØAðq&àñq&rd   c                ó$  aaaaaaa	€ \        4       oV P                  pV'       d
   V^,          pSP                  VRR7       ^ RIo^ RIHo ^ RIHo ^ RIHo	 \        V P                  .4      oVVVVVVV	3R loS! V 4       SP                  4       # )z¦
Given a triton kernel function pointer collect the transitive closure of
its dependencies

epilogue_fusion: Optional[(fused epilogue node, modified kerel src code)]
T©ÚstripN)ÚJITFunction)Ú	constexprr   c           	     óò  <€ \        R  \        P                  ! V P                  4       4       4      pV P                  P                  P                  R/ 4      pV P                  P                  P                   EFú  pVS9   d   K  W0P                  P                  9   g   K)  V P                  P                  V,          p\        VS	4      '       dZ   S
P                  4        S
P                  R4       S
P                  VP                  RR7       SP                  V4       S! V4       K±  \        SR4      '       Edx   \        VSP                  P                   P"                  4      '       EdG   RVP                  9   d˜   RS9  d‘   VP                  P                  P                  R4      pV'       dd   \        VR4      '       dR   VP$                  P'                  R4      '       d1   S
P                  R	VP$                   R
24       SP                  R4       S
P                  4        S
P                  R4       S
P                  VP                  RR7       W4P                  P(                  8w  d+   S
P                  V RVP                  P(                   24       SP                  V4       S! V4       EK;  \        V\*        \,        \.        S34      '       dÎ   S
P                  4        \        VS4      '       d   RVP0                  : R2pMV: pVP                  V4      ;p'       dS   \        V\2        4      '       d   RVP$                   RVP(                   2pMRV: 2pS
P                  V V RV 24       MS
P                  V RV 24       SP                  V4       EK*  W19   g   EK3  VR8w  g   EK=  \        VR4      '       g   EKR  VP$                  P'                  R4      '       g   EKv  \        VS4      '       d"   S
P                  V RVP4                   24       M@\        VR4      '       d/   S
P                  R	VP$                   RVP(                   RV 24       SP                  V4       EKý  	  R# )c              3  ó`   "  € T F$  pVP                   R 8X  g   K  VP                  x € K&  	  R# 5i)ÚLOAD_GLOBALN)ÚopnameÚargval)rx   Úinsts   & rb   rz   Ú^user_defined_triton_kernel_transitive_closure_source_code.<locals>.traverse.<locals>.<genexpr>�  s)   é € ð '
á3�Ø�{‰{˜mÑ+ô ˆD�KŽKÛ3ùs   ‚.š.r•   z@triton.jitTr&  Úconstexpr_functionr   r’   rî   zfrom z import dtype as dtypez@triton.constexpr_functionú = ztl.constexpr(Ú)ú: Ú.Útlz = tl.r‘   z import z as N)r"   ÚdisÚBytecodeÚfnÚ__globals__ÚgetÚ__code__Úco_namesr   Únewlinerö   ÚspliceÚsrcr  r„   ÚruntimeÚjitÚConstexprFunctionr’   Ú
startswithr‘   r�   r‹   Úboolr†   r…   rÐ   )Ú
cur_kernelÚunqualified_loadsÚglobal_annotationsÚsymbol_nameÚsymbolÚdtype_symbolÚ
symbol_strÚ
annotationÚannotation_coder(  Úcompile_wrapperr)  Úsymbols_includedÚtraverserî   Útriton_dtypes   &        €€€€€€€rb   rQ  ÚKuser_defined_triton_kernel_transitive_closure_source_code.<locals>.traverse‹  s»  ø€ ô
 'ñ '
äŸš Z§]¡]Ô3ó'
ó 
Ðð
 (Ÿ]™]×6Ñ6×:Ñ:Ð;LÈbÓQÐØ%Ÿ=™=×1Ñ1×:Õ:ˆKØÐ.Ô.ÙØŸm™m×7Ñ7Ö7Ø#Ÿ™×2Ñ2°;Õ?�Ü˜f k×2Ò2Ø#×+Ñ+Ô-Ø#×-Ñ-¨mÔ<Ø#×*Ñ*¨6¯:©:¸TÐ*ÔBØ$×(Ñ(¨Ô5Ù˜VÖ$Ü˜VÐ%9×:Ó:¼zØØ—N‘N×&Ñ&×8Ñ8÷@ó @ð
  &§*¡*Ô,°Ð@PÔ1PØ'-§y¡y×'<Ñ'<×'@Ñ'@ÀÓ'I˜ç(Ü '¨°l× CÒ CØ ,× 7Ñ 7× BÑ BÀ8× LÒ Là+×5Ñ5Ø"'¨×(?Ñ(?Ð'@Ð@VÐ Wôð -×0Ñ0°Ô9Ø#×+Ñ+Ô-Ø#×-Ñ-Ð.JÔKØ#×*Ñ*¨6¯:©:¸TÐ*ÔBØ"§i¡i×&8Ñ&8Ô8Ø'×1Ñ1Ø*˜m¨3¨v¯y©y×/AÑ/AÐ.BÐCôð %×(Ñ(¨Ô5Ù˜V×$Ü ¬¬c´4¸Ð(C×DÒDØ#×+Ñ+Ô-Ü! &¨)×4Ò4Ø'4°V·\±\Ñ4DÀAÐ%F™
à(.¡z˜
Ø%7×%;Ñ%;¸KÓ%HÐH�zÖHÜ% j´$×7Ò7à"$ Z×%:Ñ%:Ð$;¸1¸Z×=PÑ=PÐ<QÐ Rñ ,ð 13°:±.Ð.A˜OØ'×1Ñ1Ø*˜m¨OÐ+<¸CÀ
¸|ÐLõð (×1Ñ1°[°MÀÀZÀLÐ2QÔRØ$×(Ñ(¨×5à×4Ø# t×+Ü ¨×5Õ5ð ×)Ñ)×4Ñ4°X×>Õ>ô " &¨,×7Ò7Ø'×1Ñ1°[°MÀÈÏÉÀ}Ð2UÕVÜ  ¨×4Ò4Ø'×1Ñ1Ø# F×$5Ñ$5Ð#6°h¸v¿¹Ð>OÈtÐT_ÐS`Ðaôð %×(Ñ(¨×5óW ;rd   )r?   r@  r?  rî   r(  Útriton.languager)  Útriton.language.corer   r"   r‘   r  )
Úkernelr#  Ú
kernel_srcr(  rO  r)  rP  rQ  rî   rR  s
   && @@@@@@@rb   Ú9user_defined_triton_kernel_transitive_closure_source_coderX  s  s~   þ€ ô %Ó&€OØ—‘€JßØ$ QÕ'ˆ
Ø×Ñ˜:¨TÐÔ2ó Ý"Ý)Ý:ô " 6§?¡?Ð"3Ó4Ð÷V6ó V6ñp ˆVÔØ×#Ñ#Ó%Ð%rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r@  r‹   r_   r`   )ra   s   "rb   rc   rc   ç  s   € ÷ +ñ +°#ð +¸#ñ +rd   c                óš   € V P                  RR4      p \        P                  '       d   V P                  RR4      p V P                  RR4      # )zMEscape src for a '''...''' literal, optionally nested in an r"""...""" block.Ú\z\\z"""z\"\"\"z'''z\'\'\')Úreplacer*   Úcpp_wrapper)r@  s   &rb   Ú(_escape_triton_kernel_source_for_wrapperr^  ç  sA   € à
�+‰+�d˜FÓ
#€CÜ××Ðð �k‰k˜% Ó-ˆØ�;‰;�u˜kÓ*Ð*rd   c                  ó2   € ] tR tRt$ R]R&   R]R&   R tRtR# )	ÚSymbolicCallArgiò  úsympy.SymbolÚinnerr^   Ú
inner_exprc                	ó,   € \        V P                  4      # ru   )r‹   rb  ©Úselfs   &rb   Ú__str__ÚSymbolicCallArg.__str__ø  s   € Ü�4—:‘:‹Ðrd   r`   N)r‘   r’   r“   r”   r•   rg  r–   r`   rd   rb   r`  r`  ò  s   ‡ àÓàÓörd   r`  c                  óp   a € ] tR tRtV 3R ltR R ltR R ltR R ltR	 R
 ltR R lt	R R lt
RtV ;t# )ÚMemoryPlanningStateiü  c                	ó®   <€ \         SV `  4        \        P                  ! \        4      V n        \        P                  ! \        4      V n        ^ V n        R# ©r   N)ÚsuperÚ__init__ÚcollectionsÚdefaultdictr‚   Ú
reuse_poolÚcomm_buffer_reuse_poolÚtotal_allocated_buffer_size)rf  Ú	__class__s   &€rb   rn  ÚMemoryPlanningState.__init__ý  sG   ø€ Ü‰ÑÔô ×#Ò#¤DÓ)ð 	Œô ×#Ò#¤DÓ)ð 	Ô#ð 12ˆÖ(rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r  rš   r_   rE  r`   )ra   s   "rb   rc   Ú MemoryPlanningState.__annotate__	  s   € ÷ 4ñ 4 ð 4¨Tñ 4rd   c                	óL   € \        V P                  P                  VR 4      4      # ru   )rE  rq  r;  ©rf  r  s   &&rb   Ú__contains__Ú MemoryPlanningState.__contains__	  s   € Ü�D—O‘O×'Ñ'¨¨TÓ2Ó3Ð3rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r  rš   r_   ÚFreeIfNotReusedLiner`   )ra   s   "rb   rc   rw    s   € ÷ ñ �xð Ð$7ñ rd   c                	ó‚   € V P                   V,          P                  4       pVP                  '       d   \        R 4      hV# )z0expected popped reuse pool item to not be reused)rq  ÚpopÚ	is_reusedr²   ©rf  r  rÙ   s   && rb   r  ÚMemoryPlanningState.pop  s3   € Ø�‰˜sÕ#×'Ñ'Ó)ˆØ�>�>ˆ>Ü Ð!SÓTÐTØˆrd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   r  rš   rÙ   r}  r_   ÚNoner`   )ra   s   "rb   rc   rw    s"   € ÷ *ñ *˜ð *Ð(;ð *Àñ *rd   c                	ó„   € VP                   '       d   \        R 4      hV P                  V,          P                  V4       R# )z0expected pushed reuse pool item to not be reusedN)r€  r²   rq  rÈ   r�  s   &&&rb   ÚpushÚMemoryPlanningState.push  s.   € Ø�>�>ˆ>Ü Ð!SÓTÐTØ�‰˜Õ×#Ñ# DÖ)rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r  r­   r_   rE  r`   )ra   s   "rb   rc   rw    s   € ÷ @ñ @Ð(:ð @¸tñ @rd   c                	óL   € \        V P                  P                  VR 4      4      # ru   )rE  rr  r;  ry  s   &&rb   Úcomm_buffer_containsÚ(MemoryPlanningState.comm_buffer_contains  s    € Ü�D×/Ñ/×3Ñ3°C¸Ó>Ó?Ð?rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r  r­   r_   r}  r`   )ra   s   "rb   rc   rw    s   € ÷ ñ Ð#5ð Ð:Mñ rd   c                	ó‚   € V P                   V,          P                  4       pVP                  '       d   \        R 4      hV# )z<expected popped comm buffer reuse pool item to not be reused)rr  r  r€  r²   r�  s   && rb   Úcomm_buffer_popÚ#MemoryPlanningState.comm_buffer_pop  s;   € Ø×*Ñ*¨3Õ/×3Ñ3Ó5ˆØ�>�>ˆ>Ü ØNóð ð ˆrd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   r  r­   rÙ   r}  r_   r„  r`   )ra   s   "rb   rc   rw  "  s$   € ÷ 6ñ 6Ø%ð6Ø-@ð6à	ñ6rd   c                	ó„   € VP                   '       d   \        R 4      hV P                  V,          P                  V4       R# )z<expected pushed comm buffer reuse pool item to not be reusedN)r€  r²   rr  rÈ   r�  s   &&&rb   Úcomm_buffer_pushÚ$MemoryPlanningState.comm_buffer_push"  s8   € ð �>�>ˆ>Ü ØNóð ð 	×#Ñ# CÕ(×/Ñ/°Ö5rd   )rr  rq  rs  )r‘   r’   r“   r”   rn  rz  r  r†  rŠ  rŽ  r’  r–   Ú__classcell__©rt  s   @rb   rj  rj  ü  s,   ø† õ
2õ4õõ*õ
@õ÷6ó 6rd   rj  c                  ó"   € ] tR tRtR R ltRtR# )ÚHasWriteLinei,  c               ó    € V ^8„  d   QhRRRR/# )r\   rô   z$LineContext | DeferredLineBase | strr_   r„  r`   )ra   s   "rb   rc   ÚHasWriteLine.__annotate__-  s   € ×PÑPÐBÐPÀtÑPrd   c                	ó   € R # ru   r`   ©rf  rô   s   &&rb   rö   ÚHasWriteLine.writeline-  s   € ÉSrd   r`   N)r‘   r’   r“   r”   rö   r–   r`   rd   rb   r—  r—  ,  s   † ßPÑPrd   r—  c                  ó.   € ] tR tRtR R ltR R ltRtR# )rˆ   i0  c               ó    € V ^8„  d   QhRRRR/# )r\   Úcoder   r_   r„  r`   )ra   s   "rb   rc   ÚWrapperLine.__annotate__1  s   € ÷ Vñ V˜Cð V Dñ Vrd   c                	ó0   € \        R \        V 4       24      h)z#Codegen not yet supported for type ©ÚNotImplementedErrorr…   ©rf  rŸ  s   &&rb   ÚcodegenÚWrapperLine.codegen1  s   € Ü!Ð$GÌÈTË
À|Ð"TÓUÐUrd   c               ó    € V ^8„  d   QhRRRR/# ©r\   Ú	converterrZ   r_   ÚFxConversionFuncr`   )ra   s   "rb   rc   r   4  s   € ÷ Yñ Y Kð YÐ4Dñ Yrd   c                	ó0   € \        R \        V 4       24      h)z&FX codegen not yet supported for type r¢  ©rf  r©  s   &&rb   Ú
codegen_fxÚWrapperLine.codegen_fx4  s   € Ü!Ð$JÌ4ÐPTË:È,Ð"WÓXÐXrd   r`   N©r‘   r’   r“   r”   r¥  r­  r–   r`   rd   rb   rˆ   rˆ   0  s   † õV÷Yñ Yrd   c                  óP   € ] tR tRt$ R]R&   R]R&   R R ltR R	 ltR
 R ltRtR# )ÚEnterSubgraphLinei8  ÚPythonWrapperCodegenrÓ   rW   rœ   c               ó   € V ^8„  d   QhRR/# ©r\   r_   r„  r`   )ra   s   "rb   rc   ÚEnterSubgraphLine.__annotate__=  s   € ÷ Fñ F˜tñ Frd   c                	ód   € V P                   P                  V P                   P                  4       R # ru   )rÓ   Úpush_computed_sizesÚcomputed_sizesre  s   &rb   Ú__post_init__ÚEnterSubgraphLine.__post_init__=  s   € Ø�‰×(Ñ(¨¯©×)DÑ)DÖErd   c               ó    € V ^8„  d   QhRRRR/# ©r\   rŸ  r?   r_   r„  r`   )ra   s   "rb   rc   rµ  @  s   € ÷ ñ ˜Nð ¨tñ rd   c                	óp   € V P                   P                  V P                  4       VP                  4        R # ru   )rÓ   Úpush_codegened_graphrœ   Ú	do_indentr¤  s   &&rb   r¥  ÚEnterSubgraphLine.codegen@  s"   € Ø�‰×)Ñ)¨$¯*©*Ô5Ø�‰Örd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rµ  D  s   € ÷ 2ñ 2 Kð 2Ð4Dñ 2rd   c                	ó   € VP                   # ru   )Ú_generate_enter_subgraphr¬  s   &&rb   r­  ÚEnterSubgraphLine.codegen_fxD  s   € Ø×1Ñ1Ð1rd   r`   N©	r‘   r’   r“   r”   r•   r¹  r¥  r­  r–   r`   rd   rb   r±  r±  8  s#   ‡ à!Ó!ØÓõFõ÷2ñ 2rd   r±  c                  óN   € ] tR tRt$ R]R&   R]R&   R R lt]R R	 l4       tR
tR# )Ú
SwitchLineiH  r²  rÓ   z	ir.Switchr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   ÚSwitchLine.__annotate__M  s   € ÷ >ñ >˜Nð >¨tñ >rd   c                	ó   € \        R 4      h)zOnly supports FX codegen©r£  r¤  s   &&rb   r¥  ÚSwitchLine.codegenM  s   € Ü!Ð"<Ó=Ð=rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rÉ  Q  s   € ÷ *ñ *˜kð *Ð.>ñ *rd   c                	ó   € V P                   # ru   )Ú_generate_switch©r©  s   &rb   r­  ÚSwitchLine.codegen_fxP  s   € à×)Ñ)Ð)rd   r`   N©	r‘   r’   r“   r”   r•   r¥  Ústaticmethodr­  r–   r`   rd   rb   rÇ  rÇ  H  s&   ‡ à!Ó!Ø
ƒOõ>ð ô*ó ô*rd   rÇ  c                  óD   € ] tR tRt$ R]R&   R R lt]R R l4       tRtR	# )
ÚCommentLineiU  rB   rô   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   ÚCommentLine.__annotate__Y  s   € ÷ "ñ "˜Nð "¨tñ "rd   c                	ó<   € VP                  V P                  4       R # ru   )rö   rô   r¤  s   &&rb   r¥  ÚCommentLine.codegenY  s   € Ø�‰�t—y‘yÖ!rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r×  ]  s   € ÷ +ñ +˜kð +Ð.>ñ +rd   c                	ó   € V P                   # ru   )Ú_generate_commentrÐ  s   &rb   r­  ÚCommentLine.codegen_fx\  s   € à×*Ñ*Ð*rd   r`   NrÒ  r`   rd   rb   rÕ  rÕ  U  s!   ‡ à
Óõ"ð ô+ó ô+rd   rÕ  c                  óN   € ] tR tRt$ R]R&   R]R&   R R lt]R R	 l4       tR
tR# )ÚDynamicScalarLineia  r²  rÓ   zir.DynamicScalarr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   ÚDynamicScalarLine.__annotate__f  s   € ÷ 8ñ 8˜Nð 8¨tñ 8rd   c                	óP   € V P                   P                  V P                  4       R # ru   )rÓ   Ú_codegen_dynamic_scalarr˜   r¤  s   &&rb   r¥  ÚDynamicScalarLine.codegenf  s   € Ø�‰×,Ñ,¨T¯Y©YÖ7rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rá  j  s   € ÷ 2ñ 2˜kð 2Ð.>ñ 2rd   c                	ó   € V P                   # ru   )Ú_generate_dynamic_scalarrÐ  s   &rb   r­  ÚDynamicScalarLine.codegen_fxi  s   € à×1Ñ1Ð1rd   r`   NrÒ  r`   rd   rb   rß  rß  a  s'   ‡ à!Ó!Ø
Óõ8ð ô2ó ô2rd   rß  c                  óF   € ] tR tRt$ R]R&   R R ltR R ltR R	 ltR
tR# )ÚExitSubgraphLinein  r²  rÓ   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   ÚExitSubgraphLine.__annotate__r  s   € ÷ Hñ H˜tñ Hrd   c                	óX   € V P                   P                  4       V P                   n        R # ru   )rÓ   Úpop_computed_sizesr¸  re  s   &rb   r¹  ÚExitSubgraphLine.__post_init__r  s   € Ø&*§l¡l×&EÑ&EÓ&Gˆ�‰Ö#rd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   rì  u  s   € ÷ ñ ˜Nð ¨tñ rd   c                	óZ   € V P                   P                  4        VP                  4        R # ru   )rÓ   Úpop_codegened_graphÚdo_unindentr¤  s   &&rb   r¥  ÚExitSubgraphLine.codegenu  s   € Ø�‰×(Ñ(Ô*Ø×ÑÖrd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rì  y  s   € ÷ 1ñ 1 Kð 1Ð4Dñ 1rd   c                	ó   € VP                   # ru   )Ú_generate_exit_subgraphr¬  s   &&rb   r­  ÚExitSubgraphLine.codegen_fxy  s   € Ø×0Ñ0Ð0rd   r`   NrÅ  r`   rd   rb   rê  rê  n  s   ‡ à!Ó!õHõ÷1ñ 1rd   rê  c               ó    € V ^8„  d   QhRRRR/# )r\   Ú
device_idxr�   r_   z	int | strr`   )ra   s   "rb   rc   rc   ‚  s   € ÷ 
Cñ 
C Sð 
C¨Yñ 
Crd   c                ó2   € \        4       '       d   \        # T # )aö  Device index to emit in the wrapper. Under compile-on-one-rank, the call()-local
runtime current-device variable (so the wrapper is byte-identical across ranks);
otherwise the literal compile-time index.

Safe to ignore ``device_idx`` under CooR: make_fx rejects any operand that is not the
current accelerator (_coor_check_current_accelerator), so every ``device_idx`` reaching
codegen is already the current device's index; redirecting it to the runtime variable
just follows each rank's own device.
)r   Ú_COOR_DEVICE_IDX_VAR)rú  s   &rb   Ú_coor_device_idx_refrý  ‚  s   € ô $1§?¢?ÔÐB¸
ÐBrd   c                  óD   € ] tR tRt$ R]R&   R]R&   R R ltR R	 ltR
tR# )ÚEnterDeviceContextManagerLinei�  r�   rú  ú
int | NoneÚlast_seen_device_guard_indexc               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú*EnterDeviceContextManagerLine.__annotate__”  s   € ÷ &?ñ &?˜Nð &?¨tñ &?rd   c                	ó¤  € \         P                  P                  '       EdL   VP                  R 4       V P                  f„   VP                  \         P                  P                  P                  4        R24       VP                  \         P                  P                  P                  4        RV P                   R24       R# \         P                  P                  '       d'   V P                  V P                  8w  d   \        R4      h\         P                  P                  '       d!   \         P                  P                  '       d"   VP                  RV P                   R24       R# R# \        4       '       dD   \        pVP                  V R\         P                  P                  P!                  4        24       MV P                  pVP                  R\         P                  P                  P#                  V4       R	24       VP%                  4        VP                  \         P                  P                  P'                  V4      4       R# )
Ú
Nz) stream_guard(stream, this->device_idx_);z device_guard(z);z4AOTInductor only supports running on one CUDA devicezdevice_guard.set_index(r2  úwith Ú:)rH   rœ   r]  rö   r  Úwriteline_aotÚ
device_opsÚcpp_aoti_stream_guardÚwriteline_jitÚcpp_aoti_device_guardrú  Úaot_moder²   Úis_dual_wrapper_moder   rü  Úcurrent_device_idx_exprÚdevice_guardr¿  Ú
set_device)rf  rŸ  Úidxs   && rb   r¥  Ú%EnterDeviceContextManagerLine.codegen”  s°  € Ü�7‰7××ÑØ�N‰N˜4Ô ð ×0Ñ0Ò8Ø×"Ñ"Ü—w‘w×)Ñ)×?Ñ?ÓAÐBÐBkÐlôð ×"Ñ"Ü—w‘w×)Ñ)×?Ñ?ÓAÐBÀ.ÐQU×Q`ÑQ`ÐPaÐacÐdöô —7‘7×#×#Ð#ð ×8Ñ8¸D¿O¹OÔKÜ,ØRóð ô —w‘w×'×'Ð'¬1¯7©7×+G×+GÐ+GØ×&Ñ&Ð)@ÀÇÁÐ@QÐQSÐ'TÖUñ ,Hô
 �Šô +�Ø—‘Ø�e˜3œqŸw™w×1Ñ1×IÑIÓKÐLÐMõð —o‘o�Ø�N‰N˜U¤1§7¡7×#5Ñ#5×#BÑ#BÀ3Ó#GÐ"HÈÐJÔKØ�N‰NÔØ�N‰Nœ1Ÿ7™7×-Ñ-×8Ñ8¸Ó=Ö>rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r  ¼  s   € ÷ @ñ @ Kð @Ð4Dñ @rd   c                	ó   € VP                   # ru   )Ú&_generate_enter_device_context_managerr¬  s   &&rb   r­  Ú(EnterDeviceContextManagerLine.codegen_fx¼  s   € Ø×?Ñ?Ð?rd   r`   N©r‘   r’   r“   r”   r•   r¥  r­  r–   r`   rd   rb   rÿ  rÿ  �  s   ‡ àƒOØ",Ó,õ&?÷P@ñ @rd   rÿ  c                  ó.   € ] tR tRtR R ltR R ltRtR# )ÚExitDeviceContextManagerLineiÀ  c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú)ExitDeviceContextManagerLine.__annotate__Á  s   € ÷ ñ ˜Nð ¨tñ rd   c                	ój   € \         P                  P                  '       g   VP                  4        R # R # ru   ©rH   rœ   r]  ró  r¤  s   &&rb   r¥  Ú$ExitDeviceContextManagerLine.codegenÁ  s#   € Ü�w‰w×"×"Ð"Ø×ÑÖñ #rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r  Å  s   € ÷ ?ñ ? Kð ?Ð4Dñ ?rd   c                	ó   € VP                   # ru   )Ú%_generate_exit_device_context_managerr¬  s   &&rb   r­  Ú'ExitDeviceContextManagerLine.codegen_fxÅ  s   € Ø×>Ñ>Ð>rd   r`   Nr¯  r`   rd   rb   r  r  À  s   † õ÷?ñ ?rd   r  c                  óD   € ] tR tRt$ R]R&   R]R&   R R ltR R	 ltR
tR# )ÚExternKernelAllocLineiÉ  r²  rÓ   úir.ExternKernelAllocr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú"ExternKernelAllocLine.__annotate__Î  s   € ÷ Kñ K˜Nð K¨tñ Krd   c                	ó®   € V P                   p. VP                  4       OVP                  4       OpV P                  P	                  V P                   V4       R # ru   )r˜   Úcodegen_argsÚcodegen_kwargsrÓ   Ú$_generate_extern_kernel_alloc_helper)rf  rŸ  r˜   Úargss   &&  rb   r¥  ÚExternKernelAllocLine.codegenÎ  sD   € Ø�y‰yˆØ=�×"Ñ"Ó$Ð= t×':Ñ':Ó'<Ð=ˆØ�‰×9Ñ9¸$¿)¹)ÀTÖJrd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r(  Ó  s   € ÷ 7ñ 7 Kð 7Ð4Dñ 7rd   c                	ó   € VP                   # ru   )Ú_generate_extern_kernel_allocr¬  s   &&rb   r­  Ú ExternKernelAllocLine.codegen_fxÓ  s   € Ø×6Ñ6Ð6rd   r`   Nr  r`   rd   rb   r%  r%  É  s   ‡ à!Ó!Ø
ÓõK÷
7ñ 7rd   r%  c                  óD   € ] tR tRt$ R]R&   R]R&   R R ltR R	 ltR
tR# )ÚExternKernelOutLinei×  r²  rÓ   úir.ExternKernelOutr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú ExternKernelOutLine.__annotate__Ü  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c           	     	ón  € V P                   p. VP                  4       OVP                  R R7      OpVP                  4       p\        P
                  P                  '       d   VP                  R8X  d   RpMVP                  4       pVP                  4       ;p'       d   VP                  M\        P
                  P                  pV P                  P                  TVP                  4       VP                  '       d   VP                  P                  4       MRVVV P                   P                  4       4       R# )T)Úskip_outztorch::inductor::_mm_plus_mmÚaoti_torch__mm_plus_mm_outN)r˜   r*  r+  Úget_kernel_namerH   rœ   r]  Úcpp_kernel_nameÚ
get_devicer…   Údevice_typerÓ   Ú"_generate_extern_kernel_out_helperÚcodegen_referenceÚoutput_viewÚget_stack_traces)rf  rŸ  r˜   r-  Úkernel_nameÚdr�   s   &&     rb   r¥  ÚExternKernelOutLine.codegenÜ  sï   € Ø�y‰yˆØJ�×"Ñ"Ó$ÐJ t×':Ñ':ÀDÐ':Ó'IÐJˆØ×*Ñ*Ó,ˆä�G‰G××ÐØ×$Ñ$Ð(FÔFð 7‰Kà×.Ñ.Ó0ˆKØ!%§¡Ó!2Ð2˜AÖ2�—’¼¿¹×9LÑ9LˆØ�‰×7Ñ7ØØ×"Ñ"Ó$Ø48×4D×4DÐ4DˆD×Ñ×.Ñ.Ô0È$ØØØ�I‰I×&Ñ&Ó(ö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r7  ò  ó   € ÷ 5ñ 5 Kð 5Ð4Dñ 5rd   c                	ó   € VP                   # ru   )Ú_generate_extern_kernel_outr¬  s   &&rb   r­  ÚExternKernelOutLine.codegen_fxò  ó   € Ø×4Ñ4Ð4rd   r`   Nr  r`   rd   rb   r4  r4  ×  s   ‡ à!Ó!Ø
Óõ
÷,5ñ 5rd   r4  c                  ó<   € ] tR tRt$ RtR]R&   R]R&   R R ltR	tR
# )ÚExternKernelMultiOutLineiö  z²Codegen line for multi-output .out() variant calls.

Generates a kernel call with pre-allocated output buffers passed as
keyword arguments. E.g. kernel(x, out0=buf0, out1=buf1).
r²  rÓ   úir.ExternKernelMultiOutr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú%ExternKernelMultiOutLine.__annotate__  s   € ÷ 
Pñ 
P˜Nð 
P¨tñ 
Prd   c           	     	óŠ  € V P                   pVP                  4       p. VP                  4       OVP                  4       Op\	        VP
                  VP                  4       F)  w  rVVP                  V R VP                  4        24       K+  	  VP                  VP                  4        RV RRP                  V4       R24       R# )Ú=r2  Ú(ú, r3  N)r˜   r;  r*  r+  r  Úout_arg_namesÚout_variant_output_nodesrÈ   rž   rö   r  )rf  rŸ  r˜   rC  r-  Úout_nameÚout_nodes   &&     rb   r¥  Ú ExternKernelMultiOutLine.codegen  s±   € Ø�y‰yˆØ×*Ñ*Ó,ˆà=�×"Ñ"Ó$Ð= t×':Ñ':Ó'<Ð=ˆÜ"%Ø×Ñ × =Ñ =ö#
ÑˆHð �K‰K˜8˜* A h×&7Ñ&7Ó&9Ð%:Ð;Ö<ñ#
ð
 	�‰˜$Ÿ-™-›/Ð*¨#¨k¨]¸!¸D¿I¹IÀd»OÐ;LÈAÐNÖOrd   r`   N©r‘   r’   r“   r”   Ú__doc__r•   r¥  r–   r`   rd   rb   rM  rM  ö  s!   ‡ ñð "Ó!Ø
!Ó!÷
Pñ 
Prd   rM  c                  óD   € ] tR tRt$ R]R&   R]R&   R R ltR R	 ltR
tR# )ÚFreeLinei  r²  rÓ   úBufferLike | ir.TorchBindObjectr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   ÚFreeLine.__annotate__  s   € ÷ Añ A˜Nð A¨tñ Ard   c                	ó*  € V P                   P                  4       \        P                  P                  9   d(   \        R V P                   P                  4        R24      hVP                  V P                  P                  V P                   4      4       R# ©úbuffer ú was removed but is being freedN)	r˜   rž   rH   rœ   Úremoved_buffersr²   rö   rÓ   Úmake_buffer_freer¤  s   &&rb   r¥  ÚFreeLine.codegen  sj   € Ø�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ü Ø˜$Ÿ)™)×,Ñ,Ó.Ð/Ð/NÐOóð ð 	�‰�t—|‘|×4Ñ4°T·Y±YÓ?Ö@rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r`    ó   € ÷ (ñ ( Kð (Ð4Dñ (rd   c                	ó   € VP                   # ru   )Ú_generate_freer¬  s   &&rb   r­  ÚFreeLine.codegen_fx  ó   € Ø×'Ñ'Ð'rd   r`   Nr  r`   rd   rb   r]  r]    s   ‡ à!Ó!Ø
)Ó)õA÷(ñ (rd   r]  c                  ó¶   € ] tR tRt$ 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tR]R&   R R ltR R ltRtR# )ÚKernelCallLinei  r²  rÓ   r‹   rC  ztuple[Any, ...]Ú	call_argsÚraw_keysÚraw_argsú	list[str]Ú	arg_typesrE  rî   r2   Útriton_metazdict[str, Any] | NoneÚinductor_metaútorch.devicer�   Ú
graph_namerÕ   Nr   Úcurrent_stream_idxc               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   ÚKernelCallLine.__annotate__.  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c                	óD  € V P                   P                  V P                  V P                  V P                  V P
                  V P                  V P                  V P                  V P                  V P                  V P                  V P                  V P                  R 7       R# ))
rî   rt  rq  rr  ru  rv  r�   rx  rÕ   ry  N)rÓ   Ú_generate_kernel_call_helperrC  rp  rî   rt  rq  rr  ru  rv  r�   rx  rÕ   ry  r¤  s   &&rb   r¥  ÚKernelCallLine.codegen.  sw   € Ø�‰×1Ñ1Ø×ÑØ�N‰NØ—;‘;Ø—n‘nØ—]‘]Ø—]‘]Ø×(Ñ(Ø×,Ñ,Ø—;‘;Ø—‘Ø!%×!:Ñ!:Ø#×6Ñ6ð 	2ö 	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r{  >  ó   € ÷ /ñ / Kð /Ð4Dñ /rd   c                	ó   € VP                   # ru   )Ú_generate_kernel_callr¬  s   &&rb   r­  ÚKernelCallLine.codegen_fx>  ó   € Ø×.Ñ.Ð.rd   r`   )	r‘   r’   r“   r”   r•   ry  r¥  r­  r–   r`   rd   rb   ro  ro    sb   ‡ à!Ó!ØÓØÓØÓØÓØÓØƒLØÓØ(Ó(ØÓØƒOØÓØ%)Ð˜
Ó)õ
÷ /ñ /rd   ro  c                  óx   € ] tR tRt$ R]R&   R]R&   R]R&   RtR]R	&   R
tR]R&   RtR]R&   R R ltR R lt	Rt
R# )ÚKernelDefinitionLineiB  r²  rÓ   r‹   rC  Úkernel_bodyNrÖ   ÚmetadataTrE  ÚgpuÚcpp_definitionc               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú!KernelDefinitionLine.__annotate__K  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c                	óª   € V P                   P                  V P                  V P                  V P                  V P
                  V P                  R 7       R# ©)rˆ  r‰  rŠ  N)rÓ   Ú_define_kernel_helperrC  r‡  rˆ  r‰  rŠ  r¤  s   &&rb   r¥  ÚKernelDefinitionLine.codegenK  sB   € Ø�‰×*Ñ*Ø×ÑØ×ÑØ—]‘]Ø—‘Ø×.Ñ.ð 	+ö 	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rŒ  T  rG  rd   c                	ó   € VP                   # ru   )Ú_generate_kernel_definitionr¬  s   &&rb   r­  ÚKernelDefinitionLine.codegen_fxT  rK  rd   r`   )r‘   r’   r“   r”   r•   rˆ  r‰  rŠ  r¥  r­  r–   r`   rd   rb   r†  r†  B  sA   ‡ à!Ó!ØÓØÓØ€HˆjÓØ€CˆÓØ!%€N�JÓ%õ
÷5ñ 5rd   r†  c                  óF   € ] tR tRt$ R]R&   R R ltR R ltR R	 ltR
tR# )ÚMemoryPlanningLineiX  r²  rÓ   c               ó    € V ^8„  d   QhRRRR/# ©r\   Ústaterj  r_   r–  r`   )ra   s   "rb   rc   ÚMemoryPlanningLine.__annotate__\  s   € ÷ ñ Ð-ð Ð2Dñ rd   c                ó   € V # )zFirst pass to find reuser`   ©rf  r™  s   &&rb   ÚplanÚMemoryPlanningLine.plan\  s   € àˆrd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   rš  `  s   € ÷ )ñ )˜Nð )¨tñ )rd   c                ó   € R# )zSecond pass to output codeNr`   r¤  s   &&rb   r¥  ÚMemoryPlanningLine.codegen`  s   ‚ rd   c               ó   € V ^8„  d   QhRR/# ©r\   r_   r‹   r`   )ra   s   "rb   rc   rš  c  s   € ÷ ;ñ ;˜ñ ;rd   c                ó€  € . p\         P                  ! V 4       Fy  pVP                  R8X  d   K  \        WP                  4      pTP	                  VP                   RVP
                  \        P                  J d   VP                  4       MT 24       K{  	  \        V 4      P                   RRP                  V4       R2# )z6
Emits a string representation that fits on one line.
rÓ   rR  rS  rT  r3  )ÚdataclassesÚfieldsrÐ   rw   rÈ   r…   r,   rÁ   rž   r‘   r  )rf  r-  ry   Úvals   &   rb   rg  ÚMemoryPlanningLine.__str__c  sš   € ð ˆÜ ×'Ò'¨Ö-ˆEØ�z‰z˜YÔ&ÙÜ˜$§
¡
Ó+ˆCØ�K‰KØ—:‘:�,˜a°%·*±*ÄÇ	Á	Ó2I §¡¤ÈsÐSÐTöñ	 .ô �t“*×%Ñ%Ð& a¨¯	©	°$«Ð'8¸Ð:Ð:rd   r`   N)	r‘   r’   r“   r”   r•   r�  r¥  rg  r–   r`   rd   rb   r–  r–  X  s   ‡ à!Ó!õõ)÷;ñ ;rd   r–  c                  ó€   a € ] tR tRt$ Rt^tR]R&   ]P                  ! ]	R7      t
R]R&   RtR	]R
&   R V 3R lltRtV ;t# )Ú+EnterDeviceContextManagerWithStreamInfoLineir  u  Enter a device context and retrieve user stream objects.

Attributes:
    num_streams: Number of streams (determined by user annotations on nodes).
    stream_idx_to_user_obj_idx: Maps stream_idx â†’ user_object_index for
        retrieving user stream objects via _get_stream_by_index.
r�   Únum_streams)Údefault_factoryúdict[int, int]Ústream_idx_to_user_obj_idxTrE  Úsetup_stream_cachec               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú8EnterDeviceContextManagerWithStreamInfoLine.__annotate__€  s   € ÷ ñ ˜Nð ¨tñ rd   c                ó†  <€ \         P                  P                  '       dP   \        SV `  V4       \         P                  P
                  P                  WP                  V P                  4       R# \        SV `  V4       V P                  '       dA   VP                  \         R\         P                  P                  P                  4        24       V P                  ^8”  d_   \        ^V P                  4       FB  pV P                  V,          pVP                  \        P                   ! VR7       RV R24       KD  	  R# R# )z5Generate context switching and stream retrieval code.r2  ©Ú
stream_idxz = _get_stream_by_index(r3  N)rH   rœ   r]  rm  r¥  Úwrapper_codeÚcodegen_stream_info_prologuer«  r®  r¯  rö   r3   r	  Úcurrent_streamÚranger5   ra   )rf  rŸ  ÚiÚuser_obj_idxrt  s   &&  €rb   r¥  Ú3EnterDeviceContextManagerWithStreamInfoLine.codegen€  sý   ø€ ä�7‰7××ÐÜ‰G‰O˜DÔ!Ü�G‰G× Ñ ×=Ñ=Ø×&Ñ&¨×(GÑ(Göô ‰G‰O˜DÔ!à×&×&Ð&Ø—‘Ü%Ð& c¬!¯'©'×*<Ñ*<×*KÑ*KÓ*MÐ)NÐOôð ×Ñ !Ô#Ü˜q $×"2Ñ"2Ö3�AØ#'×#BÑ#BÀ1Õ#E�LØ—N‘NÜ/×6Ò6À!ÔDÐEð F2Ø2>°¸qðBöó 4ñ $rd   r`   )r‘   r’   r“   r”   r[  r«  r•   r¥  ry   r€   r®  r¯  r¥  r–   r”  r•  s   @rb   rª  rª  r  s@   ø‡ ñð €K�ÓØ1<×1BÒ1BÐSWÔ1XÐ ÓXØ#Ð˜Ó#÷ö rd   rª  c                  ó6   € ] tR tRt$ Rt^tR]R&   R R ltRtR# )	Ú*ExitDeviceContextManagerWithStreamInfoLinei˜  zoExit a device context.

Attributes:
    num_streams: Number of streams that were allocated (must match Enter).
r�   r«  c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú7ExitDeviceContextManagerWithStreamInfoLine.__annotate__¢  s   € ÷ ñ ˜Nð ¨tñ rd   c                ój   € \         P                  P                  '       g   VP                  4        R# R# )zGenerate context exit code.Nr  r¤  s   &&rb   r¥  Ú2ExitDeviceContextManagerWithStreamInfoLine.codegen¢  s#   € ä�w‰w×"×"Ð"Ø×ÑÖñ #rd   r`   N)	r‘   r’   r“   r”   r[  r«  r•   r¥  r–   r`   rd   rb   r½  r½  ˜  s   ‡ ñð €K�Ó÷ñ rd   r½  c                  ó2   € ] tR tRt$ RtR]R&   R R ltRtR# )	ÚEnterCudaStreamContextLinei¨  z�Enter a context executed by the respective device stream.

Attributes:
    stream_idx: The index number for the entering device stream context.
r�   r´  c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú'EnterCudaStreamContextLine.__annotate__²  s   € ÷ Rñ R˜Nð R¨tñ Rrd   c                	óî   € \        \        P                  R R4      pVf;   VP                  R\	        V P
                  4       R24       VP                  4        R# VP                  WP
                  4       R# )rµ  Nr  r  )rw   rH   rœ   rö   r9   r´  r¿  Ú!codegen_enter_cuda_stream_context©rf  rŸ  rµ  s   && rb   r¥  Ú"EnterCudaStreamContextLine.codegen²  sU   € ÜœqŸw™w¨¸Ó=ˆØÒØ�N‰N˜U¤?°4·?±?Ó#CÐ"DÀAÐFÔGØ�N‰NÖà×:Ñ:¸4ÇÁÖQrd   r`   NrZ  r`   rd   rb   rÃ  rÃ  ¨  s   ‡ ñð ƒO÷Rñ Rrd   rÃ  c                  ó&   € ] tR tRtRtR R ltRtR# )ÚExitCudaStreamContextLinei»  z1Generate code to exit the current stream context.c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú&ExitCudaStreamContextLine.__annotate__¿  s   € ÷ @ñ @˜Nð @¨tñ @rd   c                	óŠ   € \        \        P                  R R4      pVf   VP                  4        R# VP	                  V4       R# )rµ  N)rw   rH   rœ   ró  Ú codegen_exit_cuda_stream_contextrÈ  s   && rb   r¥  Ú!ExitCudaStreamContextLine.codegen¿  s4   € ÜœqŸw™w¨¸Ó=ˆØÒØ×ÑÖà×9Ñ9¸$Ö?rd   r`   N©r‘   r’   r“   r”   r[  r¥  r–   r`   rd   rb   rË  rË  »  s   † á;÷@ñ @rd   rË  c                  ó<   € ] tR tRt$ RtR]R&   R]R&   R R ltRtR	# )
ÚEnterCudaMemPoolContextLineiÇ  zCEnter a torch.cuda.use_mem_pool context for scheduler node codegen.r�   Úmempool_indexÚdevice_indexc               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú(EnterCudaMemPoolContextLine.__annotate__Î  s   € ÷ ñ ˜Nð ¨tñ rd   c                	ó~   € VP                  R V P                   RV P                   R24       VP                  4        R# )z:with torch.cuda.use_mem_pool(get_external_object_by_index(z
), device=z):N)rö   rÔ  rÕ  r¿  r¤  s   &&rb   r¥  Ú#EnterCudaMemPoolContextLine.codegenÎ  sD   € Ø�‰ð,Ø,0×,>Ñ,>Ð+?ð @Ø×'Ñ'Ð(¨ð,ô	
ð
 	�‰Örd   r`   NrZ  r`   rd   rb   rÓ  rÓ  Ç  s   ‡ áMàÓØÓ÷ñ rd   rÓ  c                  ó&   € ] tR tRtRtR R ltRtR# )ÚExitCudaMemPoolContextLinei×  z0Exit the active torch.cuda.use_mem_pool context.c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú'ExitCudaMemPoolContextLine.__annotate__Ú  s   € ÷ ñ ˜Nð ¨tñ rd   c                	ó&   € VP                  4        R # ru   ©ró  r¤  s   &&rb   r¥  Ú"ExitCudaMemPoolContextLine.codegenÚ  ó   € Ø×ÑÖrd   r`   NrÑ  r`   rd   rb   rÛ  rÛ  ×  s   † Ù:÷ñ rd   rÛ  c                  ó@   € ] tR tRtR tR R ltR R ltR R ltR	tR
# )ÚEfficientPeakEstimateiÞ  c                	ó’  € ^RI HpHp \        P                  P
                  P                  p\        \        P                  P                  P                  4       4      p\        \        P                  P                  4       4      pV! W44      pV! VVV4      w  V n        p^RIHp V! V\        P                  \         ^ 4      V n        R# )r\   )Úestimate_peak_memoryÚget_freeable_input_buf)ÚSegmentedTreeN)Úmemoryrå  ræ  rH   rœ   r    Únodesr"   Úgraph_inputsÚkeysÚget_output_namesÚoverall_peak_memoryÚsegmented_treerç  Úoperatorr  Úmax)	rf  rå  ræ  Úscheduler_nodesrê  Úgraph_outputsÚnames_to_freeable_bufsÚpeak_by_scheduler_noderç  s	   &        rb   rn  ÚEfficientPeakEstimate.__init__ß  s”   € ßIäŸ'™'×+Ñ+×1Ñ1ˆÜ!¤!§'¡'×"6Ñ"6×";Ñ";Ó"=Ó>ˆÜ"¤1§7¡7×#;Ñ#;Ó#=Ó>ˆÙ!7¸Ó!VÐÙ;OØØ"Øó<
Ñ8ˆÔ Ð"8õ 	2á+Ø"¤H§L¡L´#°qó
ˆÖrd   c               ó    € V ^8„  d   QhRRRR/# )r\   r˜   r™   r_   r�   r`   )ra   s   "rb   rc   Ú"EfficientPeakEstimate.__annotate__ò  s   € ÷ -ñ -˜jð -¨Sñ -rd   c                	óÎ   € \         P                  P                  P                  \         P                  P	                  V4      ^ R7      \        VP                  4       4      ,          # )r   ©Úfallback)rH   rœ   r¥   Úoptimization_hintr�   r>   r¤   ©rf  r˜   s   &&rb   Ú	_get_sizeÚEfficientPeakEstimate._get_sizeò  sL   € Ü�w‰w×Ñ×1Ñ1Ü�G‰G×/Ñ/°Ó5Àð 2ó 
ä˜4Ÿ>™>Ó+Ó,õ-ð 	-rd   c               ó    € V ^8„  d   QhRRRR/# ©r\   Úline_ar}  Úline_bÚAllocateLiner`   )ra   s   "rb   rc   r÷  ÷  s   € ÷ 
ñ 
Ð#6ð 
Àñ 
rd   c                	ó~   € V P                   P                  VP                  ^,           VP                  ^,
          4      # rf   )rî  Úsummarize_rangeÚscheduler_node_index©rf  r  r  s   &&&rb   Úpeak_betweenÚ"EfficientPeakEstimate.peak_between÷  s6   € Ø×"Ñ"×2Ñ2Ø×'Ñ'¨!Õ+¨V×-HÑ-HÈ1Õ-Ló
ð 	
rd   c               ó    € V ^8„  d   QhRRRR/# r   r`   )ra   s   "rb   rc   r÷  ü  s   € ÷ 
ñ 
Ð*=ð 
À|ñ 
rd   c                	óþ   € VP                   ^,           VP                   8X  d   R# V P                  P                  VP                   ^,           VP                   ^,
          V P                  VP                  4      4       R# ©rg   N)r  rî  Úupdate_rangerý  r˜   r  s   &&&rb   Úupdate_peak_betweenÚ)EfficientPeakEstimate.update_peak_betweenü  s^   € Ø×&Ñ&¨Õ*¨f×.IÑ.IÔIÙØ×Ñ×(Ñ(Ø×'Ñ'¨!Õ+Ø×'Ñ'¨!Õ+Ø�N‰N˜6Ÿ;™;Ó'ö	
rd   )rí  rî  N)	r‘   r’   r“   r”   rn  rý  r  r  r–   r`   rd   rb   rã  rã  Þ  s   † ò
õ&-õ

÷

ñ 
rd   rã  c                  óv   € ] tR tRt$ RtR]R&   RtR]R&   R tR	 R
 ltR R lt	R R lt
R R ltR R ltRtR# )r  i  z6Represents a buffer allocation during memory planning.r™   r˜   FrE  Úcomm_bufferc                	ó  € \         P                  P                  P                  f   \	        R4      h\         P                  P                  P
                  P                  \         P                  P                  P                  4      V n        R # ©Nz)expected scheduler.current_node to be set©rH   rœ   r    Úcurrent_noder²   ré  Úindexr  re  s   &rb   r¹  ÚAllocateLine.__post_init__  óZ   € Ü�7‰7×Ñ×)Ñ)Ò1Ü Ð!LÓMÐMÜ$%§G¡G×$5Ñ$5×$;Ñ$;×$AÑ$AÜ�G‰G×Ñ×*Ñ*ó%
ˆÖ!rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   Ú	free_liner}  rÃ   r�   r_   rE  r`   )ra   s   "rb   rc   ÚAllocateLine.__annotate__  s#   € ÷ 6ñ 6Ð-@ð 6Èð 6ÐPTñ 6rd   c                	ó  € V P                   '       d   R # VP                  ^,           V P                  8X  d   R # V P                  P                  P                  pV P                  P                  P                  W4      pW$,           pWS8*  # )T)r  r  rÓ   Úestimate_peakrí  r  )rf  r  rÃ   rí  Úpeak_memory_in_rangeÚnew_peak_memorys   &&&   rb   Úshould_reuse_bufferÚ AllocateLine.should_reuse_buffer  so   € Ø××ÐÙØ×)Ñ)¨AÕ-°×1JÑ1JÔJÙØ"Ÿl™l×8Ñ8×LÑLÐØ#Ÿ|™|×9Ñ9×FÑFÀyÓWÐØÕ5ˆØÑ5Ð5rd   c               ó    € V ^8„  d   QhRRRR/# r˜  r`   )ra   s   "rb   rc   r    s   € ÷ &ñ &Ð-ð &Ð2Dñ &rd   c           	     	ó‚  € V P                   P                  4       \        P                  P                  9   d   \        V P                  4      # V P                  '       dŠ   \        V P                   4      p\        P                  '       d]   VP                  V4      '       dF   VP                  V4      pR Vn        \        V P                  VP                   V P                   R R7      # V # \        V P                   4      p\        P                  '       Ed   W!9   Ed   VP!                  V4      p\        P                  P"                  P%                  \        P                  P'                  V P                   4      ^ R7      \)        V P                   P+                  4       4      ,          pV P-                  W44      '       dX   R Vn        V P                  P.                  P1                  W04       \        V P                  VP                   V P                   4      # VP3                  W#4       V # V P                   P5                  4       P6                  R8X  dp   V P                  P9                  V P                   4      pVeG   V;P:                  \=        \>        P@                  ! \B        PD                  V^4      4      ,          un        V # )T©r  rù  Úcpu)#r˜   rž   rH   rœ   re  ÚNullLinerÓ   r  r¶   r*   Úallow_buffer_reuserŠ  rŽ  r€  Ú	ReuseLiner«   r  r¥   rû  r�   r>   r¤   r   r  r  r†  r£   r…   Ústatic_shape_for_buffer_or_noners  r�   Ú	functoolsÚreducerï  Úmul)rf  r™  r  r  rÃ   Ústatic_shapes   &&    rb   r�  ÚAllocateLine.plan  sò  € Ø�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ü˜DŸL™LÓ)Ð)à××Ðä'¨¯	©	Ó2ˆCÜ×(×(Ð(¨U×-GÑ-GÈ×-LÒ-LØ!×1Ñ1°#Ó6�	Ø&*�	Ô#Ü Ø—L‘L )§.¡.°$·)±)Èôð ð ˆKô ˜tŸy™yÓ)ˆÜ×$×$Ñ$¨­ØŸ	™	 #›ˆIÜ—7‘7×#Ñ#×5Ñ5Ü—‘×3Ñ3°D·I±IÓ>Èð 6ó ä˜tŸy™y×2Ñ2Ó4Ó5õ6ˆDð ×'Ñ'¨	×8Ò8Ø&*�	Ô#Ø—‘×*Ñ*×>Ñ>¸yÔOÜ  §¡¨y¯~©~¸t¿y¹yÓIÐIà—
‘
˜3Ô*Ø�à�9‰9×(Ñ(Ó*×/Ñ/°5Ô8ØŸ<™<×GÑGÈÏ	É	ÓRˆLØÒ'Ø×1Ò1´SÜ×$Ò$¤X§\¡\°<ÀÓCó6õ Õ1ð ˆrd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   r  F  s   € ÷ Kñ K˜Nð K¨tñ Krd   c                	ót  € V P                   P                  4       \        P                  P                  9   d(   \        R V P                   P                  4        R24      hV P                  '       d   V P                  V4       R# VP                  V P                  P                  V P                   4      4       R# rb  )r˜   rž   rH   rœ   re  r²   r  Ú_codegen_comm_bufferrö   rÓ   Úmake_buffer_allocationr¤  s   &&rb   r¥  ÚAllocateLine.codegenF  s„   € Ø�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ü Ø˜$Ÿ)™)×,Ñ,Ó.Ð/Ð/NÐOóð ð ××ÐØ×%Ñ% dÖ+à�N‰N˜4Ÿ<™<×>Ñ>¸t¿y¹yÓIÖJrd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   r  P  s   € ÷ %ñ %¨ð %¸Dñ %rd   c                ó  € V P                   P                  4       pV P                   P                  4       pVe   VP                  f   \	        RV 24      hV P                   P                  4       p\        V P                   P                  4       4      p\        V P                   P                  4       4      pV P                   P                  4       p\        V\        P                  4      '       g   \	        R\        V4       24      hVP                  pVP                  p	V\        P                   P"                  8X  d­   \%        4       '       d.   R\&        P(                  P*                  P-                  4        R2p
MRVP                   R2p
V RV P.                  P1                  V4       R	V P.                  P1                  V4       R	V R	V
 R
V	 R\2        P4                  ! ^ R4       R2pM\7        RV 24      hVP9                  V4       R# )z*Generate allocation code for comm buffers.Nz9Comm buffer requires a valid CUDA device with index, got r¯   ztorch.device("cuda", r3  ztorch.device("cuda:ú")z = empty_strided_p2p(rT  z, group_name="z", alloc_id=zUnsupported comm buffer type: l   ÿÿÿÿ )r˜   rž   r=  r  r²   r¤   rƒ   Úget_sizeÚ
get_strider°   r   r,   r±   r…   r³   r´   ÚCommBufferTypeÚSYMM_MEMr   rH   rœ   r	  r  rÓ   Úcodegen_shape_tupleÚrandomÚrandintr£  rö   )rf  rŸ  rÐ   r�   r   ÚshaperÄ   rµ   r³   r´   Ú
device_argrô   s   &&          rb   r1  Ú!AllocateLine._codegen_comm_bufferP  sÊ  € à�y‰y×!Ñ!Ó#ˆØ—‘×%Ñ%Ó'ˆØÒ" v§|¡|Ò'?Ü ØKÈFÈ8ÐTóð ð —	‘	×#Ñ#Ó%ˆÜ�d—i‘i×(Ñ(Ó*Ó+ˆÜ�t—y‘y×+Ñ+Ó-Ó.ˆØ—‘×*Ñ*Ó,ˆÜ˜&¤"×"5Ñ"5×6Ò6Ü Ð#BÄ4ÈÃ<À.Ð!QÓRÐRØ!×2Ñ2ÐØ×&Ñ&ˆ
àœr×0Ñ0×9Ñ9Ô9ô �ŠØ4´Q·W±W×5GÑ5G×5_Ñ5_Ó5aÐ4bÐbcÐd‘
à2°6·<±<°.ÀÐC�
à�&Ð-Ø—<‘<×3Ñ3°EÓ:Ð;¸2Ø—<‘<×3Ñ3°FÓ;Ð<¸BØ�'˜Ø�,ð Ø)˜lð +Ü"ŸNšN¨1¨iÓ8Ð9¸ð<ñ ô &Ø0Ð1AÐ0BÐCóð ð 	�‰�tÖrd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r  w  s   € ÷ ,ñ , Kð ,Ð4Dñ ,rd   c                	óV   € V P                   '       d   VP                  # VP                  # ru   )r  Ú_generate_comm_buffer_allocateÚ_generate_allocater¬  s   &&rb   r­  ÚAllocateLine.codegen_fxw  s&   € Ø××ÐØ×;Ñ;Ð;Ø×+Ñ+Ð+rd   ©r  N)r‘   r’   r“   r”   r[  r•   r  r¹  r   r�  r¥  r1  r­  r–   r`   rd   rb   r  r    s;   ‡ á@à
ÓØ€K�Óò
õ6õ&õPKõ%÷N,ñ ,rd   r  c                  óh   € ] tR tRt$ R]R&   RtR]R&   RtR]R&   R tR	 R
 ltR R lt	R R lt
RtR# )r}  i}  r™   r˜   FrE  r€  r  c                	ó  € \         P                  P                  P                  f   \	        R4      h\         P                  P                  P
                  P                  \         P                  P                  P                  4      V n        R # r  r  re  s   &rb   r¹  Ú!FreeIfNotReusedLine.__post_init__ƒ  r  rd   c               ó    € V ^8„  d   QhRRRR/# r˜  r`   )ra   s   "rb   rc   Ú FreeIfNotReusedLine.__annotate__Š  s   € ÷ ñ Ð-ð Ð2Dñ rd   c                	ó„  € \        V P                  P                  4       4      ^ 8”  d   V # \        V P                  P                  \
        P                  4      '       d   V # V P                  '       d   \        R4      hV P                  P                  4       \        P                  P                  9   d   \        V P                  4      # \        P                   '       da   V P"                  '       d)   \%        V P                  4      pVP'                  W 4       V # \)        V P                  4      pVP+                  W 4       V # )r   zexpected line to not be reused)r  r˜   Úget_inputs_that_alias_outputr   rµ   r,   ÚMultiOutputLayoutr€  r²   rž   rH   rœ   re  r&  rÓ   r*   r'  r  r¶   r’  r«   r†  )rf  r™  r  s   && rb   r�  ÚFreeIfNotReusedLine.planŠ  sâ   € Üˆt�y‰y×5Ñ5Ó7Ó8¸1Ô<ØˆKÜ�d—i‘i×&Ñ&¬×(<Ñ(<×=Ò=ØˆKØ�>�>ˆ>Ü Ð!AÓBÐBØ�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ü˜DŸL™LÓ)Ð)Ü×$×$Ð$Ø××Ðä+¨D¯I©IÓ6�Ø×&Ñ& sÔ1ð ˆô ' t§y¡yÓ1�Ø—
‘
˜3Ô%Øˆrd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   rK  �  s   € ÷ %ñ %˜Nð %¨tñ %rd   c                	óv  € V P                   P                  4       \        P                  P                  9   d(   \        R V P                   P                  4        R24      hV P                  '       gÉ   V P                  P                  V P                   4      pV P                  '       d   V P                   P                  4       p\        V\        P                  4      '       g   \        R\        V4       24      hVP                  V RVP                   P"                   R24       R# VP                  V4       R# R# )rc  rd  r¯   z # z buffer freeN)r˜   rž   rH   rœ   re  r²   r€  rÓ   rf  r  r°   r   r,   r±   r…   rö   r³   r†   )rf  rŸ  rô   rµ   s   &&  rb   r¥  ÚFreeIfNotReusedLine.codegen�  sñ   € Ø�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ü Ø˜$Ÿ)™)×,Ñ,Ó.Ð/Ð/NÐOóð ð �~�~ˆ~Ø—<‘<×0Ñ0°·±Ó;ˆDØ××ÐØŸ™×2Ñ2Ó4�Ü! &¬"×*=Ñ*=×>Ò>Ü(Ø9¼$¸v»,¸ÐHóð ð —‘ $  s¨6×+BÑ+B×+HÑ+HÐ*IÈÐVÖWà—‘˜tÖ$ñ rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rK  ®  s   € ÷ 6ñ 6 Kð 6Ð4Dñ 6rd   c                	óV   € V P                   '       d   VP                  # VP                  # ru   )r  Ú_generate_comm_buffer_freeÚ_generate_free_if_not_reusedr¬  s   &&rb   r­  ÚFreeIfNotReusedLine.codegen_fx®  s&   € Ø××ÐØ×7Ñ7Ð7Ø×5Ñ5Ð5rd   rF  N)r‘   r’   r“   r”   r•   r€  r  r¹  r�  r¥  r­  r–   r`   rd   rb   r}  r}  }  s5   ‡ à
ÓØ€IˆtÓØ€K�Óò
õõ&%÷"6ñ 6rd   r}  c                  óZ   € ] tR tRt$ R]R&   R]R&   R]R&   R R ltR	 R
 ltR R ltRtR# )ÚReinterpretLinei´  r™   r˜   Ú	reused_asz	ir.Layoutrµ   c               ó    € V ^8„  d   QhRRRR/# r˜  r`   )ra   s   "rb   rc   ÚReinterpretLine.__annotate__º  s   € ÷ ñ Ð-ð Ð2Dñ rd   c                	ó   € V # ru   r`   rœ  s   &&rb   r�  ÚReinterpretLine.planº  s   € Øˆrd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   r\  ½  s   € ÷ 	
ñ 	
˜Nð 	
¨tñ 	
rd   c                	óê  € \        V P                  \        P                  4      '       g"   \	        R \        V P                  4       24      h\        V P                  P                  \        P                  4      '       g,   \	        R\        V P                  P                  4       24      hV P                  P                  V P                  P                  4       V P                  P                  4       R# )zexpected NonOwningLayout, got zexpected ReinterpretView, got N)r   rµ   r,   ÚNonOwningLayoutr²   r…   Úviewr/   rÓ   Úcodegen_deferred_allocationrZ  rž   r¤  s   &&rb   r¥  ÚReinterpretLine.codegen½  sª   € Ü˜$Ÿ+™+¤r×'9Ñ'9×:Ò:Ü Ð#AÄ$ÀtÇ{Á{ÓBSÐATÐ!UÓVÐVÜ˜$Ÿ+™+×*Ñ*¬B×,>Ñ,>×?Ò?Ü Ø0´°d·k±k×6FÑ6FÓ1GÐ0HÐIóð ð 	�‰×0Ñ0Ø�N‰N×#Ñ#Ó% t§{¡{×'7Ñ'7ö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r\  È  r€  rd   c                	ó   € VP                   # ru   )Ú_generate_reinterpretr¬  s   &&rb   r­  ÚReinterpretLine.codegen_fxÈ  r„  rd   r`   N)	r‘   r’   r“   r”   r•   r�  r¥  r­  r–   r`   rd   rb   rY  rY  ´  s(   ‡ à
ÓØÓØÓõõ	
÷/ñ /rd   rY  c                  ól   € ] tR tRt$ R]R&   R]R&   RtR]R&   RtR]R	&   R
 R ltR R ltR R lt	Rt
R# )r(  iÌ  r™   r˜   rZ  TrE  Ú
delete_oldFr  c               ó    € V ^8„  d   QhRRRR/# r˜  r`   )ra   s   "rb   rc   ÚReuseLine.__annotate__Ó  s   € ÷ ñ Ð-ð Ð2Dñ rd   c                	ó  € V P                   P                  4       \        P                  P                  9   dt   V P
                  P                  4       \        P                  P                  9  d(   \        R V P
                  P                  4        R24      h\        V P                  4      # V P
                  P                  4       \        P                  P                  9   d(   \        R V P
                  P                  4        R24      hV # )úreuse target z expected to be removedú was removed unexpectedly)	r˜   rž   rH   rœ   re  rZ  r²   r&  rÓ   rœ  s   &&rb   r�  ÚReuseLine.planÓ  sÇ   € Ø�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ø�~‰~×&Ñ&Ó(´·±×0GÑ0GÔGÜ$Ø# D§N¡N×$;Ñ$;Ó$=Ð#>Ð>UÐVóð ô ˜DŸL™LÓ)Ð)Ø�>‰>×"Ñ"Ó$¬¯©×(?Ñ(?Ô?Ü Ø §¡× 7Ñ 7Ó 9Ð:Ð:SÐTóð ð ˆrd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   rl  à  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c                	ó  € V P                   P                  4       \        P                  P                  9   d(   \        R V P                   P                  4        R24      hV P                  P                  4       \        P                  P                  9   d(   \        RV P                  P                  4        R24      hVP                  V P                  P                  V P                   V P                  V P                  4      4       R# )rc  rd  rn  ro  N)r˜   rž   rH   rœ   re  r²   rZ  rö   rÓ   Úmake_buffer_reuserj  r¤  s   &&rb   r¥  ÚReuseLine.codegenà  sÂ   € Ø�9‰9×ÑÓ¤1§7¡7×#:Ñ#:Ô:Ü Ø˜$Ÿ)™)×,Ñ,Ó.Ð/Ð/NÐOóð ð �>‰>×"Ñ"Ó$¬¯©×(?Ñ(?Ô?Ü Ø §¡× 7Ñ 7Ó 9Ð:Ð:SÐTóð ð 	�‰Ø�L‰L×*Ñ*¨4¯9©9°d·n±nÀdÇoÁoÓVö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rl  í  s   € ÷ )ñ ) Kð )Ð4Dñ )rd   c                	ó   € VP                   # ru   )Ú_generate_reuser¬  s   &&rb   r­  ÚReuseLine.codegen_fxí  s   € Ø×(Ñ(Ð(rd   r`   N)r‘   r’   r“   r”   r•   rj  r  r�  r¥  r­  r–   r`   rd   rb   r(  r(  Ì  s6   ‡ à
ÓØÓØ€J�ÓØ€K�Óõõ
÷)ñ )rd   r(  c                  ó"   € ] tR tRtR R ltRtR# )r&  iñ  c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   ÚNullLine.__annotate__ò  ri  rd   c                	ó   € VP                   # ru   )Ú_generate_nullr¬  s   &&rb   r­  ÚNullLine.codegen_fxò  rm  rd   r`   N)r‘   r’   r“   r”   r­  r–   r`   rd   rb   r&  r&  ñ  s   † ÷(ñ (rd   r&  c                  ó\   € ] tR tRt$ RtR]R&   R]R&   R]R&   R]R	&   R
 R ltR R ltRtR# )ÚMultiOutputLineiö  zM
Given a MultiOutputLayout buffer, indexes actual buffer(s) from the result.
r²  rÓ   r‹   Úresult_nameÚarg_nameúSequence[Any]Úindicesc               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   ÚMultiOutputLine.__annotate__  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c                	óê   a a€ VV 3R  loS! S P                   S P                  4      pVP                  S P                  P                   S P
                   RV S P                  P                   24       R# )c                ó°  <€ \        V4      ^ 8”  dÅ   V^ ,          w  r#\        V\        4      '       d   S! V  RV R2VR,          4      # \        V\        4      '       d@   SP                  P                  V SP                  \        V4      4      pS! WAR,          4      # \        V\        4      '       d   S! V  RV R2VR,          4      # \        RV4      hV # )r   Ú[Ú]ºrg   NNz['z']znon supported index type: )
r  Ú
issubclassr‚   rƒ   rÓ   Úcodegen_tuple_accessr�  r‹   r€   r²   )Úbasenamer„  Úityper¹  Útuple_accessÚcodegen_list_tuple_accessrf  s   &&   €€rb   r‘  Ú:MultiOutputLine.codegen.<locals>.codegen_list_tuple_access  sÈ   ø€ Ü�7‹|˜aÔØ" 1�:‘�Ü˜e¤T×*Ò*Ù4¸°zÀÀ1À#ÀQÐ5GÈÐQSÍÓUÐUÜ ¤u×-Ò-à#'§<¡<×#DÑ#DØ  $×"2Ñ"2´C¸³Fó$�Lñ 5°\È2Å;ÓOÐOÜ ¤t×,Ò,Ù4¸°zÀÀAÀ3ÀbÐ5IÈ7ÐSUÍ;ÓWÐWä(Ð)EÀuÓMÐMà�rd   r2  N)r‚  r„  rö   rÓ   Údeclarer�  Úending)rf  rŸ  r†   r‘  s   f& @rb   r¥  ÚMultiOutputLine.codegen  s]   ù€ ö	 ñ$ *¨$¯-©-¸¿¹ÓFˆØ�‰Ø�|‰|×#Ñ#Ð$ T×%5Ñ%5Ð$6°c¸%¸ÀÇÁ×ATÑATÐ@UÐVö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r†    s   € ÷ 0ñ 0 Kð 0Ð4Dñ 0rd   c                	ó   € VP                   # ru   )Ú_generate_multi_outputr¬  s   &&rb   r­  ÚMultiOutputLine.codegen_fx  s   € Ø×/Ñ/Ð/rd   r`   N)	r‘   r’   r“   r”   r[  r•   r¥  r­  r–   r`   rd   rb   r€  r€  ö  s/   ‡ ñð "Ó!ØÓØƒMØÓõ
÷00ñ 0rd   r€  c                  óN   € ] tR tRt$ R]R&   R]R&   R]R&   R R	 ltR
 R ltRtR# )ÚIndexPutFallbackLinei  r²  rÓ   úir.IndexPutFallbackr˜   zlist[ir.IRNode | None]r„  c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú!IndexPutFallbackLine.__annotate__#  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c                	óÈ  € V P                   p\        P                  ! VP                  4      '       g   \	        R 4      hR VP                  R,           4       w  r4V P
                   Uu. uF1  pV'       d   VP                  4       MV P                  P                  NK3  	  ppV P                  P                  ! VP                  4       W6V.VP                  4       O5!   R# u upi )ú*expected node.inputs to be a node sequencec              3  ó@   "  € T F  qP                  4       x € K  	  R # 5iru   ©r@  ©rx   Úts   & rb   rz   Ú/IndexPutFallbackLine.codegen.<locals>.<genexpr>'  s   é € ÐF±o°×*Ñ*×,Ð,³oùó   ‚:Nr\   NN)r˜   r,   Úis_node_sequencer�   r²   r„  r@  rÓ   Únone_strÚ_generate_index_put_fallbackr;  Úcodegen_const_args)rf  rŸ  r˜   r  Úvaluesr  r„  s   &&     rb   r¥  ÚIndexPutFallbackLine.codegen#  sº   € Ø�y‰yˆÜ×"Ò" 4§;¡;×/Ò/Ü Ð!MÓNÐNÙF°d·k±kÀ"¶oÓF‰ˆð —|’|ó
á#�÷ (+ˆC×!Ñ!Ô#°·±×0EÑ0EÒEÙ#ð 	ð 
ð
 	�‰×1Ò1Ø× Ñ Ó" A°ð	
Ø9=×9PÑ9PÓ9Rö	
ùò
s   Á(CÁ5*Cc               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rž  1  s   € ÷ 6ñ 6 Kð 6Ð4Dñ 6rd   c                	ó   € VP                   # ru   )r©  r¬  s   &&rb   r­  ÚIndexPutFallbackLine.codegen_fx1  s   € Ø×5Ñ5Ð5rd   r`   Nr  r`   rd   rb   r›  r›    s#   ‡ à!Ó!Ø
ÓØ#Ó#õ
÷6ñ 6rd   r›  c                  óD   € ] tR tRt$ R]R&   R]R&   R R ltR R	 ltR
tR# )ÚScatterFallbackLinei5  r²  rÓ   úir.ScatterFallbackr˜   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú ScatterFallbackLine.__annotate__:  s   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c           
     	ó‚  € V P                   p\        P                  ! VP                  4      '       g   \	        R 4      hVP
                  '       d   R VP                   4       w  r4pM(R VP                   4       w  r4VP                  ^,          pVP                  4       ;p'       d   VP                  M\        P                  P                  pV P                  P                  VW2P                  ^ ,          WE.VP                  VP                  VP
                  VP                   R,          VP#                  4       V4       R# )r   c              3  ó@   "  € T F  qP                  4       x € K  	  R # 5iru   r¢  r£  s   & rb   rz   Ú.ScatterFallbackLine.codegen.<locals>.<genexpr>?  s   é € ÐJ¹k¸×2Ñ2×4Ð4»kùr¦  c              3  ó@   "  € T F  qP                  4       x € K  	  R # 5iru   r¢  r£  s   & rb   rz   r·  A  s   é € ÐE¹°A×-Ñ-×/Ð/»ùr¦  r+  N)r˜   r,   r§  r�   r²   Úsrc_is_tensorÚconstant_argsr=  r…   rH   rœ   r>  rÓ   Ú_generate_scatter_fallbackr<  Úpython_kernel_namer  r+  )rf  rŸ  r˜   r  r  r@  rD  r�   s   &&      rb   r¥  ÚScatterFallbackLine.codegen:  sî   € Ø�y‰yˆÜ×"Ò" 4§;¡;×/Ò/Ü Ð!MÓNÐNØ××ÐÙJ¸d¿kºkÓJ‰OˆQ‘sáE¸¿ºÓE‰JˆQØ×$Ñ$ QÕ'ˆCØ!%§¡Ó!2Ð2˜AÖ2�—’¼¿¹×9LÑ9LˆØ�‰×/Ñ/ØØ×"Ñ" 1Õ% uÐ2Ø× Ñ Ø×#Ñ#Ø×ÑØ�K‰K˜Õ!Ø×ÑÓ!Øö		
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   r´  O  s   € ÷ 4ñ 4 Kð 4Ð4Dñ 4rd   c                	ó   € VP                   # ru   )r»  r¬  s   &&rb   r­  ÚScatterFallbackLine.codegen_fxO  s   € Ø×3Ñ3Ð3rd   r`   Nr  r`   rd   rb   r±  r±  5  s   ‡ à!Ó!Ø
Óõ
÷*4ñ 4rd   r±  c                  óN   € ] tR tRt$ R]R&   R]R&   R]R&   R R	 ltR
 R ltRtR# )ÚSymbolicCallArgLineiS  r²  rÓ   r`  ÚargrW   rœ   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú SymbolicCallArgLine.__annotate__Y  s   € ÷ Nñ N˜Nð N¨tñ Nrd   c                	óf   € V P                   P                  V P                  V P                  4       R # ru   )rÓ   Ú"_generate_symbolic_call_arg_helperrÃ  rœ   r¤  s   &&rb   r¥  ÚSymbolicCallArgLine.codegenY  s   € Ø�‰×7Ñ7¸¿¹À$Ç*Á*ÖMrd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rÅ  \  rG  rd   c                	ó   € VP                   # ru   )Ú_generate_symbolic_call_argr¬  s   &&rb   r­  ÚSymbolicCallArgLine.codegen_fx\  rK  rd   r`   Nr  r`   rd   rb   rÂ  rÂ  S  s$   ‡ à!Ó!Ø	ÓØÓõN÷5ñ 5rd   rÂ  c                  óX   € ] tR tRt$ R]R&   R]R&   R]R&   R]R	&   R
 R ltR R ltRtR# )ÚUnbackedSymbolDefsLinei`  r²  rÓ   r‹   Úoutput_namer   Úoutputsú)dict[sympy.Symbol, pytree.KeyPath] | NoneÚunbacked_bindingsc               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú#UnbackedSymbolDefsLine.__annotate__g  ó   € ÷ 
ñ 
˜Nð 
¨tñ 
rd   c                	ó|   € V P                   P                  V P                  V P                  V P                  4       R # ru   )rÓ   Ú)_codegen_unbacked_symbol_defs_for_outputsrÏ  rÐ  rÒ  r¤  s   &&rb   r¥  ÚUnbackedSymbolDefsLine.codegeng  s+   € Ø�‰×>Ñ>Ø×Ñ˜dŸl™l¨D×,BÑ,Bö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rÔ  l  s   € ÷ 8ñ 8 Kð 8Ð4Dñ 8rd   c                	ó   € VP                   # ru   )Ú_generate_unbacked_symbol_defsr¬  s   &&rb   r­  Ú!UnbackedSymbolDefsLine.codegen_fxl  s   € Ø×7Ñ7Ð7rd   r`   Nr  r`   rd   rb   rÎ  rÎ  `  s(   ‡ à!Ó!ØÓØƒLØ@Ó@õ
÷
8ñ 8rd   rÎ  c                  ó~   € ] tR tRt$ R]R&   R]R&   R]R&   R]R&   RtR]R	&   R
tR]R&   R R lt]R R l4       t	Rt
R
# )ÚAssertSizeStrideLineip  r²  rÓ   r‹   rÐ   rÃ   rÄ   ÚinputÚop_nameNútorch.dtype | Noner   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú!AssertSizeStrideLine.__annotate__y  rÕ  rd   c                	ó¨   € V P                   P                  WP                  V P                  V P                  V P
                  V P                  4       R # ru   )rÓ   Ú_codegen_assert_size_striderÐ   rÃ   rÄ   rà  r   r¤  s   &&rb   r¥  ÚAssertSizeStrideLine.codegeny  s3   € Ø�‰×0Ñ0Ø—)‘)˜TŸY™Y¨¯©°T·\±\À4Ç:Á:ö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   rã    ó   € ÷ 6ñ 6˜kð 6Ð.>ñ 6rd   c                	ó   € V P                   # ru   ©Ú_generate_assert_size_striderÐ  s   &rb   r­  ÚAssertSizeStrideLine.codegen_fx~  ó   € à×5Ñ5Ð5rd   r`   )r‘   r’   r“   r”   r•   rà  r   r¥  rÓ  r­  r–   r`   rd   rb   rÞ  rÞ  p  sE   ‡ à!Ó!Ø
ƒIØ
ƒIØƒKØ€GˆSÓØ $€EÐÓ$õ
ð
 ô6ó ô6rd   rÞ  c                  ó\   € ] tR tRt$ R]R&   R]R&   RtR]R&   R	 R
 lt]R R l4       tRt	R# )ÚGroupedAssertSizeStrideLineiƒ  r²  rÓ   úlist[tuple[str, str, str]]Úassertsrß  r‹   rà  c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú(GroupedAssertSizeStrideLine.__annotate__‰  rÕ  rd   c                	óf   € V P                   P                  WP                  V P                  4       R # ru   )rÓ   Ú#_codegen_assert_size_stride_groupedrñ  rà  r¤  s   &&rb   r¥  Ú#GroupedAssertSizeStrideLine.codegen‰  s!   € Ø�‰×8Ñ8Ø—,‘, §¡ö	
rd   c               ó    € V ^8„  d   QhRRRR/# r¨  r`   )ra   s   "rb   rc   ró  �  rè  rd   c                	ó   € V P                   # ru   rê  rÐ  s   &rb   r­  Ú&GroupedAssertSizeStrideLine.codegen_fxŽ  rí  rd   r`   N)
r‘   r’   r“   r”   r•   rà  r¥  rÓ  r­  r–   r`   rd   rb   rï  rï  ƒ  s1   ‡ à!Ó!Ø'Ó'Ø€GˆSÓõ
ð
 ô6ó ô6rd   rï  c                  óF   € ] tR tRt$ RtR]R&   R]R&   R]R&   R R	 ltR
tR# )ÚAssertDivByZeroLinei“  zûDeferred AOTI runtime check that a sizevar divisor is non-zero.

Queued from CppWrapperCpu.codegen_cpp_sizevar during the build phase
and replayed at the position of use so the check lands after the
declarations of any unbacked symbols it references.
r²  rÓ   r‹   Údivisorrà  c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   Ú AssertDivByZeroLine.__annotate__   s   € ÷ Sñ S˜Nð S¨tñ Srd   c                	óf   € V P                   P                  WP                  V P                  4       R # ru   )rÓ   Ú_codegen_assert_div_by_zerorü  rà  r¤  s   &&rb   r¥  ÚAssertDivByZeroLine.codegen   s   € Ø�‰×0Ñ0°·|±|ÀTÇ\Á\ÖRrd   r`   NrZ  r`   rd   rb   rû  rû  “  s%   ‡ ñð "Ó!ØƒLØƒL÷Sñ Srd   rû  c                  ób  a € ] tR tRt$ RtRtR]R&   V 3R lt]ER&R R	 ll4       t	R
 R lt
R R ltR R ltR R ltR R lt]R R l4       tR R lt]R R l4       tR R lt]R R l4       tR R ltR  R! ltR" R# ltR$ R% ltR& R' ltR( R) ltR* R+ ltR, R- ltR. R/ ltR0 R1 ltR2 R3 ltR4 R5 lt ER&R6 R7 llt!ER&R8 R9 llt"R: R; lt#R< R= lt$R> R? lt%R@ RA lt&RB RC lt'ER'RD RE llt(RF RG lt)RH RI lt*RJ t+RK t,RL t-RM t.RN t/RO RP lt0ER(RQ RR llt1RS RT lt2RU RV lt3RW RX lt4RY RZ lt5R[ R\ lt6R] R^ lt7R_ R` lt8Ra Rb lt9Rc Rd lt:Re Rf lt;Rg Rh lt<Ri Rj lt=Rk Rl lt>Rm Rn lt?Ro Rp lt@Rq tARr Rs ltBER&Rt Ru lltCER)Rv ltDER)Rw ltEER)Rx ltFRy tGRz R{ ltHR| tIR} R~ ltJR R€ ltKR� tLR‚ Rƒ ltMR„ tNR… R† ltO]PP¢                  R‡ Rˆ l4       tRR‰ RŠ ltSR‹ tTRŒ tUR� tVRŽ tWR� R� ltXR‘ R’ ltYER&R“ R” lltZR• R– lt[R— R˜ lt\R™ t]Rš R› lt^Rœ t_R�R/Rž RŸ llt`R�R/R  R¡ lltaR¢ R£ ltbR¤ R¥ ltcR¦ R§ ltdR¨ R© lteRª R« ltfER&R¬ R­ lltgR® R¯ lthR° R± ltiR² R³ ltjR´ tkRµ tlR¶ tmR· tnR¸ R¹ ltoRº tpR» tqER*R¼ R½ lltr]ER&R¾ R¿ ll4       tsER*RÀ RÁ llttRÂ RÃ ltu]vRÄ 4       twRÅ RÆ ltxER&RÇ RÈ lltyRÉ RÊ ltzRË RÌ lt{RÍ RÎ lt|RÏ t}RÐ t~RÑ tRÒ t€RÓ t�RÔ t‚RÕ tƒRÖ t„R× t…RØ t†ER&RÙ lt‡RÚ tˆRÛRRÜRRÝRRÞRRßRRàRRáRRâR/Rã Rä llt‰RÛRRÜRRÝRRÞRRßRRàRRáRRåRæRâRRçR/
Rè Ré lltŠRê t‹Rë tŒRì t�ER&Rí ltŽRî Rï lt�]Rð 4       t�ER+Rñ lt‘Rò t’ER,Ró lt“Rô Rõ lt”Rö R÷ lt•Rø Rù lt–ER&Rú Rû llt—Rü Rý lt˜Rþ Rÿ lt™ER  ER ltšER t›ER&ER ltœER t�ER ER ltžER tŸER ER	 lt ER
 ER lt¡ER t¢ER ER lt£ER ER lt¤ER ER lt¥ER t¦ER t§ER t¨ER t©ER tªER t«ER ER lt¬ER t­]ER 4       t®]ER 4       t¯]ER 4       t°]ER 4       t±]ER  4       t²ER! ER" lt³ER# t´ER$ tµER%t¶V ;t·# (-  r²  i¨  z:
Generate outer wrapper in Python that calls the kernels.
TrE  Úsupports_cachingc                	ó‚  <a € \         SS `  4        R S n        / S n        \	        4       S n        \	        4       S n        \	        4       S n        / S n        \        4       S n
        / S n        \        4       S n        \        4       S n        \        4       S n        \        4       S n        \        4       S n        \        4       S n        \        4       S n        \        4       S n        \        4       S n        \	        4       S n        / S n        ^ S n        / S n        \	        4       S n        . S n        RS n        RS n        RS n        RS n         RS n!        \D        PF                  PH                  '       d   RMRS n%        \D        PF                  PH                  '       d   RMRS n&        R S n'        R S n(        RS n)        / S n*        \	        4       S n+        \	        4       S n,        R S n-        S P]                  4        . S n/        . S n0        S Pc                  4        \e        S 4      '       g   S Pg                  4        S Pi                  4        \D        PF                  Pj                  '       gC   \D        PF                  Pl                  Po                  4        F  w  rS Pq                  W4       K  	  \        \r        ,          ! 4       S n:        \        \r        ,          ! 4       S n;        / S n<        \z        P|                  ! R 4      ! S P~                  4      S n?        \z        P€                  R V 3R ll4       pVS nA        / S nB        \	        4       S nC        \‰        4       S nE        \	        4       S nF        / S nG        \‘        \’        P”                  P–                  \’        P”                  P˜                  R	7      S nM        . S nN        \	        4       S nO        R # )
NÚ Ú#r„  z
std::move(r3  Tc               ó    € V ^8„  d   QhRRRR/# )r\   rô   r‹   r_   r„  r`   )ra   s   "rb   rc   Ú3PythonWrapperCodegen.__init__.<locals>.__annotate__  s   € ÷ 	;ñ 	; #ð 	;¨$ñ 	;rd   c                ó¸   <€ SP                   P                  V 4       \        P                  P                  '       d   SP
                  P                  V 4       R # R # ru   )Úimportsrö   r*   rî   rï   rø   )rô   rf  s   &€rb   Úadd_import_onceÚ6PythonWrapperCodegen.__init__.<locals>.add_import_once
  s>   ø€ à�L‰L×"Ñ" 4Ô(Ü�}‰}×5×5Ð5Ø×*Ñ*×4Ñ4°TÖ:ñ 6rd   )Údebug_printer_levelÚuse_array_ref)Prm  rn  Ú_last_default_stream_deviceÚ_pending_input_assertsr"   Ú_pending_alignment_copiesÚ_multistream_alignment_copiesÚ_alignment_orig_savedÚ_alignment_aligned_streamr   Ú_names_iterÚargs_to_buffersr?   r
  rC   ÚheaderÚprefixÚsuffixÚkernel_declarationsÚwrapper_callÚkernel_autotune_defsrø   Úsubgraph_definitionsr÷   Úkernel_autotune_example_argsÚkernel_autotune_tmp_arg_idxÚsrc_to_kernelÚkernel_numel_exprÚlinesr“  Údeclare_maybe_referencer”  Úcommentr¨  rH   rœ   r]  Ú
move_beginÚmove_endr  r   Úsupports_intermediate_hooksÚuser_defined_kernel_cacheÚunbacked_symbol_declsr¸  Úlauncher_fn_nameÚset_launcher_fn_nameÚcodegened_graph_stackÚcomputed_sizes_stackÚwrite_headerr@   Úwrite_prefixÚ!write_kernel_autotune_defs_headerr  Úconstant_reprsr�   Úwrite_constantÚ
BufferNameÚ	allocatedÚfreedÚreusesr*  Ú	lru_cacheÚwrite_get_raw_streamÚcacher  Ú_metasÚ
_meta_varsr   Úmulti_kernel_stateÚalready_codegened_subgraphsÚallocated_workspacesr   r*   Úaot_inductorÚ debug_intermediate_value_printerÚallow_stack_allocationÚdebug_printerÚadditional_filesÚexternal_kernel_libs)rf  rÐ   Úhashedr  rt  s   f   €rb   rn  ÚPythonWrapperCodegen.__init__¯  s   ù€ Ü‰ÑÔØ7;ˆÔ(ØBDˆÔ#Ü:D»,ˆÔ&ô ?I»lˆÔ*Ü6@³lˆÔ"Ø9;ˆÔ&Ü*/«'ˆÔð ð 	Ôô &Ó'ˆŒÜ)Ó+ˆŒÜ)Ó+ˆŒÜ$Ó&ˆŒÜ#1Ó#3ˆÔ Ü/Ó1ˆÔÜ$2Ó$4ˆÔ!Ü%3Ó%5ˆÔ"Ü$2Ó$4ˆÔ!Ü6@³lˆÔ"ð IKˆÔ)Ø01ˆÔ(ð .0ˆÔÜHRËˆÔØ!#ˆŒ
ØˆŒØ')ˆÔ$ØˆŒØˆŒØˆŒÜ*+¯'©'×*=×*=Ð*=™,À2ˆŒÜ Ÿw™w×2×2Ð2™¸ˆŒØ8<ˆÔ)ð PTˆÔ&Ø+/ˆÔ(ð ð 	Ô&ô ‹Lð 	Ô"ô 9C»ˆÔØ $ˆÔà×!Ñ!Ô#ð &(ˆÔ"Ø$&ˆÔ!à×ÑÔä2°4×8Ò8à×ÑÔà×.Ñ.Ô0ä�w‰w××ÐÜ !§¡× 6Ñ 6× <Ñ <Ö >‘�à×#Ñ# DÖ1ñ !?ô $¤JÖ/Ó1ˆŒÜ¤
Ö+Ó-ˆŒ
ð 57ˆŒä$-×$7Ò$7¸Ô$=Ø×%Ñ%ó%
ˆÔ!ô 
�‰÷	;ó 
ð	;ð
  /ˆÔØ&(ˆŒÜ+5«<ˆŒÜ"2Ó"4ˆÔÜ<F»LˆÔ(Ø46ˆÔ!ô 1Ü &× 3Ñ 3× TÑ TÜ ×-Ñ-×DÑDô
ˆÔð !#ˆÔô
 6@³\ˆÖ!rd   Nc               ó(   € V ^8„  d   QhRRRRRRRR/# )	r\   Úis_subgraphrE  Úsubgraph_namerÖ   Úparent_wrapperrÔ   Úpartition_signaturesú!ir.GraphPartitionSignature | Noner`   )ra   s   "rb   rc   Ú!PythonWrapperCodegen.__annotate__&  s1   € ÷ &ñ &Øð&à!ð&ð 4ð&ð @ñ	&rd   c                	óz   € V '       d+   Vf   \        R4      hVf   \        R4      h\        WV4      # \        4       # )Nz expected subgraph_name to be setz!expected parent_wrapper to be set)r²   ÚSubgraphPythonWrapperCodegenr²  )rH  rI  rJ  rK  s   &&&&rb   ÚcreateÚPythonWrapperCodegen.create%  sI   € ÷ ØÒ$Ü$Ð%GÓHÐHØÒ%Ü$Ð%HÓIÐIÜ/ØÐ/Cóð ô $Ó%Ð%rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  6  s   € ÷ 'ñ ' dñ 'rd   c                	ó   € R V n         R# )ÚcallN)r*  re  s   &rb   r+  Ú)PythonWrapperCodegen.set_launcher_fn_name6  s
   € à &ˆÖrd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÐ   r‹   rE  r_   r„  r`   )ra   s   "rb   rc   rM  :  s!   € ÷ <ñ < 3ð <°ð <¸ñ <rd   c                	óF   € V P                   P                  V R V 24       R# )z = None  # N)r  rö   )rf  rÐ   rE  s   &&&rb   r2  Ú#PythonWrapperCodegen.write_constant:  s   € Ø�‰×Ñ   k°&°Ð:Ö;rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  =  s   € ÷ pBñ pB˜dñ pBrd   c                ó   € \         P                  P                  P                  4       pRpVe   VP                  e   RVP                   2pRp\        \        P                  P                  4      ^ 8”  d   RpM6\         P                  P                  P                  P                  '       d   RpT P                  P                  RV R\        P                   RV R	2R
R7       V P                   P                  RR
R7        ^ RIHp V P                   P                  RR
R7       \        P*                  '       d   V P                   P-                  R4       \        P.                  P0                  '       Ed    V P                   P-                  R4       V P                   P-                  R4       V P                   P-                  R4       V P                   P-                  R4       V P                   P-                  R4       V P                   P-                  R4       V P                   P-                  R4       \        P.                  P2                  ;'       g)    \4        P6                  P9                  \;        4       R4      pV P                   P-                  RV R24       RV R2pRV R2p\        P.                  P<                  p\        P.                  P>                  p	\        P.                  P@                  p
V P                   P-                  R4       V P                   P                  RV RV R V	 R!V
 R"2	4       V P                   P-                  R#V R$24       V P                   P-                  R%4       R# R#   \&        \(        3 d     ELgi ; i)&z>Write the header section of the generated Python wrapper code.r  Nz
# AOT ID: zRfrom torch._inductor.codegen.debug_utils import _print_debugging_tensor_value_infozFfrom torch._inductor.runtime.debug_utils import tracked_empty_strided
z
                aH  
                from ctypes import c_void_p, c_long, c_int
                import torch
                import math
                import random
                import os
                import tempfile
                from math import inf, nan
                from cmath import nanj
                from torch._inductor.hooks import run_intermediate_hooks
                from torch._inductor.utils import maybe_profile
                from torch._inductor.codegen.memory_planning import _align as align
                from torch import device, empty_strided
                from zq import AsyncCompile
                from torch._inductor.select_algorithm import extern_kernels
                z
            Tr&  aö  
                aten = torch.ops.aten
                inductor_ops = torch.ops.inductor
                _quantized = torch.ops._quantized
                assert_size_stride = torch._C._dynamo.guards.assert_size_stride
                assert_size_stride_grouped = torch._C._dynamo.guards.assert_size_stride_grouped
                assert_alignment = torch._C._dynamo.guards.assert_alignment
                empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
                empty_strided_cpu_pinned = torch._C._dynamo.guards._empty_strided_cpu_pinned
                empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
                empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
                empty_strided_mtia = torch._C._dynamo.guards._empty_strided_mtia
                reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
                alloc_from_pool = torch.ops.inductor._alloc_from_pool
                async_compile = AsyncCompile()
            )Ú_SymmetricMemoryzs
                empty_strided_p2p = torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p
                zfrom torch.cuda import nvtxz import triton.profiler as protonz%import triton.profiler.language as plzCfrom triton.profiler.hooks import HookManager as _ProtonHookManagerzimport tritonzimport atexitz	import oszetriton.set_allocator(lambda size, align, stream: torch.empty(size, dtype=torch.uint8, device='cuda'))Úprotonzos.makedirs("z", exist_ok=True)zos.path.join("z", "inductor")z", "inductor.chrome_trace")z^from torch._inductor.runtime.proton_utils import process_proton_trace as _proton_process_tracez�
                def _proton_finalize_and_postprocess():
                    proton.finalize()
                    _trace_path = zµ
                    if os.path.exists(_trace_path):
                        _proton_process_trace(
                            _trace_path,
                            group_by_sm=z0,
                            split_invocations=z0,
                            per_cta_occupancy=z,,
                        )
                z5if not _ProtonHookManager.active_hooks: proton.start(z], backend="instrumentation", data="trace"); atexit.register(_proton_finalize_and_postprocess)úpl.enable_semantic("triton"))!ÚtorchÚ_guardsÚTracingContextÚtry_getÚaot_graph_namer�   r*   r?  r@  Ú	_inductorÚtest_configsÚtrack_memory_lifecycler
  r?  r)   r‘   r  Útorch._C._distributed_c10dr[  ÚAttributeErrorÚImportErrorÚannotate_trainingrö   rî   Úproton_profilingÚproton_output_dirÚosÚpathr  r   Úproton_group_by_smÚproton_split_invocationsÚproton_per_cta_occupancy)rf  ÚcontextÚaot_config_commentÚinductor_debug_utilsr[  Ú
output_dirÚproton_nameÚ
trace_pathÚgroup_by_smÚsplit_invocationsÚper_cta_occupancys   &          rb   r.  Ú!PythonWrapperCodegen.write_header=  sB  € ä—-‘-×.Ñ.×6Ñ6Ó8ˆØÐØÒ 7×#9Ñ#9Ò#EØ#-¨g×.DÑ.DÐ-EÐ!FÐØ!ÐÜŒv×"Ñ"×CÑCÓDÀqÔHØ#wÑ Ü�_‰_×#Ñ#×0Ñ0×G×GÐGØ#lÐ à�‰×ÑðØ#Ð$ð %ô $×,Ñ,Ð-ð .à%Ð&ð 'ð!ð$ ð' 	ô 	
ð* 	�‰×Ñðð  ð# 	ô 	
ð&	õ Dà�K‰K×Ñðð ð	 ô ô ×#×#Ð#Ø�K‰K×!Ñ!Ð"?Ô@Ü�=‰=×)×)Ñ)Ø�K‰K×!Ñ!Ð"DÔEØ�K‰K×!Ñ!Ð"IÔJØ�K‰K×!Ñ!ØUôð �K‰K×!Ñ! /Ô2Ø�K‰K×!Ñ! /Ô2Ø�K‰K×!Ñ! +Ô.Ø�K‰K×!Ñ!ðGôô  Ÿ™×8Ñ8÷ ð ¼B¿G¹G¿L¹LÜ“ ó=ˆJð �K‰K×!Ñ! M°*°Ð=NÐ"OÔPØ*¨:¨,°nÐEˆKØ)¨*¨Ð5PÐQˆJÜ Ÿ-™-×:Ñ:ˆKÜ &§¡× FÑ FÐÜ &§¡× FÑ FÐØ�K‰K×!Ñ!Øpôð �K‰K×Ñð#ð $. ,ð /)ð *5¨ð 6/Ø/@Ð.Að B/Ø/@Ð.Að Bðôð �K‰K×!Ñ!ð Ø +˜}ð -DðDôð
 �K‰K×!Ñ!Ð"@ÖAñ[ *øô	 ¤Ð,ô 	Úð	ús   Ä#M8 Í8NÎNc               ó   € V ^8„  d   QhRR/# )r\   r  r‹   r`   )ra   s   "rb   rc   rM  ¯  s   € ÷ ñ ¨3ñ rd   c                	ó   € R # ru   r`   )rf  r  s   &&rb   Úinclude_extra_headerÚ)PythonWrapperCodegen.include_extra_header¯  ó   € Ùrd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  ²  s   € ÷ ñ °4ñ rd   c                	óØ   € V P                   P                  R \        P                   R24        ^ RIHp V P                   P                  RRR7       R#   \        \        3 d     R# i ; i)a3  
                import torch
                from math import inf, nan
                from torch._dynamo.testing import rand_strided
                from torch._dynamo.utils import preserve_rng_state
                from torch._inductor.select_algorithm import AlgorithmSelectorCache
                from aH   import AsyncCompile

                async_compile = AsyncCompile()
                generate_example_value = AlgorithmSelectorCache.generate_example_value
                empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
                empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
            )Ú_cuda_getCurrentRawStreamzU
                get_raw_stream = torch._C._cuda_getCurrentRawStream
                Tr&  N)r  r?  r)   r‘   Útorch._Cr‚  rh  rg  )rf  r‚  s   & rb   r0  Ú6PythonWrapperCodegen.write_kernel_autotune_defs_header²  sw   € Ø×!Ñ!×(Ñ(ðô $×,Ñ,Ð-ð .ðô	
ð 
	Ý:à×%Ñ%×,Ñ,ðð ð	 -ö øô œ^Ð,ô 	Úð	ús   ¯#A ÁA)Á(A)c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  Ð  s   € ÷ ñ ¨$ñ rd   c                	ó*  € R \         P                   R2p\        P                  P                  '       d^   V P
                  P                  V4       V P
                  P                  \        P                  P                  P                  R4      4       \        P                  P                  '       gb   V P                  P                  VRR7       V P                  P                  \        P                  P                  P                  R4      4       R# R# )zU
            import triton
            import triton.language as tl
            from z+ import start_graph, end_graph
            Úget_raw_streamTr&  N)r0   r‘   r*   rî   rï   rø   r?  rö   rH   rœ   r	  Úimport_get_raw_stream_asr]  r
  ©rf  Ú
import_strs   & rb   Úwrite_triton_header_onceÚ-PythonWrapperCodegen.write_triton_header_onceÏ  sÍ   € ðô $×,Ñ,Ð-ð .ðˆ
ô
 �=‰=×1×1Ð1Ø×&Ñ&×-Ñ-¨jÔ9Ø×&Ñ&×0Ñ0Ü—‘×"Ñ"×;Ñ;Ð<LÓMôô �w‰w×"×"Ð"Ø�L‰L×Ñ 
°$ÐÔ7Ø�L‰L×"Ñ"Ü—‘×"Ñ"×;Ñ;Ð<LÓMöñ #rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  á  s   € ÷ 	Bñ 	B¨Tñ 	Brd   c                	óÐ  € \         P                  P                  P                  R 4      p\        P
                  P                  '       d=   V P                  P                  V4      '       g   V P                  P                  V4       \         P                  P                  '       gA   V P                  P                  V4      '       g   V P                  P                  V4       R# R# R# )r‡  N)rH   rœ   r	  rˆ  r*   rî   rï   rø   Úcontainsrö   r]  r
  )rf  Úimport_get_raw_stream_strs   & rb   Úwrite_get_raw_stream_headerÚ0PythonWrapperCodegen.write_get_raw_stream_headerá  s£   € Ü$%§G¡G×$6Ñ$6×$OÑ$OØó%
Ð!ô �=‰=×1×1Ð1Ø×-Ñ-×6Ñ6Ð7P×QÒQØ×*Ñ*×4Ñ4Ð5NÔOÜ�w‰w×"×"Ð"Ø—<‘<×(Ñ(Ð)B×CÒCØ—‘×&Ñ&Ð'@ÖAñ Dñ #rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  í  s   € ÷ +ñ +°$ñ +rd   c                	ó&   € V P                  4        R # ru   )r‘  re  s   &rb   Ú write_get_raw_stream_header_onceÚ5PythonWrapperCodegen.write_get_raw_stream_header_onceì  s   € à×(Ñ(Ö*rd   c               ó    € V ^8„  d   QhRRRR/# )r\   ÚmetaÚTritonMetaParamsr_   r‹   r`   )ra   s   "rb   rc   rM  ð  s   € ÷ !ñ !Ð"2ð !°sñ !rd   c                	ó   € \        V4      pWP                  9  d¢   R \        V P                  4       2pW P                  V&   V P                  P	                  V RV 24       \
        P                  P                  '       d<   V P                  P	                  V RV 24       V P                  P                  V4       V P                  V,          # )r˜  r2  )Úreprr:  r  r  rö   r*   rî   rï   rø   r;  r  )rf  r˜  Úvars   && rb   Úadd_meta_onceÚ"PythonWrapperCodegen.add_meta_onceð  s¤   € ä�D‹zˆØ—{‘{Ô"Øœ˜TŸ[™[Ó)Ð*Ð+ˆCà #�K‰K˜ÑØ�K‰K×!Ñ! S E¨¨T¨FÐ"3Ô4Ü�}‰}×5×5Ð5Ø×*Ñ*×4Ñ4¸°u¸CÀ¸vÐ5FÔGØ—‘×#Ñ# CÔ(à�{‰{˜4Õ Ð rd   c               ó   € V ^8„  d   QhRR/# ©r\   r_   rs  r`   )ra   s   "rb   rc   rM  ÿ  s   € ÷ 
ñ 
 ñ 
rd   c                	ó|   € V P                  4        Uu. uF  qP                  V P                  4      NK  	  up# u upi ru   )Úget_graph_outputsr@  r  ©rf  r  s   & rb   Úget_output_refsÚ$PythonWrapperCodegen.get_output_refsþ  s?   € ð =A×<RÑ<RÔ<Tó
Ù<T°q×Ñ × 1Ñ 1Ö2Ñ<Tñ
ð 	
ùò 
s   “#9c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM    s   € ÷ ñ  $ñ rd   c                	ó   € R # ru   r`   re  s   &rb   Úmark_output_typeÚ%PythonWrapperCodegen.mark_output_type  ó   € Ùrd   c               ó   € V ^8„  d   QhRR/# )r\   r_   z9dict[str, ir.TensorBox | ir.TorchBindObject | sympy.Expr]r`   )ra   s   "rb   rc   rM    s   € ÷ $ñ $à	Bñ$rd   c                	ó6   € \         P                  P                  # ru   )rH   rœ   rê  re  s   &rb   Úget_graph_inputsÚ%PythonWrapperCodegen.get_graph_inputs  s   € ô �w‰w×#Ñ#Ð#rd   c               ó   € V ^8„  d   QhRR/# ©r\   r_   zlist[IRNode]r`   )ra   s   "rb   rc   rM    s   € ÷ %ñ % <ñ %rd   c                	ó6   € \         P                  P                  # ru   )rH   rœ   rò  re  s   &rb   r¢  Ú&PythonWrapperCodegen.get_graph_outputs  s   € Ü�w‰w×$Ñ$Ð$rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM    s   € ÷ ?ñ ?¨Dñ ?rd   c                	ó  € V P                  4       P                  4        Fä  w  r\        V\        P                  \
        P                  \
        P                  \
        P                  34      '       d   KU  V\        P                  P                  9  d   Kv  \        VP                  4       4      ^ 8X  d   K–  V P                  VP                  4       4      pV P                  VP                  4       4      pW43V P                   V&   Kæ  	  R# rl  )r­  r�   r   rh   ÚBasicr,   ÚTorchBindObjectÚGeneratorStateÚOpaqueObjectStaterH   rœ   Úgraph_input_namesrD   r7  r;  r8  r  )rf  rÐ   ÚbufrÃ   rÄ   s   &    rb   Úcodegen_input_size_assertsÚ/PythonWrapperCodegen.codegen_input_size_asserts  sÊ   € Ø×.Ñ.Ó0×6Ñ6Ö8‰IˆDÜØä—K‘KÜ×&Ñ&Ü×%Ñ%Ü×(Ñ(ð	÷ò ñ ð œ1Ÿ7™7×4Ñ4Ô4Ùô ˜SŸ\™\›^Ó,°Ô1ÙØ×+Ñ+¨C¯L©L«NÓ;ˆDØ×-Ñ-¨c¯n©nÓ.>Ó?ˆFØ15°ˆD×'Ñ'¨Ó-ó+ 9rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  '  s   € ÷ (ñ (¨4ñ (rd   c                	ór  € V P                   P                  R 4       V P                  4       P                  4        Fy  w  r\	        V\
        P                  \        P                  34      '       d   K7  RV R2pV P                   P                  V4       RV R2pV P                   P                  V4       K{  	  R# )z(# make sure graph inputs are not nan/infzassert not z.isnan().any().item()z.isinf().any().item()N)	r  rö   r­  r�   r   rh   rµ  r,   r¶  )rf  rÐ   rº  rô   s   &   rb   Úcodegen_input_nan_assertsÚ.PythonWrapperCodegen.codegen_input_nan_asserts'  s”   € Ø�‰×ÑÐHÔIØ×.Ñ.Ó0×6Ñ6Ö8‰IˆDÜ˜#¤§¡¬R×-?Ñ-?Ð@×AÒAÙØ   Ð&;Ð<ˆDØ�K‰K×!Ñ! $Ô'Ø   Ð&;Ð<ˆDØ�K‰K×!Ñ! $Ö'ó 9rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  1  s   € ÷ 
ñ 
¨$ñ 
rd   c                	ó<   € V P                   P                  R 4       R# )zV

            async_compile.wait(globals())
            del async_compile
            N)r  r?  re  s   &rb   Úwrite_async_compile_waitÚ-PythonWrapperCodegen.write_async_compile_wait1  s   € Ø�‰×Ñðö	
rd   c               ó   € V ^8„  d   QhRR/# )r\   Úinput_namesrs  r`   )ra   s   "rb   rc   rM  :  s   € ÷ .ñ . iñ .rd   c                	óÌ   € R P                  V4      p\        V4      ^8X  d
   VR,          pV P                  P                  V R24       V P                  P                  R4       R# )rT  Ú,z = argszargs.clear()N)r  r  r  rö   )rf  rÆ  Úlhss   && rb   Ú
write_argsÚPythonWrapperCodegen.write_args:  sP   € Ø�i‰i˜Ó$ˆÜˆ{Ó˜qÔ Ø�3�JˆCØ�‰×Ñ   W˜oÔ.Ø�‰×Ñ˜nÖ-rd   c               ó   € V ^8„  d   QhRR/# ©r\   r_   r�   r`   )ra   s   "rb   rc   rM  A  s   € ÷ ñ °3ñ rd   c                	óÆ   € \         P                  '       d    V P                  P                  R 4       ^pV# V P                  P                  RV P                   R24       ^pV# )aÙ  
                class Runner:
                    def __init__(self, partitions):
                        self.partitions = partitions

                    def recursively_apply_fns(self, fns):
                        new_callables = []
                        for fn, c in zip(fns, self.partitions):
                            new_callables.append(fn(c))
                        self.partitions = new_callables

                    def call(self, args):
                z
                def z(args):
                )r*   Úgraph_partitionr  r?  r*  ©rf  Úprefix_indents   & rb   Ú!write_launcher_fn_call_get_indentÚ6PythonWrapperCodegen.write_launcher_fn_call_get_indentA  sp   € Ü×!×!Ð!Ø�K‰K×Ñðôð ˆMð Ðð �K‰K×ÑðØ×*Ñ*Ð+ð ,ðôð
 ˆMàÐrd   c               ó   € V ^8„  d   QhRR/# r   r`   )ra   s   "rb   rc   rM  ]  s   € ÷ )ñ ) yñ )rd   c                	ó6   € \         P                  P                  # ru   )rH   rœ   r¹  re  s   &rb   Úget_graph_input_namesÚ*PythonWrapperCodegen.get_graph_input_names]  s   € Ü�w‰w×(Ñ(Ð(rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  `  s   € ÷ :ñ :˜dñ :rd   c                	óè  € V P                   f   \        R4      hV P                  4        V P                  4       pV P                  P                  V4      ;_uu_ 4        \        P                  P                  '       d   V P                  V P                  4       \        P                  P                  4       p\        P                  '       d    V P                  P                  RV R24       V P                  4       ;p'       d   V P!                  V4       V P#                  4        \%        4       '       d   \'        V 4      '       d   V P)                  4        R R R 4       R #   + '       g   i     R # ; i)Nz#expected launcher_fn_name to be setz0training_annotation = nvtx._device_range_start('ú'))r*  r²   rÃ  rÒ  r  r  r*   rî   Údebug_sync_graphÚgenerate_debug_syncrH   rœ   Úget_training_phaseri  rö   rÖ  rÊ  Úcodegen_inputsrA   r@   Ú"codegen_input_size_and_nan_asserts)rf  rÑ  Úphaser¹  s   &   rb   r/  Ú!PythonWrapperCodegen.write_prefix`  s  € Ø× Ñ Ò(Ü Ð!FÓGÐGØ×%Ñ%Ô'Ø×>Ñ>Ó@ˆà�[‰[×Ñ ×.Õ.Ü�}‰}×-×-Ð-Ø×(Ñ(¨¯©Ô5Ü—G‘G×.Ñ.Ó0ˆEÜ×'×'Ð'Ø—‘×%Ñ%ØFÀuÀgÈRÐPôð %)×$>Ñ$>Ó$@Ð@Ð Ö@Ø—‘Ð 1Ô2à×ÑÔ!ô
 -×.Ò.Ü<¸T×BÒBà×7Ñ7Ô9÷) /×.×.Ò.ús   ÁA/E Ã6E ÄAE Å E1	c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  |  s   € ÷ -ñ -°Dñ -rd   c                	ó¢   € \         P                  '       d   V P                  4        \         P                  '       d   V P	                  4        R # R # ru   )r*   Úsize_assertsr»  Únan_assertsr¿  re  s   &rb   rß  Ú7PythonWrapperCodegen.codegen_input_size_and_nan_asserts|  s7   € Ü××ÐØ×+Ñ+Ô-Ü××ÐØ×*Ñ*Ö,ñ rd   c               ó    € V ^8„  d   QhRRRR/# )r\   rÆ  úIterable[str]r_   r„  r`   )ra   s   "rb   rc   rM  †  s   € ÷ 
Lñ 
L¸-ð 
LÈDñ 
Lrd   c                	óJ  € . pV FE  pW0P                   9   g   K  V P                   P                  V4      w  rEVP                  W4V34       KG  	  \        V4      ^8X  d"   V^ ,          w  r4pV P	                  W4VR4       R# \        V4      ^8”  d   V P                  VR4       R# R# )rg   rß  N)r  r  rÈ   r  Úwrite_assert_size_strideÚ write_assert_size_stride_grouped)rf  rÆ  Úgrouped_assertsrÐ   rÃ   rÄ   s   &&    rb   Úcodegen_deferred_input_assertsÚ3PythonWrapperCodegen.codegen_deferred_input_asserts†  sš   € Ø68ˆÛˆDØ×2Ñ2Ö2Ø#×:Ñ:×>Ñ>¸tÓD‘�Ø×&Ñ&¨°FÐ';Ö<ñ  ô ˆÓ 1Ô$Ø!0°Õ!3ÑˆD˜Ø×)Ñ)¨$°f¸gÖFÜ�Ó! AÔ%Ø×1Ñ1°/À7ÖKñ &rd   c               ó0   € V ^8„  d   QhRRRRRRRRRRRR	/# )
r\   rÐ   r‹   rÃ   rÄ   rà  r   rá  r_   r„  r`   )ra   s   "rb   rc   rM  ’  sN   € ÷ 	Wñ 	Wàð	Wð ð	Wð ð		Wð
 ð	Wð "ð	Wð 
ñ	Wrd   c           
     ó>   € V P                  \        WW#WE4      4       R# )z7Queue an assert_size_stride for emission during replay.N)rö   rÞ  )rf  rÐ   rÃ   rÄ   rà  r   s   &&&&&&rb   rê  Ú-PythonWrapperCodegen.write_assert_size_stride’  s   € ð 	�‰Ô+¨D¸ÀgÓUÖVrd   c               ó4   € V ^8„  d   Qh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„  r`   )ra   s   "rb   rc   rM  �  sP   € ÷ ñ àðð ðð ð	ð
 ðð ðð "ðð 
ñrd   c                ó¶   € Vf"   VP                  RV RV RV RV: R2	4       R# V P                  R4       VP                  RV RV RV RV RV: R24       R# )z¥Emit one assert_size_stride line to `code` (replay-phase target).

Subclasses override to change the emitted form (e.g., C++ assert with
an AOTI runtime env guard).
Nzassert_size_stride(rT  r3  zHfrom torch._inductor.runtime.runtime_utils import assert_tensor_metadatazassert_tensor_metadata()rö   r  )rf  rŸ  rÐ   rÃ   rÄ   rà  r   s   &&&&&&&rb   rå  Ú0PythonWrapperCodegen._codegen_assert_size_stride�  sv   € ð Š=Ø�N‰NÐ0°°°b¸¸¸bÀÀÈÈ7É+ÐUVÐWÖXà× Ñ ØZôð �N‰NØ)¨$¨¨r°$°°r¸&¸ÀÀEÀ7È"ÈWÉKÐWXÐYörd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   Údivisor_strr‹   rà  r_   r„  r`   )ra   s   "rb   rc   rM  µ  s&   € ÷ Hñ H°Cð HÀ#ð HÈ$ñ Hrd   c                ó<   € V P                  \        WV4      4       R# )z¾Queue a div-by-zero AOTI check for emission during replay.

codegen_cpp_sizevar is called in both build and replay phases;
set_writeline routes the WrapperLine appropriately in either case.
N)rö   rû  )rf  rö  rà  s   &&&rb   Úwrite_assert_div_by_zeroÚ-PythonWrapperCodegen.write_assert_div_by_zeroµ  s   € ð 	�‰Ô*¨4¸gÓFÖGrd   c               ó(   € V ^8„  d   QhRRRRRRRR/# )r\   rŸ  r?   rö  r‹   rà  r_   r„  r`   )ra   s   "rb   rc   rM  ½  s,   € ÷ 

ñ 

Ø"ð

Ø14ð

Ø?Bð

à	ñ

rd   c                ó   € \        R4      h)zÁEmit one div-by-zero AOTI check to `code` (replay-phase target).

Only emitted by C++ wrappers; the Python wrapper relies on Python's
native ZeroDivisionError instead of pre-checking divisors.
z6AOTI div-by-zero check is only emitted by C++ wrappersrË  )rf  rŸ  rö  rà  s   &&&&rb   r   Ú0PythonWrapperCodegen._codegen_assert_div_by_zero½  s   € ô "ØDó
ð 	
rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rñ  rð  rà  r‹   r_   r„  r`   )ra   s   "rb   rc   rM  É  s)   € ÷ Lñ LØ1ðLØ<?ðLà	ñLrd   c                ó<   € V P                  \        WV4      4       R# )z>Queue a grouped assert_size_stride for emission during replay.N)rö   rï  )rf  rñ  rà  s   &&&rb   rë  Ú5PythonWrapperCodegen.write_assert_size_stride_groupedÉ  s   € ð 	�‰Ô2°4À'ÓJÖKrd   c               ó(   € V ^8„  d   QhRRRRRRRR/# )	r\   rŸ  r?   rñ  rð  rà  r‹   r_   r„  r`   )ra   s   "rb   rc   rM  Ï  s,   € ÷ 
ñ 
Ø"ð
Ø-Gð
ØRUð
à	ñ
rd   c                	óÔ   € R P                  R V 4       4      pR P                  R V 4       4      pR P                  R V 4       4      pVP                  RV RV RV RV: R2	4       R# )	rT  c              3  ó,   "  € T F
  w  p q!x € K  	  R # 5iru   r`   )rx   rÐ   Ú_s   &  rb   rz   ÚKPythonWrapperCodegen._codegen_assert_size_stride_grouped.<locals>.<genexpr>Ò  s   é € Ð9±¡: 4¨¨Aœ$³ùó   ‚c              3  ó*   "  € T F	  w  rqx € K  	  R # 5iru   r`   )rx   r  rÃ   s   &  rb   rz   r  Ó  s   é € Ð9±¡: 1¨Aœ$³ùs   ‚c              3  ó,   "  € T F
  w   rVx € K  	  R # 5iru   r`   )rx   r  rÄ   s   &  rb   rz   r  Ô  s   é € Ð?±w¡| q¨!�F³wùr  zassert_size_stride_grouped((z), (z), r3  N)r  rö   )rf  rŸ  rñ  rà  ÚnamesÚsizesÚstridess   &&&&   rb   rõ  Ú8PythonWrapperCodegen._codegen_assert_size_stride_groupedÏ  sk   € ð —	‘	Ñ9±Ó9Ó9ˆØ—	‘	Ñ9±Ó9Ó9ˆØ—)‘)Ñ?±wÓ?Ó?ˆØ�‰Ø*¨5¨'°°e°W¸DÀÀ	ÈÈWÉKÐWXÐYö	
rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  Ù  s   € ÷ ñ °ñ rd   c                óú  € \         P                  P                  '       d   R# \         P                  P                  pV'       g   R# \	        \         P                  P
                  4      pV FG  pW29  g   K  \         P                  P                  V,          pV P                  P                  V4       KI  	  V P                  '       d3   R\         P                  n	        V P                  P                  R4       R# R# )zyPopulate pending alignment copies for non-mutated inputs.
Called from the scheduler after mutated_input_idxs is computed.NTz6from torch._C._dynamo.guards import copy_if_misaligned)rH   rœ   r]  Úinputs_to_checkr"   Úmutated_input_idxsr¹  r  r  Ú_defers_input_alignmentr
  rö   )rf  r  Úmutated_idxsr  rÐ   s   &    rb   Úregister_alignment_check_inputsÚ4PythonWrapperCodegen.register_alignment_check_inputsÙ  s®   € ô �7‰7××ÐÙÜŸ'™'×1Ñ1ˆßÙô "¤!§'¡'×"<Ñ"<Ó=ˆÛ"ˆCØÖ&Ü—w‘w×0Ñ0°Õ5�Ø×.Ñ.×2Ñ2°4Ö8ñ #ð ×)×)Ð)Ø.2ŒA�G‰GÔ+Ø�L‰L×"Ñ"ØHöñ *rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÆ  rè  r©   r�   r_   r„  r`   )ra   s   "rb   rc   rM  î  s$   € ÷ #>ñ #>Ø(ð#>Ø25ð#>à	ñ#>rd   c                ó  € \         P                  P                  '       d   R# V Fá  pW0P                  9   d5   V P                  P	                  V4       V P                  V RV R24       KG  W0P                  9   g   KY  V P                  P                  V4      V8X  d   K{  W0P                  9  d2   V P                  P                  V4       V P                  V RV 24       V P                  V RV R24       W P                  V&   Kã  	  R# )aú  Emit alignment checks/clones for the given stream.

An input with an alignment constraint needs that input to be
potentially reallocated. This emits the proper copy_if_misaligned calls
to perform that realignment.

For inputs read on a simple stream this is trivial: we emit immediately
before the first reader. Inputs read on multiple streams get 1 copy per
stream, each reading from a preserved version of the original input.
This ensures we do not need to worry about ordering between multiple
streams.
Nz = copy_if_misaligned(r3  z_orig = z_orig))rH   rœ   r]  r  Údiscardrö   r  r  r;  r  r  )rf  rÆ  r©   rÐ   s   &&& rb   Ú!codegen_deferred_alignment_copiesÚ6PythonWrapperCodegen.codegen_deferred_alignment_copiesî  sæ   € ô �7‰7××ÐÙÛˆDØ×5Ñ5Ô5Ø×.Ñ.×6Ñ6°tÔ<Ø—‘ $ Ð'=¸d¸VÀ1ÐEÖFØ×;Ñ;Ö;ð ×1Ñ1×5Ñ5°dÓ;¸vÔEÙØ×9Ñ9Ô9ð ×.Ñ.×2Ñ2°4Ô8Ø—N‘N d V¨8°D°6Ð#:Ô;Ø—‘ $ Ð'=¸d¸VÀ6ÐJÔKØ7=×.Ñ.¨tÓ4ó%  rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r  rè  r_   r„  r`   )ra   s   "rb   rc   rM    s   € ÷ =ñ =°ð =À$ñ =rd   c                ó¨   € V FK  pW P                   9   g   K  V P                   P                  V4       V P                  P                  V4       KM  	  R# )a  Reclassify inputs read on more than one stream so their alignment
copy is emitted per consuming stream (see codegen_deferred_alignment_copies)
rather than once at the first reader by line order -- which would place
the only copy inside one branch's stream and race the other branches.N)r  r  r  r  )rf  r  rÐ   s   && rb   Úmark_multistream_alignmentÚ/PythonWrapperCodegen.mark_multistream_alignment  sB   € ó
 ˆDØ×5Ñ5Ö5Ø×.Ñ.×6Ñ6°tÔ<Ø×2Ñ2×6Ñ6°tÖ<ó rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rú  r�   rx  r‹   r_   r`   )ra   s   "rb   rc   rM     s!   € ÷ ñ ¨sð Àð Èñ rd   c                	ó¶  € V P                  4        \        V4      p\        P                  P                  '       d~   \        4       '       d)   \        P                  P                  P                  4       MTpV P                  P                  V R V R24       \        P                  P                  '       d   V# V P                  V R \        V4       R24       V# )ú = get_raw_stream(r3  )r‘  r8   r*   rî   rï   r   rH   rœ   r	  r  rø   rö   r]  rý  )rf  rú  rx  rÐ   Úautotune_idxs   &&&  rb   r8  Ú)PythonWrapperCodegen.write_get_raw_stream   s¶   € Ø×(Ñ(Ô*Ü" :Ó.ˆÜ�=‰=×1×1Ð1ô
 !—?’?ô —‘×"Ñ"×:Ñ:Ô<àð ð
 ×&Ñ&×0Ñ0Ø�&Ð*¨<¨.¸Ð:ôô �w‰w×"×"Ð"à�Ø�‰˜$˜Ð1Ô2FÀzÓ2RÐ1SÐSTÐUÔVØˆrd   c                	ó(   € V P                   R,          # )rg   éÿÿÿÿ)r,  re  s   &rb   Úget_codegened_graphÚ(PythonWrapperCodegen.get_codegened_graph4  s   € Ø×)Ñ)¨"Õ-Ð-rd   c                	ó<   € V P                   P                  V4       R # ru   )r,  rÈ   )rf  rœ   s   &&rb   r¾  Ú)PythonWrapperCodegen.push_codegened_graph7  s   € Ø×"Ñ"×)Ñ)¨%Ö0rd   c                	ó6   € V P                   P                  4       # ru   )r,  r  re  s   &rb   rò  Ú(PythonWrapperCodegen.pop_codegened_graph:  s   € Ø×)Ñ)×-Ñ-Ó/Ð/rd   c                	óP   € ^ RI Hp V P                  P                  V! V4      4      # )r   )Údeepcopy)Úcopyr+  r-  rÈ   )rf  r¸  r+  s   && rb   r·  Ú(PythonWrapperCodegen.push_computed_sizes=  s!   € Ý!à×(Ñ(×/Ñ/±¸Ó0HÓIÐIrd   c                	ó6   € V P                   P                  4       # ru   )r-  r  re  s   &rb   rî  Ú'PythonWrapperCodegen.pop_computed_sizesB  s   € Ø×(Ñ(×,Ñ,Ó.Ð.rd   c               ó   € V ^8„  d   QhRR/# r£  r`   )ra   s   "rb   rc   rM  E  s   € ÷ +ñ + Cñ +rd   c                	ó.   € \        V P                  4       # ru   )Únextr  re  s   &rb   Únext_kernel_suffixÚ'PythonWrapperCodegen.next_kernel_suffixE  s   € Ü�t×'Ñ'Ó(Ð)Ð*rd   c               ó(   € V ^8„  d   QhRRRRRRRR/# )r\   rú  r�   r«  r®  zdict[int, int] | Noner_   r„  r`   )ra   s   "rb   rc   rM  H  s2   € ÷ 5-ñ 5-àð5-ð ð5-ð %:ð	5-ð
 
ñ5-rd   c           
     	óö  € V^8”  d­   Vf   \        R4      h\        P                  P                  '       g?   RpV P                  P                  V4      '       g   V P                  P                  V4       V P                  V8g  pWn        V P                  \        VV P                  VVVR7      4       M$V P                  \        WP                  4      4       \        P                  P                  '       dö   V P                  4        \        4       '       d)   \        P                  P                   P#                  4       MTpV P$                  P                  R\        P                  P                   P'                  V4       R24       V P$                  P)                  4        \+        V 4      '       d   V P-                  4        V P$                  P                  \/        V4       RV R24       Wn	        W n        R# )	rg   Nz-expected stream_idx_to_user_obj_idx to be setz@from torch._dynamo.variables.streams import _get_stream_by_index)r¯  r  r  r  r3  )r²   rH   rœ   r]  r
  r�  rö   r  rª  r  rÿ  r*   rî   rï   r‹  r   r	  r  rø   r  r¿  r@   r‘  r8   Ú_num_streams)rf  rú  r«  r®  Úimport_liner¯  r   s   &&&&   rb   Úcodegen_device_guard_enterÚ/PythonWrapperCodegen.codegen_device_guard_enterH  s¡  € ð ˜Œ?Ø)Ò1Ü$Ð%TÓUÐUÜ—7‘7×&×&Ð&àVð ð —|‘|×,Ñ,¨[×9Ò9Ø—L‘L×*Ñ*¨;Ô7Ø!%×!AÑ!AÀZÑ!OÐØ/9Ô,Ø�N‰NÜ;ØØ×5Ñ5ØØ.Ø'9ôõð �N‰NÜ-Ø× AÑ Aóôô
 �=‰=×1×1Ð1à×)Ñ)Ô+ô
 !—?’?ô —‘×"Ñ"×:Ñ:Ô<àð ð
 ×&Ñ&×0Ñ0ØœŸ™×*Ñ*×7Ñ7¸ÓEÐFÀaÐHôð ×&Ñ&×0Ñ0Ô2Ü2°4×8Ò8à×0Ñ0Ô2Ø×&Ñ&×0Ñ0Ü& zÓ2Ð3Ð3EÀlÀ^ÐSTÐUôð -7Ô)Ø!,Örd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM    s   € ÷ 5ñ 5¨4ñ 5rd   c                	ó@  € \        V R 4      '       d7   V P                  ^8”  d&   V P                  \        V P                  4      4       MV P                  \	        4       4       \
        P                  P                  '       d   V P                  P                  4        R# R# )r7  N)
r„   r7  rö   r½  r  r*   rî   rï   rø   ró  re  s   &rb   Úcodegen_device_guard_exitÚ.PythonWrapperCodegen.codegen_device_guard_exit  so   € Ü�4˜×(Ò(¨T×->Ñ->ÀÔ-BØ�N‰NÜ:¸4×;LÑ;LÓMõð �N‰NÔ7Ó9Ô:Ü�=‰=×1×1Ð1Ø×&Ñ&×2Ñ2Ö4ñ 2rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r´  r�   r_   rÃ  r`   )ra   s   "rb   rc   rM  ‰  s   € ÷ ñ àðð 
$ñrd   c                óÄ   € \         P                  P                  P                  ;pe   \	        RV R\        V4       24      h\        VR7      pV P                  V4       V# )zíGenerate data structure for entering a CUDA Stream context.

Args:
    stream_idx: The index number of the entering CUDA Stream context.

Raises:
    ValueError: If this function is called while the previous stream context isn't exited.
z!Nested stream context switching: z -> r³  )rH   rœ   r    Úcurrent_stream_nameÚ
ValueErrorr9   rÃ  rö   )rf  r´  rA  Úctx_entrances   &&  rb   Úcodegen_cuda_stream_enterÚ.PythonWrapperCodegen.codegen_cuda_stream_enter‰  sf   € ô $%§7¡7×#4Ñ#4×#HÑ#HÐHÐÒUÜØ3Ð4GÐ3HÈÜ" :Ó.Ð/ð1óð ô 2¸ZÔHˆØ�‰�|Ô$ØÐrd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  Ÿ  s   € ÷ 4ñ 4¨$ñ 4rd   c                ó8   € V P                  \        4       4       R# )z:Generate data structure for exiting a CUDA Stream context.N)rö   rË  re  s   &rb   Úcodegen_cuda_stream_exitÚ-PythonWrapperCodegen.codegen_cuda_stream_exitŸ  s   € à�‰Ô0Ó2Ö3rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rŸ  r?   r´  r�   r_   r„  r`   )ra   s   "rb   rc   rM  £  s$   € ÷ ñ Ø"ðØ03ðà	ñrd   c                	ób   € VP                  R \        V4       R24       VP                  4        R# )r  r  N)rö   r9   r¿  )rf  rŸ  r´  s   &&&rb   rÇ  Ú6PythonWrapperCodegen.codegen_enter_cuda_stream_context£  s)   € ð 	�‰˜œ¨zÓ:Ð;¸1Ð=Ô>Ø�‰Örd   c               ó    € V ^8„  d   QhRRRR/# r¼  r`   )ra   s   "rb   rc   rM  ©  s   € ÷ ñ °^ð Èñ rd   c                	ó&   € VP                  4        R # ru   rß  r¤  s   &&rb   rÏ  Ú5PythonWrapperCodegen.codegen_exit_cuda_stream_context©  rá  rd   c               ó(   € V ^8„  d   QhRRRRRRRR/# )	r\   rŸ  r?   r«  r�   r®  r­  r_   r„  r`   )ra   s   "rb   rc   rM  ¬  s2   € ÷ "ñ "àð"ð ð"ð %3ð	"ð
 
ñ"rd   c                	ó   € \         hru   rË  )rf  rŸ  r«  r®  s   &&&&rb   r¶  Ú1PythonWrapperCodegen.codegen_stream_info_prologue¬  s
   € ô "Ð!rd   c               ó    € V ^8„  d   QhRRRR/# )r\   rª   ztuple[int, int]r_   r„  r`   )ra   s   "rb   rc   rM  ´  s   € ÷ Qñ Q°/ð QÀdñ Qrd   c                ó  € \         P                  P                  '       d   \        R4      hRpV P                  P                  V4      '       g   V P                  P                  V4       Vw  r4V P                  \        W44      4       R# )z>Generate a CUDA MemPool context around code that may allocate.ú,CUDA MemPool contexts require Python wrapperzLfrom torch._dynamo.graph_bytecode_inputs import get_external_object_by_indexN)rH   rœ   r]  r²   r
  r�  rö   rÓ  )rf  rª   r8  rÔ  rÕ  s   &&   rb   Úcodegen_cuda_mempool_enterÚ/PythonWrapperCodegen.codegen_cuda_mempool_enter´  sk   € ä�7‰7××ÐÜ Ð!OÓPÐPð+ð 	ð �|‰|×$Ñ$ [×1Ò1Ø�L‰L×"Ñ" ;Ô/Ø&-Ñ#ˆØ�‰Ô2°=ÓOÖPrd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  Á  s   € ÷ 5ñ 5¨4ñ 5rd   c                óŽ   € \         P                  P                  '       d   \        R4      hV P	                  \        4       4       R# )zEGenerate data structure for exiting the current CUDA MemPool context.rU  N)rH   rœ   r]  r²   rö   rÛ  re  s   &rb   Úcodegen_cuda_mempool_exitÚ.PythonWrapperCodegen.codegen_cuda_mempool_exitÁ  s.   € ä�7‰7××ÐÜ Ð!OÓPÐPØ�‰Ô1Ó3Ö4rd   c               ó    € V ^8„  d   QhRRRR/# )r\   Úoutput_refsrs  r_   r„  r`   )ra   s   "rb   rc   rM  Ç  s   € ÷ 5ñ 5¨9ð 5¸ñ 5rd   c                	óÔ  € V'       EdD   \         P                  '       dô   V P                  P                  R RP	                  V4      ,           R,           4       V P                  P                  R4       V P                  P                  4        V P                  P                  R4       V P                  P                  4        V P                  P                  R4       V P                  P                  R4       V P                  P                  ^4       V P                  P                  RRP	                  V4      ,           R,           4       R	# V P                  P                  R4       R	# )
zreturn_vars = (rT  ú, )zfor var in return_vars:z!if isinstance(var, torch.Tensor):z#assert not var.isnan().any().item()z#assert not var.isinf().any().item()zreturn (z	return ()N)r*   rå  r  rö   r  r¿  ró  )rf  r]  s   &&rb   Úgenerate_returnÚ$PythonWrapperCodegen.generate_returnÇ  s  € ßˆ;Ü×!×!Ð!Ø×!Ñ!×+Ñ+Ø%¨¯	©	°+Ó(>Õ>ÀÕFôð ×!Ñ!×+Ñ+Ð,EÔFØ×!Ñ!×+Ñ+Ô-Ø×!Ñ!×+Ñ+Ð,OÔPØ×!Ñ!×+Ñ+Ô-Ø×!Ñ!×+Ñ+Ð,QÔRØ×!Ñ!×+Ñ+Ð,QÔRØ×!Ñ!×-Ñ-¨aÔ0à×Ñ×'Ñ'¨
°T·Y±Y¸{Ó5KÕ(KÈeÕ(SÖTà×Ñ×'Ñ'¨Ö4rd   c               ó    € V ^8„  d   QhRRRR/# ©r\   Úresultr?   r_   r„  r`   )ra   s   "rb   rc   rM  Ù  s   € ÷ ñ ¨^ð Àñ rd   c                	ó   € R # ru   r`   ©rf  rd  s   &&rb   Úgenerate_before_suffixÚ+PythonWrapperCodegen.generate_before_suffixÙ  rª  rd   c               ó    € V ^8„  d   QhRRRR/# rc  r`   )ra   s   "rb   rc   rM  Ü  s   € ÷ ñ ¨Nð ¸tñ rd   c                	óÜ   € \         P                  '       dV   R P                  V P                  4      \	        V P                  4      ^8X  d   RMR,           pVP                  RV R24       R# R# )rT  rÈ  r  z-
                runner = Runner(partitions=[z{])
                call = runner.call
                recursively_apply_fns = runner.recursively_apply_fns
                N)r*   rÏ  r  Úall_partition_namesr  r?  )rf  rd  Úall_partition_name_lists   && rb   Úgenerate_after_suffixÚ*PythonWrapperCodegen.generate_after_suffixÜ  sh   € Ü×!×!Ð!Ø&*§i¡i°×0HÑ0HÓ&IÜ˜4×3Ñ3Ó4¸Ô9‘¸rõ'Ð#ð �M‰Mð-Ø-DÐ,Eð Fðöñ "rd   c               ó    € V ^8„  d   QhRRRR/# rc  r`   )ra   s   "rb   rc   rM  ê  s   € ÷ ñ  >ð °dñ rd   c                	ó   € R # ru   r`   rf  s   &&rb   Úgenerate_endÚ!PythonWrapperCodegen.generate_endê  rª  rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r˜   zir.FallbackKernelr_   r„  r`   )ra   s   "rb   rc   rM  í  s   € ÷ :ñ :Ð->ð :À4ñ :rd   c                	óÎ   € VP                   pVe;   V\        9   d0   \        V,          P                  pVe   V! WP                  4       R # V P                  \	        W4      4       R # ru   )r¼  rP   Úpythonrö   r%  )rf  r˜   rà  Úcustom_codegens   &&  rb   Úgenerate_fallback_kernelÚ-PythonWrapperCodegen.generate_fallback_kernelí  sR   € à×)Ñ)ˆØÒ 7Ô.JÔ#JÜ9¸'ÕB×IÑIˆNØÒ)Ù˜t§^¡^Ô4ÙØ�‰Ô,¨TÓ8Ö9rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r˜   rN  r_   r„  r`   )ra   s   "rb   rc   rM  ÷  s   € ÷ 4ñ 4Ð5Lð 4ÐQUñ 4rd   c                ó"  € VP                    F  pV P                  V4       K  	  V P                  \        W4      4       VP                    F@  p\	        VP
                  \        P                  4      '       g   K/  VP                  V 4       KB  	  R# )z7Generate .out() call with pre-allocated output buffers.N)	rV  Úcodegen_allocationrö   rM  r   rµ   r,   ÚLayoutÚcodegen_size_asserts)rf  r˜   rX  s   && rb   Ú generate_extern_kernel_multi_outÚ5PythonWrapperCodegen.generate_extern_kernel_multi_out÷  se   € à×5Ô5ˆHØ×#Ñ# HÖ-ñ 6à�‰Ô/°Ó;Ô<Ø×5Ô5ˆHÜ˜(Ÿ/™/¬2¯9©9×5Ô5Ø×-Ñ-¨dÖ3ó 6rd   c               ó   € V ^8„  d   QhRR/# )r\   r˜   r&  r`   )ra   s   "rb   rc   rM   	  s   € ÷ ,ñ ,Ð1Eñ ,rd   c                	óÖ   € VP                  V 4       V P                  \        W4      4       \        VP                  \
        P                  4      '       d   VP                  V 4       R # R # ru   )Úcodegen_commentrö   r%  r   rµ   r,   r|  r}  rü  s   &&rb   Úgenerate_extern_kernel_allocÚ1PythonWrapperCodegen.generate_extern_kernel_alloc 	  sJ   € Ø×Ñ˜TÔ"Ø�‰Ô,¨TÓ8Ô9Ü�d—k‘k¤2§9¡9×-Ò-Ø×%Ñ% dÖ+ñ .rd   c           
     	ób  € \        VP                  \        P                  4      pVP	                  4       pVP                  4       pVP                  4       pV P                  p\        P                  '       d   R V9   d   RV 2p\        P                  P                  '       d!   V P                  W`P                  4       4      pMTpV'       d7   V P                  V P                   V RRP!                  V4       RV 24       R# V P                  V P                   V RV RRP!                  V4       RV 24       V P"                  '       dd   \        P$                  '       dL   VeF   \&        R,          R;;,          ^,          uu&   V P                  R	VP(                  : RV R24       R# R# R# R# )
Úview_as_complexú.clone()rS  rT  r3  r2  NÚinductorÚintermediate_hookszrun_intermediate_hooks()r   rµ   r,   Ú
NoneLayoutrž   Úget_origin_noder;  r”  r*   Úmemory_planningÚtraceÚprovenance_tracking_to_timelineÚ$define_extern_kernel_profile_wrapperr3  rö   r“  r  r'  Úgenerate_intermediate_hooksr   rÐ   )	rf  Úextern_kernelr-  Ú	no_returnrÏ  Úorigin_noderC  r”  Úwrapper_names	   &&&      rb   r,  Ú9PythonWrapperCodegen._generate_extern_kernel_alloc_helper	  sx  € ô ˜}×3Ñ3´R·]±]ÓCˆ	Ø#×,Ñ,Ó.ˆØ#×3Ñ3Ó5ˆØ#×3Ñ3Ó5ˆØ—‘ˆÜ×!×!Ð!Ð&7¸;Ô&Fð   ˜xÐ(ˆFä�<‰<×7×7Ð7Ø×DÑDØ×4Ñ4Ó6ó‰Lð 'ˆLçØ�N‰N˜dŸl™l˜^¨L¨>¸¸4¿9¹9ÀT»?Ð:KÈ1ÈVÈHÐUÖVà�N‰NØ—<‘<�.  ¨S°°¸aÀÇ	Á	È$ÃÐ?PÐPQÐRXÐQYÐZôð ×0×0Ð0Ü×6×6Ð6ØÒ+ä˜Õ$Ð%9×:¸aÕ?Ó:Ø—‘Ø-¨k×.>Ñ.>Ñ-AÀÀKÀ=ÐPQÐRöñ ,ñ 7ñ 1rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r˜   r5  r_   r„  r`   )ra   s   "rb   rc   rM  *	  s   € ÷ 8ñ 8à ð8ð 
ñ8rd   c                	ó\   € VP                  V 4       V P                  \        W4      4       R # ru   )r‚  rö   r4  rü  s   &&rb   Úgenerate_extern_kernel_outÚ/PythonWrapperCodegen.generate_extern_kernel_out*	  s$   € ð 	×Ñ˜TÔ"Ø�‰Ô*¨4Ó6Ö7rd   c               ó4   € V ^8„  d   QhRRRRRRRRRRR	R
RR/# )r\   rV  r‹   ÚoutÚout_viewrÖ   r-  rs  r�   Ústack_traceszOrderedSet[str] | Noner_   r„  r`   )ra   s   "rb   rc   rM  1	  sY   € ÷ Añ AàðAð ðAð ð	Að
 ðAð ðAð -ðAð 
ñArd   c           	     	óÐ  € \         P                  P                  P                  pVP	                  WAR R R4       TP                  RT;'       g    T 24       T;_uu_ 4        \        P                  P                  '       d!   V P                  WP                  4       4      pMTpV P                  V RRP                  V4       R24       R R R 4       R #   + '       g   i     R # ; i)NÚexternzout=rS  rT  r3  )rH   rœ   rµ  rB  Úset_printer_argsrÈ   r*   r�  rŽ  r�  r3  rö   r  )	rf  rV  r›  rœ  r-  r�   r�  Údebug_printer_managerr”  s	   &&&&&&&  rb   r?  Ú7PythonWrapperCodegen._generate_extern_kernel_out_helper1	  s²   € ô !"§¡× 4Ñ 4× BÑ BÐØ×.Ñ.¨t¸TÀ4ÈÔRØ�‰�d˜8Ÿ?˜? sÐ+Ð,Ô-ß"Ò"Ü�|‰|×;×;Ð;Ø#×HÑHØ×3Ñ3Ó5ó ‘ð  &�Ø�N‰N˜l˜^¨1¨T¯Y©Y°t«_Ð,=¸QÐ?Ô@÷ #×"×"Ò"ús   Á!A)CÃC%	c                	ó$  a € VP                   pVP                  pV'       dS   \        P                  P                  P                  V4      p\        P                  P                  P                  V4      pVP                  P                  4        R 2pRP                  V 3R lV 4       4      pRP                  V 3R lV 4       4      p\        P                  S VP                  4      pRpV RVP                   R2pV RV RV RV 2p	V RV	 R2p
V
# )	z.data_ptr()rT  c              3  óP   <"  € T F  p\         P                  SV4      x € K  	  R # 5iru   ©r²  Úval_to_arg_str©rx   Údimrf  s   & €rb   rz   ÚRPythonWrapperCodegen._generate_tma_descriptor_call_experimental.<locals>.<genexpr>P	  s$   øé € ÐXÑSWÈCÔ-×<Ñ<¸TÀ3×GÐGÓSWùó   ƒ#&c              3  óP   <"  € T F  p\         P                  SV4      x € K  	  R # 5iru   r¥  r§  s   & €rb   rz   r©  Q	  s%   øé € ð 
ÙFP¸sÔ ×/Ñ/°°c×:Ð:Ãjùrª  z$triton.tools.experimental_descriptorz.create_Úd_tma_descriptorrS  r3  )ÚdimsÚ
block_dimsrH   rœ   r¥   rû  Úoptimization_hintsÚtensorr@  r  r²  r¦  Úelement_sizeÚrank)rf  ÚdescÚapply_size_hintsr­  r®  Úptrr±  r  r9  r-  rT  s   f&&        rb   Ú*_generate_tma_descriptor_call_experimentalÚ?PythonWrapperCodegen._generate_tma_descriptor_call_experimentalG	  sþ   ø€ Ø�y‰yˆØ—_‘_ˆ
ßÜ—7‘7×#Ñ#×5Ñ5°dÓ;ˆDÜŸ™×)Ñ)×<Ñ<¸ZÓHˆJà—‘×.Ñ.Ó0Ð1°Ð=ˆà�y‰yÔXÑSWÓXÓXˆØ—Y‘Yô 
ÙFPó
ó 
ˆ
ô ,×:Ñ:¸4À×ARÑARÓSˆØ7ˆØˆx�x §	¡	˜{Ð*:Ð;ˆØ��b˜˜˜b  ¨B¨|¨nÐ=ˆØ��Q�t�f˜AˆˆØˆrd   c                	óÜ   € VP                   pV'       d*   \        P                  P                  P	                  V4      pR pV R2pVP
                  P                  4        RV 2pV RV R2pV# )z/triton.tools.tensor_descriptor.TensorDescriptorz.from_tensorrT  rS  r3  )Úblock_shaperH   rœ   r¥   r¯  r°  r@  )rf  r³  r´  r¹  r  r9  r-  rT  s   &&&     rb   Ú$_generate_tma_descriptor_call_stableÚ9PythonWrapperCodegen._generate_tma_descriptor_call_stable[	  so   € Ø×&Ñ&ˆßÜŸ'™'×*Ñ*×=Ñ=¸kÓJˆKàBˆØˆx�|Ð$ˆØ—+‘+×/Ñ/Ó1Ð2°"°[°MÐBˆØ��Q�t�f˜AˆˆØˆrd   c                	óô   € \        V\        P                  4      '       d   V P                  W4      # \        V\        P                  4      '       g   \        R \        V4       24      hV P                  W4      # )z"expected TMADescriptorStable, got )r   r,   ÚTMADescriptorExperimentalr¶  ÚTMADescriptorStabler²   r…   rº  )rf  r³  r´  s   &&&rb   Ú_generate_tma_descriptor_callÚ2PythonWrapperCodegen._generate_tma_descriptor_callf	  sh   € Ü�dœB×8Ñ8×9Ò9Ø×BÑBØóð ô ˜d¤B×$:Ñ$:×;Ò;Ü$Ð'IÌ$ÈtË*ÈÐ%VÓWÐWØ×<Ñ<¸TÓTÐTrd   c                	ó„   € V P                  V4      pVP                   R V V P                   2pV P                  V4       R# ©r2  N)r¿  rÐ   r”  rö   )rf  r³  rT  rô   s   &&  rb   Úgenerate_tma_descriptorÚ,PythonWrapperCodegen.generate_tma_descriptorp	  s:   € Ø×1Ñ1°$Ó7ˆØ—)‘)�˜C ˜v d§k¡k ]Ð3ˆØ�‰�tÖrd   c               ó   € V ^8„  d   QhRR/# )r\   r˜   r²  r`   )ra   s   "rb   rc   rM  u	  s   € ÷ 8ñ 8Ð.@ñ 8rd   c                	ó:   € V P                  \        W4      4       R # ru   )rö   r±  rü  s   &&rb   Úgenerate_scatter_fallbackÚ.PythonWrapperCodegen.generate_scatter_fallbacku	  s   € Ø�‰Ô*¨4Ó6Ö7rd   c	                	ó®  € Tp	\         P                  P                  '       d    V P                  W@P	                  4       4      pV R RP                  \        \        V4      4       2p
V	P                  R4      '       d"   V
RP                  R.V,           4      ,          p
MV'       d   V
R\        V4       2,          p
V
R,          p
V P                  V
4       R# )rS  rÈ  zaten.scatter_reducerT  r  z	, reduce=r3  N)r*   r�  rŽ  r�  r3  r  Úmapr‹   rD  r›  rö   )rf  rù   r�   r<  r¼  r¹  r+  r  r�   Úorig_python_kernel_namerô   s   &&&&&&&&&  rb   r»  Ú/PythonWrapperCodegen._generate_scatter_fallbackx	  s³   € ð #5ÐÜ�<‰<×7×7Ð7Ø!%×!JÑ!JØ"×$;Ñ$;Ó$=ó"Ðð %Ð% Q s§x¡x´´C¸Ó0@Ó'AÐ&BÐCˆØ"×-Ñ-Ð.C×DÒDØ�D—I‘I˜r˜d V�mÓ,Õ,‰DßØ�i¤ V£˜~Ð.Õ.ˆDØ��ˆØ�‰�tÖrd   c               ó    € V ^8„  d   QhRRRR/# )r\   rC  r‹   r  r`   )ra   s   "rb   rc   rM  �	  s   € ÷ #ñ #Àð #ÈSñ #rd   c                óž  € \         P                  ! RRRV RV 24      p\        P                  ! W4       V P                  P                  RV R24       V P                  P                  4       ;_uu_ 4        V P                  P                  R4       V P                  P                  RV R24       R	R	R	4       V#   + '       g   i     T# ; i)
zHWrap extern calls so profiler events use names with provenance metadata.z\Wr  Úextern_kernels_rü   z(*args, **kwargs):zE# Extra indirection is only enabled for profiler timeline provenance.rý   z(*args, **kwargs)N)ÚreÚsubÚinductor_debugÚalias_kernel_provenancer  r?  r  )rf  rC  r  Úwrapper_kernel_names   &&& rb   r�  Ú9PythonWrapperCodegen.define_extern_kernel_profile_wrapper�	  s¹   € ä ŸfšfØØØ˜k˜]¨!¨F¨8Ð4ó
Ðô
 	×.Ò.¨{ÔPØ�‰×Ñ˜TÐ"5Ð!6Ð6HÐIÔJØ�[‰[×Ñ×!Õ!Ø�K‰K×ÑØWôð �K‰K×Ñ ¨¨Ð5FÐGÔH÷	 "ð
 #Ð"÷ "Ö!ð
 #Ð"ús   Á6;B;Â;C	c               ó    € V ^8„  d   QhRRRR/# )r\   r˜   rœ  r_   r„  r`   )ra   s   "rb   rc   rM   	  s   € ÷ Bñ BÐ0Cð BÈñ Brd   c                	ó®  € . pVP                   R ,          p\        V4      p\        VP                  4       F€  w  rVVP                  V,          eV   \	        V4      p\        V\        P                  4      '       g   \        R\        V4       24      hVP                  V4       Ko  VP                  R4       K‚  	  V P                  \        WV4      4       R# )ºr\   NNNzexpected IRNode, got )r�   ÚiterÚ	enumerater„  r2  r   r,   r.   r²   r…   rÈ   rö   r›  )rf  r˜   r„  Úvalid_indicesÚiter_valid_indicesr¹  r  r  s   &&      rb   Úgenerate_index_put_fallbackÚ0PythonWrapperCodegen.generate_index_put_fallback 	  s£   € à*,ˆØŸ™ B�ˆÜ! -Ó0ÐÜ˜dŸl™lÖ+‰DˆAØ�|‰|˜A�Ò*ÜÐ/Ó0�Ü! %¬¯©×3Ò3Ü(Ð+@ÄÀeÃÀÐ)NÓOÐOØ—‘˜uÖ%à—‘˜tÖ$ñ ,ð 	�‰Ô+¨D¸Ó@ÖArd   c                	óx   € R RP                  V4       R2pW&WE.pV P                  V P                  W4      4       R# )r‰  rT  rŠ  N)r  rö   Úwrap_kernel_call)rf  rV  r  r„  r«  Ú
accumulateÚindices_strr-  s   &&&&&&  rb   r©  Ú1PythonWrapperCodegen._generate_index_put_fallback°	  s;   € Ø˜$Ÿ)™) GÓ,Ð-¨QÐ/ˆØ Ð3ˆØ�‰�t×,Ñ,¨VÓ:Ö;rd   c               ó4   € V ^8„  d   QhRRRRRRRRRR	R
RRR/# )r\   Úbuf_namer‹   r¼  Úget_argszCallable[[], Sequence[str]]Úop_overloadz6torch._ops.OpOverload | torch._ops.HigherOrderOperatorrr  rƒ  rÐ  zSequence[ir.Buffer]r_   r„  r`   )ra   s   "rb   rc   rM  µ	  sZ   € ÷ 	Vñ 	Vàð	Vð  ð	Vð .ð		Vð
 Lð	Vð  ð	Vð %ð	Vð 
ñ	Vrd   c           
     	ób   € V P                  V R V RRP                  V! 4       4       R24       R# )r2  rS  rT  r3  N)rö   r  )rf  rå  r¼  ræ  rç  rr  rÐ  s   &&&&&&&rb   Ú,generate_fallback_kernel_with_runtime_lookupÚAPythonWrapperCodegen.generate_fallback_kernel_with_runtime_lookupµ	  s2   € ð 	�‰˜(˜ 3Ð'9Ð&:¸!¸D¿I¹IÁhÃjÓ<QÐ;RÐRSÐTÖUrd   c                	óˆ   € \        R 4      ;_uu_ 4        V P                  V4      uuRRR4       #   + '       g   i     R# ; i)zPythonWrapperCodegen.generateN)r   Ú	_generate©rf  Úis_inferences   &&rb   ÚgenerateÚPythonWrapperCodegen.generateÀ	  s)   € ÜÐ9×:Õ:Ø—>‘> ,Ó/÷ ;×:×:Ó:ús	   ”0°A	c               ó   € V ^8„  d   QhRR/# rÍ  r`   )ra   s   "rb   rc   rM  Ä	  s   € ÷ ñ ¨ñ rd   c                	ó6   € \         P                  '       d   ^# ^# )r\   )r*   rÏ  re  s   &rb   Úget_wrapper_call_indentÚ,PythonWrapperCodegen.get_wrapper_call_indentÄ	  s   € Ü×!×!Ð!Ùárd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   Útargetr?   rö   úCallable[..., None]r_   zIterator[Callable[..., None]]r`   )ra   s   "rb   rc   rM  Ë	  s(   € ÷ !ñ !àð!ð 'ð!ð 
'ñ	!rd   c              #  óv   aa"  € V P                   pR VV3R llp W@n         Vx € W0n         R#   Y0n         i ; i5i)z–Temporarily redirect self.writeline to write into `target`.

Plain strings go through `writeline`.
WrapperLines are codegen'd into `target` directly.
c               ó    € V ^8„  d   QhRRRR/# )r\   rô   z2LineContext | DeferredLineBase | WrapperLine | strr_   r„  r`   )ra   s   "rb   rc   Ú8PythonWrapperCodegen.set_writeline.<locals>.__annotate__×	  s   € ÷ 	 ñ 	 ÐMð 	 ÐRVñ 	 rd   c                ój   <€ \        V \        4      '       d   V P                  S4       R # S! V 4       R # ru   )r   rˆ   r¥  )rô   rö  rö   s   &€€rb   ÚdispatchÚ4PythonWrapperCodegen.set_writeline.<locals>.dispatch×	  s#   ø€ Ü˜$¤×,Ò,à—‘˜VÖ$á˜$–rd   N©rö   )rf  rö  rö   Úoldrü  s   &ff  rb   Úset_writelineÚ"PythonWrapperCodegen.set_writelineÊ	  s4   ùé € ð �n‰nˆ÷	 ñ 	 ð	!Ø%ŒNØŠNà ŽNø˜S�Nüs   „9œ
. ¦9®6¶9c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  ä	  s   € ÷ ,ñ ,¨$ñ ,rd   c                	óâ   € V P                   P                  p\        P                  P                  '       d   V P
                  P                  V4       R # V P                  P                  V4       R # ru   )r<  Úkernel_defsr*   rî   rï   r  r?  r  )rf  r  s   & rb   Ú_write_multi_kernel_defsÚ-PythonWrapperCodegen._write_multi_kernel_defsä	  sI   € Ø×-Ñ-×9Ñ9ˆÜ�=‰=×1×1Ð1Ø×%Ñ%×,Ñ,¨[Ö9à�K‰K×Ñ˜{Ö+rd   c           	     	ó"  € \         P                  '       d   V P                  4        \        P                  ! 4       ;_uu_ 4       pVP                  V P                  P                  4       4       \         P                  '       d   V P                  V4       \         P                  '       d   V P                  4        V P                  V4       \         P                  P                  '       d1   \         P                  P                  '       g   V P                  4        V P!                  V P                  V P                  P"                  4      ;_uu_ 4        V P$                   FQ  p\'        V\(        4      '       d   VP+                  V P                  4       K6  V P                  P#                  V4       KS  	  R R R 4       V P-                  4        V P/                  4       pV P1                  4        \         P                  P2                  '       d   V P5                  V P                  4       \         P                  '       d   V P7                  4        \         P                  P8                  '       d   V P;                  4        \         P                  P                  '       d1   \         P                  P                  '       g   V P=                  4        \         P                  P                  '       d   V P?                  4        \         P@                  '       d2   \         PB                  '       g   V P                  P#                  R4       V PE                  V4       R R R 4       \G        4       pVPI                  V PJ                  4       VP#                  R4       VPI                  V PL                  4       \N        PP                  PR                  '       dk   \N        PP                  PB                  '       dK   \N        PP                  PT                  '       d+   \N        PP                  PV                  '       g   \G        4       pVPI                  V PX                  4       V P[                  4        VPI                  V P\                  4       V P_                  4       pVP                  V4      ;_uu_ 4        VPI                  V P                  4       R R R 4       V Pa                  V4       VPI                  V Pb                  4       V Pe                  V4       V Pg                  V4       V Pi                  V4       VPk                  4       V Pl                  Pk                  4       3#   + '       g   i     EL©; i  + '       g   i     EL; i  + '       g   i     LÁ; i)Nz+nvtx._device_range_end(training_annotation)r  )7r*   Úprofile_bandwidthr‹  r  Ú	ExitStackÚenter_contextr  r  Úprofiler_mark_wrapper_callÚ#generate_profiler_mark_wrapper_callÚgenerate_start_graphÚrun_wrapper_ir_passesrî   Ústore_cubinrï   Ú!generate_reset_kernel_saved_flagsr   rö   r"  r   rˆ   r¥  r  r¤  r¨  rÛ  rÜ  Úgenerate_end_graphrj  Úgenerate_proton_finalizeÚ generate_save_uncompiled_kernelsÚgenerate_and_run_autotune_blockri  r]  r`  r?   r?  r
  r  rH   rœ   r  Úis_const_graphr  r  Úfinalize_prefixr  ró  rg  r  rm  rq  Úadd_benchmark_harnessÚgetvaluewithlinemapr  )rf  rî  Ústackrô   r]  rd  Úwrapper_call_indents   &&     rb   rì  ÚPythonWrapperCodegen._generateë	  s¯  € Ü×#×#Ð#Ø×)Ñ)Ô+ä×!Ò!×#Ô# uØ×Ñ × 1Ñ 1× 8Ñ 8Ó :Ô;Ü×0×0Ð0Ø×8Ñ8¸Ô?Ü×'×'Ð'Ø×)Ñ)Ô+à×&Ñ& |Ô4ä�}‰}×(×(Ð(´·±×1W×1WÐ1WØ×6Ñ6Ô8ð ×#Ñ# D×$5Ñ$5°t×7HÑ7H×7RÑ7R×SÕSØ ŸJœJ�DÜ! $¬×4Ò4àŸ™ T×%6Ñ%6Ö7à×)Ñ)×3Ñ3°DÖ9ñ '÷ Tð ×)Ñ)Ô+à×.Ñ.Ó0ˆKØ×!Ñ!Ô#Ü�}‰}×-×-Ð-Ø×(Ñ(¨×):Ñ):Ô;ä×'×'Ð'Ø×'Ñ'Ô)ä�}‰}×-×-Ð-Ø×-Ñ-Ô/ä�}‰}×(×(Ð(´·±×1W×1WÐ1WØ×5Ñ5Ô7ä�}‰}×5×5Ð5Ø×4Ñ4Ô6ô ×'×'Ð'´×0B×0BÐ0BØ×!Ñ!×+Ñ+ØAôð × Ñ  Ô-÷] $ôb  Ó!ˆØ�‰�d—l‘lÔ#Ø×Ñ˜ÔØ�‰�d—k‘kÔ"ô
 �G‰G××ÐÜ—‘×#×#Ð#Ü—‘×&×&Ð&Ü—G‘G×0×0Ð0ä#Ó%ˆFð 	�‰�d×/Ñ/Ô0Ø×ÑÔØ�‰�d—k‘kÔ"à"×:Ñ:Ó<Ðà�]‰]Ð.×/Õ/Ø�M‰M˜$×+Ñ+Ô,÷ 0ð 	×#Ñ# FÔ+Ø�‰�d—k‘kÔ"Ø×"Ñ" 6Ô*à×Ñ˜&Ô!à×"Ñ" 6Ô*ð ×&Ñ&Ó(Ø×$Ñ$×8Ñ8Ó:ð
ð 	
÷I T×SÐSú÷ $×#Ð#ú÷N 0×/úsw   ÁA&U*Â+A U*Ã,U*ÄAU*ÅA"U	Æ6AU*È0U*È?/U*É//U*ÊU*Ê?/U*Ë/%U*ÌU*Ì+,U*Ò)U>ÕU'Õ!	U*Õ*U;	Õ>V	c                ó´  a€ V P                   P                  R4       / p\        P                  P                  '       dg   \
        P                  P                  '       dG   \        \
        P                  P                  4       UUu/ uF  w  r#V P                  V4      VbK  	  pppV P                   P                  4       R,           V P                  P                  4       ,           o\        P                  \        P                  8X  dv   \         P"                  ! \%        4       RRR7      ;_uu_ 4       pVP'                  SP)                  R4      4       VP*                  pRRR4       \        P,                  ! RX4       \/        R	R
 V3R lR7        \1        SV4       R# u uppi   + '       g   i     LO; i  \2         d   p\5        RT 24      ThRp?ii ; i)z¢
Compose self.kernel_autotune_defs and self.kernel_autotune_calls into a single block of
code and execute it to trigger Triton kernel compilation and auto-tuning
zQ
            async_compile.wait(globals())
            del async_compile
        r  z.pyF)Údirr  Údeletezutf-8NzAuto-tuning code written to %sÚartifactc                 ó   € R RRR/# )rÐ   Ú&inductor_autotune_at_compile_time_codeÚencodingÚstringr`   r`   rd   rb   r  ÚFPythonWrapperCodegen.generate_and_run_autotune_block.<locals>.<lambda>j
  s   € ØÐ@Ø˜Hñ!rd   c                 ó   <€ S # ru   r`   )Útuning_codes   €rb   r  r$  n
  s   ø€ ™{rd   )Úmetadata_fnÚ
payload_fnz%Failed to run autotuning code block: )r  r?  r*   rî   rï   rH   rœ   Úautotuning_inputsrÚ  Úget_autotuning_input_namer  rø   r-   ÚlevelÚloggingÚDEBUGÚtempfileÚNamedTemporaryFiler   ÚwriteÚencoderÐ   r+   r   ÚexecÚ	ExceptionÚRuntimeError)rf  Úscoper  r}   ÚfÚ	file_pathÚer&  s   &      @rb   r  Ú4PythonWrapperCodegen.generate_and_run_autotune_blockF
  s�  ø€ ð
 	×!Ñ!×(Ñ(ðô	
ð ˆÜ�=‰=×1×1Ð1´a·g±g×6O×6OÐ6Oô (¬¯©×(AÑ(AÔBôáB‘F�Cð ×.Ñ.¨sÓ3°QÒ6ÙBð ñ ð
 ×%Ñ%×.Ñ.Ó0Øõà×(Ñ(×1Ñ1Ó3õ4ð 	ô
 × Ñ ¤G§M¡MÔ1ô ×,Ò,Ü“K¨°e÷õ àØ—‘˜×*Ñ*¨7Ó3Ô4ØŸF™F�	÷	ô
 ×!Ò!Ø0Øôô 	Øñô +õ	
ð	SÜ�˜eÖ$ùó?÷÷ ûô( ô 	SÜÐ!FÀqÀcÐJÓKÐQRÐRûð	Sús*   ÂFÄ1-F%ÆF8 Æ%F5	Æ8GÇGÇGc                	ó^   € ^RI Hp V! V 4      P                  V P                  4      V n        R# )rg   )ÚMemoryPlannerN)rŒ  r;  r�  r"  )rf  r;  s   & rb   Úmemory_planÚ PythonWrapperCodegen.memory_planv
  s    € Ý2á" 4Ó(×-Ñ-¨d¯j©jÓ9ˆŽ
rd   c                	ó`  € V P                  4       p\        P                  P                  V4      pV P                  '       dp   \        V P                  R,          \        4      '       dI   V P                  R,          P                  P                  V9  d   V P                  P                  4        K�  \        4       .p. p. p\        \        V P                  4      4       EF^  pV P                  V,          p\        V\        4      '       d(   VP                  VR,          4      V P                  V&   KT  \        V\        4      '       d�   VP                  \        4       4       \         P"                  '       d^   VP                  V P$                  4       \        P&                  ! VP                  4      ;_uu_ 4        \)        4       V n        RRR4       Kö  Kø  \        V\*        4      '       g   EK  VP                  VP                  4       4       \         P"                  '       g   EKI  VP                  4       V n        EKa  	  VP                  VP                  4       4       \        V4      ^ 8w  d   \-        R\        V4       24      h\/        R V 4       4      pR#   + '       g   i     EKÏ  ; i)rg   Nz+expected no remaining planning states, got c              3  ó8   "  € T F  qP                   x € K  	  R # 5iru   )rs  )rx   Úss   & rb   rz   Ú9PythonWrapperCodegen.memory_plan_reuse.<locals>.<genexpr>£
  s   é € ð +
Ù3G¨a×)Ö)Ó3Gùr~   r#  )r¢  rH   rœ   Ú_get_output_namesr"  r   r–  r˜   rÐ   r  rj  r¸  r  r�  r±  rÈ   r*   r'  r  Úset_graph_handlerrã  rê  r²   Úsum)	rf  rÐ  Ú	out_namesÚplanning_statesÚpast_planning_statesÚpeak_estimate_stackr¹  rô   Ú_total_allocated_buffer_sizes	   &        rb   Úmemory_plan_reuseÚ&PythonWrapperCodegen.memory_plan_reuse{
  sñ  € Ø×(Ñ(Ó*ˆÜ—G‘G×-Ñ-¨gÓ6ˆ	ð �J�JˆJÜ˜4Ÿ:™: b�>Ô+=×>Ò>à—
‘
˜2•×#Ñ#×(Ñ(°	Ô9ð �J‰J�N‰NÖô /Ó0Ð1ˆØ!ÐØ;=ÐÜ”s˜4Ÿ:™:“×'ˆAØ—:‘:˜a•=ˆDÜ˜$Ô 2×3Ò3Ø $§	¡	¨/¸"Õ*=Ó >�—
‘
˜1“Ü˜DÔ"3×4Ò4Ø×&Ñ&Ô':Ó'<Ô=Ü×,×,Ð,Ø'×.Ñ.¨t×/AÑ/AÔBÜ×,Ò,¨T¯Z©Z×8Õ8Ü-BÓ-D˜Ô*÷ 9Ñ8ñ -ô ˜DÔ"2×3Õ3Ø$×+Ñ+¨O×,?Ñ,?Ó,AÔBÜ×,×,Ó,Ø)<×)@Ñ)@Ó)B�D×&ñ (ð 	×#Ñ# O×$7Ñ$7Ó$9Ô:ÜˆÓ 1Ô$Ü Ø=¼cÀ/Ó>RÐ=SÐTóð ô (+ñ +
Ù3Gó+
ó (
Ò$÷ 9×8Ñ8ús   Æ:JÊJ-c               ó   € V ^8„  d   QhRR/# )r\   rî  rE  r`   )ra   s   "rb   rc   rM  §
  s   € ÷ %ñ %°$ñ %rd   c                	óÐ   € V'       d)   \         P                  '       d   V P                  4        R # \         P                  '       d   \	        4       V n        V P                  4        R # ru   )r*   rŒ  r<  r'  rã  r  rJ  rí  s   &&rb   r  Ú*PythonWrapperCodegen.run_wrapper_ir_passes§
  s?   € çœF×2×2Ð2Ø×ÑÖä×(×(Ð(Ü%:Ó%<�Ô"Ø×"Ñ"Ö$rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   Úsymra  Ú
bound_varsúOrderedSet[sympy.Symbol]r_   r„  r`   )ra   s   "rb   rc   rM  °
  s(   € ÷ $$ñ $$àð$$ð -ð$$ð 
ñ	$$rd   c                	ó´  € \        \        P                  R R4      p\        VRR4      pVf   R# R R lpVP                  P	                  4        EF  w  rgWq8X  dv   \        V\        P                  4      '       dV   Wb9  dP   V! V4      '       dB   V! V4      '       d4   V P                  P                  V RV 24       VP                  V4       K�  Wa8X  g   K‰  \        V\        P                  4      '       g   K«  Wr9  g   K³  V! V4      '       g   KÃ  V! V4      '       g   KÓ  V P                  P                  V RV 24       VP                  V4       EK  	  R# )r¥   NÚ	shape_envc               ó    € V ^8„  d   QhRRRR/# )r\   r@  ra  r_   rE  r`   )ra   s   "rb   rc   ÚIPythonWrapperCodegen.maybe_emit_replacement_aliases.<locals>.__annotate__¿
  s   € ÷ 	Sñ 	S ð 	S°ñ 	Srd   c                ó`   € \        V \        P                  \        P                  34      '       * # ru   )r'   r(   ÚUNBACKED_INTÚUNBACKED_FLOAT)r@  s   &rb   Úis_backed_symbolÚMPythonWrapperCodegen.maybe_emit_replacement_aliases.<locals>.is_backed_symbol¿
  s#   € Ü% a¬$×*;Ñ*;¼T×=PÑ=PÐ)QÓRÔRÐRrd   r2  )rw   rH   rœ   Úreplacementsr�   r   rh   ÚSymbolr  rö   r  )rf  rP  rQ  r¥   rT  rZ  r@  Útgts   &&&     rb   Úmaybe_emit_replacement_aliasesÚ3PythonWrapperCodegen.maybe_emit_replacement_aliases°
  s  € ô œ1Ÿ7™7 J°Ó5ˆÜ˜H k°4Ó8ˆ	ØÒÙõ	Sð "×.Ñ.×4Ñ4×6‰HˆCà”
Ü˜s¤E§L¡L×1Ò1ØÔ)Ù$ S×)Ò)Ù$ S×)Ò)à—‘×%Ñ%¨¨¨S°°Ð&6Ô7Ø—‘˜sÖ#à–
Ü˜s¤E§L¡L×1Ô1ØÖ)Ù$ S×)Ô)Ù$ S×)Ô)à—‘×%Ñ%¨¨¨S°°Ð&6Ô7Ø—‘˜s×#ó% 7rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÐ   r‹   r†   zir.TensorBox | sympy.ExprrQ  rR  r`   )ra   s   "rb   rc   rM  Ö
  s-   € ÷ #Kñ #Kàð#Kð )ð#Kð -ñ	#Krd   c                	óŠ  € \        V\        P                  4      '       Ed   Tp\        P                  P
                  P                  V4      p\        V\        P                  4      '       dI   W#9  dC   V P                  P                  V R V 24       VP                  V4       V P                  W#4       \        V\        P                  4      '       d:   WS9  d2   V P                  P                  V R V 24       VP                  V4       R# \        V\        P                  4      '       g   R# R# \        V\        P                  \        P                  \        P                  \        P                   34      '       d   R# \"        P$                  P&                  P(                  '       d   R# \+        R\-        V4       24      h)r2  NzUnknown value type: )r   rh   r   rH   rœ   r¥   r¦   r]  r  rö   r  r_  r,   rÆ   r¶  r·  r¸  r^  rc  r*   rÏ  r²   r…   )rf  rÐ   r†   rQ  Údeferred_symbol_assignmentsÚ	raw_values   &&&&& rb   Úcodegen_input_symbol_assignmentÚ4PythonWrapperCodegen.codegen_input_symbol_assignmentÖ
  sI  € ô �eœUŸZ™Z×(Ó(ØˆIÜ—G‘G×$Ñ$×-Ñ-¨iÓ8ˆEÜ˜%¤§¡×.Ò.°5Ô3JØ—‘×%Ñ%¨¨¨s°4°&Ð&9Ô:Ø—‘˜uÔ%Ø×3Ñ3°EÔFÜ˜)¤U§\¡\×2Ò2ØÔ.Ø—K‘K×)Ñ)¨Y¨K°s¸4¸&Ð*AÔBØ—N‘N 9Ô-ÙÜ˜e¤U§\¡\×2Ò2Ùñ 3äØä—‘Ü×"Ñ"Ü×!Ñ!Ü×$Ñ$ð	÷
ò 
ñ ä�‰×%Ñ%×5×5Ð5Ùä$Ð';¼DÀ»K¸=Ð%IÓJÐJrd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  û
  s   € ÷ Kñ KÐQUñ Krd   c                	óp   € V'       d.   . pR pV F  pV! V4      ;'       g    TpK  	  V'       g   R# TpK5  R# )FNr`   )rf  rc  Ú next_deferred_symbol_assignmentsÚprogressÚ
assignments   &&   rb   Ú"_retry_deferred_symbol_assignmentsÚ7PythonWrapperCodegen._retry_deferred_symbol_assignmentsû
  s>   € ß)Ø/1Ð,ØˆHÛ9�
Ù%Ð&FÓG×SÐSÈ8’ñ :çÙØ*JÒ'ñ *rd   c               ó0   € V ^8„  d   QhRRRRRRRRR	R
RR/# )r\   rP  ra  Ú
input_namer‹   ÚkindúLiteral['size', 'stride']r¨  r�   rQ  rR  r_   r„  r`   )ra   s   "rb   rc   rM    sF   € ÷ =ñ =àð=ð ð=ð (ð	=ð
 ð=ð -ð=ð 
ñ=rd   c                	óú   € W9   d   R # VR8X  d&   V P                   P                  V RV RV R24       M$V P                   P                  V RV RV R24       VP                  V4       V P                  W4       R # )NrÃ   r2  z.size()[rŠ  z
.stride()[)r  rö   r  r_  )rf  rP  ro  rp  r¨  rQ  s   &&&&&&rb   Úbind_input_symbolÚ&PythonWrapperCodegen.bind_input_symbol  sx   € ð ÔÙØ�6Œ>Ø�K‰K×!Ñ! S E¨¨Z¨L¸ÀÀÀQÐ"GÕHà�K‰K×!Ñ! S E¨¨Z¨L¸
À3À%ÀqÐ"IÔJØ�‰�sÔØ×+Ñ+¨CÖ<rd   c           	     ó€  a aaaa€ \         \        P                  ,          ! 4       oS P                  4       pVP	                  4        UUu. uF*  w  r#\        V\        P                  4      '       g   K'  W#3NK,  	  uppVP	                  4        UUu. uF*  w  r#\        V\        P                  4      '       d   K'  W#3NK,  	  upp,           p. pV F  w  rgS P                  WgSV4       K  	  S P                  V4       \        P                  V 3R l4       o\        P                  V 3R l4       oR
R VVV VV3R lllo\        P                  P                  P	                  4        Fx  w  pw  ršp\        P                  P                  P                  V4      p\        V\        P                  4      '       g   KT  VS9   d   K]  W‘9  d   Ke  S P                  W‰W«S4       Kz  	  V EF  w  r—\        V\         P"                  4      '       g   K(  V	\        P                  P$                  9   p\&        P(                  ;'       d'    T;'       d    \+        VP-                  4       4      ^ 8g  p\/        VP-                  4       4       F`  w  r¾V'       d   S! VV	RVV4       \        P                  P                  P                  V4      pV'       g   KL  Wþ8w  g   KT  S! VV	RVV4       Kb  	  \/        VP1                  4       4       Fb  w  ppV'       d   S! VV	RVV4       \        P                  P                  P                  V4      pV'       g   KM  VV8w  g   KV  S! VV	RVV4       Kd  	  EK‚  	  S P                  V4       R R	 lpV F0  w  r—\        V\         P"                  4      '       g   K'  V! W—S4       K2  	  R# u uppi u uppi )z8Assign input symbolic shapes to wrapper-local variables.c                óP   <€ SP                   P                  V  R V  R24       V  R2# )z_size = z.size()Ú_size©r  rö   ©rÐ   rf  s   &€rb   ÚsizeofÚ3PythonWrapperCodegen.codegen_inputs.<locals>.sizeof+  s.   ø€ à�K‰K×!Ñ! T F¨(°4°&¸Ð"@ÔAØ�V˜5�>Ð!rd   c                óP   <€ SP                   P                  V  R V  R24       V  R2# )z
_stride = z	.stride()Ú_striderx  ry  s   &€rb   ÚstrideofÚ5PythonWrapperCodegen.codegen_inputs.<locals>.strideof0  s/   ø€ à�K‰K×!Ñ! T F¨*°T°F¸)Ð"DÔEØ�V˜7Ð#Ð#rd   Nc          
     ó,   € V ^8„  d   QhRRRRRRRRR	R
/# )r\   r]   zsympy.Symbol | sympy.Exprro  r‹   rp  rq  r¨  r�   r_   rE  r`   )ra   s   "rb   rc   Ú9PythonWrapperCodegen.codegen_inputs.<locals>.__annotate__5  s:   € ÷ 6	ñ 6	Ø+ð6	àð6	ð ,ð6	ð ð	6	ð ñ6	rd   c                óî  <a aaa€ SR 8X  d   SMSp\        S \        P                  4      '       dY   S S9   d   R# SP                  P	                  S  RV! S4       RS R24       SP                  S 4       SP                  S S4       R# \        S \        P                  4      '       g   R# S P                   Uu. uF  qfS9  g   K  VNK  	  pp\        V4      ^8w  d2   \        V4      ^8”  d    Ve   VVV VV3R lpVP                  V4       R# ^ RIHp	 VP                  4       p
V! S4      p\        P                  ! V RS 2RR	7      pV	! \        P                  ! S V4      V
4      pVf   R# SP                  P	                  V RV RS R24       \        V^,          4      pSP                  P	                  V
 R\!        V4       24       SP                  V
4       SP                  V
S4       R# u upi )
rÃ   Fr2  r‰  rŠ  Tc                ó   <€ S! SSSSV 4      # ru   r`   )rc  Úbind_input_expr_symbolsr¨  r]   ro  rp  s   &€€€€€rb   ÚretryÚSPythonWrapperCodegen.codegen_inputs.<locals>.bind_input_expr_symbols.<locals>.retryQ  s   ø€ Ù6Ø Ø&Ø ØØ7ó ð rd   )Ú	try_solver  )Úinteger)r   rh   r]  r  rö   r  r_  r   Úfree_symbolsr  rÈ   Útorch.utils._sympy.solver‡  r  ÚEqrn   Úpexpr)r]   ro  rp  r¨  rc  Úname_fnrP  Úundefined_symbolsr…  r‡  Úfree_symbolÚ	base_nameÚ
dim_symbolÚsolutionÚsolution_exprr„  rQ  rf  rz  r~  s   ffff&          €€€€€rb   r„  ÚDPythonWrapperCodegen.codegen_inputs.<locals>.bind_input_expr_symbols5  sÊ  ü€ ð !%¨¤‘f°HˆGÜ˜$¤§¡×-Ò-Ø˜:Ô%Ù Ø—‘×%Ñ%¨¨¨c±'¸*Ó2EÐ1FÀaÈÀuÈAÐ&NÔOØ—‘˜tÔ$Ø×3Ñ3°D¸*ÔEÙä˜d¤E§J¡J×/Ò/Ùð  $×0Ò0ó!Ù0˜¸zÑ4I—�Ñ0ð ð !ô Ð$Ó%¨Ô*äÐ)Ó*¨QÔ.Ø3Ò?÷ñ ð 0×6Ñ6°uÔ=Ùå:à+×/Ñ/Ó1ˆKÙ 
Ó+ˆIÜŸš¨¨°1°S°EÐ&:ÀDÔIˆJÙ ¤§¢¨$°
Ó!;¸[ÓIˆHØÒÙà�K‰K×!Ñ! Z L°°I°;¸aÀ¸uÀAÐ"FÔGÜDÀXÈaÅ[ÓQˆMØ�K‰K×!Ñ! [ M°´U¸=Ó5IÐ4JÐ"KÔLØ�N‰N˜;Ô'Ø×/Ñ/°¸ZÔHÙùòG!s   Â9G2ÃG2rÃ   rÄ   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   ro  r‹   r†   zir.TensorBoxrQ  rR  r`   )ra   s   "rb   rc   r�  ¤  s&   € ÷ 	ñ 	Øð	àð	ð 1ñ	rd   c                óÌ  € V \         P                  P                  9   p\        P                  ;'       d'    T;'       d    \        VP                  4       4      ^ 8g  p\        P                  ! VP                  4       VP                  4       .4       FË  p\        V\        4      '       d!   \        V\        P                  4      '       d   K;  \         P                  P                  P                  V4      p\        V\        4      '       g   K|  VP                    Uu. uF  pWb9  g   K  V'       g   K  VNK  	  pp\#        V4      ^ 8”  g   K»  \%        RV RV R24      h	  R# u upi )r   zFor z, expected z to have been codegen-ed.N)rH   rœ   r¹  r*   rä  rD   r7  r   Úfrom_iterabler8  r   r   rh   r]  r¥   r¦   r‰  r  r²   )ro  r†   rQ  Úis_named_graph_inputÚverify_assert_symbolsr]   rP  rŽ  s   &&&     rb   Ú_verify_input_symbol_assignmentÚLPythonWrapperCodegen.codegen_inputs.<locals>._verify_input_symbol_assignment¤  s3  € ð
 $.´·±×1JÑ1JÑ#JÐ Ü$*×$7Ñ$7÷ %ð %Ø$×MÐM¬°u·~±~Ó7GÓ)HÈAÑ)Mð "ô ×+Ò+¨U¯^©^Ó-=¸u×?OÑ?OÓ?QÐ,RÖS�Ü! $¬×-Ò-´¸DÄ%Ç,Á,×1OÒ1OÙÜ—w‘w×'Ñ'×0Ñ0°Ó6�Ü! $¬×-Ò-Ùð  $×0Ò0ó%á0˜ØÑ,ô ä1F÷ �CÙ0ð "ð %ô
 Ð(Ó)¨AÖ-Ü(Ø˜t˜f KÐ0AÐ/BÐB[Ð\óð ó Tùò%s   ÄE!Ä'E!Ä1E!ru   )r"   rh   r]  r­  r�   r   re  rl  r*  r9  rH   rœ   Úsymbolic_input_sourcesr¥   r¦   rs  r,   rÆ   r¹  r*   rä  rD   r7  rÚ  r8  )rf  rê  r‡   r}   r�   rc  rÐ   r†   rP  ro  rp  r¨  r˜  Úbind_assert_symbolsÚraw_sizerÃ   Ú
raw_striderÄ   rš  r„  rQ  rz  r~  s   f                  @@@@rb   rÞ  Ú#PythonWrapperCodegen.codegen_inputs  s&  ü€ ä¤§¡Ö-Ó/ˆ
ð ×,Ñ,Ó.ˆà+×1Ñ1Ô3ô
Ù3‘t�q´zÀ!ÄUÇ\Á\×7RŒFˆQ‹FÑ3ò
à ,× 2Ñ 2Ô 4ÔXÑ 4™˜¼JÀqÌ%Ï,É,×<WŒVˆa‹VÑ 4ÒXõYˆð ')Ð#Û!‰KˆDØ×0Ñ0Ø˜ZÐ)Döñ "ð 	×/Ñ/Ð0KÔLä	�‰ô	"ó 
ð	"ô 
�‰ô	$ó 
ð	$÷6	ö 6	ôp -.¯G©G×,JÑ,J×,PÑ,PÖ,RÑ(ˆCÑ(�* CÜ—'‘'×"Ñ"×+Ñ+¨CÓ0ˆCÜ˜c¤5§<¡<×0Ò0ÙØ�jÔ ÙØÔ-ÙØ×"Ñ" 3°D¸zÖJñ -Sô "(ÑˆJÜ˜e¤R§\¡\×2Ò2ÙØ#-´·±×1JÑ1JÑ#JÐ Ü"(×"5Ñ"5÷ #ð #Ø$×MÐM¬°u·~±~Ó7GÓ)HÈAÑ)Mð  ô "+¨5¯>©>Ó+;Ö!<‘�ß&Ù+Ø Ø"ØØØ3ôô —w‘w×'Ñ'×0Ñ0°Ó:�ß&Ñ&¨4Ö+;Ù+ØØ"ØØØ3öñ "=ô$ $-¨U×-=Ñ-=Ó-?Ö#@‘��Zß&Ù+Ø"Ø"Ø ØØ3ôô Ÿ™×)Ñ)×2Ñ2°:Ó>�ß&Ñ&¨6°ZÖ+?Ù+ØØ"Ø ØØ3öô $Añ3 "(ðV 	×/Ñ/Ð0KÔLõ	ó: "(ÑˆJÜ˜e¤R§\¡\×2Ò2ÙÙ+¨J¸zÖJó "(ùóA
ùãXs   Á	$N4Á2N4Â$N:Â8N:c               ó   € V ^8„  d   QhRR/# )r\   rP  ra  r`   )ra   s   "rb   rc   rM  Æ  s   € ÷ Dñ D¨ñ Drd   c                	ó¦  € \        V\        P                  4      '       d±   \        V\        P
                  4      '       d�   WP                  9   d   R # V P                  P                  V4       \        P                  P                  P                  V,          p\        W4      pV P                  \        W\        P                  4      4       R # R # R # ru   )r   rh   r]  r'   r(   ÚPRECOMPUTED_SIZEr¸  r  rH   rœ   r¥   Úinv_precomputed_replacementsr`  rö   rÂ  )rf  rP  r]   rÃ  s   &&  rb   Úensure_size_computedÚ)PythonWrapperCodegen.ensure_size_computedÆ  s�   € Ü�cœ5Ÿ<™<×(Ò(¬^¸CÄ×AVÑAV×-WÒ-WØ×)Ñ)Ô)ÙØ×Ñ×#Ñ# CÔ(Ü—7‘7×#Ñ#×@Ñ@ÀÕEˆDÜ! #Ó,ˆCØ�N‰NÔ.¨t¼!¿'¹'ÓBÖCñ .XÑ(rd   c                	ó   € R # ru   r`   re  s   &rb   r  Ú$PythonWrapperCodegen.finalize_prefixÏ  r  rd   r¦   c               ó$   € V ^8„  d   QhRRRRRR/# ©r\   r  r   r¦   rE  r_   r‹   r`   )ra   s   "rb   rc   rM  Ò  s&   € ÷ Wñ W Tð W¸ð WÈñ Wrd   c               	ó   € \        R 4      h)z8codegen_cpp_sizevar is only implemented for cpp_wrapper!)r4  ©rf  r  r¦   s   &&$rb   Úcodegen_cpp_sizevarÚ(PythonWrapperCodegen.codegen_cpp_sizevarÒ  s   € ÜÐUÓVÐVrd   c               ó$   € V ^8„  d   QhRRRRRR/# rª  r`   )ra   s   "rb   rc   rM  Õ  s!   € ÷ +ñ +¨ð +¸4ð +È3ñ +rd   c               	ó   € \        WR 7      # ))r¦   )rŒ  r¬  s   &&$rb   Úcodegen_python_sizevarÚ+PythonWrapperCodegen.codegen_python_sizevarÕ  s   € Ü�QÔ*Ð*rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r  r   r_   r‹   r`   )ra   s   "rb   rc   rM  Ø  s   € ÷ .ñ . ð .¨#ñ .rd   c                	ó$   € V P                  V4      # ru   )r±  r£  s   &&rb   Úcodegen_sizevarÚ$PythonWrapperCodegen.codegen_sizevarØ  s   € Ø×*Ñ*¨1Ó-Ð-rd   c               ó(   € V ^8„  d   QhRRRRRRRR/# )r\   rŽ  r‹   rÐ   r  r_   r`   )ra   s   "rb   rc   rM  Û  s(   € ÷ &ñ &¨Sð &¸ð &ÀCð &ÈCñ &rd   c                	ó   € V R V R2# )r‰  rŠ  r`   )rf  rŽ  rÐ   r  s   &&&&rb   r�  Ú)PythonWrapperCodegen.codegen_tuple_accessÛ  s   € Ø�˜1˜U˜G 1Ð%Ð%rd   c               ó    € V ^8„  d   QhRRRR/# ©r\   r>  zSequence[Expr]r_   r‹   r`   )ra   s   "rb   rc   rM  Þ  s   € ÷ 'ñ '°ð 'À3ñ 'rd   c                	óº   € . \        V P                  V4      Op\        V4      ^ 8X  d   R# \        V4      ^8X  d   RV^ ,           R2# RRP                  V4       R2# )r   ú()rS  r_  rT  r3  )rÊ  r±  r  r  )rf  r>  Úpartss   && rb   rí   Ú/PythonWrapperCodegen.codegen_python_shape_tupleÞ  s^   € Ø:”#�d×1Ñ1°5Ó9Ð:ˆÜˆu‹:˜Œ?ÙÜˆu‹:˜Œ?Ø�u˜Q•x�j Ð$Ð$Ø�4—9‘9˜UÓ#Ð$ AÐ&Ð&rd   c               ó    € V ^8„  d   QhRRRR/# r»  r`   )ra   s   "rb   rc   rM  æ  s   € ÷ 6ñ 6¨ð 6¸Cñ 6rd   c                	ó$   € V P                  V4      # ru   )rí   )rf  r>  s   &&rb   r;  Ú(PythonWrapperCodegen.codegen_shape_tupleæ  s   € Ø×.Ñ.¨uÓ5Ð5rd   c               ó   € V ^8„  d   QhRR/# )r\   r_   ztuple[str, list[str]]r`   )ra   s   "rb   rc   rM  é  s   € ÷ ñ à	ñrd   c                	ó°   € R P                  RP                  V\        V4      \        V4      V P	                  V4      V P	                  V4      .4      4      . 3# )zalloc_from_pool({})rT  )ra   r  rŒ  r‹   rí   )rf  rÐ   rÅ   r   r>  rÄ   s   &&&&&&rb   Úcodegen_alloc_from_poolÚ,PythonWrapperCodegen.codegen_alloc_from_poolé  s_   € ð %×+Ñ+Ø�I‰IàÜ˜&“MÜ˜“JØ×3Ñ3°EÓ:Ø×3Ñ3°FÓ;ðóó

ð ð
ð 
	rd   c               ó    € V ^8„  d   QhRRRR/# )r\   rö   r÷  r_   r‹   r`   )ra   s   "rb   rc   rM  ø  s$   € ÷ <Pñ <Pð 'ð<Pð 
ñ<Prd   c                	ó¼  a € \        V4      w  rxršpV 3R  lpVP                  4       pT;'       d    WI8H  pV'       d   W'8H  ;'       d    W88H  pT
pMdW!P                  P                  8H  ;'       d9    W1P                  P                  8H  ;'       d    WAP                  P
                  8H  pVP                  pV'       d   Ve   VV8w  d
   RV RV R2# V # V! WÒW4VV4      # )c           	     ó²   <€ S
P                  V4      pS
P                  V4      pS
P                  V4      pR V  RV RV RV R2	p	Ve   WE8w  d
   RV	 RV R2# V	# )zreinterpret_tensor(rT  r3  úaten.view.dtype()rí   rµ  )rÐ   Útgt_sizeÚ
tgt_strideÚ
tgt_offsetÚ
cast_dtypeÚ
base_dtyper@  ÚstÚoffr]   rf  s   &&&&&&    €rb   Úapply_reinterpretÚHPythonWrapperCodegen.codegen_reinterpret_view.<locals>.apply_reinterpret  sz   ø€ ð ×/Ñ/°Ó9ˆAØ×0Ñ0°Ó<ˆBØ×&Ñ& zÓ2ˆCØ(¨¨¨b°°°2°b°T¸¸C¸5ÀÐBˆDØÒ%¨*Ô*BØ)¨$¨¨r°*°¸QÐ?Ð?ØˆKrd   rÊ  rT  r3  )rÎ   rž   rµ   rÃ   rÄ   rÅ   r   )rf  rÉ   rÃ   rÄ   rÅ   rö   r   Úd_sizeÚd_strideÚd_offsetÚd_dtypeÚcollapsiblerÒ  rÐ   Ú	collapsedÚsame_layoutrÏ  s   f&&&&&&          rb   Úcodegen_reinterpret_viewÚ-PythonWrapperCodegen.codegen_reinterpret_viewø  sè   ø€ ô8 ,¨DÓ1ñ 	9ˆ˜(¨[õ		ð �}‰}‹ˆØ×6Ð6 FÑ$6ˆ	ßØ™.×?Ð?¨VÑ-?ˆKØ ‰Jð Ÿ™×(Ñ(Ñ(÷ 1ð 1ØŸk™k×0Ñ0Ñ0÷1ð 1àŸk™k×0Ñ0Ñ0ð ð
 Ÿ™ˆJçØÒ  U¨jÔ%8Ø)¨$¨¨r°%°¸Ð:Ð:Ø�VÐá  ¨V¸UÀJÓOÐOrd   c               ó   € V ^8„  d   QhRR/# )r\   Únon_blockingz
bool | strr`   )ra   s   "rb   rc   rM  6  s   € ÷ >ñ >¸*ñ >rd   c                	ó:   € V P                  V R V RV R24       R# )z.copy_(rT  r3  Nrþ  )rf  r@  ÚdstrÞ  s   &&&&rb   Úcodegen_device_copyÚ(PythonWrapperCodegen.codegen_device_copy6  s!   € Ø�‰˜#˜˜g c U¨"¨\¨N¸!Ð<Ö=rd   c               ó    € V ^8„  d   QhRRRR/# )r\   rŒ   r‹   r_   r„  r`   )ra   s   "rb   rc   rM  9  s   € ÷ 5ñ 5¨ð 5°ñ 5rd   c                	óˆ   € R V 2pV P                  V R24       V P                  V R24       V P                  V R24       R# )Ú_d2h_event_z = torch.Event()z	.record()z.synchronize()Nrþ  )rf  rŒ   Ú	event_vars   && rb   Úsync_d2h_copyÚ"PythonWrapperCodegen.sync_d2h_copy9  sH   € Ø! + Ð/ˆ	Ø�‰˜)˜Ð$4Ð5Ô6Ø�‰˜)˜ IÐ.Ô/Ø�‰˜)˜ NÐ3Ö4rd   c               ó   € V ^8„  d   QhRR/# )r\   r˜   zir.MultiOutputr`   )ra   s   "rb   rc   rM  ?  s   € ÷ Sñ S¨ñ Srd   c                	ó’   € VP                  4       pVP                  ^ 4      pV P                  \        WW1P                  4      4       R# rl  )rž   ro  rö   r€  r„  )rf  r˜   r�  r‚  s   &&  rb   Úcodegen_multi_outputÚ)PythonWrapperCodegen.codegen_multi_output?  s2   € Ø—m‘m“oˆØ—?‘? 1Ó%ˆØ�‰” t¸(ÇLÁLÓQÖRrd   c           
     	óx  € VP                    R VP                   RVP                    RVP                    2pV'       d   RVP                   RV R2pV P                  VP                   RVP                   R VP
                   RV R24       V P                  P                  \        VP                  4      4       R	# )
ú + ú if z
 < 0 else zmax(0, min(rT  z))r2  z * (r3  N)	r  rÃ   rö   Úunbacked_offset_symbolÚbase_offsetÚbase_dim_strider)  r  r‹   )rf  r˜   ÚclampÚ	index_strs   &&& rb   Úcodegen_dynamic_select_indexÚ1PythonWrapperCodegen.codegen_dynamic_select_indexD  s«   € Ø—z‘z�l # d§i¡i [°°T·Z±Z°LÀ
È4Ï:É:È,ÐWˆ	ßØ% d§i¡i [°°9°+¸RÐ@ˆIØ�‰Ø×*Ñ*Ð+¨3¨t×/?Ñ/?Ð.@ÀÀD×DXÑDXÐCYÐY]Ð^gÐ]hÐhiÐjô	
ð 	×"Ñ"×&Ñ&¤s¨4×+FÑ+FÓ'GÖHrd   c                	óš  a a€ VV 3R  lpV 3R lpSP                   pV! SP                  4      pV! SP                  4      pS P                  V RV 24       S P                  V RV 24       V! V R2V R2SP                  4      pS P                  V RV R24       S P
                  P                  \        SP                   4      4       R# )	c           	     ó  <€ SP                  \        ^ \        V SP                  4      4      4      pSP                  \        ^ \        V SP                  ,           SP                  4      4      4      pSP                  V 4      pV RV RV 2# )r   rï  z >= 0 else )rµ  r$   r%   rÃ   )r  ÚposÚnegÚx_condr˜   rf  s   &   €€rb   Úclamp_indexÚDPythonWrapperCodegen.codegen_dynamic_slice_size.<locals>.clamp_indexO  sv   ø€ Ø×&Ñ&¤s¨1¬c°!°T·Y±YÓ.?Ó'@ÓAˆCØ×&Ñ&¤s¨1¬c°!°d·i±iµ-ÀÇÁÓ.KÓ'LÓMˆCØ×)Ñ)¨!Ó,ˆFØ�U˜$˜v˜h k°#°Ð7Ð7rd   c                ó^   <€ V^8X  d   V RV  2# SP                  V4      pRV RV  RV RV 2# )rg   z - rS  rî  z	 - 1) // )rµ  )Ú	start_varÚend_varÚstepÚstep_rf  s   &&& €rb   Úcodegen_with_stepÚJPythonWrapperCodegen.codegen_dynamic_slice_size.<locals>.codegen_with_stepU  sJ   ø€ Ø�qŒyØ!˜ # i [Ð1Ð1Ø×(Ñ(¨Ó.ˆEØ�w�i˜s 9 +¨S°°°yÀÀÐHÐHrd   z	_start = z_end = Ú_startÚ_endz
 = max(0, r3  N)Úunbacked_size_symbolÚstartÚendrö   r  r)  r  r‹   )rf  r˜   rü  r  rP  r  r	  Ú	with_steps   ff      rb   Úcodegen_dynamic_slice_sizeÚ/PythonWrapperCodegen.codegen_dynamic_slice_sizeN  s»   ù€ ö	8õ	Ið ×'Ñ'ˆÙ˜DŸJ™JÓ'ˆÙ˜$Ÿ(™(Ó#ˆØ�‰˜#˜˜i¨ wÐ/Ô0Ø�‰˜#˜˜g c UÐ+Ô,Ù%¨¨¨V n¸¸¸T°lÀDÇIÁIÓNˆ	Ø�‰˜#˜˜j¨¨°1Ð5Ô6Ø×"Ñ"×&Ñ&¤s¨4×+DÑ+DÓ'EÖFrd   c                	ó:   € V P                  \        W4      4       R # ru   )rö   rß  rü  s   &&rb   Úcodegen_dynamic_scalarÚ+PythonWrapperCodegen.codegen_dynamic_scalare  s   € Ø�‰Ô(¨Ó4Ö5rd   c                	ó  € R  VP                    4       w  p\        VP                  4      ^ 8X  d$   V P                  VP                   RV R24       EM�\        VP                  4      ^8X  dK   \        VP                  ^ ,          \        4      '       d$   V P                  VP                   RV R24       EM+\        VP                  4      ^8X  dú   \        VP                  ^ ,          \        4      '       dÓ   V P                  VP                   RV R24       V P                  RVP                   RVP                  ^ ,          P                   RVP                   R	VP                  ^ ,          P                   R
2	4       V P                  VP                   RVP                   RVP                  ^ ,          P                   24       M\        RVP                   24      hV P                  VP                  4        R24       R# )c              3  ó@   "  € T F  qP                  4       x € K  	  R # 5iru   r¢  r£  s   & rb   rz   Ú?PythonWrapperCodegen._codegen_dynamic_scalar.<locals>.<genexpr>i  s   é € Ð>±+¨Q×&Ñ&×(Ð(³+ùr¦  r2  ú.item()z = 1 if z.item() else 0z_undivided = zassert z_undivided % z
 == 0, f'{z_undivided} not divisible by Ú'z_undivided // úunrecognized keypath ú = NoneN)r�   r  Úkeypathrö   rP  r   r   r   rü  r²   rž   )rf  r˜   rÉ   s   && rb   rã  Ú,PythonWrapperCodegen._codegen_dynamic_scalarh  sŒ  € Ù>°$·+²+Ó>‰ˆÜˆt�|‰|Ó Ô!Ø�N‰N˜dŸh™h˜Z s¨4¨&°Ð8Ö9Ü�—‘Ó !Ô#¬
°4·<±<Àµ?ÄM×(RÒ(RØ�N‰N˜dŸh™h˜Z x°¨v°^ÐDÖEÜ�—‘Ó !Ô#¬
°4·<±<Àµ?ÄK×(PÒ(PØ�N‰N˜dŸh™h˜Z }°T°F¸'ÐBÔCØ�N‰NØ˜$Ÿ(™(˜ =°·±¸aµ×1HÑ1HÐ0Ið JØ—x‘x�jÐ >¸t¿|¹|ÈA½×?VÑ?VÐ>WÐWXðZôð �N‰NØ—8‘8�*˜C §¡˜z¨¸¿¹ÀQ½×8OÑ8OÐ7PÐQõô !Ð#8¸¿¹¸Ð!GÓHÐHð 	�‰˜$Ÿ-™-›/Ð*¨'Ð2Ö3rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r�   rw  r_   r‹   r`   )ra   s   "rb   rc   rM  }  s   € ÷ ñ ¨Lð ¸Sñ rd   c                ó¶   € V P                   pVf#   \        4       '       d   \        4       MRpW n         VRJ d   \        V4      # \        W4       VP                  # )aæ  Render a device for the benchmark harness rank-agnostically: the bare device type
(no rank-specific index), so it is byte-identical across ranks. cpu/meta have no index
and pass through unchanged; a non-current accelerator violates the single-device
invariant and is refused. Under compile-on-one-rank the current accelerator is read
once per compilation and cached on the wrapper; once cached, the per-input/buffer
harness loops repeat neither the enabled check nor the cudaGetDevice.
F)r   r   r‹   r   r…   )rf  r�   rÌ   s   && rb   Ú_coor_device_type_strÚ*PythonWrapperCodegen._coor_device_type_str}  sN   € ð ×,Ñ,ˆØŠ;Ü1>·²Ô+Ô-ÀeˆCØ-0Ô*Ø�%‹<Ü�v“;ÐÜ'¨Ô4Ø�{‰{Ðrd   c                ó’  a a€ VV 3R lpV3R lpV3R lp/ p/ p\         P                  P                  pVEe   \        \         P                  P                  P                  4       V4       FÖ  w  r‰\        V	\        P                  4      '       g   K'  V	P                  4       p
V
P                  4       p\        V
4      pVP                  V4      pVf5   \        R\        V4       2VV	P                  V	P                   / R7      pWÕV&   \#        V	P$                  4      \#        V	P'                  4       4      3VP(                  V&   KØ  	  SP+                  . RO4       SP-                  4       ;_uu_ 4        SP/                  RRR	7       \         P                  P0                  P3                  4        FT  w  rŽSP5                  R
V 24       V! WŽP7                  4       VP'                  4       VP                  VP                   4       KV  	  \        \         P                  P8                  4      ^ 8”  d_   SP5                  R4       \         P                  P8                  P3                  4        F!  w  r�SP5                  R
V 24       V! W�4       K#  	  VP;                  4        FÛ  p\        VP(                  4      ^8  d   K  VP                  \        P<                  P?                  VP                   4      ,          pSP5                  VP@                   RV RS PC                  VP                  4       RVP                    R24       VP(                  P3                  4        F  w  pw  ppVP@                  VV3Wh&   K  	  KÝ  	  \         P                  P                  P3                  4        EFÅ  w  rŽ\        V\D        PF                  4      '       dK   \        \         P                  PH                  PJ                  P                  VR4      \L        4      '       d   Kp  \        V\N        PP                  4      '       d   SP5                  V R24       K¦  \        V\D        PR                  4      '       d5   V! V\         P                  PH                  PU                  V^*R7      4       Kú  \        V\D        PV                  4      '       d   V! VR4       EK&  \        V\N        PX                  4      '       d_   \[        4       '       d)   \         P                  P\                  P_                  4       MVP                  P`                  pV! VRV R24       EK¤  \        V\N        Pb                  4      '       d   SP5                  V R24       EKÛ  W†9   dK   Wh,          w  pppSP5                  V RV RS Pe                  V4       RS Pe                  V4       R24       EK+  \         P                  PH                  Pg                  VPi                  4       ^*R7      p\         P                  PH                  Pg                  VPk                  4       ^*R7      pV! VVVVPm                  4       VPo                  4       4       EKÈ  	  SP5                  RRPq                  \         P                  P                  P                  4       4       R24       RRR4       SP+                  . RO4       SP-                  4       ;_uu_ 4        SP/                  R\         P                  Pr                   R2RR	7       RRR4       R#   + '       g   i     Lu; i  + '       g   i     R# ; i)z2Write out codegen for benchmarking the output codec                ó¢   <€ SP                  V  R SP                  V4       RSP                  V4       RSP                  V4       RV R2
4       R# )z = rand_strided(rT  ú
, device='ú	', dtype=r3  N)rö   rí   r  )rÐ   r>  rÄ   r�   r   rù   rf  s   &&&&&€€rb   Úadd_fake_inputÚFPythonWrapperCodegen.benchmark_compiled_module.<locals>.add_fake_input‘  sa   ø€ Ø×ÑØ�&Ð(Ø×2Ñ2°5Ó9Ð:¸"Ø×2Ñ2°6Ó:Ð;ð <Ø×5Ñ5°fÓ=Ð>¸iÈÀwÈaðQörd   c                ó4   <€ SP                  V  R V 24       R# rÂ  rþ  )rÐ   r§  rù   s   &&€rb   Úadd_expr_inputÚFPythonWrapperCodegen.benchmark_compiled_module.<locals>.add_expr_input™  s   ø€ Ø×Ñ ˜v S¨¨Ð.Ö/rd   c                óB  <€ Vf   SP                  V  R24       R # ^ R Ip SP                  V  RVP                  V4      : R24       R #   \        \        TP
                  3 d9   pSP                  R\        T4       RT  R\        T4       R24        R p?R # R p?ii ; i)Nr  z = pickle.loads(r3  z.raise TypeError("Failed to pickle opaque type z for variable r4  r6  )rö   ÚpickleÚdumpsÚ	TypeErrorrg  ÚPicklingErrorr…   r‹   )rÐ   r†   r'  r8  rù   s   &&  €rb   Úadd_torchbind_inputÚKPythonWrapperCodegen.benchmark_compiled_module.<locals>.add_torchbind_inputœ  s¨   ø€ ØŠ}Ø× Ñ  D 6¨Ð!1Ô2ÙãðØ× Ñ  D 6Ð)9¸&¿,¹,ÀuÓ:MÑ9PÐPQÐ!RÖSøÜœ~¨v×/CÑ/CÐDô Ø× Ñ ØDÄTÈ%Ã[ÀMÐQ_Ð`dÐ_eÐegÔhkÐlmÓhnÐgoÐoqÐr÷ò ûðús   ¡'A
 Á
BÁ&-BÂBNÚ_shared_storage_)rŒ   rŽ   r�   r   r�   zP
                from torch._dynamo.testing import rand_strided
                Tr&  zglobal zimport picklez = rand_strided((z,), (1,), device='r   r3  r  rù  Fztorch.cuda.default_generators[z].graphsafe_get_state()z = torch.as_strided(rT  zreturn [rŠ  zÄ
                from torch._inductor.utils import print_performance
                fn = lambda: call(list(args))
                return print_performance(fn, times=times, repeat=repeat, device='z')
                )r  r  zdef get_args():)r  r  z9def benchmark_compiled_module(args, times=10, repeat=10):):rH   rœ   Úexample_inputsr  rê  rë  r   r^  ÚTensorÚuntyped_storagerŽ   r!   r;  rŠ   r  r�   r   r‚   r>  rÄ   r�   Ú
writelinesr  r?  Ú	constantsr�   rö   rÃ   Útorchbind_constantsr«  Ú_utilsÚ_element_sizerŒ   r  rh   r]  r¥   Úbacked_var_to_valr&   r,   r¶  r   rû  rµ  r·  r   r	  r  r  r¸  rí   r¯  r7  r8  r=  r¤   r  r>  )rf  rù   r!  r$  r+  Ústorage_groupsÚaliased_input_specsr.  rÐ   ÚexÚstoragerŽ   Ústorage_keyÚgroupr†   Útorchbind_objÚnumelr>  rÄ   Úgen_idxrå  s   ff                   rb   Úbenchmark_compiled_moduleÚ.PythonWrapperCodegen.benchmark_compiled_moduleŽ  sþ  ù€ ö	õ	0õ	ð  GIˆØ?AÐÜŸ™×/Ñ/ˆàÓ%Ü¤§¡× 4Ñ 4× 9Ñ 9Ó ;¸^ÖL‘�Ü! "¤e§l¡l×3Ò3ÙØ×,Ñ,Ó.�Ø Ÿ™Ó)�Ü,¨WÓ5�Ø&×*Ñ*¨;Ó7�Ø’=Ü1Ø&6´s¸>Ó7JÐ6KÐ$LØ%Ø!Ÿy™yØ Ÿh™hØ!ô�Eð 38 ;Ñ/ä&*¨2¯8©8£n´d¸2¿9¹9»;Ó6GÐ%H�—‘˜TÓ"ñ# Mð( 	×ÑÒ5Ô6Ø�]‰]�_�_Ø�M‰Mðð ð	 ô ô  !Ÿw™w×0Ñ0×6Ñ6Ö8‘�ð × Ñ  7¨4¨&Ð!1Ô2ÙØŸ*™*›,¨¯©«¸¿¹ÀeÇkÁköñ	  9ô ”1—7‘7×.Ñ.Ó/°!Ô3Ø× Ñ  Ô1Ü+,¯7©7×+FÑ+F×+LÑ+LÖ+NÑ'�Dð ×$Ñ$ w¨t¨fÐ%5Ô6Ù'¨Ö<ñ	 ,Oð (×.Ñ.Ö0�Ü�u—|‘|Ó$ qÔ(ÙØŸ™¬¯©×(BÑ(BÀ5Ç;Á;Ó(OÕO�Ø× Ñ Ø×(Ñ(Ð)Ð):¸5¸'ÐASÐTX×TnÑTnÐot×o{Ño{ÓT|ÐS}ð  ~Gð  HM÷  HSñ  HSð  GTð  TUð  Vôð .3¯\©\×-?Ñ-?Ö-AÑ)�D™/˜5 &Ø16×1BÑ1BÀEÈ6Ð0RÐ'Ó-ó .Bñ 1ô  !Ÿw™w×3Ñ3×9Ñ9×;‘�Ü˜e¤U§\¡\×2Ò2´zÜ—G‘G×$Ñ$×6Ñ6×:Ñ:¸5À$ÓGÌ÷8ò 8ñ Ü˜e¤R×%7Ñ%7×8Ò8Ø×$Ñ$¨ v¨WÐ%5Ö6Ü ¤u§z¡z×2Ò2ñ
 #ØœaŸg™g×.Ñ.×@Ñ@ÀÐQSÐ@ÓTöô   ¤u§{¡{×3Ò3ñ # 4¨×/Ü ¤r×'8Ñ'8×9Ò9ô
 )Ÿ?š?ô Ÿ™×*Ñ*×BÑBÔDà"Ÿ\™\×/Ñ/ð ñ
 #ØØ8¸¸	ÐAXÐY÷ô   ¤r×';Ñ';×<Ò<Ø×$Ñ$¨ v¨WÐ%5×6ØÔ0Ø.AÕ.GÑ+�H˜e VØ×$Ñ$Ø˜&Ð 4°X°J¸bØ×:Ñ:¸5ÓAÐBÀ"Ø×:Ñ:¸6ÓBÐCÀ1ðF÷ô ŸG™G×,Ñ,×?Ñ?ØŸ™Ó(°2ð @ó �Eô ŸW™W×-Ñ-×@Ñ@Ø×(Ñ(Ó*°Rð Aó �Fñ #ØØØØ×(Ñ(Ó*ØŸ™Ó)÷ñi  <ðx ×Ñ˜x¨¯	©	´!·'±'×2FÑ2F×2KÑ2KÓ2MÓ(NÐ'OÈqÐQÔR÷ ðD 	×ÑÚQô	
ð �]‰]�_�_Ø�M‰MðRô ST×RYÑRY×ReÑReÐQfð gðð
 ð ô ÷ ‰_÷K �_ú÷J �_ˆ_ús   Å0U\"Û(0\5Ü"\2	Ü5]	c           	     ó(  € \         P                  '       g   R# V P                  V4       VP                  . RO4       VP	                  4       ;_uu_ 4        VP                  RRR\        4        R2.4       RRR4       R#   + '       g   i     R# ; i)z<
Append a benchmark harness to generated code for debugging
NzBfrom torch._inductor.wrapper_benchmark import compiled_module_mainzargs = get_args()zcompiled_module_main('zU', lambda times, repeat: benchmark_compiled_module(args, times=times, repeat=repeat)))r  r  zif __name__ == "__main__":)r*   Úbenchmark_harnessr@  r1  r  r=   ©rf  rù   s   &&rb   r  Ú*PythonWrapperCodegen.add_benchmark_harness5  sy   € ô ×'×'Ð'Ùà×&Ñ& vÔ.à×ÑÒ@ÔAØ�]‰]�_�_Ø×ÑàXØ'à0Ô1CÓ1EÐ0Fð Gmð mð	ô	÷ �_�_Š_ús   Á!B Â B	c          
     ó,   € V ^8„  d   QhRRRRRRRRRR/# ©	r\   rC  r‹   r‡  rˆ  rÖ   r‰  rE  rŠ  r`   )ra   s   "rb   rc   rM  K  s<   € ÷ 
ñ 
àð
ð ð
ð ð	
ð
 ð
ð #ñ
rd   c                	óF   € V P                  \        V VVVVVR 7      4       R# rŽ  )rö   r†  )rf  rC  r‡  rˆ  r‰  rŠ  s   &&&&&&rb   Údefine_kernelÚ"PythonWrapperCodegen.define_kernelK  s*   € ð 	�‰Ü ØØØØ!ØØ-ôö		
rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rC  r‹   r‡  rˆ  rÖ   r`   )ra   s   "rb   rc   rM  _  s$   € ÷ 	ñ 	Øð	Ø'*ð	Ø6@ñ	rd   c                	óØ   € \         P                  P                  '       d1   V'       d)   \        P                  ! R RV\        P
                  R7      pV'       d   V R2MRpRV V  RV 2pV# )z^// z# )Úflagsr  r  z

r2  )r*   rî   rï   rÐ  rÑ  Ú	MULTILINE)rC  r‡  rˆ  Úmetadata_commentÚbodys   &&&  rb   Ú_format_kernel_definitionÚ.PythonWrapperCodegen._format_kernel_definition^  s]   € ô �=‰=×1×1Ð1·hô —v’v˜g t¨X¼R¿\¹\ÔJˆHß.6˜h˜Z r™?¸BÐØÐ&Ð'¨ }°C¸°}ÐEˆØˆrd   c          
     ó,   € V ^8„  d   QhRRRRRRRRRR/# rG  r`   )ra   s   "rb   rc   rM  j  s<   € ÷ !ñ !àð!ð ð!ð ð	!ð
 ð!ð #ñ!rd   c                	óR  € \         P                  P                  '       dY   V'       dQ   V P                  WVR 7      pV P                  P                  V4       \        P                  P                  '       d   R# V P                  WVR 7      pV P                  P                  V4       R# ))rˆ  N)
r*   rî   rï   rQ  r  r?  rH   rœ   r]  r  )rf  rC  r‡  rˆ  r‰  rŠ  rP  s   &&&&&& rb   r�  Ú*PythonWrapperCodegen._define_kernel_helperj  s‡   € ô �=‰=×1×1Ð1·cØ×1Ñ1Ø°8ð 2ó ˆDð ×%Ñ%×,Ñ,¨TÔ2Ü�w‰w×"×"Ð"áà×-Ñ-Ø¨xð .ó 
ˆð 	�‰×Ñ˜4Ö rd   c               ó   € V ^8„  d   QhRR/# )r\   rÐ   r‹   r`   )ra   s   "rb   rc   rM  €  s   € ÷ >ñ >°ñ >rd   c                	óP   € V P                   P                  VP                  4       R # ru   )r  r?  r†   )rf  rÐ   Úsubgraph_codes   &&&rb   Údefine_subgraph_launcher_fnÚ0PythonWrapperCodegen.define_subgraph_launcher_fn€  s   € Ø×!Ñ!×(Ñ(¨×)<Ñ)<Ö=rd   c                	ó   € ^RI Hp V# )rg   )ÚTritonKernel)rî   r\  )Úclsr\  s   & rb   Ú_get_triton_info_kernel_clsÚ0PythonWrapperCodegen._get_triton_info_kernel_clsƒ  s   € õ 	)àÐrd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÒ   zlist[list[int | sympy.Expr]]r#  r$  Úlaunch_kwargsztuple[str, ...]r`   )ra   s   "rb   rc   rM  �  s4   € ÷ UJñ UJð ,ðUJð >ðUJð 'ñUJrd   c	                óÊ  aa;a<a=a>a?a@€ ^ RI Hp	 ^RIHp
 ^RIHp ^RIHpHpH	p ^RI
Ho;HpHpHpHp VP                   p. o@/ o>. o=. pV=V@3R lo<R<V;V<V>V3R llpVP"                   Uu. uF  pVP$                  NK  	  ppVP"                   Uu. uF#  pVP&                  '       g   K  VP(                  NK%  	  pp\+        V4       EFú  w  ppVV9   d   V! VS;! VR7      R	R
7       K"  VS9  d   K+  SV,          pSV,          f   V! VS;! VR7      R	R7       KS  \-        V\.        P0                  4      '       dc   \-        V\.        P2                  4      '       d(   RVP4                  VP6                  P9                  4       3MR=w  pppV! VV! VVVVR7      4       KÕ  \-        V\.        P:                  4      '       d2   V! VV! VVP=                  4       VP9                  4       R7      4       EK&  \-        V\.        P>                  4      '       dQ   V! VV! VVP@                  P=                  4       VP9                  4       VPB                  PD                  R7      4       EK–  \-        V\F        \H        PJ                  34      ;'       d+    \L        PN                  PP                  PS                  V^4      pV! VV! VV4      VR7       EKý  	  \U        S@RS=VPV                   U u. uF  p \Y        V 4      NK  	  up R7      p!\L        PN                  P[                  4       p"\\        P^                  ! V"4      p#RV!RV#R/ S>C\`        Pc                  V^4      CR\e        S@S=R>R7      ./p$V'       d   \g        V4      V$R&   V'       d   \g        V4      V$R&   V U%u/ uF  p%V%S9   g   K  V%SV%,          bK  	  p&p%V&'       dp   V	! V"4      Pi                  V"4       V
f   \k        R4      hV
! V#Pl                  V#Pn                  V#Pp                  4      p'V! V'V&VPV                  4      p(V('       d   V(V$R&   \s        V4      ^8X  d6   VPu                  4       p). \w        \H        Px                  V^ ,          4      Op*EM"R V?3R llp+/ o?V U,u. uF  p,. \w        V+V,4      ONK  	  pp,V'       d   \s        V4      \s        V4      8X  g%   \k        R\s        V4       R\s        V4       R 24      h. p-\{        \}        Wb4      R! R	R"7       FU  w  p,p.V-P                  R#V! V.4      R$. \w        \€        V,4      OR%. \w        \‚        V,4      OR&. \w        \€        V,4      O/4       KW  	  R'VP                   R(V-R). \w        \„        S?P‡                  4       4      O/p). S?P‰                  4       Op*V'       d   V U/u. uF  p/VV/,          NK  	  up/V)R*&   \‹        VPŒ                  4      .p0\s        V4      ^ 8”  d[   SP‡                  4        FF  p\-        V\.        P:                  \.        P>                  34      '       d   K5  V0P                  V4       KH  	  V0P                  \…        V$4      4       V0P�                  \…        V)4      4       Ve0   V0P                  V^ ,          P=                  4       V^,          34       \g        V04      p0V0V P�                  9   d   V P�                  V0,          w  p%p$p1V%V$V1V*3# V R+\s        V P�                  4       2p%V%P“                  R,4      '       d!   V%P•                  R,4      '       g
   V%R-,          p%\—        4       p2\˜        Pš                  Pœ                  '       d   V2PŸ                  R.V%: R/24       MV2PŸ                  R.V: R/24       V%V)R0&   V P¡                  4       p3V)P£                  V3P¥                  4       4       V2P§                  V3P©                  4       4       \˜        Pš                  Pª                  '       d   V2PŸ                  R14       \­        V$4      p4V2P§                  R2. \w        WÂ4      O: R3V): R4V4: R524       \¯        W4      p5\˜        Pš                  Pœ                  '       d   V5P±                  R6V R72R6V% R724      p5\³        V54      p5V2P§                  V54       \L        PN                  P[                  4       p6V2PŸ                  R8V6Pl                   R924       \´        P¶                  ! VPŒ                  4      w  p7p8\´        P¸                  ! VPŒ                  4      p9R:V9 R;V8 2p:V P»                  V%V2P½                  4       V:4       V%V$V)3V P�                  V0&   V%V$V)V*3# u upi u upi u up i u up%i u up,i u up/i )?aÀ  Codegen a user-defined Triton kernel and return its cache entry.

Emits the ``async_compile.triton(...)`` wrapper, assigns a graph-unique
name (with a leading dunder stripped to avoid Python class-based name
mangling at the call site), and records the kernel in
``user_defined_kernel_cache``. Returns ``(name, triton_meta,
inductor_meta, extra_launcher_call_args)``; subsequent calls with the
same ``cache_key`` reuse the previously assigned name.
)Úget_interface_for_device)Ú	GPUTarget)Ú%try_filter_backend_options_for_target)Úconfig_to_dictÚ	FixedGridÚPrecomputedGrid)ÚConstexprArgÚKernelArgTypeÚSizeArgÚ	TensorArgÚTMADescriptorArgc                óL   <€ SP                  V4       SP                  V 4       R # ru   )rÈ   )r  rÃ  Úarg_indicesÚ	signatures   &&€€rb   Úadd_to_signatureÚPPythonWrapperCodegen.define_user_defined_triton_kernel.<locals>.add_to_signature¸  s   ø€ Ø×Ñ˜SÔ!Ø×Ñ˜sÖ#rd   c                ó4  <€ V'       dN   \        4       '       d	   S! W4       VP                  S9   d#   SVP                  ,          SVP                  &   R# R# VP                  S9  d   \        R VP                   R24      hV'       dE   \        4       '       d   S! V S! VP                  R7      4       MS! W4       ^SVP                  &   R# V'       d<   \        4       '       d   S! V S! VP                  R7      4       RSVP                  &   R# S! W4       R# )z	expected z
 in kwargs©rÐ   N)rG   rÐ   r²   )	r  rÃ  Úis_constexprÚequals_1Úequals_noneri  rq  r2  r  s	   &&&&&€€€€rb   Úadd_argÚGPythonWrapperCodegen.define_user_defined_triton_kernel.<locals>.add_arg¼  s×   ø€ ßÜ1×3Ò3ñ % SÔ.à—8‘8˜vÔ%ð +1°·±Õ*:�I˜cŸh™hÓ'ñ &ð —8‘8 6Ô)Ü(¨9°S·X±X°J¸jÐ)IÓJÐJçÜ5×7Ò7ñ
 )¨©lÀÇÁÔ.IÕJá(¨Ô2Ø*+�I˜cŸh™hÓ'ß Ü5×7Ò7ñ )¨©lÀÇÁÔ.IÔJØ*.�I˜cŸh™hÓ'á$ SÖ.rd   rt  T)ru  N)rw  Ústable)rÐ   Úapi_typer¹  r   )rÐ   Úbufferr   )rÐ   r|  r   rÅ   )rv  )Ú
size_dtyper„  Úargdefsrp  r�   r2  rÑ   )r„  Úpointer_range_overrideÚrestore_valueÚreset_to_zeroz!expected GPUTarget to be not NoneÚbackend_optionsc               ó    € V ^8„  d   QhRRRR/# )r\   r]   rÚ   r_   r^   r`   )ra   s   "rb   rc   ÚLPythonWrapperCodegen.define_user_defined_triton_kernel.<locals>.__annotate__a  s   € ÷ +ñ +Ð0@ð +ÀZñ +rd   c                ó¬  <€ \        V \        P                  4      '       dr   . V P                  OpV'       g   V # VP	                  \
        R 7       V F1  pVS9   d   K  \        P                  ! R\        S4       24      SV&   K3  	  \        V S4      # \        V \        4      '       g   \        R\        V 4       24      h\        P                  ! V 4      # ))r  Ú_launcher_szexpected int, got )r   rh   r   r‰  Úsortr‹   r]  r  rF   r�   r²   r…   rÝ   )r]   ÚsymbolsrP  Úextra_launcher_argss   &  €rb   Úrename_sizes_for_launcherÚYPythonWrapperCodegen.define_user_defined_triton_kernel.<locals>.rename_sizes_for_launchera  s¸   ø€ Ü˜d¤E§J¡J×/Ò/Ø2 × 1Ñ 1Ð2�Gß"Ø#˜Ø—L‘L¤S�LÔ)Û&˜ØÐ"5Ô5Ù$Ü38·<²<Ø)¬#Ð.AÓ*BÐ)CÐDó4Ð+¨CÓ0ñ  'ô & dÐ,?Ó@Ð@Ü! $¬×,Ò,Ü(Ð+=¼dÀ4»j¸\Ð)JÓKÐKÜ—}’} TÓ*Ð*rd   z/expected non-empty grids matching configs, got z grids and z configsc                ó:   € \        V ^,          P                  4      # rf   r   r  s   &rb   r  ÚHPythonWrapperCodegen.define_user_defined_triton_kernel.<locals>.<lambda>{  s   € ´3°q¸µt·{±{Ô3Crd   r  r*   ru  ÚcppÚpython_slowÚ	grid_typeÚprecomputed_gridsr‰  Údeclared_constexpr_namesr  Ú__r‹  zasync_compile.triton(z, '''rC  r]  zG
            @triton_heuristics.user_autotune(
                configs=z ,
                inductor_meta=z,
                triton_meta=z{,
                filename=__file__,
                custom_kernel=True,
            )
            @triton.jit
            rü   rS  z''', device_str='rÚ  z# Original path: r  )FFF)ÚexperimentalNNr`   )_Útorch._dynamo.device_interfacerc  Úruntime.triton_compatrd  Úruntime.triton_helpersre  Úruntime.triton_heuristicsrf  rg  rh  Úcommonri  rj  rk  rl  rm  r‘   ÚparamsrÐ   ru  ÚnumrÚ  r   r,   ÚTMADescriptorr¾  r¹  r°  r¤   rÁ   rž   r/   rÉ   rµ   rÅ   r�   rh   rÝ   rH   rœ   r¥   Ústatically_known_equalsrS   Ú	arg_namesrI   Úget_current_device_or_throwr1   rP  r€   ÚfromkeysrQ   rƒ   Úraise_if_triton_unavailabler²   r…   ÚccÚwarp_size_or_defaultr  Úsetup_grid_as_argsrÊ  rj   r  r  rÈ   rŒ  rO   r‹   r«  rë  Úidr9  Úextendr(  rD  Úendswithr?   r*   rî   Úunique_user_kernel_namesrö   r^  ÚupdateÚinductor_meta_commonr?  Úgen_common_triton_importsrj  rv   rX  r\  r^  ÚinspectÚgetsourcelinesÚgetsourcefilerI  r  )Arf  rV  rÑ   r  Úrestore_value_argsÚreset_to_zero_argsrÒ   r#  ra  rc  rd  re  rf  rg  rh  rj  rk  rl  rm  Úoriginal_nameÚequal_to_1_argsrx  Úprž  Ú
constexprsr  r  rÃ  r{  r¹  r   rv  r  Útriton_signaturer�   Údevice_propsru  rÐ   Úbackend_option_candidatesrö  Úfiltered_backend_optionsrv  Úextra_launcher_call_argsrŠ  rá   r‘  Úcfgr¹  Ú	cache_keyÚcached_inductor_metarO  Útriton_info_kernel_clsÚsanitized_triton_metarW  Úcurrent_devicer  ÚlinenoÚsrcfilerˆ  ri  rq  ro  r2  r‰  rp  sA   &&&f&&&&&                                                  @@@@@@rb   Ú!define_user_defined_triton_kernelÚ6PythonWrapperCodegen.define_user_defined_triton_kernel�  sÝ  þ€ õ( 	Lå5ÝR÷	
ñ 	
÷
	
õ 	
ð Ÿ™ˆØ)+ˆ	Ø$&ˆ	Ø!#ˆØ%'ˆö	$÷#	/ò #	/ðJ &,§]¢]Ó3¡] �Q—V”V¡]ˆ	Ð3Ø%+§]¢]ÓE¡] °a·nµn”e�a—e”e¡]ˆ
ÐEÜ! )×,‰HˆC�Ø�jÔ Ù˜™\¨sÔ3À$ÕGÙà˜&Ô Ùà˜•+ˆCà�c�{Ò"Ù˜™\¨sÔ3À×Fä˜c¤2×#3Ñ#3×4Ò4ô & c¬2×+AÑ+A×BÒBð " 3§?¡?°C·J±J×4HÑ4HÓ4JÑKà9ñ 1�H˜k¨5ñ
 ØÙ(Ø!$Ø%-Ø(3Ø"'ô	öô   ¤R§Y¡Y×/Ò/ÙØÙ!Ø!$Ø#&§<¡<£>Ø"%§-¡-£/ô÷ô   ¤R×%7Ñ%7×8Ò8ñ ØÙ!Ø!$Ø#&§8¡8×#4Ñ#4Ó#6Ø"%§-¡-£/Ø#&§:¡:×#4Ñ#4ô	÷ô  *Øœc¤5§=¡=Ð1ó ÷  ð  äŸ'™'×*Ñ*×BÑBØØóð ñ ˜C¡¨¨cÓ!2¸X×FÐFñs -ôv -ØØØØ)/×)9Ò)9Ó:Ñ)9 A”W˜Q–ZÑ)9Ñ:ô	
Ðô —‘×4Ñ4Ó6ˆÜ'×.Ò.¨vÓ6ˆàÐ)Ø�lð ð Øðä—-‘- °Ó3ðð ÜØØ'Ø+-ôðð#
ˆ÷. Ü+0Ð1CÓ+DˆK˜Ñ(çÜ+0Ð1CÓ+DˆK˜Ñ(ñ ,9ó%
Ù+8 4¸DÀF¹NÔˆD�&˜•,Ò©=ð 	"ð %
÷ %Ù$ VÓ,×HÑHÈÔPØÒ Ü$Ð%HÓIÐIÙØ×!Ñ!Ø—‘Ø×1Ñ1óˆFñ (MØÐ1°6×3CÑ3Có(Ð$÷ (Ø1I�Ð-Ñ.äˆu‹:˜Œ?à,5×,HÑ,HÓ,JˆMØ'F¬¬U¯]©]¸EÀ!½HÓ)EÐ'FÒ$÷+ð +ð" EGÐÙINÓOÉÀÐ<”sÐ4°dÓ;Ó<ÉˆEÐOçœc %›j¬C°«LÔ8Ü$ØEÄcÈ%ÃjÀ\ÐQ\Ô]`ÐahÓ]iÐ\jÐjrÐsóð ð !#ÐÜ#Ü�EÓ#Ñ)CÈT÷‘	��cð "×(Ñ(à ¡.°Ó"5Ø Ð"5¤C¬¨tÓ$4Ð"5ØÐ2¤¤U¨DÓ!1Ð2Ø%Ð':¬¬U°DÓ)9Ð':ð	öñð ˜_×5Ñ5Ø#Ð%6Ø%Ð'P¬¬SÐ2E×2LÑ2LÓ2NÓ)OÐ'PðˆMð
 (EÐ)<×)AÑ)AÓ)CÐ'DÐ$çá&0ó9Ù&0 �	˜!—�¡jñ9ˆMÐ4Ñ5ô
 ˜VŸY™Y›-˜ˆ	Üˆw‹<˜!ÔØ—}‘}–�ä! #¬¯	©	´2×3EÑ3EÐ'F×GÔGØ×$Ñ$ SÖ)ñ 'ð 	×Ñœ˜[Ó)Ô*Ø×Ñœ˜]Ó+Ô,àÒ&Ø×Ñ˜o¨aÕ0×9Ñ9Ó;¸_ÈQÕ=OÐPÔQä˜)Ó$ˆ	Ø˜×6Ñ6Ô6Ø6:×6TÑ6TØõ7Ñ3ˆD�+Ð3ð ØØ$Ø(ð	ð ð  � ¤# d×&DÑ&DÓ"EÐ!FÐGˆð �?‰?˜4× Ò ¨¯©°t×)<Ò)<Ø˜•8ˆDä(Ó*ˆÜ�=‰=×1×1Ð1Ø×%Ñ%Ð(=¸d¹XÀUÐ&KÕLà×%Ñ%Ð(=¸mÑ=NÈeÐ&TÔUà'+ˆ�mÑ$Ø!%×!AÑ!AÓ!CÐØ×ÑÐ3×HÑHÓJÔKà×ÑÐ5×OÑOÓQÔRÜ�=‰=×)×)Ð)Ø×%Ñ%Ð&DÔEô !3°;Ó ?ÐØ×Ñðà8œ3˜~Ó7Ð8Ñ;ð <Ø,Ñ/ð 0Ø2Ñ5ð 6ð		ô	
ô OØó
ˆ
ô �=‰=×1×1Ð1à#×+Ñ+¨d°=°/ÀÐ,CÀtÈDÈ6ÐQRÀ^ÓTˆJÜ=¸jÓIˆ
Ø×Ñ˜zÔ*äŸ™×<Ñ<Ó>ˆØ×!Ñ!Ð$5°n×6IÑ6IÐ5JÈ"Ð"MÔNÜ×*Ò*¨6¯9©9Ó5‰	ˆˆ6Ü×'Ò'¨¯	©	Ó2ˆØ& w i¨q°°Ð9ˆØ×ÑØØ×$Ñ$Ó&Øô	
ð 6:¸;ÈÐ4Vˆ×&Ñ& yÑ1Ø�[ -Ð1IÐIÐIùòC 4ùÚEùò@ ;ùòB%
ùò` Pùò49s0   Á&eÂeÂ$eËeÎ	eÎeÑ.eÖe c               ó    € V ^8„  d   QhRRRR/# )r\   rC  r‹   r  rÖ   r`   )ra   s   "rb   rc   rM  ä  s   € ÷ ñ ¨sð À*ñ rd   c                	ó  € V R VP                    R2pVe   VR V 2,          p\        P                  ! VRRR7      p\        WRP                  4      pVR8H  pV'       g*   V P                  \        W\        P                  4      4       V# )r  r>  T)Ú
is_integerÚis_positiver  )	r  rh   r]  r`  r>  rö   rÂ  rH   rœ   )rf  rC  Útreer  Úsym_namerP  rÃ  Úis_benchmark_kernels   &&&&    rb   Úgenerate_numel_exprÚ(PythonWrapperCodegen.generate_numel_exprä  s{   € Ø!�] ! D§K¡K =°Ð6ˆØÒØ˜!˜F˜8˜Õ$ˆHÜ�lŠl˜8°À$ÔGˆô ˜c§:¡:Ó.ˆà)¨RÑ/Ðß"Ø�N‰NÔ.¨t¼!¿'¹'ÓBÔCàˆ
rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÃ  r`  rœ   rW   r_   r„  r`   )ra   s   "rb   rc   rM  ù  s)   € ÷ Añ AØ"ðAØ+8ðAà	ñArd   c                	ól   € V P                  VP                   R \        VP                  4       24       R# rÂ  )rö   rb  rŒ  rc  )rf  rÃ  rœ   s   &&&rb   rÇ  Ú7PythonWrapperCodegen._generate_symbolic_call_arg_helperù  s)   € ð 	�‰˜#Ÿ)™)˜ C¬¨c¯n©nÓ(=Ð'>Ð?Ö@rd   c               ó   € V ^8„  d   QhRR/# ©r\   ÚwsrM   r`   )ra   s   "rb   rc   rM  þ  s   € ÷ )ñ )°ñ )rd   c                	ó  € VP                  4       p\        W4      pVP                  \        P                  8X  d   V P                  V4       EMPVP                  \        P                  8X  d3   V P                  V4       V P                  V P                  V4      4       MÿVP                  \        P                  8X  dÌ   V P                  P                  V4      pV'       di   \        V\        4      '       d!   \        VP                  \        4      '       g   \        R 4      h\        P                  ! VP                  V4      Vn        MUV P                  V4       V P                  V P                  V4      4       W0P                  V&   M\        VP                  4      h\         P"                  P$                  '       dÌ   V P&                  P                  \(        P+                  V VVP,                  VP.                  \0        P2                  P4                  P7                  VP8                  4      3RR7      4       VP                  \        P                  8w  d1   V P&                  P                  \(        P                  W4      4       R# R# R# )z=expected prior to be an AllocateLine with a WorkspaceArg node)r>  rÄ   Nrf   )rž   r  Ú	zero_moderN   ÚUNINITIALIZEDrö   ÚZERO_ON_CALLÚmake_zero_bufferÚZERO_PER_GRAPHr>  r;  r   r˜   rM   r²   Úmaximumr*   rî   rï   rø   r²  Úmake_allocationr�   r   rH   rœ   r¥   rû  r   )rf  rÒ  rÐ   rô   Úpriors   &&   rb   Úgenerate_workspace_allocationÚ2PythonWrapperCodegen.generate_workspace_allocationþ  sÎ  € Ø�{‰{‹}ˆÜ˜DÓ%ˆØ�<‰<Ô,×:Ñ:Ô:Ø�N‰N˜4Ö Ø�\‰\Ô.×;Ñ;Ô;Ø�N‰N˜4Ô Ø�N‰N˜4×0Ñ0°Ó6Õ7Ø�\‰\Ô.×=Ñ=Ô=Ø×-Ñ-×1Ñ1°$Ó7ˆEßä˜u¤l×3Ò3Ü" 5§:¡:¬|×<Ò<ä(ØWóð ô *×1Ò1°%·*±*¸bÓA�•
à—‘˜tÔ$Ø—‘˜t×4Ñ4°TÓ:Ô;Ø26×)Ñ)¨$Ò/ä  §¡Ó.Ð.ä�=‰=×1×1Ð1Ø×&Ñ&×0Ñ0Ü$×4Ñ4ØØØ—I‘IØ—H‘HÜŸ7™7×+Ñ+×=Ñ=¸b¿h¹hÓGÐIØð 5ó ô	ð �|‰|Ô0×>Ñ>Ô>Ø×*Ñ*×4Ñ4Ü(×9Ñ9¸$ÓEöñ ?ñ 2rd   c               ó   € V ^8„  d   QhRR/# rÑ  r`   )ra   s   "rb   rc   rM  )  s   € ÷ :ñ :°,ñ :rd   c                	ó|   € VP                   \        P                  8w  d   V P                  \	        W4      4       R # R # ru   )rÔ  rN   rØ  rö   r}  )rf  rÒ  s   &&rb   Úgenerate_workspace_deallocationÚ4PythonWrapperCodegen.generate_workspace_deallocation)  s,   € Ø�<‰<Ô,×;Ñ;Ô;Ø�N‰NÔ.¨tÓ8Ö9ñ <rd   c                	ó$   € V R V P                    2# )z.zero_())r”  )rf  rÐ   s   &&rb   r×  Ú%PythonWrapperCodegen.make_zero_buffer-  s   € Ø��x §¡˜}Ð-Ð-rd   c                	óH   € V R RP                  V4       RV P                   2# )rS  rT  r3  )r  r”  )rf  rÐ   rp  s   &&&rb   rà  Ú%PythonWrapperCodegen.wrap_kernel_call0  s'   € Ø��q˜Ÿ™ 9Ó-Ð.¨a°·±¨}Ð=Ð=rd   c                	ót   € VP                  \        P                  P                  P	                  4       4       R # ru   )rö   rH   rœ   r	  Úsynchronize©rf  r|  s   &&rb   rÜ  Ú(PythonWrapperCodegen.generate_debug_sync3  s#   € Ø×ÑœŸ™×+Ñ+×7Ñ7Ó9Ö:rd   c                	óü   € V P                   P                  R 4       V P                   P                  R\        P                  P                   R24       VP                  V P                   P                  4       4       R# )z*from torch.profiler import record_functionzwith record_function('graph_z_inductor_wrapper_call'):N)r  rö   rH   rœ   Úgraph_idr
  r  )rf  r  s   &&rb   r  Ú8PythonWrapperCodegen.generate_profiler_mark_wrapper_call6  sb   € Ø×Ñ×#Ñ#Ð$PÔQØ×Ñ×#Ñ#Ø*¬1¯7©7×+;Ñ+;Ð*<Ð<UÐVô	
ð 	×Ñ˜D×-Ñ-×4Ñ4Ó6Ö7rd   c                	ó<   € V P                   P                  R 4       R# )zstart_graph()N)r  rö   re  s   &rb   r  Ú)PythonWrapperCodegen.generate_start_graph=  s   € Ø×Ñ×#Ñ# OÖ4rd   c                	ób   € V P                   P                  R \        P                  : R24       R# )z
end_graph(r3  N)r  rö   r*   Úprofile_bandwidth_outputre  s   &rb   r  Ú'PythonWrapperCodegen.generate_end_graph@  s'   € Ø×Ñ×#Ñ# j´×1PÑ1PÑ0SÐSTÐ$UÖVrd   c                óˆ   € V P                   P                  \        P                  P                  P                  4       4       R# )z<Synchronize GPU to ensure proton captures all kernel events.N)r  rö   rH   rœ   r	  rç  re  s   &rb   r  Ú-PythonWrapperCodegen.generate_proton_finalizeC  s)   € à×Ñ×#Ñ#¤A§G¡G×$6Ñ$6×$BÑ$BÓ$DÖErd   c                	ó`   € V P                   P                  R \        P                   R24       R# )úU
            for kernel in globals().values():
                if isinstance(kernel, z‰.CachingAutotuner):
                    kernel.cuda_kernel_saved = False
                    kernel.cpu_kernel_saved = False
            N©r  r?  r0   r‘   re  s   &rb   r  Ú6PythonWrapperCodegen.generate_reset_kernel_saved_flagsG  s2   € Ø×Ñ× Ñ ð'ä'8×'AÑ'AÐ&Bð Cðö	
rd   c                ó`   € V P                   P                  R\        P                   R24       R# )aX  
Precompile and save the per-config kernel artifacts (CUBINs on GPU,
``.so`` + launcher ``.so`` on CPU) for Triton kernels that haven't
been precompiled and saved as a side effect of running the generated
JIT model (Python wrapper). This can happen when the model contains
control flow: only one pass through the control flow operators covers
the kernels that are saved, the remaining kernels are not launched,
hence not saved. The main purpose of this codegen is to compile and
save the Triton kernels outside the active control flow path for
subsequent AOTInductor code generation and compilation.
rõ  aé  .CachingAutotuner):
                    if kernel.device_props.type == "cpu":
                        if not kernel.cpu_kernel_saved:
                            if len(kernel.launchers) == 0:
                                kernel.precompile()
                            kernel.save_cpu_kernel(launcher=kernel.launchers[0])
                    else:
                        if not kernel.cuda_kernel_saved:
                            if len(kernel.launchers) == 0:
                                kernel.precompile()
                            kernel.save_gpu_kernel(
                                stream="stream",  # use dummy stream
                                launcher=kernel.launchers[0],
                            )
            Nrö  re  s   &rb   r  Ú5PythonWrapperCodegen.generate_save_uncompiled_kernelsQ  s4   € ð 	×Ñ× Ñ ð'ä'8×'AÑ'AÐ&Bð Cðö	
rd   c                	ó@   € R  pV Uu. uF
  q2! V4      NK  	  up# u upi )c                ó,  € \        V \        4      '       d   \        V 4      '       d
   V R ,           # T # \        V \        \        \
        \        34      '       d   \        V 4      # \        \        P                  P                  P                  V 4      4      # )r  )r   r‹   rR   r�   ÚfloatrE  r`  rŒ  rH   rœ   r¥   r¦   )rÃ  s   &rb   Úwrap_argÚAPythonWrapperCodegen.prepare_triton_kernel_call.<locals>.wrap_argr  sg   € Ü˜#œs×#Ò#ä*BÀ3×*GÒ*G�s˜Y•ÐPÈSÐPÜ˜C¤#¤u¬d´OÐ!D×EÒEÜ˜3“x�äœQŸW™W×-Ñ-×6Ñ6°sÓ;Ó<Ð<rd   r`   )rf  rp  rý  rÃ  s   &&  rb   Úprepare_triton_kernel_callÚ/PythonWrapperCodegen.prepare_triton_kernel_callq  s%   € ò	=ñ *3Ó3© #�˜–©Ñ3Ð3ùÒ3s   ˆc                	óF  a € \        V\        4      '       Ed   \        V\        P                  4      '       d3   VP	                  4       P                  4       pS P                  V,          pMpS P                  P                  V4      '       d   TpS P                  V,          pM9Vf   \        R4      hRS P                   2pTpS ;P                  ^,          un	        Vf   \        RV 24      h\        P                  P                  P                  VP                  4       4      p\        P                  P                  P                  \        P                  P                  V4      4      p\        P                  P                  P                  VP!                  4       4      pVP#                  4       p	VP%                  4       p
\        P                  P                  P'                  VP)                  4       P*                  4      pRV RV R\-        V	4       RV
 RV RV R2pS P.                  P1                  V R	V 24       \        V\        P                  4      '       d6   S P3                  VR
R7      pTpS P.                  P1                  V R	V 24       V# \5        V\6        P8                  4      '       g   \        V\:        4      '       dÚ   \        V\<        4      '       d   VS P>                  9   d   V# Vf   R# Tp\        V\:        4      '       d   VP@                  pV\        P                  P                  PB                  9   d,   \        P                  P                  PB                  V,          p\=        \        P                  P                  P'                  V4      4      # \        V\<        \D        \F        \H        34      '       d   \=        V4      # \        V\J        4      '       d    RRPM                  V 3R lV 4       4       R2# \O        R\Q        V4       24      h)NzBV.graph.get_buffer(arg) and raw_arg can't be None at the same timeÚtmp_arg_z Failed to find a buffer for arg zgenerate_example_value(rT  z, 'z', r3  r2  T)r³  r´  r„  r‰  c              3  óZ   <"  € T F   pSP                  V\        V4      4      x € K"  	  R # 5iru   ré   )rx   Úarf  s   & €rb   rz   ÚBPythonWrapperCodegen.generate_example_arg_value.<locals>.<genexpr>º  s(   øé € Ð ZÑVYÐQR ×!@Ñ!@ÀÄDÈÃG×!LÐ!LÓVYùrë   rŠ  zUnsupported type ))r   Útorch_dtyper,   rœ  Ú
get_tensorrž   r  r;  r²   r  rH   rœ   r¥   r¯  r7  Úget_allocation_sizer8  r=  r¤   rû  rÂ   rÅ   r7   rø   rö   r¿  rŒ  rh   rµ  r`  r‹   r;  rc  r¤  r�   rü  rE  r‚   r  r£  r…   )rf  rÃ  Úarg_typeÚraw_argrå  rº  rÃ   Úallocation_sizerÄ   r�   r   rÅ   r†   s   f&&&         rb   rê   Ú/PythonWrapperCodegen.generate_example_arg_value}  sQ  ø€ Ü�h¤×,Ó,Ü˜'¤2×#3Ñ#3×4Ò4à"×-Ñ-Ó/×8Ñ8Ó:�Ø×*Ñ*¨3Õ/‘Ø×%Ñ%×)Ñ)¨#×.Ò.Ø�Ø×*Ñ*¨3Õ/‘à’?Ü(Ø\óð ð & d×&FÑ&FÐ%GÐH�Ø�Ø×0Ò0°AÕ5Õ0àŠ{Ü$Ð'GÈÀuÐ%MÓNÐNÜ—7‘7×#Ñ#×6Ñ6°s·|±|³~ÓFˆDÜŸg™g×.Ñ.×AÑAÜ—‘×+Ñ+¨CÓ0óˆOô —W‘W×%Ñ%×8Ñ8¸¿¹Ó9IÓJˆFà—^‘^Ó%ˆFØ—M‘M“OˆEÜ—W‘W×%Ñ%×7Ñ7¸¿¹Ó8H×8OÑ8OÓPˆFØ-¨d¨V°2°f°X¸SÄÐQWÓAXÐ@YÐY\Ð]bÐ\cÐceÐflÐemÐmoÐpð  pAð  ABð  CˆEØ×&Ñ&×0Ñ0°H°:¸SÀÀÐ1HÔIä˜'¤2×#3Ñ#3×4Ò4ð ×:Ñ:Ø Ø%)ð ;ó �ð �Ø×*Ñ*×4Ñ4¸°zÀÀUÀGÐ5LÔMàˆOÜ˜¤%§+¡+×.Ò.´*¸SÄ/×2RÒ2Rä˜#œs×#Ò#Ø˜$Ÿ/™/Ô)Ø�JØ’?Ù!Ø�Ü˜#œ×/Ò/Ø—n‘n�Ø”a—g‘g×&Ñ&×CÑCÔCÜ—g‘g×&Ñ&×CÑCÀCÕH�ä”q—w‘w×'Ñ'×9Ñ9¸#Ó>Ó?Ð?ä˜œc¤3¬¬tÐ4×5Ò5Ü�s“8ˆOÜ˜œT×"Ò"Ø�t—y‘yÔ ZÑVYÓ ZÓZÐ[Ð[\Ð]Ð]ä%Ð(9¼$¸s»)¸Ð&EÓFÐFrd   c                	ó˜   a € \        V\        4      '       d*   R RP                  V 3R lV 4       4      ,           R,           # \        V4      # )r‰  rT  c              3  óF   <"  € T F  pSP                  V4      x € K  	  R # 5iru   )Ú_grid_dim_str)rx   rÙ   rf  s   & €rb   rz   Ú5PythonWrapperCodegen._grid_dim_str.<locals>.<genexpr>Á  s   øé € ÐRÁ\¸T × 2Ñ 2°4× 8Ð 8Ã\ùs   ƒ!rŠ  )r   r‚   r  rŒ  )rf  Úgrid_per_dims   f&rb   r  Ú"PythonWrapperCodegen._grid_dim_str¾  s?   ø€ Ü�l¤D×)Ò)à�d—i‘iÔRÁ\ÓRÓRÕRÐUXÕXðô ˜Ó&Ð&rd   r�   rî   rt  rq  rr  ru  rv  rÕ   c               ó    € V ^8„  d   QhRRRR/# ©r\   rC  r‹   ru  zTritonMeta | Noner`   )ra   s   "rb   rc   rM  Æ  s   € ÷ 6
ñ 6
àð6
ð 'ñ6
rd   c               óè  € V P                   P                  V Uu/ uF;  p\        V\        4      '       g   K  V\        P
                  P                  V4      bK=  	  up4       T;'       g    \        P
                  P                  4       p\        P
                  P                  P                  pV P                  \        V VVVVVVVV	V\        P
                  P                  V
VR7      4       R# u upi )z®
Generates kernel call code.

triton: Defines whether the backend uses Triton for codegen. Otherwise it uses the CUDA language when gpu=True,
        and C++ when gpu=False.
)rC  rp  rq  rr  rt  rî   ru  rv  r�   rx  rÕ   ry  N)r  r©  r   r‹   rH   rœ   Útry_get_bufferrŸ  r    ry  rö   ro  rÐ   )rf  rC  rp  r�   rî   rt  rq  rr  ru  rv  rÕ   rÃ  ry  s   &&&$$$$$$$$  rb   Úgenerate_kernel_callÚ)PythonWrapperCodegen.generate_kernel_callÆ  sÏ   € ð. 	×Ñ×#Ñ#ñ %óá$�CÜ˜c¤3×'ô 1�”Q—W‘W×+Ñ+¨CÓ0Ò0Ù$ñô	
ð ×@Ð@œ1Ÿ7™7×>Ñ>Ó@ˆÜŸW™W×.Ñ.×AÑAÐØ�‰ÜØØ'Ø#à!à!à#Øà'Ø+ØÜŸ7™7Ÿ<™<à%9Ø#5ô%ö	
ùòs
   šC/·$C/rx  r  ry  c               ó    € V ^8„  d   QhRRRR/# r  r`   )ra   s   "rb   rc   rM  þ  s"   € ÷ X(ñ X(àðX(ð 'ñX(rd   c       
   
     	ó.  a aa € T;'       g    \         P                  P                  4       pV'       g“   VP                  R9  d‚   VP                  R8X  d$   S P	                  S P                  SV4      4       R# VP                  R8X  d$   S P	                  S P                  SV4      4       R# \        RVP                   R24      hS P                  V4      pRP                  V4      pVeI   V\        8w  d>   S P                  4        RpS P	                  V R\        VP                  4       R	24       M!\        P                  S VP                  V
4      pV'       g&   R
V R	2pS P	                  S RS RV RV R	24       R# S P                  4        \         P"                  P$                  '       EdÛ   SS P&                  9  EdÉ   Ve   \)        V4      \)        V4      8X  g   \+        R4      hRo V'       dK   \         P                  P,                  '       d+   \         P                  P,                  P/                  VR4      o \         P0                  P2                  pR VV 3R llpV V 3R lp. p. pVf7   Ve   \+        R4      hR.\)        V4      ,          pR.\)        V4      ,          pM$\)        V4      \)        V4      8w  d   \+        R4      h/ p\5        \7        W%Wg4      4       EF	  w  pw  ppppRp\9        V\:        4      '       d%   R\;        V4      9   d   VP=                  R4      w  ppRpS '       d    VS 9   d   S P?                  S V,          4      pV'       dU   Tp\9        V\@        4      '       g<   \C        V\D        PF                  4      '       g   \9        V\H        4      '       d   VVV&   MðVR8X  d   V! WgVV4      '       d   VV,          pMÏ\9        V\@        4      '       d§   \J        PL                  ! RV4      '       d   TpMlV'       d&   S PO                  VVV4      pVPQ                  V4       M?VS PR                  9  d   S PO                  VVV4      pMS PR                  V,          ^ ,          pV'       g   VS3S PR                  V&   MS PO                  VVV4      p\9        V\:        4      '       d   \U        V4      '       d
   VR,          pTPQ                  Vf   TMV RV 24       EK  	  \W        4       '       d)   \         P                  PX                  P[                  4       MVP                  pS P\                  P	                  R\         P                  PX                  P_                  V4       R24       S P\                  Pa                  4        S P\                  P	                  S RRP                  V4       RV R	24       S P\                  Pc                  4        V'       d7   V'       d/   S P\                  P	                  RRP                  V4       24       M.V'       g'   S P\                  P	                  \e        RVR4      4       S P&                  Pg                  S4       \         P                  Ph                  '       d   R# \         P                  Pj                  Pl                  pVPo                  VSVR4       V;_uu_ 4        S P	                  S RV RV R	24       RRR4       S P                  4        R#   + '       g   i     L"; i)Úcudar%  Úmpszdevice z nyiNrT  Ú
raw_streamr  r3  z	c_void_p(r5  rS  z$call_args and arg_types do not matchc               ó   € V ^8„  d   QhRR/# r£  r`   )ra   s   "rb   rc   ÚGPythonWrapperCodegen._generate_kernel_call_helper.<locals>.__annotate__A  s   € ÷ 	ñ 	°ñ 	rd   c                 óº   <€ SP                   P                  4        U Uu. uF  w  rVS8X  g   K  V NK  	  pp pV'       d   RRP                  V4       R2# R# u upp i )z;Returns del for tensors whose last consumer is this kernel.údel rT  r  r  )r  r«  r  )r°  ÚknÚtensors_to_deleterC  rf  s      €€rb   Úget_autotune_deletion_callÚUPythonWrapperCodegen._generate_kernel_call_helper.<locals>.get_autotune_deletion_callA  sg   ø€ ð '+×&GÑ&G×&NÑ&NÔ&Pô%á&P™
˜Ø˜[Ñ(÷ �FÙ&Pð "ñ %÷
 %Ø! $§)¡)Ð,=Ó">Ð!?¸rÐBÐBÙùó%s
   ŸA¯Ac                ó²  <€ W,          pWC9   d   R# \        \        W4      4       F›  w  pw  rgWR8X  g   \        V\        4      '       g   K&  RpS'       d    VS9   d   SP	                  SV,          4      pVR8X  d   KX   VP                  4       p	\        V	P                  4       F  w  r«W´8X  g   K  V RV
 R2W4&     R# 	  K�  	  R#   \         d     K¯  i ; i)z­We try to infer raw_arg (i.e. raw_args[idx]) from remaining raw_args.
This is particularly useful for jagged cases, where the dimension is often
being passed in as an input.Tr  z.shape[rŠ  F)rÚ  r  r   r.   r*  rÂ   rÃ   r£  )rq  rr  r  Úreused_argsÚ
target_argr¹  Úraw_keyr
  Útriton_inputrµ   r¨  r@  Úautotune_argsrf  s   &&&&        €€rb   Úinfer_arg_by_inputsÚNPythonWrapperCodegen._generate_kernel_call_helper.<locals>.infer_arg_by_inputsL  sÝ   ø€ ð
 &�]�
ØÔ,Ùä-6´s¸8Ó7NÖ-OÑ)�AÑ)˜Ø”x¤z°'¼6×'BÒ'BÙ à#%�Lß$¨°MÔ)AØ'+×'EÑ'EØ)¨'Õ2ó(˜ð $ rÔ)Ù ð	!Ø!(×!3Ñ!3Ó!5˜Ü&/°·±Ö&<™F˜CØ žØ=I¸NÈ'ÐRUÐQVÐVWÐ:X Ñ 7Û'+ó '=ñ .Pñ, øô	 /ô !ò !ð!ús   Á?-CÂ1CÂ?CÃCÃCzkeys are not None but args arez#call_args and raw_args do not matchrR  r  z^(workspace|semaphore)r  r  r  z.run(z	, stream=r!  z
<del_call>)r  Úxpu)8rH   rœ   rŸ  r…   rö   rà  r4  rÿ  r  r4   r‘  rý  r  r²  r8  r‹  r*   rî   rï   r÷   r  r²   Úautotuning_mappingr;  r?  Úautotune_per_kernel_allocrÚ  r  r   r‹   Úsplitr*  r  rŒ  rh   rµ  r`  rÐ  Úmatchrê   rÈ   r  rR   r   r	  r  rø   r  r¿  ró  r<   r  r]  rµ  rB  r   )!rf  rC  rp  r�   rî   rt  rq  rr  ru  rv  rx  rÕ   ry  Úcall_args_strÚstream_nameÚ
stream_ptrÚ_per_kernelr$  r,  Úall_argsÚtensor_arg_strsr'  r¹  rÃ  r	  r)  r
  r  r*  Úarg_strÚ	guard_idxr¡  r+  s!   ff&$$$$$$$$$$                   @rb   r}  Ú1PythonWrapperCodegen._generate_kernel_call_helperþ  sþ  ú€ ð  ×@Ð@œ1Ÿ7™7×>Ñ>Ó@ˆß˜&Ÿ+™+¨_Ô<Ø�{‰{˜eÔ#Ø—‘˜t×4Ñ4°[À)ÓLÔMñ ð —‘ Ô%à—‘˜t×4Ñ4°[À)ÓLÔMñ ô # W¨V¯[©[¨M¸Ð#>Ó?Ð?ð ×7Ñ7¸	ÓBˆØŸ	™	 -Ó0ˆØÒ)Ð.@ÔDVÔ.Vð ×,Ñ,Ô.Ø&ˆKØ�N‰NØ�-Ð1Ô2FÀvÇ|Á|Ó2TÐ1UÐUVÐWõô /×CÑCØ�f—l‘l JóˆK÷ Ø$ [ M°Ð3ˆJØ�N‰NØ�-˜q  ¨Q¨}¨o¸RÀ
¸|È1ÐMôñ à×%Ñ%Ô'ô �M‰M×2×2Ñ2Ø 4×#=Ñ#=Õ=ð Ò)¬c°)«nÄÀIÃÔ.NÜ$Ð%KÓLÐLà ˆMß#¬¯©×(B×(BÐ(BÜ !§¡× :Ñ :× >Ñ >Ø(¨$ó!�ô !×-Ñ-×GÑGˆK÷	ñ 	öðB ˆHØ ˆOØÒàÒ'Ü(Ð)IÓJÐJØ ˜6¤C¨	£NÕ2�Ø ˜6¤C¨	£NÕ2‘ä�x“=¤C¨	£NÔ2Ü(Ð)NÓOÐOàˆKÜ8Aä�I¨(Ó=÷9Ñ4�Ñ4�C˜ 7¨Gð �Ü˜c¤3×'Ò'¨C´3°s³8¬Oà"Ÿy™y¨›~‘H�C˜à+/�ß  W°Ô%=Ø#'×#AÑ#AØ% gÕ.ó$�L÷  Ø*�GÜ% h´×<Ò<Ü" 8¬U¯[©[×9Ò9Ü% c¬?×;Ò;à/6˜ GÑ,øØ ”]Ñ':Ø¨¨;÷(ò (ð
 *¨'Õ2‘GÜ ¬+×6Ò6ô —x’xÐ 9¸3×?Ò?Ø"%™ß$Ø"&×"AÑ"AØ ¨7ó#˜ð (×.Ñ.¨wÕ7Ø D×$EÑ$EÔEØ"&×"AÑ"AØ ¨7ó#™ð #'×"CÑ"CÀCÕ"HÈÕ"K˜ß&ØBIÈ;ÐAW˜×9Ñ9¸#Ñ>øà"×=Ñ=¸cÀ8ÈWÓU�Gä˜c¤3×'Ò'Ô,DÀS×,IÒ,IØ˜yÕ(�GØ—‘¨3ª;¡¸s¸eÀ1ÀWÀIÐ<N×Oñg9ôt !—?’?ô —‘×"Ñ"×:Ñ:Ô<à—\‘\ð ð
 ×&Ñ&×0Ñ0ØœŸ™×*Ñ*×7Ñ7¸	ÓBÐCÀ1ÐEôð ×&Ñ&×0Ñ0Ô2Ø×&Ñ&×0Ñ0Ø�-˜u T§Y¡Y¨xÓ%8Ð$9¸À;À-ÈqÐQôð ×&Ñ&×2Ñ2Ô4çŸØ×*Ñ*×4Ñ4Ø˜4Ÿ9™9 _Ó5Ð6Ð7õ÷ !Ø×*Ñ*×4Ñ4Ü$Ø$Ø2Ø$óôð ×&Ñ&×*Ñ*¨;Ô7Ü�w‰w×"×"Ð"áô !"§¡× 4Ñ 4× BÑ BÐØ×.Ñ.¨y¸+ÀyÐRVÔWß"Ò"Ø�N‰N˜k˜]¨%°¨¸iÈÀ}ÐTUÐVÔW÷ #à×%Ñ%Ö'÷ #×"ús   Û\Ü\	c                	ó<   € V P                   P                  V4       R # ru   )r"  rÈ   r›  s   &&rb   rö   ÚPythonWrapperCodegen.writelineØ  s   € Ø�
‰
×Ñ˜$Örd   c                	ó:   € V F  pV P                  V4       K  	  R # ru   rþ  )rf  r"  rô   s   && rb   r1  ÚPythonWrapperCodegen.writelinesÛ  s   € ÛˆDØ�N‰N˜4Ö ó rd   c                	óN   € V P                   P                  \        V4      4       R # ru   )r"  rÈ   rB   )rf  Úctxs   &&rb   r
  Ú"PythonWrapperCodegen.enter_contextß  s   € Ø�
‰
×Ñœ+ cÓ*Ö+rd   c                	óB  a a	€ ^ RI Hp V! 4       '       d   ^ RIp\        V\        4      '       d    \        VP                  P                  4      # \        V\        P                  4      '       d   \        V4      # \        V\        \        34      '       dC   \        P                   ! R R4      4       o	\        \        V4      ! V	V 3R lV 4       4      4      # \        V\         P"                  P$                  4      '       d   \'        V4      # \        V\(        P*                  \(        P,                  \.        34      '       d   VP1                  4       # V! 4       '       d2   \        VXP2                  P4                  4      '       d   \        V4      # \        V\(        P6                  \(        P8                  34      '       d   VP1                  4       # \;        \        V4      4      '       dF   \=        V4      w  rVVP?                  4        F"  w  rxV\@        PB                  PD                  V&   K$  	  V# \        V\         PF                  4      '       dH   \I        4       '       d8   VPJ                  e*   \        \         PF                  ! VP                  4      4      # \        V4      # )r   )Úhas_triton_packageNc                  ó(   € ] tR tRt$ R]R&   R tRtR# )Ú1PythonWrapperCodegen.val_to_arg_str.<locals>.Shimiî  r   Úrefc                	ó   € V P                   # ru   )rG  re  s   &rb   Ú__repr__Ú:PythonWrapperCodegen.val_to_arg_str.<locals>.Shim.__repr__ò  s   € ØŸ8™8�Ord   r`   N)r‘   r’   r“   r”   r•   rI  r–   r`   rd   rb   ÚShimrF  î  s   ‡ à“ö$rd   rK  c              3  ó\   <"  € T F!  pS! \         P                  SV4      4      x € K#  	  R # 5iru   r¥  )rx   r  rK  rf  s   & €€rb   rz   Ú6PythonWrapperCodegen.val_to_arg_str.<locals>.<genexpr>÷  s)   øé € ÐVÑTUÈq™Ô1×@Ñ@ÀÀqÓI×JÐJÓTUùs   ƒ),)&Útorch.utils._tritonrD  rî   r   r   rŒ  r˜   r]   rh   r   rƒ   r‚   r¥  Ú	dataclassr›  r…   r^  Ú_opsÚ
OpOverloadr    r,   rÁ   Ú
MutableBoxr/   r@  Úlanguager   r·  r¸  r   r   r�   rH   rœ   Úopaque_value_type_classesr�   r   r  )
rf  r@  Útype_rD  rî   Úobj_reprÚopaque_typesÚnr¤  rK  s
   f&&      @rb   r¦  Ú#PythonWrapperCodegen.val_to_arg_strâ  sÍ  ù€ Ý:á×ÒÛä�aœ×"Ò"Ü˜Ÿ™Ÿ™Ó%Ð%Ü˜œ5Ÿ:™:×&Ò&Ü˜“8ˆOÜ˜œE¤4˜=×)Ò)ä×"Ñ"÷$ð $ó #ð$ô Ü�Q”ÕVÑTUÓVÓVóð ô ˜œ5Ÿ:™:×0Ñ0×1Ò1Ü& qÓ)Ð)Ü˜œBŸI™I¤r§}¡}´oÐF×GÒGØ×&Ñ&Ó(Ð(Ù×!Ò!¤j°°F·O±O×4IÑ4I×&JÒ&JÜ˜“7ˆNÜ˜œB×-Ñ-¬r×/CÑ/CÐD×EÒEØ×&Ñ&Ó(Ð(Ü$¤T¨!£W×-Ò-Ü%8¸Ó%;Ñ"ˆHØ$×*Ñ*Ö,‘�Ø78”—‘×1Ñ1°!Ó4ñ -àˆOÜ˜œ5Ÿ<™<×(Ò(¬]¯_ª_ÀÇÁÒATô
 œŸš Q§V¡VÓ,Ó-Ð-ä˜“7ˆNrd   c               ó   € V ^8„  d   QhRR/# )r\   r|  r™   r`   )ra   s   "rb   rc   rM    s   € ÷ 	
ñ 	
¨Zñ 	
rd   c           	     	óZ  € VP                  4       pVP                  4       p\        VP                  4       4      p\        \        P
                  P                  V4      4      p\        VP                  4       4      pVP                  4       pV P                  VP                  4       W#WFWW4      # ru   )r=  r¤   rƒ   r7  rH   rœ   r  r8  Úget_is_pinnedrÚ  rž   )rf  r|  r�   r   r>  Úallocation_shaperÄ   Ú	is_pinneds   &&      rb   r2  Ú+PythonWrapperCodegen.make_buffer_allocation  sŽ   € Ø×"Ñ"Ó$ˆØ× Ñ Ó"ˆÜ�f—o‘oÓ'Ó(ˆÜ ¤§¡×!<Ñ!<¸VÓ!DÓEÐÜ�v×(Ñ(Ó*Ó+ˆØ×(Ñ(Ó*ˆ	Ø×#Ñ#Ø�O‰OÓ˜v¨eÐ=Mó
ð 	
rd   c                	óˆ   € R p\         P                  P                  '       g    V P                  P	                  VRR7       R# R# )zi
            from torch._inductor.runtime.debug_utils import check_memory_step, track_tensor
            Tr&  N)rH   rœ   r]  r
  r?  r‰  s   & rb   Ú"write_memory_track_allocation_onceÚ7PythonWrapperCodegen.write_memory_track_allocation_once  s7   € ðˆ
ô �w‰w×"×"Ð"Ø�L‰L×Ñ 
°$ÐÖ7ñ #rd   c                	ó*  € Vf   TpV P                  V4      pV P                  V4      p	V P                  V4      p
\        P                  P                  P                  P
                  '       d    V RV	 RV
 RV RVP                   RV R2pMpVP                  R8X  d   V'       d   V RV	 RV
 RV R	2pMHVP                  R9   d   V R
VP                   RV	 RV
 RV R	2
pMV RV	 RV
 RVP                   RV R	2
pW‰8w  d   VRV RV
 R	2,           pV# )Nz = tracked_empty_strided(rT  z, dtype=r  z	', name='rÚ  r%  z = empty_strided_cpu_pinned(r3  z = empty_strided_rS  z = empty_strided(r   z.as_strided()r%  r  r.  Úmtia)rí   r^  rc  r*   rd  re  r…   )rf  rÐ   r�   r   r>  rÄ   r]  r^  r;  Úcodegen_allocation_shape_tupleÚcodegen_stride_tupler›  s   &&&&&&&&    rb   rÚ  Ú$PythonWrapperCodegen.make_allocation#  sˆ  € ð Ò#Ø$Ðà"×=Ñ=¸eÓDÐØ)-×)HÑ)HØó*
Ð&ð  $×>Ñ>¸vÓFÐÜ�?‰?×!Ñ!×.Ñ.×E×EÐEà�&Ð1Ø1Ð2°"Ø'Ð(ð )Ø˜ð  Ø!Ÿ;™;˜-ð (Ø˜˜bð"ñ ð �[‰[˜EÔ!§ià�&Ð4Ø1Ð2°"Ø'Ð(¨Ø�'˜ðñ ð �[‰[Ð:Ô:ð �&Ð)¨&¯+©+¨°aØ1Ð2°"Ø'Ð(¨Ø�'˜ðñ ð �&Ð)Ø1Ð2°"Ø'Ð(ð )Ø!Ÿ;™;˜- y°°°qð:ð ð Ô@à˜,Ð':Ð&;¸2Ð>RÐ=SÐSTÐUÕUˆCØˆ
rd   c                	ó:   € V P                  \        V4      4       R # ru   )rö   rÕ  r›  s   &&rb   Úmake_commentÚ!PythonWrapperCodegen.make_commentS  s   € Ø�‰”{ 4Ó(Ö)rd   c           	     	ó`   € V P                    V R V V P                   RV P                   RV 2	# )r2  ú  Ú )r“  r”  r$  )rf  Únew_nameÚold_namer$  s   &&&&rb   Úmake_tensor_aliasÚ&PythonWrapperCodegen.make_tensor_aliasV  s6   € Ø—,‘,� ˜z¨¨X¨J°t·{±{°mÀ2ÀdÇlÁlÀ^ÐSTÐU\ÐT]Ð^Ð^rd   c               ó   € V ^8„  d   QhRR/# )r\   r|  r^  r`   )ra   s   "rb   rc   rM  Y  s   € ÷ *ñ *Ð'Fñ *rd   c                	ó(   € R VP                  4        2# )r!  )rž   rè  s   &&rb   rf  Ú%PythonWrapperCodegen.make_buffer_freeY  s   € Ø�f—o‘oÓ'Ð(Ð)Ð)rd   c               ó   € V ^8„  d   QhRR/# )r\   Únames_to_delrs  r`   )ra   s   "rb   rc   rM  \  s   € ÷ Añ A¨yñ Ard   c                	ó8   € R RP                  R V 4       4       2# )r!  rT  c              3  ó$   "  € T F  qx € K  	  R # 5iru   r`   )rx   rÐ   s   & rb   rz   Ú:PythonWrapperCodegen.make_free_by_names.<locals>.<genexpr>]  s   é € Ð>±¨¤³ùs   ‚)r  )rf  rv  s   &&rb   Úmake_free_by_namesÚ'PythonWrapperCodegen.make_free_by_names\  s   € Ø�d—i‘iÑ>±Ó>Ó>Ð?Ð@Ð@rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   ro  r‹   rn  Údel_liner`   )ra   s   "rb   rc   rM  _  s'   € ÷ uñ u°3ð uÀ#ð uÐQTñ urd   c           	     	ó`   € V P                    V R V V V P                   RV P                   R2	# )r2  rl  ú reuse)r#  r”  r$  )rf  ro  rn  r}  s   &&&&rb   Úcodegen_exact_buffer_reuseÚ/PythonWrapperCodegen.codegen_exact_buffer_reuse_  s@   € Ø×.Ñ.Ð/°¨z¸¸X¸JÀxÀjÐQU×Q\ÑQ\ÐP]Ð]_Ð`d×`lÑ`lÐ_mÐmsÐtÐtrd   c               ó   € V ^8„  d   QhRR/# )r\   Údebug_handler   r`   )ra   s   "rb   rc   rM  b  s   € ÷ ñ ð !ñrd   c                	óX   € Ve&   V P                  V P                   RV RV 24       R # R # )Nz [Provenance debug handles] r  )rö   r$  )rf  rC  rƒ  s   &&&rb   Úwrite_provenance_debug_handleÚ2PythonWrapperCodegen.write_provenance_debug_handleb  s4   € ð
 Ò#Ø�N‰NØ—<‘<�.Ð <¸[¸MÈÈ<È.ÐYöñ $rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÿ  r™   Únewrj  rE  r`   )ra   s   "rb   rc   rM  l  s&   € ÷ `ñ ` Zð `°jð `Èdñ `rd   c                	óâ  € VP                  4       VP                  4       8w  d.   \        R VP                  4        RVP                  4        24      hVP                  4       pVP                  4       pRpV\        P                  P                  4       9  d   V'       d   RV P                  V4       2pVP                  4       VP                  4       8X  d6   VP                  4       VP                  4       8X  d   V P                  WEV4      # V P                  WP                  4       VP                  4       ^ V P                  P                  4      pV P                   V RV V RV P                   R2# )zexpected matching dtypes, got r
  Ú;z; r2  rl  r  )r¤   r²   rž   rH   rœ   rì  rf  r7  r8  r€  rÛ  r  rö   r“  r$  )rf  rÿ  rˆ  rj  ro  rn  r}  Úreinterpret_views   &&&&    rb   rs  Ú&PythonWrapperCodegen.make_buffer_reusel  s*  € Ø�=‰=‹?˜cŸm™m›oÔ-Ü Ø0°·±³Ð0AÀÀsÇ}Á}ÃÐFWÐXóð ð —<‘<“>ˆØ—<‘<“>ˆØˆØœ1Ÿ7™7×3Ñ3Ó5Ô5¿*Ø˜D×1Ñ1°#Ó6Ð7Ð8ˆHà�<‰<‹>˜SŸ\™\›^Ô+°·±Ó0@ÀCÇNÁNÓDTÔ0TØ×2Ñ2°8ÀxÓPÐPà×8Ñ8Ø—‘“ §¡Ó!1°1°d×6GÑ6G×6QÑ6Qó
Ðð —,‘,� ˜z¨Ð-=Ð,>¸x¸jÈÈ4Ï<É<È.ÐX^Ð_Ð_rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rÐ   r‹   rb  zir.ReinterpretViewr_   r„  r`   )ra   s   "rb   rc   rM    s#   € ÷ 
ñ 
°ð 
Ð;Mð 
ÐRVñ 
rd   c                	ó®   € V P                  \        VV P                   V R VP                  4        V P                   RV P
                   R24      4       R# )r2  rl  z aliasN)rö   rK   r“  r@  r”  r$  )rf  rÐ   rb  s   &&&rb   rc  Ú0PythonWrapperCodegen.codegen_deferred_allocation  sS   € Ø�‰ÜØØ—<‘<�.   c¨$×*@Ñ*@Ó*BÐ)CÀDÇKÁKÀ=ÐPRÐSW×S_ÑS_ÐR`Ð`fÐgóö	
rd   c               ó   € V ^8„  d   QhRR/# ©r\   r|  ú	ir.Bufferr`   )ra   s   "rb   rc   rM  ‡  s   € ÷ 53ñ 53¨ñ 53rd   c                	ó&  a€ VP                  4       pV\        P                  P                  9   gP   W P                  9   g@   \        V\        P                  \        P                  \        P                  34      '       d   R # V P                  P                  V4       \        VP                  4       \        P                  \        P                  34      '       d   VP                  4       '       g   R # VP                  4       p\        V\        P                   4      '       d   R # \        V\        P"                  4      '       d   R # \        V\        P$                  4      '       EdR   \        VP&                  \        P(                  4      '       g/   \+        R\-        VP&                  4       RVP&                   24      hVP&                  P.                  p\        V\        P0                  4      '       g   \+        \-        V4      4      hVP.                  p\        V\        P2                  \        P(                  34      '       g   \+        \-        V4      4      h\        V\        P(                  4      '       d   R V3R lloS! V4      pV P5                  V4       V P7                  \9        WW4      4       R # \        V\        P:                  4      '       d   V P7                  \=        WRR7      4       R # V P7                  \=        W4      4       R # )Nzunexpected r4  c               ó   € V ^8„  d   QhRR/# )r\   r_   r’  r`   )ra   s   "rb   rc   Ú=PythonWrapperCodegen.codegen_allocation.<locals>.__annotate__ª  s   € ÷ "ñ "¨Iñ "rd   c                ó@  <€ \        V \        P                  4      '       d   S! V P                  4       4      # \        V \        P                  4      '       d   S! V P
                  4      # \        V \        P                  4      '       g   \        \        V 4      4      hV # ru   )	r   r,   ÚBaseViewÚunwrap_viewrR  rÉ   rÁ   r²   r…   )rö  Úunwrap_viewss   &€rb   r™  Ú=PythonWrapperCodegen.codegen_allocation.<locals>.unwrap_viewsª  sm   ø€ Ü! &¬"¯+©+×6Ò6Ù+¨F×,>Ñ,>Ó,@ÓAÐAÜ! &¬"¯-©-×8Ò8Ù+¨F¯K©KÓ8Ð8Ü% f¬b¯i©i×8Ò8Ü,¬T°&«\Ó:Ð:Ø!�Mrd   Tr$  )rž   rH   rœ   re  r4  r   r,   ÚDonatedBufferÚSubgraphBufferÚInputBufferr  Úget_defining_opÚExternKernelAllocÚMultiOutputÚshould_allocater°   ÚMutationLayoutSHOULDREMOVErŠ  ra  rb  r/   r²   r…   rÉ   rÇ   rÁ   r{  rö   rY  r±   r  )rf  r|  rÐ   rµ   ÚboxÚinput_bufferr™  s   &&    @rb   r{  Ú'PythonWrapperCodegen.codegen_allocation‡  s  ø€ Ø�‰Ó ˆð ”A—G‘G×+Ñ+Ô+Ø—~‘~Ô%Ü˜&¤2×#3Ñ#3´R×5FÑ5FÌÏÉÐ"W×XÒXáØ�‰×Ñ˜4Ô äØ×&Ñ&Ó(Ü×%Ñ%¤r§~¡~Ð6÷ò ð ×*Ñ*×,Ò,áà×'Ñ'Ó)ˆÜ�fœb×;Ñ;×<Ò<ÙÜ�fœbŸm™m×,Ò,ÙÜ�fœb×0Ñ0×1Ó1Ü˜fŸk™k¬2×+=Ñ+=×>Ò>Ü$ {´4¸¿¹Ó3DÐ2EÀRÈÏÉÀ}Ð%UÓVÐVØ—+‘+×"Ñ"ˆCÜ˜c¤2§=¡=×1Ò1Ü$¤T¨#£YÓ/Ð/ØŸ8™8ˆLÜ˜l¬R¯Y©Y¼×8JÑ8JÐ,K×LÒLÜ$¤T¨,Ó%7Ó8Ð8Ü˜,¬×(:Ñ(:×;Ò;÷"ð "ñ  ,¨LÓ9�Ø×#Ñ# LÔ1Ø�N‰Nœ?¨4¸vÓNÔOÙä�fœb×1Ñ1×2Ò2Ø�N‰Nœ<¨À$ÔGÔHÙà�‰”| DÓ1Ö2rd   c                	ó�  € VP                  4       p\        V\        P                  \        P                  34      '       dl   V P                  \        W4      4       W P                  9   d@   V P                  P                  V4       V P                  V P                  V R 2.4      4       R# \        VP                  4       \        P                  4      '       d   V P                  \        WRR7      4       R# V P                  V4      '       g   R# V P                  P                  V4       V P                  \        W4      4       R# )Ú_origNTr$  )rž   r   r,   r�  r¶  rö   r]  r  r  rz  r°   r±   r}  Ú	can_reuser5  r  )rf  r|  rÐ   s   && rb   Úcodegen_freeÚ!PythonWrapperCodegen.codegen_free¾  së   € Ø�‰Ó ˆô �fœrŸ~™~¬r×/AÑ/AÐB×CÒCØ�N‰Nœ8 DÓ1Ô2ð
 ×1Ñ1Ô1Ø×*Ñ*×2Ñ2°4Ô8Ø—‘˜t×6Ñ6¸4¸&À¸Ð7GÓHÔIÙä�f×,Ñ,Ó.´×0CÑ0C×DÒDð �N‰NÔ.¨tÈÔNÔOÙà�~‰~˜f×%Ò%ÙØ�
‰
�‰�tÔà�‰Ô*¨4Ó8Ö9rd   c                	óF  € VP                  4       pV\        P                  P                  9   ;'       gé    V\        P                  P                  9   ;'       dG    \        \        P                  P                  V,          \        P                  4      '       * ;'       g~    V\        P                  P                  9   ;'       gY    V\        P                  P                  9   ;'       g4    V\        P                  P                  9   ;'       g    W0P                  9   '       * # ru   )rž   rH   rœ   re  rê  r   Úgraph_inputs_originalr,   r›  r2  r3  Únever_reuse_buffersr5  )rf  r¤  Úoutput_bufferrÐ   s   &&& rb   r¨  ÚPythonWrapperCodegen.can_reuseÙ  sé   € Ø×$Ñ$Ó&ˆà”A—G‘G×+Ñ+Ñ+÷ 
"ð 
"àœŸ™×,Ñ,Ñ,÷ ð Ü"Ü—G‘G×1Ñ1°$Õ7¼×9IÑ9Ióô ÷
"ð 
"ð ”q—w‘w×(Ñ(Ñ(÷
"ð 
"ð ”q—w‘w×2Ñ2Ñ2÷
"ð 
"ð ”q—w‘w×2Ñ2Ñ2÷
"ð 
"ð —z‘zÑ!ô
ð 	
rd   c                	ó°   € VP                  4       V P                  9   ;'       d3    V P                  VP                  4       ,          VP                  4       8H  # ru   )rž   r6  )rf  r|  Úreused_buffers   &&&rb   Ú	did_reuseÚPythonWrapperCodegen.did_reuseé  sI   € ð �O‰OÓ §¡Ñ,÷ Kð KØ—‘˜FŸO™OÓ-Õ.°-×2HÑ2HÓ2JÑJð	
rd   c               ó    € V ^8„  d   QhRRRR/# )r\   r¤  r’  r®  r`   )ra   s   "rb   rc   rM  ñ  s   € ÷ 	Eñ 	E°)ð 	EÈIñ 	Erd   c                	ó�  € \        W4      '       g   \        R 4      hV P                  V4       V P                  P	                  VP                  4       4       V P                  P	                  VP                  4       4       VP                  4       V P                  VP                  4       &   V P                  \        WV4      4       R# )z6expected input_buffer and output_buffer sizes to matchN)
r¿   r²   r{  r5  r  rž   r4  r6  rö   r(  )rf  r¤  r®  s   &&&rb   Úcodegen_inplace_reuseÚ*PythonWrapperCodegen.codegen_inplace_reuseñ  s”   € Ü$ \×AÒAÜ ØHóð ð 	×Ñ Ô-Ø�
‰
�‰�|×,Ñ,Ó.Ô/Ø�‰×Ñ˜=×1Ñ1Ó3Ô4Ø0<×0EÑ0EÓ0Gˆ�‰�M×*Ñ*Ó,Ñ-Ø�‰”y °]ÓCÖDrd   c                	ó˜   € \        V4      pW P                  9   d   V# V P                  P                  V4       V P                  V,           # ru   )r‹   r)  r  r“  )rf  rJ  rÐ   s   && rb   Úcodegen_unbacked_symbol_declÚ1PythonWrapperCodegen.codegen_unbacked_symbol_declü  sA   € Ü�6‹{ˆØ×-Ñ-Ô-ØˆKð ×&Ñ&×*Ñ*¨4Ô0Ø—<‘< $Õ&Ð&rd   c               ó(   € V ^8„  d   QhRRRRRRRR/# ©	r\   rÏ  r‹   rÐ  r   rÒ  rÑ  r_   r„  r`   )ra   s   "rb   rc   rM    s3   € ÷ 
ñ 
àð
ð ð
ð Eð	
ð
 
ñ
rd   c                	ó˜   € \        \        P                  P                  P                  V4      pV P                  \        WW#4      4       R # ru   )r   rH   rœ   r¥   rT  rö   rÎ  )rf  rÏ  rÐ  rÒ  s   &&&&rb   Ú(codegen_unbacked_symbol_defs_for_outputsÚ=PythonWrapperCodegen.codegen_unbacked_symbol_defs_for_outputs  s=   € ô 6Ü�G‰G×Ñ×&Ñ&Ð(9ó
Ðð 	�‰Ü" 4°gÓQö	
rd   c               ó(   € V ^8„  d   QhRRRRRRRR/# r¼  r`   )ra   s   "rb   rc   rM    s9   € ÷ Nñ NàðNð ðNð Eð	Nð
 
ñNrd   c                	óæ   aaaa€ V'       g   R # VP                  4        FN  w  poR V3R lloVVVV3R lpV P                  V P                  V4       RV! 4        V P                   24       KP  	  R # )Nc               ó    € V ^8„  d   QhRRRR/# )r\   r]   r‹   r  zpytree.KeyPathr`   )ra   s   "rb   rc   ÚTPythonWrapperCodegen._codegen_unbacked_symbol_defs_for_outputs.<locals>.__annotate__%  s   € ÷ Lñ L˜ð L ~ñ Lrd   c                ó  <€ VR8X  d   V # \        V4      ^8¼  d€   \        V^ ,          \        4      '       dc   \        V^,          \        P                  4      '       d<   S! V  RV^ ,          P
                   RV^,          P                   R2VR,          4      # \        V^ ,          \        4      '       d(   S! V  RV^ ,          P
                   R2VR,          4      # \        V^ ,          \        P                  4      '       dp   \        P                  P                  '       d)   S! RV^ ,          P                   RV  R2VR,          4      # S! V  R	V^ ,          P                   R
2VR,          4      # \        V^ ,          \        4      '       dp   \        P                  P                  '       d)   S! RV  RV^ ,          P                   R2VR,          4      # S! V  RV^ ,          P                   R2VR,          4      # \        RV 24      h)r\   r5  rS  r3  rØ  r½  r‹  z	std::get<z>(r‰  rŠ  z / z.__floordiv__(r  r`   )r  r   r   ÚpytreeÚSequenceKeyrÐ   r  rH   rœ   r]  r   rü  r²   )r]   r  Úgos   &&€rb   rÇ  ÚJPythonWrapperCodegen._codegen_unbacked_symbol_defs_for_outputs.<locals>.go%  sÆ  ø€ Ø˜b”=Ø�Kô ˜“L AÔ%Ü" 7¨1¥:¬}×=Ò=Ü" 7¨1¥:¬v×/AÑ/A×BÒBáØ˜&  '¨!¥*§/¡/Ð!2°!°G¸AµJ·N±NÐ3CÀ1ÐEÀwÈrÅ{óð ô   ¨¥
¬M×:Ò:Ù   a¨°­
¯©Ð'8¸Ð;¸WÀR½[ÓIÐIÜ ¨¥
¬F×,>Ñ,>×?Ò?ô Ÿ7™7×.×.Ð.ñ ˜Y w¨q¥z§~¡~Ð&6°b¸¸¸aÐ@À'È"Å+ÓNðñ   4 &¨¨'°!­*¯.©.Ð)9¸Ð ;¸WÀR½[ÓIðô
   ¨¥
¬K×8Ò8ä—w‘w×*×*Ð*Ù! A d V¨3¨w°q­z×/AÑ/AÐ.BÀ!Ð"DÀgÈbÅkÓRÐRÙ   n°W¸QµZ×5GÑ5GÐ4HÈÐJÈGÐTVÍKÓXÐXä(Ð+@ÀÀ	Ð)JÓKÐKrd   c                 óH  <€ \         P                  P                  '       dú   \        S4      ^8X  do   S^ ,          p S! S^ ,          P	                  4       \        V \        P                  4      '       d(   \        V P                  4      ^ 8w  d   SR,          4      # S4      # \        S^ ,          \        P                  4      '       g   \        R\        S^ ,          4       24      hS! SS^ ,          P                  ,          P	                  4       SR,          4      # S! SS4      # )rg   r‹  zexpected SequenceKey, got )rH   rœ   r]  r  rž   r   r,   r   r„  rÅ  rÆ  r²   r…   r  )r›  rÇ  r  rÏ  rÐ  s    €€€€rb   Úgo_outerÚPPythonWrapperCodegen._codegen_unbacked_symbol_defs_for_outputs.<locals>.go_outerD  só   ø€ Ü—7‘7×&×&Ð&ô
 ˜7“| qÔ(Ø% a�j˜ñ  "Ø# A�J×/Ñ/Ó1ä)¨#¬r¯~©~×>Ò>Ä3ÀsÇ{Á{ÓCSÐWXÔCXð $ B�Kó ð ð ")ó	 ð ô  *¨'°!­*´f×6HÑ6H×IÒIÜ"0Ø"<¼TÀ'È!Å*Ó=MÐ<NÐ Oó#ð ñ  " '¨'°!­*¯.©.Õ"9×"BÑ"BÓ"DÀgÈbÅkÓRÐRá˜k¨7Ó3Ð3rd   r2  )r�   rö   r¹  r”  )rf  rÏ  rÐ  rÒ  r@  rÊ  rÇ  r  s   &ff&  @@rb   r×  Ú>PythonWrapperCodegen._codegen_unbacked_symbol_defs_for_outputs  sm   û€ ÷ !Ùð ,×1Ñ1Ö3‰JˆAˆw÷
Lð L÷>4ð 4ð4 �N‰NØ×4Ñ4°QÓ7Ð8¸¹H»J¸<ÈÏÉÀ}ÐUöó} 4rd   c           	     	óÌ  a aaa€ VV V3R  lpVV V3R lp S P                  SP                  4       \        SRR4      pV'       d   \        P                  ! V4      M\
        P                  ! 4       pT;_uu_ 4        S P                  S P                   RSP                   24       V! 4        \        P                  p\        P                  ! SP                  4      ;_uu_ 4        SP                  P                  VR7       RRR4       V! 4        RRR4       S P                  4        R#   + '       g   i     L1; i  + '       g   i     L5; i  S P                  4        i ; i)c                 ó~  <€ \        SP                  P                  4      \        S4      8w  d8   \        R \        SP                  P                  4       R\        S4       24      h\	        SP                  P                  S4       F3  w  rSP                  SP                   V  RV SP                   24       K5  	  R# )z$expected matching input counts, got rþ   r2  N)r  rœ   rê  r²   r  rö   r“  r”  )Úinner_inputÚouter_inputÚouter_inputsrf  Úsubgraphs     €€€rb   Ú_codegen_subgraph_prefixÚSPythonWrapperCodegen.codegen_subgraph_by_inlining.<locals>._codegen_subgraph_prefixp  s¤   ø€ Ü�8—>‘>×.Ñ.Ó/´3°|Ó3DÔDÜ$Ø:¼3¸x¿~¹~×?ZÑ?ZÓ;[Ð:\Ð\`ÔadÐeqÓarÐ`sÐtóð ô -0Ø—‘×+Ñ+¨\ö-Ñ(�ð —‘Ø—|‘|�n [ M°°[°MÀ$Ç+Á+ÀÐOöó-rd   c                 ó‚  <€ \        SP                  P                  4      \        S4      8w  d8   \        R \        SP                  P                  4       R\        S4       24      h\	        SP                  P                  S4       F5  w  rSP                  V RV P                  4        SP                   24       K7  	  R# )z%expected matching output counts, got rþ   r2  N)r  rœ   rò  r²   r  rö   r@  r”  )Úinner_outputÚouter_outputÚouter_outputsrf  rÒ  s     €€€rb   Ú_codegen_subgraph_suffixÚSPythonWrapperCodegen.codegen_subgraph_by_inlining.<locals>._codegen_subgraph_suffix|  s¦   ø€ Ü�8—>‘>×/Ñ/Ó0´C¸Ó4FÔFÜ$Ø;¼CÀÇÁ×@\Ñ@\Ó<]Ð;^Ð^bÔcfÐgtÓcuÐbvÐwóð ô /2Ø—‘×,Ñ,¨mö/Ñ*�ð —‘Ø#�n C¨×(FÑ(FÓ(HÐ'IÈ$Ï+É+ÈÐWöó/rd   Úinductor_config_patchesNú subgraph: )Úparent_graph)r¾  rœ   rw   r*   Úpatchr  r  rö   r$  rÐ   rH   rC  Úcodegen_subgraphrò  )	rf  rÒ  rÑ  rØ  rÓ  rÙ  rÛ  rA  rÝ  s	   ffff     rb   Úcodegen_subgraph_by_inliningÚ1PythonWrapperCodegen.codegen_subgraph_by_inliningb  s   û€ ÷
	÷
	ð	'Ø×%Ñ% h§n¡nÔ5ô '.¨hÐ8QÐSWÓ&XÐ#÷ +ô —’Ð4Ô5ä×+Ò+Ó-ð ÷
 ’Ø—‘ $§,¡, ¨{¸8¿=¹=¸/ÐJÔKÙ(Ô*Ü Ÿw™w�Ü×(Ò(¨¯©×8Õ8Ø—N‘N×3Ñ3Ø%1ð 4ô ÷ 9ñ )Ô*÷ ð ×$Ñ$Ö&÷ 9×8ú÷	 —ûð ×$Ñ$Õ&úsB   –/E Á5E Á;A*D>Ã%D+	ÄD>ÄE Ä+D;Ä6D>Ä>E	Å	E ÅE#c               ó    € V ^8„  d   QhRRRR/# )r\   Úpartition_idr�   rK  zir.GraphPartitionSignaturer`   )ra   s   "rb   rc   rM  Ÿ  s   € ÷ <ñ <àð<ð 9ñ<rd   c           	     ó   € VP                   pVP                  p\        VP                  4       4      VP                   Uu. uF  qUP
                  NK  	  up,           pRP                  V4      \        V4      ^8X  d   RMR,           pV Uu. uF  qˆP                  4       NK  	  p	pRP                  V	4      \        V4      ^8X  d   RMR,           p
V P                  RV RV R24       VP                  4        UUu. uF  w  r¼V'       g   K  VNK  	  pppV'       d$   V P                  RRP                  V4       24       V P                  RV
 R	V R
V R24       V P                  RV R24       R# u upi u upi u uppi )z'Generate code to call a graph partitionrT  rÈ  r  Ú	partitionú	_args = [rŠ  r!  rS  z) = self.partitions[z](partitionú_args)zdel partitionÚ_argsN)Úinput_deallocationÚoutput_nodesr‚   rë  Úsymbol_inputsrÐ   r  r  rž   rö   r�   )rf  rã  rK  ré  rê  Úsymbol_inputrÆ  r�   r˜   Úoutput_namesrÐ  rÐ   Ú
deallocaterv  s   &&&           rb   Úcodegen_partition_callÚ+PythonWrapperCodegen.codegen_partition_callŸ  sp  € ð 2×DÑDÐØ+×8Ñ8ˆäÐ-×2Ñ2Ó4Ó5Ø2F×2TÒ2Tó9
Ù2T ,×ÔÑ2Tñ9
õ 
ˆð —‘˜;Ó'´#°kÓ2BÀaÔ2G©3ÈRÕPˆá4@ÓA±L¨DŸ™ž±LˆÐAØ—)‘)˜LÓ)´C¸Ó4EÈÔ4J©SÐPRÕSˆð 	�‰˜ < .°	¸&¸ÀÐCÔDð *<×)AÑ)AÔ)Cô
Ù)CÑ%�TÄz�DˆDÑ)Cð 	ñ 
÷ Ø�N‰N˜T $§)¡)¨LÓ"9Ð!:Ð;Ô<ð 	�‰Ø�ˆyÐ,¨\¨N¸+ÀlÀ^ÐSYÐZô	
ð 	�‰˜ | n°EÐ:Ö;ùò-9
ùò Bùó
s   ¿E0Â
E5Ã9E:Ä
E:c               ó    € V ^8„  d   QhRRRR/# )r\   rã  r�   r_   r‹   r`   )ra   s   "rb   rc   rM  À  s   € ÷ +ñ +¨sð +°sñ +rd   c                	ó   € R V 2# )Ú
partition_r`   )rf  rã  s   &&rb   Úget_partition_nameÚ'PythonWrapperCodegen.get_partition_nameÀ  s   € Ø˜L˜>Ð*Ð*rd   c               ó   € V ^8„  d   QhRR/# )r\   Únum_partitionsr�   r`   )ra   s   "rb   rc   rM  Ã  s   € ÷ 
ñ 
°cñ 
rd   c                	ó`   € \        \        V P                  \        V4      4      4      V n        R # ru   )r‚   rÊ  rô  r¸  rk  )rf  r÷  s   &&rb   Úset_all_partition_namesÚ,PythonWrapperCodegen.set_all_partition_namesÃ  s$   € Ü#'Ü�×'Ñ'¬¨~Ó)>Ó?ó$
ˆÖ rd   c           	     	óŠ  € R P                  V4      \        V4      ^8X  d   RMR,           pR P                  V4      \        V4      ^8X  d   RMR,           pV P                  VP                  P                   RV R24       V P                  RV RVP                  P                   RVP                  P                   R24       R# )	rT  rÈ  r  ræ  rŠ  rS  z) = rç  N)r  r  rö   rœ   rÐ   )rf  rÒ  rÑ  Úouter_flattened_outputsÚouter_output_namesÚouter_input_namess   &&&&  rb   Ú,codegen_subgraph_call_with_flattened_outputsÚAPythonWrapperCodegen.codegen_subgraph_call_with_flattened_outputsÈ  s¹   € ð "ŸY™YÐ'>Ó?ÜÐ.Ó/°1Ô4‰C¸"õ
Ðð !ŸI™I lÓ3Ü�|Ó$¨Ô)‰C¨rõ
Ðð 	�‰˜(Ÿ.™.×-Ñ-Ð.¨iÐ8IÐ7JÈ!ÐLÔMð 	�‰ØÐ"Ð# 4¨¯©×(;Ñ(;Ð'<¸A¸h¿n¹n×>QÑ>QÐ=RÐRXÐYö	
rd   c                	ó„  € R P                  V4      \        V4      ^8X  d   RMR,           pV P                  VP                  P                   RV R24       \
        P                  P                  P                  4        V P                  V RVP                  P                   RVP                  P                   R24       R# )	rT  rÈ  r  ræ  rŠ  r2  rS  rç  N)r  r  rö   rœ   rÐ   rH   r    Úfree_buffers)rf  rÒ  rÑ  Úouter_buffer_namerþ  s   &&&& rb   Úcodegen_subgraph_callÚ*PythonWrapperCodegen.codegen_subgraph_callÚ  s¢   € à ŸI™I lÓ3Ü�|Ó$¨Ô)‰C¨rõ
Ðð 	�‰˜(Ÿ.™.×-Ñ-Ð.¨iÐ8IÐ7JÈ!ÐLÔMô 	
�‰×Ñ×&Ñ&Ô(ð 	�‰Ø Ð!  X§^¡^×%8Ñ%8Ð$9¸¸8¿>¹>×;NÑ;NÐ:OÈvÐVö	
rd   c                	ó  € V P                  VP                  4       V P                  R 4       V P                  V P                   RVP                   24       \
        P                  pVP                  VP                  n        VP                  VP                  n        VP                  P                  V P                  9  Ed   \        VRR4      pV'       d   \        P                  ! V4      M\        P                  ! 4       pT;_uu_ 4        \        P                  ! RR4      ;_uu_ 4        \
        P                  ! VP                  4      ;_uu_ 4        VP                  P                  4       w  rVRRR4       RRR4       RRR4       VP                  P                  pV P                  P!                  V4       V P#                  VX4       R# R#   + '       g   i     Lg; i  + '       g   i     Lr; i  + '       g   i     L}; i)r  rÜ  rÛ  NrÏ  F)r¾  rœ   ri  r$  rÐ   rH   r]  Ú
fx_wrapperr=  rw   r*   rÞ  r  r  rC  r¥  r  rY  )rf  rÒ  rÝ  rÛ  rA  rX  r  rI  s   &&      rb   Úcodegen_subgraph_commonÚ,PythonWrapperCodegen.codegen_subgraph_commonë  si  € Ø×!Ñ! (§.¡.Ô1Ø×Ñ˜"ÔØ×Ñ˜TŸ\™\˜N¨+°h·m±m°_ÐEÔFä—w‘wˆØ%1×%=Ñ%=ˆ�‰Ô"Ø$0×$;Ñ$;ˆ�‰Ô!à�>‰>×Ñ d×&FÑ&FÕFô '.¨hÐ8QÐSWÓ&XÐ#÷ +ô —’Ð4Ô5ä×+Ò+Ó-ð ÷ ’Ü—’Ð.°×6Õ6Ü×#Ò# H§N¡N×3Õ3ð $,§>¡>×#9Ñ#9Ó#;Ñ �÷ 4÷ 7÷ ð %ŸN™N×/Ñ/ˆMØ×,Ñ,×0Ñ0°Ô?Ø×,Ñ,¨]¸MÖJñ- G÷  4×3ú÷ 7×6ú÷ —ús<   Ä G8Ä-)G%	ÅGÅ3G%	Å;G8ÇG"ÇG%	Ç%G5Ç0G8Ç8H	c                	óL   € V P                  V4       V P                  WV4       R # ru   )r	  rÿ  )rf  rÒ  rÑ  rü  s   &&&&rb   Ú'codegen_subgraph_with_flattened_outputsÚ<PythonWrapperCodegen.codegen_subgraph_with_flattened_outputs  s&   € ð 	×$Ñ$ XÔ.Ø×9Ñ9ØÐ$;ö	
rd   c                	óL   € V P                  V4       V P                  WV4       R # ru   )r	  r	  )rf  rÒ  rÑ  r	  s   &&&&rb   rß  Ú%PythonWrapperCodegen.codegen_subgraph  s#   € ð 	×$Ñ$ XÔ.Ø×"Ñ" 8Ð;LÖMrd   c                	óô  € VP                  4       pV P                  V R \        VP                  4       24       VP                   Uu. uF  q3P                  4       NK  	  pp\        P                  P                  '       dQ   \        \        VP                  4      4       Uu. uF
  qR RV R2NK  	  ppV P                  VP                  WF4       R# V P                  VP                  WB4       R# u upi u upi )ú = [None] * r‰  rŠ  N)rž   rö   r  rÐ  r�   r@  rH   rœ   r  r¸  rà  rÒ  rß  )rf  Úinvoke_subgraphrÐ   rº  rÑ  r¹  rØ  s   &&     rb   Úcodegen_invoke_subgraphÚ,PythonWrapperCodegen.codegen_invoke_subgraph  sÛ   € Ø×'Ñ'Ó)ˆà�‰˜$˜˜|¬C°×0GÑ0GÓ,HÐ+IÐJÔKØ;J×;QÒ;QÓRÑ;Q°C×-Ñ-Ö/Ñ;QˆÐRä�7‰7××Ðä(-¬c°/×2IÑ2IÓ.JÔ(KóÙ(K 1�&˜˜!˜˜A“Ñ(Kð ð ð ×-Ñ-Ø×(Ñ(¨,öð ×!Ñ! /×":Ñ":¸LÖOùò Sùòs   ÁC0Â"C5c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rM  *  s   € ÷ %5ñ %5 dñ %5rd   c                	ó®  a aaa€ SP                  4       oSP                   Uu. uF  q"P                  4       NK  	  upoSP                  P                  4       pS R 2p\	        SP                  \
        P                  4      '       g   S P                  V RV R24       MS P                  V RV R24       S P                  S R\        SP                  4       24       R VVVV 3R llp\        SP                  4      p\        SP                  4       F7  w  rxV^ 8X  d	   RRV R	2r©MWv^,
          8  d   R
RV RV 2r©MRRr©V! WšV4       K9  	  R# u upi )Ú	_selectorz = int(z.item())r3  r	  c               ó$   € V ^8„  d   QhRRRRRR/# )r\   Úkeywordr‹   Ú	conditionr_   r„  r`   )ra   s   "rb   rc   Ú9PythonWrapperCodegen.codegen_switch.<locals>.__annotate__=  s!   € ÷ 	3ñ 	3 #ð 	3°#ð 	3À$ñ 	3rd   c                ó°  <€ SP                  V  V R 24       SP                  \        SVP                  4      4       \        P                  P                  '       dH   \        \        SP                  4      4       Uu. uF  pS RV R2NK  	  ppSP                  VSV4       MSP                  VSS4       SP                  \        S4      4       R# u upi )r  r‰  rŠ  N)rö   r±  rœ   rH   r  r¸  r  rÐ  rà  rß  rê  )	r	  r	  Úbranchr¹  rØ  rÐ   r˜   rÑ  rf  s	   &&&  €€€€rb   Ú_emit_branchÚ9PythonWrapperCodegen.codegen_switch.<locals>._emit_branch=  s®   ø€ Ø�N‰N˜g˜Y y k°Ð3Ô4Ø�N‰NÔ,¨T°6·<±<Ó@ÔAÜ�w‰w××ÐÜ9>¼sÀ4Ç<Á<Ó?PÔ9QÓ RÑ9Q°A D 6¨¨1¨#¨Q£Ñ9Q�Ð RØ×1Ñ1°&¸,ÈÕVà×%Ñ% f¨l¸DÔAØ�N‰NÔ+¨DÓ1Ö2ùò	 !Ss   Á=CÚifrm  z == 0Úelifz == Úelser  N)rž   Úoperandsr@  Úselectorr   r,   ÚShapeAsConstantBufferrö   r  rÐ  ÚbranchesrÚ  )rf  r˜   rº  Úselector_refÚselector_varr	  Únum_branchesÚb_idxr	  r	  r	  rÐ   rÑ  s   ff         @@rb   Úcodegen_switchÚ#PythonWrapperCodegen.codegen_switch*  s5  û€ Ø�}‰}‹ˆà;?¿=º=ÓI¹=°C×-Ñ-Ö/¹=ÑIˆà—}‘}×6Ñ6Ó8ˆà˜˜yÐ)ˆÜ˜$Ÿ-™-¬×)AÑ)A×BÒBð �N‰N˜l˜^¨7°<°.ÀÐIÕJð �N‰N˜l˜^¨7°<°.ÀÐBÔCà�‰˜$˜˜|¬C°·±Ó,=Ð+>Ð?Ô@÷	3ó 	3ô ˜4Ÿ=™=Ó)ˆÜ& t§}¡}Ö5‰MˆEØ˜ŒzØ%)¨Q¨|¨n¸EÐ+B™Ø¨Õ)Ô)Ø%+¨q°°¸dÀ5À'Ð-J™à%+¨R˜Ù˜¨VÖ4ó 6ùò7 Js   £Ec                óZ	  a € V 3R lpVP                  4       pVP                   Uu. uF  qUP                  4       NK  	  ppVP                   Uu. uF  qUP                  4       NK  	  pp\	        V4      pS P                  V R\	        V4       24       V'       d!   S P                  V R\	        V4       R24       \        V4       F  w  ršS P                  V RV	 RV
 24       K   	  . \        \	        V4      4       U	u. uF
  q” RV	 R2NK  	  up	OVOpV R2.p\        V4      pVR	\	        V4       pV! VP                  W¼4       S P                  R
V^ ,           24       S P                  R4       V'       dy   \        V4       Fh  w  rŸS P                  \        S VP                  P                  4      4       S P                  V RV	 RV R24       S P                  \        S 4      4       Kj  	  Mw\        V4       Fh  w  rŸS P                  \        S VP                  P                  4      4       S P                  V RV	 RV R24       S P                  \        S 4      4       Kj  	  S P                  R4       S P                  \        S VP                  P                  4      4       V! VP                  WÞ4       S P                  \        S 4      4       V'       dˆ   S P                  \        S VP                  P                  4      4       \        \	        V4      4       F&  p	S P                  V RW˜,            RV RV	 R24       K(  	  S P                  \        S 4      4       S P                  \        S VP                  P                  4      4       V! VP                  W¼4       S P                  \        S 4      4       S P                  RV^ ,           24       V'       d¹   S P                  R4       \        \	        V4      4       F�  p	S P                  RV RW˜,            R24       S P                  \        S VP                  P                  4      4       S P                  V RV	 RV RW˜,            R24       S P                  \        S 4      4       K�  	  R	# R	# u upi u upi u up	i )z1while_loop is codegened as a host side while_loopc                ó”   <€ \         P                  P                  '       d   SP                  WV4       R# SP	                  WV4       R# )z3Helper method to deduplicate subgraph codegen logicN)rH   rœ   r  rà  r	  )rÒ  rÑ  rØ  rf  s   &&&€rb   rß  ÚAPythonWrapperCodegen.codegen_while_loop.<locals>.codegen_subgraphT  s6   ø€ ä�w‰w××ÐØ×1Ñ1°(È-ÖXà×<Ñ<Ø¨Mörd   r	  z.extend([[] for _ in range(z)])r‰  z] = rŠ  Ú_cond_resultNzshould_loop = zif not should_loop:z.unsqueeze(0).clone()r‡  zwhile should_loop:z	].append(z])z    should_loop = z%# Stack outputs after loop completionzif len(z]) > 0:z] = torch.stack(z	], dim=0))rž   Úcarried_inputsr@  Úadditional_inputsr  rö   rÚ  r¸  r‚   Úcond_subgraphr±  Úbody_subgraphrœ   rê  )rf  Ú
while_loopÚstack_outputrß  rÐ   rº  Úouter_carried_inputsÚouter_additional_inputsÚ
ckp_offsetr¹  ÚinpÚcond_outer_inputsÚcond_outer_outputsÚbody_outer_inputsÚbody_outer_outputsÚcarried_inputs   f&&             rb   Úcodegen_while_loopÚ'PythonWrapperCodegen.codegen_while_loopQ  s]  ø€ õ	ð ×"Ñ"Ó$ˆà/9×/HÒ/Hó 
Ù/H¨×!Ñ!Ö#Ñ/Hð 	ð  
ð 0:×/KÒ/Kó#
Ù/K¨×!Ñ!Ö#Ñ/Kð 	 ð #
ô Ð-Ó.ˆ
Ø�‰˜$˜˜|¬CÐ0DÓ,EÐ+FÐGÔHßØ�N‰NØ�&Ð3´CÐ8LÓ4MÐ3NÈcÐRôô  Ð 4Ö5‰FˆAà�N‰N˜d˜V 1 Q C t¨C¨5Ð1Ö2ñ 6ð
Ü&+¬CÐ0DÓ,EÔ&FÓGÑ&F ��a˜�s˜!‹nÑ&FÑGð
à$ð
Ðð "&  lÐ3Ð4ÐÜ Øó
Ðð /Ð/J´Ð5IÓ1JÐKÐáØ×$Ñ$Ð&7ô	
ð 	�‰˜Ð(:¸1Õ(=Ð'>Ð?Ô@Ø�‰Ð,Ô-ßä$-Ð.BÖ$CÑ �Ø—‘Ô0°°z×7OÑ7O×7UÑ7UÓVÔWØ—‘ $  q¨¨¨4°¨Ð>SÐTÔUØ—‘Ô/°Ó5Ö6ò %Dô
 %.Ð.BÖ$CÑ �Ø—‘Ô0°°z×7OÑ7O×7UÑ7UÓVÔWØ—‘ $  q¨¨¨4°¨¸hÐGÔHØ—‘Ô/°Ó5Ö6ñ %Dð
 	�‰Ð+Ô,à�‰Ô(¨¨z×/GÑ/G×/MÑ/MÓNÔOÙØ×$Ñ$Ð&7ô	
ð 	�‰Ô'¨Ó-Ô.÷ Ø�N‰NÔ,¨T°:×3KÑ3K×3QÑ3QÓRÔSÜœ3Ð3Ó4Ö5�Ø—‘ $  q¨­Ð(8¸	À$ÀÀqÈÈÈ2ÐNÖOñ 6à�N‰NÔ+¨DÓ1Ô2ð 	�‰Ô(¨¨z×/GÑ/G×/MÑ/MÓNÔOÙØ×$Ñ$Ð&7ô	
ð 	�‰Ô'¨Ó-Ô.Ø�‰Ð+Ð,>¸qÕ,AÐ+BÐCÔD÷ Ø�N‰NÐBÔCÜœ3Ð3Ó4Ö5�Ø—‘ ¨¨¨a°µÐ/?¸wÐGÔHØ—‘Ô0°°z×7OÑ7O×7UÑ7UÓVÔWØ—‘Ø�f˜A˜a˜SÐ 0°°°a¸½Ð7GÀyÐQôð —‘Ô/°Ó5Ö6ó 6ñ ùòQ 
ùò#
ùò  Hs   ¦RÁR#Ã>R(c                	óö   €  \        V R R4      '       d   R# \        V \        4      '       d   V # \        P                  P
                  P                  V 4      pVf   V# \        V4      #   \         d     R# i ; i)r‰  N)rw   r   r�   rH   rœ   Ú
_shape_envÚ_maybe_evaluate_staticr3  )r  r§  s   & rb   Ústatically_known_int_or_noneÚ1PythonWrapperCodegen.statically_known_int_or_none°  sl   € ð	Ü�q˜.¨$×/Ò/ñ Ü˜!œS×!Ò!Ø�Ü—'‘'×$Ñ$×;Ñ;¸AÓ>ˆCØŠ{Ø�
Ü�s“8ˆOøÜô 	Úð	ús!   ‚A) —A) ¯.A) Á
A) Á)A8Á7A8c                	óv   € . pV  F0  p\         P                  V4      pVf    R # VP                  V4       K2  	  V# ru   )r²  rD	  rÈ   )Úlstrd  r  r›  s   &   rb   Ú%statically_known_list_of_ints_or_noneÚ:PythonWrapperCodegen.statically_known_list_of_ints_or_noneÀ  s<   € àˆÛˆAÜ&×CÑCÀAÓFˆCØŠ{ÚØ�M‰M˜#Öñ	 ð
 ˆrd   c                	ó0   € \         P                  V 4      R J# ru   )r²  rH	  )rG	  s   &rb   Ú is_statically_known_list_of_intsÚ5PythonWrapperCodegen.is_statically_known_list_of_intsÊ  s   € ô !×FÑFÀsÓKÐSWÐWð	
rd   c                	óH   € \         P                  V P                  4       4      # ru   )r²  rH	  r7  ©r|  s   &rb   r)  Ú4PythonWrapperCodegen.static_shape_for_buffer_or_noneÐ  s   € ä#×IÑIØ�O‰OÓó
ð 	
rd   c                	ó0   € \         P                  V 4      R J# ru   )r²  r)  rN	  s   &rb   Ú!can_prove_buffer_has_static_shapeÚ6PythonWrapperCodegen.can_prove_buffer_has_static_shapeÖ  s   € ä#×CÑCÀFÓKÐSWÐWÐWrd   c               ó    € V ^8„  d   QhRRRR/# )r\   rC  r‹   Únode_schedulez*Sequence[BaseSchedulerNode] | ExternKernelr`   )ra   s   "rb   rc   rM  Ú  s   € ÷ ñ àðð Bñrd   c                	ó   € R # ru   r`   )rf  rC  rT	  s   &&&rb   Úwrite_kernel_context_guardÚ/PythonWrapperCodegen.write_kernel_context_guardÚ  s   € ñ
 	rd   c                ó   € R# )z,
Mark the beginning of kernel context guard
Nr`   re  s   &rb   Ú write_kernel_context_guard_beginÚ5PythonWrapperCodegen.write_kernel_context_guard_beginá  ó   € ñ 	rd   c                ó   € R# )z&
Mark the end of kernel context guard
Nr`   re  s   &rb   Úwrite_kernel_context_guard_endÚ3PythonWrapperCodegen.write_kernel_context_guard_endé  r[	  rd   )8r  r  r   r  r;  r:  r  r  r7  r  r  r  rC  rk  r4  r>  r=  r  r,  r$  r¸  r-  rB  r“  r#  r”  r  rD  r5  r  r
  rø   r  r  r÷   r  r  r!  r  r*  r"  r%  r&  r<  r¨  r  r6  r   r  r  r'  r)  r(  r  r8  rö   ru   )r   r  )F)NTN)NF)r  )¸r‘   r’   r“   r”   r[  r  r•   rn  rÓ  rP  r+  r2  r.  r}  r0  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  r  r8  r$  r¾  rò  r·  rî  r3  r9  r=  rD  rH  rÇ  rÏ  r¶  rV  rZ  r`  rg  rm  rq  rw  r~  rƒ  r,  r˜  r?  r¶  rº  r¿  rÃ  rÇ  r»  r�  rÝ  r©  ré  rï  ró  r  Úcontextmanagerr   r  rì  r  r<  rJ  r  r_  re  rl  rs  rÞ  r¥  r  r­  r±  rµ  r�  rí   r;  rÅ  rÛ  rá  rç  rë  rõ  r  r  rã  r  r@  r  rI  rQ  r�  rY  Úclassmethodr^  rÂ  rË  rÇ  rÜ  rà  r×  rà  rÜ  r  r  r  r  r  r  rÿ  rê   r  r  r}  rö   r1  r
  r¦  r2  ra  rÚ  ri  rp  rf  rz  r€  r…  rs  rc  r{  r©  r¨  r²  r¶  r¹  r¾  r×  rà  rï  rô  rù  rÿ  r	  r	  r	  rß  r	  r*	  r?	  rD	  rH	  rK	  r)  rQ	  rV	  rY	  r]	  r–   r”  r•  s   @rb   r²  r²  ¨  sˆ  ø‡ ñð "Ð�dÓ!õtBðl ÷&ó ð&õ 'õ<õpBõdõð: ôó ðõ"	Bð ô+ó ð+õ!ð ô
ó ð
õ
õ$õ
%õ?õ0(õ
õ.õõ8)õ:õ8-õ
L÷	Wð 	W÷ð õ0Hõ

õLõ
õ÷*#>ð #>õJ=õò(.ò1ò0òJò
/õ+÷5-ð 5-õn5õõ,4õõõ"õQõ5õ5õ$õõõ:õ4õ,ò"õH8÷Að Aõ,õ(	õUòõ
8òõ0#õ Bò <õ
	Vò0õð ×Ñô!ó ð!õ2,òY
òv.Sò`:ò
*
õX%õ$$÷L#Kð #KõJKõ=ò"nKõ`DòðW¸t÷ Wð+À$÷ +õ.õ&õ'õ6õ÷<Pð <Põ|>õ5õSò
IòGò.6ò4õ*ò"eòN÷,
ð 
ð& ÷	ó ð	÷!ð !õ,>ð ñó ðõUJ÷n
ð õ*Aõ
)õV:ò.ò>ò;ò8ò5òWòFò
ò
ò@
4õ?GòB'ð6
ð
 ð6
ð ð6
ð ð6
ð ð6
ð ð6
ð *.ð6
ð ð6
ð "÷6
ðpX(ð
 ðX(ð ðX(ð ðX(ð ðX(ð ðX(ð *.ðX(ð ðX(ð ðX(ð "ðX(ð  ÷X(òt ò!ò,õ+õ\	
ð ñ8ó ð8õ.ò`*õ_õ*õAõu÷ð õ`õ&
÷53ón:ö6
ó 
÷	Eó'÷
÷Nó`;'÷z<÷B+÷
ó

ó$
ó"KóB
óNóP÷ %5óN]7ð~ òó ðð òó ðð ò
ó ð
ð
 ò
ó ð
ð
 òXó ðX÷ó÷ò rd   r²  c                  ó6  a € ] tR tRtRtR*R V 3R llltR R ltR R ltR	 tR
 t	R t
R R ltR R ltR R ltR R ltR R ltR R ltR R ltR V 3R lltR*R R lltR R ltR  R! lt]R" R# l4       t]R$ R% l4       t]R& R' l4       tR( tR)tV ;t# )+rO  iò  zü
A wrapper codegen that generates code for a subgraph. For most of the
methods, we rely on the implementation in the PythonWrapperCodegen. But we
override a few functions to produce cleaner code (like avoiding writing
imports twice in the output code)
c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rI  r‹   rJ  r²  rK  rL  r`   )ra   s   "rb   rc   Ú)SubgraphPythonWrapperCodegen.__annotate__ú  s.   € ÷ Hñ HàðHð -ðHð @ñ	Hrd   c                	óð   <€ Wn         W n        W0n        \        SV `  4        V P                  4       pVP                  V n        VP                  V n        VP                  V n        VP                  V n	        R # ru   )
rI  rJ  rK  rm  rn  Úget_root_graphr  rø   r   r(  )rf  rI  rJ  rK  Úrootrt  s   &&&& €rb   rn  Ú%SubgraphPythonWrapperCodegen.__init__ú  sh   ø€ ð +ÔØ,ÔØ$8Ô!ä‰ÑÔà×"Ñ"Ó$ˆà$(×$=Ñ$=ˆÔ!Ø%)×%?Ñ%?ˆÔ"à!×/Ñ/ˆÔà)-×)GÑ)GˆÖ&rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rc	    s   € ÷ 3ñ 3 dñ 3rd   c                	ó(   € V P                   V n        R # ru   )rI  r*  re  s   &rb   r+  Ú1SubgraphPythonWrapperCodegen.set_launcher_fn_name  s   € ð !%× 2Ñ 2ˆÖrd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rc	    s   € ÷ ñ ˜dñ rd   c                	ó   € R # ru   r`   re  s   &rb   r.  Ú)SubgraphPythonWrapperCodegen.write_header  r  rd   c                	ó   € R # ru   r`   rD  s   &&rb   r  Ú2SubgraphPythonWrapperCodegen.add_benchmark_harness  r  rd   c                	ó   € R # ru   r`   rD  s   &&rb   r@  Ú6SubgraphPythonWrapperCodegen.benchmark_compiled_module  r  rd   c                	ó   € R # ru   r`   re  s   &rb   rÃ  Ú5SubgraphPythonWrapperCodegen.write_async_compile_wait   r  rd   c               ó   € V ^8„  d   QhRR/# r£  r`   )ra   s   "rb   rc   rc	  #  s   € ÷ 8ñ 8 Cñ 8rd   c                	ó6   € V P                   P                  4       # ru   )rJ  r3  re  s   &rb   r3  Ú/SubgraphPythonWrapperCodegen.next_kernel_suffix#  s   € à×"Ñ"×5Ñ5Ó7Ð7rd   c               ó    € V ^8„  d   QhRRRR/# rc  r`   )ra   s   "rb   rc   rc	  '  s   € ÷ ñ ¨Nð ¸tñ rd   c                	ó   € R # ru   r`   rf  s   &&rb   rm  Ú2SubgraphPythonWrapperCodegen.generate_after_suffix'  rª  rd   c               ó   € V ^8„  d   QhRR/# rÍ  r`   )ra   s   "rb   rc   rc	  *  s   € ÷ ñ °3ñ rd   c                	ó\   € V P                   P                  R V P                   R24       ^pV# )z
            def z(args):
            )r  r?  r*  rÐ  s   & rb   rÒ  Ú>SubgraphPythonWrapperCodegen.write_launcher_fn_call_get_indent*  s<   € Ø�‰×ÑðØ×&Ñ&Ð'ð (ðô	
ð
 ˆØÐrd   c               ó   € V ^8„  d   QhRR/# rÍ  r`   )ra   s   "rb   rc   rc	  3  s   € ÷ ñ ¨ñ rd   c                	ó   € ^# rf   r`   re  s   &rb   ró  Ú4SubgraphPythonWrapperCodegen.get_wrapper_call_indent3  s   € Ùrd   c               ó   € V ^8„  d   QhRR/# )r\   r_   z@dict[str, ir.TensorBox | ir.TorchBindObject | sympy.Expr | None]r`   )ra   s   "rb   rc   rc	  6  s   € ÷ 	ñ 	à	Iñ	rd   c                	óÞ   € V P                   ;p'       d:   VP                  VP                   Uu/ uF  p\        V4      VbK  	  up,          pV# \        P
                  P                  pV# u upi ru   )rK  Úinput_nodesrë  r‹   rH   rœ   rê  )rf  rp  r@  r�   s   &   rb   r­  Ú-SubgraphPythonWrapperCodegen.get_graph_inputs6  sm   € ð ×1Ñ1Ð1ˆ9Ö1Ø×*Ñ*Ø#,×#:Ò#:ó.Ù#:˜a”�A“˜’	Ñ#:ñ.õ ˆFð
 ˆô —W‘W×)Ñ)ˆFØˆùò.s   ®A*c               ó   € V ^8„  d   QhRR/# r   r`   )ra   s   "rb   rc   rc	  A  s   € ÷ ñ  yñ rd   c                	ó
  € V P                   ;p'       dP   \        VP                  P                  4       4      VP                   Uu. uF  q"P
                  NK  	  up,           pV# \        P                  P                  pV# u upi ru   )	rK  r‚   r‚	  rë  rë  rÐ   rH   rœ   r¹  )rf  rp  rì  r  s   &   rb   rÖ  Ú2SubgraphPythonWrapperCodegen.get_graph_input_namesA  su   € Ø×1Ñ1Ð1ˆ9Ö1Ü˜×.Ñ.×3Ñ3Ó5Ó6Ø6?×6MÒ6Mó:Ù6M l×!Ô!Ñ6Mñ:õ ˆEð
 ˆô —G‘G×-Ñ-ˆEØˆùò:s   ÁB c               ó   € V ^8„  d   QhRR/# r°  r`   )ra   s   "rb   rc   rc	  J  s   € ÷ ñ  <ñ rd   c                	ó~   € V P                   ;p'       d   VP                  pV# \        P                  P                  pV# ru   )rK  rê  rH   rœ   rò  )rf  rp  rÐ  s   &  rb   r¢  Ú.SubgraphPythonWrapperCodegen.get_graph_outputsJ  s;   € Ø×1Ñ1Ð1ˆ9Ö1Ø×,Ñ,ˆGð ˆô —g‘g×+Ñ+ˆGØˆrd   c               ó   € V ^8„  d   QhRR/# r‘  r`   )ra   s   "rb   rc   rc	  Q  s   € ÷ +ñ +¨ñ +rd   c                	ó’   <€ VP                  4       pV P                  ;p'       d   W#P                  9   d   R # \        SV `  V4       R # ru   )rž   rK  r‚	  rm  r{  )rf  r|  rÐ   rp  rt  s   &&  €rb   r{  Ú/SubgraphPythonWrapperCodegen.codegen_allocationQ  s=   ø€ Ø�‰Ó ˆØ×2Ñ2Ð2ˆIÖ2¸×@UÑ@UÔ8Uñ ä‰Ñ" 6Ö*rd   c               ó$   € V ^8„  d   QhRRRRRR/# )r\   rú  r�   rœ   zGraphLowering | Noner_   r‹   r`   )ra   s   "rb   rc   rc	  [  s$   € ÷ ñ ØðØ&:ðà	ñrd   c                	óT  € V P                  4        \        P                  P                  P                  ;peE   V R2pV P                  V R\        P                  P                  P                  V4       24       V# \        V4      pV P                  V R\        V4       R24       V# )NÚ_rawr2  r  r3  )
r‹  rH   rœ   r    rA  rö   r	  Ústream_handler8   rý  )rf  rú  rœ   rA  rÐ   s   &&&  rb   Ú_write_get_raw_streamÚ2SubgraphPythonWrapperCodegen._write_get_raw_stream[  s¨   € ð 	×%Ñ%Ô'ä#$§7¡7×#4Ñ#4×#HÑ#HÐHÐÒUØ)Ð*¨$Ð/ˆDØ�N‰NØ�&˜œAŸG™G×.Ñ.×<Ñ<Ð=PÓQÐRÐSôð ˆô	 ' zÓ2ˆDØ�N‰NØ�&Ð*Ô+?À
Ó+KÐ*LÈAÐNôð ˆrd   c               ó    € V ^8„  d   QhRRRR/# )r\   Úpost_grad_graph_idr�   r_   r„  r`   )ra   s   "rb   rc   rc	  l  s   € ÷ Tñ TÀð TÈñ Trd   c                ó0   € V P                  RV R24       R# )z#Generate NVTX range push for graph.z"torch.cuda.nvtx.range_push('graph rÚ  Nrþ  )rf  r”	  s   &&rb   Úcodegen_graph_nvtx_range_pushÚ:SubgraphPythonWrapperCodegen.codegen_graph_nvtx_range_pushl  s   € à�‰Ð;Ð<NÐ;OÈrÐRÖSrd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rc	  p  s   € ÷ 6ñ 6¨dñ 6rd   c                ó(   € V P                  R4       R# )z"Generate NVTX range pop for graph.ztorch.cuda.nvtx.range_pop()Nrþ  re  s   &rb   Úcodegen_graph_nvtx_range_popÚ9SubgraphPythonWrapperCodegen.codegen_graph_nvtx_range_popp  s   € à�‰Ð4Ö5rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rc	  u  s   € ÷ 7ñ 7¨$ñ 7rd   c                	ó:   € V P                   P                  4        R # ru   )rJ  r‹  re  s   &rb   r‹  Ú5SubgraphPythonWrapperCodegen.write_triton_header_oncet  s   € ð 	×Ñ×4Ñ4Ö6rd   c               ó   € V ^8„  d   QhRR/# r´  r`   )ra   s   "rb   rc   rc	  ~  s   € ÷ ?ñ ?°$ñ ?rd   c                	ó:   € V P                   P                  4        R # ru   )rJ  r•  re  s   &rb   r•  Ú=SubgraphPythonWrapperCodegen.write_get_raw_stream_header_once}  s   € ð 	×Ñ×<Ñ<Ö>rd   c               ó   € V ^8„  d   QhRR/# )r\   r_   r²  r`   )ra   s   "rb   rc   rc	  ˆ  s   € ÷ ñ Ð 4ñ rd   c                	ó¬   € T p\        V\        4      '       d   VP                  pK$  \        V\        4      '       g   \	        R \        V4       24      hV# )z#expected PythonWrapperCodegen, got )r   rO  rJ  r²  r²   r…   )rf  rf	  s   & rb   re	  Ú+SubgraphPythonWrapperCodegen.get_root_graph‡  sJ   € àDHˆÜ˜Ô;×<Ò<Ø×&Ñ&ŠDä˜$Ô 4×5Ò5Ü Ð#FÄtÈDÃzÀlÐ!SÓTÐTØˆrd   c                	ó   € R # ru   r`   re  s   &rb   r  Ú<SubgraphPythonWrapperCodegen.generate_and_run_autotune_block‘  s   € árd   )rø   r  r*  rJ  rK  r   rI  r(  ru   )r‘   r’   r“   r”   r[  rn  r+  r.  r  r@  rÃ  r3  rm  rÒ  ró  r­  rÖ  r¢  r{  r‘	  r–	  rš	  r:   r‹  r•  re	  r  r–   r”  r•  s   @rb   rO  rO  ò  s°   ø† ñ÷Hò Hõ.3õòòòõ8õõõõ	õõ÷+ð +÷õ"Tõ6ð ô7ó ð7ð ô?ó ð?ð ôó ð÷ð rd   rO  )NNru   )×Ú
__future__r   ro  r  r¥  r7  r*  r¬  r,  rï  rl  r<  rÐ  r.  Úcollections.abcr   Úenumr   Ú	itertoolsr   r   Útypingr   r	   r
   r   rh   r   r^  Ú
torch._opsÚtorch.utils._pytreeÚutilsÚ_pytreerÅ  r   r  Útorch._dynamo.utilsr   r   r   Ú#torch._inductor.codegen.debug_utilsr   Ú$torch._inductor.codegen.multi_kernelr   Ú%torch._inductor.runtime.runtime_utilsr   Útorch._library.opaque_objectr   r   Útorch._loggingr   Ú"torch.fx.experimental.proxy_tensorr   r   r   Ú%torch.fx.experimental.symbolic_shapesr   r   r   r   r   Útorch.fx.noder    Ú torch.multiprocessing.reductionsr!   Útorch.utils._ordered_setr"   Útorch.utils._sympy.functionsr#   r$   r%   Ú torch.utils._sympy.singleton_intr&   Útorch.utils._sympy.symbolr'   r(   r  r)   r*   r+   rÒ  r,   Ú	codecacher-   r.   r/   rA  r0   Úruntime.hintsr1   r2   Ústream_constantsr3   r4   r5   Ústream_utilsr6   r7   r8   r9   r:   r;   r<   r=   r>   r?   r@   rA   rB   rC   rD   rE   rF   rG   ÚvirtualizedrH   r™  rI   rJ   rK   rL   rM   rN   Ú	cpp_utilsrO   Úcustom_extern_kernel_codegenrP   Útriton_utilsrQ   rR   rS   rT   rU   rV   rî   rœ   rW   rX   r    rY   Úwrapper_fxirrZ   Ú	getLoggerr‘   ÚlogÚdoprintrŒ  rn   rv   rƒ   r�   r‹   rE  r�   rš   r­   rÁ   r™   rª  rO  rŠ   r«   r¶   r¿   rÎ   r€   r™  râ   r!  rX  r^  r`  rj  r—  rˆ   r±  rÇ  rÕ  rß  rê  rü  rý  rÿ  r  r%  r4  rM  r]  ro  r†  r–  rª  r½  rÃ  rË  rÓ  rÛ  rã  r  r}  rY  r(  r&  r€  r›  r±  rÂ  rÎ  rÞ  rï  rû  r3  ÚLiner²  rO  r`   rd   rb   Ú<module>rË	     sÑ  ðå "ã Û Û Û 
Û Û Û Û Û 	Û Û 	Û Ý $Ý ß "ß 8Ó 8ã Ý ã Û ß $Ð $Ý &ß EÑ EÝ CÝ AÝ ;ß UÝ +÷ñ ÷
õ õ .Ý ;Ý /ß ;Ñ ;Ý 9ß :ç AÓ AÝ 'ß (Ý 'ß 8ß WÑ W÷ó ÷÷ ÷ ÷ õ  ÷÷ õ Ý Fß PÑ P÷ ß<Ñ<ãå%Ý!Ý-Ý)ð ×Ó˜Ó!€á‹×Ò€õ,õð$ �—’˜uŸ{™{¨C°°s¸EÀ#ÀsÀ(½OÈdÕ<RÐRÕS€Ø˜5Ÿ<š<¨¯©°cÐ;NÐPSÐSÕTÐ Ø�YŠY˜Õ%€
Ø˜]˜O¨TÐ1Õ2Ð ð ×Ò÷3ð 3ó ð3õõ$õò>QðF ˜˜S˜•>Ð à	ˆ#�—
‘
Õ
˜CÐ
Õ  8Ð-=Ð,>ÀÀcÈ3ÀhÅÐ,OÕ#PÕPð ÷
o&÷dq&õh+ð ×Ò÷ð ó ð÷-6ñ -6ô`Q�8ô Q÷Yñ Yð ×Òô2˜ó 2ó ð2ð ×Òô	*�ó 	*ó ð	*ð ×Òô+�+ó +ó ð+ð ×Òô	2˜ó 	2ó ð	2ð ×Òô1�{ó 1ó ð1ð  +Ð õ
Cð ×Òô-@ Kó -@ó ð-@ô`? ;ô ?ð ×Òô
7˜Kó 
7ó ð
7ð ×Òô5˜+ó 5ó ð5ð< ×ÒôP˜{ó Pó ðPð. ×Òô(ˆ{ó (ó ð(ð ×Òô /�[ó  /ó ð /ðF ×Òô5˜;ó 5ó ð5ð* ×Òô;˜ó ;ó ð;ð2 ×Òô"Ð2Oó "ó ð"ðJ ×ÒôÐ1Mó ó ðð ×ÒôR ó Ró ðRð$ ×Òô@ ó @ó ð@ð ×Òô +ó ó ðô ô ÷%
ñ %
ðP ×Òôs,Ð%ó s,ó ðs,ðl ×Òô36Ð,ó 36ó ð36ðl ×Òô/Ð(ó /ó ð/ð. ×Òô!)Ð"ó !)ó ð!)ôH(Ð!ô (ð
 ×Òô#0�kó #0ó ð#0ðL ×Òô6˜;ó 6ó ð6ð. ×Òô4˜+ó 4ó ð4ð: ×Òô	5˜+ó 	5ó ð	5ð ×Òô8˜[ó 8ó ð8ð ×Òô6˜;ó 6ó ð6ð$ ×Òô6 +ó 6ó ð6ð ×ÒôS˜+ó Só ðSð  €
Ø˜KÕ'€ôG9˜7ô G9ôTraÐ#7ö ard   