
    wg $                     T   d dl Zd dlmZ d dlmZ d dlmZ d dl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 d dlZ e	e	e	       e
e
e
       eee      gZej                  ej                  ej                  fZd	Zd
Z ed       G d de             Zedk(  r ej2                          yy)    N)
namedtuple)product)	vectorize)cudaint32float32float64)CudaAPIErrordriver)skip_on_cudasim)CUDATestCase)CF)   d   i  z&ufunc API unsupported in the simulatorc                   v    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y)TestCUDAVectorizeiAB c                 v    t        t        d      d        }d}d} |||      }| j                  |||z          y )Nr   targetc                     | |z   S N abs     k/home/mcse/projects/flask/flask-venv/lib/python3.12/site-packages/numba/cuda/tests/cudapy/test_vectorize.py
vector_addz1TestCUDAVectorize.test_scalar.<locals>.vector_add.       q5L    g333333?gffffff@)r   
signaturesassertEqual)selfr   r   r   cs        r   test_scalarzTestCUDAVectorize.test_scalar,   sG    	:f	-	 
.	 q!AE"r    c                 z   t        t        d      d        }t        D ]  }t        j                  t        j
                  j                  | j                        |      }t        j                  ||      } |||      }t        j                  j                  ||       | j                  |j                  |        y )Nr   r   c                     | |z   S r   r   r   s     r   r   z-TestCUDAVectorize.test_1d.<locals>.vector_add9   r   r    dtype)r   r!   dtypesnparrayrandomNaddtestingassert_allcloser"   r)   )r#   r   tydataexpectedactuals         r   test_1dzTestCUDAVectorize.test_1d7   s    	:f	-	 
.	  	/B88BII,,TVV4B?DvvdD)Hd+FJJ&&x8V\\2.	/r    c                    t        t        d      d        }t        j                         }t        D ]  }t        j                  t
        j                  j                  | j                        |      }t        j                  ||      } ||||      }|j                         }t        j                  ||      }t
        j                  j                  ||       | j                  |j                  |        y )Nr   r   c                     | |z   S r   r   r   s     r   r   z3TestCUDAVectorize.test_1d_async.<locals>.vector_addF   r   r    r(   stream)r   r!   r   r:   r*   r+   r,   r-   r.   	to_devicecopy_to_hostr/   r0   r1   r"   r)   )	r#   r   r:   r2   r3   device_datadresultr5   r4   s	            r   test_1d_asynczTestCUDAVectorize.test_1d_asyncD   s    	:f	-	 
.	  
	/B88BII,,TVV4B?D..v6K k&IG))+FvvdD)HJJ&&x8V\\2.
	/r    c                    t        t        d      d        }t        t        dd      t        t
              D ]  \  }}}d|z  }t        j                  j                  |      j                  |      }t        j                  |j                  |      }||z   } |||      }	t        j                  j                  ||	       | j                  |	j                  |        y )Nr   r   c                     | |z   S r   r   r   s     r   r   z-TestCUDAVectorize.test_nd.<locals>.vector_addZ   r   r       r   )   )order)r   r!   r   ranger*   ordersr+   r-   astyper,   Tr0   r1   r"   r)   )
r#   r   ndr)   rD   shaper3   data2r4   r5   s
             r   test_ndzTestCUDAVectorize.test_ndX   s    	:f	-	 
.	 !(aVV D 	2Bu2IE99##E*11%8DHHTVV51Ee|He,FJJ&&x8V\\51	2r    c                    t        t        d      d        }t        j                  dt        j                        }t        j                  dt        j                        }||z   }t        j
                  |      } ||||       t        j                  j                  ||       | j                  |j                  |j                         y )Nr   r   c                     | |z   S r   r   r   s     r   r   z5TestCUDAVectorize.test_output_arg.<locals>.vector_addi   r   r    
   r(   out)
r   r!   r+   aranger   
empty_liker0   r1   r"   r)   )r#   r   ABr4   r5   s         r   test_output_argz!TestCUDAVectorize.test_output_argh   s    	:f	-	 
.	 IIb

+IIb

+q5q!1aV$


""8V46r    c                 n   t        t        d      d        }t        j                  }t        D ]  }t        j
                  ||      }t        j                  j                  |      }|j                  |      }t        j                  j                  ||       | j                  ||j                          y )Nr   r   c                     | |z   S r   r   r   s     r   r   z1TestCUDAVectorize.test_reduce.<locals>.vector_addx   r   r    r(   )r   r!   r+   r   input_sizesrR   r/   reducer0   r1   r"   r)   )r#   r   r)   nxr4   r5   s          r   test_reducezTestCUDAVectorize.test_reducew   s    	:f	-	 
.	  		2A		!5)Avv}}Q'H&&q)FJJ&&x8
 UFLL1		2r    c                    t        t        d      d        }t        j                         }t        j
                  }t        D ]  }t	        j                  ||      }t        j                  j                  |      }t        j                  ||      }|j                  ||      }t        j                  j                  ||       | j                  ||j                          y )Nr   r   c                     | |z   S r   r   r   s     r   r   z7TestCUDAVectorize.test_reduce_async.<locals>.vector_add   r   r    r(   r9   )r   r!   r   r:   r+   r   rY   rR   r/   rZ   r;   r0   r1   r"   r)   )	r#   r   r:   r)   r[   r\   r4   dxr5   s	            r   test_reduce_asyncz#TestCUDAVectorize.test_reduce_async   s    	:f	-	 
.	  	2A		!5)Avv}}Q'H6*B&&r&&9FJJ&&x8UFLL1	2r    c                 n   t        t        d      d        }d}t        j                  |t        j                        }t        j                  |      }||z   } |||      j                         }t        j                  j                  ||       | j                  |j                  |j                         y )Nr   r   c                     | |z   S r   r   r   s     r   r   z:TestCUDAVectorize.test_manual_transfer.<locals>.vector_add   r   r    rO   r(   )r   r!   r+   rR   r   r   r;   r<   r0   assert_equalr"   r)   r#   r   r[   r\   r`   r4   r5   s          r   test_manual_transferz&TestCUDAVectorize.test_manual_transfer   s    	:f	-	 
.	 IIarxx(^^Aq5Ar"//1


&16r    c                    t        t        d      d        }d}t        j                  |t        j                        j                  dd      }t        j                  |      } ||||       ||z   }|j                         }t        j                  j                  ||       | j                  |j                  |j                         y )	Nr   r   c                     | |z   S r   r   r   s     r   r   z:TestCUDAVectorize.test_ufunc_output_2d.<locals>.vector_add   r   r    rO   r(         rP   )r   r!   r+   rR   r   reshaper   r;   r<   r0   rd   r"   r)   re   s          r   test_ufunc_output_2dz&TestCUDAVectorize.test_ufunc_output_2d   s    	:f	-	 
.	 IIarxx(00A6^^A2rr"q5"


&16r    c                     t        t        d      d        } |||      }t        j                  j	                  t        j
                  |      t        j
                  |      z   |       y )Nr   r   c                     | |z   S r   r   r   s     r   r   z5TestCUDAVectorize.check_tuple_arg.<locals>.vector_add   r   r    )r   r!   r+   r0   rd   asarray)r#   r   r   r   rs        r   check_tuple_argz!TestCUDAVectorize.check_tuple_arg   sR    	:f	-	 
.	 q!




1

1 =qAr    c                 0    d}d}| j                  ||       y )N)      ?       @      @)      @      @      @)rq   )r#   r   r   s      r   test_tuple_argz TestCUDAVectorize.test_tuple_arg   s    Q"r    c                 l    t        dd      } |ddd      } |ddd      }| j                  ||       y )	NPointr\   yzrs   rt   ru   rv   rw   rx   r   rq   r#   r{   r   r   s       r   test_namedtuple_argz%TestCUDAVectorize.test_namedtuple_arg   s:    7O4C3#&C3#&Q"r    c                     t        j                  dt         j                        }||dz   f}|dz   |dz   f}| j                  ||       y )NrO   r(   rB   ri   )r+   rR   r   rq   )r#   arrr   r   s       r   test_tuple_of_array_argz)TestCUDAVectorize.test_tuple_of_array_arg   sE    ii"((+#'N1WcAgQ"r    c                     t        dd      } |ddd       |ddd      f} |d	d
d       |ddd      f}| j                  ||       y )Nr{   r|   rs   rt   ru   g      ?g      @g      @rv   rw   rx   g      @g      @g      @r   r   s       r   test_tuple_of_namedtuple_argz.TestCUDAVectorize.test_tuple_of_namedtuple_arg   sV    7O4SC3's)CDSC3's)CDQ"r    c                    t        j                  dt         j                        }|dz   }t        j                  dt         j                        dz  }|dz   }t        dd      } |||      } |||      }| j	                  ||       y )NrO   r(   ri   rB   Points)xsys)r+   rR   r   r   rq   )r#   xs1ys1xs2ys2r   r   r   s           r   test_namedtuple_of_array_argz.TestCUDAVectorize.test_namedtuple_of_array_arg   su    ii"((+Agii"((+a/AgHl3cc"cc"Q"r    c                 b    t        dd      d        }| j                  |j                  d       y )Nzf8(f8)r   r   c                     | dz  S )Nri   r   )r\   s    r   barz2TestCUDAVectorize.test_name_attribute.<locals>.bar   s    6Mr    r   )r   r"   __name__)r#   r   s     r   test_name_attributez%TestCUDAVectorize.test_name_attribute   s1    	8F	+	 
,	 	u-r    c                 F   t         j                  j                  dddd      j                  t         j                        }t        j                  |      }d }t        t        dd       }t        t        dd       }t        t        d|       t        t        d|       | j                  t        d      5  |j                          d d d        | j                  t        d      5  t        j                  dg       d d d        	 t        dgd	
      d        } ||       |t        t        d|       nt        `|t        t        d|       y t        `y # 1 sw Y   xY w# 1 sw Y   gxY w# |t        t        d|       nt        `|t        t        d|       w t        `w xY w)NrB      @   c                      t        dd      )Ni  Transfer not allowed)r
   )argskwargss     r   raising_transferzLTestCUDAVectorize.test_no_transfer_for_device_data.<locals>.raising_transfer   s    s$:;;r    cuMemcpyHtoDcuMemcpyDtoHr   zfloat32(float32)r   r   c                     | dz   S )Nrs   r   )noises    r   funcz@TestCUDAVectorize.test_no_transfer_for_device_data.<locals>.func  s    s{"r    )r+   r-   randnrG   r   r   r;   getattrr   setattrassertRaisesRegexr
   r<   r   r   r   )r#   r   r   old_HtoDold_DtoHr   s         r    test_no_transfer_for_device_dataz2TestCUDAVectorize.test_no_transfer_for_device_data   sr    		1b"-44RZZ@u%	<
 6>486>48(89(89 ##L2HI 	! 	! ##L2HI 	 NNA3	 	( *+F;# <# K
 #9'#9'5	! 	!	  	   #9'#9's$   2E!EE* EE'*6F N)r   
__module____qualname__r.   r%   r6   r?   rL   rV   r]   ra   rf   rl   rq   ry   r   r   r   r   r   r   r   r    r   r   r   %   s_    
 	A	#//(2 72$2$77B#
####.0(r    r   __main__)numpyr+   collectionsr   	itertoolsr   numbar   r   r   r   r	   numba.cuda.cudadrv.driverr
   r   numba.cuda.testingr   r   unittestr!   r*   rF   rY   r   r   mainr   r    r   <module>r      s     "   / / : . +  E5!gw'gw')
 **bjj"((	+ 
 $ 9:q( q( ;q(h zHMMO r    