在Relay中注冊新TVM算子
在本文件中，將介紹在Relay中注冊新TVM算子所需的步驟。將以添加累積算子的PR為例。PR本身建立在另一個PR的基礎上，該PR添加了一個累積和運算。
注冊新算子需要幾個步驟：

1. Add an attribute node declaring fixed arguments which are known at compile time
2. Write a type relation for your operation to integrate into Relay’s type system.
3. Use the RELAY_REGISTER_OP macro in C++ to register the operator’s arity, type, and other hints for the compiler
4. Write how the operator is computed
5. Register the compute, schedule with the relay operator
6. Define a C++ function to produce a call node for the operator and registering a Python API hook for the function
7. Wrapping the above Python API hook in a neater interface
8. Writing tests for the new relay operator

詳細過程

1. Add an attribute node declaring fixed arguments which are known at compile time
   屬性是在編譯時指定的固定參數。卷積算子的stride和伸縮，屬于卷積算子的屬性節點中的字段。
   屬性應該在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（）存在類似的接口。
   因此，在include/tvm/relay/attrs/transform.h中定義屬性時，選擇操作的坐標軸、累積數據類型和獨占性，作為struct結構體的適當字段。
   /*! \brief Attributes used in cumsum and cumprod operator */
   struct ScanopAttrs : public tvm::AttrsNode {
   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());
   TVM_ATTR_FIELD(dtype).describe(“Output data type”).set_default(NullValue());
   TVM_ATTR_FIELD(exclusive)
   .describe(“The first element is not included”)
   .set_default(Bool(false));
   }
   };

2. Writing a Type Relation
   為了允許在注冊算子時具有靈活性，以及在Relay中表達類型時，具有更大的表達能力和粒度，使用輸入和輸出類型之間的關系輸入算子。這些關系表示為函數，這些函數接受輸入類型和輸出類型列表（這些類型中的任何一種都可能不完整），并返回滿足該關系的輸入和輸出類型列表。這包括可在編譯時靜態確定的shape信息。本質上，算子的關系除了計算輸出類型外，還可以強制執行所有必要的類型規則（即通過檢查輸入類型）。
   累積積和算子的類型關系，可在src/relay/op/tensor/transform.cc中找到：
   TVM_REGISTER_NODE_TYPE(ScanopAttrs);
   bool ScanopRel(const Array& 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();
   if (data == nullptr) {
   ICHECK(types[0].as())
   << "Scanop: expect input type to be TensorType but get " << types[0];
   return false;
   }

   const auto* param = attrs.as();

   auto dtype = param->dtype;
   if (dtype.is_void()) {
   dtype = data->dtype;
   }

   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;
   }

3. Relating the Arity and Attributes to an Operation
   然后，注冊新算子的名稱，調用接口對其進行注釋。C++中的RELAY_REGISTER_OP宏，允許開發人員指定Relay中的算子的以下信息：
   ? Arity (number of arguments)
   ? Names and descriptions for positional arguments
   ? Support level (1 indicates an internal intrinsic; higher numbers indicate less integral or externally supported operators)
   ? A type relation for the operator
   ? Other annotations useful when optimizing the operation.
   Once again we add this to 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”, 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”, kOpaque);
在本例中，TOpPattern是對編譯器的一個關于算子所執行的計算模式的提示，這對于融合算子可能很有用。kOpaque說明，TVM不要費心嘗試融合這個算子。
4. Defining the Compute of the Operation
雖然現在已經為算子定義了接口，仍然需要定義如何執行累計和與積的實際計算。
編寫此代碼超出了本文的范圍。現在，假設有一個經過良好測試的算子計算實現。有關如何執行此算子的更多詳細信息，建議查閱有關張量表達式、TVM算子清單（topi）的教程，并查看python/TVM/topi/scan.py和python/TVM/topi/cuda/scan.py中的，gpu版本中的示例累積和與積實現。在累積和與積運算的情況下，直接在TIR中寫入內容，這是張量表達式和topi將lower into降低到的表示形式。
5. Hooking up Compute and Strategy with Relay
實現了計算功能之后，現在需要粘到Relay算子上。在TVM中，不僅要定義計算，還要定義算子的調度。strategy是一種選擇要使用的計算和調度的方法。例如，對于二維卷積，可能認識到正在進行深度卷積，分派到更高效的計算和調度。然而，在例子中，除了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),
wrap_topi_schedule(topi.generic.schedule_extern),
name=“cumsum.generic”,
)
return strategy

@override_native_generic_func(“cumprod_strategy”)
def cumprod_strategy(attrs, inputs, out_type, target):
“”“cumprod generic strategy”""
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_scanop(topi.cumprod),
wrap_topi_schedule(topi.generic.schedule_extern),
name=“cumprod.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

@cumprod_strategy.register([“cuda”, “gpu”])
def cumprod_strategy_cuda(attrs, inputs, out_type, target):
“”“cumprod cuda strategy”""
strategy = _op.OpStrategy()
strategy.add_implementation(
wrap_compute_scanop(topi.cuda.cumprod),
wrap_topi_schedule(topi.cuda.schedule_scan),
name=“cumprod.cuda”,
)
return strategy
在每個strategy中，定義了編寫的計算和要在add_implementation()中使用的調度。最后，將strategy與python/tvm/relay/op/_transform.py中定義的Relay算子鏈接并進行計算：

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)

cumprod

@_reg.register_compute(“cumprod”)
def compute_cumprod(attrs, inputs, output_type):
“”“Compute definition of cumprod”""
return [topi.cumprod(inputs[0], attrs.axis, attrs.dtype, attrs.exclusive)]

_reg.register_strategy(“cumprod”, strategy.cumprod_strategy)
_reg.register_shape_func(“cumprod”, False, elemwise_shape_func)

shape函數用于確定給定動態shape張量的輸出shape。在這種情況下，告訴TVM輸出shape將與輸入shape相同。
6. Creating a Relay Call Node and Exposing a Python Hoo
現在有一個工作算子，現在只需要通過Relay正確地調用節點。這一步只需要編寫一個函數，將參數作為Relay表達式。，傳遞給算子，并將調用節點返回給算子（即，應放置在RelayAST中的節點，在該節點中，算子將被調用）。
目前不支持調用屬性和類型參數（最后兩個字段），使用Op:：Get從算子注冊表獲取算子信息，并將參數傳遞給調用節點就足夠了，如下所示。在src/relay/op/tensor/transform.cc中：
Expr MakeCumsum(Expr data, Integer axis, DataType dtype, Bool exclusive) {
auto attrs = make_object();
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);

Expr MakeCumprod(Expr data, Integer axis, DataType dtype, Bool exclusive) {
auto attrs = make_object();
attrs->dtype = dtype;
attrs->axis = axis;
attrs->exclusive = exclusive;
static const Op& op = Op::Get(“cumprod”);
return Call(op, {data}, Attrs(attrs), {});
}

TVM_REGISTER_GLOBAL(“relay.op._make.cumsum”).set_body_typed(MakeCumprod);
Where TVM_REGISTER_GLOBAL exposes the MakeCumsum and MakeCumprod functions in Python via relay.op._make.cumsum(…) and relay.op._make.cumsum(…).
7. Including a Cleaner Python API Hook
通常，Relay中的約定是，通過TVM_REGISTER_GLOBAL導出的函數，應該封裝在單獨的Python函數中，而不是直接在Python中調用。對于算子，在
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)
注意，這些Python包裝器，也可能是向算子提供更簡單接口的好機會。例如，concat算子，被注冊為只使用一個算子，即一個具有要連接的張量的元組，但是Python包裝器，將張量作為參數，并在生成調用節點之前，組合成一個元組：
def concat(*args):
“”"Concatenate the input tensors along the zero axis.

Parameters
----------
args: list of TensorReturns
-------
tensor: The concatenated tensor.
"""
tup = Tuple(list(args))
return _make.concat(tup)

8. Writing Unit Tests!
   一些單元測試示例，可以在tests/python/relay/test_op_level3.py中找到，用于累積總和與乘積算子。
   其它
   梯度算子
   梯度算子對于編寫Relay中的可微程序非常重要。雖然Relay的autodiff算法可以區分一流的語言結構，但算子是不透明的。由于Relay無法查看實現，因此必須提供明確的差異化規則。
   Python和C++都可以用來編寫梯度算子，但是，例子集中在Python上，因為更常用。
   在Python中添加梯度
   Python梯度算子的集合可以在Python/tvm/relay/op/_tensor_grad.py中找到。將介紹兩個具有代表性的示例：sigmoid和multiply。
   @register_gradient(“sigmoid”)
   def sigmoid_grad(orig, grad):
   “”“Returns [grad * sigmoid(x) * (1 - sigmoid(x))].”""
   return [grad * orig * (ones_like(orig) - orig)]

這里的輸入是原始算子orig和要累加的梯度。返回的是一個列表，其中第i個索引處的元素是算子相對于算子第i個輸入的導數。通常，梯度將返回一個列表，其中包含的元素數量與基本算子的輸入數量相同。
進一步分析這個定義之前，首先應該回顧一下sigmoid函數的導數：
上面的定義看起來類似于數學定義，但有一個重要的補充，將在下面描述。
術語orig*（類似于（orig）-orig）直接匹配導數，因為這里的orig是sigmoid函數，但不只是對如何計算這個函數的梯度感興趣。將這個梯度與其它梯度組合起來，這樣就可以在整個程序中累積梯度。

這就是梯度術語的意義所在。在表達式gradorig（one_like（orig）-orig）中，乘以grad，表示如何使用到目前為止的梯度合成導數。
現在，考慮乘法，一個稍微有趣的例子：
@register_gradient(“multiply”)
def multiply_grad(orig, grad):
“”“Returns [grad * y, grad * x]”""
x, y = orig.args
return [collapse_sum_like(grad * y, x),
collapse_sum_like(grad * x, y)]

在本例中，返回的列表中有兩個元素，因為multiply是一個二進制算子。回想一下，如果偏導數是

有一個乘法所需的步驟，對于sigmoid不是必需的，因為乘法具有廣播語義。由于梯度的shape可能與輸入的shape不匹配，使用collapse_sum_like來獲取梯度grad * 項的內容，并使shape與要區分的輸入的shape匹配。
Adding a Gradient in C++
在C++中添加一個梯度，類似于在Python中添加，但是，用于注冊的接口略有不同。
首先，確保包含src/relay/transforms/pattern_utils.h。提供了用于在RelayAST中創建節點的 helper函數。然后，類似于Python示例的方式，定義梯度：
tvm::Array MultiplyGrad(const Expr& orig_call, const Expr& output_grad) {
const Call& call = orig_call.Downcast();
return { CollapseSumLike(Multiply(output_grad, call.args[1]), call.args[0]),
CollapseSumLike(Multiply(output_grad, call.args[0]), call.args[1]) };
}
注意，在C++中，不能使用Python中的算子重載，并且需要進行downcast，實現更加冗長。即使如此，也可以很容易地驗證這個定義，是否反映了Python中的早期示例。
現在，不需要使用Python裝飾器，而是需要在基礎算子的注冊末尾，添加一個對“FPrimalGradient”的set_attr調用，以便注冊梯度。
RELAY_REGISTER_OP(“multiply”)
// …
// Set other attributes
// …
.set_attr(“FPrimalGradient”, MultiplyGrad);

參考鏈接：
https://tvm.apache.org/docs/dev/relay_add_op.html

總結

以上是生活随笔為你收集整理的在Relay中注册新TVM算子的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。