o
    MeA                     @   s  d dl Z d dlZd dlZd dlZd dlZddlmZ ddlmZ ddlm	Z	m
Z
 ddlmZmZmZ ddlmZ g d	ZejjZejjZejZejZejZd
d Zdd Zdd Zdd Zdd Zdd Zdd ZG dd deZ G dd deZ!G dd deZ"G dd deZ#dS )     N   )compat   )	framework)_get_paddle_place_get_paddle_place_list)cuda_places
cpu_places
xpu_places)core)CompiledProgramExecutionStrategyBuildStrategyIpuCompiledProgramIpuStrategyc                 C   s   t  }||  |S N)r   ZPlaceZ	set_place)placep r   ED:\Projects\ConvertPro\env\Lib\site-packages\paddle/fluid/compiler.py
_place_obj&   s   
r   c                 C   s4   | r| nt  }| jD ]
}|jdv r dS qdS )N)sendrecvTF)r   default_main_programglobal_blockopstype)main_programmainopr   r   r   _is_pserver_mode,   s   
r    c                 C   s>   |   D ]}| r| d ur|  dr dS qdS )NZ_gradTF)nodesZis_opr   r   endswith)graphnoder   r   r   _has_backward_op5   s   r%   c                 C   sV   g }t |  jD ]\}}|jdkr|| q	|d d d D ]	}|  | qd S )Nfeed)	enumerater   r   r   append
_remove_op)programZpop_idxir   indexr   r   r   _prune_feed_ops=   s   r.   c                 C   sP   | j D ]"}tj}tjjj}| |jv r%t| |	  t|kr% dS qdS )NTF)
r   r   Zop_proto_and_checker_makerZOpRoleZOptimizeZkOpRoleVarAttrNameZ
attr_namesintZ	all_attrsZkOpRoleAttrName)blockr   Zop_makeroptimizer   r   r   _has_optimize_opF   s   

r2   c                 C   sH   | st  } |  jD ]}|jdkr!| |d}t|r! dS qdS )NZconditional_block_grad	sub_blockTF)r   r   r   r   r   r0   Z_block_attr_idr2   )r+   r   r3   r   r   r   _has_optimizer_in_control_flowP   s   
r4   c                 C   s@   |   }|j|d }|d u rdS t|ddpt|dd}| S )NTZ_is_distributedFis_distributed)r   varsgetgetattr)r+   var_namer0   varr5   r   r   r   _should_broadcast_or_not_exists\   s   r;   c                   @   s`   e Zd ZdZdddZ					dddZdd Zd	d
 ZdddZdd Z	dd Z
dd ZdS )r   aZ  
    :api_attr: Static Graph
    
    The CompiledProgram is used to transform a program or graph for
    various optimizations according to the configuration of build_strategy,
    for example, the operators' fusion in the computation graph, memory
    optimization during the execution of the computation graph, etc.
    For more information about build_strategy, please refer to
    :code:`paddle.static.BuildStrategy`.

    Args:
        program_or_graph (Graph|Program): This argument is the Program or Graph
            being executed.
        build_strategy(BuildStrategy): This argument is used to compile the
            program or graph with the specified options, such as operators' fusion
            in the computational graph and memory optimization during the execution
            of the computational graph. For more information about build_strategy,
            please refer to :code:`paddle.static.BuildStrategy`. The default is None.

    Returns:
        CompiledProgram

    Example:
        .. code-block:: python

            import numpy
            import paddle
            import paddle.static as static

            paddle.enable_static()

            place = paddle.CUDAPlace(0) # paddle.CPUPlace()
            exe = static.Executor(place)

            data = static.data(name='X', shape=[None, 1], dtype='float32')
            hidden = static.nn.fc(x=data, size=10)
            loss = paddle.mean(hidden)
            paddle.optimizer.SGD(learning_rate=0.01).minimize(loss)

            exe.run(static.default_startup_program())
            compiled_prog = static.CompiledProgram(
                static.default_main_program())

            x = numpy.random.random(size=(10, 1)).astype('float32')
            loss_data, = exe.run(compiled_prog,
                                feed={"X": x},
                                fetch_list=[loss.name])
    Nc                 C   s   t |tjr|| _d | _nt |tjr"t| t|j| _|| _nt	dt
| d | _d | _d | _d| _d| _d| _d | _d | _d | _|| _d | _d S )Nz[The type of program_to_graph parameter is wrong, expected Graph or Program, but received %sF)
isinstancer   Graph_graph_programr   Programr.   desc	TypeErrorr   _scope_place	_executor	_compiled_is_data_parallel_is_inference
_loss_name_share_vars_from_places_build_strategy_exec_strategy)selfZprogram_or_graphbuild_strategyr   r   r   __init__   s.   
zCompiledProgram.__init__c                 C   s   | j rJ d| jrJ dd| _ |dur|| _|| _|| _|| _t|ttfr.t	|| _
nt|| _
t| jrA| jdusAJ d| j
durSt| j
ttfsS| j
g| _
| S )a  
        This interface is used to transform the input Program or Graph to a multi-graph
        to run the model in data parallel mode. Users can use the build_strategy and
        exec_strategy to set some optimizations that can be applied during the construction
        and computation of the Graph, such as reducing the number of AllReduce operations,
        specifying the size of the thread pool used in the computation Graph running the model,
        and so on. 
        
        .. note::
            If build_strategy is specified when building CompiledProgram and calling 
            with_data_parallel, build_strategy in CompiledProgram will be overwritten, therefore, 
            if it is data parallel training, it is recommended to set build_strategy when calling 
            with_data_parallel interface.

        Args:
            loss_name (str): This parameter is the name of the loss Tensor of the model.
                **Note: If it is model training, you must set loss_name, otherwise the
                result may be problematic**. The default is None.
            build_strategy(BuildStrategy): This parameter is used to compile the
                program or graph with the specified options, such as operators' fusion
                in the computational graph and memory optimization during the execution
                of the computational graph. For more information about build_strategy,
                please refer to :code:`fluid.BuildStrategy`. The default is None.
            exec_strategy(ExecutionStrategy): exec_strategy specifies the options that can
                be changed when running the current model, such as the thread pool size.
                For more information about exec_strategy, please refer to :code:`fluid.ExecutionStrategy`.
                The default is None.
            share_vars_from(CompiledProgram): If share_vars_from is set, the current
                CompiledProgram will share the parameter value with the CompiledProgram
                specified by share_vars_from. This parameter needs to be set when model testing
                is required during model training, and the data parallel mode is used for
                training and testing. Since CompiledProgram will only distribute parameter
                Tensors to other devices when it is first executed, the CompiledProgram
                specified by share_vars_from must be run before the current CompiledProgram.
                The default is None.
            places(list(CUDAPlace)|list(CPUPlace)|list(str)|None): This parameter specifies the device
                on which the model is running. If you want to run on GPU0 and GPU1, places are
                [fluid.CUDAPlace(0), fluid.CUDAPlace(1)]; if you want to run with 2 CPUs, places are
                [fluid.CPUPlace()] * 2. If the parameter is not set, i.e. the parameter is None,
                the available device will be obtained from the environment variable when the model
                is executed: If the GPU is used, the currently available device ID is obtained
                from the environment variable FLAGS_selected_gpus or CUDA_VISIBLE_DEVICES when
                the model is executed; CPU, when the model is executed, the currently available
                CPU number is obtained from the environment variable CPU_NUM. For example,
                export CPU_NUM=4, if the environment variable is not set, the executor will
                add the variable to the environment variable and set its value to 1.
                The default is None. If ``places`` is the list of string, the string in the list
                can be ``cpu``, ``gpu:x``, where ``x`` is the index of the GPUs. 

        Returns:
            CompiledProgram

        Example:
            .. code-block:: python

                import numpy
                import os
                import paddle
                import paddle.static as static

                paddle.enable_static()

                use_cuda = True
                place = paddle.CUDAPlace(0) if use_cuda else paddle.CPUPlace()
                parallel_places = [paddle.CUDAPlace(0), paddle.CUDAPlace(1)] if use_cuda else [paddle.CPUPlace()] * 2

                # NOTE: If you use CPU to run the program, you need
                # to specify the CPU_NUM, otherwise, paddle will use
                # all the number of the logic core as the CPU_NUM,
                # in that case, the batch size of the input should be
                # greater than CPU_NUM, if not, the process will be
                # failed by an exception.
                if not use_cuda:
                    os.environ['CPU_NUM'] = str(2)

                exe = static.Executor(place)

                data = static.data(name='X', shape=[None, 1], dtype='float32')
                hidden = static.nn.fc(x=data, size=10)
                loss = paddle.mean(hidden)

                test_program = static.default_main_program().clone(for_test=True)
                paddle.optimizer.SGD(learning_rate=0.01).minimize(loss)

                exe.run(static.default_startup_program())
                compiled_train_prog = static.CompiledProgram(
                    static.default_main_program()).with_data_parallel(
                            loss_name=loss.name, places=parallel_places)
                # NOTE: if not set share_vars_from=compiled_train_prog,
                # the parameters used in test process are different with 
                # the parameters used by train process
                compiled_test_prog = static.CompiledProgram(
                    test_program).with_data_parallel(
                            share_vars_from=compiled_train_prog,
                            places=parallel_places)

                train_data = numpy.random.random(size=(10, 1)).astype('float32')
                loss_data, = exe.run(compiled_train_prog,
                                feed={"X": train_data},
                                fetch_list=[loss.name])
                test_data = numpy.random.random(size=(10, 1)).astype('float32')
                loss_data, = exe.run(compiled_test_prog,
                                feed={"X": test_data},
                                fetch_list=[loss.name])
        z5Already compiled with parallel, cannot be recompiled.z5Cannot compile with both data parallel and inference.TNzpThe loss name of CompiledProgram is None. The loss name should be set if CompiledProgram contains backward part.)rG   rH   rL   rM   rI   rJ   r<   listtupler   rK   r   r%   r>   )rN   Z	loss_namerO   Zexec_strategyZshare_vars_fromplacesr   r   r   with_data_parallel   s    o



z"CompiledProgram.with_data_parallelc                 C   sH   | j rJ d| jrJ dtt|tt|tgsJ d| _|| _| S )z Add inference optimize

        Args:
            config: instance of `NativeConfig` or `AnalysisConfig` to create predictor
        Returns:
            self
        z4Cannot compile with both data parallel and inferencez6Already compiled with inference, cannot be recompiled.T)rG   rH   anyr<   InferNativeConfigInferAnalysisConfig_infer_config)rN   configr   r   r   _with_inference_optimize:  s   
z(CompiledProgram._with_inference_optimizec                 C   s   t d)NzFSubclass of CompiledProgram should implement _with_distributed method.)NotImplementedErrorrN   r   r   r   _with_distributedM  s   z!CompiledProgram._with_distributedc              
   C   sH  | j r%|rtjd | j jstd| j jd u rtd| j j | _n|d us-J dg | _t	|t
sCt	|tsCJ dt|| jd u rLt | _t| j| j_| jd u r\t | _|| j_| jjdkr| jjtjkrvt|d | j_n| jjtjkrd| j_nt|d	 | j_d
t v rt d
 r| jjdkrtd d| j_| jjdkr| jsJ d| jr| jjdkr| jj r| jj }| jjt|ksJ d|| j_!| jr| jj"| j_#| jj$| j_%| jj&| j_'| jj(rd| j_)| jd ur'| jj*r'| jjtjksJ d| jjt| dksJ d| jj+tj,j-ks#J dd| j_.g | _/| j01 D ]A}|2 ro|3 d uro|3 4 ro|3  tj5j6j7krot89|: }| jd urot;| j|ro| j/<t89|:  q/tt=t>|}tt?| j/| _/| j/@  tA|| j/| jBrt89| jBntCDd| jE| j| j| j| j0S )Nz*share_vars_from is set, scope is ignored.
zHThe shared Program is not data parallel, cannot share variables from it.zSThe shared Program is not compiled and executed, so there is no variables to share. zPCurrently , The places type can only be list or tuple, but the input type is {}.r      r   r   ZFLAGS_use_cinnzwAt present, when CINN is turned on, each process can only contain one thread, so reset the number of threads to 1 here.z~If you use multi-trainer to train the model, you should use the data parallel model, i.e. calling with_data_parallel function.z5The trainer numbers is not equal to endpoint numbers.Tz%DGC only used under CUDA environment.z.DGC is not avaliable for single card training.zADGC                 only can be used for AllReduce BuildStrategy.F)FrJ   sysstderrwriterG   
ValueErrorrE   Zlocal_scopesZ_local_scopesr<   rR   rQ   formatr   rL   r   r    r?   Zis_distributionrM   r   Z_use_deviceZnum_threads
DeviceTypeCUDAlenXPUr   globalswarningswarnZnum_trainersZ_trainers_endpointsZtrainers_endpointsZ_nccl_comm_numZnccl_comm_numZ_use_hierarchical_allreduceZuse_hierarchical_allreduceZ$_hierarchical_allreduce_inter_nranksZ#hierarchical_allreduce_inter_nranksZsync_batch_normZenable_sequential_executionZ_enable_dgcZreduce_strategyZReduceStrategyZ	AllReduceZfuse_all_reduce_opsZ_persistable_varsr>   r!   Zis_varr:   persistableZVarDescZVarTypeZRAWcptZto_textnamer;   r)   mapr   setsortParallelExecutorrI   sixurC   )rN   rS   
use_devicescopeZtpsr$   rn   r   r   r   _compile_data_parallelR  s   



&
z&CompiledProgram._compile_data_parallelc                 C   s   t | jS r   )r   Zcreate_paddle_predictorrX   r\   r   r   r   _compile_inference  s   z"CompiledProgram._compile_inferencec                 C   s   | j r|r| j|krtd|r| j|std| S d| _ || _|| _| jr/|  | _| S | jr<| 	| j| j
| _
n| jg| _
t| j
dkrQt| jrQtdt| jtjr\tj}nt| jtjrgtj}ntj}| j|| j| j
d| _| S )a  Compile the program based on the configs.

        Args:
            scope: The variables (resources) that are associated with
               this compiled program.
            place: The location that the compiled program will be run on.

        Returns:
            self
        z,Cannot compile program with different scope.z,Cannot compile program with different place.Tr   zTIf optimizer is used in control flow, training on multi-places is not supported now.)ru   rv   rS   )rF   rC   rc   rD   Z_equalsrH   rx   rE   rG   _get_placesrK   rg   r4   r?   r[   r<   r   	CUDAPlacere   rf   XPUPlacerh   ZCPUrw   )rN   rv   r   ru   r   r   r   _compile  s>   

zCompiledProgram._compilec                 C   sn   |d u}|r|D ]}|  |  ksJ dqnt|tjr"t }nt|tjr,t }nt }|s5J d|S )Nz7Place type not match. You may set wrong type of places.zNo places for execution.)_typer<   r   rz   r   r{   r
   r	   )rN   r   Z
place_listZhas_set_placer   r   r   r   ry     s   zCompiledProgram._get_placesr   )NNNNN)__name__
__module____qualname____doc__rP   rT   rZ   r]   rw   rx   r|   ry   r   r   r   r   r   f   s     
1
 	
f1r   c                   @   s\   e Zd ZdZg Zdd Ze	dddZedd Zed	d
 Z	edd Z
edd ZdS )IpuDynamicPatcherz1
    Patcher for IPU dynamic2static support.
    c                 C   s   d S r   r   r\   r   r   r   rP     s   zIpuDynamicPatcher.__init__Nc              
      s  ddl m} ddlm  ddlm} ddlm} ddlj	}j
}j}j | fdd	}	jr>|jd
d jD ]}
|
j }|
  }|| qA|redd |dd D ndd |D dd |D jr|	 _j}|j D ][\}}| D ]R\}}|jj|jd|j|j|jdd}| |}|| |jj!||ddd W d   n1 sw   Y  "|j }|  }|| ||j| |< qq|fdd}| }|_S )zD
        Convert the ConcreteProgram to IPUConcreteProgram.
        r   )switch_to_static_graph)backward)Constant)device_guardr   Nc                     s<   j }  | j j }j }|| | S r   )r   Zgradients_with_optimizer
_optimizerstaticExecutorCPUPlaceZdefault_startup_programrun)r+   exestartup_program)r   concrete_programipu_strategypaddler   r   append_backward_desc  s   

zHIpuDynamicPatcher.convert_concrete_program.<locals>.append_backward_descZfloat16)dtypec                 S      g | ]	}|d ur|j qS r   rn   .0elemr   r   r   
<listcomp>4      z>IpuDynamicPatcher.convert_concrete_program.<locals>.<listcomp>r   c                 S   r   r   r   r   r   r   r   r   6  r   c                 S   s   g | ]}|j qS r   r   r   r   r   r   r   7  s    T)rn   rl   r   r   shapeZbelong_to_optimizerg        )value)Zinitializerc                     sf   j r%jj } dh| _jjj j| dd}jjj  j|d t	 jd
}|S )NZcumsumF)Zuse_fp16_guard)to_fp16_var_names)r   rv   )enable_fp16r   ampZCustomOpListsZunsupported_listZcast_model_to_fp16r   Zcast_parameters_to_fp16r   r   compile)Zamp_listr   r+   )r   	feed_list
fetch_listr   r   rv   r   r   func_compileP  s,   z@IpuDynamicPatcher.convert_concrete_program.<locals>.func_compile)#Zfluid.dygraph.baser   Zfluidr   Zfluid.initializerr   Zfluid.frameworkr   r   inputsoutputsr   r   global_scoper   to
parametersr:   rn   Z
get_tensorr   Z_share_data_withis_trainingr   r   Z_accumulatorsitemshelperZcreate_global_variabler   r   r   Z_get_device_for_paramZset_variable_initializerZfind_var)r   r   class_instancer   r   r   r   r   r   r   Zparam_or_bufferZparam_or_buffer_tensorZ
src_tensor	optimizerkv
param_nameZvar_tmpr:   ZdeviceZparam_or_lr_tensorZoptim_tensorr   r   r   )r   r   r   r   r   r   rv   r   convert_concrete_program  st   





z*IpuDynamicPatcher.convert_concrete_programc                    sx   ddl m} ddl m  ddlm ddl m ddlm |j} fdd}t	|d	| t
j|d	|g d
S )z Monkey patch ProgramCache discriptor to support dynamic2static in IPU.

        Args:
            ipu_strategy: The ipu_strategy used in dynamic graph.

        Returns:
            None
        r   )ProgramCache)CacheKey)logging_utils)MAX_TRACED_PROGRAM_COUNT)partial_program_fromc                    s   t | stdt|j t|}|| _|| jvsjr\|| jv r'd | jr2js2d | 	|\}}t
||j}||f| j|< t| j}|kr\d| | j| S )Nz.type(item) should be CacheKey, but received %sz3ipu_strategy chances detected. Please sync weights.zadynamic2static on IPU doesn't support mutiple caches. Please make suredynamic inputs is not used.zCurrent traced program number: {} > `max_tracing_count`:{}. Too much cached programs will bring expensive overhead. The reason may be: (1) passing tensors with different shapes, (2) passing python objects instead of tensors.)r<   rc   r   r~   hashZ_recent_keyZ_cachesneed_compilerk   Z_build_oncer   r   r   rg   rd   )rN   itemZitem_idr   _Zcurrent_tracing_countr   r   r   r   r   r   r   patch_getterz  sB   




z;IpuDynamicPatcher.patch_program_cache.<locals>.patch_getter__getitem__N)Z2fluid.dygraph.dygraph_to_static.program_translatorr   r   Zfluid.dygraph.dygraph_to_staticr   r   Z/fluid.dygraph.dygraph_to_static.partial_programr   r   setattrr   patcher_cacher)   )r   r   Z
old_getterr   r   r   r   patch_program_cacheh  s   
 z%IpuDynamicPatcher.patch_program_cachec                    sD   ddl m} |jd fdd	}t|d| tj|dg d S )Nr   LRSchedulerc                    s   | |   d| ji d S )Nlr)set_optionsZlast_lr)rN   epochr   Zold_stepr   r   
patch_step  s   
z8IpuDynamicPatcher.patch_lr_scheduler.<locals>.patch_stepstepr   )paddle.optimizer.lrr   r   r   r   r   r)   )r   r   r   r   r   r   patch_lr_scheduler  s
   z$IpuDynamicPatcher.patch_lr_schedulerc                 C   s   t |  t |  d S r   )r   r   r   )r   r   r   r   register_patch  s   
z IpuDynamicPatcher.register_patchc                  C   s"   t jD ]\} }}t| || qd S r   )r   r   r   )modulekeyattrr   r   r   release_patch  s   zIpuDynamicPatcher.release_patchr   )r~   r   r   r   r   rP   staticmethodr   r   r   r   r   r   r   r   r   r     s    `
5

r   c                   @   s   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d Z				d)ddZ					d*ddZ
d+ddZ			d,ddZdd Zdd Zdd Zdd  Zed!d" Zed#d$ Zed%d& Zed'd( ZdS )-r   a~  
    Help users precisely control the graph building in :code:`paddle.static.IpuCompiledProgram` .

    Returns:
        The IpuStrategy instance.

    Examples:
        .. code-block:: python
	
            # required: ipu

            import paddle
            import paddle.static as static

            paddle.enable_static()

            ipu_strategy = static.IpuStrategy()
    c                 C   sn   t  r"t  | _dddddd}| j| d| _g | _d| _ntdddl	m
} | r5|   d S d S )	Nr   r   )Zon_chipZuse_replicated_tensor_sharding)Zlocation_optimizerZ+accumulation_and_replication_reduction_typeZ4mean_accumulation_and_replication_reduction_strategyFTz\Can not use IpuStrategy in non IPU compiled environment, please re-compile with WITH_IPU=ON.in_dynamic_mode)r   is_compiled_with_ipur   _ipu_strategyr   has_custom_opscustom_op_namesr   RuntimeErrorr   r   r   )rN   Zdefault_optionsr   r   r   r   rP     s&   

zIpuStrategy.__init__c                 C   s   t |  dS )a  
        Register patchs function to support dynamic to static on IPU. This operation would break the dy2static functionality on CPU.
        Use `release_patch` to release the patch.

        Examples:
            .. code-block:: python
	
                # required: ipu

                import paddle
                import paddle.static as static

                ipu_strategy = static.IpuStrategy()

                ipu_strategy.register_patch()
        N)r   r   r\   r   r   r   r     s   zIpuStrategy.register_patchc                 C   s   t   dS )a@  
        Release the registered IPU functions.

        Examples:
            .. code-block:: python
	
                # required: ipu

                import paddle
                import paddle.static as static

                ipu_strategy = static.IpuStrategy()

                ipu_strategy.release_patch()
        N)r   r   r\   r   r   r   r     s   zIpuStrategy.release_patchc                 C   s:   ddl m} | r|| _| |}| j| dS td)a  
        Set optimizer to ipu_strategy in dynamic mode.

          Args:
              optimizer (Optimizer): Optimizer to be used in training.
              
          Returns:
              None.

          Examples:
              .. code-block:: python
	
                  # required: ipu

                  import paddle
                  import paddle.static as static

                  linear = paddle.nn.Linear(10, 10)
                  optimizer = paddle.optimizer.SGD(learning_rate=0.01,
                                                   parameters=linear.parameters())
                  ipu_strategy = static.IpuStrategy()
                  ipu_strategy.set_optimizer(optimizer)
        r   r   z,Only needs to set optimizer in dynamic mode.N)r   r   r   parse_optimizerr   r   r   )rN   r   r   optimizer_attrsr   r   r   set_optimizer	  s   
zIpuStrategy.set_optimizerc                    s4    fdd}|g}ddi}|D ]}| |  q|S )a  
        Parse optimizer attributes for IPU dynamic to static support. Currently only support parse lr.

          Args:
              optimizer (Optimizer): Optimizer to be parsed.
              
          Returns:
              Dict.

          Examples:
              .. code-block:: python
	
                  # required: ipu

                  import paddle
                  import paddle.static as static

                  linear = paddle.nn.Linear(10, 10)
                  optimizer = paddle.optimizer.SGD(learning_rate=0.01,
                                                   parameters=linear.parameters())
                  ipu_strategy = static.IpuStrategy()
                  attrs = ipu_strategy.parse_optimizer(optimizer)
        c                     s>   ddl m}  t jtrd jiS t j| rd  iS d S )Nr   r   r   )r   r   r<   Z_learning_ratefloatr   r   r   r   get_lrB  s   
z+IpuStrategy.parse_optimizer.<locals>.get_lrZ
is_dynamicT)update)rN   r   r   Zattr_fnr   fnr   r   r   r   )  s   zIpuStrategy.parse_optimizerr   TFc                 C   s0   |dkr
|r
t d||||d}| | dS )a=  
        Set graph configuration to the IpuStrategy instance.

        Args:
            num_ipus (int, optional): Number of IPU devices. Default 1, which means only use 1 IPU.
            is_training (bool, optional): True is training graph, False is inference graph. Default True, which means is training mode.
            batch_size (int, optional): The batch-size in the graph. Used to make the graph batch-size fixed,
                if the batch-size in the graph is dynamic. Default 1, which means the batch-size would be set 1, if the batch-size is dynamice.
            enable_manual_shard (bool, optional): Enable graph sharding or not. Only if num_ipus > 1, enable_manual_shard is able to be set True. 
                Default False, which means disabled.    
            
        Returns:
            None.

        Examples:
            .. code-block:: python
	
                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.set_graph_config(num_ipus=1,
                                            is_training=True,
                                            micro_batch_size=1,
                                            enable_manual_shard=False)
        r   zAOnly if num_ipus > 1, enable_manual_shard is able to be set True.)num_ipusr   micro_batch_sizeenable_manual_shardN)r   r   )rN   r   r   r   r   optionsr   r   r   set_graph_configO  s   #zIpuStrategy.set_graph_configc                 C   s6   |  d}|s|rtd||||d}| | dS )a  
        Set pipelining configuration to the IpuStrategy instance. Used to optimize the throughput performance.

        Args:
            enable_pipelining (bool, optional): Enable data pipelining between subgraphs. Only if enable_manual_shard=True, enable_pipelining is able to be set True. 
                Default False, which means disabled.
            batches_per_step (int, optional): Set the batches per run in data pipelining mode. Only if enable_pipelining=True, batches_per_step is able to be set > 1.
                Default 1, which means no data pipelining.
            enable_gradient_accumulation (bool, optional): Enable to accumulate gradients before updating the weights in training mode. Only if enable_pipelining=True,
                enable_gradient_accumulation is able to be set True. Default False, which means no gradient accumulation. 
            accumulation_factor (int, optional): Specify the number of micro-batches to accumulate 
                before applying the varUpdate. Default 1, which means disable the accumulation.
        
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.set_pipelining_config(enable_pipelining=False,
                                                    batches_per_step=1,
                                                    enable_gradient_accumulation=False,
                                                    accumulation_factor=1)
        r   zKOnly if enable_manual_shard=True, enable_pipelining is able to be set True.)enable_pipeliningbatches_per_stepenable_gradient_accumulationaccumulation_factorN)
get_optionr   r   )rN   r   r   r   r   r   r   r   r   r   set_pipelining_config~  s   
%z!IpuStrategy.set_pipelining_configc                 C   s   d|i}|  | dS )a  
        Set half computation configuration to the IpuStrategy instance. Used to optimize the performance.

        Args:
            enable_fp16 (bool, optional): Enable FLOAT16 mode and transform FLOAT32 to FLOAT16. Default False, which means disable FLOAT16 mode.
        
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.set_precision_config(enable_fp16=False)
        r   N)r   )rN   r   r   r   r   r   set_precision_config  s   z IpuStrategy.set_precision_configN
custom.opsc                 C   sH   |du r|}||||d}|  d|i | j| | js"d| _dS dS )a  
        Add a mapping to use popart custom ops running on the IPU.

        Args:
            paddle_op(str): the name of custom op in paddle.

            popart_op(str): the name of custom op in popart.

            domain(str): domain name of custom op in popart.

            version(int): version of custom op in popart.
        
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.add_custom_op('paddle_relu', 'popart_relu')
        N)	paddle_op	popart_opdomainversion	custom_opT)r   r   r)   r   )rN   r   r   r   r   r   r   r   r   add_custom_op  s   !
zIpuStrategy.add_custom_opc                 C   s,   | j | dh}| | rd| _dS dS )a  
        Set options from dict.

        Args:
            options(dict): dict of options.
        
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                options = {'num_ipus':1, 'enable_fp16': True}
                ipu_strategy.set_options(options)
        r   TN)r   r   keysr   )rN   r   Zrecompile_white_listr   r   r   r     s
   
zIpuStrategy.set_optionsc                 C   s   | j |d S )a  
        Get option.

        Args:
            option(str): name of option.
        
        Returns:
            option value.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                num_ipus = ipu_strategy.get_option('num_ipus')
        r   )r   r   )rN   optionr   r   r   r        zIpuStrategy.get_optionc                 C      | j | dS )a  
        Enable PopART pattern to optimize the graph.

        Args:
            pattern(string): the name of the pattern.
        
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.enable_pattern("ViewSimplifyPattern")
        N)r   enable_patternrN   patternr   r   r   r   1  r   zIpuStrategy.enable_patternc                 C   r   )a  
        Disable PopART pattern.

        Args:
            pattern(string): the name of the pattern.
        
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.disable_pattern("ViewSimplifyPattern")
        N)r   disable_patternr   r   r   r   r   J  r   zIpuStrategy.disable_patternc                 C   
   |  dS )zJ
        Get the number of IPU devices from IpuStrategy instance.
        r   r   r\   r   r   r   r   c     
zIpuStrategy.num_ipusc                 C   r   )zU
        Get the boolean of training or inference from IpuStrategy instance.
        r   r   r\   r   r   r   r   j  r   zIpuStrategy.is_trainingc                 C   r   )zX
        Get the boolean of enable pipelining or not from IpuStrategy instance.
        r   r   r\   r   r   r   r   q  r   zIpuStrategy.enable_pipeliningc                 C   r   )zS
        Get the boolean of float16 mode or not from IpuStrategy instance.
        r   r   r\   r   r   r   r   x  r   zIpuStrategy.enable_fp16)r   Tr   F)Fr   Fr   )F)Nr   r   )r~   r   r   r   rP   r   r   r   r   r   r   r   r   r   r   r   r   propertyr   r   r   r   r   r   r   r   r     sD     '
0

2
.


r   c                   @   s"   e Zd ZdZdddZdd ZdS )r   aB  
    The IpuCompiledProgram is used to transform a program to a ipu-target program,
    such as forward graph extraction, computing graph transformation, useless scale Ops clean, etc.

    Args:
        program(Program, optional): This parameter represents the :code:`Program`
            to be executed. Default is None, which means the program will be set to 
            the default program :code:`paddle.static.default_main_program()` .
        scope(Scope, optional): The scope used to run this program, you can switch
            it to different scope. Default is None, which means use the global 
            scope :code:`paddle.static.global_scope()` .
        ipu_strategy(IpuStrategy, optional): This argument is used to build the program with the
            specified options, such as half computation, training or inference session, the number of IPUs, etc.
            Default is None, which means build the program based on the default `ipu_strategy`. 

    Returns:
        IpuCompiledProgram

    Example:
        .. code-block:: python
	
            # required: ipu

            import paddle
            import paddle.static as static

            paddle.enable_static()

            a = static.data(name='data', shape=[None, 1], dtype='int32')
            b = a + 1
            main_prog = static.default_main_program()
            
            ipu_strategy = static.IpuStrategy()
            ipu_strategy.set_graph_config(num_ipus=1, is_training=True, micro_batch_size=1)
            ipu_strategy.set_pipelining_config(enable_pipelining=False, batches_per_step=1, enable_gradient_accumulation=False, accumulation_factor=1)
            ipu_strategy.set_precision_config(enable_fp16=False)
            
            ipu_compiled_program = static.IpuCompiledProgram(
                main_prog,
                ipu_strategy=ipu_strategy)
    Nc                 C   s   t  std|d u rt }t|tjstdt| || _	d| _
|d ur,|| _n
dd l}|j | _|d ur>|| _nt | _|jrLt|j| _nd| _t j | _d S )NzECan not use this function since PaddlePaddle is not compiled with IPUz:The type of program is wrong, expected Program, but got %sFr   r   )r   r   rc   r   r   r<   r@   rB   r   r?   rF   rC   r   r   r   r   r   r   rp   r   _custom_op_namesZ
IpuBackendZget_instance_backend)rN   r+   rv   r   r   r   r   r   rP     s2   zIpuCompiledProgram.__init__c                 C   sR  | j | j | j | jj | j }g }t|jD ]\}}|j	
d |jdks/|jdkr4|| q|ddd D ]}|| q<dD ]}||rR|| qF| jj	  t| jj	| _| jjryddg}	|	D ]}
t|
}|| j qkg d	}	|	D ]}
t|
}|
d
kr|d| || j qtd}| jr|d| j || j g d}	|	D ]}
t|
}|d| |d| || j qtd}t }|d| || j tj|}t| jdr| jjj }| jj|_| j }|j!| |j_"| }|D ]}|#|}|j	$d q
t|ds#| j|_%d| j_&|S )a.  
        This interface is used to compile the input Program to a program
        to run the model on the ipu.
        
        Args:
            feed_list(list): This parameter represents the input Tensors of the model.

            fetch_list(list): This parameter represents the Tensors that need to be returned
                after the model.

        Returns:
            Program

        Example:
            .. code-block:: python
    	
                # required: ipu
    
                import paddle
                import paddle.static as static
    
                paddle.enable_static()
    
                a = static.data(name='data', shape=[None, 1], dtype='int32')
                b = a + 1
                main_prog = static.default_main_program()

                ipu_strategy = static.IpuStrategy()
                ipu_strategy.set_graph_config(num_ipus=1, is_training=True, micro_batch_size=1)
                ipu_strategy.set_pipelining_config(enable_pipelining=False, batches_per_step=1, enable_gradient_accumulation=False, accumulation_factor=1)
                ipu_strategy.set_precision_config(enable_fp16=False)
                
                program = static.IpuCompiledProgram(
                    main_prog,
                    ipu_strategy=ipu_strategy).compile([a.name], [b.name])
        Fr&   fetchNr'   )r&   r   Zoptimizer_extract_passZoptimizer_state_align_pass)Zforward_graph_extract_passinfer_shape_passZavg_shard_passZdelete_scale_op_passr  r   Zpopart_canonicalization_passZ
custom_ops)Zipu_inplace_passZipu_graph_builder_passZipu_runtime_replacer_passr   Zgraph_to_program_passr+   lr_shedulerorg_program)'r   Z	set_scoperC   Zset_ipu_strategyr   r?   r   r(   r   rA   Zset_is_targetr   r)   r*   Zhas_varZ_remove_varflushr   r=   r>   r   Zget_passapplyrp   r   ZProgramDescZset_not_ownedr   r@   Z_construct_from_deschasattrr  Z	_var_namer6   Zlr_varr:   Zset_need_check_feedr  r   )rN   r   r   r   Zneed_to_remove_op_indexr,   r   r-   r:   ZpassesZ	pass_nameZa_passZconvert_passrA   r+   Zlr_var_nameZprogram_global_blockZ	feed_nameZfeed_varr   r   r   r     st   %












zIpuCompiledProgram.compile)NNN)r~   r   r   r   rP   r   r   r   r   r   r     s    
*$r   )$multiprocessingosrs   r`   rj   r^   r   rm   r   r   r   r   r	   r
   r   __all__rr   r   r   ZNativeConfigrV   ZAnalysisConfigrW   re   r   r    r%   r.   r2   r4   r;   objectr   r   r   r   r   r   r   r   <module>   sD   		

    ;   L