def delete_transpose_of_correlation(self):
nodes_need_to_delete_for_special_ops = []
nodes_need_to_insert_aster_special_ops = []
nodes_need_to_merge_for_special_ops = []
for node in self._dev_graph.nodes:
if node.op.type == NNDCT_OP.MEAN and not node.node_attr(
node.op.AttrName.KEEP_DIMS
) and self._dev_graph.parents(node):
pn = self._dev_graph.parents(node)[0]
if pn.in_tensors and _is_permute_op(
pn) and self._dev_graph.parents(pn):
gpn = self._dev_graph.parents(pn)[0]
if gpn.op.type in [
NNDCT_OP.CORRELATION1D_ELEMWISE,
NNDCT_OP.CORRELATION2D_ELEMWISE
] and node.out_tensors[0].ndim and gpn.out_tensors[
0].ndim == 5 and node.out_tensors[0].ndim == 4:
nodes_need_to_delete_for_special_ops.append(pn)
node.transpose_in_order = tuple(
_find_swim_order(5))
node.transpose_out_order = tuple(
_find_swim_order(4))
special_ops_fn[node.op.type](
node, node.transpose_in_order)
nodes_need_to_insert_aster_special_ops.append(node)
index = 0
for node in nodes_need_to_insert_aster_special_ops:
cn = self._dev_graph.children(node)[0]
node_name = "_".join([node.name, "sink_transpose", f"{index}"])
op = base_op.Permute(NNDCT_OP.PERMUTE)
new_node = Node(node_name,
op=op,
dtype=node.dtype,
in_quant_part=node.in_quant_part)
new_node.set_node_attr(new_node.op.AttrName.ORDER,
tuple(_find_sink_order(4)))
self._dev_graph.insert_node_between_nodes(new_node, node, cn)
nodes_need_to_merge_for_special_ops.append(new_node)
index += 1
for node in nodes_need_to_delete_for_special_ops:
self._dev_graph.remove_node(node)
source_nodes = []
for node in self._dev_graph.nodes:
if not node.in_tensors:
source_nodes.append(node)
transpose_group = []
reserverd_nodes = []
visited = []
for source in nodes_need_to_merge_for_special_ops:
merge_father_and_child(source, visited, transpose_group,
reserverd_nodes)
nodes_need_to_merge_for_special_ops = [
node for node in nodes_need_to_merge_for_special_ops
if node not in reserverd_nodes
]
for node in reserverd_nodes:
order = node.node_attr(node.op.AttrName.ORDER)
keep_order = True
if any([index != dim for index, dim in enumerate(order)]):
keep_order = False
if keep_order:
nodes_need_to_merge_for_special_ops.append(node)
for node in nodes_need_to_merge_for_special_ops:
self._dev_graph.remove_node(node)
merge_brothers(reserverd_nodes)
def layout_tranform(self):
"""layout_transform TORCH(NCHW) -> XIR(NHWC)"""
custom2xir = GLOBAL_MAP.get_ele(NNDCT_KEYS.CUSTOM_TO_XIR_LIST)
if custom2xir is None:
custom2xir = []
def _find_swim_order(ndim):
return {
2: [0, 1],
3: [0, 2, 1],
4: [0, 2, 3, 1],
5: [0, 3, 4, 2, 1]
}[ndim]
def _find_sink_order(ndim):
return {
2: [0, 1],
3: [0, 2, 1],
4: [0, 3, 1, 2],
5: [0, 4, 3, 1, 2]
}[ndim]
def _is_dim_transparent(node):
return node.in_tensors[0].ndim and node.out_tensors[
0].ndim and node.in_tensors[0].ndim == node.out_tensors[0].ndim
def _is_shape_transparent(node):
return node.in_tensors[0].shape and node.out_tensors[
0].shape and node.in_tensors[0].shape == node.out_tensors[
0].shape
def _have_special_layout(node):
return node.out_tensors[0].ndim and node.out_tensors[0].ndim >= 3
def _is_custom_op(node):
return isinstance(
node.op, base_op.CustomOp) and node.op.type not in custom2xir
def _is_permute_op(node):
return isinstance(node.op, base_op.Permute)
def _is_terminate_op(node):
return node.op.type == NNDCT_OP.RETURN
implicit_ops = [
NNDCT_OP.CONV2D, NNDCT_OP.DEPTHWISE_CONV2D,
NNDCT_OP.DEPTHWISE_CONVTRANSPOSE2D, NNDCT_OP.CONVTRANSPOSE2D,
NNDCT_OP.MAX_POOL, NNDCT_OP.AVG_POOL, NNDCT_OP.ADAPTIVEAVGPOOL2D,
NNDCT_OP.INTERPOLATE, NNDCT_OP.UP_SAMPLING, NNDCT_OP.RESIZE,
NNDCT_OP.BATCH_NORM, NNDCT_OP.MAX_POOL1D, NNDCT_OP.CONV1D,
NNDCT_OP.CONV3D, NNDCT_OP.DEPTHWISE_CONV3D,
NNDCT_OP.DEPTHWISE_CONVTRANSPOSE3D, NNDCT_OP.CONVTRANSPOSE3D,
NNDCT_OP.PIXEL_SHUFFLE, NNDCT_OP.PIXEL_UNSHUFFLE,
NNDCT_OP.RESIZE_3D, NNDCT_OP.RESIZE_NEAREST_3D, NNDCT_OP.REORG,
NNDCT_OP.CORRELATION1D_ELEMWISE, NNDCT_OP.CORRELATION2D_ELEMWISE,
NNDCT_OP.COST_VOLUME
]
special_ops_fn = {
NNDCT_OP.RESHAPE: shape_attr_transform_fn,
NNDCT_OP.CONCAT: axis_attr_transform_fn,
NNDCT_OP.STRIDED_SLICE: slice_attr_transform_fn,
NNDCT_OP.SUM: reduce_op_attr_transform_fn,
NNDCT_OP.MAX: reduce_op_attr_transform_fn,
NNDCT_OP.MEAN: reduce_op_attr_transform_fn,
NNDCT_OP.SHAPE: axis_attr_transform_fn,
NNDCT_OP.SOFTMAX: axis_attr_transform_fn,
NNDCT_OP.ZEROS: shape_attr_transform_fn,
}
# collect insert point for transpose
insert_pos = []
for node in self._dev_graph.nodes:
if node.op.type in implicit_ops:
insert_pos.append(node)
swim_transpose = defaultdict(list)
swim_in_transpose = defaultdict(list)
sink_transpose = defaultdict(list)
for node in insert_pos:
tranpose_out_order = tuple(
_find_swim_order(node.out_tensors[0].ndim))
swim_transpose[tranpose_out_order].append(node)
tranpose_in_order = tuple(_find_swim_order(
node.in_tensors[0].ndim))
swim_in_transpose[node] = tranpose_in_order
tranpose_out_order = tuple(
_find_sink_order(node.out_tensors[0].ndim))
sink_transpose[tranpose_out_order].append(node)
nodes_need_to_remove = []
transpose_insert_between_swim = defaultdict(list)
visited = []
# swim_transpose_order, nodes = next(iter(swim_transpose.items()))
for swim_transpose_order, nodes in swim_transpose.items():
for insert_node in nodes:
q = deque()
q.append(insert_node)
visited.append(insert_node)
insert_node.transpose_out_order = swim_transpose_order
insert_node.transpose_in_order = swim_in_transpose[insert_node]
while len(q) > 0:
node = q.popleft()
for pn in self._dev_graph.parents(node):
if pn not in visited:
if not _have_special_layout(
pn) or pn.op.type in implicit_ops:
continue
elif pn.op.type in [
NNDCT_OP.INPUT, NNDCT_OP.QUANT_STUB,
NNDCT_OP.CONST, NNDCT_OP.ZEROS
] or _is_dim_transparent(pn) and (
not _is_permute_op(pn)) and (
not _is_custom_op(pn)):
pn.transpose_out_order = node.transpose_in_order
pn.transpose_in_order = pn.transpose_out_order
if pn.op.type in special_ops_fn:
special_ops_fn[pn.op.type](
pn, pn.transpose_out_order)
q.append(pn)
visited.append(pn)
else:
# pn.transpose_out_order = [0, 2, 3, 1]
transpose_insert_between_swim[
swim_transpose_order].append((pn, node))
index = 0
for transpose_order, node_pairs in transpose_insert_between_swim.items(
):
for pn, cn in node_pairs:
node_name = "_".join([pn.name, "swim_transpose", f"{index}"])
op = base_op.Permute(NNDCT_OP.PERMUTE)
new_node = Node(node_name,
op=op,
dtype=pn.dtype,
in_quant_part=pn.in_quant_part)
new_node.set_node_attr(new_node.op.AttrName.ORDER,
list(transpose_order))
self._dev_graph.insert_node_between_nodes(new_node, pn, cn)
nodes_need_to_remove.append(new_node)
index += 1
if transpose_insert_between_swim:
self._dev_graph.reconnect_nodes()
# debug
# print("#####swim######")
# for node in self._dev_graph.nodes:
# print(node.op.type, node.name, node.transpose_out_order)
transpose_insert_between_sink = defaultdict(list)
visited = []
for node in self._dev_graph.nodes:
if node.transpose_out_order:
nodes = sink_transpose[tuple(
_find_sink_order(len(node.transpose_out_order)))]
if node not in nodes:
nodes.append(node)
for sink_transpose_order, nodes in sink_transpose.items():
for insert_node in nodes:
if insert_node not in visited:
q = deque()
q.append(insert_node)
visited.append(insert_node)
while len(q) > 0:
node = q.popleft()
for cn in self._dev_graph.children(node):
if cn not in visited:
if cn.op.type in implicit_ops or _is_terminate_op(
cn):
continue
elif cn.op.type == NNDCT_OP.SHAPE:
visited.append(cn)
if node.transpose_out_order:
special_ops_fn[cn.op.type](
cn, node.transpose_out_order)
continue
elif cn.transpose_out_order:
q.append(cn)
visited.append(cn)
elif _is_dim_transparent(cn) and (
not _is_permute_op(cn)) and (
not _is_custom_op(cn)):
cn.transpose_in_order = node.transpose_out_order
cn.transpose_out_order = cn.transpose_in_order
q.append(cn)
visited.append(cn)
if cn.op.type in special_ops_fn:
special_ops_fn[cn.op.type](
cn, cn.transpose_out_order)
else:
transpose_insert_between_sink[
sink_transpose_order].append(
(node, cn))
index = 0
for transpose_order, node_pairs in transpose_insert_between_sink.items(
):
for pn, cn in node_pairs:
node_name = "_".join([pn.name, "sink_transpose", f"{index}"])
op = base_op.Permute(NNDCT_OP.PERMUTE)
new_node = Node(node_name,
op=op,
dtype=pn.dtype,
in_quant_part=cn.in_quant_part)
new_node.set_node_attr(new_node.op.AttrName.ORDER,
list(transpose_order))
self._dev_graph.insert_node_between_nodes(new_node, pn, cn)
nodes_need_to_remove.append(new_node)
index += 1
if transpose_insert_between_sink:
self._dev_graph.reconnect_nodes()
# debug
# print("#####sink######")
# for node in self._dev_graph.nodes:
# print(node.op.type, node.name, node.transpose_out_order)
neighbor_broadcast = {}
for node in self._dev_graph.nodes:
if len(node.in_nodes) <= 1 or node in implicit_ops:
continue
if all([
node.transpose_out_order is None
for node in self._dev_graph.parents(node)
]) or all([
node.transpose_out_order is not None
for node in self._dev_graph.parents(node)
]):
continue
#if node.out_tensors[0].dtype != "float32":
# continue
transpose_order = None
for pn in self._dev_graph.parents(node):
transpose_order = pn.transpose_out_order
if transpose_order is not None:
break
neighbor_broadcast[node] = transpose_order
have_neighbors = False
for node, transpose_order in neighbor_broadcast.items():
index = 0
for pn in self._dev_graph.parents(node):
if pn.transpose_out_order is None and pn.out_tensors[
0].ndim and node.out_tensors[0].ndim and pn.out_tensors[
0].ndim == node.out_tensors[0].ndim:
# pn.transpose_out_order = node.transpose_out_order
node_name = "_".join(
[node.name, "neighbor_transpose", f"{index}"])
op = base_op.Permute(NNDCT_OP.PERMUTE)
new_node = Node(node_name,
op=op,
dtype=node.dtype,
in_quant_part=pn.in_quant_part)
new_node.set_node_attr(new_node.op.AttrName.ORDER,
list(transpose_order))
self._dev_graph.insert_node_between_nodes(
new_node, pn, node)
index += 1
nodes_need_to_remove.append(new_node)
have_neighbors = True
if have_neighbors:
self._dev_graph.reconnect_nodes()
# Debug
# print("####neightbor######")
# for node in self._dev_graph.nodes:
# print(node.op.type, node.name, node.transpose_out_order)
# remove consecutive transpose
def merge_father_and_child(node, visited, transpose_group,
reserverd_nodes):
visited.append(node)
if _is_permute_op(node):
if node.out_nodes and all([
_is_permute_op(cn)
for cn in self._dev_graph.children(node)
]):
transpose_group.append(node)
else:
transpose_group.append(node)
order = []
reserved_trans = None
for trans in transpose_group:
if trans not in nodes_need_to_remove:
reserved_trans = trans
if not order:
order = trans.node_attr(trans.op.AttrName.ORDER)
else:
new_order = len(order) * [None]
tmp_order = trans.node_attr(
trans.op.AttrName.ORDER)
for i in range(len(order)):
t_i = tmp_order[i]
new_order[i] = order[t_i]
order = new_order
if reserved_trans is None:
reserved_trans = transpose_group[-1]
reserved_trans.set_node_attr(
reserved_trans.op.AttrName.ORDER, order)
reserverd_nodes.append(reserved_trans)
transpose_group.clear()
for cn in self._dev_graph.children(node):
if cn not in visited:
merge_father_and_child(cn, visited, transpose_group,
reserverd_nodes)
def merge_brothers(reserverd_nodes):
remove_nodes = []
for node in self._dev_graph.nodes:
if len(node.out_nodes) > 1 and all([
_is_permute_op(cn)
for cn in self._dev_graph.children(node)
]):
need_merge = True
order = None
for trans_node in self._dev_graph.children(node):
if order is not None:
if order != trans_node.node_attr(
trans_node.op.AttrName.ORDER):
need_merge = False
break
else:
order = trans_node.node_attr(
trans_node.op.AttrName.ORDER)
if need_merge:
reserverd_node = None
for trans_node in self._dev_graph.children(node):
if trans_node not in nodes_need_to_remove:
reserverd_node = trans_node
if reserverd_node is None:
reserverd_node = self._dev_graph.children(node)[0]
for trans_node in self._dev_graph.children(node):
if trans_node is not reserverd_node and trans_node in reserverd_nodes:
remove_nodes.append(trans_node)
out_tensor = trans_node.out_tensors[0]
out_tensor.replace_uses_with(
reserverd_node.out_tensors[0])
for node in remove_nodes:
node.destroy()
if remove_nodes:
self._dev_graph.reconnect_nodes()
source_nodes = []
for node in self._dev_graph.nodes:
if not node.in_tensors:
source_nodes.append(node)
transpose_group = []
reserverd_nodes = []
visited = []
for source in source_nodes:
merge_father_and_child(source, visited, transpose_group,
reserverd_nodes)
nodes_need_to_remove = [
node for node in nodes_need_to_remove
if node not in reserverd_nodes
]
for node in reserverd_nodes:
order = node.node_attr(node.op.AttrName.ORDER)
keep_order = True
if any([index != dim for index, dim in enumerate(order)]):
keep_order = False
if keep_order:
nodes_need_to_remove.append(node)
for node in nodes_need_to_remove:
self._dev_graph.remove_node(node)
merge_brothers(reserverd_nodes)
# debug
# print("#####finalize######")
# for node in self._dev_graph.nodes:
# print(node.op.type, node.name, node.transpose_out_order)
def delete_transpose_of_correlation(self):
nodes_need_to_delete_for_special_ops = []
nodes_need_to_insert_aster_special_ops = []
nodes_need_to_merge_for_special_ops = []
for node in self._dev_graph.nodes:
if node.op.type == NNDCT_OP.MEAN and not node.node_attr(
node.op.AttrName.KEEP_DIMS
) and self._dev_graph.parents(node):
pn = self._dev_graph.parents(node)[0]
if pn.in_tensors and _is_permute_op(
pn) and self._dev_graph.parents(pn):
gpn = self._dev_graph.parents(pn)[0]
if gpn.op.type in [
NNDCT_OP.CORRELATION1D_ELEMWISE,
NNDCT_OP.CORRELATION2D_ELEMWISE
] and node.out_tensors[0].ndim and gpn.out_tensors[
0].ndim == 5 and node.out_tensors[0].ndim == 4:
nodes_need_to_delete_for_special_ops.append(pn)
node.transpose_in_order = tuple(
_find_swim_order(5))
node.transpose_out_order = tuple(
_find_swim_order(4))
special_ops_fn[node.op.type](
node, node.transpose_in_order)
nodes_need_to_insert_aster_special_ops.append(node)
index = 0
for node in nodes_need_to_insert_aster_special_ops:
cn = self._dev_graph.children(node)[0]
node_name = "_".join([node.name, "sink_transpose", f"{index}"])
op = base_op.Permute(NNDCT_OP.PERMUTE)
new_node = Node(node_name,
op=op,
dtype=node.dtype,
in_quant_part=node.in_quant_part)
new_node.set_node_attr(new_node.op.AttrName.ORDER,
tuple(_find_sink_order(4)))
self._dev_graph.insert_node_between_nodes(new_node, node, cn)
nodes_need_to_merge_for_special_ops.append(new_node)
index += 1
for node in nodes_need_to_delete_for_special_ops:
self._dev_graph.remove_node(node)
source_nodes = []
for node in self._dev_graph.nodes:
if not node.in_tensors:
source_nodes.append(node)
transpose_group = []
reserverd_nodes = []
visited = []
for source in nodes_need_to_merge_for_special_ops:
merge_father_and_child(source, visited, transpose_group,
reserverd_nodes)
nodes_need_to_merge_for_special_ops = [
node for node in nodes_need_to_merge_for_special_ops
if node not in reserverd_nodes
]
for node in reserverd_nodes:
order = node.node_attr(node.op.AttrName.ORDER)
keep_order = True
if any([index != dim for index, dim in enumerate(order)]):
keep_order = False
if keep_order:
nodes_need_to_merge_for_special_ops.append(node)
for node in nodes_need_to_merge_for_special_ops:
self._dev_graph.remove_node(node)
merge_brothers(reserverd_nodes)
delete_transpose_of_correlation(self)
