
    çg                     T    d dl mZmZmZ  G d de          Z G d de          ZdS )    )ABCMetaabstractmethodabstractclassmethodc                   B    e Zd Zed             Zed             ZddZdS )
DriverBasec                     d S N selfs    R/var/www/html/ai-engine/env/lib/python3.11/site-packages/triton/backends/driver.py	is_activezDriverBase.is_active           c                     d S r	   r
   r   s    r   get_current_targetzDriverBase.get_current_target
   r   r   returnNc                     d S r	   r
   r   s    r   __init__zDriverBase.__init__   s    r   )r   N)__name__
__module____qualname__r   r   r   r   r   r
   r   r   r   r      s]             ^     r   r   )	metaclassc                       e Zd Zd Zd ZdS )	GPUDriverc                     dd l j        j        | _        	 ddlm} || _        n# t          $ r fd| _        Y nw xY wj        j        | _        j        j	        | _
        d S )Nr   )_cuda_getCurrentRawStreamc                 B    j                             |           j        S r	   )cudacurrent_streamcuda_stream)idxtorchs    r   <lambda>z$GPUDriver.__init__.<locals>.<lambda>   s    %*2K2KC2P2P2\ r   )r#   r   get_device_capabilitytorch._Cr   get_current_streamImportErrorcurrent_deviceget_current_device
set_deviceset_current_device)r   r   r#   s     @r   r   zGPUDriver.__init__   s    %*Z%E"	]::::::&?D## 	] 	] 	]&\&\&\&\D###	]"'*";"'*"7s   & ==c                     |S r	   r
   )r   tensormaps_infoargss      r   assemble_tensormap_to_argz#GPUDriver.assemble_tensormap_to_arg!   s    r   N)r   r   r   r   r0   r
   r   r   r   r      s2        
8 
8 
8    r   r   N)abcr   r   r   r   r   r
   r   r   <module>r2      s    < < < < < < < < < <    7        
     r   