U
    ,dP                     @   s   d dl mZ d dl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mZmZmZmZmZ d dlmZ d dlZd dlmZmZ d dlZeG d	d
 d
ZedG dd deZedkre  dS )    )List)	dataclassfield)cudafloat32)compile_ptx_for_current_device)cossintanexploglog10log2pow)truedivN)CUDATestCaseskip_on_cudasimc                   @   sx   e Zd ZU eedZee ed< eedZ	ee ed< eedZ
ee ed< eedZee ed< eeedddZd	S )
FastMathCriterion)default_factoryfast_expectedfast_unexpectedprec_expectedprec_unexpected)testfastprecc                    s|   | t fdd| jD  | t fdd| jD  | tfdd| jD  | tfdd| jD  d S )Nc                 3   s   | ]}| kV  qd S N .0ir   r   I/tmp/pip-unpacked-wheel-eu7e0c37/numba/cuda/tests/cudapy/test_fastmath.py	<genexpr>   s     z*FastMathCriterion.check.<locals>.<genexpr>c                 3   s   | ]}| kV  qd S r   r   r   r!   r   r"   r#      s     c                 3   s   | ]}| kV  qd S r   r   r   r   r   r"   r#      s     c                 3   s   | ]}| kV  qd S r   r   r   r$   r   r"   r#      s     )
assertTrueallr   r   r   r   )selfr   r   r   r   )r   r   r"   check   s    zFastMathCriterion.checkN)__name__
__module____qualname__r   listr   r   str__annotations__r   r   r   r   r(   r   r   r   r"   r      s
   
r   z4Fastmath and PTX inspection not available on cudasimc                   @   s   e Zd Zdd ZedddZed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ejdd ZdS )TestFastMathOptionc           
      C   sx   t j||dd|}t j||d|}|| |||| t|||dd\}}t|||d\}	}|| ||	 d S )NT)devicefastmathr0   )r   jitr(   inspect_asmr   )
r'   Zpyfuncsigr0   	criterionfastverprecverZfastptx_Zprecptxr   r   r"   _test_fast_math_common   s&         
  
z)TestFastMathOption._test_fast_math_common)r6   c                    sP    fdd} fdd}| j |td d d tfd|d | j |tfd|d d S )	Nc                    s    || d< d S Nr   r   )rxopr   r"   kernel/   s    z8TestFastMathOption._test_fast_math_unary.<locals>.kernelc                    s    | S r   r   )r=   r>   r   r"   device_function2   s    zATestFastMathOption._test_fast_math_unary.<locals>.device_function   Fr0   r6   Tr:   r   )r'   r?   r6   r@   rA   r   r>   r"   _test_fast_math_unary.   s          z(TestFastMathOption._test_fast_math_unaryc                    sT    fdd} fdd}| j |td d d ttfd|d | j |ttfd|d d S )	Nc                    s    ||| d< d S r;   r   r<   r=   yr>   r   r"   r@   =   s    z9TestFastMathOption._test_fast_math_binary.<locals>.kernelc                    s
    | |S r   r   )r=   rG   r>   r   r"   r0   @   s    z9TestFastMathOption._test_fast_math_binary.<locals>.devicerB   FrC   TrD   )r'   r?   r6   r@   r0   r   r>   r"   _test_fast_math_binary<   s         z)TestFastMathOption._test_fast_math_binaryc                 C   s   |  ttdgdgd d S )Ncos.approx.ftz.f32 r   r   )rE   r   r   r'   r   r   r"   	test_cosfK   s    zTestFastMathOption.test_cosfc                 C   s   |  ttdgdgd d S )Nsin.approx.ftz.f32 rJ   )rE   r	   r   rK   r   r   r"   	test_sinfT   s    zTestFastMathOption.test_sinfc                 C   s    |  ttdddgdgd d S )NrM   rI   div.approx.ftz.f32 rJ   )rE   r
   r   rK   r   r   r"   	test_tanf]   s    zTestFastMathOption.test_tanfc                 C   s   |  ttdgdgd d S )Nzfma.rn.f32 )r   r   )rE   r   r   rK   r   r   r"   	test_expfg   s    zTestFastMathOption.test_expfc                 C   s   |  ttddgdgd d S )Nlg2.approx.ftz.f32 Z
0f3F317218rJ   )rE   r   r   rK   r   r   r"   	test_logfp   s     zTestFastMathOption.test_logfc                 C   s   |  ttddgdgd d S )NrR   Z
0f3E9A209BrJ   )rE   r   r   rK   r   r   r"   test_log10fy   s     zTestFastMathOption.test_log10fc                 C   s   |  ttdgdgd d S NrR   rJ   )rE   r   r   rK   r   r   r"   
test_log2f   s     zTestFastMathOption.test_log2fc                 C   s   |  ttdgdgd d S rU   )rH   r   r   rK   r   r   r"   	test_powf   s     zTestFastMathOption.test_powfc              	   C   s$   |  ttdgdgdgdgd d S )NrO   
div.rn.f32)r   r   r   r   )rH   r   r   rK   r   r   r"   	test_divf   s     zTestFastMathOption.test_divfc              	   C   s   dd }t d d d t t f}tj|ddd|}tj|dd|}d}tj|tj d}| t |d|f |d	d
 W 5 Q R X z|d|f |d	d
 W n tk
r   | d Y nX d S )Nc                 S   s   || | d< d S r;   r   rF   r   r   r"   f10   s    z3TestFastMathOption.test_divf_exception.<locals>.f10rB   T)r1   debug)r[   
   )Zdtypeg      $@g        z5Divide in fastmath should not throw ZeroDivisionError)r   r   r3   npemptyassertRaisesZeroDivisionErrorZfail)r'   rZ   r5   r7   r8   ZnelemZaryr   r   r"   test_divf_exception   s    z&TestFastMathOption.test_divf_exceptionc                    s   t jddddd   fdd}td d d tf}t j|dd	|}t ||}| d
|| | d|| | d
|| | d|| d S )Nzfloat32(float32, float32)Tr2   c                 S   s   | | S r   r   )abr   r   r"   foo   s    z@TestFastMathOption.test_device_fastmath_propagation.<locals>.fooc                    s&   t d}|| jk r" ||| |< d S )NrB   )r   Zgridsize)Zarrvalr    rd   r   r"   bar   s    

z@TestFastMathOption.test_device_fastmath_propagation.<locals>.barrB   )r1   zdiv.approx.f32rX   zdiv.full.f32)r   r3   r   ZassertInr4   ZassertNotIn)r'   rh   r5   r7   r8   r   rg   r"    test_device_fastmath_propagation   s    
z3TestFastMathOption.test_device_fastmath_propagationN)r)   r*   r+   r:   r   rE   rH   rL   rN   rP   rQ   rS   rT   rV   rW   rY   ra   unittestZexpectedFailureri   r   r   r   r"   r/      s   		
			
r/   __main__)typingr   Zdataclassesr   r   Znumbar   r   Znumba.cuda.compilerr   mathr   r	   r
   r   r   r   r   r   operatorr   Znumpyr]   Znumba.cuda.testingr   r   rj   r   r/   r)   mainr   r   r   r"   <module>   s   ( 1