mindspore 2.1.0__cp39-none-any.whl → 2.2.11__cp39-none-any.whl

mindspore/_akg/akg/utils/tbe_codegen_utils.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+# coding: utf-8
+# Copyright 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+import os
+import logging
+from akg.utils.util import write_code
+import json
+
+logging.getLogger().setLevel(logging.INFO)
+
+def create_directory(dir):
+    if not os.path.exists(dir):
+        os.mkdir(dir)
+
+def set_workspace_for_json(json_path, workspace_dict):
+    if workspace_dict is None:
+        return
+    with open(json_path, 'r') as f:
+        js_info = json.loads(f.read())
+    js_info["workspace"] = workspace_dict
+    write_code(js_info, json_path)
+
+
+def copy_to_akg_kernel_meta(kernel_name, postfixs, workspace_dict=None):
+    akg_kernel_mate_str = os.getenv("KERNEL_META_DIR",default="akg_kernel_meta")
+    source = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', './'))
+    import shutil
+    target = source + "/" + akg_kernel_mate_str + "/" + kernel_name
+    source = source + "/" + "kernel_meta/" + kernel_name
+    if source == target:
+        return True
+    akg_kernel_mate_str = os.path.abspath(akg_kernel_mate_str)
+    create_directory(akg_kernel_mate_str)
+    for postfix in postfixs:
+        if os.path.exists(source + postfix):
+            try:
+                shutil.move(source + postfix, target + postfix)
+                if postfix == ".json" and workspace_dict is not None:
+                    set_workspace_for_json(target + postfix, workspace_dict)
+            except IOError as e:
+                logging.error("Unable to move file. {}".format(e))
+            except Exception as e:
+                logging.error("Unexpected error:", e)
+        else:
+            logging.info("Move {} fail, exit.".format(source + postfix))
+            return False
+    return True
+
+
+def clean_env():
+    import gc
+    import sys
+
+    imported_modules = set(sys.modules.keys())
+    for obj_key in imported_modules:
+        if "conda" in obj_key:
+            continue
+        if "akg" in obj_key or "topi" in obj_key or "tvm" in obj_key:
+            del sys.modules[obj_key]
+            try:
+                del globals()[obj_key]
+            except KeyError:
+                pass
+            try:
+                del locals()[obj_key]
+            except KeyError:
+                pass
+
+    gc.collect()
+
+
+def auto_init_soc(ascend_type):
+    from tbe.common.platform import set_current_compile_soc_info
+    set_current_compile_soc_info(ascend_type)
+
+
+def build_npu_for_akg(kernel_name,
+                      stmt=None,
+                      arg_list=None,
+                      is_dynamic=False,
+                      cfg=None,
+                      simple_mode=False):
+    import tbe
+    from tbe.tvm.tir import transform
+    from tbe.tvm.driver.cce_build_module import _count_time, generate_cce_code
+    from tbe.common.buildcfg import set_current_build_config
+    from tbe.common.buildcfg.buildcfg_mapping import dynamic_shape, disable_vectorize, tik, enable_const_fold, \
+        dynamic_tik, instrument_bound_checkers, tbe_workspace_size_list_length
+    
+    set_current_build_config(tbe_workspace_size_list_length,
+                             tbe.tvm.runtime.cce_runtime.tbe_workspace_size_list_length())
+
+    if stmt is None or arg_list is None:
+        raise ValueError("No json, exit.")
+
+    func = tbe.tvm.tir.PrimFunc(arg_list, stmt)
+    mod = tbe.tvm.IRModule({kernel_name : func})
+    # _static_lower_phase_0
+    mod = transform.InjectSocVersion()(mod)
+    mod = transform.DeduceOpPlatform()(mod)
+    mod = transform.EmitInsn()(mod)
+
+    # phase 1 _static_lower_phase_emit_insn
+    mod = transform.InjectMultiCoreSync()(mod)
+    mod = transform.SplitCoproc()(mod)
+    mod = transform.SequenceSprInsn()(mod)
+
+    # phase 2
+    mod = transform.TikDoubleBufferSupport()(mod)
+    mod = transform.InjectPipeBuffer()(mod)
+    mod = transform.OptimizeDMA()(mod)
+    mod = transform.SubstituteInstr()(mod)
+    mod = transform.InjectAccessPtrMSG()(mod)
+    mod = transform.InjectPipe()(mod)
+    mod = transform.DeSequenceSprInsn()(mod)
+    mod = transform.CanonicalSimplify()(mod)
+    mod = transform.SetSPROptimizer()(mod)
+    if cfg[enable_const_fold]:
+        mod = transform.ConstantFolding()(mod)
+    if not simple_mode:
+        mod = transform.LoopPartition()(mod)
+    if cfg[disable_vectorize]:
+        mod = transform.SkipVectorize()(mod)
+    else:
+        mod = transform.VectorizeLoop()(mod)
+    mod = transform.InjectVirtualThread()(mod)
+
+    # phase 3 _static_lower_phase_3
+    mod = transform.StorageRewriteCCE()(mod)
+    mod = transform.ReorderProcess()(mod)
+    if cfg[tik] and cfg[dynamic_tik]:
+        mod = transform.TikDynamicShapeAllocMem()(mod)
+    mod = transform.UnrollLoop()(mod)
+
+    mod = transform.AutoFuseBuffer()(mod)
+    mod = transform.SetCacheMode()(mod)
+    mod = transform.Simplify()(mod)
+    mod = transform.GMConflictElimination()(mod)
+    mod = transform.MarkScalarCoreType()(mod)
+
+    # phase 4 _static_lower_phase_4
+    mod = transform.JumpInstructionElimination()(mod)
+    mod = transform.InjectSync()(mod)
+    mod = transform.PackIntrinArgConfig()(mod)
+    mod = transform.RemoveAccessPtrMSG()(mod)
+    mod = transform.Simplify()(mod)
+    mod = transform.GmAddrPrompt()(mod)
+    mod = transform.InsertCheckInvalidAccessOfDDR()(mod)
+    mod = transform.RemoveNoOp()(mod)
+    mod = transform.DeviceMark()(mod)
+    if cfg[instrument_bound_checkers]:
+        mod = transform.InstrumentBoundCheckers()(mod)
+    mod = transform.ConvertFloorDivToTruncDiv()(mod)
+    mod = transform.BuildVirtualCore()(mod)
+
+    _count_time(mod)
+    mod = transform.SplitCoreCode()(mod)
+    generate_cce_code(mod, "cce", None)
+
+
+def build_tbe_codegen(kernel_name, stmt_json, arg_json, attr, ascend_type=None):
+    import sys
+    copy_modules = sys.modules.copy()
+    clean_env()
+
+    print("build_cce_for_akg")
+    import tbe
+    from tbe.common.buildcfg.default_buildcfg import cce_default_static_build_config
+    from tbe.common.buildcfg.ascend import AscendPassContext
+    from tbe.common.buildcfg.buildcfg_mapping import dump_cce_code, save_temp_cce_file, disable_vectorize, \
+        instrument_bound_checkers, partition_const_loop, auto_unroll_max_step, auto_unroll_max_depth, \
+        auto_unroll_max_extent, unroll_explicit, dynamic_shape, enable_multicore_sync_with_atomic, \
+        kernel_meta_parent_dir
+    cfg = cce_default_static_build_config.copy()
+    cfg[dump_cce_code] = False
+    cfg[save_temp_cce_file] = True
+    cfg[disable_vectorize] = False
+    cfg[instrument_bound_checkers] = False
+    cfg[partition_const_loop] = False
+    cfg[auto_unroll_max_step] = 0
+    cfg[auto_unroll_max_depth] = 8
+    cfg[auto_unroll_max_extent] = 0
+    cfg[unroll_explicit] = True
+    cfg[dynamic_shape] = False
+    cfg[enable_multicore_sync_with_atomic] = True
+    cfg[kernel_meta_parent_dir] = os.path.realpath(os.getenv('MS_COMPILER_CACHE_PATH', './'))
+    if ascend_type is None:
+        ascend_type = "Ascend910"
+    auto_init_soc(ascend_type)
+
+    stmt = tbe.tvm.ir.load_json(stmt_json)
+    arg_list = []
+    for buff in arg_json:
+        arg_list.append(tbe.tvm.ir.load_json(buff))
+
+    is_dynamic = attr.get("dynamic", False)
+    workspace_dict = attr.get("workspace", None)
+    with AscendPassContext(config=cfg):
+        build_npu_for_akg(kernel_name,
+                          stmt,
+                          arg_list,
+                          is_dynamic=is_dynamic,
+                          cfg=cfg)
+    postfixs = [".o", ".cce", ".json"]
+    is_success = copy_to_akg_kernel_meta(kernel_name, postfixs, workspace_dict)
+    sys.modules = copy_modules
+    return is_success

mindspore/_akg/akg/utils/util.py

@@ -19,6 +19,7 @@ import os
 import json
 import types
 import akg.tvm
+from akg.utils.format_transform import get_shape
 
 
 def parse_int_const(value):
@@ -28,6 +29,49 @@ def parse_int_const(value):
         return value.value
     return None
 
+def parse_workspace_array(workspace):
+    total_bytes = []
+    workspace_list = get_shape(workspace)
+    for ws in workspace_list:
+        total_bytes.append(parse_int_const(ws))
+
+    if total_bytes is None:
+        return None
+
+    workspace_dict = {
+        "num": len(total_bytes),
+        "size": total_bytes
+    }
+    return workspace_dict
+
+def parse_workspace_map(attrs):
+    if not isinstance(attrs, akg.tvm.container.Map):
+        return None
+    if "workspace" not in attrs:
+        return None
+
+    workspace_map = attrs["workspace"]
+    if not isinstance(workspace_map, akg.tvm.container.Map):
+        return None
+    if "num" not in workspace_map or "size" not in workspace_map:
+        return None
+
+    worksapce_num =  parse_int_const(workspace_map["num"])
+    if not isinstance(workspace_map["size"], akg.tvm.container.Array):
+        return None
+
+    workspace_size = []
+    tt = get_shape(workspace_map["size"])
+    for item in tt:
+        workspace_size.append(parse_int_const(item))
+    if len(workspace_size) != worksapce_num:
+        return None
+
+    workspace_dict = {
+        "num": worksapce_num,
+        "size": workspace_size
+    }
+    return workspace_dict
 
 def parse_workspace(workspace):
     if not isinstance(workspace, akg.tvm.container.Map):
@@ -35,7 +79,10 @@ def parse_workspace(workspace):
 
     total_bytes = 0
     if "total_bytes" in workspace:
-        total_bytes = parse_int_const(workspace["total_bytes"])
+        ws = workspace["total_bytes"]
+        if isinstance(ws, akg.tvm.container.Array):
+            return parse_workspace_array(ws)
+        total_bytes = parse_int_const(ws)
 
     if total_bytes is None or total_bytes == 0:
         return None
@@ -63,3 +110,11 @@ def write_code(js_dict, fname):
     with os.fdopen(os.open(fname, os.O_WRONLY | os.O_CREAT, 0o400), 'w') as f:
         json.dump(js_dict, f, sort_keys=True, indent=4, separators=(',', ':'))
 
+def get_ascend_type(desc):
+    if "target_info" not in desc.keys():
+        return None
+
+    target_info_type = desc["target_info"]
+    if target_info_type.get("arch"):
+        return target_info_type.get("arch")
+    return None

mindspore/_check_jit_forbidden_api.py

@@ -96,9 +96,11 @@ def is_invalid_or_jit_forbidden_method(obj, obj_type, attr):
     if not hasattr(obj, attr):
         raise AttributeError(f"'{obj_type}' object has no attribute '{attr}'")
     method = getattr(obj, attr)
-    if not hasattr(method, "__module__"):
+    if not hasattr(method, "__module__") or method.__module__ is None:
         return False
     method_info = method.__module__ + '.' + method.__qualname__
     return method_info in _jit_forbidden_method
 
 add_jit_forbidden_module("mindspore.common.initializer")
+add_jit_forbidden_module("mindspore.context")
+add_jit_forbidden_module("mindspore.log")

mindspore/_checkparam.py

@@ -84,21 +84,21 @@ def _check_inc_rel(val, lower, upper, rel):
 def _format_str_one_value(value, rel):
     """format string"""
     if rel == EQ:
-        return "= {}".format(value)
+        return f"= {value}"
     if rel == NE:
-        return "!= {}".format(value)
+        return f"!= {value}"
     if rel == LT:
-        return "< {}".format(value)
+        return f"< {value}"
     if rel == LE:
-        return "<= {}".format(value)
+        return f"<= {value}"
     if rel == GT:
-        return "> {}".format(value)
+        return f"> {value}"
     if rel == GE:
-        return ">= {}".format(value)
+        return f">= {value}"
     if rel == IN:
-        return "in {}".format(value)
+        return f"in {value}"
     if rel == NOT_IN:
-        return "not in {}".format(value)
+        return f"not in {value}"
 
     return ""
 
@@ -106,13 +106,13 @@ def _format_str_one_value(value, rel):
 def _format_str_two_value(val1, val2, rel):
     """format string"""
     if rel == INC_NEITHER:
-        return "({}, {})".format(val1, val2)
+        return f"({val1}, {val2})"
     if rel == INC_LEFT:
-        return "[{}, {})".format(val1, val2)
+        return f"[{val1}, {val2})"
     if rel == INC_RIGHT:
-        return "({}, {}]".format(val1, val2)
+        return f"({val1}, {val2}]"
     if rel == INC_BOTH:
-        return "[{}, {}]".format(val1, val2)
+        return f"[{val1}, {val2}]"
 
     return ""
 
@@ -556,8 +556,7 @@ def check_str_by_regular(target, reg=None, flag=re.ASCII, prim_name=None):
         reg = r"^\w+[0-9a-zA-Z\_\.]*$"
     if re.match(reg, target, flag) is None:
         prim_name = f"For '{prim_name}', the" if prim_name else "The"
-        raise ValueError("{} '{}' is illegal, it must be match regular'{}' by flags'{}.'".format(
-            prim_name, target, reg, flag))
+        raise ValueError(f"{prim_name} '{target}' is illegal, it must be match regular'{reg}' by flags'{flag}.'")
     return True
 
 
@@ -565,11 +564,10 @@ def check_str_by_regular(target, reg=None, flag=re.ASCII, prim_name=None):
 def check_str_and_none_by_regular(target, reg=None, flag=re.ASCII, prim_name=None):
     if reg is None:
         # Named string regular expression
-        reg = r"^\w*[0-9a-zA-Z\_\.]*$"
+        reg = r"^\w*[0-9a-zA-Z\_\.\-]*$"
     if re.match(reg, target, flag) is None:
         prim_name = f"For '{prim_name}', the" if prim_name else "The"
-        raise ValueError("{} '{}' is illegal, it must be match regular'{}' by flags'{}.'".format(
-            prim_name, target, reg, flag))
+        raise ValueError(f"{prim_name} '{target}' is illegal, it must be match regular'{reg}' by flags'{flag}.'")
     return True
 
 
@@ -585,8 +583,7 @@ def check_file_name_by_regular(target, reg=None, prim_name=None):
         reg = r"^[0-9a-zA-Z@\_\-\.\:\/\\]+$"
     if re.match(reg, target) is None:
         prim_name = f"For '{prim_name}', the" if prim_name else "The"
-        raise ValueError("{} '{}' is illegal, it must be match regular '{}'.".format(
-            prim_name, target, reg))
+        raise ValueError(f"{prim_name} '{target}' is illegal, it must be match regular '{reg}'.")
 
     return True
 
@@ -802,6 +799,9 @@ def check_transpose_axis(axes, ndim):
         # if only one argument provided, it must be tuple or list
         if isinstance(perm, list):
             perm = tuple(perm)
+        elif isinstance(perm, int):
+            perm = (perm,)
+            _check_dim()
         else:
             if not isinstance(perm, tuple):
                 raise TypeError(f"For Tensor.transpose, the parameter 'axes' must be a tuple/list, " \
@@ -959,11 +959,6 @@ def is_stub_tensor(tensor):
     return hasattr(tensor, "stub")
 
 
-def is_pack_tensor(tensor):
-    """Whether it is a PackTensor."""
-    return hasattr(tensor, "__pack__")
-
-
 def expanded_shape(ndim, axis_size, axis):
     """
     Returns a shape with size = 1 for all dimensions
@@ -984,8 +979,8 @@ def infer_out_shape(*shapes):
     def _check(items, max_size, shapes):
         for item in items:
             if item not in (1, max_size):
-                raise ValueError(f'For Tensor, the dimension on each axis must be 1 or the max on the axis' \
-                                 f'to support broadcast, but got shapes {shapes,}')
+                raise ValueError(f'For Tensor, the dimension on each axis must be 1 or the max value on the axis' \
+                                 f'to support broadcasting, but got shapes {shapes,}')
     shape_out = ()
     max_len = max([len(it) for it in shapes])
     for i in range(max_len):
@@ -1269,7 +1264,7 @@ def check_input_data(*data, data_class):
             if not ret:
                 data_class_str = tuple(i.__name__ if hasattr(i, '__name__') else i for i in data_class) if isinstance(
                     data_class, (tuple, list)) else (data_class if data_class is None else data_class.__name__)
-                raise TypeError(f'The type of input data must be in the Union({data_class_str}, ' \
+                raise TypeError(f'The types of input data must be in the Union({data_class_str}, ' \
                                 f'tuple[{data_class_str}], list[{data_class_str}], dict[{data_class_str}]), ' \
                                 f'but got type {item if item is None else type(item).__name__}.')
 
@@ -1314,8 +1309,7 @@ def args_type_check(*type_args, **type_kwargs):
             for name, value in argument_dict.items():
                 if name in bound_types:
                     if value is not None and not isinstance(value, bound_types[name]):
-                        raise TypeError("The parameter '{}' must be {}, but got {}"
-                                        .format(name, bound_types[name], type(value)))
+                        raise TypeError(f"The parameter '{name}' must be {bound_types[name]}, but got {type(value)}")
             return func(*args, **kwargs)
 
         return wrapper

mindspore/_extends/graph_kernel/__init__.py

@@ -14,5 +14,4 @@
 # ============================================================================
 """init"""
 from .splitter import split_with_json
-from .expander import get_op_expander, get_expander_op_list
 from .parallel_estimate import estimate_calculation_amount, estimate_ops