ir.py 21.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
# Copyright (c) 2021 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 copy
16 17 18 19 20 21
import inspect
from os import path
import paddle
from . import core, unique_name
from .framework import _apply_pass, OpProtoHolder

22
from .proto import framework_pb2
23 24 25 26 27 28
try:
    from .proto import pass_desc_pb2
except ModuleNotFoundError:
    import sys
    sys.path.append(path.join(path.dirname(__file__), 'proto'))
    from .proto import pass_desc_pb2
29 30 31 32 33 34 35 36 37 38


def get_data_vars(program):
    data_vars = []
    for var_name, var in program.global_block().vars.items():
        if var.is_data:
            data_vars.append(var_name)
    return data_vars


39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
def _update_grad_persistable(main_program):
    grad_merge_attr_name = "grad_merge_cond_name"
    op_role_var_attr_name = core.op_proto_and_checker_maker.kOpRoleVarAttrName()
    has_grad_merge = False
    has_persistable_grad_var = False
    grad_vars = []
    for block_id in range(main_program.num_blocks):
        block = main_program.block(block_id)
        for op in block.ops:
            if grad_merge_attr_name in op.attr_names:
                has_grad_merge = True

            if op_role_var_attr_name not in op.attr_names:
                continue

            p_g = op.attr(op_role_var_attr_name)
            for g in p_g[1::2]:
                g_var = block._find_var_recursive(g)
                if g_var is None:
                    continue
                grad_vars.append(g_var)
                if g_var.persistable:
                    has_persistable_grad_var = True

    if has_grad_merge and has_persistable_grad_var:
        for g_var in grad_vars:
            g_var.persistable = True


68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85
def apply_build_strategy(main_program, startup_program, build_strategy,
                         pass_attrs):
    def update_attr(attrs, attr_types, name, value, typ=None):
        if name not in attrs:
            attrs[name] = value
        if typ:
            attr_types[name] = typ

    def apply_pass(name):
        attrs = dict(pass_attrs)
        attr_types = {}
        update_attr(attrs, attr_types, "nranks", 1, "size_t")
        update_attr(attrs, attr_types, "use_cuda", False, "bool")
        # TODO(zjl): how to skip fetch variables ?
        update_attr(attrs, attr_types, "mem_opt_skip_vars",
                    get_data_vars(main_program), "list[str]")
        _apply_pass(main_program, startup_program, name, attrs, attr_types)

86
    _update_grad_persistable(main_program)
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107
    use_cuda = pass_attrs.get("use_cuda", False)
    build_strategy = build_strategy._copy()
    if build_strategy.sync_batch_norm:
        apply_pass("sync_batch_norm_pass")
        build_strategy.sync_batch_norm = False
    if build_strategy.fuse_relu_depthwise_conv and use_cuda:
        apply_pass("fuse_relu_depthwise_conv_pass")
        build_strategy.fuse_relu_depthwise_conv = False
    if build_strategy.fuse_bn_act_ops and use_cuda:
        apply_pass("fuse_bn_act_pass")
        build_strategy.fuse_bn_act_ops = False
    if build_strategy.fuse_bn_add_act_ops and use_cuda:
        apply_pass("fuse_bn_add_act_pass")
        build_strategy.fuse_bn_add_act_ops = False
    if build_strategy.enable_auto_fusion and use_cuda:
        apply_pass("fusion_group_pass")
        build_strategy.enable_auto_fusion = False
    if build_strategy.fuse_elewise_add_act_ops:
        apply_pass("fuse_elewise_add_act_pass")
        build_strategy.fuse_elewise_add_act_ops = False
    if build_strategy.fuse_all_optimizer_ops:
108 109 110 111 112 113
        apply_pass([
            "coalesce_grad_tensor_pass",
            "fuse_adam_op_pass",
            "fuse_sgd_op_pass",
            "fuse_momentum_op_pass",
        ])
114 115 116 117 118 119 120 121 122 123 124 125 126 127
        build_strategy.fuse_all_optimizer_ops = False
    # TODO(zjl): support fuse all reduce ops
    if build_strategy.cache_runtime_context:
        apply_pass("runtime_context_cache_pass")
        build_strategy.cache_runtime_context = False
    if build_strategy.enable_addto and use_cuda:
        # NOTE: how to get fetch vars to skip memory optimization?  
        apply_pass("inplace_addto_op_pass")
        build_strategy.enable_addto = False
    if build_strategy.enable_inplace:
        apply_pass("buffer_shared_inplace_pass")
        build_strategy.enable_inplace = False
    build_strategy._clear_finalized()
    return build_strategy
128 129 130


class RegisterPassHelper(object):
131 132
    _register_helpers = list()

133 134 135
    def __init__(self, pass_pairs, pass_type=str(), input_specs=dict()):
        self._pass_type = pass_type
        self._pass_pairs = pass_pairs
136 137
        self._input_specs = input_specs
        RegisterPassHelper._register_helpers.append(self)
138 139 140 141 142 143 144 145

    def _get_args_from_func(self, func):
        args = list()
        arg_specs = inspect.getfullargspec(func)
        for arg_name in arg_specs.args:
            input_spec = self._input_specs.get(arg_name)
            if isinstance(input_spec, paddle.static.InputSpec):
                args.append(
146
                    PassDesc.VarHelper(arg_name, input_spec.shape,
147 148 149 150
                                       input_spec.dtype))
            elif isinstance(input_spec, paddle.ParamAttr):
                args.append(paddle.ParamAttr(arg_name))
            else:
151
                args.append(PassDesc.VarHelper(arg_name, [-1]))
152 153
        return args

154 155
    def _prune_program_desc(self, ops):
        for op_desc in ops:
156 157 158 159
            default_attrs = core.get_op_attrs_default_value(
                paddle.compat.to_bytes(op_desc.type))
            remove_attrs = list()
            for attr in op_desc.attrs:
160
                # attr must not in
161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
                if attr.name not in [
                        "op_namescope", "op_callstack", "op_device"
                ]:
                    attr_list_fields = attr.ListFields()
                    # attr format must be: name, type, value
                    if len(attr_list_fields) == 3:
                        attr_value = attr.ListFields()[-1][-1]
                        default_attr_value = default_attrs.get(attr.name)
                        # value must not default
                        if default_attr_value != attr_value:
                            continue
                remove_attrs.append(attr)
            for attr in remove_attrs:
                op_desc.attrs.remove(attr)

176
    def _func_to_program_desc(self, func, ops):
177 178 179 180 181
        vars = list()
        program = paddle.static.Program()
        startup_program = paddle.static.Program()
        with paddle.static.program_guard(program, startup_program):
            args = self._get_args_from_func(func)
182
            vars.extend(args)
183 184 185 186 187
            outs = func(*args)
            if not isinstance(outs, (list, tuple)):
                outs = [outs]
            for out in outs:
                if isinstance(out, PassDesc.OpHelper):
188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238
                    op_outs = out.Outputs()
                    if len(op_outs) != 1:
                        raise ValueError(
                            "Operator '{}' has multiple outputs, please specify one output variable.".
                            format(out._type))
                    for op_out in op_outs.values():
                        vars.extend(op_out)
                else:
                    vars.append(out)
        block_desc = program.current_block().desc
        for i in range(block_desc.op_size()):
            ops.add().ParseFromString(block_desc.op(i).serialize_to_string())
        self._prune_program_desc(ops)
        return vars, program.current_block().ops

    def _convert_vars_to_pass_desc(self, patterns, replaces, desc):
        for (pattern, replace) in zip(patterns, replaces):
            # Convert maps of inputs and outputs.
            var_map = desc.var_maps.add()
            var_map.pattern_var = pattern.name
            var_map.replace_var = replace.name
            conditions = desc.var_attr_conditions
            # Convert shape condition.
            if pattern.name in self._input_specs:
                condition = conditions.add()
                pattern.Attr("shape")._to_pass_desc_attr(condition.attr)
                condition.condition_value.name = ""
                condition.condition_value.type = framework_pb2.AttrType.LONGS
                condition.condition_value.longs.extend(pattern.shape)
                condition.type = pass_desc_pb2.PassDesc.ConditionType.kEQ
            # Convert attr conditions.
            if PassDesc.VarHelper == pattern.__class__:
                for attr in pattern._attrs.values():
                    if attr._condition is not None:
                        conditions.append(attr._condition)
                    conditions.extend(
                        [e._condition for e in attr._elements if e._condition])

    def _convert_ops_to_pass_desc(self, patterns, replaces, desc):
        for replace in replaces:
            if isinstance(replace, PassDesc.OpHelper):
                for attr in replace._attrs.values():
                    # Convert attr maps.
                    mapped = attr._mapped
                    if inspect.isfunction(mapped):
                        mapped = mapped(patterns)
                    attr_map = desc.op_attr_maps.add()
                    mapped._to_pass_desc_attr(attr_map.pattern_attr)
                    attr._to_pass_desc_attr(attr_map.replace_attr)
                    if mapped._operation is not None:
                        attr_map.operation.CopyFrom(mapped._operation)
239 240 241 242 243 244 245

    def SerializeMultiPassDesc(self):
        switch_static_mode = paddle.in_dynamic_mode()
        if switch_static_mode:
            paddle.enable_static()
        multi_pass_desc = pass_desc_pb2.MultiPassDesc()
        multi_pass_desc.pass_type = self._pass_type
246 247
        # Traverse all pass pairs and convert them to PassDesc data.
        # Here need to add cache in the future. 
248 249
        for (pattern, replace) in self._pass_pairs:
            pass_desc = multi_pass_desc.pass_descs.add()
250 251 252 253 254 255 256 257
            # Convert ProgramDescs of pattern and replace subgraphs.
            pattern_vars, pattern_ops = self._func_to_program_desc(
                pattern, pass_desc.pattern)
            replace_vars, replace_ops = self._func_to_program_desc(
                replace, pass_desc.replace)
            self._convert_vars_to_pass_desc(pattern_vars, replace_vars,
                                            pass_desc)
            self._convert_ops_to_pass_desc(pattern_ops, replace_ops, pass_desc)
258 259 260 261 262 263 264
        if switch_static_mode:
            paddle.disable_static()
        return multi_pass_desc.SerializeToString()


class PassDesc(object):
    class AttrHelper(object):
265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
        def __init__(self, obj, name, element_index=None):
            self._obj = obj
            self._name = name
            self._operation_type = None
            self._element_index = element_index
            self._elements = list()
            self._operation = None
            self._condition = None
            self._mapped = None

        def __getitem__(self, index):
            element = PassDesc.AttrHelper(
                self._obj, self._name, element_index=index)
            self._elements.append(element)
            return element

        def _to_pass_desc_attr(self, pass_desc_attr):
            if isinstance(self._obj, PassDesc.VarHelper):
                pass_desc_attr.role = pass_desc_pb2.PassDesc.RoleType.kVariable
                pass_desc_attr.var_name = self._obj.name
            else:
                pass_desc_attr.role = pass_desc_pb2.PassDesc.RoleType.kOperator
                pass_desc_attr.op_index = self._obj._index
            pass_desc_attr.name = self._name
            if self._operation_type is not None:
                pass_desc_attr.operation = self._operation_type
            if self._element_index is not None:
                pass_desc_attr.element_index = self._element_index

        def _to_op_desc_attr(self, value, op_desc_attr):
            op_desc_attr.name = ""
            if isinstance(value, int):
                op_desc_attr.type = framework_pb2.AttrType.INT
                op_desc_attr.i = value
            else:
                raise NotImplementedError("Unimplemented transform operation.")

        def _clone_with_operation(self, type, value=None):
            attr = PassDesc.AttrHelper(self._obj, self._name,
                                       self._element_index)
            self._elements.append(attr)
            if value is None:
                attr._operation_type = type
                return attr
            operation = pass_desc_pb2.PassDesc.Operation()
            operation.type = type
            if isinstance(value, PassDesc.AttrHelper):
                value._to_pass_desc_attr(operation.attr)
            else:
                self._to_op_desc_attr(value, operation.value)
            attr._operation = operation
            attr._operation_type = self._operation_type
            return attr

        def __sub__(self, value):
            return self._clone_with_operation(
                pass_desc_pb2.PassDesc.OperationType.kSub, value)

        def __add__(self, value):
            return self._clone_with_operation(
                pass_desc_pb2.PassDesc.OperationType.kAdd, value)

        def Size(self):
            return self._clone_with_operation(
                pass_desc_pb2.PassDesc.OperationType.kSize)

        def _set_with_condition(self, type, value):
            condition = pass_desc_pb2.PassDesc.AttrCondition()
            self._to_pass_desc_attr(condition.attr)
            condition.type = type
            if isinstance(value, PassDesc.AttrHelper):
                value._to_pass_desc_attr(condition.condition_attr)
            else:
                self._to_op_desc_attr(value, condition.condition_value)
            self._condition = condition

        def EQ(self, value):
            self._set_with_condition(pass_desc_pb2.PassDesc.ConditionType.kEQ,
                                     value)

        def MappedPattern(self, var=None, op=None, index=0, name=None):
            if all([var, op]):
                raise ValueError("Only mapped one of which var or op.")

            def mapped_var(pattern_ops):
                raise NotImplementedError(
                    "Mapping to variable is not implemented.")

            def mapped_op(pattern_ops):
                ops = [o for o in pattern_ops if o._type == op]
                if len(ops) <= index:
                    raise ValueError(
                        "Index '{}' of operator '{}' is incorrect.".format(
                            index, op))
                return PassDesc.AttrHelper(ops[index], name)

            self._mapped = mapped_op if var is None else mapped_var

    class VarHelper(paddle.static.Variable):
        def __init__(self, *args, **kwargs):
            block = paddle.static.default_main_program().current_block()
            self._var = paddle.static.data(*args, **kwargs)
            self._attrs = dict()

        def __getattr__(self, name):
            return getattr(self._var, name)

        def Attr(self, name):
            attr = self._attrs.get(name)
            if attr is None:
                attr = PassDesc.AttrHelper(self, name)
                self._attrs[name] = attr
            return attr
378 379 380 381 382 383 384 385 386 387 388

    class OpHelper(object):
        def __init__(self, type=None):
            self._type = type

        def __getattr__(self, name):
            op = PassDesc.OpHelper(name)
            op.Init()
            return op

        def __call__(self, *args, **kwargs):
389 390 391
            if len(args) > 0:
                raise ValueError(
                    "Each input argument needs to specify a parameter name.")
392
            for (in_name, in_args) in kwargs.items():
393 394 395 396 397
                op_input = self._inputs.get(in_name)
                if op_input is None:
                    raise ValueError(
                        "Operator '{}' does not have input named '{}'.".format(
                            self._type, in_name))
398 399 400 401 402 403 404 405 406
                if isinstance(in_args, (list, tuple)):
                    if len(in_args) == 0:
                        raise ValueError(
                            "Input '{}' of operator '{}' cannot be empty.".
                            format(in_name, self._type))
                else:
                    in_args = [in_args]
                for in_arg in in_args:
                    if isinstance(in_arg, PassDesc.OpHelper):
407 408 409 410 411 412 413
                        op_outs = in_arg.Outputs()
                        if len(op_outs) != 1:
                            raise ValueError(
                                "The size of outputs of operator '{}' is not equal 1, please specify one output variable.".
                                format(in_arg._type))
                        for op_out in op_outs.values():
                            op_input.extend(op_out)
414
                    else:
415 416 417 418 419 420 421
                        op_input.append(in_arg)
                self._desc.set_input(in_name, [i.name for i in op_input])
            block = paddle.static.default_main_program().current_block()
            for out_name, op_output in self._outputs.items():
                op_output_name = unique_name.generate(self._type)
                op_output.append(block.create_var(name=op_output_name))
                self._desc.set_output(out_name, [op_output_name])
422 423 424 425
            return self

        def Init(self):
            block = paddle.static.default_main_program().current_block()
426 427
            self._proto = OpProtoHolder.instance().op_proto_map.get(self._type)
            if self._proto is None:
428 429 430
                raise AttributeError(
                    "type object 'OpHelper' has no attribute '{}'".format(
                        self._type))
431 432 433 434 435 436
            self._index = len(block.ops)
            self._desc = block.desc.append_op()
            self._desc.set_type(self._type)
            self._attrs = dict()
            self._inputs = {i.name: list() for i in self._proto.inputs}
            self._outputs = {o.name: list() for o in self._proto.outputs}
437 438 439 440
            block.ops.append(self)

        def Attr(self, name):
            attr = self._attrs.get(name)
441 442 443
            if attr is None:
                attr = PassDesc.AttrHelper(self, name)
                self._attrs[name] = attr
444 445 446
            return attr

        def SetAttr(self, name, value):
447 448 449 450
            if isinstance(value, PassDesc.AttrHelper):
                self.Attr(name)._mapped = value
            else:
                self._desc._set_attr(name, value)
451

452 453 454 455 456 457 458
        def Output(self, name):
            output = self._outputs.get(name)
            if output is None:
                raise ValueError(
                    "Operator '{}' does not have output named '{}'.".format(
                        self._type, name))
            return output
459 460

        def Outputs(self):
461
            return self._outputs
462 463 464 465

    OP = OpHelper()


466
def RegisterPass(function=None, input_specs=dict()):
467 468 469 470 471 472 473 474
    """
    The function decorator of Register Pass. Decorator @RegisterPass handles
    the function and register it into a core.Pass instance. Use name of function
    as Pass type.

    Args:
        function (callable): The function with return of callable pair(s) that
            represents the pattern subgraph and the replace subgraph.
475
        input_specs (dict[str, InputSpec]): Dict of InputSpec to specific the shape/dtype
476 477 478
            information of Tensor. Some operators limit the shape and dtype of datas when
            create subgraph with Paddle APIs. So user need specify InputSpec of data to
            ensure create a correctly subgraph. Of course, this argument is not limited to
479
            matching subgraph. The default is dict().
480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520

    Returns:
        callables: Callable pair(s).

    Examples:
        .. code-block:: python

        import paddle
        from paddle.fluid.ir import RegisterPass

        @RegisterPass
        def multi_add_to_addn():
            def pattern(x, y, z):
                return paddle.add(paddle.add(x, y), z)
            def replace(x, y, z):
                return paddle.add_n([x, y, z])
            return pattern, replace
    """

    def _is_pass_pair(check_pair):
        if isinstance(check_pair, (list, tuple)):
            if len(check_pair) == 2:
                if all(map(inspect.isfunction, check_pair)):
                    return True
        return False

    def decorated(python_func):
        pass_type = python_func.__name__
        signature = inspect.signature(python_func)
        if len(signature.parameters) > 0:
            raise NotImplementedError(
                "Pass function with parameter is not supported now.")
        elif len(signature.parameters) == 0:
            pass_pairs = python_func()
            if _is_pass_pair(pass_pairs):
                pass_pairs = [pass_pairs]
            elif not all(map(_is_pass_pair, pass_pairs)):
                raise ValueError(
                    "Return value of Pass function must be (callable, callable)."
                )
            helper = RegisterPassHelper(pass_pairs, pass_type, input_specs)
521
            core.register_pass(pass_type, helper.SerializeMultiPassDesc)
522 523 524 525 526 527
        return python_func

    if inspect.isfunction(function):
        return decorated(function)

    return decorated