【TVM】在TVM Relay中创建一个自定义操作符

# 1. Introduction
* 如题，[tvm官方文档介绍](https://tvm.apache.org/docs/dev/how_to/relay_add_op.html) ，通过一个已有示例描述了创建自定义Op的过程。过程非常详细，但是仍然存在一些不利于理解的部分。包含一些其他部分的信息交叉，本文以一个自定义的过程进行了实现，并完成相关tvm的编译。并且针对文档中一些不好理解的部分做一个描述和分析。

# 2. Description
* 整个创建自定义运算被分解成了8个步骤：
	1. `定义属性节点`（Defining an Attribute Node）：首先，需要为操作符定义属性节点。属性节点用来描述操作符的参数和配置。这些属性将被用来在计算过程中传递和使用。
	2. `编写类型关系`（Writing a Type Relation）：接下来，需要为操作符定义类型关系。这包括输入和输出的类型检查和推断逻辑，以确保操作符能够处理正确类型的数据。
	3. `将参数个数和属性关联到操作符`（Relating the Arity and Attributes to an Operation）：在这一步骤中，要将操作符的参数个数（arity）和属性与具体的操作符实现关联起来。需要通过宏或函数来注册操作符，并说明其参数个数、参数名及描述、支持级别、类型关系和其他优化相关的注释。
	4. `定义操作的计算逻辑`（Defining the Compute of the Operation）：这里需要定义操作符的具体计算逻辑。这一步涉及如何实现操作符的计算，比如矩阵乘法、加法等基础运算的实现细节。
	5. `将计算逻辑和策略与 Relay 关联`（Hooking up Compute and Strategy with Relay）：将定义好的计算逻辑和策略与 Relay 系统进行集成。这使得 Relay 能够识别和使用新定义的操作符，并能够在模型优化和执行过程中应用该操作符。
	6. `创建 Relay 调用节点并暴露 Python 接口`（Creating a Relay Call Node and Exposing a Python Hook）：为操作符创建一个 Relay 调用节点，并提供 Python 接口，使得用户能够通过 Python 代码调用和使用这个新操作符。
	7. `包含更清晰的 Python API 接口`（Including a Cleaner Python API Hook）：提供更清晰和简化的 Python API 接口，以便于用户在 Python 环境中更方便地使用和调用这个新操作符。
	8. `编写单元测试`（Writing Unit Tests!）：最后，编写单元测试来验证新操作符的正确性。这些测试可以帮助检测和修复潜在的问题，确保操作符在各种情况下都能正常工作。

![](/media/202406/2024-06-14_112353_9682000.02863005381158179.png)

接下来我们基于一个完整的自定义示例来完成整个过程。

# 3. Operations

## 3.0. Pre-processing

* 克隆整个tvm工程`git clone https://github.com/apache/tvm`

* 由于是一个cmake工程，创建build目录，对tvm的相关依赖进行开关设置，并安装相关依赖。

* 可以看到tvm的整个工程源码，能够在机器上正确完成编译
![](/media/202406/2024-06-14_113219_5624530.06783700503654433.png)

## 3.1. Defining an Attribute Node
``` c++
/*! \brief Attributes used in test operator */
// 定义一个叫Test的Attributes没有任何操作，只有一个参数叫tested，用于定义是否执行了Test
struct TestAttrs : public tvm::AttrsNode<TestAttrs> {
  bool tested;
  TVM_DECLARE_ATTRS(TestAttrs, "relay.attrs.TestAttrs") {
    TVM_ATTR_FIELD(tested).describe("Whether the test elements are tested").set_default(false);
  }
};

```

## 3.2. Writing a Type Relation

* 不做任何处理只是过一个Op

``` c++
TVM_REGISTER_NODE_TYPE(TestAttrs);

bool TestopRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) {
    // num_input: 1 (data)
    ICHECK_EQ(num_inputs, 1) << "Expects one input";
    // 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 (nullptr == data) {
        ICHECK(types[0].as<IncompleteTypeNode>()) << "Testop: expect unput type to be TensorType but get "<< types[0];
        return false;
    }

const auto* param = attrs.as<TestAttrs>();

auto tested = param->tested;
    if (!tested) {
        tested = true;
    }

return true;
}
```

## 3.3. Relating the Arity and Attributes to an Operation

* 将具体数值和参数与操作做关联的过程比较好理解，但是整体各个参数的含义例如level理解起来会有点晦涩。

```shell
在 Relay 中，support level 是一个用于描述操作符（operator）或运算符在系统中的重要性和支持程度的标志。

它可以帮助开发人员和优化工具确定某个操作符在代码生成、优化或调试过程中的优先级和处理方式。

使用 support level 的场景
优化器决策：

优化器在优化计算图时，可以根据操作符的 support level 来决定是否进行某些特定优化。
基础操作符（低 support level）通常需要进行较多的优化处理，因为它们的使用频率高，性能影响大。
高 support level 的操作符可能在优化过程中被优先级较低的处理，或者作为优化的目标被替换成更高效的实现。
调试和分析：

在调试和分析过程中，可以根据 support level 快速了解操作符的优先级和重要性。
开发人员可以根据 support level 的不同，有选择地查看和调试某些操作符，帮助快速定位问题和优化性能。
代码生成：

在生成代码（例如生成低级别的机器码或硬件描述代码）时，support level 可以帮助确定生成的优先顺序和资源分配。
基础操作符可能被优先处理并映射到高效的硬件指令或加速器，而高 support level 的操作符则可能使用通用的实现。
文档和描述：

在操作符的文档中，support level 可以提供有关该操作符在系统中的角色和使用建议。
用户可以根据 support level 判断某个操作符是否适合他们的应用场景。
实例
假设在一个神经网络框架中，我们有以下操作符和它们的 support level：

Level 1: add（加法），multiply（乘法）等。这些操作符是核心计算的基础，被广泛应用在各种网络层和优化中。
Level 2: relu（激活函数），softmax（输出层函数）等。虽然重要，但它们的使用场景较为特定。
Level 3: depthwise_conv2d（深度卷积），这种操作符可能只在特定的网络结构或优化路径中才会使用。
Level 4: bilinear_resize（双线性插值），这种操作符可能用于特定的图像处理任务，在一般神经网络计算中较少使用。
```

``` c++
RELAY_REGISTER_OP("testtest")
    .describe(R"doc(this is the describe of the doc)doc" TVM_ADD_FILELINE)
    .set_num_inputs(1)
    .add_argument("data", "Tensor", "The input tensor.")
    .set_support_level(3)
    .add_type_rel("Testtest", TestopRel)
```

## 3.4. Defining the Compute of the Operation

## 3.5. Hooking up Compute and Strategy with Relay

## 3.6. Creating a Relay Call Node and Exposing a Python Hook

``` c++
Expr MakeTesttest(Expr data, bool tested) {
    auto attrs = make_object<TestAttrs>();
    attrs->tested = tested;
    static const Op& op = Op::Get("testtest");
    return Call(op, {data}, Attrs(attrs), {});
}

TVM_REGISTER_GLOBAL("relay.op._make.testtest").set_body_typed(MakeTesttest);
```
## 3.7. Including a Cleaner Python API Hook

* 在 `python/tvm/relay/op/transform.py` 中添加

``` py
def testtest(data, tested=None):
    return _make.testtest(data, tested)

```

## 3.8. Writing Unit Tests!

* `tests/python/relay/test_op_level3.py` 中添加测试用例

# Annexe
* https://tvm.apache.org/docs/dev/how_to/relay_add_op.html