Ë
    ÏÍ:jG ã                   ó‚  — d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlmZ d dl	m
Z
 d dlmZmZmZmZmZmZmZmZmZ ddlmZmZ ddlmZmZ ddlmZmZmZ dd	lmZm Z m!Z! dd
l"m#Z#m$Z$ ddl%m&Z& ddl'm(Z(m)Z)m*Z* ddl+m,Z,m-Z-m.Z. d„ Z/d„ Z0dede1fd„Z2dede1fd„Z3dede1fd„Z4dede1fd„Z5dede1fd„Z6d„ Z7d„ Z8d„ Z9dee    fd„Z:deejv                     dee!   fd„Z<e1e= e>d«      hZ? G d „ d!«      Z@ G d"„ d#e j‚                  «      ZB G d$„ d%«      ZC ed&¬'«       G d(„ d)«      «       ZD G d*„ d+e j‚                  «      ZEd-d,„ZFy).é    N)Ú	dataclass)Ú
ModuleType)	ÚAnyÚCallableÚDictÚOptionalÚTupleÚTypeÚUnionÚIterableÚListé   )ÚknobsÚlanguage)ÚirÚgluon_ir)Ú	constexprÚ	str_to_tyÚtensor)Ú_unwrap_if_constexprÚ
base_valueÚ	base_type)Úget_jit_fn_file_lineÚget_full_name)ÚJITFunction)Úfind_paths_ifÚget_iterable_pathÚset_iterable_pathé   )ÚCompilationErrorÚCompileTimeAssertionFailureÚUnsupportedLanguageConstructc                 óZ   — d}t        j                  || «      st        d|› d| › �| «      ‚| S )Nz^[a-zA-Z_][a-zA-Z0-9_]*$zinvalid z identifier: )ÚreÚmatchr    )ÚnameÚtypeÚpatterns      ús/home/mcse/projects/srt_converter/srt-converter-venv/lib/python3.12/site-packages/triton/compiler/code_generator.pyÚcheck_identifier_legalityr*      s4   € Ø)€GÜ�8‰8�G˜TÔ"Ü ¨$¨¨}¸T¸FÐCÀTÓJÐJØ€Kó    c                 ó†  — dj                  |D �cg c]  }|j                  «       ‘Œ c}«      }dj                  t        |«      D �cg c]  }|› dt        ||   «      › �‘Œ c}«      }|j	                  dd«      }|j	                  dd«      }|j	                  dd«      j	                  dd«      }| › d	|› d	|› �}|S c c}w c c}w )
NÚ_Úcú.Ú_d_ú'Ú_sq_ú[ú]Ú__)ÚjoinÚmangleÚsortedÚreprÚreplace)r&   Úarg_tysÚ	constantsÚtyÚmangled_arg_namesÚiÚmangled_constantsÚrets           r)   Ú	mangle_fnrB      sÄ   € àŸ™¸Ö!@°" "§)¡)¥+Ò!@ÓAÐØŸ™ÄvÈiÓGXÖ!YÀ! Q C q¬¨i¸©lÓ);Ð(<Ò"=Ò!YÓZÐØ)×1Ñ1°#°uÓ=ÐØ)×1Ñ1°#°vÓ>Ðà)×1Ñ1°#°sÓ;×CÑCÀCÈÓMÐØˆF�"Ð&Ð' rÐ*;Ð)<Ð
=€CØ€Jùò "AùÚ!Ys   �B9ÁB>ÚoÚreturnc                 ó"   — t        | t        «      S ©N)Ú
isinstancer   ©rC   s    r)   Ú_is_triton_valuerI   +   s   € Ü�aœÓ$Ð$r+   c                 ó"   — t        | t        «      S rF   )rG   r   rH   s    r)   Ú_is_triton_tensorrK   /   s   € Ü�aœÓ Ð r+   c                 ól   — | d u xs/ t        | t        t        j                  j                  t
        f«      S rF   )rG   r   r   ÚcoreÚdtyper   rH   s    r)   Ú_is_constexprrO   3   s)   € Ø�ˆ9ÒTœ
 1¤y´(·-±-×2EÑ2EÄ{Ð&SÓTÐTr+   c                 ó†   — t        | «      xr5 | j                  j                  «       xr | j                  j                  dk7  S ©Nr   )rK   r'   Úis_blockÚnumelrH   s    r)   Ú_is_non_scalar_tensorrT   7   s/   € Ü˜QÓÒM Q§V¡V§_¡_Ó%6Ò%L¸1¿6¹6¿<¹<È1Ñ;LÐMr+   c                 ó.   — t        | t        t        f«      S rF   )rG   ÚlistÚtuplerH   s    r)   Ú_is_list_likerX   ;   s   € Ü�aœ$¤˜Ó'Ð'r+   c                 óà   — |j                   rbt        |«      D ]S  \  }}t        |«      rŒt        |«      sŒt	        |j
                  | d|j                  › d|j                  |   › d|› �«      ‚ y y )Nz	Function z= is marked noinline, but was called with non-scalar argument ú:)ÚnoinlineÚ	enumeraterO   rT   r"   ÚsrcÚ__name__Ú	arg_names)ÚnodeÚfnÚargsÚidxÚargs        r)   Ú_check_fn_argsre   ?   s‚   € Ø	‡{‚{Ü! $›ò 	‰HˆC�Ü  Õ%Ô*?ÀÕ*DÜ2Ø—F‘F˜DØ §¡˜}Ð,iÐjl×jvÑjvÐwzÑj{Ði|Ð|}ð  Bð  ~Cð  Dóð ñ	ð r+   c                 ób   — t        | t        «      xr t        | t        «      xr t	        | d«      S )NÚ_fields)rG   r'   Ú
issubclassrW   Úhasattr)Úvals    r)   Ú_is_namedtuplerk   I   s'   € Ü�cœ4Ó ÒW¤Z°´UÓ%;ÒWÄÈÈYÓ@WÐWr+   c                 óÒ  — t        t        | «      «      r| j                  }nDt        | t        j
                  «      r| j                  j                  }nJ dt        | «      › �«       ‚| D �cg c]
  } ||«      ‘Œ }}|D �cg c]  }|€t        |«      n|‘Œ }}|D �cg c]  }|j                  ‘Œ }}t	        j
                  |t	        j                  ||«      «      S c c}w c c}w c c}w )NzUnsupported type )	rk   r'   rg   rG   r   rW   Úfieldsr   Ú
tuple_type)Úvaluera   rm   ÚvÚvalsÚtypess         r)   Ú_apply_to_tuple_valuesrs   M   sÁ   € Ü”d˜5“kÔ"Ø—‘‰Ü	�Eœ8Ÿ>™>Ô	*Ø—‘×"Ñ"‰à7Ð)¬$¨u«+¨Ð7Ó7ˆuà Ö!�a‰Bˆq�EÐ!€DÐ!Ø6:Ö;°˜A˜IŒI�aŒL¨1Ñ,Ð;€DÐ;Ø!Ö"˜ˆQ�V‹VÐ"€EÐ"Ü�>‰>˜$¤× 3Ñ 3°E¸6Ó BÓCÐCùò "ùÚ;ùÚ"s   Á*CÁ?CÂC$Úvaluesc                 ó:   — g }| D ]  }|j                  |«       Œ |S rF   )Ú_flatten_ir)rt   Úhandlesrp   s      r)   Úflatten_values_to_irrx   [   s'   € Ø€GØò ˆØ	�‰�gÕðà€Nr+   rw   rr   c              #   óp   K  — d}|D ]  }|j                  | |«      \  }}|–— Œ |t        | «      k(  sJ ‚y ­w)Nr   )Ú_unflatten_irÚlen)rw   rr   Úcursorr=   ro   s        r)   Úunflatten_ir_valuesr}   b   sI   è ø€ Ø€FØò ˆØ×(Ñ(¨°&Ó9‰ˆˆvØ‹ðð ”S˜“\Ò!Ð!Ñ!ùs   ‚46c                   ó   — e Zd Zd„ Zd„ Zd„ Zy)Úenter_sub_regionc                 ó   — || _         y rF   )Ú	generator)Úselfr�   s     r)   Ú__init__zenter_sub_region.__init__o   s	   € Ø"ˆ�r+   c                 óœ  — | j                   j                  j                  «       | _        | j                   j                  j                  «       | _        i | j                   _        | j                   j                  j                  «       | _        | j                   j                  j                  «       | _
        | j                  | j                  fS rF   )r�   ÚlscopeÚcopyÚliveinsÚ
local_defsÚ	prev_defsÚbuilderÚget_insertion_blockÚinsert_blockÚget_insertion_pointÚinsert_point)r‚   s    r)   Ú	__enter__zenter_sub_region.__enter__r   s�   € à—~‘~×,Ñ,×1Ñ1Ó3ˆŒØŸ™×2Ñ2×7Ñ7Ó9ˆŒØ$&ˆ�‰Ô!Ø ŸN™N×2Ñ2×FÑFÓHˆÔØ ŸN™N×2Ñ2×FÑFÓHˆÔØ�|‰|˜T×.Ñ.Ð.Ð.r+   c                 óÎ   — | j                   j                  j                  | j                  «       | j                  | j                   _        | j                  | j                   _        y rF   )r�   rŠ   Úrestore_insertion_pointrŽ   r‡   r…   r‰   rˆ   )r‚   rb   Úkwargss      r)   Ú__exit__zenter_sub_region.__exit__{   s@   € Ø�‰×Ñ×6Ñ6°t×7HÑ7HÔIØ $§¡ˆ�‰ÔØ$(§N¡Nˆ�‰Õ!r+   N)r^   Ú
__module__Ú__qualname__rƒ   r�   r“   © r+   r)   r   r   m   s   „ ò#ò/ó3r+   r   c                   óž  — e Zd Zd„ Zdefd„Zdefd„Zdefd„Zdej                  defd„Z
dej                  defd„Zdej                  defd	„Zdej                  defd
„Zdej"                  defd„Zdej&                  defd„Zdej*                  defd„Zdej.                  defd„Zdej2                  defd„Zdej6                  defd„Zy)ÚContainsReturnCheckerc                 ó   — || _         y rF   )Úgscope)r‚   rš   s     r)   rƒ   zContainsReturnChecker.__init__„   s	   € Øˆ�r+   rD   c                 ó,   ‡ — t        ˆ fd„|D «       «      S )Nc              3   ó@   •K  — | ]  }‰j                  |«      –— Œ y ­wrF   ©Úvisit)Ú.0Úsr‚   s     €r)   ú	<genexpr>z5ContainsReturnChecker._visit_stmts.<locals>.<genexpr>ˆ   s   øè ø€ Ò/ Q�4—:‘:˜a—=Ñ/ùó   ƒ)Úany)r‚   Úbodys   ` r)   Ú_visit_stmtsz"ContainsReturnChecker._visit_stmts‡   s   ø€ ÜÓ/¨$Ô/Ó/Ð/r+   c                  ó   — y©NFr–   )r‚   ra   s     r)   Ú_visit_functionz%ContainsReturnChecker._visit_functionŠ   s   € ð r+   c                 ó4  — d}t        j                  |«      D ]}  \  }}t        |t        «      r8|D ]2  }t        |t         j                  «      sŒ|xs | j                  |«      }Œ4 ŒNt        |t         j                  «      sŒi|xs | j                  |«      }Œ |S r§   )ÚastÚiter_fieldsrG   rV   ÚASTrž   )r‚   r`   rA   r-   ro   Úitems         r)   Úgeneric_visitz#ContainsReturnChecker.generic_visit�   s‡   € ØˆÜŸ™¨Ó-ò 	/‰HˆAˆuÜ˜%¤Ô&Ø!ò 6�DÜ! $¬¯©Õ0Ø!Ò5 T§Z¡Z°Ó%5™ñ6ô ˜E¤3§7¡7Õ+ØÒ.˜TŸZ™Z¨Ó.‘ð	/ð ˆ
r+   r`   c                 óZ  — t        |j                  t        j                  «      rm|j                  j                  | j
                  v rJ| j
                  |j                  j                     }t        ||j                  «      }| j                  |«      S y| j                  |j                  «      S r§   )
rG   ro   rª   ÚNameÚidrš   ÚgetattrÚattrr¨   rž   )r‚   r`   ro   ra   s       r)   Úvisit_Attributez%ContainsReturnChecker.visit_Attribute›   sw   € ô �d—j‘j¤#§(¡(Ô+Ø�z‰z�}‰} §¡Ñ+ØŸ™ D§J¡J§M¡MÑ2�Ü˜U D§I¡IÓ.�Ø×+Ñ+¨BÓ/Ð/ØØ�z‰z˜$Ÿ*™*Ó%Ð%r+   c                 óÔ   — t        |j                  «      t        j                  u ry|j                  | j
                  v r*| j
                  |j                     }| j                  |«      S yr§   )r'   Úctxrª   ÚStorer±   rš   r¨   )r‚   r`   ra   s      r)   Ú
visit_Namez ContainsReturnChecker.visit_Nameª   sO   € Ü�—‘‹>œSŸY™YÑ&ØØ�7‰7�d—k‘kÑ!Ø—‘˜TŸW™WÑ%ˆBØ×'Ñ'¨Ó+Ð+Ør+   c                  ó   — y©NTr–   ©r‚   r`   s     r)   Úvisit_Returnz"ContainsReturnChecker.visit_Return²   ó   € Ør+   c                  ó   — yr§   r–   r»   s     r)   Úvisit_Assignz"ContainsReturnChecker.visit_Assignµ   ó   € ð r+   c                  ó   — yr§   r–   r»   s     r)   Úvisit_AugAssignz%ContainsReturnChecker.visit_AugAssignº   rÀ   r+   c                 ó8   — | j                  |j                  «      S rF   ©r¥   r¤   r»   s     r)   Úvisit_Modulez"ContainsReturnChecker.visit_Module¿   ó   € Ø× Ñ  §¡Ó+Ð+r+   c                 ó8   — | j                  |j                  «      S rF   rÄ   r»   s     r)   Úvisit_FunctionDefz'ContainsReturnChecker.visit_FunctionDefÂ   rÆ   r+   c                 ó’   — | j                  |j                  «      }|j                  r|xs | j                  |j                  «      }|S rF   )r¥   r¤   Úorelse)r‚   r`   rA   s      r)   Úvisit_IfzContainsReturnChecker.visit_IfÅ   s=   € ð ×Ñ §	¡	Ó*ˆØ�;Š;ØÒ7˜×*Ñ*¨4¯;©;Ó7ˆCØˆ
r+   c                 ór   — | j                  |j                  «      xs | j                  |j                  «      S rF   )rž   r¤   rÊ   r»   s     r)   Úvisit_IfExpz!ContainsReturnChecker.visit_IfExpÒ   s'   € Ø�z‰z˜$Ÿ)™)Ó$Ò?¨¯
©
°4·;±;Ó(?Ð?r+   c                 ó8   — | j                  |j                  «      S rF   )rž   Úfuncr»   s     r)   Ú
visit_Callz ContainsReturnChecker.visit_CallÕ   s   € Ø�z‰z˜$Ÿ)™)Ó$Ð$r+   N)r^   r”   r•   rƒ   Úboolr¥   r¨   r®   rª   Ú	Attributer´   r°   r¸   ÚReturnr¼   ÚAssignr¿   Ú	AugAssignrÂ   ÚModulerÅ   ÚFunctionDefrÈ   ÚIfrË   ÚIfExprÍ   ÚCallrÐ   r–   r+   r)   r˜   r˜   ‚   s  „ òð0 Dó 0ð Tó ð	 Tó 	ð& C§M¡Mð &°dó &ð˜sŸx™xð ¨Dó ð §¡ð °ó ð §¡ð °ó ð
 C§M¡Mð °dó ð
, §¡ð ,°ó ,ð, c§o¡oð ,¸$ó ,ð˜SŸV™Vð ¨ó ð@ §	¡	ð @¨dó @ð%˜sŸx™xð %¨Dô %r+   r˜   c                   ó¾   — e Zd Zd„ Zdej
                  dee   deej                     fd„Z	dej
                  deej                     fd„Z
dej
                  fd„Zd„ Zy	)
ÚASTFunctionc                 ó<   — || _         || _        || _        || _        y rF   )Ú	ret_typesÚ	arg_typesr<   Úattrs)r‚   rÞ   rß   r<   rà   s        r)   rƒ   zASTFunction.__init__Û   s   € Ø"ˆŒØ"ˆŒØ"ˆŒØˆ�
r+   rŠ   rr   rD   c                 óB   — g }|D ]  }|€Œ|j                  ||«       Œ |S rF   )Ú_flatten_ir_types)r‚   rŠ   rr   Úir_typesr=   s        r)   Úflatten_ir_typeszASTFunction.flatten_ir_typesá   s5   € ØˆØò 	4ˆBØˆzØØ× Ñ  ¨(Õ3ð	4ð ˆr+   c                 ó:   — | j                  || j                  «      S rF   )rä   rÞ   )r‚   rŠ   s     r)   Úreturn_types_irzASTFunction.return_types_iré   s   € Ø×$Ñ$ W¨d¯n©nÓ=Ð=r+   c                 ó  ‡ — ˆ fd„}t        t        ‰ j                  |«      «      }|D �cg c]  }t        ‰ j                  |«      ‘Œ }}‰ j	                  ||«      }‰ j                  |«      }|j                  ||«      S c c}w )Nc                 ó,   •— | ‰j                   vxr |d uS rF   ©r<   ©Úpathr-   r‚   s     €r)   ú<lambda>z'ASTFunction.serialize.<locals>.<lambda>ï   ó   ø€  ¨T¯^©^Ð!;Ò!MÀÈÀ€ r+   )rV   r   rß   r   rä   ræ   Úget_function_ty)r‚   rŠ   Úis_valÚ	val_pathsrë   rß   Úarg_types_irÚret_types_irs   `       r)   Ú	serializezASTFunction.serializeì   sz   ø€ ó NˆÜœ t§~¡~°vÓ>Ó?ˆ	ØIRÖSÀÔ& t§~¡~°tÕ<ÐSˆ	ÐSØ×,Ñ,¨W°iÓ@ˆØ×+Ñ+¨GÓ4ˆØ×&Ñ& |°\ÓBÐBùò Ts   ªA>c                 óp  ‡ ‡— ˆfd„Š ‰‰ j                   «      }ˆ fd„}t        t        ‰ j                   |«      «      }d}t        |j	                  «       «      D �cg c]  }|j                  |«      ‘Œ }}|D ]s  }t        ‰ j                   |«      }	‰ j                  j                  |g «      }
|
D ]  \  }}|j                  |||«       Œ |	j                  ||«      \  }}t        |||«       Œu ‰ j                  }|j                  «       D ]%  \  }}t        ||t        j                  |«      «       Œ' |S c c}w )Nc                 óÖ   •— t        | t        t        t        j                  f«      r*t        j                  | D �cg c]
  } ‰|«      ‘Œ c}| «      S t        j
                  d «      S c c}w rF   )rG   rV   rW   r   rn   r   )r=   ÚxÚmake_templates     €r)   r÷   z.ASTFunction.deserialize.<locals>.make_templateø   sR   ø€ Ü˜"œt¤U¬H×,?Ñ,?Ð@ÔAÜ—~‘~ÀÖ&D¸A¡}°QÕ'7Ò&DÀbÓIÐIÜ×%Ñ% dÓ+Ð+ùò 'Es   ºA&c                 ó,   •— | ‰j                   vxr |d uS rF   ré   rê   s     €r)   rì   z)ASTFunction.deserialize.<locals>.<lambda>þ   rí   r+   r   )rß   rV   r   ÚrangeÚget_num_argsrb   r   rà   ÚgetÚset_arg_attrrz   r   r<   Úitemsr   r   )r‚   ra   rq   rï   rð   r|   r?   rw   rë   r=   Ú
attr_specsÚ	attr_nameÚattr_valrj   r<   r÷   s   `              @r)   ÚdeserializezASTFunction.deserializeö   s"  ù€ ô	,ñ
 ˜TŸ^™^Ó,ˆÛMˆÜœ t§~¡~°vÓ>Ó?ˆ	àˆÜ',¨R¯_©_Ó->Ó'?Ö@ !�2—7‘7˜1•:Ð@ˆÐ@Øò 	/ˆDÜ" 4§>¡>°4Ó8ˆBàŸ™Ÿ™¨¨bÓ1ˆJØ'1ò =Ñ#�	˜8Ø—‘ ¨	°8Õ<ð=ð ×*Ñ*¨7°FÓ;‰KˆC�Ü˜d D¨#Õ.ð	/ð —N‘Nˆ	Ø"Ÿ™Ó*ò 	C‰IˆD�#Ü˜d D¬(×*<Ñ*<¸SÓ*AÕBð	Càˆùò As   ÁD3N)r^   r”   r•   rƒ   r   rŠ   r   r   r'   rä   ræ   ró   r  r–   r+   r)   rÜ   rÜ   Ù   sj   „ òð¨¯
©
ð ¸4À	¹?ð ÈtÐTV×T[ÑT[É}ó ð> r§z¡zð >°d¸2¿7¹7±mó >ðC §¡ó Cór+   rÜ   T)Úfrozenc                   ó"   — e Zd ZU eed<   eed<   y)ÚBoundJITMethodÚ__self__Ú__func__N)r^   r”   r•   r   Ú__annotations__r   r–   r+   r)   r  r    s   … àÓØÔr+   r  c                   ó  ‡— e Zd ZU 	 	 dgdedee   dee   fd„Zee	e
eeeeefD � �ci c]  }|j                   |“Œ c}} Zeeef   ed<   ej'                  dej*                  j,                  fdej.                  fd	ej0                  ff«       d
„ Zd„ Zd„ Zdedeeef   ddfd„Zd„ Z d„ Z!d„ Z"d„ Z#d„ Z$d„ Z%defd„Z&d„ Z'd„ Z(d„ Z)d„ Z*d„ Z+d„ Z,d„ Z-d„ Z.d „ Z/d!„ Z0d"„ Z1d#„ Z2d$„ Z3e4jj                  d%e4jl                  d&e4jn                  d'e4jp                  d(e4jr                  d)e4jt                  d*e4jv                  d+e4jx                  d,e4jz                  d-e4j|                  d.e4j~                  d/e4j€                  d0iZAeeBe4j†                     ef   ed1<   d2„ ZDd3„ ZEd4„ ZFd5„ ZGd6„ ZHd7„ ZId8„ ZJe4j–                  d9e4j˜                  d:e4jš                  d;e4jœ                  d<e4jž                  d=e4j                   d>iZQeeBe4j¤                     ef   ed?<   d@„ ZSe4j¨                  dAe4jª                  dBe4j¬                  dCe4j®                  dDiZXeeBe4j²                     ef   edE<   dF„ ZZdG„ Z[dH„ Z\dI„ Z]dJ„ Z^dK„ Z_dL„ Z`dM„ ZadN„ ZbdO„ ZcdP„ Zddeeeef   fdQ„ZfdefdR„ZgdSefdT„ZhdU„ ZidV„ ZjdWe4jÖ                  fdX„Zle4jÚ                  dYe4jÜ                  dZiZoeeBe4jà                     ef   ed[<   d\„ Zqd]„ Zrd^„ Zsd_„ Ztˆfd`„Zuda„ ZvdWe4jî                  ddfdb„Zxdc„ ZydddelzmZ{ ej*                  jø                  exej*                  jú                   eye~«      e{jø                  exe{jú                   eye~«      e eye«      e eye«      iZee€e�e4jî                  gef   f   edf<   ˆxZ‚S c c}} w )hÚCodeGeneratorNÚjit_fnÚfunction_typesÚ	file_namec                 ó’  — || _         |j                  «       r8ddlm} t	        j
                  |«      | _         || j                  «      | _        n7ddlm	} t        j                  |«      | _         || j                  «      | _        || _        |dz
  | _        | j                  j                  ||d«       || j                  _        || j                  _        |€i n|| j                  _        |	€| j                  j#                  «       n|	| _        |€i n|| _        || _        i | _        |j-                  «       D ]†  \  }}t/        |t0        «      r*|j3                  |j4                  |«      | j*                  |<   Œ@t7        |dd«      }||v r't7        ||   |j4                  «      | j*                  |<   Œx|| j*                  |<   Œˆ i | _        || _        |
r#||j=                  d«      dz   d  }t?        |d«      }|| _         |
| _!        d | _"        || _#        g | _$        d | _%        i | _&        | jO                  «       | _(        d | _)        d	| _*        y )
Nr   )ÚGluonSemantic)ÚTritonSemanticr   r”   Ú r/   ÚfunctionF)+ÚcontextÚis_gluonÚ,triton.experimental.gluon.language._semanticr  r   ÚGluonOpBuilderrŠ   ÚsemanticÚtriton.language.semanticr  r   r  Ú
begin_lineÚset_locÚoptionsÚcodegen_fnsÚ
module_mapÚcreate_moduleÚmoduleÚfunction_ret_typesÚ	prototyperš   rý   rG   r   rû   r^   r²   r…   r
  Úrfindr*   Úfunction_nameÚ	is_kernelÚcur_noder[   Ú	scf_stackÚret_typerˆ   Ú_define_name_lookupÚdereference_namera   Úvisiting_arg_default_value)r‚   r  r   rš   r"  r
  r  r  r  r  r#  r  r[   r  r  r  r  Úkrp   Úmodule_names                       r)   rƒ   zCodeGenerator.__init__  sý  € ð ˆŒØ�?‰?ÔÝRÜ#×2Ñ2°7Ó;ˆDŒLÙ)¨$¯,©,Ó7ˆD�Må?ÜŸ:™: gÓ.ˆDŒLÙ*¨4¯<©<Ó8ˆDŒMØ"ˆŒà$ q™.ˆŒØ�‰×Ñ˜Y¨
°AÔ6Ø&ˆ�‰Ôð $/ˆ�‰Ô Ø(2Ð(:¡"À
ˆ�‰ÔØ6<°n�d—l‘l×0Ñ0Ô2È&ˆŒØ(6Ð(>¡"ÀNˆÔØ"ˆŒàˆŒØ—L‘L“Nò 		#‰DˆAˆqÜ˜!œZÔ(Ø!+§¡°·
±
¸AÓ!>�—‘˜A‘Øä! ! \°2Ó6ˆKØ˜jÑ(Ü!(¨°KÑ)@À!Ç*Á*Ó!M�—‘˜A’à!"�—‘˜A’ð		#ð ˆŒØˆŒáØ)¨-×*=Ñ*=¸cÓ*BÀQÑ*FÐ*GÐHˆMÜ5°mÀZÓPˆMØ*ˆÔØ"ˆŒØˆŒØ ˆŒØˆŒØˆŒð .0ˆŒØ6:×6NÑ6NÓ6PˆÔØˆŒð +0ˆÕ'r+   Úbuiltin_namespaceÚprintÚminÚmaxc                 óD   — t        | j                  j                  ||«      S rF   )r"   r
  r]   )r‚   r`   Úmessages      r)   Ú_unsupportedzCodeGenerator._unsupported`  s   € Ü+¨D¯K©K¯O©O¸TÀ7ÓKÐKr+   c                 ór   — t        «       }| j                  j                  ||«      }||u ryt        |«      ryy)NFT)Úobjectrš   rû   rO   )r‚   r&   Úabsent_markerrj   s       r)   Ú_is_constexpr_globalz"CodeGenerator._is_constexpr_globalc  s6   € Ü›ˆØ�k‰k�o‰o˜d MÓ2ˆØ�-ÑØä˜ÔØàr+   c                 ó|   ‡ ‡‡‡— dt         fˆ fd„Šdt         fˆ fd„Št        «       Šdt         dt        fˆˆˆˆ fd„}|S )Nr&   c                 ó<   •— ‰j                   j                  | |«      S rF   )r…   rû   )r&   Úabsentr‚   s     €r)   Úlocal_lookupz7CodeGenerator._define_name_lookup.<locals>.local_lookupp  s   ø€ à—;‘;—?‘? 4¨Ó0Ð0r+   c                 ób  •— ‰j                   j                  | |«      }t        ||u | ‰j                  v t	        |«      t
        u t        |t        «      t        |dd«      t        |dd«      t        |dd«      j                  d«      t        |dd«      j                  d«      t        |t        j                  «      t        |«      ‰j                  | «      ‰j                  t        j                   j"                  g«      r|S t%        t'        j(                  d| › d	�«      j+                  d
d«      «      ‚)NÚ__triton_builtin__FÚ__triton_aggregate__r”   r  ztriton.languagez"triton.experimental.gluon.languagez.                Cannot access global variable aà   from within @jit'ed
                function. Triton kernels can only access global variables that
                are instanstiated as constexpr (`x = triton.language.constexpr(42)`). Note that this is different from
                annotating a variable as constexpr (`x: triton.language.constexpr = 42`), which is not supported.  Alternatively, set the
                envvar TRITON_ALLOW_NON_CONSTEXPR_GLOBALS=1, but we do not
                promise to support this forever.ú
ú )rš   rû   r£   r,  r'   r   rG   r   r²   Ú
startswithr   rN   rk   r6  r)  r   ÚcompilationÚallow_non_constexpr_globalsÚ	NameErrorÚtextwrapÚdedentr:   )r&   r9  rj   r‚   s      €r)   Úglobal_lookupz8CodeGenerator._define_name_lookup.<locals>.global_lookupt  s  ø€ Ø—+‘+—/‘/ $¨Ó/ˆCô Ø˜6�MØ˜D×2Ñ2Ð2Ü˜“I¤Ð+Ü˜s¤KÓ0Ü˜CÐ!5°uÓ=Ü˜CÐ!7¸Ó?Ü˜C ¨rÓ2×=Ñ=Ð>OÓPÜ˜C ¨rÓ2×=Ñ=Ð>bÓcÜ˜s¤H§N¡NÓ3Ü" 3Ó'Ø×-Ñ-¨dÓ3ð ×3Ñ3Ü×%Ñ%×AÑAð!ô ð$ �
ÜÜ—‘ð %/Ø/3¨fð 51ð!4ó 5÷ 6=±W¸TÀ3Ó5GóIð Ir+   rD   c                 ó‚   •— ‰}‰‰‰j                   j                  fD ]  } || |«      }||usŒ|c S  t        | › d�«      ‚)Nz is not defined)r,  rû   rC  )r&   r9  Úlookup_functionro   r5  rF  r:  r‚   s       €€€€r)   Úname_lookupz6CodeGenerator._define_name_lookup.<locals>.name_lookup—  sW   ø€ Ø"ˆFØ#/°À×@VÑ@V×@ZÑ@ZÐ#Zò !�Ù'¨¨fÓ5�Ø Ò&Ø ’Lð!ô ˜t˜f OÐ4Ó5Ð5r+   )Ústrr4  r   )r‚   rI  r5  rF  r:  s   ` @@@r)   r'  z!CodeGenerator._define_name_lookupn  sG   û€ ð	1œsõ 	1ð	I¤õ 	IôB ›ˆð	6œcð 	6¤c÷ 	6ð 	6ð Ðr+   r&   ro   rD   c                 ó@   — || j                   |<   || j                  |<   y)zâ This function:
            called by visit_Assign() & visit_FunctionDef() to store left value (lvalue)
        1. record local defined name (FIXME: should consider control flow)
        2. store tensor in self.lvalue
        N)r…   rˆ   )r‚   r&   ro   s      r)   Ú	set_valuezCodeGenerator.set_value¡  s    € ð "ˆ�‰�DÑØ %ˆ�‰˜Òr+   c                 ór   — | j                   j                  «       }| j                   j                  «       }||fS rF   )rŠ   Úget_locr�   )r‚   ÚlocÚips      r)   Ú_get_insertion_point_and_locz*CodeGenerator._get_insertion_point_and_locª  s1   € ð �l‰l×"Ñ"Ó$ˆØ�\‰\×-Ñ-Ó/ˆØ�3ˆwˆr+   c                 óp   — | j                   j                  |«       | j                   j                  |«       y rF   )rŠ   r‘   r  )r‚   rP  rO  s      r)   Ú_set_insertion_point_and_locz*CodeGenerator._set_insertion_point_and_loc²  s&   € Ø�‰×,Ñ,¨RÔ0Ø�‰×Ñ˜SÕ!r+   c                 óˆ   — t        |«      s|g}|D ]/  }| j                  |«       t        |t        j                  «      sŒ/ y  y rF   )rX   rž   rG   rª   rÓ   )r‚   ÚstmtsÚstmts      r)   Úvisit_compound_statementz&CodeGenerator.visit_compound_statement¹  s?   € ä˜UÔ#Ø�GˆEØò 	ˆDØ�J‰J�tÔô ˜$¤§
¡
Õ+Ùñ	r+   c                 óD   — t         j                  j                  | |«       y rF   ©rª   ÚNodeVisitorr®   r»   s     r)   rÅ   zCodeGenerator.visit_ModuleÄ  ó   € Ü�‰×%Ñ% d¨DÕ1r+   c                 óÆ   — | j                  |j                  «      }|�J ‚t        j                  |j                  D �cg c]  }| j                  |«      ‘Œ c}«      }|S c c}w rF   )rž   r¶   r   rW   Úelts)r‚   r`   r¶   Úeltr]  s        r)   Ú
visit_ListzCodeGenerator.visit_ListÇ  sN   € Ø�j‰j˜Ÿ™Ó"ˆØˆ{Ðˆ{Ü�~‰~¸$¿)¹)ÖD°3˜tŸz™z¨#�ÒDÓEˆØˆùò Es   ½Ac                 ó&  ‡ ‡— ‰ j                  |j                  «      }g }ˆˆ fd„Š ‰|«      }|€t        j                  }nCt	        |t        j
                  j                  «      sJ ‚|j                  |«       |j                  }‰ j                  j                  |«       ‰ j                  €|‰ _        n*‰ j                  |k7  rt        d‰ j                  › d|› �«      ‚‰ j                  j                  «       }‰ j                  j                  |«       y )Nc                 óÔ   •— t        | t        j                  «      rt        | ‰«      S t        | t        j                  t
        t        f«      r‰j                  j                  | «      S | S rF   )	rG   r   rW   rs   r   ÚintÚfloatr  Ú	to_tensor)ro   Údecayr‚   s    €€r)   re  z)CodeGenerator.visit_Return.<locals>.decayÒ  sO   ø€ Ü˜%¤§¡Ô0Ü-¨e°UÓ;Ð;Ü˜E¤H×$6Ñ$6¼¼UÐ#CÔDØ—}‘}×.Ñ.¨uÓ5Ð5ØˆLr+   zInconsistent return types: ú and )rž   ro   r   ÚvoidrG   rM   r   rv   r'   rŠ   rA   r&  Ú	TypeErrorÚcreate_blockÚset_insertion_point_to_end)r‚   r`   Ú	ret_valuerw   Úret_tyÚpost_ret_blockre  s   `     @r)   r¼   zCodeGenerator.visit_ReturnÎ  sÝ   ù€ Ø—J‘J˜tŸz™zÓ*ˆ	Øˆõ	ñ ˜)Ó$ˆ	àÐÜ—]‘]‰Fä˜i¬¯©×)AÑ)AÔBÐBÐBØ×!Ñ! 'Ô*Ø—^‘^ˆFØ�‰×Ñ˜Ô!Ø�=‰=Ð Ø"ˆD�MØ�]‰]˜fÒ$ÜÐ9¸$¿-¹-¸ÈÈfÈXÐVÓWÐWð Ÿ™×2Ñ2Ó4ˆØ�‰×/Ñ/°Õ?r+   c                 óœ   — | j                  |j                  «      }t        |t        j                  j
                  «      sJ ‚|j                  S rF   )rž   ro   rG   r   rM   rW   rt   )r‚   r`   rb   s      r)   Úvisit_StarredzCodeGenerator.visit_Starredì  s8   € Ø�z‰z˜$Ÿ*™*Ó%ˆÜ˜$¤§¡× 3Ñ 3Ô4Ð4Ð4Ø�{‰{Ðr+   c                 ó¤  — | j                  |j                  «      \  }}| j                  r| j                  |d«      ‚t	        |j                  j
                  d d d…   «      D ]Ä  \  }}|j                  j                  | dz
     }|j                  }|j                  }t        j                  |t        j                  «       ¬«      }	|€t        j                  |	g|¬«      }
nt        j                  |	||¬«      }
	 | j                  rJ ‚d| _        | j                  |
«       d| _        ŒÆ | j                  rd	nd
}| j                  j!                  | j"                  «      }| j"                  j%                  | j&                  | j(                  ||| j*                  «      | _        | j&                  j-                  | j                  «       | j                  j/                  «       }| j                  j1                  | j                  «      }t3        ||«      D ]  \  }}| j5                  ||«       Œ | j"                  j7                  «       }| j"                  j9                  |«       | j;                  |j<                  «       | j"                  j7                  «       j?                  «       rJ ‚| j@                  �| j@                  tB        jD                  k(  r2tB        jD                  | _         | j"                  jG                  g «       �n	tI        | j@                  tB        jJ                  «      r&| j@                  jL                  | j                  _'        n| j@                  g| j                  _'        | j                  jQ                  | j                  j!                  | j"                  «      «       | j"                  jG                  | j                  jS                  | j"                  «      D �cg c]  }| j"                  jU                  |«      ‘Œ c}«       | j                  jW                  «        |r| j"                  jY                  |«       y y # d| _        w xY wc c}w )Nz,nested function definition is not supported.éÿÿÿÿr   )r±   r¶   ©Útargetsro   )Útargetro   Ú
annotationTFÚpublicÚprivate)-rž   rb   ra   r2  r\   Údefaultsru  rd   rª   r°   r·   rÔ   Ú	AnnAssignr)  r#  r   ró   rŠ   Úget_or_insert_functionr  r"  r[   Ú	push_backÚadd_entry_blockr  ÚziprL  r‹   Úset_insertion_point_to_startrW  r¤   Úhas_terminatorr&  r   rg  rA   rG   rn   rr   rÞ   Ú
reset_typeræ   Úcreate_poisonÚfinalizerj  )r‚   r`   r_   Úkwarg_namesr?   Údefault_valueÚarg_noderu  r&   Ú	st_targetÚ	init_nodeÚ
visibilityÚfn_tyÚentryÚ
arg_valuesÚarg_nameÚ	arg_valueÚ	insert_ptr=   s                      r)   rÈ   zCodeGenerator.visit_FunctionDefñ  s)  € Ø!%§¡¨D¯I©IÓ!6Ñˆ	�;Ø�7Š7Ø×#Ñ# DÐ*XÓYÐYä )¨$¯)©)×*<Ñ*<¹T¸r¸TÑ*BÓ Cò 	8ÑˆAˆ}Ø—y‘y—~‘~ q b¨1¡fÑ-ˆHØ!×,Ñ,ˆJØ—<‘<ˆDÜŸ™ D¬c¯i©i«kÔ:ˆIØÐ!ÜŸJ™J°	¨{À-ÔP‘	äŸM™M°À-Ð\fÔg�	ð8Ø×:Ò:Ð:Ð:Ø26�Ô/Ø—
‘
˜9Ô%à27�Õ/ð	8ð" "&§¢‘X°Yˆ
Ø—‘×(Ñ(¨¯©Ó6ˆØ—,‘,×5Ñ5°d·k±kÀ4×CUÑCUÐW\Ð^hÐjn×jwÑjwÓxˆŒØ�‰×Ñ˜dŸg™gÔ&Ø—‘×'Ñ'Ó)ˆØ—^‘^×/Ñ/°·±Ó8ˆ
ä#& y°*Ó#=ò 	0ÑˆH�iØ�N‰N˜8 YÕ/ð	0à—L‘L×4Ñ4Ó6ˆ	Ø�‰×1Ñ1°%Ô8à×%Ñ% d§i¡iÔ0ð —<‘<×3Ñ3Ó5×DÑDÔFÐFÐFØ�=‰=Ð  D§M¡M´X·]±]Ò$BÜ$ŸM™MˆDŒMØ�L‰L×Ñ˜RÖ ä˜$Ÿ-™-¬×)<Ñ)<Ô=Ø+/¯=©=×+>Ñ+>�—‘Õ(à,0¯M©M¨?�—‘Ô(Ø�G‰G×Ñ˜tŸ~™~×7Ñ7¸¿¹ÓEÔFØ�L‰L×ÑÀtÇ~Á~×GeÑGeÐfj×frÑfrÓGsÖtÀ˜dŸl™l×8Ñ8¸Õ<ÒtÔuØ�‰×ÑÔáØ�L‰L×3Ñ3°IÕ>ð øð? 38�Õ/üò8 us   Ã:&QÏ"QÑ	Q
c                 ó�   — g }|j                   D ]  }|| j                  |«      gz  }Œ | j                  |j                  «      }||fS rF   )rb   rž   Úkwarg)r‚   r`   r_   rd   rƒ  s        r)   Úvisit_argumentszCodeGenerator.visit_arguments&  sL   € Øˆ	Ø—9‘9ò 	+ˆCØ˜$Ÿ*™* S›/Ð*Ñ*‰Ið	+à—j‘j §¡Ó,ˆØ˜+Ð%Ð%r+   c                 óZ   — t         j                  j                  | |«       |j                  S rF   )rª   rZ  r®   rd   r»   s     r)   Ú	visit_argzCodeGenerator.visit_arg-  s   € Ü�‰×%Ñ% d¨DÔ1Ø�x‰xˆr+   c                 ób  — | j                  |j                  «      }| j                  |j                  «      }| j                  |j                  «      }|t        k(  rE|| j
                  v rt        |› d�«      ‚t	        |«      }|| j
                  |<   | j
                  |   S | j                  |«      S )Nz4 is already defined. constexpr cannot be reassigned.)rž   ru  rt  ro   r   r…   Ú
ValueErrorr¿   )r‚   r`   ru  rt  ro   s        r)   Úvisit_AnnAssignzCodeGenerator.visit_AnnAssign1  s£   € à—Z‘Z §¡Ó0ˆ
Ø—‘˜DŸK™KÓ(ˆØ—
‘
˜4Ÿ:™:Ó&ˆàœÒ"Ø˜Ÿ™Ñ$Ü  F 8ð ,Dð "Eó Fð Fä˜eÓ$ˆEØ"'ˆD�K‰K˜ÑØ—;‘;˜vÑ&Ð&à× Ñ  Ó&Ð&r+   c                 ój  — t        |j                  t        j                  «      sJ ‚t        |t        j                  «      r| j                  ||«      S t        |t        j                  «      r=t        |j                  «      D ]$  \  }}| j                  ||j                  |   «       Œ& y t        |t        j                  «      r3| j                  |j                  «      }t        ||j                  |«       y t        |t        j                   «      sJ ‚| j#                  | j                  |«      |«       y rF   )rG   r¶   rª   r·   Ú	SubscriptÚvisit_Subscript_Storer	   r\   r]  ÚassignTargetrt   rÒ   rž   ro   Úsetattrr³   r°   rL  )r‚   rt  ro   r?   Úbases        r)   rš  zCodeGenerator.assignTargetA  sÞ   € Ü˜&Ÿ*™*¤c§i¡iÔ0Ð0Ð0Ü�fœcŸm™mÔ,Ø×-Ñ-¨f°eÓ<Ð<Ü�fœcŸi™iÔ(Ü& v§{¡{Ó3ò ;‘	��6Ø×!Ñ! &¨%¯,©,°q©/Õ:ð;àÜ�fœcŸm™mÔ,Ø—:‘:˜fŸl™lÓ+ˆDÜ�D˜&Ÿ+™+ uÔ-ØÜ˜&¤#§(¡(Ô+Ð+Ð+Ø�‰�t—z‘z &Ó)¨5Õ1r+   c                 ó  ‡ ‡— ˆˆ fd„Š ‰‰ j                  |j                  «      «      }t        |t        j                  «      r|j
                  gn|j                  }t        |«      dk(  sJ ‚‰ j                  |d   |«       y )Nc                 ó  •— t        | t        j                  «      rt        | ‰«      S t        j                  t        j                  f}t        | «      } | �2t        | «      s't        | |«      s‰j                  j                  | «      } | S rF   )	rG   r   rW   rs   rN   r   rI   r  rd  )ro   Únative_nontensor_typesÚ_sanitize_valuer‚   s     €€r)   r   z3CodeGenerator.visit_Assign.<locals>._sanitize_valueR  sm   ø€ Ü˜%¤§¡Ô0Ü-¨e°_ÓEÐEÜ&.§n¡n´h·n±nÐ%EÐ"Ü(¨Ó/ˆEØÐ Ü$ UÔ+Ü˜uÐ&<Ô=ØŸ™×/Ñ/°Ó6�ØˆLr+   r   r   )	rž   ro   rG   rª   ry  rt  rs  r{   rš  )r‚   r`   rt   rs  r   s   `   @r)   r¿   zCodeGenerator.visit_AssignP  se   ù€ õ		ñ ! §¡¨D¯J©JÓ!7Ó8ˆÜ#-¨d´C·M±MÔ#B�4—;‘;‘-ÈÏÉˆÜ�7‹|˜qÒ Ð Ð Ø×Ñ˜' !™* fÕ-r+   c                 óP  — t        j                  |j                  «      }t        j                  «       |_        t        j                  ||j                  |j                  «      }t        j                  |j                  g|¬«      }| j                  |«       | j                  |«      S )Nrr  )r†   Údeepcopyrt  rª   ÚLoadr¶   ÚBinOpÚopro   rÔ   rž   )r‚   r`   ÚlhsÚrhsÚassigns        r)   rÂ   zCodeGenerator.visit_AugAssignb  sk   € Ü�m‰m˜DŸK™KÓ(ˆÜ—(‘(“*ˆŒÜ�i‰i˜˜TŸW™W d§j¡jÓ1ˆÜ—‘ T§[¡[ M¸Ô=ˆØ�
‰
�6ÔØ�z‰z˜#‹Ðr+   c                 óš   — t        |j                  «      t        j                  u r|j                  S | j                  |j                  «      S rF   )r'   r¶   rª   r·   r±   r(  r»   s     r)   r¸   zCodeGenerator.visit_Namej  s4   € Ü�—‘‹>œSŸY™YÑ&Ø—7‘7ˆNØ×$Ñ$ T§W¡WÓ-Ð-r+   c                 óD   — t         j                  j                  | |«       y rF   rY  r»   s     r)   Úvisit_StorezCodeGenerator.visit_Storeo  r[  r+   c                 óD   — t         j                  j                  | |«       y rF   rY  r»   s     r)   Ú
visit_LoadzCodeGenerator.visit_Loadr  r[  r+   c                 ó†   — |j                   D �cg c]  }| j                  |«      ‘Œ }}t        j                  |«      S c c}w rF   )r]  rž   r   rW   )r‚   r`   rö   rb   s       r)   Úvisit_TuplezCodeGenerator.visit_Tupleu  s5   € Ø'+§y¡yÖ1 !�—
‘
˜1•Ð1ˆÐ1Ü�~‰~˜dÓ#Ð#ùò 2s   �>c                 ón  — t        |«      r t        ||«      || j                  ¬«      S t        |«      r5t        j                  dd|«      } t        ||«      || j                  ¬«      S t        |t        t        j                  f«      st        |t        «      rt        |«      } t        ||«      |«      S )N©Ú	_semanticz__(.*)__z__r\1__)	rK   r²   r  r$   ÚsubrG   r   r   rW   )r‚   Úmethod_namer¦  r§  Úreverse_method_names        r)   Ú_apply_binary_methodz"CodeGenerator._apply_binary_methody  s”   € ä˜SÔ!Ø,”7˜3 Ó,¨S¸D¿M¹MÔJÐJÜ˜SÔ!Ü"$§&¡&¨°jÀ+Ó"NÐØ4”7˜3Ð 3Ó4°SÀDÇMÁMÔRÐRÜ˜#¤	¬8¯>©>Ð:Ô;Ä
È3ÔPYÔ@ZÜ˜C“.ˆCØ(Œw�s˜KÓ(¨Ó-Ð-r+   c                 ó^  — | j                  |j                  «      }| j                  |j                  «      }| j                  j	                  t        |j                  «      «      }|€5| j                  |dj                  |j                  j                  «      «      ‚| j                  |||«      S )Nz8AST binary operator '{}' is not (currently) implemented.)rž   ÚleftÚrightÚ_method_name_for_bin_oprû   r'   r¥  r2  Úformatr^   r¶  )r‚   r`   r¦  r§  r´  s        r)   Úvisit_BinOpzCodeGenerator.visit_BinOp„  s‘   € Ø�j‰j˜Ÿ™Ó#ˆØ�j‰j˜Ÿ™Ó$ˆØ×2Ñ2×6Ñ6´t¸D¿G¹G³}ÓEˆØÐØ×#Ñ# DØ$^×$eÑ$eÐfj×fmÑfm×fvÑfvÓ$wóyð yà×(Ñ(¨°c¸3Ó?Ð?r+   Ú__add__Ú__sub__Ú__mul__Ú__truediv__Ú__floordiv__Ú__mod__Ú__pow__Ú
__lshift__Ú
__rshift__Ú__and__Ú__or__Ú__xor__rº  c                 ó˜  — | j                   j                  |«       | j                  |j                  «       | j                   j	                  «       }| j
                  j                  «       }i }|j                  r†| j                   j                  |«       |j                  «       | _        i | _        | j                  |j                  «       | j
                  j                  «       }| j                   j	                  «       }g }|D ]©  }|df|dffD ]a  \  }	}
||	v sŒt        |	|   «      t        ||   «      k(  }|r |	|   j                  ||   j                  k(  rŒJJ d|› d||   › d|
› d|	|   › �«       ‚ ||v s||v r|j                  |«       ||v r||vr||   ||<   ||v sŒ�||vsŒ¢||   ||<   Œ« t        |j                  «       |j                  «       z  «      D ]i  }||v rŒ||   }|j                  }||   }|j                  }t        |«      t        |«      k(  }|r||k(  sJ d|› d|› d	|› d
�«       ‚|j                  |«       Œk |||||fS )NÚthenÚelsezinitial value for `z` is of type z
, but the z block redefines it as zMismatched type for z between then block (z) and else block (ú))rŠ   r~  rW  r¤   r‹   rˆ   r†   rÊ   r…   r'   Úappendr8   Úkeys)r‚   r`   r‡   Ú
then_blockÚ
else_blockÚ	then_defsÚ	else_defsÚnamesr&   ÚdefsÚ
block_nameÚ
type_equalÚthen_valÚthen_tyÚelse_valÚelse_tys                   r)   Úvisit_then_else_blocksz$CodeGenerator.visit_then_else_blocksœ  s~  € à�‰×1Ñ1°*Ô=Ø×%Ñ% d§i¡iÔ0Ø—\‘\×5Ñ5Ó7ˆ
Ø—O‘O×(Ñ(Ó*ˆ	àˆ	Ø�;Š;Ø�L‰L×5Ñ5°jÔAØ!Ÿ,™,›.ˆDŒKØ ˆDŒOØ×)Ñ)¨$¯+©+Ô6ØŸ™×,Ñ,Ó.ˆIØŸ™×9Ñ9Ó;ˆJð ˆàò 	0ˆDà&/°Ð%8¸9ÀfÐ:MÐ$Nò SÑ ��jØ˜4’<Ü!% d¨4¡jÓ!1´T¸'À$¹-Ó5HÑ!H�JÙ%¨$¨t©*¯/©/¸WÀT¹]×=OÑ=OÓ*Oð SØ-¨d¨V°=ÀÈÁÀð P#Ø#- ,Ð.EÀdÈ4ÁjÀ\ðSóSÐOðSð �yÑ  D¨IÑ$5Ø—‘˜TÔ"à�yÑ  T°Ñ%:Ø")¨$¡-�	˜$‘à�yÒ  T°Ò%:Ø")¨$¡-�	˜$’ð	0ô$ ˜9Ÿ>™>Ó+¨i¯n©nÓ.>Ñ>Ó?ò 	ˆDØ�u‰}ØØ  ‘ˆHØ—m‘mˆGØ  ‘ˆHØ—m‘mˆGÜ˜h›¬4°«>Ñ9ˆJÙ '¨WÒ"4ð .Ø& t fÐ,AÀ'Àð K#Ø#* )¨1ð.ó.Ð4ð �L‰L˜Õð	ð ˜) Z°¸UÐBÐBr+   c                 ó@  ‡‡— t        | «      5 }|\  }}| j                  j                  «       }| j                  j                  «       }| j                  j                  |«       | j                  j	                  |j
                  ||«       | j                  ||||«      \  ŠŠ}}}| j                  j                  «       }	| j                  j                  |«       |j                  «       rJ |› «       ‚t        ˆfd„|D «       «      }
| j                  j                  |	|
«       | j                  j                  |«       |j                  «       rJ |› «       ‚t        ˆfd„|D «       «      }| j                  j                  |	|«       t        |
«      t        |«      k(  sJ ‚t        |
|«      D ];  \  }}|j                  «       }||j                  «       k(  sJ ‚|	j                  |«       Œ= 	 d d d «       | j                  j                  	«       t        t        
«      «      D �cg c]  }|	j!                  |«      ‘Œ }}D �cg c]  }‰|   j"                  ‘Œ }}t%        ||«      }t        ||«      D ]  \  }}| j'                  ||«       Œ y # 1 sw Y   Œ£xY wc c}w c c}w )Nc              3   ó(   •K  — | ]	  }‰|   –— Œ y ­wrF   r–   ©rŸ   r&   rÑ  s     €r)   r¡   z3CodeGenerator.visit_if_top_level.<locals>.<genexpr>à  ó   øè ø€ Ò/RÀD°	¸$µÑ/Rùó   ƒc              3   ó(   •K  — | ]	  }‰|   –— Œ y ­wrF   r–   ©rŸ   r&   rÒ  s     €r)   r¡   z3CodeGenerator.visit_if_top_level.<locals>.<genexpr>å  rß  rà  )r   rŠ   ri  rj  Úcreate_cond_branchÚhandlerÛ  r  rx   Úcreate_branchr{   r}  Úget_typeÚadd_argumentr~  rù   rd   r'   r}   rL  )r‚   Úcondr`   Úsrr‡   Úip_blockrÏ  rÐ  rÓ  Úendif_blockÚthen_handlesÚelse_handlesÚthen_hÚelse_hr=   r?   Úres_handlesr&   rr   Ú
new_valuesÚ	new_valuerÒ  rÑ  s                        @@r)   Úvisit_if_top_levelz CodeGenerator.visit_if_top_levelÐ  sI  ù€ Ü˜dÓ#ð 	- rØ "ÑˆG�XØŸ™×2Ñ2Ó4ˆJØŸ™×2Ñ2Ó4ˆJà�L‰L×3Ñ3°HÔ=Ø�L‰L×+Ñ+¨D¯K©K¸ÀZÔPð ×+Ñ+¨D°'¸:ÀzÓRñ @ˆI�y *¨j¸%ð Ÿ,™,×3Ñ3Ó5ˆKà�L‰L×3Ñ3°JÔ?Ø!×0Ñ0Ô2ÐC°z°lÓCÐ2Ü/Ó/RÈEÔ/RÓRˆLØ�L‰L×&Ñ& {°LÔAà�L‰L×3Ñ3°JÔ?Ø!×0Ñ0Ô2ÐC°z°lÓCÐ2Ü/Ó/RÈEÔ/RÓRˆLØ�L‰L×&Ñ& {°LÔAÜ�|Ó$¬¨LÓ(9Ò9Ð9Ð9Ü"% l°LÓ"Aò -‘�˜Ø—_‘_Ó&�Ø˜VŸ_™_Ó.Ò.Ð.Ð.Ø×(Ñ(¨Õ,ñ-÷/	-ð: 	�‰×1Ñ1°+Ô>ä38¼¸\Ó9JÓ3KÖL¨a�{—‘ qÕ)ÐLˆÐLØ27Ö8¨$�˜4‘×%Ó%Ð8ˆÐ8Ü(¨°eÓ<ˆ
Ü" 5¨*Ó5ò 	,‰OˆD�)Ø�N‰N˜4 Õ+ñ	,÷E	-ð 	-üò> MùÚ8s   ŽGJ
È!JÈ?JÊ
Jc                 óB  ‡‡— t        | «      5 }|\  }}| j                  «       \  }}| j                  j                  «       }|j                  r| j                  j                  «       nd }	| j                  ||||	«      \  ŠŠ}}	}
t        ˆfd„|
D «       «      }| j                  ||«       | j                  j                  |D �cg c]  }|j                  «       ‘Œ c}|j                  d«      }|j                  |j                  «       «       | j                  j                  |j                  «       «       t        |
«      dkD  r| j                  j                  |«       |j                  s|j!                  «       }	n|	j                  |j!                  «       «       | j                  j                  |j!                  «       «       t        |
«      dkD  r/t        ˆfd„|
D «       «      }| j                  j                  |«       d d d «       t#        t        «      «      D �cg c]  }j%                  |«      ‘Œ }}
D �cg c]  }‰|   j&                  ‘Œ }}t)        ||«      }t+        |
|«      D ]  \  }}| j-                  ||«       Œ y c c}w # 1 sw Y   Œ�xY wc c}w c c}w )Nc              3   ó(   •K  — | ]	  }‰|   –— Œ y ­wrF   r–   rÞ  s     €r)   r¡   z-CodeGenerator.visit_if_scf.<locals>.<genexpr>   rß  rà  Tr   c              3   ó(   •K  — | ]	  }‰|   –— Œ y ­wrF   r–   râ  s     €r)   r¡   z-CodeGenerator.visit_if_scf.<locals>.<genexpr>  s   øè ø€ Ò3VÈ°I¸dµOÑ3Vùrà  )r   rQ  rŠ   ri  rÊ   rÛ  rx   rS  Úcreate_if_opræ  rä  Úmerge_block_beforeÚget_then_blockrj  r{   Úcreate_yield_opÚget_else_blockrù   Ú
get_resultr'   r}   r}  rL  )r‚   rè  r`   ré  r‡   r-   rP  Úlast_locrÏ  rÐ  rÓ  rì  ÚhÚif_oprí  r?   rð  r&   rr   rñ  rò  rÒ  rÑ  s                        @@r)   Úvisit_if_scfzCodeGenerator.visit_if_scf÷  s-  ù€ Ü˜dÓ#ð 	; rØ‰JˆG�QØ×<Ñ<Ó>‰LˆB�ØŸ™×2Ñ2Ó4ˆJØ8<¿º˜Ÿ™×2Ñ2Ô4ÈˆJà×+Ñ+¨D°'¸:ÀzÓRñ @ˆI�y *¨j¸%ô 0Ó/RÈEÔ/RÓRˆLØ×-Ñ-¨b°(Ô;Ø—L‘L×-Ñ-À\Ö.RÀ¨q¯z©z­|Ò.RÐTX×T_ÑT_ÐaeÓfˆEØ×)Ñ)¨%×*>Ñ*>Ó*@ÔAØ�L‰L×3Ñ3°E×4HÑ4HÓ4JÔKÜ�5‹z˜AŠ~Ø—‘×,Ñ,¨\Ô:Ø—;’;Ø"×1Ñ1Ó3‘
à×-Ñ-¨e×.BÑ.BÓ.DÔEØ�L‰L×3Ñ3°E×4HÑ4HÓ4JÔKÜ�5‹z˜AŠ~Ü3Ó3VÐPUÔ3VÓV�Ø—‘×,Ñ,¨\Ô:÷-	;ô0 5:¼#¸lÓ:KÓ4LÖM¨q�u×'Ñ'¨Õ*ÐMˆÐMØ27Ö8¨$�˜4‘×%Ó%Ð8ˆÐ8Ü(¨°eÓ<ˆ
Ü" 5¨*Ó5ò 	,‰OˆD�)Ø�N‰N˜4 Õ+ñ	,ùò# /S÷	;ð 	;üò0 NùÚ8s+   ŽB4JÃJ
ÃD&JÈJÈ;JÊJÊJc           	      óØ  — | j                  |j                  «      }t        |«      �r-t        |«      r| j	                  |d«      ‚|j
                  j                  «       rat        j                  dt        j                  |j                  «      z  «       t        j                  j                  || j                  | ¬«      }|j                  t        j                   | j                  ¬«      }t#        | j$                  «      j                  |«      }|r1| j&                  r| j	                  |d«      ‚| j)                  ||«       y | j+                  ||«       y t-        |«      }t        |«      t.        vrO| j	                  |dj1                  dj3                  d„ t.        D «       «      t        |«      j4                  «      «      ‚|r|j6                  n|j8                  }| j;                  |«       y )	Nz=Boolean value of Tensor with more than one value is ambiguousziIf conditional called with multidimensional Tensor instead of scalar; please use "if (%s).item()" instead)r²  Ú
_generatorr±  zÔCannot have `return` statements inside `while` or `for` statements in triton (note that this also applies to `return` statements that are inside functions transitively called from within `while`/`for` statements)úO`if` conditionals can only accept values of type {{{}}}, not objects of type {}ú, c              3   ó4   K  — | ]  }|j                   –— Œ y ­wrF   ©r^   ©rŸ   r-   s     r)   r¡   z)CodeGenerator.visit_If.<locals>.<genexpr>2  ó   è ø€ Ò!G° !§*¥*Ñ!Gùó   ‚)rž   ÚtestrK   rT   r2  r'   rR   ÚwarningsÚwarnrª   Úunparser   rM   Ú_unsplatr  ÚtoÚint1r˜   rš   r%  ró  r   r   Ú_condition_typesr»  r6   r^   r¤   rÊ   rW  )r‚   r`   rè  Úcontains_returnÚactive_blocks        r)   rË   zCodeGenerator.visit_If  sŠ  € Ø�z‰z˜$Ÿ)™)Ó$ˆä˜TÕ"Ü$ TÔ*Ø×'Ñ'¨Ð.mÓnÐnØ�y‰y×!Ñ!Ô#Ü—‘ð BÜ—k‘k $§)¡)Ó,ñ-ô.ô  —}‘}×-Ñ-¨d¸d¿m¹mÐX\Ð-Ó]�Ø—7‘7œ8Ÿ=™=°D·M±M�7ÓBˆDÜ3°D·K±KÓ@×FÑFÀtÓLˆOÙØ—>’>Ø×+Ñ+Øð TóUð Uð ×'Ñ'¨¨dÕ3à×!Ñ! $¨Õ-ä'¨Ó-ˆDä�D‹zÔ!1Ñ1Ø×'Ñ'ØÐk×rÑrØŸ	™	Ñ!GÔ6FÔ!GÓGÜ˜T›
×+Ñ+ó-ó.ð .ñ
 )-˜4Ÿ9š9°$·+±+ˆLØ×)Ñ)¨,Õ7r+   c           	      ó8  — | j                  |j                  «      }t        |«      �rE|j                  t        j
                  | j                  ¬«      }t        | «      5  | j                  «       \  }}| j                  j                  «       }| j                  j                  |«       | j                  j                  | j                  |j                  «      «      }| j                  j                  «       }| j                  j                  «       }| j                  j                  |«       | j                  j                  | j                  |j                  «      «      }| j                  j                  «       }| j!                  ||«       |j"                  |j"                  k(  s!J d|j"                  › d|j"                  › �«       ‚|j"                  }	|	t        j$                  k7  r|	j'                  | j                  «      gng }
| j                  j)                  |
|j*                  d«      }|j-                  |j/                  «       «       |
rO| j                  j1                  |j/                  «       «       | j                  j3                  |j*                  g«       | j                  j1                  |j/                  «       «       |j-                  |j5                  «       «       |
rO| j                  j1                  |j5                  «       «       | j                  j3                  |j*                  g«       |
r/t        j6                  j9                  |j;                  d«      |	«      nd cd d d «       S t=        |«      }t#        |«      t>        vrO| jA                  |djC                  djE                  d„ t>        D «       «      t#        |«      jF                  «      «      ‚|r| j                  |j                  «      S | j                  |j                  «      S # 1 sw Y   y xY w)	Nr±  zATernary expression with dynamic condition has inconsistent types rf  Tr   r  r  c              3   ó4   K  — | ]  }|j                   –— Œ y ­wrF   r  r  s     r)   r¡   z,CodeGenerator.visit_IfExp.<locals>.<genexpr>f  r  r	  )$rž   r
  rK   r  r   r  r  r   rQ  rŠ   ri  r~  rd  r¤   r‹   rÊ   rS  r'   rg  Úto_irr÷  rä  rø  rù  rj  rú  rû  rM   r   rü  r   r  r2  r»  r6   r^   )r‚   r`   rè  rP  rý  rÏ  r×  rÐ  rÙ  r&  Úret_type_irrÿ  s               r)   rÍ   zCodeGenerator.visit_IfExp8  s  € Ø�z‰z˜$Ÿ)™)Ó$ˆÜ˜TÕ"Ø—7‘7œ8Ÿ=™=°D·M±M�7ÓBˆDä! $Ó'ñ !dØ#×@Ñ@ÓB‘��Hà!Ÿ\™\×6Ñ6Ó8�
Ø—‘×9Ñ9¸*ÔEØŸ=™=×2Ñ2°4·:±:¸d¿i¹iÓ3HÓI�Ø!Ÿ\™\×=Ñ=Ó?�
à!Ÿ\™\×6Ñ6Ó8�
Ø—‘×9Ñ9¸*ÔEð  Ÿ=™=×2Ñ2°4·:±:¸d¿k¹kÓ3JÓK�Ø!Ÿ\™\×=Ñ=Ó?�
à×1Ñ1°"°hÔ?à—}‘}¨¯©Ò5ð |ØWÐX`×XeÑXeÐWfÐfkÐlt×lyÑlyÐkzÐ{ó|Ð5à#Ÿ=™=�à@HÌHÏMÉMÒ@Y˜xŸ~™~¨d¯l©lÓ;Ñ<Ð_a�ØŸ™×1Ñ1°+¸t¿{¹{ÈDÓQ�Ø×-Ñ-¨e×.BÑ.BÓ.DÔEÙØ—L‘L×;Ñ;¸E×<PÑ<PÓ<RÔSØ—L‘L×0Ñ0°(·/±/Ð1BÔCà—‘×7Ñ7¸×8LÑ8LÓ8NÔOØ×-Ñ-¨e×.BÑ.BÓ.DÔEÙØ—L‘L×;Ñ;¸E×<PÑ<PÓ<RÔSØ—L‘L×0Ñ0°(·/±/Ð1BÔCÙNY”x—}‘}×+Ñ+¨E×,<Ñ,<¸QÓ,?ÀÔJÐ_c÷C!dñ !dôF (¨Ó-ˆDô �D‹zÔ!1Ñ1Ø×'Ñ'ØÐk×rÑrØŸ	™	Ñ!GÔ6FÔ!GÓGÜ˜T›
×+Ñ+ó-ó.ð .ñ Ø—z‘z $§)¡)Ó,Ð,à—z‘z $§+¡+Ó.Ð.÷]!dð !dús   ÁLPÐPc                  ó   — y rF   r–   r»   s     r)   Ú
visit_PasszCodeGenerator.visit_Passm  s   € Ør+   c                 óô  — t        |j                  «      dk(  rt        |j                  «      dk(  s| j                  |d«      ‚| j	                  |j
                  «      }| j	                  |j                  d   «      }t        |«      }t        |«      }t        |j                  d   «      t        j                  u rt        ||u «      S t        |j                  d   «      t        j                  u rt        ||u«      S | j                  j                  t        |j                  d   «      «      }|€8| j                  |dj                  |j                  d   j                  «      «      ‚| j!                  |||«      S )Nr   z1simultaneous multiple comparison is not supportedr   z<AST comparison operator '{}' is not (currently) implemented.)r{   ÚcomparatorsÚopsr2  rž   r¸  r   r'   rª   ÚIsr   ÚIsNotÚ_method_name_for_comp_oprû   r»  r^   r¶  )r‚   r`   r¦  r§  Ú	lhs_valueÚ	rhs_valuer´  s          r)   Úvisit_ComparezCodeGenerator.visit_Comparep  s@  € Ü�D×$Ñ$Ó%¨Ò*¬s°4·8±8«}ÀÒ/AØ×#Ñ# DÐ*]Ó^Ð^Ø�j‰j˜Ÿ™Ó#ˆØ�j‰j˜×)Ñ)¨!Ñ,Ó-ˆÜ(¨Ó-ˆ	Ü(¨Ó-ˆ	Ü�—‘˜‘Ó¤§¡Ñ&Ü˜Y¨)Ð3Ó4Ð4Ü�—‘˜‘Ó¤§	¡	Ñ)Ü˜Y¨iÐ7Ó8Ð8Ø×3Ñ3×7Ñ7¼¸T¿X¹XÀa¹[Ó8IÓJˆØÐØ×#Ñ#ØÐT×[Ñ[Ð\`×\dÑ\dÐefÑ\g×\pÑ\pÓqósð sà×(Ñ(¨°c¸3Ó?Ð?r+   Ú__eq__Ú__ne__Ú__lt__Ú__le__Ú__gt__Ú__ge__r  c           
      óö  — | j                  |j                  «      }| j                  j                  t	        |j
                  «      «      }|€*| j                  |d|j
                  j                  › d�«      ‚t        |«      r t        ||«      | j                  ¬«      S 	  t        ||«      «       S # t        $ r? |dk(  rt        | «      cY S | j                  |d|› dt	        |«      j                  › �«      ‚w xY w)NzAST unary operator 'z!' is not (currently) implemented.r±  Ú__not__z)' is not (currently) implemented on type )rž   ÚoperandÚ_method_name_for_unary_oprû   r'   r¥  r2  r^   rK   r²   r  ÚAttributeErrorr   )r‚   r`   r+  ra   s       r)   Úvisit_UnaryOpzCodeGenerator.visit_UnaryOp…  sô   € Ø—*‘*˜TŸ\™\Ó*ˆØ×+Ñ+×/Ñ/´°T·W±W³Ó>ˆØˆ:Ø×#Ñ# DÐ,@ÀÇÁ×AQÑAQÐ@RÐRsÐ*tÓuÐuÜ˜WÔ%Ø'”7˜7 BÓ'°$·-±-Ô@Ð@ð	tØ'”7˜7 BÓ'Ó)Ð)øÜò 	tØ�YŠÜ  W Ó-Ò-Ø×#Ñ#ØÐ,¨R¨DÐ0YÔZ^Ð_fÓZg×ZpÑZpÐYqÐrótð tð	tús   ÂB0 Â0C8Ã,C8Ú__neg__Ú__pos__r*  Ú
__invert__r,  c                 ó2  — t        |«      sJ d|› d�«       ‚t        |«      sJ d|› d�«       ‚t        |«      t        |«      u sJ d|› d�«       ‚t        |«      r?|j                  |j                  k(  s%J d|› d|j                  › d|j                  › d	�«       ‚y y )
Nzcannot reassign constxpr z in the loopzcannot reasign constexpr zLoop carried variable z changed typezLoop-carried variable z has initial type z but is re-assigned to z: in loop! Please make sure that the type stays consistent.)rI   r'   rK   )r‚   r&   Úloop_valÚlive_vals       r)   Ú_verify_loop_carried_variablez+CodeGenerator._verify_loop_carried_variable˜  s¹   € Ü Ô)ÐYÐ-FÀtÀfÈLÐ+YÓYÐ)Ü Ô)ÐYÐ-FÀtÀfÈLÐ+YÓYÐ)Ü�H‹~¤ h£Ñ/Ð]Ð3IÈ$ÈÈ}Ð1]Ó]Ð/Ü$ XÔ.°(·-±-À8Ç=Á=Ò2Pð 	@Ø$ T FÐ*<¸X¿]¹]¸Oð L%Ø%-§]¡] Oð 4?ð@ó	@ÐPÐ2PÐ.r+   c                 ó8   — | j                  |j                  «      S rF   )rž   Úcontext_exprr»   s     r)   Úvisit_withitemzCodeGenerator.visit_withitem¡  s   € Ø�z‰z˜$×+Ñ+Ó,Ð,r+   c                 óÜ  — t        |j                  «      dk(  sJ ‚|j                  d   j                  }| j                  |j                  «      }|t
        j                  k(  r_|j                  D �cg c]  }| j                  |«      ‘Œ }} ||d| j                  iŽ5  | j                  |j                  «       d d d «       y | j                  |j                  «       y c c}w # 1 sw Y   y xY w)Nr   r   Ú_builder)r{   rý   r7  rž   rÏ   r   Ú
async_taskrb   rŠ   rW  r¤   )r‚   r`   r  ÚwithitemClassrd   rb   s         r)   Ú
visit_WithzCodeGenerator.visit_With¤  sÀ   € Ü�4—:‘:‹ !Ò#Ð#Ð#Ø—*‘*˜Q‘-×,Ñ,ˆØŸ
™
 7§<¡<Ó0ˆØœH×/Ñ/Ò/Ø/6¯|©|Ö<¨�D—J‘J˜s•OÐ<ˆDÐ<Ù Ð<¨t¯|©|Ñ<ñ 9Ø×-Ñ-¨d¯i©iÔ8÷9ð 9ð ×)Ñ)¨$¯)©)Õ4ùò	 =÷9ð 9ús   Á0CÂC"Ã"C+c                 ó$	  — t        | «      5 }|\  }}| j                  «       \  }}| j                  j                  «       }| j                  j	                  |«       | j
                  j                  |«       | j                  |j                  «       | j
                  j                  «        | j                  }|j                  «        g }	g }
|D ]F  }||v sŒ||   }||   }| j                  |||«       |	j                  |«       |
j                  |«       ŒH t        |
«      }|D �cg c]  }|j                  «       ‘Œ }}|
D �cg c]  }|j                  ‘Œ }}| j!                  ||«       | j                  j#                  ||«      }| j                  j%                  |j'                  «       |«      }| j                  j	                  |«       t)        t+        |«      «      D �cg c]  }|j-                  |«      ‘Œ }}t/        ||«      }t1        |	|«      D ]#  \  }}|| j2                  |<   || j                  |<   Œ% | j5                  |j6                  «      }| j                  j9                  |«       | j                  j;                  |j<                  |«       | j                  j%                  |j?                  «       |«      }| j                  j	                  |«       t)        t+        |«      «      D �cg c]  }|j-                  |«      ‘Œ }}t/        ||«      }t1        |	|«      D ]#  \  }}|| j2                  |<   || j                  |<   Œ% | j
                  j                  |«       | j                  |j                  «       | j
                  j                  «        | j                  }g }|D ]  }||v sŒ||   jA                  |«       Œ | j                  jC                  |«       d d d «       t)        t+        «      «      D �cg c]  }jE                  |«      ‘Œ }}t/        |«      }t1        	|«      D ]#  \  }} | | j2                  |<   | | j                  |<   Œ% |jF                  D ]  }!J d«       ‚ y c c}w c c}w c c}w c c}w # 1 sw Y   Œ£xY wc c}w )NzNot implemented)'r   rQ  rŠ   ri  r~  r%  rÍ  rW  r¤   Úpoprˆ   Úeraser5  rx   ræ  r'   rS  Úcreate_while_opÚcreate_block_with_parentÚ
get_beforerù   r{   rd   r}   r}  r…   rž   r
  rj  Úcreate_condition_oprä  Ú	get_afterrv   rú  rü  rÊ   rª   rZ  r®   )"r‚   r`   ré  r‡   rŒ   rP  rý  ÚdummyÚ	loop_defsrÓ  Ú	init_argsr&   r3  r4  Úinit_handlesrþ  Úinit_tysÚaÚinit_fe_tysÚwhile_opÚbefore_blockr?   Ú
block_argsÚcondition_argsrj   rè  Úafter_blockÚbody_handlesÚ	body_argsÚyieldsÚresult_handlesÚresult_valsÚnew_defrV  s"                                     r)   Úvisit_WhilezCodeGenerator.visit_While¯  s  € Ü˜dÓ#ð @	1 rØ$&Ñ!ˆG�\Ø×<Ñ<Ó>‰LˆB�ð —L‘L×-Ñ-Ó/ˆEØ�L‰L×5Ñ5°eÔ<Ø�N‰N×!Ñ! $Ô'Ø×)Ñ)¨$¯)©)Ô4Ø�N‰N×ÑÔ ØŸ™ˆIØ�K‰KŒMð ˆEØˆIØ!ò 	/�Ø˜7’?à(¨™�HØ& t™}�HØ×6Ñ6°t¸XÀxÔPð —L‘L Ô&Ø×$Ñ$ XÕ.ð	/ô 0°	Ó:ˆLØ.:Ö;¨˜Ÿ
™
�Ð;ˆHÐ;Ø+4Ö5 a˜1Ÿ6›6Ð5ˆKÐ5Ø×-Ñ-¨b°(Ô;Ø—|‘|×3Ñ3°H¸lÓKˆHàŸ<™<×@Ñ@À×ATÑATÓAVÐX`ÓaˆLØ�L‰L×5Ñ5°lÔCÜ7<¼SÀÓ=NÓ7OÖP°!˜,×*Ñ*¨1Õ-ÐPˆJÐPÜ0°¸[ÓIˆNÜ  ¨Ó7ò ,‘	��cØ$'�—‘˜DÑ!Ø(+�—‘ Ò%ð,ð —:‘:˜dŸi™iÓ(ˆDØ�L‰L×3Ñ3°LÔAà�L‰L×,Ñ,¨T¯[©[¸*ÔEàŸ,™,×?Ñ?À×@RÑ@RÓ@TÐV^Ó_ˆKð �L‰L×5Ñ5°kÔBÜ8=¼cÀ,Ó>OÓ8PÖQ°1˜KŸO™O¨AÕ.ÐQˆLÐQÜ+¨L¸+ÓFˆIÜ  ¨	Ó2ò ,‘	��cØ$'�—‘˜DÑ!Ø(+�—‘ Ò%ð,ð �N‰N×!Ñ! $Ô'Ø×)Ñ)¨$¯)©)Ô4Ø�N‰N×ÑÔ ØŸ™ˆIØˆFØ!ò 8�Ø˜7’?Ø˜d‘O×/Ñ/°Õ7ð8ð �L‰L×(Ñ(¨Ô0÷A@	1ôF ;@ÄÀLÓ@QÓ:RÖS°Q˜(×-Ñ-¨aÕ0ÐSˆÐSÜ)¨.¸+ÓFˆÜ  ¨Ó4ò 	,‰MˆD�'Ø 'ˆD�K‰K˜ÑØ$+ˆD�O‰O˜DÒ!ð	,ð —K‘Kò 	6ˆDØ+Ð+Ó+�5ñ	6ùòY <ùÚ5ùò Qùò R÷e@	1ð @	1üòF TsX   ŒCRÃARÄ$Q-Ä;RÅQ2ÅBRÇQ7Ç7C7RË.Q<ÌB'RÎ.1RÏ=RÑ-RÒR
c                 ó  — t        |j                  t        j                  «      sJ ‚| j	                  |j
                  «      }| j	                  |j                  «      }t        |«      r|j                  || j                  ¬«      S ||   S )Nr±  )
rG   r¶   rª   r£  rž   ro   ÚslicerK   Ú__getitem__r  )r‚   r`   r¦  Úslicess       r)   Úvisit_Subscript_Loadz"CodeGenerator.visit_Subscript_Loadý  sg   € Ü˜$Ÿ(™(¤C§H¡HÔ-Ð-Ð-Ø�j‰j˜Ÿ™Ó$ˆØ—‘˜DŸJ™JÓ'ˆÜ˜SÔ!Ø—?‘? 6°T·]±]�?ÓCÐCØ�6‰{Ðr+   c                 ó  — t        |j                  t        j                  «      sJ ‚| j	                  |j
                  «      }| j	                  |j                  «      }t        |t        j                  «      sJ ‚|j                  ||«       y rF   )
rG   r¶   rª   r·   rž   ro   rZ  r   rW   Ú__setitem__)r‚   r`   ro   r¦  r\  s        r)   r™  z#CodeGenerator.visit_Subscript_Store  s`   € Ü˜$Ÿ(™(¤C§I¡IÔ.Ð.Ð.Ø�j‰j˜Ÿ™Ó$ˆØ—‘˜DŸJ™JÓ'ˆÜ˜#œxŸ~™~Ô.Ð.Ð.Ø�‰˜ Õ&r+   c                 ó$   — | j                  |«      S rF   )r]  r»   s     r)   Úvisit_SubscriptzCodeGenerator.visit_Subscript  s   € Ø×(Ñ(¨Ó.Ð.r+   c                 ó^   — |j                   D �cg c]  }| j                  |«      ‘Œ c}S c c}w rF   )Údimsrž   )r‚   r`   Údims      r)   Úvisit_ExtSlicezCodeGenerator.visit_ExtSlice  s"   € Ø+/¯9©9Ö5 C�—
‘
˜3•Ò5Ð5ùÒ5s   �*c           	      óÞ  ‡ — ‰ j                  |j                  j                  «      }|j                  j                  D �cg c]  }‰ j                  |«      ‘Œ }}t	        ˆ fd„|j                  j
                  D «       «      }|t        j                  k(  rÑ ||i |¤Ž}t        |j                  j                  |j                  j                  |j                  j                  «      }|D ]z  }t        |«      ‰ j                  |j                  j                   <   ‰ j#                  |j$                  «       |j&                  D ]"  }	t(        j*                  j-                  ‰ |	«       Œ$ Œ| y d }
d }d}d}d}|t        j                  u ri ||i |¤Ž}|j                  }|j                  }|j                  }|j.                  }
|j0                  }|j2                  }|j4                  }|j6                  }n½|t        u rªt9        |«      dkD  r|d   n#‰ j                  t)        j:                  d«      «      }t9        |«      dkD  r|d   n'‰ j                  |j                  j                  d   «      }t9        |«      dkD  r|d   n#‰ j                  t)        j:                  d«      «      }nt=        d«      ‚d}t?        |«      r+|j                  dk  rt        |j                   «      }d}||}}‰ j@                  jC                  |«      }‰ j@                  jC                  |«      }‰ j@                  jC                  |«      }|jD                  jG                  «       r4|jD                  jG                  «       r|jD                  jG                  «       s3tI        d|jD                  › d	|jD                  › d	|jD                  › d
�«      ‚‰ j@                  jK                  |jD                  |jD                  «      }‰ j@                  jK                  ||jD                  «      }|jM                  ‰ jN                  «      }|jP                  t        jR                  jD                  jT                  jV                  k(  }|jX                  }|jX                  }|jX                  }‰ jN                  j[                  |||«      }‰ jN                  j[                  |||«      }‰ jN                  j[                  |||«      }‰ jN                  j]                  |«      }‰ j_                  |j                  j                   t        jR                  ja                  ||«      «       tc        ‰ «      5 }|\  }}‰ je                  «       \  }}‰ jN                  jg                  «       }‰ jN                  ji                  |«       ‰ jj                  jm                  |«       ‰ j#                  |j$                  «       ‰ jj                  jo                  «        |jq                  «        g }g }g }‰ jr                  D ]a  } | |v sŒ‰ jr                  |    }!||    }"‰ ju                  | |!|"«       |jm                  | «       |jm                  |"«       |jm                  |!«       Œc ‰ jw                  ||«       ty        |«      }#|D �$cg c]  }$|$jz                  ‘Œ }%}$‰ jN                  j}                  ||||#«      }&t        |
«      �+|&j�                  d‰ jN                  jƒ                  |
«      «       t        |«      �+|&j�                  d‰ jN                  jƒ                  |«      «       |r*|&j�                  d‰ jN                  j…                  «       «       |r*|&j�                  d‰ jN                  j…                  «       «       |r*|&j�                  d‰ jN                  j…                  «       «       ‰ jj                  jm                  |«       |&j‡                  d«      }'‰ jN                  ji                  |'«       |j‰                  «       ‰ _        i ‰ _9        t        t9        |#«      «      D �cg c]  }|'j‹                  |dz   «      ‘Œ }(}t�        |(|%«      })t�        ||)«      D ]  \  } }*‰ j_                  | |*«       Œ ‰ j#                  |j$                  «       ‰ jj                  jo                  «        g }‰ jr                  D ]R  } | |v sŒ‰ jr                  |    }+t‘        |+t        «      r‰ j@                  jC                  |+«      }+|jm                  |+«       ŒT t9        |«      dkD  r&ty        |«      },‰ jN                  j“                  |,«       |'j•                  «       }-|-j—                  «       dk(  sJ d«       ‚‰ jN                  ji                  |'«       |&j™                  «       }|r8‰ jN                  j›                  ||«      }‰ jN                  j�                  ||«      }‰ j                  |j                  j                      jX                  jŸ                  |«       ‰ j_                  |j                  j                   t        jR                  ja                  ||«      «       d d d «       t        t9        #«      «      D �cg c]  }&j¡                  |«      ‘Œ }.}t�        |.%«      }/t�        |/«      D ]  \  } }*‰ j_                  | |*«       Œ |j&                  D ]  }	J d«       ‚ y c c}w c c}$w c c}w # 1 sw Y   Œ’xY wc c}w )Nc              3   ó@   •K  — | ]  }‰j                  |«      –— Œ y ­wrF   r�   ©rŸ   Úkeywordr‚   s     €r)   r¡   z*CodeGenerator.visit_For.<locals>.<genexpr>  s   øè ø€ ÒQ°7˜4Ÿ:™: g×.ÑQùr¢   Fr   r   r   zAOnly `range` and `static_range` iterators are currently supportedTz0For loop bounds and step must all be ints, are (r  rÌ  ztt.num_stagesztt.loop_unroll_factorztt.disallow_acc_multi_bufferz
tt.flattenztt.warp_specializez7We use SCF, so the loop body should only have one blockz)Don't know what to do with else after for)Qrž   ÚiterrÏ   rb   ÚdictÚkeywordsr   Ústatic_rangerù   Ústartro   ÚendÚstepr   r…   rt  r±   rW  r¤   rÊ   rª   rZ  r®   Ú
num_stagesÚloop_unroll_factorÚdisallow_acc_multi_bufferÚflattenÚwarp_specializer{   ÚNumÚRuntimeErrorrO   r  rd  rN   Úis_intrh  Úinteger_promote_implr  rŠ   Úint_signednessrM   Ú
SIGNEDNESSÚSIGNEDrä  Úcreate_int_castr�  rL  r   r   rQ  ri  r~  r%  rÍ  r?  r@  rˆ   r5  rS  rx   r'   Úcreate_for_opr   Úset_attrÚget_int32_attrÚget_unit_attrÚget_bodyr†   rd   r}   r}  rG   rú  Ú
get_parentÚsizeÚget_induction_varÚ
create_subÚ
create_addÚreplace_all_uses_withrü  )0r‚   r`   ÚIteratorClassrd   Ú	iter_argsÚiter_kwargsÚiteratorrm  r?   rV  rq  rr  rs  rt  ru  ÚlbÚubrp  Únegative_stepÚiv_typeÚ
iv_ir_typeÚiv_is_signedÚivré  r‡   rŒ   rP  rý  ÚblockrH  rT  rÓ  r&   r3  r4  rI  rp   rJ  Úfor_opÚfor_op_bodyÚblock_handlesrO  rj   ÚlocalÚyield_handlesÚfor_op_regionrU  Úresult_valuess0   `                                               r)   Ú	visit_ForzCodeGenerator.visit_For  sÈ  ø€ ØŸ
™
 4§9¡9§>¡>Ó2ˆØ04·	±	·±Ö?¨�T—Z‘Z •_Ð?ˆ	Ð?ÜÓQ¸d¿i¹i×>PÑ>PÔQÓQˆØœH×1Ñ1Ò1Ù$ iÐ?°;Ñ?ˆHÜ  §¡×!5Ñ!5°x·|±|×7IÑ7IÈ8Ï=É=×K^ÑK^Ó_ˆLØ!ò >�Ü.7¸«l�—‘˜DŸK™KŸN™NÑ+Ø×-Ñ-¨d¯i©iÔ8Ø ŸK™Kò >�DÜ—O‘O×1Ñ1°$¸Õ=ñ>ð>ð
 Øˆ
Ø!ÐØ$)Ð!ØˆØˆØœHŸN™NÑ*Ù$ iÐ?°;Ñ?ˆHð —‘ˆBØ—‘ˆBØ—=‘=ˆDØ!×,Ñ,ˆJØ!)×!<Ñ!<ÐØ(0×(JÑ(JÐ%Ø×&Ñ&ˆGØ&×6Ñ6‰OØœeÑ#ô "% Y£°!Ò!3�˜1’¸¿¹ÄCÇGÁGÈAÃJÓ9OˆBÜ!$ Y£°!Ò!3�˜1’¸¿¹ÀDÇIÁIÇNÁNÐSTÑDUÓ9VˆBÜ#& y£>°AÒ#5�9˜Q’<¸4¿:¹:ÄcÇgÁgÈaÃjÓ;Q‰DäÐbÓcÐcàˆÜ˜Ô 4§:¡:°¢>Ü˜dŸj™j˜[Ó)ˆDØ ˆMØ˜�ˆBØ�]‰]×$Ñ$ RÓ(ˆØ�]‰]×$Ñ$ RÓ(ˆØ�}‰}×&Ñ& tÓ,ˆà�x‰x�‰Ô ¨¯©¯©Ô(9ÀÇÁ×ARÑARÔATÜÐNÈrÏxÉxÈjÐXZÐ[]×[cÑ[cÐZdÐdfÐgk×gqÑgqÐfrÐrsÐtÓuÐuØ—-‘-×4Ñ4°R·X±X¸r¿x¹xÓHˆØ—-‘-×4Ñ4°W¸d¿j¹jÓIˆØ—]‘] 4§<¡<Ó0ˆ
Ø×-Ñ-´·±×1DÑ1D×1OÑ1O×1VÑ1VÑVˆà�Y‰YˆØ�Y‰YˆØ�{‰{ˆà�\‰\×)Ñ)¨"¨j¸,ÓGˆØ�\‰\×)Ñ)¨"¨j¸,ÓGˆØ�|‰|×+Ñ+¨D°*¸lÓKˆà�\‰\×'Ñ'¨
Ó3ˆØ�‰�t—{‘{—~‘~¤x§}¡}×';Ñ';¸BÀÓ'HÔIä˜dÓ#ð N	N rØ$&Ñ!ˆG�\Ø×<Ñ<Ó>‰LˆB�ð —L‘L×-Ñ-Ó/ˆEØ�L‰L×5Ñ5°eÔ<à�N‰N×!Ñ! $Ô'Ø×)Ñ)¨$¯)©)Ô4Ø�N‰N×ÑÔ Ø�K‰KŒMð ˆIØˆFØˆEØŸ™ò ,�Ø˜7’?Ø#Ÿ™¨tÑ4�HØ& t™}�HØ×6Ñ6°t¸XÀxÔPà—L‘L Ô&Ø×$Ñ$ XÔ.Ø—M‘M (Õ+ð,ð ×-Ñ-¨b°(Ô;Ü/°	Ó:ˆLØ(1Ö2 1˜Ÿ›Ð2ˆHÐ2Ø—\‘\×/Ñ/°°B¸¸lÓKˆFÜ# JÓ/Ð;Ø—‘ °·±×1LÑ1LÈZÓ1XÔYÜ#Ð$6Ó7ÐCØ—‘Ð 7¸¿¹×9TÑ9TÐUgÓ9hÔiÙ(Ø—‘Ð >ÀÇÁ×@ZÑ@ZÓ@\Ô]ÙØ—‘ ¨d¯l©l×.HÑ.HÓ.JÔKÙØ—‘Ð 4°d·l±l×6PÑ6PÓ6RÔSà�N‰N×!Ñ! $Ô'Ø Ÿ/™/¨!Ó,ˆKØ�L‰L×5Ñ5°kÔBà!Ÿ,™,›.ˆDŒKØ ˆDŒOÜ=BÄ3À|ÓCTÓ=UÖV¸˜[Ÿ_™_¨Q°©UÕ3ÐVˆMÐVÜ,¨]¸HÓEˆJÜ  ¨
Ó3ò *‘	��cØ—‘˜t SÕ)ð*à×)Ñ)¨$¯)©)Ô4Ø�N‰N×ÑÔ ØˆFØŸ™ò )�Ø˜7’?Ø ŸO™O¨DÑ1�EÜ! %¬Ô3Ø $§¡× 7Ñ 7¸Ó >˜Ø—M‘M %Õ(ð)ô �6‹{˜QŠÜ 4°VÓ <�Ø—‘×,Ñ,¨]Ô;Ø'×2Ñ2Ó4ˆMØ ×%Ñ%Ó'¨1Ò,ÐgÐ.gÓgÐ,ð �L‰L×5Ñ5°kÔBØ×)Ñ)Ó+ˆBÙØ—\‘\×,Ñ,¨R°Ó4�Ø—\‘\×,Ñ,¨R°Ó4�Ø�K‰K˜Ÿ™Ÿ™Ñ'×.Ñ.×DÑDÀRÔHØ�N‰N˜4Ÿ;™;Ÿ>™>¬8¯=©=×+?Ñ+?ÀÀGÓ+LÔM÷]N	Nôb 9>¼cÀ,Ó>OÓ8PÖQ°1˜&×+Ñ+¨AÕ.ÐQˆÐQÜ+¨N¸HÓEˆÜ˜U MÓ2ò 	&‰IˆD�#Ø�N‰N˜4 Õ%ð	&ð —K‘Kò 	6ˆDØEÐEÓE�5ñ	6ùòq @ùòB 3ùò& W÷eN	Nð N	Nüòb RsF   ¿kÕCkØ
A=kÚkÚF	kà#kà>A>kâ=Fké+k*ë
këk'c                 óÒ   — | j                  |j                  «      }| j                  |j                  «      }| j                  |j                  «      }t	        j
                  |||«      S rF   )rž   ÚlowerÚupperrp  r   rZ  )r‚   r`   rž  rŸ  rp  s        r)   Úvisit_SlicezCodeGenerator.visit_Slice°  sK   € Ø—
‘
˜4Ÿ:™:Ó&ˆØ—
‘
˜4Ÿ:™:Ó&ˆØ�z‰z˜$Ÿ)™)Ó$ˆÜ�~‰~˜e U¨DÓ1Ð1r+   c                 ó8   — | j                  |j                  «      S rF   )rž   ro   r»   s     r)   Úvisit_IndexzCodeGenerator.visit_Index¶  s   € Ø�z‰z˜$Ÿ*™*Ó%Ð%r+   c                 óP   — |j                   | j                  |j                  «      fS rF   )rd   rž   ro   r»   s     r)   Úvisit_keywordzCodeGenerator.visit_keyword¹  s   € Ø�x‰x˜Ÿ™ D§J¡JÓ/Ð/Ð/r+   c                 óâ   — | j                  |j                  «      }|j                  �| j                  |j                  «      nd}t        j                  j                  ||| j                  ¬«      S )Nr  r±  )rž   r
  Úmsgr   rM   Údevice_assertr  )r‚   r`   r
  r¦  s       r)   Úvisit_AssertzCodeGenerator.visit_Assert¼  sQ   € Ø�z‰z˜$Ÿ)™)Ó$ˆØ&*§h¡hÐ&:ˆd�j‰j˜Ÿ™Ô"ÀˆÜ�}‰}×*Ñ*¨4°ÀÇÁÐ*ÓNÐNr+   ra   c                 ó`  — t        j                  |j                  g|¢­i |¤Ž}|j                  D �cg c]  }||   ‘Œ	 }}t	        |«      D ]W  \  }}t        |t        j                  t        t        t        t        f«      sŒ6t        j                  j                  |«      ||<   ŒY t        |d„ «      }|D �ci c]  }|t        ||«      “Œ }}t        |d„ «      }	|	D �cg c]  }t        ||«      ‘Œ }
}t!        t#        |«      |
D �cg c]  }|j$                  ‘Œ c}|«      }| j&                  j)                  |«      �sNt+        |«      \  }}|D �cg c]Y  }|�/t        |t        t        t        j                  j                  f«      rt        j                  j                  n|j$                  ‘Œ[ }}t-        g ||t/        «       «      }t1        | j2                  ||j5                  «       | j&                  ||| j6                  |j8                  ||| j:                  j<                  | j:                  j>                  | j:                  j@                  ¬«      }	 |jC                  |jE                  «       «       |jV                  }|| j6                  |<   n| j6                  |   }| j&                  jY                  |«      }t[        |
«      }
| j:                  j]                  ||
«      }|t        j^                  k(  ry ta        |jc                  «       «      D �cg c]  }|je                  |«      ‘Œ }}tg        ti        ||g«      «      S c c}w c c}w c c}w c c}w c c}w # tF        $ rL}tH        jJ                  jL                  r‚ tO        | jP                  jR                  | jT                  d «      |‚d }~ww xY wc c}w )Nc                 ó   — t        |«      S rF   ©rO   ©r-   rö   s     r)   rì   z0CodeGenerator.call_JitFunction.<locals>.<lambda>Ç  s   € ´MÀ!Ó4D€ r+   c                 ó   — t        |«       S rF   r«  r¬  s     r)   rì   z0CodeGenerator.call_JitFunction.<locals>.<lambda>É  s   € ¼ÀqÓ9IÐ5I€ r+   )
r  r
  r"  r  r[   r  r  r  r  r  )5ÚinspectÚgetcallargsra   r_   r\   rG   r   rN   rc  rb  rÑ   r   rM   r   r   r   rB   r   r'   r  Úhas_functionr   rÜ   rk  r	  r  Úget_capture_scoper  r[   rŠ   r  r  r  rž   ÚparseÚ	Exceptionr   rA  Úfront_end_debuggingr    r
  r]   r$  r&  Úget_functionrx   Úcallrg  rù   Úget_num_resultsrü  Únextr}   )r‚   ra   rb   r’   r&   r?   rd   Úargs_cstrë   Ú	args_pathÚargs_valÚfn_namer  r  rß   r   r�   ÚeÚcallee_ret_typeÚsymbolÚcall_oprw   s                         r)   Úcall_JitFunctionzCodeGenerator.call_JitFunctionÁ  s$  € Ü×"Ñ" 2§5¡5Ð:¨4Ò:°6Ñ:ˆØ')§|¡|Ö4˜t��T“
Ð4ˆÐ4Ü “oò 	7‰FˆAˆsÜ˜#¤§¡´´s¼DÄ+ÐNÕOÜ"Ÿ-™-×1Ñ1°#Ó6��Q’ð	7ô ! Ñ'DÓEˆØDLÖM¸D�DÔ+¨D°$Ó7Ñ7ÐMˆÐMÜ! $Ñ(IÓJˆ	Ø>GÖH°dÔ% d¨DÕ1ÐHˆÐHäœM¨"Ó-ÀHÖ/M¸S°·³Ò/MÈxÓXˆà�{‰{×'Ñ'¨Õ0ä$8¸Ó$<Ñ!ˆI�zð  öð ð ,/¨;¼*ÀSÜFJÌCÔQY×Q^ÑQ^×QdÑQdÐEeô;g”—‘×'Ò'Ølo×ltÑltñuðˆIð ô
 $ B¨	°8¼T»VÓDˆIÜ% d§l¡l°I¸r×?SÑ?SÓ?UÐ^b×^iÑ^iÐrtØ4;ÈD×LcÑLcØ/1¯{©{ÀiÐ\fØ.2¯l©l×.BÑ.BÐPT×P\ÑP\×PhÑPhØ15·±×1HÑ1Hô	JˆIð
TØ—‘ §¡£
Ô+ð (×0Ñ0ˆOØ/>ˆD×#Ñ# GÒ,à"×5Ñ5°gÑ>ˆOØ—‘×)Ñ)¨'Ó2ˆÜ'¨Ó1ˆØ—,‘,×#Ñ# F¨HÓ5ˆØœhŸm™mÒ+ØÜ27¸×8OÑ8OÓ8QÓ2RÖS¨Q�7×%Ñ% aÕ(ÐSˆÐSÜÔ'¨°/Ð1BÓCÓDÐDùòW 5ùò
 NùâHùâ/Mùò
øô ò Tä×$Ñ$×8Ò8ØÜ& t§{¡{§¡¸¿¹ÀtÓLÐRSÐSûð	Tüò Ts=   ²L:Â6L?ÃMÄM	
ÅAMÉM Ì
N+Í	N(ÍAN#Î#N(c                 ó6  ‡ — t        ‰ j                  |j                  «      «      }t        |t        «      s&‰ j
                  j                  |«      }|�	 |‰ |«      S t        |dd«      }|r�t        |dd«      rtdt        j                  |j                  «      z  g}t        |t        «      r|j                  |«       t        ‰ j                  j                  |dj                  |«      «      ‚t!        ˆ fd„|j"                  D «       «      }|j$                  D �cg c]  }‰ j                  |«      ‘Œ }}t'        t(        j*                  j-                  d„ |D «       «      «      }t        |t        «      r(|j/                  d|j0                  «       |j2                  }t        |t4        «      r t7        |||«       ‰ j9                  |||«      S t;        |d	«      rt=        |j0                  «      st>        j@                  jC                  |«      rhd
‰ jD                  i}	tG        jH                  |«      }
d|
jJ                  v r‰ |	d<   	  ||i |	¤|¤Ž}t        |tL        «      rt?        jL                  |«      }|S |‰ jV                  jY                  «       v rt[        t         |«      } ||i |¤Ž}t]        t_        |«      «      rta        |d„ «      S |S c c}w # tN        $ rB}tP        jR                  jT                  r‚ t        ‰ j                  j                  |d «      |‚d }~ww xY w)NÚ_must_use_resultFÚ
_is_unusedz#The result of %s is not being used.r?  c              3   ó@   •K  — | ]  }‰j                  |«      –— Œ y ­wrF   r�   rh  s     €r)   r¡   z+CodeGenerator.visit_Call.<locals>.<genexpr>þ  s   øè ø€ ÒD¨7�4—:‘:˜g×&ÑDùr¢   c              3   óF   K  — | ]  }t        |t        «      r|n|g–— Œ y ­wrF   )rG   rV   )rŸ   rö   s     r)   r¡   z+CodeGenerator.visit_Call.<locals>.<genexpr>   s#   è ø€ Ò1bÐXY´zÀ!ÄTÔ7J±!ÐQRÐPSÓ2SÑ1bùs   ‚!r   r  r²  r  c                 ó   — | S rF   r–   )rö   s    r)   rì   z*CodeGenerator.visit_Call.<locals>.<lambda>   s   € °Q€ r+   )1r   rž   rÏ   rG   r  Ú statically_implemented_functionsrû   r²   rª   r  rJ  rÍ  r    r
  r]   r6   rk  rl  rb   rV   Ú	itertoolsÚchainÚfrom_iterableÚinsertr  r  r   re   rÁ  ri   rI   r   rM   Ú
is_builtinr  r®  Ú	signatureÚ
parametersrW   r³  r   rA  r´  r,  rt   Úmaprk   r'   rs   )r‚   r`   ra   Ústatic_implementationÚmurÚerror_messageÚkwsrd   rb   Úextra_kwargsÚsigrA   r½  s   `            r)   rÐ   zCodeGenerator.visit_Callð  s‘  ø€ Ü! $§*¡*¨T¯Y©YÓ"7Ó8ˆÜ˜"œnÔ-Ø$(×$IÑ$I×$MÑ$MÈbÓ$QÐ!Ø$Ð0Ù,¨T°4Ó8Ð8ä�bÐ,¨eÓ4ˆÙ”7˜4 ¨uÔ5ØBÄSÇ[Á[ÐQU×QZÑQZÓE[Ñ[Ð\ˆMÜ˜#œsÔ#Ø×$Ñ$ SÔ)Ü" 4§;¡;§?¡?°D¸#¿(¹(À=Ó:QÓRÐRäÓD°d·m±mÔDÓDˆØ+/¯9©9Ö5 C�—
‘
˜3•Ð5ˆÐ5Ü”I—O‘O×1Ñ1Ñ1bÐ]aÔ1bÓbÓcˆÜ�bœ.Ô)Ø�K‰K˜˜2Ÿ;™;Ô'Ø—‘ˆBÜ�bœ+Ô&Ü˜4  TÔ*Ø×(Ñ(¨¨T°3Ó7Ð7Ü�B˜
Ô#Ô(8¸¿¹Ô(EÌ(Ï-É-×JbÑJbÐceÔJfØ'¨¯©Ð7ˆLÜ×#Ñ# BÓ'ˆCØ˜sŸ~™~Ñ-Ø-1�˜\Ñ*ðKÙ˜$Ð6 ,Ð6°#Ñ6�ä˜c¤5Ô)Ü"Ÿ.™.¨Ó-�CØ�
ð �×'Ñ'×.Ñ.Ó0Ñ0ÜÔ+¨TÓ2ˆDÙ�$Ð˜#ÑˆÜ;IÌ$ÈsË)Ô;TÔ% c©;Ó7Ð]ÐZ]Ð]ùòC 6øô& ò 	KÜ×$Ñ$×8Ò8Øô ' t§{¡{§¡¸¸dÓCÈÐJûð	Kús   ÄKÉ 0K Ë	LË=LÌLc                 ó,   — t        |j                  «      S rF   )r   ro   r»   s     r)   Úvisit_ConstantzCodeGenerator.visit_Constant"  s   € Ü˜Ÿ™Ó$Ð$r+   r`   c           
      óª  — | j                   j                  t        |j                  «      «      }|€5| j	                  |dj                  |j                  j                  «      «      ‚g }|j                  D ]É  }| j                  |«      }t        |«      s't        |«      }|du r	|dk(  r|c S |du sŒ<|dk(  sŒB|c S |j                  j                  «       rYt        |dd «      }|�|| j                  z  }t        j                  dt         | j"                  |t%        j&                  |«      ¬«       |j)                  |«       ŒË t+        |«      d	k(  r|j)                  «       t+        |«      d
k\  rS|j-                  «       }|j-                  «       }	| j/                  ||	|«      }
|j)                  |
«       t+        |«      d
k\  rŒSt+        |«      dk(  sJ ‚|d	   S )Nz9AST boolean operator '{}' is not (currently) implemented.FÚlogical_andTÚ
logical_orÚlinenozeLogical operators 'and' and 'or' are deprecated for non-scalar tensors; please use '&' or '|' instead)ÚcategoryÚfilenamerÜ  Úsourcer   r   r   )Ú_method_name_for_bool_oprû   r'   r¥  r2  r»  r^   rt   rž   rK   rÑ   rR   r²   r  r  Úwarn_explicitÚUserWarningr  rª   r  rÍ  r{   r?  r¶  )r‚   r`   r´  Únontrivial_valuesÚsubnodero   ÚbvrÜ  r§  r¦  Úress              r)   Úvisit_BoolOpzCodeGenerator.visit_BoolOp%  sÆ  € Ø×3Ñ3×7Ñ7¼¸T¿W¹W»ÓFˆØÐØ×#Ñ#ØÐQ×XÑXÐY]×Y`ÑY`×YiÑYiÓjólð lð Ðà—{‘{ò 	0ˆGð —J‘J˜wÓ'ˆEÜ$ UÔ+ä˜%“[�Ø˜%‘K k°]Ò&Bà ’LØ˜$’J [°LÓ%@à ’Lð —:‘:×&Ñ&Ô(Ü$ T¨8°TÓ:�FØÐ)Ø $§/¡/Ñ1˜Ü×*Ñ*ØÜ!,Ø!%§¡Ø%Ü"Ÿ{™{¨4Ó0õð "×(Ñ(¨Õ/ð9	0ô< Ð Ó! QÒ&ð ×$Ñ$ UÔ+äÐ#Ó$¨Ò)Ø#×'Ñ'Ó)ˆCØ#×'Ñ'Ó)ˆCØ×+Ñ+¨K¸¸cÓBˆCØ×$Ñ$ SÔ)ô	 Ð#Ó$¨Ó)ô Ð$Ó%¨Ò*Ð*Ð*Ø  Ñ#Ð#r+   rÚ  rÛ  rà  c                 óx  — | j                  |j                  «      }t        |«      r+|j                  dk(  r| j                  j                  |d«      S t        |t        «      r|j                  dk7  r|j                  }t        ||j                  «      }t        |«      rt        |t        «      rt        ||«      S |S )NÚT)r   r   ro   )rž   ro   rK   r³   r  ÚpermuterG   r   r²   rI   r   r  )r‚   r`   r¦  r³   s       r)   r´   zCodeGenerator.visit_Attribute[  s�   € Ø�j‰j˜Ÿ™Ó$ˆÜ˜SÔ! d§i¡i°3Ò&6Ø—=‘=×(Ñ(¨¨fÓ5Ð5ä�cœ9Ô%¨$¯)©)°wÒ*>Ø—)‘)ˆCÜ�s˜DŸI™IÓ&ˆÜ˜CÔ ¤Z°´kÔ%BÜ! # tÓ,Ð,Øˆr+   c                 óf   — d|j                   _        t        j                  j	                  | |«       y rº   )ro   rÄ  rª   rZ  r®   r»   s     r)   Ú
visit_ExprzCodeGenerator.visit_Exprg  s"   € Ø $ˆ�
‰
ÔÜ�‰×%Ñ% d¨DÕ1r+   c                  ó   — y rF   r–   r»   s     r)   Úvisit_NoneTypezCodeGenerator.visit_NoneTypek  r½   r+   c           
      óp  — t        |j                  «      }t        |«      D �]  \  }}t        |t        j
                  «      rt        |j                  «      ||<   Œ:t        |t        j                  «      rŽ|j                  }| j                  |j                  «      }t        |«      s'| j                  |dt        t        |«      «      z   «      ‚|dk  rdndt        |«      z   dz   j                  |j                  «      ||<   Œât!        dj                  t        |«      «      «      ‚ dj#                  |«      S )Nz^Cannot evaluate f-string containing non-constexpr conversion values, found conversion of type r   z{}z{!ú}z:encountered unexpected node of type {} in a JoinedStr noder  )rV   rt   r\   rG   rª   ÚConstantrJ  ro   ÚFormattedValueÚ
conversionrž   rO   r2  r'   Úchrr»  ÚAssertionErrorr6   )r‚   r`   rt   r?   ro   Úconversion_codeÚ	evaluateds          r)   Úvisit_JoinedStrzCodeGenerator.visit_JoinedStrn  s  € Ü�d—k‘kÓ"ˆÜ! &Ó)ó 	w‰HˆAˆuÜ˜%¤§¡Ô.Ü §¡Ó,��q’	Ü˜E¤3×#5Ñ#5Ô6Ø"'×"2Ñ"2�Ø ŸJ™J u§{¡{Ó3�	Ü$ YÔ/Ø×+Ñ+ØØxÜœd 9›oÓ.ñ/ó0ð 0ð &5°qÒ%8™T¸dÄSÈÓEYÑ>YÐ\_Ñ>_×gÑgÐhq×hwÑhwÓx��q’	ä$Ð%a×%hÑ%hÔimÐnsÓitÓ%uÓvÐvð	wð �w‰w�v‹Ðr+   c           	      óT  •— |€y t        j                  «       5  t        j                  dt        «       t        j                  dt        «       | j
                  }| j                  j                  «       }|| _        t        |d«      rnt        |d«      rb| j                  j                  | j                  | j                  |j                  z   |j                  «       | j                  j                  «       }	 t        ‰| �=  |«      }|r"|| _        | j                  j                  |«       |cd d d «       S # t         $ r ‚ t"        $ rU}t$        j&                  j(                  r‚ t!        | j*                  j,                  | j
                  t/        |«      «      d ‚d }~ww xY w# 1 sw Y   y xY w)NÚignorerÜ  Ú
col_offset)r  Úcatch_warningsÚsimplefilterÚDeprecationWarningÚPendingDeprecationWarningr$  rŠ   rN  ri   r  r  r  rÜ  rû  Úsuperrž   r    r³  r   rA  r´  r
  r]   r9   )r‚   r`   Ú	last_noderý  rA   r½  Ú	__class__s         €r)   rž   zCodeGenerator.visit€  sT  ø€ Øˆ<ØÜ×$Ñ$Ó&ñ 	ô ×!Ñ! (Ô,>Ô?Ü×!Ñ! (Ô,EÔFØŸ™ˆIØ—|‘|×+Ñ+Ó-ˆHØ ˆDŒMÜ�t˜XÔ&¬7°4¸Ô+FØ—‘×$Ñ$ T§^¡^°T·_±_ÀtÇ{Á{Ñ5RÐTX×TcÑTcÔdØŸ<™<×/Ñ/Ó1�ð	ZÜ‘g‘m DÓ)�ñ Ø )�”Ø—‘×$Ñ$ XÔ.Ø÷5	ñ 	øô $ò ØÜò ZÜ×$Ñ$×8Ò8Øô ' t§{¡{§¡¸¿¹ÄtÈAÃwÓOÐUYÐYûðZú÷	ð 	ús1   ™CFÃ6D4Ä%FÄ4FÅAFÆFÆFÆF'c                 ój   — | j                  |dj                  t        |«      j                  «      «      ‚)Nzunsupported AST node type: {})r2  r»  r'   r^   r»   s     r)   r®   zCodeGenerator.generic_visitŸ  s,   € Ø×Ñ Ð&E×&LÑ&LÌTÐRVËZ×M`ÑM`Ó&aÓbÐbr+   c                 ó  — t        |j                  «      }d|cxk  rdk  r"n t        d«      ‚t        |j                  «      rt        d«      ‚t	        | j                  |j                  d   «      «      }t        |t        «      st        d«      ‚|sQ|dk(  rd}n	 | j                  |j                  d   «      }t        | j                  j                  |t	        |«      «      ‚y # t        $ r}dt        |«      z   dz   }Y d }~ŒJd }~ww xY w)	Nr   r   z=`static_assert` requires one or two positional arguments onlyzqAssertion condition could not be determined at compile-time. Make sure that it depends only on `constexpr` valuesr   r  z'<failed to evaluate assertion message: ú>)r{   rb   rl  rh  r   rž   rG   rÑ   ÚNotImplementedErrorr³  r9   r!   r
  r]   )r‚   r`   Ú	arg_countÚpassedr1  r½  s         r)   Úexecute_static_assertz#CodeGenerator.execute_static_assert¢  sò   € Ü˜Ÿ	™	“Nˆ	Ø�IÔ" Ô"ÜÐ[Ó\Ð\ô (+¨4¯=©=Ô'9ÜÐ[Ó\Ð\ä% d§j¡j°·±¸1±Ó&>Ó?ˆÜ˜&¤$Ô'Ü%ð Dóð ñ Ø˜AŠ~Ø‘ðXØ"Ÿj™j¨¯©°1©Ó6�Gô .¨d¯k©k¯o©o¸tÔEYÐZaÓEbÓcÐcØøô	 !ò XØGÌ$ÈqË'ÑQÐTWÑW•GûðXús   ÂC$ Ã$	DÃ-DÄDc                 ó4   ‡ — dt         j                  fˆ fd„}|S )Nr`   c                 ó  •‡ — ˆ fd„|j                   D «       D ��ci c]  \  }}|t        |«      “Œ }}}|j                  D �cg c]  }t        ‰ j                  |«      «      ‘Œ }}t	         ‰|i |¤Ž«      S c c}}w c c}w )Nc              3   ó@   •K  — | ]  }‰j                  |«      –— Œ y ­wrF   r�   rh  s     €r)   r¡   z=CodeGenerator.static_executor.<locals>.ret.<locals>.<genexpr>½  s   øè ø€ Ò#U¸G D§J¡J¨w×$7Ñ#Uùr¢   )rl  r   rb   rž   r   )r‚   r`   r&   ro   rÔ  rd   rb   Ú	python_fns   `      €r)   rA   z*CodeGenerator.static_executor.<locals>.retº  s„   ù€ ó $VÀtÇ}Á}Ô#U÷á�D˜%ð Ô*¨5Ó1Ñ1ðˆCñ ð FJÇYÁYÖO¸cÔ(¨¯©°C«Õ9ÐOˆDÐOÜ™Y¨Ð4°Ñ4Ó5Ð5ùóùò Ps   ›A6Á!A<)rª   rÚ   )r  rA   s   ` r)   Ústatic_executorzCodeGenerator.static_executor¸  s   ø€ ð	6œCŸH™Hõ 	6ð ˆ
r+   r   )r   rÈ  )NFNFNr   )ƒr^   r”   r•   r   r   r   rJ  rƒ   r{   rV   rù   rc  rb  rG   r²   ri   r,  r   r  Úupdater   rM   Údevice_printÚminimumÚmaximumr2  r6  r'  r   r   r   rL  rQ  rS  rW  rÅ   r_  r¼   ro  rÈ   r‘  r“  r–  rš  r¿   rÂ   r¸   r«  r­  r¯  r¶  r¼  rª   ÚAddÚSubÚMultÚDivÚFloorDivÚModÚPowÚLShiftÚRShiftÚBitAndÚBitOrÚBitXorrº  r
   ÚoperatorrÛ  ró  r   rË   rÍ   r  r"  ÚEqÚNotEqÚLtÚLtEÚGtÚGtEr  Úcmpopr.  ÚUSubÚUAddÚNotÚInvertr,  Úunaryopr5  r8  r=  rX  r]  r™  ra  re  rœ  r   r¢  r	   r¤  r¨  rÁ  rÐ   rØ  ÚBoolOprç  ÚAndÚOrrà  Úboolopr´   rì  rî  rø  rž   r®   rÚ   r	  r  Úexperimental.gluonÚttglÚstatic_assertÚstatic_printr-  rÈ  r4  r   Ú__classcell__)rŸ   r-   r  s   00@r)   r	  r	    s?  ø… ð `eØ=>ñ90È+ð 90Ø?GÈ¹~ð90à$ S™Mó90ðz �t˜U E¨3°
¸GÀWÐM÷)àð 	
�
‰
�A‰ó)Ð�t˜C ˜H‘~ó ð ×ÑØ	�(—-‘-×,Ñ,Ð-Ø	�× Ñ Ð!Ø	�× Ñ Ð!ðô òLò	ò1ðf&˜cð &¨%°
¸IÐ0EÑ*Fð &È4ó &òò"ò	ò2òò@ð< Só ò
3?òj&òò'ò 2ò.ò$ò.ò
2ò2ò$ò	.ò@ð 	�‰�Ø�‰�Ø�‰�)Ø�‰�Ø�‰�nØ�‰�Ø�‰�Ø�
‰
�LØ�
‰
�LØ�
‰
�IØ�	‰	�8Ø�
‰
�Ið>Ð˜T $ s§|¡|Ñ"4°cÐ"9Ñ:ó ò2Còh$,òN,ò> 8òD3/òjò@ð$ 	�‰�˜#Ÿ)™) X¨s¯v©v°xÀÇÁÈ(ÐTW×TZÑTZÐ\dÐfi×fmÑfmÐowð<Ð˜d 4¨¯	©	¡?°CÐ#7Ñ8ó òtð  	�‰�)˜SŸX™X y°#·'±'¸9ÀcÇjÁjÐR^ð?Ð˜t D¨¯©Ñ$5°sÐ$:Ñ;ó ò@ò-ò	5òL6ò\ò'ò/ò6ò\6ò|2ò&ð0 U¨3°¨8¡_ó 0ðO Có Oð
-E ;ó -Eò^0^òd%ð2$ §¡ó 2$ðh >A¿W¹WÀmÐUX×U[ÑU[Ð]iÐ<jÐ˜d 4¨¯
©
Ñ#3°SÐ#8Ñ9Ójò
ò2òòô$ò>cð¨#¯(©(ð °tó ò,
õ 6à�‰×#Ñ#Ð%:Ø�‰×"Ñ"¡O°EÓ$:Ø×ÑÐ1Ø×Ñ™?¨5Ó1Ø‰_˜SÓ!Ø‰_˜SÓ!ðQÐ$ d¨6°8¸S¿X¹X¸JÈ¸OÑ3LÐ+LÑ&M÷ ùó_#)s   §M;r	  c                 óÔ  — d gt        | j                  «      z  }|j                  j                  «       D ].  \  }}	| j                  j	                  |«      }
t        |	«      ||
<   Œ0 t        g ||j                  |j                  «      }t        | «      \  }}ddl
m} t        d„ |j                  «      }|D �ci c]"  }| j                  |d      |j                  |   “Œ$ }}|j                  }  |dddg«      ||«      }t        ||| j                  «       | j                  |«      | d||||||¬«      }|j!                  | j#                  «       «       |j$                  }||_        |S c c}w )	Nr   )Ú
namedtuplec                 ó   — t        | «      dk(  S rQ   )r{   )rp   s    r)   rì   zast_to_ttir.<locals>.<lambda>Ø  s   € œc !›f¨™k€ r+   ÚSpecializationProxyr<   rÎ  T)
rš   r"  r
  r#  r  r  r  r  r  r  )r{   r_   rÎ  rý   Úindexr   rÜ   r<   rà   r   Úcollectionsr6  Úfilterr	  r±  r9   rž   r²  r  r  )ra   r]   r  r  r  r  r  rß   r*  rp   rc   r   r  r  r6  Úleavesr?   r<   rÎ  Úproxyr�   rA   s                         r)   Úast_to_ttirr>  Ï  sW  € Ø�œ˜RŸ\™\Ó*Ñ*€IØ—‘×#Ñ#Ó%ò &‰ˆˆ1Ø�l‰l× Ñ  Ó#ˆÜ" 1›ˆ	�#Šð&ô ˜B 	¨3¯=©=¸#¿)¹)ÓD€IÜ0°Ó4Ñ€Iˆzå&ÜÑ)¨3¯=©=Ó9€FØ?EÖF¸!�—‘˜a ™dÑ# S§]¡]°1Ñ%5Ñ5ÐF€IÐFØ—‘€IØI‰JÐ,¨{¸KÐ.HÓIÈ)ÐU^Ó_€EÜ˜g y¸×9MÑ9MÓ9OÐ_a×_fÑ_fÐglÓ_mØ%'°4À9ÐYcÐmtØ*5À*ÐU[ô]€Ið ‡O�O�B—H‘H“JÔØ
×
Ñ
€Cà€C„KØ€Jùò Gs   Â6'E%rF   )Grª   r†   r®  r$   r  rD  rÉ  Údataclassesr   rr   r   Útypingr   r   r   r   r	   r
   r   r   r   r  r   r   Ú_C.libtritonr   r   r   r   r   Úlanguage.corer   r   r   Úruntime.jitr   r   Úruntimer   Ú_utilsr   r   r   Úerrorsr    r!   r"   r*   rB   rÑ   rI   rK   rO   rT   rX   re   rk   rs   rx   ro   r}   rb  r'   r  r   rZ  r˜   rÜ   r  r	  r>  r–   r+   r)   ú<module>rG     sn  ðÛ 
Û Û Û 	Û Û Û Ý !Ý ß T× TÕ Tç ß 'ß 3Ñ 3ß GÑ Gß =å !ß HÑ Hç aÑ aòò	ð%˜ð % ó %ð!˜ð ! ó !ðU�Sð U˜Tó UðN˜Sð N Tó Nð(�Sð (˜Tó (òòXòDð ¨*Ñ!5ó ð"  b§h¡h¡ð "¸¸Y¹ó "ð ˜#™t D›zÐ*Ð ÷3ñ 3ô*T%˜CŸO™Oô T%÷n7ñ 7ñt �$Ô÷ð ó ðô
s�C—O‘Oô sôl%r+   