MLIR多层编译框架实现全同态加密的讨论

发布时间:2021-02-27 11:10

摘要

MLIR[1]是谷歌推出的开源编译框架。本文基于论文:SyFER-MLIR : Integrating Fully Homomorphic Encryption Into the MLIR Compiler Framework [2],介绍将全同态加密技术集成到MLIR编译框架中,达到“一次加密,多设备编译”的目的。这个实例对于如何使用MLIR,具有非常好的启发性。本文分为两部分,第一部分是MLIR相关技术介绍,另一部分是介绍同态加密技术及其在MLIR中的集成。

 

MLIR简介

MLIR是谷歌团队针对异构系统提出的开源编译框架,旨在打造具有可复用性和可扩展性的编译器基础设施。正如其名,MLIR主要特点在于多层级表达,在不同层级间实现转换,优化以及代码生成。MLIR可以针对领域专用架构架构(DSA:Domain Specific Architecture),比如现在出现的很多ML加速器,快速搭建编译器;同时也能够连接现有的编译技术,比如LLVM。MLIR源起于Google Tensor Flow框架的编译生态中存在的诸多挑战[3],如图1所示。

 

图1 Tensor Flow 编译生态[3]

 

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]。

 

图2 MLIR 语法格式

 

图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表示加密方法。

2. 客户端将加密的数据发送到服务端。

3. 当客户端想对自己的数据执行一个函数操作f(),会将操作的名称发送给服务端,比如图中的query。

4. 服务端接收到命令后,在加密的数据上进行f()操作。

5. 服务端返回f()操作的结果。

6. 客户端通过解密的方法获得原始数据的计算结果。
在上文关于MLIR技术介绍中提到了在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层级可以根据不同的硬件设备产生不同的代码,这样带来了“一次加密,多设备编译”的便捷性。

图6 GSW和B/FVLowering的详细过程

 

从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本身具有的特性。

 

图7 多项式乘加的优化

 

本文的核心在于将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

上一个: 深度学习在反演成像中的最新进展

下一个: 面向大规模图计算的系统优化(3)

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多