mindspore 2.4.10__cp311-none-any.whl → 2.5.0__cp311-none-any.whl

mindspore/ops_generate/pyboost_overload_functions_cpp_generator.py

@@ -0,0 +1,344 @@
+# Copyright 2024 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.
+# ============================================================================
+"""
+This module defines the PyboostFunctionsGenerator class for generating C++ functions for PyBoost operations.
+
+The generator processes operator prototypes and constructs the necessary function definitions, including
+conversions for optional parameters and tensor arguments. It generates the registration code and includes
+the necessary header files for the generated functions.
+"""
+
+import os
+import template
+from template import Template
+import gen_constants as K
+import pyboost_utils
+import op_api_proto
+from gen_utils import save_file
+from op_proto import OpProto
+from op_template_parser import OpTemplateParser
+from base_generator import BaseGenerator
+
+
+class PyboostOverloadFunctionsGenerator(BaseGenerator):
+    """
+    Generates PyBoost overload functions cpp code based on operator prototypes.
+
+    This class processes operator prototypes (`op_protos`) to create the necessary C++ function definitions for
+    PyBoost operations. It constructs function bodies, handles optional value conversions, and generates
+    registration code and header inclusions.
+    """
+
+    def __init__(self):
+        self.PYBOOST_OVERLOAD_FUNCTIONS_TEMPLATE = template.PYBOOST_OVERLOAD_FUNCTIONS_CC_TEMPLATE
+        self.PYBOOST_MINT_CLASS_DEF = template.PYBOOST_MINT_CLASS_DEF
+        self.PYBOOST_OVERLOAD_MINT_CLASS_DEF = template.PYBOOST_OVERLOAD_MINT_CLASS_DEF
+        self.TENSOR_FUNC_UT_BODY = template.TENSOR_FUNC_UT_BODY
+        self.TENSOR_FUNC_UT_OVERLOAD_BODY = template.TENSOR_FUNC_UT_OVERLOAD_BODY
+
+        self.single_case_template = Template(
+            'case ${case_id}:\n'
+            '  ${device_dispatcher}\n'
+            '  break;\n'
+        )
+        self.device_dispatcher_template = Template(
+            'if (backend == kAscendDevice || backend == kDavinciDevice) {\n'
+            '  ${ascend_dispatcher}\n'
+            '} else if (backend == kCPUDevice) {\n'
+            '  ${cpu_dispatcher}\n'
+            '} else if (backend == kGPUDevice) {\n'
+            '  ${gpu_dispatcher}\n'
+            '} else {\n'
+            '  MS_LOG(ERROR) << "Device target is not supported!";\n'
+            '  return py::none();\n'
+            '}'
+        )
+        self.pyboost_return_template = Template(
+            '${arg_handler_processor}\n'
+            'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
+            'return mindspore::tensor::ToPython(${pyboost_base_func_name}_Base(${prim_name}, arg_list));\n'
+        )
+        self.callback_python_template = Template(
+            'MS_LOG(INFO) << "Callback python method: ${py_method}";\n'
+            'py::function fn = python_adapter::GetPyFn(\"mindspore.ops.tensor_method\", \"${py_method}\");\n'
+            'py::object res = fn(*args, **kwargs);\n'
+            'return res;\n'
+        )
+        self.pybind_register_template = Template(
+            '(void)py::class_<${cpp_func_name}Functional, Functional, std::shared_ptr<${cpp_func_name}Functional>>\n'
+            '  (*m, "${cpp_func_name}Functional_")\n'
+            '  .def("__call__", &${cpp_func_name}Functional::Call, "Call ${cpp_func_name} functional.");\n'
+            'm->attr("_${mint_func_name}_instance") = ${mint_func_name}_instance;'
+        )
+        self.callback_python_in_ut_template = Template(
+            'MS_LOG(INFO) << "Callback python method in UT: ${py_method}";\n'
+            'fn = python_adapter::GetPyFn(\"mindspore.ops.tensor_method\", \"${py_method}\");\n'
+            'res = fn(*args, **kwargs);\n'
+            'break;\n'
+        )
+        self.single_case_in_ut_template = Template(
+            'case ${case_id}:\n'
+            '  ${device_dispatcher}\n'
+        )
+
+    def generate(self, work_path, op_protos, mint_func_protos_data, alias_func_mapping):
+        """
+        Generates the C++ PyBoost functions and writes them to the specified files.
+
+        This method processes a list of operator prototypes (`op_protos`), extracting necessary information
+        such as operator names, arguments, and conversion types. It constructs the function definitions, includes,
+        and registration code. The generated content is saved to the specified path as a C++ source file.
+
+        Args:
+            work_path (str): The file path where the generated files will be saved.
+            op_protos (list): A list of operator prototypes containing information about the operators to be processed.
+            mint_func_protos_data (dict): A dict of tensor prototypes containing device-related information.
+            alias_func_mapping (dict): A dict mapping from api name to its alias api name.
+
+        Returns:
+            None
+        """
+
+        mint_classes_def_list = []
+
+        _, single_mint_func_data, overload_mint_func_data = op_api_proto.categorize_func_data(mint_func_protos_data)
+        single_func_call_body_list, single_cpp_class_name_list = (
+            self._get_single_func_call_body_list(single_mint_func_data))
+        overload_func_call_body_list, overload_cpp_class_name_list = (
+            self._get_overload_func_call_body_list(overload_mint_func_data))
+
+        mint_classes_def_list.extend(single_func_call_body_list)
+        mint_classes_def_list.extend(overload_func_call_body_list)
+
+        cpp_class_name_list = single_cpp_class_name_list + overload_cpp_class_name_list
+        mint_classes_reg_list = (
+            self._get_mint_func_reg_list(single_mint_func_data, overload_mint_func_data, cpp_class_name_list))
+
+        pyboost_overload_file_str = (
+            self.PYBOOST_OVERLOAD_FUNCTIONS_TEMPLATE.replace(mint_func_classes_def=mint_classes_def_list,
+                                                             pybind_register_code=mint_classes_reg_list))
+        save_path = os.path.join(work_path, K.PIPELINE_PYBOOST_FUNC_GEN_PATH)
+        file_name = "pyboost_overload_functions.cc"
+        save_file(save_path, file_name, pyboost_overload_file_str)
+
+    def _get_single_func_call_body_list(self, single_op_func_data):
+        """
+        Generates the list of call body strings for single operation functions.
+
+        Args:
+            single_op_func_data (dict): Dictionary of tensor function prototypes with only one definition.
+
+        Returns:
+            func_call_body_list (list): Updated str list for generating C++ function call bodies.
+            cpp_class_name_list (list): The list of non-overloaded c++ functional classes' names.
+        """
+        func_call_body_list, cpp_class_name_list = [], []
+        for _, func_proto in single_op_func_data.items():
+            func_name = func_proto.func_name
+            class_name = func_proto.op_proto.op_class.name
+            device_dispatcher_str = self._get_device_dispatchers_str(func_proto)
+            signature_str = self._generate_single_signature_str(
+                func_proto.op_proto, func_proto.kw_only_args, func_proto.varargs)
+            op_args = func_proto.op_proto.op_args
+            max_size = len(op_args)
+            ut_body = self.TENSOR_FUNC_UT_BODY.replace(py_method=func_proto.py_method)
+            func_call_body_list.append(self.PYBOOST_MINT_CLASS_DEF.replace(
+                class_name=class_name,
+                func_name=func_name,
+                device_dispatcher=device_dispatcher_str,
+                signatures=signature_str,
+                max_args=max_size,
+                ut_body=ut_body))
+            cpp_class_name_list.append(class_name)
+        return func_call_body_list, cpp_class_name_list
+
+    def _get_overload_func_call_body_list(self, overload_op_func_data):
+        """
+        Generates the list of call body strings for overloaded operation functions.
+
+        Args:
+            overload_op_func_data (dict): Dictionary of tensor function prototypes with overloaded definitions.
+
+        Returns:
+            func_call_body_list (list): Updated str list for generating C++ function call bodies.
+            cpp_class_name_list (list): The list of overloaded c++ functional classes' names.
+        """
+        func_call_body_list, cpp_class_name_list = [], []
+        for func_api_name, func_protos in overload_op_func_data.items():
+            func_call_body_list.append(
+                self._get_overload_func_call_str(func_api_name, func_protos, cpp_class_name_list))
+        return func_call_body_list, cpp_class_name_list
+
+    def _get_overload_func_call_str(self, func_api_name, func_protos, cpp_class_name_list):
+        """
+        Generates C++ call body string for overloaded tensor functions.
+
+        Args:
+            func_api_name (str): Name of the function API.
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated call body string for the overloaded functions.
+        """
+        signatures_str = self._generate_func_signatures_str(func_protos)
+        dispatch_cases = self._get_dispatch_cases(func_protos)
+        ut_dispatch_cases = self._get_ut_dispatch_cases(func_protos)
+        ut_overload_body = self.TENSOR_FUNC_UT_OVERLOAD_BODY.replace(ut_dispatch_cases=ut_dispatch_cases)
+
+        max_size = 0
+        for tensor_proto in func_protos:
+            op_proto = tensor_proto.op_proto
+            op_args = op_proto.op_args
+            max_size = max(len(op_args), max_size)
+        cpp_func_name = pyboost_utils.format_func_api_name(func_api_name)
+        cpp_class_name_list.append(cpp_func_name)
+        overload_func_call_str = self.PYBOOST_OVERLOAD_MINT_CLASS_DEF.replace(cpp_func_name=cpp_func_name,
+                                                                              func_name=func_api_name,
+                                                                              signatures=signatures_str,
+                                                                              dispatch_cases=dispatch_cases,
+                                                                              max_args=max_size,
+                                                                              ut_overload_body=ut_overload_body)
+        return overload_func_call_str
+
+    def _generate_func_signatures_str(self, func_protos) -> str:
+        """
+        Generates function signatures as a string from the given prototypes.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated function signatures string.
+        """
+        sig_str = ''
+        first_sig = True
+        for tensor_proto in func_protos:
+            op_proto = tensor_proto.op_proto
+            if not first_sig:
+                sig_str += ',\n'
+            first_sig = False
+            sig_str += self._generate_single_signature_str(op_proto, tensor_proto.kw_only_args, tensor_proto.varargs)
+        return sig_str
+
+    def _generate_single_signature_str(self, op_proto: OpProto, kw_only_args, varargs) -> str:
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.generate_signature_str(kw_only_args, varargs, is_tensor_api=False)
+
+    def _get_dispatch_cases(self, func_protos):
+        """
+        Generates C++ switch-case statements for dispatching tensor function calls.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated switch-case dispatch statements.
+        """
+        dispatch_cases_str = ''
+        for idx, func_proto in enumerate(func_protos):
+            device_dispatcher_str = self._get_device_dispatchers_str(func_proto)
+            dispatch_cases_str += self.single_case_template.replace(case_id=idx,
+                                                                    device_dispatcher=device_dispatcher_str)
+        dispatch_cases_str += 'default:\n'
+        dispatch_cases_str += '  return py::none();'
+        return dispatch_cases_str
+
+    def _get_ut_dispatch_cases(self, func_protos):
+        """
+        Generates C++ switch-case statements for dispatching tensor function calls.
+
+        Args:
+            func_protos (list): List of TensorFuncProto objects representing the function prototypes.
+
+        Returns:
+            str: Generated switch-case dispatch statements.
+        """
+        dispatch_cases_str = ''
+        for idx, func_proto in enumerate(func_protos):
+            device_dispatcher_str = self.callback_python_in_ut_template.replace(py_method=func_proto.py_method)
+            dispatch_cases_str += self.single_case_in_ut_template.replace(case_id=idx,
+                                                                          device_dispatcher=device_dispatcher_str)
+        dispatch_cases_str += 'default:\n'
+        dispatch_cases_str += '  res = py::none();'
+        return dispatch_cases_str
+
+    def _get_device_dispatchers_str(self, func_proto):
+        """
+        Generates device-specific dispatch strings for the given function prototype.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate dispatch strings for.
+
+        Returns:
+            str: Generated device-specific dispatch string.
+        """
+        ascend_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'ascend')
+        cpu_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'cpu')
+        gpu_dispatcher_str = self._get_single_device_dispatcher_str(func_proto, 'gpu')
+        device_dispatcher_str = self.device_dispatcher_template.replace(ascend_dispatcher=ascend_dispatcher_str,
+                                                                        cpu_dispatcher=cpu_dispatcher_str,
+                                                                        gpu_dispatcher=gpu_dispatcher_str)
+        return device_dispatcher_str
+
+    def _get_single_device_dispatcher_str(self, func_proto, device):
+        """
+        Generates the dispatch string for a specific device.
+
+        Args:
+            func_proto (TensorFuncProto): Function prototype to generate the dispatcher for.
+            device (str): Device type ('ascend', 'cpu', 'gpu').
+
+        Returns:
+            str: Generated device dispatcher string.
+        """
+        func_proto_device = getattr(func_proto, device)
+        if func_proto_device == 'pyboost':
+            arg_handler_processor_str = self._get_arg_handler_processor(func_proto.func_name, func_proto.op_proto)
+            op_parser = OpTemplateParser(func_proto.op_proto)
+            op_pyboost_func_name = op_parser.get_pyboost_func_name()
+            prim_name = f"prim::kPrim{func_proto.op_proto.op_class.name}"
+            return self.pyboost_return_template.replace(arg_handler_processor=arg_handler_processor_str,
+                                                        class_name=func_proto.op_proto.op_class.name,
+                                                        prim_name=prim_name,
+                                                        pyboost_base_func_name=op_pyboost_func_name)
+        if func_proto_device == 'py_method':
+            return self.callback_python_template.replace(py_method=func_proto.py_method)
+
+        raise TypeError("Only support pyboost or python_method.")
+
+    def _get_arg_handler_processor(self, func_name, op_proto):
+        op_parser = OpTemplateParser(op_proto)
+        return op_parser.get_arg_handler_processor(func_name, op_proto, is_tensor_api=False)
+
+    def _get_mint_func_reg_list(self, single_mint_func_data, overload_mint_func_data, cpp_class_names):
+        """
+        Generates the list of pybind definition strings for mint functions.
+
+        Args:
+            single_mint_func_data (dict): Dictionary of single mint function data.
+            overload_mint_func_data (dict): Dictionary of overload mint function data.
+            cpp_class_names (list): List of C++ class names.
+
+        Returns:
+            list: list of strs for generating pybind definitions of mint functions' API.
+        """
+        # the order of single_mint_func_data/overload_mint_func_data matters
+        mint_func_names = list(single_mint_func_data.keys()) + list(overload_mint_func_data.keys())
+
+        mint_func_reg_list = []
+        for mint_func_name, cpp_func_name in zip(mint_func_names, cpp_class_names):
+            mint_func_reg_list.append(self.pybind_register_template.replace(mint_func_name=mint_func_name,
+                                                                            cpp_func_name=cpp_func_name))
+        return mint_func_reg_list

mindspore/ops_generate/pyboost_utils.py

@@ -17,6 +17,7 @@
 import os
 import logging
 from gen_utils import safe_load_yaml
+from op_proto import OpProto
 
 
 def is_optional_param(op_arg):
@@ -163,11 +164,11 @@ def tuple_input_to_cpp_type(dtype: str):
         'tuple[float]': 'float',
         'tuple[bool]': 'bool',
         'tuple[str]': 'string',
-        'tuple[tensor]': 'TensorPtr',
+        'tuple[tensor]': 'mindspore::tensor::TensorPtr',
         'list[int]': 'int64_t',
         'list[float]': 'float',
         'list[bool]': 'bool',
-        'list[tensor]': 'TensorPtr',
+        'list[tensor]': 'mindspore::tensor::TensorPtr',
     }
     return types_map.get(dtype)
 
@@ -187,14 +188,14 @@ def get_input_dtype(dtype: str, optional):
     Convert type
     """
     # add more type here
-    value_tuple = 'ValueTuplePtr'
+    value_tuple = 'mindspore::ValueTuplePtr'
     type_convert = {
-        'int': 'Int64ImmPtr',
-        'float': 'FP32ImmPtr',
-        'bool': 'BoolImmPtr',
-        'number': 'ScalarPtr',
-        'str': 'StringImmPtr',
-        'tensor': 'BaseTensorPtr',
+        'int': 'mindspore::Int64ImmPtr',
+        'float': 'mindspore::FP32ImmPtr',
+        'bool': 'mindspore::BoolImmPtr',
+        'number': 'mindspore::ScalarPtr',
+        'str': 'mindspore::StringImmPtr',
+        'tensor': 'mindspore::tensor::BaseTensorPtr',
         'tuple[int]': value_tuple,
         'tuple[float]': value_tuple,
         'tuple[bool]': value_tuple,
@@ -204,14 +205,14 @@ def get_input_dtype(dtype: str, optional):
         'list[bool]': value_tuple,
         'list[tensor]': value_tuple,
     }
-    value_tuple_optional = 'std::optional<ValueTuplePtr>'
+    value_tuple_optional = 'std::optional<mindspore::ValueTuplePtr>'
     optional_type_convert = {
-        'int': 'std::optional<Int64ImmPtr>',
-        'float': 'std::optional<FP32ImmPtr>',
-        'bool': 'std::optional<BoolImmPtr>',
-        'number': 'std::optional<ScalarPtr>',
-        'str': 'std::optional<StringImmPtr>',
-        'tensor': 'std::optional<BaseTensorPtr>',
+        'int': 'std::optional<mindspore::Int64ImmPtr>',
+        'float': 'std::optional<mindspore::FP32ImmPtr>',
+        'bool': 'std::optional<mindspore::BoolImmPtr>',
+        'number': 'std::optional<mindspore::ScalarPtr>',
+        'str': 'std::optional<mindspore::StringImmPtr>',
+        'tensor': 'std::optional<mindspore::tensor::BaseTensorPtr>',
         'tuple[int]': value_tuple_optional,
         'tuple[float]': value_tuple_optional,
         'tuple[bool]': value_tuple_optional,
@@ -239,9 +240,9 @@ def get_return_type(dtype: str):
     """
     # add more type here
     type_convert = {
-        'tuple[tensor]': 'std::vector<tensor::TensorPtr>',
-        'list[tensor]': 'std::vector<tensor::TensorPtr>',
-        'tensor': 'tensor::TensorPtr',
+        'tuple[tensor]': 'std::vector<mindspore::tensor::TensorPtr>',
+        'list[tensor]': 'std::vector<mindspore::tensor::TensorPtr>',
+        'tensor': 'mindspore::tensor::TensorPtr',
     }
     if dtype in type_convert:
         return type_convert[dtype]
@@ -266,11 +267,10 @@ def get_op_name(operator_name, class_def):
     """
     Get op name for python class Primitive or c++ OpDef name.
     """
+    if class_def:
+        return class_def
+
     class_name = ''.join(word.capitalize() for word in operator_name.split('_'))
-    if class_def is not None:
-        item = class_def.get("name")
-        if item is not None:
-            class_name = item
     return class_name
 
 
@@ -278,10 +278,6 @@ def get_pyboost_name(operator_name):
     return 'pyboost_' + operator_name
 
 
-def convert_python_func_name_to_c(func_name: str) -> str:
-    return ''.join(word.capitalize() for word in func_name.split('_'))
-
-
 def get_const_number_convert(arg_name, op_arg):
     cpp_type = number_input_to_cpp_type(op_arg.arg_dtype)
     if op_arg.is_type_id:
@@ -297,8 +293,8 @@ def get_tuple_input_convert(arg_name, arg_type):
     :return:
     """
     cpp_type = tuple_input_to_cpp_type(arg_type)
-    if cpp_type == "TensorPtr":
-        cpp_type = "BaseTensorPtr"
+    if cpp_type == "mindspore::tensor::TensorPtr":
+        cpp_type = "mindspore::tensor::BaseTensorPtr"
     return f"std::vector<{cpp_type}> {arg_name}_vector = ConvertValueTupleToVector<{cpp_type}>({arg_name});\n"
 
 
@@ -311,12 +307,52 @@ def is_pyboost_enable(operator_data):
     return False
 
 
+def format_func_api_name(func_api_name):
+    """
+    Converts a snake_case string to PascalCase format with the first letter capitalized.
+    Additionally, it preserves the trailing underscore. In special cases, such as double
+    underscore names (e.g., __add__), it converts them into PascalCase.
+
+    Args:
+        func_api_name (str): The input snake_case string.
+
+    Returns:
+        str: The converted PascalCase string.
+    """
+    # Check if the string ends with '_'
+    is_one_underscore = func_api_name.endswith('_')
+
+    # Check if it is a double-underscore name (special method names)
+    is_double_underscore = func_api_name.startswith('__') and func_api_name.endswith('__')
+
+    # If it is a double-underscore name, remove the leading and trailing underscores
+    if is_double_underscore:
+        func_api_name = func_api_name[2:-2]
+
+    # If the original name ends with '_' but is not a double-underscore name, remove the trailing '_'
+    if is_one_underscore and not is_double_underscore:
+        func_api_name = func_api_name[:-1]
+
+    # Convert snake_case to PascalCase
+    formatted_func_api_name = ''.join(x.capitalize() for x in func_api_name.split('_'))
+
+    # If the original name ends with '_' but is not a double-underscore name, append the trailing underscore
+    if is_one_underscore and not is_double_underscore:
+        formatted_func_api_name += '_'
+
+    # If the original name is a double-underscore name, add a 'Magic' suffix.
+    if is_double_underscore:
+        formatted_func_api_name += 'Magic'
+
+    return formatted_func_api_name
+
+
 def convert_types(inputs):
     '''convert type to acl type'''
     inputs_dtypes = {}
     flag = False
     for i in inputs:
-        inputs_dtypes[i] = inputs.get(i).get('dtype')
+        inputs_dtypes[i] = i.arg_dtype
         if inputs_dtypes[i] != 'tensor':
             flag = True
         if 'tuple' in inputs_dtypes[i]:
@@ -338,16 +374,39 @@ def convert_types(inputs):
     return inputs_dtypes, flag
 
 
-def get_dtypes(op_yaml):
+def get_dtypes(op_proto: OpProto):
     """get op inputs and outputs dtypes"""
-    inputs = op_yaml.get('args')
-    outputs = op_yaml.get('returns')
+    inputs = op_proto.op_args
+    outputs = op_proto.op_returns
     inputs_dtypes, flag_in = convert_types(inputs)
     outputs_dtypes, flag_out = convert_types(outputs)
     none_tensor_exist = (flag_in or flag_out)
     return inputs_dtypes, outputs_dtypes, none_tensor_exist
 
 
+def merge_strings_by_chunk_size(string_list, chunk_size=50):
+    """
+    Merges a list of strings into smaller chunks, with each chunk having a specified maximum size.
+
+    Args:
+        string_list (list of str): A list of strings to be merged.
+        chunk_size (int, optional): The maximum size of each merged chunk. Defaults to 50.
+
+    Returns:
+        list of str: A list of merged strings, where each string contains up to `chunk_size` characters.
+
+    Example:
+        >>> strings = ["Hello", "world", "this", "is", "a", "test"]
+        >>> merge_strings_by_chunk_size(strings, chunk_size=2)
+        ['Helloworld', 'thisis', 'atest']
+    """
+    merged_strings = [
+        "".join(string_list[i:i + chunk_size])  # Merge the current grouped string
+        for i in range(0, len(string_list), chunk_size)
+    ]
+    return merged_strings
+
+
 class AclnnUtils:
     """
     aclnn utils