本文為tvm 教程的翻譯版。這部分介紹了如何在tvm中添加新的relay算子，具體的是以一個累乘（cumprod）算子為例進行介紹。
新增relay算子基本是下面幾個步驟：
- 定義新增算子的屬性節點（Attribute Node），聲明在編譯時已知的固定參數
- 為新增算子編寫類型關系，以集成到relay的類型系統中
- 使用C++
RELAY_REGISTER_OP宏，為新增算子注冊生命參數數量、類型、提示信息 - 算子的compute
- 注冊算子的compute、schedule
- 定義C++函數，為新增算子生成調用節點，並為該函數注冊 Python API hook
- 將上面的 Python API hook 封裝成簡潔的調用方式
- 為新的relay 算子編寫測試
新增算子的屬性節點
算子屬性是編譯期已知的參數。以卷積算子為例，strid、dilation就屬於卷積算子的屬性。這部分算子屬性定義在include/tvm/relay/attrs/下。
最終來說，我們期望定義有如下屬性說明的算子，其python側的接口如下所示
def cumprod(data, axis=None, dtype=None, exclusive=None):
"""Numpy style cumprod op. Return the cumulative inclusive product of the elements along
a given axis.
Parameters
----------
data : relay.Expr
The input data to the operator.
axis : int, optional
Axis along which the cumulative product is computed. The default (None) is to compute
the cumprod over the flattened array.
dtype : string, optional
Type of the returned array and of the accumulator in which the elements are multiplied.
If dtype is not specified, it defaults to the dtype of data.
exclusive : bool, optional
If true will return exclusive product in which the first element is not
included. In other terms, if true, the j-th output element would be
the product of the first (j-1) elements. Otherwise, it would be the product of
the first j elements. The product of zero elements will be 1.
Returns
-------
result : relay.Expr
The result has the same size as data, and the same shape as data if axis is not None.
If axis is None, the result is a 1-d array.
"""
.cumsum()有類似的接口。
因此，在定義我們新增算子（cumprod）屬性時，需要選擇操作的軸、數據類型和排他性作為屬性字段。include/tvm/relay/attrs/transform.h
ScanopAttrs 這里定義了對累加、累乘等操作的屬性定義。對累乘來說就不需要額外定義了。
/*! \brief Attributes used in cumsum and cumprod operator */
struct ScanopAttrs : public tvm::AttrsNode<ScanopAttrs> {
Integer axis;
DataType dtype;
Bool exclusive = Bool(false);
TVM_DECLARE_ATTRS(ScanopAttrs, "relay.attrs.ScanopAttrs") {
TVM_ATTR_FIELD(axis).describe("The axis to operate over").set_default(NullValue<Integer>());
TVM_ATTR_FIELD(dtype).describe("Output data type").set_default(NullValue<DataType>());
TVM_ATTR_FIELD(exclusive)
.describe("The first element is not included")
.set_default(Bool(false));
}
};
但是如果是其他的算子，需要自己定義相應的屬性節點。如BiasAdd就需要單獨定義
struct BiasAddAttrs : public tvm::AttrsNode<BiasAddAttrs> {
int axis;
TVM_DECLARE_ATTRS(BiasAddAttrs, "relay.attrs.BiasAddAttrs") {
TVM_ATTR_FIELD(axis).describe("The axis to add the bias").set_default(1);
}
};
類型推導 Type Relation
為了算子注冊的靈活性以及relay算子有更好的泛化能力，relay算子通過輸入輸出之間的類型關系來實例化。
這些關系通過一系列的函數進行表示（這些函數是以算子輸入輸出類型為參數，返回滿足類型關系的輸入輸出列表）， 、、？
這包括編譯期已知的輸入輸出的shape 信息
本質上，算子relation除了推到輸出類型外，還能夠強制指定類型規則（檢查輸入類型）。
然后就是官網教程的給的例子src/relay/op/tensor/transform.cc。這里依舊是ScanopAttrs
TVM_REGISTER_NODE_TYPE(ScanopAttrs);
bool ScanopRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) {
// types: [data, output]
ICHECK_EQ(types.size(), 2) << "Expects two types, one for the input and another for the output";
const auto* data = types[0].as<TensorTypeNode>(); //輸入的tensor信息
if (data == nullptr) {
ICHECK(types[0].as<IncompleteTypeNode>())
<< "Scanop: expect input type to be TensorType but get " << types[0];
return false;
}
const auto* param = attrs.as<ScanopAttrs>(); //算子屬性
auto dtype = param->dtype;
if (dtype.is_void()) {
dtype = data->dtype;
}
//設置輸出tensor屬性
if (param->axis.defined()) {
reporter->Assign(types[1], TensorType(data->shape, dtype));
} else {
auto prod = data->shape[0];
for (size_t i = 1; i < data->shape.size(); ++i) {
prod = prod * data->shape[i];
}
reporter->Assign(types[1], TensorType({prod}, dtype));
}
return true;
}
從上面的例子可以看出 XXXOpRel 的主要功能是根據輸入類型確定輸出類型。特別的， TensorType的構造函數可以看出，需要指定輸出的shape信息，這部分主要目的就是infershape和infertype。
關聯算子的參數數目、屬性
這一步的操作，為自定義算子注冊算子名稱，通過調用接口增加算子注釋。這里需要用到C++的宏RELAY_REGISTER_OP
涉及的參數含義如下：
- Arity（參數數量）
- 位置參數的名稱和描述
- 支持級別（1 表示內部實現;較高的數字表示較少的內部支持或外部支持的算子）
- 算子的類型關系
- 優化算子時有用的其他注釋。
src/relay/op/tensor/transform.cc
RELAY_REGISTER_OP("cumsum")
.describe(
R"doc(Return the cumulative sum of the elements along a given axis.)doc" TVM_ADD_FILELINE)
.set_num_inputs(1)
.add_argument("data", "Tensor", "The input tensor.")
.set_support_level(3)
.add_type_rel("Cumsum", ScanopRel)
.set_attr<TOpPattern>("TOpPattern", kOpaque);
RELAY_REGISTER_OP("cumprod")
.describe(
R"doc(Return the cumulative product of the elements along a given axis.)doc" TVM_ADD_FILELINE)
.set_num_inputs(1)
.add_argument("data", "Tensor", "The input tensor.")
.set_support_level(3)
.add_type_rel("Cumprod", ScanopRel)
.set_attr<TOpPattern>("TOpPattern", kOpaque);// 不融合
注：set_attr<TOpPattern>("TOpPattern", );此處表示融合算子是，跳過此算子。
編寫的算子compute
到現在，我們已經實現了算子的接口，但是還缺少算子的compute邏輯。這部分內容超出了這個教程的范圍。
對於cumprod和cumsum，CPU實現可以參考python/tvm/topi/scan.py，GPU實現可以參考python/tvm/topi/cuda/scan.py。
這里這兩個的實現，直接在TIR基礎上實現得到的。
def scanop(
data: tvm.te.Tensor,
binop: Callable[["tvm.Expr", "tvm.Expr"], "tvm.Expr"],
identity_value: "tvm.Expr",
op_name: str,
axis: Optional[int] = None,
dtype: Optional[str] = None,
exclusive: Optional[bool] = None,
) -> tvm.te.Tensor:
if dtype is None or dtype == "":
dtype = data.dtype
if exclusive is None:
exclusive = False
def maybe_cast(x):
if dtype != data.dtype:
return cast(x, dtype)
return x
axis_mul_before = 1
axis_mul_after = 1
if axis is None:
axis = 0
cumsum_axis_len = prod(data.shape)
shape = (cumsum_axis_len,)
else:
if not isinstance(axis, int):
axis = get_const_int(axis)
shape = data.shape
cumsum_axis_len = shape[axis]
if axis < 0:
axis = len(shape) + axis
for i, value in enumerate(shape, 0):
if i < axis:
axis_mul_before *= value
elif i > axis:
axis_mul_after *= value
def gen_ir(data_buf, out_buf):
ib = ir_builder.create()
data_buf = ib.buffer_ptr(data_buf)
out_buf = ib.buffer_ptr(out_buf)
with ib.for_range(0, axis_mul_before * axis_mul_after, "fused", kind="parallel") as fused:
i = fused // axis_mul_after
j = fused % axis_mul_after
base_idx = i * cumsum_axis_len * axis_mul_after + j
if exclusive:
out_buf[base_idx] = cast(identity_value, dtype)
else:
out_buf[base_idx] = maybe_cast(data_buf[base_idx])
with ib.for_range(0, cumsum_axis_len - 1, "_k") as _k:
k = _k + 1
cur_idx = base_idx + k * axis_mul_after
prev_idx = base_idx + (k - 1) * axis_mul_after
if exclusive:
out_buf[cur_idx] = binop(out_buf[prev_idx], maybe_cast(data_buf[prev_idx]))
else:
out_buf[cur_idx] = binop(out_buf[prev_idx], maybe_cast(data_buf[cur_idx]))
return ib.get()
out_buf = decl_buffer(shape, dtype, "out_buf")
return extern(
[shape],
[data],
lambda ins, outs: gen_ir(ins[0], outs[0]),
dtype=dtype,
out_buffers=[out_buf],
name=op_name,
tag=op_name,
)
def cumsum(
data: tvm.te.Tensor,
axis: Optional[int] = None,
dtype: Optional[int] = None,
exclusive: Optional[bool] = None,
) -> tvm.te.Tensor:
return scanop(
data=data,
binop=generic.add,
identity_value=0,
op_name="cumsum_generic",
axis=axis,
dtype=dtype,
exclusive=exclusive,
)
注冊算子的compute、schedule
在實現了算子compute邏輯以后，需要與我們實現的算子接口綁定在一起。在TVM中，這就需要不僅實現算子的compute接口，還要實現對應的schedule。而strategy就是對compute選擇合適的schedule。
以卷積算子為例，算子編譯時，可能會發現這是一個depthwise卷積，進而去選擇更高效的schedule實現。
一般情況下，僅僅考慮CPU、GPU版本即可。
python/tvm/relay/op/strategy/generic.py python/tvm/relay/op/strategy/cuda.py
def wrap_compute_scanop(topi_compute):
"""Wrap scanop style topi compute"""
def _compute_scanop(attrs, inputs, _):
return [topi_compute(inputs[0], attrs.axis, attrs.dtype, attrs.exclusive)]
return _compute_scanop
@override_native_generic_func("cumsum_strategy")
def cumsum_strategy(attrs, inputs, out_type, target):
"""cumsum generic strategy"""
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_scanop(topi.cumsum), #上面寫的compute
wrap_topi_schedule(topi.generic.schedule_extern),
name="cumsum.generic",
)
return strategy
@cumsum_strategy.register(["cuda", "gpu"])
def cumsum_strategy_cuda(attrs, inputs, out_type, target):
"""cumsum cuda strategy"""
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_scanop(topi.cuda.cumsum),
wrap_topi_schedule(topi.cuda.schedule_scan),
name="cumsum.cuda",
)
return strategy
對於每個strategy，與對應的compute、schedule通過add_implementation關聯起來。
這里的shape_func時對輸入時動態shape廠家推導有用。
# cumsum
@_reg.register_compute("cumsum")
def compute_cumsum(attrs, inputs, output_type):
"""Compute definition of cumsum"""
return [topi.cumsum(inputs[0], attrs.axis, attrs.dtype, attrs.exclusive)]
_reg.register_strategy("cumsum", strategy.cumsum_strategy)
_reg.register_shape_func("cumsum", False, elemwise_shape_func)
定義C++函數，為新增算子生成調用節點，並為該函數注冊 Python API hook
現在我們有一個可以調用的relay算子了，下一步就是如何通過relay call node調用。這就需要實現一個函數，傳遞相應的參數給對於的relay算子，並且返回對應算子的Call Node（這個算子最終在Relay表達式的AST里面）。
當前不支持直接調用 Attrs和參數。所以需要在函數中構造對應的AttrsNode，傳遞給對應的Call Node。
Expr MakeCumsum(Expr data, Integer axis, DataType dtype, Bool exclusive) {
auto attrs = make_object<ScanopAttrs>();
attrs->dtype = dtype;
attrs->axis = axis;
attrs->exclusive = exclusive;
static const Op& op = Op::Get("cumsum");
return Call(op, {data}, Attrs(attrs), {});
}
TVM_REGISTER_GLOBAL("relay.op._make.cumsum").set_body_typed(MakeCumsum);
Op::Get("cumsum")的實現如下。具體怎么注冊到OpRegistry的，TODO
const Op& Op::Get(const String& name) {
const OpRegEntry* reg = OpRegistry::Global()->Get(name);
ICHECK(reg != nullptr) << "AttributeError: Operator " << name << " is not registered";
return reg->op();
}
這里看一下Call的實現，實際上是得到一個call Node，里面保存了算子及其屬性信息。
Call::Call(Expr op, Array<Expr> args, Attrs attrs, Array<Type> type_args, Span span) {
ObjectPtr<CallNode> n = make_object<CallNode>();
n->op = std::move(op);
n->args = std::move(args);
n->attrs = std::move(attrs);
n->type_args = std::move(type_args);
n->span = std::move(span);
data_ = std::move(n);
}
Op::Get src/relay/op/tensor/transform.cc
相關接口暴露到python側，是通過.TVM_REGISTER_GLOBAL MakeCumsum MakeCumprod relay.op._make.cumsum(...) relay.op._make.cumsum(...)實現的。
細節TODO
將上面的 Python API hook 封裝成簡潔的調用方式
為更方便的使用，通常的做法是構造單獨的函數，因此最好封裝成更簡潔的python接口。教程的例子，定義在
TVM_REGISTER_GLOBAL python/tvm/relay/op/transform.py
def cumsum(data, axis=None, dtype=None, exclusive=None):
return _make.cumsum(data, axis, dtype, exclusive)
def cumprod(data, axis=None, dtype=None, exclusive=None):
return _make.cumprod(data, axis, dtype, exclusive)
特別的，如果不定參數的，需要包成Tuple形式進行傳遞。
def concat(*args):
"""Concatenate the input tensors along the zero axis.
Parameters
----------
args: list of Tensor
Returns
-------
tensor: The concatenated tensor.
"""
tup = Tuple(list(args))
return _make.concat(tup)
為新的relay 算子編寫測試
參考 tests/python/relay/test_op_level3.py