U
    hâËdb"  ã                   @   sB  d dl Zd dlmZ d dlmZmZmZmZm	Z	m
Z
mZmZmZ d dlmZ d dlmZ d dlmZmZmZ d dlmZmZ d dlZd dl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,G d,d-„ d-eƒZ-e.d.k�r>e /¡  dS )/é    N)Úcompile_ptx)	Úf2Úi1Úi2Úi4Úi8Úu1Úu2Úu4Úu8)Úcuda)Útypes)ÚCUDATestCaseÚskip_on_cudasimÚskip_unless_cc_53)Úfloat16Úfloat32c                 C   s   t | ƒS ©N)Úfloat©Úx© r   ú]/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/numba/cuda/tests/cudapy/test_casting.pyÚnative_cast   s    r   c                 C   s
   t  | ¡S r   )ÚnpÚint8r   r   r   r   Úto_int8   s    r   c                 C   s
   t  | ¡S r   )r   Úint16r   r   r   r   Úto_int16   s    r   c                 C   s
   t  | ¡S r   )r   Úint32r   r   r   r   Úto_int32   s    r    c                 C   s
   t  | ¡S r   )r   Úint64r   r   r   r   Úto_int64   s    r"   c                 C   s
   t  | ¡S r   )r   Úuint8r   r   r   r   Úto_uint8"   s    r$   c                 C   s
   t  | ¡S r   )r   Úuint16r   r   r   r   Ú	to_uint16&   s    r&   c                 C   s
   t  | ¡S r   )r   Úuint32r   r   r   r   Ú	to_uint32*   s    r(   c                 C   s
   t  | ¡S r   )r   Úuint64r   r   r   r   Ú	to_uint64.   s    r*   c                 C   s   t  | ¡t  d¡ S )Ng      à?©r   r   r   r   r   r   Ú
to_float162   s    r,   c                 C   s   t  | ¡t  d¡ S ©Né   )r   r   r   r   r   r   Ú
to_float328   s    r/   c                 C   s   t  | ¡t  d¡ S r-   )r   Úfloat64r   r   r   r   Ú
to_float64<   s    r1   c                 C   s
   t  | ¡S r   )r   Ú	complex64r   r   r   r   Úto_complex64@   s    r3   c                 C   s
   t  | ¡S r   )r   Ú
complex128r   r   r   r   Úto_complex128D   s    r5   c                 C   s   t j t | ¡d¡S r-   ©r   Zfp16Zhmulr   r   r   r   r   r   Úcuda_int_literal_to_float16O   s    r7   c                 C   s   t  | ¡t  d¡ S r-   r+   r   r   r   r   Ú reference_int_literal_to_float16U   s    r8   c                 C   s   t j t | ¡d¡S ©Ng      @r6   r   r   r   r   Úcuda_float_literal_to_float16Y   s    r:   c                 C   s   t  | ¡t  d¡ S r9   r+   r   r   r   r   Ú"reference_float_literal_to_float16_   s    r;   c                   @   s¼   e Zd Zdd„ Zedd„ ƒZedƒdd„ ƒZedd	„ ƒZedƒd
d„ ƒZ	edd„ ƒZ
edd„ ƒZedƒdd„ ƒZedƒdd„ ƒZedd„ ƒZedƒdd„ ƒZedd„ ƒZedƒdd„ ƒZdS )ÚTestCastingc                    s6   t jdd�|ƒ‰t j‡fdd„ƒ‰ ‡ ‡‡fdd„}|S )NT©Zdevicec                    s   ˆ | d ƒ|d< d S )Nr   r   )ÚargÚres)Úwrapped_funcr   r   Úcuda_wrapper_fng   s    z4TestCasting._create_wrapped.<locals>.cuda_wrapper_fnc                    s:   t jdˆd�}| |d< t jdˆd�}ˆ d ||ƒ |d S )Né   )Zdtyper   )rB   rB   )r   Zzeros)r>   ZargarrayZresarray)rA   ÚintypeÚouttyper   r   Ú
wrapper_fnk   s
    z/TestCasting._create_wrapped.<locals>.wrapper_fn)r   Zjit)ÚselfÚpyfuncrC   rD   rE   r   )rA   rC   rD   r@   r   Ú_create_wrappedd   s
    zTestCasting._create_wrappedc                 C   sÊ   t tttf}tjtjtjtjf}tj	tj
tjf}t||ƒD ]Š\}}|D ]|}| j||d��d |  |||¡}|  |dƒ|dƒ¡ |  |dƒtdƒ¡ |  |dƒ|dƒ¡ |  |dƒtdƒ¡ W 5 Q R X qFq:d S )N©ÚfromtyÚtotyçš™™™™™(@çš™™™™™(À)r   r   r    r"   r   r   r   r   r!   r   r   r0   ÚzipÚsubTestrH   ÚassertEqualÚint©rF   ÚpyfuncsÚtotysÚfromtysrG   rK   rJ   Úcfuncr   r   r   Útest_float_to_intt   s    zTestCasting.test_float_to_intz(Compilation unsupported in the simulatorc                 C   sP   t tttf}d}t||ƒD ]0\}}t|tfdd�\}}|  d|› d�|¡ qd S )N©é   é   é    é@   Tr=   z	cvt.rni.sú.f16)r   r   r    r"   rN   r   r   ÚassertIn©rF   rS   ÚsizesrG   ÚsizeÚptxÚ_r   r   r   Útest_float16_to_int_ptxƒ   s
    z#TestCasting.test_float16_to_int_ptxc                 C   s¢   t tttf}tjtjtjtjf}tj	tj
tjf}t||ƒD ]b\}}|D ]T}| j||d��< |  |||¡}|  |dƒ|dƒ¡ |  |dƒtdƒ¡ W 5 Q R X qFq:d S )NrI   rL   )r   r   r    r"   r   r#   r%   r'   r)   r   r   r0   rN   rO   rH   rP   rQ   rR   r   r   r   Útest_float_to_uintŒ   s    zTestCasting.test_float_to_uintc                 C   sP   t tttf}d}t||ƒD ]0\}}t|tfdd�\}}|  d|› d�|¡ qd S )NrX   Tr=   z	cvt.rni.ur]   )r$   r&   r(   r*   rN   r   r   r^   r_   r   r   r   Útest_float16_to_uint_ptx™   s
    z$TestCasting.test_float16_to_uint_ptxc              
   C   sn   t ttf}tjtjtjf}t||ƒD ]D\}}| j|d��* |  	|tj
|¡}|  |dƒ|dƒ¡ W 5 Q R X q$d S )N)rK   éA  )r,   r/   r1   r   r   r   r0   rN   rO   rH   r!   rP   )rF   rS   rT   rG   rK   rV   r   r   r   Útest_int_to_float¢   s    
zTestCasting.test_int_to_floatc              
   C   sf   t tf}ttf}t||ƒD ]F\}}| j|d��, |  |tjtj¡}|  	|dƒ|dƒ¡ W 5 Q R X qd S )N)Úfuncrg   )
r7   r:   r8   r;   rN   rO   rH   r   r   rP   )rF   Z	cudafuncsZ	hostfuncsZcudafuncZhostfuncrV   r   r   r   Útest_literal_to_float16¬   s    ÿÿz#TestCasting.test_literal_to_float16c                 C   sN   t tttf}d}t||ƒD ].\}}tt|fdd�\}}|  d|› �|¡ qd S )NrX   Tr=   zcvt.rn.f16.s)r   r   r   r   rN   r   r,   r^   ©rF   rU   r`   Útyra   rb   rc   r   r   r   Útest_int_to_float16_ptx¸   s
    z#TestCasting.test_int_to_float16_ptxc                 C   sN   t tttf}d}t||ƒD ].\}}tt|fdd�\}}|  d|› �|¡ qd S )NrX   Tr=   zcvt.rn.f16.u)r   r	   r
   r   rN   r   r,   r^   rk   r   r   r   Útest_uint_to_float16_ptxÁ   s
    z$TestCasting.test_uint_to_float16_ptxc              
   C   sª   t ttf}tjtjtjf}t t	||ƒ|¡D ]x\\}}}| j
||d��X |  |||¡}tjj|dƒ|dƒ|dƒ dd� tjj|dƒ|dƒ|dƒ dd� W 5 Q R X q,d S )NrI   rL   r.   ga2U0*©3?)ZrtolrM   )r,   r/   r1   r   r   r   r0   Ú	itertoolsÚproductrN   rO   rH   ÚtestingÚassert_allclose)rF   rS   ZtysrG   rJ   rK   rV   r   r   r   Útest_float_to_floatÊ   s    
 ÿ ÿzTestCasting.test_float_to_floatc                 C   sL   t tf}d}t||ƒD ]0\}}t|tfdd�\}}|  d|› d�|¡ qd S )N)Zf32Zf64Tr=   zcvt.r]   )r/   r1   rN   r   r   r^   )rF   rS   Z	postfixesrG   Zpostfixrb   rc   r   r   r   Útest_float16_to_float_ptxÛ   s
    z%TestCasting.test_float16_to_float_ptxc                 C   s¦   t tf}tjtjf}tjtjtjf}t||ƒD ]r\}}|D ]d}| j	||d��L |  
|||¡}tj |dƒ||dƒƒ¡ tj |dƒ||dƒƒd ¡ W 5 Q R X q:q.d S )NrI   g®Gáz®	@g®Gáz®	Ày                )r3   r5   r   r2   r4   r   r   r0   rN   rO   rH   rq   rr   rR   r   r   r   Útest_float_to_complexä   s    
ÿÿz!TestCasting.test_float_to_complexc                 C   sH   t jttfdd�\}}|  d|¡ t jttfdd�\}}|  d|¡ d S )NTr=   zst.f32zst.u16)r   r   r   r   r^   r   )rF   Zfloat32_ptxrc   Zfloat16_ptxr   r   r   Útest_native_cast÷   s    zTestCasting.test_native_castN)Ú__name__Ú
__module__Ú__qualname__rH   r   rW   r   rd   re   rf   rh   rj   rm   rn   rs   rt   ru   rv   r   r   r   r   r<   c   s2   




	





r<   Ú__main__)0Únumpyr   Z
numba.cudar   Znumba.core.typesr   r   r   r   r   r   r	   r
   r   Znumbar   Z
numba.corer   Znumba.cuda.testingr   r   r   Znumba.typesr   r   ro   Zunittestr   r   r   r    r"   r$   r&   r(   r*   r,   r/   r1   r3   r5   r7   r8   r:   r;   r<   rw   Úmainr   r   r   r   Ú<module>   s<   , 
