【源碼研讀】MLIR Dialect 分層設計

以「疑問 - 求解」的形式來組織調研，此處記錄整個過程。

1. MLIR 中的 Dialect 是「分層」設計的麼？

先問是不是，再談爲什麼。從 LLVM 社區可以看出，至少在做 Codegen 時，是採用了「分層」的思想來逐步 Lowering 的（具體見下圖）。MLIR 爲編譯優化而生，分層 Lowering 是比較符合設計直覺的，在多硬件、多場景的 Dialect 擴展性上具有天然優勢。

在 MLIR 的設計初衷裏，「擴展性」是非常強烈的需求考量，要能夠縱向支撐用戶不同的場景，也要橫向支撐不同硬件接入進來，低成本地編譯優化。如此，MLIR 社區纔有可能做大做強。

另外，在 MLIR 核心開發者 ZhangLei 的技術博客裏，也提到了 MLIR 的宏觀圖景，原文表述如下：

從類型的角度，恰當分層的軟件棧需要支持對張量、buffer、向量、標量等進行建模以及一步步分解和遞降（即 Lowering）。從操作的角度，我們需要計算和控制流。控制流可以是顯式的基礎塊跳轉，也可以內含於結構化操作之中。

從上圖來對比飛槳現有的新 IR 體系設計，以及考慮未來新增的「分層」Dialect：

TF、TFLite 層對應於新 IR 體系中的 PaddleDialect ? 即組網編譯期高層表示（與主框架耦合緊密，放到 paddle 目錄下）
MHLO、TOSA 層面是爲了「屏蔽非常多的框架」信息，起到「沙漏」的作用，以形成某種協調的「規範」（個人理解類似ONNX角色，是第一層屏蔽的協議）
也對應於PaddleDialect，但僅對應於「規範的」Type和Operation子集。對於飛槳而言，其實應該要考慮規範化這一層，而不是僅僅滿足於支持自己的框架
中間層 Dialect 設計的優勢在於：「獨立且靈活」。
- 原文描述： 高層和低層的 dialect 通常處於 MLIR 系統的邊界，所以需要準確地描述某一 MLIR 之外的標的。中間層的 dialect 則沒有這樣的限制，所以中間層的 dialect 具有更大的設計空間以及更高的設計靈活性。傳統的中間表示，如 LLVM IR 或者 SPIR-V，通常都是完整 (complete) 的；它們包含所需的所有指令來表示整個 CPU 後者 GPU 程序。相較而言，中間層的 dialect 則可以認爲是部分 (partial) 中間表示。這種組織結構有助於解耦 (decoupling) 和可組合性—我們可以通過混用這一層的不同 dialect 來表示原始模型，同時不同 dialect 可以獨立發展演進。這些dialect 有的用來表示計算或者負載 (payload)，有的則表示控制流或者某種結構 (structure)。
linalg dialect 是用以表示結構的重要 dialect 之一，linalg dialect 的文檔中還有許多其他不錯的設計考慮值得一讀（TODO）
- 分爲兩大類：generic op和named op
- 其 op 既可以操作 Tensor，也可以操作 buffer，兩者分別對應於 MLIR 中的 tensor and memref 類型。兩者皆是高維的抽象，並都可以支持 dynamic shape。
- MHLO、TOSA Dialect 可以轉換爲 linalg Dialect，這種轉換會保持在張量這一抽象層級，所以其目的並非遞降，而是爲接下來的轉換做準備。
- 從文檔裏給出的代碼來看，linalg 更像是爲編譯器而生的，比較類似CINN的Lowering，但作者明確說了，這一層並未 Lowering？

**初步結論：*8 無論是從 Type 體系的宏觀圖景，還是 Codegen 時的流程，都能捕捉到 MLIR 在設計上有「分層設計」的思想。

2. 從TF、Torch、LLVM 源碼中來看，各個 Dialect 的組織形式是什麼樣子的？

2.1 TF 中的 Dialect

倉庫地址：https://github.com/tensorflow/mlir-hlo

從上述倉庫的首頁介紹來看，mlir-hlo 有兩個託管地址：

TF 主倉庫：tensorflow/compiler/xla/mlir_hlo ，援引的mlir_hlo 獨立倉庫，隨 TF 編包
獨立倉庫：https://github.com/tensorflow/mlir-hlo ，爲了方便開發者不依賴 TF 龐大的包

子問題：mlir-hlo 是什麼？TF 已經有 XLA+HLO，爲什麼還要孵化 mlir-hlo 呢？

以下是參閱了官網資料後的一些「依據」收集，和粗淺的理解。

① 用於 XLA 風格編譯的 MLIR 方言
其定義了三種方言來支持使用 MLIR 的類似 HLO 的Compile pipline：

chlo：“客戶端”HLO 方言，旨在更接近前端（包括隱式廣播語義）。
mhlo：“元”-HLO 方言；與類似xla_hlo，但具有動態形狀支持的擴展。（是基於 Tensor 層面）
lmhlo：“late”-“meta”-HLO，是執行緩衝區（即 buffer）分配後的IR。在 XLA 中，緩衝區分配是一個輔助數據結構，用於跟蹤這些信息，而這種單獨的方言在 IR 中將其具體化。（是基於Buffer層面的）

個人理解，基本流程是：hlo → chlo → mhlo → Linalg → bufferization → lmhlo

從代碼目錄組織上，也能看出 mhlo、lhlo、thlo（這是啥？）

以 mhlo 目錄中的 CMakeLists.txt 和 hlo_ops.h/.cc 爲例，評估 MHLODialect 的編譯依賴。

Dialect 定義中的頭文件依賴比較清晰，均爲 LLVM 和 mlir 底層目錄的依賴
CMakeLists.txt 中也沒有明顯發現對於 TF 主框架數據結構的依賴（可能跟 mlir-hlo 獨立倉庫有關係）

其他的一些疑問待求解：

除了 lhlo，還有單獨的一個 lhlo_gpu 目錄，定義了 LmhloGpuDialect 方言，這個是做什麼的？

2.2 Torch 中的 Dialect

倉庫地址：https://github.com/llvm/torch-mlir

目前 torch-mlir 是放在 llvm 組織下的，但尚未正式化到 llvm project 裏，也是一個「中間代理人」的角色，支撐Pytorch、JAX 模型能夠藉助此模塊流暢地接入到 MLIR 體系下，然後複用 MLIR 的多硬件編譯器優化的能力。

從目錄結構上，Dialect 下只有兩個目錄，分別是 Torch 和TorchConversion:

對於 Torch Dialect，從TorchOps.td可以看出其Operation的定義與 TorchScript 有者千絲萬縷的聯繫，其中定義瞭如下核心的 Operation:

NnModuleOp，NnModuleTerminatorOp、SlotOp：表示 torch.nn.Module
ClassTypeOp，ClassTypeTerminatorOp：用於輔助建模 torch.nn.Module
MethodOp，AttrOp，GlobalSlotOp，GlobalSlotModuleInitializerOp，GlobalSlotGetOp，GlobalSlotSetOp
PrimListUnpackOp，PrimTupleConstructOp，PrimCreateObjectOp，PrimGetAttrOp等
PrimLoopOp，PrimLoopConditionOp，PrimIfOp，PrimIfYieldOp
PrimLoadOp，PrimEnterOp，PrimExitOp
ConstantStrOp，ConstantDeviceOp，ConstantIntOp，ConstantNumberOp
DerefineOp，OperatorOp，LinearParamsCreateOp， ValueTensorLiteralOp，CopyToValueTensorOp
OverwriteTensorContentsOp，PromoteDtypesOp
ShapeCalculateOp，ShapeCalculateYieldOp，ShapeCalculateYieldShapesOp
DtypeCalculateOp，DtypeCalculateYieldOp，DtypeCalculateYieldDtypesOp

舉一個栗子，熟悉 TorchScript 的同學也許能很快地理解爲什麼需要上面這些 Operation 的定義：

 // 用於理解 NnModuleOp，SlotOp的作用
    %2 = torch.nn_module {
      torch.slot "b", %bool_true : !torch.bool
      torch.slot "i", %int3 : !torch.int
      torch.slot "f", %float : !torch.float
      torch.slot "t", %t : !torch.tensor
      torch.slot "submodule", %1 : !torch.nn.Module
    } : !torch.nn.Module<"my_class_name">
    
    
    // 用於理解 ClassTypeOp 的作用
    // A simple empty torch.class_type, with corresponding torch.nn_module.
    torch.class_type @empty {}
    %submodule = torch.nn_module {} : !torch.nn.Module<"empty">

    // A class type with many members.
    torch.class_type @test {
      torch.attr "b" : !torch.bool
      torch.attr "i" : !torch.int
      torch.attr "f" : !torch.float
      torch.attr "t" : !torch.tensor
      torch.attr "submodule" : !torch.nn.Module<"empty">
      torch.method "method", @f
    }
    torch.nn_module {
      // These must match the order and names in the `torch.class_type`.
      torch.slot "b", %bool_true : !torch.bool
      torch.slot "i", %int3 : !torch.int
      torch.slot "f", %float : !torch.float
      torch.slot "t", %t : !torch.tensor
      torch.slot "submodule", %submodule : !torch.nn.Module<"empty">
    } : !torch.nn.Module<"test">

也許從「問題空間」的角度也更好理解：如何解決用戶持有 TorchScript 得到的模型表示後的 MLIR 轉換問題，需要如何定義一個「算子集合」來承接表示 TorchScript 中的每個「操作」Operation？
另外，Torch Dialect 下 Transforms目錄下定義了非常多的「轉換」規則：

此時，我們再看 TorchConversion Dialect 目錄下的源碼，和其定位，如下是 TorchConversionBase.td 裏描述。可以看出這是一個「轉換相關」的Dialect。

這裏似乎與 TF MHLO 不同。在 MHLO 裏每層的 Dialect 定義是有明確的「表示含義」的，所有的「轉換」操作是放在 transforms 裏的
而 Torch 的 Transforms目錄放的是 Pass（也與轉換有關係），但這裏卻單獨定義了一個「轉換」相關的 Dialect （有點費解）

def TorchConversion_Dialect : Dialect {
  // `torch_conversion` is too verbose.
  let name = "torch_c";
  let cppNamespace = "::mlir::torch::TorchConversion";
  let description = [{
    This dialect contains ops and transforms for converting from the Torch
    backend contract to the linalg-on-tensors backend contract.

    This mainly consists of converting ops and types from `torch` dialect
    to the mix of dialects of the linalg-on-tensors backend contract, such as
    tensor ops being converted linalg-on-tensors and `!torch.vtensor` being
    converted to the builtin `tensor` type.
  }];

  let hasConstantMaterializer = 1;
}

在 TorchConversionOps.td 裏定義了轉換 Dialect 相關的Ops：