o
    i;                      @   s:  d dl mZmZmZ d dlmZ d dlmZmZ d dl	m
Z
 d dlmZmZm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Zd Zd!Zd"Z d#Z!d$Z"d%Z#d&Z$d'Z%d(Z&d)Z'd*Z(d+Z)d,Z*eed-d.d/d0 Z+eed1d.d2d3 Z,eed-d.d4d5 Z-eed1d.d6d7 Z.eed-d.d8d9 Z/eed-d.d:d; Z0eed1d.d<d= Z1eed1d.d>d? Z2eed-d.d@dA Z3eed1d.dBdC Z4eed-d.dDdE Z5eed1d.dFdG Z6eed-d.dHdI Z7eed1d.dJdK Z8edLG dMdN dNeZ9e:dOkre;  dS dS )P    )cudanjittypes)TypingError)overloadoverload_attribute)typeof)CUDATestCaseskip_on_cudasimunittestNc                   C      d S N r   r   r   b/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/numba/cuda/tests/cudapy/test_overload.pygeneric_func_1      r   c                   C   r   r   r   r   r   r   r   cuda_func_1   r   r   c                   C   r   r   r   r   r   r   r   generic_func_2   r   r   c                   C   r   r   r   r   r   r   r   cuda_func_2   r   r   c                   C   r   r   r   r   r   r   r   generic_calls_generic   r   r   c                   C   r   r   r   r   r   r   r   generic_calls_cuda   r   r   c                   C   r   r   r   r   r   r   r   cuda_calls_generic#   r   r   c                   C   r   r   r   r   r   r   r   cuda_calls_cuda'   r   r   c                   C   r   r   r   r   r   r   r   target_overloaded+   r   r   c                   C   r   r   r   r   r   r   r   generic_calls_target_overloaded/   r   r   c                   C   r   r   r   r   r   r   r   cuda_calls_target_overloaded3   r   r   c                   C   r   r   r   r   r   r   r   )target_overloaded_calls_target_overloaded7   r   r                                    %   )   +   generictargetc                 C      dd }|S )Nc                 S      | d  t 9  < d S Nr   )GENERIC_FUNCTION_1xr   r   r   implT      zol_generic_func_1.<locals>.implr   r3   r4   r   r   r   ol_generic_func_1R      r7   r   c                 C   r.   )Nc                 S   r/   r0   )CUDA_FUNCTION_1r2   r   r   r   r4   [   r5   zol_cuda_func_1.<locals>.implr   r6   r   r   r   ol_cuda_func_1Y   r8   r:   c                 C   r.   )Nc                 S   r/   r0   )GENERIC_FUNCTION_2r2   r   r   r   r4   b   r5   zol_generic_func_2.<locals>.implr   r6   r   r   r   ol_generic_func_2`   r8   r<   c                 C   r.   )Nc                 S   r/   r0   )CUDA_FUNCTION_2r2   r   r   r   r4   i   r5   zol_cuda_func.<locals>.implr   r6   r   r   r   ol_cuda_funcg   r8   r>   c                 C   r.   )Nc                 S      | d  t 9  < t|  d S r0   )GENERIC_CALLS_GENERICr   r2   r   r   r   r4   p      z&ol_generic_calls_generic.<locals>.implr   r6   r   r   r   ol_generic_calls_genericn      rB   c                 C   r.   )Nc                 S   r?   r0   )GENERIC_CALLS_CUDAr   r2   r   r   r   r4   x   rA   z#ol_generic_calls_cuda.<locals>.implr   r6   r   r   r   ol_generic_calls_cudav   rC   rE   c                 C   r.   )Nc                 S   r?   r0   )CUDA_CALLS_GENERICr   r2   r   r   r   r4      rA   z#ol_cuda_calls_generic.<locals>.implr   r6   r   r   r   ol_cuda_calls_generic~   rC   rG   c                 C   r.   )Nc                 S   r?   r0   )CUDA_CALLS_CUDAr   r2   r   r   r   r4      rA   z ol_cuda_calls_cuda.<locals>.implr   r6   r   r   r   ol_cuda_calls_cuda   rC   rI   c                 C   r.   )Nc                 S   r/   r0   )GENERIC_TARGET_OLr2   r   r   r   r4      r5   z*ol_target_overloaded_generic.<locals>.implr   r6   r   r   r   ol_target_overloaded_generic   r8   rK   c                 C   r.   )Nc                 S   r/   r0   )CUDA_TARGET_OLr2   r   r   r   r4      r5   z'ol_target_overloaded_cuda.<locals>.implr   r6   r   r   r   ol_target_overloaded_cuda   r8   rM   c                 C   r.   )Nc                 S   r?   r0   )GENERIC_CALLS_TARGET_OLr   r2   r   r   r   r4      rA   z0ol_generic_calls_target_overloaded.<locals>.implr   r6   r   r   r   "ol_generic_calls_target_overloaded   rC   rO   c                 C   r.   )Nc                 S   r?   r0   )CUDA_CALLS_TARGET_OLr   r2   r   r   r   r4      rA   z-ol_cuda_calls_target_overloaded.<locals>.implr   r6   r   r   r   ol_cuda_calls_target_overloaded   rC   rQ   c                 C   r.   )Nc                 S   r?   r0   )!GENERIC_TARGET_OL_CALLS_TARGET_OLr   r2   r   r   r   r4      rA   z8ol_generic_calls_target_overloaded_generic.<locals>.implr   r6   r   r   r   *ol_generic_calls_target_overloaded_generic   rC   rS   c                 C   r.   )Nc                 S   r?   r0   )CUDA_TARGET_OL_CALLS_TARGET_OLr   r2   r   r   r   r4      rA   z5ol_generic_calls_target_overloaded_cuda.<locals>.implr   r6   r   r   r   'ol_generic_calls_target_overloaded_cuda   rC   rU   z$Overloading not supported in cudasimc                   @   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d Zdd Zdd Zdd Zdd Zdd  Zd!S )"TestOverloadc                 C   s6   t jdt jd}t|d | | |d | d S )N   dtype)rW   rW   r   )nponesint32r   jitassertEqualselfkernelexpectedr3   r   r   r   check_overload   s   zTestOverload.check_overloadc                 C   s0   t jdt jd}t|| | |d | d S )NrW   rX   r   )rZ   r[   r\   r   r^   r_   r   r   r   check_overload_cpu   s   zTestOverload.check_overload_cpuc                 C      dd }t }| || d S )Nc                 S      t |  d S r   )r   r2   r   r   r   ra         z)TestOverload.test_generic.<locals>.kernel)r1   rc   r`   ra   rb   r   r   r   test_generic      zTestOverload.test_genericc                 C   re   )Nc                 S   rf   r   )r   r2   r   r   r   ra      rg   z&TestOverload.test_cuda.<locals>.kernel)r9   rc   rh   r   r   r   	test_cuda   rj   zTestOverload.test_cudac                 C       dd }t t }| || d S )Nc                 S      t |  t|  d S r   )r   r   r2   r   r   r   ra         z2TestOverload.test_generic_and_cuda.<locals>.kernel)r1   r9   rc   rh   r   r   r   test_generic_and_cuda      z"TestOverload.test_generic_and_cudac                 C   rl   )Nc                 S   rm   r   )r   r   r2   r   r   r   ra      rn   z8TestOverload.test_call_two_generic_calls.<locals>.kernel)r1   r;   rc   rh   r   r   r   test_call_two_generic_calls   rp   z(TestOverload.test_call_two_generic_callsc                 C   rl   )Nc                 S   rm   r   )r   r   r2   r   r   r   ra      rn   z5TestOverload.test_call_two_cuda_calls.<locals>.kernel)r9   r=   rc   rh   r   r   r   test_call_two_cuda_calls   rp   z%TestOverload.test_call_two_cuda_callsc                 C   rl   )Nc                 S   rf   r   )r   r2   r   r   r   ra      rg   z7TestOverload.test_generic_calls_generic.<locals>.kernel)r@   r1   rc   rh   r   r   r   test_generic_calls_generic      z'TestOverload.test_generic_calls_genericc                 C   rl   )Nc                 S   rf   r   )r   r2   r   r   r   ra      rg   z4TestOverload.test_generic_calls_cuda.<locals>.kernel)rD   r9   rc   rh   r   r   r   test_generic_calls_cuda   rt   z$TestOverload.test_generic_calls_cudac                 C   rl   )Nc                 S   rf   r   )r   r2   r   r   r   ra      rg   z4TestOverload.test_cuda_calls_generic.<locals>.kernel)rF   r1   rc   rh   r   r   r   test_cuda_calls_generic   rt   z$TestOverload.test_cuda_calls_genericc                 C   rl   )Nc                 S   rf   r   )r   r2   r   r   r   ra     rg   z1TestOverload.test_cuda_calls_cuda.<locals>.kernel)rH   r9   rc   rh   r   r   r   test_cuda_calls_cuda  rt   z!TestOverload.test_cuda_calls_cudac                 C   re   )Nc                 S   rf   r   )r   r2   r   r   r   ra     rg   z8TestOverload.test_call_target_overloaded.<locals>.kernel)rL   rc   rh   r   r   r   test_call_target_overloaded
  rj   z(TestOverload.test_call_target_overloadedc                 C   rl   )Nc                 S   rf   r   )r   r2   r   r   r   ra     rg   zATestOverload.test_generic_calls_target_overloaded.<locals>.kernel)rN   rL   rc   rh   r   r   r   $test_generic_calls_target_overloaded  rt   z1TestOverload.test_generic_calls_target_overloadedc                 C   rl   )Nc                 S   rf   r   )r   r2   r   r   r   ra     rg   z>TestOverload.test_cuda_calls_target_overloaded.<locals>.kernel)rP   rL   rc   rh   r   r   r   !test_cuda_calls_target_overloaded  rt   z.TestOverload.test_cuda_calls_target_overloadedc                 C   s4   dd }t t }| || tt }| || d S )Nc                 S   rf   r   )r   r2   r   r   r   ra      rg   zKTestOverload.test_target_overloaded_calls_target_overloaded.<locals>.kernel)rT   rL   rc   rR   rJ   rd   rh   r   r   r   .test_target_overloaded_calls_target_overloaded  s
   z;TestOverload.test_target_overloaded_calls_target_overloadedc                 C   s   |   \}}t| }t|ddddd }| td tt|dd }W d    n1 s2w   Y  t	t
tjd d d	 |d
d }d S )N	cuda_onlyr   r,   c                 S   r.   )Nc                 S   s   dS )N*   r   )objr   r   r   imp1  r   zTTestOverload.test_overload_attribute_target.<locals>.ov_dummy_cuda_attr.<locals>.impr   )r~   r   r   r   r   ov_dummy_cuda_attr/  rC   zGTestOverload.test_overload_attribute_target.<locals>.ov_dummy_cuda_attrzUnknown attribute 'cuda_only'c                 S   s   | j S r   r|   r2   r   r   r   illegal_target_attr_use:  s   zLTestOverload.test_overload_attribute_target.<locals>.illegal_target_attr_userW   c                 S   s   |j | d< d S r0   r   )resdummyr   r   r   cuda_target_attr_useA  s   zITestOverload.test_overload_attribute_target.<locals>.cuda_target_attr_use)make_dummy_typer   r   assertRaisesRegexr   r   r   int64r   r]   void)r`   MyDummyMyDummyTypemydummy_typer   r   r   r   r   r   test_overload_attribute_target+  s   

	z+TestOverload.test_overload_attribute_targetN)__name__
__module____qualname__rc   rd   ri   rk   ro   rq   rr   rs   ru   rv   rw   rx   ry   rz   r{   r   r   r   r   r   rV      s"    rV   __main__)<numbar   r   r   numba.core.errorsr   numba.core.extendingr   r   numba.core.typing.typeofr   numba.cuda.testingr	   r
   r   numpyrZ   r   r   r   r   r   r   r   r   r   r   r   r   r1   r9   r;   r=   r@   rD   rF   rH   rJ   rL   rN   rP   rR   rT   r7   r:   r<   r>   rB   rE   rG   rI   rK   rM   rO   rQ   rS   rU   rV   r   mainr   r   r   r   <module>   s    	



























 

