Ë
    1²jí  ã            	       ó2  — d dl Z d dlZd dlmZ d dlmZmZmZmZ ddl	m	Z	 ddl
mZmZ d dlmZ  e j                   «       Zeeeee   ee   f   Z G d„ d	«      Z G d
„ de«      Zddœdedeeeef      dee   fd„Zdddœdee   deeeef      dee   fd„Zy)é    N©Úwraps)ÚOptionalÚUnionÚAnyÚSequenceé   )Úflags)Útransform_tensor_metricsÚset_metric_kernels)Úprotonc                   óT   — e Zd ZdZddedeeeef      ddfd„Zd„ Z	d„ Z
d	„ Zd
„ Zd„ Zy)Úscopea  
    A context manager and decorator for entering and exiting a scope.

    Usage:
        context manager:
        ```python
        with proton.scope("test0", {metric_name: metric_value}):
            foo[1,](x, y)
        ```

        decorator:
        ```python
        @proton.scope("test0", {metric_name: metric_value})
        def foo(x, y):
            ...
        ```

    Args:
        name (str): The name of the scope.
        metrics (dict[str, float], optional): The metrics of the scope. Default is None.
    NÚnameÚmetricsÚreturnc                 ó.   — || _         || _        d | _        y ©N©r   r   Úid)Úselfr   r   s      úT/var/www/html/ia-enhanced/venv/lib/python3.12/site-packages/triton/profiler/scope.pyÚ__init__zscope.__init__&   s   € ØˆŒ	ØˆŒØˆ�ó    c                 ó@  — t         j                  sy t        j                  «       | _        t        j
                  | j                  | j                  «       | j                  r>t        «        t        j                  | j                  gt        | j                  «      ¢­Ž  y y r   )r
   Úprofiling_onÚ	libprotonÚrecord_scoper   Úenter_scoper   r   r   Úadd_metricsr   ©r   s    r   Ú_enter_scopezscope._enter_scope+   si   € Ü×!Ò!ØÜ×(Ñ(Ó*ˆŒÜ×Ñ˜dŸg™g t§y¡yÔ1Ø�<Š<ÜÔ Ü×!Ñ! $§'¡'ÐSÔ,DÀTÇ\Á\Ó,RÔSð r   c                 ó’   — t         j                  r| j                  €y t        j                  | j                  | j
                  «       y r   )r
   r   r   r   Ú
exit_scoper   r!   s    r   Ú_exit_scopezscope._exit_scope4   s.   € Ü×!Ò! T§W¡W _ØÜ×Ñ˜TŸW™W d§i¡iÕ0r   c                 ó&   — | j                  «        | S r   )r"   r!   s    r   Ú	__enter__zscope.__enter__9   s   € Ø×ÑÔØˆr   c                 ó$   — | j                  «        y r   )r%   )r   Úexc_typeÚ	exc_valueÚ	tracebacks       r   Ú__exit__zscope.__exit__=   s   € Ø×ÑÕr   c                 ó2   ‡ ‡— t        ‰«      ˆˆ fd„«       }|S )Nc                  ó€   •— ‰j                  «        	  ‰| i |¤Ž‰j                  «        S # ‰j                  «        w xY wr   )r"   r%   )ÚargsÚkwargsÚfuncr   s     €€r   Úwrapperzscope.__call__.<locals>.wrapperB   s:   ø€ à×ÑÔð#Ù˜TÐ, VÑ,à× Ñ Õ"ø�× Ñ Õ"ús   “+ «=r   )r   r1   r2   s   `` r   Ú__call__zscope.__call__@   s    ù€ ä	ˆt‹ô	#ó 
ð	#ð ˆr   r   )Ú__name__Ú
__module__Ú__qualname__Ú__doc__Ústrr   Údictr   r   r"   r%   r'   r,   r3   © r   r   r   r      sH   „ ññ,˜Sð ¨8°D¸¸c¸±NÑ+Cð Ètó ò
Tò1ò
òó
r   r   c                   óV   ‡ — e Zd ZdZd	dedeeeef      ddfˆ fd„Zˆ fd„Z	ˆ fd„Z
ˆ xZS )
Úcpu_timed_scopezÈ
    A scope that measures elapsed time (cpu_time).

    Args:
        name (str): The name of the scope.
        metrics (dict[str, float], optional): Additional metrics to add. Default is None.
    Nr   r   r   c                 óX   •— t         ‰| �  ||«       d | _        |rd|v rt        d«      ‚y y )NÚcpu_timez'The metric name 'cpu_time' is reserved.)Úsuperr   Ú
start_timeÚ
ValueError)r   r   r   Ú	__class__s      €r   r   zcpu_timed_scope.__init__V   s6   ø€ Ü‰Ñ˜˜wÔ'ØˆŒÙ�z WÑ,ÜÐFÓGÐGð -ˆ7r   c                 óv   •— t         j                  sy t        j                  «       | _        t
        ‰| �  «        y r   )r
   r   ÚtimeÚtime_nsr@   r?   r"   )r   rB   s    €r   r"   zcpu_timed_scope._enter_scope\   s'   ø€ Ü×!Ò!ØÜŸ,™,›.ˆŒÜ‰ÑÕr   c                 óâ   •— t         j                  sy | j                  �Ct        j                  «       | j                  z
  }t        j                  | j                  d|i«       t        ‰| �%  «        y )Nzcpu_time (ns)(exc))
r
   r   r@   rD   rE   r   r    r   r?   r%   )r   r>   rB   s     €r   r%   zcpu_timed_scope._exit_scopeb   sR   ø€ Ü×!Ò!ØØ�?‰?Ð&Ü—|‘|“~¨¯©Ñ7ˆHÜ×!Ñ! $§'¡'Ð,@À(Ð+KÔLÜ‰ÑÕr   r   )r4   r5   r6   r7   r8   r   r9   r   r   r"   r%   Ú__classcell__)rB   s   @r   r<   r<   M   sB   ø„ ññH˜Sð H¨8°D¸¸c¸±NÑ+Cð HÈtõ Hô÷ð r   r<   )r   r   r   r   c                óH  — t         j                  sy t        j                  «       }t	        t
        dg «      t
        _        t
        j                  j                  || f«       t        j                  || «       |r)t        «        t        j                  |gt        |«      ¢­Ž  |S )NÚscopes)r
   r   r   r   ÚgetattrÚthread_local_scopesrI   Úappendr   r   r    r   r   s      r   r   r   k   s~   € Ü×ÒØÜ	×	Ñ	Ó	!€BÜ!(Ô)<¸hÈÓ!KÔÔÜ×Ñ×%Ñ% r¨4 jÔ1Ü×Ñ˜"˜dÔ#ÙÜÔÜ×Ñ˜bÐEÔ#;¸GÓ#DÓEØ€Ir   c                ó$  — t         j                  sy t        j                  j	                  «       \  }}| r| |k7  rt        d| › d|› �«      ‚| s|} |r)t        «        t        j                  |gt        |«      ¢­Ž  t        j                  || «       |S )NzScope name mismatch: z != )r
   r   rK   rI   ÚpoprA   r   r   r    r   r$   )r   r   r   Úpopped_names       r   r$   r$   x   sˆ   € ä×ÒØÜ)×0Ñ0×4Ñ4Ó6�O€BˆÙ�˜Ò#ÜÐ0°°°d¸;¸-ÐHÓIÐIÙØˆÙÜÔÜ×Ñ˜bÐEÔ#;¸GÓ#DÓEÜ×Ñ˜˜TÔ"Ø€Ir   r   )Ú	threadingrD   Ú	functoolsr   Útypingr   r   r   r   r
   Úmetricr   r   Útriton._C.libprotonr   r   ÚlocalrK   ÚfloatÚintÚMetricValueTyper   r<   r8   r9   r   r$   r:   r   r   ú<module>rY      sË   ðÛ Û Ý ß 1Ó 1å ß @Ý 3à%�i—o‘oÓ'Ð à˜˜s H¨U¡O°X¸c±]ÐBÑC€÷;ñ ;ô|�eô ð< CGò 
�cð 
 x°°S¸#°X±Ñ'?ð 
È8ÐTWÉ=ó 
ðÐRVò �X˜c‘]ð °xÀÀSÈ#ÀXÁÑ7Oð Ð[cÐdgÑ[hô r   