compiler.py 52.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import multiprocessing
import os
X
polish  
Xin Pan 已提交
17
import sys
18
import warnings
X
Xin Pan 已提交
19
from . import framework
20
from .framework import _get_paddle_place, _get_paddle_place_list
21
from .framework import cuda_places, cpu_places, xpu_places
22 23
from . import core

J
jianghaicheng 已提交
24
__all__ = [
25 26 27 28 29
    'CompiledProgram',
    'ExecutionStrategy',
    'BuildStrategy',
    'IpuCompiledProgram',
    'IpuStrategy',
J
jianghaicheng 已提交
30
]
X
Xin Pan 已提交
31

32 33
ExecutionStrategy = core.ParallelExecutor.ExecutionStrategy
BuildStrategy = core.ParallelExecutor.BuildStrategy
F
flame 已提交
34 35
InferNativeConfig = core.NativeConfig
InferAnalysisConfig = core.AnalysisConfig
36
DeviceType = core.DeviceType
37 38 39 40 41 42 43 44


def _place_obj(place):
    p = core.Place()
    p.set_place(place)
    return p


45
def _is_pserver_mode(main_program):
46
    main = main_program if main_program else framework.default_main_program()
47 48 49 50 51 52
    for op in main.global_block().ops:
        if op.type in ["send", "recv"]:
            return True
    return False


C
chengduo 已提交
53 54
def _has_backward_op(graph):
    for node in graph.nodes():
55 56 57 58 59
        if (
            node.is_op()
            and node.op() is not None
            and node.op().type().endswith("_grad")
        ):
C
chengduo 已提交
60 61 62 63
            return True
    return False


64 65 66 67
def _prune_feed_ops(program):
    # prune the feed ops in the program.
    pop_idx = []
    for i, op in enumerate(program.global_block().ops):
68 69
        if op.type == "feed":
            pop_idx.append(i)
70 71 72 73
    for index in pop_idx[::-1]:
        program.global_block()._remove_op(index)


74 75 76 77
def _has_optimize_op(block):
    for op in block.ops:
        op_maker = core.op_proto_and_checker_maker
        optimize = core.op_proto_and_checker_maker.OpRole.Optimize
78 79 80
        if op_maker.kOpRoleVarAttrName() in op.attr_names and int(
            op.all_attrs()[op_maker.kOpRoleAttrName()]
        ) == int(optimize):
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
            return True
    return False


def _has_optimizer_in_control_flow(program):
    if not program:
        program = framework.default_main_program()
    for op in program.global_block().ops:
        if op.type == "conditional_block_grad":
            sub_block = program.block(op._block_attr_id("sub_block"))
            if _has_optimize_op(sub_block):
                return True

    return False


97 98 99 100 101 102
def _should_broadcast_or_not_exists(program, var_name):
    block = program.global_block()
    var = block.vars.get(var_name, None)
    if var is None:
        return True
    is_distributed = getattr(var, '_is_distributed', False) or getattr(
103 104
        var, 'is_distributed', False
    )
105 106 107
    return not is_distributed


X
polish  
Xin Pan 已提交
108
class CompiledProgram(object):
X
polish  
Xin Pan 已提交
109
    """
110
    :api_attr: Static Graph
111

C
chengduo 已提交
112 113 114 115 116
    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
117
    :code:`paddle.static.BuildStrategy`.
X
polish  
Xin Pan 已提交
118

C
chengduo 已提交
119
    Args:
120
        program_or_graph (Graph|Program): This argument is the Program or Graph
C
chengduo 已提交
121
            being executed.
122
        build_strategy(BuildStrategy): This argument is used to compile the
C
chengduo 已提交
123 124 125
            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,
126
            please refer to :code:`paddle.static.BuildStrategy`. The default is None.
X
Xin Pan 已提交
127

C
chengduo 已提交
128 129
    Returns:
        CompiledProgram
X
polish  
Xin Pan 已提交
130 131

    Example:
X
Xin Pan 已提交
132
        .. code-block:: python
133

134 135 136
            import numpy
            import paddle
            import paddle.static as static
137

138
            paddle.enable_static()
139

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

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

148 149 150 151 152 153 154 155
            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])
X
polish  
Xin Pan 已提交
156 157
    """

C
chengduo 已提交
158
    def __init__(self, program_or_graph, build_strategy=None):
X
Xin Pan 已提交
159 160
        if isinstance(program_or_graph, core.Graph):
            self._graph = program_or_graph
161
            # don't not create a new program here.
X
Xin Pan 已提交
162 163
            self._program = None
        elif isinstance(program_or_graph, framework.Program):
164
            _prune_feed_ops(program_or_graph)
X
Xin Pan 已提交
165 166 167
            self._graph = core.Graph(program_or_graph.desc)
            self._program = program_or_graph
        else:
168 169
            raise TypeError(
                "The type of program_to_graph parameter is wrong, expected Graph or Program, but received %s"
170 171
                % type(program_or_graph)
            )
X
Xin Pan 已提交
172

X
polish  
Xin Pan 已提交
173 174 175
        self._scope = None
        self._place = None
        self._executor = None
176 177
        self._compiled = False
        self._is_data_parallel = False
F
flame 已提交
178
        self._is_inference = False
C
chengduo 已提交
179 180 181 182 183
        self._loss_name = None
        self._share_vars_from = None
        self._places = None
        self._build_strategy = build_strategy
        self._exec_strategy = None
184

185 186 187 188 189 190 191 192
    def with_data_parallel(
        self,
        loss_name=None,
        build_strategy=None,
        exec_strategy=None,
        share_vars_from=None,
        places=None,
    ):
C
chengduo 已提交
193 194 195 196 197 198
        """
        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,
199 200
        and so on.

201
        .. note::
202 203 204
            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
205
            with_data_parallel interface.
C
chengduo 已提交
206 207

        Args:
208
            loss_name (str): This parameter is the name of the loss Tensor of the model.
C
chengduo 已提交
209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224
                **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
225
                Tensors to other devices when it is first executed, the CompiledProgram
C
chengduo 已提交
226 227
                specified by share_vars_from must be run before the current CompiledProgram.
                The default is None.
228
            places(list(CUDAPlace)|list(CPUPlace)|list(str)|None): This parameter specifies the device
C
chengduo 已提交
229 230 231 232 233 234 235 236 237 238
                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.
239
                The default is None. If ``places`` is the list of string, the string in the list
240
                can be ``cpu``, ``gpu:x``, where ``x`` is the index of the GPUs.
C
chengduo 已提交
241 242 243

        Returns:
            CompiledProgram
X
Xin Pan 已提交
244

245 246 247
        Example:
            .. code-block:: python

248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270
                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')
271
                hidden = static.nn.fc(x=data, size=10)
272 273 274 275 276 277 278 279 280 281
                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,
282
                # the parameters used in test process are different with
283 284 285 286 287 288 289 290
                # 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,
291 292
                                feed={"X": train_data},
                                fetch_list=[loss.name])
293 294
                test_data = numpy.random.random(size=(10, 1)).astype('float32')
                loss_data, = exe.run(compiled_test_prog,
295 296
                                feed={"X": test_data},
                                fetch_list=[loss.name])
X
Xin Pan 已提交
297
        """
298 299 300 301 302 303
        assert (
            not self._is_data_parallel
        ), "Already compiled with parallel, cannot be recompiled."
        assert (
            not self._is_inference
        ), "Cannot compile with both data parallel and inference."
304
        self._is_data_parallel = True
C
chengduo 已提交
305 306 307 308
        # FIXME(zcd): Currently, the build_strategy can be set during creating
        # CompiledProgram or calling with_data_parallel, and it may be confusing,
        # but in the long run, we should set up build_strategy only when creating
        # CompiledProgram, and exec_strategy should be deprecated.
309 310
        if build_strategy is not None:
            self._build_strategy = build_strategy
311 312
        self._exec_strategy = exec_strategy
        self._loss_name = loss_name
X
polish  
Xin Pan 已提交
313
        self._share_vars_from = share_vars_from
314 315 316 317
        if isinstance(places, (list, tuple)):
            self._places = _get_paddle_place_list(places)
        else:
            self._places = _get_paddle_place(places)
C
chengduo 已提交
318 319

        if _has_backward_op(self._graph):
320 321 322
            assert (
                self._loss_name is not None
            ), "The loss name of CompiledProgram is None. The loss name should be set if CompiledProgram contains backward part."
C
chengduo 已提交
323 324 325 326 327

        if self._places is not None:
            if not isinstance(self._places, (list, tuple)):
                self._places = [self._places]

328 329
        return self

F
flame 已提交
330
    def _with_inference_optimize(self, config):
331
        """Add inference optimize
F
flame 已提交
332 333 334 335 336 337

        Args:
            config: instance of `NativeConfig` or `AnalysisConfig` to create predictor
        Returns:
            self
        """
338 339 340 341 342 343 344 345 346 347 348 349 350
        assert (
            not self._is_data_parallel
        ), "Cannot compile with both data parallel and inference"
        assert (
            not self._is_inference
        ), "Already compiled with inference, cannot be recompiled."

        assert any(
            [
                isinstance(config, InferNativeConfig),
                isinstance(config, InferAnalysisConfig),
            ]
        )
F
flame 已提交
351 352 353
        self._is_inference = True
        self._infer_config = config
        return self
X
polish  
Xin Pan 已提交
354

F
flame 已提交
355
    def _with_distributed(self):
356 357 358
        raise NotImplementedError(
            "Subclass of CompiledProgram should implement _with_distributed method."
        )
X
polish  
Xin Pan 已提交
359

360
    def _compile_data_parallel(self, places, use_device, scope=None):
X
polish  
Xin Pan 已提交
361
        if self._share_vars_from:
362
            if scope:
X
polish  
Xin Pan 已提交
363 364
                sys.stderr.write("share_vars_from is set, scope is ignored.\n")
            if not self._share_vars_from._is_data_parallel:
365 366
                raise ValueError(
                    "The shared Program is not data parallel, cannot "
367 368
                    "share variables from it."
                )
X
polish  
Xin Pan 已提交
369 370
            if self._share_vars_from._executor is None:
                raise ValueError(
371
                    "The shared Program is not compiled and executed, so there is no "
372 373
                    "variables to share."
                )
X
polish  
Xin Pan 已提交
374 375
            self._local_scopes = self._share_vars_from._executor.local_scopes()
        else:
376
            assert scope is not None, ""
X
polish  
Xin Pan 已提交
377
            self._local_scopes = []
378

379 380 381 382 383
        assert isinstance(places, tuple) or isinstance(
            places, list
        ), "Currently , The places type can only be list or tuple, but the input type is {}.".format(
            type(places)
        )
C
chengduo 已提交
384 385 386 387 388 389 390

        if self._build_strategy is None:
            self._build_strategy = BuildStrategy()
        self._build_strategy.is_distribution = _is_pserver_mode(self._program)

        if self._exec_strategy is None:
            self._exec_strategy = ExecutionStrategy()
391
        self._exec_strategy._use_device = use_device
392 393

        if self._exec_strategy.num_threads == 0:
394
            if self._exec_strategy._use_device == DeviceType.CUDA:
395 396
                # Experiments on se-resnext shows that too many threads hurt
                # performance. Worth tunning for other models in the future.
C
chengduo 已提交
397
                self._exec_strategy.num_threads = len(places) * 4
398
            elif self._exec_strategy._use_device == DeviceType.XPU:
399 400
                # Currently only single thread is supported in Kunlun XPU.
                self._exec_strategy.num_threads = 1
401
            else:
C
chengduo 已提交
402 403
                self._exec_strategy.num_threads = len(places) * 2

404 405 406 407 408 409 410 411 412
        if (
            "FLAGS_use_cinn" in core.globals()
            and core.globals()["FLAGS_use_cinn"]
            and self._exec_strategy.num_threads != 1
        ):
            warnings.warn(
                "At present, when CINN is turned on, each process can "
                "only contain one thread, so reset the number of threads to 1 here."
            )
413 414
            self._exec_strategy.num_threads = 1

C
chengduo 已提交
415
        if self._build_strategy.num_trainers > 1:
416 417
            assert self._is_data_parallel, (
                "If you use multi-trainer to train the model, you should use "
C
chengduo 已提交
418
                "the data parallel model, i.e. calling with_data_parallel function."
419
            )
420

X
Xin Pan 已提交
421 422
        # TODO(wuyi): trainer endpoings should be passed in through
        # build_strategy, not program.xxx.
423
        # TODO(gongwb): let user to set them once.
424 425 426 427 428
        if (
            self._program
            and self._build_strategy.num_trainers > 1
            and self._program._trainers_endpoints
        ):
X
Xin Pan 已提交
429
            tps = self._program._trainers_endpoints
D
dzhwinter 已提交
430

431
            assert self._build_strategy.num_trainers == len(
432 433
                tps
            ), "The trainer numbers is not equal to endpoint numbers."
X
Xin Pan 已提交
434 435
            self._build_strategy.trainers_endpoints = tps

436 437
        if self._program:
            self._build_strategy.nccl_comm_num = self._program._nccl_comm_num
438 439 440 441 442 443
            self._build_strategy.use_hierarchical_allreduce = (
                self._program._use_hierarchical_allreduce
            )
            self._build_strategy.hierarchical_allreduce_inter_nranks = (
                self._program._hierarchical_allreduce_inter_nranks
            )
444

Q
qingqing01 已提交
445 446 447
        if self._build_strategy.sync_batch_norm:
            self._build_strategy.enable_sequential_execution = True

448
        if self._program is not None and self._program._enable_dgc:
449 450 451 452 453 454 455 456 457 458
            assert (
                self._exec_strategy._use_device == DeviceType.CUDA
            ), "DGC only used under CUDA environment."
            assert (
                self._build_strategy.num_trainers * len(places) > 1
            ), "DGC is not avaliable for single card training."
            assert (
                self._build_strategy.reduce_strategy
                == BuildStrategy.ReduceStrategy.AllReduce
            ), "DGC \
459
                only can be used for AllReduce BuildStrategy."
460 461 462 463

            # DGC doesn't support fuse for now, close fuse.
            self._build_strategy.fuse_all_reduce_ops = False

X
Xin Pan 已提交
464
        self._persistable_vars = []
Z
Zhen Wang 已提交
465
        for node in self._graph.nodes():
466 467 468 469 470 471
            if (
                node.is_var()
                and node.var() is not None
                and node.var().persistable()
                and node.var().type() != core.VarDesc.VarType.RAW
            ):
472
                name = node.name()
473 474 475 476
                if (
                    self._program is not None
                    and _should_broadcast_or_not_exists(self._program, name)
                ):
477
                    self._persistable_vars.append(node.name())
478

C
chengduo 已提交
479 480
        places = list(map(_place_obj, places))

Y
Yan Xu 已提交
481 482 483 484 485 486
        # ParallelExecutor would broadcast all the parameters during initializing.
        # The parameters of each process should be in the same ordered for the data-parallelism
        # distributed training to keep the broadcast correct.
        self._persistable_vars = list(set(self._persistable_vars))
        self._persistable_vars.sort()

487 488 489 490 491 492 493 494 495 496
        return core.ParallelExecutor(
            places,
            self._persistable_vars,
            self._loss_name if self._loss_name else '',
            self._scope,
            self._local_scopes,
            self._exec_strategy,
            self._build_strategy,
            self._graph,
        )
497

F
flame 已提交
498 499 500
    def _compile_inference(self):
        return core.create_paddle_predictor(self._infer_config)

501
    def _compile(self, scope, place):
X
Xin Pan 已提交
502 503 504 505 506 507 508 509 510 511
        """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
        """
512
        if self._compiled:
X
polish  
Xin Pan 已提交
513
            if scope and self._scope != scope:
514
                raise ValueError("Cannot compile program with different scope.")
S
sneaxiy 已提交
515
            if place and not self._place._equals(place):
516
                raise ValueError("Cannot compile program with different place.")
517
            return self
X
fix  
Xin Pan 已提交
518
        self._compiled = True
519 520 521

        self._scope = scope
        self._place = place
C
chengduo 已提交
522 523

        if self._is_inference:
F
flame 已提交
524
            self._executor = self._compile_inference()
525
        else:
C
chengduo 已提交
526 527 528 529
            if self._is_data_parallel:
                self._places = self._get_places(self._place, self._places)
            else:
                self._places = [self._place]
530 531 532 533

            # Todo(liym27):If optimizer is used in control flow,
            #  training on multi-places is not supported now, will
            #  be supported later.
534 535 536
            if len(self._places) > 1 and _has_optimizer_in_control_flow(
                self._program
            ):
537 538
                raise NotImplementedError(
                    "If optimizer is used in control flow, "
539 540
                    "training on multi-places is not supported now."
                )
541
            if isinstance(self._place, core.CUDAPlace):
542
                use_device = DeviceType.CUDA
543
            elif isinstance(self._place, core.XPUPlace):
544
                use_device = DeviceType.XPU
545
            else:
546
                use_device = DeviceType.CPU
547 548 549
            self._executor = self._compile_data_parallel(
                use_device=use_device, scope=self._scope, places=self._places
            )
550
        return self
C
chengduo 已提交
551 552

    def _get_places(self, place, place_list):
553
        has_set_place = place_list is not None
C
chengduo 已提交
554 555
        if has_set_place:
            for p in place_list:
556 557 558
                assert (
                    p._type() == place._type()
                ), "Place type not match. You may set wrong type of places."
C
chengduo 已提交
559
        else:
560 561 562 563 564 565
            if isinstance(place, core.CUDAPlace):
                place_list = cuda_places()
            elif isinstance(place, core.XPUPlace):
                place_list = xpu_places()
            else:
                place_list = cpu_places()
566
        assert place_list, "No places for execution."
C
chengduo 已提交
567
        return place_list
J
jianghaicheng 已提交
568 569


570 571 572 573 574 575 576 577 578 579 580
class IpuDynamicPatcher(object):
    """
    Patcher for IPU dynamic2static support.
    """

    patcher_cache = []

    def __init__(self):
        pass

    @staticmethod
581 582 583
    def convert_concrete_program(
        ipu_strategy, concrete_program, class_instance=None
    ):
584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618
        """
        Convert the ConcreteProgram to IPUConcreteProgram.
        """
        from ..fluid.dygraph.base import switch_to_static_graph
        from ..fluid import backward
        from ..fluid.initializer import Constant
        from ..fluid.framework import device_guard
        import paddle

        inputs = concrete_program.inputs
        outputs = concrete_program.outputs
        startup_program = concrete_program.startup_program

        scope = paddle.static.global_scope()

        @switch_to_static_graph
        def append_backward_desc():
            program = concrete_program.main_program

            # backward with optimizer to add backward graph to program
            backward.gradients_with_optimizer(program, ipu_strategy._optimizer)

            # initialize backward parameters
            exe = paddle.static.Executor(paddle.CPUPlace())
            startup_program = paddle.static.default_startup_program()
            exe.run(startup_program)

            return program

        if ipu_strategy.enable_fp16:
            class_instance.to(dtype="float16")

        # copy the bias and filters
        for param_or_buffer in concrete_program.parameters:
            param_or_buffer_tensor = scope.var(
619 620
                param_or_buffer.name
            ).get_tensor()
621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642
            src_tensor = param_or_buffer.value().get_tensor()
            param_or_buffer_tensor._share_data_with(src_tensor)

        # TODO(czr): feed and fetch list needs to consider more type
        if class_instance:
            feed_list = [elem.name for elem in inputs[1:] if elem is not None]
        else:
            feed_list = [elem.name for elem in inputs if elem is not None]
        fetch_list = [elem.name for elem in outputs]

        if ipu_strategy.is_training:
            concrete_program.main_program = append_backward_desc()
            # copy optimizer parameters
            optimizer = ipu_strategy._optimizer
            for k, v in optimizer._accumulators.items():
                for param_name, var_tmp in v.items():
                    var = optimizer.helper.create_global_variable(
                        name=var_tmp.name,
                        persistable=True,
                        dtype=var_tmp.dtype,
                        type=var_tmp.type,
                        shape=var_tmp.shape,
643 644
                        belong_to_optimizer=True,
                    )
645 646 647
                    device = optimizer._get_device_for_param(param_name)
                    with device_guard(device):
                        optimizer.helper.set_variable_initializer(
648 649
                            var, initializer=Constant(value=0.0)
                        )
650
                    param_or_lr_tensor = scope.find_var(
651 652
                        var_tmp.name
                    ).get_tensor()
653 654 655 656 657 658 659 660 661 662 663 664
                    optim_tensor = var.value().get_tensor()
                    param_or_lr_tensor._share_data_with(optim_tensor)
                    optimizer._accumulators[k][param_name] = var

        @switch_to_static_graph
        def func_compile():
            if ipu_strategy.enable_fp16:
                amp_list = paddle.static.amp.CustomOpLists()
                amp_list.unsupported_list = {"cumsum"}
                to_fp16_var_names = paddle.static.amp.cast_model_to_fp16(
                    concrete_program.main_program,
                    amp_list,
665 666
                    use_fp16_guard=False,
                )
667 668 669
                paddle.static.amp.cast_parameters_to_fp16(
                    paddle.CPUPlace(),
                    concrete_program.main_program,
670 671 672 673 674 675 676 677
                    to_fp16_var_names=to_fp16_var_names,
                )

            program = IpuCompiledProgram(
                concrete_program.main_program,
                ipu_strategy=ipu_strategy,
                scope=scope,
            ).compile(feed_list, fetch_list)
678 679 680 681 682 683 684 685
            return program

        main_program = func_compile()
        concrete_program.main_program = main_program
        return concrete_program

    @staticmethod
    def patch_program_cache(ipu_strategy):
686
        """Monkey patch ProgramCache discriptor to support dynamic2static in IPU.
687 688 689 690 691 692 693

        Args:
            ipu_strategy: The ipu_strategy used in dynamic graph.

        Returns:
            None
        """
694 695 696 697 698 699
        from ..fluid.dygraph.dygraph_to_static.program_translator import (
            ProgramCache,
        )
        from ..fluid.dygraph.dygraph_to_static.program_translator import (
            CacheKey,
        )
700
        from ..fluid.dygraph.dygraph_to_static import logging_utils
701 702 703 704 705 706
        from ..fluid.dygraph.dygraph_to_static.program_translator import (
            MAX_TRACED_PROGRAM_COUNT,
        )
        from ..fluid.dygraph.dygraph_to_static.partial_program import (
            partial_program_from,
        )
707 708 709 710 711 712

        old_getter = ProgramCache.__getitem__

        def patch_getter(self, item):
            if not isinstance(item, CacheKey):
                raise ValueError(
713 714 715
                    'type(item) should be CacheKey, but received %s'
                    % type(item).__name__
                )
716 717 718 719 720
            item_id = hash(item)
            self._recent_key = item_id
            if item_id not in self._caches or ipu_strategy.need_compile:
                if item_id in self._caches:
                    logging_utils.warn(
721 722
                        "ipu_strategy chances detected. Please sync weights."
                    )
723 724 725
                if self._caches and not ipu_strategy.need_compile:
                    logging_utils.warn(
                        "dynamic2static on IPU doesn't support mutiple caches. Please make sure"
726 727
                        "dynamic inputs is not used."
                    )
728 729
                concrete_program, _ = self._build_once(item)
                concrete_program = IpuDynamicPatcher.convert_concrete_program(
730 731
                    ipu_strategy, concrete_program, item.class_instance
                )
732

733 734 735 736
                self._caches[item_id] = (
                    concrete_program,
                    partial_program_from(concrete_program),
                )
737 738 739 740 741
                # Note: raise warnings if number of traced program is more than `max_tracing_count`
                current_tracing_count = len(self._caches)
                if current_tracing_count > MAX_TRACED_PROGRAM_COUNT:
                    logging_utils.warn(
                        "Current traced program number: {} > `max_tracing_count`:{}. Too much cached programs will bring expensive overhead. "
742 743 744 745
                        "The reason may be: (1) passing tensors with different shapes, (2) passing python objects instead of tensors.".format(
                            current_tracing_count, MAX_TRACED_PROGRAM_COUNT
                        )
                    )
746 747 748 749 750

            return self._caches[item_id]

        setattr(ProgramCache, '__getitem__', patch_getter)
        IpuDynamicPatcher.patcher_cache.append(
751 752
            [ProgramCache, '__getitem__', old_getter]
        )
753 754 755 756

    @staticmethod
    def patch_lr_scheduler(ipu_strategy):
        from paddle.optimizer.lr import LRScheduler
757

758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779
        # For IPU dynamic graph usage, lr_var is not synced in executor as static mode do.
        # Manually set lr to ipu_strategy to update the lr.
        old_step = LRScheduler.step

        def patch_step(self, epoch=None):
            old_step(self, epoch)
            ipu_strategy.set_options({"lr": self.last_lr})

        setattr(LRScheduler, 'step', patch_step)
        IpuDynamicPatcher.patcher_cache.append([LRScheduler, 'step', old_step])

    @staticmethod
    def register_patch(ipu_strategy):
        IpuDynamicPatcher.patch_program_cache(ipu_strategy)
        IpuDynamicPatcher.patch_lr_scheduler(ipu_strategy)

    @staticmethod
    def release_patch():
        for module, key, attr in IpuDynamicPatcher.patcher_cache:
            setattr(module, key, attr)


J
jianghaicheng 已提交
780 781 782 783 784 785 786 787 788
class IpuStrategy(object):
    """
    Help users precisely control the graph building in :code:`paddle.static.IpuCompiledProgram` .

    Returns:
        The IpuStrategy instance.

    Examples:
        .. code-block:: python
789

J
jianghaicheng 已提交
790 791 792 793 794 795
            # required: ipu

            import paddle
            import paddle.static as static

            paddle.enable_static()
796

J
jianghaicheng 已提交
797 798 799 800 801 802
            ipu_strategy = static.IpuStrategy()
    """

    def __init__(self):
        if core.is_compiled_with_ipu():
            self._ipu_strategy = core.IpuStrategy()
803 804 805 806 807
            default_options = {
                'location_optimizer': {
                    'on_chip': 0,
                    'use_replicated_tensor_sharding': 1,
                },  # set optimizer location
808 809
                'accumulation_and_replication_reduction_type': 1,  # popart::ReductionType::Mean
                'mean_accumulation_and_replication_reduction_strategy': 1,  # popart::MeanReductionStrategy::Post
810 811 812 813
            }
            self._ipu_strategy.set_options(default_options)
            self.has_custom_ops = False
            self.custom_op_names = []
814
            self.need_compile = True
J
jianghaicheng 已提交
815 816 817 818
        else:
            raise RuntimeError(
                "Can not use IpuStrategy in non IPU compiled environment, please re-compile with WITH_IPU=ON."
            )
819
        from paddle import in_dynamic_mode
820

821 822 823 824 825 826 827 828 829 830
        if in_dynamic_mode():
            self.register_patch()

    def register_patch(self):
        """
        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
831

832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848
                # required: ipu

                import paddle
                import paddle.static as static

                ipu_strategy = static.IpuStrategy()

                ipu_strategy.register_patch()
        """
        IpuDynamicPatcher.register_patch(self)

    def release_patch(self):
        """
        Release the registered IPU functions.

        Examples:
            .. code-block:: python
849

850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866
                # required: ipu

                import paddle
                import paddle.static as static

                ipu_strategy = static.IpuStrategy()

                ipu_strategy.release_patch()
        """
        IpuDynamicPatcher.release_patch()

    def set_optimizer(self, optimizer):
        """
        Set optimizer to ipu_strategy in dynamic mode.

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

868 869 870 871 872
          Returns:
              None.

          Examples:
              .. code-block:: python
873

874 875 876 877 878 879 880 881 882 883 884 885
                  # 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)
        """
        from paddle import in_dynamic_mode
886

887 888 889 890 891 892 893 894 895 896 897 898 899
        if in_dynamic_mode():
            self._optimizer = optimizer
            optimizer_attrs = self.parse_optimizer(optimizer)
            self._ipu_strategy.set_options(optimizer_attrs)
        else:
            raise RuntimeError("Only needs to set optimizer in dynamic mode.")

    def parse_optimizer(self, optimizer):
        """
        Parse optimizer attributes for IPU dynamic to static support. Currently only support parse lr.

          Args:
              optimizer (Optimizer): Optimizer to be parsed.
900

901 902 903 904 905
          Returns:
              Dict.

          Examples:
              .. code-block:: python
906

907 908 909 910 911 912 913 914 915 916 917 918 919 920
                  # 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)
        """

        def get_lr():
            from paddle.optimizer.lr import LRScheduler
921

922 923 924 925 926 927 928 929 930 931
            if isinstance(optimizer._learning_rate, float):
                return {"lr": optimizer._learning_rate}
            elif isinstance(optimizer._learning_rate, LRScheduler):
                return {"lr": optimizer._learning_rate()}

        attr_fn = [get_lr]
        optimizer_attrs = {"is_dynamic": True}
        for fn in attr_fn:
            optimizer_attrs.update(fn())
        return optimizer_attrs
J
jianghaicheng 已提交
932

933 934 935 936 937 938 939
    def set_graph_config(
        self,
        num_ipus=1,
        is_training=True,
        micro_batch_size=1,
        enable_manual_shard=False,
    ):
J
jianghaicheng 已提交
940 941 942 943 944 945 946 947
        """
        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.
948 949 950
            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.

J
jianghaicheng 已提交
951 952 953 954 955
        Returns:
            None.

        Examples:
            .. code-block:: python
956

J
jianghaicheng 已提交
957 958 959 960 961 962
                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()
963

J
jianghaicheng 已提交
964
                ipu_strategy = static.IpuStrategy()
965
                ipu_strategy.set_graph_config(num_ipus=1,
J
jianghaicheng 已提交
966
                                            is_training=True,
A
Allen Guo 已提交
967
                                            micro_batch_size=1,
968
                                            enable_manual_shard=False)
J
jianghaicheng 已提交
969
        """
970
        if num_ipus == 1 and enable_manual_shard:
J
jianghaicheng 已提交
971 972 973
            raise RuntimeError(
                "Only if num_ipus > 1, enable_manual_shard is able to be set True."
            )
974 975 976
        options = {
            'num_ipus': num_ipus,
            'is_training': is_training,
A
Allen Guo 已提交
977
            'micro_batch_size': micro_batch_size,
978 979 980 981
            'enable_manual_shard': enable_manual_shard,
        }
        self.set_options(options)

982 983 984 985 986 987 988
    def set_pipelining_config(
        self,
        enable_pipelining=False,
        batches_per_step=1,
        enable_gradient_accumulation=False,
        accumulation_factor=1,
    ):
J
jianghaicheng 已提交
989 990 991 992
        """
        Set pipelining configuration to the IpuStrategy instance. Used to optimize the throughput performance.

        Args:
993
            enable_pipelining (bool, optional): Enable data pipelining between subgraphs. Only if enable_manual_shard=True, enable_pipelining is able to be set True.
J
jianghaicheng 已提交
994 995 996
                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.
A
Allen Guo 已提交
997
            enable_gradient_accumulation (bool, optional): Enable to accumulate gradients before updating the weights in training mode. Only if enable_pipelining=True,
998 999
                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
J
jianghaicheng 已提交
1000
                before applying the varUpdate. Default 1, which means disable the accumulation.
1001

J
jianghaicheng 已提交
1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
1016 1017
                ipu_strategy.set_pipelining_config(enable_pipelining=False,
                                                    batches_per_step=1,
A
Allen Guo 已提交
1018
                                                    enable_gradient_accumulation=False,
1019
                                                    accumulation_factor=1)
J
jianghaicheng 已提交
1020
        """
1021 1022
        enable_manual_shard = self.get_option('enable_manual_shard')
        if not enable_manual_shard and enable_pipelining:
J
jianghaicheng 已提交
1023 1024 1025
            raise RuntimeError(
                "Only if enable_manual_shard=True, enable_pipelining is able to be set True."
            )
1026 1027 1028
        options = {
            'enable_pipelining': enable_pipelining,
            'batches_per_step': batches_per_step,
A
Allen Guo 已提交
1029
            'enable_gradient_accumulation': enable_gradient_accumulation,
1030 1031 1032 1033 1034
            'accumulation_factor': accumulation_factor,
        }
        self.set_options(options)

    def set_precision_config(self, enable_fp16=False):
J
jianghaicheng 已提交
1035 1036 1037 1038 1039
        """
        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.
1040

J
jianghaicheng 已提交
1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
1055 1056
                ipu_strategy.set_precision_config(enable_fp16=False)
        """
1057 1058 1059
        options = {
            'enable_fp16': enable_fp16,
        }
1060 1061
        self.set_options(options)

1062 1063 1064
    def add_custom_op(
        self, paddle_op, popart_op=None, domain='custom.ops', version=1
    ):
J
jianghaicheng 已提交
1065
        """
1066
        Add a mapping to use popart custom ops running on the IPU.
J
jianghaicheng 已提交
1067

1068 1069
        Args:
            paddle_op(str): the name of custom op in paddle.
J
jianghaicheng 已提交
1070

1071
            popart_op(str): the name of custom op in popart.
J
jianghaicheng 已提交
1072

1073
            domain(str): domain name of custom op in popart.
J
jianghaicheng 已提交
1074

1075
            version(int): version of custom op in popart.
1076

1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091
        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')
J
jianghaicheng 已提交
1092
        """
1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106
        if popart_op is None:
            popart_op = paddle_op
        custom_op = {
            'paddle_op': paddle_op,
            'popart_op': popart_op,
            'domain': domain,
            'version': version,
        }
        self.set_options({'custom_op': custom_op})
        self.custom_op_names.append(paddle_op)
        if not self.has_custom_ops:
            self.has_custom_ops = True

    def set_options(self, options):
J
jianghaicheng 已提交
1107
        """
1108
        Set options from dict.
J
jianghaicheng 已提交
1109

1110 1111
        Args:
            options(dict): dict of options.
1112

1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128
        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)
J
jianghaicheng 已提交
1129
        """
1130
        self._ipu_strategy.set_options(options)
1131 1132 1133 1134
        # check whether to recompile program with updated ipu options.
        recompile_white_list = {'lr'}
        if options.keys() - recompile_white_list:
            self.need_compile = True
J
jianghaicheng 已提交
1135

1136
    def get_option(self, option):
J
jianghaicheng 已提交
1137
        """
1138 1139 1140 1141
        Get option.

        Args:
            option(str): name of option.
1142

1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157
        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')
J
jianghaicheng 已提交
1158
        """
1159
        return self._ipu_strategy.get_option(option)['value']
J
jianghaicheng 已提交
1160

A
Allen Guo 已提交
1161 1162 1163 1164 1165 1166
    def enable_pattern(self, pattern):
        """
        Enable PopART pattern to optimize the graph.

        Args:
            pattern(string): the name of the pattern.
1167

A
Allen Guo 已提交
1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191
        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")
        """
        self._ipu_strategy.enable_pattern(pattern)

    def disable_pattern(self, pattern):
        """
        Disable PopART pattern.

        Args:
            pattern(string): the name of the pattern.
1192

A
Allen Guo 已提交
1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210
        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")
        """
        self._ipu_strategy.disable_pattern(pattern)

J
jianghaicheng 已提交
1211
    @property
1212
    def num_ipus(self):
J
jianghaicheng 已提交
1213
        """
1214
        Get the number of IPU devices from IpuStrategy instance.
J
jianghaicheng 已提交
1215
        """
1216
        return self.get_option('num_ipus')
J
jianghaicheng 已提交
1217 1218

    @property
1219
    def is_training(self):
J
jianghaicheng 已提交
1220
        """
1221
        Get the boolean of training or inference from IpuStrategy instance.
J
jianghaicheng 已提交
1222
        """
1223
        return self.get_option('is_training')
J
jianghaicheng 已提交
1224 1225

    @property
1226
    def enable_pipelining(self):
J
jianghaicheng 已提交
1227
        """
1228
        Get the boolean of enable pipelining or not from IpuStrategy instance.
J
jianghaicheng 已提交
1229
        """
1230
        return self.get_option('enable_pipelining')
J
jianghaicheng 已提交
1231 1232 1233 1234 1235 1236

    @property
    def enable_fp16(self):
        """
        Get the boolean of float16 mode or not from IpuStrategy instance.
        """
1237
        return self.get_option('enable_fp16')
J
jianghaicheng 已提交
1238 1239 1240 1241 1242 1243 1244 1245 1246


class IpuCompiledProgram(object):
    """
    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`
1247
            to be executed. Default is None, which means the program will be set to
J
jianghaicheng 已提交
1248 1249
            the default program :code:`paddle.static.default_main_program()` .
        scope(Scope, optional): The scope used to run this program, you can switch
1250
            it to different scope. Default is None, which means use the global
J
jianghaicheng 已提交
1251 1252 1253
            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.
1254
            Default is None, which means build the program based on the default `ipu_strategy`.
J
jianghaicheng 已提交
1255 1256 1257 1258 1259 1260

    Returns:
        IpuCompiledProgram

    Example:
        .. code-block:: python
1261

J
jianghaicheng 已提交
1262 1263 1264 1265 1266 1267 1268 1269 1270 1271
            # 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()
1272

J
jianghaicheng 已提交
1273
            ipu_strategy = static.IpuStrategy()
A
Allen Guo 已提交
1274 1275
            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)
1276
            ipu_strategy.set_precision_config(enable_fp16=False)
1277

J
jianghaicheng 已提交
1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289
            ipu_compiled_program = static.IpuCompiledProgram(
                main_prog,
                ipu_strategy=ipu_strategy)
    """

    def __init__(self, program=None, scope=None, ipu_strategy=None):
        if not core.is_compiled_with_ipu():
            raise ValueError(
                "Can not use this function since PaddlePaddle is not compiled with IPU"
            )

        if program is None:
1290
            program = framework.default_main_program()
J
jianghaicheng 已提交
1291 1292 1293

        if not isinstance(program, framework.Program):
            raise TypeError(
1294 1295 1296
                "The type of program is wrong, expected Program, but got %s"
                % type(program)
            )
J
jianghaicheng 已提交
1297 1298 1299 1300 1301 1302 1303

        self._program = program
        self._compiled = False

        if scope is not None:
            self._scope = scope
        else:
1304 1305
            # import here to avoiding confused
            import paddle
1306

J
jianghaicheng 已提交
1307 1308 1309
            self._scope = paddle.static.global_scope()

        if ipu_strategy is not None:
1310
            self._ipu_strategy = ipu_strategy
J
jianghaicheng 已提交
1311
        else:
1312
            self._ipu_strategy = IpuStrategy()
J
jianghaicheng 已提交
1313

1314 1315 1316 1317 1318 1319
        if ipu_strategy.has_custom_ops:
            self._custom_op_names = set(ipu_strategy.custom_op_names)
        else:
            self._custom_op_names = ()

        self._backend = core.IpuBackend.get_instance()
J
jianghaicheng 已提交
1320 1321 1322 1323 1324

    def compile(self, feed_list, fetch_list):
        """
        This interface is used to compile the input Program to a program
        to run the model on the ipu.
1325

J
jianghaicheng 已提交
1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336
        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
1337

J
jianghaicheng 已提交
1338
                # required: ipu
1339

J
jianghaicheng 已提交
1340 1341
                import paddle
                import paddle.static as static
1342

J
jianghaicheng 已提交
1343
                paddle.enable_static()
1344

J
jianghaicheng 已提交
1345 1346 1347 1348 1349
                a = static.data(name='data', shape=[None, 1], dtype='int32')
                b = a + 1
                main_prog = static.default_main_program()

                ipu_strategy = static.IpuStrategy()
A
Allen Guo 已提交
1350 1351
                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)
1352
                ipu_strategy.set_precision_config(enable_fp16=False)
1353

J
jianghaicheng 已提交
1354 1355 1356 1357
                program = static.IpuCompiledProgram(
                    main_prog,
                    ipu_strategy=ipu_strategy).compile([a.name], [b.name])
        """
1358 1359 1360
        self._backend.set_scope(self._scope)
        self._backend.set_ipu_strategy(self._ipu_strategy._ipu_strategy)

J
jianghaicheng 已提交
1361 1362 1363 1364 1365
        # feed and fetch doesn't have corresponding popart op, so we rm both here
        global_block = self._program.global_block()
        need_to_remove_op_index = []
        for i, op in enumerate(global_block.ops):
            op.desc.set_is_target(False)
1366
            if op.type == 'feed' or op.type == 'fetch':
J
jianghaicheng 已提交
1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378
                need_to_remove_op_index.append(i)

        for index in need_to_remove_op_index[::-1]:
            global_block._remove_op(index)

        for var in ['feed', 'fetch']:
            if global_block.has_var(var):
                global_block._remove_var(var)

        self._program.desc.flush()
        self._graph = core.Graph(self._program.desc)

1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414
        if self._ipu_strategy.is_training:
            passes = [
                'optimizer_extract_pass',
                'optimizer_state_align_pass',
            ]
            for pass_name in passes:
                a_pass = core.get_pass(pass_name)
                a_pass.apply(self._graph)

        passes = [
            'forward_graph_extract_pass',
            'infer_shape_pass',
            'avg_shard_pass',
            'delete_scale_op_pass',
        ]
        for pass_name in passes:
            a_pass = core.get_pass(pass_name)
            if pass_name == 'infer_shape_pass':
                a_pass.set('feed_list', feed_list)
            a_pass.apply(self._graph)

        a_pass = core.get_pass('popart_canonicalization_pass')
        if self._custom_op_names:
            a_pass.set('custom_ops', self._custom_op_names)
        a_pass.apply(self._graph)

        passes = [
            'ipu_inplace_pass',
            'ipu_graph_builder_pass',
            'ipu_runtime_replacer_pass',
        ]
        for pass_name in passes:
            a_pass = core.get_pass(pass_name)
            a_pass.set('feed_list', feed_list)
            a_pass.set('fetch_list', fetch_list)
            a_pass.apply(self._graph)
J
jianghaicheng 已提交
1415 1416 1417 1418 1419 1420 1421 1422 1423 1424 1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443

        convert_pass = core.get_pass('graph_to_program_pass')
        desc = core.ProgramDesc()
        convert_pass.set_not_owned('program', desc)
        convert_pass.apply(self._graph)
        program = framework.Program._construct_from_desc(desc)

        if hasattr(self._program, 'lr_sheduler'):
            # how to share var between two different block ?
            lr_var_name = self._program.lr_sheduler._var_name

            program.lr_sheduler = self._program.lr_sheduler
            # Program.clone will clone lr_sheduler, so i set lr_var as
            # lr_sheduler attribute
            global_block = self._program.global_block()
            program.lr_sheduler.lr_var = global_block.vars[lr_var_name]

        # with popart, we need to support batches_per_step, what means
        # the shape of feed_var and feed_tensor(maybe numpy array) will
        # mismatch, so we set need_check_feed to False. Thus we can avoid
        # modify logic of run.
        program_global_block = program.global_block()
        for feed_name in feed_list:
            feed_var = program_global_block.var(feed_name)
            feed_var.desc.set_need_check_feed(False)

        if not hasattr(program, 'org_program'):
            program.org_program = self._program

1444 1445
        self._ipu_strategy.need_compile = False

J
jianghaicheng 已提交
1446
        return program