
    Χg                         d dl Z d dlZ e j        d          defd            Z e j        d          defd            Z e j        d          d             Z e j        d          d             Zd Zd Z	dS )	    Nreturnc                  R    	 ddl m}  | d uS # t          $ r Y dS t          $ r Y dS w xY w)Nr   
triton_keyF)triton.compiler.compilerr   ImportErrorRuntimeErrorr   s    O/var/www/html/ai-engine/env/lib/python3.11/site-packages/torch/utils/_triton.pyhas_triton_packager      s_    777777%%   uu   uus   	 
&	&&c                  d    ddl m d } d }| |dfd} |            ot                      S )Nr   )get_interface_for_devicec                 F    | j                                         j        dk    S )N   )Workerget_device_propertiesmajordevice_interfaces    r
   cuda_extra_checkz$has_triton.<locals>.cuda_extra_check   s    &<<>>DII    c                     dS )NT r   s    r
   _return_truez has_triton.<locals>._return_true   s    tr   )cudaxpuc                                                       D ]2\  } } |           }|                                r ||          r dS 3dS )NTF)itemsis_available)deviceextra_checkr   r   triton_supported_devicess      r
    is_device_compatible_with_tritonz4has_triton.<locals>.is_device_compatible_with_triton   sg    #;#A#A#C#C 	 	FK77??,,.. ;;?O3P3P ttur   )torch._dynamo.device_interfacer   r   )r   r   r"   r   r!   s      @@r
   
has_tritonr$      s    GGGGGGJ J J   )9NN      ,+--F2D2F2FFr   c                  b    ddl m}  ddlm} |j                                        } | |          S )Nr   )make_backend)driver)r   r&   triton.runtime.driverr'   activeget_current_target)r&   r'   targets      r
   triton_backendr,   (   sI    555555,,,,,,]--//F<r   c                     ddl m}  t                      } |              d|                                 }t	          j        |                    d                                                                                    S )Nr   r   -zutf-8)	r   r   r,   hashhashlibsha256encode	hexdigestupper)r   backendkeys      r
   triton_hash_with_backendr7   1   sv    333333GZ\\
,
,GLLNN
,
,C >#**W--..88::@@BBBr   c                     | j                             d          rd| j         dd          z   }n4| j                             d          rd| j         dd          z   }n| j         }d|z   S )Nfpfloat   bfbfloatztriton.language.)name
startswith)dtypesuffixs     r
   dtype_to_stringrB   <   sp    zT"" 5:abb>)			t	$	$ EJqrrN*&&r   c                  2    dd l } d | j        j        _        d S )Nr   c                      t          |           S )N)rB   )selfs    r
   <lambda>z)patch_triton_dtype_repr.<locals>.<lambda>M   s    /$2G2G r   )tritonlanguager@   __repr__)rG   s    r
   patch_triton_dtype_reprrJ   F   s$    MMM &H%GFO"""r   )
	functoolsr0   	lru_cacheboolr   r$   r,   r7   rB   rJ   r   r   r
   <module>rN      s        TD     TGD G G G G* T      TC C C' ' 'H H H H Hr   