如何给TPU-MLIR添加新的算子

众所周知，一个完整的模型实际上是由一系列算子组成的，所以如果我们想让编译器更通用，那么支持尽可能多的算子就是一个绕不开的工作。

这样无论算子是来自onnx、Caffe、PyTorch中的哪个框架，我们都可以在TPU-MLIR中找到对应的算子来表达。

首先，要添加一个新的算子，我们就需要像前端转换那一集里提到的先进行算子定义。

在 MLIR 中，您可以直接使用 TableGen 工具来完成定义工作，而不是自己实现所有包含每个算子的输入、输出和属性的 cpp 模板。

在 TPU-MLIR 中，不同 dialect 的算子定义在不同的 td 文件中，这些算子将在编译器 build 时注册在相应的 Dialect 下。

但是定义部分只是生成了模板，也就是说，我们的编译器还不知道这个算子会对输入张量做什么处理，所以我们需要通过实现相应目录下的 inference 方法来完成这部分工作。

在 Top dialect 中，除了 inference 接口，我们还需要为每个算子实现是 FLOPs 和 Shape 接口。 前者用于计算浮点运算量，后者用于在输出 shape 未知的情况下推理出输出 shape。

在 MLIR 中，我们有 RankedTensorType 和 UnRankedTensorType。

这些接口的声明是在 td 文件中被要求的，所以所有从 Top_Op 类派生的算子都需要声明这些接口。

同样，我们还必须为每个 Tpu 算子实现 inference 接口。 由于我们可以直接从 Top 算子获取 FLOPs 和 Shape 信息，所以这里不需要再实现这些接口。

由于 Top 和 Tpu 算子是在 CPU 上做推理工作，所以有时我们会把推理工作交给 oneDNN，一个跨平台的神经网络库，主要用于提高 CPU 上的推理性能。 不过这部分我就不再细说，如果大家有兴趣的话，我们可以再做一个视频来介绍一下。

所以如果大家想了解 oneDNN 的话，记得在视频底下留言让我们知道。

我们知道，TPU 算子最终会被用于不同硬件的代码生成，所以对于这个 Dialect 中的算子，需要为每个硬件实现额外的接口。

其中 LocalGenInterface 用于应用了 LayerGroup 的算子，而 没有应用 LayerGroup 的算子则会使用 GlobalGenInterface。 所以你会看到所有的算子中都有 GlobalGenInterface，但只有其中一部分算子实现了 LocalGen。

在 GlobalGen 中，张量在 Global Memory 里，因此我们需要做的是准备后端 API 所需的所有参数，例如算子的属性以及输入和输出张量的 Global 地址。

对于 LocalGen，张量位于 Local Memory 中，这意味着它已经完成了将 tensor 从 Global 到 Local Mmeory 的搬运，因此我们需要调用 local 的后端 API。 此外，在量化的情况下，有时我们需要计算缓冲区大小以存储中间结果。 这是因为中间结果通常以更高位的数据类型存储。 比如在 int8 量化中，我们需要先将计算结果存储为 int16 或者 int32 数据，然后再重新量化回 int8。

完成定义和接口实现工作后，还有一件需要完成的事情就是 lowering。

在 TopToTpu pass 中，我们需要应用算子转换的 Pattern set，这需要我们为每个硬件中的每个算子实现转换 Pattern。

一共要做 3 步，首先，在头文件中声明 Lowering pattern。 接着，实现该 Pattern， 然后将其添加到 Pattern set 中。

如本例所示，我们在实现 Pattern 部分主要要做的是将当前的 Top op 替换为对应的 Tpu op，并根据指定的量化模式设置该 op 的 Type。

至此，添加新算子的工作就完成了。

审核编辑：汤梓红