如何在 Imperative Runtime 中添加算子#

我们以添加 add 算子为例进行说明。

注册 Op Trait#

在 imperative/src/impl/ops/ 目录下对应的文件中（例如 misc.cpp ）：

定义好 apply_on_var_node 方法；
调用宏 OP_TRAIT_REG.

namespace { namespace add {
auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
  auto&& op = static_cast<const Add&>(def);
  mgb_assert(inputs.size() == 2);
  return opr::Add::make(inputs[0], inputs[1], op.param());
}

OP_TRAIT_REG(Add, Add)
          .apply_on_var_node(apply_on_var_node)
          .fallback();
}}  // add

Note

可以参考 dnn/src/common 中的文件组织，确定对应的文件。

在 src/core/include/megbrain/ir/ops.td 中添加你新写的算子。
```
def Add: MgbHashableOp<"Add", [AddParam]>;
```

添加算子#

Warning

下面的流程中使用 add_example.py 文件作为举例，然而在实际开发时，新增算子（以及测试）最好按照对应的分类放在指定文件（如 elemwise.py ）中。可以参考已有的分类逻辑进行判断，如果依旧不清楚如何分类，可以向核心开发人员进行询问。

Note

文档字符串写法请参考 Python 文档字符串风格指南。

Functional 实现#

向 imperative/python/megengine/functional 目录添加文件 add_example.py:

from typing import Sequence

from ..core._imperative_rt.core2 import apply
from ..core.ops import builtin
from ..tensor import Tensor


def add_example(a: Tensor, b: Tensor, m: int=0) -> Tensor:
    """
    add example
    """
    op = builtin.Add(m)
    return apply(op, a, b)

在 imperative/python/megengine/functional/__init__.py 文件中进行导入：
```
from .add_example import add_example
```

在 imperative/python/test/unit/functional/ 对应文件中添加测试：

def test_add_example():
   a = np.random.randn(2, 2, 2).astype(np.float32)
   b = np.random.randn(2, 2, 2).astype(np.float32)
   res = a + b + 3
   c = F.add_example(tensor(a), tensor(b), m=3)[0]
   np.testing.assert_allclose(c.numpy(), res, atol=1e-5)

Note

如果有需要，还应该提供对应的 Module 实现。

Module 实现#

向 imperative/python/megengine/module 目录添加文件 add_example.py:

from abc import abstractmethod
from typing import Tuple, Union

from ..functional import (
    add_example,
)
from .module import Module

class AddExample(Module):
    def __init__(
        self, m: int,
    ):
        super().__init__()
        self.m = m

    def forward(self, a, b):
        return add_example(a, b, self.m)

在 imperative/python/megengine/module/__init__.py 文件中进行导入：
```
from .add_example import AddExample
```
在 imperative/python/test/unit/module/ 对应文件中添加测试。

编译和测试#

假设我们需要对 imperative/python/test/unit/functional/test_functional.py 中的 test_add_example 进行测试。

确保执行了 make develop 命令：

mkdir -p build
cd build
cmake .. -DMGE_WITH_CUDA=OFF -DMGE_WITH_TEST=ON -DMGE_BUILD_IMPERATIVE_RT=ON
make -j$(nproc)
make develop

Note

设置 DMGE_WITH_CUDA=ON 将开启 CUDA 来进行测试。

设置 PYTHONPATH:

export PYTHONPATH="$(git rev-parse --show-toplevel)/imperative/python"

使用 pytest 进行测试：

pytest unit/functional/test_functional.py -k "test_add_example"

Note

需要按照 imperative/python 下的各个 requires 文件中的要求安装所需要的对应版本软件；
编辑完 .py 文件之后请使用 imperative/python/scripts/format.sh 进行格式化;
想要执行 CUDA 测试，请确保你已经配置好了 CUDA 环境。