mindspore.ops.CustomOpBuilder

class mindspore.ops.CustomOpBuilder(name, sources, backend=None, include_paths=None, cflags=None, ldflags=None, **kwargs)

CustomOpBuilder is used to initialize and configure custom operators for MindSpore. Users can define and load custom operator modules through this class and apply them to the network.

In most cases, users only need to provide the source files and additional compilation options in the constructor and call the load method to complete the compilation and loading of the operator. If users have specific customization requirements, they can inherit this class and override certain methods. It is important to note that if methods are overridden, some parameters passed to the constructor may be ignored.

Warning

This is an experimental API that is subject to change.

Parameters

• name (str) – The unique name of the custom operator module, used to identify the operator.

• sources (Union[list[str], tuple[str], str]) – The source file(s) of the custom operator. It can be a single file path or a list of file paths.

• backend (str, optional) – The target backend for the operator, such as "CPU" or "Ascend". Default: None.

• include_paths (Union[list[str], tuple[str], str], optional) – Additionally included paths needed during compilation. Default: None.

• cflags (str, optional) – Extra C++ compiler flags to be used during compilation. Default: None.

• ldflags (str, optional) – Extra linker flags to be used during linking. Default: None.

• kwargs (dict, optional) –

  Additional keyword arguments for future extensions or specific custom requirements.

  • build_dir (str, optional): The directory used to generate the operator build files. If this argument is set, the provided path will be used directly. If not set, a subdirectory named after the operator's name will be created under the path specified by the environment variable MS_COMPILER_CACHE_PATH (defaulting to "./kernel_meta"), and the files will be placed in this subdirectory. Default: None.

  • enable_atb (bool, optional): Whether to call ATB (Ascend Transformer Boost) operator. If set to True, the backend must be Ascend or left empty. Default: False.

  • enable_asdsip (bool, optional): Whether to call ASDSIP (Ascend SiP Boost) operator. If set to True, the backend must be Ascend or left empty. Default: False.

  • op_def (Union[list[str], tuple[str], str], optional): Path(s) to the operator definition file(s) (YAML format). When using custom operators in graph mode, this parameter is mandatory. It can be a single file path string or a list of file path strings. Default: None.

  • op_doc (Union[list[str], tuple[str], str], optional): Path(s) to the operator documentation file(s) (YAML format). This parameter is optional and used to provide additional documentation for the operator. It can be a single file path string or a list of file path strings. Default: None.

Note

• If the backend argument is provided, additional default flags will be automatically added to the compilation and linking steps to support the operator's target backend. The default options can be referenced in the implementation of the get_cflags and get_ldflags methods in the CustomOpBuilder.

• The sources argument must point to valid source files for the custom operator.

Supported Platforms:

Ascend CPU

Examples

>>> from mindspore import ops
>>> builder = ops.CustomOpBuilder(
...     name="custom_op_cpu",
...     sources="custom_ops_impl/pybind_op_cpu.cpp",
...     backend="CPU"
... )
>>> my_ops = builder.load()

build()

Build the custom operator module.

This method generates a dynamic library file for the custom operator based on the provided source files, include paths, compilation flags, and link flags.

Returns

str, The path to the compiled module.

get_cflags()

Get the C++ compiler flags for building the custom operator.

Returns

list[str], A list of C++ compiler flags.

get_include_paths()

Get the include paths required for compiling the custom operator.

Returns

list[str], A list of include paths.

get_ldflags()

Get the linker flags for building the custom operator.

Returns

list[str], A list of linker flags.

get_sources()

Get the source files for the custom operator.

Returns

str or list[str], The source file(s) for the operator.

load()

Build and load the custom operator module.

Returns

Module, The loaded custom operator module.