o
    hTG                     @   s*  d dl Z d dlZd dlmZ d dlm  mZ ejjZ	d dl
m  m  mZ d dlm  m  m  mZ d dlm  m  m  mZ d dlm  m  m  m  mZ d dlm  m  m  mZ d dlm  m  mZ d dlm  m  mZ d dlm  m  m  mZ d dlmZ d dl m  m!  m"  m#Z# d dl$m  m!  m%Z% ddl&m'Z' d dl(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z. de*e/e-e' f fddZ0d	e*e/e-e' f d
e'de,e/ fddZ1d	e*e/e-e' f d
e'de,e' ddfddZ2de*e/e-e' f fddZ3de*e/e-e' f fddZ4dS )    N)get_native_backend_config   )NSNodeTargetType)CallableDictListOptionalSetTuplereturnc                  C   sF  g t jht jht jhtjhtjhtjht jhtj	ht j
tjht jtjjjht jtjjjht jtjht jtjht jtjht jhtjtjhtjhtjtjhtjt jddtjht j tj!ht j"tj#ht j$tj%htj&ddt j'tj&ht j(ht j)ht j*ht j+ht j,htj-htj.htj/ht j0ht j1ht j2ht j3ht j4ht j5ht j6ht j7ht j8ht j9ht j:tj;htj<htj=htj>htj?htj@htjAht jBtjCht jDtjEht jFtjGtjGddhddtjHt jIht jJtjKtjLhtjMhtjNhtjOhtjPhtjQhtjRhtjShtjThtjUhtjVhtjWhtjXhtjYhtj jZj[ht j\htj]htj^ht j_ht j`taj`htjbtcjbht jdhtjeht jfhtjghtjhh} ti }t jt jjj	jkfg}|jlm D ]E\}}|}tn|totpfr|d	 }tn|totpfs|jqd ur|r||jqf |jsd ur|r||jsf |jtd ur|r||jtf qtujvtujwtujxtujyfD ]}|m D ]\}}|r||f qqtujztuj{tuj|fD ]}	|	m D ]\}\}
}|r||
f |r||f qqtuj}m D ]\}\}
}|r||
f |r||f qtuj~tujtjfD ]}|m D ]\}}|r||f q?q9tjfD ]}|m D ]\}}|r||f qXqR|D ]!\}}| D ]}||v s{||v r|| ||  nqoqii }d
}| D ]}t|}|d7 }|||< q|S )Nrelurelu_sigmoidsigmoid_tanh_tanhhardsigmoid_hardsigmoidr   r   )nnConv1dConv2dConv3dFconv1dconv2dconv3dLinearlinear	AvgPool1dtorch
avg_pool1d	AvgPool2d_C_nn
avg_pool2d	AvgPool3d
avg_pool3dAdaptiveAvgPool1dadaptive_avg_pool1dAdaptiveAvgPool2dadaptive_avg_pool2dAdaptiveAvgPool3dadaptive_avg_pool3dLSTMaddoperatorcatmulr   ReLU	MaxPool1d
max_pool1d	MaxPool2d
max_pool2d	MaxPool3d
max_pool3dr   SigmoidBatchNorm2dBatchNorm3dConvTranspose1dConvTranspose2dConvTranspose3dconv_transpose1dconv_transpose2dconv_transpose3dELU	EmbeddingEmbeddingBag	GroupNorm	HardswishInstanceNorm1dInstanceNorm2dInstanceNorm3d	LayerNorm	LeakyReLUReLU6relu6elu	hardswish
group_norminstance_norm
layer_norm
leaky_reluSiLUsiluMishmishTanhr   r   HardsigmoidHardtanhhardtanh	hardtanh_floordiv	unsqueezestacksqueezesortrepeat_interleaveminmeanmax	transposeflattenclampchunk
functionalinterpolateDropoutdropoutmatmulSoftmaxPReLUnnqprelutoqPixelShufflepixel_shufflePixelUnshufflepixel_unshufflenarrowr   modulesNonDynamicallyQuantizableLinear!_pattern_complex_format_to_configitems
isinstancelisttuplefused_moduleappend
qat_modulereference_quantized_module_lower_to_native_backendSTATIC_LOWER_MODULE_MAPDYNAMIC_LOWER_MODULE_MAPWEIGHT_ONLY_LOWER_MODULE_MAP SPECIAL_PATTERN_LOWER_MODULE_MAPSTATIC_LOWER_FUSED_MODULE_MAP(STATIC_LOWER_FUSED_MODULE_TWO_INPUTS_MAPDYNAMIC_LOWER_FUSED_MODULE_MAPSTATIC_LOWER_FUNCTIONAL_MAPQBIN_OP_MAPPINGQBIN_RELU_OP_MAPPINGquantization_mappings,DEFAULT_FLOAT_TO_QUANTIZED_OPERATOR_MAPPINGS%DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGSstr)sets_of_related_opsbackend_confignew_connectionspatternconfigfirst_elementsource_to_targetsourcetargetsource_to_double_targettarget1target2item1item2set_of_related_ops base_name_to_sets_of_related_opscounter	base_name r   M/var/www/html/ai/venv/lib/python3.10/site-packages/torch/ao/ns/fx/mappings.py$get_base_name_to_sets_of_related_ops   s  	#',049=BFKSW[`hkoruy|                 "  &  *  /  3  7  ;  ?  C  G  L  Q  Y  `  f  j  n  r  v  z  ~                                   #    '    +    /    4    9    <    @    C    G    M	

r   r   opc                 C   s&   |   D ]\}}||v r|  S qd S )N)r}   )r   r   r   r   r   r   r   get_base_name_for_op  s
   r   
related_opc                 C   sr   |d ur|   D ]}||v r||  d S qt| dd}t|| v r0|d7 }t|| v s&|h| t|< d S )Nz was not foundr   r   )valuesr/   AssertionErrorr   )r   r   r   r   r   r   r   r   add_op_to_sets_of_related_ops  s   
r   c                  C   s(  t jt jt jt jtjt jt jt j	t j
t jt jt jt jtjtjtjtjtjt jh} t }tjtjtjtjtjtjtjtjtjtjtjtj	tj
tjtjtjh}h t jt jtjt jtjt jtjtjt j t j!t jt j"t j#t j$t j%t j&t j't j(t j)t j*tj+tj,j-j.tj,j-j/tjtj0tj1tj2tj3tj4tj5tj6tj7tj8tj9tj:tj;tj<tj}h t=j>t?j>t@j>tAj>tj=jBjjCt=jDt=jEt=jFt?jDt?jEt?jFt?jGt?jHt=jItAjIt=jJt=jKt=jLt=jMt=jNt=jOt=jPt=jQt=jRt=jSt=jTt=jUt=jVt=jWt=jXt=jYt=jZt=j[t=j\t]j^t]j_t]j`t]jat]jbt]jct]jdt]jet]jft]jgthjethjfthjgthjithjjthjkthj`thjathjbthjcthjdtljct]jmt]jnt]jot]jp}h tqj>tqjDtqjEtqjFtqjJtqjKtqjLtqjMtqjNtqjPtqjRtqjStqjTtqjUtqjVtqjWtqjGtqjHtqjLtqj[tqj\trj^trj_trj`trjatrjbtrjctrjmtrjntrjotrjp}t=jst=jtt=jut=jvt=jwt=jxt=jyt=jzt=j{t=j|t=jLt=j}t=j~t=jt=jt=jt=jt=jt=jXh}h d}| |||||||dS )N>   r   r   r   r   r   r   r   r   )funs_io_type_fp32funs_io_type_fp16funs_io_type_int8funs_io_type_fp32_or_int8mods_io_type_fp32mods_io_type_int8mods_io_type_fp32_or_int8meths_io_type_fp32_or_int8)r   r   r   r   r   r    r1   rO   rP   rR   rS   rT   rn   rV   rX   r0   r/   r2   sumrs   setrt   linear_reluconv1d_reluconv2d_reluconv3d_relur   r   r   r   r^   r)   r+   r-   r\   r]   rl   r5   r7   r9   rN   rv   rx   r!   r#   r$   r%   r'   rj   ri   rh   rg   rf   re   rd   ry   rc   rb   ra   r`   r_   r   r   nnqatnnqatdnnqdrz   r{   r   r   r   rD   rE   r.   r;   r<   rm   r=   r>   r?   rC   rF   rH   rI   rJ   rK   rG   rL   rM   rU   rW   rp   rq   nniBNReLU2dBNReLU3d
ConvReLU1d
ConvReLU2d
ConvReLU3d
LinearReLU
LinearBn1dConvBn1dConvBn2dConvBn3dnniqatConvBnReLU1dConvBnReLU2dConvBnReLU3dnniqdLinearLeakyReLU
LinearTanh	ConvAdd2dConvAddReLU2drr   nniqr3   rY   r:   rZ   r(   r*   r,   r   r"   r&   r[   Identityr4   r6   r8   ru   rw   )FUNS_IO_TYPE_FP32FUNS_IO_TYPE_FP16FUNS_IO_TYPE_INT8FUNS_IO_TYPE_FP32_OR_INT8MODS_IO_TYPE_FP32MODS_IO_TYPE_INT8MODS_IO_TYPE_FP32_OR_INT8METHS_IO_TYPE_FP32_OR_INT8r   r   r   get_node_type_to_io_type_map  s  	
 !"#$%&)
	
 !"#$%&'()*+,-./0123456789:;<=>A	
#r   c                  C   s(   t jtjh} tjh}h d}| ||dS )N>   tore   sizeviewrj   ri   numelshapedetachrepeatdetach_permutereshaperesize_ra   squeeze_rg   r_   
contiguous
dequantize
unsqueeze_rc   )funs_unmatchablemods_unmatchablemeths_unmatchable)r    quantize_per_tensorr0   getitemr   r   )FUNS_UNMATCHABLEMODS_UNMATCHABLEMETHS_UNMATCHABLEr   r   r   get_unmatchable_types_map  s   r   )5r0   r    torch.nnr   torch.nn.functionalrk   r   ops	quantizedrt   torch.ao.nn.quantizedaorr   torch.ao.nn.quantized.dynamicdynamicr   torch.ao.nn.intrinsic.quantized	intrinsicr   'torch.ao.nn.intrinsic.quantized.dynamicr   torch.ao.nn.intrinsic.qatqatr   torch.ao.nn.intrinsicr   torch.ao.nn.qatr   torch.ao.nn.qat.dynamicr   $torch.ao.quantization.backend_configr   1torch.ao.quantization.fx._lower_to_native_backendquantizationfxr   +torch.ao.quantization.quantization_mappingsr   ns_typesr   typingr   r   r   r   r	   r
   r   r   r   r   r   r   r   r   r   r   <module>   sR    $    3


 j