o
    + iūq  ć                   @   sh  d dl Z d dlZd dlZd dlZejj ”  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T ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddlT ddl T ddl!m"Z" ddl#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1 e+ Z2ee3ej4d	d
Z5G dd dZ6dS )é    N)Śpir)Śclear_shape_infoŚget_value_shape_range_info)Ś
get_loggeré   )Ś*)Śconverter_registry)ŚRefitManagerŚ	RefitRoleŚTensorRTConfigManagerŚTensorRTConstantManagerŚall_ops_into_trtŚget_cache_pathŚget_trt_versionŚget_trt_version_listŚis_shape_tensorŚ	map_dtypeŚremove_duplicate_valueŚset_dynamic_rangeŚweight_to_tensorŚzero_dims_to_one_dimsz&%(asctime)s-%(levelname)s: %(message)s)Śfmtc                   @   s6   e Zd ZdddZdd Zdd Zdd	 Zd
d ZdS )ŚPaddleToTensorRTConverterNc           
      C   s   || _ || _|| _t | _t | _| ”  ” }i }|D ]}| 	”  
” d }| j  |” ” }| j |||” qi | _i | _d| _t d”}	|	 d d” d S )NŚparameter_namer   zlibnvinfer_plugin.soŚ )ŚscopeŚprogramŚ
trt_configr   Śconstant_managerr	   Śrefit_managerŚglobal_blockZall_parametersŚget_defining_opŚattrsŚvarŚ
get_tensorŚset_constant_valueŚ
input_infoZtrt_output_value_mapŚ
engine_numŚctypesŚCDLLZinitLibNvInferPlugins)
ŚselfZpaddle_programr   r   ŚparamsZ
param_dictŚvŚnameZweight_tensorZtrt_plugin_lib© r.   śe/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/paddle/tensorrt/converter.pyŚ__init__I   s    
z"PaddleToTensorRTConverter.__init__c                    s¼   t t| ” j}i }i  g }dd }|D ]9}| ” D ]}| |j< |||j< ||r0| |” q| ” D ]}| ” }	|	 	” sHt
 d|	 ” q5|	||	j< q5q fdd| ” D }
|
|fS )Nc                 S   s$   |   ” D ]}| ” dkr dS qdS )Nścf.yieldTF)Śall_used_opsr-   )ŚvalueŚopr.   r.   r/   Z__is_output_valuee   s
   ’zNPaddleToTensorRTConverter.find_graph_inputs_outputs.<locals>.__is_output_valuezSkipping uninitialized source: c                    s   g | ]
\}}| vr|qS r.   r.   )Ś.0Zvalue_idr3   ©Śoutput_valuesr.   r/   Ś
<listcomp>{   s
    žzGPaddleToTensorRTConverter.find_graph_inputs_outputs.<locals>.<listcomp>)ŚnextŚiterŚblocksŚopsŚresultsŚidŚappendŚoperandsŚsourceŚinitializedŚ_loggerŚwarningŚitems)r*   Śgroup_opŚ
operationsZ
all_valuesZgraph_output_valuesZ+_PaddleToTensorRTConverter__is_output_valuer4   ŚresultŚoperandrA   Śinput_valuesr.   r6   r/   Śfind_graph_inputs_outputs^   s.   


ś
	žz3PaddleToTensorRTConverter.find_graph_inputs_outputsc           Z      C   sō  ddl m} t| j}| jd ur| jjrt d| ”  ” | jjs*t d| ” t	t
| ” j}|  |”\}}t t tjj””}| dttjj> ”}	| ” }
i }i }i }i }i }i }i }g }g }g }|j}| jd urs| jj}g }|D ]"}| ” }| ” dkr| ” D ]}| ” }| |” qqw| |” qwt|}|D ]Ž}| ” }| ” dkrų| ” d }| |” t  | j! "|””}| jj#rņ|j$}t %|”}|	 &||”}||_| 'd”||j(< | j) *| 'd”j|” | j) +||t,j-” q |||j(< q | ” d	krN| ” d
 } | j. /| ” 0” }!| j! 1| |!|” t  | j! "| ””}!| jj#rH|j$}t %|”}|	 &||!”}| |_| 'd”||j(< | j) *| 'd”j|!” q |!||j(< q |j$}"t2|j3j}#d|j( }$t4|"dkrgdg}"|	j5|$|#|"d}%| |$” | |” |%||j(< q |D ](}&|& ” dks|& ” dkrqg }'|& ” D ]}| ” }| 6” s¬|' d ” q| 7”  8” }(|(d urĀt4|( 9” dkrĀq| ”  ” })|)dkrg }*| ”  ” D ]+}+|+ ” },|,j(}-|-|v rśt:|	|,||- |& ” }.t;|	|.}.|* |.” qÕt<|- d|' |*” q|j(}/|/|v r't:|	|||/ |& ” }.t;|	|.}.|' |.” qt<|/ d|j=|j>j=kr;t?|&|' |  @|	|&|'”}0g }1tA|& B” D ]?\}2}3|3 C” rt4|3 7”  8”  9” dkrg|1 |3” qJ|3 D” }4|4D ]}5|5 ” dkr|5 B” }6|1 E|6” qmqJ|1 |3” qJtA|1D ]\}2}3|2t4|0k r¢|0|2 ||3j(< qd ||3j(< qqtA|D ]D\}7}||j( }8| ”  ” }9|9dksÉ|9d	krĖqÆ|8j}$| jjsąt d| d| ”  ” g }:g };g }<g }=g }>g }?| ” }@|@ ” dkr8|@ Fd” ”  ” dkr8| jG|j( d }:| jG|j( d };| jG|j( d }<|8jHr7| jG|j( d }=| jG|j( d }>| jG|j( d }?n@tI|dtJjKjLjMjN}:tI|dtJjKjLjMjO};tI|dtJjKjLjMjP}<|8jHrxtI|dtJjKjLjMjN}=tI|dtJjKjLjMjO}>tI|dtJjKjLjMjP}?|8jHs¬| jjs¢t d| d|: ” t d| d|; ” t d| d|< ” |
jQ|$|:|;|<d n/| jjsŅt d| d|= ” t d | d|> ” t d!| d|? ” |
jR|$|=|>|?d |:||$< |;||$< |<||$< |=||$< |>||$< |?||$< qÆg }Ag }Bg }CtSt4|D ]„}D||D }E||Ej( }F|Fd u r"|B d"” |A g ” |C d ” q|	 T|F” |B |Fj” |A |Ej$” |C |Ej3” | U|D” V” rCqg }:g };g }<t4|Ej$dkrotI|EdtJjKjLjMjN}:tI|EdtJjKjLjMjO};tI|EdtJjKjLjMjP}<g }=g }>g }?tH|ErtI|EdtJjKjLjMjN}=tI|EdtJjKjLjMjO}>tI|EdtJjKjLjMjP}?|:|;|<|=|>|?d#| jG|Ej(< q| W” }G| jr»| jj#r»|G XtjYjZ” |G [|
” t\d d$ksÕt\d d$krŚt\d d%krŚ| jj]|G_^|G _tj`ja| jjb” |j=|jcj=kr|jdrż|G XtjYjc” t d&” nnt ed'” nh|j=|jfj=krFt\d d(kr/|jgr)th|d)r)|G XtjYjf” t d*” nBt ed+” n<|jdr@|G XtjYjc” t ed,” n+t ed'” n%|j=|j>j=kr[|G XtjYj>” t d-” n| jd urkt d.| jj ” t\d d$kst\d d$krt\d d/krt\d/ dkr| jd ur| jjr|G XtjYji” | j|	|G”}H|Hd usØJ d0tJjKjk l” }I||I_m||I_n||I_o||I_p||I_q||I_r| jjs|I_stt| ju}J| jjsrŚ|JsŚt d1” d|I_s| jj#rķ| jj#|I_#||I_v| j) w” |I_xty|}Kttz {|K |d2”” }” d3d4 }Lt~| jj}M|M d5|L d6| j d7}Nt|Nd8}O|O |H” W d    n	1 s(w   Y  |M d5|L d6| j d9}Pt|Pd:}O|O |K” W d    n	1 sOw   Y  |N|I_tJj ”  tJjjL |”r t |” tJj ||I||B|A|Cd"”}QtSt4|QD ]R}D| U|D” V” rq|||D }R|Q|D }S| jG|Rj( d }T| jG|Rj( d }U| jG|Rj( d }V| jG|Rj( d }W| jG|Rj( d }X| jG|Rj( d }Y|T|U|V|W|X|Yd#| jG|Sj(< q|W d    n1 sŚw   Y  W d    |QS W d    |QS 1 sów   Y  |QS );Nr   )ŚPrecisionModezforce_fp32_ops: śstart process zbuiltin.combineśbuiltin.parameterr   r   śbuiltin.constantr3   Zinput_)r-   ŚdtypeŚshapezbuiltin.splitz! not found in value_to_trt_tensorzset shape of z	, op is: zpd_op.tensorrt_engineŚ	min_shapeŚ	opt_shapeŚ	max_shapeŚ	min_valueŚ	opt_valueŚ	max_valueFTzset min_shape of z as zset opt_shape of zset max_shape of )ŚminŚoptŚmaxzset min_value of shape input: zset opt_value of shape input: zset max_value of shape input: r   )rR   rS   rT   rU   rV   rW   é   é   zRun Paddle-TRT FP16 modez8Hardware does not support FP16. Continuing in FP32 mode.é	   Zplateform_has_fast_bf16zRun Paddle-TRT BF16 modez8Hardware does not support BF16. Continuing in FP32 mode.zKBecause the version of TensorRT is less than 9.0, run  Paddle-TRT FP16 modezRun Paddle-TRT INT8 modezDefault precision mode é   z<Failed to build engine. please see ERROR log from trt.Loggerz_You have enabled CudaGraph, but not the entire graph offload to trt, now return to normal mode.zutf-8é   i įõz/engine_Ś_z.trtŚwbz.pirŚw)ZexportrL   r   r   Zops_run_floatrC   ŚinfoZget_force_fp32_opsŚdisable_logglingr9   r:   r;   r<   rK   ŚtrtZBuilderŚLoggerŚERRORZcreate_networkŚintZNetworkDefinitionCreationFlagZEXPLICIT_BATCHZcreate_optimization_profileZFP32Śprecision_moder!   r-   r@   rA   r?   r   r"   ZWeightsr   Zget_constant_valueZrefit_params_pathrQ   ZDimsZadd_constantZ
get_outputr>   r   Zset_trt_weight_tensorZset_mappingr
   ZCONSTANTr   r#   r$   r%   r   rP   ŚlenZ	add_inputrB   ŚtypeZas_vec_typeŚas_listr   r   ŚRuntimeErrorr3   ZINT8r   ŚconvertŚ	enumerater=   Z
is_combiner2   ŚextendZoperand_sourcer&   r   r   ŚpaddleŚbaseŚcoreZ	ShapeModeZkMINZkOPTZkMAXŚ	set_shapeZset_shape_inputŚrangeZmark_outputrH   Ś	use_emptyZcreate_builder_configZset_flagZBuilderFlagZREFITZadd_optimization_profileŚversion_listZoptimization_levelZbuilder_optimization_levelZset_memory_pool_limitZMemoryPoolTypeŚ	WORKSPACEZworkspace_sizeZFP16Zplatform_has_fast_fp16rD   ZBF16Zplatform_has_fast_bfp16ŚhasattrZPREFER_PRECISION_CONSTRAINTSZbuild_serialized_networkŚ	libpaddleZTRTEngineParamsZmin_input_shapeZmax_input_shapeZoptim_input_shapeZmin_shape_tensorZmax_shape_tensorZoptim_shape_tensorZuse_cuda_graphr   r   Śrefit_param_nameZget_all_mappingsZrefit_param_names2trt_namesŚstrŚhashlibŚsha256ŚencodeŚ	hexdigestr   Zsave_model_dirr'   ŚopenŚwriteZengine_serialized_dataZ	pir_utilsZIrGuardr   Śprogram_guardZset_insertion_pointŚ_C_opsZtensorrt_engine)Zr*   r   rF   rL   Ztrt_managerrG   rJ   r7   ZbuilderŚnetworkZprofileZvalue_to_trt_tensorZmin_shape_mapZopt_shape_mapZmax_shape_mapZmin_value_mapZopt_value_mapZmax_value_mapZinput_namesZnew_input_valuesr{   ri   Zorigin_input_valuer3   Zdefining_oprI   rA   Ś
param_nameŚweightZpaddle_shapeZ	trt_shapeZconstant_layerŚconstant_value_nameZconstant_tensorrQ   rP   Z
input_nameZinput_tensorr4   r@   Zvec_typeZdefine_op_nameZoperand_listZcombined_operandZcombined_sourceZcombined_source_idZtrt_input_tensorZ	source_idZtrt_outsr=   ŚidxrH   Zused_opsZuse_opZsplit_outputsŚiZ	trt_inputZdefining_op_namerR   rS   rT   rU   rV   rW   Zvalue_define_opZ
out_shapesZ	out_namesZ	out_typesZ	out_indexZresult_valueZoutput_tensorŚconfigZ
trt_engineZ
trt_paramsZall_nodes_offload_to_trtZ	group_strZengine_nameZ
CACHE_ROOTZ
CACHE_FILEŚfZPIR_DUMP_FILEŚoutZ	ori_valueŚcurrent_valueZorin_min_shapeZorin_opt_shapeZorin_max_shapeZorin_min_valueZorin_opt_valueZorin_max_valuer.   r.   r/   Śconvert_subgraph_to_trt   s>  
’
ž

’
’’’ž’’’

’’’






ü’’
ü’’



żü



’’’’’’’’
’
’’’’





’’’
’’’ś	
’’’
’’’’’’

’

’’’’’’
ł
śō*ō!ß!ß!z1PaddleToTensorRTConverter.convert_subgraph_to_trtc                 C   s`   t  }| ” }|dv rd S t ||”}|d u rtd| d||||}t|tjr.|fS |S )N)r1   zConverter for z not implemented.)r   r-   r   ŚgetŚNotImplementedErrorŚ
isinstancere   ZITensor)r*   r   Z	paddle_opZinputsZtrt_versionZop_nameZconverter_funcZoutsr.   r.   r/   rn   U  s   
’z!PaddleToTensorRTConverter.convertc              
   C   s>  | j  ” jD ]H}| ” dks| ” dkrN| jjs"t d| ”  ” |  jd7  _|  	| j |”}| 
” }tt|D ]}||  || ” q:| j  ”  |” qd}| j  ” jD ]Ā}| ” dkr~| ” d }|snd|vrnd	}qW| 
” d
  ” r~| j  ”  |” | ” dkr| 
” d
  ” r| j  ”  |” qW| 
” d
 }| ” d }t tjjj|j ”}	| j |” ” }
tj|
|	d ” }t|ttfrČ|g}tjj | j ”2 tj j!j "”  t|dkrļtj# $dg|d
 |jt %d
””}ntj# &||jt %d
””}W d    n	1 sw   Y  | |g” | j  ”  |” qWt'  d S )Nzcinn_op.groupzbuiltin.grouprM   r   FrN   r   Zconstant_foldingTr   rO   r3   )rP   )(r   r    r<   r-   r   rd   rC   rc   r'   r   r=   ru   rj   Zreplace_all_uses_withZ	remove_opr"   rv   ŚnprP   rq   r   rs   Z_PADDLE_PIR_DTYPE_2_NUMPY_DTYPEr   r#   r$   ŚarrayŚtolistr   rh   Śfloatr   rr   rz   Zreset_insertion_point_to_startr   ŚfullZ	CUDAPlaceZfull_int_arrayr   )r*   r4   Znew_outZorin_out_valuesZo_iZsave_one_parameterr   Zconstant_resultr   Z	out_dtypeZtensor_dataZconstant_arrayZ
full_valuer.   r.   r/   Śconvert_program_to_trtf  s   ’’’’’ž’žüżö
z0PaddleToTensorRTConverter.convert_program_to_trt)N)Ś__name__Ś
__module__Ś__qualname__r0   rK   r   rn   r   r.   r.   r.   r/   r   H   s    
%   Ur   )7r(   r}   Śloggingrq   rr   rs   Zregister_paddle_pluginZtensorrtre   r   Zpaddle.base.corer   r   Zpaddle.base.log_helperr   Zimpls.activationZimpls.attributeZimpls.commonZ
impls.convZimpls.creationZimpls.einsumZimpls.inputZimpls.linalgZimpls.logicZimpls.manipulationZ
impls.mathZ
impls.normZ	impls.opsZimpls.othersZimpls.poolingZimpls.searchZ
impls.statZimpls.visionŚregisterr   Śutilr	   r
   r   r   r   r   r   r   r   r   r   r   r   r   rw   r   ŚINFOrC   r   r.   r.   r.   r/   Ś<module>   sF   @’