o
    ۾i                     @   sX   d dl Zd dlmZ d dlmZ d dlmZmZ G dd deZ	e
dkr*e  dS dS )    N)driver)cuda)unittestContextResettingTestCasec                   @   s,   e Zd Zdd Zdd Zdd Zdd Zd	S )
TestHostAllocc                 C   s   d}t  j|dd}ttj}tj||j ||d}d}t	||| | 
t||k || t|}t|||j | 
t||k | 
t||k d S )N    T)mapped)shapedtypebuffer   )r   current_contextmemhostallocnpr
   uint8ndarrayitemsizer   device_memset
assertTrueallfill
empty_likedevice_to_hostsize)selfnmemr
   arymagicrecv r    \/home/ubuntu/.local/lib/python3.10/site-packages/numba/cuda/tests/cudadrv/test_host_alloc.pytest_host_alloc_driver   s   

z$TestHostAlloc.test_host_alloc_driverc                 C   s|   t jdtjd}|d | t|dk t |}t	|dt
| | t|dk || | t|dk d S N
   r
   {   r   )r   pinned_arrayr   uint32r   r   r   	to_devicer   r   device_memory_sizecopy_to_host)r   r   devaryr    r    r!   test_host_alloc_pinned   s   


z$TestHostAlloc.test_host_alloc_pinnedc                 C   sl   t jdtjd}|d | t|dk t|dt	| | t|dk | t
|dkdk d S r#   )r   mapped_arrayr   r(   r   r   r   r   r   r*   sumr   r   r    r    r!   test_host_alloc_mapped(   s   
z$TestHostAlloc.test_host_alloc_mappedc                 C   s  t jdtjdt jdtjdfD ]|}td|d d < | t|d dk | t|d d d dk | t|dk dk | t|dkdk | t|dkd	k | t|dkd
k | t|d dk | t|d dk | t|d dk | t|d dk qd S )Nr$   r%      7      d               i     g       @g     6@)r   r.   r   r(   r'   ranger   r/   r0   r    r    r!   test_host_operators0   s   z!TestHostAlloc.test_host_operatorsN)__name__
__module____qualname__r"   r-   r1   r<   r    r    r    r!   r      s
    
r   __main__)numpyr   numba.cuda.cudadrvr   numbar   numba.cuda.testingr   r   r   r=   mainr    r    r    r!   <module>   s    9