Apache TVM 是一個深度的深度學習編譯框架，適用于 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 →https://tvm.hyper.ai/

本文檔將以添加?cumulative product?算子的 PR（基于?cumulative sum?算子 PR）為例，介紹在 Relay 中注冊一個新的 TVM 算子所需的步驟。

注冊一個新的算子需要如下幾個步驟：

1. 添加一個屬性節點，聲明在編譯時已知的固定參數
2. 為算子編寫一個類型關系，以整合到 Relay 的類型系統中
3. 使用 C++ 中的?RELAY_REGISTER_OP?宏，為編譯器注冊算子的數量、類型和其他提示
4. 編寫算子的計算方式
5. 用 Relay 注冊算子和 schedule
6. 定義一個為算子產生調用節點的 C++ 函數，并為該函數注冊一個 Python API hook
7. 將上述 Python API hook 放在一個更簡潔的接口中
8. 為新的 Relay 算子編寫測試

1. 定義屬性節點

屬性是在編譯時已知的固定參數。卷積算子的步長和擴張可能屬于卷積算子屬性節點字段的恰當示例。

屬性應在文件夾?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.
    參數
    ----------
    data : relay.Expr 類型
        算子的輸入數據。
    axis : int 類型，可選
        Axis along which the cumulative product is computed. The default (None) is to compute the cumprod over the flattened array.
    dtype : string 類型，可選
        Type of the returned array and of the accumulator in which the elements are multiplied.
        如果 dtype 沒有被指定, 那么它默認為 data 的 dtype。
    exclusive : bool 類型，可選
        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.
    返回
    -------
    result : relay.Expr 類型
        如果 axis 不為空的話，結果的大小和形狀和 data 一樣。
        如果 axis 為空的話, 結果是一個一維數組。
    """

cumsum()?存在類似的接口。

因此，在?include/tvm/relay/attrs/transform.h?中定義屬性時，可以選擇算子的 axis、accumulation dtype 及 exclusivity 作為結構體的合適字段。

/*! 用在 cumsum 和 cumprod 算子中的簡單屬性 */
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));
  }
};

2. 編寫類型關系

為了提高注冊算子的靈活性，在 Relay 中表示類型時更突出，算子使用輸入和輸出類型之間的關系進行類型化。這些關系被表示為函數，它接收一個輸入類型和輸出類型的列表（這些類型中的任何一個都可能是不完整的），然后返回一個滿足關系的輸入和輸出類型的列表，包括可以在編譯時靜態確定的形狀信息?；旧?，一個算子的關系除了計算輸出類型外，還可以執行所有必要的類型化規則（即通過檢查輸入類型）。

在?src/relay/op/tensor/transform.cc?中可以找到 cumulative product 與 cumulative product 算子的類型關系。

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

    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. 將參數數量和屬性與算子關聯起來

注冊新算子的名稱，并為其添加調用接口的注解。C++ 中的?RELAY_REGISTER_OP?宏允許開發者在 Relay 中指定一個算子的以下信息：

• 參數數量
• 位置參數的名稱和描述
• 支持級別（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);

在這種情況下，TOpPattern?是對編譯器關于算子執行的計算模式的提示，這對于融合算子可能很有用。kOpaque?提示 TVM 無需融合這個算子。

4. 定義算子的計算

為算子定義接口后，仍需定義如何執行 cumulative sum 和 cumulative product 的實際計算。

假設算子計算的實現方式，經過了多輪測試且表現良好。推薦查看?張量表達式教程、TVM 算子清單（topi）、python/tvm/topi/scan.py?中 cumulative sum 及 cumulative product 相關實現案例，以及?python/tvm/topi/cuda/scan.py?中的 GPU 版本。在 cumulative sum 及 cumulative product 算子中，可以直接用?TIR，張量表達式及 topi 降級后表示為 TIR。

5. 將計算（compute）和策略（strategy）與 Relay 關聯起來

實現計算函數后，需要將其與 Relay 算子粘合在一起。在 TVM 中，這意味著不僅要定義 computation，還要定義算子的 schedule。策略決定使用哪種 computation 及 schedule。例如，對于二維卷積，識別出這屬于一種深度卷積后，最終將其分配給一個更有效的 computation 和 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 基本策略"""
    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 基本策略"""
    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 = _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 = _op.OpStrategy()
    strategy.add_implementation(
        wrap_compute_scanop(topi.cuda.cumprod),
        wrap_topi_schedule(topi.cuda.schedule_scan),
        name="cumprod.cuda",
    )
    return strategy

每個策略都定義了寫入的 compute 以及在?add_implementation()?中使用的 schedule。最后，將 strategy 和 compute 與python/tvm/relay/op/_transform.py?中定義的 Relay 算子關聯起來。

# cumsum
@_reg.register_compute("cumsum")
def compute_cumsum(attrs, inputs, output_type):
    """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):
    """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 函數用于確定 output shape，給定一個動態 shaped tensor。在這種情況下，TVM 的 output shape 與 input shape 保持一致。

6. 創建 Relay 調用節點并提供 Python Hook

現在已經有了一個可以運行的算子，接下來只需通過一個 Relay 調用節點（Relay Call Node）正確地調用即可。這一步需要簡單地編寫一個函數，接收算子的參數（作為 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<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);

Expr MakeCumprod(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("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(...)?分別暴露（expose）Python 中的?MakeCumsum?和?MakeCumprod?函數。

7. 包含一個更簡潔的 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 wrapper 也可能為算子提供更簡潔的接口。例如?concat?算子被注冊為只接受一個算子（即一個帶有要連接的張量的元組），但是 Python wrapper 將張量作為參數，并在產生調用節點之前將它們組合成一個元組。

def concat(*args):
    """圍繞零軸連接輸入張量。

    參數
    ----------
    args: Tensor 列表

    返回
    -------
    tensor: 連接的張量。
    """
    tup = Tuple(list(args))
    return _make.concat(tup)

8. 編寫單元測試

更多用于 cumulative sum 和 cumulative product 算子的單元測試示例，請查看?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):
    """返回 [grad * sigmoid(x) * (1 - sigmoid(x))]."""
    return [grad * orig * (ones_like(orig) - orig)]

這里的輸入是原始算子?orig?以及梯度算子?grad，返回是一個列表，其中第 i 個索引的元素，是算子相對于算子第 i 個輸入的導數。通常，梯度算子將返回一個列表，其元素的個數和基礎算子（base operator）的輸入一樣多。

進一步分析這個定義之前，先回憶一下 sigmoid 函數的導數：?σ/?x=σ(x)(1?σ(x))。上面的定義看起來類似于數學定義，但有一個重要的補充：

術語?orig * (ones_like(orig) - orig)?直接匹配導數，因為這里的?orig?是 sigmoid 函數。除了要了解如何計算該函數的梯度之外，還要掌握該梯度與其他梯度組合的方法，即在整個程序中累積梯度。這就是?grad?的作用。在表達式?grad * orig * (ones_like(orig) - orig)中，乘以?grad?指定了到目前為止如何用梯度組合導數。

接下來請看?multiply的示例：

@register_gradient("multiply")
def multiply_grad(orig, grad):
    """返回 [grad * y, grad * x]"""
    x, y = orig.args
    return [collapse_sum_like(grad * y, x),
            collapse_sum_like(grad * x, y)]

在此示例中，返回列表中有兩個元素，因為?multiply?是二元運算符（binary operator）。如果 f(x,y) = xy，偏導數是 ?f / ?x = y 和 ?f / ?y = x。

與?sigmoid?相比，multiply?需要一個額外的步驟，因為?multiply?具有廣播語義（broadcasting semantics）。由于?grad?的 shape 可能與輸入 shape 不匹配，所以我們使用?collapse_sum_like?來獲取?grad * <var>?項的內容，并使其 shape 與做微分的輸入 shape 相匹配。

在 C++ 中添加梯度算子

在 C++ 中添加梯度算子的方法，與在 Python 中添加梯度算子類似，但注冊的接口略有不同。

首先，確保?src/relay/transforms/pattern_utils.h?被包含在內。它提供了用于在 Relay AST 中創建節點的輔助函數。定義梯度算子的方式與 Python 類似：

tvm::Array<Expr> MultiplyGrad(const Expr& orig_call, const Expr& output_grad) {
    const Call& call = orig_call.Downcast<Call>();
    return { CollapseSumLike(Multiply(output_grad, call.args[1]), call.args[0]),
             CollapseSumLike(Multiply(output_grad, call.args[0]), call.args[1]) };
}

注意，在 C++ 中不能使用與 Python 相同的運算符重載（operator overloading），而是需要向下轉換，因此實現更加冗長。即便如此，我們仍然可以輕易地驗證這個定義反映了 Python 中先前的例子。

要注冊梯度算子，這里無需使用 Python 修飾器，只需要在基礎算子注冊的末尾添加?set_attr?調用 “FPrimalGradient” 即可。

RELAY_REGISTER_OP("multiply")
    // ...
    // 設置其他屬性
    // ...
    .set_attr<FPrimalGradient>("FPrimalGradient", MultiplyGrad);