{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 自定义算子（基于Custom表达）\n",
"\n",
"[](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.3.0rc2/tutorials/experts/zh_cn/operation/mindspore_op_custom.ipynb) [](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.3.0rc2/tutorials/experts/zh_cn/operation/mindspore_op_custom.py) [](https://gitee.com/mindspore/docs/blob/r2.3.0rc2/tutorials/experts/source_zh_cn/operation/op_custom.ipynb)\n",
"\n",
"## 概述\n",
"\n",
"当开发网络遇到内置算子不足以满足需求时，你可以利用MindSpore的Python API中的[Custom](https://www.mindspore.cn/docs/zh-CN/r2.3.0rc2/api_python/ops/mindspore.ops.Custom.html#mindspore-ops-custom)原语方便快捷地进行不同类型自定义算子的定义和使用。\n",
"\n",
"传统的添加一个自定义算子的方式，需要完成算子原语注册、算子实现、算子信息注册三部分工作。\n",
"\n",
"其中：\n",
"\n",
"- 算子原语：定义了算子在网络中的前端接口原型，也是组成网络模型的基础单元，主要包括算子的名称、属性（可选）、输入输出名称、输出shape推理方法、输出数据类型推理方法等信息。\n",
"- 算子实现：在Python侧定义函数（JIT自定义算子）或C++侧定义类（GPU和CPU自定义算子），描述算子内部计算逻辑的实现。\n",
"- 算子信息：描述自定义算子的基本信息，如算子名称、支持的输入输出数据类型、支持的输入输出数据格式和属性等。它是后端做算子选择和映射时的依据。\n",
"\n",
"相比于传统自定义算子方式，基于`Custom`原语自定义算子具有如下优势：\n",
"\n",
"- 不同的自定义算子对应的算子原语都是`Custom`原语，无需对每个自定义算子定义一个相应的算子原语。上述提到的三部分工作可以在网络脚本中以统一的接口进行实现，并作为网络表达的一部分，不需要对MindSpore框架进行侵入式修改和重新编译。\n",
"- 实现了不同方式自定义算子的接口和使用统一，方便网络开发者根据需要灵活选用不同的自定义方式。\n",
"- 新增支持hybrid等自定义算子方式，并且可以跨平台使用。\n",
"\n",
"基于[Custom](https://www.mindspore.cn/docs/zh-CN/r2.3.0rc2/api_python/ops/mindspore.ops.Custom.html#mindspore-ops-custom)原语的自定义算子支持的算子开发方式包括：hybrid、aot、pyfunc、julia、akg。\n",
"\n",
"不同的算子开发方式差异如下：\n",
"\n",
"| 算子开发方式 | 开发语言 | 编译方式 | 支持平台 | 推荐场景 |\n",
"|:-------|:------------------|:------ |:------ |:------------------------|\n",
"| [pyfunc](#自定义算子入门) | Python | N/A | `CPU` | 快速算法验证、需要与Python进行交互等场景 |\n",
"| [hybrid](#hybrid类型的自定义算子开发) | MindSpore HYBRID DSL | JIT | `GPU` `CPU` | 全平台通用开发和快速验证 |\n",
"| [akg](#akg类型的自定义算子开发) | MindSpore AKG DSL | JIT | `GPU` | 用于开发验证场景，不建议普通用户使用 |\n",
"| [aot](#aot类型的自定义算子开发) | C/C++/CUDA | AOT | `GPU` `CPU` | 高性能手写、对接调用第三方算子库场景 |\n",
"| [julia](#julia类型的自定义算子开发) | Julia | N/A | `CPU` | 科学计算场景、需要使用Julia编程等场景 |\n",
"\n",
"> - DSL全称是Domain Specific Language。\n",
"> - AOT（Ahead Of Time）编译方式指的是，算子实现函数需提前被编译为动态链接库，然后在网络运行时由框架自动调用；JIT（Just In Time）编译方式则不需要提前编译算子实现函数，而是在网络编译或运行期间被框架直接编译。\n",
"> - 为了区别自定义算子的类型和编译方式，下面的文中用aot指代自定义算子的类型，用AOT指代自定义算子的编译方式。\n",
"\n",
"不同平台的不同场景下的推荐开发方式如下：\n",
"\n",
"- GPU: hybrid（通用场景），aot（基于CUDA的高性能实现）；\n",
"- CPU: hybrid（通用场景），aot（基于C++的高性能实现）。\n",
"\n",
"不同的开发方式使用不同的开发语言实现算子计算逻辑，但是自定义算子的开发流程是一致的，包括算子实现、算子输出shape和数据类型推理和算子信息注册（可选）。网络开发者可以根据需要选用不同的自定义算子开发方式。下面分别介绍这几种自定义算子开发方式，每种开发方式均提供示例。\n",
"\n",
"> 更多示例可参考MindSpore源码中[tests/st/ops/graph_kernel/custom](https://gitee.com/mindspore/mindspore/tree/v2.3.0-rc2/tests/st/ops/graph_kernel/custom)下的用例。\n",
"\n",
"## 自定义算子入门\n",
"\n",
"为了帮助用户快速入门自定义算子，这里以pyfunc类型自定义算子为例帮助用户理解自定义算子的定义流程。下面基于pyfunc模式定义一个实现sin计算的自定义算子。pyfunc类型的自定义算子使用原生Python语法定义算子实现函数，描述算子内部计算逻辑的实现。网络运行时框架会自动调用此函数。为了表达自定义算子的计算，我们写一个基于numpy的计算正弦函数的Python原生函数。"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"def sin_by_numpy(x):\n",
" return np.sin(x)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"然后我们要定义两个函数，一个是张量形状的推导函数（infer_shape），另一个是张量数据类型的推导函数（infer_dtype）。这里要注意：\n",
"\n",
"- 张量形状的推导函数是输入张量的形状；\n",
"- 张量数据类型的推导函数是输入张量的数据类型。"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"def infer_shape(x):\n",
"\n",
" # 1. 这里的输入x是算子输入张量的形状\n",
" # 2. sin函数是逐元素计算，输入的形状和输出的一样\n",
" return x\n",
"\n",
"def infer_dtype(x):\n",
"\n",
" # 1. 这里的输入x是算子输入张量的数据类型\n",
" # 2. sin函数输入的数据类型和输出的一样\n",
" return x\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"下面我们用上面的函数自定义一个算子，其输入包括\n",
"\n",
"- func：自定义算子的函数表达，这里我们用`sin_by_numpy`函数；\n",
"- out_shape: 输出形状的推导函数，这里我们用`infer_shape`函数；\n",
"- out_dtype: 输出数据类型的推导函数，这里我们用`infer_dtype`函数；\n",
"- func_type: 自定义算子类型，这里我们用`\"pyfunc\"`。"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"from mindspore import ops\n",
"\n",
"sin_by_numpy_op = ops.Custom(func=sin_by_numpy, # 这里填入自定义算子的函数表达\n",
" out_shape=infer_shape, # 这里填入输出形状的推导函数\n",
" out_dtype=infer_dtype, # 这里填入输出数据类型的推导函数\n",
" func_type=\"pyfunc\" # 这里填入自定义算子类型\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"加上其他环境依赖依赖和算子调用语句，我们获得完整的自定义算子用例如下。"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[0. 0.841471 0.19866933 0.29552022 0.38941833]\n"
]
}
],
"source": [
"import numpy as np\n",
"import mindspore as ms\n",
"from mindspore import ops\n",
"\n",
"ms.set_context(mode=ms.GRAPH_MODE, device_target=\"CPU\")\n",
"\n",
"def sin_by_numpy(x):\n",
" return np.sin(x)\n",
"\n",
"def infer_shape(x):\n",
" return x\n",
"\n",
"def infer_dtype(x):\n",
" return x\n",
"\n",
"sin_by_numpy_op = ops.Custom(func=sin_by_numpy,\n",
" out_shape=infer_shape,\n",
" out_dtype=infer_dtype,\n",
" func_type=\"pyfunc\")\n",
"input_tensor = ms.Tensor([0, 1, 0.2, 0.3, 0.4], dtype=ms.float32)\n",
"result_cus = sin_by_numpy_op(input_tensor)\n",
"print(result_cus)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"我们可以得到结果为，即上面输入对应的sin值。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"如此我们完成一个pyfunc类型自定义算子的定义。对于更多完整的pyfunc类型自定义算子的例子，参见MindSpore源码中的[用例](https://gitee.com/mindspore/mindspore/blob/v2.3.0-rc2/tests/st/ops/graph_kernel/custom/test_custom_pyfunc.py)。\n",
"\n",
"## 采用JIT编译的自定义算子\n",
"\n",
"JIT（Just In Time）指算子在网络编译或运行期间被框架直接编译。用户可以直接用Python脚本在网络脚本中直接定义此种类型的自定义算子，然后根据算子和后端类型调用对应算子编译器自动编译。此种类型的自定义算子定义方便，而且有着更好的后端适应性。\n",
"\n",
"### Hybrid类型的自定义算子开发\n",
"\n",
"Hybrid类型的自定义算子是自定义算子的默认定义类型。通过使用Hybrid类型的自定义算子，用户可以用类Python的语法描述算子计算逻辑，且无需关注MindSpore框架对于算子定义的工程细节，让用户专注于算法本身。\n",
"\n",
"Hybrid类型的自定义算子使用[MindSpore Hybrid DSL](https://www.mindspore.cn/tutorials/experts/zh-CN/r2.3.0rc2/operation/ms_kernel.html#语法规则)描述算子内部计算逻辑的实现。用MindSpore Hybrid DSL定义的函数可以被[AKG算子编译器](https://gitee.com/mindspore/akg)解析进行JIT编译生成高效算子，在大规模模型的训练推理中使用。同时，用MindSpore Hybrid DSL定义的函数可以当做一个`numpy`函数直接调用，方便用户调试的同时也可以灵活的切换到[pyfunc 类型的自定义算子](#自定义算子入门一个例子)，做到一次开发，多个模式多个平台多个场景复用的自定义算子表达。\n",
"\n",
"下面用例(test_custom_hybrid.py)介绍hybrid类型的自定义算子开发流程，其中自定义算子实现两个输入张量相加的功能。\n",
"值得注意的是，Hybrid类型的自定义算子采取源码变换的方式打通MindSpore的图编译器和算子编译器，用户可以直接使用MindSpore Hybrid DSL提供的关键词，例如下面的`output_tensor`，而无需引入对应Python函数。更多MindSpore Hybrid DSL关键词的介绍，参见[MindSpore Hybrid DSL关键词](https://www.mindspore.cn/tutorials/experts/zh-CN/r2.3.0rc2/operation/ms_kernel.html#关键词)。"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[[2. 2.]\n",
" [4. 4.]]\n"
]
}
],
"source": [
"import numpy as np\n",
"from mindspore import ops\n",
"import mindspore as ms\n",
"from mindspore.ops import kernel\n",
"\n",
"ms.set_context(device_target=\"CPU\")\n",
"\n",
"# 算子实现，Hybrid DSL\n",
"@kernel\n",
"def add(a, b):\n",
" c = output_tensor(a.shape, a.dtype)\n",
" for i0 in range(a.shape[0]):\n",
" for i1 in range(a.shape[1]):\n",
" c[i0, i1] = a[i0, i1] + b[i0, i1]\n",
" return c\n",
"\n",
"if __name__ == \"__main__\":\n",
" # 定义hybrid类型的自定义算子(Custom的默认模式)\n",
" op = ops.Custom(add)\n",
"\n",
" x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)\n",
" x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)\n",
" output = op(ms.Tensor(x0), ms.Tensor(x1))\n",
" print(output)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"本例中，有如下几点需要说明：\n",
"\n",
"- Hybrid类型是Custom的默认类型。\n",
"- Hybrid类型自定义算子的输入必须是一个带有[@kernel](https://www.mindspore.cn/docs/zh-CN/r2.3.0rc2/api_python/ops/mindspore.ops.kernel.html)的函数。\n",
"- Hybrid类型自定义算子定义时可以使用自带的自动shape/dtype推导函数，也可以手动输入shape/dtype推导函数。\n",
"\n",
"执行用例："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"python test_custom_hybrid.py\n",
"```\n",
"\n",
"执行结果：\n",
"\n",
"```text\n",
"[[2. 2.]\n",
" [4. 4.]]\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"对于更多完整的hybrid类型自定义算子的例子，参见MindSpore源码中的[用例](https://gitee.com/mindspore/mindspore/blob/v2.3.0-rc2/tests/st/ops/graph_kernel/custom/test_ms_kernel.py)。\n",
"\n",
"### akg类型的自定义算子开发\n",
"\n",
"akg类型的自定义算子使用[MindSpore AKG](https://gitee.com/mindspore/akg)算子DSL，描述算子内部计算逻辑的实现。MindSpore AKG是基于TVM（Tensor Virtual Machine）和Polyhedral技术的算子开发和编译框架，支持Hybrid、IR builder和TVM compute等多种类型的算子DSL。\n",
"\n",
"算子输出shape和数据类型推理可以通过定义Python函数实现，描述算子输出shape和数据类型的推导逻辑。\n",
"\n",
"若算子包含属性或者只支持特定的输入输出数据类型或数据格式，则需要注册算子信息，算子信息生成方式请参考[算子信息注册](https://www.mindspore.cn/tutorials/experts/zh-CN/r2.3.0rc2/operation/op_custom_adv.html#算子信息注册)。若未注册算子信息，在后端做算子选择和映射的时候，将会从当前算子的输入中推导算子信息。\n",
"\n",
"下面以test_custom_akg.py为例介绍akg类型的自定义算子开发流程，其中自定义算子实现两个输入张量相加的功能。\n",
"\n",
"test_custom_akg.py内容："
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[[2. 2.]\n",
" [4. 4.]]\n"
]
}
],
"source": [
"import numpy as np\n",
"import mindspore as ms\n",
"import mindspore.ops as ops\n",
"\n",
"ms.set_context(device_target=\"CPU\")\n",
"\n",
"# 算子实现，Hybrid DSL\n",
"def add(a, b):\n",
" c = output_tensor(a.shape, a.dtype)\n",
" for i0 in range(a.shape[0]):\n",
" for i1 in range(a.shape[1]):\n",
" c[i0, i1] = a[i0, i1] + b[i0, i1]\n",
" return c\n",
"\n",
"if __name__ == \"__main__\":\n",
" # 定义akg类型的自定义算子\n",
" op = ops.Custom(add, out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type=\"akg\")\n",
"\n",
" x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)\n",
" x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)\n",
" output = op(ms.Tensor(x0), ms.Tensor(x1))\n",
" print(output)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"本例中，有如下几点需要说明：\n",
"\n",
"- `set_context(device_target=\"GPU\")`表示算子运行在GPU平台。\n",
"- 用Python lambda函数定义输出shape和数据类型推理函数，并分别传给`Custom`原语的`out_shape`和`out_dtype`参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。\n",
"- 未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。\n",
"\n",
"执行用例："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"python test_custom_akg.py\n",
"```\n",
"\n",
"执行结果：\n",
"\n",
"```text\n",
"[[2. 2.]\n",
" [4. 4.]]\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"对于更多完整的akg类型自定义算子的例子，参见MindSpore源码中的[用例](https://gitee.com/mindspore/mindspore/blob/v2.3.0-rc2/tests/st/ops/graph_kernel/custom/test_custom_akg.py)。\n",
"\n",
"## 采用AOT编译的自定义算子\n",
"\n",
"AOT类型的自定义算子指用户事先把算子编译成二进制文件后接入网络。通常用户通过C/C++/CUDA等编程语言手工优化算子实现，并把算子以动态库的形式接入MindSpore加速网络。如此，用户可以针对算子进行极致优化，发挥对应后端硬件的极致性能。这里我们会介绍AOT类型自定义算子的一些基础知识，对于AOT类型自定义算子的更多用法和功能，请参见[AOT类型自定义算子进阶用法](https://www.mindspore.cn/tutorials/experts/zh-CN/r2.3.0rc2/operation/op_custom_aot.html)\n",
"\n",
"### aot类型的自定义算子开发\n",
"\n",
"aot类型的自定义算子采用AOT编译方式，要求网络开发者基于特定接口，手写算子实现函数对应的源码文件，并提前将源码文件编译为动态链接库，然后在网络运行时框架会自动调用执行动态链接库中的函数。在算子实现的开发语言方面，GPU平台支持CUDA，CPU平台支持C和C++。源码文件中的算子实现函数的接口规范如下："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"extern \"C\" int CustomFunc(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"其中，函数名`CustomFunc`可替换成任意有效函数名。返回值为int类型，约定0表示正常退出，非0表示发生异常。参数列表的含义如下：\n",
"\n",
"- nparam (int): 输入输出总数。比如算子有2个输入，1个输出，则nparam的值为3。\n",
"- params (void \\*\\*): 输入输出指针数组。比如算子有2个输入，1个输出，params[0]指向第一个输入数据，params[1]指向第二个输入数据，params[2]指向输出数据。\n",
"- ndims (int \\*): 输入输出shape维度数组。比如params[i]是个shape[1024, 1024]的张量，则ndims[i]的值为2。\n",
"- shapes (int64_t \\*\\*): 输入输出shape数组。比如params[i]是个shape[1024, 1024]的张量，则shapes[i][0]的值为1024，shapes[i][1]的值为1024。\n",
"- dtypes (const char \\*\\*): 输入输出数据类型数组。dtypes里的元素取值可为：\"float32\", \"float16\", \"float\", \"float64\", \"int\", \"int8\", \"int16\", \"int32\", \"int64\", \"uint\", \"uint8\", \"uint16\", \"uint32\", \"uint64\", \"bool\"。\n",
"- stream (void \\*): CUDA流指针，仅定义GPU算子实现时需要。\n",
"- extra (void \\*): 用于后续扩展。\n",
"\n",
"在Python脚本中，`Custom`接口中的`func`输入的格式为`Path_To_Func:CustomFunc`，其中`CustomFunc`为上面函数的名字，而`Path_To_Func`为对应函数源文件或者二进制库的地址。\n",
"\n",
"> - MindSpore识别自动编译的方式为文件名后缀。为了使用自动编译功能，请使用后缀为`cpp`、`cc`或者`cu`的源文件。其他情况MindSpore将处理为二进制库的路径；\n",
"> - 为了防止恶意第三方库篡改，请在环境变量`MS_CUSTOM_AOT_WHITE_LIST`设置合法第三方库的路径。只有在`MS_CUSTOM_AOT_WHITE_LIST`设置的目录及其子目录下文件才会被自定义算子调用。\n",
"\n",
"算子输出shape和数据类型推理可以通过定义Python函数实现，描述算子输出shape和数据类型的推导逻辑。\n",
"\n",
"若自定义算子只支持特定的输入输出数据类型，则需要定义算子信息，算子信息生成方式请参考[算子信息注册](https://www.mindspore.cn/tutorials/experts/zh-CN/r2.3.0rc2/operation/op_custom_adv.html#算子信息注册)。\n",
"\n",
"下面通过例子介绍GPU平台和CPU平台上aot类型的自定义算子开发流程，其中自定义算子实现两个输入张量相加的功能。\n",
"\n",
"#### GPU示例\n",
"\n",
"使用CUDA语言，编写算子实现的源码文件add.cu：\n",
"\n",
"```text\n",
"#define THREADS 1024\n",
"__global__ void CustomAddKernel(float *input1, float *input2, float *output, size_t size) {\n",
" auto idx = blockIdx.x * THREADS + threadIdx.x;\n",
" if (idx < size) {\n",
" output[idx] = input1[idx] + input2[idx];\n",
" }\n",
"}\n",
"\n",
"extern \"C\" int CustomAdd(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream,\n",
" void *extra) {\n",
" cudaStream_t custream = static_cast<cudaStream_t>(stream);\n",
" if (nparam != 3) return 1;\n",
" void *input1 = params[0];\n",
" void *input2 = params[1];\n",
" void *output = params[2];\n",
" size_t size = 1;\n",
"\n",
" for (int i = 0; i < ndims[2]; i++) {\n",
" size *= shapes[2][i];\n",
" }\n",
" int n = size / THREADS;\n",
" for (int i = 0; i < nparam; i++) {\n",
" if (strcmp(dtypes[i], \"float32\") != 0) {\n",
" return 2;\n",
" }\n",
" }\n",
" CustomAddKernel<<<n + 1, THREADS, 0, custream>>>(static_cast<float *>(input1), static_cast<float *>(input2),\n",
" static_cast<float *>(output), size);\n",
" return 0;\n",
"}\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"将add.cu编译成动态库add.so："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"nvcc --shared -Xcompiler -fPIC -o add.so add.cu\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"编写测试用例test_custom_aot.py：\n",
"\n",
"```text\n",
"import numpy as np\n",
"import mindspore as ms\n",
"import mindspore.ops as ops\n",
"\n",
"ms.set_context(device_target=\"GPU\")\n",
"\n",
"if __name__ == \"__main__\":\n",
" # 定义aot类型的自定义算子\n",
" op = ops.Custom(\"./add.so:CustomAdd\", out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type=\"aot\")\n",
"\n",
" x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)\n",
" x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)\n",
" output = op(ms.Tensor(x0), ms.Tensor(x1))\n",
" print(output)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"本例中，有如下几点需要说明：\n",
"\n",
"- 本例中需要将test_custom_aot.py和add.so放置在同一目录下，若add.so在其他目录，则需要将`Custom`第一个参数里路径修改为add.so的绝对路径。\n",
"- 用Python lambda函数定义输出shape和数据类型推理函数，并分别传给`Custom`原语的`out_shape`和`out_dtype`参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。\n",
"- 未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。\n",
"\n",
"执行用例："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"python test_custom_aot.py\n",
"```\n",
"\n",
"执行结果：\n",
"\n",
"```text\n",
"[[2. 2.]\n",
" [4. 4.]]\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### CPU示例\n",
"\n",
"使用C或者C++语言，编写算子实现的源码文件add.cc：\n",
"\n",
"```text\n",
"#include <string.h>\n",
"using size_t = decltype(sizeof(int));\n",
"using int64_t = decltype(sizeof(long));\n",
"\n",
"extern \"C\" int CustomAdd(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra) {\n",
" if (nparam != 3) return 1;\n",
" float *input1 = static_cast<float *>(params[0]);\n",
" float *input2 = static_cast<float *>(params[1]);\n",
" float *output = static_cast<float *>(params[2]);\n",
" size_t size = 1;\n",
" for (int i = 0; i < ndims[2]; i++) {\n",
" size *= shapes[2][i];\n",
" }\n",
" for (int i = 0; i < nparam; i++) {\n",
" if (strcmp(dtypes[i], \"float32\") != 0) {\n",
" return 2;\n",
" }\n",
" }\n",
" for (int i = 0; i < size; i++) {\n",
" output[i] = input1[i] + input2[i];\n",
" }\n",
" return 0;\n",
"}\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"将add.cc编译成动态库add.so："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"g++ --shared -fPIC -o add.so add.cc\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"编写测试用例test_custom_aot.py：\n",
"\n",
"```text\n",
"import numpy as np\n",
"import mindspore as ms\n",
"import mindspore.ops as ops\n",
"\n",
"ms.set_context(device_target=\"CPU\")\n",
"\n",
"if __name__ == \"__main__\":\n",
" # 定义aot类型的自定义算子\n",
" op = ops.Custom(\"./add.so:CustomAdd\", out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type=\"aot\")\n",
"\n",
" x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)\n",
" x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)\n",
" output = op(ms.Tensor(x0), ms.Tensor(x1))\n",
" print(output)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"本例中，有如下几点需要说明：\n",
"\n",
"- 本例中需要将test_custom_aot.py和add.so放置在同一目录下，若add.so在其他目录，则需要将`Custom`第一个参数里路径修改为add.so的绝对路径。\n",
"- 用Python lambda函数定义输出shape和数据类型推理函数，并分别传给`Custom`原语的`out_shape`和`out_dtype`参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。\n",
"- 未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。\n",
"\n",
"执行用例："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"python test_custom_aot.py\n",
"```\n",
"\n",
"执行结果：\n",
"\n",
"```text\n",
"[[2. 2.]\n",
" [4. 4.]]\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"对于更多完整的aot类型自定义算子的例子，参见MindSpore源码中的[用例](https://gitee.com/mindspore/mindspore/blob/v2.3.0-rc2/tests/st/ops/graph_kernel/custom/test_custom_aot.py)。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 自定义算子接入第三方前端\n",
"\n",
"作为MindSpore未来的发展方向之一，AI和科学计算的融合越来越受到业界的重视。MindSpore自定义算子基于自身表达的灵活性，也在科学计算方面做出了探索：把面向HPC的编程前端以自定义算子的方式接入MindSpore。\n",
"\n",
"### julia类型的自定义算子开发\n",
"\n",
"Julia是一种速度快且使用简单的高级通用编程语言，最初设计用于科学计算领域，而由于其高效而实用的特性，近些年来越来越受到用户的青睐，逐步迈向主流编程语言。\n",
"julia类型的自定义算子使用Julia语法定义算子实现函数，描述算子内部计算逻辑的实现。网络运行时框架会自动调用执行相应的Julia函数。\n",
"\n",
"算子输出shape和数据类型推导可以通过定义Python函数实现，描述算子输出shape和数据类型的推导逻辑。\n",
"\n",
"若自定义算子只支持特定的输入输出数据类型，则需要定义算子信息，算子信息生成方式请参考[算子信息注册](https://www.mindspore.cn/tutorials/experts/zh-CN/r2.3.0rc2/operation/op_custom_adv.html#算子信息注册)。\n",
"\n",
"下面以两个输入张量相加为例，介绍julia类型的自定义算子开发流程:\n",
"\n",
"首先，用户需要通过单独文件实现Julia函数，如(add.jl)：\n",
"\n",
"```text\n",
"# add.jl\n",
"module Add\n",
"# inputs: x, y, output: z, output should use .= to inplace assign\n",
"function add(x, y, z)\n",
" z .= x + y\n",
"end\n",
"end\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"其次，在网络脚本中通过自定义算子方式引用上面所写的Julia函数，以test_custom_julia.py为例：\n",
"\n",
"```text\n",
"import numpy as np\n",
"import mindspore as ms\n",
"import mindspore.ops as ops\n",
"\n",
"ms.set_context(device_target=\"CPU\")\n",
"\n",
"if __name__ == \"__main__\":\n",
" # 定义julia类型的自定义算子\n",
" op = ops.Custom(\"./add.jl:Add:add\", out_shape=lambda x, _: x, out_dtype=lambda x, _: x, func_type=\"julia\")\n",
" x0 = np.array([[0.0, 0.0], [1.0, 1.0]]).astype(np.float32)\n",
" x1 = np.array([[2.0, 2.0], [3.0, 3.0]]).astype(np.float32)\n",
" output = op(ms.Tensor(x0), ms.Tensor(x1))\n",
" print(output)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"本例中，有如下几点需要说明：\n",
"\n",
"- 用Python lambda函数定义输出shape和数据类型推理函数，并分别传给`Custom`原语的`out_shape`和`out_dtype`参数。本例中lambda函数表明输出shape和数据类型和第一个输入张量的信息相同。\n",
"- 未注册算子信息，所以自定义算子的算子信息将会从算子输入中推理。\n",
"\n",
"执行用例："
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```text\n",
"python test_custom_julia.py\n",
"```\n",
"\n",
"执行结果：\n",
"\n",
"```text\n",
"[[2. 2.]\n",
" [4. 4.]]\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"注意事项：\n",
"\n",
"1. 用户需确保下载正确版本的Julia，即version>=1.6.0。\n",
"2. 由于运行时调用的Julia C api是从`libjulia.so`中获取的，因此需要用户设置`julia/lib`到`LD_LIBRARY_PATH`，以julia-1.6.5为例:\n",
"\n",
" ```bash\n",
" # download julia-1.6.5\n",
" wget https://julialang-s3.julialang.org/bin/linux/x64/1.6/julia-1.6.5-linux-x86_64.tar.gz\n",
" # for arm server\n",
" # wget https://julialang-s3.julialang.org/bin/linux/aarch64/1.6/julia-1.6.5-linux-aarch64.tar.gz\n",
" # extract file\n",
" tar xvf julia-1.6.5-linux-x86_64.tar.gz\n",
" # if $JULIA_DIR not exist\n",
" export LD_LIBRARY_PATH=$PWD/julia-1.6.5/lib:$LD_LIBRARY_PATH\n",
" # else\n",
" export LD_LIBRARY_PATH=$JULIA_DIR/lib:$LD_LIBRARY_PATH\n",
" ```\n",
"\n",
"3. `Custom` 第一个入参指定用户书写的Julia函数需按照`file_name:module_name:func_name`格式指定，`file_name`需包含文件路径，建议使用绝对路径。\n",
"4. Julia代码文件需包含`module`, `module`内包含`function`，且`module`/`function`都以`end`结束。\n",
"5. Julia函数的输入输出顺序需与算子的输入输出顺序一致。\n",
"6. Julia函数的最终输出，即kernel output的赋值需要使用`.=`，否则结果无法写入内存。\n",
"7. Julia代码支持[Julia](https://docs.julialang.org/en/v1/)的常用语法，用户需自行保证语法正确，函数可正确执行。\n",
"8. 用户想在Julia文件内使用Julia的第三方软件包，需自行下载对应软件以确保能正确调用，可以通过 `import pkg; pkg.add(\"somepkg\")`进行安装。\n",
"9. `julia array`在内存上是`column major`排列的，而`numpy array`是`row major`排列的，如果Julia和numpy做比较，非elemwise计算需考虑内存排布。在Julia函数中，可以通过如下代码示例进行`numpy array`和`julia array`的相互转换:\n",
"\n",
" ```julia\n",
" function change_input_to_row_major(x)\n",
" return permutedims(reshape(x, reverse(size(x))), length(size(x)):-1:1)\n",
" end\n",
"\n",
" function change_output_to_row_major(x)\n",
" return reshape(permutedims(x, length(size(x)):-1:1), size(x))\n",
" end\n",
" ```\n",
"\n",
" 以矩阵乘为例：\n",
"\n",
" ```julia\n",
" # julia array is column-major, numpy array is row-major\n",
" # user should change julia or numpy's layout to keep same behavior\n",
" #= EXAMPLE\n",
" A[2,3] B[3,4] C[2,4]\n",
" NUMPY:\n",
" [[1, 2, 3] [[1, 2, 3, 4] [[38, 44, 50, 56]\n",
" [4, 5, 6]] [5, 6, 7, 8] [83, 98, 113,128]]\n",
" [9,10,11,12]]\n",
" JULIA:\n",
" change_input_to_row_major:\n",
" 1.inputs read numpy data from memory:\n",
" [[1, 3, 5] [[1, 4, 7,10]\n",
" [2, 4, 6]] [2, 5, 8,11]\n",
" [3, 6, 9,12]]\n",
" 2.inputs after reshape(reverse(shape)):\n",
" [[1, 4] [[1, 5, 9]\n",
" [2, 5] [2, 6,10]\n",
" [3, 6]] [3, 7,11]\n",
" [4, 8,12]]\n",
" 3.inputs after transpose/permutedims:\n",
" [[1, 2, 3] [[1, 2, 3, 4] [[38, 44, 50, 56]\n",
" [4, 5, 6]] [5, 6, 7, 8] [83, 98, 113,128]]\n",
" [9,10,11,12]]\n",
" change_output_to_row_major:\n",
" 1.output after transpose/permutedims:\n",
" [[38, 83]\n",
" [44, 98]\n",
" [50,113]\n",
" [56,128]\n",
" 2.output after reshape:\n",
" [[38, 50, 83, 113]\n",
" [44, 56, 98, 128]]\n",
" 3.output read numpy data from memory:\n",
" [[38, 44, 50, 56]\n",
" [83, 98,113, 128]]\n",
" =#\n",
" function foo!(x, y, z)\n",
" x = change_input_to_row_major(x)\n",
" y = change_input_to_row_major(y)\n",
" z .= gemm(x, y, z)\n",
" z .= change_output_to_row_major(z)\n",
" end\n",
" ```\n",
"\n",
"对于更多完整的jullia类型自定义算子的例子，参见MindSpore源码中的[用例](https://gitee.com/mindspore/mindspore/blob/v2.3.0-rc2/tests/st/ops/graph_kernel/custom/test_custom_julia.py)。"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "MindSpore",
"language": "python",
"name": "mindspore"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.5-final"
}
},
"nbformat": 4,
"nbformat_minor": 5
}