timekiller
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

remove split resahpe output

This commit is contained in:
Yang Jiao 2022-03-31 15:51:09 +08:00
parent fcb0d72e38
commit 35dc0aff10
2 changed files with 1 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
akg

@ -1 +1 @@
Subproject commit 14f0688d5dd8b92c782c43bf2028f73eafa4c31b
Subproject commit 0ddab6d9cad4c4c9faf72585e5a52faf79f55e4d

38
mindspore/python/mindspore/_extends/graph_kernel/model/graph_split.py
View File

@ -532,47 +532,9 @@ class GraphSplitByPattern:
"""Split graph by pattern"""
self.pattern_fuse()
self.recompute_fuse()
# The reshape should not be output node
# Note: after this function, the input output relation is not maintained.
self.split_output_reshapes()
subgraphs, graphmodes = self.to_subgraphs()
return subgraphs, graphmodes
def split_output_reshapes(self):
"""Force split the output Reshapes into other new area"""
def _remove_output_reshape(reshape_ops, other_ops):
def _run():
for op in reshape_ops:
if any((to_op in other_ops for to_op in op.output.to_ops)):
reshape_ops.remove(op)
other_ops.append(op)
return True
return False
while _run():
pass
new_areas = []
for area in self.areas:
reshape_ops = list(op for op in area.ops if PrimLib.iter_type(op) == PrimLib.RESHAPE)
other_ops = list(op for op in area.ops if op not in reshape_ops)
if not other_ops or not reshape_ops:
continue
# remove the output reshape from "reshape_ops" and add it into "other_ops"
_remove_output_reshape(reshape_ops, other_ops)
if not reshape_ops:
continue
for op in reshape_ops:
a = self.Area(op, False, 0, self.reach_tab)
self.set_default_mode(a)
new_areas.append(a)
area.ops = other_ops
if len(other_ops) == 1:
self.set_default_mode(area)
if new_areas:
self.areas += new_areas
def set_recompute(self, dom_area, ops, user_area):
"""set the recompute area and connect with other areas"""
self.recom_area.recompute_ops.extend(ops)