mindspore 1.10.0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

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
@@ -97,25 +119,21 @@ def _compile_aot(file):
         Custom.compiled_bin.append(func_path)
 
         if file.endswith("cpp") or file.endswith("cc"):
-            cmd = ["g++", "-std=c++17", "--shared", "-fPIC"]
+            cmd = ["g++", "-std=c++17", "--shared", "-fPIC", "-D_GLIBCXX_USE_CXX11_ABI=0"]
             cmd += [include_file, "-o", func_path, file]
         elif file.endswith("cu"):
             cmd = ["nvcc"]
             cmd += ["--shared", "-Xcompiler", "-fPIC", "-O3", "-gencode", "arch=compute_70, code=sm_70"]
-            cmd += ["--use_fast_math", "--expt-relaxed-constexpr", "--expt-extended-lambda"]
-
-            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:
-                cmd += ["-gencode", "arch=compute_80,code=sm_80"]
+            cmd += ["--use_fast_math", "--expt-relaxed-constexpr"]
+            cmd += ["-D_GLIBCXX_USE_CXX11_ABI=0"]
+
+            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))
@@ -141,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/r1.10/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:
@@ -166,7 +184,7 @@ class Custom(ops.PrimitiveWithInfer):
               1. A AKG operator implementation function, which can use ir builder/tvm compute/hybrid grammar.
               2. A TBE operator implementation function.
               3. A pure python function
-              4. An ms_kernel decorated function written by the Hybrid DSL.
+              4. An kernel decorated function written by the Hybrid DSL.
 
             - str: If func is of str type, then str should be a path of file along with a function name.
               This could be used when func_type is "aot" or "julia".
@@ -317,10 +335,10 @@ class Custom(ops.PrimitiveWithInfer):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> import mindspore.ops as ops
         >>> import numpy as np
-        >>> from mindspore.ops import CustomRegOp, custom_info_register, DataType, ms_kernel
-        >>> from mindspore.common import dtype as mstype
+        >>> from mindspore import Tensor, ops
+        >>> from mindspore.ops import CustomRegOp, custom_info_register, DataType, kernel
+        >>> from mindspore import dtype as mstype
         >>> from mindspore.nn import Cell
         >>> input_x = Tensor(np.ones([16, 16]).astype(np.float32))
         >>> input_y = Tensor(np.ones([16, 16]).astype(np.float32))
@@ -329,8 +347,8 @@ class Custom(ops.PrimitiveWithInfer):
         >>> # This is the default func_type in Custom,
         >>> # and both out_shape and out_dtype can be None(default value).
         >>> # In this case, the input func must be a function written in the Hybrid DSL
-        >>> # and decorated by @ms_kernel.
-        >>> @ms_kernel
+        >>> # and decorated by @kernel.
+        >>> @kernel
         ... def add_script(a, b):
         ...     c = output_tensor(a.shape, a.dtype)
         ...     for i0 in range(a.shape[0]):
@@ -436,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")
@@ -453,7 +474,7 @@ class Custom(ops.PrimitiveWithInfer):
         self._is_ms_kernel = False
 
         self._check_func()
-        self._update_func_info()
+        self._update_func_info(reg_info)
         self.add_prim_attr("func_name", self.func_name)
         self.add_prim_attr("uniq_name", self.uniq_name)
         if self.func_type == "hybrid":
@@ -468,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)
 
@@ -497,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:
@@ -533,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 {}"
@@ -551,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):
@@ -592,17 +625,17 @@ class Custom(ops.PrimitiveWithInfer):
             self._check_julia_func()
         elif self.func_type == "hybrid":
             if not hasattr(self.func, "ms_kernel_flag"):
-                raise TypeError("{}, 'func' must be a function decorated by ms_kernel".format(self.log_prefix))
+                raise TypeError("{}, 'func' must be a function decorated by kernel".format(self.log_prefix))
             self._is_ms_kernel = True
             self._func_compile_attrs = getattr(self.func, "compile_attrs", {})
         elif self.func_type == "akg":
             if hasattr(self.func, "ms_kernel_flag"):
                 logger.warning("{}. To have a better user experience, the mode hybrid is suggested "
-                               "for the input function with decorator @ms_kernel. "
+                               "for the input function with decorator @kernel. "
                                "To enable this mode, set the 'func_type' to be \"hybrid\"".format(self.log_prefix))
         elif self.func_type == "pyfunc":
             if hasattr(self.func, "ms_kernel_flag"):
-                logger.warning("{}. Now you are using the function with decorator @ms_kernel in the mode pyfunc. "
+                logger.warning("{}. Now you are using the function with decorator @kernel in the mode pyfunc. "
                                "The kernel will be executed as a native python function, which might lead to "
                                "low efficiency. To accelerate the kernel, set the 'func_type' to be \"hybrid\""
                                .format(self.log_prefix))
@@ -611,7 +644,64 @@ class Custom(ops.PrimitiveWithInfer):
                 raise TypeError("{}, 'func' must be of type function, but got {}"
                                 .format(self.log_prefix, type(self.func)))
 
-    def _update_func_info(self):
+    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):
             # For the func_type other then hybrid, get the original function if func is decorated
@@ -626,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)
@@ -660,11 +739,38 @@ class Custom(ops.PrimitiveWithInfer):
             # func name
             self.func_name = self.func
             # uniq func name
-            self.uniq_name = self.name + "_" + self.func_name
+            prefix = self.name
+            if reg_info is None:
+                reg_info = {}
+            reg_info_list = self._get_expanded_list(reg_info)
+            for reg_info_item in reg_info_list:
+                if not isinstance(reg_info_item, (str, dict)):
+                    continue
+                if isinstance(reg_info_item, str):
+                    reg_info_item = json.loads(reg_info_item)
+                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 {}"
                             .format(self.name, type(self.func)))
 
+    def _update_reg_attrs(self, reg_info):
+        """Update op attrs in reg_info."""
+        for _, item in enumerate(reg_info.get("outputs", [])):
+            output_name_list = []
+            if isinstance(item, dict) and item.get("name"):
+                output_name_list.append(item.get("name"))
+            self.add_prim_attr("output_names", output_name_list)
+
+        if isinstance(reg_info.get("op_name"), str):
+            self.add_prim_attr("reg_op_name", reg_info.get("op_name"))
+
+        if self.func_type == "aot":
+            if reg_info.get("attr") is not None and isinstance(reg_info["attr"], list):
+                for item in reg_info["attr"]:
+                    if isinstance(item, dict) and item.get("value") is not None:
+                        self.add_prim_attr(item["name"], item["value"])
+
     def _register_info(self, info):
         """Register reg_info."""
         reg_info = info
@@ -687,18 +793,11 @@ class Custom(ops.PrimitiveWithInfer):
                     new_dtype_format.append(i + (DataType.I32_Default,))
                 reg_info["dtype_format"] = new_dtype_format
 
-            for _, item in enumerate(reg_info.get("outputs", [])):
-                output_name_list = []
-                if isinstance(item, dict) and item.get("name"):
-                    output_name_list.append(item.get("name"))
-                self.add_prim_attr("output_names", output_name_list)
-
-            if isinstance(reg_info.get("op_name"), str):
-                self.add_prim_attr("reg_op_name", reg_info.get("op_name"))
+            self._update_reg_attrs(reg_info)
 
             target = self._get_target(reg_info)
             # Reg info for func is only registered once for a certain target
-            if self._has_registered(target, reg_info):
+            if self._has_registered(target):
                 continue
             # Register
             reg_info = self._reformat_reg_info(reg_info, target)
@@ -710,7 +809,7 @@ class Custom(ops.PrimitiveWithInfer):
                                  "'custom_info_register' to bind it to 'func' if 'func' is a function."
                                  .format(self.log_prefix))
             self._save_attr(reg_info)
-            self._save_register_status(target, reg_info)
+            self._save_register_status(target)
 
     def _get_expanded_list(self, data):
         """Recursive function to parse elements in list or tuple."""
@@ -724,43 +823,34 @@ class Custom(ops.PrimitiveWithInfer):
             data_list.append(data)
         return data_list
 
-    def _get_registered_targets(self, reg_info=None):
+    def _get_registered_targets(self):
         """Get the registered targets of func."""
         targets = []
-        if reg_info is None:
-            reg_info = {}
         if callable(self.func):
             targets = getattr(self.func, "registered_targets", [])
         elif isinstance(self.func, str):
-            if isinstance(reg_info.get("op_name"), str):
-                reg_op_name = reg_info.get("op_name")
-            else:
-                reg_op_name = self.func
-            targets = Custom.registered_func.get(reg_op_name, [])
+            targets = Custom.registered_func.get(self.uniq_name, [])
         if not isinstance(targets, list):
             targets = [targets]
         return targets
 
-    def _has_registered(self, target, reg_info):
+    def _has_registered(self, target):
         """Check if registration information is registered in target."""
-        registered_targets = self._get_registered_targets(reg_info)
+        registered_targets = self._get_registered_targets()
         return target in registered_targets
 
-    def _save_register_status(self, target, reg_info):
+    def _save_register_status(self, target):
         """Save registration status for target."""
         if callable(self.func):
             registered_targets = getattr(self.func, "registered_targets", [])
             registered_targets.append(target)
             setattr(self.func, "registered_targets", registered_targets)
         elif isinstance(self.func, str):
-            if isinstance(reg_info.get("op_name"), str):
-                reg_op_name = reg_info.get("op_name")
+            func_name = self.uniq_name
+            if isinstance(Custom.registered_func.get(func_name), list):
+                Custom.registered_func.get(func_name).append(target)
             else:
-                reg_op_name = self.func
-            if isinstance(Custom.registered_func.get(reg_op_name), list):
-                Custom.registered_func.get(reg_op_name).append(target)
-            else:
-                Custom.registered_func[reg_op_name] = [target]
+                Custom.registered_func[func_name] = [target]
 
     def _get_op_name(self, reg_info):
         if self.func_type == "aicpu":
@@ -786,20 +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))
-        if self.func_type == "aot":
-            if reg_info.get("attr") is not None and isinstance(reg_info["attr"], list):
-                for item in reg_info["attr"]:
-                    if isinstance(item, dict) and item.get("value") is not None:
-                        self.add_prim_attr(item["name"], item["value"])
         return reg_info
 
     def _get_target(self, reg_info):
@@ -833,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):
@@ -854,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.
@@ -882,7 +969,7 @@ class Custom(ops.PrimitiveWithInfer):
         """Add primitive_target to primitive's attr."""
         registered_targets = self._get_registered_targets()
         if self.func_type == "pyfunc":
-            self.add_prim_attr("primitive_target", "CPU")
+            self.set_device("CPU")
             if registered_targets and registered_targets != ["CPU"]:
                 logger.warning("{}, only supports CPU platform, but got registered target {}. "
                                "We will run it on CPU".format(self.log_prefix, registered_targets))
@@ -890,11 +977,11 @@ class Custom(ops.PrimitiveWithInfer):
             if len(registered_targets) != 1:
                 logger.info("{}, target will be set according to context.".format(self.log_prefix))
             elif registered_targets == ["GPU"]:
-                self.add_prim_attr("primitive_target", "GPU")
+                self.set_device("GPU")
             elif registered_targets == ["CPU"]:
-                self.add_prim_attr("primitive_target", "CPU")
+                self.set_device("CPU")
         elif self.func_type == "julia":
-            self.add_prim_attr("primitive_target", "CPU")
+            self.set_device("CPU")
             device_target = context.get_context('device_target')
             if device_target == "CPU":
                 pass
@@ -962,12 +1049,12 @@ class Custom(ops.PrimitiveWithInfer):
                                      for idx, val in enumerate(symbols)]
 
             if any(i[1] != -1 for i in inplace_assign_output):
-                value = " ".join((str(j) for i in inplace_assign_output for j in i))
-                self.add_prim_attr("inplace_assign_output", value)
+                self.add_prim_attr("inplace_assign_output", " ".join(
+                    (str(j) for i in inplace_assign_output for j in i)))
 
     def _auto_infer(self, *args):
         """
-        the automatic infer function for functions with @ms_kernel decorator
+        the automatic infer function for functions with @kernel decorator
         """
         fake_input = []
         enable_infer_value = True