memory_optimization_transpiler.py 10.6 KB
Newer Older
D
dzhwinter 已提交
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# 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.

15 16 17 18 19
from collections import defaultdict
import framework
from framework import Program, default_main_program, Parameter, Variable
import backward
from backward import _rename_arg_
20 21 22 23 24 25 26 27 28 29 30
from . import core

dtype_to_size = {
    core.DataType.FP16: 2,
    core.DataType.FP32: 4,
    core.DataType.FP64: 8,
    core.DataType.INT16: 2,
    core.DataType.INT32: 4,
    core.DataType.INT64: 8,
    core.DataType.BOOL: 1
}
31 32 33


class ControlFlowGraph(object):
34
    def __init__(self, Program, ops, forward_num, skip_opt):
35
        self._program = Program
36 37 38 39
        self._ops = ops
        self._forward_num = forward_num
        self._successors = defaultdict(set)
        self._presuccessors = defaultdict(set)
40 41 42 43
        self._uses = defaultdict(set)
        self._defs = defaultdict(set)
        self._live_in = defaultdict(set)
        self._live_out = defaultdict(set)
44
        self._skip_opt = skip_opt
45 46 47 48 49 50

    def _add_connections(self, connections):
        for node1, node2 in connections:
            self._add(node1, node2)

    def _add(self, node1, node2):
51 52
        self._successors[node1].add(node2)
        self._presuccessors[node2].add(node1)
53 54

    def _build_graph(self):
55
        self.op_size = len(self._ops)
56 57 58
        op_node_connections = [(i, i + 1) for i in range(self.op_size - 1)]
        self._add_connections(op_node_connections)
        for i in range(self.op_size):
59 60
            self._uses[i].update(self._ops[i].input_arg_names())
            self._defs[i].update(self._ops[i].output_arg_names())
61

62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
    def _update_graph(self, old_name, new_name, begin_idx=0):
        for i in range(begin_idx, self.op_size):
            if old_name in self._uses[i]:
                self._uses[i].remove(old_name)
                self._uses[i].add(new_name)
            if old_name in self._defs[i]:
                self._defs[i].remove(old_name)
                self._defs[i].add(new_name)
            if old_name in self._live_in[i]:
                self._live_in[i].remove(old_name)
                self._live_out[i].add(new_name)
            if old_name in self._live_out[i]:
                self._live_out[i].remove(old_name)
                self._live_out[i].add(new_name)

77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
    def _reach_fixed_point(self, live_in, live_out):
        if len(live_in) != len(self._live_in):
            return False
        if len(live_out) != len(self._live_out):
            return False
        for i in range(self.op_size):
            if live_in[i] != self._live_in[i]:
                return False
        for i in range(self.op_size):
            if live_out[i] != self._live_out[i]:
                return False
        return True

    def _dataflow_analyze(self):
        self._build_graph()
        live_in = defaultdict(set)
        live_out = defaultdict(set)
        while True:
Q
QI JUN 已提交
95
            for i in range(self.op_size, 0, -1):
96 97
                live_in[i] = set(self._live_in[i])
                live_out[i] = set(self._live_out[i])
98
                for s in self._successors[i]:
99
                    self._live_out[i] |= self._live_in[s]
Q
QI JUN 已提交
100 101
                self._live_in[i] = self._uses[i] | (
                    self._live_out[i] - self._defs[i])
102 103 104 105 106 107 108
            if self._reach_fixed_point(live_in, live_out):
                break

    def _get_diff(self, a, b):
        u = a & b
        return a - u, b - u

109 110 111 112 113 114 115 116 117 118 119 120
    def _has_var(self, block_desc, var_name, is_forward):
        if is_forward:
            return block_desc.has_var(str(var_name))
        else:
            return block_desc.has_var_recursive(str(var_name))

    def _find_var(self, block_desc, var_name, is_forward):
        if is_forward:
            return block_desc.find_var(str(var_name))
        else:
            return block_desc.find_var_recursive(str(var_name))

121
    def memory_optimize(self):
122 123 124 125 126 127 128 129 130 131 132
        def check_var_validity(block_desc, x, is_forward):
            if str(x) == "@EMPTY@":
                return False
            if not self._has_var(block_desc, x, is_forward):
                return False
            if self._find_var(block_desc, x, is_forward).persistable():
                return False
            if self._find_var(
                    block_desc, x,
                    is_forward).type() != core.VarDesc.VarType.LOD_TENSOR:
                return False
133 134 135 136
            if x in self._skip_opt:
                return False
            if not self._find_var(block_desc, x, is_forward).shape():
                return False
137 138
            return True

139 140 141 142
        self._build_graph()
        self._dataflow_analyze()
        self.pool = []
        for i in range(self.op_size):
143 144 145 146 147
            op = self._ops[i]
            if op.type() == "while" or op.type() == "while_grad":
                continue
            block_desc = op.block()
            is_forward = i < self._forward_num
148
            if self.pool:
149 150 151 152 153 154 155
                defs_can_optimize = filter(
                    lambda x: check_var_validity(block_desc, x, is_forward),
                    self._defs[i])
                out_pair = [
                    (x, self._find_var(block_desc, x, is_forward).shape())
                    for x in defs_can_optimize
                ]
156
                for x, x_shape in out_pair:
157 158 159
                    # If x is both in uses and defs, it can not be optimized!
                    if x in self._uses[i]:
                        continue
160 161 162 163 164 165 166 167 168
                    for index, cache_pair in enumerate(self.pool):
                        cache_var = cache_pair[0]
                        cache_shape = cache_pair[1]
                        if x_shape == cache_shape:
                            if self._has_var(block_desc, cache_var, is_forward):
                                x_dtype = self._find_var(block_desc, x,
                                                         is_forward).dtype()
                                cache_dtype = self._find_var(
                                    block_desc, cache_var, is_forward).dtype()
169 170 171
                                # TODO(qijun): actually, we should compare dtype_to_size[x_dtype]
                                # and dtype_to_size[cache_dtype]
                                if x_dtype == cache_dtype:
172 173 174 175 176 177
                                    print(("Hit Cache !!!! cache pool index "
                                           "is %d, var name is %s, "
                                           "cached var name is %s, "
                                           "var shape is %s ") %
                                          (index, x, cache_var,
                                           str(cache_shape)))
178
                                    self.pool.pop(index)
179 180
                                    if x == cache_var:
                                        break
181
                                    _rename_arg_(
182 183 184 185
                                        self._ops, x, cache_var, begin_idx=i)
                                    self._program.block(block_desc.id).var(
                                        str(x)).desc = self._find_var(
                                            block_desc, cache_var, is_forward)
186 187 188
                                    self._update_graph(
                                        x, cache_var, begin_idx=i)
                                    break
189 190 191 192

            in_diff, out_diff = self._get_diff(self._live_in[i],
                                               self._live_out[i])
            can_optimize = filter(
193
                lambda x: check_var_validity(block_desc, x, is_forward),
194 195 196
                in_diff)
            if can_optimize:
                for var_name in can_optimize:
197 198 199 200 201 202 203 204 205 206
                    self.pool.append((var_name, self._find_var(
                        block_desc, var_name, is_forward).shape()))


def get_cfgs(input_program):
    ops_list = []
    pdesc = input_program.get_desc()
    block_desc = pdesc.block(0)
    op_size = block_desc.op_size()
    # Get global block ops
207 208
    ops_list.append(
        ([block_desc.op(i) for i in range(op_size)], op_size, set()))
209 210 211

    while_sub_block_ids = []
    while_grad_sub_block_ids = []
212
    while_block_id_pair = []
213
    while_op_dict = {}
214 215 216 217 218

    for i in range(op_size):
        op = block_desc.op(i)
        if op.type() == "while":
            while_sub_block_ids.append(op.attr("sub_block").id)
219
            while_op_dict[op.attr("sub_block").id] = op
220 221
        elif op.type() == "while_grad":
            while_grad_sub_block_ids.append(op.attr("sub_block").id)
222
            while_op_dict[op.attr("sub_block").id] = op
223 224 225 226 227

    # Find while/while_grad block pair
    for grad_id in while_grad_sub_block_ids:
        parent_id = pdesc.block(grad_id).parent
        if parent_id in while_sub_block_ids:
228
            while_block_id_pair.append((parent_id, grad_id))
229 230 231
            while_sub_block_ids.remove(parent_id)

    # Get while/while_grad block ops
232
    for parent_id, grad_id in while_block_id_pair:
233 234 235 236 237 238 239 240 241 242 243
        while_block_ops = []
        while_block = pdesc.block(parent_id)
        while_block_op_size = while_block.op_size()
        for i in range(while_block_op_size):
            while_block_ops.append(while_block.op(i))

        while_grad_block = pdesc.block(grad_id)
        while_grad_block_op_size = while_grad_block.op_size()
        for i in range(while_grad_block_op_size):
            while_block_ops.append(while_grad_block.op(i))

244 245 246 247
        while_op_output = set()
        while_op_output.update(while_op_dict[parent_id].output_arg_names())
        while_op_output.update(while_op_dict[grad_id].output_arg_names())

248
        ops_list.append((while_block_ops, while_block_op_size, while_op_output))
249 250 251 252 253 254 255 256 257

    # Process rest while block ops
    for parent_id in while_sub_block_ids:
        while_block_ops = []
        while_block = pdesc.block(parent_id)
        while_block_op_size = while_block.op_size()
        for i in range(while_block_op_size):
            while_block_ops.append(while_block.op(i))

258 259 260 261
        while_op_output = set()
        while_op_output.update(while_op_dict[parent_id].output_arg_names())

        ops_list.append((while_block_ops, while_block_op_size, while_op_output))
262

263 264 265 266
    cfgs = [
        ControlFlowGraph(input_program, ops, forward_num, skip_opt)
        for ops, forward_num, skip_opt in ops_list
    ]
267
    return cfgs
268 269 270


def memory_optimize(input_program):
271 272 273
    cfgs = get_cfgs(input_program)
    for cfg in cfgs:
        cfg.memory_optimize()