compiler.py 52.1 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


108
class CompiledProgram:
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
class IpuDynamicPatcher:
571 572 573 574 575 576 577 578 579 580
    """
    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
        from paddle.jit.dy2static.program_translator import (
695
            CacheKey,
696
            ProgramCache,
697 698
            MAX_TRACED_PROGRAM_COUNT,
        )
699
        from paddle.jit.dy2static import logging_utils
700
        from paddle.jit.dy2static.partial_program import (
701 702
            partial_program_from,
        )
703 704 705 706 707 708

        old_getter = ProgramCache.__getitem__

        def patch_getter(self, item):
            if not isinstance(item, CacheKey):
                raise ValueError(
709 710 711
                    'type(item) should be CacheKey, but received %s'
                    % type(item).__name__
                )
712 713 714 715 716
            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(
717 718
                        "ipu_strategy chances detected. Please sync weights."
                    )
719 720 721
                if self._caches and not ipu_strategy.need_compile:
                    logging_utils.warn(
                        "dynamic2static on IPU doesn't support mutiple caches. Please make sure"
722 723
                        "dynamic inputs is not used."
                    )
724 725
                concrete_program, _ = self._build_once(item)
                concrete_program = IpuDynamicPatcher.convert_concrete_program(
726 727
                    ipu_strategy, concrete_program, item.class_instance
                )
728

729 730 731 732
                self._caches[item_id] = (
                    concrete_program,
                    partial_program_from(concrete_program),
                )
733 734 735 736 737
                # 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. "
738 739 740 741
                        "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
                        )
                    )
742 743 744 745 746

            return self._caches[item_id]

        setattr(ProgramCache, '__getitem__', patch_getter)
        IpuDynamicPatcher.patcher_cache.append(
747 748
            [ProgramCache, '__getitem__', old_getter]
        )
749 750 751 752

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

754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775
        # 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)


776
class IpuStrategy:
J
jianghaicheng 已提交
777 778 779 780 781 782 783 784
    """
    Help users precisely control the graph building in :code:`paddle.static.IpuCompiledProgram` .

    Returns:
        The IpuStrategy instance.

    Examples:
        .. code-block:: python
785

J
jianghaicheng 已提交
786 787 788 789 790 791
            # required: ipu

            import paddle
            import paddle.static as static

            paddle.enable_static()
792

J
jianghaicheng 已提交
793 794 795 796 797 798
            ipu_strategy = static.IpuStrategy()
    """

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

817 818 819 820 821 822 823 824 825 826
        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
827

828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844
                # 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
845

846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862
                # 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.
863

864 865 866 867 868
          Returns:
              None.

          Examples:
              .. code-block:: python
869

870 871 872 873 874 875 876 877 878 879 880 881
                  # 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
882

883 884 885 886 887 888 889 890 891 892 893 894 895
        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.
896

897 898 899 900 901
          Returns:
              Dict.

          Examples:
              .. code-block:: python
902

903 904 905 906 907 908 909 910 911 912 913 914 915 916
                  # 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
917

918 919 920 921 922 923 924 925 926 927
            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 已提交
928

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

        Examples:
            .. code-block:: python
952

J
jianghaicheng 已提交
953 954 955 956 957 958
                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()
959

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

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

        Args:
989
            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 已提交
990 991 992
                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 已提交
993
            enable_gradient_accumulation (bool, optional): Enable to accumulate gradients before updating the weights in training mode. Only if enable_pipelining=True,
994 995
                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 已提交
996
                before applying the varUpdate. Default 1, which means disable the accumulation.
997

J
jianghaicheng 已提交
998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

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

    def set_precision_config(self, enable_fp16=False):
J
jianghaicheng 已提交
1031 1032 1033 1034 1035
        """
        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.
1036

J
jianghaicheng 已提交
1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050
        Returns:
            None.

        Examples:
            .. code-block:: python

                # required: ipu

                import paddle
                import paddle.static as static

                paddle.enable_static()

                ipu_strategy = static.IpuStrategy()
1051 1052
                ipu_strategy.set_precision_config(enable_fp16=False)
        """
1053 1054 1055
        options = {
            'enable_fp16': enable_fp16,
        }
1056 1057
        self.set_options(options)

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

1064 1065
        Args:
            paddle_op(str): the name of custom op in paddle.
J
jianghaicheng 已提交
1066

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

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

1071
            version(int): version of custom op in popart.
1072

1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087
        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 已提交
1088
        """
1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102
        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 已提交
1103
        """
1104
        Set options from dict.
J
jianghaicheng 已提交
1105

1106 1107
        Args:
            options(dict): dict of options.
1108

1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124
        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 已提交
1125
        """
1126
        self._ipu_strategy.set_options(options)
1127 1128 1129 1130
        # 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 已提交
1131

1132
    def get_option(self, option):
J
jianghaicheng 已提交
1133
        """
1134 1135 1136 1137
        Get option.

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

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

A
Allen Guo 已提交
1157 1158 1159 1160 1161 1162
    def enable_pattern(self, pattern):
        """
        Enable PopART pattern to optimize the graph.

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

A
Allen Guo 已提交
1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187
        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.
1188

A
Allen Guo 已提交
1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206
        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 已提交
1207
    @property
1208
    def num_ipus(self):
J
jianghaicheng 已提交
1209
        """
1210
        Get the number of IPU devices from IpuStrategy instance.
J
jianghaicheng 已提交
1211
        """
1212
        return self.get_option('num_ipus')
J
jianghaicheng 已提交
1213 1214

    @property
1215
    def is_training(self):
J
jianghaicheng 已提交
1216
        """
1217
        Get the boolean of training or inference from IpuStrategy instance.
J
jianghaicheng 已提交
1218
        """
1219
        return self.get_option('is_training')
J
jianghaicheng 已提交
1220 1221

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

    @property
    def enable_fp16(self):
        """
        Get the boolean of float16 mode or not from IpuStrategy instance.
        """
1233
        return self.get_option('enable_fp16')
J
jianghaicheng 已提交
1234 1235


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

    Returns:
        IpuCompiledProgram

    Example:
        .. code-block:: python
1257

J
jianghaicheng 已提交
1258 1259 1260 1261 1262 1263 1264 1265 1266 1267
            # 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()
1268

J
jianghaicheng 已提交
1269
            ipu_strategy = static.IpuStrategy()
A
Allen Guo 已提交
1270 1271
            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)
1272
            ipu_strategy.set_precision_config(enable_fp16=False)
1273

J
jianghaicheng 已提交
1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285
            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:
1286
            program = framework.default_main_program()
J
jianghaicheng 已提交
1287 1288 1289

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

        self._program = program
        self._compiled = False

        if scope is not None:
            self._scope = scope
        else:
1300 1301
            # import here to avoiding confused
            import paddle
1302

J
jianghaicheng 已提交
1303 1304 1305
            self._scope = paddle.static.global_scope()

        if ipu_strategy is not None:
1306
            self._ipu_strategy = ipu_strategy
J
jianghaicheng 已提交
1307
        else:
1308
            self._ipu_strategy = IpuStrategy()
J
jianghaicheng 已提交
1309

1310 1311 1312 1313 1314 1315
        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 已提交
1316 1317 1318 1319 1320

    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.
1321

J
jianghaicheng 已提交
1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332
        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
1333

J
jianghaicheng 已提交
1334
                # required: ipu
1335

J
jianghaicheng 已提交
1336 1337
                import paddle
                import paddle.static as static
1338

J
jianghaicheng 已提交
1339
                paddle.enable_static()
1340

J
jianghaicheng 已提交
1341 1342 1343 1344 1345
                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 已提交
1346 1347
                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)
1348
                ipu_strategy.set_precision_config(enable_fp16=False)
1349

J
jianghaicheng 已提交
1350 1351 1352 1353
                program = static.IpuCompiledProgram(
                    main_prog,
                    ipu_strategy=ipu_strategy).compile([a.name], [b.name])
        """
1354 1355 1356
        self._backend.set_scope(self._scope)
        self._backend.set_ipu_strategy(self._ipu_strategy._ipu_strategy)

J
jianghaicheng 已提交
1357 1358 1359 1360 1361
        # 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)
1362
            if op.type == 'feed' or op.type == 'fetch':
J
jianghaicheng 已提交
1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374
                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)

1375 1376 1377 1378 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
        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 已提交
1411 1412 1413 1414 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

        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

1440 1441
        self._ipu_strategy.need_compile = False

J
jianghaicheng 已提交
1442
        return program