摘要
MLIR[1]是谷歌推出的開源編譯框架。本文基於論文：SyFER-MLIR : Integrating Fully Homomorphic Encryption Into the MLIR Compiler Framework [2]，介紹將全同態加密技術整合到MLIR編譯框架中，達到“一次加密，多裝置編譯”的目的。這個例項對於如何使用MLIR，具有非常好的啟發性。本文分為兩部分，第一部分是MLIR相關技術介紹，另一部分是介紹同態加密技術及其在MLIR中的整合。
MLIR是谷歌團隊針對異構系統提出的開源編譯框架，旨在打造具有可複用性和可擴充套件性的編譯器基礎設施。正如其名，MLIR主要特點在於多層級表達，在不同層級間實現轉換，最佳化以及程式碼生成。MLIR可以針對領域專用架構架構（DSA：Domain Specific Architecture），比如現在出現的很多ML加速器，快速搭建編譯器；同時也能夠連線現有的編譯技術，比如LLVM。MLIR源起於Google Tensor Flow框架的編譯生態中存在的諸多挑戰[3]，如圖1所示。
MLIR的核心是為構建現代編譯器提供豐富的基礎元件，這些基礎元件用於不同層級的中間表示的轉化及最佳化中。不同於傳統編譯器透過一層的中間表達（IR: Intermediate Representation）直接翻譯為可執行指令，MLIR透過多層級的中間表達實現轉換，在每層級的IR表示中，MLIR提供豐富的最佳化元件進行最佳化。使用者也可以自己開發最佳化元件，最終翻譯為可執行指令，對映到不同的裝置。MLIR的設計理念使得編譯器整體框架從抽象到具體結構清晰，層次分明，任務明確。MLIR專案文件中的幾個“Rationale”文章是理解其設計思想的很好參考。
MLIR將中間表達也稱之為方言(dialect)，方言都符合相同的語法格式，如圖2。在MLIR語法中，採用的基於SSA(Static Single Assignment)的資料結構，Operation（下文翻譯為操作）是一等公民(The first class citizen)，每個Op都有自己的識別符號(Identifier)，Op接受輸入引數(Value)同時也可以返回引數，Op具有屬性(Attribute)，屬性在編譯階段需要保持為常量，用於在編譯階段告知編譯器Op的基本資訊。具體可以參考後文[5]。
圖3 MLIR 程式碼生成的流水線(Codegen Pipeline)
圖3表示的是MLIR方言不斷Lowering的流水線，圖中左側部分表示在每層方言中進行的最佳化操作，中間部分表示方言之間的轉換。前端的輸入來自TensorFlow，對應於編譯器的結構，可以看成為MLIR的前端。在方言的不斷轉換的過程中，體現的是MLIR的重要特點progressive lowering。以下以圖3為例介紹MLIR的lowering生成程式碼的過程。
1.TF的描述轉換為HLO(High-Level Optimizer)方言[6]，HLO實際上是XLA（Accelerated Linear Algebra）的IR[7]表示，裡面的Op基本和XLA一一對應。
2.HLO轉換為LHLO(Late HLO)[6]，HLO和LHLO的區別在於HLO注重的是tensor的表達，不考慮到記憶體的分配，比如tensor<8x32x16xfp32>，僅僅表示為tensor，沒有具體的記憶體資訊，LHLO會為tensor開闢記憶體空間，也就是圖中的buffer assign，buffer assign相當於傳統編譯器中的記憶體分配。
3.LHLO被轉換為Linalg(Linear Algebra) 方言[8]和Affine方言[9]，Linalg為線性代數方言，Affine是針對多面體編譯的方言。在此層方言表達中完成tiling展開，for-loop迴圈最佳化操作等。
4.針對不同的硬體裝置產生程式碼，在此層轉換成LLVM IR，然後再透過LLVM的codegen生成程式碼。
MLIR提供了一個比較靈活和規範的框架，相應的，我們也看到很多基於MLIR的工作，逐漸形成一定的生態，甚至超出了ML編譯器的範疇。但是，對於MLIR的爭論和質疑也一直存在。一個近期的討論可以參考知乎問題“如何評價MLIR專案中Linalg Dialect的設計思想？”（https://www.zhihu.com/question/442964082），這裡不再贅述。
拋開這些爭議，從實際的應用上，MLIR的出現確實起到了輔助編譯器開發的作用。大家可以更快地搭建解決特定問題的編譯工具，或者將一些新的特色引入到編譯器中，比如下面將要討論的基於MLIR編譯器框架實現同態加密，非常具有啟發性。
同態加密（HE: Homomorphic Encryption)[10]是一項直接在加密資料上進行運算而無需先解密資料再運算的加密技術，最終透過解密計算結果就可以得到在非加密狀態下執行的結果。同態加密技術廣泛應用在涉及到敏感資料處理的場景，比如醫療資料分析和基因分析等。
圖4 客戶端-服務端的同態加密場景（C代表客戶端，S代表服務端）[11]
下面以圖4的客戶端-服務端同態加密使用場景介紹同態加密解密過程：
1. 客戶端將資料進行同態加密處理得到加密資料Enc(m)，m表示資訊，Enc表示加密方法。
3. 當客戶端想對自己的資料執行一個函式操作f()，會將操作的名稱傳送給服務端，比如圖中的query。
4. 服務端接收到命令後，在加密的資料上進行f()操作。
6. 客戶端透過解密的方法獲得原始資料的計算結果。
在上文關於MLIR技術介紹中提到了在MLIR的程式碼生成的流程中採用不同層級的中間表達和最佳化，那麼也理應可以在中間表達和最佳化中引入關於同態加密的技術。
同態加密分為4種[11]，本文所分析的論文針對的是全同態加密(FHE:Full Homomorphic Encryption)，以下統稱為同態加密。常見的FHE方法包括Brakerski-Gentry-Vaikuntanathan(BGV)[12]，Fan-Vercauten(B/FV)[13] 和Cheon-Kim-Kim-Song(CKKS)[14] 等。針對這些FHE方法，現有的實現是透過呼叫庫的方式實現，比如MicrosoftSEAL[15]。但是目前的HFE庫存在幾點侷限：1.只能進行有限的跨操作最佳化(cross-operation optimization)，跨操作最佳化可以理解為移除掉不必要的指令或者調整操作順序，參考後文圖7。庫雖然提供了一系列的原語(primitives)，這些原語操作都是單獨最佳化，不能進行跨操作最佳化; 2.不能進行重寫操作(rewrite)以及高層次的最佳化，比如NAND(x,x,parameters)不能重寫為NOT(x,parameters) ; 3. 不能進行低層次的最佳化，比如巢狀迴圈最佳化和迴圈排程; 4.缺少模組化，這樣導致更改加密方法時，需要重寫程式碼。
針對以上基於FHE庫實現方式存在的缺點，論文中提出了利用MLIR的編譯框架產生加密的程式，其核心思想在於將FHE庫改寫為方言，類似前文提及到的將XLA改為HLO，目前實現的是Gentry-Sahai-Waters(GSW)[16]和B/FV[13]，具體流程如下圖5所示。
圖5 GSW和B/FVprogressive lowering pipeline的簡易框圖
正常的MLIR編譯流程如圖5中的左圖所示，基於GSW和BFV的加密編譯流程如圖5中中圖和右圖所示。普通程式碼作為輸入，採用MLIR的progressive lowering的功能將程式碼轉換為較低級別的IR，這些IR在經過自開發的lowering規則轉換為加密的IR，同態加密的方言轉換過程為Linalg->Affine/SCF(StructureControl Flow)[16] ->Standard -> GSW/BFVIR，實現了同態加密的IR然後按照正常的MLIR編譯流程編譯，最後在LLVM層級可以根據不同的硬體裝置產生不同的程式碼，這樣帶來了“一次加密，多裝置編譯”的便捷性。
從GSW和BFV向下lowering的詳細過程如圖6所示，GSW首先lowering到S-GSW(Simple GSW)、Standard以及LinAlg方言，S-GSW然後lowering到Affine，SCF和standard方言中，最後Affine，Standard，LinAlg和SCF都lowering到LLVM。圖中的箭頭表示的Pass，實現的作用是從一種方言向另一種方言的轉換，顏色越深表示實現難度越大，藍色是難度最低的，MLIR的框架中提供相關的Pass實現，使用者也可以自己開發Pass。可惜文章沒有給出GSW和S-GSW以及BFV和S-BFV之間的區別。
在同態加密的方法中，每個加密的函式操作能表示為一個DAG圖（Directed Acyclic Graph），每個節點表示一個原語，每條邊表示一個原語的輸出作為另外一個原語的輸入。
針對本文選擇的GSW和BVF，它們的原語都很簡單，GSW的原語是二進位制的閘電路(Binary Gate)，BVF的原語是加法和乘法，簡單的原語極大地降低了最佳化的難度。
在Lowering的過程中，論文使用了三類最佳化技術。第一，使用重寫機制完成最佳化操作，重寫操作可以簡單地理解為使用一種表示方式替代一種或者多種的表示，目的在於使用高效能的Op替代低效能的Op，比如前文提到的使用NOT(x,parameters)替代NAND(x,x,parameters)，此外還提及到了複雜的重寫機制，比如修剪程式碼，刪除常量和最佳化兩輸入和三輸入的子圖，但文中沒給出具體的例項說明。第二，實現跨操作（cross-operation）最佳化，跨操作最佳化的目的是刪除不必要的指令，比如文中舉例合併數論轉換(NTT:Number Theoretic Transform)加法的操作，如圖7所示。第三，針對不同的硬體進行迴圈最佳化，這也是MLIR本身具有的特性。
本文的核心在於將FHE兩種的庫GSW和BFV改寫為MLIR中的方言，並實現了各個層級的轉換和最佳化，最終將同態加密技術整合到了MLIR編譯框架中。整體論文的實現難度相對較低，也比較遺憾文中沒給出實驗結果，但作者的思路是值得肯定的。
在AI應用領域中，對於端側敏感資料的處理往往採用本地私有化部署的方式，在端側完成推理，但端側裝置的效能有時會成為瓶頸。或許存在下面兩種可能，在端裝置上使用集成了同態加密的MLIR編譯器，將編譯後的指令傳到雲端執行，然後雲端將執行結果返回給端裝置；雲端部署同態加密MLIR編譯器，端裝置將同態加密的資料傳送到雲端處理，雲端部署編譯器在於能夠適配不同的加密方法，具有更高的編譯效能。
如果讀者對其他同態加密的編譯技術感興趣，可以參考論文[18]。
由於水平有限，文中存在不足的地方請各位讀者批評指正，也歡迎大家一起參與我們的討論。
[1] MLIR:https://mlir.llvm.org/
[2] Govindarajan, S. and W. Moses. “SyFER-MLIR: Integrating Fully Homomorphic Encryption Into the MLIR Compiler Framework.” (2020).
[3] https://llvm.org/devmtg/2019-04/slides/Keynote-ShpeismanLattner-MLIR.pdf
[4] https://www.tensorflow.org/mlir/overview
[5] Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: A compiler infrastructure for the end of moore’s law, 2020
[6] https://github.com/tensorflow/mlir-hlo
[7] https://www.tensorflow.org/mlir/xla_gpu_codegen
[8] https://mlir.llvm.org/docs/Rationale/RationaleLinalgDialect/
[9] https://mlir.llvm.org/docs/Dialects/Affine/
[10] https://en.wikipedia.org/wiki/Homomorphic_encryption
[11] A. Acar, H. Aksu, A. S. Uluagac, and M. Conti. A survey on homomorphic encryption schemes: Theory and implementation. ACM Computing Surveys (CSUR), 51(4):1–35, 2018.
[12] Z. Brakerski, C. Gentry, and V. Vaikuntanathan. (leveled) fully homomorphic encryption without bootstrapping. ACM Transactions on Computation Theory (TOCT), 6(3):1–36, 2014.
[13] J. Fan and F. Vercauteren. Somewhat practical fully homomorphic encryption. IACR Cryptol. ePrint Arch., 2012:144, 2012.
[14.J. H. Cheon, A. Kim, M. Kim, and Y. Song. Homomorphic encryption for arithmetic of approximate numbers. In International Conference on the Theory and Application of Cryptology and Information Security, pages 409–437. Springer, 2017.
[15
]
H. Chen, K. Laine, and R. Player. Simple encrypted arithmetic libraryseal v2. 1. In International Conference on Financial Cryptography and Data Security, pages 3–18. Springer, 2017.
[16] C. Gentry, A. Sahai, and B. Waters. Homomorphic encryption from learning with errors: Conceptually-simpler, asymptotically-faster, attribute-based. In Annual Cryptology Conference, pages 75–92. Springer, 2013.
[17] https://mlir.llvm.org/docs/Dialects/SCFDialect/
[18] https://arxiv.org/pdf/2101.07078.pdf
壁仞科技研究院作為壁仞科技的前沿研究部門，旨在研究新型智慧計算系統的關鍵技術，重點關注新型架構，先進編譯技術和設計方法學，並將逐漸拓展研究方向，探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則，將積極投入各類產學研合作並參與開源社群的建設，為相關領域的技術進步做出自己的貢獻。
