U
    ?h                     @   s   d dl mZ d dlZd dlmZmZmZ d dlm	Z	m
Z
mZmZmZmZ ejdd Zejdd Zejd	d
 Zdd ZeedG dd de
Zedkre	  dS )    )print_functionN)configcudaint32)unittestCUDATestCaseskip_on_cudasimskip_unless_cc_60skip_if_cudadevrt_missingskip_if_mvc_enabledc                 C   s   t j  d| d< d S Ng      ?r   )r   cg	this_gridA r   a/var/www/html/venv/lib/python3.8/site-packages/numba/cuda/tests/cudapy/test_cooperative_groups.pyr      s    
r   c                 C   s   t j }|  d| d< d S r   )r   r   r   sync)r   gr   r   r   
sync_group   s    
r   c                 C   s   t d| d< d S N   r   )r   gridr   r   r   r   no_sync   s    r   c                 C   sl   t d}t j }| jd }| jd }td|D ]4}|| d }| |d |f d | ||f< |  q2d S r   )r   r   r   r   shaperanger   )Mcolr   rowscolsrowZoppositer   r   r   sequential_rows   s    



r!   zCG not supported with MVCc                   @   st   e Zd Zedd Zeeddd Zedd Zeeddd	 Zed
dd Z	edd Z
edd ZdS )TestCudaCooperativeGroupsc                 C   s6   t jdt jd}td | | t |d d d S Nr   Z
fill_valuer   r   r   zValue was not set)npfullnanr   assertFalseisnanselfr   r   r   r   test_this_grid4   s    z(TestCudaCooperativeGroups.test_this_gridzFSimulator doesn't differentiate between normal and cooperative kernelsc                 C   s@   t jdt jd}td | tj D ]\}}| |j q&d S Nr   r$   r%   )r&   r'   r(   r   	overloadsitems
assertTruecooperativer,   r   keyoverloadr   r   r   test_this_grid_is_cooperative<   s    z7TestCudaCooperativeGroups.test_this_grid_is_cooperativec                 C   s6   t jdt jd}td | | t |d d d S r#   )r&   r'   r(   r   r)   r*   r+   r   r   r   test_sync_groupG   s    z)TestCudaCooperativeGroups.test_sync_groupc                 C   s@   t jdt jd}td | tj D ]\}}| |j q&d S r.   )r&   r'   r(   r   r/   r0   r1   r2   r3   r   r   r   test_sync_group_is_cooperativeO   s    z8TestCudaCooperativeGroups.test_sync_group_is_cooperativez$Simulator does not implement linkingc                 C   sZ   t jdt jd}td | tj D ].\}}| |j |jj	D ]}| 
d| qBq&dS )z
        We should only mark a kernel as cooperative and link cudadevrt if the
        kernel uses grid sync. Here we ensure that one that doesn't use grid
        synsync isn't marked as such.
        r   r$   r%   Z	cudadevrtN)r&   r'   r(   r   r/   r0   r)   r2   Z_codelibraryZ_linking_filesZassertNotIn)r,   r   r4   r5   linkr   r   r   ,test_false_cooperative_doesnt_link_cudadevrtY   s    zFTestCudaCooperativeGroups.test_false_cooperative_doesnt_link_cudadevrtc           
      C   s   t jrd}nd}tj|tjd}d}|jd | }td d d d df f}t|t}|j	| }|
|}||kr~td |||f | tt|d |d dfj}	tj||	 d S )N)    r;   )   r<   )Zdtyper;   r   z1GPU cannot support enough cooperative grid blocksr   )r   ZENABLE_CUDASIMr&   Zzerosr   r   r   jitr!   r/   max_cooperative_grid_blocksr   skipZtileZarangeTZtestingZassert_equal)
r,   r   r   ZblockdimZgriddimsigc_sequential_rowsr5   mb	referencer   r   r   test_sync_at_matrix_rowh   s    


 z1TestCudaCooperativeGroups.test_sync_at_matrix_rowc                 C   sj   t d d d d df f}t|t}|j| }|d}|d}|d}| || | || d S )Nr      )   rG   )rG      rH   )r   r   r=   r!   r/   r>   ZassertEqual)r,   rA   rB   r5   Zblocks1dZblocks2dZblocks3dr   r   r    test_max_cooperative_grid_blocks   s    



z:TestCudaCooperativeGroups.test_max_cooperative_grid_blocksN)__name__
__module____qualname__r	   r-   r   r6   r7   r8   r:   rE   rI   r   r   r   r   r"   1   s    



r"   __main__)
__future__r   numpyr&   Znumbar   r   r   Znumba.cuda.testingr   r   r   r	   r
   r   r=   r   r   r   r!   r"   rJ   mainr   r   r   r   <module>   s    


_