
    zjA                     0    d dl mZ ddlmZ d Zd Zd ZdS )   )_apply_pass   )corec                     g }|                                  j                                        D ]!\  }}|j        r|                    |           "|S N)global_blockvarsitemsis_dataappend)program	data_varsvar_namevars       c/lsinfo/ai/hellotax_ai/data_center/backend/venv/lib/python3.11/site-packages/paddle/framework/ir.pyget_data_varsr      sZ    I --//4::<< ' '#; 	'X&&&    c                    d}t           j                                        }d}d}g }t          | j                  D ]}|                     |          }|j        D ]p}||j        v rd}||j        vr|                    |          }	|	dd d         D ]8}
|	                    |
          }||
                    |           |j        rd}9q|r|r|D ]}d|_        
d S d S d S )Ngrad_merge_cond_nameFTr   r   )r   op_proto_and_checker_makerkOpRoleVarAttrNamerange
num_blocksblockops
attr_namesattr_find_var_recursiver   persistable)main_programgrad_merge_attr_nameop_role_var_attr_namehas_grad_mergehas_persistable_grad_var	grad_varsblock_idr   opp_ggg_vars               r   _update_grad_persistabler+      sM   1 ;NNPPN$I,122 4 4""8,,) 	4 	4B#r}44!%$BM99''/00CAY 4 411!44=  '''$ 4/3,4	4   %2 % 	% 	%E $E% % % %	% 	%r   c                     dd fd}t                                          dd          }|                                }|j        r |d           d|_        |j        r|r |d           d|_        |j        r |d           d|_        |j        r|r |d           d|_        |j        r|r |d	           d|_        |j        r|r |d
           d|_        |j	        r |d           d|_	        |j
        r |d           d|_
        |j        r |d           d|_        |j        r |g d           d|_        |j        r |d           d|_        |                                 |S )Nc                 *    || vr|| |<   |r|||<   d S d S r    )attrs
attr_typesnamevaluetyps        r   update_attrz)apply_build_strategy.<locals>.update_attr;   s8    uE$K 	#"Jt	# 	#r   c                     t                    }i } ||ddd            ||ddd            ||dt                    d           t          | ||           d S )	Nnranksr   size_tuse_cudaFboolmem_opt_skip_varsz	list[str])dictr   r   )r1   r/   r0   r    
pass_attrsstartup_programr4   s      r   
apply_passz(apply_build_strategy.<locals>.apply_passA   s    Z  
E:xH===E:z5&AAA,''	
 	
 	
 	L/4
KKKKKr   r8   Fsync_batch_norm_passfuse_relu_depthwise_conv_passfuse_resunit_passfuse_bn_act_passfuse_bn_add_act_passfusion_group_passfuse_gemm_epilogue_passfuse_dot_product_attention_passfuse_elewise_add_act_pass)coalesce_grad_tensor_passfuse_adam_op_passfuse_sgd_op_passfuse_momentum_op_passruntime_context_cache_passr   )r+   get_copysync_batch_normfuse_relu_depthwise_convfuse_resunitfuse_bn_act_opsfuse_bn_add_act_opsenable_auto_fusionfuse_gemm_epiloguefuse_dot_product_attentionfuse_elewise_add_act_opsfuse_all_optimizer_opscache_runtime_context_clear_finalized)r    r=   build_strategyr<   r>   r8   r4   s   `` `  @r   apply_build_strategyr\   8   sR   # # # #L L L L L L L L \***~~j%00H#))++N% /
)***).&. 88 8
233327/" ,
&'''&+#% /( /
%&&&).&) 3h 3
)***-2*( 2X 2
&''',1)( 2
,---,1)0 :
4555491. 8
.///27/, 	6
  	
 	
 	
 16-+ 5
/000/4,##%%%r   N)base.frameworkr    r   r   r+   r\   r.   r   r   <module>r_      sg    ) ( ( ( ( (        % % %:E E E E Er   