{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "500b6c20",
   "metadata": {},
   "source": [
    "# 自定义算子注册\n",
    "\n",
    "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.3.0rc2/resource/_static/logo_notebook.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.3.0rc2/tutorials/experts/zh_cn/operation/mindspore_op_custom_adv.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.3.0rc2/resource/_static/logo_download_code.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.3.0rc2/tutorials/experts/zh_cn/operation/mindspore_op_custom_adv.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.3.0rc2/resource/_static/logo_source.svg)](https://gitee.com/mindspore/docs/blob/r2.3.0rc2/tutorials/experts/source_zh_cn/operation/op_custom_adv.ipynb)\n",
    "\n",
    "## 算子信息注册\n",
    "\n",
    "算子信息主要描述了算子实现函数所支持的输入输出类型、输入输出数据格式、属性和target（平台信息），它是后端做算子选择和映射时的依据。它通过[CustomRegOp](https://www.mindspore.cn/docs/zh-CN/r2.3.0rc2/api_python/ops/mindspore.ops.CustomRegOp.html#mindspore-ops-customregop)接口定义，通过[custom_info_register](https://www.mindspore.cn/docs/zh-CN/r2.3.0rc2/api_python/ops/mindspore.ops.custom_info_register.html#mindspore-ops-custom-info-register)装饰器或者[Custom](https://www.mindspore.cn/docs/zh-CN/r2.3.0rc2/api_python/ops/mindspore.ops.Custom.html#mindspore-ops-custom)原语构造函数中的`reg_info`参数，实现算子信息与算子实现函数的绑定，并最终注册到MindSpore C++侧的算子信息库。`reg_info`参数优先级高于`custom_info_register`装饰器。\n",
    "\n",
    "算子信息中的target的值可以为\"Ascend\"或\"GPU\"或\"CPU\"，描述的是算子实现函数在当前target上所支持的输入输出类型、输入输出数据格式和属性等信息，对于同一个算子实现函数，其在不同target上支持的输入输出类型可能不一致，所以通过target进行区分。算子信息在同一target下只会被注册一次。\n",
    "\n",
    "> - 算子信息中定义输入输出信息的个数和顺序、算子实现函数中的输入输出信息的个数和顺序，两者要完全一致。\n",
    "> - 对于akg类型的自定义算子，若算子存在属性输入，则必须注册算子信息，算子信息中的属性名称与算子实现函数中使用的属性名称要一致；对于tbe类型的自定义算子，当前必须注册算子信息；对于aot类型的自定义算子，由于算子实现函数需要预先编译成动态库，所以无法通过装饰器方式绑定算子信息，只能通过`reg_info`参数传入算子信息。\n",
    "> - 若自定义算子只支持特定的输入输出数据类型或数据格式，则需要注册算子信息，以便在后端做算子选择时进行数据类型和数据格式的检查。对于不提供算子信息的情况，则在后端做算子选择和映射的时候，将会从当前算子的输入中推导信息。\n",
    "\n",
    "## 定义算子反向传播函数\n",
    "\n",
    "如果算子要支持自动微分，需要定义其反向传播函数（bprop），然后将bprop函数传入`Custom`原语构造函数的`bprop`参数。你需要在bprop中描述利用正向输入、正向输出和输出梯度得到输入梯度的反向计算逻辑。反向计算逻辑可以使用内置算子或自定义Custom算子。\n",
    "\n",
    "定义算子反向传播函数时需注意以下几点：\n",
    "\n",
    "- bprop函数的入参顺序约定为正向的输入、正向的输出、输出梯度。若算子为多输出算子，正向输出和输出梯度将以元组的形式提供。\n",
    "- bprop函数的返回值形式约定为输入梯度组成的元组，元组中元素的顺序与正向输入参数顺序一致。即使只有一个输入梯度，返回值也要求是元组的形式。\n",
    "\n",
    "下面test_grad.py为例，展示反向传播函数的用法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b7a95c3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 2.  8. 18.]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import mindspore as ms\n",
    "from mindspore.nn import Cell\n",
    "import mindspore.ops as ops\n",
    "\n",
    "ms.set_context(mode=ms.GRAPH_MODE, device_target=\"GPU\")\n",
    "\n",
    "# 自定义算子正向实现\n",
    "def square(x):\n",
    "    y = output_tensor(x.shape, x.dtype)\n",
    "    for i0 in range(x.shape[0]):\n",
    "        y[i0] = y[i0] * y[i0]\n",
    "    return y\n",
    "\n",
    "# 自定义算子反向实现\n",
    "def square_grad(x, dout):\n",
    "    dx = output_tensor(x.shape, x.dtype)\n",
    "    for i0 in range(x.shape[0]):\n",
    "        dx[i0] = 2.0 * x[i0]\n",
    "    for i0 in range(x.shape[0]):\n",
    "        dx[i0] = dx[i0] * dout[i0]\n",
    "    return dx\n",
    "\n",
    "# 反向传播函数\n",
    "def bprop():\n",
    "    op = ops.Custom(square_grad, lambda x, _: x, lambda x, _: x, func_type=\"akg\")\n",
    "\n",
    "    def custom_bprop(x, out, dout):\n",
    "        dx = op(x, dout)\n",
    "        return (dx,)\n",
    "\n",
    "    return custom_bprop\n",
    "\n",
    "class Net(Cell):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        # 定义akg类型的自定义算子，并提供反向传播函数\n",
    "        self.op = ops.Custom(square, lambda x: x, lambda x: x, bprop=bprop(), func_type=\"akg\")\n",
    "\n",
    "    def construct(self, x):\n",
    "        return self.op(x)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    x = np.array([1.0, 4.0, 9.0]).astype(np.float32)\n",
    "    dx = ms.grad(Net())(ms.Tensor(x))\n",
    "    print(dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b83531",
   "metadata": {},
   "source": [
    "其中：\n",
    "\n",
    "- 反向传播函数中使用是的akg类型的自定义算子，算子定义与使用需要分开，即自定义算子在`custom_bprop`函数外面定义，在`custom_bprop`函数内部使用。\n",
    "\n",
    "执行用例："
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7668ea78",
   "metadata": {},
   "source": [
    "python test_grad.py\n",
    "\n",
    "执行结果：\n",
    "\n",
    "[ 2.  8. 18.]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db293aae",
   "metadata": {},
   "source": [
    "> 更多示例可参考MindSpore源码中[tests/st/ops/graph_kernel/custom](https://gitee.com/mindspore/mindspore/tree/v2.3.0-rc2/tests/st/ops/graph_kernel/custom)下的用例。"
   ]
  }
 ],
 "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}