mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/_extends/builtin_operations.py

@@ -20,6 +20,7 @@ from mindspore.ops import composite as C
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.common.dtype import dtype_to_nptype, get_py_obj_dtype
+from mindspore._c_expression import Tensor as CTensor
 
 
 def ScalarAdd(x, y):
@@ -103,8 +104,11 @@ def zeros_like_tensor(x):
 
 def OnesLike(x):
     """Implement `oneslike`."""
-    x = x.asnumpy()
-    value = Tensor(np.ones(x.shape).astype(x.dtype))
+    if isinstance(x, (Tensor, CTensor)):
+        x = x.asnumpy()
+        value = Tensor(np.ones(x.shape).astype(x.dtype))
+    else:
+        value = Tensor(1.0)
     return value
 
 
@@ -176,11 +180,35 @@ def tuple_to_array(x):
     return Tensor(np.array(x))
 
 
-def stop_gradient(x):
-    """Implement `stop_gradient`."""
+def StopGradient(x):
+    """Implement `StopGradient`."""
     return x
 
 
+def TensorShape(x):
+    """Implement `TensorShape`."""
+    return Tensor(F.shape(x))
+
+
+def DynamicBroadcastGradientArgs(x, y):
+    """Implement `DynamicBroadcastGradientArgs`."""
+    return -1, -1
+
+
+def StridedSlice(x, begin, end, stride):
+    """Implement `StridedSlice`."""
+    if isinstance(x, Tensor):
+        x = x.asnumpy()
+        ret = x
+        for i in range(len(x)):
+            ret[i:] = x[begin[i], end[i], stride[i]]
+        return Tensor(ret)
+    ret = x
+    for i in range(len(x)):
+        ret[i:] = x[begin[i], end[i], stride[i]]
+    return ret
+
+
 hyper_map = C.HyperMap()
 
 

mindspore/_extends/graph_kernel/model/graph_split.py

@@ -15,9 +15,7 @@
 """Cost model splitter"""
 from functools import reduce as prod_reduce
 from functools import partial
-from mindspore import log as logger
 from .model import PrimLib, Graph, Tensor, Operator
-from .model import DataFormat as DF
 
 
 def tensor_size(tensor):
@@ -175,7 +173,7 @@ class ReshapeElimChecker:
         def _get_remap_axis(in_shape, out_shape):
             rin, rout = [], []
             in_prod, out_prod, out_idx = 1, out_shape[-1], -1
-            in_ext, out_ext = -len(in_shape)-1, -len(out_shape)-1
+            in_ext, out_ext = -len(in_shape) - 1, -len(out_shape) - 1
             for in_idx in range(-1, in_ext, -1):
                 in_prod = in_prod * in_shape[in_idx]
                 while out_prod < in_prod:
@@ -191,6 +189,7 @@ class ReshapeElimChecker:
             if out_idx > out_ext:
                 rout.extend([i for i in range(out_idx, out_ext, -1)])
             return rin, rout
+
         remap_in, remap_out = _get_remap_axis(reshape.inputs[0].shape, reshape.output.shape)
         self.exc_fwd = self._collect_exc_ops(reshape, remap_in, True)
         self.exc_bwd = self._collect_exc_ops(reshape, remap_out, False)
@@ -198,10 +197,11 @@ class ReshapeElimChecker:
     @staticmethod
     def _collect_exc_ops(reshape, remap_axis, is_fwd):
         """collect exclude ops of reshape"""
+
         def _propagate(remap, src, des):
             out_remap = []
             src_prod, des_prod, des_idx = 1, 1, 0
-            for src_idx in range(-1, -len(src)-1, -1):
+            for src_idx in range(-1, -len(src) - 1, -1):
                 src_prod = src_prod * src[src_idx]
                 if src_idx in remap:
                     while des_prod < src_prod:
@@ -256,6 +256,7 @@ class ReshapeElimChecker:
                         push_stack(op, remap)
                     else:
                         exc_ops.add(op)
+
         exc_ops, stack, visited = set(), [], {reshape}
         if is_fwd:
             _visit_fwd(reshape, remap_axis)
@@ -296,7 +297,7 @@ class ReduceOutFuseChecker:
             for to in op.output.to_ops:
                 idx = to.inputs.index(op.output)
                 if PrimLib.iter_type(to) > PrimLib.ELEMWISE or \
-                   tensor_size(to.inputs[idx]) != tensor_size(to.output):
+                        tensor_size(to.inputs[idx]) != tensor_size(to.output):
                     self.output_excluded.add(to)
                 else:
                     recursion_stack.append(to)
@@ -384,7 +385,7 @@ class GraphSplitByPattern:
             self.unique_id = unique_id
             self.reach_tab = reach_tab
             self.checkers = []
-            if self.pattern == PrimLib.RESHAPE and init_op.inputs: # reshape's input may be empty (const value)
+            if self.pattern == PrimLib.RESHAPE and init_op.inputs:  # reshape's input may be empty (const value)
                 self.checkers.append(ReshapeElimChecker(init_op))
             elif self.pattern == PrimLib.REDUCE:
                 self.checkers.append(ReduceOutFuseChecker(init_op))
@@ -430,6 +431,7 @@ class GraphSplitByPattern:
 
         def fuse_confirm(self, area):
             """confirm if area can be fused"""
+
             def _check(a, b, res, fwd):
                 for checker in a.checkers:
                     r = checker.check(b.ops, fwd)
@@ -441,6 +443,7 @@ class GraphSplitByPattern:
             def _commit(a, res):
                 for i, checker in enumerate(a.checkers):
                     checker.commit(res[i])
+
             res1, res2 = [], []
             if not _check(self, area, res1, True) or not _check(area, self, res2, False):
                 return False
@@ -645,7 +648,7 @@ class GraphSplitByPattern:
         area.mode = self.get_default_mode(area.ops[0])
 
     @staticmethod
-    def limit_area_size(dominant, fuse_areas, limit_size=200):
+    def limit_area_size(dominant, fuse_areas, limit_size):
         """Remove some areas if the size is too large"""
         area_sizes = map(lambda area: len(area.ops), fuse_areas)
         dom_size = len(dominant.ops)
@@ -671,7 +674,7 @@ class GraphSplitByPattern:
                     continue
                 fuse_areas, is_forward = result
                 if not is_stitch:
-                    fuse_areas = self.limit_area_size(dominant, fuse_areas)
+                    fuse_areas = self.limit_area_size(dominant, fuse_areas, self.flags['composite_op_limit_size'])
                     if not fuse_areas:
                         continue
                 changed = False
@@ -703,6 +706,7 @@ class GraphSplitByPattern:
 
     def hfuse(self, selector):
         """Fuse horizontal areas with same input tensor"""
+
         def _do_fuse(areas):
             for i in range(len(areas) - 1):
                 dom = areas[i]
@@ -748,7 +752,7 @@ class GraphSplitByPattern:
             result = selector(dominant)
             if result and result[0]:
                 fuse_areas, _ = result
-                fuse_areas = self.limit_area_size(dominant, fuse_areas)
+                fuse_areas = self.limit_area_size(dominant, fuse_areas, self.flags['composite_op_limit_size'])
                 if not fuse_areas:
                     continue
                 if fuse_areas[0] in [self.recom_area, user] and user.fuse_confirm(self.recom_area):
@@ -788,9 +792,47 @@ class GraphSplitByPattern:
         self.pattern_fuse()
         if self.enable_recompute:
             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 recompute_fuse(self):
         """find recompute regions and copy them out to new Areas"""
 
@@ -959,10 +1001,21 @@ class GraphSplitGpu(GraphSplitByPattern):
                 return reduce_size >= 1024
             return True
 
+        def _may_multi_filter(dom_ops):
+            count = 1
+            stack = [dom_ops[0]]
+            while stack:
+                op = stack.pop()
+                for t in op.inputs:
+                    if t.op and t.op in dom_ops:
+                        count = count + 1
+                        stack.append(t.op)
+            return count < len(dom_ops)
+
         def _reduce_output(dom):
             if dom.pattern != PrimLib.REDUCE:
                 return []
-            if reduce_nums(dom.ops) > 1:
+            if _may_multi_filter(dom.ops):
                 return []
             if _is_atomic_add_available(dom):
                 return []
@@ -1122,6 +1175,7 @@ class GraphSplitGpu(GraphSplitByPattern):
 
         def _broadcast_tot(dom):
             """Fuse rule for TensorScatterAdd and UnsortedSegmentSum."""
+
             def _same_input(op1, op2):
                 return bool(set(op1.inputs) & set(op2.inputs))
 
@@ -1182,7 +1236,7 @@ class GraphSplitGpu(GraphSplitByPattern):
             if dom.is_output or len(to_ops) != 1 or len(dom.out_relations) != 1:
                 return []
             to_area = list(dom.out_relations.keys())[0]
-            if (to_area.pattern > PrimLib.REDUCE and to_area.dom_op().prim not in injective_ops) or \
+            if (to_area.pattern >= PrimLib.REDUCE and to_area.dom_op().prim not in injective_ops) or \
                     to_ops[0] not in to_area.ops:
                 return []
             if len(to_area.ops) > self.TRANSPOSE_FUSE_DEPTH:
@@ -1214,6 +1268,7 @@ class GraphSplitGpu(GraphSplitByPattern):
         def _link_csr(dom):
             def _same_input(op1, op2):
                 return bool(set(op1.inputs.copy()) & set(op2.inputs.copy()))
+
             fuse_arg = {"CSRReduceSum": slice(1, 3), "CSRGather": slice(2, 3)}
             arg_idx = fuse_arg.get(dom.dom_op().prim, -1)
             if arg_idx == -1:
@@ -1275,220 +1330,9 @@ class GraphSplitGpu(GraphSplitByPattern):
             _fuse_once(fuse_func)
 
 
-class GraphSplitAscend(GraphSplitByPattern):
-    """Graph splitter"""
-    BROADCAST_FUSE_DEPTH = 6
-
-    def __init__(self, graph, flags):
-        super().__init__(graph, flags)
-        self.reduce_fuse_depth = 10 if self.reduce_fuse_depth < 0 else self.reduce_fuse_depth
-
-    def get_default_mode(self, op):
-        """Get default mode for Ascend"""
-
-        def _dtype_same(tensors):
-            dtype = tensors[0].dtype
-            for tensor_ in tensors:
-                if tensor_.dtype != dtype:
-                    return False
-            return True
-
-        if op.prim == "MatMul":
-            if op.inputs[0].dtype == "float16" and not _dtype_same(op.inputs):
-                return self.Area.MODE_COMPOSITE
-        if op.prim in ("Tile", "BroadcastTo", "ExpandDims"):
-            return self.Area.MODE_COMPOSITE
-        return self.Area.MODE_BASIC
-
-    def pattern_fuse(self, fuse_func=None):
-        """fuse Areas by pattern"""
-
-        def _likely_multicore(dom):
-            op = dom.dom_op()
-            iter_size = tensor_size(op.output if not PrimLib.is_reduce(op) else op.inputs[0])
-            return iter_size > 1024
-
-        def _broadcast_pat_exclude(dom, a, r):
-            if _likely_multicore(a) and (dom.is_output or len(dom.ops) > self.BROADCAST_FUSE_DEPTH):
-                return True
-            return a.pattern > PrimLib.REDUCE or r > PrimLib.BROADCAST
-
-        def _broadcast_bwd_depth(dom):
-            if dom.pattern not in (PrimLib.ELEMWISE, PrimLib.BROADCAST) or len(dom.out_relations) != 1:
-                return []
-            a, r = list(dom.out_relations.items())[0]
-            if _broadcast_pat_exclude(dom, a, r) or len(a.in_relations) != 1:
-                return []
-            return [a], False
-
-        def _broadcast_bwd_width(dom):
-            if dom.pattern not in (PrimLib.ELEMWISE, PrimLib.BROADCAST):
-                return []
-            fused = []
-            for a, r in dom.out_relations.items():
-                if _broadcast_pat_exclude(dom, a, r) or not dom.check_acyclic(a) or \
-                        (fused and tensor_size(fused[0].dom_op().output) != tensor_size(a.dom_op().output)):
-                    return []
-                fused.append(a)
-            return fused, False
-
-        def _reduce_pat_exclude(dom, a, r):
-            if len(a.ops) > self.reduce_fuse_depth:
-                return True
-            if r == PrimLib.BROADCAST and _likely_multicore(dom) and \
-                    (dom.is_output or len(dom.ops) > self.BROADCAST_FUSE_DEPTH):
-                return True
-            return a.pattern > PrimLib.BROADCAST or r > PrimLib.REDUCE
-
-        def _reduce_depth(dom):
-            if dom.pattern != PrimLib.REDUCE or len(dom.in_relations) != 1:
-                return []
-            a, r = list(dom.in_relations.items())[0]
-            if _reduce_pat_exclude(dom, a, r) or len(a.out_relations) != 1:
-                return []
-            return [a], True
-
-        def _reduce_width(dom):
-            if dom.pattern != PrimLib.REDUCE:
-                return []
-            fused = []
-            for a, r in dom.in_relations.items():
-                if not _reduce_pat_exclude(dom, a, r) and a.check_acyclic(dom):
-                    fused.append(a)
-            return fused, True
-
-        def _matmul_depth(dom):
-            if dom.dom_op().prim != "MatMul" and dom.dom_op().prim != "BatchMatMul":
-                return []
-            fused = []
-            for a, _ in dom.out_relations.items():
-                if (((a.dom_op().prim == "AddN" or a.dom_op().prim == "Add" or a.dom_op().prim == "Cast")
-                     and dom.dom_op().prim == "MatMul")
-                        or (a.pattern == PrimLib.ELEMWISE and dom.dom_op().prim == "BatchMatMul")) \
-                        and a.check_acyclic(dom):
-                    fused.append(a)
-            return fused, False
-
-        def _reduce_output(dom):
-            if dom.pattern != PrimLib.REDUCE:
-                return []
-            op_attrs = dom.dom_op().attrs
-            if not op_attrs.get('reduce_output_fuse'):
-                return []
-            fused = []
-            for a, r in dom.out_relations.items():
-                if a.pattern <= PrimLib.BROADCAST and r <= PrimLib.BROADCAST and \
-                        dom.check_acyclic(a):
-                    fused.append(a)
-            return fused, False
-
-        def _reduce_stitch(dom):
-            if dom.pattern != PrimLib.REDUCE:
-                return []
-            if tensor_size(dom.ops[0].output) == 1:
-                return []
-            if tensor_size(dom.ops[0].inputs[0]) < 32 * 16 * 16:
-                return []
-
-            fused = []
-            for a, r in dom.out_relations.items():
-                if not may_stitch(dom, a, r, 32, 32 * 16 * 16):
-                    continue
-                if a.pattern == PrimLib.REDUCE:
-                    if a.ops[0].attrs['reduce_axis'] == dom.ops[0].attrs['reduce_axis']:
-                        dom.stitch_info.stitch_ops.add(dom.ops[0].output.name)
-                        fused.append(a)
-                elif a.pattern <= PrimLib.BROADCAST:
-                    dom.stitch_info.stitch_ops.add(dom.ops[0].output.name)
-                    fused.append(a)
-            return fused, False
-
-        def _transdata_pattern_support(dom, a):
-            transdata_op = dom.dom_op()
-
-            # Currently, if transdata has the pad, it is not used to fuse
-            def _has_pad():
-                res = False
-                input_shape = transdata_op.inputs[0].shape
-                output_shape = transdata_op.output.shape
-                cube_size = 16
-                for dim in input_shape[-2:]:
-                    if dim % cube_size != 0:
-                        res = True
-                for dim in output_shape[-2:]:
-                    if dim % cube_size != 0:
-                        res = True
-                return res
-
-            has_pad = _has_pad()
-            if has_pad:
-                return False
-
-            if a.dom_op().prim == "MatMul" and len(dom.ops) == 1:
-                return True
-
-            # reshape/elewise/broadcast + transdata
-            if a.pattern <= PrimLib.BROADCAST and len(dom.ops) == 1:
-                op_attrs = dom.dom_op().attrs
-                if 'src_format' not in op_attrs.keys() \
-                        or 'dst_format' not in op_attrs.keys():
-                    logger.error("For 'TransData', can not find the attr 'src_format' or 'dst_format'")
-                    return False
-                src_format, dst_format = op_attrs['src_format'], op_attrs['dst_format']
-                if src_format == DF.FRAC_NZ and dst_format in (DF.DEFAULT, DF.NCHW):
-                    return True
-                # For the Default/NCHW to FRAC_NZ, currently only the Cast+Transdata is supported
-                if src_format in (DF.DEFAULT, DF.NCHW) and dst_format == DF.FRAC_NZ \
-                        and len(a.ops) == 1 and a.dom_op().prim == "Cast" and not a.is_output:
-                    return True
-            return False
-
-        def _transdata(dom):
-            if dom.dom_op().prim != "TransData":
-                return []
-            fused = []
-            for a, _ in dom.in_relations.items():
-                if _transdata_pattern_support(dom, a) and a.check_acyclic(dom):
-                    fused.append(a)
-            return fused, True
-
-        def _fuse_loop():
-            self.fuse(CommonPattern.reshape)
-            self.fuse(CommonPattern.assign)
-            self.fuse(CommonPattern.elemwise_depth)
-            self.fuse(CommonPattern.elemwise_width)
-            self.fuse(CommonPattern.broadcast_depth)
-            self.fuse(CommonPattern.broadcast_width)
-            self.fuse(_reduce_depth)
-            self.fuse(_reduce_width)
-            self.fuse(_broadcast_bwd_depth)
-            self.fuse(_broadcast_bwd_width)
-            self.fuse(_matmul_depth)
-            self.fuse(_reduce_output)
-            if self.enable_stitch_fusion:
-                self.fuse(_reduce_stitch, True)
-            self.fuse(_transdata)
-
-        def _fuse_once(fuse_func):
-            if fuse_func(CommonPattern.reshape) or \
-                    fuse_func(CommonPattern.elemwise_depth) or fuse_func(CommonPattern.elemwise_width) or \
-                    fuse_func(CommonPattern.broadcast_depth) or fuse_func(CommonPattern.broadcast_width) or \
-                    fuse_func(_reduce_depth) or fuse_func(_reduce_width) or \
-                    fuse_func(_broadcast_bwd_depth) or fuse_func(_broadcast_bwd_width) or \
-                    fuse_func(_matmul_depth) or fuse_func(_reduce_output) or fuse_func(_transdata):
-                pass
-
-        if fuse_func is None:
-            _fuse_loop()
-        else:
-            _fuse_once(fuse_func)
-
-
 def split(graph, target, flags):
     """Split graph"""
     result = None
     if target == "cuda":
         result = GraphSplitGpu(graph, flags).split()
-    else:
-        result = GraphSplitAscend(graph, flags).split()
     return result

mindspore/_extends/parallel_compile/akg_compiler/akg_process.py

@@ -21,7 +21,7 @@ from multiprocessing import Pool, cpu_count
 from mindspore import log as logger
 from mindspore._extends.parallel_compile.akg_compiler.get_file_path import get_akg_path
 from mindspore._extends.parallel_compile.akg_compiler.util import get_ascend_compile_dirs, create_compile_dirs, \
-    get_log_level, update_attr, select_best, print_compile_log, check_tbe_support
+    get_log_level, update_attr, select_best, print_compile_log, check_tbe_support, get_kernel_meta_parent_dir
 
 
 def _compile_akg_task_default(json_strs, attrs):
@@ -35,6 +35,7 @@ def _compile_akg_task_default(json_strs, attrs):
     sys.path.insert(0, get_akg_path())
     p = __import__("akg", globals(), locals(), ['ms'], 0)
     func = getattr(p.ms, "compilewithjson")
+    os.environ["MS_COMPILER_CACHE_PATH"] = get_kernel_meta_parent_dir(attrs)
 
     for json_str in json_strs:
         res = func(json_str, attrs)
@@ -66,12 +67,12 @@ def _compile_akg_task_ascend(json_strs, attrs):
         return
     log_level = get_log_level(attrs)
     compiler = os.path.join(os.path.split(os.path.realpath(__file__))[0], "compiler.py")
-    compile_dirs = get_ascend_compile_dirs()
+    compile_dirs = get_ascend_compile_dirs(attrs)
     kernel_meta_dir = compile_dirs.get("kernel_meta_dir")
     akg_compile_dir = compile_dirs.get("akg_compile_dir")
     tbe_compile_dir = compile_dirs.get("tbe_compile_dir")
     composite_graph_dir = compile_dirs.get("composite_graph_dir")
-    attrs = update_attr(attrs, {"dump_composite_graph": composite_graph_dir})
+    attrs = update_attr(attrs, {"dump_composite_graph": composite_graph_dir, "optimize_for_tbe": True})
     for json_str in json_strs:
         json_desc = json.loads(json_str)
         op_name = json_desc["op"]
@@ -87,11 +88,13 @@ def _compile_akg_task_ascend(json_strs, attrs):
         _compile_subprocess(compiler, akg_compile_dir, info_path, "AKG", attrs, compile_log, log_level)
 
         # Load composite optimized json str and compile it with TBE
-        tbe_support = check_tbe_support(json_desc)
-        if tbe_support:
-            composite_graph_path = os.path.join(composite_graph_dir, op_name + ".info")
-            if not os.path.isfile(composite_graph_path):
-                composite_graph_path = info_path
+        composite_graph_path = os.path.join(composite_graph_dir, op_name + ".info")
+        if not os.path.isfile(composite_graph_path):
+            composite_graph_path = info_path
+        with open(composite_graph_path, 'r') as f:
+            composite_graph = f.read()
+        if "buffer_stitch" not in json_desc and "parallel_fusion" not in json_desc and \
+                check_tbe_support(json.loads(composite_graph)):
             _compile_subprocess(compiler, tbe_compile_dir, composite_graph_path, "TBE", attrs, compile_log, log_level)
 
         print_compile_log(compile_log)
@@ -149,7 +152,7 @@ class AkgProcess:
                              "not be zero.")
         args = list((arg, attrs) for arg in self.args)
         if self.platform == "ASCEND":
-            create_compile_dirs(get_ascend_compile_dirs())
+            create_compile_dirs(get_ascend_compile_dirs(attrs))
             with Pool(processes=self.process_num) as pool:
                 res = pool.starmap_async(_compile_akg_task_ascend, args)
                 res.get(timeout=self.wait_time)