Ë
    g^(hÇU ã                   óD	  — U d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlZd dl	m
Z
 d dlmZ d dlmZmZmZmZ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 d dlmZ d dlmZmZm Z  d d	l!m"Z"m#Z#m$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/ ddl0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7 ddl8m9Z9m:Z:m;Z;m<Z<m=Z=m>Z>m?Z?m@Z@mAZAmBZBmCZCmDZD ddlEmFZFmGZGmHZHmIZI ddlJmKZKmLZLmMZMmNZNmOZOmPZPmQZQmRZRmSZSmTZTmUZUmVZV ddlWmXZXmYZYmZZZm[Z[m\Z\m]Z]m^Z^m_Z_m`Z`maZambZbmcZcmdZdmeZe ejÌ                  dk(  Zg ejÐ                  d«      d„ «       ZiejÔ                  j×                  eld«      Zm eg d¢«      Zndddddddd d!d!d"œ
Zo eg d#¢«      Zpd$d%d&d'd(d)d*d+d,d-d.œ
Zqd/d0d1œZrejæ                  ejè                  gZuejì                  ejî                  ejæ                  ejè                  ejð                  ejò                  ejô                  ejö                  ejø                  g	Z}e~ejþ                     e€d2<   ejî                  ejæ                  ejè                  ejò                  ejô                  gZ�e~ejþ                     e€d3<   d4„ Z‚d5„ Zƒ	 	 dwd6ee�j                     fd7„Z…d8„ Z†d9eRd:e�j                  d;eˆd<e�j                  d=e�j                  d>eLfd?„Z‰d@eeˆeNf   dAeˆdBeˆdCejþ                  dDeeˆeŠf   f
dE„Z‹dFeRdGeˆdHeˆfdI„ZŒejÐ                  d6e�j                  dJe�j                  fdK„«       Z�ejÐ                  d6e�j                  dJe�j                  dLeŠfdM„«       ZŽejÐ                  	 dxd6e�j                  dJe�j                  dLeeŠ   fdN„«       Z�e�j                    G dO„ dP«      «       Z‘ G dQ„ dRe5«      Z’ G dS„ dT«      Z“ G dU„ dVeU«      Z”e”�j+                  dW«        G dX„ dYe”«      Z–e–�j+                  dZ«       e–�j/                  «         G d[„ d\e–«      Z˜ G d]„ d^eS«      Z™ G d_„ d`e™«      Zš G da„ dbeš«      Z›dce/d>eœeejþ                     exf   fdd„Z� G de„ df«      Zž G dg„ dhe™«      ZŸ G di„ dje™«      Z  G dk„ dle«      Z¡ G dm„ dne2«      Z¢ G do„ dp«      Z£ G dq„ dr«      Z¤e�j                    G ds„ dt«      «       Z¥e�j                    G du„ dv«      «       Z¦y)yé    N)ÚSequence)ÚEnum)ÚAnyÚCallableÚcastÚOptionalÚUnion)Údependencies)Úis_float_dtypeÚis_integer_dtype)Ú
OrderedSet)ÚCeilDivÚFloorDivÚModularIndexing)Úfree_symbol_is_typeÚsymbol_is_typeÚSymTé   )Úcountersé   )Ú	codecacheÚconfigÚcpp_builderÚcpu_vec_isaÚirÚmetrics)ÚLoopBody)ÚBaseSchedulerNodeÚBaseSchedulingÚExternKernelSchedulerNodeÚForeachKernelSchedulerNodeÚFusedSchedulerNodeÚ	SchedulerÚSchedulerNode)Úcache_on_selfÚget_bounds_index_exprÚget_fused_kernel_nameÚhas_free_symbolsÚis_multi_outputs_templateÚis_welford_reductionÚparallel_num_threadsÚPlaceholderÚsympy_index_symbolÚsympy_index_symbol_with_prefixÚsympy_productÚ
sympy_subs)ÚNullKernelHandlerÚopsÚOpsValueÚVé   )ÚBackendFeatureÚBracesBufferÚCSEÚCSEVariableÚDataTypePropagationÚDeferredLineÚDTYPE_TO_COMPUTATION_DTYPEÚIndentedBufferÚKernelÚ
KernelArgsÚOpOverridesÚOptimizationContext)Ú_get_dtype_from_loopbodiesÚ_get_loop_bodyÚcexprÚcexpr_indexÚcodegen_randÚCppCSEVariableÚDTYPE_TO_CPPÚ
INDEX_TYPEÚLocalBufferContextÚmay_unify_binary_op_mask_typeÚpromote_argsÚ(template_fusion_with_epilogues_supportedÚunify_mask_base_typeÚvalue_to_cppÚwin32c                  ó   — t         rdS dS )Nz__declspec(dllexport)Ú ©Ú_IS_WINDOWS© ó    úY/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/torch/_inductor/codegen/cpp.pyÚget_export_declarationrX   W   s   € å&1Ð"Ð9°rÐ9rV   Úschedule)ú+Ú*ú^ú||ÚminÚmaxrZ   r[   r\   r^   r_   ÚargminÚargmaxr]   Úwelford)
ÚsumÚprodÚxor_sumr^   r_   r`   ra   ÚanyÚwelford_reduceÚwelford_combine)
r_   r^   rc   rd   re   rg   rh   r`   ra   rf   z
at::TensorÚlongÚdoubleÚboolzstd::stringzc10::ScalarTypezat::MemoryFormatz
at::Layoutz
at::Devicez
at::Scalar)
ÚTensorÚintÚfloatrk   ÚstrÚ
ScalarTypeÚMemoryFormatÚLayoutÚDeviceÚnumberzstd::vectorzstd::optional)ÚListr   ÚVECTORIZABLE_DTYPESÚMASKED_VECTORIZABLE_DTYPESc                 óŽ  — |t         v rt        j                  }| dv ry| dk(  ry| dv r|t        |   }|t        j                  k(  r| dv rt        t        j
                     }t        |«      rd|› d�nd	|› d
�}t        |«      rd	|› d�nd	|› d�}| dv r|n|}| dv r|S d|› d|› d�S t        | «      rdt        |   › d�S t        | «      ‚)N)re   rc   rf   r   rd   r5   )r_   ra   r^   r`   ©r`   ra   ú-std::numeric_limits<ú>::infinity()ústd::numeric_limits<ú>::min()ú>::max())r_   ra   )r_   r^   úIndexValue<z>{0, ú}úWelford<ú>())	ÚDTYPE_LOWP_FPÚtorchÚfloat32rH   rk   rn   r   r*   ÚAssertionError)Úreduction_typeÚdtypeÚcdtypeÚmin_varÚmax_varÚinit_vars         rW   Úreduction_initr�   ¦   s  € Ø”Ñô —‘ˆØÐ2Ñ2ØØ˜ÒØØÐ;Ñ;Ü˜eÑ$ˆØ”E—J‘JÒ >Ð5IÑ#IÜ!¤%§+¡+Ñ.ˆFô ˜eÔ$ð $ F 8¨=Ñ9à'¨ x¨xÐ8ð 	ô ˜eÔ$ð # 6 (¨-Ñ8à'¨ x¨xÐ8ð 	ð
 -Ð0AÑA‘7Àwˆð  Ñ/ð ð	
ð ˜v˜h f¨X¨J°bÐ9ð	
ô
 ˜NÔ+Øœ, uÑ-Ð.¨cÐ2Ð2Ü
˜Ó
(Ð(rV   c                 ó°   — t         t        |      }t        | «      rd|› d�S | dv r0|t        j                  k(  rt         t        j
                     }d|› d�S |S )Nr�   ú>ry   r   )rH   r<   r*   r„   rk   rn   )r‡   rˆ   Úscalar_types      rW   Úreduction_acc_typer‘   È   sa   € ÜÔ9¸%Ñ@ÑA€KÜ˜NÔ+Ø˜+˜ aÐ(Ð(ØÐ-Ñ-Ø”E—J‘JÒÜ&¤u§{¡{Ñ3ˆKØ˜[˜M¨Ð+Ð+ØÐrV   Úindexc           	      ó>  — |t         j                  k(  }| dk(  r|rdnd}|› d|› d|› �S | dk(  r|› d|› �S | dk(  r|› d|› �S | d	k(  r|› d
|› �S | dv r| › d|› d|› d�S | dk(  r	d|› d|› d�S | dk(  r6t        |t        «      r|\  }}}	nt	        | |«      \  }}}	d|› d|› d|› d|	› d�	S | dv rkt        |d«      rD|j                  t         j                  k(  r'|j                  s|�| › d|› d|› d|› d�S | › d|› d|› d�S |�| › d|› d|› d|› d�S | › d|› d|› d�S t        | «      ‚)Nrc   ú|rZ   ú rd   ú * re   ú ^ rf   ú || )r^   r_   z_propagate_nan(ú, ú)rg   úwelford_combine(rh   ú, {ú})ry   rˆ   z	_combine(z, static_cast<float>(ú), ú)))	r„   rk   Ú
isinstanceÚtupleÚreduction_projectÚhasattrrˆ   Úis_vecr†   )
r‡   ÚvarÚ
next_valuer’   Ú	src_dtypeÚis_boolÚconjunctionÚmeanÚm2Úweights
             rW   Úreduction_combiner­   Ó   sÞ  € ð œ5Ÿ:™:Ñ%€GØ˜ÒÙ$‘c¨#ˆØ��a˜�} A j \Ð2Ð2Ø˜ÒØ��c˜*˜Ð&Ð&Ø˜Ò"Ø��c˜*˜Ð&Ð&Ø˜ÒØ��d˜:˜,Ð'Ð'Ø˜Ñ'Ø Ð! °°°R¸
°|À1ÐEÐEØÐ)Ò)Ø! #  b¨¨°AÐ6Ð6ØÐ*Ò*Ü�j¤%Ô(Ø)ÑˆD�"‘fä0°ÀÓLÑˆD�"�fØ! #  d¨4¨&°°2°$°b¸¸ÀÐDÐDØÐ-Ñ-ä�J Ô(Ø× Ñ ¤E§J¡JÒ.Ø×%Ò%àÐ Ø(Ð)¨°3°%Ð7LÈZÈLÐX[Ð\aÐ[bÐbcÐdÐdð &Ð& i°¨uÐ4IÈ*ÈÐUWÐXðð ÐØ$Ð% Y¨s¨e°2°j°\ÀÀEÀ7È!ÐLÐLà$Ð% Y¨s¨e°2°j°\ÀÐCÐCÜ
˜Ó
(Ð(rV   c                 óJ   — t        | «      r|› d�|› d�|› d�fS | dv r|› d�S |S )Nz.meanz.m2z.weightry   z.index)r*   )r‡   Úaccs     rW   r¢   r¢     sC   € Ü˜NÔ+Ø��eˆ}   S˜k¨c¨U°'¨?Ð:Ð:Ø	Ð/Ñ	/Ø��fˆ~ÐØ€JrV   ÚcodeÚiter_varÚnew_iter_varÚ
loop_startÚloop_endÚreturnc                 ób  — t        «       }t        j                  «       5 }|j                  dt        › d|› dt        |«      › d�|› dt        |«      › d|› d�z   «       |j                  |j                  «       «       t        | j                  «      D ]…  \  }}t        |t        t        f«      sJ ‚d}	t        |t        «      r|j                  }	|j                  }t        j                   d	|› z   d	z   |› |«      }
|	rt        |	|
«      }
|j                  |
«       Œ‡ 	 ddd«       |S # 1 sw Y   |S xY w)
aú  
    f(iter_var) is transformed to f(new_iter_var) under the inner loop
      \/
    for (new_iter_var = loop_start; new_iter_var < loop_end; new_iter_var++) {
        f(new_iter_var)
    }
    Please be careful while using this function,
    as the variable defined in f(iter_var) will be invalid outside the for loop.
    For example:
    auto tmp0 = in_ptr[x0]; ->
    for (new_x0 = start; new_x0 < end; new_x0++){
        auto tmp0 = in_ptr[new_x0];
    }
    The tmp0 is invalid outside the loop.
    zfor (r•   ú = ú;ú < ú; ú++)Nú\b)r7   Ú
contextlibÚ	ExitStackÚ	writelinerI   rE   Úenter_contextÚindentÚ	enumerateÚ_linesr    ro   r;   ÚnameÚlineÚreÚsub)r°   r±   r²   r³   r´   Útransformed_codeÚstackÚ_rÅ   Údeferred_nameÚnew_lines              rW   Úmove_code_under_inner_looprÍ   	  s?  € ô, $“~ÐÜ	×	Ñ	Ó	ð 1 5Ø×"Ñ"Ø”J�<˜q  ¨c´+¸jÓ2IÐ1JÈ!ÐLØ�˜c¤+¨hÓ"7Ð!8¸¸<¸.ÈÐLñMô	
ð 	×ÑÐ,×3Ñ3Ó5Ô6Ü  §¡Ó-ò 	1‰GˆAˆtÜØäÜ ðôð ð ð !ˆMÜ˜$¤Ô-Ø $§	¡	�Ø—y‘y�Ü—v‘v˜e¨¨
Ñ3°eÑ;À¸~ÐPTÓUˆHÙÜ'¨°xÓ@�Ø×&Ñ& xÕ0ñ	1÷1ð, Ð÷-1ð, Ðús   ŸC:D$Ä$D.Úacc_varÚacc_typer‡   rˆ   Úlenc                 óâ   — t        «       }t        j                  «       rd| › d|› d|› d�n
|› d| › d|› d�}|j                  |› «       |j	                  d|› d	�d
d| › d |||«      › d�dg«       |S )a˜  
    MSVC don't support dynamic array(VLA). So we use std::unique_ptr here.
    Ref: https://stackoverflow.com/questions/56555406/creating-dynamic-sized-array-using-msvc-c-compiler
    MSVC is the only one compiler without VLA. support. Since MSVC can't get good performance here.
    We just use unique_ptr make it works on MSVC.
    For other compilers, we continue to use VLA to get best performence.
    úauto z_arr = std::make_unique<z[]>(ú);r•   ú_arr[ú];úfor (int i = 0; i < ú; i++)ú{ú    z
_arr[i] = r¸   r€   )r=   r   Ú
is_msvc_clr¿   Ú
writelines)rÎ   rÏ   r‡   rˆ   rÐ   Úinit_fnÚcode_bufferÚacc_decls           rW   Úreduction_prefix_arrayrß   9  s©   € ô !Ó"€Kô ×!Ñ!Ô#ð �ˆyÐ0°°
¸$¸s¸eÀ2ÑFàˆz˜˜7˜) 5¨¨¨RÐ0ð ð
 ×Ñ˜X˜JÔ(Ø×Ñà" 3 % vÐ.ØØ�7�)˜:¡g¨n¸eÓ&DÐ%EÀQÐGØð		
ôð ÐrV   ÚbufferrÄ   Únew_namec                 óD  — t        | j                  «      D ]ˆ  \  }}t        |t        t        f«      sJ ‚t        |t        «      r/t        j                  d|› z   dz   |› |j                  «      |_        Œ]t        j                  d|› z   dz   |› |«      | j                  |<   ŒŠ y )Nr¼   )rÂ   rÃ   r    ro   r;   rÆ   rÇ   rÅ   )rà   rÄ   rá   ÚirÅ   s        rW   Úreplace_acc_namerä   Z  sœ   € Ü˜VŸ]™]Ó+ò V‰ˆˆ4ÜØäÜðô
ð 	
ð 
ô �dœLÔ)ÜŸ™˜u¨$¨Ñ0°5Ñ8¸X¸JÈÏÉÓSˆD�Iä!Ÿv™v e°°Ñ&7¸%Ñ&?ÀHÀ:ÐPTÓUˆF�M‰M˜!ÒñVrV   r¥   c                 ó®   — | j                  |«      st        j                  j                  S ||dz   i}t	        | |«      }t        j
                  || z
  «      S ©Nr5   )ÚhasÚsympyÚSÚZeror0   Úsimplify)r’   r¥   ÚreplacementÚ	new_indexs       rW   Ú	stride_atrî   i  sJ   € à�9‰9�SŒ>ô �w‰w�|‰|ÐØ˜˜a™�.€KÜ˜5 +Ó.€IÜ�>‰>˜) eÑ+Ó,Ð,rV   Ú
vec_lengthc                 ó¦  ‡‡‡‡	— dŠdŠ	ˆˆˆfd„}ˆ	ˆˆfd„}| }t        j                  dd¬«      }| j                  t        «      r| j	                  t        ‰|«      |«      } t        j                  dd¬«      }| j                  t
        «      r| j	                  t        ‰||«      |«      } t        j                  | «      } | |k7  rt        | ‰‰«      S | S )a±  
    Simplifies the index expression within the range of a vectorized loop.
    Given a vectorized loop variable `var` in the range of a loop with `vec_length`,
    this function transforms the `index` into an equivalent form. It handles
    simplifications for cases where `var` can be expressed as `vec_length * a + b`,
    where `b` ranges from 0 to `vec_length - 1`. The function reduces occurrences
    of `FloorDiv` and `ModularIndexing` in the `index` with best-effort optimizations.

    NOTE:
    The simplified index expression is intended for analysis purposes only, not
    for code generation. It replaces `FloorDiv` and `ModularIndexing` with free variables
    which are not dependent on the loop variable `var` in the vectorized range. Check
    https://github.com/pytorch/pytorch/pull/117221#discussion_r1449746217 for more details.

    Examples:
    1. If `var` is `x3` and `vec_length` is 16, and `x3 = 16*a + b`, then
       `FloorDiv(x3, div)` or `ModularIndexing(x3, div, mod)` becomes a free variable
       when `div` is divisible by 16.
    2. `ModularIndexing(x3, 1, mod)` can be simplified to `x3 + c` where `c` is a free
       variable when `mod` is divisible by 16.
    r   c                 ó�   •— t        ‰| «      }t        j                  | ‰«      ‰k(  rt        j                  ‰› d‰› �«      }‰dz  Š|S )NÚ_div_cr5   )r   rè   ÚgcdÚSymbol)ÚdivisorÚresultÚdiv_freevar_idr¥   rï   s     €€€rW   Úvisit_indexing_divz7simplify_index_in_vec_range.<locals>.visit_indexing_div�  sK   ø€ ä˜#˜wÓ'ˆÜ�9‰9�W˜jÓ)¨ZÒ7Ü—\‘\ S E¨°Ð/?Ð"@ÓAˆFØ˜aÑˆNØˆrV   c                 ó  •— t        ‰| |«      }t        j                  | ‰«      ‰k(  r!t        j                  ‰› d‰› �«      }‰dz  Š|S | dk(  r;t        j                  |‰«      ‰k(  r"‰t        j                  ‰› d‰› �«      z   }‰dz  Š|S )NÚ_mod_cr5   )r   rè   ró   rô   )rõ   Úmodulusrö   Úmod_freevar_idr¥   rï   s      €€€rW   Úvisit_modular_indexingz;simplify_index_in_vec_range.<locals>.visit_modular_indexing˜  s˜   ø€ ä   g¨wÓ7ˆÜ�9‰9�W˜jÓ)¨ZÒ7Ü—\‘\ S E¨°Ð/?Ð"@ÓAˆFØ˜aÑˆNð ˆð ˜Š\œeŸi™i¨°Ó<À
ÒJØœ5Ÿ<™<¨3¨%¨v°nÐ5EÐ(FÓGÑGˆFØ˜aÑˆNØˆrV   rõ   T)Úintegerrû   )rè   ÚWildrç   r   Úreplacer   rë   Úsimplify_index_in_vec_range)
r’   r¥   rï   rø   rý   Úoriginal_indexÚdivÚmodr÷   rü   s
    ``     @@rW   r  r  u  s´   û€ ð0 €NØ€Nöö	ð €Nä
�*‰*�Y¨Ô
-€CØ‡y�y”ÔØ—‘œh s¨CÓ0Ð2DÓEˆä
�*‰*�Y¨Ô
-€CØ‡y�y”Ô!Ø—‘œo¨c°3¸Ó<Ð>TÓUˆä�N‰N˜5Ó!€EØ�ÒÜ*¨5°#°zÓBÐBà€LrV   c                 ó8   — |rt        | ||«      } t        | |«      S ©N)r  rî   )r’   r¥   rï   s      rW   Ústride_at_vec_ranger  ´  s"   € ñ Ü+¨E°3¸
ÓCˆÜ�U˜CÓ Ð rV   c                   ó&   — e Zd ZU dZeed<   eed<   y)ÚParallelDepthz{
    A class representing parallel depth.
    Includes the starting depth of parallelism and the depth of parallelism.
    Úparallel_depthÚstart_depthN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__rm   Ú__annotations__rU   rV   rW   r	  r	  ½  s   … ñð
 ÓØÔrV   r	  c                   ód   ‡ — e Zd Zededefd„«       Zdddeeee	f      fˆ fd„Z
d„ Zd	„ Zd
„ Zˆ xZS )ÚOuterLoopFusedSchedulerNodeÚnode1Únode2c                 óž  — |j                   |j                   u sJ ‚t        d„ ||fD «       «      sJ ‚t        d„ ||fD «       «      rt | |j                   t        |«      t        u rt        |j                  «       «      n|gt        |«      t        u r t        |j                  «       «      z   |«      S |gz   |«      S  | |j                   ||g|«      S )Nc              3   óT   K  — | ]   }t        |«      t        t        t        fv –— Œ" y ­wr  )Útyper  r$   r"   ©Ú.0Únodes     rW   ú	<genexpr>z3OuterLoopFusedSchedulerNode.fuse.<locals>.<genexpr>Î  s1   è ø€ ò 
ð ô �‹Jä+ÜÜ"ðôñ
ùó   ‚&(c              3   ó>   K  — | ]  }t        |«      t        u –— Œ y ­wr  ©r  r  r  s     rW   r  z3OuterLoopFusedSchedulerNode.fuse.<locals>.<genexpr>×  ó   è ø€ ÒT¸TŒt�D‹zÔ8Ô8ÑTùó   ‚)Ú	schedulerÚallrf   r  r  ÚlistÚget_outer_nodes)Úclsr  r  Úouter_loop_fusion_depths       rW   Úfusez OuterLoopFusedSchedulerNode.fuseÉ  s÷   € ð �‰ %§/¡/Ñ1Ð1Ð1Üñ 
ð  ˜ô
ô 
ð 	
ð 
ô ÑTÀeÈUÀ^ÔTÔTÙØ—‘ô ˜E“{Ô&AÑAô ˜×.Ñ.Ó0Ô1ð ðô ˜E“{Ô&AÑAô ˜×.Ñ.Ó0Ó1ñð (ó!ð ð ðñð (ó!ð ñ& �u—‘¨°¨Ð8OÓPÐPrV   r!  r#   Úouter_fused_nodesc                 óè   •— || _         || _        g }| j                   D ]B  }t        |t        t        f«      sJ ‚|j                  t        |j                  «       «      «       ŒD t        ‰| �%  ||«       y r  )
r(  r&  r    r$   r"   Úextendr#  Ú	get_nodesÚsuperÚ__init__)Úselfr!  r(  r&  Úflatten_snodesÚ_nodeÚ	__class__s         €rW   r-  z$OuterLoopFusedSchedulerNode.__init__í  su   ø€ ð ð 	Ôð (?ˆÔ$ØˆØ×+Ñ+ò 	;ˆEÜ˜e¤mÔ5GÐ%HÔIÐIÐIØ×!Ñ!¤$ u§¡Ó'8Ó"9Õ:ð	;ô 	‰Ñ˜ NÕ3rV   c                 ó   — | j                   S r  )r(  ©r.  s    rW   r$  z+OuterLoopFusedSchedulerNode.get_outer_nodesý  s   € Ø×%Ñ%Ð%rV   c           
      ó8  ‡— dt         dt         dt        dt        dt        f
ˆfd„Št        t	        |«      dz
  «      D ]0  }||   j
                  }||dz      j
                  } ‰|||d«      rŒ0 y	 |D ]¦  }t        j                  t        j                  |j                  d | «      }t	        |j                  «      |kD  sŒMt        |t        j                  «      sŒht        |j                  |   t        j                  «      sŒ�|d
z  |j                  |   k  sŒ¦ y	 y)NÚleft_loop_nestÚright_loop_nestÚloop_fusion_depthÚcurrent_checking_depthrµ   c                 ó^  •‡‡— | j                   sJ ‚|j                   sJ ‚| j                   |   Š|j                   |   Šg d¢}t        ˆˆfd„|D «       «      sy|dk\  sJ ‚|dz
  x}dkD  rE|dz   }|t        | j                   «      k  sJ ‚|t        |j                   «      k  sJ ‚ ‰| |||«      syy)N)r¥   ÚsizeÚoffsetÚstepsc              3   óP   •K  — | ]  }t        ‰|«      t        ‰|«      k(  –— Œ y ­wr  )Úgetattr)r  Úattr_compareÚleft_loop_levelÚright_loop_levels     €€rW   r  zaOuterLoopFusedSchedulerNode.check_outer_fusion_loop_level_attr.<locals>._inner.<locals>.<genexpr>  s2   øè ø€ ò ð %ô ˜O¨\Ó:ÜÐ/°Ó>õ?ñùs   ƒ#&Fr5   r   T)Úloopsr"  rÐ   )r5  r6  r7  r8  Úouter_loops_attr_compare_listr@  rA  Ú_inners        @@€rW   rD  zNOuterLoopFusedSchedulerNode.check_outer_fusion_loop_level_attr.<locals>._inner	  sç   ú€ ð "×'Ò'Ð'Ð'Ø"×(Ò(Ð(Ð(Ø,×2Ñ2Ð3IÑJˆOØ.×4Ñ4Ð5KÑLÐò-Ð)ô ô ð )Fôô ð à$¨Ò)Ð)Ð)Ø%6¸Ñ%:Ð:Ð!¸aÒ?à)?À!Ñ)CÐ&Ø-´°N×4HÑ4HÓ0IÒIÐIÐIØ-´°O×4IÑ4IÓ0JÒJÐJÐJÙØ"Ø#Ø%Ø*ô	ð !àrV   r5   r   Fé,  T)ÚLoopNestrm   rk   ÚrangerÐ   Ú	loop_nestÚ	functoolsÚreduceÚoperatorÚmulÚrangesr    rè   ÚInteger)	r.  Úcpp_kernel_proxy_listr&  Úidxr5  r6  Úcpp_kernel_proxyÚouter_rangesrD  s	           @rW   Ú"check_outer_fusion_loop_level_attrz>OuterLoopFusedSchedulerNode.check_outer_fusion_loop_level_attr   s<  ø€ ð(	Ü$ð(	ä%ð(	ô  #ð(	ô %(ð	(	ô
 õ(	ôT œÐ2Ó3°aÑ7Ó8ò 		ˆCØ2°3Ñ7×AÑAˆNØ3°C¸!±GÑ<×FÑFˆOÙØØØ'Øõ	ñ ð		ð !6ò 	ÐÜ$×+Ñ+Ü—‘Ø ×'Ñ'Ð(@Ð)@ÐAóˆLô Ð$×+Ñ+Ó,Ð/FÓFÜ˜|¬U¯]©]Õ;ÜØ$×+Ñ+Ð,CÑDÜ—M‘Mõð ! 3Ñ&Ø"×)Ñ)Ð*AÑBóCñ ð)	ð, rV   c                 óD  — |d   j                   }t        |«      }|D �cg c]'  }|j                  j                  | j                  «      ‘Œ) c}|_        |d   }||j                  _        |j                  j                  d | j                   |j                  _        |S c c}w ©Nr   )Úkernel_groupÚOuterLoopFusedKernelrH  Úfrom_loop_levelr&  ÚinnerÚkernelrB  )r.  rO  rV  Úouter_loop_fused_kernelÚproxyÚouter_fused_proxys         rW   Úmerge_outer_fusion_kernelsz6OuterLoopFusedSchedulerNode.merge_outer_fusion_kernelsV  s¦   € ð -¨QÑ/×<Ñ<ˆÜ"6°|Ó"DÐð /ö)
àð �O‰O×+Ñ+¨D×,HÑ,HÕIò)
ÐÔ%ð 2°!Ñ4ÐØ-DÐ×#Ñ#Ô*Ø,=×,GÑ,G×,MÑ,MØ*ˆd×*Ñ*ð-
Ð×#Ñ#Ô)ð !Ð ùò)
s   Ÿ,B)r  r  r  Úclassmethodr   r'  r#  r	   r"   r$   r-  r$  rS  r^  Ú__classcell__©r1  s   @rW   r  r  È  sb   ø„ Øð!QØ%ð!QØ.?ò!Qó ð!QðF4àð4ð   Ð&8¸-Ð&GÑ HÑIõ4ò &òTöl!rV   r  c                   ó2   — e Zd Zddefd„Zd„ Zd„ Zd„ Zd„ Zy)	ÚRecordOptimizationContextÚ	func_namec                 ó.   — || _         d | _        d | _        y r  )rd  Úcurrent_nodeÚopt_ctx)r.  rd  s     rW   r-  z"RecordOptimizationContext.__init__i  s   € Ø"ˆŒØ59ˆÔØ6:ˆ�rV   c                 óÒ  — t         j                  sJ ‚t         j                  j                  sJ ‚t         j                  j                  | _        | j                  €J ‚t        j                  | j                  j
                  v r-| j                  j
                  t        j                     | _        nt        «       | _        | j                  €J ‚| j                  | j                  _        | S r  )	r4   Úinterpreterrf  rA   ÚkeyÚmetarg  rd  Úops_namer3  s    rW   Ú	__enter__z#RecordOptimizationContext.__enter__n  s­   € Ü�}Š}Ðˆ}Ü�}‰}×)Ò)Ð)Ð)äŸM™M×6Ñ6ˆÔØ× Ñ Ð,Ð,Ð,Ü×"Ñ" d×&7Ñ&7×&<Ñ&<Ñ<Ø×,Ñ,×1Ñ1Ô2E×2IÑ2IÑJˆD�Lä.Ó0ˆDŒLØ�|‰|Ð'Ð'Ð'Ø $§¡ˆ�‰ÔØˆrV   c                 óž   — | j                   sJ ‚| j                  sJ ‚| j                  | j                   j                  t        j                  <   y r  )rf  rg  rk  rA   rj  ©r.  Úexc_typeÚexc_valÚexc_tbs       rW   Ú__exit__z"RecordOptimizationContext.__exit__|  s>   € Ø× Ò Ð Ð Ø�|Š|Ðˆ|Ø:>¿,¹,ˆ×Ñ×ÑÔ2×6Ñ6Ò7rV   c                 ó   — | j                   S r  )rg  r3  s    rW   Úget_opt_ctxz%RecordOptimizationContext.get_opt_ctx�  s   € Ø�|‰|ÐrV   c                 ó6   — | j                   sJ ‚| j                   S r  )rf  r3  s    rW   Úget_fx_nodez%RecordOptimizationContext.get_fx_node„  s   € Ø× Ò Ð Ð Ø× Ñ Ð rV   N)rR   )	r  r  r  ro   r-  rm  rs  ru  rw  rU   rV   rW   rc  rc  h  s#   „ ñ; #ó ;ò
òGò
ó!rV   rc  c                   ó  — e Zd ZdZed„ «       Zed„ «       Zed„ «       ZedLd„«       Zed„ «       Z	ed„ «       Z
ed	„ «       Zed
„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Z ed„ «       Z!ed „ «       Z"ed!„ «       Z#ed"„ «       Z$ed#„ «       Z%ed$„ «       Z&ed%„ «       Z'ed&„ «       Z(ed'„ «       Z)ed(„ «       Z*ed)„ «       Z+ed*„ «       Z,ed+„ «       Z-ed,„ «       Z.ed-„ «       Z/ed.„ «       Z0ed/„ «       Z1ed0„ «       Z2ed1„ «       Z3ed2„ «       Z4ed3„ «       Z5ed4„ «       Z6ed5„ «       Z7ed6„ «       Z8ed7„ «       Z9ed8„ «       Z:ed9„ «       Z;ed:„ «       Z<ed;„ «       Z=ed<„ «       Z>ed=„ «       Z?ed>„ «       Z@ed?„ «       ZAed@„ «       ZBedA„ «       ZCedB„ «       ZDedC„ «       ZEedD„ «       ZFedEeGj�                  dFeGj�                  fdG„«       ZIedEeGj�                  dFeGj�                  fdH„«       ZJedEeGj�                  dFeGj�                  fdI„«       ZKedJ„ «       ZLedK„ «       ZMy)MÚCppOverrideszMap element-wise ops to C++c                 ó   — d| › d| › d|› d�S )Nú	decltype(ú)(ú + rš   rU   ©ÚaÚbs     rW   ÚaddzCppOverrides.addŒ  ó   € à˜1˜#˜R ˜s # a S¨Ð*Ð*rV   c                 ó   — d| › d| › d|› d�S )Nr{  r|  ú - rš   rU   r~  s     rW   rÇ   zCppOverrides.sub�  r‚  rV   c                 ó   — d| › d| › d|› d�S )Nr{  r|  r–   rš   rU   r~  s     rW   rL  zCppOverrides.mul”  r‚  rV   Nc                 ó¶  — t        | t        «      sJ ‚|€| j                  }t        j                  j                  | ||«      }t        j                  j                  j                  t        j                  j                  |«      }|j                  d| |fd|i«       |t        v r6|t        j                  k(  r#	 t        j                  j                  | |||«       |S )NÚto_dtyper§   )r    rG   rˆ   r4   rZ  Úget_to_dtype_exprÚcseÚgenerateÚcomputeÚupdate_on_argsrƒ   r„   rn   Úcache_dtype_convert)Úxrˆ   r§   Úuse_compute_typesÚexprÚcsevars         rW   r‡  zCppOverrides.to_dtype˜  s¯   € ä˜!œ^Ô,Ð,Ð,ØÐØŸ™ˆIÜ�x‰x×)Ñ)¨!¨U°IÓ>ˆÜ—‘—‘×&Ñ&¤q§x¡x×'7Ñ'7¸Ó>ˆØ×Ñ˜j¨1¨e¨*°{ÀIÐ6NÔOØ”MÑ! i´5·;±;Ò&>ðô> �H‰H×(Ñ(¨¨I°v¸uÔEØˆrV   c                 óT   — |t         v sJ |› dt        › d�«       ‚dt         |   › d| › d�S )Nz missing from z.DTYPE_TO_CPPzc10::bit_cast<ú>(rš   )rH   r  )rŽ  rˆ   r§   s      rW   Úto_dtype_bitcastzCppOverrides.to_dtype_bitcastÃ  s=   € àœÑ$ÐU¨¨¨~¼h¸ZÀ}Ð&UÓUÐ$Ø¤¨UÑ 3Ð4°B°q°c¸Ð;Ð;rV   c                 ó   — d| › d�S )Nz	std::abs(rš   rU   ©rŽ  s    rW   ÚabszCppOverrides.absÈ  ó   € à˜1˜#˜QÐÐrV   c                 ó   — d| › d�S )Nz	std::sin(rš   rU   r–  s    rW   ÚsinzCppOverrides.sinÌ  r˜  rV   c                 ó   — d| › d�S )Nz	std::cos(rš   rU   r–  s    rW   ÚcoszCppOverrides.cosÐ  r˜  rV   c                 ó   — d| › d| › d�S )Nr{  z)(-rš   rU   r–  s    rW   ÚnegzCppOverrides.negÔ  ó   € à˜1˜#˜S   1Ð%Ð%rV   c                 ó   — d| › d�S )Nz	std::exp(rš   rU   r–  s    rW   ÚexpzCppOverrides.expØ  s   € ð ˜1˜#˜QÐÐrV   c                 ó   — d| › d�S )Nz
std::exp2(rš   rU   r–  s    rW   Úexp2zCppOverrides.exp2Ý  ó   € à˜A˜3˜aÐ Ð rV   c                 ó   — d| › d�S )Nzstd::expm1(rš   rU   r–  s    rW   Úexpm1zCppOverrides.expm1á  ó   € à˜Q˜C˜qÐ!Ð!rV   c                 ó   — d| › d�S )Nz	std::erf(rš   rU   r–  s    rW   ÚerfzCppOverrides.erfå  r˜  rV   c                 ó   — d| › d�S )Nz
std::erfc(rš   rU   r–  s    rW   ÚerfczCppOverrides.erfcé  r¤  rV   c                 ó   — d| › d�S )Nzcalc_erfinv(rš   rU   r–  s    rW   ÚerfinvzCppOverrides.erfinví  ó   € à˜a˜S Ð"Ð"rV   c                 ó   — d| › d�S )Nz
std::sqrt(rš   rU   r–  s    rW   ÚsqrtzCppOverrides.sqrtñ  r¤  rV   c                 ó   — d| › d�S )Nz1 / std::sqrt(rš   rU   r–  s    rW   ÚrsqrtzCppOverrides.rsqrtõ  s   € à ˜s !Ð$Ð$rV   c                 ó|   — t         j                  j                  }|dk(  r| › d| › d�S |€d| › d�S t        d|›�«      ‚)NÚaccuracyú + decltype(ú)(1)zstd::log1p(rš   ú8unrecognized config cpp.inject_log1p_bug_TESTING_ONLY = ©r   ÚcppÚinject_log1p_bug_TESTING_ONLYr†   ©rŽ  Úbugs     rW   Úlog1pzCppOverrides.log1pù  sW   € ä�j‰j×6Ñ6ˆØ�*ÒØ�S˜ Q C tÐ,Ð,Øˆ[Ø    1Ð%Ð%ä ØJÈ3È'ÐRóð rV   c                 ó   — d| › d�S )Nz	std::tan(rš   rU   r–  s    rW   ÚtanzCppOverrides.tan  r˜  rV   c                 ó   — d| › d�S )Nz
std::tanh(rš   rU   r–  s    rW   ÚtanhzCppOverrides.tanh	  r¤  rV   c                 ó&   — t         rd| › d�S d| › d�S )z§
        On windows std::signbit only support float type.
        Ref: https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/signbit?view=msvc-170
        z std::signbit(static_cast<float>(rŸ   zstd::signbit(rš   rS   r–  s    rW   ÚsignbitzCppOverrides.signbit  s-   € õ ð /¨q¨c°Ð4ð	
ð !   1Ð%ð	
rV   c                 ó   — d| › d|› d�S )Nz	std::pow(r™   rš   rU   r~  s     rW   ÚpowzCppOverrides.pow  s   € à˜1˜#˜R ˜s !Ð$Ð$rV   c                 ó   — d| › d�S )Nz	std::log(rš   rU   r–  s    rW   ÚlogzCppOverrides.log  r˜  rV   c                 ó   — d| › d�S )Nzstd::nearbyint(rš   rU   r–  s    rW   ÚroundzCppOverrides.round!  s   € à    1Ð%Ð%rV   c                 ó   — d| › d�S )Nzstd::floor(rš   rU   r–  s    rW   ÚfloorzCppOverrides.floor%  r§  rV   c                 óH   — | › d|› �}| › d|› �}d| › d|› d|› d|› d|› d|› d	�S )
Nú / ú % ú((z
 < 0) != (z	 < 0) ? (z != 0 ? z - 1 : z) : rš   rU   )r  r€  ÚquotÚrems       rW   ÚfloordivzCppOverrides.floordiv)  sR   € ð ��C˜�sˆ|ˆØ��3�q�cˆlˆØ�A�3�j   9¨S¨E°¸$¸¸wÀtÀfÈDÐQUÐPVÐVWÐXÐXrV   c                 ó   — d| › d�S )Nz
std::ceil(rš   rU   r–  s    rW   ÚceilzCppOverrides.ceil0  r¤  rV   c                 ó   — d| › d�S )Nzstd::trunc(rš   rU   r–  s    rW   ÚtrunczCppOverrides.trunc4  r§  rV   c                 ó   — | › d|› �S ©NrÍ  rU   r~  s     rW   ÚtruncdivzCppOverrides.truncdiv8  s   € ð ��C˜�sˆ|ÐrV   c                 ó   — d| › d|› d�S )Nz
std::fmod(r™   rš   rU   r~  s     rW   ÚfmodzCppOverrides.fmod=  s   € à˜A˜3˜b   1Ð%Ð%rV   c                 ó   — d| › d�S )Nzstd::isinf(rš   rU   r–  s    rW   ÚisinfzCppOverrides.isinfA  r§  rV   c                 ó   — d| › d�S )Nzstd::isnan(rš   rU   r–  s    rW   ÚisnanzCppOverrides.isnanE  r§  rV   c                 ó   — d| › d�S )Nzstd::lgamma(rš   rU   r–  s    rW   ÚlgammazCppOverrides.lgammaI  r®  rV   c                 ó   — d| › d�S )Nz
std::acos(rš   rU   r–  s    rW   ÚacoszCppOverrides.acosM  r¤  rV   c                 ó   — d| › d�S )Nzstd::acosh(rš   rU   r–  s    rW   ÚacoshzCppOverrides.acoshQ  r§  rV   c                 ó   — d| › d�S )Nz
std::cosh(rš   rU   r–  s    rW   ÚcoshzCppOverrides.coshU  r¤  rV   c                 ó   — d| › d�S )Nz
std::sinh(rš   rU   r–  s    rW   ÚsinhzCppOverrides.sinhY  r¤  rV   c                 ó   — d| › d�S )Nz
std::asin(rš   rU   r–  s    rW   ÚasinzCppOverrides.asin]  r¤  rV   c                 ó   — d| › d�S )Nzstd::asinh(rš   rU   r–  s    rW   ÚasinhzCppOverrides.asinha  r§  rV   c                 ó   — d| › d|› d�S )Nzstd::atan2(r™   rš   rU   ©rŽ  Úys     rW   Úatan2zCppOverrides.atan2e  ó   € à˜Q˜C˜r !  AÐ&Ð&rV   c                 ó   — d| › d�S )Nz
std::atan(rš   rU   r–  s    rW   ÚatanzCppOverrides.atani  r¤  rV   c                 ó   — d| › d�S )Nzstd::atanh(rš   rU   r–  s    rW   ÚatanhzCppOverrides.atanhm  r§  rV   c                 ó   — d| › d|› d�S )Nzstd::copysign(r™   rš   rU   rï  s     rW   ÚcopysignzCppOverrides.copysignq  s   € à ˜s " Q C qÐ)Ð)rV   c           	      óŒ  — d| › d�d| › d�f}t        d„ |D «       «      rt        d„ |D «       «      S t        «       }t        j                  j
                  j                  t        j                  ¬«      }t        j                  j
                  j                  | j                  ¬«      }|j                  d|› d�«       |j                  d	|› d
| › d|› d�«       t        j                  j                  j                  |«       ||f}t        ||«      D ]/  \  }}t        j                  j
                  j                  ||«       Œ1 ||fS )Núfrexp(ú)[0]ú)[1]c              3   ór   K  — | ]/  }t         j                  j                  j                  |«      d u–— Œ1 y ­wr  ©r4   rZ  r‰  Útry_get©r  Ú	cache_keys     rW   r  z%CppOverrides.frexp.<locals>.<genexpr>x  ó(   è ø€ ÒW¸yŒq�x‰x�|‰|×#Ñ# IÓ.°dÔ:ÑWùó   ‚57c              3   ón   K  — | ]-  }t         j                  j                  j                  |«      –— Œ/ y ­wr  rþ  r   s     rW   r  z%CppOverrides.frexp.<locals>.<genexpr>y  ó#   è ø€ ÒU¸YœŸ™Ÿ™×-Ñ-¨i×8ÑUùó   ‚35©rˆ   zint32_t r¸   rÒ   z = std::frexp(ú, &rÓ   )r"  r¡   r7   r4   rZ  r‰  Únewvarr„   Úint32rˆ   r¿   r‹  ÚspliceÚzipÚput)rŽ  Ú
cache_keysr°   ÚexponentÚmantissaÚcse_varsr  Úcse_vars           rW   ÚfrexpzCppOverrides.frexpu  s  € à˜a˜S Ð%¨°¨s°$Ð'7Ð7ˆ
ÜÑWÈJÔWÔWÜÑUÈ*ÔUÓUÐUä‹~ˆÜ—8‘8—<‘<×&Ñ&¬U¯[©[Ð&Ó9ˆÜ—8‘8—<‘<×&Ñ&¨Q¯W©WÐ&Ó5ˆØ�‰˜ ( ¨1Ð-Ô.Ø�‰˜˜x˜j¨°q°c¸¸X¸JÀbÐIÔJÜ	�‰×Ñ×Ñ Ô%Ø˜hÐ'ˆÜ"% j°(Ó";ò 	1ÑˆI�wÜ�H‰H�L‰L×Ñ˜Y¨Õ0ð	1à˜Ð!Ð!rV   c                 ó   — d| › d|› d�S )Nzstd::hypot(r™   rš   rU   rï  s     rW   ÚhypotzCppOverrides.hypot†  rò  rV   c                 ó   — d| › d�S )Nzstd::log10(rš   rU   r–  s    rW   Úlog10zCppOverrides.log10Š  r§  rV   c                 ó   — d| › d�S )Nz
std::log2(rš   rU   r–  s    rW   Úlog2zCppOverrides.log2Ž  r¤  rV   c                 ó   — d| › d|› d�S )Nzstd::nextafter(r™   rš   rU   rï  s     rW   Ú	nextafterzCppOverrides.nextafter’  s   € à    2 a S¨Ð*Ð*rV   c                 ó¢   — t         j                  j                  }|dk(  ry|dk(  r| › d�S |dk(  r| › d| › d�S |€	d| › d	| › d
�S t        d|›�«      ‚)NÚcompile_errorúcompile error!Úruntime_errorú	; throw 1r´  rµ  r¶  z	std::max(ú, decltype(ú)(0))ú7unrecognized config cpp.inject_relu_bug_TESTING_ONLY = ©r   r¹  Úinject_relu_bug_TESTING_ONLYr†   r»  s     rW   ÚreluzCppOverrides.relu–  s|   € ä�j‰j×5Ñ5ˆØ�/Ò!Ø#Ø�OÒ#Ø�S˜	�?Ð"Ø�JÒØ�S˜ Q C tÐ,Ð,Øˆ[Ø˜q˜c ¨Q¨C¨uÐ5Ð5ä ØIÈ#ÈÐQóð rV   c                 ó   — d| › d|› d�S )Nzmin_propagate_nan(r™   rš   rU   r~  s     rW   ÚminimumzCppOverrides.minimum¦  ó   € à# A 3 b¨¨¨1Ð-Ð-rV   c                 ó   — d| › d|› d�S )Nzmax_propagate_nan(r™   rš   rU   r~  s     rW   ÚmaximumzCppOverrides.maximumª  r)  rV   c                 ó   — | › d|› d|› �S )Nú ? ú : rU   )r  r€  Úcs      rW   ÚwherezCppOverrides.where®  s   € à��C˜�s˜#˜a˜SÐ!Ð!rV   c                 ó   — d| › d|› d�S )Nzmod(r™   rš   rU   r~  s     rW   r  zCppOverrides.mod²  s   € à�a�S˜˜1˜#˜QÐÐrV   c                 ó(   — t        | t        |   «      S r  )rO   rH   )Úvalrˆ   s     rW   ÚconstantzCppOverrides.constant¶  s   € ä˜C¤¨eÑ!4Ó5Ð5rV   c                 ó  — t        t        j                  j                  | «      «      }t        j                  j                  j                  t        j                  j                  |t        | «      ¬«      }t        j                  ||«      S )N©Úbounds)
rD   r4   rZ  Úrename_indexingr‰  rŠ  r‹  r&   r2   r‡  )r�  rˆ   Úidx_strr¥   s       rW   Ú
index_exprzCppOverrides.index_exprº  sb   € äœŸ™×0Ñ0°Ó6Ó7ˆÜ�h‰h�l‰l×#Ñ#Ü�H‰H×Ñ˜gÔ.CÀDÓ.Ið $ó 
ˆô �|‰|˜C Ó'Ð'rV   c                 ó"  — t        «       }t        j                  j                  j	                  «       }|j                  d|› d�«       t        j                  j                  |«      5  |j                  «       5   |«       }|j                  d|› d�«       d d d «       d d d «       |j                  d«       t        j                  j                  j                  |«       t        |d|› d�«      }| › d|› d|› �S # 1 sw Y   ŒexY w# 1 sw Y   ŒixY w)	NrÒ   ú = [&]úreturn r¸   r{  z())r-  z() : )r7   r4   rZ  r‰  r	  r¿   Úswap_buffersrÁ   r‹  r  rO   )ÚmaskÚbodyÚotherr°   Úbody_varrö   Ú
other_codes          rW   ÚmaskedzCppOverrides.maskedÂ  sê   € ä‹~ˆô —8‘8—<‘<×&Ñ&Ó(ˆØ�‰˜˜x˜j¨Ð/Ô0Ü�X‰X×"Ñ" 4Ó(ñ 	0¨$¯+©+«-ñ 	0Ù“VˆFØ�N‰N˜W V H¨AÐ.Ô/÷	0÷ 	0ð 	�‰�sÔÜ	�‰×Ñ×Ñ Ô%ô " %¨9°X°J¸cÐ)BÓCˆ
Ø��s˜8˜* E¨*¨Ð6Ð6÷	0ð 	0ú÷ 	0ð 	0ús$   Á'DÁ8C9ÂDÃ9D	Ã>DÄDc                 ó   — | › d|› �S )Nú && rU   r~  s     rW   Úlogical_andzCppOverrides.logical_andÓ  ó   € à��D˜˜ˆ}ÐrV   c                 ó   — d| › �S )Nú!rU   ©r  s    rW   Úlogical_notzCppOverrides.logical_not×  ó   € à�1�#ˆwˆrV   c                 ó   — | › d|› �S )Nr˜   rU   r~  s     rW   Ú
logical_orzCppOverrides.logical_orÛ  rH  rV   c                 ó   — | › d|› �S )Nú != rU   r~  s     rW   Úlogical_xorzCppOverrides.logical_xorß  rH  rV   c                 ó   — d| › d| › d|› d�S )Nr{  r|  ú & rš   rU   r~  s     rW   Úbitwise_andzCppOverrides.bitwise_andã  r‚  rV   c                 ó   — d| › d| › d�S )Nr{  z)(~rš   rU   rK  s    rW   Úbitwise_notzCppOverrides.bitwise_notç  rŸ  rV   c                 ó   — d| › d| › d|› d�S )Nr{  r|  ú | rš   rU   r~  s     rW   Ú
bitwise_orzCppOverrides.bitwise_orë  r‚  rV   c                 ó   — d| › d| › d|› d�S )Nr{  r|  r—   rš   rU   r~  s     rW   Úbitwise_xorzCppOverrides.bitwise_xorï  r‚  rV   c                 óä  — t        «       }|j                  d«       |j                  «       5  t        | j                     }|j                  d|› d|› d�«       |j                  d|› d|› d|› d�«       |j                  «       5  |j                  d	| › d
�«       d d d «       |j                  d	| › d|› d| › d|› d�	«       d d d «       |j                  d«       |S # 1 sw Y   ŒBxY w# 1 sw Y   Œ(xY w)Nú[&]()úconstexpr decltype(ú) max_shift = sizeof(z) * CHAR_BIT;ú$if ((static_cast<std::make_signed_t<ú>>(ú) < 0) || (ú >= max_shift))úreturn decltype(z)(0);z#)(static_cast<std::make_unsigned_t<z) << rÓ   ú()©r7   r¿   rÁ   rH   rˆ   ©r  r€  r°   Úscalar_ts       rW   Úbitwise_left_shiftzCppOverrides.bitwise_left_shiftó  s  € ä‹~ˆØ�‰�wÔØ�[‰[‹]ñ 	Ü# A§G¡GÑ,ˆHØ�N‰NØ% a SÐ(=¸h¸ZÀ}ÐUôð �N‰NØ6°x°jÀÀAÀ3ÀkÐRSÐQTÐTcÐdôð —‘“ñ <Ø—‘Ð!1°!°°EÐ:Ô;÷<à�N‰NØ" 1 #Ð%HÈÈ
ÐRUÐVWÐUXÐX]Ð^_Ð]`Ð`bÐcô÷	ð 	�‰�tÔØˆ÷<ð <ú÷	ð 	ús$   ¬AC&ÂCÂ&C&ÃC#	ÃC&Ã&C/c           
      óê  — t        «       }|j                  d«       |j                  «       5  t        | j                     }|j                  d|› d|› d|› d�«       |j                  d|› d|› d|› d	�«       |j                  «       5  |j                  d
| › d| › d�«       d d d «       |j                  d
| › d| › d|› d�«       d d d «       |j                  d«       |S # 1 sw Y   Œ?xY w# 1 sw Y   Œ(xY w)Nr^  r_  r`  z ) * CHAR_BIT - std::is_signed_v<z>;ra  rb  rc  rd  re  r|  z >> max_shift);ú >> rÓ   rf  rg  rh  s       rW   Úbitwise_right_shiftz CppOverrides.bitwise_right_shift  s  € ä‹~ˆØ�‰�wÔØ�[‰[‹]ñ 
	AÜ# A§G¡GÑ,ˆHØ�N‰NØ% a SÐ(=¸h¸ZÐGgÐhpÐgqÐqsÐtôð �N‰NØ6°x°jÀÀAÀ3ÀkÐRSÐQTÐTcÐdôð —‘“ñ KØ—‘Ð!1°!°°B°q°c¸ÐIÔJ÷Kà�N‰NÐ-¨a¨S°°1°#°T¸!¸¸BÐ?Ô@÷
	Að 	�‰�tÔØˆ÷	Kð Kú÷
	Að 
	Aús$   ¬AC)ÂCÂ#C)ÃC&	Ã"C)Ã)C2Úseedr;  c                 ó   — d| › d|› d�S )Nznormalized_rand_cpu(r™   rš   rU   ©rn  r;  s     rW   ÚrandzCppOverrides.rand  s   € à% d V¨2¨f¨X°QÐ7Ð7rV   c                 ó   — d| › d|› d�S )Nz
randn_cpu(r™   rš   rU   rp  s     rW   ÚrandnzCppOverrides.randn  s   € à˜D˜6  F 8¨1Ð-Ð-rV   c           	      ó    — d| › d|› d|› d|› d�	S )Nzrandint64_cpu(r™   rš   rU   )rn  r;  ÚlowÚhighs       rW   Ú	randint64zCppOverrides.randint64!  s#   € à ˜v R¨ x¨r°#°°b¸¸¸aÐ@Ð@rV   c                 ó   — d| › d| › d| › d�S )Nr{  z)(1) / (decltype(z)(1) + std::exp(-rŸ   rU   r–  s    rW   ÚsigmoidzCppOverrides.sigmoid%  s   € à˜1˜#Ð.¨q¨cÐ1BÀ1À#ÀRÐHÐHrV   c           
      óN  — t        «       }d| › d�}d| › d�}|j                  d«       |j                  «       5  |j                  d| › d|› d|› d�«       |j                  d	| › d
|› d|› d�«       |j                  d«       d d d «       |j                  d«       |S # 1 sw Y   ŒxY w)Nr{  ú)(0)r¶  r^  úauto left = z > 0 ? r.  r¸   úauto right = z < 0 ? úreturn left - right;rf  ©r7   r¿   rÁ   )rŽ  r°   Úscalar_zeroÚ
scalar_ones       rW   ÚsignzCppOverrides.sign)  s²   € ä‹~ˆØ! !  DÐ)ˆØ    4Ð(ˆ
Ø�‰�wÔØ�[‰[‹]ñ 	3Ø�N‰N˜\¨!¨¨G°J°<¸sÀ;À-ÈqÐQÔRØ�N‰N˜]¨1¨#¨W°Z°LÀÀKÀ=ÐPQÐRÔSØ�N‰NÐ1Ô2÷	3ð 	�‰�tÔØˆ÷	3ð 	3ús   ¸ABÂB$©NT)Nr  r  r  r  Ústaticmethodr�  rÇ   rL  r‡  r”  r—  rš  rœ  rž  r¡  r£  r¦  r©  r«  r­  r°  r²  r½  r¿  rÁ  rÃ  rÅ  rÇ  rÉ  rË  rÒ  rÔ  rÖ  rÙ  rÛ  rÝ  rß  rá  rã  rå  rç  ré  rë  rí  rñ  rô  rö  rø  r  r  r  r  r  r&  r(  r+  r0  r  r4  r:  rD  rG  rL  rO  rR  rU  rW  rZ  r\  rj  rm  rè   ÚExprrq  rs  rw  ry  r‚  rU   rV   rW   ry  ry  ‰  sÚ  „ Ù%àñ+ó ð+ð ñ+ó ð+ð ñ+ó ð+ð ò(ó ð(ðT ñ<ó ð<ð ñ ó ð ð ñ ó ð ð ñ ó ð ð ñ&ó ð&ð ñ ó ð ð ñ!ó ð!ð ñ"ó ð"ð ñ ó ð ð ñ!ó ð!ð ñ#ó ð#ð ñ!ó ð!ð ñ%ó ð%ð ñ	ó ð	ð ñ ó ð ð ñ!ó ð!ð ñ	
ó ð	
ð ñ%ó ð%ð ñ ó ð ð ñ&ó ð&ð ñ"ó ð"ð ñYó ðYð ñ!ó ð!ð ñ"ó ð"ð ñó ðð ñ&ó ð&ð ñ"ó ð"ð ñ"ó ð"ð ñ#ó ð#ð ñ!ó ð!ð ñ"ó ð"ð ñ!ó ð!ð ñ!ó ð!ð ñ!ó ð!ð ñ"ó ð"ð ñ'ó ð'ð ñ!ó ð!ð ñ"ó ð"ð ñ*ó ð*ð ñ"ó ð"ð  ñ'ó ð'ð ñ"ó ð"ð ñ!ó ð!ð ñ+ó ð+ð ñó ðð ñ.ó ð.ð ñ.ó ð.ð ñ"ó ð"ð ñ ó ð ð ñ6ó ð6ð ñ(ó ð(ð ñ7ó ð7ð  ñó ðð ñó ðð ñó ðð ñó ðð ñ+ó ð+ð ñ&ó ð&ð ñ+ó ð+ð ñ+ó ð+ð ñó ðð& ñó ðð" ð8�5—:‘:ð 8 u§z¡zò 8ó ð8ð ð.�E—J‘Jð .¨¯
©
ò .ó ð.ð ðA˜Ÿ
™
ð A¨E¯J©Jò Aó ðAð ñIó ðIð ñ
ó ñ
rV   ry  r¹  c                   ó  ‡ — e Zd ZdZˆ fd„Zed„ «       Zed„ «       Zed„ «       Zed„ «       Z	ed„ «       Z
ed„ «       Zed	„ «       Zed
„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Zed„ «       Z ed„ «       Z!ed„ «       Z"ed „ «       Z#ed!„ «       Z$ed"„ «       Z%ed#„ «       Z&ed$„ «       Z'ed%„ «       Z(ed&„ «       Z)ed'„ «       Z*ed(„ «       Z+ed)„ «       Z,ed*„ «       Z-ed+„ «       Z.ed,„ «       Z/ed-„ «       Z0ed.„ «       Z1ed/„ «       Z2ed0„ «       Z3ed1„ «       Z4ed2„ «       Z5ed3„ «       Z6ed4„ «       Z7ed5„ «       Z8ed6„ «       Z9ed7„ «       Z:ed8„ «       Z;ed9„ «       Z<ed:„ «       Z=ed;„ «       Z>ed<„ «       Z?ed=„ «       Z@ed>„ «       ZAed?„ «       ZBed@„ «       ZCedA„ «       ZDedB„ «       ZEedC„ «       ZFedD„ «       ZGedE„ «       ZHedF„ «       ZIedG„ «       ZJedH„ «       ZKedI„ «       ZLedJ„ «       ZMedRdK„«       ZNedL„ «       ZOedM„ «       ZPedN„ «       ZQedO„ «       ZReSdP„ «       ZTeSdQ„ «       ZUˆ xZVS )SÚCppVecOverridesz.Map element-wise ops to aten vectorization C++c                 óì   •‡— t         ‰| �  | «      Šˆˆfd„}t        t        «      j	                  «       D ]<  \  }}t        |dd «      t        k(  sŒ|dvsŒ t        ‰| ||j                  «      «       Œ> ‰S )Nc                 ó   •‡ — ˆˆ ˆfd„}|S )Nc                  óð  •— | D �cg c]@  }t        |t        t        j                  f«      st        |t        «      r|j
                  s|‘ŒB }}| D �cg c]   }t        |t        «      r|j
                  r|‘Œ" }}t        | «      }|rÉ|rÇg }| D ]À  }t        |t        t        j                  f«      r�t        |t        j                  «      r1|j                  s%t        j                  |t        j                  «      }n$t        j                  |t        j                  «      }t        |t        «      r|j                  n|}|j                  |«       ŒÂ |r>t!        |«      dk(  rt#        |«      }n$‰
t$        j&                  k(  rt#        |dd  «      |dd  |r�|r›t        t(        j*                  t,        «      sJ ‚|D �cg c]p  }t        |t        «      r\|j
                  sP‰
t$        j.                  t$        j0                  t$        j2                  fvrt(        j*                  j5                  |«      n|‘Œr }}|r ‰
|i |¤ŽS t7        t$        ‰«      }t9        |‰
j:                  «      }|€J ‚ || i |¤ŽS c c}w c c}w c c}w )Nr   r5   )r    rm   rè   r…  rG   r¤   r#  Ú	is_numberr2   r:  r„   Úint64r4  r3   ÚvalueÚappendrÐ   rL   r‡  r0  r4   rZ  ÚCppVecKernelrq  rs  rw  Ú	broadcastr,  r>  r  )ÚargsÚkwargsÚargÚscalarsÚvectorsÚnew_argsÚnew_argÚ
scalar_opsÚscalar_funcr1  Úfuncr.  s            €€€rW   Úwrapperz6CppVecOverrides.__new__.<locals>.wrap.<locals>.wrapperM  s8  ø€ ð  $öàÜ! #¬¬U¯Z©ZÐ'8Ô9Ü" 3¬Ô7ÀÇ
Â
ò ð�ð ð  $öàÜ! #¤~Ô6¸3¿:º:ò ð�ð ô
   ›:�Ù™wØ!�HØ#ò -˜Ü% c¬C´·±Ð+<Ô=Ü)¨#¬u¯z©zÔ:À3Ç=Â=Ü&)§n¡n°S¼%¿+¹+Ó&F¡ä&)§l¡l°3¼¿¹Ó&D Ü/9¸#¼xÔ/H #§)¢)Èc˜CØ Ÿ™¨Õ,ð-ñ ô ˜8“}¨Ò)Ü#/°Ó#9™Ø¤×!6Ñ!6Ò6Ü'3°H¸Q¸R°LÓ'A˜  ˜ñ ™wÜ%¤a§h¡h´Ô=Ð=Ð=ð  (0ö ð $ô !+¨7´NÔ CØ(/¯ªØ$(ä$3×$8Ñ$8Ü$3×$9Ñ$9Ü$3×$=Ñ$=ð("ñ%"ô	 ŸH™H×.Ñ.¨wÔ7ð ")ñ)ð �Hð  ñ$ Ù Ð4¨VÑ4Ð4ô "'¤¸Ó!=�JÜ")¨*°d·m±mÓ"D�KØ&Ð2Ð2Ð2Ù&¨Ð7°Ñ7Ð7ùòùòùò@ s   †AI)Á%I.Æ6A5I3rU   )rš  r›  r1  r.  s   ` €€rW   Úwrapz%CppVecOverrides.__new__.<locals>.wrap@  s   ù€ ö@8ðD ˆNrV   r1  )rD  r:  )	r,  Ú__new__Úvarsr‡  Úitemsr>  r„  ÚsetattrÚ__func__)r%  r‘  Úkargsrœ  rÄ   Úmethodr.  r1  s         @€rW   r�  zCppVecOverrides.__new__=  sx   ù€ Ü‰w‰˜sÓ#ˆõO	ôb !¤Ó1×7Ñ7Ó9ò 	;‰LˆD�&Ü�v˜{¨DÓ1´\ÓAÀdð Sò Gô ˜˜d¡D¨¯©Ó$9Õ:ð	;ð ˆrV   c                 ó   — | › d|› �S )Nr}  rU   r~  s     rW   r�  zCppVecOverrides.addš  ó   € à��C˜�sˆ|ÐrV   c                 ó   — | › d|› �S )Nr„  rU   r~  s     rW   rÇ   zCppVecOverrides.subž  r¥  rV   c                 ó   — | › d|› �S ©Nr–   rU   r~  s     rW   rL  zCppVecOverrides.mul¢  r¥  rV   c                 ó   — | › d|› �S rØ  rU   r~  s     rW   ÚtruedivzCppVecOverrides.truediv¦  r¥  rV   c                 ó   — | › d�S )Nz.abs()rU   r–  s    rW   r—  zCppVecOverrides.absª  ó   € à��Fˆ|ÐrV   c                 ó   — | › d�S )Nz.sin()rU   r–  s    rW   rš  zCppVecOverrides.sin®  r¬  rV   c                 ó   — | › d�S )Nz.cos()rU   r–  s    rW   rœ  zCppVecOverrides.cos²  r¬  rV   c                 ó   — | › d�S )Nz.exp()rU   r–  s    rW   r¡  zCppVecOverrides.exp¶  r¬  rV   c                 ó   — | › d�S )Nz.exp2()rU   r–  s    rW   r£  zCppVecOverrides.exp2º  ó   € à��Gˆ}ÐrV   c                 ó   — d| › d�}| › d|› �S )Nr{  r¶  z	.exp() - rU   )rŽ  Úvec_ones     rW   r¦  zCppVecOverrides.expm1¾  s#   € ð ˜a˜S Ð%ˆØ��I˜g˜YÐ'Ð'rV   c                 ó   — | › d�S )Nz.erf()rU   r–  s    rW   r©  zCppVecOverrides.erfÄ  r¬  rV   c                 ó   — | › d�S )Nz.erfc()rU   r–  s    rW   r«  zCppVecOverrides.erfcÈ  r±  rV   c                 ó   — | › d�S )Nz	.erfinv()rU   r–  s    rW   r­  zCppVecOverrides.erfinvÌ  ó   € à��IˆÐrV   c                 ó   — | › d�S )Nz.sqrt()rU   r–  s    rW   r°  zCppVecOverrides.sqrtÐ  r±  rV   c                 óæ   — t        t        j                  t        «      sJ ‚t        | t        «      sJ ‚| j
                  €J ‚t        j                  j                  | j
                  «      › d| › d|› d�S )Nú(ú == rš   ©r    r4   rZ  r�  rG   rˆ   Ú_get_mask_typerï  s     rW   ÚeqzCppVecOverrides.eqÔ  óc   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü˜!œ^Ô,Ð,Ð,Ø�w‰wÐ"Ð"Ð"Ü—(‘(×)Ñ)¨!¯'©'Ó2Ð3°1°Q°C°t¸A¸3¸aÐ@Ð@rV   c                 ó¾  — t        t        j                  t        «      sJ ‚t        | t        «      sJ ‚| j
                  t        j                  k(  rO|j
                  t        j                  k(  sJ ‚t        t        j                  j                  | |f«      \  }}|› d|› �S | j
                  €J ‚t        j                  j                  | j
                  «      › d| › d|› d�S )NrQ  rº  rš   )r    r4   rZ  r�  rG   rˆ   r„   rk   rN   r‹  r½  )rŽ  rð  Úx_castÚy_casts       rW   ÚnezCppVecOverrides.neÛ  s¸   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü˜!œ^Ô,Ð,Ð,Ø�7‰7”e—j‘jÒ Ø—7‘7œeŸj™jÒ(Ð(Ð(Ü1´!·(±(×2BÑ2BÀQÈÀFÓK‰NˆF�FØ�X˜T & Ð*Ð*à—7‘7Ð&Ð&Ð&Ü—h‘h×-Ñ-¨a¯g©gÓ6Ð7°q¸¸¸4À¸sÀ!ÐDÐDrV   c                 óæ   — t        t        j                  t        «      sJ ‚t        | t        «      sJ ‚| j
                  €J ‚t        j                  j                  | j
                  «      › d| › d|› d�S )Nrº  r¹   rš   r¼  rï  s     rW   ÚltzCppVecOverrides.ltç  óc   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü˜!œ^Ô,Ð,Ð,Ø�w‰wÐ"Ð"Ð"Ü—(‘(×)Ñ)¨!¯'©'Ó2Ð3°1°Q°C°s¸1¸#¸QÐ?Ð?rV   c                 óæ   — t        t        j                  t        «      sJ ‚t        | t        «      sJ ‚| j
                  €J ‚t        j                  j                  | j
                  «      › d| › d|› d�S )Nrº  z > rš   r¼  rï  s     rW   ÚgtzCppVecOverrides.gtî  rÆ  rV   c                 óæ   — t        t        j                  t        «      sJ ‚t        | t        «      sJ ‚| j
                  €J ‚t        j                  j                  | j
                  «      › d| › d|› d�S )Nrº  ú <= rš   r¼  rï  s     rW   ÚlezCppVecOverrides.leõ  r¿  rV   c                 óæ   — t        t        j                  t        «      sJ ‚t        | t        «      sJ ‚| j
                  €J ‚t        j                  j                  | j
                  «      › d| › d|› d�S )Nrº  ú >= rš   r¼  rï  s     rW   ÚgezCppVecOverrides.geü  r¿  rV   c                 ó   — | › d|› �S ©NrT  rU   rï  s     rW   Úand_zCppVecOverrides.and_  r¥  rV   c                 ó   — | › d�S )Nz.rsqrt()rU   r–  s    rW   r²  zCppVecOverrides.rsqrt  ó   € à��Hˆ~ÐrV   c                 ó   — | › d|› d�S )Nz.pow(rš   rU   r~  s     rW   rÅ  zCppVecOverrides.pow  s   € à��E˜!˜˜AˆÐrV   c                 ó   — | › d�S )Nz.log()rU   r–  s    rW   rÇ  zCppVecOverrides.log  r¬  rV   c                 ó   — | › d�S )Nz.round()rU   r–  s    rW   rÉ  zCppVecOverrides.round  rÓ  rV   c                 ó   — | › d�S )Nz.floor()rU   r–  s    rW   rË  zCppVecOverrides.floor  rÓ  rV   c                 ó   — | › d�S )Nz.ceil()rU   r–  s    rW   rÔ  zCppVecOverrides.ceil  r±  rV   c                 ó   — | › d�S )Nz.trunc()rU   r–  s    rW   rÖ  zCppVecOverrides.trunc  rÓ  rV   c                 ó   — | › d|› d�S )Nz.fmod(rš   rU   r~  s     rW   rÛ  zCppVecOverrides.fmod#  s   € à��F˜1˜#˜QÐÐrV   c                 ó   — | › d�S )Nz	.lgamma()rU   r–  s    rW   rá  zCppVecOverrides.lgamma'  r·  rV   c                 ó.   — t        | |«      \  } }| › d|› �S rÐ  ©rK   r~  s     rW   rG  zCppVecOverrides.logical_and+  ó"   € ä,¨Q°Ó2‰ˆˆ1Ø��C˜�sˆ|ÐrV   c                 ó   — d| › �S ©Nú~rU   rK  s    rW   rL  zCppVecOverrides.logical_not0  rM  rV   c                 ó.   — t        | |«      \  } }| › d|› �S ©NrY  rÝ  r~  s     rW   rO  zCppVecOverrides.logical_or4  rÞ  rV   c                 ó.   — t        | |«      \  } }| › d|› �S ©Nr—   rÝ  r~  s     rW   rR  zCppVecOverrides.logical_xor9  rÞ  rV   c                 ó.   — t        | |«      \  } }| › d|› �S rÐ  rÝ  r~  s     rW   rU  zCppVecOverrides.bitwise_and>  rÞ  rV   c                 ó   — d| › �S rà  rU   rK  s    rW   rW  zCppVecOverrides.bitwise_notC  rM  rV   c                 ó.   — t        | |«      \  } }| › d|› �S rã  rÝ  r~  s     rW   rZ  zCppVecOverrides.bitwise_orG  rÞ  rV   c                 ó.   — t        | |«      \  } }| › d|› �S rå  rÝ  r~  s     rW   r\  zCppVecOverrides.bitwise_xorL  rÞ  rV   c                 ó   — | › d|› �S )Nz << rU   r~  s     rW   rj  z"CppVecOverrides.bitwise_left_shiftQ  rH  rV   c                 ó   — | › d|› �S )Nrl  rU   r~  s     rW   rm  z#CppVecOverrides.bitwise_right_shiftU  rH  rV   c                 ó„   — t        t        j                  t        «      sJ ‚t        j                  j	                  | |«      › S r  )r    r4   rZ  r�  Úload)rÄ   r;  s     rW   Ú	load_seedzCppVecOverrides.load_seedY  s.   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü—(‘(—-‘-  fÓ-Ð.Ð/rV   c                 ó|   — t        t        j                  t        «      sJ ‚t	        «       }d| › d�}t        |||«      S )Nz)result[offset_idx] = normalized_rand_cpu(ú, offset[offset_idx]);©r    r4   rZ  r�  r7   rF   ©rn  r;  r°   Úrand_functions       rW   rq  zCppVecOverrides.rand^  s@   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü‹~ˆà7¸°vÐ=SÐTð 	ô ˜F D¨-Ó8Ð8rV   c                 ó|   — t        t        j                  t        «      sJ ‚t	        «       }d| › d�}t        |||«      S )Nzresult[offset_idx] = randn_cpu(rð  rñ  rò  s       rW   rs  zCppVecOverrides.randng  s;   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü‹~ˆØ9¸$¸Ð?UÐVˆÜ˜F D¨-Ó8Ð8rV   c                 ó¦   — t        t        j                  t        «      sJ ‚t	        «       }d| › d|› d|› d�}t        |||t        j                  «      S )Nz#result[offset_idx] = randint64_cpu(z, offset[offset_idx], r™   rÓ   )r    r4   rZ  r�  r7   rF   r„   rŒ  )rn  r;  ru  rv  r°   ró  s         rW   rw  zCppVecOverrides.randint64n  sS   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü‹~ˆØ=¸d¸VÐCYÐZ]ÐY^Ð^`ÐaeÐ`fÐfhÐiˆÜ˜F D¨-¼¿¹ÓEÐErV   c                 ó~   — | j                   |j                   k(  sJ d«       ‚| › dt        j                  | |«      › d|› �S )Nz;remainder vec implementation expect the same inputs' dtype.z - (z) * )rˆ   r‡  rÒ  r~  s     rW   Ú	remainderzCppVecOverrides.remainderu  sI   € à�w‰w˜!Ÿ'™'Ò!ð 	
ØIó	
Ð!ð ��Dœ×1Ñ1°!°QÓ7Ð8¸¸Q¸CÐ@Ð@rV   c                 ó   — | › d�S )Nz.tan()rU   rK  s    rW   r¿  zCppVecOverrides.tan|  r¬  rV   c                 ó   — | › d�S )Nz.tanh()rU   rK  s    rW   rÁ  zCppVecOverrides.tanh€  r±  rV   c                 ó   — | › d�S )Nz.reciprocal()rU   rK  s    rW   Ú
reciprocalzCppVecOverrides.reciprocal„  s   € à��MÐ"Ð"rV   c                 ó   — | › d�S )Nz.atan()rU   r–  s    rW   rô  zCppVecOverrides.atanˆ  r±  rV   c                 ó   — | › d�S )Nz.acos()rU   r–  s    rW   rã  zCppVecOverrides.acosŒ  r±  rV   c                 ó   — | › d�S )Nz.asin()rU   r–  s    rW   rë  zCppVecOverrides.asin�  r±  rV   c                 ó   — | › d�S )Nz.cosh()rU   r–  s    rW   rç  zCppVecOverrides.cosh”  r±  rV   c                 ó   — | › d�S )Nz.sinh()rU   r–  s    rW   ré  zCppVecOverrides.sinh˜  r±  rV   c                 ó   — | › d�S )Nz.log10()rU   r–  s    rW   r  zCppVecOverrides.log10œ  rÓ  rV   c                 ó   — | › d�S )Nz.log2()rU   r–  s    rW   r  zCppVecOverrides.log2   r±  rV   c                 ó   — | › d|› d�S )Nz.nextafter(rš   rU   rï  s     rW   r  zCppVecOverrides.nextafter¤  s   € à��K ˜s !Ð$Ð$rV   c                 ó   — | › d|› d�S )Nz
.copysign(rš   rU   r~  s     rW   rø  zCppVecOverrides.copysign¨  s   € à��J˜q˜c Ð#Ð#rV   c                 ó   — | › d|› d�S )Nz.atan2(rš   rU   r~  s     rW   rñ  zCppVecOverrides.atan2¬  ó   € à��G˜A˜3˜aÐ Ð rV   c                 ó   — | › d|› d�S )Nz.hypot(rš   rU   r~  s     rW   r  zCppVecOverrides.hypot°  r  rV   c           
      ó<   — d| › d�}d| › d�}|› d|› d| › d|› d| › d�
S )	Nr{  r¶  z)(0.5)z * ((r}  z)/(r„  z)).log()rU   )rŽ  r³  Úvec_one_halfs      rW   rö  zCppVecOverrides.atanh´  sE   € ð ˜a˜S Ð%ˆØ" 1 # VÐ,ˆØ�˜u W I¨S°°°3°w°i¸sÀ1À#ÀXÐNÐNrV   c                 ó   — | › d�S )Nz.asinh()rU   r–  s    rW   rí  zCppVecOverrides.asinh»  rÓ  rV   c                 ó   — | › d�S )Nz.acosh()rU   r–  s    rW   rå  zCppVecOverrides.acosh¿  rÓ  rV   c                 ó¢   — t         j                  j                  }|dk(  ry|dk(  r| › d�S |dk(  r| › d| › d�S |€	d| › d	| › d
�S t        d|›�«      ‚)Nr  r  r  r   r´  rµ  r¶  zat::vec::clamp_min(r!  r"  r#  r$  r»  s     rW   r&  zCppVecOverrides.reluÃ  s|   € ä�j‰j×5Ñ5ˆØ�/Ò!Ø#Ø�OÒ#Ø�S˜	�?Ð"Ø�JÒØ�S˜ Q C tÐ,Ð,Øˆ[Ø(¨¨¨;°q°c¸Ð?Ð?ä ØIÈ#ÈÐQóð rV   c                 ó   — d| › d| › d| › d�S )Nr{  z)(1)/(decltype(z)(1) + z.neg().exp())rU   r–  s    rW   ry  zCppVecOverrides.sigmoidÔ  s   € à˜1˜#˜_¨Q¨C¨w°q°c¸ÐGÐGrV   c                 ó   — | › d�S )Nz.neg()rU   r–  s    rW   rž  zCppVecOverrides.negØ  r¬  rV   c                 óÒ  — t        | j                  «      r)| j                  |j                  k(  sJ d«       ‚d| › d|› d�S t        d„ | |fD «       «      sJ ‚d| › d�}t        j                  j                  |j                  «      dk  r,|› ddt        j                  j                  z  dz
  › d	|› d
|› d�}| › d|› �}d| › d|› d|› d�}d| › d|› d|› d|› d�	}|› d|› d|› d|› d
|› d|› d�S )NzDdiv_floor_floating_vec implementation expect the same inputs' dtype.zdiv_floor_floating_vec(r™   rš   c              3   óF   K  — | ]  }t        |j                  «      –— Œ y ­wr  )r   rˆ   )r  Úitems     rW   r  z+CppVecOverrides.floordiv.<locals>.<genexpr>ä  s   è ø€ ÒG¸Ô'¨¯
©
×3ÑGùó   ‚!r{  r5   ú::blend<r“  ú(1), rÍ  rº  rÎ  rQ  z(0))rÏ  r¹   z	(0)) != (z(0)))z	::blendv(r„  rT  )r   rˆ   r"  r4   rZ  Ú_get_raw_num_vectorsÚtiling_factor)r  r€  Ú_trÐ  Úhas_remÚis_negs         rW   rÒ  zCppVecOverrides.floordivÜ  s1  € ä˜!Ÿ'™'Ô"Ø—7‘7˜aŸg™gÒ%ð ØVóÐ%ð -¨Q¨C¨r°!°°AÐ6Ð6äÑGÀÀAÀÔGÔGÐGÐGà˜Q˜C˜qÐ!ˆBÜ�x‰x×,Ñ,¨Q¯W©WÓ5¸Ò9à�d˜( A¬¯©×)?Ñ)?Ñ$?À1Ñ#DÐ"EÀRÈÀtÈ5ÐQRÐPSÐSTÐU�Ø�S˜˜A˜3�<ˆDØ˜!˜˜C ˜s $ r d¨$Ð/ˆGØ˜!˜˜C ˜t 9¨Q¨C¨s°2°$°eÐ<ˆFØ�T˜ 4 &¨¨4¨&°°B°4°u¸W¸IÀSÈÈÐPQÐRÐRrV   c                 óÌ   — t         j                  j                  |j                  «      dk  r2d|› d�}|› ddt         j                  j                  z  dz
  › d|› d|› d�}| › d|› �S )Nr5   r{  rš   r  r“  r  rÍ  )r4   rZ  r  rˆ   r  )r  r€  r  s      rW   rÙ  zCppVecOverrides.truncdivï  sp   € ô �8‰8×(Ñ(¨¯©Ó1°AÒ5à˜Q˜C˜qÐ!ˆBØ�$�h ¤Q§X¡X×%;Ñ%;Ñ ;¸qÑ@ÐAÀÀBÀ4ÀuÈQÈCÈqÐQˆAØ��C˜�sˆ|ÐrV   c                 óì   — | j                   t        j                  k(  rO|j                   t        j                  k(  sJ ‚t        t        j
                  j                  | |f«      \  }}|› d|› �S d| › d|› d�S )NrT  úat::vec::minimum(r™   rš   ©rˆ   r„   rk   rN   r4   rZ  r‹  ©r  r€  Úa_castÚb_casts       rW   r(  zCppVecOverrides.minimumø  ól   € à�7‰7”e—j‘jÒ Ø—7‘7œeŸj™jÒ(Ð(Ð(Ü1´!·(±(×2BÑ2BÀQÈÀFÓK‰NˆF�FØ�X˜S  Ð)Ð)à& q c¨¨A¨3¨aÐ0Ð0rV   c                 óì   — | j                   t        j                  k(  rO|j                   t        j                  k(  sJ ‚t        t        j
                  j                  | |f«      \  }}|› d|› �S d| › d|› d�S )NrY  úat::vec::maximum(r™   rš   r  r  s       rW   r+  zCppVecOverrides.maximum  r!  rV   c                 ó   — | › d| › �S r¨  rU   rK  s    rW   ÚsquarezCppVecOverrides.square
  r¥  rV   c                 óœ  — t        t        j                  t        «      sJ ‚|j                  t
        j                  k(  rY|j                  t
        j                  k(  sJ ‚t        t        j                  j                  | ||f«      \  }}}d|› d|› d|› d|› d�	S d|› d|› d|› dt        j                  j                  | |j                  «      › d�	S )Nr{  ú
)::blendv(r™   rš   )
r    r4   rZ  r�  rˆ   r„   rk   rN   r‹  Ú_get_mask_cast)r  r€  r/  Úblendv_aÚblendv_bÚblendv_cs         rW   r0  zCppVecOverrides.where  sÁ   € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ø�7‰7”e—j‘jÒ Ø—7‘7œeŸj™jÒ(Ð(Ð(Ü+?Ü—‘× Ñ  1 a¨ )ó,Ñ(ˆH�h ð ˜x˜j¨
°8°*¸B¸x¸jÈÈ8È*ÐTUÐVÐVà˜q˜c ¨A¨3¨b°°°2´a·h±h×6MÑ6MÈaÐQR×QXÑQXÓ6YÐ5ZÐZ[Ð\Ð\rV   c                 ó~  — t        «       }d| › d�}d| › d�}d| › d|› d|› d|› d| › d�}d| › d|› d|› d| › d|› d�}|j                  d«       |j                  «       5  |j                  d	|› d
�«       |j                  d|› d
�«       |j                  d«       d d d «       |j                  d«       |S # 1 sw Y   ŒxY w)Nr{  r{  r¶  r'  r™   r¹   rš   r^  r|  r¸   r}  r~  rf  r  )rŽ  r°   Úvec_zeror³  Úblendv_lÚblendv_rs         rW   r‚  zCppVecOverrides.sign  sè   € ä‹~ˆØ˜q˜c Ð&ˆØ˜a˜S Ð%ˆØ˜q˜c ¨H¨:°R¸°yÀÀ8À*ÈCÐPQÈsÐRSÐTˆØ˜q˜c ¨H¨:°R¸°yÀÀ1À#ÀSÈÈ
ÐRSÐTˆØ�‰�wÔØ�[‰[‹]ñ 	3Ø�N‰N˜\¨(¨°1Ð5Ô6Ø�N‰N˜]¨8¨*°AÐ6Ô7Ø�N‰NÐ1Ô2÷	3ð 	�‰�tÔØˆ÷	3ð 	3ús   Á<B3Â3B<c           
      óæ  — |t         j                  t         j                  t         j                  t         j                  t         j
                  t         j                  t         j                  t         j                  t         j                  f	v sJ t        › d|› �«       ‚t        | t        «      sJ ‚| j                  }t        j                  j!                  | ||«      }t        j                  j"                  j%                  t        j                  j&                  |«      }|j)                  d| |fd|i«       |t*        v r5|t         j                  k(  r"t        j                  j-                  | |||«       |S )Nz does not support r‡  r§   )r„   rk   Úfloat64rn   Úbfloat16Úfloat16Úuint8Úint8r
  rŒ  r  r    rG   rˆ   r4   rZ  rˆ  r‰  rŠ  r‹  rŒ  rƒ   r�  )rŽ  rˆ   r§   Úuse_compute_dtypesr�  r‘  s         rW   r‡  zCppVecOverrides.to_dtype)  s  € àÜ�J‰JÜ�M‰MÜ�K‰KÜ�N‰NÜ�M‰MÜ�K‰KÜ�J‰JÜ�K‰KÜ�K‰Kð

ñ 

ð 
	2ô ˆZÐ)¨%¨Ð1ó
	2ð 

ô ˜!œ^Ô,Ð,Ð,Ø—G‘Gˆ	Ü�x‰x×)Ñ)¨!¨U°IÓ>ˆÜ—‘—‘×&Ñ&¤q§x¡x×'7Ñ'7¸Ó>ˆØ×Ñ˜j¨1¨e¨*°{ÀIÐ6NÔOØ”MÑ! i´5·;±;Ò&>Ü�H‰H×(Ñ(¨¨I°v¸uÔEØˆrV   c                 óz   — t         j                  j                  }|dk(  r| › d| › d�S |€| › d�S t        d|›�«      ‚)Nr´  rµ  r¶  z.log1p()r·  r¸  r»  s     rW   r½  zCppVecOverrides.log1p?  sT   € ä�j‰j×6Ñ6ˆØ�*ÒØ�S˜ Q C tÐ,Ð,Øˆ[Ø�S˜�>Ð!ä ØJÈ3È'ÐRóð rV   c                 óŒ	  ‡— t        t        j                  t        «      sJ ‚t	        «       }t        j                  j
                  j                  «       }t        j                  j                  | «      5 }|j                  d|› d�«       t        j                  j                  |«      5  |j                  «       5   |«       }|j                  d|› d�«       d d d «       d d d «       d d d «       |j                  d«       t        j                  j                  j                  |«       j                  Š|› d�}ˆfd„}|j                  r|}	n ||«      }	t        |t         ‰   «      }
 ||
«      }t        t"        «      sJ |«       ‚|j                  �rt	        «       }|j                  d«       t        j                  j                  |«      5  |j                  «       5  |j                  d|› d	�«       |j                  «       5  |j                  d|› d�«       d d d «       |j                  d
«       |j                  «       5  t        j                  j
                  j%                  t        j                  j                  |	«      }t        j                  j
                  j%                  t        j                  j                  |«      }t        |t"        «      sJ |«       ‚t        |t"        «      sJ |«       ‚‰|_        ‰|_        t        j                  j&                  }|j                  d|j)                  |||«      › d�«       d d d «       d d d «       d d d «       |j                  d«       t        j                  j
                  j%                  t        j                  j                  |«      }n¡|j                  rKt        j                  j
                  j%                  t        j                  j                  | › d|	› d|› �«      }nJt        j                  j
                  j%                  t        j                  j                  | › d|› d|
› �«      }|j+                  d| |||fi «       |S # 1 sw Y   �Œ—xY w# 1 sw Y   �ŒœxY w# 1 sw Y   �Œ¡xY w# 1 sw Y   �ŒyxY w# 1 sw Y   �Œ\xY w# 1 sw Y   �ŒaxY w# 1 sw Y   �ŒfxY w)NrÒ   r<  r=  r¸   rf  c                 ó¼   •— ‰t         j                  k(  r$t        j                  j	                  «       › d| › d�S t        j                  j                  ‰«      › d| › d�S )Nú::from(rš   rº  )r„   rk   r4   rZ  r½  Ú_get_vec_type)r°   rˆ   s    €rW   Úmaskify_or_vecifyz1CppVecOverrides.masked.<locals>.maskify_or_vecify[  s]   ø€ ð œEŸJ™JÒ&ô —8‘8×*Ñ*Ó,Ð-¨W°T°F¸!Ð<ðô Ÿ™×.Ñ.¨uÓ5Ð6°a¸°v¸QÐ?ðrV   ú[&]úif (z.all_zero())Úelser-  r.  rD  )r    r4   rZ  r�  r7   r‰  r	  rD  r¿   r>  rÁ   r‹  r  rˆ   r¤   rO   rH   rG   rŠ  Ú	overridesr0  rŒ  )r?  r@  rA  r°   r¥   Únew_maskrö   Ú	body_coder<  Úbody_code_vecrC  Úother_code_vecÚbody_vec_varÚother_vec_varr@  r‘  rˆ   s                   @rW   rD  zCppVecOverrides.maskedK  sÊ  ø€ äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü‹~ˆÜ�h‰h�l‰l×!Ñ!Ó#ˆÜ�X‰X�_‰_˜TÓ"ð 	4 hØ�N‰N˜U 3 % vÐ.Ô/Ü—‘×&Ñ& tÓ,ñ 4¨d¯k©k«mñ 4Ù›�Ø—‘ ¨¨°Ð2Ô3÷4÷ 4÷	4ð
 	�‰�sÔÜ	�‰×Ñ×Ñ Ô%à—‘ˆØ�e˜2�Jˆ	ô	ð �=Š=Ø%‰Má-¨iÓ8ˆMÜ! %¬°eÑ)<Ó=ˆ
á*¨:Ó6ˆÜ˜(¤NÔ3Ð=°XÓ=Ð3Ø�?‹?Ü“>ˆDØ�N‰N˜5Ô!Ü—‘×&Ñ& tÓ,ñ ¨d¯k©k«mñ Ø—‘  h Z¨|Ð<Ô=Ø—[‘[“]ñ @Ø—N‘N W¨^Ð,<¸AÐ#>Ô?÷@à—‘˜vÔ&Ø—[‘[“]ñ ä#$§8¡8§<¡<×#8Ñ#8ÜŸ™×(Ñ(Ø%ó$�Lô %&§H¡H§L¡L×$9Ñ$9ÜŸ™×(Ñ(Ø&ó%�Mô & l´NÔCÐQÀ\ÓQÐCÜ% m´^ÔDÐSÀmÓSÐDØ).�LÔ&Ø*/�MÔ'äŸ™×*Ñ*ð ð —N‘NØ! )§/¡/°(¸LÈ-Ó"XÐ!YÐYZÐ[ô÷#÷÷ ð2 �N‰N˜4Ô Ü—X‘X—\‘\×*Ñ*Ü—‘× Ñ Øó‰Fð �]Š]Ü—X‘X—\‘\×*Ñ*Ü—‘× Ñ  T F¨#¨m¨_¸CÀÐ?OÐ"Pó‰Fô —X‘X—\‘\×*Ñ*Ü—‘× Ñ  T F¨#¨i¨[¸¸J¸<Ð"HóˆFð
 	×Ñ˜h¨¨t°U¸FÐ(CÀRÔHØˆ÷I4ñ 4ú÷ 4ñ 4ú÷	4ñ 	4ú÷>@ñ @ú÷ñ ú÷ñ ú÷ ñ ús‘   Á35RÂ(Q8Â9Q+ÃQ8ÃRÇR9Ç&R,ÈRÈ)R,ÉDRÍR,ÍR9Ñ+Q5Ñ0Q8Ñ8R	Ñ=RÒRÒRÒR,ÒR)Ò$R,Ò,R6	Ò1R9Ò9Sc                 óX  — t        t        j                  t        «      sJ ‚t        j                  j	                  | «      }t        j                  j
                  t        j                  j                     }t        j                  j                  ||«      }|dk(  rt        j                  | |«      S |�©t        j                  j                  j                  t        j                  j                  t        |«      t        | «      ¬«      }t        j                   ||«      }t        |t"        «      r|j$                  }t        j                  j'                  ||«      }n:t        j                  j)                  d ||t        j                  j                  «      }|j+                  d| |fi «       |S )Nr   r6  r:  )r    r4   rZ  r�  r8  ÚitervarsÚ
tiling_idxÚ_try_get_const_stridery  r:  r‰  rŠ  r‹  rD   r&   r2   r‡  r3   r�  ÚarangeÚ_load_or_store_non_contiguousrŒ  )r�  rˆ   r’   Ú
tiling_varÚstriderP  r�  r‘  s           rW   r:  zCppVecOverrides.index_expr˜  s3  € äœ!Ÿ(™(¤LÔ1Ð1Ð1Ü—‘×(Ñ(¨Ó.ˆÜ—X‘X×&Ñ&¤q§x¡x×':Ñ':Ñ;ˆ
Ü—‘×/Ñ/°°zÓBˆØ�QŠ;Ü×*Ñ*¨4°Ó7Ð7ØÐÜ—(‘(—,‘,×'Ñ'Ü—‘× Ñ ¤%¨£,Ô7LÈTÓ7Rð (ó ˆCô —L‘L  eÓ,ˆEÜ˜%¤Ô*ØŸ™�Ü—X‘X—_‘_ U¨FÓ3‰Fä—X‘X×;Ñ;Ø�e˜U¤A§H¡H×$4Ñ$4óˆFð 	×Ñ˜l¨T°5¨M¸2Ô>ØˆrV   c           
      óx  — d| › d�d| › d�f}t        d„ |D «       «      rt        d„ |D «       «      S t        | j                     }t        j
                  j                  rt        j
                  j                  nt        j
                  j                  }t        «       }t        j
                  j                  j                  t        j                  ¬«      }t        j
                  j                  j                  | j                  ¬«      }|j                  d| fi ¬«       |j                  d| fi ¬«       t        j
                  j                  | j                  «      }|d	k(  rd
|› d�nd|› d|› d�}|j                  |d	k(  rd|› d�nd|› d|› d�«       |j                  |› d|› d�«       |j                  d«       |j!                  «       5  |j                  d|› dt        j
                  j                  › d�«       |j                  | › dt#        |«      › d�«       |j                  dt        j
                  j                  › d�«       |j                  d|› dt        j
                  j                  › d�«       |j                  dt#        |«      › d�«       |j!                  «       5  |j                  d«       d d d «       |j                  |d	k(  r|› dt#        |«      › d�n|› d|› d t#        |«      › d�«       |j                  |› d!|› d"t#        |«      › d�«       d d d «       |j                  d#«       t        j
                  j$                  j'                  |«       ||f}	t)        ||	«      D ]/  \  }
}t        j
                  j                  j+                  |
|«       Œ1 ||fS # 1 sw Y   ŒíxY w# 1 sw Y   Œ•xY w)$Nrú  rû  rü  c              3   ór   K  — | ]/  }t         j                  j                  j                  |«      d u–— Œ1 y ­wr  rþ  r   s     rW   r  z(CppVecOverrides.frexp.<locals>.<genexpr>²  r  r  c              3   ón   K  — | ]-  }t         j                  j                  j                  |«      –— Œ/ y ­wr  rþ  r   s     rW   r  z(CppVecOverrides.frexp.<locals>.<genexpr>³  r  r  r  r  )r’  r5   úat::vec::Vectorized<r�   úat::vec::VectorizedN<r™   zat::vec::Vectorized<int32_t> r¸   zat::vec::VectorizedN<int32_t, ú> r•   r^  ú__at_align__ std::array<ú	> tmpbuf;ú.store(tmpbuf.data(), rÓ   z!__at_align__ std::array<int32_t, z> tmpbuf_exponent;z> tmpbuf_mantissa;rÖ   r×   z@tmpbuf_mantissa[i] = std::frexp(tmpbuf[i], &tmpbuf_exponent[i]);z? = at::vec::Vectorized<int32_t>::loadu(tmpbuf_exponent.data(), z! = at::vec::VectorizedN<int32_t, z!>::loadu(tmpbuf_exponent.data(), r·   z ::loadu(tmpbuf_mantissa.data(), z();)r"  r¡   rH   rˆ   r4   rZ  Ú	tail_sizer  r7   r‰  r	  r„   r
  rŒ  Ú_get_num_vectorsr¿   rÁ   rE   r‹  r  r  r  )rŽ  r  r‰   r:  r°   r  r  Ún_vecÚ
mantissa_tr  r  r  s               rW   r  zCppVecOverrides.frexp¯  s{  € à˜a˜S Ð%¨°¨s°$Ð'7Ð7ˆ
ÜÑWÈJÔWÔWÜÑUÈ*ÔUÓUÐUä˜aŸg™gÑ&ˆÜ%&§X¡X×%7Ò%7Œq�x‰x×!Ò!¼Q¿X¹X×=SÑ=SˆÜ‹~ˆÜ—8‘8—<‘<×&Ñ&¬U¯[©[Ð&Ó9ˆÜ—8‘8—<‘<×&Ñ&¨Q¯W©WÐ&Ó5ˆØ×Ñ ¨!¨°bÐÔ9Ø×Ñ ¨!¨°bÐÔ9Ü—‘×)Ñ)¨!¯'©'Ó2ˆð ˜Šzð # 6 (¨!Ñ,à(¨¨°°5°'¸Ð;ð 	ð
 	�‰à˜Šzð ,¨H¨:°QÑ7à1°%°¸¸8¸*ÀAÐFô	
ð
 	�‰˜*˜ Q x j°Ð2Ô3Ø�‰�wÔØ�[‰[‹]ñ 	Ø�N‰NØ*¨6¨(°"´Q·X±X×5KÑ5KÐ4LÈIÐVôð �N‰N˜a˜SÐ 6´{À4Ó7HÐ6IÈÐLÔMØ�N‰NØ3´A·H±H×4JÑ4JÐ3KÐK]Ð^ôð �N‰NØ*¨6¨(°"´Q·X±X×5KÑ5KÐ4LÐL^Ð_ôð �N‰NÐ1´+¸dÓ2CÐ1DÀFÐKÔLØ—‘“ñ Ø—‘ØVô÷ð �N‰Nà˜A’:ð �*Ð[Ô\gÐhlÓ\mÐ[nÐnpÑqà �zÐ!BÀ5À'ÐIjÔkvÐw{Ók|Ðj}Ð}ð  Aôð
 �N‰NØ�*˜C 
˜|Ð+KÌKÐX\ÓL]ÐK^Ð^`Ðaô÷+	ð0 	�‰�uÔÜ	�‰×Ñ×Ñ Ô%Ø˜hÐ'ˆÜ"% j°(Ó";ò 	1ÑˆI�wÜ�H‰H�L‰L×Ñ˜Y¨Õ0ð	1à˜Ð!Ð!÷#ð ú÷	ð 	ús&   Ç
CN0Ê&N$Ê8A$N0Î$N-	Î)N0Î0N9c                 ó   ‡— ˆfd„}|S )Nc                  ó  •— |rJ ‚t         j                  }t        |t        «      sJ ‚t	        «       }|j                  d«       | d   j                  }|j                  |«      }|j                  r|j                  n|j                  }g }t        |   }‰j                  dv }	|	rdn|}
‰j                  dk(  rt        | d      n|
}
|j                  «       5  t        | «      D ]£  \  }}t        |t        «      r}|j                  sJ ‚|j                  |k(  sJ ‚|j                  d|› d|j                  › d	|› d
�«       |j                  |› d|› dt!        |«      › d�«       |j#                  d|› d�«       Œ“|j#                  |«       Œ¥ |j                  d|
› d|j                  › d�«        ‰|Ž }|j                  dt!        |«      › d�«       |j                  «       5  |j                  d|› d
�«       d d d «       |	r|j                  rJ ‚d}d|› d|› d�}n#dt!        |«      › �}|dk(  rd|
› d�}n	d|
› d|› d�}|j                  d|› d|› d�«       d d d «       |j                  d«       |S # 1 sw Y   Œ{xY w# 1 sw Y   Œ(xY w) Nr^  r   )rÝ  rß  rÃ  rk   r”  éþÿÿÿrU  r™   z> tmpbufr¸   z.store(tmpbufz	.data(), rÓ   Útmpbufz[i]z> tmpbuf_out;rÖ   r×   ztmpbuf_out[i] = ztmpbuf_out.data()úat::vec::VecMask<ú,z>::fromztmpbuf_out.data(), r5   rR  z>::loaduz at::vec::VectorizedN<r=  rº  rf  )r4   rZ  r    r�  r7   r¿   rˆ   rY  rX  r  rH   r  rÁ   rÂ   rG   r¤   rE   rŽ  )r‘  r’  rZ  r°   Ú	vec_dtyperZ  r:  Úscalar_argsr‰   Úoutput_maskÚoctypeÚargidxr“  ÚresÚ	load_argsÚload_fnr™  s                   €rW   rY  z)CppVecOverrides._scalarize.<locals>.innerê  sÁ  ø€ ÙÐ�:Ü—X‘XˆFÜ˜f¤lÔ3Ð3Ð3Ü“>ˆDØ�N‰N˜7Ô#Ø˜Q™Ÿ™ˆIØ×+Ñ+¨IÓ6ˆEØ'-×'7Ò'7�6×#Ò#¸V×=QÑ=QˆDØˆKÜ! )Ñ,ˆFØ%×.Ñ.ð 3ð ˆKñ
  +‘V°ˆFð  ×(Ñ(Ð,>Ò>ô ˜T "™XÒ&àð ð
 —‘“ñ BÜ#,¨T£?ò 0‘K�F˜CÜ! #¤~Ô6Ø"ŸzšzÐ)˜zØ"Ÿy™y¨IÒ5Ð5Ð5ØŸ™Ø6°v°h¸bÀ×AUÑAUÐ@VÐV^Ð_eÐ^fÐfgÐhôð Ÿ™Ø"˜e =°°¸	Ä+ÈdÓBSÐATÐTVÐWôð $×*Ñ*¨V°F°8¸3Ð+?Õ@à#×*Ñ*¨3Õ/ð0ð —‘Ø.¨v¨h°b¸×9MÑ9MÐ8NÈmÐ\ôñ " ;Ð/�Ø—‘Ð!5´kÀ$Ó6GÐ5HÈÐOÔPØ—[‘[“]ñ >Ø—N‘NÐ%5°c°U¸!Ð#<Ô=÷>áØ%×/Ò/Ð/Ð/Ø 3�IØ 1°&°¸¸5¸'ÀÐI‘Gà"5´kÀ$Ó6GÐ5HÐ I�IØ ’zØ$8¸¸ÀÐ"I™à$:¸6¸(À"ÀUÀGÈ8Ð"T˜Ø—‘ ¨¨	°°9°+¸RÐ@ÔA÷?Bð@ �N‰N˜4Ô ØˆK÷>ð >ú÷'Bð Bús&   ÃDI6ÇI*Ç0AI6É*I3	É/I6É6I?rU   )r%  r™  rY  s    ` rW   Ú
_scalarizezCppVecOverrides._scalarizeè  s   ø€ ô7	ðr ˆrV   c                 ó  — t        t        «      }t        t        «      j                  «       D ]S  \  }}t	        |t
        «      sŒ||vsŒ| j                  |j                  «      }||_        t        | |t        |«      «       ŒU y r  )
rž  r‡  ry  rŸ  r    r„  rj  r¡  r  r   )r%  Úvec_varsrÄ   r£  rš  s        rW   Ú_initialize_scalarizez%CppVecOverrides._initialize_scalarize%  sh   € äœÓ(ˆÜ ¤Ó.×4Ñ4Ó6ò 	7‰LˆD�&Ü˜&¤,Õ/°DÀÒ4HØ—~‘~ f§o¡oÓ6�Ø $�”Ü˜˜T¤<°Ó#5Õ6ñ		7rV   rƒ  )Wr  r  r  r  r�  r„  r�  rÇ   rL  rª  r—  rš  rœ  r¡  r£  r¦  r©  r«  r­  r°  r¾  rÃ  rÅ  rÈ  rË  rÎ  rÑ  r²  rÅ  rÇ  rÉ  rË  rÔ  rÖ  rÛ  rá  rG  rL  rO  rR  rU  rW  rZ  r\  rj  rm  rî  rq  rs  rw  r÷  r¿  rÁ  rû  rô  rã  rë  rç  ré  r  r  r  rø  rñ  r  rö  rí  rå  r&  ry  rž  rÒ  rÙ  r(  r+  r%  r0  r‚  r‡  r½  rD  r:  r  r_  rj  rm  r`  ra  s   @rW   r‡  r‡  :  sV  ø„ Ù8ô[ðz ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñ(ó ð(ð
 ñó ðð ñó ðð ñó ðð ñó ðð ñAó ðAð ñ	Eó ð	Eð ñ@ó ð@ð ñ@ó ð@ð ñAó ðAð ñAó ðAð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñ ó ð ð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñ0ó ð0ð ñ9ó ð9ð ñ9ó ð9ð ñFó ðFð ñAó ðAð ñó ðð ñó ðð ñ#ó ð#ð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñó ðð ñ%ó ð%ð ñ$ó ð$ð ñ!ó ð!ð ñ!ó ð!ð ñOó ðOð ñó ðð ñó ðð ñó ðð  ñHó ðHð ñó ðð ñSó ðSð$ ñó ðð ñ1ó ð1ð ñ1ó ð1ð ñó ðð ñ	]ó ð	]ð ñó ðð òó ðð* ñ	ó ð	ð ñJó ðJðX ñó ðð, ñ6"ó ð6"ðp ñ:ó ð:ðx ñ7ó ô7rV   r‡  Úcppvecc                   ó   — e Zd Zed„ «       Zy)ÚCppTile2DOverridesc                 ó¬   — t        t        j                  t        «      sJ ‚t        j                  j	                  | «      } t
        j                  | |«      S r  )r    r4   rZ  ÚCppTile2DKernelÚtransform_indexingr‡  r:  )r�  rˆ   s     rW   r:  zCppTile2DOverrides.index_expr4  s=   € äœ!Ÿ(™(¤OÔ4Ð4Ð4Ü�x‰x×*Ñ*¨4Ó0ˆÜ×)Ñ)¨$°Ó6Ð6rV   N)r  r  r  r„  r:  rU   rV   rW   rp  rp  3  s   „ Øñ7ó ñ7rV   rp  c                   óœ  ‡ — e Zd ZeZeZdZdZˆ fd„Z	e
efd„Zd„ Zd.dee   fd„Zej$                  d„ «       Z	 d/d	ej*                  fd
„Zd	ej*                  defd„Zd	ej*                  dej2                  fd„Zd	ej*                  dej2                  fd„Zd„ Zdej*                  dej*                  dedefd„Zded	ej*                  fd„Zd.d„Z de!e"ef   dedede#jH                  fd„Z%d.dee&   fd„Z'd„ Z(d „ Z)d!„ Z*d"„ Z+d#„ Z,d$„ Z-e.defd%„«       Z/d&„ Z0ej$                  d'„ «       Z1d(„ Z2d)„ Z3d*„ Z4	 	 d0ded+ee   d,eej2                     fd-„Z5ˆ xZ6S )1Ú	CppKernelrÒ   r¸   c                 óh  •— t         ‰| �  |«       i | _        g | _        d | _        g | _        g | _        d | _        t        «       | _	        g | _
        t        «       | _        t        «       | _        t        «       | _        t        «       | _        t        «       | _        d| _        t        «       | _        t%        | j&                  | j(                  d¬«      | _        t%        | j&                  | j(                  d¬«      | _        t        «       | _        t        «       | _        || _        i | _        g | _        y )NFÚtmp_acc)Úname_prefixÚwrecps)r,  r-  Úactive_rangesÚinner_itervarsÚcall_rangesrM  rH  Úreduction_depthr=   Úreduction_prefixÚreduction_prefix_generatorsÚreduction_suffixÚparallel_reduction_prefixÚparallel_reduction_suffixÚlocal_reduction_initÚlocal_reduction_storesÚis_reductionÚnon_parallel_reduction_prefixr8   Únewvar_prefixÚsuffixÚreduction_cseÚweight_recps_cseÚpreloadsÚ
poststoresÚnum_threadsÚreduction_omp_decÚreduction_var_names)r.  r‘  r�  r1  s      €rW   r-  zCppKernel.__init__A  s
  ø€ Ü‰Ñ˜Ôð HJˆÔð 35ˆÔØ=AˆÔØ(*ˆŒØ,.ˆŒØ#ˆÔÜ .Ó 0ˆÔð <>ˆÔ(Ü .Ó 0ˆÔÜ)7Ó)9ˆÔ&Ü)7Ó)9ˆÔ&Ü$2Ó$4ˆÔ!Ü&4Ó&6ˆÔ#Ø!ˆÔÜ-;Ó-=ˆÔ*Ü  ×!3Ñ!3°T·[±[ÈiÔXˆÔÜ #Ø×Ñ §¡¸ô!
ˆÔô 'Ó(ˆŒÜ(Ó*ˆŒØ&ˆÔØ=?ˆÔØ.0ˆÕ rV   c                 ó8  — t         j                  j                  r'| j                  s| j                  j	                  d«       |› d�}t         j                  j                  rdn	t        «       }|› d�}	| j                  j	                  |› d|› d |||«      › d�«       | j                  j                  t        ||||||«      «       | j                  j	                  |	› d|› d�«       | j                  j                  d|› d	�d
d|› d ||||	|¬«      › d�dg«       y )Nz(int max_threads = omp_get_max_threads();Ú_localÚmax_threadsz	_arr[tid]r•   r·   r¸   zfor (int tid = 0; tid < z; tid++)rØ   rÙ   ©r§   r€   )r   r¹  Údynamic_threadsr�  r¿   r+   rƒ  r  rß   r„  r‚  rÛ   )
r.  r¯   rÏ   r‡   rˆ   Úreduction_combine_fnÚreduction_init_fnÚ	acc_localr�  Úacc_local_in_arrays
             rW   Ú_gen_parallel_reduction_buffersz)CppKernel._gen_parallel_reduction_buffersf  s=  € ô �:‰:×%Ò%¨d×.LÒ.LØ×*Ñ*×4Ñ4Ø:ôð �e˜6�Nˆ	ä#ŸZ™Z×7Ò7‰MÔ=QÓ=Sð 	ð !$˜u IÐ.ÐØ×!Ñ!×+Ñ+Øˆj˜˜)˜ CÑ(9¸.È%Ó(PÐ'QÐQRÐSô	
ð 	×&Ñ&×-Ñ-Ü"ØØØØØØ!óô		
ð 	×#Ñ#×-Ñ-Ð1CÐ0DÀCÈ	À{ÐRSÐ.TÔUØ×&Ñ&×1Ñ1à*¨;¨-°xÐ@ØØ�s�e˜3Ñ3°NÀCÐI[ÐglÔmÐnÐnoÐpØð	õ	
rV   c                 óZ   — | j                   D ]  }t        | j                  ||› d�«       Œ y )Nr‘  )r�  rä   Ústores)r.  Úvar_names     rW   Ú%update_stores_with_parallel_reductionz/CppKernel.update_stores_with_parallel_reduction�  s0   € Ø×0Ñ0ò 	IˆHÜ˜TŸ[™[¨(°x°jÀÐ4GÕHñ	IrV   r°   c                 ó”  — |�J ‚t        «       }t        j                  «       5 }t        | d«      rK|j	                  | j
                  «       | j                  |«       |j                  |j                  «       «       |j	                  | j                  «       |j	                  | j                  «       |j	                  | j                  «       d d d «       t        | d«      r|j	                  | j                  «       | j                  r5| j                  D ]&  }| j                  |   \  }}t        |||› d�||«      }Œ( |S # 1 sw Y   ŒsxY w)NÚcodegen_inner_loopsÚ_tail)r7   r½   r¾   r£   r  r‹  rŸ  rÀ   rÁ   Úloadsr‹  r›  rŒ  r{  rz  rÍ   )r.  r°   rÉ   rP  ÚstartÚends         rW   Úgen_bodyzCppKernel.gen_body“  s  € Øˆ|Ðˆ|Ü‹~ˆÜ×!Ñ!Ó#ð 	% uÜ�tÐ2Ô3Ø—‘˜DŸM™MÔ*Ø×(Ñ(¨Ô.Ø×#Ñ# D§K¡K£MÔ2Ø�K‰K˜Ÿ
™
Ô#Ø�K‰K˜Ÿ™Ô%Ø�K‰K˜Ÿ™Ô$÷	%ô �4Ð.Ô/Ø�K‰K˜Ÿ™Ô(à×ÒØ×*Ñ*ò X�Ø!×/Ñ/°Ñ4‘
��sÜ1°$¸À¸uÀE¸]ÈEÐSVÓW‘ðXð ˆ÷	%ð 	%ús   £B)D>Ä>Ec              #   ó  K  — | j                   }|rYt        j                  ||«      }t        |t        «      r3|j
                  }t        |t        «      sJ ‚t        j                  |_	        || _         	 |–— || _         y# || _         w xY w­w)z>Context manager to add an additional mask to loads and stores.N)
Ú
_load_maskr2   rÑ  r    r3   r�  rG   r„   rk   rˆ   )r.  r?  Úpriors      rW   rD  zCppKernel.masked§  ss   è ø€ ð —‘ˆÙÜ—8‘8˜D %Ó(ˆDÜ˜$¤Ô)Ø—z‘z�Ü! $¬Ô7Ð7Ð7ô #ŸZ™Z�”
àˆŒð	$ØŠJà#ˆD�Oø˜eˆD�Oüs   ‚A/B
Á2A> Á6B
Á>	BÂB
r’   c                 óP   — | j                   |   }|||z  |z   i}t        ||«      }|S r  )rH  r0   )r.  r’   ÚscaleÚitervar_idxr;  r¥   rì   rí   s           rW   Úscale_index_with_offsetz!CppKernel.scale_index_with_offsetº  s7   € ð �m‰m˜KÑ(ˆØ˜C %™K¨&Ñ0Ð1ˆÜ˜u kÓ2ˆ	ØÐrV   rµ   c                 ó6   — t        | j                  |«      «      S )z¦
        Convert an index expr to a string that can be used in cpp code.
        e.g. a sympy expression "s2" may actually appear as "ks1" in the cpp kernel.
        )rD   r8  ©r.  r’   s     rW   Úindex_to_strzCppKernel.index_to_strÂ  s   € ô
 �T×)Ñ)¨%Ó0Ó1Ð1rV   Úitervarc                 óD   ‡ ‡— t        ˆˆ fd„|j                  D «       «      S )z]
        Check if an index has free symbol CppCSEVariable that depends on `itervar`.
        c              3   ó(  •K  — | ]‰  }|j                   ‰j                  j                  v ret        ‰j                  j                  |j                      t        «      r4‰j                  j                  |j                      j                  ‰«      –— Œ‹ y ­wr  )rÄ   r‰  Úvarname_mapr    rG   Ú
depends_on)r  Úsr¯  r.  s     €€rW   r  z6CppKernel.index_indirect_depends_on.<locals>.<genexpr>Í  sj   øè ø€ ò 
àØ�v‰v˜Ÿ™×-Ñ-Ñ-Ü˜4Ÿ8™8×/Ñ/°·±Ñ7¼ÔHð �H‰H× Ñ  §¡Ñ(×3Ñ3°G×<ñ
ùs   ƒBB)rf   Úfree_symbols©r.  r’   r¯  s   ` `rW   Úindex_indirect_depends_onz#CppKernel.index_indirect_depends_onÉ  s%   ù€ ô ô 
à×'Ñ'ô
ó 
ð 	
rV   c                 óF   — ||j                   v xs | j                  ||«      S r  )rµ  r·  r¶  s      rW   Úindex_depends_onzCppKernel.index_depends_onÔ  s,   € Ø˜%×,Ñ,Ð,ò 
°×0NÑ0NØ�7ó1
ð 	
rV   c                 óT   — t        t        | j                  | j                  «      «      S r  )Údictr  rH  rM  r3  s    rW   Ú
var_rangeszCppKernel.var_rangesÙ  s   € Ü”C˜Ÿ™ t§{¡{Ó3Ó4Ð4rV   r�  r:  ÚlowerÚupperc                 ó  — |s|sy t        |t        j                  «      }|rIt        j                  |t
        j                  «      j                  }t        j                  j                  }n‰t        j                  j                  }	 | j                  t        j                  _
        t        j                  |t
        j                  «      j                  }|t        j                  _
        | j                  }|r.t        j                  j                  | j                  |«      «      nd }	| j                  ||rdnd |	| j                  «      }
| j                   j#                  ||
d¬«       y # |t        j                  _
        w xY w)NÚ0F)Ú
assignment)r   r   ÚTMPr2   r:  r„   rŒ  r�  r4   rZ  r‹  r¡  Úsexprr8  Úindirect_assertr¦  r‰  rŠ  )r.  r�  r:  r½  r¾  Úindirectr‘  rà   Úprior_computeÚsize_strrÅ   s              rW   Úcheck_boundszCppKernel.check_boundsÜ  s
  € ñ ™Øä& t¬T¯X©XÓ6ˆÙä—^‘^ D¬%¯+©+Ó6×<Ñ<ˆFÜ—X‘X×%Ñ%‰Fô ŸH™H×,Ñ,ˆMð1Ø#'§:¡:”—‘Ô ÜŸ™¨¬e¯k©kÓ:×@Ñ@�à#0”—‘Ô Ø—Z‘ZˆFáAF”1—8‘8—>‘> $×"6Ñ"6°tÓ"<Ô=ÈDˆà×#Ñ#Ø™5‘C d¨H°d·o±oó
ˆð 	�‰×Ñ˜& $°5ÐÕ9øð $1”—‘Õ ús   ÂAE' Å'E>rÄ   c                 óø   — | j                   j                  |«      }| j                  |«      }|› dt        |«      › d�}| j                  j                  | j                  |«      }|j                  d| ||fi «       |S )Nú[ú]rí  )r‘  Úinputr8  rE   r‰  rŠ  r¡  rŒ  )r.  rÄ   r’   r¥   rÅ   r‘  s         rW   rí  zCppKernel.loadü  st   € Ø�i‰i�o‰o˜dÓ#ˆØ×$Ñ$ UÓ+ˆØ��aœ EÓ*Ð+¨1Ð-ˆØ—‘×"Ñ" 4§:¡:¨tÓ4ˆØ×Ñ˜f t¨T°5Ð&9¸2Ô>ØˆrV   c                 ó  — d|v sJ ‚| j                   j                  |«      }| j                  |«      }|€|› dt        |«      › d|› d�}n–|dk(  rƒt        j
                  j                  s$| j                  dk(  r|› dt        |«      › d|› d�}nSt        j                  j                  |«      }dt        |   › d	|› d
�}d|› dt        |«      › d|› d�}nt        d|› �«      ‚| j                  j                  t        ||«      «       y )NÚbufrÊ  ú] = r¸   Ú
atomic_addr5   z] += zstatic_cast<r“  rš   zatomic_add(&z], rÓ   ústore mode=)r‘  Úoutputr8  rE   r   r¹  r”  r�  r4   ÚgraphÚ	get_dtyperH   ÚNotImplementedErrorr›  r¿   r;   )r.  rÄ   r’   r�  Úmoder¥   rÅ   rˆ   s           rW   ÚstorezCppKernel.store  s  € Ø˜‰}Ðˆ}Ø�i‰i×Ñ˜tÓ$ˆØ×$Ñ$ UÓ+ˆØˆ<Ø�U˜!œK¨Ó.Ð/¨t°E°7¸!Ð<‰DØ�\Ò!Ü—:‘:×-Ò-°$×2BÑ2BÀaÒ2GØ˜˜a¤¨EÓ 2Ð3°5¸¸¸qÐA‘äŸ™×)Ñ)¨$Ó/�à&¤|°EÑ':Ð&;¸2¸e¸WÀAÐF�Ø% c U¨!¬K¸Ó,>Ð+?¸sÀ5À'ÈÐL‘ä%¨°D°6Ð&:Ó;Ð;Ø�‰×Ñœl¨4°Ó6Õ7rV   r¯   rÏ   Úrtyperˆ   c                 ó@   ‡‡‡‡‡— ddt         t           fˆˆˆˆˆfd„}|S )Nr:  c                 óL   •— | €‰› d‰› d ‰‰‰«      › d�S t        ‰‰‰‰| ‰«      S )Nr•   r·   r¸   )rß   )r:  r¯   rÏ   rˆ   rÜ   rØ  s    €€€€€rW   rY  z.CppKernel._gen_reduction_prefix.<locals>.inner$  sH   ø€ Øˆ|Ø"˜ 1 S E¨©W°U¸EÓ-BÐ,CÀ1ÐEÐEä-ØØØØØØóð rV   r  )r   rm   )r.  r¯   rÏ   rØ  rˆ   rÜ   rY  s    ````` rW   Ú_gen_reduction_prefixzCppKernel._gen_reduction_prefix  s   ü€ ñ	œ¤™÷ 	ñ 	ð ˆrV   c                 óh   — | j                   D ]#  }| j                  j                   ||«      «       Œ% y r  )r  r~  r  )r.  r:  Úgen_fns      rW   Úfinalize_reduction_prefixz#CppKernel.finalize_reduction_prefix3  s/   € Ø×6Ñ6ò 	7ˆFØ×!Ñ!×(Ñ(©°«Õ6ñ	7rV   c           
      ó^  — |dv }|||f}|| j                   j                  v r| j                   j                  |   S | j                   j                  | j                  d|› �d¬«      }| j                  j                  |› «       d| _        |r|n|}t        ||«      }	| j                  j                  | j                  ||	||t        «      «       | j                  €J ‚| j                  | j                     }
t        | j                  dz   t        | j                  «      «      D ]$  }|
| j                  |   z  | j                  |   z   }
Œ& | j                   j#                  |› dt%        ||||
«      › d�«       | j'                  ||	||«       t)        ||«      }|| j                   j                  |<   |S )	N©ra   r`   ú
reduction F©ÚwriteTr5   r·   r¸   )r‰  Úreduction_cacherŠ  r¡  r�  rŽ  r…  r‘   r  rÛ  r�   r}  rH  rG  rÐ   rM  r›  r¿   r­   r™  r¢   )r.  rˆ   r§   r‡   r�  Úargmax_or_argminÚreduction_keyr¯   Ú
init_dtyperÏ   r’   rã   rö   s                rW   Ú	reductionzCppKernel.reduction7  s²  € Ø)Ð-AÐAÐØ! >°5Ð8ˆØ˜D×.Ñ.×>Ñ>Ñ>Ø×%Ñ%×5Ñ5°mÑDÐDà× Ñ ×)Ñ)Ø�J‰J˜* ] OÐ4¸Eð *ó 
ˆð 	× Ñ ×'Ñ'¨3¨%Ô1Ø ˆÔÙ"2‘Y¸ˆ
Ü% n°jÓAˆØ×(Ñ(×/Ñ/Ø×&Ñ&Ø�X˜~¨z¼>óô	
ð
 ×#Ñ#Ð/Ð/Ð/Ø—‘˜d×2Ñ2Ñ3ˆÜ�t×+Ñ+¨aÑ/´°T·]±]Ó1CÓDò 	>ˆAØ˜DŸK™K¨™NÑ*¨T¯]©]¸1Ñ-=Ñ=‰Eð	>à�‰×ÑØˆe�3Ô(¨¸¸eÀUÓKÐLÈAÐNô	
ð 	×,Ñ,¨S°(¸NÈJÔWÜ" >°3Ó7ˆØ<Bˆ×Ñ×*Ñ*¨=Ñ9ØˆrV   c                 óÊ   — | j                  |«      }| j                  j                  |«      }| j                  j	                  t        ||› dt        |«      › d|› d�«      «       y )NrÊ  rÏ  r¸   )r8  r‘  rÒ  r€  r¿   r;   rE   )r.  rÄ   r’   r�  r¥   s        rW   Ústore_reductionzCppKernel.store_reductionV  s]   € Ø×$Ñ$ UÓ+ˆØ�i‰i×Ñ˜tÓ$ˆØ×Ñ×'Ñ'Ü˜ #  a¬°EÓ(:Ð';¸4À¸wÀaÐHÓIõ	
rV   c                 óª  — | j                   ri| j                   t        |«      t        |«      z   k(  s+J | j                   › dt        |«      › dt        |«      › �«       ‚| j                  t        |«      k(  s¥J ‚t        |«      t        |«      z   | _         | j                   D �cg c]  }| j	                  |«      ‘Œ c}| _        t        t        | j
                  «      «      D �cg c]  }t        t        j                  |«      ‘Œ c}| _
        t        |«      | _        | j                  d | j                   | j                  | j                  d  fS c c}w c c}w )Nr»  r}  )r|  r¡   r}  rÐ   r8  rM  rG  r.   r   ÚXBLOCKrH  )r.  ÚlengthsÚreduction_lengthsrŽ  Úns        rW   Ú
set_rangeszCppKernel.set_ranges]  s2  € Ø×ÒØ×#Ñ#¤u¨W£~¼Ð>OÓ8PÑ'PÒPð Ø×#Ñ#Ð$ D¬¨w«Ð(8¸¼EÐBSÓ<TÐ;UÐVóÐPð ×'Ñ'¬3¨w«<Ò7Ð7Ð7ä$ W›~´Ð6GÓ0HÑHˆDÔØ<@×<LÑ<LÖM°q˜4×/Ñ/°Õ2ÒMˆDŒKô œs 4§;¡;Ó/Ó0öàô /¬t¯{©{¸AÕ>òˆDŒMô $' w£<ˆDÔ à�M‰MÐ0˜D×0Ñ0Ð1Ø�M‰M˜$×.Ñ.Ð0Ð1ð
ð 	
ùò Nùòs   Â EÃ !Ec                 óš   — | j                   €J ‚t        j                  j                  j	                  t        | j                   «      d¬«      S )Né    ©Úfallback)r|  r4   rÓ  ÚsizevarsÚ	size_hintr/   r3  s    rW   rö  zCppKernel.size_hintp  sF   € Ø×ÑÐ+Ð+Ð+Ü�w‰w×Ñ×)Ñ)Ü˜$×*Ñ*Ó+°dð *ó 
ð 	
rV   c                 ó^  ‡‡‡‡‡‡‡‡‡‡— t        | t        «      sJ ‚t        «       Š| j                  €J ‚t        |j                  t
        «      r+|j                  j                  |j                  «       ‰«      Šn | j                  |j                  «       ‰«      Š|j                  d uxr# |j                  ‰j                     j                  Št        j                  «       5 }‰j                  r6‰r‰j                  «        n‰j                  ‰«       |j!                  ‰«       n4‰dkD  r/‰j#                  «       r|j%                  ‰j'                  «       «       dt(        fˆˆˆfd„Šdd„Š	 ddt(        dt*        fˆˆˆˆˆˆfd„Šddt(        dt*        fˆˆfd„Š	 	 ddt(        dt*        dt,        fˆˆfd	„Š|j%                  ‰j'                  «       «       t        |j                  t
        «      �r t        t.        j0                  t2        «      �rt.        j0                  j4                  rçt.        j0                  j4                  }|j7                  «       D ]º  }t9        |j;                  «       j<                  D �cg c]  }| j?                  |«      ‘Œ c}«      }t@        |j;                  «       jB                     }	d
|	› dtE        |«      › d�}
|jG                  «       }‰jI                  d|	› d|› d|
› d�«       ‰jI                  |	› d|› d|› d�«       Œ¼  ‰|«       d d d «       y c c}w # 1 sw Y   y xY w)Nr5   Ú
_loop_nestc                 ó¤  •‡ — ˆ ˆfd„}‰ j                  «       }t        |t        «      r|j                  D ]
  Š  ‰‰ «       Œ y t        |t        «      sJ ‚‰ j
                  � |«       r|j                  «        t        j                  «       5 }|j                  ‰j                  «       «       |j                  ‰«       d d d «       y # 1 sw Y   y xY w)Nc                  ó†   •— ‰j                   sJ ‚‰j                   ‰j                     } | j                  xr | j                  S r  )rB  r  r…  Úparallel)Úrootrø  Ú	par_depths    €€rW   Úis_parallel_reductionzOCppKernel.codegen_loops_impl.<locals>.gen_kernel.<locals>.is_parallel_reduction”  s=   ø€ Ø%×+Ò+Ð+Ð+Ø%×+Ñ+¨I×,AÑ,AÑB�DØ×,Ñ,Ò>°·±Ð>rV   )Ú
get_kernelr    rW  rY  ÚCppKernelProxyrB  r�  r½   r¾   rÀ   rÁ   r¤  )rø  rþ  rZ  rÉ   r°   Úgen_loop_nestrý  s   `   €€€rW   Ú
gen_kernelz0CppKernel.codegen_loops_impl.<locals>.gen_kernel“  s±   ù€ õ?ð
 $×.Ñ.Ó0�Ü˜fÔ&:Ô;Ø&,§l¡lò 2˜
Ù% jÕ1ñ2ô & f¬nÔ=Ð=Ð=Ø!×'Ñ'Ð3Ñ8MÔ8OØ×DÑDÔFÜ#×-Ñ-Ó/ð .°5Ø×+Ñ+¨D¯K©K«MÔ:ØŸ™¨Ô-÷.÷ .ñ .ús   Â1CÃCc                 ó¤   — |r| j                   }|r| j                  |z   }|S | j                  }|r|| j                  z   }|S || j                  z   }|S r  )r€  r‚  r~  r�  r†  )rZ  rû  Ú	is_suffixrˆ  Úprefixs        rW   Úget_reduction_prefix_suffixzACppKernel.codegen_loops_impl.<locals>.get_reduction_prefix_suffix¥  sd   € ÙØ#×4Ñ4�FÙØ!'×!AÑ!AÀFÑ!J˜Ø!�Mà#×4Ñ4�FÙØ!'¨&×*JÑ*JÑ!J˜ð "�Mð "(¨&×*NÑ*NÑ!N˜Ø!�MrV   Údepthc                 ó  •— | j                  «       }| j                  sJ ‚| j                  |   }t        j                  «       5 }|j                  rI|sG ‰	||j
                  d¬«      }|r|j                  ‰j                  «       «       ‰j                  |«       ‰
rR|j
                  rF‰j                  ‰«       |j                  r)|j                  sJ ‚‰j                  |j                  «        ‰| |«       ‰
rC|j
                  r7|j                  r‰j                  |j                  «       ‰j                  «        |j                  r&|s$‰j                   ‰	||j
                  d¬«      «       d d d «       y # 1 sw Y   y xY w)NF)r  T)rÿ  rB  r½   r¾   r…  rû  rÀ   rÁ   r  rƒ  r„  Úclose)rø  r  Úin_reductionrZ  ÚloopÚstack_outerr~  r°   Úgen_loop_atr  Úis_reduction_loopÚthreadsÚworksharings          €€€€€€rW   Úgen_loop_with_reductionz=CppKernel.codegen_loops_impl.<locals>.gen_loop_with_reduction³  sA  ø€ ð $×.Ñ.Ó0�Ø!×'Ò'Ð'Ð'Ø!×'Ñ'¨Ñ.�Ü×)Ñ)Ó+ð ¨{Ø×(Ò(±Ù+FØ" D§M¡M¸Uô,Ð(ñ ,Ø'×5Ñ5°d·k±k³mÔDØŸ™Ð$4Ô5Ù(¨T¯]ª]Ø#×,Ñ,¨WÔ5Ø!×6Ò6Ø#)×#@Ò#@Ð@Ð#@Ø ŸK™K¨×(CÑ(CÔDá 
¨EÔ2á(¨T¯]ª]Ø!×8Ò8Ø ŸK™K¨×(EÑ(EÔFØ#×)Ñ)Ô+Ø×(Ò(±ØŸ™Ù7Ø &¨¯©Àôô÷+÷ ñ ús   ÁD*E6Å6E?c                 óX  •— t        j                  «       5 }| j                  sJ ‚| j                  |   }|j                  «       }|€
	 d d d «       y ‰j	                  |«       |j                  ‰j                  «       «        ‰| |dz   |j                  «       d d d «       y # 1 sw Y   y xY wræ   )r½   r¾   rB  ÚlinesrÛ   rÀ   rÁ   r…  )rø  r  rÉ   r  Ú
loop_linesr°   r  s        €€rW   r  z1CppKernel.codegen_loops_impl.<locals>.gen_loop_atÔ  s¡   ø€ Ü×)Ñ)Ó+ð L¨uØ%×+Ò+Ð+Ð+Ø%×+Ñ+¨EÑ2�DØ!%§¡£�JØ!Ð)Ø÷Lð Lð —O‘O JÔ/Ø×'Ñ'¨¯©«Ô6Ù! *¨e°a©i¸×9JÑ9JÔK÷L÷ Lñ Lús   –1B ÁAB Â B)r
  c                 ót   •— | j                   �|t        | j                   «      k(  r	 ‰| «       y  ‰| ||«       y r  )rB  rÐ   )rø  r  r
  r  r  s      €€rW   r  z3CppKernel.codegen_loops_impl.<locals>.gen_loop_nestß  s7   ø€ ð
 ×#Ñ#Ð+¨u¼¸J×<LÑ<LÓ8MÒ/MÙ˜zÕ*á+¨J¸¸|ÕLrV   zstd::make_unique<z []>(rš   zstd::unique_ptr<z	 []> buf_r·   r¸   z* z = buf_z.get();)FF)r   F)r   )%r    r   r+   r|  rZ  rW  Údecide_parallel_depthÚmax_parallel_depthrB  r  r…  r½   r¾   r
  r	  rû  Úmark_parallelÚsinglerÀ   rÁ   rF  rm   rk   r4   Úlocal_buffer_contextrJ   Úlocal_buffersÚvaluesr/   Ú
get_layoutr:  r8  rH   rˆ   rD   Úget_namer  )r.  rH  r°   r  rÉ   r  Úlocal_bufferÚsize_valÚlocal_buf_sizeÚlocal_buf_dtypeÚallocateÚlocal_buffer_namer  r  r  r  r  r  rý  r  s     ``        @@@@@@@@rW   Úcodegen_loops_implzCppKernel.codegen_loops_implv  s  ÿù€ Ü˜$¤Ô/Ð/Ð/Ü&Ó(ˆØ×ÑÐ+Ð+Ð+Ü�i×&Ñ&Ô(<Ô=Ø!×(Ñ(×>Ñ>Ø×,Ñ,Ó.°ó‰Ið ×2Ñ2Ø×,Ñ,Ó.°óˆIð
 �O‰O 4Ð'ò DØ—‘ 	× 5Ñ 5Ñ6×CÑCð 	ô ×!Ñ!Ó#ð |	% uØ×'Ò'Ù$à×%Ñ%Õ'à×(Ñ(¨Ô1Ø×'Ñ'¨	Õ2Ø˜1’Ø×%Ñ%Ô'Ø×'Ñ'¨¯©«Ô6ð.¤x÷ .ó$"ð DIñÜ$ðÜ-0÷ò ñB	L¬ð 	L¼ö 	Lð Ø%*ñMÜ$ðMäðMô #öMð ×Ñ §¡£Ô.ô ˜9×+Ñ+Ô-AÕBÜœq×5Ñ5Ô7IÕJÜ×*Ñ*×8Ò8ô !"× 6Ñ 6× DÑ D�Ø$1×$8Ñ$8Ó$:ò �Lä%2ð -9×,CÑ,CÓ,E×,JÑ,Jöà (ð !×0Ñ0°Õ:òó&�Nô '3°<×3JÑ3JÓ3L×3RÑ3RÑ&S�OØ!2°?Ð2CÀ5ÌÈ~ÓI^ÐH_Ð_`Ða�HØ(4×(=Ñ(=Ó(?Ð%Ø—K‘KØ*¨?Ð*;¸9ÐEVÐDWÐWZÐ[cÐZdÐdeÐfôð —K‘KØ*Ð+¨2Ð.?Ð-@ÀÐHYÐGZÐZaÐbõðñ" ˜)Ô$÷y|	%ð |	%ùò\÷]|	%ð |	%ús    Ã!FL#É4LÊB	L#ÌL#Ì#L,c                 óT   — t         j                  | «      }| j                  |||«       y r  )rF  Úbuildr%  )r.  r°   r  rH  s       rW   Úcodegen_loopszCppKernel.codegen_loops	  s"   € Ü—N‘N 4Ó(ˆ	Ø×Ñ 	¨4°Õ=rV   c                 ó:   — t         j                  j                  ryy)NÚAOTI_TORCH_CHECKÚTORCH_CHECK)r4   rÓ  Úaot_moder3  s    rW   Úassert_functionzCppKernel.assert_function		  s   € ä�7‰7×ÒØ%à rV   c                 ó  — | j                   €J ‚| j                   |j                  |j                  |j                  z    }| j                  «       }d}d}|D ]m  }t        j
                  j                  j                  |d¬«      }|d|z  k\  s||k(  r n3||z  t        j                  j                  k  r n|dz  }||z  }||z  }Œo t        j                  j                  r|dk(  rt        |«      dkD  rd}t        ||j                  ¬«      S )Nr5   r   rò  ró  r   ©r
  r  )r|  r  r
  rö  r4   rÓ  rõ  r   r¹  Úmin_chunk_sizer”  rÐ   r	  )	r.  r  r  rM  ÚseqÚparr  r�  Úhints	            rW   r  zCppKernel.decide_parallel_depth	  s  € Ø×ÑÐ+Ð+Ð+Ø×!Ñ!Ø×*Ñ*Ø"×.Ñ.Ð1C×1RÑ1RÑRð
ˆð
 �n‰nÓˆØˆØˆØò 		ˆDÜ—7‘7×#Ñ#×-Ñ-¨d¸TÐ-ÓBˆDØ�a˜'‘kÒ! S¨G¢^ÙØ�g‰~¤§
¡
× 9Ñ 9Ò9áØ�Q‰JˆEØ�4‰KˆCØ�4‰K‰Cð		ô �:‰:×%Ò%¨%°1ª*¼¸V»ÀqºØˆEÜØ Ð.@×.LÑ.Lô
ð 	
rV   c              #   ó  K  — | j                   | j                  | j                  | j                  f}t	        «       | _         t	        «       | _        t	        «       | _        | j                  j                  «       | _        d –— | j                  j                  | j                   «       | j                  j                  | j                  «       | j                  j                  | j                  «       |\  | _         | _        | _        | _        y ­wr  )r¡  r‹  r›  r‰  r=   Úcloner€  r  )r.  r§  s     rW   Úwrite_to_suffixzCppKernel.write_to_suffix-	  sµ   è ø€ à—‘˜TŸ\™\¨4¯;©;¸¿¹ÐAˆÜ#Ó%ˆŒ
Ü%Ó'ˆŒÜ$Ó&ˆŒØ—8‘8—>‘>Ó#ˆŒÛØ×Ñ×$Ñ$ T§Z¡ZÔ0Ø×Ñ×$Ñ$ T§\¡\Ô2Ø×Ñ×$Ñ$ T§[¡[Ô1Ø<AÑ9ˆŒ�T”\ 4¤;°µùs   ‚D
Dc                 ó   — t        |i |¤ŽS r  )rG   )r.  r‘  r’  s      rW   Úcreate_cse_varzCppKernel.create_cse_var:	  s   € Ü˜tÐ. vÑ.Ð.rV   c                 ó"   — dt         |   › d|› d�S )Nzc10::convert<r“  rš   )rH   )r.  Úsrcrˆ   r§   s       rW   rˆ  zCppKernel.get_to_dtype_expr=	  s   € Øœ|¨EÑ2Ð3°2°c°U¸!Ð<Ð<rV   c                 ób   — | j                  |||«      }| j                  j                  ||«       y r  )rˆ  r‰  r  )r.  ÚdstÚ	dst_dtyper:  r§   r�  s         rW   r�  zCppKernel.cache_dtype_convert@	  s(   € Ø×%Ñ% c¨9°iÓ@ˆØ�‰�‰�T˜3ÕrV   r  r¥   c                 óV  ‡ ‡
— |€d}‰ j                   syg Š
ˆ
ˆ fd„}|�-|‰ j                   v sJ ‚‰ j                   |   \  }} ||||«      s4y‰ j                   j                  «       D ]  \  }}|\  }} ||||«      rŒ y dj                  ‰
«      }	|	r|j                  d|› d|	› d�«       yy)	NrR   Tc                 ó8  •— | |k(  ryd }t        ‰j                  «      D ]  \  }}||k(  sŒ|} n t        ‰«      t        k(  r|r| dk(  r|‰j                  |   k(  rd}‰j                  |› dt        | «      › �«       ‰j                  |› dt        |«      › �«       y)NFr   r5   rÍ  r¹   T)rÂ   rH  r  ru  rM  rŽ  rE   )r¢  r£  r¥   Úvar_idrã   Ú_varÚ
conditionsr.  s         €€rW   Úgenz)CppKernel.codegen_conditions.<locals>.genP	  s¬   ø€ Ø˜Š|ØØˆFÜ$ T§]¡]Ó3ò ‘��4Ø˜$“;Ø�FÙðô
 �T“
œiÒ'ÙØ˜Q’JØ˜4Ÿ;™; vÑ.Ò.à�Ø×Ñ   T¬+°eÓ*<Ð)=Ð>Ô?Ø×Ñ   S¬°SÓ)9Ð(:Ð;Ô<ØrV   FrF  zif(rº  rŸ   )rz  rŸ  Újoinr¿   )r.  r°   r  r¥   rC  r¢  r£  rA  Ú_rangeÚjoined_conditionsrB  s   `         @rW   Úcodegen_conditionszCppKernel.codegen_conditionsD	  sÔ   ù€ ð ˆ>ØˆFØ×!Ò!ØØˆ
õ	ð& ˆ?Ø˜$×,Ñ,Ñ,Ð,Ð,Ø×+Ñ+¨CÑ0‰JˆE�3Ù�u˜c 3Ô'Øà $× 2Ñ 2× 8Ñ 8Ó :ò !‘��fØ#‘
��sÙ˜5 # tÕ,Ù ð!ð #ŸK™K¨
Ó3ÐÙØ�N‰N˜S  ¨Ð*;Ð)<¸BÐ?Ô@ØàrV   r  )r5   éÿÿÿÿr   ©NN)7r  r  r  ry  r@  rD   rÃ  r‡  rˆ  r-  r­   r�   r™  r�  r   r7   r¤  r½   ÚcontextmanagerrD  rè   r…  r«  ro   r®  rô   r·  r¹  r¼  rk   rÈ  rí  r×  r	   r9   r„   rˆ   rÛ  rm   rÞ  rè  rê  rð  rö  r%  r(  Úpropertyr-  r  r6  r8  rˆ  r�  rG  r`  ra  s   @rW   ru  ru  ;  s   ø„ Ø€IØ€EØ€MØ€Fô#1ðV /Ø(ó'
òRIñ˜X lÑ3ó ð( ×Ññ$ó ð$ð& BCñØ—Z‘Zóð2 %§*¡*ð 2°ó 2ð	
¨u¯z©zð 	
ÀEÇLÁLó 	
ð
 e§j¡jð 
¸5¿<¹<ó 
ò
5ð:à�j‰jð:ð �j‰jð:ð ð	:ð
 ó:ð@˜ð  U§Z¡Zó ó8ð$à�; Ð#Ñ$ðð ðð ð	ð
 �{‰{óñ:7¨h°s©mó 7òò>
ò
ò&
òM%ò^>ð ð! ò !ó ð!ò
ð: ×Ññ
Bó ð
Bò/ò=ò ð !%Ø&*ñ	.àð.ð ˜‘ð.ð �e—l‘lÑ#÷	.rV   ru  c                   óJ  ‡ — e Zd ZeZ	 d)ˆ fd„	Zdej                  dej                  fd„Z	de
j                  defd„Zde
j                  defd	„Zde
j                  defd
„Ze
j                  fde
j                  defd„Zdede
j                  defd„Z	 d)dedej                  de
j                  dee   fd„Z	 	 	 d*dee   dej                  de
j                  dee   deeeef      dedee   fd„Zdedej                  fˆ fd„Z	 d+deeef   dedej                  de
j                  def
d„Zd)d„Zd„ Zd„ Zdedefd„Z dedej                  defd „Z!d!„ Z"d"„ Z#d)d#„Z$ddde
jJ                  fdeej                     d$ee   d%ee
j                     fd&„Z&d)ˆ fd'„	Z'ˆ fd(„Z(ˆ xZ)S ),r�  Nc                 óØ   •— t         ‰| �  ||«       t        j                  «       | _        | j                  sJ ‚|dkD  sJ d«       ‚|| _        || _        || _        |r|| _        y || _        y )Nr   z0Expect pass in Non-Zero tiling_factor explicitly)	r,  r-  r   Úpick_vec_isaÚvec_isar  rI  rX  Ú	num_elems)r.  r‘  r�  r  rI  rX  r1  s         €rW   r-  zCppVecKernel.__init__x	  si   ø€ ô 	‰Ñ˜˜{Ô+Ü"×/Ñ/Ó1ˆŒØ�|Š|Ðˆ|Ø˜qÒ ÐTÐ"TÓTÐ Ø*ˆÔØ$ˆŒØ"ˆŒÙ&/˜ˆ�°]ˆ�rV   r’   r¯  c                 óì   ‡ — ‰ j                  ||«      ry ˆ fd„|j                  D «       D ]"  }t        |t        «      sJ ‚|j                  sŒ" y  t        ||‰ j                  «      }|j                  r|S d S )Nc              3   ó˜   •K  — | ]A  }t        |t        j                  «      r%‰j                  j                  |j
                     –— ŒC y ­wr  ©r   r   rÂ  r‰  r²  rÄ   ©r  r´  r.  s     €rW   r  z5CppVecKernel._try_get_const_stride.<locals>.<genexpr>Œ	  s:   øè ø€ ò 
àÜ˜a¤§¡Ô*ð �H‰H× Ñ  §¡Õ(ñ
ùó   ƒAA
)r·  rµ  r    rG   r¤   r  r  r‹  )r.  r’   r¯  Úindirect_varrN  s   `    rW   rJ  z"CppVecKernel._try_get_const_stride‰	  s|   ø€ Ø×)Ñ)¨%°Ô9Øó
à×'Ñ'ô
ò 	ˆLô
 ˜l¬NÔ;Ð;Ð;Ø×"Ó"Ùð	ô % U¨G°T×5GÑ5GÓHˆØ×)Ò)ˆvÐ3¨tÐ3rV   rˆ   rµ   c                 ó¨   — t        j                  | j                  |j                  z  dz  | j                  j                  «       z  «      }|dk\  sJ ‚|S )Né   r5   )ÚmathrÔ  r  ÚitemsizerO  Ú	bit_width©r.  rˆ   Únum_vectorss      rW   rY  zCppVecKernel._get_num_vectors—	  sO   € Ü—i‘iØ×Ñ §¡Ñ/°!Ñ3°d·l±l×6LÑ6LÓ6NÑNó
ˆð ˜aÒÐÐØÐrV   c                 óp   — | j                   |j                  z  dz  | j                  j                  «       z  S )NrX  )r  rZ  rO  r[  )r.  rˆ   s     rW   r  z!CppVecKernel._get_raw_num_vectorsž	  s0   € ð ×!Ñ! E§N¡NÑ2°QÑ6¸¿¹×9OÑ9OÓ9QÑQÐQrV   c                 óh   — | j                  |«      }|dk(  rdt        |   › d�S dt        |   › d|› d�S )Nr5   rR  r�   rS  ra  )rY  rH   r\  s      rW   r;  zCppVecKernel._get_vec_type£	  sJ   € Ø×+Ñ+¨EÓ2ˆØ˜!ÒØ)¬,°uÑ*=Ð)>¸aÐ@Ð@à*¬<¸Ñ+>Ð*?¸qÀÀÈQÐOÐOrV   c                 ól   — |t         j                  k(  ry| j                  |«      }dt        |   › d|› d�S )NrR   r`  ra  r�   )r„   rk   rY  rH   r\  s      rW   r½  zCppVecKernel._get_mask_typeª	  s<   € Ø”E—J‘JÒØØ×+Ñ+¨EÓ2ˆØ"¤<°Ñ#6Ð"7°q¸¸ÀQÐGÐGrV   r?  c                 ó¢   — |j                   t        j                  k(  sJ t        |«      «       ‚| j	                  |«      }|› dt
        |   › d|› d�S )Nz.template cast<ra  r‚   )rˆ   r„   rk   ÚreprrY  rH   )r.  r?  rˆ   r]  s       rW   r(  zCppVecKernel._get_mask_cast°	  sP   € Ø�z‰zœUŸZ™ZÒ'Ð3¬¨d«Ó3Ð'Ø×+Ñ+¨EÓ2ˆØ��¤|°EÑ':Ð&;¸1¸[¸MÈÐMÐMrV   r¥   Ú	load_maskc                 óæ  — t         |   }| j                  |«      }d}|rS|j                  s&| j                  t        j
                  «      › d|› d�}n!| j                  |t        j
                  «      › }|dk7  r|› dt        |«      › �n|}|t        j                  k(  r| j                  «       › d|› d�}	|	S |r|› d|› d|› d|› d�n,| j                  |«      › d	|› d
t        | j                  «      › d�}	|	S )a  
        Get a load line str that loads a vector from `var` at `index` of type `dtype`.
        If `load_mask` is not None, we do a masked load accordingly.
        Notes on the `dtype`:
        1. We always load `self.tiling_factor` number of elements regardless of the `dtype`.
           It means we load half of the vector lanes for 16-bit data types and quarter of the
           vector lanes for 8-bit data types.
        2. `torch.bool` and `torch.uint8` could mean masks and we load them as float mask vectors.
        Nr:  rš   r   r}  z.template loadu<ra  r“  ú::loadu(r™   )rH   rY  r¤   r½  r„   rn   r(  rE   rk   r;  rP  )
r.  r¥   r’   rˆ   rc  Úcpp_typer]  Úload_mask_strÚloadbufrÅ   s
             rW   Ú_get_vec_load_linezCppVecKernel._get_vec_load_lineµ	  s  € ô    Ñ&ˆØ×+Ñ+¨EÓ2ˆØˆÙØ×#Ò#à#'×#6Ñ#6´u·{±{Ó#CÐ"DÀGÈIÈ;ÐVWÐ X‘à#'×#6Ñ#6°yÄ%Ç+Á+Ó#NÐ"O�Ø5:¸a²Z�S�E˜œ[¨Ó/Ð0Ñ1ÀSˆØ”E—J‘JÒà×)Ñ)Ó+Ð,¨G°G°9¸AÐ>ˆDð ˆñ !ð !�/Ð!1°(°¸1¸[¸MÈÈGÈ9ÐTUÑVà×*Ñ*¨5Ó1Ð2°(¸7¸)À2ÄkÐRV×R`ÑR`ÓFaÐEbÐbcÐdð ð
 ˆrV   Frà   Ústore_valueÚ
accu_storec                 óZ  ‡ ‡‡‡— |r	|€J d«       ‚|r|sJ ‚‰€‰ j                   Šdt        j                  dt        fˆ fd„Šdt        j                  dt        fˆ fd„Šdt        dt        fˆˆˆˆ fd„}t        «       }|j                  d	«       |j                  «       5   ‰|«      }	 ‰|«      }
d
t        |   › d|
› d�}|j                  |«       |r |j                  |› dt        |	«      › d�«       t        ‰ j                  ‰ j                     › d�«      }i }ˆ fd„|j                  D «       D ]4  }t        |t        «      sJ ‚|j                  sŒ" ||«      }|› d|› d�||<   Œ6 ‰ j!                  |‰ j                  |¬«      }d}‰ j"                  �l|rJ d«       ‚t        ‰ j"                  t        «      sJ ‰ j"                  «       ‚‰ j"                  j                  r‰ j"                  › d|› d�}n‰ j"                  › d�}t%        j&                  «       r|j                  d‰ j(                  › �«       n|j                  d‰ j(                  › �«       |j                  d|› d�|› dt        ‰ j*                  «      › d�z   |› d�z   «       |j                  «       5  t-        j.                  «       5 }t        |«      }|D ]#  }t1        j2                  d|› z   dz   ||   |«      }Œ% |�|› d|› d�n|› }|r4|j                  d |› d�«       |j5                  |j                  «       «       |r!|rd!nd"}|j                  |› d#|› d$|› d%�«       n|j                  d&|› d'|› d(�«       ddd«       ddd«       |s(‰ j7                  d)d*|«      }|j                  d+|› d(�«       ddd«       |j                  d,«       |r#|j                  d(«       ‰j9                  |«       y‰ j:                  j=                  ‰|«      }t        |t        «      sJ ‚d-|_        |S # 1 sw Y   Œ°xY w# 1 sw Y   Œ´xY w# 1 sw Y   ŒŽxY w).a÷  
        Load or store a vector in a non-contiguous way. The vector is initialized from an array that is
        filled in an inner loop over the tiling factor.
        :param var: buffer to load from or store to, i.e. `var[transformed(index)]`. If None, we load the index
                    as index expression, i.e. `transformed(index)`.
        :param index: index into the `var` or the index expression by its own if `var` is None.
                      The `index` could contain indirect indexing or the tiling itervar. When used in
                      the inner loop, the index is transformed as follows:
                      1. the index is linearized along the tiling dim.
                      2. the indirect indexing vector variables are transformed into arrays over the tiling dim.
        :param dtype: data type of `var` or `index` if `var` is None.
        :param buffer: the code buffer to write the generated code to. If None, we write to `self.loads`.
        :param store_value: the value to store. If None, we load the vector.
        :param accu_store: whether accumulate the store_value to store_ptr. If True, a store_value should be provided
        :return: a CppCSEVariable that represents the loaded vector or None if it is a store.
        Nzstore var must be providedrˆ   rµ   c                 ór   •— | j                   dk  r‰j                  d| j                   z  z  S ‰j                  S ©Né   )rZ  rP  ©rˆ   r.  s    €rW   Úget_result_sizezCCppVecKernel._load_or_store_non_contiguous.<locals>.get_result_sizeù	  s1   ø€ Ø�~‰~ Ò!Ø—~‘~¨¨e¯n©nÑ)<Ñ=Ð=à—~‘~Ð%rV   c                 ór   •— | j                   dk  r‰j                  d| j                   z  z  S ‰j                  S rn  )rZ  r  rp  s    €rW   Úget_tiling_sizezCCppVecKernel._load_or_store_non_contiguous.<locals>.get_tiling_sizeÿ	  s5   ø€ Ø�~‰~ Ò!Ø×)Ñ)¨Q°%·.±.Ñ-@ÑAÐAà×)Ñ)Ð)rV   Úvec_varc                 óL  •— | j                   sJ ‚t        «       }|j                  d«       |j                  «       5  | j                  }|€J ‚|t
        j                  k(  rt
        j                  } ‰|«      } ‰	|«      }|j                  dt        |   › d|› d�«       | › dt        |«      › d�}|j                  |«       |j                  d«       d d d «       |j                  d«       ‰
j                  j                  ‰|«      }t        |t        «      sJ ‚|S # 1 sw Y   ŒJxY w)	Nr=  rU  r™   rV  rW  rÓ   zreturn tmpbuf;rf  )r¤   r7   r¿   rÁ   rˆ   r„   rk   rn   rH   rE   r‰  rŠ  r    rG   )rt  r°   rb  Úresult_sizeÚtiling_sizerÅ   r‘  rà   rq  rs  r.  s          €€€€rW   Úvec_to_arrayz@CppVecKernel._load_or_store_non_contiguous.<locals>.vec_to_array
  s  ø€ Ø—>’>Ð!�>Ü“>ˆDØ�N‰N˜5Ô!Ø—‘“ñ 1Ø#ŸM™M�	Ø Ð,Ð,Ð,Ø¤§
¡
Ò*Ü %§¡�IÙ-¨iÓ8�Ù-¨iÓ8�Ø—‘Ø.¬|¸IÑ/FÐ.GÀrÈ+ÈÐV_Ð`ôð "˜Ð"8¼À[Ó9QÐ8RÐRTÐU�Ø—‘˜tÔ$Ø—‘Ð/Ô0÷1ð �N‰N˜4Ô Ø—X‘X×&Ñ& v¨tÓ4ˆFÜ˜f¤nÔ5Ð5Ð5ØˆM÷!1ð 1ús   »BDÄD#r=  rU  r™   rV  rW  rÓ   rD  c              3   ó˜   •K  — | ]A  }t        |t        j                  «      r%‰j                  j                  |j
                     –— ŒC y ­wr  rS  rT  s     €rW   r  z=CppVecKernel._load_or_store_non_contiguous.<locals>.<genexpr>,
  s:   øè ø€ ò !àÜ! !¤T§X¡XÔ.ð —‘×$Ñ$ Q§V¡VÕ,ñ!ùrU  rÊ  rË  ©rª  r;  zunexpected store with load maskz.is_masked(rš   z != 0z#pragma GCC unroll z#pragma unroll ú
for (long ú = 0; r¹   rº   r»   r¼   r>  ú+=ú=r•   z tmpbuf[rÕ   ztmpbuf[rÏ  r¸   ztmpbuf.data()r   r=  rf  T)r¡  r„   rˆ   rm   rG   r7   r¿   rÁ   rH   rE   r-   rH  rI  rµ  r    r¤   r«  r¦  r   Úis_gccr  rP  r½   r¾   rÆ   rÇ   rÀ   ri  r  r‰  rŠ  )r.  r¥   r’   rˆ   rà   rj  rk  rx  r°   rv  rw  Úresult_declareÚitervar_innerÚreplacementsrV  Ú	array_varrc  rÉ   Úindex_cÚrhsr©   Ú	load_liner‘  rq  rs  s   `   `                  @@rW   rL  z*CppVecKernel._load_or_store_non_contiguousÚ	  s[  û€ ñ2  # /ÐOÐ3OÓOÐ1ÙÙÐ�;Øˆ>Ø—Z‘ZˆFð	&¤5§;¡;ð 	&´3õ 	&ð	*¤5§;¡;ð 	*´3õ 	*ð	¤.ð 	´^÷ 	ð 	ô, ‹~ˆØ�‰�uÔØ�[‰[‹]ñ ?	7Ù)¨%Ó0ˆKÙ)¨%Ó0ˆKà*¬<¸Ñ+>Ð*?¸rÀ+ÀÈiÐXð ð �N‰N˜>Ô*ÙØ—‘Ø"�mÐ#9¼+ÀkÓ:RÐ9SÐSUÐVôô /Ø—=‘= §¡Ñ1Ð2°&Ð9óˆMð ˆLó!à×+Ñ+ô!ò Q�ô
 " ,´Ô?Ð?Ð?Ø×&Ó&Ù ,¨\Ó :�IØ4=°;¸aÀ¸ÈaÐ1P�L Ò.ðQð ×0Ñ0Ø 4§?¡?¸=ð 1ó ˆEð ˆIØ�‰Ð*Ù&ÐIÐ(IÓI�Ü! $§/¡/´>ÔBÐSÀDÇOÁOÓSÐBØ—?‘?×)Ò)Ø#'§?¡?Ð"3°;¸}¸oÈQÐ O‘Ià#'§?¡?Ð"3°5Ð 9�IÜ×!Ñ!Ô#Ø—‘Ð!4°T×5GÑ5GÐ4HÐIÕJà—‘ °×1CÑ1CÐ0DÐEÔFØ�N‰NØ˜]˜O¨6Ð2Ø"�O 3¤{°4·>±>Ó'BÐ&CÀ2ÐFñGà"�O 3Ð'ñ(ôð
 —‘“ñ H¤
× 4Ñ 4Ó 6ð H¸%Ü% eÓ,�Ø$0ò �LÜ Ÿf™fØ < .Ñ1°EÑ9Ø$ \Ñ2Øó‘Gðð .1¨_˜˜˜Q˜w˜i qÑ)ÀWÀI�ÙØ—N‘N T¨)¨°AÐ#6Ô7Ø×'Ñ'¨¯©«Ô6ÙÙ*4¡$¸#�KØ—N‘N c U¨!¨K¨=¸ÀÀÈrÐ#RÕSà—N‘N W¨]¨O¸4À¸uÀAÐ#FÔG÷!H÷ Hñ" Ø ×3Ñ3°OÀQÈÓN�	Ø—‘ ¨¨°1Ð5Ô6÷?	7ð@ 	�‰�tÔÙØ�N‰N˜3ÔØ�M‰M˜$ÔØà—X‘X×&Ñ& v¨tÓ4ˆFÜ˜f¤nÔ5Ð5Ð5Ø ˆFŒMØˆM÷;Hð Hú÷ Hð Hú÷Y?	7ð ?	7úsE   ÂB1P!ÅEP!ÊPÊ(B2P	ÍPÍ"2P!Ð	PÐPÐP	ÐP!Ð!P*rÄ   c                 ó6  •— | j                   j                  |«      }| j                  |«      }t        j                  j                  |«      }| j                  | j                     }| j                  ||«      }|dk(  rt        ‰	| �)  ||«      S |dk(  rE| j                  |||| j                  «      }| j                  j                  | j                  |«      }n| j!                  |||«      }t#        |t$        «      sJ ‚|j'                  d| ||fi «       d|_        |S )Nr   r5   rí  T)r‘  rÌ  r8  r4   rÓ  rÔ  rH  rI  rJ  r,  rí  ri  r¦  r‰  rŠ  r¡  rL  r    rG   rŒ  r¤   )
r.  rÄ   r’   r¥   rˆ   rM  rN  rÅ   r‘  r1  s
            €rW   rí  zCppVecKernel.loadh
  sù   ø€ Ø�i‰i�o‰o˜dÓ#ˆØ×$Ñ$ UÓ+ˆÜ—‘×!Ñ! $Ó'ˆØ—]‘] 4§?¡?Ñ3ˆ
Ø×+Ñ+¨E°:Ó>ˆØ�QŠ;ä‘7‘<  eÓ,Ð,Ø�qŠ[à×*Ñ*¨3°°u¸d¿o¹oÓNˆDØ—X‘X×&Ñ& t§z¡z°4Ó8‰Fà×7Ñ7¸¸UÀEÓJˆFÜ˜&¤.Ô1Ð1Ð1Ø×Ñ˜f t¨T°5Ð&9¸2Ô>ØˆŒØˆrV   r�  c           	      óØ  — t        |t        «      s#t        |t        «      r|j                  sJ |«       ‚| j                  | j
                     }|› dt        |«      › �}| j                  ||«      }t        «       }	|dk(  rÕ|rl|t        j                  k(  r#| j                  €| j                  |«      › d|› d�n,| j                  |«      › d|› dt        | j                  «      › d�}
d|› d|
› d�}|t        j                  k(  r%| j                  €|	j                  |› d|› d�«       |	S |	j                  |› d|› dt        | j                  «      › d�«       |	S | j                  ||||	||¬	«       |	S )
a2  
        Get a store line buffer that stores `value` into `var` at `index` of `dtype`. It handles
        both contiguous and non-contiguous store cases.
        :param value: Vectorized type templaterized on `dtype`.
        :param var: buffer to store into.
        :index: index into the `var`.
        r}  r5   re  rš   r™   rº  ú.store(rÓ   )rà   rj  rk  )r    ro   rG   r¤   rH  rI  rE   rJ  r=   r„   rn   rX  r;  rP  r¿   rL  )r.  r�  r¥   r’   rˆ   rk  rM  Úvar_exprrN  r°   rí  s              rW   Ú_get_store_linezCppVecKernel._get_store_line|
  s‘  € ô" ˜%¤Ô%Ü�uœnÔ-°%·,²,ð	àó	ð 
ð —]‘] 4§?¡?Ñ3ˆ
Ø�U˜#œk¨%Ó0Ð1Ð2ˆØ×+Ñ+¨E°:Ó>ˆÜÓˆØ�QŠ;Ùð ¤§¡Ò+°·±Ð0Fð ×)Ñ)¨%Ó0Ð1°¸(¸À1ÑEà ×.Ñ.¨uÓ5Ð6°h¸x¸jÈÌ;ÐW[×WeÑWeÓKfÐJgÐghÐið ð
 ˜E˜7 # d V¨1Ð-�ØœŸ™Ò#¨¯©Ð(>Ø—‘ % ¨°¨z¸Ð<Ô=ð ˆð —‘Ø�g˜W X J¨b´¸T¿^¹^Ó1LÐ0MÈRÐPôð ˆð ×.Ñ.Ø�U˜E¨$¸EÈjð /ô ð ˆrV   c                 ó  ‡— d‰v sJ ‚t        |t        «      sJ |«       ‚|j                  s| j                  |«      }| j                  j                  ‰«      }| j                  |«      }t        j                  j                  ‰«      }|€B| j                  ||||«      }| j                  j                  |j                  ˆfd„«      «       y |dk(  �r.t        j                  j                   sT| j"                  dk(  rE| j                  |› |||d¬«      }| j                  j                  |j                  ˆfd„«      «       y | j%                  |«      }| j%                  t&        j(                  «      }	t*        |   }
t-        j.                  |t&        j(                  «      j0                  }t        |t        «      r|j                  sJ ‚d|
› d	|	› d	|› d
|› d	|› d	|› d�}| j                  j3                  t5        ‰|«      «       y t7        d|› �«      ‚)NrÎ  c                 ó   •— t        ‰| «      S r  ©r;   ©rŽ  rÄ   s    €rW   ú<lambda>z$CppVecKernel.store.<locals>.<lambda>³
  s   ø€ ´,¸tÀQÓ2G€ rV   rÐ  r5   T)rk  c                 ó   •— t        ‰| «      S r  rŽ  r�  s    €rW   r�  z$CppVecKernel.store.<locals>.<lambda>½
  s   ø€ ´lÀ4ÈÓ6K€ rV   zatomic_add_vec<r™   r“  rÓ   rÑ  )r    rG   r¤   r�  r‘  rÒ  r8  r4   rÓ  rÔ  r‹  r›  r  Úmapr   r¹  r”  r�  rY  r„   rŒ  rH   r2   r:  r�  r¿   r;   rÕ  )r.  rÄ   r’   r�  rÖ  r¥   rˆ   r°   Ún_srcÚn_idxr‰   rÅ   s    `          rW   r×  zCppVecKernel.store¨
  sÇ  ø€ Ø˜‰}Ðˆ}Ü˜%¤Ô0Ð7°%Ó7Ð0Ø�|Š|à—N‘N 5Ó)ˆEØ�i‰i×Ñ˜tÓ$ˆØ×$Ñ$ UÓ+ˆÜ—‘×!Ñ! $Ó'ˆØˆ<Ø×'Ñ'¨¨s°E¸5ÓAˆDØ�K‰K×Ñ˜tŸx™xÓ(GÓHÕIØ�\Ó!Ü—:‘:×-Ò-°$×2BÑ2BÀaÒ2GØ×+Ñ+Ø�gØØØØ#ð ,ó �ð —‘×"Ñ" 4§8¡8Ó,KÓ#LÕMà×-Ñ-¨eÓ4�Ø×-Ñ-¬e¯k©kÓ:�Ü% eÑ,�ÜŸ™ u¬e¯k©kÓ:×@Ñ@�Ü! %¬Ô8¸U¿\º\ÐIÐIØ(¨¨°°5°'¸¸E¸7À"ÀSÀEÈÈEÈ7ÐRTÐUZÐT[Ð[]Ð^�Ø—‘×%Ñ%¤l°4¸Ó&>Õ?ä%¨°D°6Ð&:Ó;Ð;rV   c           
      ó†  — |t         v sJ ‚|dv }| j                  | j                  k\  }|r|n|}t        |t        «      sJ |«       ‚|j
                  s| j                  |«      }|||f}|| j                  j                  v r| j                  j                  |   S d}	|	› dt        |   › d�}
t        ||«      }| j                  ||«      }| j                  j                  | j                  d|› �d¬«      }t        |t        «      sJ ‚|› d�}d	|› �}| xj                  |› ||gz  c_        d
| _        | j                   j#                  | j%                  ||||t&        «      «       | j                   j#                  | j%                  ||||| j(                  «      «       t+        j,                  d„ | j.                  | j                  d  «      }|dk(  �r| j                  €J ‚| j                   j#                  | j%                  ||||| j(                  «      «       t+        j,                  d„ | j.                  | j                  d  «      }| j                  | j                  k\  r| j0                  nd}t3        ||«      | _        | j4                  | j6                  j                  vrã| j6                  j                  | j8                  d| j4                  › �d¬«      | _        | j:                  | j6                  j                  | j4                  <   | j<                  j?                  | jA                  |«      «       tB        jD                  jF                  rdn	tI        «       }| jJ                  j?                  | jA                  ||«      «       n(| j6                  j                  | j4                     | _        | jL                  r|n|}| jN                  j?                  |› d| jQ                  |||d
«      › d�«       n´| j                  €J ‚| jR                  | j                     }tU        | j                  dz   tW        | jR                  «      «      D ]$  }|| j.                  |   z  | jR                  |   z   }Œ& ||||dœ}| jN                  j?                  |› d | jP                  ||fi |¤Ž› d�«       | jY                  ||||| jP                  | j(                  ¬«       | jY                  ||||tZ        t&        ¬«       |dk(  r+| jY                  ||||| jP                  | j(                  ¬«       |t\        j^                  k(  }|�r1ta        |«      rS| jc                  |«      dv sJ d«       ‚d|› d�}d|› d�}| jd                  j?                  |› dt[        |||«      › d�«       n¢|r	|› d|› d�}n—|r|dv rd|› d�}nŠ|dk(  sJ ‚|› d�}n}d| jQ                  |dd «      z   d!z   }|t\        j^                  k(  }|rt\        jf                  n|}d"t        |   › d�}
d#t        |   › d$| jc                  |«      › d�}|› d%|
› d&|
› d'|› d$|› d�
}| jd                  j?                  |› dt[        ||||¬(«      › d�«       |}n>|}ta        |«      r1d	|› �}| jd                  j?                  |› dt[        |||«      › d�«       ti        ||«      }|| j                  j                  |<   |S ))Nrà  zat::vecz::Vectorized<r�   rá  Frâ  Ú_vecÚmasked_Tc                 ó   — | |z  S r  rU   rï  s     rW   r�  z(CppVecKernel.reduction.<locals>.<lambda>õ
  s
   € ˜˜Q™€ rV   rg   c                 ó   — | |z  S r  rU   rï  s     rW   r�  z(CppVecKernel.reduction.<locals>.<lambda>  s
   € ˜Q ™U€ rV   r5   r’  r·   r¸   )r¦   r’   Úhorizontal_reductionr§   )r•  r–  )r5   r   z4Welford reduction does not support VectorizedN (N>2)zwelford_vec_reduce_all(rš   z_vec_reduce_all()rf   rc   r_   rJ  z.all_zero()r^   z.all_masked()z	{ return rŽ  rð  z; }rR  zat::vec::vec_reduce_all<r™   z([](z& x, z& y) r“  )5ÚVECTORIZABLE_RTYPESrI  r}  r    rG   r¤   r�  r‰  rä  rH   r‘   Úreduction_acc_type_vecrŠ  r¡  r�  r…  r  rŽ  rÛ  r�   Úreduction_init_vecrI  rJ  rM  r  r   Úweight_recp_vec_rangerŠ  r‹  Úweight_recps_valr†  r¿   Úwelford_weight_reciprocal_vecr   r¹  r”  r+   rƒ  rX  r›  Úreduction_combine_vecrH  rG  rÐ   r™  r­   r„   rk   r*   rY  r€  rn   r¢   ) r.  rˆ   r§   r‡   r�  rå  rš  rç  ræ  Úvec_nsÚvecrÏ   Úacc_type_vecr¯   Úacc_vecÚmasked_acc_vecÚreduction_sizeÚreduction_factorr�  Úacc_vec_r’   rã   r’  r¨   r¦   Úmasked_next_valueÚreduce_all_bodyrb  Úvec_reduce_all_funcÚtmpvarÚmasked_tmpvarrö   s                                    rW   rè  zCppVecKernel.reductionÉ
  s…  € ð Ô!4Ñ4Ð4Ð4Ø)Ð-AÐAÐØ#Ÿ™°$×2FÑ2FÑFÐÙ"2‘Y¸ˆ
Ü˜%¤Ô0Ð7°%Ó7Ð0à�|Š|Ø—N‘N 5Ó)ˆEà! >°5Ð8ˆØ˜D×.Ñ.×>Ñ>Ñ>Ø×%Ñ%×5Ñ5°mÑDÐDàˆØ�˜¤l°5Ñ&9Ð%:¸!Ð<ˆÜ% n°jÓAˆØ×2Ñ2°>À:ÓNˆà× Ñ ×)Ñ)Ø�J‰J˜* ] OÐ4¸Eð *ó 
ˆô ˜#œ~Ô.Ð.Ð.Ø�E˜�,ˆØ" 7 )Ð,ˆØ× Ò ¨ u¨w¸Ð$GÑGÕ Ø ˆÔØ×(Ñ(×/Ñ/Ø×&Ñ&Ø�X˜~¨z¼>óô	
ð
 	×(Ñ(×/Ñ/Ø×&Ñ&ØØØØØ×'Ñ'óô	
ô #×)Ñ)Ù §¡¨D×,@Ñ,@Ð,BÐ Có
ˆð Ð-Ó-à×'Ñ'Ð3Ð3Ð3à×,Ñ,×3Ñ3Ø×*Ñ*Ø"Ø Ø"ØØ×+Ñ+óôô '×-Ñ-Ù" D§K¡K°×0DÑ0DÐ0FÐ$GóˆNð '+§o¡o¸×9MÑ9MÒ&M�×"Ò"ÐSTð ô *2°.ÐBRÓ)SˆDÔ&Ø×)Ñ)°×1FÑ1F×1VÑ1VÑVØ(,×(=Ñ(=×(FÑ(FØ—L‘L J¨t×/IÑ/IÐ.JÐ"KÐSXð )Gó )�Ô%ð ×)Ñ)ð ×%Ñ%×5Ñ5°d×6PÑ6PÑQð ×2Ñ2×<Ñ<Ø×6Ñ6°uÓ=ôô —z‘z×1Ò1ñ "ä-Ó/ð ð
 ×)Ñ)×3Ñ3Ø×6Ñ6°u¸kÓJõð )-×(=Ñ(=×(MÑ(MØ×.Ñ.ñ)�Ô%ð *.¯ª‘~¸WˆHØ�K‰K×!Ñ!Ø�*˜C × :Ñ :¸>È8ÐUZÐ\`Ó aÐbÐbcÐdõð ×'Ñ'Ð3Ð3Ð3Ø—M‘M $×"6Ñ"6Ñ7ˆEÜ˜4×/Ñ/°!Ñ3´S¸¿¹Ó5GÓHò B�Ø §¡¨A¡Ñ.°·±¸qÑ1AÑA‘ðBð $ØØ(<Ø&ñ	ˆFð �K‰K×!Ñ!Ø�)˜3Ð9˜t×9Ñ9¸.È'Ñ\ÐU[Ñ\Ð]Ð]^Ð_ôð 	×,Ñ,ØØØØØ!%×!;Ñ!;Ø"×5Ñ5ð 	-ô 	
ð 	×,Ñ,ØØØØÜ!2Ü,ð 	-ô 	
ð Ð-Ò-à×0Ñ0ØØØØØ%)×%?Ñ%?Ø"&×"9Ñ"9ð 1ô ð œ5Ÿ:™:Ñ%ˆÚä# NÔ3Ø×,Ñ,¨UÓ3ð 8ñ ð Jð JóJð ð  7°w°i¸qÐA�
Ø&=¸nÐ=MÈQÐ$OÐ!Ø×%Ñ%×/Ñ/Ø�e˜3Ô0°ÀÐFWÓXÐYÐYZÐ[õñ "Ø .Ð/Ð/?À¸yÈÐJ‘
ÙØ!ð &ñ ð
 $% W I¨[Ð!9‘Jà)¨UÒ2Ð2Ð2Ø$+ 9¨MÐ!:‘Jð  Ø×0Ñ0°ÀÀcÓJñKàñð  ð
  ¤5§:¡:Ñ-�á+2œEŸKšK¸�	Ø,¬\¸)Ñ-DÐ,EÀQÐG�Ø(@ÄÈiÑAXÐ@YÐY[Ð\`×\qÑ\qÐr{Ó\|Ð[}Ð}~Ð&Ð#Ø 3Ð4°D¸¸¸UÀ3À%ÀuÈ_ÐL]Ð]_Ð`gÐ_hÐhiÐj�
à×!Ñ!×+Ñ+Ø�%�sÔ,¨^¸SÀ*ÐXaÔbÐcÐcdÐeôð ‰FàˆFÜ# NÔ3Ø")¨&¨Ð 2�Ø×%Ñ%×/Ñ/Ø�h˜cÔ"3°NÀFÈMÓ"ZÐ![Ð[\Ð]ôô # >°6Ó:ˆØ<Bˆ×Ñ×*Ñ*¨=Ñ9ØˆrV   c                 ó&  ‡— | j                  |«      }| j                  j                  ‰«      }t        j                  j                  ‰«      }|j                  r%|t        j                  k(  r|nt        j                  nt        j                  }t        j                  j                  |«      }t        j                  j                  |«      }t        «       }	| j                  | j                  k\  r.|	j!                  |› dt#        |«      › dt$        |   › d|› d�«       nÈ||k7  r t$        |   › d|› �}
|t        j&                  k(  r&|› d| j                  t        j&                  «      › d�}n?||cxk(  rdk(  rn nd	t$        |   › d|› d
�}n d	t$        |   › d|› dt$        |   › d|› d|› d
�}|	j!                  d|
› d|› d�«       |
}|	j)                  | j+                  ||||«      «       | j,                  j)                  |	j/                  ˆfd„«      «       y )NrÊ  z] = static_cast<r“  rÓ   rÊ   z.template cast<bool,r‚   r5   úat::vec::convert<rš   ra  rÒ   r·   r¸   c                 ó   •— t        ‰| «      S r  rŽ  r�  s    €rW   r�  z.CppVecKernel.store_reduction.<locals>.<lambda>ª  s   ø€ ¼ÀTÈ1Ó8M€ rV   )r8  r‘  rÒ  r4   rÓ  rÔ  Úis_floating_pointr„   rj   rn   rŒ  rZ  rY  r=   rI  r}  r¿   rE   rH   rk   r  r‹  r€  r’  )r.  rÄ   r’   r�  r¥   Ú	out_dtyperˆ   Úout_num_vectorsÚsrc_num_vectorsr°   Úconverted_valueÚconverts    `          rW   rê  zCppVecKernel.store_reduction†  s  ø€ Ø×$Ñ$ UÓ+ˆØ�i‰i×Ñ˜tÓ$ˆÜ—G‘G×%Ñ% dÓ+ˆ	ð ×*Ò*ð $¤u§|¡|Ò3‰Y¼¿ºä—‘ð 	ô
 Ÿ(™(×3Ñ3°IÓ>ˆÜŸ(™(×3Ñ3°EÓ:ˆÜÓˆØ�?‰?˜d×2Ñ2Ò2à�N‰NØ�%�qœ UÓ+Ð,Ð,<¼\È)Ñ=TÐ<UÐUWÐX]ÐW^Ð^`Ðaõð
 ˜EÒ!Ü%1°)Ñ%<Ð$=¸Q¸u¸gÐ"F�Ø¤§
¡
Ò*Ø!& Ð';¸D×<QÑ<QÔRW×R\ÑR\Ó<]Ð;^Ð^aÐb‘Gà&¨/Ô>¸QÕ>à/´¸YÑ0GÐ/HÈÈ5È'ÐQRÐSñ  ð
 0´¸YÑ0GÐ/HÈØ.Ð/¨q´¸eÑ1DÐ0EÀQÀÐFWÐWYÐZ_ÐY`Ð`aðcð  ð —‘  Ð&7°s¸7¸)À1ÐEÔFØ'�Ø�K‰K˜×,Ñ,¨U°C¸À	ÓJÔKØ×Ñ×$Ñ$ T§X¡XÓ.MÓ%NÕOrV   Ú
scalar_varc                 ó  — |j                   rJ ‚|j                  t        j                  k(  rE| j                  j                  | j                  | j                  «       › d|j                  › d�«      }n]|j                  €J ‚| j                  j                  | j                  | j                  |j                  «      › d|j                  › d�«      }t        |t        «      sJ ‚|j                  |_        |j                  |_        d|_         |S )Nr:  rš   rº  T)r¤   rˆ   r„   rk   r‰  rŠ  r‹  r½  rÄ   r;  r    rG   Údependent_itervars)r.  r¸  rt  s      rW   r�  zCppVecKernel.broadcast¬  së   € Ø×$Ò$Ð$Ð$Ø×ÑœuŸz™zÒ)Ø—h‘h×'Ñ'Ø—‘ ×!4Ñ!4Ó!6Ð 7°w¸z¿¹Ð>OÈqÐQó‰Gð ×#Ñ#Ð/Ð/Ð/Ø—h‘h×'Ñ'Ø—‘Ø×%Ñ% j×&6Ñ&6Ó7Ð8¸¸*¿/¹/Ð9JÈ!ÐLóˆGô ˜'¤>Ô2Ð2Ð2Ø"×(Ñ(ˆŒØ%/×%BÑ%BˆÔ"ØˆŒØˆrV   rN  c           	      ó"  — |j                   rJ ‚|j                  €J ‚| j                  j                  | j                  | j                  |j                  «      › d|› d|› d�«      }t        |t        «      sJ ‚|j                  |_        d|_         |S )Nz	::arange(r™   rš   T)r¤   rˆ   r‰  rŠ  r‹  r;  r    rG   )r.  r’   rN  r‘  s       rW   rK  zCppVecKernel.arange¾  s‰   € Ø—<’<ÐÐØ�{‰{Ð&Ð&Ð&Ø—‘×"Ñ"Ø�L‰LØ×!Ñ! %§+¡+Ó.Ð/¨y¸¸¸rÀ&ÀÈÐKó
ˆô ˜&¤.Ô1Ð1Ð1Ø—{‘{ˆŒØˆŒØˆrV   c                 óÈ  — t         |   }| j                  |«      }t        |«      rd|› d�S |dv rWt        |   }| j	                  ||«      }|dk(  rt        |«      rd|› d�nd|› d�}nt        |«      rd|› d�nd|› d	�}|› d
|› d�S |dk(  r| j                  «       › d�S t        ||«      }|› d
|› d�}	|t        j                  k(  r|dv sJ ‚| j                  «       › d|› d�S |	S )Nr�   r‚   ry   r`   r|   r{   r~   rz   r}   rº  rš   rf   z	::from(0))r^   r_   rc   r:  )
r<   r;  r*   rH   rœ  r   r½  r�   r„   rk   )
r.  r‡   rˆ   r�   Úvec_typer‰   rÏ   r3  Úscalar_initÚvec_inits
             rW   r�  zCppVecKernel.reduction_init_vecÊ  sA  € Ü0°Ñ7ˆØ×%Ñ% kÓ2ˆä Ô/Ø˜h˜Z sÐ+Ð+àÐ1Ñ1Ü! +Ñ.ˆFØ×2Ñ2°>À5ÓIˆHØ Ò)ô & eÔ,ð +¨6¨(°-Ñ@à/°¨x°xÐ@ñ ô & eÔ,ð ,¨F¨8°=ÑAà/°¨x°xÐ@ð ð
 �Z˜q   QÐ'Ð'à˜UÒ"Ø×)Ñ)Ó+Ð,¨IÐ6Ð6ä$ ^°UÓ;ˆØ�Z˜q  ¨QÐ/ˆØ”E—J‘JÒØ!Ð%:Ñ:Ð:Ð:Ø×)Ñ)Ó+Ð,¨G°K°=ÀÐBÐBØˆrV   c                 ó¦  — t         |   }| j                  |«      }t        |«      rd|› d�S |dv rw| j                  |«      }| j                  t        j
                  «      }|t        j                  k(  r!dt        t        j                     › d|› d|› d�S dt        |   › d|› d|› d�S |t        j                  k(  r|dv sJ ‚| j                  «       › S |S )Nr�   r�   ry   zIndexValueVec<r™   )r^   r_   rf   rc   )
r<   r;  r*   rY  r„   rŒ  rk   rH   rn   r½  )r.  r‡   rˆ   r�   r½  r“  r”  s          rW   rœ  z#CppVecKernel.reduction_acc_type_vecì  sæ   € Ü0°Ñ7ˆØ×%Ñ% kÓ2ˆÜ Ô/Ø˜h˜Z qÐ)Ð)ØÐ1Ñ1Ø×)Ñ)¨+Ó6ˆEØ×)Ñ)¬%¯+©+Ó6ˆEØœŸ
™
Ò"Ø'¬´U·[±[Ñ(AÐ'BÀ"ÀUÀGÈ2ÈeÈWÐTUÐVÐVØ#¤L°Ñ$=Ð#>¸bÀÀÀrÈ%ÈÐPQÐRÐRØ”E—J‘JÒØ!Ð%AÑAÐAÐAØ×)Ñ)Ó+Ð,Ð-ØˆrV   c                 óª   — |rt        | j                  |«      n| j                  }t        |«      }d| j                  |«      › d| j                  › d|› d�S )Nzstatic WeightRecp<rT  rº  rÓ   )r   rž  rE   r;  rŸ  )r.  rˆ   r�  Úvec_num_range_threadÚvec_num_range_thread_exprs        rW   r   z*CppVecKernel.welford_weight_reciprocal_vecü  sk   € ñ ô �D×.Ñ.°Ô<à×+Ñ+ð 	ô
 %0Ð0DÓ$EÐ!à  ×!3Ñ!3°EÓ!:Ð ;¸2¸d×>SÑ>SÐ=TØØ(Ð)Øðð	
rV   rš  r§   c                 ó  — |t         j                  k(  }|dk(  r=| j                  rd|› d|› dt        | j                  «      › d�S |r|› d|› �S d|› d|› d�S |dk(  r=| j                  rd|› d|› dt        | j                  «      › d�S |r|› d	|› �S d
|› d|› d�S |dk(  r;| j                  rd|› d|› dt        | j                  «      › d�S |rdnd}	|› d|	› d|› �S |dk(  r2| j                  rd|› d|› dt        | j                  «      › d�S |› d|› �S |dk(  r2| j                  rd|› d|› dt        | j                  «      › d�S |› d|› �S |dk(  r„|rN| j                  r,d|› d|› dt        | j                  «      › d| j                  › d�	S d|› d|› d| j                  › d�S | j                  rd|› d|› dt        | j                  «      › d�S d|› d|› d�S |dk(  rgt        |t        «      r|\  }
}}nt        ||«      \  }
}}| j                  r%d|› d|
› d|› d|› dt        | j                  «      › d�S d|› d|
› d|› d|› d�	S |dv rä|€J ‚t        |   }|t         j                  k(  rt        t         j                     }| j                  |«      }| j                  t         j                  «      }d}d}|�%|€J ‚dt        |«      j                  «       › �}d|› �}| j                  r.|› d|› d|› d|› |› d |› d|› |› dt        | j                  «      › d�S |› d|› d|› d|› |› d |› d|› |› d�S |d!k(  r]t        |t        «      rF|j                  t         j                  k(  sJ ‚t!        t"        j$                  j&                  |f«      \  }|› d|› �S t(        ‚)"Nr_   zmax_masked_reduce(r™   rš   rY  r#  r^   zmin_masked_reduce(rT  r  rc   zsum_masked_reduce(r”   rZ   r•   rd   zprod_masked_reduce(r–   re   zxor_sum_masked_reduce(r—   rg   r›   r  rh   rœ   z}, r�   ry   rR   z_combine_vec<r“  rf   )r„   rk   rX  rE   rŸ  r    r¡   r¢   rH   rn   rY  rŒ  ro   r½  rG   rˆ   rN   r4   rZ  r‹  rÕ  )r.  r‡   r¥   r¦   Úuse_weight_recpsr’   rš  r§   r¨   r©   rª   r«   r¬   r‰   r“  r”  Út_extraÚ	arg_extras                     rW   r¡  z"CppVecKernel.reduction_combine_vec
  sš  € ð œuŸz™zÑ)ˆØ˜UÒ"Ø�~Š~Ø+¨C¨5°°:°,¸bÄÈTÏ^É^ÓA\Ð@]Ð]^Ð_Ð_ñ ð �e˜3˜z˜lÐ+ðð -¨S¨E°°J°<¸qÐAðð
 ˜uÒ$Ø�~Š~Ø+¨C¨5°°:°,¸bÄÈTÏ^É^ÓA\Ð@]Ð]^Ð_Ð_ñ ð �e˜3˜z˜lÐ+ðð -¨S¨E°°J°<¸qÐAðð
 ˜uÒ$Ø�~Š~Ø+¨C¨5°°:°,¸bÄÈTÏ^É^ÓA\Ð@]Ð]^Ð_Ð_á%,™c°#�Ø˜˜a ˜}¨A¨j¨\Ð:Ð:Ø˜vÒ%Ø�~Š~Ø,¨S¨E°°J°<¸rÄ+ÈdÏnÉnÓB]ÐA^Ð^_Ð`Ð`à˜˜c * Ð.Ð.Ø˜yÒ(Ø�~Š~Ø/°¨u°B°z°lÀ"Ä[ÐQU×Q_ÑQ_ÓE`ÐDaÐabÐcÐcà˜˜c * Ð.Ð.ØÐ/Ò/ÙØ—>’>Ø-¨c¨U°"°Z°LÀÄ;ÈtÏ~É~ÓC^ÐB_Ð_bÐcg×cxÑcxÐbyÐyzÐ{Ð{à-¨c¨U°"°Z°LÀÀD×DYÑDYÐCZÐZ[Ð\Ð\à—>’>Ø-¨c¨U°"°Z°LÀÄ;ÈtÏ~É~ÓC^ÐB_Ð_`ÐaÐaà-¨c¨U°"°Z°LÀÐBÐBØÐ0Ò0Ü˜*¤eÔ,à#-Ñ ��b™&ô $5°^ÀZÓ#PÑ ��b˜&Ø�~Š~Ø)¨#¨¨d°4°&¸¸2¸$¸bÀÀÈÌ[ÐY]×YgÑYgÓMhÐLiÐijÐkÐkà)¨#¨¨d°4°&¸¸2¸$¸bÀÀÈÐLÐLØÐ3Ñ3ØÐ(Ð(Ð(Ü! )Ñ,ˆFØœEŸJ™JÒ&Ü%¤e§k¡kÑ2�Ø×)Ñ)¨)Ó4ˆEØ×)Ñ)¬%¯+©+Ó6ˆEØˆGØˆIØÐ Ø+Ð7Ð7Ð7ØœsÐ#7Ó8×>Ñ>Ó@ÐAÐB�Ø   ˜L�	Ø�~Š~à%Ð& m°F°8¸2¸e¸WÀBÀuÀgÈgÈYð WØ�u˜B˜z˜l¨9¨+°R¼ÀDÇNÁNÓ8SÐ7TÐTUðWðð
 )Ð)¨°v°h¸bÀÀÀrÈ%ÈÐQXÐPYÐY[Ð\_Ð[`Ð`bÐcmÐbnÐoxÐnyÐyzÐ{Ð{Ø˜uÒ$Ü˜*¤nÔ5Ø!×'Ñ'¬5¯:©:Ò5Ð5Ð5Ü 4´Q·X±X×5EÑ5EÈ
À}Ó U‘�Ø�U˜#˜j˜\Ð*Ð*ä%Ð%rV   c           	      ó€  •— t        |t        «      sJ ‚|j                  €J ‚|j                  s4t        |t        «      r|j                  rd|› d�}t        ‰	| �  ||||«      S |}|}|r!| j                  |j                  «      › d|› d�}|r!| j                  |j                  «      › d|› d�}|r|rd|› d|› d|› d|› d�	}|› d|› d|› �}n#|r|› d|› �}|› d|› �}n|sJ ‚|› d|› �}|› d|› �}| j                  |j                  «      › d|› d�}|r6|j                  s!| j                  |j                  «      › d|› d�}d|› d|› d�}| j                  rS| j                  |j                  «      › d| j                  |j                  «      › d	|› d
t        | j                  «      › d�}d|› d�}| j                  › d|› d|› d�S )Nrº  z).all_masked()rš   rÊ  z) & (r¹   z) | ~(z::set(z::from(1), (rž   z, "index out of bounds: z"))r    rG   rˆ   r¤   r,  rÄ  r;  r½  rX  rE   r-  )
r.  r¥   r½  r¾  r?  Úlower_scalarÚupper_scalarÚcondÚ
cond_printr1  s
            €rW   rÄ  zCppVecKernel.indirect_asserti  s$  ø€ Ü˜#œ~Ô.Ð.Ð.Ø�y‰yÐ$Ð$Ð$Ø�zŠzÜ˜$¤Ô/°D·K²KØ˜4˜& Ð/�Ü‘7Ñ*¨3°°u¸dÓCÐCØˆØˆÙØ×)Ñ)¨#¯)©)Ó4Ð5°Q°u°g¸QÐ?ˆEÙØ×)Ñ)¨#¯)©)Ó4Ð5°Q°u°g¸QÐ?ˆEÙ‘UØ�u�g˜T #  e¨C¨5°°E°7¸!Ð<ˆDØ(˜>¨¨c¨U°#°l°^ÐD‰JÙØ�W˜D  Ð&ˆDØ(˜>¨¨c¨UÐ3‰JáˆL�5Ø�U˜#˜e˜WÐ%ˆDØ˜5  L >Ð2ˆJØ×%Ñ% c§i¡iÓ0Ð1°°4°&¸Ð:ˆÙØ—;’;Ø×-Ñ-¨c¯i©iÓ8Ð9¸¸4¸&ÀÐB�à�t�f˜F 4 &¨Ð*ˆDØ�>Š>à×&Ñ& s§y¡yÓ1Ð2°&¸×9LÑ9LÈSÏYÉYÓ9WÐ8Xð YØ�V˜3œ{¨4¯>©>Ó:Ð;¸1ð>ð ð �4�&˜Ð'ˆØ×&Ñ&Ð' q¨¨Ð.FÀzÀlÐRTÐUÐUrV   c                 ó*  •— t        |t        «      sJ ‚|j                  st        ‰	| �  |||«      S t
        |   }| j                  |«      }t
        |   }| j                  |«      }d|› d�}|t        j                  k7  r2|t        j                  k(  r| j                  |«      › d|› d|› d|› d�}|S |t        j                  k(  r |t        j                  k7  r|› d|› d|› d�}|S ||k7  r+||cxk(  rdk(  rn nd	|› d|› d�}|S d	|› d|› d|› d|› d|› d�}|S )
Nrº  rš   z::from<ra  r“  z.to<r‚   r5   r°  )
r    rG   r¤   r,  rˆ  rH   rY  r„   rk   r½  )
r.  r:  rˆ   r§   Úsrc_cpp_typerµ  Údst_cpp_typeÚdst_num_vectorsr�  r1  s
            €rW   rˆ  zCppVecKernel.get_to_dtype_exprŽ  s\  ø€ Ü˜#œ~Ô.Ð.Ð.Ø�zŠzÜ‘7Ñ,¨S°%¸ÓCÐCÜ# IÑ.ˆØ×/Ñ/°	Ó:ˆÜ# EÑ*ˆØ×/Ñ/°Ó6ˆØ�3�%�qˆzˆØœŸ
™
Ò" u´·
±
Ò':Ø×)Ñ)¨)Ó4Ð5°W¸\¸NÈ!ÈOÐK\Ð\^Ð_bÐ^cÐcdÐeˆDð ˆð œ%Ÿ*™*Ò$¨´%·*±*Ò)<Ø�U˜$˜|˜n¨A¨oÐ->¸cÐBˆDð ˆð ˜%ÒØ /Ô6°QÕ6Ø*¨<¨.¸¸3¸%¸qÐA�ð ˆð +¨<¨.¸¸/Ð9JÈ!ÈLÈ>ÐYZÐ[jÐZkÐkmÐnqÐmrÐrsÐt�ØˆrV   r  )NNF)F)*r  r  r  r‡  r@  r-  rè   r…  rô   rJ  r„   rˆ   rm   rY  rn   r  ro   r;  r½  rG   r(  r   ri  r=   r	   rk   rL  rí  r‹  r×  rè  rê  r�  rK  r�  rœ  r   r…   r¡  rÄ  rˆ  r`  ra  s   @rW   r�  r�  u	  s˜  ø„ Ø€Ið õCð"4¨5¯:©:ð 4ÀÇÁó 4ð e§k¡kð °có ðR¨%¯+©+ð R¸%ó Rð
P 5§;¡;ð P°3ó Pð 38·+±+ñ H E§K¡Kð HÀ#ó HðN >ð N¸%¿+¹+ð NÈ#ó Nð /3ñ#àð#ð �z‰zð#ð �{‰{ð	#ð
 ˜NÑ+ó#ðT ,0Ø<@Ø ñLà�c‰]ðLð �z‰zðLð �{‰{ð	Lð
 ˜Ñ(ðLð ˜e C¨Ð$7Ñ8Ñ9ðLð ðLð 
�.Ñ	!óLð\˜ð  U§Z¡Zõ ð4 !ñ*à�S˜.Ð(Ñ)ð*ð ð*ð �z‰zð	*ð
 �{‰{ð*ð ó*óX<òB{òz$PðL Nð °~ó ð$
˜Nð 
°E·L±Lð 
À^ó 
ò òDó 
ð& Ø(,Ø/3Ø+0¯=©=ñ]&ð ˜Ÿ™Ñ%ð]&ð ' t™nð]&ð ˜EŸK™KÑ(ó]&õ~#V÷Jð rV   r�  c                   ó¾   ‡ — e Zd ZdZeZ	 	 dˆ fd„	Zd„ Zd„ Z	 dd„Z	de
dej                  fˆ fd„Zdˆ fd	„	Zd
„ Zˆ fd„Zdej                  dej                  fd„Zˆ xZS )rr  an  
    A vector kernel that handles the 2d tiles with the tile size defined in `tiling_factor` on
    the inner-most loop level and one of the outer loop level (`outer_tiling_idx`). When the data
    tile is accessed in a contiguous way from the outer loop axis, a transposition is applied on the
    tile to make the access contiguous from the inner-most loop axis. Then, the same vectorization
    logic from its parent `CppVecKernel` is leveraged for load/store/compute. The transposed tile load
    and store are generated into kernel.preloads and kernel.poststores buffers.

    The loop structure looks like below:
    for ...
      for i_outer ...
        for ...
          for inner_most ...
            // generated by CppTile2DKernel
            float tmp0[16*16]; at::vec::transpose_mxn<...>(tmp0, in_ptr0 + ..., ...); // into kernel.preloads
            float tmp1[16*16]; // into kernel.preloads
            for i_inner ... { // the kernel inner loop
              vectorized loads/compute/stores (e.g., load tmp0, store tmp1) // into kernel.loads/compute/stores
            }
            at::vec::transpose_mxn(out_ptr0 + ..., tmp1, ...) // into kernel.poststores
          for inner_most ... (tail)
            // generated by CppVecKernel
            ...
      for i_outer ... (tail)
        for ...
          for ...
            // generated by CppKernel
            ...
    c                 ó–   •— t         ‰| �  ||||d   |«       || _        || _        || _        |r|n|| _        |r|n|| _        d| _        y )Nr5   T)r,  r-  Útiling_indicesÚinner_tail_sizeÚouter_tail_sizeÚinner_num_elemsÚouter_num_elemsÚinner_is_tiling_idx)r.  r‘  r�  r  rÓ  rÔ  rÕ  r1  s          €rW   r-  zCppTile2DKernel.__init__Ä  s`   ø€ ô 	‰ÑØØØØ˜1ÑØô	
ð -ˆÔØ.ˆÔØ.ˆÔÙ2A™À}ˆÔÙ2A™À}ˆÔØ#'ˆÕ rV   c                 óL   — t        | j                  | j                     › d�«      S )NrD  )r-   rH  Ú	outer_idxr3  s    rW   Úinner_itervarzCppTile2DKernel.inner_itervarÛ  s"   € Ü! T§]¡]°4·>±>Ñ%BÐ$CÀ6Ð"JÓKÐKrV   c                 ób  — | j                   | j                     }| j                   | j                     }t        ||| j                  «      }t        ||| j                  «      }| j
                  d u xr@ |dk(  xr9 |j                  |«      xr& |j                  |«       xr |j                  |«       S ræ   )rH  rÚ  rI  r  r  r¦  rç   )r.  r’   Ú	outer_varÚ	inner_varÚouter_strideÚinner_strides         rW   Úneed_vec_transposez"CppTile2DKernel.need_vec_transposeÞ  s­   € Ø—M‘M $§.¡.Ñ1ˆ	Ø—M‘M $§/¡/Ñ2ˆ	Ü*¨5°)¸T×=OÑ=OÓPˆÜ*¨5°)¸T×=OÑ=OÓPˆà�O‰O˜tÐ#ò 0Ø Ñ!ò0à—	‘	˜)Ó$ò0ð !×$Ñ$ YÓ/Ð/ò0ð !×$Ñ$ YÓ/Ð/ð	
rV   c                 ón  — t         j                  j                  |«      }| j                  }|› dt	        |«      › �}d}	t	        t        || j                  | j                     | j                  «      «      › }
t	        | j                  «      › }|r|	|}	}||
}}
d}| j                  |z  r| j                  | j                  }}n| j                  | j                  }}|r|dk(  rdnd}t        |t        j                  «      r|j                  r&t        |t        j                  «      rA|j                  s5dt         |   › d|› d	|› d
|
› d
|	› d
|› d
t	        |«      › d
t	        |«      › d�}n4dt         |   › dt	        |«      › dt	        |«      › d|› d	|› d
|
› d
|	› d
|› d�}|r| j"                  j%                  «       }na| j"                  j'                  |«      s)| j"                  j)                  | j*                  |d¬«      }nd}| j"                  j-                  |«      }|r>t         |   }d|› d|› d�}|› d|› d|› d|› d|› d�
}| j*                  j/                  |«       |j1                  dt3        |«      «      }|r'| j4                  j/                  t7        ||«      «       |S | j*                  j/                  |«       |S )Nr}  Ú__place_holder__TrÐ  ÚtrueÚfalseztranspose_mxn<ra  r“  r™   rÓ   Frâ  zalignas(std::max(std::size_t(z), alignof(z)))r•   rÊ  r[   rÕ   )r4   rÓ  rÔ  r  rE   r  rH  rI  rP  rØ  rÖ  r×  r    rè   r…  r‹  rH   r‰  r	  ÚcontainsrŠ  r‹  Úgetr¿   r   ro   rŒ  r;   )r.  rÄ   r¥   r’   Úis_storeÚ
store_moderˆ   Úfactorr:  r<  Úld_srcÚld_dstÚneed_defineÚMÚNrÐ  Úload_or_storeÚtile_varÚ	cpp_dtypeÚalignasÚdefine_lines                        rW   Úgen_transposed_tile_load_storez.CppTile2DKernel.gen_transposed_tile_load_storeë  sÆ  € ô —‘×!Ñ! $Ó'ˆØ×#Ñ#ˆØ��Sœ UÓ+Ð,Ð-ˆØ ˆÜÔ 3°E¸4¿=¹=ÈÏÉÑ;YÐ[_×[mÑ[mÓ nÓoÐpˆÜ §¡Ó/Ð0ˆÙØ˜C�ˆCØ# V�FˆFàˆØ×#Ñ# hÒ.Ø×'Ñ'¨×)=Ñ)=ˆq‰Að ×$Ñ$Ø×$Ñ$ð ˆAñ !)¨j¸LÒ.H‘VÐPWˆ
Ü�qœ%Ÿ*™*Ô%¨a¯kªkÜ�qœ%Ÿ*™*Ô%¨a¯kªkð !¤¨eÑ!4Ð 5°Q°z°lð CØ�5˜˜6˜( " S E¨¨F¨8°2´kÀ!³nÐ5EÀRÌÐTUËÐGWÐWYð[ñ ð !¤¨eÑ!4Ð 5°Q´{À1³~Ð6FÀaÌÐTUËÐGWÐWXÐYcÐXdð eØ�5˜˜6˜( " S E¨¨F¨8°2ð7ð ñ Ø—x‘x—‘Ó(‰HØ—‘×"Ñ" =Ô1Ø—x‘x×(Ñ(¨¯©¸ÈUÐ(ÓS‰HàˆKØ—x‘x—|‘| MÓ2ˆHáÜ$ UÑ+ˆIð 6°f°X¸[ÈÈÐSVÐWˆGØ$˜I Q y k°°8°*¸A¸f¸XÀQÀvÀhÈbÐQˆKØ�M‰M×#Ñ# KÔ0à%×-Ñ-Ð.@Ä#ÀhÃ-ÓPˆÙØ�O‰O×%Ñ%¤l°4¸Ó&GÔHð ˆð �M‰M×#Ñ# MÔ2àˆrV   rÄ   r’   c                 óV  •— | j                   j                  |«      }| j                  |«      }| j                  «       }| j	                  |«      r»| j                  |||d¬«      }|› dt        || j                  z  «      › �}t        j                  j                  |«      }| j                  |d|«      }| j                  j                  | j                  |«      }	|	j                  d| ||fi «       t!        |	t"        «      sJ ‚d|	_        |	S | j'                  |«      }
t(        ‰| �U  ||
«      S )NF)rè  r}  r   rí  T)r‘  rÌ  r8  rÛ  rá  rõ  rE   rP  r4   rÓ  rÔ  ri  r‰  rŠ  r¡  rŒ  r    rG   r¤   rs  r,  rí  )r.  rÄ   r’   r¥   rY  rñ  rh  rˆ   rÅ   r‘  rí   r1  s              €rW   rí  zCppTile2DKernel.load'  s  ø€ Ø�i‰i�o‰o˜dÓ#ˆØ×$Ñ$ UÓ+ˆà×"Ñ"Ó$ˆØ×"Ñ" 5Ô)Ø×:Ñ:Ø�c˜5¨5ð ;ó ˆHð "˜
 #¤k°%¸$¿.¹.Ñ2HÓ&IÐ%JÐKˆGÜ—G‘G×%Ñ% dÓ+ˆEØ×*Ñ*¨7°A°uÓ=ˆDØ—X‘X×&Ñ& t§z¡z°4Ó8ˆFØ×!Ñ! &¨4°°uÐ*=¸rÔBÜ˜f¤nÔ5Ð5Ð5Ø ˆFŒMØˆMà×/Ñ/°Ó6ˆIÜ‘7‘<  iÓ0Ð0rV   c                 óþ  •— d|v sJ ‚t        |t        «      sJ |«       ‚|j                  s| j                  |«      }| j                  j                  |«      }| j                  «       }| j                  |«      }| j                  |«      rÒ| j                  |||d|¬«      }|› dt        || j                  z  «      › �}| j                  sFt        j                  j                  |«      t         t"        j$                  t"        j&                  gz   v r|› d|› dt        | j                  «      › d�}	n|› d|› d�}	| j(                  j+                  t-        ||	«      «       y | j/                  |«      }
t0        ‰| �e  ||
||«       y )NrÎ  T)rè  ré  r}  r‰  r™   rÓ   )r    rG   r¤   r�  r‘  rÒ  rÛ  r8  rá  rõ  rE   rP  rX  r4   rÓ  rÔ  rƒ   r„   r4  r5  r›  r¿   r;   rs  r,  r×  )r.  rÄ   r’   r�  rÖ  r¥   rY  rñ  ÚstorebufrÅ   rí   r1  s              €rW   r×  zCppTile2DKernel.store=  sd  ø€ Ø˜‰}Ðˆ}Ü˜%¤Ô0Ð7°%Ó7Ð0Ø�|Š|à—N‘N 5Ó)ˆEà�i‰i×Ñ˜tÓ$ˆà×"Ñ"Ó$ˆØ×$Ñ$ UÓ+ˆØ×"Ñ" 5Ô)Ø×:Ñ:Ø�c˜5¨4¸Dð ;ó ˆHð #˜ 3¤{°5¸4¿>¹>Ñ3IÓ'JÐ&KÐLˆHØ�~Š~¤§¡×!2Ñ!2°4Ó!8¼MÜ—‘Ü—
‘
ðMñ =ñ "ð  ˜ ¨ z°´KÀÇÁÓ4OÐ3PÐPRÐS‘à˜ ¨ z°Ð4�Ø�K‰K×!Ñ!¤,¨t°TÓ":Õ;à×/Ñ/°Ó6ˆIÜ‰G‰M˜$ 	¨5°$Õ7rV   c                 ó  — | j                  «       }| j                  r2|j                  d|› d|› dt        | j                  «      › d|› d�	«       y |j                  d|› d|› dt        | j
                  «      › d|› d�	«       y )Nr{  r|  r¹   rº   r»   )rÛ  rØ  r¿   rE   r×  rÖ  )r.  r°   rY  s      rW   rŸ  z#CppTile2DKernel.codegen_inner_loopsZ  s�   € Ø×"Ñ"Ó$ˆØ×#Ò#Ø�N‰NØ˜U˜G 6¨%¨°´KÀ×@TÑ@TÓ4UÐ3VÐVXÐY^ÐX_Ð_bÐcõð �N‰NØ˜U˜G 6¨%¨°´KÀ×@TÑ@TÓ4UÐ3VÐVXÐY^ÐX_Ð_bÐcõrV   c                 óœ  •— t         ‰| �  ||«      }| j                  d   | j                  k  r| j                  nt	        | j                  «      \  | _        | _        | j                  | j                  d   k(  r+| j                  | _        | j                  | _
        d| _        |S | j                  | _        | j                  | _
        d| _        |S )Nr5   r   FT)r,  rð  rÓ  r}  ÚreversedrÚ  rI  rÕ  rX  r×  rP  rØ  rÔ  rÖ  )r.  ÚgroupÚreduction_grouprž  r1  s       €rW   rð  zCppTile2DKernel.set_rangese  sÀ   ø€ Ü‰wÑ! %¨Ó9ˆð ×"Ñ" 1Ñ%¨×(<Ñ(<Ò<ð ×Òä˜$×-Ñ-Ó.ñ 	(ˆŒ˜œð
 �?‰?˜d×1Ñ1°!Ñ4Ò4Ø!×1Ñ1ˆDŒNØ!×1Ñ1ˆDŒNØ',ˆDÔ$ð
 ˆð "×1Ñ1ˆDŒNØ!×1Ñ1ˆDŒNØ'+ˆDÔ$ØˆrV   rµ   c                 óZ   — | j                  || j                  | j                  «       ¬«      S )Nrz  )r«  rÚ  rÛ  r­  s     rW   rs  z"CppTile2DKernel.transform_indexingw  s0   € Ø×+Ñ+ØØŸ™Ø×%Ñ%Ó'ð ,ó 
ð 	
rV   rI  r  )r  r  r  r  rp  r@  r-  rÛ  rá  rõ  ro   rè   r…  rí  r×  rŸ  rð  rs  r`  ra  s   @rW   rr  rr  £  ss   ø„ ñð< #€Ið Øõ(ò.Lò
ð 6:ó:ðx1˜ð 1 U§Z¡Zõ 1õ,8ò:	ôð$
¨¯
©
ð 
°u·z±z÷ 
rV   rr  Ú_bodyc                 ój  — | j                   gt        | j                  j                  «       «      z   }d}d}|D ]ö  }|j                  j
                  D ]Û  }|j                  dk(  s|j                  dv rŒ!|j                  dvrd}t        |d«      r�|j                  r‘t        j                  |j                  v sJ ‚|j                  t        j                     }|j                  r|j                  t        vrd}Œ¥|�&||j                  k7  sŒ·t        j                  d«       ŒÍ|j                  }ŒÚd}ŒÝ Œø ||fS )	zÞ
    Returns the low precision data type (torch.float16/torch.bfloat16) contained in the nodes
    and if all the nodes can codegen with this data type without converting to float.
    Otherwise returns None and True.
    NFÚplaceholder)Ú	get_indexr:  )rí  r×  r—  rž  rÒ  Trk  z.bf16 and fp16 are mixed in the scheduler node.)Ú
root_blockr#  Ú	subblocksr  rÓ  ÚnodesÚopÚtargetr£   rk  rA   rj  rˆ   rƒ   ÚwarningsÚwarn)rÿ  Ú
sub_blocksÚ_lowp_fp_typeÚ	_use_fp32Ú	sub_blockr0  rg  s          rW   Úget_loop_body_lowp_fpr    s$  € ð ×"Ñ"Ð#¤d¨5¯?©?×+AÑ+AÓ+CÓ&DÑD€Jà+/€MØ€IØò !ˆ	Ø—_‘_×*Ñ*ò 	!ˆEØ�x‰x˜=Ò(¨E¯L©Lð =ñ -ð ð �|‰|ð $ñ ð !�	ä�u˜fÔ%¨%¯*ª*Ü*×.Ñ.°%·*±*Ñ<Ð<Ð<Ø/4¯z©zÔ:M×:QÑ:QÑ/R�Ø—}’}¨¯©¼]Ñ(JØ $‘IØ"Ð.Ø$¨¯©Ó5Ü Ÿ™Ð&VÕWà$+§M¡M‘Mà ‘	ñ9	!ð!ð> ˜)Ð#Ð#rV   c                   óJ   ‡ — e Zd ZdZˆ fd„Zdeee   ee   f   fd„Zd„ Z	ˆ xZ
S )ÚTilingSelectz˜
    Implement the heuristic to select the tiling factors and tiling indices.
    In the future, we can implement advanced heuristic in a subclass.
    c                 ó"   •— t         ‰| �  «        y r  )r,  r-  ©r.  r1  s    €rW   r-  zTilingSelect.__init__±  s   ø€ Ü‰ÑÕrV   rµ   c           	      ó”  ‡#— t        |«      }t        |«      }|sJ ‚t        d„ |D «       «      rg g fS t        j                  }t        |d   «      d   Š#‰#rt        ˆ#fd„|dd  D «       «      r‰#}t        j                  «       j                  |¬«      }| j                  |||«      }|�r™t        |d„ ¬«      \  }}	t        |«      t        |	«      z   }
t        j                  j                  �rVd„ }d	„ }d
„ }t!        t#        |
«      «      D �cg c]  }t%        t&        j(                  |«      ‘Œ }}t#        |«      }|d | ||d  }}i }i }|D �]n  }|j*                  gt-        |j.                  j1                  «       «      z   }|D �]4  }|j2                  j4                  D �]  }|j6                  dv rŠ|j6                  dk(  rdnd}|j8                  j;                  ||f«      |j<                  |   j<                  d      } |||«      r4 |||||«      }|j6                  dk(  r|€n|dvr ||j6                  |«       t?        |j6                  t@        «      sŒ·|j6                  jC                  d«      rŒÓ|j6                  dv rŒâ|j6                  |vrd||j6                  <   �Œ||j6                  xx   dz  cc<   �Œ �Œ7 �Œq tE        |j1                  «       «      }tE        |j1                  «       «      }d}d}||k\  s|dkD  r||z  |k\  rg g fS |	s9|r7t#        |«      dk(  r)tG        ||d      g«      s||d      |dz  k  r	|dk  rg g fS |tH        v rÑt        j                  «       j                  |¬«      } |D ]¨  }!|!dk  r|!t#        |
«      z   }!|!dk  s|!t#        |
«      k\  rŒ*tG        |
«      retJ        j2                  jL                  jO                  |
|!   d¬«      }"|"| k  sŒitJ        j2                  jL                  jQ                  |"| «       | dz  } n|
|!   | k  sŒ£| dz  } n t#        |«      dk(  r|g|fS t#        |«      dk(  r||g|fS g g fS c c}w )Nc              3   ó,   K  — | ]  }|t         v–— Œ y ­wr  )rv   ©r  rˆ   s     rW   r  z-TilingSelect.select_tiling.<locals>.<genexpr>½  s   è ø€ ÒH°EˆuÔ/Ô/ÑHùó   ‚r   c              3   ó@   •K  — | ]  }t        |«      d    ‰k(  –— Œ y­w)r   N)r  )r  Ú	loop_bodyÚ_lowp_fp_dtypes     €rW   r  z-TilingSelect.select_tiling.<locals>.<genexpr>Á  s(   øè ø€ ò "
àô # 9Ó-¨aÑ0°NÕBñ"
ùó   ƒr5   r  c                 ó   — t        | d   «      S ræ   ©rÐ   ©Úsizess    rW   r�  z,TilingSelect.select_tiling.<locals>.<lambda>Î  s   € ´#°e¸A±h³-€ rV   ©rj  c                 óL   — ||d      }t        | ||«      }|j                  r|S d S rU  )r  r‹  )r’   rH  r  rÓ  r¯  rN  s         rW   Ú_try_get_stridez3TilingSelect.select_tiling.<locals>._try_get_strideÔ  s4   € ð ' ~°aÑ'8Ñ9�GÜ0°¸ÀÓO�FØ%+×%5Ò%5˜6Ð?¸4Ð?rV   c                 ó2   — | |vrd|| <   y || xx   dz  cc<   y ræ   rU   )Ú	node_nameÚnon_contig_indexing_op_counters     rW   Ú_update_negative_op_countz=TilingSelect.select_tiling.<locals>._update_negative_op_countÞ  s(   € ð !Ð(FÑFØDEÐ6°yÒAà6°yÓAÀQÑFÔArV   c                 ó–   — t        |«      dk(  xr: t        | «      dkD  xr* |d   dk\  r|d   n|d   t        | «      z   t        | «      k  S ©Nr5   r   r  )rH  rÓ  s     rW   Ú_is_valid_indicesz5TilingSelect.select_tiling.<locals>._is_valid_indicesæ  sb   € ô
 ˜NÓ+¨qÑ0ò (Ü ›M¨AÑ-ò(ð  .¨aÑ0°AÒ5ð +¨1Ò-à!/°Ñ!2´S¸³]Ñ!Bä˜h›-ñ(ð	rV   )r:  rí  r×  r:  r   ©r   r5   Úmasked_subblock)r2   rÒ  r4  r  g¸…ëQ¸¾?é#   ro  é
   ró  ))rC   rB   rf   r„   rn   r  r"  r   rN  Ú	nelementsÚ_select_tiling_indicesr_   r¡   r   r¹  Úenable_tiling_heuristicsrG  rÐ   r.   r   rì  r  r#  r  r  rÓ  r  r  r@  Úindexing_from_argsr‘  r    ro   Ú
startswithrc   r(   rƒ   r4   rõ  rö  Úguard_lt)$r.  Úfn_listÚvar_sizes_listÚloop_bodiesÚ
all_dtypesrˆ   r  rÓ  rü  rý  r|  r!  r%  r(  rï  rH  r}  rž  Úreduction_varsÚ
op_counterr$  rÿ  r
  r  r0  Úarg_idxr’   rN  Úop_numÚnon_contig_indexing_op_numÚratio_thresholdÚquantity_thresholdÚfactor_lowpÚtiling_indiceÚ
call_ranger  s$                                      @rW   Úselect_tilingzTilingSelect.select_tiling´  s´  ø€ ô % WÓ-ˆÜ/°Ó<ˆ
ÙÐˆzÜÑH¸ZÔHÔHØ�r�6ˆMÜ—‘ˆÜ.¨{¸1©~Ó>¸qÑAˆÙœcó "
à(¨¨˜_ô"
ô 
ð #ˆEä#×0Ñ0Ó2×<Ñ<À5Ð<ÓIˆØ×4Ñ4Ø�^ ]ó
ˆò Ü%(ØÑ$?ô&Ñ"ˆE�?ô   ›,¬¨Ó)?Ñ?ˆKä�z‰z×2Ó2ò@òGòô" #¤3 {Ó#3Ó4öàô 3´4·;±;ÀÕBð�ð ô #& e£*�àÐ-˜oÐ.Ø˜_Ð-Ð.ð %�ð .0�
ð BDÐ.Ø(ó B�EØ"'×"2Ñ"2Ð!3´d¸5¿?¹?×;QÑ;QÓ;SÓ6TÑ!T�JØ%/ó B˜	Ø%.§_¡_×%:Ñ%:ó B˜EØ$Ÿ|™|Ð/NÑNà/4¯|©|¸|Ò/K©!ÐQR Ø(1¯©×(IÑ(IØ%)¨>Ð$:ó)"à"'§*¡*¨WÑ"5×":Ñ":¸1Ñ"=ñ)? ñ $5°X¸~Ô#NÙ-<Ø(-¨x¸Èó.& Fð
 ,1¯<©<¸<Ò+Gð )/©à-3¸6Ñ-Aá(AØ,1¯L©LÐ:Xô)*ô  *¨%¯,©,¼Õ<Ø %§¡× 7Ñ 7Ð8IÕ JØ#(§<¡<Ø#Mò$Nð $)§<¡<°zÑ#AØ?@ J¨u¯|©|Ó$<à$.¨u¯|©|Ó$<ÀÑ$AÕ$<ò7BòBðBô@ ˜Z×.Ñ.Ó0Ó1�Ü-0Ø2×9Ñ9Ó;ó.Ð*ð #'�Ø%'Ð"Ø-Ð1CÒCØ˜Q’JØ2°VÑ;¸ÒNð
 ˜r˜6�Mñ (ÙÜ˜NÓ+¨qÒ0Ü,à! .°Ñ"3Ñ4ðôð
 ˜n¨QÑ/Ñ0°=À1Ñ3DÒDØ šð ˜r˜6�MàœÑ%ô *×6Ñ6Ó8×BÑBÈÐBÓO�Ø%3ò �MØ$ qÒ(Ø(5¼¸KÓ8HÑ(H˜Ø$ qÒ(¨M¼SÀÓ=MÒ,MØ Ü'¨Ô4Ü%&§W¡W×%5Ñ%5×%?Ñ%?Ø'¨Ñ6Àð &@ó &˜
ð &¨Ó3ÜŸG™G×,Ñ,×5Ñ5°jÀ+ÔNØ,7¸1Ñ,<˜MÙ!Ø$ ]Ñ3°kÓAØ(3°qÑ(8˜Ùðô" �>Ó" aÒ'Ø%�¨Ð6Ð6Ü�>Ó" aÒ'Ø% }Ð5°~ÐEÐEØ�2ˆvˆùòQs   Ä
!Qc           	      ó¸  — g }t        ||«      D ]`  \  }}t        j                  |g|¢­Ž }|t        j                  |j
                  |j                  «      D �cg c]  }|j                  ‘Œ c}z  }Œb t        t           «       }	g }
t        t           «       }t        t           «       }|D �]  }|j                  D ]õ  }t        j                  d|j                  «      sŒ$t        |||«      }|dk(  rŒ7|dk(  rO|	j                  t        |j                  dd  «      «       |
j!                  t        |j                  dd  «      «       Œ‹t#        d„ |j                  D «       «      r(|j                  t        |j                  dd  «      «       ŒÏ|j                  t        |j                  dd  «      «       Œ÷ �Œ
 |	|z
  |z
  }t%        |d„ ¬«      \  }}t'        |«      t'        |«      z   }t'        |	«      dk(  r|dz
  gS |rt)        |«      dd  S |	|z  |z
  }t)        |	«      }t'        |«      dk(  r|d   |v r|d   |dz
  k(  r|S t)        ||
j*                  ¬«      dd  S c c}w )	Nz^d\d+$r   r5   c              3   óP   K  — | ]  }t        |t        j                  «      –— Œ  y ­wr  )r   r   ÚSIZE©r  r´  s     rW   r  z6TilingSelect._select_tiling_indices.<locals>.<genexpr>w  s   è ø€ ÒS¸!œ¨¬4¯9©9×5ÑSùó   ‚$&c                 ó   — t        | d   «      S ræ   r  r  s    rW   r�  z5TilingSelect._select_tiling_indices.<locals>.<lambda>|  ó   € ÄsÈ5ÐQRÉ8Ã}€ rV   r  rH  r   )r  r
   Úextract_read_writesÚ	itertoolsÚchainÚreadsÚwritesr’   r   rm   rµ  rÆ   ÚsearchrÄ   r  r�  rŽ  r"  r_   rÐ   ÚsortedÚcount)r.  r3  r4  r  Ú	all_indexÚfnÚ	var_sizesÚrwÚdepÚcontig_varsÚcontig_vars_listÚnon_contig_stride_constÚnon_contig_stride_otherr’   r¥   rN  Úcontig_onlyrü  rý  Únum_itervarsÚcontig_and_const_strideÚcontig_vars_sorteds                         rW   r.  z#TilingSelect._select_tiling_indices_  sV  € ð ˆ	Ü  ¨.Ó9ò 	U‰MˆB�	Ü×1Ñ1°"ÐA°yÒAˆBØ¬y¯©¸r¿x¹xÈÏÉÓ/SÖT¨˜#Ÿ)›)ÒTÑT‰Ið	Uô !¤‘oÓ'ˆØÐÜ",¬S¡/Ó"3ÐÜ",¬S¡/Ó"3ÐØó 	CˆEØ×)Ñ)ò C�Ü—y‘y ¨C¯H©HÔ5ØÜ,¨U°C¸ÓG�Ø˜Q’;ØØ˜q’[Ø—O‘O¤C¨¯©°°¨Ó$5Ô6Ø$×+Ñ+¬C°·±¸¸°Ó,=Õ>ÜÑS¸v×?RÑ?RÔSÔSØ+×/Ñ/´°C·H±H¸Q¸R°LÓ0AÕBà+×/Ñ/´°C·H±H¸Q¸R°LÓ0AÕBòCð	Cð "Ð$;Ñ;Ð>UÑUˆÜ!$ ^Ñ9TÔ!UÑˆˆÜ˜5“z¤C¨Ó$8Ñ8ˆÜˆ{Ó˜qÒ à  1Ñ$Ð%Ð%ÙÜ˜+Ó& r sÐ+Ð+àÐ1Ñ1Ø#ñ#$Ðô $ KÓ0ÐäÐ"Ó# qÒ(Ø" 2Ñ&Ð*AÑAØ" 2Ñ&¨,¸Ñ*:Ò:à%Ð%ÜÐ(Ð.>×.DÑ.DÔEÀbÀcÐJÐJùòK Us   ÁI)r  r  r  r  r-  r¡   r#  rm   rA  r.  r`  ra  s   @rW   r  r  «  s7   ø„ ñô
ðið 
ˆt�C‰y˜$˜s™)Ð#Ñ	$ó	iöV.KrV   r  c                   ó–   ‡ — e Zd Zˆ fd„Zd„ Zdefd„Zdefd„Zd„ Z	d„ Z
d	„ Zd
ee   fd„Zd„ Zd„ Zddee   fd„Zdeded   fd„Zˆ xZS )r   c                 óÌ   •— t         ‰| �  |j                  |j                  j                  «       || _        d | _        d | _        t        j                  «       | _
        g | _        y r  )r,  r-  r‘  Úwsr�  rV  rH  r|  r   rN  Úpicked_vec_isaÚkernels©r.  rV  r1  s     €rW   r-  zCppKernelProxy.__init__‘  sQ   ø€ Ü‰Ñ˜×*Ñ*¨L¯O©O×,GÑ,GÔHØ(ˆÔØˆŒØˆÔÜ2=×2JÑ2JÓ2LˆÔØ(*ˆ�rV   c                 ó`   — |D ])  }t        |t        «      sJ ‚t        j                  |«       Œ+ y r  )r    r$   r:   Úpropagate_scheduler_node)r.  r  r0  s      rW   Údata_type_propagationz$CppKernelProxy.data_type_propagation™  s1   € Øò 	@ˆEÜ˜e¤]Ô3Ð3Ð3Ü×8Ñ8¸Õ?ñ	@rV   Úscheduler_nodec                 óÌ   — t        |j                  t        «      syt        j                  |«       t        |j                  «      d   d uxr t        |j                  «      d    S )NTr   r5   )r    rÿ  r   r:   re  r  )r.  rg  s     rW   Úis_lowp_fp_schedulerz#CppKernelProxy.is_lowp_fp_schedulerŸ  s\   € Ü˜.×.Ñ.´Ô9Øä×4Ñ4°^ÔDä! .×"6Ñ"6Ó7¸Ñ:À$ÐFò CÜ)¨.×*>Ñ*>Ó?ÀÑBÐBð	
rV   r  c                 óÔ   — dt         j                  j                  fd„}|j                  gt	        |j
                  j                  «       «      z   }|D ]  } ||j                  «       Œ y )NÚ	sub_graphc                 ól  ‡‡‡‡‡‡‡‡— dt         j                  j                  dt        t         j                     fd„Šdt         j                  j                  dt        t         j                     fd„Šdt         j                  j                  dt         j                  fˆfd„Šdt         j                  j                  dt         j                  fˆfd„Šdt         j                  j                  dt         j                  fˆˆfd„}t        | j                  «      }g Š|D �]9  }|j                  d	v r¬ ‰|«      xŠt        v rœt        ˆˆfd
„|j                  D «       «      rŒB|j                  d   }| j                  |«      5  | j                  d||t         j                  f¬«      Š|j                  ‰ˆfd„«       t         xj"                  dz  c_        d d d «       Œ¾|j                  dk(  r¨ ‰|«      xŠt        v r˜|j                  \  }}}}} ||‰«      rŒùt$        j&                  j)                  |«      Š| j+                  |«      5  | j                  d||‰f¬«      Š|j-                  |‰«       t         xj"                  dz  c_        d d d «       �Œu|j                  dk(  r’|j                  \  }Š}}	}
|t        v s�Œ ‰t         j                  t         j.                  t         j0                  t         j2                  fv sJ ‚|‰t        v rt         j                  n‰t         j                  |	|
f|_        �Œ|j                  dk(  r`|j                  d   t        v rK|j                  \  }}
Št        ˆˆfd„|j                  D «       «      r�Œk||
t         j                  f|_        �Œ…|j                  dk(  rq|j                  d   t        v r\|j                  \  }}Št        ˆˆfd„|j                  D «       «      r�ŒÚ‰j5                  |«       ||t         j                  f|_        �Œ|j                  dk(  �r%|j                  \  }}Š}|t        v rd |||«      s[| j+                  |«      5  | j                  d|||f¬«      Š|j-                  |‰«       t         xj"                  dz  c_        d d d «       ‰t        v s�Œœt        ˆˆfd„|j                  D «       «      r�Œ½|j                  d   }| j                  |«      5  | j                  d||t         j                  f¬«      Š|j                  ‰ˆfd„«       t         xj"                  dz  c_        d d d «       �Œ:�Œ< dt         j                  j6                  fˆfd„} || «       y # 1 sw Y   �ŒpxY w# 1 sw Y   �Œ}xY w# 1 sw Y   Œ÷xY w# 1 sw Y   �Œ–xY w)Nr  rµ   c                 óZ  — | j                   dk(  r,t        j                  j                  | j                  d   «      S | j                   dk(  r| j                  d   S | j                   dk(  rCt        | j                  «      dkD  r| j                  d   S | j                  j                  dd«      S y)	z6Get input dtype for nodes that may consumes lowp fp dtr×  r5   r”  rH  r‡  r   r§   N)r  r4   rÓ  rÔ  r‘  rÐ   r’  rç  ©r  s    rW   Úget_input_dtypez]CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.get_input_dtype«  sˆ   € à—;‘; 'Ò)ÜŸ7™7×,Ñ,¨T¯Y©Y°q©\Ó:Ð:Ø—[‘[Ð$6Ò6ØŸ9™9 R™=Ð(Ø—[‘[ JÒ.Ü˜4Ÿ9™9“~¨Ò)Ø#Ÿy™y¨™|Ð+à#Ÿ{™{Ÿ™¨{¸DÓAÐAàrV   c                 ó$  — | j                   dk(  rFt        | j                  «      dk(  sJ ‚t        j                  j                  | j                  d   «      S | j                   dv r| j                  d   S | j                   dk(  r| j                  d   S y)	z6Get output dtype for nodes that may produce lowp fp dtrí  r   r5   )r‡  r4  r:  rH  r”  r   N)r  rÐ   r‘  r4   rÓ  rÔ  rn  s    rW   Úget_output_dtypez^CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.get_output_dtype¹  sz   € à—;‘; &Ò(Ü˜tŸy™y›>¨QÒ.Ð.Ð.ÜŸ7™7×,Ñ,¨T¯Y©Y°q©\Ó:Ð:Ø—[‘[Ð$JÑJØŸ9™9 R™=Ð(Ø—[‘[Ð$6Ò6ØŸ9™9 Q™<Ð'àrV   Údtc                 ó.   •— |t         v sJ ‚ ‰| «      |k(  S )z]Check if the given node produces output with expected low precision floating point data type.)rƒ   )r  rr  rq  s     €rW   Úis_lowp_fp_sourcez_CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.is_lowp_fp_sourceÅ  s!   ø€ àœ]Ñ*Ð*Ð*Ù'¨Ó-°Ñ3Ð3rV   c                 óX   •— |t         v sJ ‚ ‰| «      x}r||k(  S | j                  dk(  ryy)zZCheck if the given node accept input with expected low precision floating point data type.r‡  TF)rƒ   r  )r  rr  Úinput_dtypero  s      €rW   Úis_lowp_fp_sinkz]CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.is_lowp_fp_sinkÊ  s>   ø€ àœ]Ñ*Ð*Ð*Ù"1°$Ó"7Ð7�;Ð7Ø&¨"Ñ,Ð,Ø—[‘[ JÒ.àà rV   c                 óZ   •‡—  ‰| ‰«      xr t        ˆˆfd„| j                  D «       «      S )z¥Check if the node is a lowp fp sources which are all directly fed to ops that accepts lowp fp input
                thus no need to promote to float
                c              3   ó0   •K  — | ]  } ‰|‰«      –— Œ y ­wr  rU   ©r  Úuserrr  rw  s     €€rW   r  z}CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.is_lowp_fp_source_no_promote.<locals>.<genexpr>Ù  s   øè ø€ ò ;Ø26‘O D¨"×-ñ;ùó   ƒ©r"  Úusers)r  rr  rw  rt  s    `€€rW   Úis_lowp_fp_source_no_promotezjCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.is_lowp_fp_source_no_promoteÕ  s1   ù€ ñ )¨¨rÓ2ò ´sô ;Ø:>¿*¹*ô;ó 8ð rV   )rí  r:  c              3   ó0   •K  — | ]  } ‰|‰«      –— Œ y ­wr  rU   rz  s     €€rW   r  zWCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.<genexpr>å  ó   øè ø€ ÒM¸™?¨4°×4ÑMùr|  r   r‡  ©r‘  c                 ó   •— | ‰uS r  rU   ©rï  Úto_type_nodes    €rW   r�  zVCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.<lambda>í  s   ø€ °A¸\Ð4I€ rV   r5   r×  rè  r4  rH  c              3   ó0   •K  — | ]  } ‰|‰«      –— Œ y ­wr  rU   rz  s     €€rW   r  zWCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.<genexpr>  r�  r|  c              3   ó0   •K  — | ]  } ‰|‰«      –— Œ y ­wr  rU   rz  s     €€rW   r  zWCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.<genexpr>"  r�  r|  r”  c              3   ó0   •K  — | ]  } ‰|‰«      –— Œ y ­wr  rU   )r  r{  rˆ   rw  s     €€rW   r  zWCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.<genexpr>J  s   øè ø€ ÒUÀ¡°°e× <ÑUùr|  c                 ó   •— | ‰uS r  rU   r„  s    €rW   r�  zVCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.<lambda>R  s   ø€ ¸AÀ\Ð<Q€ rV   rk  c                 óV   •— dt         j                  j                  fˆfd„} || «       y )Nrk  c                 óB  •‡— dt         j                  j                  fd„}| j                  D �cg c]  }|j                  dk(  sŒ|‘Œ }}|D �cg c]  } ||«      sŒ||j
                  i‘Œ }}|D ]†  }|j                  «       D ]q  \  Š}‰| j                  v sŒt        ˆfd„|D «       «      s‰‰v sŒ.t        d„ |D «       «      sŒA‰j                  d   }‰j                  |«       | j                  ‰«       Œs Œˆ | j                  €| j                  «        y y c c}w c c}w )NÚto_nodec                 ó:   — t        d„ | j                  D «       «      S )Nc              3   ó:   K  — | ]  }|j                   d k(  –— Œ y­w)r‡  N)r  ©r  Úusrs     rW   r  z®CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.eliminate_to_dtype.<locals>._eliminate_duplicate_to_node.<locals>._used_by_to.<locals>.<genexpr>a  s   è ø€ Ò"UÀ 3§:¡:°Õ#;Ñ"Uùs   ‚r}  )rŒ  s    rW   Ú_used_by_toz›CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.eliminate_to_dtype.<locals>._eliminate_duplicate_to_node.<locals>._used_by_to`  s   € Ü"Ñ"UÀwÇ}Á}Ô"UÓUÐUrV   r‡  c              3   ó\   •K  — | ]#  }|j                   d    ‰j                   d    k(  –— Œ% y­w©rH  Nr‚  )r  r�  r  s     €rW   r  z™CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.eliminate_to_dtype.<locals>._eliminate_duplicate_to_node.<locals>.<genexpr>l  s&   øè ø€ Ò#SÀc C§H¡H¨R¡L°D·I±I¸b±MÕ$AÑ#Sùs   ƒ),c              3   óF   K  — | ]  }|j                   d    t        v –— Œ y­wr“  )r‘  rƒ   r�  s     rW   r  z™CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.eliminate_to_dtype.<locals>._eliminate_duplicate_to_node.<locals>.<genexpr>o  s"   è ø€ ò ,&ØJM¨¯©°©¼Ô(Eñ,&ùr  rH  )r„   ÚfxÚNoder  r  r~  rŸ  r"  Úall_input_nodesÚreplace_all_uses_withÚ
erase_nodeÚowning_moduleÚlint)	rk  r‘  r  Úall_to_nodesÚall_to_nodes_and_usersÚ
node_usersr~  Úval_nodeÚto_lowp_fp_legalized_nodess	     `     €rW   Ú_eliminate_duplicate_to_nodez†CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.eliminate_to_dtype.<locals>._eliminate_duplicate_to_nodeY  s-  ù€ ðV¬U¯X©X¯]©]ó Vð *3¯©ö$Ø!%¸D¿K¹KÈ:Ó<Ušð$�Lð $ð 8Dö.Ø/3Á{ÐSWÕGX˜˜tŸz™zÒ*ð.Ð*ð .ð '=ò ;˜
Ø+5×+;Ñ+;Ó+=ò ;™K˜D %Ø# y§¡Ò6Ü #Ó#SÈUÔ#SÔ Sà$(Ð,FÒ$FÜ(+ñ ,&ØQVô,&õ )&ð
 ,0×+?Ñ+?ÀÑ+C Ø $× :Ñ :¸8Ô DØ )× 4Ñ 4°TÕ :ñ;ð;ð, !×.Ñ.Ð6Ø!Ÿ™Õ(ð 7ùò9$ùò.s   ¯DÁDÁDÁD)r„   r•  ÚGraph)rk  r¡  r   s     €rW   Úeliminate_to_dtypez`CppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtype.<locals>.eliminate_to_dtypeX  s"   ø€ ð')¼E¿H¹H¿N¹Nõ ')ñR -¨YÕ7rV   )r„   r•  r–  r   rˆ   r#  r  r  rƒ   r"  r~  r‘  Úinserting_afterÚcall_methodrn   r˜  r   Úcpp_to_dtype_countr4   rÓ  rÔ  Úinserting_beforeÚreplace_input_withr2  r3  rŒ  rŽ  r¢  )rk  r  Úsub_graph_nodesr0  r2   rÄ   rÊ   Ú	value_varr§   r‡   r�  rŽ  r£  rr  rˆ   ro  rq  rw  rt  r   r…  s                @@@@@@@@rW   Úadd_to_dtypezDCppKernelProxy.legalize_lowp_fp_dtype_loopbody.<locals>.add_to_dtypeª  s  ÿ€ ð ¤e§h¡h§m¡mð  ¼ÄÇÁÑ8Mó  ð
 ¤u§x¡x§}¡}ð 
 ¼Ä%Ç+Á+Ñ9Nó 
 ð4¬¯©¯©ð 4¼5¿;¹;õ 4ð
	!¤e§h¡h§m¡mð 	!¼¿¹õ 	!ð´5·8±8·=±=ð ÄeÇkÁkö ô # 9§?¡?Ó3ˆOØ)+Ð&Ø(ó w�à—L‘LÐ$:Ñ:Ù/°Ó6Ð6˜¼=ÑHô ÔMÀÇÁÔMÔMØ ØŸ*™* Q™-�CØ"×2Ñ2°5Ó9ñ 8Ø'0×'<Ñ'<Ø&¨c°5¼%¿+¹+Ð-Fð (=ó (˜ð ×3Ñ3Ø(Ó*Iôô  ×2Ò2°aÑ7Õ2÷8ð 8ð —L‘L GÒ+Ù.¨uÓ5Ð5˜¼-ÑGà16·±Ñ.�C˜˜q )¨QÙ3°I¸rÔBØ ÜŸG™G×-Ñ-¨dÓ3�EØ"×3Ñ3°EÓ:ñ 8Ø'0×'<Ñ'<Ø&¨c°9¸eÐ-Dð (=ó (˜ð ×0Ñ0°¸LÔIÜ×2Ò2°aÑ7Õ2÷8ñ 8ð —\‘\ [Ò0ð Ÿ
™
ñØØØ!Ø&Øà ¤MÓ1ð  %Ü!ŸK™KÜ!ŸN™NÜ!ŸM™MÜ!ŸK™Kð	)ñ  ð ð  ð  Ø+0´MÑ+AœEŸKšKÀuÜ!ŸK™KØ*Ø!ð&˜ž
ð —\‘\ ZÒ/°E·J±J¸r±NÄmÑ4Sà',§z¡zÑ$�S˜% ÜÔMÀÇÁÔMÔMÙ Ø"% u¬e¯k©kÐ!:�E–JØ—\‘\ ZÒ/°E·J±J¸r±NÄmÑ4Sà#(§:¡:‘L�S˜!˜RÜÔMÀÇÁÔMÔMÙ ð /×5Ñ5°eÔ<Ø"% q¬%¯+©+Ð!6�E–JØ—\‘\Ð%7Ó7Ø9>¿¹Ñ6�S˜) U¨Ið !¤MÑ1ñ  <¸IÀyÔQØ!*×!;Ñ!;¸EÓ!Bñ @Ø/8×/DÑ/DØ$.°c¸9ÀiÐ5Pð 0Eó 0" ð !&× 8Ñ 8¸ÀLÔ QÜ '× :Ò :¸aÑ ?Õ :÷@ð ¤Ó-ô  ÔUÈÏÉÔUÖUà"'§*¡*¨Q¡-˜CØ!*×!:Ñ!:¸5Ó!Añ @Ø/8×/DÑ/DØ$.°c¸5Ä%Ç+Á+Ð5Nð 0Eó 0" ð !&× ;Ñ ;Ø$0Ó2Qô!"ô !(× :Ò :¸aÑ ?Õ :÷@ñ @ñ ðowðr*8¬e¯h©h¯n©nõ *8ñX ˜yÕ)÷y8ñ 8ú÷ 8ñ 8ú÷J@ð @ú÷ @ñ @ús4   Å?AVÉAVÑ%AVÓ:AV)ÖV	ÖV	ÖV&	Ö)V3	)r„   r•  r¢  r  r#  r  r  rÓ  )r.  r  r«  r
  r  s        rW   Úlegalize_lowp_fp_dtype_loopbodyz.CppKernelProxy.legalize_lowp_fp_dtype_loopbody©  s]   € ðZ	*¤E§H¡H§N¡Nó Z	*ðx  ×*Ñ*Ð+¬d°9×3FÑ3F×3MÑ3MÓ3OÓ.PÑPˆ
Ø#ò 	*ˆIÙ˜Ÿ™Õ)ñ	*rV   c                 óª  ‡ — t        ˆ fd„|D «       «      rÛ|D ]Õ  }|j                  j                  gt        |j                  j                  j                  «       «      z   }|D ]‰  }|j                  j                  D ]n  }|j                  dv sŒ|j                  sJ ‚t        j                  |j                  v sJ ‚|j                  t        j                     }|j                  t        v rŒnJ ‚ Œ‹ Œ× y |D ]^  }t        |t        «      sJ ‚t        |j                  t         «      sJ ‚|j                  }|j#                  «       rŒN‰ j%                  |«       Œ` y )Nc              3   ód   •K  — | ]'  }t        |t        «      xr ‰j                  |«      –— Œ) y ­wr  )r    r$   ri  )r  r0  r.  s     €rW   r  z8CppKernelProxy.legalize_lowp_fp_dtype.<locals>.<genexpr>‹  s3   øè ø€ ò 
àô �uœmÓ,ÒQ°×1JÑ1JÈ5Ó1QÓQñ
ùó   ƒ-0)rí  r×  )r"  rÿ  r  r#  r  r  rÓ  r  r  rk  rA   rj  rˆ   rƒ   r    r$   r   Úis_memory_copyr¬  )r.  r  r0  r
  r  Úfx_noderg  r@  s   `       rW   Úlegalize_lowp_fp_dtypez%CppKernelProxy.legalize_lowp_fp_dtypeŠ  sD  ø€ Üó 
àô
ô 
ð
 ò B�Ø#Ÿk™k×4Ñ4Ð5¼Ø—K‘K×)Ñ)×0Ñ0Ó2ó9ñ �
ð ",ò B�IØ#,§?¡?×#8Ñ#8ò B˜Ø"Ÿ>™>Ð->Ò>Ø#*§<¢<Ð/ <Ü#6×#:Ñ#:¸g¿l¹lÑ#JÐJÐ#JØ;B¿<¹<Ü 3× 7Ñ 7ñ<˜Gð $+§=¡=´MÒ#AÐAÐ#AñBñBð	Bð àò 	;ˆEÜ˜e¤]Ô3Ð3Ð3Ü˜eŸk™k¬8Ô4Ð4Ð4Ø"Ÿ[™[ˆDØ×&Ñ&Õ(Ø×4Ñ4°TÕ:ñ	;rV   c           	      ó^  ‡‡‡ ‡!‡"‡#— t        ‰«      t        ‰«      k(  sJ ‚| j                  Š!t        ‰d„ ¬«      \  Š Š"| j                  ‰ ‰"«       ˆ!ˆ#fd„}ˆˆ ˆ"ˆfd„Š# |t        «      }t
        j                  xj                  |j                  z  c_        t
        j                  xj                  |j                  z  c_        t        j                  |«      | _        | j                  r| j                  s6|g| _        | j                  dd «       | j                  j!                  | «       y t"        j$                  j&                  j)                  d¬«      5  t+        «       }|j-                  ‰‰«      \  }}t        |«      t        |«      k(  sJ ‚d}t/        t1        ‰«      «      }	t3        d„ |	D «       «      rd}d}
d }|r{d}|d	   }|d
z   }t        | j                  j4                  «      |kD  rM| j                  j4                  |   j6                  }| j                  j4                  |   j6                  }|xr | }
t        |«      d
k(  �rt8        xj:                  d
z  c_        | j                  j=                  |d	   |d	   ¬«      } |t>        |d	   |d	   «      }|j@                  |jB                  z
  }|jD                  d	|jB                  fi|_#        t&        jH                  jJ                  r|r |t>        |d	   |d	   |«      }n|}|jD                  g|_&        |jD                  |jB                  |j@                  fi|_#        ||g| _        |}�n¡t        |«      dk(  �rŠ|d
   t        | j                  «      d
z
  k(  r|d	   |d
   k(  sJ ‚t8        xj:                  dz  c_        | j                  j=                  |d	   |d	   ¬«      }d	|jB                  f|jB                  |j@                  fdœ}|j@                  |jB                  z
  }| j                  j=                  |d
   |d	   ¬«      }d	|jB                  f|jB                  |j@                  fdœ}|j@                  |jB                  z
  } |tN        |d	   |«      }|jD                  |d   |jD                  |d   i|_#        g }t&        jH                  jJ                  rh|rfdD ]`  \  }}|dk(  r|nd }|dk(  r|nd } |tN        |d	   |||«      }|jD                  ||   |jD                  ||   i|_#        |jQ                  |«       Œb n¸ |t>        |d	   |d	   «      }|jD                  |d   |jD                  |d   i|_#        |jD                  g|_&        |jQ                  |«       |jD                  |d   |jD                  d	|j@                  fi|_#        |jD                  |jD                  g|_&        |jQ                  |«       |g|z   | _        |}n|g| _        | j                  |
|«       | j                  j!                  | «       d d d «       y # 1 sw Y   y xY w)Nc                 ó   — t        | d   «      S ræ   r  r  s    rW   r�  z2CppKernelProxy.codegen_functions.<locals>.<lambda>«  rH  rV   r  c                 óœ   •—  ‰j                   | g|¢­Ž 5 }t        xj                  dz  c_         ‰|«       |cd d d «       S # 1 sw Y   y xY wræ   )Ú
new_kernelr   Úgenerated_kernel_count)r%  r‘  rZ  rV  Úruns      €€rW   Úcodegen_kernelz8CppKernelProxy.codegen_functions.<locals>.codegen_kernel¯  sL   ø€ Ø(�×(Ñ(¨Ð4¨tÒ4ð ¸ô ×.Ò.°!Ñ3Õ.á�F”Ø÷÷ ò ús   •#AÁAc           	      óT  •— | j                  ‰‰«      \  }}d}t        ‰‰	«      D ]u  \  }}|‰‰ft        t        j                  ‰‰«      «      dffv r|rJ ‚ |||«       Œ;d}|‰dfk(  sJ d|› d‰› d‰› �«       ‚| j                  «       5   ||d«       d d d «       Œw y # 1 sw Y   Œ‚xY w)NFrU   Tzunexpected group: rQ  r™   )rð  r  r¡   rJ  rK  r6  )
rZ  rž  r7  Ú	in_suffixrR  rS  r3  rü  rý  r4  s
         €€€€rW   r¸  z-CppKernelProxy.codegen_functions.<locals>.run¸  sí   ø€ Ø#)×#4Ñ#4°U¸OÓ#LÑ ˆD�.ØˆIÜ!$ W¨nÓ!=ò %‘��IØØ˜OÐ,Üœ9Ÿ?™?¨5°/ÓBÓCÀRÐHð!ñ ñ  )Ð(˜=Ù�t˜^Õ,à $�IØ$ØØð)ò ð Vð ,¨I¨;°d¸5¸'ÀÀOÐCTÐUóVð ð
  ×/Ñ/Ó1ñ %Ù˜4 œ÷%ð %ñ%÷%ð %ús   Â	
BÂB'	F©Úinplace_buffersTc              3   ó,   K  — | ]  }|t         v–— Œ y ­wr  )rw   r  s     rW   r  z3CppKernelProxy.codegen_functions.<locals>.<genexpr>ä  s   è ø€ ÒS¸u�5Ô :Ô:ÑSùr  r   r5   )rê  r   ©ÚmainÚtailrÀ  )r¿  )rÁ  rÀ  )rÁ  rÁ  rÁ  ))rÐ   rV  r_   rð  ru  r4   rÓ  Úremoved_buffersÚinplaced_to_removerF  r'  rH  ra  rH  rb  Úaggregate_reduction_buffersÚ
set_kernelr„   Ú	_inductorr   Úpatchr  rA  rB   rC   rf   rB  r…  r   Úgenerated_cpp_vec_kernel_countÚtiler�  r:  Ú
tiled_sizer¥   rz  r¹  Úenable_loop_tail_vecr{  rr  rŽ  )$r.  r3  r4  r¹  Úscalar_kernelÚtiling_selectÚtiling_factorsrÓ  Úcould_masked_vecr6  Ú_inner_loop_reduction_outer_notÚ_outer_loopÚinner_loop_reductionÚouter_loop_levelÚinner_loop_levelÚouter_loop_reductionr  Ú
vec_kernelrX  Útail_kernelÚ
outer_looprR  rÕ  Ú
inner_loopÚinner_rangesrÔ  Útile2d_kernelÚouter_rÚinner_rÚ_inner_tail_sizeÚ_outer_tail_sizerZ  rü  rV  rý  r¸  s$    ``                             @@@@rW   Úcodegen_functionsz CppKernelProxy.codegen_functions¨  s  ý€ Ü�7‹|œs >Ó2Ò2Ð2Ð2Ø×(Ñ(ˆÜ!$ ^Ñ9TÔ!UÑˆˆà�‰˜˜Ô/õ	÷	%ñ( '¤yÓ1ˆÜ	�‰×Ò =×#@Ñ#@Ñ@ÕÜ	�‰×"Ò" m×&FÑ&FÑFÕ"Ü!Ÿ™¨Ó6ˆŒà×"Ò"¨$¯-ª-Ø)˜?ˆDŒLØ×,Ñ,¨U°DÔ9Ø�N‰N×%Ñ% dÔ+Øô �_‰_×#Ñ#×)Ñ)¸%Ð)Ó@ñ C	,Ü(›NˆMØ-:×-HÑ-HØ˜ó.Ñ*ˆN˜Nô �~Ó&¬#¨nÓ*=Ò=Ð=Ð=à#ÐÜ3´NÀ7Ó4KÓLˆJÜÑSÈ
ÔSÔSà#(Ð à.3Ð+ØˆKÙØ',Ð$Ø#1°!Ñ#4Ð Ø#3°aÑ#7Ð Ü�t—~‘~×+Ñ+Ó,Ð/?Ò?Ø+/¯>©>×+?Ñ+?Ø(ñ,ç"‘lð )ð ,0¯>©>×+?Ñ+?Ø(ñ,ç"‘lð )ð -ÒIÐ5IÐ1Ið 4ô �>Ó" aÓ'Ü×6Ò6¸!Ñ;Õ6Ø—~‘~×*Ñ*¨>¸!Ñ+<À^ÐTUÑEVÐ*ÓW�Ù+Ü  .°Ñ"3°^ÀAÑ5Fó�
ð !ŸI™I¨¯©Ñ7�	Ø,0¯H©H°q¸$¿/¹/Ð6JÐ+K�
Ô(Ü—:‘:×2Ò2Ñ7GÙ"0Ü$Ø& qÑ)Ø& qÑ)Ø!ó	#‘Kð #0�KØ48·H±H°:�MÔ0Ø-1¯X©X¸¿¹ÈÏÉÐ7SÐ,T�Ô)Ø *¨KÐ8�”Ø"’Ü�^Ó$¨Ó)à" 1Ñ%¬¨T¯]©]Ó);¸aÑ)?Ò?Ø& qÑ)¨^¸AÑ->Ò>ðð?ô ×6Ò6¸!Ñ;Õ6Ø!Ÿ^™^×0Ñ0Ø" 1Ñ%¨n¸QÑ.?ð 1ó �
ð  
× 5Ñ 5Ð6Ø'×2Ñ2°J·O±OÐDñ �ð #-§/¡/°J×4IÑ4IÑ"I�Ø!Ÿ^™^×0Ñ0Ø" 1Ñ%¨n¸QÑ.?ð 1ó �
ð  
× 5Ñ 5Ð6Ø'×2Ñ2°J·O±OÐDñ �ð #-§/¡/°J×4IÑ4IÑ"I�Ù .Ü#Ø" 1Ñ%Ø"ó!�ð —N‘N L°Ñ$8Ø—N‘N L°Ñ$8ð/�Ô+ð !�Ü—:‘:×2Ò2Ñ7Gð-ò 3Ñ(˜ ð 07¸&Ò/@™OÀdð )ð 07¸&Ò/@™OÀdð )ñ "0Ü+Ø*¨1Ñ-Ø*Ø,Ø,ó"˜ð 'ŸN™N¨L¸Ñ,AØ&ŸN™N¨L¸Ñ,Að0˜Ô,ð $×*Ñ*¨6Õ2ñ-3ñ0 "0Ü$ n°QÑ&7¸ÈÑ9Jó"�Jð #Ÿ™¨°VÑ(<Ø"Ÿ™¨°VÑ(<ð0�JÔ,ð 2<·±Ð0@�JÔ-Ø×&Ñ& zÔ2à"Ÿ™¨°VÑ(<Ø"Ÿ™¨¨J¯O©OÐ(<ð3�MÔ/ð 5?·N±NÀJÇNÁNÐ3S�MÔ0Ø×&Ñ& }Ô5Ø -˜°Ñ<�”Ø(‘à -˜�”Ø×,Ñ,Ø/°ôð �N‰N×%Ñ% dÔ+÷GC	,÷ C	,ñ C	,ús   ÅSX#Ø#X,c                 ó‚   — |D ](  }| j                  |«       t        j                  |«       Œ* | j                  ||«       y r  )r¬  r:   Úpropagate_loopbodyrà  )r.  r5  r4  r@  s       rW   Úcodegen_loop_bodiesz"CppKernelProxy.codegen_loop_bodies`  s?   € Øò 	9ˆDØ×0Ñ0°Ô6Ü×2Ñ2°4Õ8ð	9ð 	×Ñ˜{¨NÕ;rV   r  c                 óÎ  — | j                  |«       | j                  |«       t        |«      dk\  sJ ‚d„ }|D �cg c]  }t        j                  ||«      ‘Œ }}t        t        j                  t        «      r2t        j                  j                  rd„ }|D �cg c]
  } ||«      ‘Œ }}|D �cg c]  }|j                  d   ‘Œ }}| j                  ||«       y c c}w c c}w c c}w )Nr5   c                 ó¾   — | j                  «        | j                  «        t        t        j                  t
        «      r | j                  |Ž S | j                  |«      S r  )Údecide_inplace_updateÚmark_runr    r4   rZ  r1   rÿ  Úcodegen)r  Ú
index_varss     rW   rR  z(CppKernelProxy.codegen_nodes.<locals>.fnl  sF   € Ø×&Ñ&Ô(Ø�M‰MŒOÜœ!Ÿ(™(Ô$5Ô6Ø!�t—z‘z :Ð.Ð.à—|‘| JÓ/Ð/rV   c                 óR   — t         j                  j                  | «      }| |_        |S r  )r4   r  Úlocalize_functionÚoriginal_fn)rR  Ú
wrapped_fns     rW   Úwrap_fnz-CppKernelProxy.codegen_nodes.<locals>.wrap_fn{  s+   € Ü×3Ñ3×EÑEØó�
ð *,�
Ô&Ø!Ð!rV   )r²  rf  rÐ   rI  Úpartialr    r4   r  rJ   r  rü  rà  )r.  r  rR  r  r3  rî  r4  s          rW   Úcodegen_nodeszCppKernelProxy.codegen_nodesf  sÏ   € à×#Ñ# EÔ*Ø×"Ñ" 5Ô)Ü�5‹z˜QŠÐˆò	0ð <AÖA°4”9×$Ñ$ R¨Õ.ÐAˆÐAô ”q×-Ñ-Ô/AÔBÜ×&Ñ&×4Ò4ò"ð .5Ö5 r‘w˜r•{Ð5ˆGÐ5à49Ö:¨D˜$Ÿ*™* Q›-Ð:ˆÐ:Ø×Ñ˜w¨Õ7ùò# Bùò 6ùâ:s   ºCÂCÂ-C"c                 ó>   — | j                  | j                  ||«       y r  )r%  rH  )r.  r°   r  s      rW   r(  zCppKernelProxy.codegen_loops‡  s   € Ø×Ñ §¡°°kÕBrV   c                 óF   — | j                   D ]  }|j                  «        Œ y r  )rb  r�  ©r.  rZ  s     rW   r�  z4CppKernelProxy.update_stores_with_parallel_reductionŠ  s!   € Ø—l‘lò 	;ˆFØ×8Ñ8Õ:ñ	;rV   r°   c                 ó(  — |€J ‚d}| j                   D ]q  }t        j                  «       5 }|j                  ||«      r@d}|j	                  |j                  «       «       |j                  |j                  «       «       d d d «       Œs y # 1 sw Y   Œ~xY w)NÚ
C10_LIKELYÚC10_UNLIKELY)rb  r½   r¾   rG  rÀ   rÁ   r  r¤  )r.  r°   Ú	if_prefixrZ  rÉ   s        rW   r¤  zCppKernelProxy.gen_bodyŽ  sŠ   € ØÐÐÐØ ˆ	Ø—l‘lò 	3ˆFÜ×%Ñ%Ó'ð 3¨5Ø×,Ñ,¨T°9Ô=Ø .�IØ×'Ñ'¨¯©«Ô6Ø—K‘K §¡Ó 1Ô2÷	3ð 3ñ	3÷3ð 3ús   ªABÂB	Úinner_loop_reduction_outer_notrØ  Ú	LoopLevelc                 ót  ‡ — dˆ fd„}‰ j                   d   }|r|sJ ‚ ||«       nZ|j                  «        ‰ j                  j                  |j                  «       ‰ j                  j                  |j                  «       ‰ j
                  j                  |j
                  «       ‰ j                  j                  |j                  «       ‰ j                  j                  |j                  «       ‰ j                  j                  |j                  «       ‰ j                  j                  |j                  «       y )Nc           
      ó¶  •— t        ‰j                  «      dk\  sJ ‚‰j                  d   }‰j                  d   }t        |t        «      sJ ‚t	        |«      t
        k(  r^|j                  |j                  «       |j                  «        ‰j                  j                  |j                  |j                  z   «       n5|j                  «        ‰j                  j                  |j                  «       t        «       }t        j                  «       5 }|j                  |d| j                  «      r:|j                  |j!                  «       «       |j                  |j"                  «       d d d «       t        j                  «       5 }|j                  |d| j                  «      �r	|j                  |j!                  «       «       t	        |«      t
        k(  r½|j$                  }|D ]X  }|› d| j                  › dt'        | j(                  «      › d�}t+        |j,                  ||«       t+        |j"                  ||«       ŒZ |j                  t/        |j"                  | j                  | j                  › d	�| j(                  | j0                  «      «       n|j                  |j"                  «       d d d «       |‰_        y # 1 sw Y   �ŒVxY w# 1 sw Y   |‰_        y xY w)
Nr   r   rH  rõ  rö  rÔ   z_tail - rË  r   )rÐ   rb  r    r�  r  ru  rÞ  r  r~  r  r7   r½   r¾   rG  r¥   rÀ   rÁ   r€  r�  rE   rÊ  rä   r›  rÍ   r:  )	rØ  Úmain_loop_kernelÚtail_loop_kernelÚ
suffix_bufrÉ   r7  rÄ   rá   r.  s	           €rW   Ú!aggregate_reduction_prefix_suffixzUCppKernelProxy.aggregate_reduction_buffers.<locals>.aggregate_reduction_prefix_suffixž  sŒ  ø€ Ü�t—|‘|Ó$¨Ò)Ð)Ð)Ø#Ÿ|™|¨A™ÐØ#Ÿ|™|¨BÑ/ÐÜÐ.´Ô=Ð=Ð=ô Ð$Ó%¬Ò2ð !×:Ñ:Ø$×2Ñ2ôð !×:Ñ:Ô<Ø×%Ñ%×,Ñ,Ø$×5Ñ5Ø&×7Ñ7ñ8õð
 !×:Ñ:Ô<Ø×%Ñ%×,Ñ,Ð-=×-NÑ-NÔOô &›ˆJÜ×%Ñ%Ó'ð I¨5Ø#×6Ñ6Ø ¨j¯n©nôð ×'Ñ'¨
×(9Ñ(9Ó(;Ô<Ø×%Ñ%Ð&6×&GÑ&GÔH÷Iô ×%Ñ%Ó'ð M¨5Ø#×6Ñ6Ø °
·±õð ×'Ñ'¨
×(9Ñ(9Ó(;Ô<ÜÐ,Ó-´Ò:Ø)9×)MÑ)M˜Ø$2ò ˜DØ*.¨¨u°Z·^±^Ð4DÀHÌ[ÐYc×YnÑYnÓMoÐLpÐpqÐ'r˜HÜ,Ð-=×-DÑ-DÀdÈHÔUÜ,Ø 0× AÑ AÀ4Èõðð #×)Ñ)Ü6Ø 0× AÑ AØ *§¡Ø#-§>¡>Ð"2°%Ð 8Ø *× 5Ñ 5Ø *§¡óõð #×)Ñ)Ð*:×*KÑ*KÔL÷/Mð0 %/ˆDÕ!÷=Iñ Iú÷Mð0 %/ˆDÕ!ús   ÄAJ;ÆD(KÊ;KËKr   )rØ  rù  )
rb  rÞ  r~  r  r€  r�  r‚  rƒ  r„  r†  )r.  rø  rØ  rÿ  Úmain_kernels   `    rW   rÄ  z*CppKernelProxy.aggregate_reduction_buffers˜  sð   ø€ õ6	/ðp —l‘l 1‘oˆÙ)ÙÐ�:Ù-¨jÕ9à×1Ñ1Ô3Ø×!Ñ!×(Ñ(¨×)EÑ)EÔFØ×!Ñ!×(Ñ(¨×)EÑ)EÔFØ×&Ñ&×-Ñ-¨k×.SÑ.SÔTØ×&Ñ&×-Ñ-¨k×.SÑ.SÔTØ×!Ñ!×(Ñ(¨×)IÑ)IÔJØ×#Ñ#×*Ñ*¨;×+MÑ+MÔNØ×*Ñ*×1Ñ1Ø×5Ñ5õ	
rV   r  )r  r  r  r-  rf  r$   ri  r   r¬  r²  rà  rã  r#  rð  r(  r�  r   r7   r¤  rk   rÄ  r`  ra  s   @rW   r   r   �  s…   ø„ ô+ò@ð
°=ó 
ð_*¸ó _*òB;ò<v,òp<ð8 4¨Ñ#6ó 8òBCò;ñ3˜X lÑ3ó 3ðL
Ø.2ðL
Ø@HÈÑ@U÷L
rV   r   c                   ó$   ‡ — e Zd Zˆ fd„Zd„ Zˆ xZS )rW  c                 óp   •— t         ‰| �  |j                  |j                  j                  «       g | _        y r  )r,  r-  r‘  r`  r�  rY  rc  s     €rW   r-  zOuterLoopFusedKernel.__init__è  s)   ø€ Ü‰Ñ˜×*Ñ*¨L¯O©O×,GÑ,GÔHØ%'ˆ�
rV   c           
      ó   — g }| j                   D �cg c]  }|j                  «       ‘Œ }}|D ]h  }|j                  }|€J ‚|j                  |j	                  t        t        |«      |j                  z
  |j                  ¬«      |«      j                  «       Œj t        t        |j                  t        |«      «      |j                  ¬«      S c c}w )Nr/  )rY  rÿ  r|  rŽ  r  r	  rÐ   r  r
  r^   r_   )r.  r  r  Úkernels_parallel_depthrH  Únested_kernelsrZ  r|  s           rW   r  z*OuterLoopFusedKernel.decide_parallel_depthì  sÞ   € Ø!#Ðà48·J±Jö+
Ø'0ˆI× Ñ Õ"ð+
ˆð +
ð %ò 	ˆFð !×,Ñ,ˆKØÐ*Ð*Ð*Ø"×)Ñ)Ø×,Ñ,Ü!ä Ó,Ð/A×/MÑ/MÑMà$6×$BÑ$Bô	ð ó÷ !‘.õ
ð	ô  ÜØ"×1Ñ1´3Ð7MÓ3Nóð +×6Ñ6ô	
ð 	
ùò)+
s   ‘C)r  r  r  r-  r  r`  ra  s   @rW   rW  rW  ç  s   ø„ ô(ö
rV   rW  c                   ó   — e Zd ZdZdZdZy)ÚReasonFusedNodesÚsame_vars_reduceÚcompatible_reductionÚcompatible_ranges_no_reductionN)r  r  r  ÚSAME_VARS_REDUCEÚCOMPATIBLE_REDUCTIONÚCOMPATIBLE_RANGES_NO_REDUCTIONrU   rV   rW   r  r  
  s   „ Ø)ÐØ1ÐØ%EÑ"rV   r  c                   ó†  ‡ — e Zd ZdZ eej                  ej                  g«      Ze	de
j                  dee   fd„«       Zˆ fd„Zdefd„Zd„ Zd	„ Zd
„ Zdee   fd„Zd„ Zd„ Zd„ Zdededefd„Zd„ Zd„ Zd„ Zd„ Zdee    fd„Z!de"fd„Z#de$e"e%e f   fd„Z&dedefd„Z'dede(e   de(e   fd„Z)d „ Z*d!„ Z+d"„ Z,d%d#„Z-d$„ Z.ˆ xZ/S )&ÚCppSchedulingiô  Údevicerµ   c                 ó   — | j                   S r  )Úbackend_features)r%  r  s     rW   Úget_backend_featuresz"CppScheduling.get_backend_features  s   € à×#Ñ#Ð#rV   c                 óV   •— t         ‰| �  |«       |r| j                  «        d| _        y ©NF)r,  r-  Úreset_kernel_groupÚ_ready_to_flush)r.  r!  r1  s     €rW   r-  zCppScheduling.__init__   s'   ø€ Ü‰Ñ˜Ô#ÙØ×#Ñ#Ô%Ø$ˆÕrV   Ústatusc                 ó   — || _         y r  ©r  )r.  r  s     rW   Ú_set_flush_statuszCppScheduling._set_flush_status&  s
   € Ø%ˆÕrV   c                 ó&   — t        d„ |D «       «      S )Nc              3   óŠ   K  — | ];  }t        t        t        j                  j                  j
                  |«      «      –— Œ= y ­wr  )r¡   r’  r4   rÓ  rõ  rë   rE  s     rW   r  z)CppScheduling.group_fn.<locals>.<genexpr>*  s,   è ø€ ÒMÀ!”Uœ3œqŸw™w×/Ñ/×8Ñ8¸!Ó<×=ÑMùs   ‚AA)r¡   )r.  r  s     rW   Úgroup_fnzCppScheduling.group_fn)  s   € ÜÑMÀuÔMÓMÐMrV   c                 ó"   — t        «       | _        y r  )ÚKernelGrouprV  r3  s    rW   r  z CppScheduling.reset_kernel_group,  s   € Ü'›MˆÕrV   c                 ó”  ‡— |j                  «       s|j                  «       rt        j                  ||«      S |j                  «       r(|j                  «       rJ ‚t	        j                  ||«      S | j                  ||«      t        j                  k(  �rçt        |t        t        f«      sJ ‚t        |t        t        f«      sJ ‚|j                  \  }\  }}|j                  \  }\  }}|dk(  r|dk(  s	J ||f«       ‚ˆfd„Št        |«      t        |«      k  r|n|}t        |t        «      sJ ‚t        |«      t        |«      k  r|n|}	 ‰|	«      }
|j                  |
¬«       |j                  \  }\  }}|j                  \  }\  }}||k(  rt	        j                  ||«      S  ‰|«      }t        |	t        «      r|	j                  |¬«       ngt        |	t        «      sJ ‚|	j                  D ]&  }t        |t        «      sJ ‚|j                  |¬«       Œ( t	        |	j                  |	j                  «      }	|j                  \  }\  }}|j                  \  }\  }}||k(  s	J ||f«       ‚t	        j                  ||«      S | j                  ||«      r't         j                  ||| j#                  ||«      «      S t	        j                  ||«      S )NrU   c                 óD  •— t        | t        «      r“t        | j                  «      dkD  sJ | j                  «       ‚d }t	        t
           «       }| j                  D ];  } ‰	|«      \  }}|€|}||k(  sJ ||| j                  f«       ‚|j                  |«       Œ= |t        |«      fS t        | t        «      sJ ‚| j                  }t        |t        j                  «      sJ ‚|j                  «       \  }}}|j                  t        |j                  j                  «       «      fS rU  )r    r"   rÐ   Úsnodesr   r   Úupdater#  r$   r  r   ÚComputedBufferÚget_default_sizes_bodyr¼  Úindexing_exprsr  )
r  r¼  r'  ÚsnodeÚvÚexprsÚcomp_bufferrÊ   r@  Úget_indexing_ranges_exprss
            €rW   r,  z5CppScheduling.fuse.<locals>.get_indexing_ranges_exprsA  s  ø€ Ü! $Ô(:Ô;Ü" 4§;¡;Ó/°!Ò3Ð@°T·[±[Ó@Ð3Ø%)˜
Ü)3´C©Ó):˜Ø%)§[¡[ò 9˜EÙ'@ÀÓ'G™H˜A˜uØ)Ð1Ø-. 
Ø#-°¢?ÐP°ZÀÀDÇKÁKÐ4PÓP ?Ø*×1Ñ1°%Õ8ð9ð  *¬4°Ó+?Ð?Ð?ä)¨$´Ô>Ð>Ð>Ø&*§i¡i˜Ü)¨+´r×7HÑ7HÔIÐIÐIØ%0×%GÑ%GÓ%I™
˜˜4 Ø#Ÿ™´°T×5HÑ5H×5OÑ5OÓ5QÓ0RÐRÐRrV   )Úextra_indexing_constraints)Ú
is_foreachr!   r'  Úis_templater"   Ú_why_fuse_nodesr  r  r    r$   rü  rÐ   Úrecompute_size_and_bodyr#  r!  Úcan_fuse_vertical_outer_loopr  Ú_get_outer_loop_fusion_depth)r.  r  r  rÊ   Úvars1Úreduce1Úvars2Úreduce2Únode_to_recompÚref_nodeÚref_indexing_constraintsÚ#node_to_recomp_indexing_constraintsr(  r,  s                @rW   r'  zCppScheduling.fuse/  sÀ  ø€ Ø×ÑÔ ×!1Ñ!1Ô!3Ü-×2Ñ2°5¸%Ó@Ð@Ø×ÑÔ Ø×(Ñ(Ô*Ð*Ð*Ü%×*Ñ*¨5°%Ó8Ð8ð ×$Ñ$ U¨EÓ2Ü#×BÑBóCô " %¬-Ô9KÐ)LÔMÐMÐMÜ! %¬-Ô9KÐ)LÔMÐMÐMà&+§k¡kÑ#�Ñ#�E˜7Ø&+§k¡kÑ#�Ñ#�E˜7Ø "’}¨°BªÐJ¸À'Ð8JÓJÐ6ôSô& +.¨e«*´s¸5³zÒ*A¡Àu�Ü! .´-Ô@Ð@Ð@ä$'¨£J´°U³Ò$;™5À�á+DÀXÓ+NÐ(à×6Ñ6Ø/Gð 7ô ð !&§¡‘�‘:�E˜1Ø %§¡‘�‘:�E˜1à˜E’>Ü-×2Ñ2°5¸%Ó@Ð@ñ 7PØ"ó7Ð3ô ˜h¬Ô6Ø×4Ñ4Ø3Vð 5õ ô & hÔ0BÔCÐCÐCØ!)§¡ò ˜Ü)¨%´Ô?Ð?Ð?Ø×5Ñ5Ø7Zð 6õ ðô
  2°(×2DÑ2DÀhÇoÁoÓV�Hà %§¡‘�‘:�E˜1Ø %§¡‘�‘:�E˜1Ø ’~Ð5¨¨u ~Ó5�~Ü)×.Ñ.¨u°eÓ<Ð<Ø×2Ñ2°5¸%Ô@Ü2×7Ñ7Ø˜5 $×"CÑ"CÀEÈ5Ó"Qóð ô *×.Ñ.¨u°eÓ<Ð<rV   c                 óþ   — |j                   \  }\  }}|j                   \  }\  }}||k(  r||k(  rt        j                  S |dk(  r|||z   k(  rt        j                  S | j	                  ||«      rt        j
                  S y )NrU   )rü  r  r  r  Ú&_can_fuse_nodes_with_compatible_rangesr  )r.  r  r  rÊ   r4  r5  r6  r7  s           rW   r0  zCppScheduling._why_fuse_nodes�  s   € Ø#Ÿk™kÑˆÑˆE�7Ø#Ÿk™kÑˆÑˆE�7à�EŠ>˜g¨Ò0Ü#×4Ñ4Ð4Ø�bŠ=˜U e¨g¡oÒ5Ü#×8Ñ8Ð8Ø×6Ñ6°u¸eÔDÜ#×BÑBÐBàrV   c                 ó<  — |j                   \  }\  }}|j                   \  }\  }}|dk(  xr |dk(  }t        j                  |«      t        j                  |«      k(  }	t        |«      dk(  xs t        |«      dk(  }
|r|	r|
syt        |«      t        |«      k  r|n|}t        |«      t        |«      k  r|n|}t	        |t
        «      ryt	        |t        «      sJ ‚t	        |j                  t        j                  «      ryt	        |j                  t        j                  «      sJ ‚|j                  j                  j                  «       }d }t	        |t
        «      ràt        t        t        df      «       }|j                   D ]Š  }t	        |j                  t        j                  «      r ndt	        |j                  t        j                  «      sJ ‚|j#                  t        |j                  j                  j                  «       «      «       ŒŒ t        |«      dk7  ryt%        t'        t)        |«      «      «      }n\t	        |t        «      sJ ‚t	        |j                  t        j                  «      sJ ‚|j                  j                  j                  «       }||k7  ryy)NrU   r5   F.T)rü  rY  rd   rÐ   r    r"   r$   r  r   ÚTemplateBufferr%  ÚdataÚget_sizer   r¡   r   r#  r�  r#  ÚnextÚiter)r.  r  r  rÊ   r4  r5  r6  r7  Úc1Úc2Úc3r8  r9  Úranges2Úranges1Ú
ranges_setr(  s                    rW   r=  z4CppScheduling._can_fuse_nodes_with_compatible_rangesŽ  s  € ð $Ÿk™kÑˆÑˆE�7Ø#Ÿk™kÑˆÑˆE�7à˜‰]Ò,˜w¨"™}ˆÜ�Y‰Y�uÓ¤§¡¨5Ó!1Ñ1ˆÜ�‹Z˜1‰_Ò/¤ E£
¨a¡ˆÙ‘r™bØä"% e£*¬s°5«zÒ"9™¸uˆÜ ›J¬¨U«Ò3‘5¸ˆô �nÔ&8Ô9Øô ˜.¬-Ô8Ð8Ð8Ü�n×)Ñ)¬2×+<Ñ+<Ô=ØÜ˜.×-Ñ-¬r×/@Ñ/@ÔAÐAÐAð !×%Ñ%×*Ñ*×3Ñ3Ó5ˆØˆÜ�hÔ 2Ô3Ü#¤E¬#¨s¨(¡OÑ4Ó6ˆJØ!Ÿ™ò B�Ü˜eŸj™j¬"×*;Ñ*;Ô<ÙÜ! %§*¡*¬b×.?Ñ.?Ô@Ð@Ð@Ø—‘œu U§Z¡Z§_¡_×%=Ñ%=Ó%?Ó@ÕAð	Bô �:‹ !Ò#Øäœ4¤ ZÓ 0Ó1Ó2‰Gä˜h¬Ô6Ð6Ð6Ü˜hŸm™m¬R×->Ñ->Ô?Ð?Ð?Ø—m‘m×(Ñ(×1Ñ1Ó3ˆGà�gÒØàrV   c                 ó´   — t        |t        t        f«      sJ ‚t        |t        t        f«      sJ ‚t        d„ ||fD «       «      ry| j	                  ||«      d uS )Nc              3   ó<   K  — | ]  }t        |t        «      –— Œ y ­wr  )r    r  r  s     rW   r  z:CppScheduling._can_fuse_horizontal_impl.<locals>.<genexpr>É  s   è ø€ ò 
Ø>BŒJ�tÔ8×9ñ
ùs   ‚F)r    r"   r$   rf   r0  ©r.  r  r  s      rW   Ú_can_fuse_horizontal_implz'CppScheduling._can_fuse_horizontal_implÆ  sd   € Ü˜%Ô"4´mÐ!DÔEÐEÐEÜ˜%Ô"4´mÐ!DÔEÐEÐEÜñ 
ØGLÈeÀnô
ô 
ð Ø×#Ñ# E¨5Ó1¸Ð=Ð=rV   c                 ó  — |j                  «       s|j                  «       ryt        |j                  «       «      t        |j                  «       «      z   t        j                  j
                  kD  ry| j                  ||«      S r  )r/  rÐ   r+  r   r¹  Úmax_horizontal_fusion_sizerM  rL  s      rW   Úcan_fuse_horizontalz!CppScheduling.can_fuse_horizontalÏ  sf   € Ø×ÑÔ %×"3Ñ"3Ô"5Øä�—‘Ó!Ó"¤S¨¯©Ó):Ó%;Ñ;Ü�j‰j×3Ñ3ò4ð à×-Ñ-¨e°UÓ;Ð;rV   r  r  c                 óp  — |j                  «       x}r¤t        |j                  t        j                  «      xr~ t        |j
                  t        j                  «      xrX t        |j
                  j                  «      dk(  xr4 |j
                  j                  d   j                  «       |j                  k(  S y)Nr5   r   F)Úget_template_noder    Úlayoutr   ÚMultiOutputLayoutr  ÚMultiOutputrÐ   Úinputsr  rÄ   )r.  r  r  Útemplate_bufs       rW   Úcan_fuse_multi_outputs_templatez-CppScheduling.can_fuse_multi_outputs_templateÚ  s›   € ð !×2Ñ2Ó4Ð4ˆ<Ð4ä˜<×.Ñ.´×0DÑ0DÓEò IÜ˜uŸz™z¬2¯>©>Ó:òIä˜Ÿ
™
×)Ñ)Ó*¨aÑ/òIð —J‘J×%Ñ% aÑ(×1Ñ1Ó3°|×7HÑ7HÑHð	ð rV   c                 óì  — d}t        d„ ||fD «       «      s|S t        |t        «      r|j                  «       d   n|}t        |t        t
        f«      sJ ‚t        |t        «      r|j                  «       d   n|}t        |t        t
        f«      sJ ‚|j                  \  }\  }}|j                  \  }\  }	}
|dk(  r|	dk(  r|dk7  r|
dk7  r|S t        d„ ||fD «       «      r'|j                  |j                  k(  r|j                  S |S t        t        |«      t        |	«      «      }|dk\  rI|d | |	d | k(  r>t        d„ ||fD «       «      r(t        |«      t        u r|n|}|j                  |k(  r|S |S |S |S )Nr   c              3   óT   K  — | ]   }t        |«      t        t        t        fv –— Œ" y ­wr  )r  r  r"   r$   r  s     rW   r  z=CppScheduling._get_outer_loop_fusion_depth.<locals>.<genexpr>è  s/   è ø€ ò 
ð ô �‹JÜ+Ô-?ÄÐOôPñ
ùr  rH  rU   c              3   ó>   K  — | ]  }t        |«      t        u –— Œ y ­wr  r  r  s     rW   r  z=CppScheduling._get_outer_loop_fusion_depth.<locals>.<genexpr>  r  r   r5   c              3   ó>   K  — | ]  }t        |«      t        u –— Œ y ­wr  r  r  s     rW   r  z=CppScheduling._get_outer_loop_fusion_depth.<locals>.<genexpr>  s   è ø€ ò Ø>B”�T“
Ô9Ô9ñùr   )r"  r    r  r$  r"   r$   rü  r&  r^   rÐ   rf   r  )r.  r  r  ÚDISABLE_OUTER_LOOP_FUSIONÚ_node1Ú_node2rÊ   r4  r5  r6  r7  r&  Ú_compare_nodes                rW   r3  z*CppScheduling._get_outer_loop_fusion_depthæ  sÇ  € Ø$%Ð!Üñ 
ð  ˜ô
ô 
ð
 -Ð,ô ˜%Ô!<Ô=ð ×!Ñ!Ó# BÒ'àð 	ô
 ˜&Ô#5´}Ð"EÔFÐFÐFô ˜%Ô!<Ô=ð ×!Ñ!Ó# AÒ&àð 	ô
 ˜&Ô#5´}Ð"EÔFÐFÐFà$Ÿl™lÑˆÑˆE�7Ø$Ÿl™lÑˆÑˆE�7Ø�BŠ;˜5 Bš;¨7°bª=¸WÈº]à,Ð,ÜÑTÀeÈUÀ^ÔTÔTð ×0Ñ0°E×4QÑ4QÒQð ×-Ñ-ðð /ðô
 #&¤c¨%£j´#°e³*Ó"=Ðà# qÒ(ØÐ.Ð.Ð/°5Ð9QÐ:QÐ3RÒRäñ ØGLÈeÀnôô ô " %›[Ô,GÑG‘EÈUð ð !×8Ñ8Ð<SÒSà2Ð2à4Ð4ð /Ð.Ø(Ð(rV   c                 ó  — |j                  «        xro |j                  «        xr\ |j                  «       |j                  z  xr= | j                  ||«      xr |j	                  «         xr | j                  ||«      dk\  S ræ   )r/  Úget_operation_namesÚ	ancestorsrM  r…  r3  rL  s      rW   r2  z*CppScheduling.can_fuse_vertical_outer_loop  s”   € à×!Ñ!Ó#Ð#ò EØ×%Ñ%Ó'Ð'òEà×)Ñ)Ó+¨e¯o©oÑ=òEð ×.Ñ.¨u°eÓ<ò -Ø×*Ñ*Ó,Ð,ðòEð ×1Ñ1°%¸Ó?À1ÑDð		
rV   c                 ó*   — | j                  ||«      ryyr'  )r2  rL  s      rW   Úget_fusion_pair_priorityz&CppScheduling.get_fusion_pair_priority(  s   € Ø×,Ñ,¨U°EÔ:ààrV   c                 ó   — |j                  «       ry|j                  «       r%t        ||g«      \  }}|j                  «        xr |S | j                  ||«      xr |j                  «        xs | j	                  ||«      S r  )r/  rM   r…  rM  r2  )r.  r  r  Útemplate_fusion_supportedrÊ   s        rW   Úcan_fuse_verticalzCppScheduling.can_fuse_vertical/  s‰   € Ø×ÑÔàØ×ÑÔÜ+SØ˜�wó,Ñ(Ð% qð ×)Ñ)Ó+Ð+ÒIÐ0IÐIà×*Ñ*¨5°%Ó8ÒUÀ×ASÑASÓAUÐ=Uò=à×.Ñ.¨u°eÓ<ð	=rV   r  c                 óè  ‡‡‡‡‡— t        d„ |D «       «      r|S dŠdŠd}d}d}d}|D �]z  }t        |j                  t        j                  «      sJ ‚|j                  j                  «       \  }}}|j                  j                  «       D �]  \  Š}	|	j                  t        «      D ]÷  Št        ˆfd„|j                  D «       «      r‰|k7  r‰}|dz  }|dkD  r|c c c S t        ‰j                  d   t        j                  j                  j                  «      sŒv‰j                  d   |j                  v sŒ’‰€Œ•t!        ˆˆfd„|j                  j                  «       D «       «      sŒÃ‰j                  d   dkD  sŒÖ‰j                  d   Š‰j                  d   Šd	}|}Œù �Œ �Œ} |s|S dŠˆˆˆfd
„}
|D ]  }||k(  sŒ	|j#                  |
¬«       Œ |D ]  }||k7  sŒ	|j#                  ‰|
¬«       Œ |S )aI  
        Apply loop split optimization.
        When one of the indexing_exprs contains a division, we eliminate the division by splitting the loop
        to avoid non-contiguous loads, subject to the following conditions:
            1. No reduction and no mudular index for all nodes.
            2. The indexing_exprs of all nodes contain only one (or more, but all the same) division,
               where the divisor is an integer and not too small (the divisor > 8), the dividend is
               one of the iter_vars, and this var, i.e. the dimension that needs to be split, is
               contiguous in all other indexing_exprs.

        For example, if the node's var_ranges: {z0: 2, z1: 9216, z2: 960} and indexing_exprs:
        {'index0': 8847360*z0 + 960*z1 + z2, 'index1': 32*z0 + (z2//30), 'index2': z2},
        we will split z2 -> 30*z2 + z3, then the node's var_ranges will be changed to
        {z0: 2, z1: 9216, z2: 32, z3: 30} and indexing_exprs will be changed to
        {'index0': 8847360*z0 + 960*z1 + 30*z2 + z3, 'index1': 32*z0 + z2, 'index2': 30*z2 + z3}.
        c              3   óÄ   K  — | ]X  }t        |j                  d    d    «      dk7  xs4 t        d„ |j                  j                  j                  «       D «       «      –— ŒZ y­w)r5   r   c              3   óF   K  — | ]  }|j                  t        «      –— Œ y ­wr  )rç   r   )r  r�  s     rW   r  z9CppScheduling.try_loop_split.<locals>.<genexpr>.<genexpr>Q  s   è ø€ ò Ø.2�—‘œ×)ñùr  N)rÐ   rü  rf   rÿ  r'  r  r  s     rW   r  z/CppScheduling.try_loop_split.<locals>.<genexpr>O  sc   è ø€ ò 
ð
 ô	 �—
‘
˜1‘˜aÑ Ó! QÑ&ò Üñ Ø6:·j±j×6OÑ6O×6VÑ6VÓ6Xôó óñ
ùs   ‚AA Nr   Fc              3   ó@   •K  — | ]  }‰j                  |«      –— Œ y ­wr  )rç   )r  r¥   Údiv_exprs     €rW   r  z/CppScheduling.try_loop_split.<locals>.<genexpr>e  s   øè ø€ ÒQ°#˜HŸL™L¨×-ÑQùr  r5   c              3   ód   •K  — | ]'  \  }}|‰k7  rt        |‰j                  d    «      dv –— Œ) y­w)r   r)  N)r  r‘  )r  Úname_Úexpr_rm  rÄ   s      €€rW   r  z/CppScheduling.try_loop_split.<locals>.<genexpr>p  s9   øè ø€ ò  á ,  uØ$¨š}ô 0°°x·}±}ÀQÑ7GÓHÈFÔRñ ùr¯  rX  Tc                 óä  •— | \  }}|\  }}|j                  ‰«      }|j                  «       }||   ‰z  ||<   |j                  |dz   ‰«       t        j                  ||d¬«      \  \  }	}
}|	j                  «       }|j                  |dz   «      }‰||   z  |z   ||<   t        j                  |||g||	|«      }‰s/|j                  t        |j                  j                  «       «      fŠ||f||	|ffS )Nr5   rð  )r  )r’   ÚcopyÚinsertr
   Úindex_vars_no_squeezeÚpopr   r   r¼  r#  r'  r  )r  r@  rž  Ú
index_sizeÚreduce_sizeré  Úreduce_varsÚ	split_idxÚnew_index_sizeÚnew_index_varsrÊ   r¼  Ú	iter_varsÚdivisor_varr-  Úsplit_numberÚ	split_vars                 €€€rW   Ú
loop_splitz0CppScheduling.try_loop_split.<locals>.loop_split‚  s  ø€ Ø&+Ñ#ˆJ˜Ø&*Ñ#ˆJ˜Ø"×(Ñ(¨Ó3ˆIØ'Ÿ_™_Ó.ˆNØ(2°9Ñ(=ÀÑ(MˆN˜9Ñ%Ø×!Ñ! )¨a¡-°Ô>Ü.:×.PÑ.PØ °Cô/Ñ+Ñˆ^˜Q ð '×+Ñ+Ó-ˆIØ#Ÿ-™-¨	°A©Ó6ˆKØ#/°)¸IÑ2FÑ#FÈÑ#TˆI�iÑ Ü—;‘;Ø�y +Ð.°
¸NÈKóˆDñ .à—O‘OÜ˜×,Ñ,×3Ñ3Ó5Ó6ð.Ð*ð
   Ð-ØØ Ð-ðð rV   )Úrecompute_sizes_body_func)r-  r�  )rf   r    r  r   r%  r&  r'  rŸ  Úfindr   r|  r‘  rè   ÚcoreÚnumbersrN  r"  r1  )r.  r  Únum_divÚ	div_expr_Ú	match_divÚmatched_noder  rÊ   Úoriginal_bodyr�  r€  rm  r-  rÄ   r~  r  s              @@@@@rW   Útry_loop_splitzCppScheduling.try_loop_split<  s
  ü€ ô& ñ 
ð
 ô
ô 
ð ˆLàˆ	ØˆØˆØˆ	Øˆ	Øˆàó 	,ˆDÜ˜dŸi™i¬×):Ñ):Ô;Ð;Ð;Ø"&§)¡)×"BÑ"BÓ"DÑˆAˆ}˜aØ+×:Ñ:×@Ñ@ÓBó ,‘
��dØ $§	¡	¬(Ó 3ò ,�HäÓQ¸×9PÑ9PÔQÔQØ$¨	Ò1à$,˜	Ø 1™˜Ø ’{Ø$žä" 8§=¡=°Ñ#3´U·Z±Z×5GÑ5G×5OÑ5OÕPØ$ŸM™M¨!Ñ,°×0GÑ0GÒGØ Ñ,Üô  à0=×0LÑ0L×0RÑ0RÓ0Tô õ ð
 %ŸM™M¨!Ñ,¨qÓ0à$,§M¡M°!Ñ$4˜	Ø'/§}¡}°QÑ'7˜Ø$(˜	Ø'+™ò/,ò,ð	,ñ< ØˆLà%)Ð"ö	ð: ò 	SˆDØ�|Ó#Ø×,Ñ,ÀzÐ,ÕRð	Sð ò 	ˆDØ�|Ó#Ø×,Ñ,Ø/IØ.8ð -õ ð	ð ˆrV   r  c                 óH  ‡‡‡	— | j                   Št        j                  }g Šg Š	t        |t        «      sJ ‚dt        fˆˆˆ	fd„} ||«      sË|t        _        ‰j                  «        ‰	j                  «        t        j                  j                  j                  d¬«      5  |j                  «       D ]X  }t        |t        t        f«      sJ ‚|j                  «       }t        ‰«      }|j                  |«       ‰j!                  ||«       ŒZ 	 ddd«       yy# 1 sw Y   yxY w)a  
        Generate the code for the outer loop fused scheduler node.
        1. Codegen with fused outer loop: depends on the analysis of
            the outer loop fused scheduler node, with or without the local buffer.
        2. If failed, fallback to standard codegen.
        r  c           	      óv  •‡ ‡‡‡‡‡— t        ‰ t        «      sJ ‚‰j                  «        ‰j                  «        dt        fd„Šg }i Št	        ˆˆ fd„‰ j                  «       D «       «      �rét        t           «       Š‰ j                  «       D �]Ä  Št        ‰t        «      sJ ‚‰j                  ‰j                  «       «       ‰j                  «       st        ‰j                  «       «      dk7  rŒb‰j                  «       d   Št	        ˆ fd„‰j                  D «       «      sŒ”‰j                   }t        |t"        j$                  «      sJ ‚|j'                  «       }‰ j(                  t         ‰‰«      «      z
  }ˆˆfd„}|j+                  «       r |«       s�Œ	t#        j,                  |j.                  |j0                  |j2                  |d |j4                  |d «      }ˆˆfd	„}d
} |||«      }	|	sEt#        j6                  |› dt        |«      › �|¬«      }	|j9                  |	«       g ‰|	j:                  <   ‰|	j:                     j9                  |«       �ŒÇ t=        ‰j>                  «      5 }
t        |«      dkD  r4|D ]/  }|j:                  €J ‚|
jA                  |‰|j:                     «       Œ1 ‰ j                  «       D ]t  }t        |tB        t        f«      sJ ‚tE        ‰«      }|jG                  |j                  «       «       ‰j9                  |«       ‰j9                  |j                  «       «       Œv ‰ jI                  ‰‰ j(                  «      sD|
jJ                  D ]+  }tL        jN                  jJ                  jQ                  |«       Œ- 	 ddd«       ytR        jT                  j9                  tS        jV                  t        ‰«      t        |
jX                  «      ¬«      «       ‰ j[                  ‰«      }‰j]                  |g t^        j`                  jc                  ‰«      ¢«       ddd«       y# 1 sw Y   yxY w)zN
            Codegen code with fused outer loop and local Buffer.
            r  c                 óÀ   — t        | t        t        f«      sJ ‚| j                  «       }t	        |d„ ¬«      j
                  \  }\  }}t        |«      t        |«      z   }|S )Nc                 ó4   — t        | j                  «       «      S r  )rm   r…  r–  s    rW   r�  z~CppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.get_call_ranges.<locals>.<lambda>Ç  s   € ¬¨Q¯^©^Ó-=Ó)>€ rV   r  )r    r$   r"   r+  r_   rü  r¡   )r  r  rÊ   rü  rý  r|  s         rW   Úget_call_rangeszlCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.get_call_rangesÃ  s`   € Ü! $¬Ô8JÐ(KÔLÐLÐLØ-1¯^©^Ó-=�Ü.1ØÑ>ô/ç‘%ñ ,�Ñ+�E˜?ô $ E›l¬U°?Ó-CÑC�Ø"Ð"rV   c              3   ó`   •K  — | ]%  }t         ‰|«      «      ‰j                  d z   k(  –— Œ' y­w)r5   N)rÐ   r&  )r  r0  r�  r  s     €€rW   r  zfCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.<genexpr>Ï  s3   øè ø€ ò àô ‘O EÓ*Ó+¨t×/KÑ/KÈaÑ/OÕOñùs   ƒ+.r5   r   c              3   óV   •K  — | ]   }|j                   ‰j                  «       v –— Œ" y ­wr  )r  r+  )r  r{  r  s     €rW   r  zfCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.<genexpr>ä  s&   øè ø€ ò Ø:>˜Ÿ	™	 T§^¡^Ó%5Ô5ñùó   ƒ&)c                  óB  •‡‡— dŠd} t        ‰j                  j                  j                  «       «      D ]  \  }}‰| |z  z  Š| |z  } Œ ‰j                  j	                  ‰j                  «       «      }ˆfd„Š ‰|«      xr t        ˆˆfd„‰j                  D «       «      S )Nr   r5   c                 ó   •— | ‰k(  S r  rU   )rŽ  Úcontiguous_index_exprs    €rW   Úis_contiguous_indexz–CppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.is_all_write_read_contiguous.<locals>.is_contiguous_indexú  s   ø€ Ø'(Ð,AÑ'AÐ ArV   c              3   óÈ   •K  — | ]Y  }t        |j                  t        «      xr9  ‰|j                  j                  j	                  ‰j                  «       «      «      –— Œ[ y ­wr  )r    r  r$   rÿ  Úget_read_exprr  )r  r{  r–  Úscheduler_buffers     €€rW   r  zŒCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.is_all_write_read_contiguous.<locals>.<genexpr>ý  s\   øè ø€ ò Qð %)ô !+¨4¯9©9´mÓ Dò !"Ù$7Ø$(§I¡I§O¡O×$AÑ$AØ(8×(AÑ(AÓ(Có%&ó%"ó!"ñQùs   ƒAA")rû  rÿ  r¼  rŸ  Úget_write_exprr  r"  r~  )rN  r¥   rG  Úwrite_index_exprr•  r–  r™  rg  s       @@€€rW   Úis_all_write_read_contiguouszyCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.is_all_write_read_contiguousî  s·   ú€ Ø45Ð1Ø%&˜FÜ.6Ø .× 4Ñ 4× ?Ñ ?× EÑ EÓ Gó/ò 0¡
  Uð !6¸À#¹Ñ EÐ 5Ø &¨%¡¡ð	0ð
 0>×/CÑ/C×/RÑ/RØ 0× 9Ñ 9Ó ;ó0Ð,ôBñ $7Ð7GÓ#Hò $ÌSô Qð -=×,BÑ,BôQó Nð rV   Nc                 ó~   •— |D ]7  }| |j                   k(  sŒt        ˆfd„‰|j                     D «       «      sŒ5|c S  y )Nc              3   óÐ   •K  — | ]]  }|j                   �Ot        ˆfd„t        j                  j                  j
                  |j                      j                  D «       «      –— Œ_ y ­w)Nc              3   óV   •K  — | ]   }|j                   j                  «       ‰v –— Œ" y ­wr  )r  r  )r  r{  Úvisited_scheduler_nodess     €rW   r  z�CppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.try_share_local_buffer.<locals>.<genexpr>.<genexpr>  s,   øè ø€ ò (&à,0ð )-¯	©	×(:Ñ(:Ó(<Ð@WÔ(Wñ(&ùr’  )rÄ   r"  r4   rÓ  r!  Úname_to_bufr~  )r  Úglobal_bufferr   s     €rW   r  z†CppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.try_share_local_buffer.<locals>.<genexpr>  s`   øè ø€ ò S"ð )6ð (5×'9Ñ'9Ð'Eô %(ó (&ä45·G±G×4EÑ4E×4QÑ4QØ,9×,>Ñ,>ñ5*ç*/©%ô	(&÷ %&ñS"ùs   ƒA#A&)rS  r"  rÄ   )Úlocal_buffer_layoutr  Ú	local_bufÚlocal_to_global_buffersr   s      €€rW   Útry_share_local_bufferzsCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf.<locals>.try_share_local_buffer  sS   ø€ Ø-:ò 5 	Ø#6¸)×:JÑ:JÓ#JÌsó S"ð :QØ(1¯©ñ:&ôS"õ P"ð ,5Ò$4ð5ð $(rV   Úlocal_buffer_datarÊ   )rÄ   rS  F)Úlocal_buffer_numberT)2r    r  Úclearr   r"  r$  r   ro   r+  r$   r�  r  r…  rÐ   Úget_outputsr~  r  r   r%  r  r&  Úis_contiguousÚFixedLayoutr  rˆ   r:  rN  ÚBufferrŽ  rÄ   rJ   r‘  Úadd_local_bufferr"   r   rð  rS  rÂ  r4   rÓ  Úremover   Ú!cpp_outer_loop_fused_inner_countsÚCppOuterLoopFusedCountr  r^  Úfinalize_kernelrJ  rK  Úfrom_iterable)r  r  r¢  Úglobal_buffer_layoutÚsize_offsetrœ  r£  r¦  Úlocal_buf_prefixÚlocal_buffer_usedÚscoper  r0  rQ  Úremoved_bufferÚouter_fusion_cpp_kernel_proxyr�  r¥  r™  rg  r   rO  rV  Ú
nodes_lists   `               @@@@@€€€rW   Ú$try_outer_loop_fusion_with_local_bufzSCppScheduling.codegen_outer_loop_node.<locals>.try_outer_loop_fusion_with_local_buf»  s
  þ€ ô ˜dÔ$?Ô@Ð@Ð@Ø!×'Ñ'Ô)Ø×ÑÔð#Ô&7ó #ð .0ˆMàBDÐ#Üô à!×1Ñ1Ó3ôõ ô +5´S©/Ó*;Ð'Ø&*§n¡nÓ&6ó [�Nä% n´mÔDÐDÐDØ+×/Ñ/°×0GÑ0GÓ0IÔJà&×3Ñ3Ô5Ü˜~×9Ñ9Ó;Ó<ÀÒAà à'5×'AÑ'AÓ'CÀAÑ'FÐ$Üó ØBR×BXÑBXôõ ð )9×(=Ñ(=˜Ü)¨-¼×9JÑ9JÔKÐKÐKØ/<×/GÑ/GÓ/IÐ,Ø&*×&BÑ&BÄSÙ+¨NÓ;óFñ '˜õð4 1×>Ñ>Ô@Ù <Ô >á$ä.0¯n©nØ0×7Ñ7Ø0×6Ñ6Ø0×5Ñ5°k°lÐCØ0×7Ñ7¸¸ÐEó	/Ð+õ(ð" ,?Ð(á,BØ/°ó-Ð)ñ  1ä02·	±	Ø(8Ð'9¸¼3¸}Ó;MÐ:NÐ%OØ':ô1Ð-ð *×0Ñ0Ð1BÔCØNPÐ3Ð4E×4JÑ4JÑKØ/Ð0A×0FÑ0FÑG×NÑNØ)öðs[ôz $ L×$5Ñ$5Ó6ð "¸%Ü�}Ó%¨Ò)Ø(5ò ˜Ø+×0Ñ0Ð<Ð<Ð<Ø×.Ñ.Ø(Ð*AÀ,×BSÑBSÑ*Tõðð
 "×1Ñ1Ó3ò 9�EÜ% eÔ.@Ä-Ð-PÔQÐQÐQÜ'5°lÓ'CÐ$Ø$×2Ñ2°5·?±?Ó3DÔEØ)×0Ñ0Ð1AÔBØ×%Ñ% e§o¡oÓ&7Õ8ð9ð ×>Ñ>Ø)¨4×+GÑ+Gôð +0×*?Ñ*?ò G˜ô Ÿ™×/Ñ/×6Ñ6°~ÕFðGð !÷+"ð "ô, ×9Ñ9×@Ñ@Ü×2Ñ2ÜÐ1Ó2Ü,/°×0CÑ0CÓ,Dôôð 15×0OÑ0OØ)ó1Ð-ð ×,Ñ,Ø1Ø@”i—o‘o×3Ñ3°JÓ?Ð@ô÷?"ðH ÷I"ðH ús   É*D!P/ÎBP/Ð/P8Fr¼  N)rV  r   rÈ  r    r  r©  r„   rÆ  r   rÇ  r$  r"   r$   r+  r   rð  r²  )
r.  r  rÈ  r¼  r0  Ú_nodesrQ  rO  rV  r»  s
          @@@rW   Úcodegen_outer_loop_nodez%CppScheduling.codegen_outer_loop_node«  s  ú€ ð ×(Ñ(ˆÜ)0×)OÑ)OÐ&Ø68ÐØ02ˆ
Ü˜$Ô ;Ô<Ð<Ð<ð_	Ô7R÷ _	ñB 4°DÔ9à5SŒGÔ2Ø!×'Ñ'Ô)Ø×ÑÔô —‘×'Ñ'×-Ñ-¸eÐ-ÓDñ KØ!×1Ñ1Ó3ò K�EÜ% eÔ.@Ä-Ð-PÔQÐQÐQØ27·/±/Ó2C�FÜ'5°lÓ'CÐ$Ø$×2Ñ2°6Ô:Ø ×0Ñ0Ð1AÀ6ÕJñK÷Kð Kð :÷Kð Kús   Â!A,DÄD!c                 óh  — | j                   }t        |t        «      r| j                  |«       nO|j	                  «       }| j                  |«      }t        |«      }|j                  |«       |j                  ||«       | j                  «       }|t        j                  kD  r| j                  d«       yy)zC
        Turn an set of pre-fused nodes into a C++ kernel.
        TN)rV  r    r  r¾  r+  rŠ  r   rð  r²  Ú_get_scheduled_num_argsr  ÚMAX_FUSED_KERNEL_ARGS_NUMr  )r.  r  rV  r  rQ  Úargs_nums         rW   Úcodegen_nodezCppScheduling.codegen_nodel  s�   € ð ×(Ñ(ˆä�dÔ7Ô8Ø×(Ñ(¨Õ.à)-¯©Ó)9ˆEØ×'Ñ'¨Ó.ˆEÜ-¨lÓ;ÐØ×*Ñ*¨5Ô1Ø×(Ñ(Ð)9¸5ÔAà×/Ñ/Ó1ˆØ”m×=Ñ=Ò=Ø×"Ñ" 4Õ(ð >rV   c                 ón   — t        |t        «      xr$ t        |j                  t        j                  «      S r  )r    r$   r  r   ÚCppTemplateBuffer)r.  r  s     rW   Úis_cpp_templatezCppScheduling.is_cpp_template‚  s,   € Ü˜$¤Ó.ò 
´:Ø�I‰I”r×+Ñ+ó4
ð 	
rV   Útemplate_nodeÚepilogue_nodesÚprologue_nodesc                 óÈ  — |rJ ‚|D �cg c]  }t        |t        t        f«      r|‘Œ }}t        d   dxx   dz  cc<   t        d   dxx   t	        |«      z  cc<   | j                  |«      sJ d«       ‚t        t        |«      }|j                  \  }\  }}|dk(  sJ ‚t        t        j                  |j                  «      }|D �cg c]  }|j                  ‘Œ }	}t        d„ |	D «       «      sJ d«       ‚d	„ }
 |
||j                  |	«      }|j                  |||	¬
«      \  }}|5  t        |j                  «      s|j                  «        |D ]  }|j                  «        Œ  |«       }ddd«       t!        j"                  |«      5  |g|¢}| j%                  ||j&                  «      }ddd«       t        |j                  «      r­t	        |j(                  «      dk(  sJ d«       ‚|j(                  d   j*                  D ]r  }t        |j                  t,        «      sJ d«       ‚t        |j                  j                  t        j.                  «      sJ d«       ‚|j                  j                  «        Œt |j1                  |«       t         j2                  xj4                  |j4                  z  c_        | j7                  «        yc c}w c c}w # 1 sw Y   �ŒfxY w# 1 sw Y   �Œ3xY w)zG
        Codegen a CPP template, possibly with fused epilogues
        ÚinductorÚcpp_templated_kernel_counterr5   Úcpp_epilogue_fusion_counterzlTemplate node passed to CppScheduler.codegen_template must be a SchedulerNode that wraps a CppTemplateBufferrU   c              3   óP   K  — | ]  }t        |t        j                  «      –— Œ  y ­wr  )r    r   r%  )r  rï  s     rW   r  z1CppScheduling.codegen_template.<locals>.<genexpr>¦  s   è ø€ ÒO¸”:˜a¤×!2Ñ!2×3ÑOùrF  z9Epilogue nodes must all be instances of ir.ComputedBufferc                 ó–   ‡— ‰sy| j                  «       |v sJ ‚|| j                  «          j                  }t        ˆfd„|D «       «       S )NFc              3   ó†   •K  — | ]8  }t        |j                  t        «      xr |j                  j                  ‰v –— Œ: y ­wr  )r    r  r   )r  r{  rÈ  s     €rW   r  zZCppScheduling.codegen_template.<locals>.template_buffer_has_other_users.<locals>.<genexpr>²  sA   øè ø€ ò ð ô ˜4Ÿ9™9Ô&7Ó8ò 5Ø—I‘I—N‘N nÐ4ó5ñùs   ƒ>A)r  r~  r"  )Útemplate_bufferÚoutputs_by_namerÈ  r~  s     ` rW   Útemplate_buffer_has_other_userszGCppScheduling.codegen_template.<locals>.template_buffer_has_other_usersª  s\   ø€ ñ "Øà"×+Ñ+Ó-°Ñ@Ð@Ð@Ø# O×$<Ñ$<Ó$>Ñ?×EÑEˆEÜó ð "ôó ð ð rV   )Ú$flag_template_buffer_has_other_usersrÈ  NzSMulti outputs template should be with 1 output template buffer of MultiOutputLayoutr   z?Multi outputs template should be with ExternKernelSchedulerNodez7Multi outputs template has multi users with MultiOutput)r    r$   r"   r   rÐ   rÆ  r   rü  r   rÅ  r  r"  rÒ  Úmake_kernel_renderr)   rç  r4   Úset_kernel_handlerÚdefine_kernelr‘  Úoutputsr~  r    rU  Úcall_kernelrÓ  rÂ  Úfree_buffers_in_scheduler)r.  rÇ  rÈ  rÉ  Úepilogue_noderÊ   ÚrnumelÚctbrï  Úepilogue_ir_nodesrÓ  rÔ  rZ  Úrenderr  Úsrc_codeÚnode_scheduleÚkernel_namer{  s                      rW   Úcodegen_templatezCppScheduling.codegen_template‡  sì  € ñ "Ð!Ð!ð
 "0ö
àÜ˜-¬-Ô9KÐ)LÔMò ð
ˆð 
ô 	�ÑÐ;Ó<ÀÑAÓ<Ü�ÑÐ:Ó;¼sÀ>Ó?RÑRÓ;Ø×#Ñ# MÔ2ð 	
Øzó	
Ð2ô œ]¨MÓ:ˆØ&×,Ñ,‰ˆ‰;ˆAˆvØ˜Š|Ðˆ|Ü$(¬×)=Ñ)=¸}×?QÑ?QÓ$Rˆà*ö;
ØˆA�F‹Fð;
Ðð ;
ô ÑOÐ=NÔOÔOð 	
ØGó	
ÐOò	ñ 0OØ�×.Ñ.Ð0Aó0
Ð,ð ×/Ñ/ØØ1UØ,ð 0ó 
‰ˆ�ð
 ñ 	 Ü,¨]×-?Ñ-?Ô@Ø×&Ñ&Ô(Ø&ò  �Ø—‘•ð á“xˆH÷	 ô ×!Ñ! &Ó)ñ 	SØ*Ð<¨^Ð<ˆMØ×,Ñ,¨X°}ÀfÇkÁkÓRˆK÷	Sô % ]×%7Ñ%7Ô8ô �}×,Ñ,Ó-°Ò2ð ØeóÐ2ð &×-Ñ-¨aÑ0×6Ñ6ò %�Ü! $§)¡)Ô-FÔGð ØUóÐGô " $§)¡)§.¡.´"·.±.ÔAð ØMóÐAð —	‘	×"Ñ"Õ$ð%ð 	×Ñ˜;¨Ô,Ü	�‰×Ò 6×#9Ñ#9Ñ9ÕØ×&Ñ&Õ(ùòS
ùò ;
÷:	 ñ 	 ú÷	Sñ 	Sús$   ‰K ÃKÄ"AK
Æ#KË
KËK!c                 ó6   — | j                   j                  «       S r  )rV  Úget_num_argsr3  s    rW   rÀ  z%CppScheduling._get_scheduled_num_argsÞ  s   € Ø× Ñ ×-Ñ-Ó/Ð/rV   c                 ó   — | j                   S r  r  r3  s    rW   Úready_to_flushzCppScheduling.ready_to_flushá  s   € Ø×#Ñ#Ð#rV   c                  ó   — y r  rU   r3  s    rW   Úcodegen_synczCppScheduling.codegen_syncä  s   € ØrV   c                 óî  — t         j                  j                  }t        j                  j
                  r$t        |t        j                  j
                  «      nd}dj                  d||j                  «       g«      }t         j                  j                  r|nd}|j                  t        t        j                  «      |«      }|j                  t        t        j                  «      |«      }|j                  dd«      }|j                  d«      }|j!                  d|«      }	|||	d	z    › d
�}
t#        «       }|€| j$                  j&                  n|}|j)                  «       \  }}}t         j                  j                  s|j+                  d|›d�«       |j-                  |d¬«       t         j                  j                  s|j+                  d«       |j/                  ||j1                  «       d|
¬«       |S )NrR   rÊ   r¹  rZ  z#pragma CMTz//z
extern "C"rš   r5   z;
zasync_compile.cpp_pybinding(z, '''T)Ústripz''')F)ÚgpuÚcpp_definition)r4   rÓ  Úwrapper_coder   r¹  Údescriptive_namesr'   rD  Únext_kernel_suffixÚcpp_wrapperr   ro   r,   ÚKERNEL_NAMEÚDESCRIPTIVE_NAMEÚrfindr‚  r=   rV  r‘  Úcpp_argdefsr¿   r  r×  Úgetvalue)r.  rà  r  Úkernel_argsr›  Ú
fused_namerâ  Úkernel_decl_nameÚ
first_charÚ	last_charÚkernel_definitionÚcompile_wrapperr‘  rÊ   Ú	arg_typess                  rW   r×  zCppScheduling.define_kernelç  s²  € Ü—'‘'×&Ñ&ˆô �z‰z×+Ò+ô " %¬¯©×)EÑ)EÔFàð 	ð
 —h‘h  z°7×3MÑ3MÓ3OÐPÓQˆÜ*+¯'©'×*=Ò*=™;À8ÐØ×#Ñ#¤C¬×(?Ñ(?Ó$@ÐBRÓSˆØ×#Ñ#¤C¬×(DÑ(DÓ$EÀ{ÓSˆð ×#Ñ# M°4Ó8ˆð —^‘^ LÓ1ˆ
Ø—M‘M # zÓ2ˆ	Ø'¨
°YÀ±]ÐCÐDÀCÐHÐä(Ó*ˆØ)4Ð)<ˆt× Ñ ×%Ò%À+ˆØ×*Ñ*Ó,‰ˆˆ1ˆiÜ�w‰w×"Ò"Ø×%Ñ%Ð(DÀYÀMÐQVÐ&WÔXØ×Ñ˜x¨tÐÔ4Ü�w‰w×"Ò"Ø×%Ñ% fÔ-Ø×ÑØØ×$Ñ$Ó&ØØ,ð	 	ô 	
ð ÐrV   c                 ó2  — | j                   j                  «       }|rZ| j                  || j                   j                  «      }| j                   j	                  t
        j                  j                  |«       | j                  «        | j                  d«       y r  )
rV  Úcodegen_groupr×  Úscheduled_nodesrÙ  r4   rÓ  rî  r  r  )r.  rà  râ  s      rW   ÚflushzCppScheduling.flush  sv   € Ø×$Ñ$×2Ñ2Ó4ˆÙØ×,Ñ,Ø˜$×+Ñ+×;Ñ;óˆKð ×Ñ×)Ñ)¬!¯'©'×*>Ñ*>ÀÔLØ×ÑÔ!Ø×Ñ˜uÕ%rV   r  )0r  r  r  rÁ  r   r6   ÚINPLACE_BUFFERSÚREDUCE_TO_SINGLE_ELEMENTr  r_  r„   r  r  r-  rk   r  r  r  r'  r   r  r0  r=  rM  rP  r   rX  r3  r2  re  rh  r#  r$   rŠ  r  r¾  r	   r"   rÃ  rÆ  r   rã  rÀ  rç  ré  r×  r  r`  ra  s   @rW   r  r    sp  ø„ ð !$ÐÙ!à×*Ñ*Ø×3Ñ3ð	
óÐð ð$¨%¯,©,ð $¸:ÀnÑ;Uò $ó ð$ô%ð&¨ó &òNò*òP=ðd¨xÐ8HÑ/Ió ò6òp>ò	<ð
Ø&ð
Ø/@ð
à	ó
ò4)òl

òò=ðm D¨Ñ$7ó mð^Kà)óKðB)àÐ/Ð1CÀ]ÐRÑSó)ð,
Ð$5ð 
¸$ó 
ð
U)à(ðU)ð !Ð!2Ñ3ðU)ð !Ð!2Ñ3ó	U)òn0ò$òó#öJ&rV   r  c                   óD   ‡ — e Zd Zˆ fd„Zd„ Zd„ Zd„ Zddefd„Zd„ Z	ˆ xZ
S )	r   c                 ó  •— t         ‰| �  «        t        «       | _        t	        «       | _        t        | j
                  «      | _        t        j                  «       | _
        | j                  j                  | j                  «       g | _        y r  )r,  r-  r?   r‘  r7   Ú
loops_codeÚWorkSharingr`  r½   r¾   rÉ   rÀ   r  r  s    €rW   r-  zKernelGroup.__init__  s^   ø€ Ü‰ÑÔÜ“LˆŒ	Ü&›.ˆŒÜ˜dŸo™oÓ.ˆŒÜ×)Ñ)Ó+ˆŒ
Ø�
‰
× Ñ  §¡Ô)Ø!ˆÕrV   c                 ó:   —  || j                   t        «       g|¢­Ž S r  )r‘  r+   )r.  r%  r‘  s      rW   r¶  zKernelGroup.new_kernel!  s   € Ù�4—9‘9Ô2Ó4Ð<°tÒ<Ð<rV   c                 ó‚   — | xj                   |z  c_         | j                  }| j                  }|j                  ||«       y r  )r  r  r`  r(  )r.  r¶  r  r°   r`  s        rW   r²  zKernelGroup.finalize_kernel$  s5   € Ø×Ò Ñ%ÕØ�‰ˆØ�W‰WˆØ× Ñ   rÕ*rV   c                 óX   — | j                   j                  «       \  }}}t        |«      }|S r  )r‘  rõ  rÐ   )r.  Úarg_defsÚ
_call_argsÚ
_arg_typesrÂ  s        rW   rå  zKernelGroup.get_num_args*  s)   € Ø+/¯9©9×+@Ñ+@Ó+BÑ(ˆ�*˜jÜ�x“=ˆØˆrV   rµ   c           	      ó.  — | j                   j                  «        | j                  syt        «       }t        j
                  j                  xr t        j                  dv }|r|j                  dg«       |j                  t        j                  «       «       |€t        t        j                  «      n|}|€t        t        j                   «      n|}| j"                  j%                  «       \  }}}dj'                  d«      j)                  |«      }t+        «       }|j                  d|› d|› d|› d	�«       |j-                  «       5  |rHt.        j0                  j2                  }	|	�d
t        |	«      z   dz   nd}
|j                  d|
|z   › d�g«       | j"                  j5                  «       D ]  \  }}|j                  d|› d|› d�«       Œ |j7                  | j8                  «       d d d «       |j;                  «       S # 1 sw Y   |j;                  «       S xY w)NrR   )ÚlinuxrP   z!#include <ATen/record_function.h>z,
é   zextern "C" z void rº  rš   Úgraph_rÊ   zRECORD_FUNCTION("z#", c10::ArrayRef<c10::IValue>({}));rÒ   r·   r¸   )rÉ   r	  r  r7   r   r¹  Úenable_kernel_profileÚsysÚplatformrÛ   r¿   r   Ú
cpp_prefixro   r,   rò  ró  r‘  rõ  ÚljustrD  rX   rÁ   r4   rÓ  Úgraph_idÚaliasesr  r  rö  )r.  rÄ   r°   r  rù  râ  r  rÊ   Úfunc_export_declr  r  ÚoldÚnews                rW   r   zKernelGroup.codegen_group/  sÞ  € Ø�
‰
×ÑÔØ×#Ò#ØÜ‹~ˆô !'§
¡
× @Ñ @ò !
ÄSÇ\Á\ð V
ð F
Ðñ !Ø�O‰OÐ@ÐAÔBØ�‰”y×+Ñ+Ó-Ô.ð <@¸<œ3œ{×6Ñ6Ô7ÈTÐØ;?¸<”cœ+×6Ñ6Ô7ÈTˆØŸ™×.Ñ.Ó0‰ˆ�!�QØ—;‘;˜r“?×'Ñ'¨Ó1ˆÜ1Ó3ÐØ�‰ØÐ*Ð+¨6Ð2BÐ1CÀ1ÀXÀJÈaÐPô	
ð
 �[‰[‹]ñ 	)Ù$ÜŸ7™7×+Ñ+�Ø;CÐ;O˜¤C¨£MÑ1°CÒ7ÐUW�Ø—‘à+¨F°[Ñ,@Ð+AÐAfÐgðôð
 !ŸI™I×-Ñ-Ó/ò 7‘��SØ—‘  s e¨3¨s¨e°1Ð5Õ6ð7à�K‰K˜Ÿ™Ô(÷	)ð �}‰}‹Ð÷	)ð �}‰}‹Ðús   ÅB G<Ç<Hc                 ój   — | j                   j                  «       \  }}}|j                  ||d|¬«       y )NF)Útritonrþ  )r‘  rõ  Úgenerate_kernel_call)r.  r›  râ  rÊ   Ú	call_argsrþ  s         rW   rÙ  zKernelGroup.call_kernelW  s7   € Ø"&§)¡)×"7Ñ"7Ó"9Ñˆˆ9�iØ×$Ñ$Ø˜¨5¸Ið 	%õ 	
rV   r  )r  r  r  r-  r¶  r²  rå  ro   r   rÙ  r`  ra  s   @rW   r   r     s)   ø„ ô"ò=ò+òñ
&¨#ó &öP
rV   r   c                   ó0   — e Zd Zd„ Zd„ Zd„ Zd„ Zd„ Zd„ Zy)r  c                 ó`   — || _         d| _        d | _        t        j                  «       | _        y r  )r°   Úin_parallelr�  r½   r¾   rÉ   )r.  r°   s     rW   r-  zWorkSharing.__init___  s)   € ØˆŒ	Ø ˆÔØˆÔÜ×)Ñ)Ó+ˆ�
rV   c                 óÖ  — | j                   r|| j                  k7  r| j                  «        | j                   s²|| _        d| _         t        j                  j
                  r| j                  j                  d«       n| j                  j                  d|› d�«       | j                  j                  | j                  j                  «       «       | j                  j                  d«       y y )NTz#pragma omp parallelz!#pragma omp parallel num_threads(rš   zint tid = omp_get_thread_num();)r#  r�  r	  r   r¹  r”  r°   r¿   rÉ   rÀ   rÁ   )r.  r  s     rW   rû  zWorkSharing.parallele  s°   € Ø×Ò ¨4×+;Ñ+;Ò ;à�J‰JŒLØ×ÒØ&ˆDÔØ#ˆDÔÜ�z‰z×)Ò)Ø—	‘	×#Ñ#Ð$:Õ;à—	‘	×#Ñ#Ð&GÈÀyÐPQÐ$RÔSØ�J‰J×$Ñ$ T§Y¡Y×%5Ñ%5Ó%7Ô8Ø�I‰I×ÑØ1õð  rV   c                 óh   — | j                   r| j                  j                  d«       | j                   S )Nz#pragma omp single)r#  r°   r¿   r3  s    rW   r  zWorkSharing.singleu  s*   € Ø×ÒØ�I‰I×ÑÐ 4Ô5Ø×ÑÐrV   c                 óF   — | j                   j                  «        d| _        y r  )rÉ   r	  r#  r3  s    rW   r	  zWorkSharing.closez  s   € Ø�
‰
×ÑÔØ ˆÕrV   c                 ó:   — | j                   j                  «        | S r  )rÉ   rm  r3  s    rW   rm  zWorkSharing.__enter__~  s   € Ø�
‰
×ÑÔØˆrV   c                 ó>   — | j                   j                  |||«       y r  )rÉ   rs  ro  s       rW   rs  zWorkSharing.__exit__‚  s   € Ø�
‰
×Ñ˜H g¨vÕ6rV   N)	r  r  r  r-  rû  r  r	  rm  rs  rU   rV   rW   r  r  ^  s    „ ò,òò  ò
!òó7rV   r  c                   ó”  — e Zd ZU dZeej                     ed<   dZeej                     ed<   ej                  j                  Zej                  ed<   ej                  j                  Zej                  ed<   ej                  j                  Zej                  ed<   dZeed<   d	Zeed
<   d	Zeed<   d	Zeed<   d	Zeed<   d„ Zd„ Zd„ Zy)rù  Nr¥   r:  r;  rÊ  r<  r   rû  FÚsimd_ompÚsimd_vecÚ	collapsedr…  c                 ój   — t        j                  «       }|r|j                  «       | _        y d| _        y rU  )r   rN  r-  Úsimd_nelements)r.  ra  s     rW   Ú__post_init__zLoopLevel.__post_init__›  s-   € ô .9×-EÑ-EÓ-GˆÙAO >×#;Ñ#;Ó#=ˆÕÐUVˆÕrV   c                 ó  — t        j                  |«      }t        | j                  | j                  «      }||_        d|_        t        |j                  |«      |z  |_        | j                  |_	        d|_
        | j                  |_        |S )NTF)rè   rN  rù  r¥   r:  r<  r+  r   rÊ  rû  r,  r…  )r.  rê  Úsympy_factorr  s       rW   rÉ  zLoopLevel.tile§  sn   € Ü—}‘} VÓ,ˆÜ˜Ÿ™ 4§9¡9Ó-ˆØ!ˆŒ
ØˆŒÜ" 4§9¡9¨lÓ;¸lÑJˆŒØŸ™ˆŒØˆŒØ ×-Ñ-ˆÔØˆrV   c                 ó   — t        | j                  «      }t        | j                  «      }t        j                  j
                  r||k(  ry | j                  r| j                  dkD  rd| j                  › d�nd}| j                  rFd}| j                  dkD  r|d| j                  › d�z  }| j                  r\|j                  dd|› �«      }nF| j                  rd}n7| j                  rd	|› �}n%| j                  st        j                  «       rd
}nd}t        › d| j                  › d|› �}| j                  › d|› �}| j                   j"                  r%| j                  › dt        | j                   «      › �}n;| j                  › dt        | j                   «      › dt        | j                   «      › d�}d|› d|› d|› d�}| j$                  s|s|gS ||gS )Nr5   zsimd simdlen(z) rR   z#pragma omp forz
 collapse(rš   z for z#pragma omp z#pragma GCC ivdepr•   r~  ú<r}  z+=(z == 0 ? 1 : zfor(rº   )rE   r;  r:  r   r¹  Úno_redundant_loopsr*  r.  rû  r   r+  r…  r   r  rI   r¥   r<  r‹  r,  )	r.  Úoffset_exprÚ	size_exprÚsimdÚline1Ú
offset_strrÇ  Ú	steps_strÚline2s	            rW   r  zLoopLevel.lines²  s½  € Ü! $§+¡+Ó.ˆÜ §	¡	Ó*ˆ	Ü�:‰:×(Ò(¨[¸IÒ-EØð �}Š} ×!4Ñ!4°qÒ!8ð ˜D×/Ñ/Ð0°Ñ3àð 	ð
 �=Š=à%ˆEØ�}‰}˜qÒ Ø˜: d§m¡m _°AÐ6Ñ6�Ø�}Š}ØŸ™ g°°t°f¨~Ó>‘Ø�]Š]Ø‰EØ�]Š]Ø" 4 &Ð)‰EØ×"Ò"¤{×'9Ñ'9Ô';Ø'‰EàˆEÜ"�| 1 T§X¡X J¨a°¨}Ð=ˆ
Ø—h‘h�Z˜q  Ð,ˆØ�:‰:×ÒØŸ8™8˜* B¤{°4·:±:Ó'>Ð&?Ð@‰Ið
 —8‘8�*˜C¤¨D¯J©JÓ 7Ð8ð 9Ü" 4§:¡:Ó.Ð/¨qð2ð ð �z�l " X J¨b°°¸1Ð=ˆØ�>Š>¡Ø�7ˆNØ�uˆ~ÐrV   )r  r  r  r¥   r   rè   r…  r  r:  ré   rê   r;  rÊ  ÚOner<  rû  rm   r*  rk   r+  r,  r…  r/  rÉ  r  rU   rV   rW   rù  rù  †  s®   … à $€Cˆ�%—*‘*Ñ	Ó$Ø!%€Dˆ(�5—:‘:Ñ
Ó%ØŸ™Ÿ™€FˆE�J‰JÓ%ð #ŸW™WŸ\™\€J�—
‘
Ó)ØŸ™Ÿ™€Eˆ5�:‰:Ó#Ø€HˆcÓØ€HˆdÓØ€HˆdÓØ€IˆtÓØ€L�$Óò
Wò	ó'rV   rù  c                   ó–   — e Zd ZU dZdZeee      ed<   dZ	ee
   ed<   ede
fd„«       Zd„ Zed„ «       Zd„ Zd	„ Zd
e
fd„Zd„ Zdefd„Zy)rF  aV  
    A loop-nest-like structure. It is built with the `build` method
    as a loop nest and then will perform loop-tiling at some depth.

    A typical case is for vectorization, where we typically do loop-tiling
    at the innermost loop level. A more complicated case is when we do
    2D tiling at both the innermost and outer levels.
    NrB  rZ  c                 ó$  — | j                   }| j                  }| j                  }|€J ‚d}t        t	        ||«      «      D ]B  \  }\  }}t        ||«      }|s|g}n|j                  |«       ||k\  sŒ2| j                  |_        ŒD t        |«      }	|	S )z4Build a LoopNest with the given `kernel` as the leafN)	rH  rM  r}  rÂ   r  rù  rŽ  r…  rF  )
rZ  rH  rM  r}  rB  Úloop_idxr¥   r:  r  rH  s
             rW   r'  zLoopNest.buildê  sŸ   € ð —?‘?ˆØ—‘ˆØ ×0Ñ0ˆØÐ*Ð*Ð*à+/ˆÜ%.¬s°8¸VÓ/DÓ%Eò 	8Ñ!ˆH‘k�s˜DÜ˜S $Ó'ˆDÙØ˜‘à—‘˜TÔ"Ø˜?Ó*Ø$*×$7Ñ$7�Õ!ð	8ô ˜U“Oˆ	ØÐrV   c                 ó,   — t        | j                  «      S r  )rk   rB  r3  s    rW   Ú__bool__zLoopNest.__bool__ÿ  s   € Ü�D—J‘JÓÐrV   c                 óê  — | j                   €t        dd¬«      S d}d}| j                   d   j                  }t        j                  d«      }| j                   D ]'  }|j                  |k7  r n||j
                  z  }|dz  }Œ) |t        | j                   «      k  rÎt        |t        j                  «      r´t        | j                   |   j
                  t        j                  «      rƒ|dz  | j                   |   j
                  k  rd|}d}| j                   |   j                  }t        |t        | j                   «      «      D ]%  }| j                   |   j                  |k7  r n|dz  }Œ' t        ||¬«      S )a‰  
        Maximal allowed depth for parallelism: All reduction or non-reduction levels.
        When the range of the first inner loop beyond the maximum parallel depth is much
        larger than the range of all outer loops within the maximum parallel depth,
        change the starting depth of parallelism to the first inner loop and recalculate
        the maximum parallel depth.
        r   r/  r5   rE  )	rB  r	  r…  rè   rN  r:  rÐ   r    rG  )r.  r  Ú	max_depthr…  Ú
loop_sizesr  rã   s          rW   r  zLoopNest.max_parallel_depth  sG  € ð �:‰:ÐÜ °¸qÔAÐAàˆØˆ	Ø—z‘z !‘}×1Ñ1ˆÜ—]‘] 1Ó%ˆ
Ø—J‘Jò 	ˆDØ× Ñ  LÒ0ÙØ# d§i¡iÑ/ˆJØ˜‰N‰Ið		ð œ˜DŸJ™J›Ò'Ü˜:¤u§}¡}Ô5Ü˜4Ÿ:™: iÑ0×5Ñ5´u·}±}ÔEØ˜SÑ  4§:¡:¨iÑ#8×#=Ñ#=Ò=à#ˆKØˆIØŸ:™: kÑ2×?Ñ?ˆLÜ˜;¬¨D¯J©J«Ó8ò �Ø—:‘:˜a‘=×-Ñ-°Ò=ÙØ˜Q‘‘	ðô ¨IÀ;ÔOÐOrV   c                 óÚ  — |j                   | j                  «       j                   k  sJ d«       ‚| j                  €J ‚t        | j                  «      |j                   k\  sJ ‚| j                  |j                     }|j                   |_        |j                  rt        xj                  dz  c_        t        |j                  dz   |j                   «      D ]  }d| j                  |   _
        Œ y )Nz?Parallel depth cannot exceed the maximal allowed parallel depthr5   T)r
  r  rB  rÐ   r  rû  r…  r   Úparallel_reduction_countrG  r,  )r.  rý  r  rã   s       rW   r  zLoopNest.mark_parallel)  sÏ   € Ø×'Ñ'¨4×+BÑ+BÓ+D×+SÑ+SÒSð 	
ØMó	
ÐSð �z‰zÐ%Ð%Ð%Ü�4—:‘:‹ )×":Ñ":Ò:Ð:Ð:Ø�z‰z˜)×/Ñ/Ñ0ˆØ!×0Ñ0ˆŒØ×ÒÜ×,Ò,°Ñ1Õ,Ü�y×,Ñ,¨qÑ0°)×2JÑ2JÓKò 	+ˆAØ&*ˆD�J‰J�q‰MÕ#ñ	+rV   c                 ó’   — | j                   sJ ‚| j                   |   j                  |«      | j                   |<   | j                   |   S )zÝ
        Do loop-tiling at the `depth` level with `factor`.
            for (x0 = 0; x0 < x0_end; x0++)
            ->
            for (x0 = 0; x0 < x0_end; x0 += factor)
        See details in Note [tiled_size].
        )rB  rÉ  )r.  r  rê  s      rW   rÉ  zLoopNest.tile6  sA   € ð �zŠzÐˆzØ ŸJ™J uÑ-×2Ñ2°6Ó:ˆ�
‰
�5ÑØ�z‰z˜%Ñ Ð rV   rµ   c                 ó6   — | j                   sJ ‚| j                   S r  ©rZ  r3  s    rW   rÿ  zLoopNest.get_kernelB  s   € Ø�{Š{Ðˆ{Ø�{‰{ÐrV   c                 ó   — || _         y r  rI  ró  s     rW   rÅ  zLoopNest.set_kernelF  s	   € Øˆ�rV   Úlevelc                 óÐ   — | j                   sJ ‚t        | j                   «      |k\  sJ ‚|t        | j                   «      k(  rd n| j                   |d  }t        || j                  «      S r  )rB  rÐ   rF  rZ  )r.  rK  rB  s      rW   rX  zLoopNest.from_loop_levelI  sV   € Ø�zŠzÐˆzÜ�4—:‘:‹ %Ò'Ð'Ð'Ø¤ T§Z¡Z£Ò0‘°d·j±jÀÀÐ6HˆÜ˜˜tŸ{™{Ó+Ð+rV   )r  r  r  r  rB  r   r#  rù  r  rZ  ru  r„  r'  rA  r%   r  r  rÉ  rÿ  rÅ  rm   rX  rU   rV   rW   rF  rF  Ü  s‡   … ñð (,€Eˆ8�D˜‘OÑ$Ó+Ø"&€FˆH�YÑÓ&àð�iò ó ðò( ð ñ$Pó ð$PòL+ò
!ð˜Ió òð, Sô ,rV   rF  rI  r  )§r½   ÚdataclassesrI  rJ  rY  rK  rÆ   r  r  Úcollections.abcr   Úenumr   Útypingr   r   r   r   r	   rè   r„   Útorch.fxÚtorch._inductorr
   Útorch._prims_commonr   r   Útorch.utils._ordered_setr   Útorch.utils._sympy.functionsr   r   r   Útorch.utils._sympy.symbolr   r   r   Ú_dynamo.utilsr   rR   r   r   r   r   r   r   r  r   r!  r   r   r    r!   r"   r#   r$   Úutilsr%   r&   r'   r(   r)   r*   r+   r,   r-   r.   r/   r0   Úvirtualizedr1   r2   r3   r4   Úcommonr6   r7   r8   r9   r:   r;   r<   r=   r>   r?   r@   rA   Ú	cpp_utilsrB   rC   rD   rE   rF   rG   rH   rI   rJ   rK   rL   rM   rN   rO   r  rT   Ú	lru_cacherX   Ú_loggingÚgetArtifactLoggerr  Úschedule_logÚNATIVE_OMP_RTYPESÚRTYPE_TO_CPPr›  ÚPYTHON_TO_CPPÚCONTAINER_PYTHON_TO_CPPr2  r3  rƒ   r1  rn   rk   r4  r5  r
  rŒ  rv   r#  rˆ   r  rw   r�   r‘   rô   r­   r¢   r…  ro   rÍ   rm   rß   rä   rî   r  r  Ú	dataclassr	  r  rc  ry  Ú_initialize_pointwise_overridesr‡  rm  rp  ru  r�  rr  r¡   r  r  r   rW  r  r  r   r  rù  rF  rU   rV   rW   ú<module>rf     s/  ðä Û Û Û Û Û Û 	Û 
Û Ý $Ý ß 7Õ 7ã ã Û Ý (ß @Ý /ß KÑ Kß OÑ Oå %ß G× GÝ  ÷÷ ñ ÷÷ ÷ ó ÷ >Ó =÷÷ ÷ ó ÷÷ ÷ ÷ ð$ �l‰l˜gÑ%€ð €×Ñ�TÓñ:ó ð:ð �~‰~×/Ñ/°¸*ÓE€áÒBÓCÐ àØØØØØØØØØ ñ€ñ !òóÐ ð  ØØØØØ#Ø&ØØØñ€ð ØñÐ ð 
‡N�NØ	‡M�Mð€ð 
‡M�MØ	‡K�KØ	‡N�NØ	‡M�MØ	‡J�JØ	‡K�KØ	‡J�JØ	‡K�KØ	‡K�Kð
*Ð �T˜%Ÿ+™+Ñ&ó 
ð 
‡K�KØ	‡N�NØ	‡M�MØ	‡K�KØ	‡J�Jð1Ð ˜D §¡Ñ-ó ò)òDð %)Øñ+)ð �E—L’LÑ!ó	+)ò\ð-Ø
ð-à�jŠjð-ð ð-ð —
’
ð	-ð
 �jŠjð-ð ó-ð`Ø�3˜Ð#Ñ$ðàðð ðð �;‰;ð	ð
 
ˆs�Cˆx‰óðBV˜^ð V°3ð VÀ#ó Vð ×Ñð-�U—Z’Zð - e§l¢lò -ó ð-ð ×Ñð; u§z¢zð ;¸¿
º
ð ;ÐPSò ;ó ð;ð| ×ÑàFJñ!Ø�:Š:ð!Ø!ŸLšLð!Ø6>¸s±mò!ó ð!ð ×Ò÷ð ó ðô]!Ð"4ô ]!÷@!ñ !ôBk�;ô kð\ × ,Ò ,¨UÔ 3ôr7�lô r7ðj × /Ò /°Ô 9Ø × %Ò %Ô 'ô7˜ô 7ôw�ô wôtk�9ô kô\Y
�lô Y
ðx)$ ð )$¨e°H¸U¿[¹[Ñ4IÈ4Ð4OÑ.Pó )$÷XbKñ bKôJT	
�Yô T	
ôn 
˜9ô  
ôFF�tô FôD&�Nô D&÷ND
ñ D
÷N%7ñ %7ðP ×Ò÷Rð Ró ðRðj ×Ò÷p,ð p,ó ñp,rV   