a
    
db[                     @   s   d dl mZmZmZ d dlZd dlZd dlZd dlZd dlZd dl	Z
d dlmZmZmZ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mZ d dlmZ G dd	 d	eZG d
d deeZe  e dkre!  dS )    )print_functionabsolute_importdivisionN)njitset_num_threadsget_num_threadsprangeconfigthreading_layerguvectorize)get_thread_id)TypingError)TestCaseskip_parfors_unsupportedtag)TestInSubprocessc                   @   s  e Zd ZdZdd Zdd Zedd Zee	e
jdk d	d
d Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd Zee	e
jdk d	dd  Zee	e
jdk d	e	ejd! d"d#d$ Zd%d& Zed'd( Zd)S )*TestNumThreadsFc                 C   s   t tj d S Nr   r	   NUMBA_NUM_THREADSself r   e/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/numba/tests/test_num_threads.pysetUp   s    zTestNumThreads.setUpc                 C   sJ   t  dkr | t||k n&t  dv r:tj|| ndsFJ dd S )NZtbb)Zomp	workqueuer   Zunreachable)r
   
assertTruenpalltestingassert_equal)r   expectedresultr   r   r   
check_mask   s
    

zTestNumThreads.check_maskc              	   C   sr   t dd }d}|tf|jtffD ]J\}}| |}|  W d    n1 sP0    Y  | |t|j q"d S )Nc                   S   s   t d d S )NZ
wrong_type)r   r   r   r   r   foo)   s    z5TestNumThreads.test_set_num_threads_type.<locals>.fooz2The number of threads specified must be an integer)r   r   py_func	TypeErrorassertRaisesassertInstr	exception)r   r$   r!   fnZerrtyZraisesr   r   r   test_set_num_threads_type&   s    
$z(TestNumThreads.test_set_num_threads_type   zNot enough CPU coresc                 C   s   t j}| t | td | t d t| | t | | t td W d    n1 sh0    Y  | t t|d  W d    n1 s0    Y  d S )Nr-   r      )r	   r   assertEqualr   r   r'   
ValueError)r   max_threadsr   r   r   _test_set_num_threads_basic3   s    &z*TestNumThreads._test_set_num_threads_basicc                 C   s|   t j}tdd }| | | td | | d t| | | | tdd }| |dd | ||| d S )Nc                   S   s   t  S r   r   r   r   r   r   get_nI   s    z=TestNumThreads._test_set_num_threads_basic_jit.<locals>.get_nr-   c                 S   s   t |  t S r   r   r   nr   r   r   	set_get_nS   s    zATestNumThreads._test_set_num_threads_basic_jit.<locals>.set_get_n)r	   r   r   r/   r   )r   r1   r4   r8   r   r   r   _test_set_num_threads_basic_jitD   s    

z.TestNumThreads._test_set_num_threads_basic_jitc                 C   s  t j}tdgdddddd }tjdtjd	}|| tj|| td
 tjdtjd	}|| tj|d
 t| tjdtjd	}|| tj|| tdgdddddd }tjdtjd	}d
|d< || tj|d
 tjdtjd	}||d< || tj|| d S )Nvoid(int64[:])(n)TparallelZnopythontargetc                 S   s   t  | d d < d S r   r3   xr   r   r   r4   `   s    zETestNumThreads._test_set_num_threads_basic_guvectorize.<locals>.get_n@KL Zdtyper-   c                 S   s   t | d  t | d d < d S Nr   r5   r6   r   r   r   r8   s   s    zITestNumThreads._test_set_num_threads_basic_guvectorize.<locals>.set_get_nr   )	r	   r   r   r   zerosint64r   r    r   )r   r1   r4   r@   r8   r   r   r   '_test_set_num_threads_basic_guvectorize[   s@    

z6TestNumThreads._test_set_num_threads_basic_guvectorizec                 C   sr   t d tdddd }tdgdddd	d
d }| }tj|d tjdtjd}|| tj|d d S )Nr-   Tr<   c                  S   s,   d} t | f}t| D ]}t ||< q|S N   )r   emptyr   r   r@   bufir   r   r   	test_func   s
    zCTestNumThreads._test_set_num_threads_outside_jit.<locals>.test_funcr:   r;   r<   r=   c                 S   s   t  | d d < d S r   r3   r?   r   r   r   test_gufunc   s    zETestNumThreads._test_set_num_threads_outside_jit.<locals>.test_gufuncrA   rC   )r   r   r   r   r   r    rE   rF   )r   rO   rP   outr@   r   r   r   !_test_set_num_threads_outside_jit   s    

z0TestNumThreads._test_set_num_threads_outside_jitc                 C   s0   t dddd }d}||}tj|| d S )NTrH   c                 S   s4   d}t |f}t|  t|D ]}t ||< q |S rI   )r   rK   r   r   r   )nthreadsr@   rM   rN   r   r   r   rO      s    zBTestNumThreads._test_set_num_threads_inside_jit.<locals>.test_funcr-   )r   r   r   r    )r   rO   maskrQ   r   r   r    _test_set_num_threads_inside_jit   s
    
z/TestNumThreads._test_set_num_threads_inside_jitc                 C   sP   t dgdddddd }tjdtjd	}d
}||d< || tj|| d S )Nr:   r;   Tr<   r=   c                 S   s   t | d  t | d d < d S rD   r5   r?   r   r   r   rO      s    zJTestNumThreads._test_set_num_threads_inside_guvectorize.<locals>.test_funcrA   rC   r-   r   )r   r   rE   rF   r   r    )r   rO   r@   rT   r   r   r   (_test_set_num_threads_inside_guvectorize   s    
z7TestNumThreads._test_set_num_threads_inside_guvectorizec                 C   s   t dtdtjd D ]}t| tdddd }| }| ||f| tdgd	dd
ddd }tj	ddtj
dd}tjdtj
d}||| | || | |tt| qd S )Nr-      r.   TrH   c                  S   s<   d} t | f}t| D ]}t ||< qtt |t fS NrB   )r   rK   r   r   lenuniquer   rL   r   r   r   rO      s
    zITestNumThreads._test_get_num_threads_truth_outside_jit.<locals>.test_funcvoid(int64[:], int64[:])(n), (m)r<   r=   c                 S   s   t  | d d < t |d< d S rD   )r   r   r@   rQ   r   r   r   rP      s    zKTestNumThreads._test_get_num_threads_truth_outside_jit.<locals>.test_gufuncrA   rC   d   iP  r.   )rangeminr	   r   r   r   r#   r   r   fullrF   reshaperE   rY   rZ   )r   rT   rO   rQ   rP   r@   r   r   r   '_test_get_num_threads_truth_outside_jit   s"    


z6TestNumThreads._test_get_num_threads_truth_outside_jitc                    s   t dtdtjd D ] tdd fdd}| }|   f| tdgd	dd
d fdd}tjddtj	d
d}tjdtj	d}||| |  | |  tt| qd S )Nr-   rW   r.   TrH   c                     sD   t   d} t| f}t| D ]}t ||< q tt|t fS rX   )r   r   rK   r   r   rY   rZ   r   rL   rT   r   r   rO      s    zHTestNumThreads._test_get_num_threads_truth_inside_jit.<locals>.test_funcr[   r\   r<   r=   c                    s$   t   t | d d < t |d< d S rD   )r   r   r   r]   rg   r   r   rP      s    zJTestNumThreads._test_get_num_threads_truth_inside_jit.<locals>.test_gufuncrA   r^   rC   r_   ra   )rb   rc   r	   r   r   r#   r   r   rd   rF   re   rE   rY   rZ   )r   rO   rQ   rP   r@   r   rg   r   &_test_get_num_threads_truth_inside_jit   s     
z5TestNumThreads._test_get_num_threads_truth_inside_jitc                    s  t  dkr| d tjd tjdtj  tdddd  fd	d
}dD ]}||}|\}}|dd\}}tj|| tj|| dt   }|dkr|	 df}	n
t
|}	tj|	| t f}
tdD ]}||
|d d f< qtj|
| qVd S )Nr   workqueue is not threadsafer.   r-   TrH   c                 S   s*   | j \}}t|D ]}t | ||f< qd S r   )shaper   r   rM   ZfidMNrN   r   r   r   
child_func  s    
z=TestNumThreads._test_nested_parallelism_1.<locals>.child_funcc                    s>   | dkrd fdd	}n| dkr:d fdd	}|S )	Nr   Fc                    s6   t dd fdd}|r*|| S || S d S )NTrH   c                    s`   d}t  f}t|  t D ]4}d|  }t| |k rL|| |t 7 }q"||fS Nr   r.   )r   rE   r   r   r   )rS   accrM   rN   
local_maskrl   rm   rn   rT   r   r   
_test_func"  s    
bTestNumThreads._test_nested_parallelism_1.<locals>.get_test.<locals>.test_func.<locals>._test_func)r   r%   )rS   r%   rs   rr   r   r   rO   !  s
    
zNTestNumThreads._test_nested_parallelism_1.<locals>.get_test.<locals>.test_funcr   c                    s    fdd}t jft jd}t jdft jd}dt   df}dg}d}|sxt||ddd	|}nt||dd
|}|||| ||fS )Nc                    sD   t  t |d  |d  k r.||d  | d  t 7  < d S rD   r5   )rp   rM   rq   )rm   rn   rS   r   r   rs   6  s
    rt   rC   r.   z%void(int64[:], int64[:, :], int64[:])z(p), (n, m), (p)Tr<   r=   Zforceobj)r   rE   rF   arangere   r   )rS   r%   rs   rM   rp   rq   sigZlayoutrr   rS   r   rO   5  s&    )F)Fr   )	test_typerO   rr   r   r   get_test  s
    z;TestNumThreads._test_nested_parallelism_1.<locals>.get_testr   r   )r%   r   )r
   skipTestr	   r   r   r   r   r    rv   re   sumrE   rb   )r   rz   ry   rO   got_accZgot_arrZexp_accZexp_arrZmath_acc_expZmath_accmath_arrrN   r   rr   r   _test_nested_parallelism_1  s.    




/
z)TestNumThreads._test_nested_parallelism_1c           	         s   t  dkr| d tjd dtj d   fdd}tjd i }dD ]>}dD ]4}|d	krl|d	krlqVt ||||||f< qVqN|d
 }| D ]}tj|| qt	 f}t
dtjD ]}|||d d f< qtj|| d S )Nr   ri   r.      c                    s   | dkrt dd}n$| dkr(t dd}n| dkr8dd }|d	d
  |dv r|dkr`t ddn$|dkrtt ddn|dkrdd  fdd}nH|dkrtdgddddn|dkrtdgddd fdd}|S )Nr<   TrH   r   Fnonec                 S   s   | S r   r   r?   r   r   r   	child_decx  s    zNTestNumThreads._test_nested_parallelism_2.<locals>.get_impl.<locals>.child_decc                 S   s2   | j \}}t| t|D ]}t | ||f< qd S r   )rj   r   r   r   rk   r   r   r   child{  s    
zJTestNumThreads._test_nested_parallelism_2.<locals>.get_impl.<locals>.childr<   r   r   c                 S   s   | S r   r   r?   r   r   r   test_dec  s    zMTestNumThreads._test_nested_parallelism_2.<locals>.get_impl.<locals>.test_decc                    sV   t  f}t|  t D ]2}d|  }|tjk r|| t |ksJ q|S )Nr.   )r   rE   r   r   r	   r   r   )rS   rM   rN   rq   )rl   rm   r   rT   r   r   rO     s    

zNTestNumThreads._test_nested_parallelism_2.<locals>.get_impl.<locals>.test_funcr   zint64[:,:], int64[:]z(n, m), (k)r=   guvectorize-objru   c                    sP    fdd}t jft jd}dt   df}||| |S )Nc                    s:   t  |d tjk r6 | |d  t |d ks6J d S rD   )r   r	   r   r   )rM   rq   )r   rS   r   r   rs     s    zbTestNumThreads._test_nested_parallelism_2.<locals>.get_impl.<locals>.test_func.<locals>._test_funcrC   r.   )r   rE   rF   rv   re   )rS   rs   rM   rq   )rl   rm   r   rT   r   rx   r   rO     s    	
r{   )
child_typery   r   rO   rl   rm   rT   )r   r   r   get_implq  s8    
z;TestNumThreads._test_nested_parallelism_2.<locals>.get_impl)r<   r   r   r   r   r   r   )r   r   )r
   r|   r	   r   r   valuesr   r   r    rE   rb   )	r   r   Z
res_arraysry   r   Zpy_arrZarrr   rN   r   r   r   _test_nested_parallelism_2f  s0    


B
z)TestNumThreads._test_nested_parallelism_2   c                    s   t  dkr| d d tdd fddtddfdd	}d
} | }|d }||\}}| || | || fdd}||\}}| || | || d S )Nr   ri   i@B TrH   c                    sD   t  }d}t|  t D ]}|d7 }t ||< q|t |fS ro   )r   rE   r   r   r   rZ   )Zlocal_nttidrp   rN   )BIGr   r   work  s    
z7TestNumThreads._test_nested_parallelism_3.<locals>.workc                    sV   t |  t| }d}t| D ](} | d \}}t|||< ||7 }q|t|fS ro   )r   r   rE   r   rY   rZ   )rS   lenstotalrN   my_acctidsr   r   r   test_func_jit  s    

z@TestNumThreads._test_nested_parallelism_3.<locals>.test_func_jitr-   r.   c                    sj   t dgdddd fdd}tj dftjd	}tj tjd	 df}||| | t|fS )
Nzint64[:], int64[:]z(n), (n)Tr<   r=   c                    s0    d \}}t ||d< | d  |7  < d S )Nr.   r   )rY   )r   r   r   r   )rS   r   r   r   test_func_guvectorize  s    zbTestNumThreads._test_nested_parallelism_3.<locals>.test_guvectorize.<locals>.test_func_guvectorizer.   rC   )r   r   rE   rF   re   r}   rZ   )rS   r   r   r   r   rx   r   test_guvectorize  s    
zCTestNumThreads._test_nested_parallelism_3.<locals>.test_guvectorize)r
   r|   r   r/   r#   )r   r   ZNTZexpected_accZexpected_thread_countr~   Zgot_tcr   r   )r   r   r   _test_nested_parallelism_3  s"    

	
z)TestNumThreads._test_nested_parallelism_3linuxz
Linux onlyc                    s   t d}tdd   fdd}d}|   tj t| |   | | }|j||fd}|	  |
  | | | d S )Nforkc                   S   s   t  S r   r3   r   r   r   r   r$     s    z8TestNumThreads._test_threadmask_across_fork.<locals>.fooc                    s   |     d S r   )put)queuer$   r   r   wrap  s    z9TestNumThreads._test_threadmask_across_fork.<locals>.wrapr.   )r>   args)multiprocessingZget_contextr   r/   r	   r   r   QueueProcessstartjoinget)r   Zforkctxr   rT   Zshared_queuepr   r   r   _test_threadmask_across_fork  s    

z+TestNumThreads._test_threadmask_across_forkc                 C   s   t tj d S r   r   r   r   r   r   tearDown$  s    zTestNumThreads.tearDownc           	      C   s   t  }d}tdddd }tdddd }| |d	 ||\}}| |d	 tj|d	 ||\}}| |d	 tj|d	 d S )
N   FrH   c                 S   s0   t  }t| }t| D ]}t  ||< q||fS r   )r   r   onesr   )sizenjit_par_false_tidresrN   r   r   r   	par_false,  s
    
zBTestNumThreads._test_get_thread_id_not_parallel.<locals>.par_falseTc                 S   s0   t  }t| }t| D ]}t  ||< q||fS r   )r   r   r   rb   )r   njit_par_true_tidr   rN   r   r   r   par_true4  s
    
zATestNumThreads._test_get_thread_id_not_parallel.<locals>.par_truer   )r   r   r/   r   r   r    )	r   Zpython_get_thread_idZcheck_array_sizer   r   r   Znjit_par_false_arrr   Znjit_par_true_arrr   r   r    _test_get_thread_id_not_parallel'  s    

z/TestNumThreads._test_get_thread_id_not_parallelN)__name__
__module____qualname__Z_numba_parallel_test_r   r#   r   r,   unittestZskipIfr	   r   r2   r9   rG   rR   rU   rV   rf   rh   r   r   r   sysplatform
startswithr   r   r   r   r   r   r   r      sZ   

'!$Wg;r   c                   @   sB   e Zd ZeZdZdd dD Zdd Zedd Z	ed	d
 Z
dS )TestNumThreadsBackendsFc                 C   s   g | ]}|t jkr|qS r   )r	   r   ).0rN   r   r   r   
<listcomp>J      z!TestNumThreadsBackends.<listcomp>)r.   r-   r   r      c                 C   s>   t j }t||d< t||d< tjddd|g}| ||S )NZNUMBA_THREADING_LAYERr   z-mznumba.runtestsz-v)osenvironcopyr)   r   
executableZrun_cmd)r   testr
   num_threadsZenv_copyZcmdliner   r   r   run_test_in_separate_processL  s
    
z3TestNumThreadsBackends.run_test_in_separate_processc           	         s^   | j }| jj}d|||f  fdd}d|dd   f }t| |td|| d S )Nz%s.%s.%sc                    st   |   \}}| jr(td||f  | d| | d|v | d|v td|}|rp| |d d S )Nzstdout:
 "%s"
 stderr:
 "%s"OKFAILERRORz\.\.\. skipped '(.*?)'r.   )	r   _DEBUGprintr(   r   researchr|   group)r   oembackendZinjected_methodr   r   r   test_templateY  s    z5TestNumThreadsBackends._inject.<locals>.test_templatez%s_%s_%s_threadsr.   Zlong_running)r   _classr   setattrr   )	clsnamer   backend_guardr   ZthemodZtheclsr   Zinjected_testr   r   r   _injectS  s    zTestNumThreadsBackends._injectc              	   C   sR   | j j D ]@}| j D ]0\}}| jD ] }|ds8q(| |||| q(qqd S )NZ_test_)r   __dict__r   backendsitemsr   r   r   )r   r   r   r   r   r   r   r   generatei  s    

zTestNumThreadsBackends.generateN)r   r   r   r   r   r   r   r   classmethodr   r   r   r   r   r   r   E  s   
r   __main__)"
__future__r   r   r   r   r   r   r   r   numpyr   Znumbar   r   r   r   r	   r
   r   Znumba.np.ufunc.parallelr   Znumba.core.errorsr   Znumba.tests.supportr   r   r   Z!numba.tests.test_parallel_backendr   r   r   r   r   mainr   r   r   r   <module>   s(   $    5.