Ë
    ŒÓ^jV  ã                   ó  — d dl mZ d dlmZmZ d dlm	Z	m
Z
 g d¢Ze
d„ «       Ze
d„ «       Ze
d„ «       Ze
d„ «       Ze G d	„ d
«      «       Ze
d„ «       Ze
d„ «       Ze	d„ «       Ze	d„ «       Ze
d„ «       Ze
d„ «       Ze
d„ «       Ze
d„ «       Zy)é    )Ú
_aggregate)Ú_coreÚ	_standard)Úconstexpr_functionÚjit)Úpack2Úunpack2ÚpackÚunpackÚfmaÚFloat2Tensorc                 óX   — t        j                  dd| |gt         j                  dd¬«      S )Nz'
        add.f32x2 $0, $1, $2;
        ú=l,l,lTé   ©ÚdtypeÚis_purer
   ©ÚttglÚinline_asm_elementwiseÚint64©ÚaÚbs     ú…/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/triton/experimental/gluon/language/nvidia/blackwell/float2.pyÚ
_add_f32x2r      ó3   € ä×&Ñ&ð	ð 	Ø	
ˆAˆÜ�j‰jØØô	ð 	ó    c                 óX   — t        j                  dd| |gt         j                  dd¬«      S )Nz'
        sub.f32x2 $0, $1, $2;
        r   Tr   r   r   r   s     r   Ú
_sub_f32x2r       r   r   c                 óX   — t        j                  dd| |gt         j                  dd¬«      S )Nz'
        mul.f32x2 $0, $1, $2;
        r   Tr   r   r   r   s     r   Ú
_mul_f32x2r"   +   r   r   c                 óZ   — t        j                  dd| ||gt         j                  dd¬«      S )Nz.
        fma.rn.f32x2 $0, $1, $2, $3;
        z=l,l,l,lTr   r   r   ©r   r   Úcs      r   Ú
_fma_f32x2r&   9   s5   € ä×&Ñ&ð	ð 	Ø	
ˆAˆqˆ	Ü�j‰jØØô	ð 	r   c                   ó°   — e Zd ZU ej                  ed<   edej                  fd„«       Zed„ «       Z	ed„ «       Z
ed„ «       Zedej                  fd„«       Zy)	r   Úvaluec                 ó   — || _         y ©N)r(   )Úselfr(   s     r   Ú__init__zFloat2Tensor.__init__K   s	   € àˆ�
r   c                 óœ   — t        j                  t        |t        «      d«       t        t	        | j
                  |j
                  «      «      S ©Nzrhs must be a Float2Tensor)r   Ústatic_assertÚ
isinstancer   r   r(   ©r+   Úrhss     r   Ú__add__zFloat2Tensor.__add__O   ó4   € ä×Ñœ: c¬<Ó8Ð:VÔWÜœJ t§z¡z°3·9±9Ó=Ó>Ð>r   c                 óœ   — t        j                  t        |t        «      d«       t        t	        | j
                  |j
                  «      «      S r.   )r   r/   r0   r   r    r(   r1   s     r   Ú__sub__zFloat2Tensor.__sub__T   r4   r   c                 óœ   — t        j                  t        |t        «      d«       t        t	        | j
                  |j
                  «      «      S r.   )r   r/   r0   r   r"   r(   r1   s     r   Ú__mul__zFloat2Tensor.__mul__Y   r4   r   Úaxisc                 ó`   — t        t        j                  | j                  |t        ¬«      «      S )N)r9   Ú
combine_fn)r   r   Úreducer(   r   )r+   r9   s     r   ÚsumzFloat2Tensor.sum^   s   € äœDŸK™K¨¯
©
¸Ì*ÔUÓVÐVr   N)Ú__name__Ú
__module__Ú__qualname__r   ÚtensorÚ__annotations__r   r,   r   r3   r6   r8   Ú	constexprr=   © r   r   r   r   G   s‡   … à�;‰;Óàð˜dŸk™kò ó ðð 	ñ?ó 	ð?ð 	ñ?ó 	ð?ð 	ñ?ó 	ð?ð 	ðW˜Ÿ™ò Wó 	ñWr   r   c                 ón   — t        j                  dd| |gt         j                  dd¬«      }t        |«      S )Nz)
        mov.b64 $0, { $1, $2 };
        z=l,r,rTr   r   )r   r   r   r   )Úx0Úx1r(   s      r   r   r   c   s=   € ä×'Ñ'ð	ð 	Ø	ˆRˆÜ�j‰jØØô	€Eô ˜ÓÐr   c                 óŠ   — t        j                  dd| j                  gt         j                  t         j                  gdd¬«      S )Nz)
        mov.b64 { $0, $1 }, $2;
        z=r,=r,lTr   r   )r   r   r(   Úfloat32)Úxs    r   r	   r	   r   s>   € ä×&Ñ&ð	ð 	Ø	
�‰ˆ	Ü�|‰|œTŸ\™\Ð*ØØô	ð 	r   c                 ó~  — | D �cg c]  }|‘Œ } }| |   dk\  s
J d|› �«       ‚| |xx   dz  cc<   | j                  |dz   d«       t        t        t        | «      «      «      }|t        |«      dz
     ||dz      c||dz   <   |t        |«      dz
  <   t	        j
                  | «      t	        j
                  |«      fS c c}w )Né   z'not enough elements to pack along axis r   )ÚinsertÚlistÚrangeÚlenr   Útuple©Úshaper9   ÚdÚpermutes       r   Ú_get_split_shaperV   €   sÆ   € àÖ�1ŠQÐ€EÐØ�‰;˜!ÒÐMÐFÀtÀfÐMÓMÐØ	ˆ$ƒK�AÑƒKØ	‡L�L�˜‘˜1ÔÜ”5œ˜U›Ó$Ó%€GØ3:¼3¸w»<È!Ñ;KÑ3LÈgÐVZÐ]^ÑV^ÑN_Ð0€GˆD�1‰HÑ�wœs 7›|¨aÑ/Ñ0Ü�:‰:�eÓœdŸj™j¨Ó1Ð1Ð1ùò s   …	B:c                 ó  — | D �cg c]  }|‘Œ } }| |xx   dz  cc<   t        t        t        | «      «      «      }|j                  |dz   t        |«      «       t	        j
                  | «      t	        j
                  |«      fS c c}w )NrL   r   )rN   rO   rP   rM   r   rQ   rR   s       r   Ú_get_join_shaperX   ‹   sn   € àÖ�1ŠQÐ€EÐØ	ˆ$ƒK�1ÑƒKÜ”5œ˜U›Ó$Ó%€GØ‡N�N�4˜!‘8œS ›\Ô*Ü�:‰:�eÓœdŸj™j¨Ó1Ð1Ð1ùò	 s   …	Bc                 ó¬   — t        | j                  |«      }  | j                  |d   Ž j                  |d   Ž j	                  «       \  }}t        ||«      S )Nr   r   )rV   rS   ÚreshaperU   Úsplitr   )rJ   r9   ÚsprF   rG   s        r   r
   r
   ”   sP   € ä)¨!¯'©'°4Ó8€BØ&ˆYˆQ�Y‰Y˜˜1™Ð×&Ñ&¨¨1©Ð.×4Ñ4Ó6�F€BˆÜ��R‹=Ðr   c                 óÎ   — | j                   j                  }t        ||«      }t        | «      \  }}  t	        j
                  ||«      j                  |d   Ž j                  |d   Ž S )Nr   r   )r(   rS   rX   r	   r   ÚjoinrU   rZ   )rJ   r9   rS   r\   rF   rG   s         r   r   r   ›   s\   € àŸG™GŸM™M€EÜ(¨°Ó5€BÜ�Q‹Z�F€BˆØ4Ð$Œ4�9‰9�R˜Ó×$Ñ$ b¨¡eÐ,×4Ñ4°b¸±eÐ<Ð<r   c                 óÜ   — t        j                  |j                  t         j                  k(  d«       t	        j
                  | j                  |t         j                  ¬«      }t        ||«      S )Nzfill_value must be a float32)r   )r   r/   r   rI   ÚstdlibÚ	full_liker(   r   )rJ   Ú
fill_valueÚfills      r   ra   ra   £   sK   € ä×Ñ�z×'Ñ'¬4¯<©<Ñ7Ð9WÔXÜ×Ñ˜AŸG™G Z´t·|±|ÔD€DÜ��tÓÐr   c                 ój   — t        t        | j                  |j                  |j                  «      «      S r*   )r   r&   r(   r$   s      r   r   r   ª   s"   € äœ
 1§7¡7¨A¯G©G°Q·W±WÓ=Ó>Ð>r   N)Útriton.language.corer   Ú	aggregateÚ"triton.experimental.gluon.languager   r   r   r`   Ú"triton.experimental.gluon._runtimer   r   Ú__all__r   r    r"   r&   r   r   r	   rV   rX   r
   r   ra   r   rD   r   r   ú<module>rj      s!  ðÝ 8ß Qß Fò€ð ñ
ó ð
ð ñ
ó ð
ð ñ
ó ð
ð ñ
ó ð
ð ÷Wð Wó ðWð6 ñó ðð ñ
ó ð
ð ñ2ó ð2ð ñ2ó ð2ð ñó ðð ñ=ó ð=ð ñó ðð ñ?ó ñ?r   