mindspore 2.0.0a0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/operations/comm_ops.py

@@ -1,4 +1,4 @@
-# Copyright 2020 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,8 +19,7 @@ from __future__ import absolute_import
 from __future__ import division
 
 from mindspore.common import Tensor
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.communication.management import get_rank, get_group_size, GlobalComm, _get_group, _host_distribute
 from mindspore.common import dtype as mstype
 from mindspore.ops.primitive import PrimitiveWithInfer, PrimitiveWithCheck, Primitive, prim_attr_register
@@ -54,11 +53,11 @@ class ReduceOp:
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with multiple devices.
 
@@ -114,11 +113,10 @@ class AllReduce(Primitive):
     Reduces the tensor data across all devices in such a way that all devices will get the same final result.
 
     Note:
-        The operation of AllReduce does not support "prod" currently.
         The tensors must have the same shape and format in all processes of the collection.
 
     Args:
-        op (str): Specifies an operation used for element-wise reductions, like sum, max, and min.
+        op (str): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
                   On the CPU, only 'sum' is supported. Default: ReduceOp.SUM.
         group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
 
@@ -132,7 +130,6 @@ class AllReduce(Primitive):
     Raises:
         TypeError: If any of `op` and `group` is not a str,
                    or fusion is not an integer, or the input's dtype is bool.
-        ValueError: If the `op` is "prod".
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -143,11 +140,11 @@ class AllReduce(Primitive):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with multiple devices.
 
@@ -196,7 +193,8 @@ class AllGather(PrimitiveWithInfer):
     Gathers tensors from the specified communication group.
 
     Note:
-        The tensors must have the same shape and format in all processes of the collection.
+        - The tensors must have the same shape and format in all processes of the collection.
+        - Currently only supports GRAPH_MODE and it should be called in Cell.
 
     Args:
         group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
@@ -222,11 +220,11 @@ class AllGather(PrimitiveWithInfer):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with 2 devices.
 
@@ -263,7 +261,7 @@ class AllGather(PrimitiveWithInfer):
         validator.check_value_type('group', _get_group(group), (str,), self.name)
         self.rank = get_rank(_get_group(group))
         self.rank_size = get_group_size(_get_group(group))
-        validator.check('rank', self.rank, 'rank_size', self.rank_size, Rel.LT, self.name)
+        validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
         self.add_prim_attr('rank_size', self.rank_size)
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('fusion', 0)
@@ -300,13 +298,14 @@ class _MiniStepAllGather(PrimitiveWithInfer):
         validator.check_value_type('group', _get_group(group), (str,), self.name)
         self.rank = get_rank(_get_group(group))
         self.rank_size = get_group_size(_get_group(group))
-        validator.check('rank', self.rank, 'rank_size', self.rank_size, Rel.LT, self.name)
+        validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
         self.add_prim_attr('rank_size', self.rank_size)
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('fusion', 1)
         self.grad_accumulation_step = grad_accumulation_step
         self.mean_flag = mean_flag
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, z_shape):
         validator.check_positive_int(len(x_shape), "x shape", self.name)
@@ -331,15 +330,17 @@ class _MicroStepAllGather(PrimitiveWithInfer):
     @prim_attr_register
     def __init__(self, group=GlobalComm.WORLD_COMM_GROUP, mean_flag=None):
         validator.check_value_type('group', _get_group(group), (str,), self.name)
-        self.rank = get_rank(_get_group(group))
-        self.rank_size = get_group_size(_get_group(group))
-        validator.check('rank', self.rank, 'rank_size', self.rank_size, Rel.LT, self.name)
-        self.add_prim_attr('rank_size', self.rank_size)
-        self.add_prim_attr('group', _get_group(group))
-        self.add_prim_attr('fusion', 1)
-        self.add_prim_attr('do_mirror', False)
-        self.mean_flag = mean_flag
-        self.add_prim_attr('order_enforce_skip', True)
+        self.rank_size = 1
+        if group != "":
+            self.rank = get_rank(_get_group(group))
+            self.rank_size = get_group_size(_get_group(group))
+            validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
+            self.add_prim_attr('rank_size', self.rank_size)
+            self.add_prim_attr('group', _get_group(group))
+            self.add_prim_attr('fusion', 1)
+            self.add_prim_attr('do_mirror', False)
+            self.mean_flag = mean_flag
+            self.add_prim_attr('order_enforce_skip', True)
 
     def infer_shape(self, x_shape, z_shape):
         validator.check_positive_int(len(x_shape), "x shape", self.name)
@@ -383,9 +384,9 @@ class _HostAllGather(PrimitiveWithInfer):
         if group is None:
             raise ValueError(f"For '{self.name}', the 'group' cannot be None, but got {group}.")
         validator.check_value_type('group', group, (tuple, list), self.name)
-        validator.check_int(len(group), 2, Rel.GE, "group size", self.name)
+        validator.check_int(len(group), 2, validator.GE, "group size", self.name)
         for r in group:
-            validator.check_int_range(r, 0, 7, Rel.INC_BOTH, "rank_id", self.name)
+            validator.check_int_range(r, 0, 7, validator.INC_BOTH, "rank_id", self.name)
             validator.check_value_type("rank_id", r, (int,), self.name)
         self.group_size = len(group)
         self.add_prim_attr('group', group)
@@ -410,7 +411,7 @@ class ReduceScatter(Primitive):
     r"""
     Reduces and scatters tensors from the specified communication group.
     For more details about it, please refer to `Distributed Set Communication Primitives - ReduceScatter \
-    <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/communicate_ops.html#reducescatter>`_ .
+    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#reducescatter>`_ .
 
     Note:
         The tensors must have the same shape and format in all processes of the collection.
@@ -441,11 +442,11 @@ class ReduceScatter(Primitive):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with 2 devices.
 
@@ -522,9 +523,9 @@ class _HostReduceScatter(PrimitiveWithInfer):
             raise ValueError(f"For '{self.name}', the 'group' cannot be None, but got {group}.")
         validator.check_value_type('op', op, (type(ReduceOp.SUM),), self.name)
         validator.check_value_type('group', group, (tuple, list), self.name)
-        validator.check_int(len(group), 2, Rel.GE, "group size", self.name)
+        validator.check_int(len(group), 2, validator.GE, "group size", self.name)
         for r in group:
-            validator.check_int_range(r, 0, 7, Rel.INC_BOTH, "rank_id", self.name)
+            validator.check_int_range(r, 0, 7, validator.INC_BOTH, "rank_id", self.name)
             validator.check_value_type("rank_id", r, (int,), self.name)
         self.op = op
         self.group_size = len(group)
@@ -563,7 +564,7 @@ class Broadcast(PrimitiveWithInfer):
         - **input_x** (tuple[Tensor]) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
 
     Outputs:
-        Tensor, has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.
+        tuple[Tensor], Tensor has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.
         The contents depend on the data of the `root_rank` device.
 
     Raises:
@@ -578,11 +579,11 @@ class Broadcast(PrimitiveWithInfer):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with multiple devices.
 
@@ -628,7 +629,7 @@ class Broadcast(PrimitiveWithInfer):
         if not isinstance(x_dtype, tuple):
             raise TypeError(f"For '{self.name}', the 'input_x' must be a tuple, but got {type(x_dtype).__name__}!")
         for _ele in x_dtype:
-            check_collective_target_dtype('x', _ele, self.name)
+            check_collective_target_dtype('input_x', _ele, self.name)
         return x_dtype
 
 
@@ -698,11 +699,11 @@ class NeighborExchange(Primitive):
         The user needs to preset
         communication environment variables before running the following example, please check the details on the
         official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r2.0.0-alpha/api_python/mindspore.ops.html#communication-operator>`_.
+        <https://www.mindspore.cn/docs/en/r2.0/api_python/mindspore.ops.html#communication-operator>`_.
 
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/zh-CN/r2.0.0-alpha/parallel/communicate_ops.html#注意事项>`_.
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to.
@@ -778,7 +779,7 @@ class AlltoAll(PrimitiveWithInfer):
     Note:
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/zh-CN/r2.0.0-alpha/parallel/communicate_ops.html#注意事项>`_.
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
 
     Args:
         split_count (int): On each process, divide blocks into split_count number.
@@ -793,11 +794,9 @@ class AlltoAll(PrimitiveWithInfer):
         Tensor. If the shape of input tensor is :math:`(x_1, x_2, ..., x_R)`, then the shape of output tensor is
         :math:`(y_1, y_2, ..., y_R)`, where:
 
-        :math:`y_{split\_dim} = x_{split\_dim} / split\_count`
-
-        :math:`y_{concat\_dim} = x_{concat\_dim} * split\_count`
-
-        :math:`y_{other} = x_{other}`.
+        - :math:`y_{split\_dim} = x_{split\_dim} / split\_count`
+        - :math:`y_{concat\_dim} = x_{concat\_dim} * split\_count`
+        - :math:`y_{other} = x_{other}`.
 
     Raises:
         TypeError: If group is not a string.
@@ -811,11 +810,11 @@ class AlltoAll(PrimitiveWithInfer):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with 8 devices.
 
@@ -885,13 +884,13 @@ class NeighborExchangeV2(Primitive):
     NeighborExchangeV2 sends data from the local rank to ranks in the `send_rank_ids`,
     as while receive data from `recv_rank_ids`. Please refer to
     `Distributed Set Communication Primitives - NeighborExchangeV2 \
-    <https://www.mindspore.cn/tutorials/experts/zh-CN/r2.0.0-alpha/parallel/communicate_ops.html#neighborexchangev2>`_
+    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#neighborexchangev2>`_
     to learn about how the data is exchanged between neighborhood devices.
 
     Note:
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/zh-CN/r2.0.0-alpha/parallel/communicate_ops.html#注意事项>`_.
+        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to. 8 rank_ids represents 8 directions, if one
@@ -929,11 +928,11 @@ class NeighborExchangeV2(Primitive):
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
             Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_ascend.html#preparations>`_
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
             for more details.
 
             For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/parallel/train_gpu.html#preparation>`_ .
+            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
 
             This example should be run with 2 devices.
 
@@ -1032,6 +1031,7 @@ class _MirrorMiniStepOperator(PrimitiveWithInfer):
         self.mean_flag = mean_flag
         self.grad_accumulation_step = grad_accumulation_step
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, z_shape):
         return x_shape
@@ -1137,6 +1137,7 @@ class _VirtualAssignAdd(PrimitiveWithInfer):
     def __init__(self):
         """Initialize _VirtualAssignAdd."""
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, y_shape):
         return x_shape
@@ -1187,6 +1188,7 @@ class _MirrorMicroStepOperator(PrimitiveWithInfer):
         self.dev_num = dev_num
         self.mean_flag = mean_flag
         self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
 
     def infer_shape(self, x_shape, z_shape):
         return x_shape

mindspore/ops/operations/control_ops.py

@@ -17,7 +17,7 @@
 
 from __future__ import absolute_import
 from mindspore.ops.primitive import PrimitiveWithInfer, prim_attr_register
-from mindspore._checkparam import Validator as validator, Rel
+from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 
 
@@ -76,7 +76,7 @@ class GeSwitch(PrimitiveWithInfer):
         raise NotImplementedError
 
     def infer_shape(self, data, pred):
-        validator.check_int_range(len(pred), 0, 1, Rel.INC_BOTH, "pred rank", self.name)
+        validator.check_int_range(len(pred), 0, 1, validator.INC_BOTH, "pred rank", self.name)
         return data, data
 
     def infer_dtype(self, data_type, pred_type):

mindspore/ops/operations/custom_ops.py

@@ -26,6 +26,7 @@ import importlib
 import platform
 import subprocess
 import numpy as np
+import mindspore as ms
 from mindspore._c_expression import Oplib, typing
 from mindspore import context
 from mindspore.common import Tensor
@@ -61,6 +62,27 @@ def _get_cache_path():
     return cache_path
 
 
+def _get_cuda_bare_metal_version():
+    """
+    Automatically get the cuda version.
+
+    Returns:
+        tuple(str), the version of cuda of the platform.ss
+    """
+    raw_output = subprocess.check_output(["nvcc", "-V"],
+                                         universal_newlines=True)
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    version_major = release[0]
+    version_idx = release_idx + 1
+    version = output[version_idx].split(".")
+    version_middle = version[1] if len(version) > 1 else 0
+    version_minor = version[2] if len(version) > 2 else 0
+
+    return int(version_major), int(version_middle), int(version_minor)
+
+
 def _compile_aot(file):
     """
     Automatically compile the source file for custom aot
@@ -105,18 +127,13 @@ def _compile_aot(file):
             cmd += ["--use_fast_math", "--expt-relaxed-constexpr"]
             cmd += ["-D_GLIBCXX_USE_CXX11_ABI=0"]
 
-            def _get_cuda_bare_metal_version():
-                raw_output = subprocess.check_output(["nvcc", "-V"],
-                                                     universal_newlines=True)
-                output = raw_output.split()
-                release_idx = output.index("release") + 1
-                release = output[release_idx].split(".")
-                version_major = release[0]
-
-                return int(version_major)
-
-            if _get_cuda_bare_metal_version() >= 11:
+            v_major, v_mid, v_minor = _get_cuda_bare_metal_version()
+            if v_major >= 11:
                 cmd += ["-gencode", "arch=compute_80,code=sm_80", "--expt-extended-lambda"]
+            elif v_major == 10 and not(v_mid >= 1 and v_minor >= 168):
+                logger.warning("The current version of nvcc, V{}.{}.{},  might have unfixed issues with std string, "
+                               "which will lead to errors in aot custom op with attrs."
+                               "The version higher than V10.1.168 is recommended".format(v_major, v_mid, v_minor))
             cmd += [include_file, "-o", func_path, file]
         else:
             raise ValueError("The source file must be a cc/cpp/cu file, but get: {}".format(file))
@@ -142,10 +159,10 @@ class Custom(ops.PrimitiveWithInfer):
     function if needed. Then these `Custom` objects can be directly used in neural networks.
     Detailed description and introduction of user-defined operators, including correct writing of parameters,
     please refer to `Custom Operators Tutorial
-    <https://www.mindspore.cn/tutorials/experts/en/r2.0.0-alpha/operation/op_custom.html>`_ .
+    <https://www.mindspore.cn/tutorials/experts/en/r2.0/operation/op_custom.html>`_ .
 
     .. warning::
-        This is an experimental prototype that is subject to change.
+        This is an experimental API that is subject to change.
 
     .. note::
         The supported platforms are determined by the input `func_type`. The supported platforms are as follows:
@@ -437,6 +454,9 @@ class Custom(ops.PrimitiveWithInfer):
     registered_func = {}
     attr_dict = {}  # Save input_names and attr_names for func.
     compiled_bin = []  # Save names for compiled bin.
+    tbe_path_checked = []  # Save paths for tbe functions which is safe to be imported as module.
+    tbe_path_failed = []  # Save paths for tbe functions which fail to be imported as module.
+    op_path_in_cache = []  # Save paths for op functions created in the cached.
 
     def __init__(self, func, out_shape=None, out_dtype=None, func_type="hybrid", bprop=None, reg_info=None):
         ops.PrimitiveWithInfer.__init__(self, "Custom")
@@ -469,15 +489,7 @@ class Custom(ops.PrimitiveWithInfer):
         self.out_shape = out_shape
         self.out_dtype = out_dtype
         self.bprop = bprop
-        self.fake_output = False
-        self.single_scalar_output = False
-        if not self.out_dtype:
-            self.fake_output = True
-        elif not self.out_shape:
-            self.single_scalar_output = True
-        self.add_prim_attr("fake_output", self.fake_output)
-        self.add_prim_attr("single_scalar_output", self.single_scalar_output)
-
+        self._update_op_attr()
         # Register info
         self._register_info(reg_info)
 
@@ -498,6 +510,7 @@ class Custom(ops.PrimitiveWithInfer):
         self._update_attr()
 
     def __infer__(self, *args):
+        """Infer function of the custom op"""
         if callable(self.out_shape):
             infer_shape = self.out_shape(*(x["shape"] for x in args))
         else:
@@ -534,7 +547,12 @@ class Custom(ops.PrimitiveWithInfer):
                 logger.warning("{}, 'out_dtype' is an empty tuple. Add a placeholder instead. "
                                "Not recommend to use it as it could be any uninitialized data.".format(self.log_prefix))
                 infer_dtype = mstype.int32
-
+        if self.func_type == "aot":
+            if infer_shape is None:
+                logger.warning("{}, 'out_shape' is None. Add a placeholder instead. "
+                               "A CPP version of infer shape function is required "
+                               "in this case.".format(self.log_prefix))
+                infer_shape = (1,)
         # after all automatic infer information fulfillment, throw error if infer_shape/infer_dtype is still None
         if not isinstance(infer_shape, (tuple, list)):
             raise TypeError("{}, 'out_shape' must be one of [tuple, list, function], but got {}"
@@ -552,8 +570,22 @@ class Custom(ops.PrimitiveWithInfer):
         return out
 
     def get_bprop(self):
+        """Get the bprop of the custom op"""
         return self.bprop
 
+    def _update_op_attr(self):
+        """Update the attrs of the custom op"""
+        if self.out_shape is None and self.func_type == "aot":
+            self.add_prim_attr("cpp_infer_shape", True)
+        self.fake_output = False
+        self.single_scalar_output = False
+        if not self.out_dtype:
+            self.fake_output = True
+        elif not self.out_shape:
+            self.single_scalar_output = True
+        self.add_prim_attr("fake_output", self.fake_output)
+        self.add_prim_attr("single_scalar_output", self.single_scalar_output)
+
     def _check_julia_func(self):
         """Check the validity of julia func"""
         if not isinstance(self.func, str):
@@ -612,6 +644,63 @@ class Custom(ops.PrimitiveWithInfer):
                 raise TypeError("{}, 'func' must be of type function, but got {}"
                                 .format(self.log_prefix, type(self.func)))
 
+    def _update_func_imply_path(self):
+        """Update op_imply_path of func"""
+        file_path = os.path.realpath(inspect.getfile(self.func))
+
+        if not self.func_type == "tbe":
+            # Custom ops with type other than tbe doesn't need to import from the path
+            # use the file path directly
+            return file_path
+        # For the custom op of type tbe, the kernel compiler will import the module from file path.
+        # we will try import in the initialization,
+        if file_path in Custom.tbe_path_checked:
+            logger.info("The file of {} has already been checked good to be imported.".format(self.func_name))
+            return file_path
+
+        if file_path not in Custom.tbe_path_failed:
+            # As a single file might include multiply functions
+            # we will not try the file path which already failed in previous trials
+            mod_spec = importlib.util.spec_from_file_location(
+                self.func_name, file_path)
+            custom_mod = importlib.util.module_from_spec(mod_spec)
+            try:
+                mod_spec.loader.exec_module(custom_mod)
+            except (ImportError, RecursionError):
+                Custom.tbe_path_failed.append(file_path)
+            else:
+                Custom.tbe_path_checked.append(file_path)
+                return file_path
+
+        # Create a new file for each tbe function
+        op_imply_path = os.path.realpath(_get_cache_path() + self.func_name + ".py")
+        if op_imply_path in Custom.op_path_in_cache:
+            logger.info("The new file of {} has already been created.".format(self.func_name))
+            return op_imply_path
+
+        logger.warning("Fail to import the original source file. Create a new source file for {}. "
+                       "The new file will not include the dependency for the op function. "
+                       "Check the definition of the function {} "
+                       "in the file: {}".format(self.func_name, self.func_name, op_imply_path))
+
+        Custom.op_path_in_cache.append(op_imply_path)
+
+        if os.path.exists(op_imply_path):
+            try:
+                os.remove(op_imply_path)
+            except FileNotFoundError:
+                logger.warning("Fail to remove the existing file. Check the definition of the function {} "
+                               "in the file: {}".format(self.func_name, op_imply_path))
+
+        with open(op_imply_path, 'at') as file:
+            if platform.system() != "Windows":
+                fcntl.flock(file.fileno(), fcntl.LOCK_EX)
+            file.seek(0, 2)
+            if file.tell() == 0:
+                file.write(self.func_source_str)
+        os.chmod(op_imply_path, stat.S_IRUSR | stat.S_IWUSR)
+        return op_imply_path
+
     def _update_func_info(self, reg_info):
         """Update information of func"""
         if callable(self.func):
@@ -627,19 +716,8 @@ class Custom(ops.PrimitiveWithInfer):
             if index != -1:
                 self.func_source_str = self.func_source_str[index:]
 
-            op_imply_path = os.path.realpath(_get_cache_path() + self.func_name + ".py")
-            if os.path.exists(op_imply_path):
-                os.remove(op_imply_path)
-            with open(op_imply_path, 'at') as file:
-                if platform.system() != "Windows":
-                    fcntl.flock(file.fileno(), fcntl.LOCK_EX)
-                file.seek(0, 2)
-                if file.tell() == 0:
-                    file.write(self.func_source_str)
-            os.chmod(op_imply_path, stat.S_IRUSR | stat.S_IWUSR)
-
-            # path of func
-            self.imply_path = op_imply_path
+            # update path of func for TBE type of custom op
+            self.imply_path = self._update_func_imply_path()
             if self._is_ms_kernel:
                 # static check for the Hybrid DSL in hybrid
                 root = ast.parse(self.func_source_str)
@@ -670,7 +748,7 @@ class Custom(ops.PrimitiveWithInfer):
                     continue
                 if isinstance(reg_info_item, str):
                     reg_info_item = json.loads(reg_info_item)
-                prefix = prefix + "_" + reg_info_item.get("op_name", "")
+                prefix = "_".join([prefix, reg_info_item.get("op_name", "")])
             self.uniq_name = prefix + "_" + self.func_name
         else:
             raise TypeError("For '{}', 'func' must be of type function or str, but got {}"
@@ -798,16 +876,15 @@ class Custom(ops.PrimitiveWithInfer):
                     if isinstance(item, dict) and item.get("value") is None:
                         reg_info["attr"][i]["value"] = "all"
             reg_info["async_flag"] = reg_info.get("async_flag", False)
-            reg_info["binfile_name"] = "%s.so" % self.func_name
+            reg_info["binfile"] = "%s.so" % self.func_name
             reg_info["compute_cost"] = reg_info.get("compute_cost", 10)
-            reg_info["kernel_name"] = self.func_name
+            reg_info["kernel"] = self.func_name
             reg_info["partial_flag"] = reg_info.get("partial_flag", True)
-            reg_info["need_check_supported"] = reg_info.get("need_check_supported", False)
+            reg_info["needCheckSupport"] = reg_info.get("need_check_supported", False)
         # Supplement necessary info for AKG if these information is missing in reg_info
         if reg_info["imply_type"] == "AKG":
             target_to_processor = {"Ascend": "AiCore", "GPU": "CUDA", "CPU": "CPU"}
             reg_info["processor"] = reg_info.get("processor", target_to_processor.get(target))
-
         return reg_info
 
     def _get_target(self, reg_info):
@@ -841,8 +918,8 @@ class Custom(ops.PrimitiveWithInfer):
                 reg_info["imply_type"].strip():
             return reg_info["imply_type"]
         # Infer imply_type from func_type
-        func_type_to_imply_type = {"hybrid": "AKG", "akg": "AKG", "tbe": "TBE", "aicpu": "AiCPU", "aot": target,
-                                   "pyfunc": target, "julia": target}
+        func_type_to_imply_type = {"hybrid": "AKG", "akg": "AKG", "tbe": "TBE", "aicpu": "AiCPU", "pyfunc": target,
+                                   "julia": target, "aot": "BiSheng" if target == "Ascend" else target}
         return func_type_to_imply_type.get(self.func_type, "AKG")
 
     def _save_attr(self, reg_info):
@@ -862,9 +939,11 @@ class Custom(ops.PrimitiveWithInfer):
         for item in tensor_inputs:
             if isinstance(item, dict) and item.get("name") is not None:
                 input_names.append(item["name"])
+        has_input_name = bool(input_names)
         for item in attr:
             if isinstance(item, dict) and item.get("name") is not None:
-                input_names.append(item["name"])
+                if has_input_name or context.get_context("mode") != ms.PYNATIVE_MODE:
+                    input_names.append(item["name"])
                 attr_names.append(item["name"])
         cur_attr = {"input_names": input_names, "attr_names": attr_names}
         # If func does not have attr, save current attr.