mindspore 2.1.0__cp39-cp39-win_amd64.whl → 2.2.11__cp39-cp39-win_amd64.whl

mindspore/_extends/graph_kernel/model/model_builder.py

@@ -13,9 +13,6 @@
 # limitations under the License.
 # ===========================================================================
 """GraphKernel model builder"""
-
-import copy
-from . import op_infer
 from .model import Tensor, Value, Operator, Graph, AlignShape
 
 
@@ -95,18 +92,6 @@ class GraphBuilder:
         node.all_inputs = inputs
         self.current.graph.add(node)
 
-    def emit(self, prim, inputs, name=None, attrs=None):
-        """Emit a new operation"""
-        if attrs is None:
-            attrs = {}
-        if isinstance(inputs, (Tensor, Value)):
-            inputs = [inputs]
-        tensor_inputs = [t for t in inputs if isinstance(t, (Tensor, Value))]
-        out_shape, out_dtype, out_format = op_infer.infer(prim, tensor_inputs, attrs)
-        output = self.tensor(out_shape, out_dtype, out_format, name)
-        self.op(prim, output, inputs, attrs)
-        return output
-
     def get(self):
         """Get graphs"""
         return self.graphs
@@ -169,15 +154,18 @@ class CompositeGraph:
             for op in desc['op_desc']:
                 inputs = [self.tensors.get(d['tensor_name'], None) for x in op['input_desc']
                           for d in x if 'value' not in d]
+                if op['name'] in ('ReduceSum', 'ReduceMax', 'ReduceMin'):
+                    axis = op['input_desc'][1][0]['value']
+                    if isinstance(axis, int):
+                        axis = [axis]
+                    if not op['attr']:
+                        attr = [{'name': 'axis', 'dtype': 'listInt', 'value': axis}]
+                        op['attr'] = attr
+                    else:
+                        op['attr'].append({'name': 'axis', 'dtype': 'listInt', 'value': axis})
                 out_desc = op['output_desc']
                 name, shape, dtype, data_format = out_desc[0]['tensor_name'], out_desc[
                     0]['shape'], out_desc[0]['data_type'], out_desc[0]['format']
-                if op['name'] == 'InplaceAssign':
-                    inputs[0].add_buddy(inputs[1])
-                    inputs[1].para_type = Tensor.PARA_OUTPUT
-                    output = inputs[2]
-                    self.tensors[name] = output
-                    continue
                 output = self.tensors.get(name, None)
                 if not output:
                     output = builder.tensor(shape, dtype, data_format, name=name)
@@ -186,46 +174,17 @@ class CompositeGraph:
         self.graph = builder.get()[0]
         self.desc = desc
 
-    def _pre_dump(self, outputs):
-        """restore name to before load"""
-        inplace_assign = {}  # y_name, output_name
-        inplace_assign_z = None
-        for op in self.desc['op_desc']:
-            if op['name'] == 'InplaceAssign':
-                inplace_assign[op['input_desc'][1][0]['tensor_name']] = op['output_desc'][0]['tensor_name']
-        if inplace_assign:
-            for t in outputs:
-                if t.name not in inplace_assign:
-                    inplace_assign_z = t
-        return inplace_assign, inplace_assign_z
 
     def dump(self, subgraph):
         """Dump Graph to json"""
         desc = {}
         inputs, outputs = subgraph.deduce_parameters()
         graph_ops = set(subgraph.ops)
-        inplace_assign, inplace_assign_z = self._pre_dump(outputs)
 
         def dump_output(t):
-            if t.name in inplace_assign:
-                z = inplace_assign_z if inplace_assign_z is not None else self.tensors.get(t.name, None)
-                return {'data_type': z.dtype, 'shape': z.shape, 'tensor_name': inplace_assign.get(t.name)}
             return {'data_type': t.dtype, 'shape': t.shape, 'tensor_name': t.name}
 
         def dump_op_desc(d):
-            if d['name'] == 'InplaceAssign':
-                y = d['input_desc'][1][0]['tensor_name']
-                if self.tensors[y].op in graph_ops:
-                    z, fake = (inplace_assign_z, False) if inplace_assign_z is not None else (self.tensors.get(y), True)
-                    inplace_desc = copy.deepcopy(d)
-                    inplace_desc['attr'] = {'name': 'fake_output', 'value': fake}
-                    z_desc, out_desc = inplace_desc['input_desc'][2][0], inplace_desc['output_desc'][0]
-                    z_desc['shape'] = z.shape
-                    z_desc['data_type'] = z.dtype
-                    z_desc['tensor_name'] = z.name
-                    out_desc['shape'] = z.shape
-                    out_desc['data_type'] = z.dtype
-                    return inplace_desc
             op = self.tensors[d['output_desc'][0]['tensor_name']].op
             if op in graph_ops or op in subgraph.recompute_ops:
                 return d

mindspore/_extends/graph_kernel/splitter.py

@@ -36,7 +36,6 @@ def split_with_json(json_str, flags_str):
         subgraphs, graph_mode = model.split(comp.graph, target, flags)
         is_multi_graph = len(subgraphs) > 1
         graph_list = list(map(comp.dump, subgraphs))
-        _reset_graphmode_for_inplaceassign(graph_list, graph_mode)
         result = {"multi_graph": is_multi_graph,
                   "graph_desc": graph_list,
                   "graph_mode": graph_mode}
@@ -51,8 +50,9 @@ def split_with_json(json_str, flags_str):
 def _load_repository(graph, flags):
     """Load repository if exists"""
     def check_repo(op, best_split, op_desc):
-        if not isinstance(best_split, dict) or "group_num" not in best_split or "graph_mode" not in best_split \
-                or "split_result" not in best_split:
+        if not isinstance(best_split, dict):
+            return False
+        if "group_num" not in best_split or "graph_mode" not in best_split or "split_result" not in best_split:
             logger.warning("The graph split repository of {} should be a dict which contains 'group_num', 'graph_mode' "
                            "and 'split_result' field, but got {}".format(op, best_split))
             return False
@@ -114,19 +114,12 @@ def _load_repository(graph, flags):
     return result
 
 
-def _reset_graphmode_for_inplaceassign(graph_list, graph_mode):
-    """Operator with InplaceAssign should always be composite op"""
-    for i, g in enumerate(graph_list):
-        if any((op['name'] == 'InplaceAssign' for op in g['op_desc'])):
-            graph_mode[i] = 'composite'
-
-
 def _dump_split_info(use_repo, graph_str, graph, subgraphs, graph_mode, graph_list):
     """Dump split info as text"""
     graph_kernel_dump_path = "graph_kernel_dump"
     utils.create_dir(graph_kernel_dump_path)
     filename = os.path.join(graph_kernel_dump_path, "graph_kernel_split_mode.%d.txt" % os.getpid())
-    with os.fdopen(os.open(filename, os.O_WRONLY | os.O_CREAT), "a+") as f:
+    with os.fdopen(os.open(filename, os.O_WRONLY | os.O_CREAT, 0o600), "a+") as f:
         f.write("********** main graph: {} **********\n".format(graph.name))
         f.write("input json:\n{}\n".format(graph_str))
         f.write("graph desc:\n{}\n".format(str(graph)))

mindspore/_extends/parallel_compile/akg_compiler/akg_process.py

@@ -113,30 +113,115 @@ def create_akg_parallel_process(process_num, wait_time, platform):
     return AkgProcess(process_num, wait_time, platform)
 
 
-class AkgProcess:
-    """akg kernel parallel process"""
+def _is_input_shape_dynamic(desc_d):
+    input_lists = desc_d.get("input_desc", [])
+    if input_lists is None:
+        return True
+    for input_desc in input_lists:
+        shape = input_desc[0].get("shape", ())
+        if -1 in shape or -2 in shape:
+            return True
+    return False
 
-    def __init__(self, process_num, wait_time, platform):
+
+def _compile_akg_v2_task_default(json_strs, attrs, driver):
+    """
+    compile func called in single process
+
+    Parameters:
+        json_strs: list. List contains multiple kernel infos, suitable for json compile api.
+    """
+    log_level = get_log_level(attrs)
+    kernel_meta_dir = os.path.join(get_kernel_meta_parent_dir(attrs), "akg_kernel_meta")
+    for json_str in json_strs:
+        json_desc = json.loads(json_str)
+        op_name = json_desc["op"]
+        info_path = os.path.join(kernel_meta_dir, op_name + ".info")
+        if not os.path.isfile(info_path):
+            raise FileNotFoundError(f"Can not compile non-existing file \"{info_path}\"")
+        # Compile json str with AKG
+        bisheng_cpp_path = os.getenv("BISHENG_CPP_PATH", default="")
+        compiler = driver(input_file=info_path, output_dir=kernel_meta_dir, bisheng_tools_dir=bisheng_cpp_path,
+                          dynamic_shape=_is_input_shape_dynamic(json_desc))
+        try:
+            compiler.compile()
+        except RuntimeError as exc:
+            if log_level == "ERROR":
+                raise ValueError(f"Compile error, json str: {json_str}! build attrs: {attrs}") from exc
+            logger.info(f"Will try to split, json str: {json_str}! build attrs: {attrs}")
+
+
+def create_akg_v2_parallel_process(process_num, wait_time, platform):
+    """
+    create Akg V2 Parallel Compiler object
+
+    Returns:
+        AKG V2 ParallelCompiler
+    """
+    return AkgV2Process(process_num, wait_time, platform)
+
+
+class AkgProcessBase:
+    """base class for akg kernel parallel process"""
+
+    def __init__(self, name, process_num, wait_time, platform):
         """
         Args:
             process_num: int. processes number
             wait_time: int. max time the function blocked
         """
         if not isinstance(process_num, int):
-            raise ValueError("AKG kernel compiling process number must be of type int, but got {} with type {}"
-                             .format(process_num, type(wait_time)))
+            raise ValueError(
+                f"{name} kernel compiling process number must be of type int"
+                ", but got {process_num} with type {type(wait_time)}")
         if not isinstance(wait_time, int):
-            raise ValueError("AKG kernel compiling wait time must be of type int, but got {} with type {}"
-                             .format(wait_time, type(wait_time)))
+            raise ValueError(
+                f"{name} kernel compiling wait time must be of type int,"
+                " but got {wait_time} with type {type(wait_time)}")
         if process_num == 0:
             process_num = 1
         max_proc_num = 16
+        self.name = name
         self.process_num = min([cpu_count(), max_proc_num, process_num])
         self.args = list([] for _ in range(self.process_num))
         self.wait_time = wait_time
         self.platform = platform
         self.argc = 0
 
+    def compile(self, attrs=None):
+        """
+        compile kernel by multi processes
+        Return:
+            True for all compile success, False for some failed.
+        """
+        del attrs
+        raise NotImplementedError
+
+    def accept_json(self, json_str):
+        """
+        accept json data before compile
+        Args:
+            json_str: str. kernel info.
+        """
+        if not isinstance(json_str, str):
+            raise ValueError(
+                f"In {self.name} kernel compiling, the kernel json must be of type str"
+                ", but got {json_str} with type { type(json_str)}")
+        self.args[self.argc % self.process_num].append(json_str)
+        self.argc += 1
+
+
+class AkgProcess(AkgProcessBase):
+    """akg kernel parallel process"""
+
+    def __init__(self, process_num, wait_time, platform):
+        """
+        Args:
+            process_num: int. processes number
+            wait_time: int. max time the function blocked
+        """
+        super(AkgProcess, self).__init__("AKG", process_num, wait_time, platform)
+
     def compile(self, attrs=None):
         """
         compile kernel by multi processes
@@ -162,14 +247,36 @@ class AkgProcess:
                 res.get(timeout=self.wait_time)
         return True
 
-    def accept_json(self, json_str):
+
+class AkgV2Process(AkgProcessBase):
+    """akg v2 kernel parallel process"""
+
+    def __init__(self, process_num, wait_time, platform):
         """
-        accept json data before compile
         Args:
-            json_str: str. kernel info.
+            process_num: int. processes number
+            wait_time: int. max time the function blocked
         """
-        if not isinstance(json_str, str):
-            raise ValueError("In AKG kernel compiling, the kernel json must be of type str, but got {} with type {}"
-                             .format(json, type(json)))
-        self.args[self.argc % self.process_num].append(json_str)
-        self.argc += 1
+        super(AkgV2Process, self).__init__("AKG V2", process_num, wait_time, platform)
+
+    def compile(self, attrs=None):
+        """
+        compile kernel by multi processes
+        Return:
+            True for all compile success, False for some failed.
+        """
+        if self.argc == 0:
+            raise ValueError("In AKG V2 kernel compiling, the number of kernel json that need to be compiled can "
+                             "not be zero.")
+        akg_v2_path = os.getenv("AKG_V2_PATH", default="")
+        if akg_v2_path == "":
+            raise ValueError(
+                "The path to akg v2 compiler is not specified. Set the path to the compiler in AKG_V2_PATH")
+        sys.path.append(akg_v2_path)
+        p = __import__("akg_v2", globals(), locals())
+        driver = getattr(p, "AkgV2Driver")
+        args = list((arg, attrs, driver) for arg in self.args)
+        with Pool(processes=self.process_num) as pool:
+            res = pool.starmap_async(_compile_akg_v2_task_default, args)
+            res.get(timeout=self.wait_time)
+        return True

mindspore/_extends/parallel_compile/akg_compiler/build_tbe_kernel.py

@@ -1,4 +1,4 @@
-# Copyright 2021 Huawei Technologies Co., Ltd
+# Copyright 2021-2023 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -25,6 +25,20 @@ from tbe.common.buildcfg import build_config
 from tbe.dsl import auto_schedule
 from tbe.dsl import build as tbe_build
 import tbe.common.context.op_context as op_context
+from impl.dynamic.add import _add_check_format, _infer_shape
+
+SHAPE = "shape"
+FORMAT = "format"
+DATA_TYPE = "data_type"
+NEW_SHAPE = "new_shape"
+ORI_SHAPE = "ori_shape"
+ORI_FORMAT = "ori_format"
+DST_TYPE = "dst_type"
+DST_ORI_SHAPE = "dst_ori_shape"
+INPUT_DESC = "input_desc"
+OUTPUT_DESC = "output_desc"
+ENABLE_VECTOR_2X = "enable_vector_2x"
+ENABLE_GROUP_INPLACE = "enable_group_inplace"
 
 
 def initialize(kernel_meta_parent_dir):
@@ -49,14 +63,14 @@ def update_config(config, op_names):
     change_type_dict = {"MatMul": (True, False),
                         "BatchMatMul": (True, False)}
     config["bool_storage_as_1bit"] = True
-    config["enable_group_inplace"] = False
-    config["enable_vector_2x"] = True
+    config[ENABLE_GROUP_INPLACE] = False
+    config[ENABLE_VECTOR_2X] = True
     for op in op_names:
         if op in bool_storage_as_1bit_oplist:
             config["bool_storage_as_1bit"] = False
         enable_group_inplace, enable_vector_2x = change_type_dict.get(op, (False, True))
-        config["enable_group_inplace"] = config["enable_group_inplace"] or enable_group_inplace
-        config["enable_vector_2x"] = config["enable_vector_2x"] and enable_vector_2x
+        config[ENABLE_GROUP_INPLACE] = config[ENABLE_GROUP_INPLACE] or enable_group_inplace
+        config[ENABLE_VECTOR_2X] = config[ENABLE_VECTOR_2X] and enable_vector_2x
 
 
 def add_new_shape(names, shapes, new_shapes, inputs):
@@ -70,11 +84,11 @@ def add_new_shape(names, shapes, new_shapes, inputs):
             continue
         if name not in inputs:
             raise RuntimeError("Can not support reshape on output tensor {}".format(name))
-        if "new_shape" not in inputs[name]:
-            inputs[name]["new_shape"] = new_shapes[i]
-        elif new_shapes[i] != inputs[name]["new_shape"]:
+        if NEW_SHAPE not in inputs[name]:
+            inputs[name][NEW_SHAPE] = new_shapes[i]
+        elif new_shapes[i] != inputs[name][NEW_SHAPE]:
             raise RuntimeError("Find different new_shape {} and {} for {}"
-                               .format(inputs[name]["new_shape"], new_shapes[i], name))
+                               .format(inputs[name][NEW_SHAPE], new_shapes[i], name))
 
 
 class TransShape:
@@ -93,22 +107,21 @@ class TransShape:
             if v.get("value") is not None:
                 continue
             names.append(k)
-            shapes.append(v["shape"])
-            ori_shapes.append(v["ori_shape"] if v.get("ori_shape") else None)
-            formats.append(v["format"])
-            ori_formats.append(v["ori_format"])
+            shapes.append(v[SHAPE])
+            ori_shapes.append(v[ORI_SHAPE] if v.get(ORI_SHAPE) else None)
+            formats.append(v[FORMAT])
+            ori_formats.append(v[ORI_FORMAT])
         if len(shapes) == 2 and len(shapes[0]) != len(shapes[1]):
-            from impl.add import _add_check_format, _infer_shape
-            format_pattern = _add_check_format({"shape": shapes[0], "format": formats[0]},
-                                               {"shape": shapes[1], "format": formats[1]})
+            format_pattern = _add_check_format({SHAPE: shapes[0], FORMAT: formats[0]},
+                                               {SHAPE: shapes[1], FORMAT: formats[1]})
             ori_shape0 = ori_shapes[0] if ori_shapes[0] is not None else infer_ori_shape(
                 shapes[0], formats[0], ori_formats[0])
             ori_shape1 = ori_shapes[1] if ori_shapes[1] is not None else infer_ori_shape(
                 shapes[1], formats[1], ori_formats[1])
             new_shapes = [None, None]
             new_shapes[0], new_shapes[1] = _infer_shape(format_pattern,
-                                                        {"shape": shapes[0], "ori_shape": ori_shape0},
-                                                        {"shape": shapes[1], "ori_shape": ori_shape1})
+                                                        {SHAPE: shapes[0], ORI_SHAPE: ori_shape0},
+                                                        {SHAPE: shapes[1], ORI_SHAPE: ori_shape1})
             new_shapes[0], new_shapes[1], _ = shape_util.broadcast_shapes(new_shapes[0], new_shapes[1],
                                                                           param_name_input1="input0",
                                                                           param_name_input2="input1")
@@ -119,7 +132,7 @@ class TransShape:
         """deal with batch_matmul."""
         for k, v in op_inputs.items():
             # batch dimension of BatchMatMul must be fused to 1D
-            shape = v["shape"]
+            shape = v[SHAPE]
             if len(shape) > 5:
                 new_shape = [functools.reduce(lambda x, y: x * y, shape[:-4])] + shape[-4:]
                 add_new_shape(k, shape, new_shape, inputs)
@@ -135,6 +148,10 @@ class TransShape:
 
 def infer_ori_shape(shape, cur_format, ori_format):
     """Given current format and shape, infer the shape with ori_format."""
+
+    def _shape_error(current_shape, current_format):
+        raise ValueError("Invalid shape {} for format {}".format(current_shape, current_format))
+
     if cur_format == ori_format:
         return shape
     default_formats = ["DefaultFormat", "ND", "NCHW"]
@@ -145,7 +162,7 @@ def infer_ori_shape(shape, cur_format, ori_format):
     if cur_format == "FRACTAL_NZ" and ori_format in default_formats:
         dims = len(shape)
         if dims < 4:
-            raise ValueError("Invalid shape {} for format {}".format(shape, cur_format))
+            _shape_error(shape, cur_format)
         ori_shape = shape[:dims - 4]
         m = shape[-3] * shape[-2]
         n = shape[-4] * shape[-1]
@@ -155,13 +172,13 @@ def infer_ori_shape(shape, cur_format, ori_format):
 
     if cur_format == "NC1HWC0" and ori_format in default_formats:
         if len(shape) != 5:
-            raise ValueError("Invalid shape {} for format {}".format(shape, cur_format))
+            _shape_error(shape, cur_format)
         ori_shape = [shape[0], shape[1] * shape[4], shape[2], shape[3]]
         return ori_shape
 
     if cur_format == "NHWC" and ori_format in default_formats:
         if len(shape) != 4:
-            raise ValueError("Invalid shape {} for format {}".format(shape, cur_format))
+            _shape_error(shape, cur_format)
         ori_shape = [shape[0], shape[3], shape[1], shape[2]]
         return ori_shape
 
@@ -202,7 +219,7 @@ def get_input_desc(input_desc):
     res = {}
     for desc in input_desc:
         for item in desc:
-            item["shape"] = [1] if not item["shape"] else item["shape"]
+            item[SHAPE] = [1] if not item[SHAPE] else item[SHAPE]
             res[item["tensor_name"]] = item
     return res
 
@@ -215,7 +232,7 @@ def get_inputs_tensor(input_desc, all_tensors):
             name = item["tensor_name"]
             if item.get("value") is not None:
                 # const value
-                all_tensors[name] = tvm.const(item["value"], item["data_type"])
+                all_tensors[name] = tvm.const(item["value"], item[DATA_TYPE])
             if all_tensors.get(name) is None:
                 raise ValueError("Tensor [{}] not found.".format(name))
             inputs.append(all_tensors[name])
@@ -237,17 +254,17 @@ def get_op_attrs(op, fusion_op_name):
     op_name = op["name"]
     op_attrs = get_attr_dict(op.get("attr"))
     if op_name == "BatchMatMul":
-        op_attrs["dst_type"] = op["output_desc"][0]["data_type"]
-        op_attrs["dst_ori_shape"] = op["output_desc"][0].get("ori_shape")
-        if op_attrs.get("dst_ori_shape") is None:
-            op_attrs["dst_ori_shape"] = infer_ori_shape(op["output_desc"][0]["shape"],
-                                                        op["output_desc"][0]["format"],
-                                                        op["output_desc"][0]["ori_format"])
+        op_attrs[DST_TYPE] = op[OUTPUT_DESC][0][DATA_TYPE]
+        op_attrs[DST_ORI_SHAPE] = op[OUTPUT_DESC][0].get(ORI_SHAPE)
+        if op_attrs.get(DST_ORI_SHAPE) is None:
+            op_attrs[DST_ORI_SHAPE] = infer_ori_shape(op[OUTPUT_DESC][0][SHAPE],
+                                                      op[OUTPUT_DESC][0][FORMAT],
+                                                      op[OUTPUT_DESC][0][ORI_FORMAT])
     elif op_name == "MatMul":
-        op_attrs["dst_type"] = op["output_desc"][0]["data_type"]
-        op_attrs["dst_format"] = op["output_desc"][0]["format"]
+        op_attrs[DST_TYPE] = op[OUTPUT_DESC][0][DATA_TYPE]
+        op_attrs["dst_format"] = op[OUTPUT_DESC][0][FORMAT]
     elif op_name == "Cast":
-        op_attrs["dst_type"] = op["output_desc"][0]["data_type"]
+        op_attrs[DST_TYPE] = op[OUTPUT_DESC][0][DATA_TYPE]
     op_attrs["fusion_op_name"] = fusion_op_name
     return op_attrs
 
@@ -256,17 +273,17 @@ def create_placeholders(inputs):
     """Create placeholders."""
     tensors = {}
     for k, v in inputs.items():
-        dtype = v["data_type"]
+        dtype = v[DATA_TYPE]
         if dtype == "bool":
             dtype = "int8"
-        shape = v["shape"]
-        if "new_shape" in v:
-            shape = v["new_shape"]
+        shape = v[SHAPE]
+        if NEW_SHAPE in v:
+            shape = v[NEW_SHAPE]
         attr = {
-            "format": v.get("format"),
+            FORMAT: v.get(FORMAT),
             "sub_format": v.get("sub_format", ""),
-            "ori_shape": v.get("ori_shape"),
-            "ori_format": v.get("ori_format"),
+            ORI_SHAPE: v.get(ORI_SHAPE),
+            ORI_FORMAT: v.get(ORI_FORMAT),
             "addr_type": v.get("addr_type", 0),
             "valid_shape": v.get("valid_shape", []),
             "slice_offset": v.get("slice_offset", []),
@@ -276,8 +293,8 @@ def create_placeholders(inputs):
             "L1_valid_size": v.get("L1_valid_size", -1),
             "range": v.get("range", [])
         }
-        if attr.get("ori_shape") is None:
-            attr["ori_shape"] = infer_ori_shape(v.get("shape"), v.get("format"), attr.get("ori_format"))
+        if attr.get(ORI_SHAPE) is None:
+            attr[ORI_SHAPE] = infer_ori_shape(v.get(SHAPE), v.get(FORMAT), attr.get(ORI_FORMAT))
         tensors[k] = tvm.placeholder(shape=shape, name=k, dtype=dtype, attrs=attr)
     return tensors
 
@@ -289,8 +306,8 @@ def same_shape(inputs):
     base_shape = -1
     for _, v in inputs.items():
         if base_shape == -1:
-            base_shape = v["shape"]
-        if v["shape"] != base_shape:
+            base_shape = v[SHAPE]
+        if v[SHAPE] != base_shape:
             return False
     return True
 
@@ -298,17 +315,17 @@ def same_shape(inputs):
 def create_input_tensors(json_dict):
     """Create input placeholders."""
     fold_dim = True
-    inputs = get_input_desc(json_dict.get("input_desc", []))
+    inputs = get_input_desc(json_dict.get(INPUT_DESC, []))
     for op in json_dict["op_desc"]:
         op_name = op["name"]
         pattern = get_op_reg_info(op_name, "pattern")
-        op_inputs = get_input_desc(op.get("input_desc", []))
+        op_inputs = get_input_desc(op.get(INPUT_DESC, []))
         TransShape.run(op_name, pattern, op_inputs, inputs)
         if pattern != OpPattern.ELEMWISE or not same_shape(op_inputs):
             fold_dim = False
     if fold_dim:
         for k, v in inputs.items():
-            shape = v["shape"]
+            shape = v[SHAPE]
             new_shape = [functools.reduce(lambda x, y: x * y, shape[:])]
             add_new_shape(k, shape, new_shape, inputs)
     return create_placeholders(inputs)
@@ -324,28 +341,28 @@ def create_fusion_op_name(op_names):
 
 
 def update_format(json_dict):
-    """Some format like DefaultFormat is not recognized in TBE, need to covert these formats."""
+    """Some format like DefaultFormat is not recognized in TBE, need to convert these formats."""
 
     def _update_input_format(input_desc):
         for desc in input_desc:
             for item in desc:
-                if item["format"] == "DefaultFormat":
-                    item["format"] = "ND"
-                if item.get("ori_format") is None or item["ori_format"] == "DefaultFormat":
-                    item["ori_format"] = "NCHW"
+                if item[FORMAT] == "DefaultFormat":
+                    item[FORMAT] = "ND"
+                if item.get(ORI_FORMAT) is None or item[ORI_FORMAT] == "DefaultFormat":
+                    item[ORI_FORMAT] = "NCHW"
 
     def _update_output_format(output_desc):
         for item in output_desc:
-            if item["format"] == "DefaultFormat":
-                item["format"] = "ND"
-            if item.get("ori_format") is None or item["ori_format"] == "DefaultFormat":
-                item["ori_format"] = "NCHW"
+            if item[FORMAT] == "DefaultFormat":
+                item[FORMAT] = "ND"
+            if item.get(ORI_FORMAT) is None or item[ORI_FORMAT] == "DefaultFormat":
+                item[ORI_FORMAT] = "NCHW"
 
-    _update_input_format(json_dict.get("input_desc", []))
-    _update_output_format(json_dict["output_desc"])
+    _update_input_format(json_dict.get(INPUT_DESC, []))
+    _update_output_format(json_dict[OUTPUT_DESC])
     for op in json_dict["op_desc"]:
-        _update_input_format(op.get("input_desc", []))
-        _update_output_format(op["output_desc"])
+        _update_input_format(op.get(INPUT_DESC, []))
+        _update_output_format(op[OUTPUT_DESC])
 
 
 def gen_args_remap(orig_inputs_name, orig_outputs_name, inputs_name, outputs_name, inplace_names):
@@ -410,7 +427,7 @@ def update_json(json_dict, inputs_name, outputs_name, inplace_names, kernel_meta
             pass
         # generate new .json
         try:
-            with os.fdopen(os.open(json_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o660), 'w') as fi:
+            with os.fdopen(os.open(json_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), 'w') as fi:
                 json.dump(json_dict, fi, sort_keys=True, indent=4, separators=(',', ':'))
         except OSError:
             pass
@@ -441,8 +458,8 @@ def build(json_str, kernel_meta_parent_dir):
     """Build kernel."""
     json_dict = json.loads(json_str)
     update_format(json_dict)
-    inputs_name = get_inputs_name(json_dict.get("input_desc", []))
-    outputs_name, inplace_names = get_outputs_info(json_dict["output_desc"])
+    inputs_name = get_inputs_name(json_dict.get(INPUT_DESC, []))
+    outputs_name, inplace_names = get_outputs_info(json_dict[OUTPUT_DESC])
     op_names = get_all_op_name(json_dict["op_desc"])
     fusion_op_name = create_fusion_op_name(op_names)
 
@@ -458,7 +475,7 @@ def build(json_str, kernel_meta_parent_dir):
         for op in json_dict["op_desc"]:
             op_name = op["name"]
             # get op input tensor
-            op_inputs = get_inputs_tensor(op.get("input_desc", []), all_tensors)
+            op_inputs = get_inputs_tensor(op.get(INPUT_DESC, []), all_tensors)
             # get op attrs
             op_attrs = get_op_attrs(op, fusion_op_name)
             # op compute
@@ -466,10 +483,10 @@ def build(json_str, kernel_meta_parent_dir):
             # update op output tensor
             if not isinstance(op_outputs, (list, tuple)):
                 op_outputs = [op_outputs]
-            if len(op["output_desc"]) != len(op_outputs):
+            if len(op[OUTPUT_DESC]) != len(op_outputs):
                 raise ValueError("len(op[\"output_desc\"] is not equal to the number of real output tensors in op[{}]: "
-                                 "{} vs {}".format(op_name, len(op["output_desc"]), len(op_outputs)))
-            for i, desc in enumerate(op["output_desc"]):
+                                 "{} vs {}".format(op_name, len(op[OUTPUT_DESC]), len(op_outputs)))
+            for i, desc in enumerate(op[OUTPUT_DESC]):
                 all_tensors[desc["tensor_name"]] = op_outputs[i]
 
         # Collect input, output tensors

mindspore/_extends/parallel_compile/akg_compiler/tbe_topi.py

@@ -310,7 +310,8 @@ def _log(x, attrs=None):
     if base <= 0 and not math.isclose(base, -1.0, rel_tol=1e-8, abs_tol=0.0):
         raise ValueError("base must be strictly positive or -1, but got {}".format(base))
     from impl.log import log_compute
-    return log_compute(x, None, base, scale, shift, kernel_name=attrs["fusion_op_name"])
+    output_desc = {"dtype": x.dtype, "shape": x.shape}
+    return log_compute(x, output_desc, base, scale, shift, kernel_name=attrs["fusion_op_name"])
 
 
 @reg_op("Maximum", pattern=OpPattern.ELEMWISE)
@@ -349,7 +350,8 @@ def _mul(x0, x1, attrs=None):
         return tbe.dsl.vmuls(x1, x0)
     x0, x1 = _broadcast(x0, x1)
     from impl.mul import mul_compute
-    return mul_compute(x0, x1, None, kernel_name=attrs["fusion_op_name"])
+    output_desc = {"dtype": x0.dtype, "shape": x0.shape}
+    return mul_compute(x0, x1, output_desc, kernel_name=attrs["fusion_op_name"])
 
 
 @reg_op("Neg", pattern=OpPattern.ELEMWISE)