U
    hdW                     @   s   d dl mZ d dlmZmZmZmZmZmZ d dl	m
Z
mZ d dlmZ d dlmZmZmZ edG dd dejZedG d	d
 d
eZedG dd dejZedkre  dS )    sqrt)cudafloat32int16int32uint32void)compile_ptxcompile_ptx_for_current_device)NVVM)skip_on_cudasimunittestCUDATestCasez(Compilation unsupported in the simulatorc                   @   s\   e Z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S )TestCompileToPTXc                 C   sl   dd }t d d  t d d  t d d  f}t||\}}| d| | d| | d| | |t d S )Nc                 S   s.   t d}|t| k r*|| ||  | |< d S )N   )r   gridlen)rxyi r   ^/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/numba/cuda/tests/cudapy/test_compiler.pyf   s    
z.TestCompileToPTX.test_global_kernel.<locals>.ffunc_retval.visible .func.visible .entry)r   r
   assertNotInassertInassertEqualr	   selfr   argsptxrestyr   r   r   test_global_kernel   s    "z#TestCompileToPTX.test_global_kernelc                 C   s   dd }t t f}t||dd\}}| d| | d| | d| | |t  ttt}t||dd\}}| |t ttt}t||dd\}}| |t d}t||dd\}}| |t d S )	Nc                 S   s   | | S Nr   r   r   r   r   r   add   s    z2TestCompileToPTX.test_device_function.<locals>.addTdevicer   r   r   zuint32(uint32, uint32))r   r
   r   r   r    r   r   r   )r"   r)   r#   r$   r%   Z	sig_int32Z	sig_int16Z
sig_stringr   r   r   test_device_function   s     

z%TestCompileToPTX.test_device_functionc                 C   s   dd }t t t t f}t||dd\}}| d| | d| | d| t||ddd\}}| d	| | d
| | d| d S )Nc                 S   s   t | | | | S r'   r   )r   r   zdr   r   r   r   <   s    z)TestCompileToPTX.test_fastmath.<locals>.fTr*   z
fma.rn.f32z
div.rn.f32zsqrt.rn.f32)r+   Zfastmathzfma.rn.ftz.f32zdiv.approx.ftz.f32zsqrt.approx.ftz.f32)r   r
   r   r!   r   r   r   test_fastmath;   s    zTestCompileToPTX.test_fastmathc                 C   s.   t  js| d | |d | |d d S )Nz$debuginfo not generated for NVVM 3.4z\.section\s+\.debug_info\.file.*test_compiler.py")r   	is_nvvm70skipTestassertRegexr"   r$   r   r   r   check_debug_infoO   s    
z!TestCompileToPTX.check_debug_infoc                 C   s*   dd }t |dddd\}}| | d S )Nc                   S   s   d S r'   r   r   r   r   r   r   b   s    z;TestCompileToPTX.test_device_function_with_debug.<locals>.fr   T)r+   debugr
   r5   r"   r   r$   r%   r   r   r   test_device_function_with_debug[   s    z0TestCompileToPTX.test_device_function_with_debugc                 C   s(   dd }t |ddd\}}| | d S )Nc                   S   s   d S r'   r   r   r   r   r   r   j   s    z2TestCompileToPTX.test_kernel_with_debug.<locals>.fr   T)r6   r7   r8   r   r   r   test_kernel_with_debugh   s    z'TestCompileToPTX.test_kernel_with_debugc                 C   s   |  |d d S )Nr0   )r3   r4   r   r   r   check_line_infop   s    z TestCompileToPTX.check_line_infoc                 C   s<   t  js| d dd }t|dddd\}}| | d S )N#lineinfo not generated for NVVM 3.4c                   S   s   d S r'   r   r   r   r   r   r   z   s    z?TestCompileToPTX.test_device_function_with_line_info.<locals>.fr   T)r+   lineinfor   r1   r2   r
   r;   r8   r   r   r   #test_device_function_with_line_infov   s
    
z4TestCompileToPTX.test_device_function_with_line_infoc                 C   s:   t  js| d dd }t|ddd\}}| | d S )Nr<   c                   S   s   d S r'   r   r   r   r   r   r      s    z6TestCompileToPTX.test_kernel_with_line_info.<locals>.fr   T)r=   r>   r8   r   r   r   test_kernel_with_line_info   s
    
z+TestCompileToPTX.test_kernel_with_line_infoc              	   C   sF   dd }|  td( t|td d d td d d f W 5 Q R X d S )Nc                 S   s   | d |d  S )Nr   r   r(   r   r   r   r      s    z5TestCompileToPTX.test_non_void_return_type.<locals>.fzmust have void return typer   )assertRaisesRegex	TypeErrorr
   r   )r"   r   r   r   r   test_non_void_return_type   s    z*TestCompileToPTX.test_non_void_return_typeN)__name__
__module____qualname__r&   r,   r/   r5   r9   r:   r;   r?   r@   rC   r   r   r   r   r   	   s   

r   c                   @   s   e Zd Zdd ZdS ) TestCompileToPTXForCurrentDevicec                 C   s`   dd }t t f}t||dd\}}t j}tjj|}d|d  |d  }| || d S )Nc                 S   s   | | S r'   r   r(   r   r   r   r)      s    zQTestCompileToPTXForCurrentDevice.test_compile_ptx_for_current_device.<locals>.addTr*   z.target sm_r   r   )	r   r   r   Zget_current_deviceZcompute_capabilityZcudadrvZnvvmZfind_closest_archr   )r"   r)   r#   r$   r%   Z	device_cccctargetr   r   r   #test_compile_ptx_for_current_device   s    
zDTestCompileToPTXForCurrentDevice.test_compile_ptx_for_current_deviceN)rD   rE   rF   rJ   r   r   r   r   rG      s   rG   c                   @   s   e Zd ZdZdd ZdS )TestCompileOnlyTestszFor tests where we can only check correctness by examining the compiler
    output rather than observing the effects of execution.c                 C   sb   dd }t |tfdd\}}d}|dD ]}d|kr*|d7 }q*d	}| ||d
| d|  d S )Nc                 S   s   t d t |  d S )N    )r   Z	nanosleep)r   r   r   r   use_nanosleep   s    
z:TestCompileOnlyTests.test_nanosleep.<locals>.use_nanosleep)   r   )rH   r   
znanosleep.u32r      zGot z" nanosleep instructions, expected )r
   r   splitr    )r"   rM   r$   r%   Znanosleep_countlineexpectedr   r   r   test_nanosleep   s    
z#TestCompileOnlyTests.test_nanosleepN)rD   rE   rF   __doc__rT   r   r   r   r   rK      s   rK   __main__N)mathr   Znumbar   r   r   r   r   r	   Z
numba.cudar
   r   Znumba.cuda.cudadrv.nvvmr   Znumba.cuda.testingr   r   r   ZTestCaser   rG   rK   rD   mainr   r   r   r   <module>   s     	