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

眾所周知，一個(gè)完整的模型實(shí)際上是由一系列算子組成的，所以如果我們想讓編譯器更通用，那么支持盡可能多的算子就是一個(gè)繞不開的工作。

這樣無論算子是來自onnx、Caffe、PyTorch中的哪個(gè)框架，我們都可以在TPU-MLIR中找到對(duì)應(yīng)的算子來表達(dá)。

首先，要添加一個(gè)新的算子，我們就需要像前端轉(zhuǎn)換那一集里提到的先進(jìn)行算子定義。

在 MLIR 中，您可以直接使用 TableGen 工具來完成定義工作，而不是自己實(shí)現(xiàn)所有包含每個(gè)算子的輸入、輸出和屬性的 cpp 模板。

在 TPU-MLIR 中，不同 dialect 的算子定義在不同的 td 文件中，這些算子將在編譯器 build 時(shí)注冊(cè)在相應(yīng)的 Dialect 下。

但是定義部分只是生成了模板，也就是說，我們的編譯器還不知道這個(gè)算子會(huì)對(duì)輸入張量做什么處理，所以我們需要通過實(shí)現(xiàn)相應(yīng)目錄下的 inference 方法來完成這部分工作。

在 Top dialect 中，除了 inference 接口，我們還需要為每個(gè)算子實(shí)現(xiàn)是 FLOPs 和 Shape 接口。 前者用于計(jì)算浮點(diǎn)運(yùn)算量，后者用于在輸出 shape 未知的情況下推理出輸出 shape。

在 MLIR 中，我們有 RankedTensorType 和 UnRankedTensorType。

這些接口的聲明是在 td 文件中被要求的，所以所有從 Top_Op 類派生的算子都需要聲明這些接口。

同樣，我們還必須為每個(gè) Tpu 算子實(shí)現(xiàn) inference 接口。 由于我們可以直接從 Top 算子獲取 FLOPs 和 Shape 信息，所以這里不需要再實(shí)現(xiàn)這些接口。

由于 Top 和 Tpu 算子是在 CPU 上做推理工作，所以有時(shí)我們會(huì)把推理工作交給 oneDNN，一個(gè)跨平臺(tái)的神經(jīng)網(wǎng)絡(luò)庫，主要用于提高 CPU 上的推理性能。 不過這部分我就不再細(xì)說，如果大家有興趣的話，我們可以再做一個(gè)視頻來介紹一下。

所以如果大家想了解 oneDNN 的話，記得在視頻底下留言讓我們知道。

我們知道，TPU 算子最終會(huì)被用于不同硬件的代碼生成，所以對(duì)于這個(gè) Dialect 中的算子，需要為每個(gè)硬件實(shí)現(xiàn)額外的接口。

其中 LocalGenInterface 用于應(yīng)用了 LayerGroup 的算子，而 沒有應(yīng)用 LayerGroup 的算子則會(huì)使用 GlobalGenInterface。 所以你會(huì)看到所有的算子中都有 GlobalGenInterface，但只有其中一部分算子實(shí)現(xiàn)了 LocalGen。

在 GlobalGen 中，張量在 Global Memory 里，因此我們需要做的是準(zhǔn)備后端 API 所需的所有參數(shù)，例如算子的屬性以及輸入和輸出張量的 Global 地址。

對(duì)于 LocalGen，張量位于 Local Memory 中，這意味著它已經(jīng)完成了將 tensor 從 Global 到 Local Mmeory 的搬運(yùn)，因此我們需要調(diào)用 local 的后端 API。 此外，在量化的情況下，有時(shí)我們需要計(jì)算緩沖區(qū)大小以存儲(chǔ)中間結(jié)果。 這是因?yàn)橹虚g結(jié)果通常以更高位的數(shù)據(jù)類型存儲(chǔ)。 比如在 int8 量化中，我們需要先將計(jì)算結(jié)果存儲(chǔ)為 int16 或者 int32 數(shù)據(jù)，然后再重新量化回 int8。

完成定義和接口實(shí)現(xiàn)工作后，還有一件需要完成的事情就是 lowering。

在 TopToTpu pass 中，我們需要應(yīng)用算子轉(zhuǎn)換的 Pattern set，這需要我們?yōu)槊總€(gè)硬件中的每個(gè)算子實(shí)現(xiàn)轉(zhuǎn)換 Pattern。

一共要做 3 步，首先，在頭文件中聲明 Lowering pattern。 接著，實(shí)現(xiàn)該 Pattern， 然后將其添加到 Pattern set 中。

如本例所示，我們?cè)趯?shí)現(xiàn) Pattern 部分主要要做的是將當(dāng)前的 Top op 替換為對(duì)應(yīng)的 Tpu op，并根據(jù)指定的量化模式設(shè)置該 op 的 Type。

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

審核編輯：湯梓紅