U
    hâËd®  ã                   @   s0  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mZmZmZmZ d dlmZmZ d dlmZ d dl
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 d dl m!Z! d d	l"m#Z# d d
l$m%Z% dd„ Z&G dd„ de'ƒZ(dd„ Z)G dd„ de'ƒZ*dd„ Z+dd„ Z,dd„ Z-dd„ Z.dd„ Z/dS )é    N)Úadd)ÚtypesÚirÚrewritesÚconfigÚir_utils)Úinfer_globalÚAbstractTemplate)Ú	signature)ÚutilsÚtyping)	Úget_call_tableÚmk_unique_varÚcompile_to_numba_irÚreplace_arg_nodesÚguardÚfind_callnameÚrequireÚ
find_constÚGuardException)ÚNumbaValueError)ÚOPERATORS_TO_BUILTINS)Únumpy_supportc                 C   s   |dkr| | S | S d S ©Nr   © )Údim_sizeÚindex_constr   r   úU/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/numba/stencils/stencilparfor.pyÚ_compute_last_ind   s    r   c                   @   sT   e Zd Zdd„ Zdd„ Zdd„ Zdd„ Zd	d
„ Zdd„ Zdd„ Z	dd„ Z
dd„ ZdS )ÚStencilPassc                 C   s.   || _ || _|| _|| _|| _|| _|| _d S ©N)Úfunc_irÚtypemapÚ	calltypesÚarray_analysisÚ	typingctxÚ	targetctxÚflags)Úselfr!   r"   r#   r$   r%   r&   r'   r   r   r   Ú__init__"   s    zStencilPass.__init__c                    s  ddl m} tˆ jjƒ\}}g }i }| ¡ D ].\}}|D ] }t||ƒr8| |¡ |||< q8q,|sddS ˆ jj ¡ D �]¦\}	}
tt	t
|
jƒƒƒD �]ˆ\}‰tˆtjƒ�rÌtˆjtjƒ�rÌˆjjdk�rÌˆjjj|k�rÌtˆjjƒ}‡fdd„ttˆjjƒƒD ƒ}ˆjj}t‡ fdd„|D ƒƒ}|D ]}t|tjƒ�rtd	ƒ‚�q| d
¡}|ˆjjj }t|ˆ j||
j|
j |ˆ j!ˆ j"ƒ\}}}|j# dd¡}ˆ  $|	||||ˆj%|||¡	}|
jd|… | |
j|d d…  |
_qŒtˆtjƒrŒtˆjtjƒrŒˆjjdkrŒt&t'ˆ jˆjƒdkrŒt (dˆj ¡ˆ_qŒqpdS )zP Finds all calls to StencilFuncs in the IR and converts them to parfor.
        r   )ÚStencilFuncNÚcallc                    s   i | ]}|ˆ j j| “qS r   )ÚvalueÚargs©Ú.0Úi)Ústmtr   r   Ú
<dictcomp>I   s      z#StencilPass.run.<locals>.<dictcomp>c                 3   s   | ]}ˆ j |j V  qd S r    )r"   Únamer.   ©r(   r   r   Ú	<genexpr>L   s     z"StencilPass.run.<locals>.<genexpr>zITuple parameters not supported for stencil kernels in parallel=True mode.ÚoutÚindex_offsetsé   )ÚstencilÚnumba))Znumba.stencils.stencilr*   r   r!   ÚblocksÚitemsÚ
isinstanceÚappendÚreversedÚlistÚ	enumerateÚbodyr   ÚAssignr,   ÚExprÚopÚfuncr3   ÚdictÚkwsÚrangeÚlenr-   Útupler   Z	BaseTupleÚ
ValueErrorÚgetÚget_stencil_irr%   ÚscopeÚlocr"   r#   ÚoptionsÚ_mk_stencil_parforÚtargetr   r   ÚConst)r(   r*   Z
call_tableÚ_Zstencil_callsZstencil_dictZcall_varnameÚ	call_listZone_callÚlabelÚblockr0   rH   Ú
input_dictÚin_argsZarg_typemapZarg_typeÚout_arrÚsfÚ
stencil_irÚrtÚarg_to_arr_dictr7   Ú	gen_nodesr   )r(   r1   r   Úrun,   sx    

ÿ
þý
ÿ
    ý

    þ(ÿ
þÿýzStencilPass.runc              	   C   s¬   |  ¡ D ]ž\}}|j}|j}g }|jD ]x}	t|	tjƒr”| ¡ }
t|
tjƒrdt|
j	tj
ƒrd|
j	jdksht‚| t |
j	j	||¡¡ | t ||¡¡ q&| |	¡ q&||_qdS )zß
        Find return statements in the IR and replace them with a SetItem
        call of the value "returned" by the kernel into the result array.
        Returns the block labels that contained return statements.
        ÚcastN)r<   rO   rP   rB   r=   r   ÚReturnÚpoprC   r,   rD   rE   ÚAssertionErrorr>   ZJump)r(   r;   Úexit_value_varÚparfor_body_exit_labelrW   rX   rO   rP   Únew_bodyr1   Z	prev_stmtr   r   r   Úreplace_return_with_setitemh   s     
ÿ
þz'StencilPass.replace_return_with_setitemc
           S         sì	  g }
|j }tjdkr4td||ˆ||||ƒ t |¡ |d }ˆj|j }t |ˆj¡\}}t 	|¡}t 
|||ˆjˆj¡ tjdkr–tdƒ t |¡ t |ˆjj|ˆj¡ tjdkrÈtdƒ t |¡ ˆj|j j‰ |j}|j}g }tˆ ƒD ].}t |tdƒ|¡}tjˆj|j< | |¡ qîˆ ||||||	¡\}}tjdk�rhtdƒ td|ƒ td	|ƒ t |¡ g }ˆj |¡}| |¡}ˆ t|ƒk�s”t‚g }g }tˆ ƒD ]f}ˆ || || |
||¡}ˆ  || |
||¡}| |¡ | |¡ | t!j"j# $|| ||d¡¡ �q¤t%| &¡ ƒd }t '||¡||< t |td
ƒ|¡} |j(ˆj| j< g }!ˆ dk�rd|d }"nPt |tdƒ|¡}"tj) *tjˆ ¡ˆj|"j< tj+ ,||¡}#t -|#|"|¡}$|! |$¡ tjdk�rÒtdƒ t |¡ t '||¡}%ˆdk�rÒˆj|j }t d¡}&t ||&|¡}'tj+ .|d|¡}(tj) *tj|j¡ˆj|&< |%j/ 0t -|(|'|¡g¡ t d¡})t ||)|¡‰d|j1k�r¬|j1d }*t2j3 3|*¡}+ˆj4 5|+|j(¡�s t6dƒ‚| (|*¡},n
| (d¡},t 7|,|¡}-|j(ˆj|)< |%j/ 0t -|-ˆ|¡g¡ t d¡}.t ||.|¡‰t!j8jj9 :|j(|j|j;¡ˆjˆj< t |tdƒ|¡}/tj< =t>¡ˆj|/j< t ?dt>|¡}0t -|0|/|¡}1|%j/ |1¡ t@ A|j(¡jBjC}2|2dk�rŒd}2tj+ .|/|2|¡}3t |tdƒ|¡}4tjD E|j(¡ˆj|4j< t -|3|4|¡}5|%j/ |5¡ t Fdt>jGˆ|'|4gˆj4ˆjˆj¡}6t |tdƒ|¡}7t -t 7d|¡|7|¡}8|6 |8¡ tjHˆj|7j< t |tdƒ|¡}9t -t 7d|¡|9|¡}:|6 |:¡ tjˆj|9j< t |tdƒ|¡};ˆj4 ItJ¡}<|<ˆj|;j< t ?dtJ|¡}=t -|=|;|¡}>|6 |>¡ ˆj4 K|<tjHfd i ¡}?tj+jL|;|7|7fd|d �}@|?ˆj|@< t |td!ƒ|¡}AtjMˆj|Aj< t -|@|A|¡}>|6 |>¡ ‡ ‡‡‡fd"d#„}Bt|jƒD ]V}C|Ag|j }D|Ag|j }E|B|<|C|||;|6||D|9d$ƒ
 |B|<|C|||;|6||E||C d%ƒ
 �q^| Nˆ|¡ |%j/ 0|6¡ �n¦d|j1k�rx|j1d }*t2j3 3|*¡}+ˆj4 5|+|j(¡�sd}FtO|Fƒ‚t |td&ƒ|¡}Aˆj4 ItJ¡}<|<ˆj|Aj< t ?dtJ|¡}=t -|=|A|¡}G|%j/ |G¡ ˆj4 K|<tjHfd i ¡}?tj+jL|Add|d �}H|?ˆj|H< t |td'ƒ|¡}ItjMˆj|Ij< t -|H|I|¡}>|%j/ |>¡ t 7| (|*¡|¡}Jt |td(ƒ|¡}K|j(ˆj|Kj< t -|J|K|¡}L|%j/ |L¡ t PˆtJddƒ|I|K|¡}M|%j/ |M¡ tQtjHˆjˆj ˆj|Ij ˆjˆj j(ƒ}?|?ˆj|M< ˆ R|| |¡ tjdk�r¤td)ƒ t |¡ t Sˆ|"| |¡}NtQtjHˆjˆj ˆj|"j ˆjˆj j(ƒˆj|N< || j/ 0|!¡ || j/ |N¡ t Td*d+¡}Ot |td(ƒ|O¡}Pt -tj7d|Od,�|P|O¡}L|| j/ |L¡ || j/ t U|P|O¡¡ t V|¡}|t%| &¡ ƒ j/ W¡  tjdk�	r td-ƒ t |¡ d.||gf}Qt!j"j# X||%|||"||QˆjY¡}R|
 |R¡ |
 t -ˆ||¡¡ |
S )/z> Converts a set of stencil kernel blocks to a parfor.
        r8   rR   r   z#stencil_blocks after copy_propagatez'stencil_blocks after removing dead codez$parfor_index_varz-stencil_blocks after replace stencil accesseszstart_lengths:zend_lengths:z$parfor_exit_valuez$parfor_index_tuple_varz.stencil_blocks after creating parfor index varNZin_arr_shapeÚshapeZzero_valÚcvalz-cval type does not match stencil return type.Zstencil_outputz	$np_g_varÚnpÚboolZbool_z$np_attr_attrÚemptyz	$none_varz$zero_index_varz$slice_func_varÚsliceé   r   ©rF   r-   rH   rP   ú$slicec
                    sš  ˆj  | tjfd i ¡}
|| }t|ttjfƒs4t‚t |t	dƒ|¡}tjˆj
|j< t|tƒrvt t ||¡||¡}nt |||¡}| |¡ tjj||	r¢||fn||fd|d�}|
ˆj|< t |t	dƒ|¡}tjˆj
|j< t |||¡}| |¡ |||< t |t	dƒ|¡}tj tjˆ ¡ˆj
|j< tj ||¡}t |||¡}| |¡ t ˆ|ˆ|¡}ttjˆj
ˆj ˆj
|j ˆj
ˆj jƒˆj|< | |¡ d S )Nrp   z$border_indr   rq   rr   z$border_index_tuple_var)r%   Úresolve_function_typer   Úintpr=   Úintr   ÚVarre   r   r"   r3   rC   rT   r>   rD   r+   r#   Úslice2_typeÚ
containersÚUniTupleÚbuild_tupleÚSetItemr
   ÚnoneÚdtype)Úslice_fn_tyÚdimrO   rP   Úslice_func_varÚstmtsZborder_indsZborder_tuple_itemsZ	other_argZother_firstÚsigÚsiZsi_varZ	si_assignÚslice_callexprZborder_slice_varÚslice_assignZborder_ind_varÚ
tuple_callÚtuple_assignÚsetitem_call©Úndimsr[   r(   Zzero_varr   r   Úhandle_borderA  sZ    
ý

ü

ÿÿ ÿ
 

ý
z5StencilPass._mk_stencil_parfor.<locals>.handle_borderTFz	$py_g_varz$slice_instz$cval_constz%stencil_blocks after replacing returnZstencilparfor_dummyéÿÿÿÿ)rP   z#stencil_blocks after adding SetItemr9   )Zr;   r   ÚDEBUG_ARRAY_OPTÚprintr   Údump_blocksr"   r3   Zcopy_propagateÚget_name_var_tableZapply_copy_propagater#   Zremove_deadr!   Ú	arg_namesÚndimrO   rP   rI   r   rv   r   r   rt   r>   Ú_replace_stencil_accessesr$   Zget_equiv_setÚ	get_shaperJ   re   Ú_get_stencil_last_indÚ_get_stencil_start_indr:   ZparforsÚparforZLoopNestÚmaxÚkeysZBlockr}   rx   ry   rD   rz   rC   ÚgetattrrB   ÚextendrQ   r   Ztypeofr%   Zcan_convertrL   rT   ÚcoreZnpytypesZArrayZlayoutÚmiscÚModulerl   ÚGlobalr   Zas_dtypeÚtypeÚ__name__Ú	functionsZNumberClassZgen_np_callrn   r|   Zresolve_value_typero   rs   r+   rw   Zinsert_equivr   ÚStaticSetItemr
   ri   r{   ZLocrc   Zsimplify_CFGrd   ZParforr'   )Sr(   rW   rZ   r[   r]   r7   rS   Úreturn_typeÚstencil_funcr_   r`   Ústencil_blocksÚin_arrZ
in_arr_typZin_cpsZout_cpsÚname_var_tablerO   rP   Úparfor_varsr0   Z
parfor_varÚstart_lengthsÚend_lengthsZ	loopnestsZ	equiv_setZin_arr_dim_sizesZ
start_indsZ	last_indsÚlast_indZ	start_indrg   rf   Zfor_replacing_retZparfor_ind_varr†   r‡   Z
init_blockZ
shape_nameZ	shape_varZshape_getattrZ	zero_namerk   Zcval_tyZtemp2Z
full_constZso_nameZdtype_g_np_varZ
dtype_g_npZdtype_g_np_assignZreturn_type_nameZdtype_np_attr_callZdtype_attr_varZdtype_attr_assignr�   Znone_varZnone_assignZzero_index_varZzero_index_assignr€   r~   Zslice_gr…   r‚   r„   Ú	slice_varr‹   r   Zstart_tuple_itemsZlast_tuple_itemsÚmsgZslice_assignedZcallexprZslice_inst_varZcval_const_valZcval_const_varZcval_const_assignZsetitemexprrˆ   Z	dummy_locZret_const_varÚpatternr—   r   r‰   r   rR   �   sN   
  ÿ

û

ÿ

ÿÿ    þ




   ÿ   ÿ

  ÿ


ÿÿ ÿ




ÿ





ýÿ
ú



þý

8÷
÷

þÿ
ÿÿ ÿ þ
þ
ÿ
 

ý


ÿ

    ÿ
zStencilPass._mk_stencil_parforc                 C   s4  |}|dk�r0t  |tdƒ|¡}tj| j|j< t|tj	ƒrRt  
t  ||¡||¡}nt  
|||¡}| |¡ t  |tdƒ|¡}tj| j|j< t  |tdƒ|¡}	t t¡}
tj |
¡}|| j|	j< t  d|
|¡}t  
||	|¡}| |¡ t j |	||gd|¡}| | jtjtjgi ¡| j|< t  
|||¡}| |¡ |S )Nr   Zstencil_const_varr¬   Zcompute_last_ind_varr   r   )r   rv   r   r   rt   r"   r3   r=   ÚnumbersÚNumberrC   rT   r>   r:   Znjitr   r¢   Z
DispatcherrŸ   rD   r+   Zget_call_typer%   r#   )r(   r   Z
end_lengthr`   rO   rP   r¬   r   Úconst_assignZg_varZ
check_funcZfunc_typZg_objZg_assignÚ
index_callÚindex_assignr   r   r   r•   ò  s>    
ÿ ÿ


 
 ÿ

z!StencilPass._get_stencil_last_indc           	      C   s�   t |tƒrtt|dƒƒS dd„ }t|i | j| jtjf| j	| j
ƒ}t|jƒdksRt‚|j ¡ d }t||gƒ ||jd d… 7 }|jd jj}|S )Nr   c                 S   s   t t| dƒƒS r   )ÚabsÚmin)Zs_lengthr   r   r   Úget_start_ind  s    z9StencilPass._get_stencil_start_ind.<locals>.get_start_indr8   éþÿÿÿ)r=   ru   rµ   r¶   r   r%   r&   r   rt   r"   r#   rJ   r;   re   Úpopitemr   rB   r,   )	r(   Ústart_lengthr`   rO   rP   r·   Úf_irrX   Zret_varr   r   r   r–     s    

  þz"StencilPass._get_stencil_start_indc              	      s°  ˆj }|d }dd„ |D ƒ}	d|jkr`|jd D ]}
|
ˆ kr0tdƒ‚q0‡ fdd„|jd D ƒ}ng }|j|krvtdƒ‚ˆj|j j}|j}|j}|jdk}|r´|dg }|dg }n d	d„ |jD ƒ}d
d„ |jD ƒ}t	 
|¡}d}| ¡ D �]¦\}}g }|jD �]Š}t|tjƒ�r@t|jtjƒ�r@|jjdk�r@|jjj|	k�sjt|tjƒ�s\t|tjƒ�rr|jj|	k�rrtdƒ‚t|tjƒ�r€t|jtjƒ�r€|jjdk�r€|jjj|	k�r€|jjj|k�r€|jj}|dk�rØ|g}n"t|dƒ�rú|j|k�rú||j }t	 |¡ˆ_‡‡fdd„|D ƒ}|�r6ˆ |t|ƒ|||¡}|�rŠt|tjƒ�s^tdd„ |D ƒƒ�rftdƒ‚ttt||ƒƒ}ttt ||ƒƒ}d}ˆ |t|ƒ|||¡}|dk�r´|d }nPt |t!dƒ|¡}t"j# $t"j%|¡ˆj|j< tj &||¡}t |||¡}| '|¡ t(‡fdd„|D ƒƒ�r0ˆj|jjj j)}nˆj|jjj }tj *|jj||¡}t+|ˆj|jjj ˆj|j ƒˆj,|< ||_| '|¡ qþ||_qê|�r¨|�s¨tdƒ‚||fS )zÄ Convert relative indexing in the stencil kernel to standard indexing
            by adding the loop index variables to the corresponding dimensions
            of the array index tuples.
        r   c                 S   s   g | ]
}|j ‘qS r   )r3   ©r/   Úxr   r   r   Ú
<listcomp>+  s     z9StencilPass._replace_stencil_accesses.<locals>.<listcomp>Zstandard_indexingz[Standard indexing requested for an array name not present in the stencil kernel definition.c                    s   g | ]}ˆ | ‘qS r   r   r¼   )r_   r   r   r¾   2  s     zYThe first argument to a stencil kernel must use relative indexing, not standard indexing.Nc                 S   s   g | ]}|d  ‘qS )r   r   r¼   r   r   r   r¾   I  s     c                 S   s   g | ]}|d  ‘qS )r8   r   r¼   r   r   r   r¾   J  s     F)ÚsetitemZstatic_setitemz?Assignments to arrays passed to stencil kernels is not allowed.)Zstatic_getitemÚgetitemr8   r3   c                    s   g | ]}t ˆˆ j|ƒ‘qS r   )Ú_get_const_index_exprr!   ©r/   Úv)r(   r]   r   r   r¾   t  s   ÿ  ÿc                 S   s   g | ]}t |tƒ ‘qS r   ©r=   ru   rÂ   r   r   r   r¾     s     z<Variable stencil index only possible with known neighborhoodTz$parfor_index_ind_varc                    s   g | ]}ˆ j |j tjk‘qS r   )r"   r3   r   rt   rÂ   r4   r   r   r¾   ˜  s   ÿz=Stencil kernel with no accesses to relatively indexed arrays.)-r;   rQ   rL   r3   r"   r’   rO   rP   Zneighborhoodr   Zget_tuple_tabler<   rB   r=   r   rC   r,   rD   rE   r{   r£   rS   ÚindexÚhasattrZbuild_definitionsZ_definitionsÚ_add_index_offsetsr@   rv   ÚanyÚmapr¶   r˜   r   r   rx   ry   rt   rz   r>   Úallr}   rÀ   r
   r#   )r(   r]   r©   rZ   r7   r¥   r_   r¦   r§   Zin_arg_namesr½   Zstandard_indexedrŠ   rO   rP   Zneed_to_calc_kernelrª   r«   Ztuple_tableZfound_relative_indexrW   rX   rh   r1   Ú
index_listÚ
index_varsZind_varr†   r‡   Zgetitem_return_typZgetitem_callr   )r_   r(   r]   r   r“   #  sê    


ÿ



ÿ
þý
ü
û
úÿ
þýü

ÿ   ÿÿÿ   ÿ

ÿÿ ÿ
ÿ
ÿÿ
ý
z%StencilPass._replace_stencil_accessesc                 C   s<  t |ƒt |ƒkst‚tdd„ || D ƒƒr:ttt||ƒƒS g }g }tt |ƒƒD �]Ü}|| }	t|	tƒrªt	 
|tdƒ|¡}	tj| j|	j< t	 t	 || |¡|	|¡}
| |
¡ || }t|tƒ�rt	 
|tdƒ|¡}tj| j|j< t	 t	 || |¡||¡}
| |
¡ t|	tƒ�s&t| j|	j tjjƒ�r\| j|j tjk�s>t‚|  |	||||¡}| |¡ qNt|tƒ�s€t| j|j tjjƒ�r¶| j|	j tjk�s˜t‚|  ||	|||¡}| |¡ qNt	 
|tdƒ|¡}tj| j|j< t	j tj|	||¡}| j tjtjtjfi ¡| j|< t	 |||¡}| |¡ | |¡ qN| |¡ |S )zw Does the actual work of adding loop index variables to the
            relative index constants or variables.
        c                 S   s   g | ]}t |tƒ‘qS r   rÄ   rÂ   r   r   r   r¾   ¸  s     z2StencilPass._add_index_offsets.<locals>.<listcomp>Úold_index_varÚ
offset_varZoffset_stencil_index)rJ   re   rÊ   r@   rÉ   r   rI   r=   ru   r   rv   r   r   rt   r"   r3   rC   rT   r>   ro   r�   Z	SliceTypeÚ_add_offset_to_slicerD   ÚbinopÚoperatorr%   rs   r#   r›   )r(   rË   r7   rh   rO   rP   Ú	out_nodesrÌ   r0   rÍ   r²   rÎ   Ú	index_varr³   r´   r   r   r   rÇ   °  s–    
 ÿ ÿ
 ÿ ÿ
ÿÿ  ÿ
ÿÿ  ÿ
 ÿ ÿ 
 ÿ


zStencilPass._add_index_offsetsc                 C   sÎ   t |tƒrBd |j|j¡}i }t|i |ƒ |d }|g}tjf}	n&dd„ }||g}| j|j	 }
|
tjf}	| j
jjj}t||| j| j|	| j| jƒ}|j ¡ \}}t||ƒ |jd jj}| |jd d… ¡ |S )NzRdef f(offset):
                return slice({} + offset, {} + offset)
            Úfc                 S   s   t | j| | j| ƒS r    )ro   ÚstartÚstop)Z	old_sliceÚoffsetr   r   r   rÔ      s    z+StencilPass._add_offset_to_slice.<locals>.fr¸   )r=   ro   ÚformatrÕ   rÖ   Úexecr   rt   r"   r3   r!   Zfunc_idrF   Ú__globals__r   r%   r&   r#   r;   r¹   r   rB   r,   r›   )r(   r­   rÎ   rÒ   rO   rP   Zf_textrÔ   r-   Zarg_typsZ
slice_typeZ_globalsr»   rU   rX   Z	new_indexr   r   r   rÏ   ô  s2    
 þ

  ÿ
z StencilPass._add_offset_to_sliceN)r¡   Ú
__module__Ú__qualname__r)   ra   ri   rR   r•   r–   r“   rÇ   rÏ   r   r   r   r   r   !   s   
<  s" Dr   c                 C   sÌ  ddl m} ddlm}	 ddlm}
 ddlm} | j 	¡ }t	 
|j¡}||_t |j¡}d|krhtdƒ‚ddlm}	 |	j}|	 ||¡�ˆ t||||ƒ}tj d|j¡ ||jj|jj|jj|jjd	ƒ\|j_|j_|j_}|
j|jj|jj|jjd
d	|jj|jjtjd� W 5 Q R X t  |t !¡ ¡}t"| #¡ ƒ}t$| #¡ ƒ}tj% &|¡ tj'dk�rft(dƒ t )|¡ i }|jj *¡ D ].\}}t+ ,|t-|ƒ|¡}|||< |||j.< �qvt /||¡ tj'dk�rÐt(dƒ t )|¡ |jj *¡ D ]\}}|||< �qÜi }| 0¡ D ]‚}|j1D ]t}t2|t+j3ƒ�rt2|j4t+j5ƒ�rtj'dk�rTt(d||j4j6|j4j.|j4j6|kƒ ||j4j6 j.||j4j.< ||j4j6 |_4�q�qütj'dk�r¨t(d|ƒ t(dƒ t )|¡ t 7|¡ ||_||  8|¡d |fS )z'get typed IR from stencil bytecode
    r   )Ú
CPUContext)Ú
cpu_target)Útype_annotations)Útype_inference_stager6   z6Cannot use the reserved word 'out' in stencil kernels.zbefore-inferenceNr   )r!   r"   r#   ZliftedZlifted_fromr-   r¤   Zhtml_outputr8   zInitial stencil_blockszAfter replace_varsrY   r_   zAfter replace arg with arr)9Znumba.core.cpurÝ   Znumba.core.registryrÞ   Znumba.core.annotationsrß   Znumba.core.typed_passesrà   Z	kernel_irÚcopyÚdeepcopyr;   r   r�   rL   Ztarget_contextZnested_contextÚDummyPipeliner   Zrewrite_registryÚapplyÚstater%   r&   r!   r-   r"   r¤   r#   ZTypeAnnotationr   ÚHTMLZadd_offset_to_labelsZ
next_labelr¶   r™   r˜   Z_the_max_labelÚupdater�   rŽ   r�   r<   r   rv   r   r3   Zreplace_varsÚvaluesrB   r=   rC   r,   ZArgrÅ   Zremove_delsZget_return_type)r\   r%   r-   rO   rP   rY   r"   r#   rÝ   rÞ   rß   rà   Zstencil_func_irr¦   r¨   r&   ÚtprU   Z	min_labelZ	max_labelZvar_dictrÃ   ÚtypZnew_varr+   Zcall_typr_   rX   r1   r   r   r   rN     sŒ    
   þøÿ


 
ÿ


rN   c                   @   s   e Zd Zdd„ ZdS )rã   c                 C   sP   ddl m} |ƒ | _|| j_|| j_|| j_|| j_d | j_d | j_d | j_	d S )Nr   )Ú	StateDict)
Znumba.core.compilerrë   rå   r%   r&   r-   r!   r"   r¤   r#   )r(   r%   r&   r-   r»   rë   r   r   r   r)   g  s    zDummyPipeline.__init__N)r¡   rÛ   rÜ   r)   r   r   r   r   rã   f  s   rã   c                 C   s   t t| ||ƒ}|dk	r|S |S )z²
    infer index_var as constant if it is of a expression form like c-1 where c
    is a constant in the outer function.
    index_var is assumed to be inside stencil kernel
    N)r   Ú_get_const_index_expr_inner)r]   r!   rÓ   Ú	const_valr   r   r   rÁ   s  s       ÿrÁ   c                 C   sr   t t|tjƒƒ tt| ||ƒ}|dk	r*|S t | |¡}tt| ||ƒ}|dk	rP|S tt	| ||ƒ}|dk	rj|S t
‚dS )zWinner constant inference function that calls constant, unary and binary
    cases.
    N)r   r=   r   rv   r   Ú_get_const_two_irsr   Zget_definitionÚ_get_const_unary_exprÚ_get_const_binary_exprr   )r]   r!   rÓ   Ú	var_constÚ	index_defr   r   r   rì     s,       ÿ   ÿrì   c                 C   s8   t t| |ƒ}|dk	r|S t t||ƒ}|dk	r0|S t‚dS )zWget constant in either of two IRs if available
    otherwise, throw GuardException
    N)r   r   r   )Zir1Zir2Úvarrñ   r   r   r   rî   –  s    rî   c                 C   sF   t t|tjƒo|jdkƒ |j}t| ||ƒ}t|j }t	d 
||¡ƒS )zQevaluate constant unary expr if possible
    otherwise, raise GuardException
    Zunaryz{}{})r   r=   r   rD   rE   r,   rì   r   ÚfnÚevalrØ   )r]   r!   rò   Z	inner_varrí   rE   r   r   r   rï   ¢  s
    
rï   c                 C   sR   t t|tjƒo|jdkƒ t| ||jƒ}t| ||jƒ}t|j	 }t
d |||¡ƒS )zRevaluate constant binary expr if possible
    otherwise, raise GuardException
    rÐ   z{}{}{})r   r=   r   rD   rE   rì   ÚlhsÚrhsr   rô   rõ   rØ   )r]   r!   rò   Zarg1Zarg2rE   r   r   r   rð   ­  s
    
rð   )0r°   rá   r   ZpytypesrÑ   r   Únumpyrl   Znumba.parfors.parforr:   Z
numba.corer   r   r   r   Znumba.core.typing.templatesr   r	   Znumba.core.typingr
   r   r   Znumba.core.ir_utilsr   r   r   r   r   r   r   r   r   Znumba.core.errorsr   Znumba.core.utilsr   Znumba.npr   r   Úobjectr   rN   rã   rÁ   rì   rî   rï   rð   r   r   r   r   Ú<module>   s8   ,     rX