
    ܙdDE                    J   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 d dlZd dl	Z
d dlZd dlmZ d dlmZmZmZmZmZmZ d dlmZmZmZmZmZmZmZmZmZmZmZ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+ d d	l,m-Z-m.Z.m/Z/m0Z0 d d
l1m2Z2  G d de)j3                  Z4d Z5d Z6 edg dd          Z7d Z8 G d de0          Z9d Z:d Z;d Z<d Z=d Z>d Z?d Z@d ZAd ZBd ZCd  ZDd! ZEd" ZFd# ZGd$ ZHd% ZId& ZJd' ZKd( ZLd) ZMd* ZNd+ ZOd, ZPd- ZQd. ZRd/ ZSd0 ZTd1 ZUd2 ZVd3 ZWd4 ZXd5 ZYdS )6    N)make_dataclass)parfor)typesirconfigcompilersigutilscgutils)add_offset_to_labelsreplace_var_namesremove_delslegalize_namesrename_labelsget_name_var_tablevisit_vars_innerget_definitionguardget_call_tableis_pureget_np_ufunc_typget_unused_var_nameis_const_callfixup_var_define_in_scopetransfer_scopefind_max_labelget_global_func_typ)	signature)lowering)ensure_parallel_support)NumbaParallelSafetyWarningNotDefinedErrorCompilerErrorInternalError)ParforLoweringBuilderc                   "     e Zd ZdZ fdZ xZS )ParforLowerzlThis is a custom lowering class that extends standard lowering so as
    to accommodate parfor.Parfor nodes.c                     t          |t          j                  rt          | |           d S t	                                          |           d S N)
isinstancer   Parfor_lower_parfor_parallelsuper
lower_instselfinst	__class__s     =lib/python3.11/site-packages/numba/parfors/parfor_lowering.pyr-   zParforLower.lower_inst1   sK    dFM** 	%"4.....GGt$$$$$    )__name__
__module____qualname____doc__r-   __classcell__r1   s   @r2   r&   r&   ,   sB        + +% % % % % % % % %r3   r&   c                 R    |j         t          | |          S  |j         | |          S r(   )lowerer_lower_parfor_parallel_std)r;   r   s     r2   r+   r+   8   s.    ~)'6:::v~gv...r3   c                    ddl m} t                       | j        j        }| j        }| j        }| j        j        }t          j        |          | j        _        t          j
        r(t          d| j        t          | j                             | j        j        }| j        }t          j
        rt          d            |j                     |j        j        }	|j        j        }
t          j
        r)t          d|j        dt          |j                             |j        j        D ]3}t          j
        rt          d|           |                     |           4|j        D ]?}||vr9||         }t+          j        |
||	          }|                     |j        |           @i }i }t2          j        j                            ||j        || j        ||           t          j
        r t          d|           t          d	|           |j        J t2          j        j                            ||j                  }|j         |j!        }}t          j
        r t          d|           t          d|           tE          |          }i }g }|dk    rC|j        j        }
|j        j        }	tG          | |
|	          }|$                    t2          j%        j&        j'        j(        tS          t2          j%        j&        j'        j(                  d          }|*                    |+                    |g           tX          j-        d          }t]          |          D ]z}||         }| j        j        |         }t+          j        |
||	          }t_          |          }|j0        }t          j
        r5t          d||||tY          j1        |          |t          |                     te          |tX          j3        j4                  r|j5        dz   } nd} |$                    tJ          j6        to          tJ          j6                  tY          j8        tX          j-        |           fdtY          j1        |          i          }!|g}"te          |tX          j3        j4                  r|*                    t*          j9        :                    |d|	          tY          j8        tX          j-        |j5                  d          }#t]          |j5                  D ]Y}$|*                    t*          j9        ;                    |#|$d
|	          tX          j-        d          }%|"<                    |%           Z|=                    |"d          }&|j>        ?                    |          }'|@                    |'tY          j1        |                    }(|+                    |!|&|(g          })|*                    |)|d          }*|*||j        <   |<                    |*           ||         jA        }+|+
te          |tX          j3        j4                  r|$                    tJ          jB        to          tJ          jB                  tY          j8        tX          j-        |j5                  |fdtY          j1        |          i          },|@                    |+|d          }-|+                    |,|#|-|(g          }.|*                    |.|d           }/|<                    |/           n|@                    |+|d           }/n|}/t          jC        rd!t          |          z   d"z   }0tY          jE        |0          }1|@                    |0|1d#          }2t+          jF        |2|gd
|	$          }3t          tX          jH        ||2j                 ||j                           | j        jI        |3<   t          d%|3           |                     |3           ||j                 }4|J                    |4          }5t          jL        ||5          }6|
M                    d&|	          }7|4||7j        <   |6| j        |7j        <   t          jN        || O                    |j                  |5'          5 }8|P                    |8jQ        |6           |R                    |*|7|/(           d
d
d
           n# 1 swxY w Y   ||jS                                        }9d)|9_T        ||jU        d         jV        j                 }:|jU        dd
         D ]};||;jV        j                 |:k    sJ d*t2          j        j        _W        	 t          | |||||9i t          |          |:|j        
  
        \  }<}=}>}?}@d+t2          j        j        _W        n# d+t2          j        j        _W        w xY wd,g|=z   }=tE          |          }AtE          |=          tE          |          z
  |Az
  }Bt          j
        rPt          d-|=           t          d.|B           t          d/|           t          d0|           t          d1|A           t          |j[        |B|A|=|>|j        |          }Ct          j
        rt          d2|C           d3 |jU        D             }Dt          j
        r%t          d4|jU                   t          d5|D           t          | |<|C|>|=|?|D||||j        |:|j        |@           |dk    rt          ||| ||           |D ]0}E|                     t+          j^        |Ej        |	6                     1|| j        _        t          j
        rt          d7           d
S d
S )8a  Lowerer that handles LLVM code generation for parfor.
    This function lowers a parfor IR node to LLVM.
    The general approach is as follows:
    1) The code from the parfor's init block is lowered normally
       in the context of the current function.
    2) The body of the parfor is transformed into a gufunc function.
    3) Code is inserted into the main function that calls do_scheduling
       to divide the iteration space for each thread, allocates
       reduction arrays, calls the gufunc function, and then invokes
       the reduction function across the reduction arrays to produce
       the final reduction values.
    r   )get_thread_countzlowerer.fndescr+   zinit_block =  zlower init_block instr = 	alias_maparg_aliasesNzparfor_redvars:zparfor_reddict:)r;   scopeloc )fobjftypeargs)rG   num_threads_var)rhstypnamereduction_info   dtype)rE   rF   rG   kwsshaperedarr_shaperedshapeonedimtuple_size_varrK   )cvalrJ   redarrinit_val)rU   rJ   rK   redtosetzres_print1 for redvar :	str_constrG   varargrC   res_print_redvarz$loop_index)intp)objindexvalnumpyTFschedzfunc_args = znum_inputs = parfor_outputs = parfor_redvars = znum_reductions = zgu_signature = c                 6    g | ]}|j         |j        |j        fS rD   )startstopstep.0ls     r2   
<listcomp>z._lower_parfor_parallel_std.<locals>.<listcomp>p  s%    HHHAGQVQV,HHHr3   zloop_nests = loop_ranges = )rC   z_lower_parfor_parallel done)_numba.np.ufunc.parallelr>   r   contexttyping_contextbuilderfndesctypemapcopyr   DEBUG_ARRAY_OPTprinttypevarmapdump
init_blockrC   rB   bodyr-   racesr   Var_alloca_varrK   numbaparforsr   find_potential_aliases_parforparamsfunc_irget_parfor_outputsredvarsreddictlenr$   bind_global_functionnpufuncparallel_iget_num_threadsr   assigncallr   r^   rangeredtyp_to_redarraytyperN   DTyper)   npytypesArrayndimemptyr   UniTupleExprgetattrstatic_getitemappendmake_tuple_variable
_typingctxresolve_value_typemake_const_variablerW   fullDEBUG_ARRAY_OPT_RUNTIMEstrStringLiteralPrintr   none	calltypesget_value_typer
   alloca_onceredefine	for_rangeloadvarstorer`   setitemflagserror_model
loop_nestsindex_variablesequential_parfor_lowering_create_gufunc_for_parfor_bodybool_create_shape_signatureget_shape_classescall_parallel_gufunc#_parfor_lowering_finalize_reductionDel)Fr;   r   r>   	typingctx	targetctxrr   orig_typemaprt   ry   rC   rB   instrracevarrvtyprvr@   rA   parfor_output_arraysparfor_redvarsparfor_reddictnredvarsredarrs
to_cleanuppfbdrget_num_threadsrH   ired_name
redvar_typredvarredarrvar_typreddtype	redarrdimglbl_np_emptysize_var_listredshape_varj	onedimvarsize_varrU   dt
empty_call
redarr_varrW   full_func_nodeinit_val_var	full_callrX   res_print_strstrconsttyplhs	res_printnum_thread_typentllvm_typealloc_loop_varnumba_ir_loop_index_varloopr   index_var_typrl   func	func_argsfunc_sigfunc_arg_typesexp_name_to_tuple_varnum_reductions
num_inputsgu_signatureloop_rangesvsF                                                                         r2   r<   r<   ?   s    988888.IIoG >)L!Y|44GN FW^0D0DEEEn$G^F &'''


C#E  Pov0#tF<M7N7NOOO"' " "! 	6-u5555!!!!< 0 0&  G$Ew,,B///IK	M66vv}g(/KQ Q Q *k9%%%m[))) =$$$ =/BB %+^V^NN 1000000 >""HGJ!||!'#%gULLL  44(:%ehn&=&OPP 5 
 
  ,,

?
44
" ' $ $ x h	[ h	[A%a(H /9JVE8S11F 3:>>M$*H% 
$!K))#))	 	 	 *en&:;; &Oa/			
 "66X&rx00N5:y99 ek(334 7  M --M *en&:;; 4$||==uz:?CC'  ,     z// 4 4A %G22<D#NN!J- !- ! !I
 "((3333 00$4 1  H
 #66x@@D**%+h:O:O*PPBM2GGJM &  J
 $.GFK j)))%h/8H#j%.*>?? % &+%?%?W.rw77!N5:zGG$ %ek(&;&;< &@ & &N $)#<#<%$' $= $ $L !&

&lL"-M !+ ! !I  %||%&'  ,    H %%h////$88%$'  9    HH "1 2$<s6{{$JS$PM"'"5m"D"DK33*'( 4  C !#sFm04#!? !? !?I:CEJ=DSX=N=DV[=Q;S ;SGN,Y7 ,i888&&y111 &o&:;O#22?CCK$0+FFN ',nn]C&H&H#4CG+01 <JGN278 "7GOOO<P,Q,QXcddd [hldj.999 *4KQYZZZ[ [ [ [ [ [ [ [ [ [ [ [ [ [ [ LEEF-a0?DEMqrr" ? ?q',->>>>>6:EM3	@
 #AVWiE2OO]FL#: #:		 				 ;@77%7???? 	I%I((NY#&:";";;nLJ 3ni(((oz***!#7888!>222!>222*  L  /... IHf6GHHHK -ov0111,,,    !||+GWno	
 	
 	

  4 426!&c2223333 *GN -+,,,,,- -s   >4b>>c	c		0f f(_ReductionInforedvar_inforedvar_namer   r   
redarr_typrW   T)frozenc           	      <   |                                 D ]\  }}|j        j        |         }|j        j        |j                 }|                    |          }	t          ||         |||||	          }
|
j        j        t          nt          } || |||
           dS )zYEmit code to finalize the reduction from the intermediate values of
    each thread.
    r   N)
itemsrs   rt   rK   r   r   r   redop_lower_trivial_inplace_binops_lower_non_trivial_reduce)r   r   r;   r   thread_count_varr   r   r   r   rW   reduce_infohandlers               r2   r   r     s     $+==?? @ @Z ^+K8
^+JO<
??;//$(5#!!!
 
 
 "-3? 101 	 	!1;????-@ @r3   c                        e Zd Z fdZ xZS ) ParforsUnexpectedReduceNodeErrorc                 P    t                                          d|            d S )Nz!Unknown reduce instruction node: )r,   __init__r.   s     r2   r  z)ParforsUnexpectedReduceNodeError.__init__  s*    CTCCDDDDDr3   )r4   r5   r6   r  r8   r9   s   @r2   r   r     sA        E E E E E E E E Er3   r   c                    |j         j        D ]}t          ||          rndt          ||j                  r@|j        j        }t          ||||          }|                    ||j	        j
                   nt          |          t          | |||j                  r nt          j        r2|j        }|                    | j         d|j         d| d|           dS dS )z+Lower trivial inplace-binop reduction.
    rT   z	: parfor z reduction  =N)r   reduce_nodes_lower_var_to_var_assign!_is_inplace_binop_and_rhs_is_initr   valuefn_emit_binop_reduce_callstorevartargetrK   r   _fix_redvar_name_ssa_mismatchr   r   print_variablerC   r4   )r   r;   r   r   r0   r  redvar_resultvarnames           r2   r   r     s&    '4  #GT22 	9.t[5LMM 	9B3G-{ M ]1ABBBB 34888 )$#.#:< < 	E	 % 
)zGGBKGGGGGG	
 	
 	
 	
 	

 
r3   c                   
 |j          d
|j        j                            
|j                   |                    |j                  }t          j        |j	        |          5 }|j
        }|j        j        D ]}t          ||          rnt          |t          j                  rjt#          
fd|                                D                       r=t'          |||          }|                    |
           |                    |           nt-          |          t/          | |||j                   r nddd           n# 1 swxY w Y   t0          j        r*|j         }	|                    | j         d|	 d|	           dS dS )zFLower non-trivial reduction such as call to `functools.reduce()`.
    #initc              3   .   K   | ]}|j         k    V  d S r(   rT   )rk   var	init_names     r2   	<genexpr>z,_lower_non_trivial_reduce.<locals>.<genexpr>  s*      NNcCH	1NNNNNNr3   Nz: parfor non-trivial reduction r  )r   rs   rt   
setdefaultr   r   rK   r
   r   rr   r`   r   r  r  r)   r   Assignany	list_vars_emit_getitem_callr
  r-   r   r  r   r   r  rC   )r   r;   r   r   num_thread_llvalr   tidr0   elemr  r  s             @r2   r   r     s    *111IN%%i1GHHH '7'<==		7?,<	=	= j+8 	 	D'66 =T29-- =NNNNT^^=M=MNNNNN=)#wDD  y111""4(((( 7t<<< -VWd'2'>@ @ %              , % 
)zEE'EEE	
 	
 	
 	
 	

 
s   %CEE	E	c                     t          |t          j                  rat          |j        t          j                  rB|                     |j        j                  }|                     ||j        j                   dS dS )zaLower Var->Var assignment.

    Returns True if-and-only-if `inst` is a Var->Var assignment.
    rT   TF)	r)   r   r  r  r~   r   rK   r
  r  )r;   r0   loadeds      r2   r  r    sj    
 $	"" z$*bf'E'E 11dk&6777t5r3   c                     d }|j         }|j        }|j        }|                    |j        j                  }|| f}t          |j        |t          j	                  }	|
                    |||	|          }
|
S )z%Emit call to ``redarr_var[idx]``
    c                     | |         S r(   rD   )rV   r`   s     r2   reducer_getitemz+_emit_getitem_call.<locals>.reducer_getitem  s    e}r3   )rr   rp   r   r   r   rK   r   r   r   r^   compile_internal)idxr;   r   r"  rr   ctxr   arg_arrrG   sigr  s              r2   r  r    s}       oG
/C'Jook49::GS>D
K*J

C
CC#tDDDKr3   c           
      4   d }d }t           j        |t           j        |t           j        |t           j        |t           j        |i|          }|j        }|j        }|j        }	|	                    |j
        j                  }
t          j        r|j
        j                            |j                  }t#          j        |j
        |gd|j                  }|j        j        }t-          t.          j        ||j
        j                 ||j                           |j        j        |<   |                    |           |	                    |j                  }||
|j        f}t-          |j        t.          j        |	|j                  }|                    ||||          }|S )z;Emit call to the ``binop`` for the reduction variable.
    c                 D    |}t          |           D ]}|||         z  }|S r(   r   thread_countrV   initcr   s        r2   reduction_addz._emit_binop_reduce_call.<locals>.reduction_add.  2    |$$ 	 	ANAAr3   c                 D    |}t          |           D ]}|||         z  }|S r(   r*  r+  s        r2   reduction_mulz._emit_binop_reduce_call.<locals>.reduction_mul4  r0  r3   Nr[   )operatoriaddisubimul	ifloordivitruedivrp   rr   r   r   r   rK   r   r   rB   getr   r   r   rC   rs   rt   r   r   r   r   r-   rW   r   uintpr#  )binopr;   r   r   r/  r2  kernelr%  rr   r   r&  init_varr   rt   arg_thread_countrG   r'  r  s                     r2   r	  r	  *  s        	}}}M= F /CoG'Jook49::G% &)/33K4KLLH((3D
 
 
	 .(.7J 6 ;<HM"/
 /
 + 	9%%%'7'<==g{';<D
Z9O C ((&#tDDMr3   c                     t          | t          j                  sdS | j        }t          |t          j                  sdS |j        dk    rdS |j        j        | dk    rdS dS )zDIs ``inst`` an inplace-binop and the RHS is the reduction init?
    Finplace_binopr  T)r)   r   r  r  r   oprI   rK   )r0   r   rI   s      r2   r  r  ^  sq     dBI&& u
*Cc27## u
v  u
w|+,,,,,u4r3   c                    | j         j        }t          |t          j                  r	 |                    |          }|j        }|j        j        }||k    }n # t          $ r ||j        j	        k    }Y nw xY w|rH||j        j	        k    r8|
                    |j        j	                  }	|                    |	|           dS dS )z5Fix reduction variable name mismatch due to SSA.
    rT   TF)r{   rB   r)   r   r  	get_exactunversioned_namer  r!   rK   r   r
  )
r   r;   r0   r   rB   reduction_varredvar_unver_nametarget_unver_nameis_same_source_varra   s
             r2   r  r  m  s     #E$	"" 	H!OOK88M !. > $ <!26G!G  	A 	A 	A "-0@!@		A  	 dk...oodk&677  ; 777t5s   A A43A4c                 z    t           j        rJt          d|||           |dd         D ],}t          d|t          |            |                     -t	          |          |z
  } fd|dd         D             }	t                      }
|	D ]}|r|D ]}|
                    |           |
rt          |
          dz   nd}|	                    d|f           |
                    |           |dz   }|
                    |           i t          d          }|
D ]}|dk    rt          |          |<   |dz  } t          |          }d	|ifd
g }g }d}d}t           j        r0t          d|           t          d|	           t          d|           t          |	|          D ]\  }}|dz   }|rt          fd|D                       }nd}||k    r<|                    t          |gt          |dd                   z                        m|                    |           ||z  }||fS )z&Create shape signature for GUFunc
    r   rM   Nargumentrt   c                 6    g | ]}|vr |           ndS )rK  )rD   )rk   r  r   r}   rt   s     r2   rm   z+_create_shape_signature.<locals>.<listcomp>  s;    jjjZ]#U:J:J  g6666PUjjjr3   r   alatest_alphac                 j    | dk    r||          S dxx         dz  cc<   t          d                   S )Nr   rO  rM   )chr)r.  	class_map
alpha_dicts     r2   
bump_alphaz+_create_shape_signature.<locals>.bump_alpha  sE    66Q<~&&&!+&&&z.1222r3   rD   rG   classesthreadcount_ordinalc              3   0   K   | ]} |          V  d S r(   rD   )rk   r.  rT  rR  s     r2   r  z*_create_shape_signature.<locals>.<genexpr>  s/      CC!ZZ955CCCCCCr3   )r   rv   rw   rx   r   setaddmaxinsertordrQ  ziptupler   list)r   r   r   rG   r   r}   rt   r   
num_inoutsrU  	class_set_class	max_classthread_num_classalphabetnrV  gu_singu_soutcountsyms_sinclsargdim_symsrS  rT  rR  s   `    ``                 @@@r2   r   r     s     Q'^T5QQQabb 	Q 	QA*aa*;*;Aw*O*O*OPPPPT^+JjjjjjjaefgfhfhaijjjGI ! ! 	! ! !a    &/6I""QINN1yl###MM) 1}MM"###I3xxH  
q&&h--9Q<q=8h-- 8,J3 3 3 3 3 FGEH :fdi!!!#%8999&& ! !S	 	CCCCCsCCCCCHHHJ MM%!4 5Xabb\8J8J JKKLLLLMM(### HHGr3   c                 ^    t          | j                  D ]\  }}t          d|d|           d S )N    r?   )	enumerater|   rw   )blockr   r0   s      r2   _print_blockrr    s@    UZ(( $ $4fad####$ $r3   c                 x    |                                  D ]$\  }}t          d|           t          |           %dS )z%Pretty-print a set of IR blocks.
    zlabel: N)r   rw   rr  )	body_dictlabelrq  s      r2   _print_bodyrv    sL     "))  uiU r3   c                 2   |                                  }t          |                                          }t          |                                          }||         j        }||         j                            t          j        ||                     |S r(   )	ru   minkeysrZ  rC   r|   r   r   Jump)	loop_bodyblocksfirst_label
last_labelrC   s        r2   wrap_loop_bodyr    su    ^^Ffkkmm$$KV[[]]##J


 C
:""27;#<#<===Mr3   c                     t          |                                           }| |         j        d d         | |         _        d S )NrM  )rZ  ry  r|   )r{  r~  s     r2   unwrap_loop_bodyr    s=    Y^^%%&&J!*:!6!;CRC!@Ijr3   c                     | |v rdS | |v r,|                     |            |                    |            dS |                     |            dS )a_  If the variable is already defined more than once, do nothing.
       Else if defined exactly once previously then transition this
       variable to the defined more than once set (remove it from
       def_once set and add to def_more set).
       Else this must be the first time we've seen this variable defined
       so add to def_once set.
    N)rY  remove)a_defdef_oncedef_mores      r2   add_to_def_once_setsr    sa     	(		UUr3   c                    |                      t          j                  }|D ]}|j        j        }t          |||           |j        }	t          |	t          j                  r.t          |	j        t          j
                  r|	j        j        ||<   t          |	t          j                  r/|	j        dk    r$|	j        j        |v r|	j        j        |	j        f||<   t          |	t          j                  rm|	j        dk    rb|	j        j        |v rT||	j        j                 \  }
}|
|v r*||
         }t!          ||          st          |
||           nt          |
||           t          |	t          j                  r`|	j        dk    rU|	j        D ]M}t          |t          j                  r|j        }||         }t'          |dd          rt          |||           NdS )a  Effect changes to the set of variables defined once or more than once
       for a single block.
       block - the block to process
       def_once - set of variable names known to be defined exactly once
       def_more - set of variable names known to be defined more than once
       getattr_taken - dict mapping variable name to tuple of object and attribute taken
       module_assigns - dict mapping variable name to the Global that they came from
    r   r   mutableFN)
find_instsr   r  r  rK   r  r  r)   Globalpytypes
ModuleTyper4   r   rA  attrr   r   rG   r~   r   )rq  r  r  getattr_takenrt   module_assignsassignments
one_assignr  rI   base_obj	base_attrbase_mod_nameargvaravtypes                  r2   compute_def_once_blockr    s    ""29--K! 0E 0E
!&UHh777c29%% 	; #)W%788 ;(+	(:u%c27## 	>)(;(;	RZ@Z@Z$'INCH#=M% c27## 	C&(8(8SX]m=[=[ #0">Hi>)) !/x 8$]I>> G )8XFFF %XxBBBc27## 	E&(8(8 ( 
E 
Efbf-- )#[F  69e44 E )8DDDa0E 0Er3   c           
         |                                  D ]o\  }}t          ||||||           |j        D ]N}t          |t          j                  r2t          |j        |||||           t          |j        |||||           OpdS )zCompute the set of variables defined exactly once in the given set of blocks
       and use the given sets for storing which variables are defined once, more than
       once and which have had a getattr call on them.
    N)	r   r  r|   r)   r   r*   r{   compute_def_once_internalr{  )	r{  r  r  r  rt   r  ru  rq  r0   s	            r2   r  r  ;  s     ")) 
v 
vu 	uh-R`aaaJ 	v 	vD$.. v&t(M[bdrsss)$.(Hm]dftuuu	v
v 
vr3   c                 r    t                      }t                      }i }i }t          | |||||           ||fS )zRCompute the set of variables defined exactly once in the given set of blocks.
    )rX  r  )r{  rt   r  r  r  r  s         r2   compute_def_oncer  M  sD     uuHuuHMNi8]GUcdddXr3   c                 r    t          | t          j                  sJ |                    | j                   | S r(   )r)   r   r~   rY  rK   )r  varsets     r2   	find_varsr  W  s3    c26"""""
JJsxJr3   c                    | j         j        |v r:|                    | df           t          j        dk    rt          d| d           dS t                      }t          | j        t          |           |
                    |          }t          j        dk    rt          d| d|d|           t          |          d	k    rt          | j        d |          rt          j        dk    r!t          d
| || j         j                            |                    |            t          || j         j                 t          j        j                  s|| j         j        gz  }dS t          |          d	k    r9|                    | df           t          j        dk    rt          d| d           n8|                    | df           t          j        dk    rt          d| d           dS )Nzstored arrayrM   Instructionz: could not be hoisted because the created array is stored.Fz_hoist_internal:zuses:zdiff:r   zWill hoist instructionT
dependencyz. could not be hoisted because of a dependency.znot purez, could not be hoisted because it isn't pure.)r  rK   r   r   rv   rw   rX  r   r  r  
differencer   r   r)   r   r   r   )	r0   dep_on_param
call_tablehoistednot_hoistedrt   stored_arraysusesdiffs	            r2   _hoist_internalr  \  s   {=((D.1222!Q&&-'cdddu55DTZD111??<((D"" $wEEE
4yyA~~'$*dJ??~!Q&&*D'$+:J2KLLLt'$+"23U^5IJJ 	/T[-..Ltt99q==l3444%**mT+[\\\j1222%**mT+YZZZ5r3   c                    |j         D ]}t          |t          j        t          j        f          r`|                     |j        j                   t          ||j	        j                 dd          r|                    |j	        j                   t          |t          j                  r.t          | ||j        |           t          | ||j        |           d S )Nr  F)r|   r)   r   StaticSetItemSetItemrY  r  rK   r   r  r   r*   find_setitems_blockr{   find_setitems_bodyr{  )setitemsitemssetrq  rt   r0   s        r2   r  r  {  s    
 L LdR-rz:;; 
	LLL)***
 wtz/EBB .TZ_---fm,, 	L(DOWMMMx4>7KKKL Lr3   c                 ^    |                                 D ]\  }}t          | |||           dS )z
      Find the arrays that are written into (goes into setitems) and the
      mutable objects (mostly arrays) that are written into other arrays
      (goes into itemsset).
    N)r   r  )r  r  r{  rt   ru  rq  s         r2   r  r    sF     ")) @ @uHhw????@ @r3   c           	      0   t          | t          j                  r{t          | j        t          j                  r\| j        j        dk    rL| j        j        j        |v r9|| j        j        j                 }|dt          gk    rt          | ||||||          S dS )Nr   r   F)
r)   r   r  r  r   rA  r   rK   r   r  )r0   r  r  r  r  rt   r  	call_lists           r2   empty_container_allocator_hoistr    s    4## H4:rw''H


**tz34	"%%"4z7#.H H H5r3   c                    t          j         |           }g }g }t          ||          \  }}t          |          \  }	}
t                      }t                      }t	          ||||           t          t          |                              |                    }t          j        dk    rt          d|d|d|d|d| 
  
         |D ]}t          |||           |                                D ]X\  }}g }|j        D ]@}t          |||	||||          rt          |t          j                  r%|j        j        |v rt'          |||	||||          rWnt          |t(          j                  rg }t          j        dk    r#t          d           |                                 |j        j        D ]k}t          |||	||||          rt          |t          j                  r$|j        j        |v rt'          |||	||||          rV|                    |           l||j        _        |                    |           B||_        Z||fS )NrM   zhoist - def_once:z	setitems:z	itemsset:zdep_on_param:zparfor_params:r   )ru   r  r   rX  r  r_  r  r   rv   rw   r  r   r|   r  r)   r   r  r  rK   r  r   r*   rz   r{   r   )parfor_paramsr{  rt   wrapped_blocksr  r  r  r  r  r  reverse_call_tabler  r  siru  rq  	new_blockr0   new_init_blockib_insts                       r2   hoistr    s   9]++LGK *)W==Hh'5n'E'E$Z#uuHuuHx9g>>>L))44X>>??L""!8[(KQY[jlx  {K  MZ  	[  	[  	[ 5 5R84444!))  u	J 	# 	#D.t\:#*K(L L 6D"),, 61AX1M1M"4z#*K(L L  	
 D&-00 6!#)Q..(OOOIIKKK#3 
3 
3G6w"G['8M M % $Wbi88 %+x77*7L*+2K(T T % %"))'2222'5$T""""

Kr3   c                 B    t          | t          j        j                   S r(   )r)   r   r   r   )redtypes    r2   redtyp_is_scalarr    s    '5>#78888r3   c                     d}t          | t          j        j                  r|| j        z  }| j        } t          j                            | |d          S )ziGo from a reducation variable type to a reduction array type used to hold
       per-worker results.
    rM   C)r)   r   r   r   r   rN   )redtypr   s     r2   r   r     sM     I&%..// V[ 	>	3777r3   c                     t          | t          j        j                  sJ t          j                            | j        | j        | j                  S )zYGiven a reduction array type, find the type of the reduction argument to the gufunc.
    )r)   r   r   r   rN   r   layout)redarraytyps    r2   redarraytype_to_sigr    sB     k5>#788888> 1;3C[EWXXXr3   c                 z    t          |           }|                                D ]\  }}||k    r||         ||<   |S )aQ   We use ir_utils.legalize_names to replace internal IR variable names
        containing illegal characters (e.g. period) with a legal character
        (underscore) so as to create legal variable names.
        The original variable names are in the typemap so we also
        need to add the legalized name to the typemap as well.
    )r   r   )namesrt   outdictxys        r2   legalize_names_with_typemapr    sH     U##G $ $166 GAJNr3   c                 ^    t          | t          j                  r| j        dk    r| j        S | S )Nr   )r)   r   ArrayCompatibler   rN   )r  s    r2   to_scalar_from_0dr    s.    !U*++ 6Q;;7NHr3   c
           
      Z,  ef t           j        dk    rt          d           |j        j        }
t          j        |j                  }t          |           t          |j	                  }d |j	        D             }|j
        }t          j        j                            ||          }| j        j        t          j        j                            | j        ||| j        j                  \  }}t)          t+          t-          |          t-          |          z
  t-          |          z
                      }t           j        dk    r|t          d|dt/          |                     t          d|dt/          |                     t          d|dt/          |                     t          d|dt/          |                     g }i }i }d	}g }|D ]}|         }t1          |t2          j                  st1          |t2          j                  r|j        }|j        }|t           j        k    sJ g }t?          |          D ]O}d
tA          |          z   }|!                    |           |!                    |           ||f||<   |dz  }||<   P|||<   |!                    |           t1          |t2          j"                  st1          |t2          j#                  r|j        }|j        }|t           j        k    sJ g }t?          |          D ]U}d
tA          |          z   }|!                    |           |!                    |           ||f||<   |dz  }||         |<   V|||<   |!                    |           |!                    |           |}t           j        dk    rt          d|dt/          |                     |	$                    t-          |                    }	|	D ])} d| z  }!tK          j&        tO          |!|
                     *tQ          |	|| j        j                   g }"g }#|D ]m}$|$dz   }%|"!                    |%           tS          |$                   }&|#!                    |&           tU          |&          }'|%v r|%         |'k    sJ h|'|%<   n||z   |"z   }t           j        dk    rlt          d|dt/          |                     t          d|dt/          |                     t          d|dt/          |                     tW          |           tY          ||z   |z             ft           j        dk    r>t          dt)          f-                                          dt/          f                     tY          |          eefd|D             }(t           j        dk    rt          dt)          e-                                          dt/          e                     t          d|(dt/          |(                     |D ]=})t          d|)           t          d|)         dt/          |)                              >fd|D             }*fd||z   D             |#z   }+t           j        dk    r t          d|*           t          d|+           t]          |f           |},ffd|D             }|}-g }d}.t?          t          |-                    D ]v}/|.rW|/t          |          k     rDt1          |*|/         t2          j/        j0                  r|!                    |-|/         dz              [|!                    |-|/                    wt]          |e           tc          |          }0te          d|0          }1t           j        dk    rt          d|dt/          |                     dtg          ti          |                    5                    d d!          z  }2t           j        rt          d"t/          |2          d|2           d#}3|3d$|2z   d%z   d&6                    |          z   d'z   z  }3tn          t          d(}4|-                                D ]"\  }5}6|5         }7|3d)f|5         z   z  }3t1          |7t2          j#                  st1          |7t2          j                  rd*}8nd}8|8rtq          tr          | j        |5          }9d+}:t           j        rt          d,|9t/          |9                     |9t1          |9tt          j;                  rV|9j<        d-k    rKtq          tr          | j        |9j=                  }:t           j        rt          d.|:t/          |:                     nWt1          |9tt          j>                  r=|9j?                 }:t           j        r$t          d.|:t/          |:          |:j?                   |:t1          |:tt          j@                  st1          |:tt          jA                  r<|:jB        };t           j        rt          d/|;t/          |;                     |;|4|:j?        <   nt1          |:t2          jC        jD                  r_|:j?        E                    d0          d	         }<t           j        r)t          d1|<|:jF        t/          |:jF                             |:jF        |4|<<   n@t           j        rt          d2           t          d3tA          |5          z   |5j                  |3d4|7jF        jH        z   d0z   z  }3t          |6|7jJ                  D ]\  }=}>|3|>d5z   f|=         z   d6z   z  }3n?|3d7d&6                    ffd8|6D                       z   z  }3t          |6          dk    r|3d6z  }3|3d9z  }3$t?          t          |                    D ]F}/|.rBt1          |*|/         t2          j/        j0                  r|3d)|-|/         z   d:z   ||/         z   d9z   z  }3Gd;}?t          |"          d	k    r|3d)|?z   d4z   z  }3|3d<z  }3t          |"|          D ]s\  }%}$|3d)f|$         z   d5z   f|%         z   d=z   |?z   d>z   z  }3t           jK        r?|3d?z  }3|3d@f|$         z   d6z   f|$         z   dAz   z  }3|3dBf|%         z   d6z   f|%         z   dAz   z  }3tt?          |          D ]X}@t?          |@dz             D ]}A|3d)z  }3|@}B|3dC|(|@         z   dDz   tA          |B          z   dEz   tA          |B|z             z   dFz   z  }3Yt           jK        rSt?          |dz             D ]}A|3d)z  }3|3dGz  }3t?          |          D ]}@|3dH|(|@         z   dIz   |(|@         z   d6z   z  }3 |3d9z  }3t?          |dz             D ]}A|3d)z  }3|3|1dJz   z  }3t          |"|          D ]V\  }%}$t           jK        r"|3dKf|$         z   d9z   z  }3|3dLf|%         z   d9z   z  }3|3d)f|%         z   d=z   |?z   dMz   f|$         z   dNz   z  }3W|3dOz  }3t           j        r=t          dPt/          |3          dN|3           t          dQ|4t/          |4                     i }Ct          |3|4|C           |C|2         }Dt           j        rt          dRt/          |D          dN|D           t          jN        |D          }Et           j        r]t          dSt/          |E                     |EO                                 t          dTt/          |                     tW          |           tc          |EjP                  }Fi }G|1gt+          fQ                                          z   |(z   }H|F-                                D ]1\  }=}$|=|Hvr(|j        jR        S                    |=|
          j?        |G|=<   2t]          |EjP        |G           t           j        r#t          dU           |EO                                 t2          j/        0                    |ddV          g|*z   }It           j        rt          dWt/          |I          dN|I           t          |EjP                  dz   }Jt          ||J          }t          |          dz   }Kt           jK        r|-                                D ]\  }L}M|M                                }N|NV                                 |Mj        }
|MjR        }O|MjW        D ]r}P|N!                    |P           t1          |Ptt          jX                  r?|PjY        j?                 t2          jZ        vrRdX[                    |PjY        j?                  }Qt3          j\        |Q          }R|OS                    dY|
          }Stu          jX        tu          j]        |Q|
Z          |S|
[          }T|R|Sj?        <   |N!                    |T           tu          j^        |S|PjY        gd+|
\          }U|N!                    |U           t          j_        j`        a                    t2          jb        |Sj?                 |PjY        j?                           }V|V| j        j        |U<   t|N||L<   t           j        rt          d]           tW          |           t          |          }Wt          |||W          \  }X}Y|EjP        t          |EjP        f                                                   }Z|ZjW        d+d^         |Xz   |ZjW        d^         gz   |Z_W        t          |           | jh        d_         }[|X|Yd`|[ji        |jj        <   t           j        rt          da           tW          |           |EjP        -                                D ]\  }L}Mt          |MjW                  D ]e\  }}Pt1          |Ptt          jX                  rD|PjY        j?        |1k    r3|Pj        }
|MjR        }Otu          jl        |O|
          }\|MjW        d+|         |\_W        |MjW        |dz   d+         |M_W        t          |f                                          }]|\!                    tu          jm        |]|
                     |-                                D ]\  }^}_t          |_|O          |EjP        |^<   t          |f                                          }`|M|EjP        |K<   |\|EjP        |L<   |EjP        |`         !                    tu          jm        |K|
                      ng t           j        r#t          db           |EO                                 t          |EjP                  |E_P        t          |EjP                   t           j        rCt          dc           |EO                                 t          dd|           t          de           |jq        }a|s"t           j        rt          df           d*|_q        t          |EjP                    G dg dht          js                  }bt          jt        |||E|It2          jb        |||bi          }c|a|_q        t          t2          jb        g|IR  }dt           j        rt          dj|d           |c|,|d|+|fS )ka~  
    Takes a parfor and creates a gufunc function for its body.
    There are two parts to this function.
    1) Code to iterate across the iteration space as defined by the schedule.
    2) The parfor body that does the work for a single point in the iteration space.
    Part 1 is created as Python text for simplicity with a sentinel assignment to mark the point
    in the IR where the parfor body should be added.
    This Python text is 'exec'ed into existence and its IR retrieved with run_frontend.
    The IR is scanned for the sentinel assignment where that basic block is split and the IR
    for the parfor body inserted.
    rM   z'starting _create_gufunc_for_parfor_bodyc                 &    g | ]}|j         j        S rD   )r   rK   rj   s     r2   rm   z2_create_gufunc_for_parfor_body.<locals>.<listcomp>   s    EEEaA$)EEEr3   zparfor_params = r?   rd   zparfor_inputs = re   r   expanded_tuple_var_z$parfor_inputs post tuple handling = zVariable %s used in parallel loop may be written to simultaneously by multiple workers and may result in non-deterministic or unintended results._arrzloop_indices = zloop_body = zparam_dict = c                      g | ]
}|         S rD   rD   )rk   r   ind_dicts     r2   rm   z2_create_gufunc_for_parfor_body.<locals>.<listcomp>  s    <<<!(1+<<<r3   zind_dict = zlegal_loop_indices = zpd = z
pd type = c                 :    g | ]}t          |                   S rD   )r  rk   r   rt   s     r2   rm   z2_create_gufunc_for_parfor_body.<locals>.<listcomp>  s&    HHHQ$WQZ00HHHr3   c                      g | ]
}|         S rD   rD   r  s     r2   rm   z2_create_gufunc_for_parfor_body.<locals>.<listcomp>  s    KKKQgajKKKr3   znew param_types:znew func_arg_types:c                      g | ]
}|         S rD   rD   )rk   r   
param_dicts     r2   rm   z2_create_gufunc_for_parfor_body.<locals>.<listcomp>  s    :::qZ]:::r3   Fparam__sentinel__zlegal parfor_params = z__numba_parfor_gufunc_%s-_zgufunc_name  zdef z(sched, z, z):
)r   r   ro  TNz	func_def:r   znamed_tuple_def:zgval:(zname:z0Didn't find definition of namedtuple for globls.zCould not find definition of z = =,z = (c                      g | ]
}|         S rD   rD   )rk   r  r  s     r2   rm   z2_create_gufunc_for_parfor_body.<locals>.<listcomp>  s    -O-O-Ojm-O-O-Or3   z)
z = np.ascontiguousarray(!ParallelAcceleratorGufuncThreadIdz*numba.np.ufunc.parallel._iget_thread_id()
[z]
z<    print("thread id =", ParallelAcceleratorGufuncThreadId)
zF    print("initial reduction value",ParallelAcceleratorGufuncThreadId,z.shape)
z>    print("reduction array",ParallelAcceleratorGufuncThreadId,zfor z in range(sched[z	], sched[z] + np.uint8(1)):
zprint("z",z = 0
zD    print("final reduction value",ParallelAcceleratorGufuncThreadId,zD    print("final reduction array",ParallelAcceleratorGufuncThreadId,z] = 
z    return None
zgufunc_txt = zglobls:zgufunc_func = zgufunc_ir dump zloop_body dump zgufunc_ir dump after renaming r  zgufunc_param_types = z{} =rZ   )r  rC   )r  r  rC   r[   zparfor loop bodyrM  parfor_diagnostics)r  r  zAfter hoistingz#gufunc_ir last dump before renamingzgufunc_ir last dumpr   rt   z(No aliases found so adding noalias flag.c                       e Zd Zd ZdS )<_create_gufunc_for_parfor_body.<locals>.ParforGufuncCompilerc                 N   ddl m} t          j        } |d          }|                    | j                  }|j                            |j                   |                    | j                  }|j                            |j                   |	                                 |gS )Nr   )PassManagerfull_parfor_gufunc)
numba.core.compiler_machineryr  r   DefaultPassBuilderdefine_parfor_gufunc_pipelinestatepassesextend/define_parfor_gufunc_nopython_lowering_pipelinefinalize)r/   r  dpbpmparfor_gufunc_passeslowering_passess         r2   define_pipelineszM_create_gufunc_for_parfor_body.<locals>.ParforGufuncCompiler.define_pipelines   s    AAAAAA-C122B#&#D#DTZ#P#P I18999!QQRVR\]]OI_3444KKMMM4Kr3   N)r4   r5   r6   r  rD   r3   r2   ParforGufuncCompilerr    s#        
	 
	 
	 
	 
	r3   r  )pipeline_classz5finished create_gufunc_for_parfor_body. kernel_sig = )ur   rv   rw   r{   rC   ru   r{  r   r   r   r   r   r   r   r   rs   rt   get_parfor_reductionsr   r   sortedr_  rX  rx   r)   r   r   NamedUniTupleri  rN   PARFOR_MAX_TUPLE_SIZEr   r   r   Tuple
NamedTupler  warningswarnr    replace_var_with_arrayr   r  rv  r  r   r   r   r   r   r   hexhashreplacejoinr   r   r   r   r   rA  r   ArgrK   r  FreeVarr  
containersBaseNamedTuplesplitinstance_classr"   r4   r]  fieldsr   execr   run_frontendrz   r|  valuesrB   r   r   r   clearr|   r  r  number_domainformatr   Constr   coretypingr   r   r  r  rx  ry  r  metadata
hoist_infoidrp  Blockrz  r   rZ  r   noaliasr   CompilerBase
compile_ir)gr;   r   rt   r   r   r   localshas_aliasesr   r}   rC   r{  
parfor_dimloop_indicesr  parfor_outputsr   r   parfor_inputstuple_expanded_parfor_inputstuple_var_to_expanded_namesexpanded_name_to_tuple_varnext_expanded_tuple_varparfor_tuple_paramspipi_typetuple_counttuple_dtypethis_var_expansionr   expanded_nametuple_typesracemsgparfor_redarrsparfor_red_arg_typesr  arrredarraytype	redarrsiglegal_loop_indicespdparam_typesr   parfor_argsparfor_params_origascontigpindexloop_body_var_tablesentinel_namegufunc_name
gufunc_txtgloblstup_var	exp_namestup_type	named_tupfunc_defnamed_tuple_defgvalnamed_tuple_namerK   
field_namegufunc_thread_id_vareachdimindent	sched_dimloclsgufunc_func	gufunc_ir	var_tablenew_var_dictreserved_namesgufunc_param_typesgufunc_stub_last_label	new_labelru  rq  r  rB   r0   strvalr   r   
assign_lhs
print_noder'  r  r  r  start_blockdiagnostics
prev_blockbody_first_labelrl   bbody_last_label	old_aliasr  kernel_func
kernel_sigr  r  sg     `                                                                                                  @@r2   r   r     s   , ""7888


C 	&*++I	V&''JEE63DEEEL MM])<<V]SSNn$G%*]%9%O%O0H&J &J"NN   	! 	!" "M "" -d=6I6IJJJ!>3^8L8LMMM -d=6I6IJJJ!>3^8L8LMMM $& "$!# .4 .4"+gu~.. *	4*WeFY2Z2Z *	4!-K!-K&">>>>>!#;'' 5 5 5<S8T8T T,33MBBB"))-888 >@G*=9'1,')4&& /A'+&&r****-- 	4GUEU1V1V 	4 "-K!-K&">>>>>!#;'' 8 8 5<S8T8T T,33MBBB"))-888=?G*=9'1,')4Q&&.@'+&&r****(//33330M""4mS$}J]J]^^^ S0011E < <=?CD 	0c::;;;;5)Wgn6NOOO N 	% 	%Flc"""-gcl;;##L111'55	'>>3<9,,,,,$GCLL "N2^CM"" -d=6I6IJJJsD4F4FGGGnid9oo>>>I -]^-KNa-acjkkJ""  ""$ $	 	 	 +<AAH<<<<|<<<""mVHNN$4$455sDNNKKK##$$		& 	& 	&
   	E 	EB'2,S$wr{2C2CDDDD IHHH-HHHKKKKK=>+IKKKNbbN"" +...#^444 i,,,K::::M:::M&MH.//00 = = 	=S''''{6*EN,@AA (  !3F!;G!CDDDD  !3F!;<<<< i***,Y77'8KLLM""$		! 	! 	! -DLL!!#s++-K Cnd;//kBBBJ &;&ii..0289 9J ''F
 :??AA : :7#fz'222
x!122 	x!455	III .	" ^W_gFFH"O% =k8T(^^<<<#x11 	KK6))&+NGOX]&[&[O- Z0/4CXCXYYY"&11 K&-hm&<O- K0/"?33_5IK K K*	:: N
;;N*0D- 9gtT$ZZ88837F?/001A1PQQ N'6';'A'A#'F'Fq'I$- Dg'7-<"?#ABBD D D 0?/MF+,) NLMMM#$Cc'll$R%,[2 2 2%("9"BBSHHJ$'	8?$C$C H H jj3.D1AACGG

H &499-O-O-O-OY-O-O-O#P#PPPJ9~~"" c!
e

M**++ N N 	N
;v#68LMM 	N6$6v$>>,-/<V/DEGLM NJ ?
>Qf33e;;
CC
 77 U USfz#.c?#%()+?@BGH 	H
) 	U[[Jdgqrugvvy||  @J  KN  @O  O  R]  ]  ]J\_ijm_nnqtt  xB  CF  xG  G  JU  U  UJ $$ . .GaK(( 	! 	!F& JJ	v)'23)* 9~~& #	#
 9%& ' '' -- 	.

 % JN++ 	! 	!F& JJh
Z(( 	i 	iG$!3G!<<uDGYZaGbbehhhJJe
 
Q''  f

-(**J77 I IS) 	}beopsettw|||Jbeopsettw|||Jfz#.&')/02<S/BDHI 	I

%%J /otJ//zBBBif...EVU###$K F[ 1 14EEE%k22I i111i111I #9#344IL#_Z  !!"$67N__&& R R	c&&!'!2!8!A!A$!L!L!QLi&555 .///...1c# # $&12  ##$$		  	  	  ,I,<==A %Y0FGGIy))A-I %  )%OO-- 	) 	)LE5

IOO)CKE
 ? ?  &&&dBI.. ?t{/08KKK  $]]4;+;<<F"'"5f"="=K..c::C!#C1P1P1P253"@ "@ "@J(3GCH%$$Z000 "$T[/A$TW!X!X!XJ$$Z000*+55ej+238+<+24;3C+DF FC <?GN,Z8(Ie  !!!I#I..N 	7NSSG["3y'7'<'<'>'>#?#?@K"',w6+:J2:N9OOKY "#78K4;8C)E )EK69%  I "(..00 # #u ,, !	 !	GAtI  #';#3}#D#Dh  XeS11
"'*RaR.
"ZA/
#&y~~'7'7#8#8  !!"'*:C"@"@AAA (oo// C CFQ*8E*B*BI$Q''"%inn&6&6"7"7.3	 +*4	 '  188GIs++- - - 3444$Y%566I	 !!! "#$$$gui!!!I ! 	><===i.///    x4    %
+- - -K EM5:;(:;;;J SEzRRRZA[[[r3   c           	          g }|j         D ]}t          |t          j                  r|j        j        | v rv|j        }|j        j        }t          j        d|          }|	                    d|          }	t          j        ||	j        <   t          j        ||	|          }
|                    |
           |	                    d|          }||j        j                 ||j        <   |                    t          j        |j        ||                     t          j        |j        |	||          }t          t          j        t          j                            ||j        j                 dd          t          j        ||j        j                           ||<   |                    |           t          |t(          j                  r0t-          | d|j        i||           t-          | |j        ||           |                    |           |S )Nr   z$const_ind_0z$valrM   r  )r|   r)   r   r  r  rK   rC   rB   r  r   r   r:  r   r  r  r   r   r   r   r^   r   r*   replace_var_with_array_internalr{   r{  )varsrq  rt   r   r  r0   rC   rB   
const_node	const_varconst_assignval_varsetitem_nodes                r2   replace_var_with_array_in_blockrs    s   I
  dBI&& 	V4;+;t+C+C(CK%E!S))J~s;;I&+kGIN#9ZC@@L\***nnVS11G$+DK,<$=GGL!RYtz7C@@AAA:dk9gsKKL&/
EN009I1JAsSSUZU_ahimitiyaz'| '|Il#\***fm,, 	V+D1do2FQZ[[[+D$.'9UUUr3   c                 h    |                                 D ]\  }}t          | |||          |_        d S r(   )r   rs  r|   )rm  r{  rt   r   ru  rq  s         r2   rl  rl  :  sD    !)) V Vu4T5'9UU

V Vr3   c                     t          | |||           | D ]D}||         }|                    |d            t          j                            |dd          ||<   Ed S )NrM   r  )rl  popr   r   r   )rm  r{  rt   r   r   el_typs         r2   r  r  >  sk    #D)WiHHH : :At^))&!S99

: :r3   c                 d'   	]^_  j         ^ j        ]ddlm}m} t
          j        rvt          d           t          d|j        |j	        |j
        |j                   t          d|           t          d|           t          d|           t          d|           t          j        |          \  }}|j                            |j        j                  }|\  }} |              |||||d	d
          }|j        }|j                                         t
          j        rt          d||           ^ fd}t)          |          }t+          |          D ]}||         \  }}} ||          } ||          }|dk    sJ  ||          }|||f||<   t
          j        rPt          d                    |          |||           t/          j        ]d                    |          |||           t2          j                            d          }t2          j                            |          }t2          j                            |          }^                    t<          j                  } ^                    t<          j                   }!t2          j                            |           }"t2          j                            |!          }#^!                    t<          j         d          }$^!                    t<          j         d          }%|%j"        }&^#                    |           }'|$                    d           |$                    d          }(t
          j        rt          d|j%                   |j%        r| })|"}*n|!})|#}*t/          j&        ]|)^!                    t<          j         |          d          }+t/          j&        ]|)^!                    t<          j         |          d          },t+          |          D ]}||         \  }}}|j"        |&k    r]'                    ||&          }|j"        |&k    r]'                    ||&          }|j"        |&k    r]'                    ||&          }](                    ||%          }])                    |]*                    |+^!                    t<          j         |          g                     ])                    |]*                    |,^!                    t<          j         |          g                     t/          j+        ]j,        t2          j        -                    |!g           d          }-t/          j+        ]j,        t2          j        -                    t2          j        .                                |!g          d          }.t/          j+        ]j,        t2          j        -                    t2          j                            t<          j        j/                  g           d          }/]0                    |/g           }0]0                    |-g           }1t/          j1        ]]2                    d|0|0"                    d                              5  t/          j        ]d|0           ^j3        4                    ]tj          d           ddd           n# 1 swxY w Y   t2          j        -                    |!|!|!|"|"g          }2t/          j+        ]j,        |2d          }3]0                    |3|0^!                    t<          j         |          |+|,g          }4]0                    |.|$g           ^!                    t<          j         |dz            }5]6                    |4|5          }6]7                    |)|6d           }7t
          j        rdnd}8t2          j        -                    |"|!|"|"|!|*| g          }9|j%        rt/          j+        ]j,        |9d!          }:nt/          j+        ]j,        |9d"          }:]0                    |:^!                    t<          j         |          |+|,|4|7^!                    t<          j        |8          g            	fd#|D             };t)          |          }<t)          |          |<z
  }=] fd$__fd%|d|=         D             |;z   }>t)          |>          }?t)          |          dz   }@t/          j&        ]|^!                    t<          j        d|?z             d&          }g }A])                    ]8                    |7|          |           |A9                    ^!                    t<          j        |'                     i }Bt+          |?          D ]	}|>|         }C||         }D||         }E]*                    |^!                    t<          j        |dz             g          }F||=k    r ^:                    |E          ^]|C          }Gt/          j;        ]|Gj<        |Ej=                  }Ht+          t)          |H                    D ]}I|A9                    |H|I                    ])                    ]8                    |Gj>        |          |F           t          |Et<          j@                  r|D|v r|EjA        t<          jB        k    r^C                    |EjA                  nt2          j                            d          }Jt/          j&        ]|J          }K])                    |C|K           ])                    ]8                    |K|          |F           |C|Kf|B|D<   |A9                    ^!                    t<          j        ^#                    |J                                ^:                    |E          ^]|C          }Gt/          j;        ]|Gj<        |Ej=                  }Ht+          t)          |H                    D ]}I|A9                    |H|I                    ])                    ]8                    |Gj>        |          |F           ||@k     rt          |Et<          jD                  r |Ej"        }L^E                    ]|C|E|L          }Cn|E}Lt          |Lt<          jF                  s^C                    |L          nt2          j                            d          }Jt/          j&        ]|J          }M])                    |C|M           nct          |Et<          jF                  s^C                    |E          nt2          j                            d          }Jt/          j&        ]|J          }M])                    ]8                    |M|          |F           i }Ng }O|(d         g}O^!                    t<          j        d|z            |N|(d         <   t)          |          t)          |>          k    sJ t)          |          t)          |          k    sJ t)          |          t)          ||z             k    sJ t)          |          t)          |j        dd                   k    sJ t          ||>|||z             D ]T\  }D}C}E}Pt          |Et<          jH        jI                  r|Ej=        t)          |P          z
  }nd}t
          j        rt          d'|Dd(|Pd)|Ed*|           |PD ]}Qt
          j        rt          d+|Dd,|E           |D|v r$^!                    t<          j        d          |N|Q<   nK ^:                    |E          ^]|C          }Gt/          j;        ]|GjJ        |Ej=                  }R|R|         |N|Q<   |Q|OvrRt
          j        r1t          d-|Qd.|           t/          j        ]|Qd/z   |N|Q                    |O9                    |Q           |dz   }Vt)          |N          dz   }St/          j&        ]| |Sd0          }R])                    |4|R           d}|OD ]}Qt
          j        rt/          j        ]|Qd/z   |N|Q                    ])                    |N|Q         ]*                    |R^!                    t<          j        |          g                     |dz   }|?dz   t)          |A          z   }Tt/          j&        ]| ^!                    t<          j        |T          d1          }U])                    ^!                    t<          j        d|z  |'z            |U           t+          |?          D ]R}|$}V]*                    |U^!                    t<          j        d|z             g          }F])                    |V|F           St+          t)          |A                    D ]Y}I]*                    |U^!                    t<          j        d|?z   |Iz             g          }F])                    |A|I         |F           Zt/          jK        |          }Wt2          j        -                    t2          j        .                                ||"|"|g          }Xt/          j+        ]j,        |X|          }Y^jL        M                    |j                   t
          j        rt/          j        ]d2|Y           ]0                    |Y||R|U|Wg           t
          j        rt/          j        ]d3|Y           ]0                    |.|1g           |BN                                D ]{\  }Z}[|[\  }C}K]*                    |K^!                    t<          j        d          g          }\])                    ]O                    |\           P                    |Z                     |^jL        M                    |j                   dS )4zF
    Adds the call to the gufunc function from the main function.
    r   )build_gufunc_wrapper_launch_threadsmake_parallel_loopzouter_sig = rn   	expr_argsexpr_arg_typesr   FT)cache
is_parforszparallel function = c                     t          | t          j                  r                    | j                  S                     t          j        |           S r(   )r)   r   r~   r   rK   get_constantr   r:  )r   rp   r;   s    r2   
load_rangez(call_parallel_gufunc.<locals>.load_rangek  sC    a   	8??16***''Q777r3   rM   z'call_parallel_gufunc loop_ranges[{}] = zloop range[{}]: %d %d (%d)
   z%Parfor has potentially negative start
dim_starts)sizerK   	dim_stopsget_parallel_chunksizerT   set_parallel_chunksizer   z<=znum_threads: %d
)z@Invalid number of threads. This likely indicates a bug in Numba.Nget_sched_size   rc   do_scheduling_signeddo_scheduling_unsignedc                 P    g | ]"}                     |         j                  #S rD   )r   rK   )rk   r  r;   
redarrdicts     r2   rm   z(call_parallel_gufunc.<locals>.<listcomp>  s,    DDDqwz!}122DDDr3   c                     | v r8|          \  }}                     |          }                    ||          }|S                      |           S )a?  Given a variable name, if that variable is not a new name
           introduced as the extracted part of a tuple then just return
           the variable loaded from its name.  However, if the variable
           does represent part of a tuple, as recognized by the name of
           the variable being present in the exp_name_to_tuple_var dict,
           then we load the original tuple var instead that we get from
           the dict and then extract the corresponding element of the
           tuple, also stored and returned to use in the dict (i.e., offset).
        )r   extract_value)r  orig_tupoffsetrI  resrr   r   r;   s        r2   load_potential_tuple_varz6call_parallel_gufunc.<locals>.load_potential_tuple_var  s]     %%%4Q7Hfooh//G''88CJ??1%%%r3   c                 &    g | ]} |          S rD   rD   )rk   r  r  s     r2   rm   z(call_parallel_gufunc.<locals>.<listcomp>	  s%    III((++IIIr3   pargszvar =zgu_sig =ztype =zi =zvar = z type = z
dim_sym = z, i = z = %d
pshapepstepszbefore calling kernel %p
zafter calling kernel %p
)Qrp   rr   ro   ry  rz  r   rv   rw   rG   return_typerecvrpysigr	   normalize_signaturelibraryget_functionrs   llvm_func_namerK   _ensure_finalizedr   r   r  r
   printfllvmliter   IntTypePointerTyper   r   r^   r:  r  rx   get_abi_sizeofrv  signedr   sextsubr   gepget_or_insert_functionmoduleFunctionTypeVoidTypebitwidthr   if_unlikelyicmp_signed	call_convreturn_user_excRuntimeErrormulallocabitcastr   
make_arrayunpack_tuplestridesr   datar)   r  rN   booleanget_data_typeOptionalcastBooleanr]  r   r   rP   get_null_valueactive_code_libraryadd_linking_libraryr   loadgetvar)`r;   cresr   	outer_sigr|  r}  r   r   r   r  r{   r   r}   r   ry  rz  rG   r  	llvm_funcsinsoutinfowrapper_namer  num_dimr   rg   rh   ri   byte_t
byte_ptr_tbyte_ptr_ptr_tintp_tuintp_t
intp_ptr_tuintp_ptr_tzerooneone_typesizeof_intp	sched_sig
sched_typesched_ptr_typer  r  get_chunksizeset_chunksizer   num_threadscurrent_chunksizeget_sched_size_fntyr  num_divisions
multiplier
sched_sizerc   
debug_flagscheduling_fntydo_schedulingr   r   ninoutsall_argsnum_argsnum_inpsarray_stridesrv_to_arg_dictrl  r  atydstaryr  r   rJ   rv_argunpacked_atyptrsig_dim_dictoccurrencesgu_sigdim_symshapesnshapes	num_stepsstepsstepsizer  fntyr  kr   only_elem_ptrrr   rp   r  s`   `        `   `                                                                               @@@r2   r   r   E  s    oGoG, , , , , , , ,  ,"###nini.Coy	0 	0 	0,,,k9%%%///nl+++ !4Y??D+))$+*DEEIIC O	4d&+> > >D9LL""$$$ :$lD9998 8 8 8 8 8 +G7^^ . .'NtT
5!!z$qyyyyz$t,A! 	.;BB1EEu  N7$B$I$I!$L$L $. . . [  ##F((00J[,,Z88N##EJ//F$$U[11G((00J+))'22KQ//D


u{A
.
.CxH((00K MM!

I M5}7KLLL %
#
$ $'"6"6K#" #"(46 6 6J #'"6"6K#" #"(35 5 5I 7^^ Q Q'NtT:!!LL11E9  <<h//D9  <<h//D{{4%%7;;((Q( ()* *	+ 	+ 	+
 	dGKK	)0)=)=ek1)M)M(NP P 	Q 	Q 	Q 	Q 2  "--%' ' 'M
 2  !5!5!7!7'CC%' ' 'M
 4  !4!4UZ5H!I!I2NN O
 ,,33K]B77		Wg&9&9$2=2B2B12E2E'G 'G 
H 
H ^ ^w 3[AAA))'<3]	^ 	^ 	^^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ #+227Wgz[e<fgg3  N LL+292F2Fu{T[2\2\2<2;2= > >M LL'''%%ek7Q;??J]J77JNN:JWNEEE,3!Jk..Wj*g~vVX XO V6w~7F<RT T T  6w~7F<TV V V LL  W& &'19m7''J
, ,-. . . EDDDDGDDDG7||H)nnx'G& & & & & & &( JIIIYxx5HIIIGSH8}}H3xx!|H!!JL    D MMM'//%44d;;;--ej+FFGGGN8__ 2A 2AqklQkk$!5!5ej!a%!H!H IJJ<<)'$$S))'7C@@C*7CKJJG3w<<(( 1 1$$WQZ0000MM'//#(J??EEEEU233 &	Ae||)u}44 ,,SY777$K//22  !,Wc::c6***goofjAA3GGG'*Fms#$$W%9%9%*gF\F\]`FaFa%b%bcccc-g((--gwDD!.wSXNNs7||,, 5 5A!((4444gooch
CCSIIII8||c5>22 '#&8L!,,wS,GGCC#&L(u}EE3w,,\:::$K//22  )'377c3'''' )em<<3w,,S111$K//22  )'377MM'//#z::C@@@@ LKQ<.K!(!5!5ej!g+!N!NL1y>>S]]****y>>S000000y>>St__,,,,y>>S!3444444!$Y%3S4Z"A "A  S#vc5>/00 	3v;;&AAA! 	M'3
FHc5!LLL 	 	G% 6hZ555e||(/(<(<UZ(K(KW%%-g((--gwDD -gsy#(KK(.q	W%{**) X,1===N7Gi,?gAVWWW""7+++AAA	( ,!#G &wXNNNFMM-(((	A  ! 	PN7Gi$7g9NOOO!7;;((
A' '($) $)	* 	* 	*
 E
 1s=111Ig22J	# #)13 3 3E MM'&&uz1w;3LMM   8__ % %kk%'"6"6uz1q5"I"I!JKKh$$$$3}%%&& - -kk$$JHq 02 234 4 	mA&,,,, !*--D;##HK$8$8$:$:%3Z%/%=> >D 
	'l	K	KB33DLAAA Bw <bAAALLdFE40111 Aw ;R@@@LL!2 3444$$&& F F1VFW-A-A%*a-P-P,QRRgll=117>>!3D3DEEEE33DLAAAAAs   =8[[[)Zru   r3  r   r  r  dataclassesr   llvmlite.irr  rb   r   r   numba.parforsr   
numba.corer   r   r   r	   r
   numba.core.ir_utilsr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   numba.core.typingr   r   numba.parfors.parforr   numba.core.errorsr    r!   r"   r#   #numba.parfors.parfor_lowering_utilsr$   Lowerr&   r+   r<   r   r   r   r   r   r  r  r	  r  r  r   rr  rv  r  r  r  r  r  r  r  r  r  r  r  r  r  r   r  r  r  r   rs  rl  r  r   rD   r3   r2   <module>r     s           & & & & & &                      E E E E E E E E E E E E E E E E                                       ( ( ' ' ' ' '       8 8 8 8 8 8            F E E E E E	% 	% 	% 	% 	%(. 	% 	% 	%/ / /R- R- R-j
       !@ !@ !@HE E E E E} E E E

 
 
<$
 $
 $
L	 	 	   1 1 1h  ! ! !FC C CJ$ $ $    A A A   <E <E <E|v v v$    
  >L L L@ @ @
 
 
2  2  2 h9 9 9
8 
8 
8Y Y Y     _\ _\ _\B  8V V V: : :oB oB oB oB oBr3   