o
    i                      @   s&  d dl Z d dlmZ d dlmZmZ ddlmZ dd Zdd	 Z	d
d Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zd d! Zd"d# Zd$d% Zd&d' Zd(d) Zd*d+ Zd,d- Zd.d/ Zd0d1 Zd2d3d4d5d6d7d8d9d:d;d<d=d>d?d@ZdAdB ZG dCdD dDe Z!dEdF Z"dS )G    N)ir)cgutilstargetconfig   )nvvmc                 C   sN   dt | d }tt|tt|t|t|f}t| ||S )NZ___numba_atomic_iZ	_cas_hack)strr   FunctionTypeIntTypePointerTyper   get_or_insert_function)lmodisizefnamefnty r   c/var/www/html/eduruby.in/lip-sync/lip-sync-env/lib/python3.10/site-packages/numba/cuda/nvvmutils.pydeclare_atomic_cas_int   s   r   c                 C   s   |  |||dd}| |dS )N	monotonicr   )ZcmpxchgZextract_value)builderr   r   Zptrcmpvaloutr   r   r   atomic_cmpxchg   s   r   c                 C   s6   d}t t  t t  dt  f}t| ||S )Nz#llvm.nvvm.atomic.load.add.f32.p0f32r   r   r   Z	FloatTyper
   r   r   r   r   r   r   r   r   declare_atomic_add_float32   s
   
r   c                 C   sP   t   }|jdkrd}nd}tt tt t f}t	| ||S )N)   r   z#llvm.nvvm.atomic.load.add.f64.p0f64Z___numba_atomic_double_add)
r   ZConfigStacktopZcompute_capabilityr   r   
DoubleTyper
   r   r   )r   flagsr   r   r   r   r   declare_atomic_add_float64   s   

r    c                 C   4   d}t t  t t  t  f}t| ||S )NZ___numba_atomic_float_subr   r   r   r   r   declare_atomic_sub_float32'   
   
r"   c                 C   r!   )NZ___numba_atomic_double_subr   r   r   r
   r   r   r   r   r   r   declare_atomic_sub_float64.   r#   r%   c                 C   :   d}t t dt t dt df}t| ||S )Nz"llvm.nvvm.atomic.load.inc.32.p0i32    r   r   r	   r
   r   r   r   r   r   r   declare_atomic_inc_int325   
   r)   c                 C   r&   )NZ___numba_atomic_u64_inc@   r(   r   r   r   r   declare_atomic_inc_int64<   r*   r,   c                 C   r&   )Nz"llvm.nvvm.atomic.load.dec.32.p0i32r'   r(   r   r   r   r   declare_atomic_dec_int32C   r*   r-   c                 C   r&   )NZ___numba_atomic_u64_decr+   r(   r   r   r   r   declare_atomic_dec_int64J   r*   r.   c                 C   r!   )NZ___numba_atomic_float_maxr   r   r   r   r   declare_atomic_max_float32Q   r#   r/   c                 C   r!   )NZ___numba_atomic_double_maxr$   r   r   r   r   declare_atomic_max_float64X   r#   r0   c                 C   r!   )NZ___numba_atomic_float_minr   r   r   r   r   declare_atomic_min_float32_   r#   r1   c                 C   r!   )NZ___numba_atomic_double_minr$   r   r   r   r   declare_atomic_min_float64f   r#   r2   c                 C   r!   )NZ___numba_atomic_float_nanmaxr   r   r   r   r   declare_atomic_nanmax_float32m   r#   r3   c                 C   r!   )NZ___numba_atomic_double_nanmaxr$   r   r   r   r   declare_atomic_nanmax_float64t   r#   r4   c                 C   r!   )NZ___numba_atomic_float_nanminr   r   r   r   r   declare_atomic_nanmin_float32{   r#   r5   c                 C   r!   )NZ___numba_atomic_double_nanminr$   r   r   r   r   declare_atomic_nanmin_float64   r#   r6   c                 C   s,   d}t t dt df}t| ||S )NZcudaCGGetIntrinsicHandler+   r'   r   r   r	   r   r   r   r   r   r    declare_cudaCGGetIntrinsicHandle   s
   
r8   c                 C   s4   d}t t dt dt df}t| ||S )NZcudaCGSynchronizer'   r+   r7   r   r   r   r   declare_cudaCGSynchronize   s
   r9   c                 C   s`   | j jj}t|dd }tj||jdtj	d}d|_
d|_||_| |ttddS )	Nzutf-8    _str)nameZ	addrspaceZinternalT   Zgeneric)Zbasic_blockfunctionmoduler   Zmake_bytearrayencodeZadd_global_variabletyper   ZADDRSPACE_CONSTANTlinkageZglobal_constantZinitializerZaddrspacecastr   r
   r	   )r   valuer   cvalglr   r   r   declare_string   s   
rF   c                 C   s8   t t d}t t d||g}t| |d}|S )Nr=   r'   vprintf)r   r
   r	   r   r   r   )r   Z	voidptrtyZ	vprintftyrG   r   r   r   declare_vprint   s   rH   zllvm.nvvm.read.ptx.sreg.tid.xzllvm.nvvm.read.ptx.sreg.tid.yzllvm.nvvm.read.ptx.sreg.tid.zzllvm.nvvm.read.ptx.sreg.ntid.xzllvm.nvvm.read.ptx.sreg.ntid.yzllvm.nvvm.read.ptx.sreg.ntid.zzllvm.nvvm.read.ptx.sreg.ctaid.xzllvm.nvvm.read.ptx.sreg.ctaid.yzllvm.nvvm.read.ptx.sreg.ctaid.zz llvm.nvvm.read.ptx.sreg.nctaid.xz llvm.nvvm.read.ptx.sreg.nctaid.yz llvm.nvvm.read.ptx.sreg.nctaid.zz llvm.nvvm.read.ptx.sreg.warpsizezllvm.nvvm.read.ptx.sreg.laneid)ztid.xztid.yztid.zzntid.xzntid.yzntid.zzctaid.xzctaid.yzctaid.zznctaid.xznctaid.yznctaid.zZwarpsizeZlaneidc                 C   s6   | j }ttdd}t||t| }| |dS )Nr'   r   )r?   r   r   r	   r   r   SREG_MAPPINGcall)r   r<   r?   r   fnr   r   r   	call_sreg   s   rL   c                   @   s<   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd ZdS )SRegBuilderc                 C   s
   || _ d S N)r   )selfr   r   r   r   __init__   s   
zSRegBuilder.__init__c                 C      t | jd| S )Nztid.%srL   r   rO   xyzr   r   r   tid      zSRegBuilder.tidc                 C   rQ   )Nzctaid.%srR   rS   r   r   r   ctaid   rV   zSRegBuilder.ctaidc                 C   rQ   )Nzntid.%srR   rS   r   r   r   ntid   rV   zSRegBuilder.ntidc                 C   rQ   )Nz	nctaid.%srR   rS   r   r   r   nctaid   rV   zSRegBuilder.nctaidc                 C   s:   |  |}| |}| |}| j| j|||}|S rN   )rU   rX   rW   r   addmul)rO   rT   rU   rX   rY   resr   r   r   getdim   s
   


zSRegBuilder.getdimN)	__name__
__module____qualname__rP   rU   rW   rX   rY   r]   r   r   r   r   rM      s    rM   c                    s@   t |   fdddD }tt|d |}|dkr|d S |S )Nc                 3   s    | ]}  |V  qd S rN   )r]   ).0rT   Zsregr   r   	<genexpr>   s    z get_global_id.<locals>.<genexpr>rT   r   r   )rM   list	itertoolsislice)r   dimitseqr   rb   r   get_global_id   s   rj   )#re   Zllvmliter   Z
numba.corer   r   Zcudadrvr   r   r   r   r    r"   r%   r)   r,   r-   r.   r/   r0   r1   r2   r3   r4   r5   r6   r8   r9   rF   rH   rI   rL   objectrM   rj   r   r   r   r   <module>   sX    	