Ë
    cQêj;  ã                   ó°   — d dl Z d dlZd dlZddlmZmZmZmZmZm	Z	 ddl
mZ ddlmZ  G d„ de«      Z	  G d	„ d
e«      Z	  G d„ de«      Z G d„ de«      Zy)é    Né   )ÚFLOAT8_TYPESÚTENSOR_NAME_QUANT_SUFFIXÚQuantizedValueÚQuantizedValueTypeÚfind_by_nameÚget_mul_nodeé   )ÚQuantOperatorBase)ÚQDQOperatorBasec                   ó$   ‡ — e Zd Zˆ fd„Zd„ Zˆ xZS )Ú	QOpMatMulc                 ó&   •— t         ‰| �  ||«       y ©N©ÚsuperÚ__init__©ÚselfÚonnx_quantizerÚ	onnx_nodeÚ	__class__s      €ún/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/onnxruntime/quantization/operators/matmul.pyr   zQOpMatMul.__init__   ó   ø€ Ü‰Ñ˜¨Õ3ó    c                 óè  — | j                   j                  | j                  «      s.t        j                  d| j                  j
                  › d�«       y| j                   j                  | j                  j                  d   «      s`| j                   j                  | j                  j                  d   «      s.t        j                  d| j                  j
                  › d�«       y| j                   j                  rw| j                   j                  | j                  j                  d   «      sEt        j                  d| j                   j                  › d| j                  j
                  › d�«       yy	)
NzIgnore MatMul ú]Fr
   r   z&Ignore MatMul due to non float inputs z%Ignore MatMul due to non constant B: ú[T)Ú	quantizerÚshould_quantize_nodeÚnodeÚloggingÚdebugÚnameÚis_float_tensorÚinputÚinfoÚq_matmul_const_b_onlyÚfind_initializer_in_pathÚgraph_scope)r   s    r   Úshould_quantizezQOpMatMul.should_quantize   s  € Ø�~‰~×2Ñ2°4·9±9Ô=Ü�M‰M˜N¨4¯9©9¯>©>Ð*:¸!Ð<Ô=Øà—‘×.Ñ.¨t¯y©y¯©¸qÑ/AÔBØ—‘×.Ñ.¨t¯y©y¯©¸qÑ/AÔBä�L‰LÐAÀ$Ç)Á)Ç.Á.ÐAQÐQRÐSÔTØð �>‰>×/Ò/Ø—>‘>×:Ñ:¸4¿9¹9¿?¹?È1Ñ;MÔNÜ—‘ÐDÀTÇ^Á^×E_ÑE_ÐD`Ð`aÐbf×bkÑbk×bpÑbpÐaqÐqrÐsÔtØØr   )Ú__name__Ú
__module__Ú__qualname__r   r+   Ú__classcell__©r   s   @r   r   r      s   ø„ ô4ör   r   c                   ó$   ‡ — e Zd Zˆ fd„Zd„ Zˆ xZS )ÚMatMulIntegerc                 ó&   •— t         ‰| �  ||«       y r   r   r   s      €r   r   zMatMulInteger.__init__/   r   r   c                 óŒ  — | j                   }|j                  dk(  sJ ‚| j                  j                  |dg«      \  }}}}| j                  j	                  |dgdd¬«      \  }}}}	|j                  |«       |j                  |«       |j                  |«       |j                  |	«       |j                  d   dz   }
|j                  r|j                  dz   nd}t        j                  j                  d	||z   |
g|«      }|j                  |«       |
d
z   }| j                  j                  |j                  d   d¬«      }t        j                  j                  d|
g|g|
dz   |¬«      }|j                  |«       t        |«      dk(  sJ ‚|r|dz   n|d   dz   |d   z   dz   }t        || j                  j                  «      }|€!t!        ||dz   |«      }|j                  |«       |j                  d   }d}|r|dz   }|j                  t!        ||g|j                  d   |«      «       | j                  xj                  |z  c_        y )NÚMatMulr   r
   T©Úreduce_rangeÚop_level_per_channelÚ_output_quantizedÚ_quantÚ r2   Ú_cast_output)Ú	mandatoryÚCastÚ_cast)Útor   Ú_scales_mulÚ_Ú_mulz:0Ú_output_scale_mul)r!   Úop_typer   Úquantize_activationÚquantize_weightÚextendÚoutputr$   ÚonnxÚhelperÚ	make_nodeÚappendÚget_tensor_typeÚlenr   Ú	new_nodesr	   )r   r!   Úquantized_input_namesÚzero_point_namesÚscale_namesÚnodesÚquantized_input_names_weightÚzero_point_names_weightÚscale_names_weightÚnodes_weightÚmatmul_integer_outputÚmatmul_integer_nameÚmatmul_integer_nodeÚcast_op_outputÚotypeÚ	cast_nodeÚscales_mul_opÚscales_mul_nodeÚscales_mul_op_outputÚoutput_scale_mul_ops                       r   ÚquantizezMatMulInteger.quantize2   sl  € Ø�y‰yˆØ�|‰|˜xÒ'Ð'Ð'ð �N‰N×.Ñ.¨t°a°SÓ9ñ	
Ø!ØØØð �N‰N×*Ñ*¨4°!°À4Ð^bÐ*Ócñ	
Ø(Ø#ØØà×$Ñ$Ð%AÔBØ×ÑÐ 7Ô8Ø×ÑÐ-Ô.Ø�‰�\Ô"à $§¡¨A¡Ð1DÑ DÐØ6:·i²i˜dŸi™i¨(Ò2ÀRÐÜ"Ÿk™k×3Ñ3ØØ!Ð$4Ñ4Ø"Ð#Øó	
Ðð 	�‰Ð(Ô)ð /°Ñ?ˆØ—‘×.Ñ.¨t¯{©{¸1©~ÈÐ.ÓNˆÜ—K‘K×)Ñ)ØØ"Ð#ØÐØ! GÑ+Øð *ó 
ˆ	ð 	�‰�YÔô �;Ó 1Ò$Ð$Ð$ñ #ð   -Ò/à˜Q‘ #Ñ%¨°A©Ñ6¸Ñ?ð 	ô ' }°d·n±n×6NÑ6NÓOˆØÐ"Ü*¨;¸ÈÑ8LÈmÓ\ˆOØ�L‰L˜Ô)à.×5Ñ5°aÑ8Ðð !ÐÙØ"5Ð8KÑ"KÐØ�‰ÜØÐ!5Ð6Ø—‘˜A‘Ø#óô	
ð 	�‰× Ò  EÑ)Ö r   ©r,   r-   r.   r   rc   r/   r0   s   @r   r2   r2   .   s   ø„ ô4öG*r   r2   c                   ó(   ‡ — e Zd Zˆ fd„Zˆ fd„Zˆ xZS )ÚQLinearMatMulc                 ó&   •— t         ‰| �  ||«       y r   r   r   s      €r   r   zQLinearMatMul.__init__‚   r   r   c                 óþ  •— | j                   }|j                  dk(  sJ ‚| j                  j                  |dg«      \  }}}}| j                  j	                  |dgdd¬«      \  }}}}	|j                  |«       |j                  |«       |j                  |«       |j                  |	«       | j                  j                  |j                  d   «      \  }
}}}}|
r|€t        ‰| �%  «       S |j                  d   t        z   }|j                  r|j                  dz   nd}g }|j                  |d   «       |j                  |d   «       |j                  |d   «       |j                  |d   «       |j                  |d   «       |j                  |d   «       |j                  |«       |j                  |«       | j                  j                  t        v rdnd}t        j                   j#                  d	||g||¬
«      }|j                  |«       t%        |j                  d   |||t&        j(                  «      }|| j                  j*                  |j                  d   <   | j                  xj,                  |z  c_        y )Nr5   r   r
   Tr6   r:   r;   zcom.microsoftrf   )Údomain)r!   rE   r   rF   rG   rH   Ú_get_quantization_paramsrI   r   rc   r   r$   rM   Úweight_qTyper   rJ   rK   rL   r   r   ÚInputÚquantized_value_maprP   )r   r!   rQ   rR   rS   rT   rU   rV   rW   rX   Ú
data_foundÚoutput_scale_nameÚoutput_zp_namerB   Úqlinear_matmul_outputÚqlinear_matmul_nameÚqlinear_matmul_inputsri   Úqlinear_matmul_nodeÚq_outputr   s                       €r   rc   zQLinearMatMul.quantize…   sk  ø€ Ø�y‰yˆØ�|‰|˜xÒ'Ð'Ð'ð �N‰N×.Ñ.¨t°a°SÓ9ñ	
Ø!ØØØð �N‰N×*Ñ*¨4°!°À4Ð^bÐ*Ócñ	
Ø(Ø#ØØà×$Ñ$Ð%AÔBØ×ÑÐ 7Ô8Ø×ÑÐ-Ô.à�‰�\Ô"ð �N‰N×3Ñ3°D·K±KÀ±NÓCñ	
ØØØØØáÐ2Ð:Ü‘7Ñ#Ó%Ð%à $§¡¨A¡Ô1IÑ IÐØ6:·i²i˜dŸi™i¨(Ò2ÀRÐà "Ðà×$Ñ$Ð%:¸1Ñ%=Ô>Ø×$Ñ$ [°¡^Ô4Ø×$Ñ$Ð%5°aÑ%8Ô9à×$Ñ$Ð%:¸1Ñ%=Ô>Ø×$Ñ$ [°¡^Ô4Ø×$Ñ$Ð%5°aÑ%8Ô9à×$Ñ$Ð%6Ô7Ø×$Ñ$ ^Ô4à$(§N¡N×$?Ñ$?Ä<Ñ$O‘ÐUWˆÜ"Ÿk™k×3Ñ3ØØ!Ø"Ð#ØØð 4ó 
Ðð 	�‰Ð(Ô)ô "Ø�K‰K˜‰NØ!ØØÜ×$Ñ$ó
ˆð >Fˆ�‰×*Ñ*¨4¯;©;°q©>Ñ:à�‰× Ò  EÑ)Ö r   rd   r0   s   @r   rf   rf   �   s   ø„ ô4÷D*ð D*r   rf   c                   ó$   ‡ — e Zd Zˆ fd„Zd„ Zˆ xZS )Ú	QDQMatMulc                 ó&   •— t         ‰| �  ||«       y r   r   r   s      €r   r   zQDQMatMul.__init__Í   r   r   c                 ó0  — | j                   }|j                  dk(  sJ ‚| j                  r|j                  }n*t	        j
                  |j                  |j                  «      }|D ]±  }t        || j                  j                  j                  «       «      rf| j                  j                  |d|j                  ¬«      \  }}|r| j                  j                  ||«       Œ{| j                  j                  |«       Œ—| j                  j                  |«       Œ³ y )Nr5   r
   )Údefault_axisrE   )r!   rE   Údisable_qdq_for_node_outputr&   Ú	itertoolsÚchainrI   r   r   ÚmodelÚinitializerÚis_tensor_per_channelÚ"quantize_weight_tensor_per_channelÚquantize_weight_tensorÚquantize_activation_tensor)r   r!   Únodes_to_iterateÚtensor_nameÚis_per_channelÚchannel_axiss         r   rc   zQDQMatMul.quantizeÐ   sß   € Ø�y‰yˆØ�|‰|˜xÒ'Ð'Ð'à×+Ò+Ø#Ÿz™zÑä(Ÿ™¨t¯z©z¸4¿;¹;ÓGÐà+ò 
	GˆKÜ˜K¨¯©×)=Ñ)=×)IÑ)IÓ)KÔLØ/3¯~©~×/SÑ/SØ¨a¸¿¹ð 0Tó 0Ñ,� ñ "Ø—N‘N×EÑEÀkÐS_Õ`à—N‘N×9Ñ9¸+ÕFà—‘×9Ñ9¸+ÕFñ
	Gr   rd   r0   s   @r   rw   rw   Ì   s   ø„ ô4öGr   rw   )r|   r"   rJ   Úquant_utilsr   r   r   r   r   r	   Úbase_operatorr   Úqdq_base_operatorr   r   r2   rf   rw   © r   r   ú<module>rŒ      sc   ðÛ Û ã ÷÷ õ -Ý .ôÐ!ô ð.ô
K*�Iô K*ð\ô
H*�Iô H*ôVG�õ Gr   