
    ܙd;                     H   d dl Zd dlmc mZ d dlmZ d dl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 d dlZd dlZd dlmZ d dlmZ efd	Z ed
           G d de                      Z ed
           G d de                      Zedk    r ej                     dS dS )    N)
namedtuple)voidint32float32float64)guvectorize)cuda)skip_on_cudasimCUDATestCase)NumbaPerformanceWarning)override_configc           
          t          t          | d d d d f         | d d d d f         | d d d d f                   gdd          d             }|S )Nz(m,n),(n,p)->(m,p)r	   targetc           
          | j         \  }}|j         \  }}t          |          D ]R}t          |          D ]@}d|||f<   t          |          D ]'}|||fxx         | ||f         |||f         z  z  cc<   (ASd S Nr   shaperange)	ABCmnpijks	            Clib/python3.11/site-packages/numba/cuda/tests/cudapy/test_gufunc.py
matmulcorez*_get_matmulcore_gufunc.<locals>.matmulcore   s     w1w1q 	1 	1A1XX 1 1!Q$q 1 1AadGGGqAw1a400GGGG11	1 	1    )r   r   )dtyper    s     r   _get_matmulcore_gufuncr#      s|    $uQQQT{E!!!QQQ$Kqqq!!!t==>%     1 1   1 r!   z&ufunc API unsupported in the simulatorc                       e 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d Zd Zd Zd Zd Zd Zd Zd Zd ZdS )TestCUDAGufuncc                    t                      }d}t          j        |dz  dz  t          j                                      |dd          }t          j        |dz  dz  t          j                                      |dd          } |||          }t          j        ||          }|                     t          j        ||                     d S N      r"      	r#   nparanger   reshapeutmatrix_multiply
assertTrueallcloseselfgufunc	matrix_ctr   r   r   Golds          r   test_gufunc_smallz TestCUDAGufunc.test_gufunc_small"   s    '))	Ii!ma'rz:::BB9aCDF FIi!ma'rz:::BB9aCDF F F1aLL!!Q''At,,-----r!   c                    t                      }d}t          j        |dz  dz  t          j                                      |dd          }t          j        |dz  dz  t          j                                      |dd          }t          j        |          } |||                                          }t          j	        ||          }| 
                    t          j        ||                     d S r'   )r#   r-   r.   r   r/   r	   	to_devicecopy_to_hostr0   r1   r2   r3   )r5   r6   r7   r   r   dBr   r8   s           r   test_gufunc_auto_transferz(TestCUDAGufunc.test_gufunc_auto_transfer0   s    '))	Ii!ma'rz:::BB9aCDF FIi!ma'rz:::BB9aCDF F ^AF1bMM&&((!!Q''At,,-----r!   c                    t                      }d}t          j        |dz  dz  t          j                                      |dd          }t          j        |dz  dz  t          j                                      |dd          } |||          }t          j        ||          }|                     t          j        ||                     d S )N  r(   r)   r*   r+   r,   r4   s          r   test_gufunczTestCUDAGufunc.test_gufunc@   s    '))	Ii!ma'rz:::BB9aCDF FIi!ma'rz:::BB9aCDF F F1aLL!!Q''At,,-----r!   c                    t                      }d}t          j        |dz  dz  t          j                                      dddd          }t          j        |dz  dz  t          j                                      dddd          } |||          }t          j        ||          }|                     t          j        ||                     d S )Nd   r(   r)   r*      r+   r,   r4   s          r   test_gufunc_hidimz TestCUDAGufunc.test_gufunc_hidimN   s    '))	Ii!ma'rz:::BB1b!QOOIi!ma'rz:::BB1b!QOOF1aLL!!Q''At,,-----r!   c                    t          t                    }t          j                            ddd          }t          j                            dd          }t          j        ||          } |||          }t          j                            ||            ||t          j	        |d                    }t          j                            ||           d S )Nr*   
      )rG      rI   )
r#   r   r-   randomrandnr0   r1   testingassert_allclosetile)r5   r6   XYgoldres1res2s          r   test_gufunc_new_axisz#TestCUDAGufunc.test_gufunc_new_axisZ   s    'g666IOOB1%%IOOAq!!!!Q''va||

""4...vaJ//00

""4.....r!   c                    t                      }d}t          j        |dz  dz  t          j                                      |dd          }t          j        |dz  dz  t          j                                      |dd          }t          j                    }t          j        ||          }t          j        ||          }t          j        d|j	        |          } |||||          }|
                    |	          }	|                                 t          j        ||          }
|                     t          j        |	|
                     d S )
Nr@   r(   r)   r*   r+   )r@   r(   r+   )r   r"   stream)outrV   )rV   )r#   r-   r.   r   r/   r	   rV   r;   device_arrayr"   r<   synchronizer0   r1   r2   r3   )r5   r6   r7   r   r   rV   dAr=   dCr   r8   s              r   test_gufunc_streamz!TestCUDAGufunc.test_gufunc_streami   s?   ')) 	Ii!ma'rz:::BB9aCDF FIi!ma'rz:::BB9aCDF F ^Av&&^Av&&\PPPVB6222OO6O**!!Q''At,,-----r!   c                 V   t          t          t          d d          t          d d                    gdd          d             }t          j        dt          j                  dz   }t          j        |          } |||           t          j                            ||           d S )	N(x)->(x)r	   r   c                 J    t          |j                  D ]}| |         ||<   d S Nr   sizer   r   r   s      r   copyz&TestCUDAGufunc.test_copy.<locals>.copy   4     16]]  t! r!   rG   r*   rI   rW   r   r   r   r-   r.   
zeros_likerL   rM   r5   rd   r   r   s       r   	test_copyzTestCUDAGufunc.test_copy   s    	d7111:wqqqz223"
$ 
$ 
$	 	
$ 
$	 Ib
+++a/M!QA

""1a(((((r!   c                 N   t          t          d d          t          d d          fgdd          d             }t          j        dt          j                  dz   }t          j        |          } |||           |                     t          j        ||                     d S )	Nr^   r	   r   c                 J    t          |j                  D ]}| |         ||<   d S r`   ra   rc   s      r   rd   z9TestCUDAGufunc.test_copy_unspecified_return.<locals>.copy   re   r!   rG   r*   rI   rf   )r   r   r-   r.   rh   r2   r3   ri   s       r   test_copy_unspecified_returnz+TestCUDAGufunc.test_copy_unspecified_return   s     
wqqqz7111:./"
$ 
$ 
$	 	
$ 
$	 Ib
+++a/M!QAAq))*****r!   c                 f   t          t          t          d d          t          d d                    gdd          d             }t          j        dt          j                  dz   }t          j        |          } |||           |                     t          j        ||                     d S )	Nr^   r	   r   c                 J    t          |j                  D ]}| |         ||<   d S r`   ra   rc   s      r   rd   z*TestCUDAGufunc.test_copy_odd.<locals>.copy   re   r!      r*   rI   rf   )r   r   r   r-   r.   rh   r2   r3   ri   s       r   test_copy_oddzTestCUDAGufunc.test_copy_odd   s    	d7111:wqqqz223"
$ 
$ 
$	 	
$ 
$	 Ib
+++a/M!QAAq))*****r!   c           	         t          t          t          d d d d f         t          d d d d f                   gdd          d             }t          j        dt          j                                      dd          d	z   }t          j        |          } |||
           |                     t          j        ||                     d S )Nz(x, y)->(x, y)r	   r   c                     t          |j        d                   D ].}t          |j        d                   D ]}| ||f         |||f<   /d S )Nr   rI   )r   r   )r   r   xys       r   copy2dz*TestCUDAGufunc.test_copy2d.<locals>.copy2d   sd     171:&& & &qwqz** & &A1gAadGG&& &r!      r*   r+      rI   rf   )	r   r   r   r-   r.   r/   rh   r2   r3   )r5   rv   r   r   s       r   test_copy2dzTestCUDAGufunc.test_copy2d   s    	d7111aaa4='!!!QQQ$-889%"
$ 
$ 
$	& 	&
$ 
$	&
 Ib
+++33Aq99A=M!qaAq))*****r!   c           	      4   t          dgdd          d             }t          j                            d                              d          }t          j                            d                              d          }t          j        |j        d                                       d          }t          d	d
          5  t          j	        d          5 } ||||           | 
                    |d         j        t                     |                     dt          |d         j                             |                     dt          |d         j                             d d d            n# 1 swxY w Y   d d d            d S # 1 swxY w Y   d S )N(void(float32[:], float32[:], float32[:])(n),(n)->(n)r	   r   c                 l    | j         d         }t          |          D ]}| |         ||         z  ||<   d S r   r   abdistlenr   s        r   numba_dist_cudazMTestCUDAGufunc.test_inefficient_launch_configuration.<locals>.numba_dist_cuda   E     '!*C3ZZ & &A$1+Q& &r!   i   r   r   CUDA_LOW_OCCUPANCY_WARNINGSrI   Trecordz	Grid sizezlow occupancy)r   r-   rJ   randastypezerosr   r   warningscatch_warningsassertEqualcategoryr   assertInstrmessage)r5   r   r   r   r   ws         r   %test_inefficient_launch_configurationz4TestCUDAGufunc.test_inefficient_launch_configuration   s   	@A#F
4 
4 
4	& 	&
4 
4	&
 INN9%%,,Y77INN9%%,,Y77x
##**955:A>> 	B 	B(555 B1d+++  10GHHHk3qt|+<+<===os1Q4</@/@AAA	B B B B B B B B B B B B B B B	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	Bs7   FBE5)F5E9	9F<E9	=FFFc                    t          dgddd          d             }t          j                            d                              d                              d	          }t          j                            d                              d                              d	          }t          j        |          }t          d
d          5  t          j	        d          5 } ||||           | 
                    t          |          d           d d d            n# 1 swxY w Y   d d d            d S # 1 swxY w Y   d S )Nr{   r|   Tr	   )nopythonr   c                 l    | j         d         }t          |          D ]}| |         ||         z  ||<   d S r   r   r~   s        r   numba_dist_cuda2zLTestCUDAGufunc.test_efficient_launch_configuration.<locals>.numba_dist_cuda2   r   r!   i   r   )i   r(   r   rI   r   r   )r   r-   rJ   r   r   r/   rh   r   r   r   r   r   )r5   r   r   r   r   r   s         r   #test_efficient_launch_configurationz2TestCUDAGufunc.test_efficient_launch_configuration   s   	@A#d6
C 
C 
C	& 	&
C 
C	&
 INN:&&--i88GK   	
INN:&&--i88GK   	
}Q:A>> 	, 	,(555 ,  At,,,  Q+++, , , , , , , , , , , , , , ,	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	, 	,s6   D7"1DD7D#	#D7&D#	'D77D;>D;c           
         d } t          t          t          d d          t          d d                    gddd          |           |                     t                    5 } t          t          t          d d          t          d d                    gddd          |           d d d            n# 1 swxY w Y   |                     dt          |j                             d S )Nc                     d S r`    r   r   s     r   fooz.TestCUDAGufunc.test_nopython_flag.<locals>.foo       Dr!   r^   r	   T)r   r   Fznopython flag must be True)r   r   r   assertRaises	TypeErrorr   r   	exception)r5   r   raisess      r   test_nopython_flagz!TestCUDAGufunc.test_nopython_flag   s?   	 	 		#T'!!!*gaaaj112Jv!	# 	# 	##&	( 	( 	( y)) 	<V7Kgaaaj'!!!*556
%7 7 77:< < <	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	< 	5s6;K7L7LMMMMMs   #AB55B9<B9c           
      r   d }|                      t                    5 } t          t          t          d d          t          d d                    gdddd          |           d d d            n# 1 swxY w Y   d}t          |j                  }|                     |d t          |                   |           |t          |          d          	                                
                    d          }d	 |D             }|                     t          d
dg          t          |                     d S )Nc                     d S r`   r   r   s     r   r   z.TestCUDAGufunc.test_invalid_flags.<locals>.foo   r   r!   r^   r	   TF)r   what1ever2z/The following target options are not supported:,c                 8    g | ]}|                     d           S )z'" )strip).0r   s     r   
<listcomp>z5TestCUDAGufunc.test_invalid_flags.<locals>.<listcomp>   s"    000Q000r!   r   r   )r   r   r   r   r   r   r   r   r   r   splitset)r5   r   r   headmsgitemss         r   test_invalid_flagsz!TestCUDAGufunc.test_invalid_flags   s   	 	 	 y)) 	EV@Kgaaaj'!!!*556
%T@ @ @@CE E E	E 	E 	E 	E 	E 	E 	E 	E 	E 	E 	E 	E 	E 	E 	E A&"##Zc$iiZ$///CIIJJ%%''--c2200%000gw/00#e**=====s   AA11A58A5c                    t          t          t          d d          t          d d                    gdd          d             }t          j        dt          j                  x}}|                     t                    5 } ||||           d d d            n# 1 swxY w Y   d}|                     t          |j	                  |           d S )	Nr^   r	   r   c                     d S r`   r   )inprW   s     r   r   z2TestCUDAGufunc.test_duplicated_output.<locals>.foo  s    Dr!   rG   r*   rf   z<cannot specify argument 'out' as both positional and keyword)
r   r   r   r-   r   r   
ValueErrorr   r   r   )r5   r   r   rW   r   r   s         r   test_duplicated_outputz%TestCUDAGufunc.test_duplicated_output  s   	d7111:wqqqz223Z	O	O	O	 	 
P	O	 HRrz2222cz** 	#fCSc""""	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# 	# MV-..44444s   BB #B c                 d   t          t          d d          t          d d          t          d d          fgdd          d             } |||          }t          j        t          j        |          t          j        |          z  d          }t          j                            ||           d S )Nz(n),(n)->()r	   r   c                 z    d}t          t          |                     D ]}|| |         ||         z  z  }||d<   d S r   )r   r   )rt   ru   rsr   s        r   	gu_reducez1TestCUDAGufunc.check_tuple_arg.<locals>.gu_reduce  sI     A3q66]] ! !QqTAaD[ AaDDDr!   rI   )axis)r   r   r-   sumasarrayrL   assert_equal)r5   r   r   r   r   expecteds         r   check_tuple_argzTestCUDAGufunc.check_tuple_arg  s    	wqqqz7111:wqqqz:;]"
$ 
$ 
$	 	
$ 
$	 IaOO6"*Q--"*Q--7a@@@

!,,,,,r!   c                 :    d}d}|                      ||           d S )N)      ?       @      @      @      @      @)      ?      @      @      @      @      @)r   r5   r   r   s      r   test_tuple_of_tuple_argz&TestCUDAGufunc.test_tuple_of_tuple_arg  s,    Q"""""r!   c                     t          dd          } |ddd           |ddd          f} |d	d
d           |ddd          f}|                     ||           d S )NPointrt   ru   zr   r   r   r   r   r   r   r   r   r   r   r   )r   r   )r5   r   r   r   s       r   test_tuple_of_namedtuple_argz+TestCUDAGufunc.test_tuple_of_namedtuple_arg!  s    7O44USC3'''USC3''')USC3'''USC3''')Q"""""r!   c                     t          j        d          t          j        d          f}t          j        d          t          j        d          f}|                     ||           d S )Nr   r   r   r   )r-   r   r   r   s      r   test_tuple_of_array_argz&TestCUDAGufunc.test_tuple_of_array_arg)  s`    Z((Z((*Z((Z((*Q"""""r!   c                 X    t                      }|                     |j        d           d S )Nr    )r#   r   __name__)r5   r6   s     r   test_gufunc_namezTestCUDAGufunc.test_gufunc_name0  s+    ')),77777r!   c           	      r   |                      t                    5 }t          t          t          d d          t          d d                    gdd          d             }d d d            n# 1 swxY w Y   t	          |j                  }|                     d|           |                     d|           d S )Nz(m)->(m)r	   r   c                     d S r`   r   )rt   ru   s     r   fz.TestCUDAGufunc.test_bad_return_type.<locals>.f6  s    r!   z+guvectorized functions cannot return valueszspecifies int32 return type)r   r   r   r   r   r   r   )r5   ter   r   s       r   test_bad_return_typez#TestCUDAGufunc.test_bad_return_type4  s    y)) 	R%aaa%(334jPPP  QP	 	 	 	 	 	 	 	 	 	 	 	 	 	 	
 ",CSIII3S99999s   AA,,A03A0c                    t          t          d d          t          d d          t          d d          fgdd          d             }t          j        d          }|                     t
                    5 } ||           d d d            n# 1 swxY w Y   t          |j                  }|                     d|           |                     d|           |                     d|           |                     t
                    5 } |||||           d d d            n# 1 swxY w Y   t          |j                  }|                     d|           |                     d|           |                     d	|           d S )
Nz(m),(m)->(m)r	   r   c                     d S r`   r   r   s      r   r   z;TestCUDAGufunc.test_incorrect_number_of_pos_args.<locals>.f?  	     Dr!   r+   %gufunc accepts 2 positional argumentszor 3 positional argumentsGot 1 positional argument.zGot 4 positional arguments.	r   r   r-   r.   r   r   r   r   r   r5   r   arrr   r   s        r   !test_incorrect_number_of_pos_argsz0TestCUDAGufunc.test_incorrect_number_of_pos_args>  s   	uQQQxqqq5845#F
4 
4 
4	 	
4 
4	 ill y)) 	RAcFFF	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 ",=sCCC137772C888 y)) 	"RAc3S!!!	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" 	" ",=sCCC137773S99999s$   5BBBD##D'*D'N)r   
__module____qualname__r9   r>   rA   rE   rT   r\   rj   rm   rq   ry   r   r   r   r   r   r   r   r   r   r   r   r   r   r!   r   r%   r%      sd       . . .. . . . . .
. 
. 
./ / /. . .0) ) )+ + ++ + ++ + +"B B B&, , ,&N N N> > >
5 
5 
5- - -# # ## # ## # #8 8 8: : :: : : : :r!   r%   c                   ,    e Zd Zd Zd Zd Zd Zd ZdS )TestMultipleOutputsc           	         t          t          t          d d          t          d d          t          d d                    gdd          d             }t          j        dt          j                  dz   }t          j        |          }t          j        |          } ||||           t          j                            ||           t          j                            ||           d S )N(x)->(x),(x)r	   r   c                 `    t          |j                  D ]}| |         ||<   | |         ||<   d S r`   ra   r   r   r   r   s       r   rd   zKTestMultipleOutputs.test_multiple_outputs_same_type_passed_in.<locals>.copy\  sA     16]]  t!t! r!   rG   r*   rI   rg   )r5   rd   r   r   r   s        r   )test_multiple_outputs_same_type_passed_inz=TestMultipleOutputs.test_multiple_outputs_same_type_passed_in[  s    	d7111:wqqqz7111:>>?#"
$ 
$ 
$	 	
$ 
$	
 Ib
+++a/M!M!Q1

""1a(((

""1a(((((r!   c           	         t          t          t          d d          t          d d          t          d d                    gdd          d             }t          j        dt          j                  dz   }t          j        |          }t          j        |          } ||||           t          j                            ||           t          j                            |dz  |           d S )	Nr   r	   r   c                 f    t          |j                  D ]}| |         ||<   | |         dz  ||<   d S Nr(   ra   r   s       r   copy_and_doublezRTestMultipleOutputs.test_multiple_outputs_distinct_values.<locals>.copy_and_doublem  E     16]]    t!tax!   r!   rG   r*   rI   r(   rg   r5   r   r   r   r   s        r   %test_multiple_outputs_distinct_valuesz9TestMultipleOutputs.test_multiple_outputs_distinct_valuesk  s    	d7111:wqqqz7111:>>?#"
$ 
$ 
$	  	 
$ 
$	 
 Ib
+++a/M!M!1a   

""1a(((

""1q5!,,,,,r!   c           	         t          t          t          d d          t          d d          t          d d                    gdd          d             }t          j        dt          j                  dz   } ||          \  }}t          j                            ||           t          j                            |dz  |           d S )	Nr   r	   r   c                 f    t          |j                  D ]}| |         ||<   | |         dz  ||<   d S r   ra   r   s       r   r   zLTestMultipleOutputs.test_multiple_output_allocation.<locals>.copy_and_double}  r   r!   rG   r*   rI   r(   )r   r   r   r-   r.   rL   rM   r   s        r   test_multiple_output_allocationz3TestMultipleOutputs.test_multiple_output_allocation|  s    	d7111:wqqqz7111:>>?#"
$ 
$ 
$	  	 
$ 
$	 
 Ib
+++a/q!!1

""1a(((

""1q5!,,,,,r!   c           	         t          t          t          d d          t          d d          t          d d                    gdd          d             }t	          j        dt          j                  dz   }t	          j        |          }t	          j        |t          j                  } ||||           t          j                            ||           t          j                            |t	          j        d          z  |           d S )	Nr   r	   r   c                 f    t          |j                  D ]}| |         ||<   | |         dz  ||<   d S )Nr   ra   r   s       r   copy_and_multiplyzJTestMultipleOutputs.test_multiple_output_dtypes.<locals>.copy_and_multiply  sE     16]] " "t!tcz!" "r!   rG   r*   rI   r   )	r   r   r   r   r-   r.   rh   rL   rM   )r5   r  r   r   r   s        r   test_multiple_output_dtypesz/TestMultipleOutputs.test_multiple_output_dtypes  s    	d58U111Xwqqqz::;#"
$ 
$ 
$	" 	"
$ 
$	"
 Ib)))A-M!M!2:...!Q"""

""1a(((

""1rz##6:::::r!   c                 *   t          t          d d          t          d d          t          d d          t          d d          fgdd          d             }t          j        d          }|                     t
                    5 } ||           d d d            n# 1 swxY w Y   t          |j                  }|                     d|           |                     d|           |                     d|           |                     t
                    5 } ||||||           d d d            n# 1 swxY w Y   t          |j                  }|                     d|           |                     d|           |                     d	|           d S )
Nz(m),(m)->(m),(m)r	   r   c                     d S r`   r   )rt   ru   r   r   s       r   r   z@TestMultipleOutputs.test_incorrect_number_of_pos_args.<locals>.f  r   r!   r+   r   zor 4 positional argumentsr   zGot 5 positional arguments.r   r   s        r   r   z5TestMultipleOutputs.test_incorrect_number_of_pos_args  s
   	uQQQxqqq58U111X>?'
8 
8 
8	 	
8 
8	 ill y)) 	RAcFFF	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 ",=sCCC137772C888 y)) 	'RAc3S#&&&	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' ",=sCCC137773S99999s$   BB"BD22D69D6N)r   r   r   r   r   r   r  r   r   r!   r   r   r   Y  s_        ) ) ) - - -"- - -; ; ;": : : : :r!   r   __main__)numpyr-   numpy.core.umath_testscoreumath_testsr0   collectionsr   numbar   r   r   r   r   r	   numba.cuda.testingr
   r   unittestr   numba.core.errorsr   numba.tests.supportr   r#   r%   r   r   mainr   r!   r   <module>r     s       # # # # # # # # # " " " " " " / / / / / / / / / / / /             < < < < < < < <   5 5 5 5 5 5 / / / / / / ")      9::v: v: v: v: v:\ v: v: ;:v:r	 9::Y: Y: Y: Y: Y:, Y: Y: ;:Y:x zHMOOOOO r!   