添加 relay 算子

添加 relay 算子#

参考：relay_add_op

属性是固定的参数，应该在编译时就知道。比如卷积算子的 stride 和 expand 是属于卷积算子属性节点的字段的适当的例子。属性一般定义在 include/tvm/relay/attrs/ 文件夹中，也可以直接通过 Python 接口定义：

register_op_attr 函数用于注册算子的属性。

参数：

  • op_name（str 类型）：算子的名称。

  • attr_key（str 类型）：属性的名称。

  • value（object 类型，可选）：要设置的值。

  • level（int 类型，可选）：优先级级别，默认为 10。

返回值：

  • 如果 value 没有指定，则返回注册函数 fregister。

  • 如果 value 指定了，则直接返回注册函数 _register(value)。

内部函数 _register 用于实际执行算子注册，它调用了 _ffi_api.RegisterOpAttr 函数来注册算子的属性。

import set_env

import tvm
from tvm import relay
from tvm.ir import register_op_attr
from tvm.relay.op import op as _op, op_attrs

op_name = "add_one"
_op.register(op_name, "doc(Add one to the given tensor.)doc")
# op_attr = register_op_attr(op_name, "testing", op_attrs.ScanopAttrs)
_op.get(op_name).set_num_inputs(1)
_op.get(op_name).add_argument("data", "Tensor", "The input tensor.")
# 添加关系函数
_op.get(op_name).add_type_rel("Identity") # "Identity", "Broadcast"
_op.register_pattern(op_name, _op.OpPattern.ELEMWISE)
_op.register_stateful(op_name, False)

_op.get(op_name).get_attr("TOpPattern")

0

x = relay.var("x", dtype="float32", shape=(10,))

custom_op = _op.get("add_one")
z = relay.Call(custom_op, [x])
print(z)

free_var %x: Tensor[(10), float32];
add_one(%x)