§
    ŠŠtj×&  ã                   ó  — d dl Z d dlZd dlmZ d dlmZ d dlmZ d dlZd dlm	Z	 d dl
Z
d dlmZmZmZ ddlmZ dd	lmZ d
dlmZmZmZ d
dlmZmZmZ d
dlmZmZ d
dlmZ  e j         e!¦  «        Z" G d„ d¦  «        Z# G d„ dee¦  «        Z$dS )é    N)ÚCallable)Úpartial)ÚAny)ÚExpr)Úbound_sympyÚSymPyValueRangeAnalysisÚValueRangesé   ©ÚPowByNatural)Úint_ooé   )ÚInterpreterShimÚLoopBodyÚLoopBodyBlock)ÚDefaultHandlerÚReductionTypeÚ	StoreMode)Úcache_on_selfÚdominated_nodes)ÚVc                   óš  — e Zd ZdZdeddfd„Zdefd„Zede	e
j        j        ee         f         fd„¦   «         Zde	eed	ef         f         de	eed	ee         f         f         fd
„Zdede	e
j        j        ee         f         dedede	eed	ef         f         dee         fd„Zdedee         dee         fd„Zdedee         fd„ZdS )Ú	BoundVarsa¢  
    Performs Value Range Analysis on LoopBody's fx graph by calling BoundVars.run()
    It exposes the ranges of the nodes in the `bounds` variable

    Note. A current limitation of this analysis is that it just works on a per-loop basis.
    We should be able to propagate the bounds between across the whole graph. This may benefit
    the case a bounded variable is returned by a kernel and fed into another.
    Ú	loop_bodyÚreturnNc                 ó  ‡— dt           t          z  dt          fd„Š|| _        ˆfd„|j                             ¦   «         D ¦   «         | _        t          d„ | j                             ¦   «         D ¦   «         ¦  «        | _        i | _	        d S )NÚvr   c                 óX   — t          | t          ¦  «        rt          | ¦  «        j        n| S ©N)Ú
isinstancer   r   Úupper)r   s    úT/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/torch/_inductor/bounds.pyÚupper_boundz'BoundVars.__init__.<locals>.upper_bound'   s%   € Ý+5°a½Ñ+>Ô+>ÐE•;˜q‘>”>Ô'Ð'ÀAÐEó    c           	      óf   •— i | ]-\  }}|t          t                   d  ‰|¦  «        dz
  ¦  «        “Œ.S )r   r   )r	   r   )Ú.0Úkr   r#   s      €r"   ú
<dictcomp>z&BoundVars.__init__.<locals>.<dictcomp>+   sK   ø€ ð !
ð !
ð !
á��1ð �{�4Ô   K K°¡N¤N°QÑ$6Ñ7Ô7ð!
ð !
ð !
r$   c              3   óX   K  — | ]%}|j         d dt          j        fv s	d|j         v ¯!|V — Œ&dS )ÚloadÚ	reductionÚmasked_subblockN)ÚtargetÚoperatorÚgetitem©r&   Únodes     r"   ú	<genexpr>z%BoundVars.__init__.<locals>.<genexpr>0   sT   è è € ð .
ð .
àØŒ{˜v {µHÔ4DÐEÐEÐEØ  D¤KÐ/Ð/ð ð 0Ð/Ð/Ð/ð	.
ð .
r$   )
r   Úintr   Ú
var_rangesÚitemsÚreplacement_valsr   Ú	get_nodesÚunbounded_varsÚ_bounds)Úselfr   r#   s     @r"   Ú__init__zBoundVars.__init__&   s¿   ø€ ð	F�4¥#™:ð 	F­#ð 	Fð 	Fð 	Fð 	Fð #ˆŒð!
ð !
ð !
ð !
à!Ô,×2Ò2Ñ4Ô4ð!
ñ !
ô !
ˆÔõ
 .ð .
ð .
àœ×0Ò0Ñ2Ô2ð.
ñ .
ô .
ñ 
ô 
ˆÔð @BˆŒˆˆr$   c           
      ó`   — | j         j        › d| j        › d| j        › d| j        › d| j        › d�
S )Nz(loop_body=z,
 replacement_vals=z, 
unbounded_vars=z, 
_bounds=ú))Ú	__class__Ú__name__r   r6   r8   r9   )r:   s    r"   Ú__repr__zBoundVars.__repr__9   s`   € àŒ~Ô&ð 'ð 'Øœð'ð 'à $Ô 5ð'ð 'ð #Ô1ð'ð 'ð ”|ð	'ð 'ð 'ð	
r$   c                 ón  — |                       | j        j        ¦  «        }| j        D ]Z}t	          |j        t          ¦  «        rd|j        vr5d|j        vr,t          t                    	                    ¦   «         | j
        |<   Œ[t          j        t          ¦   «         ¦  «        5  t          | j        j        j        |¦  «        }t"                               d| j        j        j        ¦  «         |                     t          j        ¦   «         | j
        ¬¦  «         d d d ¦  «         n# 1 swxY w Y   | j
        S )Nr,   Úset_indirectzget_bounds:
%s©Úinitial_env)Úswap_submodulesr   Ú
submodulesr8   r    r-   Ústrr	   r   Úunknownr9   r   Úset_ops_handlerÚValueRangeAnalysisr   Ú
root_blockÚgraphÚlogÚdebugÚrunÚget_ops_handler)r:   rF   r1   Úinterpreters       r"   Ú
get_boundszBoundVars.get_boundsB   sZ  € à×)Ò)¨$¬.Ô*CÑDÔDˆ
ð Ô'ð 	Að 	AˆDå˜dœk­3Ñ/Ô/ð AØ!¨¬Ð4Ð4Ø"¨$¬+Ð5Ð5å%0µÔ%6×%>Ò%>Ñ%@Ô%@�”˜TÑ"øåÔÕ1Ñ3Ô3Ñ4Ô4ð 	Kð 	KÝ)¨$¬.Ô*CÔ*IÈ:ÑVÔVˆKÝ�IŠIÐ'¨¬Ô)BÔ)HÑIÔIÐIØ�OŠO�AÔ-Ñ/Ô/¸T¼\ˆOÑJÔJÐJð	Kð 	Kð 	Kñ 	Kô 	Kð 	Kð 	Kð 	Kð 	Kð 	Kð 	Køøøð 	Kð 	Kð 	Kð 	Kð Œ|Ðs   Â"A7D%Ä%D)Ä,D)rF   .c                 óÂ  ‡ ‡— i Š|D ]×}|dk    r‰ j         ‰|<   Œd|v rZ‰ j        j        |         }dt          dt          t
          t
          gt          t                   f         fˆˆ fd„} ||¦  «        ‰|<   Œqd|v rQt          |t          d¦  «        d …         ¦  «        }‰ j        j
        |         }t          ‰ j        |¦  «        }|‰|<   ŒÆd|v sJ ‚||         ‰|<   ŒØ‰S )NÚ	get_indexr,   Úsubblockr   c                 ó   •‡ — ˆˆˆ fd„S )Nc                 ó@   •— ‰                      ‰‰j        | |‰¦  «        S r   )r,   r9   )ÚmaskÚvalueÚresultr:   rU   s     €€€r"   ú<lambda>z<BoundVars.swap_submodules.<locals>.make_fn.<locals>.<lambda>h   s$   ø€ ¨t×/CÒ/CØ  $¤,°°e¸Vñ0ô 0€ r$   © )rU   rZ   r:   s   `€€r"   Úmake_fnz*BoundVars.swap_submodules.<locals>.make_fne   s'   øø€ ðð ð ð ð ð ð r$   rB   Úscan)rT   r   Ú	subblocksr   r   r   r	   r   r3   ÚlenÚindirect_varsr   rB   )	r:   rF   ÚkeyrU   r]   ÚidxÚvarÚindirectrZ   s	   `       @r"   rE   zBoundVars.swap_submodulesU   s"  øø€ ð ?AˆØð 	.ð 	.ˆCØ�kÒ!Ð!Ø"œn��s‘�Ø" cÐ)Ð)Øœ>Ô3°CÔ8�ðÝ+ðå�s¥C˜j­+µdÔ*;Ð;Ô<ðð ð ð ð ð ð ð &˜g hÑ/Ô/��s‘�Ø 3Ð&Ð&Ý˜#�c .Ñ1Ô1Ð3Ð3Ô4Ñ5Ô5�Ø”nÔ2°3Ô7�Ý" 4Ô#4°cÑ:Ô:�Ø&��s‘�à �}�}�}�}Ø(¨œo��s‘�àˆr$   rU   ÚenvrX   rY   c                 óø   — t          |j        |¦  «        }|                     t          j        ¦   «         |¬¦  «         d„ |j        j        D ¦   «         }t          |¦  «        dk    sJ ‚|j        |d                  S )NrC   c                 ó(   — g | ]}|j         d k    ¯|‘ŒS )Úoutput)r-   r0   s     r"   ú
<listcomp>z-BoundVars.masked_subblock.<locals>.<listcomp>‚   s$   € ÐSÐSÐS˜4¸4¼;È(Ò;RÐ;R�$Ð;RÐ;RÐ;Rr$   r   r   )r   rL   rO   r   rP   Únodesr`   rf   )r:   rU   rf   rX   rY   rF   Úinterpri   s           r"   r,   zBoundVars.masked_subblockx   sv   € õ ! ¤°Ñ<Ô<ˆØ�
Š
•1Ô$Ñ&Ô&°Cˆ
Ñ8Ô8Ð8ØSÐS 8¤>Ô#7ÐSÑSÔSˆÝ�6‰{Œ{˜aÒÐÐÐð Œz˜& œ)Ô$Ð$r$   ÚoldÚnewc                 óH   — t          |t          ¦  «        sJ ‚|| j        |<   |S r   )r    r	   r6   )r:   rm   rn   s      r"   rB   zBoundVars.set_indirectˆ   s*   € Ý˜#�{Ñ+Ô+Ð+Ð+Ð+Ø%(ˆÔ˜cÑ"Øˆ
r$   Únamec                 ó    — | j         j        |         }| j                             |¦  «        }|€t	          || j        ¦  «        }|| j        |<   |S r   )r   Úindexing_exprsr6   Úgetr   )r:   rp   ÚexprÚbounds       r"   rT   zBoundVars.get_index�   sR   € ØŒ~Ô,¨TÔ2ˆØÔ%×)Ò)¨$Ñ/Ô/ˆØˆ=Ý  dÔ&;Ñ<Ô<ˆEð ',ˆÔ˜dÑ#Øˆr$   )r?   Ú
__module__Ú__qualname__Ú__doc__r   r;   rG   r@   r   ÚdictÚtorchÚfxÚNoder	   r   rR   r   r   rE   r   r,   rB   rT   r\   r$   r"   r   r      s¿  € € € € € ðð ðB (ð B¨tð Bð Bð Bð Bð&
˜#ð 
ð 
ð 
ð 
ð ð˜D ¤¤°¸DÔ0AÐ!AÔBð ð ð ñ „]ðð$!Ø˜s H¨S°#¨XÔ$6Ð6Ô7ð!à	ˆc�8˜C ¨TÔ!2Ð2Ô3Ð3Ô	4ð!ð !ð !ð !ðF%àð%ð �%”(”- ¨TÔ!2Ð2Ô3ð%ð ð	%ð
 ð%ð ˜˜h s¨C xÔ0Ð0Ô1ð%ð 
�TÔ	ð%ð %ð %ð %ð  ð ¨;°tÔ+<ð ÀÈTÔARð ð ð ð ð
	˜cð 	 k°$Ô&7ð 	ð 	ð 	ð 	ð 	ð 	r$   r   c                   ó’  — e Zd Zdd„Zedededee         fd„¦   «         Zdede	edf         de
eef         defd	„Zded
ej        dee         fd„Z	 d ded
ej        dededdf
d„Zdej        dej        dededee         f
d„Zed
edej        dee         fd„¦   «         Zed
edej        dee         fd„¦   «         Ze	 	 d!dedej        dej        dz  dedee         f
d„¦   «         Zededee         fd„¦   «         Zededee         fd„¦   «         Zedededee         fd„¦   «         Zedededee         fd„¦   «         ZdS )"rJ   r   Nc                 óN   — d| _         d}|D ]}t          | || j        ¦  «         Œd S )NrJ   )ÚxorÚlogical_andÚ
logical_orÚlogical_not)rp   ÚsetattrÚbool_handler)r:   Úboolean_operatorsÚops      r"   r;   zValueRangeAnalysis.__init__š   sD   € Ø(ˆŒ	ð
Ðð $ð 	1ð 	1ˆBÝ�D˜"˜dÔ/Ñ0Ô0Ð0Ð0ð	1ð 	1r$   ÚargsÚkwargsc                  óJ   — t          t          j        t          j        ¦  «        S r   )r	   ÚsympyÚfalseÚtrue)r‡   rˆ   s     r"   r„   zValueRangeAnalysis.bool_handler¥   s   € õ �5œ;­¬
Ñ3Ô3Ð3r$   rp   .c                 ó(   — t          j        ¦   «         S r   ©r	   rH   )r:   rp   r‡   rˆ   s       r"   Ú_defaultzValueRangeAnalysis._defaultª   s   € õ Ô"Ñ$Ô$Ð$r$   Úindexc                 ó(   — t          j        ¦   «         S r   rŽ   )r:   rp   r�   s      r"   r*   zValueRangeAnalysis.load¯   s   € ÝÔ"Ñ$Ô$Ð$r$   rY   Úmodec                 ó   — d S r   r\   )r:   rp   r�   rY   r’   s        r"   ÚstorezValueRangeAnalysis.store²   s	   € ð 	ˆr$   ÚdtypeÚ	src_dtypeÚreduction_typec                 ó(   — t          j        ¦   «         S r   rŽ   )r:   r•   r–   r—   rY   s        r"   r+   zValueRangeAnalysis.reduction·   s   € õ Ô"Ñ$Ô$Ð$r$   c                 ó\   — t          |t          ¦  «        sJ ‚|                      ||¦  «        S r   )r    r	   Úto_dtype©Úclsr�   r•   s      r"   Ú
index_exprzValueRangeAnalysis.index_exprÀ   s,   € å˜%¥Ñ-Ô-Ð-Ð-Ð-Ø�|Š|˜E 5Ñ)Ô)Ð)r$   c                 ó.   — |                       ||¦  «        S r   )r�   r›   s      r"   Ú
value_exprzValueRangeAnalysis.value_exprÅ   s   € à�~Š~˜e UÑ+Ô+Ð+r$   TÚxÚuse_compute_typesc                 ó¼  — t          j        | ¦  «        } |t          j        k    r€|                      ¦   «         rt          j        | j        dk    ¦  «        S | j        r| S d| vrt          j        t          j        ¦  «        S t          t          j	        t          j        ¦  «        S dt          dt          j        dt          j        fd„}| j        ra|                      ¦   «         r)| j        rdnd}t          j         |||¦  «        ¦  «        S t           |d|¦  «         |d|¦  «        ¦  «        S t           || j        |¦  «         || j        |¦  «        ¦  «        S )Nr   r    r•   r   c                 ó¬   — |j         rt          j        | ¦  «        S | t          t           fv r| S 	 t          j        | ¦  «        S # t
          $ r | cY S w xY wr   )Úis_floating_pointrŠ   ÚFloatr   ÚIntegerÚ	TypeError)r    r•   s     r"   Úcastz)ValueRangeAnalysis.to_dtype.<locals>.castÝ   sk   € àÔ&ð 	Ý”{ 1‘~”~Ð%à�¥& Ð)Ð)Ð)Ø�HðÝ œ=¨Ñ+Ô+Ð+øÝ ð ð ð à�H�H�Hðøøøs   °A ÁAÁAr   )r	   Úwraprz   ÚboolÚis_singletonÚlowerÚis_boolrŠ   rŒ   r‹   r   r•   r   r!   )r    r•   r–   r¡   r¨   Úvals         r"   rš   zValueRangeAnalysis.to_dtypeÉ   sU  € õ Ô˜QÑÔˆà•E”JÒÐØ�~Š~ÑÔð <Ý"Ô'¨¬°1ªÑ5Ô5Ð5Ø”ð <Ø�Ø˜!��Ý"Ô'­¬
Ñ3Ô3Ð3å"¥5¤;µ´
Ñ;Ô;Ð;ð	•Cð 	¥¤ð 	µ´
ð 	ð 	ð 	ð 	ð Œ9ð 	KØ�~Š~ÑÔð CØœ7Ð)�a�a¨�Ý"Ô'¨¨¨S°%Ñ(8Ô(8Ñ9Ô9Ð9å" 4 4¨¨5¡>¤>°4°4¸¸5±>´>ÑBÔBÐBõ ˜t˜t A¤G¨UÑ3Ô3°T°T¸!¼'À5Ñ5IÔ5IÑJÔJÐJr$   c                 ó.   — t          j        | d„ ¦  «        S )Nc                 ó"   — t          | d¦  «        S )Nr
   r   )Úys    r"   r[   z+ValueRangeAnalysis.square.<locals>.<lambda>÷   s   € ½LÈÈAÑ<NÔ<N€ r$   )r	   Úconvex_min_zero_map©r    s    r"   ÚsquarezValueRangeAnalysis.squareô   s   € õ Ô.¨qÐ2NÐ2NÑOÔOÐOr$   c                 ó@   — t          j        | t          j        ¦  «        S r   )r	   Údecreasing_mapr.   Únegr³   s    r"   r·   zValueRangeAnalysis.negù   s   € õ Ô)¨!­X¬\Ñ:Ô:Ð:r$   ÚaÚbc                 óŠ   — |                       ||¦  «        }|t          j        ¦   «         k    r|S |                      |¦  «        S r   )Útruedivr	   rH   Útrunc)rœ   r¸   r¹   r    s       r"   ÚtruncdivzValueRangeAnalysis.truncdiv  s>   € à�KŠK˜˜1ÑÔˆØ•Ô#Ñ%Ô%Ò%Ð%ØˆHà�yŠy˜‰|Œ|Ðr$   c                 óT   — |                       ||                      |¦  «        ¦  «        S r   )Úaddr·   )rœ   r¸   r¹   s      r"   ÚsubzValueRangeAnalysis.sub	  s    € à�wŠw�q˜#Ÿ'š' !™*œ*Ñ%Ô%Ð%r$   )r   Nr   )NT)r?   rv   rw   r;   Ústaticmethodr   r	   r„   rG   Útuplery   r�   rŠ   r   r*   r   r”   rz   r•   r   r+   Úclassmethodr�   rŸ   rª   rš   r´   r·   r½   rÀ   r\   r$   r"   rJ   rJ   ™   s6  € € € € € ð	1ð 	1ð 	1ð 	1ð ð4˜Cð 4¨3ð 4°;¸sÔ3Cð 4ð 4ð 4ñ „\ð4ð%˜Sð %¨¨c°3¨h¬ð %ÀÀcÈ3ÀhÄð %ÐTWð %ð %ð %ð %ð
%˜ð % U¤Zð %°KÀÔ4Dð %ð %ð %ð %ð KOðð ØðØ %¤
ðØ36ðØ>Gðà	ðð ð ð ð
%àŒ{ð%ð ”;ð%ð &ð	%ð
 ð%ð 
�SÔ	ð%ð %ð %ð %ð ð*˜sð *¨5¬;ð *¸;ÀsÔ;Kð *ð *ð *ñ „[ð*ð ð,˜sð ,¨5¬;ð ,¸;ÀsÔ;Kð ,ð ,ð ,ñ „[ð,ð ð
 )-Ø"&ð	'Kð 'KØð'KàŒ{ð'Kð ”; Ñ%ð'Kð  ð	'Kð
 
�SÔ	ð'Kð 'Kð 'Kñ „\ð'KðR ðP�#ð P˜+ cÔ*ð Pð Pð Pñ „\ðPð ð;ˆsð ;�{ 3Ô'ð ;ð ;ð ;ñ „\ð;ð ð˜ð  ð ¨°SÔ)9ð ð ð ñ „[ðð ð&�Cð &˜Cð & K°Ô$4ð &ð &ð &ñ „[ð&ð &ð &r$   rJ   )%Úloggingr.   Úcollections.abcr   Ú	functoolsr   Útypingr   rŠ   r   rz   Útorch.utils._sympy.value_rangesr   r   r	   Úutils._sympy.functionsr   Úutils._sympy.numbersr   r   r   r   r   Úops_handlerr   r   r   Úutilsr   r   Úvirtualizedr   Ú	getLoggerr?   rM   r   rJ   r\   r$   r"   ú<module>rÏ      sµ  ðØ €€€Ø €€€Ø $Ð $Ð $Ð $Ð $Ð $Ø Ð Ð Ð Ð Ð Ø Ð Ð Ð Ð Ð à €€€Ø Ð Ð Ð Ð Ð à €€€ðð ð ð ð ð ð ð ð ð ð 2Ð 1Ð 1Ð 1Ð 1Ð 1Ø )Ð )Ð )Ð )Ð )Ð )Ø ?Ð ?Ð ?Ð ?Ð ?Ð ?Ð ?Ð ?Ð ?Ð ?Ø AÐ AÐ AÐ AÐ AÐ AÐ AÐ AÐ AÐ AØ 1Ð 1Ð 1Ð 1Ð 1Ð 1Ð 1Ð 1Ø Ð Ð Ð Ð Ð ð €gÔ˜Ñ!Ô!€ðzð zð zð zð zñ zô zð zðzr&ð r&ð r&ð r&ð r&Ð0°.ñ r&ô r&ð r&ð r&ð r&r$   