
    J/Ph,                         d dl mZ d dlZd dlmZmZ d dlZd dlZ G d de          Z	e
dk    r ej                     dS dS )    )cudaN)skip_on_cudasimCUDATestCasec                   >   e Zd Z ej         eej                  dk     d          d             Z e	d          d             Z
 ej         eej                  dk     d          d             Z ej         eej                  dk     d          d             ZdS )	TestMultiGPUContext   zneed more than 1 gpusc                    t          j        d          d             }d }d}t          j        |t          j                  }t          j        |t          j                  }t           j        d         5   |d|f         ||           d d d            n# 1 swxY w Y    |||            |d|f         ||            |||           t           j        d         5  t          j        |t          j                  }t          j        |t          j                  } |d|f         ||           t           j        d         5  t          j        |t          j                  }t          j        |t          j                  }	 |d|f         ||	           d d d            n# 1 swxY w Y   d d d            n# 1 swxY w Y    |||            |||	           t          j        |t          j                  }t          j        |t          j                  } |d|f         ||            |||           d S )Nzvoid(float64[:], float64[:])c                 d    t          j        d          }||j        k     r| |         dz   ||<   d S d S N   r   gridsize)inpoutis      e/var/www/html/test/jupyter/venv/lib/python3.11/site-packages/numba/cuda/tests/cudapy/test_multigpu.pycopy_plus_1z>TestMultiGPUContext.test_multigpu_context.<locals>.copy_plus_1   s6    	!A38||Q!A |    c                 L    t           j                            | dz   |           d S r   )nptestingassert_equal)r   r   s     r   checkz8TestMultiGPUContext.test_multigpu_context.<locals>.check   s$    J##C!GS11111r       dtyper   r   )r   jitr   arangefloat64gpus)
selfr   r   NABA0B0A1B1s
             r   test_multigpu_contextz)TestMultiGPUContext.test_multigpu_context	   s   	0	1	1	$ 	$ 
2	1	$
	2 	2 	2 Iarz***Iarz***Yq\ 	$ 	$K1a###	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	$ 	aAqD!QaYq\ 	* 	*1BJ///B1BJ///BK1b"%%%1 * *Yq
333Yq
333!AqD!"b)))* * * * * * * * * * * * * * *	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	b"b"Iarz***Iarz***AqD!QasJ   5BBBA'F>AF'F>'F+	+F>.F+	/F>>GGz+Simulator does not support multiple threadsc                 n   d t          j        t          j        d                    d}d g|z  fdt	          |          D             }|D ]}|                                 |D ]}|                                 D ].}t          |t                    r|| 	                    |           /d S )Nc                     	 | 5  |                                 }d d d            n# 1 swxY w Y   t          j        |t          j        d          k              ||<   d S # t          $ r}|||<   Y d }~d S d }~ww xY w)N
   )copy_to_hostr   allr   	Exception)gpudAresultsridxarres         r   workz4TestMultiGPUContext.test_multithreaded.<locals>.work4   s    = , ,//++C, , , , , , , , , , , , , , , !#sbimm'; < <	  " " " !"s2   A %A )A )A 
A9)A44A9r-   c           	      b    g | ]+}t          j        t          j        j        |f           ,S ))targetargs)	threadingThreadr   r!   current).0r   r2   r3   r7   s     r   
<listcomp>z:TestMultiGPUContext.test_multithreaded.<locals>.<listcomp>C   sT     - - - #4ty7H797A7G H H H - - -r   )
r   	to_devicer   r   rangestartjoin
isinstanceBaseException
assertTrue)r"   nthreadsthreadsthrr2   r3   r7   s        @@@r   test_multithreadedz&TestMultiGPUContext.test_multithreaded2   s    		= 		= 		= ^BIbMM**&8#- - - - - -!(OO- - -  	 	BHHJJJJ 	 	BGGIIII 	# 	#A!]++ #""""		# 	#r   c                    t           j        d             }t          j        dt          j                  }t           j        d         5  t          j        |          }d d d            n# 1 swxY w Y   t           j        d         5  t          j        |          }d d d            n# 1 swxY w Y   t           j        d         5   |d         |d           d d d            n# 1 swxY w Y   t           j        d         5   |d         |d           d d d            n# 1 swxY w Y   t           j        d         5  t          j                            |	                                |dz              d d d            n# 1 swxY w Y   t           j        d         5  t          j                            |	                                |dz              d d d            d S # 1 swxY w Y   d S )Nc                 h    t          j        d          }|| j        k     r| |xx         |z  cc<   d S d S r   r   )r5   valr   s      r   vector_add_scalarz@TestMultiGPUContext.test_with_context.<locals>.vector_add_scalarU   s:    	!A38||A# |r   r-   r   r   r   )r   r-   r   )
r   r   r   r   float32r!   r@   r   r   r.   )r"   rO   hostarrarr1arr2s        r   test_with_contextz%TestMultiGPUContext.test_with_contextR   sY    
	 	 
	
 )Bbj111Yq\ 	+ 	+>'**D	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ Yq\ 	+ 	+>'**D	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ Yq\ 	. 	.$e$T1---	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. Yq\ 	. 	.$e$T1---	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. 	. Yq\ 	H 	HJ##D$5$5$7$7'A+GGG	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H Yq\ 	H 	HJ##D$5$5$7$7'A+GGG	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	H 	Hsl   A))A-0A-B''B+.B+C##C'*C' DD#&D#<6E>>FF6GG"%G"c                    t           j        d         5  t          j                    }|                    d          s|                     d           d d d            n# 1 swxY w Y   t          j        dt
          j                  }t           j        d         5  t          j        |          }d d d            n# 1 swxY w Y   t           j        d         5  t          j        t          j	        |                    }d d d            n# 1 swxY w Y   t           j        d         5  |
                    |           t
          j                            |                                |           d d d            d S # 1 swxY w Y   d S )Nr   r   z!Peer access between GPUs disabledr-   r   )r   r!   current_contextcan_access_peerskipTestr   r   rP   r@   
zeros_likecopy_to_devicer   r   r.   )r"   ctxrQ   rR   rS   s        r   test_with_context_peer_copyz/TestMultiGPUContext.test_with_context_peer_copyn   st   
 Yq\ 	C 	C&((C&&q)) CABBB	C 	C 	C 	C 	C 	C 	C 	C 	C 	C 	C 	C 	C 	C 	C )Bbj111 Yq\ 	+ 	+>'**D	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ 	+ Yq\ 	: 	:>"-"8"899D	: 	: 	: 	: 	: 	: 	: 	: 	: 	: 	: 	: 	: 	: 	: Yq\ 	B 	B%%% J##D$5$5$7$7AAA	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	B 	BsH   >AA!$A!B;;B?B?'DDD(AE==FFN)__name__
__module____qualname__unittestskipIflenr   r!   r*   r   rK   rT   r\    r   r   r   r      s        X_SS^^a')@AA& & BA&P _BCC# # DC#> X_SS^^a')@AAH H BAH6 X_SS^^a')@AAB B BAB B Br   r   __main__)numbar   numpyr   numba.cuda.testingr   r   r;   r`   r   r]   mainrc   r   r   <module>ri      s              < < < < < < < <     @B @B @B @B @B, @B @B @BF zHMOOOOO r   