mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.11__cp38-cp38-manylinux1_x86_64.whl

mindspore/.commit_id

@@ -1 +1 @@
-__commit_id__ = '[sha1]:5822529e,[branch]:(HEAD,origin/r2.1,r2.1)'
+__commit_id__ = '[sha1]:8c390933,[branch]:(HEAD,origin/r2.2,r2.2)'

mindspore/__init__.py

@@ -20,7 +20,7 @@ from mindspore import common, dataset, mindrecord, train, log, amp
 from mindspore import profiler, communication, numpy, parallel
 from mindspore.common import *
 from mindspore.mindrecord import *
-from mindspore.ops import _op_impl, grad, value_and_grad, vjp, jvp, jacfwd, jacrev, vmap, get_grad
+from mindspore.ops import _op_impl, grad, value_and_grad, vjp, jvp, jacfwd, jacrev, vmap, get_grad, constexpr
 from mindspore.train import *
 from mindspore.log import *
 from mindspore.context import GRAPH_MODE, PYNATIVE_MODE, set_context, get_context, set_auto_parallel_context, \
@@ -31,8 +31,10 @@ from mindspore.profiler import Profiler, EnvProfiler
 from mindspore.parallel import set_algo_parameters, get_algo_parameters, reset_algo_parameters, \
     rank_list_for_transform, transform_checkpoint_by_rank, transform_checkpoints, merge_pipeline_strategys, shard
 from mindspore.rewrite import SymbolTree, ScopedValue, Node, NodeType, TreeNodeHelper
+from mindspore.safeguard import obfuscate_ckpt, load_obf_params_into_net
 from mindspore._check_jit_forbidden_api import get_obj_module_and_name_info, is_jit_forbidden_module, \
     is_invalid_or_jit_forbidden_method
+from mindspore import _install_custom
 
 
 __all__ = ["run_check"]
@@ -43,4 +45,5 @@ __all__.extend(log.__all__)
 __all__.extend(context.__all__)
 __all__.extend(parallel.__all__)
 __all__.extend(rewrite.__all__)
+__all__.extend(safeguard.__all__)
 __all__.append("Profiler")

mindspore/_akg/akg/build_module.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # coding: utf-8
-# Copyright 2019 Huawei Technologies Co., Ltd
+# Copyright 2019-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.
@@ -28,7 +28,6 @@ from akg.tvm import _api_internal
 from akg.tvm import schedule
 
 
-tuning_spaces = None
 help_tiling_level = {
     "None": 0, "General": 1, "Candidates": 2, "Tuning": 3
 }
@@ -36,7 +35,7 @@ EMPTY_CODE = 0
 L0_DEFAULT_TILING = 1
 
 
-def dump_tiling_info(level):
+def dump_tiling_info(level, tuning_spaces = None):
     """Dump tiling info."""
     if tuning_spaces is None:
         return
@@ -107,7 +106,8 @@ def lower(sch, args, shape_params=None, name="default_function", binds=None, att
     ret = _api_internal._Lower(sch, args, shape_params, name,
                                tmp_binds, tmp_attrs, simple_mode,
                                polyhedral, tuning, target, cfg)
-
+    if tmp_attrs is None:
+        tmp_attrs = {}
     level = tmp_attrs.get("help_tiling")
     if attrs.get("use_new_space", False):
         # new space: constraints format
@@ -118,7 +118,6 @@ def lower(sch, args, shape_params=None, name="default_function", binds=None, att
         return space
     elif tuning or (level is not None and level > help_tiling_level['None']):
         level = help_tiling_level['Tuning'] if tuning else level
-        global tuning_spaces
         tuning_spaces = {}
         tuning_spaces["index"] = ret.index_table.asnumpy().tolist()
         tuning_spaces["c1_range"] = ret.c1_tile_range_table.asnumpy().tolist()
@@ -132,7 +131,7 @@ def lower(sch, args, shape_params=None, name="default_function", binds=None, att
         if level >= help_tiling_level["Candidates"]:
             tuning_spaces["tuning_space"] = ret.tiling_candidate.asnumpy().tolist()
         if not tuning:
-            dump_tiling_info(level)
+            dump_tiling_info(level, tuning_spaces)
     return ret
 
 

mindspore/_akg/akg/composite/build_module.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # coding: utf-8
-# Copyright 2020-2022 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -19,19 +19,17 @@ import os
 import json
 from collections.abc import Iterable
 import akg
+import math
 from akg import tvm
 from tvm.autotvm.env import AutotvmGlobalScope
+from akg.utils.util import parse_workspace_map
+from akg.utils.tbe_codegen_utils import build_tbe_codegen
 from akg.utils.kernel_exec import ReturnType, is_symbolic_tiling
 from .split_stitch import split_stitch_attr
-from .construct_args import get_construct_args, get_tune_construct_args, \
-    should_enable_attr, get_stmt_for_tune, add_attrs_in_segment_infos, \
-    update_attrs
 from .construct_args import ConstructType, ConstructKey
 from .construct_args import get_construct_args, get_tune_construct_args, \
     should_enable_attr, get_stmt_for_tune, add_attrs_in_segment_infos
-from .split_stitch import split_stitch_attr
-
-
+from utils.util import get_ascend_type
 def generate_trait(desc):
     """
     generate trait of kernel description
@@ -317,12 +315,12 @@ def merge_attrs(attrs_a, attrs_b):
     return attrs
 
 
-def read_repo_file(repo_file):
+def read_repo_file(repo_file, is_json_load=True):
     if not os.path.exists(repo_file):
         return {}
     with open(repo_file, 'r') as f:
-        repo = json.loads(f.read())
-    return repo
+        repo = f.read()
+    return json.loads(repo) if is_json_load else repo
 
 
 def _get_default_repository_file(process):
@@ -407,6 +405,18 @@ def _update_attrs_ascend(all_ops, attr):
     attr["multicore_loop_switch_hoist"] = "UnsortedSegmentSum" not in all_ops
     return attr
 
+def _cpp_build(attrs, process, poly, segment_tree, segment_infos):
+    if attrs.get("is_tbe_codegen"):
+        func = tvm.get_global_func("lower_composite")
+    else:
+        func = tvm.get_global_func("lower_composite_to_module")
+
+    if "ret_mode" in attrs:
+        return _build_for_tuning(attrs, func, process, segment_tree, segment_infos)
+
+    res = func(process, poly, segment_tree, segment_infos)
+    return res
+
 
 def _build_to_module(desc_s, desc_d, attrs=None, poly=True):
     """
@@ -512,11 +522,7 @@ def _build_to_module(desc_s, desc_d, attrs=None, poly=True):
     segment_tree, segment_infos = get_construct_args(desc_s, attrs, post_funcs)
     process = desc_d["process"]
 
-    func = tvm.get_global_func("lower_composite_to_module")
-    if "ret_mode" in attrs and poly:
-        return _build_for_tuning(attrs, func, process, segment_tree, segment_infos)
-    return func(process, poly, segment_tree, segment_infos)
-
+    return _cpp_build(attrs, process, poly, segment_tree, segment_infos)
 
 def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
     """
@@ -633,14 +639,125 @@ def _build_to_module_ascend(desc_s_in, desc_d_in, attr, use_repo=True):
         ConstructType.STITCH: _stitch_postprocess,
         ConstructType.NORMAL: _normal_postprocess,
     }
-    segment_tree, segment_infos = get_construct_args(desc_s_in, attr, post_funcs)
     process = desc_d_in["process"]
+    kernel_name = desc_d_in['op']
+    ascend_type = get_ascend_type(desc_d_in)
+    ascend_type_to_section = {"Ascend910A": "1.6", "Ascend310P3": "1.7",
+                              "Ascend910B1": "2.1", "Ascend910B2": "2.2", "Ascend910B3": "2.3", "Ascend910B4": "2.4"}
+    if ascend_type is not None:
+        section = ascend_type_to_section.get(ascend_type, "1.6")
+        config_func = akg.tvm.get_global_func("cce.set_product_section")
+        config_func(section)
+        if section >= "2.1":
+            attr["is_tbe_codegen"] = True
+            attr["pragma_modshift"] = True
+    segment_tree, segment_infos = get_construct_args(desc_s_in, attr, post_funcs)
 
-    func = tvm.get_global_func("lower_composite_to_module")
-    if "ret_mode" in attr:
-        return _build_for_tuning(attr, func, process, segment_tree, segment_infos)
-    return func(process, True, segment_tree, segment_infos)
-
+    if desc_d_in.get("enable_cce_lib"):
+        attr["enable_cce_lib"] = True
+        return _build_to_module_ascend_lib(desc_s_in, kernel_name)
+
+    poly = True
+    res = _cpp_build(attr, process, poly, segment_tree, segment_infos)
+    if attr.get("is_tbe_codegen"):
+        stmt_json = akg.tvm.save_json(res[0], "0.8.0")
+        args_json = []
+        for buf in res[1]:
+            args_json.append(akg.tvm.save_json(buf, "0.8.0"))
+        
+        workspace_dict = parse_workspace_map(res[2])
+        if workspace_dict is not None:
+            attr["workspace"] = workspace_dict
+
+        is_success = build_tbe_codegen(kernel_name, stmt_json, args_json, attr, ascend_type)
+        if not is_success:
+            raise TypeError("npu_inference codegen failed.")
+        return kernel_name
+    return res
+
+def _build_to_module_ascend_lib(desc_s_in, kernel_name):
+    def _get_all_shape(shapes):
+        shape_split = shapes.split(".")
+        shape_list = []
+        for shape in shape_split:
+            if "-" in shape:
+                tmp_shape = shape.split("-")[0]
+                for _ in range(shape.count("-") + 1):
+                    shape_list.append(tmp_shape)
+            else:
+                shape_list.append(shape)
+        return shape_list
+    
+    def _get_tiling_info(desc_s):
+        compute, shape, dtype = generate_trait(desc_s)
+        tiling_info = {}
+        if "MatMul" in compute:
+            trans_a = compute.split("_")[1]
+            trans_b = compute.split("_")[-1].split(".")[0]
+
+            shape_list = _get_all_shape(shape)
+            bias_flag = int(len(shape_list) > 3)
+            tensor_A = shape_list[0]
+            tensor_B = shape_list[1]
+
+            tensor_A_split = tensor_A.split("_")
+            if len(tensor_A_split) > 2:
+                batch_size = int(tensor_A.split("_")[0])
+            else:
+                batch_size = 1
+            if trans_a == "1":
+                M = int(tensor_A_split[-1])
+                K = int(tensor_A_split[-2])
+            else:
+                M = int(tensor_A_split[-2])
+                K = int(tensor_A_split[-1])
+
+            if trans_b == "1":
+                N = int(tensor_B.split("_")[-2])
+            else:
+                N = int(tensor_B.split("_")[-1])
+            tensor_A_type = str(dtype.split("-")[0])
+            tiling_info = {"batch_size":batch_size, "M": M, "N": N, "K": K, "trans_a": int(trans_a), "trans_b": int(trans_b),
+                           "tensor_A_type": tensor_A_type, "bias_flag": bias_flag, "op_type": "MatMul"}
+        elif "PagedAttention" in compute or "PagedAttentionMask" in compute:
+            shape_list = _get_all_shape(shape)
+            query = shape_list[0]
+            key_cache = shape_list[1]
+            table_shape = shape_list[3]
+
+            num_tokens = int(query.split("_")[0])
+            num_heads = int(query.split("_")[1])
+            embedding_size = int(query.split("_")[2])
+            num_blocks = int(key_cache.split("_")[0])
+            block_size = int(key_cache.split("_")[1])
+            kv_heads = int(key_cache.split("_")[2])
+
+            max_num_blocks_per_query = int(table_shape.split("_")[1])
+            tor = float(1.0 / math.sqrt(1.0 * embedding_size))
+
+            tiling_info = {"num_tokens": num_tokens, "num_heads": num_heads, "embedding_size": embedding_size, 
+                           "num_blocks": num_blocks, "block_size": block_size, "max_num_blocks_per_query": max_num_blocks_per_query,
+                           "tor": tor, "kv_heads": kv_heads, "op_type": "PagedAttention"}
+            if "PagedAttentionMask" in compute:
+                mask_shape = shape_list[5]
+                tiling_info["mask"] = list(map(int, mask_shape.split("_")))
+                tiling_info["op_type"] = "PagedAttentionMask"
+        elif "ReshapeAndCache" in compute:
+            shape_list = _get_all_shape(shape)
+            kv = shape_list[0]
+
+            num_tokens = int(kv.split("_")[0])
+            num_heads = int(kv.split("_")[1])
+            head_size = int(kv.split("_")[2])
+
+            tiling_info = {"num_tokens": num_tokens, "num_heads": num_heads, "head_size": head_size, 
+                           "op_type": "ReshapeAndCache"}
+        return tiling_info
+
+    func = tvm.get_global_func("build_cce_lib")
+    tiling_info = _get_tiling_info(json.loads(desc_s_in))
+    func(kernel_name, tiling_info, None)
+    return kernel_name
 
 def _set_backend(desc_d):
     desc_d_process = desc_d
@@ -739,4 +856,4 @@ def get_tiling_space(kernel_desc, level=1, attr=None):
         spaces['c0_mod'] = ret.c0_tile_mod_table.asnumpy().tolist()
         if level >= 2:
             spaces['tuning_space'] = ret.tiling_candidate.asnumpy().tolist()
-    return spaces
+    return spaces

mindspore/_akg/akg/composite/split_stitch.py

@@ -240,11 +240,10 @@ def parse_merged_json(desc_d, stitch_tensor_name, input_tensor_name, output_tens
     return extra_subgraph_output, list(final_output_list), final_output_within_graph
 
 
-def process_inplace_assign(op_info, inplace_assign_map, fake_output_list):
-    inplace_assign_map[op_info['output_desc'][0]['tensor_name']] = op_info['input_desc'][0][0]['tensor_name']
-    if op_info['attr'][0]['name'] == 'fake_output' and op_info['attr'][0]['value'] == 1:
-        fake_output_list.append(op_info['output_desc'][0]['tensor_name'])
-    return inplace_assign_map, fake_output_list
+def process_assign(op_info, assign_map, fake_output_list):
+    assign_map[op_info['output_desc'][0]['tensor_name']] = op_info['input_desc'][0][0]['tensor_name']
+    fake_output_list.append(op_info['output_desc'][0]['tensor_name'])
+    return assign_map, fake_output_list
 
 
 def process_each_subgraph(op_info, req_map, input_tensor_name, output_tensor_name, stitch_node_list, sg):
@@ -291,16 +290,16 @@ def process_each_subgraph(op_info, req_map, input_tensor_name, output_tensor_nam
 
 
 def collect_subgraph_info(desc_d, sub_stitch_graphs, req_map, input_tensor_name, output_tensor_name, stitch_node_list):
-    inplace_assign_map = {}
+    assign_map = {}
     fake_output_list = []
     # traversal desc_d by reverse topologically order.
     for i in range(len(desc_d['op_desc']) - 1, -1, -1):
         op_info = desc_d['op_desc'][i]
-        if op_info['name'] == "InplaceAssign":
-            inplace_assign_map, fake_output_list = process_inplace_assign(op_info, inplace_assign_map, fake_output_list)
+        if op_info['name'] == "Assign":
+            assign_map, fake_output_list = process_assign(op_info, assign_map, fake_output_list)
         for sg in sub_stitch_graphs:
             process_each_subgraph(op_info, req_map, input_tensor_name, output_tensor_name, stitch_node_list, sg)
-    return sub_stitch_graphs, inplace_assign_map, fake_output_list
+    return sub_stitch_graphs, assign_map, fake_output_list
 
 
 def sub_graph_info(sub_graph, desc_d):
@@ -391,7 +390,7 @@ def stitch_json_split(desc_d):
     for i, stitch_op in enumerate(stitch_node):
         sub_stitch_graphs.append(Graph(stitch_op))
 
-    sub_stitch_graphs, inplace_assign_map, fake_output_list = \
+    sub_stitch_graphs, assign_map, fake_output_list = \
         collect_subgraph_info(desc_d, sub_stitch_graphs, req_map,
                               input_tensor_name, complement_output, stitchnode_list)
     # reverse op order to generate topological subgraph
@@ -418,7 +417,7 @@ def stitch_json_split(desc_d):
         else:
             clean_info = reuse_map.get(fake_op)
             reuse_map.pop(fake_op)
-        clean_op_map[inplace_assign_map.get(fake_op)] = clean_info
+        clean_op_map[assign_map.get(fake_op)] = clean_info
 
     if not alloc_map:
         alloc_map['EMPTY'] = []

mindspore/_akg/akg/ms/info_version_adapt.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Huawei Technologies Co., Ltd
+# Copyright 2022-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.
@@ -26,11 +26,64 @@ v1(r1.8):
    * on cpu, it contains "arch", "system" and "feature" fields.
    * on gpu, it contains "compute_capability" and "sm_count" fields,
      the "compute_capability" already exists from v0.
+
+v2(r2.2):
+1. For all the ops, the number of inputs in akg info is equal to the number of
+real inputs, which means that no inputs are converted to attributes in the info.
+AKG needs to handle this process by itself from now on.
+2. For those inputs need to be converted to attributes, a key in "input_desc" of
+this op named "value" is guaranteed, which contains the value of this input.
 """
 
 import logging
 
 
+def convert_input_to_attr(kernel_info:dict):
+    op_info = {
+        "Reshape": [(1, "shape")],
+        "ReduceMax": [(1, "axis")],
+        "ReduceMin": [(1, "axis")],
+        "ReduceSum": [(1, "axis")],
+        "Transpose": [(1, "perm")],
+        "ExpandDims": [(1, "axis")],
+        "Tile": [(1, "multiples")],
+        "StridedSlice": [(3, "strides"), (2, "end"), (1, "begin")],
+        "OneHot": [(1, "depth")],
+        "Gather": [(2, "axis")],
+        "UnsortedSegmentSum": [(2, "num_segments")],
+        "CumSum": [(1, "axis")],
+    }
+
+    int_input_required_ops = {
+        "OneHot",
+        "UnsortedSegmentSum",
+    }
+
+    ops = kernel_info["op_desc"]
+    for op in ops:
+        op_name = op["name"]
+        if op_name in op_info:
+            attr = []
+            if op["attr"]:
+                attr = op["attr"]
+            for input_info in op_info[op_name]:
+                input_index = input_info[0]
+                input_name = input_info[1]
+                if input_index >= len(op["input_desc"]):
+                    continue
+                input_desc_i = op["input_desc"].pop(input_index)
+                input_value = input_desc_i[0]["value"]
+                input_dtype = "listInt"
+                if op_name not in int_input_required_ops and isinstance(input_value, int):
+                    input_value = [input_value]
+                if isinstance(input_value, int):
+                    input_dtype = "int"
+                attr.append(
+                    {"name": input_name, "dtype": input_dtype, "value": input_value}
+                )
+            op["attr"] = attr
+
+
 def _adapt_version_0(kernel_info: dict):
     if "compute_capability" in kernel_info:
         # move the "compute_capability" into "target_info" field.
@@ -39,8 +92,21 @@ def _adapt_version_0(kernel_info: dict):
         kernel_info.pop("compute_capability")
 
 
+def _adapt_version_1(kernel_info:dict):
+    pass
+
+
+def _adapt_version_2(kernel_info:dict):
+    if kernel_info.get("version", 0) < 2:
+        return
+    else:
+        convert_input_to_attr(kernel_info)
+
+
 _adapt_func_list = [
     _adapt_version_0,
+    _adapt_version_1,
+    _adapt_version_2,
 ]
 
 

mindspore/_akg/akg/tvm/api.py

@@ -155,21 +155,22 @@ def load_json(json_str):
     return _api_internal._load_json(json_str)
 
 
-def save_json(node):
+def save_json(node, version="0.6.0"):
     """Save tvm object as json string.
 
     Parameters
     ----------
     node : Node
         A TVM Node object to be saved.
+    version : str
+        A TVM version string, only support "0.6.0", "0.8.0" for now.
 
     Returns
     -------
     json_str : str
         Saved json string.
     """
-    return _api_internal._save_json(node)
-
+    return _api_internal._save_json(node, version)
 
 def var(name="tindex", dtype=int32):
     """Create a new variable with specified name and dtype

mindspore/_akg/akg/tvm/autotvm/__init__.py

@@ -42,8 +42,7 @@ from . import env
 #from . import tophub
 
 # some shortcuts
-from .measure import measure_option, MeasureInput, MeasureResult, MeasureErrorNo, \
-    LocalBuilder, LocalRunner, RPCRunner
+from .measure import measure_option, MeasureInput, MeasureResult, MeasureErrorNo, LocalBuilder
 from .tuner import callback
 from .task import template, get_config, create, ConfigSpace, ConfigEntity, \
     register_topi_compute, register_topi_schedule, \

mindspore/_akg/akg/tvm/autotvm/graph_tuner/base_graph_tuner.py

@@ -441,11 +441,7 @@ class BaseGraphTuner(object):
         builder = autotvm.LocalBuilder(n_parallel=n_parallel, build_func=build_func)
         runner = autotvm.LocalRunner(number=min_exec_num, repeat=1, timeout=timeout)
         if use_rpc:
-            if device_key is None:
-                raise RuntimeError("device_key need to be set to use rpc tracker mode.")
-            runner = autotvm.measure.RPCRunner(device_key, host, port, n_parallel=n_parallel,
-                                               number=min_exec_num, repeat=1,
-                                               timeout=timeout)
+            raise RuntimeError("rpc is not supported.")
         measure_option = autotvm.measure_option(builder=builder, runner=runner)
         for args in args_list:
             data, in_layout, out_layout = args

mindspore/_akg/akg/tvm/autotvm/measure/__init__.py

@@ -18,6 +18,6 @@
 
 from .measure import MeasureInput, MeasureResult, MeasureErrorNo, measure_option, \
     create_measure_batch
-from .measure_methods import LocalBuilder, LocalRunner, RPCRunner, request_remote
+from .measure_methods import LocalBuilder
 from .executor import Executor
 from .local_executor import LocalExecutor

mindspore/_akg/akg/tvm/autotvm/measure/measure.py

@@ -46,7 +46,7 @@ class MeasureResult(namedtuple("MeasureResult", ["costs", "error_no", "all_cost"
     error_no: int
         Denote error type, defined by MeasureErrorNo
     all_cost: float
-        All cost of this measure, including rpc, compilation, test runs
+        All cost of this measure, including compilation, test runs
     timestamp: float
         The absolute time stamp when we finish measurement.
     """
@@ -191,15 +191,6 @@ def measure_option(builder, runner):
     >>>         timeout=5)
     >>> )
 
-    # example setting for using remote devices
-    >>> measure_option = autotvm.measure_option(
-    >>>    builder=autotvm.LocalBuilder(),  # use all local cpu cores for compilation
-    >>>    runner=autotvm.RPCRunner(
-    >>>        'rasp3b', 'locahost', 9190, # device key, host and port of the rpc tracker
-    >>>        number=4,
-    >>>        timeout=4) # timeout of a run on the device. RPC request waiting time is excluded.
-    >>>)
-
     Note
     ----
     To make measurement results accurate, you should pick the correct value for the argument