添加新的算子，以”nn.vacc_dropout”为例
算子名称说明：
- 跟神经网络有关的，前面加上”nn.”，例如：
_nn.conv2d, n.dense,nn.relu,nn.pad_等等 其它不带”nn.”前缀的是一般的算子，例如：
_round，clip，cast, sum，power，reshape_等等1 添加relay算子
1.1 在c++中注册relay算子
一般地，注册relay算子相关的c++文件在src/relay/op/目录下的某个子目录（没有强制要求）。可以新建一个文件，也可以直接在其它文件中添加。
vacc_dropout算子对应的c++文件是src/relay/pass/vacc/hw_agnostic/convert_strides.cc，源码如下：RELAY_REGISTER_OP("nn.vacc_dropout") //算子名为"nn.vacc_dropout".describe(R"code(Applies the dropout with H/W stride = 2 to the input array.Examples::x = [[ 1., 4., 7., 10.],[ 2., 5., 8., 11.],[ 3., 6., 9., 12.]]vacc_dropout(x) = [[1., 7.],[3., 9.]]x = [[[ 1., 2., 3.],[ 4., 5., 6.],[ 7., 8., 9.]],[[ 1., 2., 3.],[ 4., 5., 6.],[ 7., 8., 9.]],[[ 1., 2., 3.],[ 4., 5., 6.],[ 7., 8., 9.]]]vacc_dropout(x) = [[[ 1., 3.],[ 7., 9.]],[[ 1., 3.],[ 7., 9.]],[[ 1., 3.],[ 7., 9.]]])code" TVM_ADD_FILELINE).set_num_inputs(1).add_argument("data", "Tensor", "The input tensor.").set_support_level(4).add_type_rel("VaccDropout", VaccDropoutRel).set_attr<TOpPattern>("TOpPattern", kInjective);
说明：注册算子使用的宏是：RELAY_REGISTER_OP
C++中的RELAY_REGISTER_OP宏允许开发人员指定以下有关Relay中算子的信息：
- set_num_inputs：设置输入参数个数
- add_argument：添加输入参数的名称,类型和描述
- set_attrs_type：设置属性类型
- set_support_level：设置支持级别（1表示内部固有；较大的数字表示不够完整或外部支持的算子）
- add_type_rel：算子的类型关系（输入与输出之间的类型关系）
- set_attr：设置算子的其它属性
上面几个函数之所以能够一个接一个的采用.的方式连续调用，是因为他们都是OpRegistry的成员函数，而且这些函数的返回值都是OpRegistry&类型的，所以能够使用.的形式继续调用它的成员函数。这些成员函数如下：
set_num_inputs
add_argument
添加输入参数的说明：名称，类型和描述。有几个输入参数，就添加几个说明。
大多数算子只有一个或2个输入，少数的有多个。例如：cast， ``reshape``，``transpose，等只有一个输入；
.set_num_inputs(1).add_argument("data", "Tensor", "The input tensor.")
add，``subtract，``divide，``multiply，``power，``mod，``right_shift，``left_shift，``maximum，``minimum，等算子有2个输入：
.set_num_inputs(2) \.add_argument("lhs", "Tensor", "The left hand side tensor.") \.add_argument("rhs", "Tensor", "The right hand side tensor.") \
set_support_level
设置支持级别（1表示内部固有；较大的数字表示不够完整或外部支持的算子）
自定义的算子，可以设置为3。
set_attrs_type
设置属性类型。如果算子的参数，除了输入参数之外，还有其它的属性参数，就应该设置属性的类型。
tvm已经定义了很多常用的属性。
c++端参考include/tvm/relay/attrs/
python端参考：python/tvm/relay/op/op_attrs.py
例如，case算子的参数，除了输入的tensor之外，还需要指定需要转换的数据类型，所以，需要一个转换类型属性（CastAttrs），该属性类至少包含一个数据类型参数。CastAttrs属性定义如下：
c++端：include/tvm/relay/attrs/transform.h
/*! \brief data type cast */struct CastAttrs : public tvm::AttrsNode<CastAttrs> {DataType dtype;TVM_DECLARE_ATTRS(CastAttrs, "relay.attrs.CastAttrs") {TVM_ATTR_FIELD(dtype).describe("Target data type");}}; // struct CastAttrs.
python端：python/tvm/relay/op/op_attrs.py
@register_relay_attr_nodeclass CastAttrs(Attrs):"""Attributes for transform.cast"""
然后，算子cast的定义就设置了该类型，
RELAY_REGISTER_OP("cast")....set_attrs_type<CastAttrs>()...
如果已定义的属性类型中，没有符合要求的，就需要自定义一个属性类型，参考：自定义属性类
add_type_rel
给算子添加输入和输出之间的类型约束关系。
需要注意的是，add_type_rel()中的第二个参数是固定格式的类型相关的函数，tvm提供了一些常用的，例如：
AffineBroadcastRel
BroadcastCompRel
BroadcastRel
ConcatenateRel
IdentityCastRel
IdentityConcatenateRel
IdentityRel
除了以上几种，还可以自定义类型关系函数，该函数的一些格式如下：
- 1 函数名以Rel结尾
- 2 函数返回值为bool型
- 3 函数的参数一共4个，依次是：
- const Array
& types, //包括输入和输出 - int num_inputs, //输入参数的个数
- const Attrs& attrs, //属性
- const TypeReporter& reporter //包含类型约束的容器
- const Array
例如，自定义的VaccDropoutRel如下：
bool VaccDropoutRel(const Array<Type>& types,int num_inputs,const Attrs& attrs,const TypeReporter& reporter) {CHECK_EQ(types.size(), 2);//1个输入+1个输出，都是TensorTypeconst auto* data = types[0].as<TensorTypeNode>();if (data == nullptr) return false;auto dshape = data->shape;auto num_axis = dshape.size();//张量的轴数(维度)，例如4维的[N,C,H,W]std::vector<int64_t> stride_vec;//步长for (size_t i = 0; i < num_axis - 2; ++i) {stride_vec.push_back(1);}stride_vec.push_back(2);//H方向步长为2stride_vec.push_back(2);//W方向步长为2//计算输出张量的形状std::vector<IndexExpr> oshape(dshape.size());for (size_t i = 0; i < num_axis; ++i) {int64_t stride_v = stride_vec[i];const int64_t* p_dim_size = as_const_int(dshape[i]);CHECK(p_dim_size)<< "vacc_dropout requires dimension to be concrete int";int64_t dim_size = p_dim_size[0];oshape[i] = make_const(dshape[i].dtype(), (dim_size + stride_v - 1) / stride_v);}reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));return true;}
set_attr
set_attr用于设置算子的其它属性。例如：
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
.set_attr
TOpPattern是比较常用的一个属性，用于在图融合时算子的级别。
.set_attr
用于注册计算(c++端)，此处设置了计算，那么在python端就不用注册计算。
1.2 定义一个C++函数为算子生成一个CallNode，并为该函数注册一个Python API钩子
在c++中调用TVM_REGISTER_GLOBAL或TVM_REGISTER_API 注册一个API函数（该函数在第3步中被Python封装成一个简洁的API），该返回一个CallNode
例如
Expr MakeVaccDropout(Expr data) {// 获取vacc_dropout算子，并生成一个该算子的CallNodestatic const Op& op = Op::Get("nn.vacc_dropout");return CallNode::make(op, {data}, Attrs{}, {});}TVM_REGISTER_API("relay.op._make.vacc_dropout").set_body_typed(MakeVaccDropout);
说明：不一定非得写一个函数MakeVaccDropout， 也可以直接在TVM_REGISTER_API中写一个匿名函数。 例如：
TVM_REGISTER_API("relay.op._make.vacc_dropout").set_body_typed<Expr(Expr)>([](Expr data) {static const Op& op = Op::Get("nn.vacc_dropout");return CallNode::make(op, {data}, Attrs{}, {});});
注意：
RELAYREGISTER_OP注册新的算子，或设置相应的算子的属性
TVM_REGISTER_API注册一个全局的API函数，等效于 TVM_REGISTER_GLOBAL
#define TVM_REGISTER_API(_OpName) TVM_REGISTER_GLOBAL(OpName)
TVM_REGISTER_API和的RELAY_REGISTER_OP进一步封装的宏还有：
RELAY_REGISTER_UNARY_OP
RELAY_REGISTER_BINARY_OP
RELAY_REGISTER_CMP_OP
/*! Quick helper macro* - Expose a positional make function to construct the node.* - Register op to the registry.** We make the decision to always only expose positional argument.* We will do rewrapping in the frontend to support language* sugars such as keyword arguments and default value.* \param OpName the name of registry.*/#define RELAY_REGISTER_UNARY_OP(OpName) \TVM_REGISTER_API("relay.op._make." OpName) \.set_body_typed<Expr(Expr)>([](Expr data) { \static const Op& op = Op::Get(OpName); \return CallNode::make(op, {data}, Attrs(), {}); \}); \RELAY_REGISTER_OP(OpName) \.set_num_inputs(1) \.add_argument("data", "Tensor", "The input tensor.") \.add_type_rel("Identity", IdentityRel) \.set_attr<TOpPattern>("TOpPattern", kElemWise) \.set_attr<TOpIsStateful>("TOpIsStateful", false) \.set_attr<FInferCorrectLayout>("FInferCorrectLayout", \ElemwiseArbitraryLayout) \/*! Quick helper macro* - Expose a positional make function to construct the node.* - Register op to the registry.** We make the decision to always only expose positional argument.* We will do rewrapping in the frontend to support language* sugars such as keyword arguments and default value.** \param OpName the name of registry.*/#define RELAY_REGISTER_BINARY_OP(OpName) \TVM_REGISTER_API("relay.op._make." OpName) \.set_body_typed<Expr(Expr, Expr)>([](Expr lhs, Expr rhs) { \static const Op& op = Op::Get(OpName); \return CallNode::make(op, {lhs, rhs}, Attrs(), {}); \}); \RELAY_REGISTER_OP(OpName) \.set_num_inputs(2) \.add_argument("lhs", "Tensor", "The left hand side tensor.") \.add_argument("rhs", "Tensor", "The right hand side tensor.") \.add_type_rel("Broadcast", BroadcastRel) \.set_attr<TOpPattern>("TOpPattern", kBroadcast) \.set_attr<TOpIsStateful>("TOpIsStateful", false) \.set_attr<FInferCorrectLayout>("FInferCorrectLayout", \BinaryBroadcastLayout)// Comparisons#define RELAY_REGISTER_CMP_OP(OpName) \TVM_REGISTER_API("relay.op._make." OpName) \.set_body_typed<Expr(Expr, Expr)>([](Expr lhs, Expr rhs) { \static const Op& op = Op::Get(OpName); \return CallNode::make(op, {lhs, rhs}, Attrs(), {}); \}); \RELAY_REGISTER_OP(OpName) \.set_num_inputs(2) \.add_argument("lhs", "Tensor", "The left hand side tensor.") \.add_argument("rhs", "Tensor", "The right hand side tensor.") \.add_type_rel("BroadcastComp", BroadcastCompRel) \.set_attr<TOpPattern>("TOpPattern", kBroadcast) \.set_attr<TOpIsStateful>("TOpIsStateful", false) \.set_attr<FInferCorrectLayout>("FInferCorrectLayout", \BinaryBroadcastLayout)
1.3 定义relay算子relay.vacc_dropout
即在Python中将第2步中公开的函数包装在更简洁的接口中
一般地，第2步中通过TVM_REGISTER_API 或 TVM_REGISTER_GLOBAL公开的函数，在python中不会被直接调用，而是被封装成一个单独的API。该文件一般在**python/tvm/relay/op/**目录下。
例如，在python/tvm/relay/op/tensor.py文件中如下定义：
def vacc_dropout(data):"""vacc_dropout dropout with H/W stride = 2 to the input arrayParameters----------data : tvm.relay.ExprThe input tensor.Returns-------result : tvm.relay.ExprThe shape tensor."""return _make.vacc_dropout(data)
该算子即为relay.vacc_dropout
说明：_make.``vacc_dropout就是在第2步中注册的全局函数relay.op._make.vacc_dropout
定义了vacc_dropout函数后，在python中会调用relay.``vacc_dropout就相当于直接调用全局函数relay.op._make.vacc_dropout。
1.4 在python端为算子注册compute, schedule
例如，在python/tvm/relay/op/_tensor.py中代码如下：
(1)注册算子的计算
# 注册算子的计算@register_compute("nn.vacc_dropout")def vacc_dropout_compute(attrs, inputs, output_type, target):assert len(inputs) == 1print("relay/op: vacc dropout compute is called")// 返回值是[]封装的list, 如果topi.nn.vacc_dropout(inputs[0])的返回值是list，就不用封装成list// 直接返回topi.nn.vacc_dropout(inputs[0])即可return [topi.nn.vacc_dropout(inputs[0])]
注意：也可以在c++端为算子注册compute，，通过.set_attr来实现，例如：
.set_attr
在c++端和python端只能选其一。
注意：
1 如果是在c++端实现计算，那么该relay算子的计算不会调用topi中的算子实现计算。
2 如果是在python端实现计算，注意该返回值的类型是列表类型，用[]封装。如果topi….中定义的算子的返回值本身就是列表，那么在本函数中就不用再包装成一个列表。
(2)注册算子的调度
# 注册算子的调度@register_schedule("nn.vacc_dropout")def vacc_dropout_schedule(attrs, outs, target):print("relay/op: vacc dropout schedule is called, output len %d"%(len(outs)))return topi.generic.schedule_extern(outs)
也可以直接通过如下方式来注册到已有的调度：
register_schedule
例如：register_schedule(“log”, schedule_broadcast)
schedule_broadcast是已定义的调度：schedule_broadcast = schedule_injective
def schedule_injective(attrs, outputs, target):"""Generic schedule for binary broadcast."""with target:return topi.generic.schedule_injective(outputs)
(3)注册pattern
register_pattern(“nn.relu”, OpPattern.ELEMWISE)
它等效于在c++中的.set_attr
自定义说明：
vacc_dropout_schedule是调度相关的，保持固定格式，只需修改下print内容即可。
vacc_dropout_compute是计算相关的，保持固定格式，需要自定义该函数，例如，topi.nn.vacc_dropout是在topi/python/topi/nn/vacc_op.py中定义的函数，见下一部分。
2 定义topi算子topi.nn.vacc_dropout
所谓通用计算，即没有指定特殊的device target （如cuda），采用cpu计算。
在第4步中注册计算中所调用的topi.nn.vacc_dropout， 定义在topi/python/topi/nn/vacc_op.py，代码如下：
@tvm.target.generic_funcdef vacc_dropout(data):"""vacc_dropout operator.Parameters----------input : tvm.Tensor4-D with shape [batch, in_channel, in_height, in_width]Returns-------output : tvm.Tensor4-D with shape [batch, in_channel, out_height, out_width]"""print("topi.nn.generic vacc_dropout")begin = []end = data.shapestrides = []dims = len(data.shape)for i in range(dims):begin.append(0)if i < dims - 2:strides.append(1)else:strides.append(2)return cpp.strided_slice(data, begin, end, strides) //调用c++中的函数，或python中的函数
说明：@tvm.target.generic_func装饰器是定义在python/tvm/target.py
3 实现该算子的在自定义硬件设备上的计算
3.1 在vacc.py中注册计算的函数
例如，在vacc/python/vacc/top/vacc.py文件中如下代码（我们的硬件叫vacc，所有算子在vacc上的计算都在vacc/python/vacc/top/vacc.py中实现）：
@autotvm.register_topi_compute(topi.nn.vacc_dropout, 'vacc', 'direct')def _declaration_vacc_dropout(cfg, data):print("vacc.top.nn.vacc_dropout compute")return vacnn.vacc_dropout(data)
说明：
@autotvm.register_topi_compute:装饰器定义在/home/xxw/dev/tvm/python/tvm/autotvm/task/topi_integration.py
@autotvm.register_topi_compute: 第1个参数对应的是topi.nn.vacc_dropout，就是在第5步中定义的topi算子，2个参数是我们的硬件是vacc。@autotvm.register_topi_compute的作用就是在使用targert是vacc时，将topi.nn.vacc_dropout转到_declaration_vacc_dropout，也即是自定义的vacc的计算。
在此处会调用vacnn实现在vacc设备的计算：vacnn.vacc_dropout(data)。
3.2 在vacnn.py中定义计算的函数
vacnn.vacc_dropout是咱们自己的硬件vacc中实现的，该文件是python/tvm/contrib/vacnn.py，源码如下：
def vacc_dropout(input):N, C, H, W = topi.util.get_const_tuple(input.shape)oH = (H + 1) // 2oW = (W + 1) // 2oshape = [N, C, oH, oW]return _api.extern(oshape, [input],lambda ins, outs: _intrin.call_packed("tvm.contrib.vacnn.vacc_dropout",ins[0],outs[0]), name="vacc_dropout")
注意：该函数最后调用了一个ID为tvm.contrib.vacnn.vacc_dropout的函数。
3.3 在vacnn.cc中注册ID为tvm.contrib.vacnn.vacc_dropout的函数
void VacnnDropout(TVMArgs args, TVMRetValue* ret) {DLTensor* src = args[0];DLTensor* dst = args[1];VACNNThreadEntry* entry_ptr = VACNNThreadEntry::ThreadLocal();vacnnHandle_t handle = entry_ptr->handle;auto* repo = entry_ptr->container_repo;auto* queue = reinterpret_cast<InstructionQueue<InstructionConvolution, InsnFieldsMatrixMani>*>(repo->GetContainer(kContainerQueue, kComputationMatrix, handle->core_num,true));vacnnTensorDescriptor_t t_desc;vacnnDataType_t src_data_type = vacnnDataType::DLTypeToVACNNType(src->dtype);for (int i = 0; i < handle->core_num; ++i) {handle->core_id = i;// Get descriptor for srcint src_node_id = GetNodeId(src->data);std::vector<int64_t> src_shape =GetTensorShapeByCore(src, src_node_id, handle->core_id);VACNN_CALL(vacnnSetTensor4dDescriptor(&t_desc, vacnnTensorFormat_t(0), src_data_type, src->ndim,static_cast<int32_t>(src_shape[0]), static_cast<int32_t>(src_shape[1]),static_cast<int32_t>(src_shape[2]),static_cast<int32_t>(src_shape[3])));// Get address for srcvoid* src_addr = GetDataAddress(src->data);// Get address for dstvoid* dst_addr = GetDataAddress(dst->data);InstructionConvolution* insn = queue->GetItem(i);VACNN_INSN_GEN(vacnnDropout(handle, insn, t_desc, src_addr, dst_addr));}// Push the queue to device runtime when instructions all generated.if (queue->FullLoaded()) {// VACNN_INSN_PUSH(pad_forward)VACNN_INSN_CACHE(dropout, VACNN_OP_DROPOUT, VACNN_ISA_CONVOLUTION)repo->ResetContainer();}}TVM_REGISTER_GLOBAL("tvm.contrib.vacnn.vacc_dropout").set_body([](TVMArgs args, TVMRetValue* ret) { VacnnDropout(args, ret); });
而在VacnnDropout中，调用的很多函数都是NNlib中的定义的函数，例如vacnnDropout。
如果在NNlib中没有相关函数的实现，可以暂时使用模拟的函数替代。例如，vacc_slice算子的NNlib中的函数没有完成，就在vacnn_sim.cc中注册了一个ID为tvm.contrib.vacnn.vacc_slice_forward_sim的函数。
TVM_REGISTER_GLOBAL("tvm.contrib.vacnn.vacc_slice_forward_sim").set_body([](TVMArgs args, TVMRetValue *ret) {DLTensor *src = args[0];DLTensor *dst = args[5];printf("fake: vacnn.vacc_slice \n");std::cout << "vacc_slice_forward src:" << src->data << std::endl;std::cout << "vacc_slice_forward dst:" << dst->data << std::endl;memcpy(dst->data, src->data, 16 * 4);printf("dst[0] is: %f\n", ((float *)(dst->data))[0]);});
然后，在vacc.py中调用的是tvm.contrib.vacnn.vacc_slice_forward_sim
def vacc_slice_forward(x, begin, end):n = len(x.shape)oshape = []for i in range(n):if i == n - 2:# H axisoshape.append(end[0] - begin[0])elif i == n - 1:# W axisoshape.append(end[1] - begin[1])else:oshape.append(x.shape[i])return _api.extern(oshape, [x],lambda ins, outs: _intrin.call_packed("tvm.contrib.vacnn.vacc_slice_forward_sim", #注意，这是一个模拟的函数，仅供测试ins[0],begin[0],begin[1],end[0],end[1],outs[0]), name="vacc_slice_forward")
4 使用和测试算子
c++
static const Op& dropout_op = Op::Get("nn.vacc_dropout");
python
def test_conv2d_1x1_s2():n, c, h, w = 4, 3, 224, 224x = relay.var("x", relay.ty.TensorType((n, c, h, w), "float32"))w1 = relay.var("w1")w2 = relay.var("w2")def before():y = relay.nn.conv2d(x, w1, strides=(2,2), kernel_size=(1, 1), channels=8)y = relay.vacc_activation('sigmoid', y, 1.0, 2.0)y = relay.nn.conv2d(y, w2, strides=(2,2), kernel_size=(1,1), channels=4)y = relay.nn.relu(y)return relay.Function([x, w1, w2], y)def expect():y = relay.nn.conv2d(x, w1, strides=(1,1), kernel_size=(1, 1), channels=8)y = relay.vacc_activation('sigmoid', y, 1.0, 2.0)y = relay.vacc_dropout(y)y = relay.nn.conv2d(y, w2, strides=(1,1), kernel_size=(1,1), channels=4)y = relay.nn.relu(y)y = relay.vacc_dropout(y)return relay.Function([x, w1, w2], y)converted = run_convert_strides(before())expected = run_infer_type(expect())assert analysis.alpha_equal(converted, expected)
…
若有收获，就点个赞吧
0 人点赞
