如何將算子添加到Relay
本文將介紹在Relay中注冊新TVM算子所需的步驟。將PR添加累積產品運算示例。PR本身建立在另一個PR的基礎上，該PR添加了一個累積和運算。
注冊新算子需要幾個步驟：
添加一個屬性節點，聲明編譯時已知的固定參數
為算子編寫一個類型關系，集成到Relay的類型系統中。
使用C++中的RELAY_REGISTER_OP宏，注冊編譯器的算子屬性、類型和其它提示
如何計算算子
向Relay算子注冊計算、調度
定義C++函數，為算子生成調用節點并注冊函數的Python API hook
將上述Python API hook包裝在一個更完備的接口中
為新Relay算子編寫測試

1. 定義屬性節點
   屬性是在編譯時已知的固定參數。卷積算子的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. 編寫類型關系
   為了允許在注冊算子時具有靈活性，以及在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. 將Arity和屬性與運算關聯
   然后，注冊新運算的名稱，用調用接口進行注釋。C++中的RELAY_REGISTER_OP宏，允許開發人員指定Relay中的算子的以下信息：
   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”, 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. 定義運算的計算
雖然現在已經為操作定義了接口，但仍然需要定義如何執行累計和與積的實際計算。
編寫此代碼超出了本問的范圍。現在，假設有一個經過良好測試的操作計算實現。有關如何執行此操作的更多詳細信息，建議查閱有關張量表達式、TVM算子清單（topi）的文件，查看python/TVM/topi/scan.py和python/TVM/topi/cuda/scan.py中的gpu版本中的示例累積和與產品實現。在累積和與積運算的情況下，直接在TIR中寫入內容，這是張量表達式和topi將降低到的表示形式。
5. 將計算和策略與Relay連接起來
實現了計算功能后，現在需要粘到Relay操作上。在TVM中，不僅要定義計算，還要定義操作的調度。策略是一種選擇要使用的計算和計劃的方法。例如，對于二維卷積，可能認識到正在進行深度卷積，因此分派到更高效的計算和調度。然而，在例子中，除了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
在每個策略中，定義了編寫的計算和要在add_implementation()中使用的調度。最后，將該策略與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. 創建Relay調用節點，開放Python hook
現在有一個工作操作，現在只需要通過Relay調用節點正確地調用。這一步只需要編寫一個函數，將參數作為Relay表達式傳遞給算子，將調用節點返回給算子（即，應該放置在Relay AST中的節點，在該節點中，算子將被調用）。
目前不支持調用屬性和類型參數（最后兩個字段），因此使用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);
其中TVM_REGISTER_GLOBAL通過relay.op._make.cumsum（…）和relay.op._make.cumsum（…），在Python中開放MakeCumsum和MakeCumprod函數。
7. 包括一個更完整的Python API Hook
通常，Realy中的約定，通過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. 單元測試！
   一些單元測試示例，可以在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\u sum\u來獲取梯度*項的內容，使shape與要區分的輸入的shape匹配。
在C++中添加梯度
在C++中添加一個梯度類似于在Python中添加一個，但是用于注冊的接口略有不同。
首先，確保包含src/relay/transforms/pattern_utils.h。提供了用于在中繼AST中創建節點的幫助器函數。然后，類似于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中的操作符重載，需要進行降維，因此實現更加冗長。即使如此，也可以很容易地驗證這個定義，是否反映了Python中的早期示例。
現在，不需要使用Python裝飾器，需要在基算子的注冊末尾，添加一個對“FPrimalGradient”的set_attr調用，以便注冊梯度。
RELAY_REGISTER_OP(“multiply”)
// …
// Set other attributes
// …
.set_attr(“FPrimalGradient”, MultiplyGrad);

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

總結

以上是生活随笔為你收集整理的如何将算子添加到Relay的全部內容，希望文章能夠幫你解決所遇到的問題。

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