§
    OŠtjx  ã                   ó  — d dl mZmZmZmZm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 d dlmZmZ d dlmZmZ  ed	¦  «        Zefd
edeg ef         fd„Zd„ Zdd„Z de _!         G d„ d¦  «        Z" G d„ dee         ¦  «        Z#dS )é    )ÚSequenceÚListÚTypeVarÚTupleÚCallableN)ÚTritonSemanticé   )Ú_core)Ú
AutoLayoutÚDistributedLayoutÚDistributedLinearLayoutÚSliceLayoutÚSharedLayoutÚCoalescedLayoutÚSharedLinearLayout)ÚGluonOpBuilderÚcompute_tmem_reg_layout)Úflatten_values_to_irÚunflatten_ir_valuesÚTensorTyÚcondÚmsg_fnc                 ó0   — | s | |¦   «         ¦  «        ‚d S ©N© )r   r   Úcategorys      új/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/triton/experimental/gluon/language/_semantic.pyÚ_checkr      s*   € Øð !Øˆh�v�v‘x”xÑ Ô Ð ð!ð !ó    c                 ó^   — t          | t          ¦  «        ot          d„ | D ¦   «         ¦  «        S )Nc              3   ó@   K  — | ]}t          |t          ¦  «        V — Œd S r   ©Ú
isinstanceÚint)Ú.0Úis     r   ú	<genexpr>z_is_int_list.<locals>.<genexpr>   s,   è è € Ð.QÐ.QÀa­z¸!½SÑ/AÔ/AÐ.QÐ.QÐ.QÐ.QÐ.QÐ.Qr   )r#   r   Úall©Úvalues    r   Ú_is_int_listr+      s.   € Ý�e�XÑ&Ô&ÐQ­3Ð.QÐ.QÈ5Ð.QÑ.QÔ.QÑ+QÔ+QÐQr   c                 óÆ  ‡‡‡‡‡‡— t          t          ‰t          ¦  «        d„ ¦  «         t          ‰dv ˆfd„¦  «         t          t          ‰t          ¦  «        ˆfd„¦  «         t          ‰dk    o‰‰dz
  z  dk    d„ ¦  «         t	          ‰¦  «        Št          t          d	„ ‰D ¦   «         ¦  «        ˆfd
„¦  «         t          ‰¦  «        }t          |dk    d„ ¦  «         |€g }‰dk    }|rdn‰Š|r'|D ]$}t          t          |¦  «        |k    d„ ¦  «         Œ%t          | ‰|‰‰|¦  «        Št          ‰d uˆˆˆfd„¦  «         |�r‰d         }	‰j        sSt          ‰j	        d         d|	dz  gk    ˆfd„¦  «         ‰j         
                    d|	dz  g¦  «         ddg‰j	        d<   n´‰j        d         d|	dz  gk    rž| j        Šdt          ‰j        ¦  «        z  }
t          |
d‰z  k    ˆfd„¦  «         ‰j        }dD ]O}t          ‰|¦  «        }t          |¦  «        D ]-\  }}|d|	dz  gk    r||         |d         c|d<   ||<   ‰c c S Œ.ŒPJ d|› �¦   «         ‚‰S )Nc                  ó   — dS )Nzinstr_variant must be a stringr   r   r   r   ú<lambda>z*_compute_tmem_reg_layout.<locals>.<lambda>   s   € Ð3S€ r   )Ú32x32bÚ16x64bÚ16x128bÚ16x256bÚ16x32bx2Ú32x32b_splitnc                  ó   •— d‰ › �S )Nzunknown instr_variant: r   )Úinstr_variants   €r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>   s   ø€ Ð<¨]Ð<Ð<€ r   c                  ó(   •— dt          ‰ ¦  «        ›�S )Nz!num_warps must be an int but got ©Útype©Ú	num_warpss   €r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>   s   ø€ Ð/fÕSWÐXaÑSbÔSbÐ/fÐ/f€ r   é   r	   r   c                  ó   — dS )Nz)num_warps must be a power of two and >= 4r   r   r   r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>   s   € ÐJu€ r   c              3   ó@   K  — | ]}t          |t          ¦  «        V — Œd S r   r"   )r%   Údims     r   r'   z+_compute_tmem_reg_layout.<locals>.<genexpr>   s,   è è € Ð5Ð5¨�z˜#�sÑ#Ô#Ð5Ð5Ð5Ð5Ð5Ð5r   c                  ó   •— d‰ › �S )Nz#shape entries must be ints but got r   ©Úshapes   €r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>   s   ø€ Ð?lÐejÐ?lÐ?l€ r   é   c                  ó   — dS )Nzexpected a 2D tensorr   r   r   r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>   s   € Ð4€ r   r4   r/   c                  ó   — dS )Nzcga_layout basis rank mismatchr   r   r   r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>(   s   € Ð/O€ r   c                  ó   •— d‰ › d‰› d‰› �S )NzTMEM layout 'z' unsupported for shape z and num_warps r   )Úatom_variantr;   rB   s   €€€r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>3   s!   ø€ Ðj <ÐjÐjÈÐjÐjÐ_hÐjÐj€ r   éÿÿÿÿc                  ó   •— d‰ › �S )NzJsplitn with 1 register requires the last lane basis to be [0, N / 2]. Got r   )Ú
layout_objs   €r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>;   s   ø€ ÐtÐhrÐtÐt€ r   é    c                  ó   •— dd‰ z  › d‰ › d�S )NzETo be able to `tmem.load` into `tl.split` you need to have more than rK   ú z¥-bit registers, as you need to use the instruction 32x32b.x1 twice. You can always load into instr_variant="32x32b" and then convert_layout to this layout otherwise.r   )Úbitwidths   €r   r.   z*_compute_tmem_reg_layout.<locals>.<lambda>B   s3   ø€ ð 3]Ø˜h™ð3]ð 3]Ø)1ð3]ð 3]ð 3]€ r   )Ú
lane_basesÚ
warp_basesFz6splitn requires at least one basis of [0, N / 2]. Got )r   r#   Ústrr$   Úlistr(   Úlenr   Ú	reg_basesrO   ÚappendÚprimitive_bitwidthÚgetattrÚ	enumerate)Ú
element_tyrB   Úlayoutr;   r6   Ú
cga_layoutÚrankÚsplitnÚbasisÚNÚnum_regrT   Ú	bases_strÚbasesr&   rG   rN   rJ   s    ` ``          @@@r   Ú_compute_tmem_reg_layoutrc      sV  øøøøøø€ Ý
�:�m¥SÑ)Ô)Ð+SÐ+SÑTÔTÐTÝ
ˆ=ÐcÐcØ<Ð<Ð<Ð<ñ>ô >ð >å
�:�i¥Ñ%Ô%Ð'fÐ'fÐ'fÐ'fÑgÔgÐgÝ
ˆ9˜Š>Ð@˜y¨I¸©MÑ:¸qÒ@ÐBuÐBuÑvÔvÐvå�‰KŒK€EÝ
�3Ð5Ð5¨uÐ5Ñ5Ô5Ñ5Ô5Ð7lÐ7lÐ7lÐ7lÑmÔmÐmÝˆu‰:Œ:€DÝ
ˆ4�1Š9Ð4Ð4Ñ5Ô5Ð5àÐØˆ
Ø˜oÒ-€FØ%Ð8�8�8¨=€Làð QØð 	Qð 	QˆEÝ•3�u‘:”: Ò%Ð'OÐ'OÑPÔPÐPÐPå(ØØØØØØñô €Jõ ˆ:˜TÐ!ØjÐjÐjÐjÐjÐjñlô lð lð ñ \Ø�!ŒHˆØÔ#ð 	\õ �:Ô(¨Ô,°°A¸±F°Ò;ØtÐtÐtÐtñvô vð vàÔ ×'Ò'¨¨A°©F¨Ñ4Ô4Ð4Ø)*¨A¨ˆJÔ! "Ñ%Ð%ØÔ! "Ô%¨!¨Q°!©V¨Ò4Ð4Ø!Ô4ˆHØ�˜ZÔ1Ñ2Ô2Ñ2ˆGÝØ˜" ™.Ò(ð +]ð +]ð +]ð +]ñ^ô ^ð ^ð #Ô,ˆIØ9ð *ð *�	Ý 
¨IÑ6Ô6�Ý )¨%Ñ 0Ô 0ð *ð *‘H�A�uØ  A¨¡F Ò+Ð+Ø27¸´(¸IÀb¼MÐ/˜	 "™ u¨Q¡xØ)Ð)Ð)Ð)Ð)Ð)ð ,ð*ð \Ð[ÐSYÐ[Ð[Ñ[Ô[�5ØÐr   Tc                   ó&   — e Zd Zdefd„Zd„ Zd„ ZdS )ÚGluonCallerContextr;   c                 ó   — || _         d S r   r:   )Úselfr;   s     r   Ú__init__zGluonCallerContext.__init__W   s   € Ø"ˆŒˆˆr   c                 ó   — d| j         › �S )NÚ_NWr:   ©rg   s    r   ÚmanglezGluonCallerContext.mangleZ   s   € Ø%�T”^Ð%Ð%Ð%r   c                 ób   — |                      d|                     | j        ¦  «        ¦  «         d S )Nzttg.num-warps)Úset_attrÚget_int32_attrr;   )rg   ÚfnÚbuilders      r   Úinitialize_calleez$GluonCallerContext.initialize_callee]   s,   € Ø
�Š�O W×%;Ò%;¸D¼NÑ%KÔ%KÑLÔLÐLÐLÐLr   N)Ú__name__Ú
__module__Ú__qualname__r$   rh   rl   rr   r   r   r   re   re   U   sR   € € € € € ð# #ð #ð #ð #ð #ð&ð &ð &ðMð Mð Mð Mð Mr   re   c            
       óÀ  ‡ — e Zd ZU ej        ZeZeed<   defd„Zd„ Z	d„ Z
dee         dee         fd„Zded	ed
efd„Zdeded
efˆ fd„Zded
eeef         fˆ fd„Zdedee         d
efˆ fd„Zdedee         d
efd„Zdeded
efˆ fd„Zˆ fd„Zdedee         defˆ fd„Zd„ Zd„ ZdAd„Zd„ Zd „ Zd!„ Zd"„ Zd#„ Z d$„ Z!d%„ Z"d&„ Z#d'„ Z$d(„ Z%d)„ Z&d*„ Z'd+„ Z(d,„ Z)d-„ Z*e+d.„ ¦   «         Z,d/e-e         d	ed0ed
eed1f         fd2„Z.d/e-e         d	ed
eed1f         fd3„Z/ded4ed5ed
efd6„Z0dededed
efd7„Z1d8ed9ed	ed
efd:„Z2d8ed
efd;„Z3d<e-e         d=e-e         fd>„Z4d?„ Z5d@„ Z6ˆ xZ7S )BÚGluonSemanticrq   c                 ó   — || _         d S r   )rq   )rg   rq   s     r   rh   zGluonSemantic.__init__g   s   € ØˆŒˆˆr   c                 óœ   — |g k    r|}n.t          j        ||| j                             |¦  «        ¦  «        }|                      ||¦  «        S r   )ÚttglÚdistributed_typerq   Úget_gluon_layout_from_tensorÚtensor)rg   ÚhandleÚ	scalar_tyrB   Útys        r   Ú_wrap_handle_infer_layoutz'GluonSemantic._wrap_handle_infer_layoutj   sK   € Ø�BŠ;ˆ;ØˆBˆBåÔ& y°%¸¼×9bÒ9bÐciÑ9jÔ9jÑkÔkˆBØ�{Š{˜6 2Ñ&Ô&Ð&r   c                 óX   — |                       |j        |j        j        |j        ¦  «        S r   )r�   r~   r9   ÚscalarrB   )rg   r}   s     r   Ú_wrap_tensor_infer_layoutz'GluonSemantic._wrap_tensor_infer_layoutq   s%   € Ø×-Ò-¨f¬m¸V¼[Ô=OÐQWÔQ]Ñ^Ô^Ð^r   Ú	lhs_shapeÚ	rhs_shapec                 ó²  — t          |¦  «        t          |¦  «        k    rt          d|› d|› �¦  «        ‚g }t          |¦  «        D ]�\  }}||         }|dk    r|                     |¦  «         Œ)|dk    s||k    r|                     |¦  «         ŒKt          dt	          |¦  «        z   dz   t	          |¦  «        z   dz   t	          |¦  «        z   ¦  «        ‚|S )Nú!Cannot broadcast, rank mismatch: ú, r	   z?Cannot make_shape_compatible: incompatible dimensions at index ú: ú and )rS   Ú
ValueErrorrX   rU   rQ   )rg   r…   r†   Ú	ret_shaper&   ÚleftÚrights          r   Ú_broadcast_shapeszGluonSemantic._broadcast_shapest   s  € Ýˆy‰>Œ>�S ™^œ^Ò+Ð+ÝÐYÀÐYÐYÈiÐYÐYÑZÔZÐZàˆ	Ý  Ñ+Ô+ð 	að 	a‰GˆAˆtØ˜a”LˆEØ�qŠyˆyØ× Ò  Ñ'Ô'Ð'Ð'Ø˜1’*�* %¨4¢- -Ø× Ò  Ñ&Ô&Ð&Ð&å ð "-Ý/2°1©v¬vñ"6Ø8<ñ"=Ý?BÀ4¹y¼yñ"IØKRñ"SÝUXÐY^ÑU_ÔU_ñ"`ñ aô að aàÐr   ÚinputÚaxisÚreturnc                 ób  ‡‡‡— d„ ‰j         D ¦   «         }|                     ‰d¦  «         ‰dk     r‰t          ‰j         ¦  «        z  Št          t	          ‰j        t          j        ¦  «        ˆfd„¦  «         ‰j        j        Št          t	          ‰t          t          t          f¦  «        ˆfd„¦  «         t          t	          ‰t          t          f¦  «        p
‰j        ‰k    ˆˆfd„¦  «         | j                             ‰j        ‰¦  «        }|                      |‰j        j        |¦  «        S )Nc                 ó6   — g | ]}t          j        |¦  «        ‘ŒS r   )rz   Ú_unwrap_if_constexpr©r%   Úxs     r   ú
<listcomp>z-GluonSemantic.expand_dims.<locals>.<listcomp>…   s#   € ÐGÐGÐG°a•TÔ.¨qÑ1Ô1ÐGÐGÐGr   r	   r   c                  ó   •— d‰ j         ›�S ©Nz=expected expand_dims input to be a distributed_type but got: r8   ©r‘   s   €r   r.   z+GluonSemantic.expand_dims.<locals>.<lambda>Œ   ó   ø€ ÐeÐW\ÔWaÐeÐe€ r   c                  ó   •— d‰ › �S )Nz;expected expand_dims input to have a SliceLayout, but got: r   ©rZ   s   €r   r.   z+GluonSemantic.expand_dims.<locals>.<lambda>�   s   ø€ Ð]ÐU[Ð]Ð]€ r   c                  ó   •— d‰ › d‰j         › �S )Nz7expected expand_dims input layout to be sliced in axis z	 but got ©r?   )r’   rZ   s   €€r   r.   z+GluonSemantic.expand_dims.<locals>.<lambda>’   s   ø€ ÐiÈdÐiÐiÐ]cÔ]gÐiÐi€ r   )rB   ÚinsertrS   r   r#   r9   rz   r{   rZ   r   r   r   r?   rq   Úcreate_expand_dimsr~   r�   rƒ   )rg   r‘   r’   Ú	dst_shaper~   rZ   s    ``  @r   Úexpand_dimszGluonSemantic.expand_dims„   s7  øøø€ ØGÐG¸5¼;ÐGÑGÔGˆ	Ø×Ò˜˜qÑ!Ô!Ð!à�!Š8ˆ8Ø•C˜œÑ$Ô$Ñ$ˆDå�z˜%œ*¥dÔ&;Ñ<Ô<ØeÐeÐeÐeñ	gô 	gð 	gà”Ô"ˆÝ�z˜&¥;µ
½OÐ"LÑMÔMØ]Ð]Ð]Ð]ñ	_ô 	_ð 	_åÝ�v¥
­OÐ<Ñ=Ô=ÐSÀÄÈtÒASØiÐiÐiÐiÐiñ	kô 	kð 	kð ”×0Ò0°´¸tÑDÔDˆØ×-Ò-¨f°e´jÔ6GÈÑSÔSÐSr   ÚaÚbc                 óØ   •— |                       ||¦  «        \  }}t          |j        g k    d„ ¦  «         t          ¦   «                              ||¦  «        }|                      |¦  «        S )Nc                  ó   — dS )NzCannot join scalars in gluonr   r   r   r   r.   z$GluonSemantic.join.<locals>.<lambda>™   s   € Ð&D€ r   )Úbroadcast_impl_valuer   rB   ÚsuperÚjoinr„   )rg   r¦   r§   r*   Ú	__class__s       €r   r¬   zGluonSemantic.join—   sa   ø€ Ø×(Ò(¨¨AÑ.Ô.‰ˆˆ1ÝˆqŒw˜"Š}ÐDÐDÑEÔEÐEÝ‘”—’˜Q Ñ"Ô"ˆØ×-Ò-¨eÑ4Ô4Ð4r   c                 ó    •— t          ¦   «                              |¦  «        \  }}|                      |¦  «        |                      |¦  «        fS r   )r«   Úsplitr„   )rg   r¦   ÚlhsÚrhsr­   s       €r   r¯   zGluonSemantic.split�   sD   ø€ Ý‘7”7—=’= Ñ#Ô#‰ˆˆSØ×-Ò-¨cÑ2Ô2°D×4RÒ4RÐSVÑ4WÔ4WÐWÐWr   Údimsc                 ór   •— t          ¦   «                              ||¦  «        }|                      |¦  «        S r   )r«   Úpermuter„   )rg   r‘   r²   r*   r­   s       €r   r´   zGluonSemantic.permute¡   s.   ø€ Ý‘”—’  tÑ,Ô,ˆØ×-Ò-¨eÑ4Ô4Ð4r   rB   c                 ó”  ‡‡‡— t          t          ‰j        t          j        ¦  «        ˆfd„¦  «         ‰j                             ¦   «         Št          t          ‰¦  «        t          ‰¦  «        k    ˆˆfd„¦  «         ‰‰k    r‰S t          ‰¦  «        D ];\  }}‰|         |k    r*|dk    r$t          d‰|         › d|› d|› d‰› d‰› �
¦  «        ‚Œ<t          j        ‰j        j	        ‰‰j        j
        ¦  «        }| j                             ‰j        |                     | j        ¦  «        ¦  «        }|                      ||¦  «        S )	Nc                  ó   •— d‰ j         ›�S r›   r8   rœ   s   €r   r.   z4GluonSemantic.broadcast_impl_shape.<locals>.<lambda>§   r�   r   c                  ó   •— d‰› d‰ › �S )Nrˆ   r‰   r   )rB   Ú	src_shapes   €€r   r.   z4GluonSemantic.broadcast_impl_shape.<locals>.<lambda>©   s   ø€ Ð5mÐYbÐ5mÐ5mÐfkÐ5mÐ5m€ r   r	   z3Cannot broadcast, the expanded size of the tensor (z ) must match the existing size (z) at non-singleton dimension rŠ   r‰   )r   r#   r9   rz   r{   Úget_block_shapesrS   rX   rŒ   rƒ   rZ   rq   Úcreate_broadcastr~   Úto_irr}   )rg   r‘   rB   r&   ÚitemÚret_tyr~   r¸   s    ``    @r   Úbroadcast_impl_shapez"GluonSemantic.broadcast_impl_shape¥   sy  øøø€ Ý�z˜%œ*¥dÔ&;Ñ<Ô<ØeÐeÐeÐeñ	gô 	gð 	gà”J×/Ò/Ñ1Ô1ˆ	Ý�s�9‰~Œ~¥ U¡¤Ò+Ð-mÐ-mÐ-mÐ-mÐ-mÑnÔnÐnØ�IÒÐØˆLÝ  Ñ+Ô+ð 	@ð 	@‰GˆAˆtØ�QŒx˜4ÒÐ D¨A¢I IÝ ð "?ÐW\Ð]^ÔW_ð "?ð "?ØCGð"?ð "?à%&ð"?ð "?à*3ð"?ð "?à7<ð"?ð "?ñ @ô @ð @øõ Ô& u¤zÔ'8¸%ÀÄÔARÑSÔSˆØ”×.Ò.¨u¬|¸V¿\º\È$Ì,Ñ=WÔ=WÑXÔXˆØ�{Š{˜6 6Ñ*Ô*Ð*r   r°   r±   c                 ó†  •‡‡	— |j         Š|j         Š	‰                     ¦   «         r‰	                     ¦   «         s"t          ¦   «                              ||¦  «        S t	          t          ‰t          j        ¦  «        ˆfd„¦  «         t	          t          ‰	t          j        ¦  «        ˆ	fd„¦  «         ‰                     ¦   «         }‰	                     ¦   «         }|  	                    ||¦  «        }t          ‰j
        t          ¦  «        }t          ‰	j
        t          ¦  «        }|r|s|                      |‰	j
        ¦  «        }nO|r|s|                      |‰j
        ¦  «        }n/‰j
        ‰	j
        k    rt          d‰j
        › d‰	j
        › �¦  «        ‚|                      ||¦  «        }|                      ||¦  «        }||fS )Nc                  ó   •— d‰ ›�S )Nz@expected broadcast left input to be a distributed_type but got: r   )Úlhs_tys   €r   r.   z4GluonSemantic.broadcast_impl_value.<locals>.<lambda>½   s   ø€ ÐdÐZ`ÐdÐd€ r   c                  ó   •— d‰ ›�S )NzAexpected broadcast right input to be a distributed_type but got: r   )Úrhs_tys   €r   r.   z4GluonSemantic.broadcast_impl_value.<locals>.<lambda>¿   s   ø€ ÐeÐ[aÐeÐe€ r   zLayout mismatch in broadcast: z vs )r9   Úis_blockr«   rª   r   r#   rz   r{   r¹   r�   rZ   r   Úset_auto_layoutrŒ   r¾   )rg   r°   r±   r…   r†   r�   Úis_lhs_autoÚis_rhs_autorÁ   rÃ   r­   s           @@€r   rª   z"GluonSemantic.broadcast_impl_valueµ   sÆ  øøø€ Ø”ˆØ”ˆà�ŠÑ Ô ð 	:¨¯ªÑ(9Ô(9ð 	:Ý‘7”7×/Ò/°°SÑ9Ô9Ð9å�z˜&¥$Ô"7Ñ8Ô8ØdÐdÐdÐdñ	fô 	fð 	få�z˜&¥$Ô"7Ñ8Ô8ØeÐeÐeÐeñ	gô 	gð 	gð ×+Ò+Ñ-Ô-ˆ	Ø×+Ò+Ñ-Ô-ˆ	Ø×*Ò*¨9°iÑ@Ô@ˆ	å  ¤µ
Ñ;Ô;ˆÝ  ¤µ
Ñ;Ô;ˆØð 	b˜{ð 	bØ×&Ò& s¨F¬MÑ:Ô:ˆCˆCØð 	b ð 	bØ×&Ò& s¨F¬MÑ:Ô:ˆCˆCØŒ]˜fœmÒ+Ð+ÝÐ`¸f¼mÐ`Ð`ÐQWÔQ^Ð`Ð`ÑaÔaÐaà×'Ò'¨¨YÑ7Ô7ˆØ×'Ò'¨¨YÑ7Ô7ˆØ�Cˆxˆr   c                 ó¸   •— ||z
  g}|€t          ¦   «         }t          j        t          j        ||¦  «        }t	          ¦   «                              |||¬¦  «        S )N)r½   )r   rz   r{   Úint32r«   Úarange)rg   ÚstartÚendrZ   rB   r½   r­   s         €r   rÊ   zGluonSemantic.arangeÒ   sN   ø€ Ø�u‘�ˆØˆ>Ý‘\”\ˆFÝÔ&¥t¤z°5¸&ÑAÔAˆÝ‰wŒw�~Š~˜e S°ˆ~Ñ8Ô8Ð8r   r¤   Úcan_reorderc                 ó˜   •— t          | d„ ¦  «         t          ¦   «                              |||¦  «        }|                      |¦  «        S )Nc                  ó   — dS )Nz%can_reorder is not supported in gluonr   r   r   r   r.   z'GluonSemantic.reshape.<locals>.<lambda>Ú   s   € Ð(O€ r   )r   r«   Úreshaper„   )rg   r‘   r¤   rÍ   r*   r­   s        €r   rÐ   zGluonSemantic.reshapeÙ   sF   ø€ Ý�;ˆÐ OÐ OÑPÔPÐPÝ‘”—’  y°+Ñ>Ô>ˆØ×-Ò-¨eÑ4Ô4Ð4r   c                 óü   — t          |¦  «        dk    r|S t          j        |j        ||¦  «        }| j                             |                     | j        ¦  «        |j        ¦  «        }t          j        ||¦  «        S )Nr   )	rS   rz   r{   Údtyperq   Úcreate_splatr»   r~   r}   )rg   r*   rB   rZ   r½   r~   s         r   ÚsplatzGluonSemantic.splatÞ   se   € Ýˆu‰:Œ:˜Š?ˆ?ØˆLÝÔ& u¤{°E¸6ÑBÔBˆØ”×*Ò*¨6¯<ª<¸¼Ñ+EÔ+EÀuÄ|ÑTÔTˆÝŒ{˜6 6Ñ*Ô*Ð*r   c                 ó|   — |                       ||¦  «        }|€t          ¦   «         }|                      |||¦  «        S r   )Úmake_scalarr   rÔ   )rg   rB   r*   rÒ   rZ   rƒ   s         r   ÚfullzGluonSemantic.fullå   s;   € Ø×!Ò! %¨Ñ/Ô/ˆØˆ>Ý‘\”\ˆFØ�zŠz˜& %¨Ñ0Ô0Ð0r   Fc                 ó�  ‡‡— |j         Št          t          ‰t          j        ¦  «        ˆfd„¦  «         t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          j        ‰j        ‰j        ‰¦  «        }|                     | j	        ¦  «        }|rw| j	         
                    ||j        ¦  «        sWt          d‰j        › d‰› d|                      ‰j        ‰j        ¦  «        › d|                      ‰‰j        ¦  «        › �¦  «        ‚| j	                             ||j        ¦  «        }t          j        ||¦  «        S )Nc                  ó   •— d‰ ›�S )Nz@expected convert_layout input to be a distributed_type but got: r   )r€   s   €r   r.   z.GluonSemantic.convert_layout.<locals>.<lambda>î   s   ø€ Ð`ÐZ\Ð`Ð`€ r   c                  ó   •— d‰ › �S ©Nz4expected 'layout' to be a DistributedLayout but got r   rŸ   s   €r   r.   z.GluonSemantic.convert_layout.<locals>.<lambda>ð   ó   ø€ ÐVÈfÐVÐV€ r   zlayout conversion from z to z) is not trivial.
The linear layouts are:
ú
)r9   r   r#   rz   r{   r   rY   rB   r»   rq   Úis_convert_layout_trivialr~   Ú	TypeErrorrZ   Úto_linear_layoutÚcreate_convert_layoutr}   )rg   r*   rZ   Úassert_trivialr½   Ú	ret_ty_irr~   r€   s     `    @r   Úconvert_layoutzGluonSemantic.convert_layoutë   si  øø€ ØŒZˆÝ�z˜"�dÔ3Ñ4Ô4Ø`Ð`Ð`Ð`ñ	bô 	bð 	bå�z˜&¥$Ô"8Ñ9Ô9ØVÐVÐVÐVñ	Xô 	Xð 	XåÔ& r¤}°b´hÀÑGÔGˆØ—L’L ¤Ñ.Ô.ˆ	Øð 	J $¤,×"HÒ"HÈÐTYÔT`Ñ"aÔ"að 	JÝð I°b´ið Ið IÀVð Ið IØ8<×8MÒ8MÈbÌiÐY[ÔYaÑ8bÔ8bðIð Ià#×4Ò4°V¸R¼XÑFÔFðIð Iñ Jô Jð Jð ”×3Ò3°I¸u¼|ÑLÔLˆÝŒ{˜6 6Ñ*Ô*Ð*r   c                 ó,  ‡‡‡— t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          t	          ‰¦  «        ˆfd„¦  «         t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          j        ‰‰‰‰¦  «        }|�9| j                             | 	                    | j        ¦  «        |j
        ¦  «        }n2| j                             | 	                    | j        ¦  «        ¦  «        }t          j        |‰‰‰‰¦  «        S )Nc                  ó   •— d‰ › �S )Nz,expected 'element_ty' to be a dtype but got r   )rY   s   €r   r.   z/GluonSemantic.allocate_shared.<locals>.<lambda>û   s   ø€ Ð;vÐjtÐ;vÐ;v€ r   c                  ó   •— d‰ › �S ©Nz1all elements of 'shape' must be integers but got r   rA   s   €r   r.   z/GluonSemantic.allocate_shared.<locals>.<lambda>ü   ó   ø€ Ð,gÐ`eÐ,gÐ,g€ r   c                  ó   •— d‰ › �S ©Nz/expected 'layout' to be a SharedLayout but got r   rŸ   s   €r   r.   z/GluonSemantic.allocate_shared.<locals>.<lambda>þ   ó   ø€ ÐQÈÐQÐQ€ r   )r   r#   rz   rÒ   r+   r   Úshared_memory_descriptor_typerq   Úcreate_local_allocr»   r~   Úshared_memory_descriptor)rg   rY   rB   rZ   r*   r€   r~   s    ```   r   Úallocate_sharedzGluonSemantic.allocate_sharedú   sÿ   øøø€ Ý�z˜*¥d¤jÑ1Ô1Ð3vÐ3vÐ3vÐ3vÑwÔwÐwÝ�|˜EÑ"Ô"Ð$gÐ$gÐ$gÐ$gÑhÔhÐhÝ�z˜&¥$Ô"3Ñ4Ô4ØQÐQÐQÐQñ	Sô 	Sð 	SåÔ/°
¸EÀ6È5ÑQÔQˆØÐØ”\×4Ò4°R·X²X¸d¼lÑ5KÔ5KÈUÌ\ÑZÔZˆFˆFà”\×4Ò4°R·X²X¸d¼lÑ5KÔ5KÑLÔLˆFÝÔ,¨V°ZÀÈÐPUÑVÔVÐVr   c                 ó4  ‡— t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          j        |j        |j        ‰¦  «        }| j                             | 	                    | j        ¦  «        |j
        ¦  «        }t          j        ||¦  «        S )Nc                  ó   •— d‰ › �S rÛ   r   rŸ   s   €r   r.   z+GluonSemantic.shared_load.<locals>.<lambda>  rÜ   r   )r   r#   rz   r   r{   rÒ   rB   rq   Úcreate_local_loadr»   r~   r}   )rg   Úmem_descrZ   r½   r~   s     `  r   Úshared_loadzGluonSemantic.shared_load  s„   ø€ Ý�z˜&¥$Ô"8Ñ9Ô9ØVÐVÐVÐVñ	Xô 	Xð 	XåÔ& x¤~°x´~ÀvÑNÔNˆØ”×/Ò/°·²¸T¼\Ñ0JÔ0JÈHÌOÑ\Ô\ˆÝŒ{˜6 6Ñ*Ô*Ð*r   c                 ó2  ‡‡— t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          ‰j        ‰j        k    ˆˆfd„¦  «         t          ‰j        ‰j        k    ˆˆfd„¦  «         | j                             ‰j        ‰j        ¦  «         d S )Nc                  ó(   •— dt          ‰ ¦  «        › �S )Nz+expected 'value' to be a tensor, but got a r8   r)   s   €r   r.   z,GluonSemantic.shared_store.<locals>.<lambda>  s   ø€ Ð7rÕeiÐjoÑepÔepÐ7rÐ7r€ r   c                  ó*   •— d‰j         › d‰ j         › d�S )Nzsource shape z and destination shape ú must matchrA   ©rô   r*   s   €€r   r.   z,GluonSemantic.shared_store.<locals>.<lambda>  ó   ø€ Ðf u¤{ÐfÐfÈ8Ì>ÐfÐfÐf€ r   c                  ó*   •— d‰j         › d‰ j         › d�S )Nzsource dtype z and destination dtype rù   ©rÒ   rú   s   €€r   r.   z,GluonSemantic.shared_store.<locals>.<lambda>  rû   r   )	r   r#   rz   r}   rB   rÒ   rq   Úcreate_local_storer~   )rg   rô   r*   s    ``r   Úshared_storezGluonSemantic.shared_store  sª   øø€ Ý�z˜%¥¤Ñ-Ô-Ð/rÐ/rÐ/rÐ/rÑsÔsÐsÝˆuŒ{˜hœnÒ,ØfÐfÐfÐfÐfñ	hô 	hð 	håˆuŒ{˜hœnÒ,ØfÐfÐfÐfÐfñ	hô 	hð 	hàŒ×'Ò'¨¬¸¼ÑFÔFÐFÐFÐFr   c                 ó¦  ‡‡‡— t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          t          ‰t          ¦  «        ˆfd„¦  «         t          t          ‰j        ¦  «        ‰j        k    ˆˆfd„¦  «         t          d‰cxk    o
‰j        k     nc ˆˆfd„¦  «         t          ‰j         	                    ¦   «         ˆfd„¦  «         t          j
        ‰j        ‰j        ‰j        j        ¦  «        }| j                             |                     | j        ¦  «        ‰j        ‰j        ‰¦  «        }t          j        ||¦  «        S )Nc                  ó(   •— dt          ‰ ¦  «        › �S ©Nz-expected 'indices' to be a tensor, but got a r8   ©Úindicess   €r   r.   z-GluonSemantic.shared_gather.<locals>.<lambda>  ó   ø€ ÐVÅtÈGÁ}Ä}ÐVÐV€ r   c                  ó(   •— dt          ‰ ¦  «        › �S ©Nz(expected 'axis' to be an int, but got a r8   ©r’   s   €r   r.   z-GluonSemantic.shared_gather.<locals>.<lambda>  ó   ø€ Ð.eÕY]Ð^bÑYcÔYcÐ.eÐ.e€ r   c                  óB   •— dt          ‰ j        ¦  «        › d‰j        › �S ©Nz*indices rank must match memdesc rank: got r‹   ©rS   rB   r\   ©r  rô   s   €€r   r.   z-GluonSemantic.shared_gather.<locals>.<lambda>  ó&   ø€ ÐiÅÀWÄ]ÑASÔASÐiÐiÐZbÔZgÐiÐi€ r   r   c                  ó   •— d‰ › d‰j         › �S ©Nzaxis z# is out of bounds for memdesc rank ©r\   ©r’   rô   s   €€r   r.   z-GluonSemantic.shared_gather.<locals>.<lambda>  ó   ø€ Ð2r¸$Ð2rÐ2rÐckÔcpÐ2rÐ2r€ r   c                  ó   •— d‰ j         › �S ©Nz%indices must have integer dtype, got rý   r  s   €r   r.   z-GluonSemantic.shared_gather.<locals>.<lambda>  ó   ø€ Ð/fÐW^ÔWdÐ/fÐ/f€ r   )r   r#   rz   r}   r$   rS   rB   r\   rÒ   Úis_intr{   r9   rZ   rq   Úcreate_local_gatherr»   r~   )rg   rô   r  r’   r½   r~   s    ```  r   Úshared_gatherzGluonSemantic.shared_gather  sT  øøø€ Ý�z˜'¥4¤;Ñ/Ô/ØVÐVÐVÐVñ	Xô 	Xð 	Xå�z˜$¥Ñ$Ô$Ð&eÐ&eÐ&eÐ&eÑfÔfÐfÝÝ�”ÑÔ (¤-Ò/ØiÐiÐiÐiÐiñ	kô 	kð 	kõ 	ˆq�DÐ(Ð(Ò(Ð(˜8œ=Ò(Ð(Ð(Ð(Ð*rÐ*rÐ*rÐ*rÐ*rÑsÔsÐsÝˆwŒ}×#Ò#Ñ%Ô%Ð'fÐ'fÐ'fÐ'fÑgÔgÐgåÔ& x¤~°w´}ÀgÄlÔFYÑZÔZˆØ”×1Ò1°&·,²,¸t¼|Ñ2LÔ2LÈhÌoÐ_fÔ_mÐosÑtÔtˆÝŒ{˜6 6Ñ*Ô*Ð*r   c                 ó8  ‡‡‡‡— t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          t          ‰t          ¦  «        ˆfd„¦  «         t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          t          ‰j        ¦  «        ‰j        k    ˆˆfd„¦  «         t          d‰cxk    o
‰j        k     nc ˆˆfd„¦  «         t          ‰j         	                    ¦   «         ˆfd„¦  «         t          ‰j        ‰j        k    ˆˆfd„¦  «         t          ‰j
        j        ‰j
        j        k    d	„ ¦  «         t          ‰j        ‰j        k    ˆˆfd
„¦  «         | j                             ‰j        ‰j        ‰j        ‰¦  «         d S )Nc                  ó(   •— dt          ‰ ¦  «        › �S r  r8   r  s   €r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>%  r  r   c                  ó(   •— dt          ‰ ¦  «        › �S r  r8   r  s   €r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>&  r	  r   c                  ó(   •— dt          ‰ ¦  «        › �S )Nz,expected 'values' to be a tensor, but got a r8   )Úvaluess   €r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>'  s   ø€ Ð8uÕgkÐlrÑgsÔgsÐ8uÐ8u€ r   c                  óB   •— dt          ‰ j        ¦  «        › d‰j        › �S r  r  r  s   €€r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>*  r  r   r   c                  ó   •— d‰ › d‰j         › �S r  r  r  s   €€r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>+  r  r   c                  ó   •— d‰ j         › �S r  rý   r  s   €r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>,  r  r   c                  ó(   •— d‰j         › d‰ j         › �S )Nz0values must have the same shape as indices: got r‹   rA   )r  r  s   €€r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>.  s   ø€ ÐlÈ&Ì,ÐlÐlÐ]dÔ]jÐlÐl€ r   c                  ó   — dS )Nz+values must have the same layout as indicesr   r   r   r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>/  s   € ÐBo€ r   c                  ó(   •— d‰j         › d‰ j         › �S )Nz=values element type must match destination element type: got r‹   rý   )rô   r  s   €€r   r.   z.GluonSemantic.shared_scatter.<locals>.<lambda>2  s    ø€ ÐwÐTZÔT`ÐwÐwÐgoÔguÐwÐw€ r   )r   r#   rz   r}   r$   rS   rB   r\   rÒ   r  r9   rZ   rq   Úcreate_local_scatterr~   )rg   rô   r  r  r’   s    ````r   Úshared_scatterzGluonSemantic.shared_scatter#  sÉ  øøøø€ Ý�z˜'¥4¤;Ñ/Ô/ØVÐVÐVÐVñ	Xô 	Xð 	Xå�z˜$¥Ñ$Ô$Ð&eÐ&eÐ&eÐ&eÑfÔfÐfÝ�z˜&¥$¤+Ñ.Ô.Ð0uÐ0uÐ0uÐ0uÑvÔvÐvÝÝ�”ÑÔ (¤-Ò/ØiÐiÐiÐiÐiñ	kô 	kð 	kõ 	ˆq�DÐ(Ð(Ò(Ð(˜8œ=Ò(Ð(Ð(Ð(Ð*rÐ*rÐ*rÐ*rÐ*rÑsÔsÐsÝˆwŒ}×#Ò#Ñ%Ô%Ð'fÐ'fÐ'fÐ'fÑgÔgÐgÝˆvŒ|˜wœ}Ò,ØlÐlÐlÐlÐlñ	nô 	nð 	nåˆvŒ{Ô! W¤\Ô%8Ò8Ð:oÐ:oÑpÔpÐpÝØŒL˜HœNÒ*ØwÐwÐwÐwÐwñ	yô 	yð 	yð 	Œ×)Ò)¨(¬/¸6¼=È'Ì.ÐZ^Ñ_Ô_Ð_Ð_Ð_r   c                 ó(  — t          |t          j        ¦  «        st          dt	          |¦  «        › �¦  «        ‚t          |t          j        ¦  «        st          dt	          |¦  «        › �¦  «        ‚|j        |j        k    r t          d|j        › d|j        › d�¦  «        ‚|j        |j        k    r t          d|j        › d|j        › d�¦  «        ‚|j        |j	        t          |j        ¦  «         d …         k    rt          d	|j        › d
|j	        › �¦  «        ‚|j                             | j        ¦  «        }|j                             | j        ¦  «        }| j                             ||t          |j        ¦  «        |j        j        ¦  «        S )NzIbank_conflicts expects the register layout to be a distributed_type, got zTbank_conflicts expects the shared layout to be a shared_memory_descriptor_type, got zregister shape z and shared shape rù   z$mismatched dtypes between register (z) and shared (z	) layoutsz,bank_conflicts NYI for subslices. Got shape z and alloc_shape )r#   rz   r{   rß   r9   rí   rB   rŒ   rY   Úalloc_shaperS   rZ   Ú_to_irrq   Úget_shared_bank_conflictsrR   rV   )rg   Údistr_tyÚ	shared_tyÚreg_attrÚshared_attrs        r   Úbank_conflictszGluonSemantic.bank_conflicts6  s·  € Ý˜(¥DÔ$9Ñ:Ô:ð 	nÝØlÕ\`ÐaiÑ\jÔ\jÐlÐlñnô nð nõ ˜)¥TÔ%GÑHÔHð 	ÝØxÕgkÐluÑgvÔgvÐxÐxñô ð ð Œ>˜Yœ_Ò,Ð,ÝÐm¨x¬~ÐmÐmÐQZÔQ`ÐmÐmÐmÑnÔnÐnØÔ 8Ô#6Ò6Ð6ÝØy°xÔ7JÐyÐyÐZcÔZnÐyÐyÐyñô ð ð Œ?˜iÔ3µS¸¼Ñ5IÔ5IÐ4IÐ4JÐ4JÔKÒKÐKÝØx¸y¼ÐxÐxÐajÔavÐxÐxñô ð ð ”?×)Ò)¨$¬,Ñ7Ô7ˆØÔ&×-Ò-¨d¬lÑ;Ô;ˆØŒ|×5Ò5°hÀÍTÐRZÔR`ÑMaÔMaØ6>Ô6IÔ6\ñ^ô ^ð 	^r   c           	      óh  ‡— ddl m}m} t          t	          ‰t
          t          ||f¦  «        ˆfd„¦  «         t	          ‰t          t          t          f¦  «        rt          j        ‰¦  «        S t          j        | j                             ‰                     | j        ¦  «        |¦  «        ¦  «        S )Nr   )ÚTensorMemoryLayoutÚTensorMemoryScalesLayoutc                  ó(   •— dt          ‰ ¦  «        › �S )NzcExpected a DistributedLayout, SharedLayout, or TensorMemoryLayout or TensorMemoryScalesLayout, got r8   rŸ   s   €r   r.   z0GluonSemantic.to_linear_layout.<locals>.<lambda>V  s&   ø€ ð AÕrvÐw}Ñr~Ôr~ð  Að  Að r   )Ú3triton.experimental.gluon.language.nvidia.blackwellr1  r2  r   r#   r   r   r   r   r   rz   Ú	constexprrq   rà   r)  )rg   rZ   rB   r1  r2  s    `   r   rà   zGluonSemantic.to_linear_layoutP  sÒ   ø€ ð	
ð 	
ð 	
ð 	
ð 	
ð 	
ð 	
ð 	
õ 	Ý�vÕ 1µ<ÐASÐUmÐnÑoÔoð rAð rAð rAð rAñ	
ô 	
ð 	
õ
 �f�zÕ+BÕDVÐWÑXÔXð 	*Ý”> &Ñ)Ô)Ð)åŒ~˜dœl×;Ò;¸F¿MºMÈ$Ì,Ñ<WÔ<WÐY^Ñ_Ô_Ñ`Ô`Ð`r   c                 óD   — | j                              |j        ¦  «         d S r   )rq   Úcreate_local_deallocr~   )rg   rô   s     r   Úshared_dealloczGluonSemantic.shared_dealloc_  s    € ØŒ×)Ò)¨(¬/Ñ:Ô:Ð:Ð:Ð:r   c                 ó’  ‡‡— ‰j         }t          t          ‰t          ¦  «        ˆfd„¦  «         t          t          |j        t
          ¦  «        ˆfd„¦  «         | j                             ‰                     | j        ¦  «        ‰j	        ¦  «        }t          j        |j        |j        ‰¦  «        }|                      ||¦  «        S )Nc                  ó   •— d‰ › �S )Nz9set_auto_layout must set to a distributed layout but got r   rŸ   s   €r   r.   z/GluonSemantic.set_auto_layout.<locals>.<lambda>e  s   ø€ Ð[ÐSYÐ[Ð[€ r   c                  ó"   •— d‰ j         j        › �S )Nz4set_auto_layout input must have auto layout but got ©r9   rZ   r)   s   €r   r.   z/GluonSemantic.set_auto_layout.<locals>.<lambda>g  s   ø€ ÐaÈeÌjÔN_ÐaÐa€ r   )r9   r   r#   r   rZ   r   rq   Úcreate_set_auto_layoutr)  r~   rz   r{   rY   rB   r}   )rg   r*   rZ   Úsrc_tyr~   Úres_tys    ``   r   rÅ   zGluonSemantic.set_auto_layoutb  s¾   øø€ Ø”ˆÝ�z˜&Õ"3Ñ4Ô4Ø[Ð[Ð[Ð[ñ	]ô 	]ð 	]å�z˜&œ-­Ñ4Ô4ØaÐaÐaÐañ	cô 	cð 	cà”×4Ò4°V·]²]À4Ä<Ñ5PÔ5PÐRWÔR^Ñ_Ô_ˆÝÔ& vÔ'8¸&¼,ÈÑOÔOˆØ�{Š{˜6 6Ñ*Ô*Ð*r   c                 ó2  ‡‡‡— t          t          ‰t          ¦  «        ˆfd„¦  «         t          t          ‰t          ¦  «        ˆfd„¦  «         t          t          ‰t          ¦  «        ˆfd„¦  «         dg|j        z  }‰|‰<   t	          |j        ¦  «        }‰|‰<   |j        }t          j        |j	        |||j
        j        ¦  «        }| j        }	|	                     |                     |	¦  «        |j        |¦  «        }
t          j        |
fi |j        ¤ŽS )Nc                  ó   •— d‰ › �S )Nz&expected 'start' to be an int but got r   )rË   s   €r   r.   z-GluonSemantic.memdesc_slice.<locals>.<lambda>m  s   ø€ Ð/_ÐX]Ð/_Ð/_€ r   c                  ó   •— d‰ › �S )Nz'expected 'length' to be an int but got r   )Úlengths   €r   r.   z-GluonSemantic.memdesc_slice.<locals>.<lambda>n  s   ø€ Ð0bÐZ`Ð0bÐ0b€ r   c                  ó   •— d‰ › �S )Nz$expected 'dim' to be an int but got r   r¡   s   €r   r.   z-GluonSemantic.memdesc_slice.<locals>.<lambda>o  s   ø€ Ð-YÐTWÐ-YÐ-Y€ r   r   )r   r#   r$   r\   rR   rB   rZ   rz   rí   rÒ   r9   r(  rq   Úcreate_memdesc_subslicer»   r~   rï   Ú__dict__)rg   rô   rË   rC  r?   ÚoffsetsrB   rZ   r€   rq   r~   s     ```      r   Úmemdesc_slicezGluonSemantic.memdesc_slicel  s  øøø€ Ý�z˜%¥Ñ%Ô%Ð'_Ð'_Ð'_Ð'_Ñ`Ô`Ð`Ý�z˜&¥#Ñ&Ô&Ð(bÐ(bÐ(bÐ(bÑcÔcÐcÝ�z˜#�sÑ#Ô#Ð%YÐ%YÐ%YÐ%YÑZÔZÐZØ�#˜œÑ%ˆØˆ�‰Ý�X”^Ñ$Ô$ˆØˆˆc‰
Ø”ˆÝÔ/°´ÀÀvÈxÌ}ÔOhÑiÔiˆØ”,ˆØ×0Ò0°·²¸'Ñ1BÔ1BÀHÄOÐU\Ñ]Ô]ˆÝÔ,¨VÐCÐC°r´{ÐCÐCÐCr   c                 ó¬  ‡— |                       ‰¦  «        Št          ‰j        t          j        k    ˆfd„¦  «         |j        dd …         }|                       ‰¦  «        j        Š|j        }t          j        |j	        |||¦  «        }| j
        }|                     |                     |¦  «        |j        ‰¦  «        }t          j        |fi |j        ¤ŽS )Nc                  ó   •— d‰ j         › �S )Nz%expected 'index' to be int32 but got r8   ©Úindexs   €r   r.   z-GluonSemantic.memdesc_index.<locals>.<lambda>|  s   ø€ Ð1eÐY^ÔYcÐ1eÐ1e€ r   r	   )Ú	to_tensorr   r9   rz   rÉ   rB   r~   rZ   rí   rÒ   rq   Úcreate_memdesc_indexr»   rï   rF  )rg   rô   rL  rB   rZ   r€   rq   r~   s     `     r   Úmemdesc_indexzGluonSemantic.memdesc_indexz  sÂ   ø€ Ø—’˜uÑ%Ô%ˆÝˆuŒz�TœZÒ'Ð)eÐ)eÐ)eÐ)eÑfÔfÐfØ”˜q˜r˜rÔ"ˆØ—’˜uÑ%Ô%Ô,ˆØ”ˆÝÔ/°´ÀÀvÈuÑUÔUˆØ”,ˆØ×-Ò-¨b¯hªh°wÑ.?Ô.?ÀÄÐRWÑXÔXˆÝÔ,¨VÐCÐC°r´{ÐCÐCÐCr   c                 óü  ‡‡‡— t          t          ‰¦  «        ˆfd„¦  «         t          t          ‰¦  «        t          ‰j        ¦  «        k    ˆˆfd„¦  «         ˆfd„‰D ¦   «         }‰j        j        Š‰d t          ‰¦  «        ‰j        z
  …         }|ˆˆfd„‰D ¦   «         z  }| j                             ‰j	        ‰¦  «        }| j         
                    |¦  «        }t          j        |‰j        |||¬¦  «        S )Nc                  ó   •— d‰ › �S )Nz1all elements of 'order' must be integers but got r   )Úorders   €r   r.   z-GluonSemantic.memdesc_trans.<locals>.<lambda>†  ré   r   c                  ó:   •— d‰ j         › dt          ‰¦  «        › d�S )Nzsource rank (z) and order length (z) must match)r\   rS   )rô   rR  s   €€r   r.   z-GluonSemantic.memdesc_trans.<locals>.<lambda>‰  s#   ø€ Ð_ H¤MÐ_Ð_ÅsÈ5ÁzÄzÐ_Ð_Ð_€ r   c                 ó*   •— g | ]}‰j         |         ‘ŒS r   rA   )r%   r&   rô   s     €r   r™   z/GluonSemantic.memdesc_trans.<locals>.<listcomp>‹  s    ø€ Ð2Ð2Ð2 q�” Ô"Ð2Ð2Ð2r   c                 óZ   •— g | ]'}‰t          ‰¦  «        ‰j        z
  d …         |         ‘Œ(S r   )rS   r\   )r%   r&   r(  rô   s     €€r   r™   z/GluonSemantic.memdesc_trans.<locals>.<listcomp>Ž  s9   ø€ Ð]Ð]Ð]ÐRS˜K­¨KÑ(8Ô(8¸8¼=Ñ(HÐ(IÐ(IÔJÈ1ÔMÐ]Ð]Ð]r   ©rY   rB   r(  rZ   )r   r+   rS   rB   r9   r(  r\   rq   Úcreate_memdesc_transr~   Úget_gluon_layout_from_memdescrz   rï   rÒ   )rg   rô   rR  rB   Únew_alloc_shaper~   rZ   r(  s    ``    @r   Úmemdesc_transzGluonSemantic.memdesc_trans…  s%  øøø€ Ý�|˜EÑ"Ô"Ð$gÐ$gÐ$gÐ$gÑhÔhÐhÝÝ�‰JŒJ�#˜hœnÑ-Ô-Ò-Ø_Ð_Ð_Ð_Ð_ñ	aô 	að 	að 3Ð2Ð2Ð2¨EÐ2Ñ2Ô2ˆØ”mÔ/ˆØ%Ð&G¥s¨;Ñ'7Ô'7¸(¼-Ñ'GÐ&GÔHˆØÐ]Ð]Ð]Ð]Ð]ÐW\Ð]Ñ]Ô]Ñ]ˆà”×2Ò2°8´?ÀEÑJÔJˆØ”×;Ò;¸FÑCÔCˆÝÔ,¨VÀÄÐV[Ø9HÐQWðYñ Yô Yð 	Yr   c                 óò  ‡‡— t          t          ‰¦  «        ˆfd„¦  «         t          t          j        ‰¦  «        t          j        ‰j        ¦  «        k    ˆˆfd„¦  «         | j                             ‰j        ‰¦  «        }| j                             |¦  «        }‰j	        j
        }t          |¦  «        ‰j        z
  }|d |…         t          ‰¦  «        z   }t          j        |‰j        ‰||¬¦  «        S )Nc                  ó   •— d‰ › �S rè   r   rA   s   €r   r.   z/GluonSemantic.memdesc_reshape.<locals>.<lambda>–  ré   r   c                  ó   •— d‰ j         › d‰› �S )Nz)memdesc_reshape total elements mismatch: z -> rA   )rô   rB   s   €€r   r.   z/GluonSemantic.memdesc_reshape.<locals>.<lambda>™  s&   ø€ ð 4Ø œð4ð 4Ø,1ð4ð 4€ r   rV  )r   r+   ÚmathÚprodrB   rq   Úcreate_memdesc_reshaper~   rX  r9   r(  rS   r\   rR   rz   rï   rÒ   )rg   rô   rB   r~   rZ   r(  Ú
prefix_lenrY  s    ``     r   Úmemdesc_reshapezGluonSemantic.memdesc_reshape•  s  øø€ Ý�|˜EÑ"Ô"Ð$gÐ$gÐ$gÐ$gÑhÔhÐhÝÝŒI�eÑÔ¥¤	¨(¬.Ñ 9Ô 9Ò9ð5ð 5ð 5ð 5ð 5ñ	
ô 	
ð 	
ð ”×4Ò4°X´_ÀeÑLÔLˆØ”×;Ò;¸FÑCÔCˆØ”mÔ/ˆÝ˜Ñ%Ô%¨¬Ñ5ˆ
Ø% k z kÔ2µT¸%±[´[Ñ@ˆåÔ,ØØ”~ØØ'Øð
ñ 
ô 
ð 	
r   c                 óÀ  ‡‡‡— t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          t	          ‰¦  «        ˆfd„¦  «         t          t          ‰t          j        ¦  «        ˆfd„¦  «         t          j        ‰‰‰‰¦  «        }| j                             | 	                    | j        ¦  «        |j
        ¦  «        }t          j        |fi |j        ¤ŽS )Nc                  ó   •— d‰ › �S )Nz'expected 'dtype' to be a dtype but got r   rý   s   €r   r.   z3GluonSemantic.memdesc_reinterpret.<locals>.<lambda>¬  s   ø€ Ð6gÐ`eÐ6gÐ6g€ r   c                  ó   •— d‰ › �S rè   r   rA   s   €r   r.   z3GluonSemantic.memdesc_reinterpret.<locals>.<lambda>­  ré   r   c                  ó   •— d‰ › �S rë   r   rŸ   s   €r   r.   z3GluonSemantic.memdesc_reinterpret.<locals>.<lambda>¯  rì   r   )r   r#   rz   rÒ   r+   r   rí   rq   Úcreate_memdesc_reinterpretr»   r~   rï   rF  )rg   rô   rÒ   rB   rZ   r€   r~   s     ```  r   Úmemdesc_reinterpretz!GluonSemantic.memdesc_reinterpret«  s×   øøø€ Ý�z˜%¥¤Ñ,Ô,Ð.gÐ.gÐ.gÐ.gÑhÔhÐhÝ�|˜EÑ"Ô"Ð$gÐ$gÐ$gÐ$gÑhÔhÐhÝ�z˜&¥$Ô"3Ñ4Ô4ØQÐQÐQÐQñ	Sô 	Sð 	SåÔ/°°u¸fÀeÑLÔLˆØ”×8Ò8¸¿ºÀ$Ä,Ñ9OÔ9OÐQYÔQ`ÑaÔaˆÝÔ,¨VÐCÐC°r´{ÐCÐCÐCr   c                 ód   — |rt          j        |||¦  «        }n|}|                      ||¦  «        S r   )rz   r{   r}   )rg   r˜   r   r�   rZ   r?  s         r   Úwrap_tensorzGluonSemantic.wrap_tensor´  s9   € Øð 	ÝÔ*¨9°iÀÑHÔHˆFˆFàˆFØ�{Š{˜1˜fÑ%Ô%Ð%r   c                 ó  ‡‡‡— | D ]2Št          t          ‰j        t          j        ¦  «        ˆfd„¦  «         Œ3d„ | D ¦   «         Š‰d         Št          t          ˆfd„‰dd …         D ¦   «         ¦  «        ˆfd„¦  «         d S )Nc                  ó   •— d‰ j         ›�S ©Nz#expected distributed_type but got: r8   )r˜   s   €r   r.   z2GluonSemantic._check_same_layout.<locals>.<lambda>¾  s   ø€ ÐFvÐlmÔlrÐFvÐFv€ r   c                 ó&   — g | ]}|j         j        ‘ŒS r   r<  r—   s     r   r™   z4GluonSemantic._check_same_layout.<locals>.<listcomp>¿  s   € Ð-Ð-Ð- Q�1”6”=Ð-Ð-Ð-r   r   c              3   ó$   •K  — | ]
}|‰k    V — Œd S r   r   )r%   ÚlÚl0s     €r   r'   z3GluonSemantic._check_same_layout.<locals>.<genexpr>Á  s'   øè è € Ð0Ð0˜q�1˜’7Ð0Ð0Ð0Ð0Ð0Ð0r   r	   c                  ó   •— d‰ › �S )Nz3Expected inputs to have matching layouts, but got: r   )Úlayoutss   €r   r.   z2GluonSemantic._check_same_layout.<locals>.<lambda>Â  s   ø€ ÐVÈWÐVÐV€ r   )r   r#   r9   rz   r{   r(   )Úxsrq  rs  r˜   s    @@@r   Ú_check_same_layoutz GluonSemantic._check_same_layout»  s²   øøø€ àð 	xð 	xˆAÝ•:˜aœf¥dÔ&;Ñ<Ô<Ð>vÐ>vÐ>vÐ>vÑwÔwÐwÐwØ-Ð-¨"Ð-Ñ-Ô-ˆØ�QŒZˆÝ�sÐ0Ð0Ð0Ð0 G¨A¨B¨B¤KÐ0Ñ0Ô0Ñ0Ô0ØVÐVÐVÐVñ	Xô 	Xð 	Xð 	Xð 	Xr   ÚinputsÚreverse.c                 óà  ‡ ‡‡‡— ‰d         j         j        Št          ‰¦  «        }| |cxk    r|k     sn J d|› d|› d�¦   «         ‚|dk     r||z  }‰D ]}|j         j        ‰k    s
J d¦   «         ‚Œ‰ j                             d„ ‰D ¦   «         ||¦  «        Š |‰¦  «         ‰                     ¦   «         sJ ‚t          ˆˆˆ ˆfd„t          t          ‰¦  «        ¦  «        D ¦   «         ¦  «        S )Nr   z
scan axis z must be < inputs rank (ú)z(all scan inputs must have the same shapec                 ó   — g | ]	}|j         ‘Œ
S r   ©r~   ©r%   Úts     r   r™   z2GluonSemantic.associative_scan.<locals>.<listcomp>Ñ  s   € Ð+EÐ+EÐ+E¸¨A¬HÐ+EÐ+EÐ+Er   c              3   óŒ   •K  — | ]>}‰                      ‰                     |¦  «        ‰|         j        j        ‰¦  «        V — Œ?d S r   ©r�   Ú
get_resultr9   rƒ   )r%   r&   rv  Úscan_oprg   rB   s     €€€€r   r'   z1GluonSemantic.associative_scan.<locals>.<genexpr>Õ  sa   øè è € ð )ð )àð ×*Ò*¨7×+=Ò+=¸aÑ+@Ô+@À&ÈÄ)Ä.ÔBWÐY^Ñ_Ô_ð)ð )ð )ð )ð )ð )r   )r9   rB   rS   rq   Úcreate_scanÚverifyÚtupleÚrange)	rg   rv  r’   Úregion_builder_fnrw  r\   r}  r�  rB   s	   ``     @@r   Úassociative_scanzGluonSemantic.associative_scanÄ  sK  øøøø€ à�q”	”Ô$ˆÝ�5‰zŒzˆàˆu˜Ð#Ð#Ò#Ð#˜tÒ#Ð#Ð#Ð#Ð#Ð%W°$Ð%WÐ%WÐPTÐ%WÐ%WÐ%WÑ#Ô#Ð#à�!Š8ˆ8Ø�D‰LˆDàð 	Uð 	UˆAØ”6”< 5Ò(Ð(Ð(Ð*TÑ(Ô(Ð(Ð(à”,×*Ò*Ð+EÐ+E¸fÐ+EÑ+EÔ+EÀtÈWÑUÔUˆØÐ˜'Ñ"Ô"Ð"Ø�~Š~ÑÔÐÐÐåð )ð )ð )ð )ð )ð )ð )å�3˜v™;œ;Ñ'Ô'ð)ñ )ô )ñ )ô )ð 	)r   c                 óz  ‡ ‡‡‡‡‡‡— ‰€t          ˆ fd„‰D ¦   «         ¦  «        ŠdŠ‰d         j        j        Št          ‰¦  «        Št	          d‰cxk    o‰k     nc ˆˆfd„¦  «         ‰                      ‰¦  «         ˆfd„t          ‰¦  «        D ¦   «         Št          ˆfd„‰D ¦   «         ¦  «        s
J d¦   «         ‚‰ j         	                    d„ ‰D ¦   «         ‰¦  «        Š |‰¦  «         ‰ 
                    ¦   «         sJ ‚t          ˆˆˆˆ fd„t          t          ‰¦  «        ¦  «        D ¦   «         ¦  «        S )	Nc              3   ó^   •K  — | ]'}‰                      ||j        j        gd ¬¦  «        V — Œ(dS )F)rÍ   N)rÐ   Únumelr*   )r%   r}  rg   s     €r   r'   z*GluonSemantic.reduction.<locals>.<genexpr>Û  s<   øè è € Ð_Ð_ÐST˜4Ÿ<š<¨¨A¬G¬M¨?È˜<ÑNÔNÐ_Ð_Ð_Ð_Ð_Ð_r   r   c                  ó   •— d‰› d‰ › �S )Nz/expected reduction axis to be in the range [0, z
) but got r   ©r’   r\   s   €€r   r.   z)GluonSemantic.reduction.<locals>.<lambda>à  s   ø€ Ð)qÐ[_Ð)qÐ)qÐkoÐ)qÐ)q€ r   c                 ó&   •— g | ]\  }}|‰k    ¯|‘ŒS r   r   )r%   r&   Úsr’   s      €r   r™   z+GluonSemantic.reduction.<locals>.<listcomp>â  s"   ø€ ÐAÐAÐA™4˜1˜a°q¸D²y°y�Q°y°y°yr   c              3   ó8   •K  — | ]}|j         j        ‰k    V — Œd S r   )r9   rB   )r%   r}  rB   s     €r   r'   z*GluonSemantic.reduction.<locals>.<genexpr>ã  s,   øè è € Ð9Ð9¨Q�1”6”< 5Ò(Ð9Ð9Ð9Ð9Ð9Ð9r   z-all reduction inputs must have the same shapec                 ó   — g | ]	}|j         ‘Œ
S r   r{  r|  s     r   r™   z+GluonSemantic.reduction.<locals>.<listcomp>å  s   € Ð/IÐ/IÐ/I¸Q°´Ð/IÐ/IÐ/Ir   c              3   óŒ   •K  — | ]>}‰                      ‰                     |¦  «        ‰|         j        j        ‰¦  «        V — Œ?d S r   r  )r%   r&   rv  Ú	reduce_opr�   rg   s     €€€€r   r'   z*GluonSemantic.reduction.<locals>.<genexpr>é  sa   øè è € ð )ð )àð ×*Ò*¨9×+?Ò+?ÀÑ+BÔ+BÀFÈ1ÄIÄNÔDYÐ[dÑeÔeð)ð )ð )ð )ð )ð )r   )r„  r9   rB   rS   r   ru  rX   r(   rq   Úcreate_reducerƒ  r…  )rg   rv  r’   r†  r\   r’  r�   rB   s   ``` @@@@r   Ú	reductionzGluonSemantic.reductionÙ  s“  øøøøøøø€ Øˆ<ÝÐ_Ð_Ð_Ð_ÐX^Ð_Ñ_Ô_Ñ_Ô_ˆFØˆDà�q”	”Ô$ˆÝ�5‰zŒzˆÝˆq�DÐÐÒÐ˜4ÒÐÐÐÐ!qÐ!qÐ!qÐ!qÐ!qÑrÔrÐrØ×Ò Ñ'Ô'Ð'ØAÐAÐAÐA¥9¨UÑ#3Ô#3ÐAÑAÔAˆ	ÝÐ9Ð9Ð9Ð9°&Ð9Ñ9Ô9Ñ9Ô9ÐjÐjÐ;jÑjÔjÐ9à”L×.Ò.Ð/IÐ/IÀ&Ð/IÑ/IÔ/IÈ4ÑPÔPˆ	ØÐ˜)Ñ$Ô$Ð$Ø×ÒÑ!Ô!Ð!Ð!Ð!åð )ð )ð )ð )ð )ð )ð )å�3˜v™;œ;Ñ'Ô'ð)ñ )ô )ñ )ô )ð 	)r   Únum_binsÚmaskc                 ó"  — t          t          |j        ¦  «        dk    d„ ¦  «         t          |j                             ¦   «         d„ ¦  «         t          |d ud„ ¦  «         |�M|                      ||¦  «        \  }}t          |j        j                             ¦   «         d„ ¦  «         |j	        }| 
                    | j        ¦  «        }| j                             |j	        |||¦  «        }|                      |t          j        |g|¦  «        S )Nr	   c                  ó   — dS )Nz histogram only supports 1D inputr   r   r   r   r.   z)GluonSemantic.histogram.<locals>.<lambda>î  s   € Ð.P€ r   c                  ó   — dS )Nz%histogram only supports integer inputr   r   r   r   r.   z)GluonSemantic.histogram.<locals>.<lambda>ï  s   € Ð-T€ r   c                  ó   — dS )Nz'histogram requires a destination layoutr   r   r   r   r.   z)GluonSemantic.histogram.<locals>.<lambda>ð  s   € Ð+T€ r   c                  ó   — dS )Nz"Mask must have boolean scalar typer   r   r   r   r.   z)GluonSemantic.histogram.<locals>.<lambda>ó  s   € Ð7[€ r   )r   rS   rB   rÒ   r  rª   r9   rƒ   Úis_boolr~   r)  rq   Úcreate_histogramrj  rz   rÉ   )rg   r‘   r•  r–  rZ   Úlayout_attrr~   s          r   Ú	histogramzGluonSemantic.histogramí  sù   € Ý�s�5”;ÑÔ 1Ò$Ð&PÐ&PÑQÔQÐQÝˆuŒ{×!Ò!Ñ#Ô#Ð%TÐ%TÑUÔUÐUÝˆv˜TÐ!Ð#TÐ#TÑUÔUÐUØÐØ×3Ò3°D¸%Ñ@Ô@‰KˆD�%Ý�4”9Ô#×+Ò+Ñ-Ô-Ð/[Ð/[Ñ\Ô\Ð\Ø”;ˆDØ—m’m D¤LÑ1Ô1ˆØ”×.Ò.¨u¬|¸XÀtÈ[ÑYÔYˆØ×Ò ­¬
°X°JÀÑGÔGÐGr   c           	      ó°  — t          |d ud„ ¦  «         t          |d„ ¦  «         t          t          |j        ¦  «        dk    d„ ¦  «         t          j        |j        j        |j        d         |j        d         z   g|¦  «        }|                      | j         	                    |j
        |j
        |                     | j        ¦  «        ¦  «        |¦  «        S )Nc                  ó   — dS )Nz!cat requires a destination layoutr   r   r   r   r.   z#GluonSemantic.cat.<locals>.<lambda>ú  s   € Ð+N€ r   c                  ó   — dS )Nz;current implementation of `cat` always may reorder elementsr   r   r   r   r.   z#GluonSemantic.cat.<locals>.<lambda>û  s   € Ð$a€ r   r	   c                  ó   — dS )Nzcat requires a rank-1 inputr   r   r   r   r.   z#GluonSemantic.cat.<locals>.<lambda>ü  s   € Ð,I€ r   r   )r   rS   rB   rz   r{   r9   rƒ   r}   rq   Ú
create_catr~   r»   )rg   r°   r±   rÍ   rZ   Úret_types         r   ÚcatzGluonSemantic.catù  s»   € Ýˆv˜TÐ!Ð#NÐ#NÑOÔOÐOÝˆ{ÐaÐaÑbÔbÐbÝ�s�3”9‰~Œ~ Ò"Ð$IÐ$IÑJÔJÐJÝÔ(¨¬¬¸3¼9ÀQ¼<È#Ì)ÐTUÌ,Ñ;VÐ:WÐY_Ñ`Ô`ˆØ�{Š{˜4œ<×2Ò2°3´:¸s¼zÈ8Ï>Ê>ÐZ^ÔZfÑKgÔKgÑhÔhÐjrÑsÔsÐsr   ÚsrcrL  c                 óf  ‡‡‡‡— t          t          ‰j        t          j        ¦  «        ˆfd„¦  «         t          t          ‰j        t          j        ¦  «        ˆfd„¦  «         t          ‰j        j                             ¦   «         ˆfd„¦  «         t          ‰j        j        ¦  «        Št          t          ‰j        j        ¦  «        ‰k    d„ ¦  «         t          ‰ ‰cxk    o‰k     nc ˆˆfd„¦  «         ‰dk     r‰‰z  Št          ‰¦  «        D ]@}|‰k    rŒ	t          ‰j        j        |         ‰j        j        |         k    ˆfd„¦  «         ŒA| j
                             ‰j        ‰j        ‰¦  «        }|                      |‰j        j        ‰j        j        ‰j        j        ¦  «        S )Nc                  ó   •— d‰ j         ›�S rm  r8   )r§  s   €r   r.   z&GluonSemantic.gather.<locals>.<lambda>  s   ø€ ÐDvÐjmÔjrÐDvÐDv€ r   c                  ó   •— d‰ j         ›�S rm  r8   rK  s   €r   r.   z&GluonSemantic.gather.<locals>.<lambda>  s   ø€ ÐK¸U¼ZÐKÐK€ r   c                  ó"   •— d‰ j         j        ›�S )Nz&expected integer scalar type but got: )r9   rƒ   rK  s   €r   r.   z&GluonSemantic.gather.<locals>.<lambda>  s   ø€ Ð3qÐ\aÔ\fÔ\mÐ3qÐ3q€ r   c                  ó   — dS )Nz0source and index tensors must have the same rankr   r   r   r   r.   z&GluonSemantic.gather.<locals>.<lambda>  s   € Ð6h€ r   c                  ó   •— d‰ › d‰› d�S )Nzgather axis z must be < source rank (ry  r   rŒ  s   €€r   r.   z&GluonSemantic.gather.<locals>.<lambda>  s   ø€ Ð-a¸DÐ-aÐ-aÐZ^Ð-aÐ-aÐ-a€ r   r   c                  ó   •— d‰ › d�S )Nz
index dim z( must match the corresponding source dimr   r  s   €r   r.   z&GluonSemantic.gather.<locals>.<lambda>  s   ø€ ÐS TÐSÐSÐS€ r   )r   r#   r9   rz   r{   rƒ   r  rS   rB   r…  rq   Úcreate_gatherr~   rj  rZ   )rg   r§  rL  r’   ÚdÚgatherr\   s    ```  @r   r±  zGluonSemantic.gather   s¶  øøøø€ Ý�z˜#œ(¥DÔ$9Ñ:Ô:Ð<vÐ<vÐ<vÐ<vÑwÔwÐwÝ�z˜%œ*¥dÔ&;Ñ<Ô<ØKÐKÐKÐKñ	Mô 	Mð 	MåˆuŒzÔ ×'Ò'Ñ)Ô)Ð+qÐ+qÐ+qÐ+qÑrÔrÐrå�3”8”>Ñ"Ô"ˆÝ�s�5”:Ô#Ñ$Ô$¨Ò,Ð.hÐ.hÑiÔiÐiÝ�ˆu˜Ð#Ð#Ò#Ð#˜tÒ#Ð#Ð#Ð#Ð%aÐ%aÐ%aÐ%aÐ%aÑbÔbÐbØ�!Š8ˆ8Ø�D‰LˆDå�t‘”ð 	ð 	ˆAØ�DŠyˆyØÝØ”
Ô  Ô# s¤x¤~°aÔ'8Ò8ØSÐSÐSÐSñô ð ð ð ”×+Ò+¨C¬J¸¼ÀdÑKÔKˆØ×Ò ¨¬¬¸¼Ô9IÈ5Ì:ÔK\Ñ]Ô]Ð]r   c                 óô   — | j                              |j        |                     | j         ¦  «        |¦  «        }t	          |j        j        ¦  «        }||xx         dz  cc<   |                      |||¦  «        S )NrC   )rq   Úcreate_fp4_to_fpr~   r»   rR   r9   rB   r�   )rg   r§  Ú	elem_typer’   ÚresultrB   s         r   Ú	fp4_to_fpzGluonSemantic.fp4_to_fp  sk   € Ø”×.Ò.¨s¬z¸9¿?º?È4Ì<Ñ;XÔ;XÐZ^Ñ_Ô_ˆÝ�S”X”^Ñ$Ô$ˆØˆdˆˆŒ�qÑˆˆ‰Ø×-Ò-¨f°iÀÑGÔGÐGr   Úworker_num_warpsÚworker_num_regsc                 ó0  ‡‡‡‡— |D ]7\  }Št          t          ‰t          t          j        f¦  «        ˆfd„¦  «         Œ8t	          |¦  «        dk    s
J d¦   «         ‚|d         \  }}t	          |¦  «        dz
  }|dd …         }	|t	          |¦  «        k    sJ d|› dt	          |¦  «        › d�¦   «         ‚|�1|t	          |¦  «        k    sJ d|› dt	          |¦  «        › d�¦   «         ‚| j        }
|
                     ¦   «         }|
                     ¦   «         }|
                     |¦  «         | 	                    ||i ¬	¦  «        }t          |g¦  «        }|
                     |¦  «         d
„ |D ¦   «         }d„ |	D ¦   «         }t          |g ¦  «        }|
                     |¦  «         |
                     ||¦  «        Š‰                     ¦   «                              |¦  «         |�‰                     |¦  «         |
                     ‰                     ¦   «         g ¦  «         |
                     ||¦  «        }d„ |D ¦   «         }dŠt+          |	¦  «        D ]Ñ\  }\  }Št-          ||         ¬¦  «        }|
                     |                     |¦  «        |¦  «        Š||         }ˆˆfd„t1          t	          |¦  «        ¦  «        D ¦   «         }t3          |d„ ‰D ¦   «         ¦  «        }| 	                    ||i |¬¦  «         |
                     ¦   «          ‰t	          |¦  «        z  ŠŒÒ|
                     ‰                     ¦   «         ¦  «         ˆfd„t1          t	          |¦  «        ¦  «        D ¦   «         }t;          t3          ||j        g¦  «        ¦  «        S )Nc                  ó(   •— dt          ‰ ¦  «        › �S )Nz9function arguments must be a tuple of arguments, but got r8   )Úargss   €r   r.   z/GluonSemantic.warp_specialize.<locals>.<lambda>   s   ø€ ÐcÕW[Ð\`ÑWaÔWaÐcÐc€ r   r	   z8expected at least one function for the default partitionr   zwarp specialize got z partitions but z warp countsz register counts)Úkwargsc                 ó6   — g | ]}|                      ¦   «         ‘ŒS r   ©Úget_type)r%   Úrs     r   r™   z1GluonSemantic.warp_specialize.<locals>.<listcomp>9  s    € Ð;Ð;Ð;¨˜Ÿ
š
™œÐ;Ð;Ð;r   c                 ó2   — g | ]\  }}t          |¦  «        ‘ŒS r   )r   )r%   Ú_r»  s      r   r™   z1GluonSemantic.warp_specialize.<locals>.<listcomp><  s%   € ÐIÐIÐI±g°a¸Õ+¨DÑ1Ô1ÐIÐIÐIr   c                 ó6   — g | ]}|                      ¦   «         ‘ŒS r   r¾  ©r%   Úargs     r   r™   z1GluonSemantic.warp_specialize.<locals>.<listcomp>H  s    € Ð9Ð9Ð9¨�S—\’\‘^”^Ð9Ð9Ð9r   r:   c                 ó@   •— g | ]}‰                      ‰|z   ¦  «        ‘ŒS r   )Úget_argument)r%   ÚjÚarg_itÚblocks     €€r   r™   z1GluonSemantic.warp_specialize.<locals>.<listcomp>N  s+   ø€ ÐXÐXÐX¸Q˜%×,Ò,¨V°a©ZÑ8Ô8ÐXÐXÐXr   c                 ó   — g | ]	}|j         ‘Œ
S r   r8   rÄ  s     r   r™   z1GluonSemantic.warp_specialize.<locals>.<listcomp>O  s   € Ð9SÐ9SÐ9SÀs¸#¼(Ð9SÐ9SÐ9Sr   )r¼  Úcaller_contextc                 ó:   •— g | ]}‰                      |¦  «        ‘ŒS r   )r€  )r%   r&   Úws_ops     €r   r™   z1GluonSemantic.warp_specialize.<locals>.<listcomp>U  s'   ø€ ÐNÐNÐN°˜×(Ò(¨Ñ+Ô+ÐNÐNÐNr   )r   r#   r„  rz   rS   rq   Úget_insertion_pointÚ	new_blockÚset_insertion_point_to_startÚcall_JitFunctionr   Úcreate_warp_yieldÚsumÚrestore_insertion_pointÚcreate_warp_specializeÚget_default_regionÚ	push_backÚset_requested_registersÚcreate_block_with_parentÚget_partition_op_holderÚ!create_warp_specialize_partitionsrX   re   Ú
get_regionr…  r   Úcreate_warp_returnÚset_insertion_point_afterÚget_operationÚnextr9   )rg   Úfunctions_and_argsr·  r¸  Ú	generatorrÂ  Údefault_partitionÚdefault_argsÚnum_partitionsÚworkersrq   Ú	insert_ptÚdefault_blockÚdefault_resultÚmlir_resultsÚresult_typesÚworker_argsÚ	mlir_argsÚpartitions_opÚ	arg_typesr&   ÚfuncrÌ  Ú
block_argsrÉ  r»  rÊ  rÎ  s                           @@@@r   Úwarp_specializezGluonSemantic.warp_specialize  s  øøøø€ à)ð 	eð 	e‰GˆAˆtÝ•:˜d¥U­D¬JÐ$7Ñ8Ô8ØcÐcÐcÐcñeô eð eð eõ Ð%Ñ&Ô&¨!Ò+Ð+Ð+Ð-gÑ+Ô+Ð+Ø*<¸QÔ*?Ñ'Ð˜<ÝÐ/Ñ0Ô0°1Ñ4ˆØ$ Q R RÔ(ˆà¥Øñ"
ô "
ò 
ð 
ð 
àe .ÐeÐeÅ#ÐFVÑBWÔBWÐeÐeÐeñ
ô 
ð 
ð Ð&Ø!¥SØñ&ô &ò ð ð àl nÐlÐlÅcÈ/ÑFZÔFZÐlÐlÐlñô ð ð ”,ˆØ×/Ò/Ñ1Ô1ˆ	ð  ×)Ò)Ñ+Ô+ˆØ×,Ò,¨]Ñ;Ô;Ð;Ø"×3Ò3Ð4EÀ|Ð\^Ð3Ñ_Ô_ˆÝ+¨^Ð,<Ñ=Ô=ˆØ×!Ò! ,Ñ/Ô/Ð/Ø;Ð;¨lÐ;Ñ;Ô;ˆð JÐIÀÐIÑIÔIˆÝ˜ RÑ(Ô(ˆ	Ø×'Ò'¨	Ñ2Ô2Ð2Ø×.Ò.¨|Ð=MÑNÔNˆØ× Ò Ñ"Ô"×,Ò,¨]Ñ;Ô;Ð;àÐ&Ø×)Ò)¨/Ñ:Ô:Ð:ð 	×(Ò(¨×)FÒ)FÑ)HÔ)HÈ"ÑMÔMÐMØ×AÒAÀ)È^Ñ\Ô\ˆØ9Ð9¨yÐ9Ñ9Ô9ˆ	ØˆÝ(¨Ñ1Ô1ð 	%ð 	%‰OˆA‰|��dÝ/Ð:JÈ1Ô:MÐNÑNÔNˆNØ×4Ò4°]×5MÒ5MÈaÑ5PÔ5PÐR[Ñ\Ô\ˆEØ# AœˆIØXÐXÐXÐXÐXÅ%ÍÈIÉÌÑBWÔBWÐXÑXÔXˆJÝ,¨ZÐ9SÐ9SÈdÐ9SÑ9SÔ9SÑTÔTˆJØ×&Ò& t¨ZÀÐSaÐ&ÑbÔbÐbØ×&Ò&Ñ(Ô(Ð(Ø•c˜)‘n”nÑ$ˆFˆFà×)Ò)¨%×*=Ò*=Ñ*?Ô*?Ñ@Ô@Ð@ØNÐNÐNÐNµU½3¸|Ñ;LÔ;LÑ5MÔ5MÐNÑNÔNˆÝÕ'¨°~Ô7JÐ6KÑLÔLÑMÔMÐMr   c                 óH   — t          j        | j        j        j        ¦  «        S r   )rz   r5  rq   ÚoptionsÚnum_ctasrk   s    r   rö  zGluonSemantic.num_ctasX  s   € ÝŒ~˜dœlÔ2Ô;Ñ<Ô<Ð<r   c                 óÊ   — |j         �:t          |j         t          ¦  «        sJ ‚t          j        |j         j        ¦  «        S t          j        | j        j        j        ¦  «        S r   )rÌ  r#   re   rz   r5  r;   rq   rõ  )rg   rã  s     r   r;   zGluonSemantic.num_warps[  sS   € ØÔ#Ð/Ý˜iÔ6Õ8JÑKÔKÐKÐKÐKÝ”> )Ô":Ô"DÑEÔEÐEÝŒ~˜dœlÔ2Ô<Ñ=Ô=Ð=r   )F)8rs   rt   ru   rz   r}   Úlangr   Ú__annotations__rh   r�   r„   r   r$   r�   r   r¥   r¬   r   r¯   r´   r¾   rª   rÊ   ÚboolrÐ   rÔ   r×   rä   rð   rõ   rÿ   r  r&  r/  rà   r8  rÅ   rH  rO  rZ  rb  rh  rj  Ústaticmethodru  r   r‡  r”  rŸ  r¦  r±  r¶  ró  rö  r;   Ú__classcell__)r­   s   @r   rw   rw   a   s  ø€ € € € € € ØŒ[€FØ€DàÐÐÑð ð ð ð ð ð'ð 'ð 'ð_ð _ð _ð¨4°¬9ð ÀÀcÄð ð ð ð ð T ð T°ð T¸ð Tð Tð Tð Tð&5�hð 5 8ð 5°ð 5ð 5ð 5ð 5ð 5ð 5ðX�xð X E¨(°HÐ*<Ô$=ð Xð Xð Xð Xð Xð Xð5˜Xð 5¨U°3¬Zð 5¸Hð 5ð 5ð 5ð 5ð 5ð 5ð+¨(ð +¸5À¼:ð +È(ð +ð +ð +ð +ð ¨ð °xð ÀHð ð ð ð ð ð ð:9ð 9ð 9ð 9ð 9ð5˜Xð 5°$°s´)ð 5È$ð 5ð 5ð 5ð 5ð 5ð 5ð
+ð +ð +ð1ð 1ð 1ð+ð +ð +ð +ð
Wð 
Wð 
Wð+ð +ð +ðGð Gð Gð+ð +ð +ð`ð `ð `ð&^ð ^ð ^ð4að að að;ð ;ð ;ð+ð +ð +ðDð Dð Dð	Dð 	Dð 	DðYð Yð Yð 
ð 
ð 
ð,Dð Dð Dð&ð &ð &ð ðXð Xñ „\ðXð) x°Ô'9ð )Àð )Ø"&ð)Ø+0°¸3°Ô+?ð)ð )ð )ð )ð*) ¨Ô 2ð )¸#ð )ÐUZÐ[cÐehÐ[hÔUið )ð )ð )ð )ð(
H˜xð 
H°3ð 
H¸hð 
HÐS[ð 
Hð 
Hð 
Hð 
Hðt�xð t hð t¸Tð tÈhð tð tð tð tð^˜(ð ^¨8ð ^¸3ð ^À8ð ^ð ^ð ^ð ^ð,H˜Xð H¸8ð Hð Hð Hð Hð:NÀHÈSÄMð :NÐdlÐmpÔdqð :Nð :Nð :Nð :Nðx=ð =ð =ð>ð >ð >ð >ð >ð >ð >r   rw   r   )$Útypingr   r   r   r   r   r^  Útriton.language.semanticr   Ú r
   rz   Ú_layoutsr   r   r   r   r   r   r   Útriton._C.libtriton.gluon_irr   r   Útriton.compiler.code_generatorr   r   r   rŒ   rú  rQ   r   r+   rc   Ú__triton_builtin__re   rw   r   r   r   ú<module>r     sâ  ðØ ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ð ;Ø €€€Ø 3Ð 3Ð 3Ð 3Ð 3Ð 3Ø Ð Ð Ð Ð Ð ð Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  Mð  MØ PÐ PÐ PÐ PÐ PÐ PÐ PÐ PØ TÐ TÐ TÐ TÐ TÐ TÐ TÐ Tàˆ7�:ÑÔ€ð <Fð !ð !�ð !˜x¨¨C¨Ô0ð !ð !ð !ð !ð
Rð Rð Rð:ð :ð :ð :ðz /3Ð Ô +ð	Mð 	Mð 	Mð 	Mð 	Mñ 	Mô 	Mð 	Mð~>ð ~>ð ~>ð ~>ð ~>�N 8Ô,ñ ~>ô ~>ð ~>ð ~>ð ~>r   