开发者必看：深度解读隐语密态计算设备 SPU
“隐语” 是开源的可信隐私计算框架，内置 MPC、TEE、同态等多种密态计算虚拟设备供灵活选择，提供丰富的联邦学习算法和差分隐私机制。
代码开源：
SPU 是 Secretflow Processing Unit 的简称，它作为隐语平台的密态计算单元，为隐语提供安全的计算服务。
1.SPU 概念理解
密态计算单元这个概念听起来比较晦涩，我们用一个实际的例子介绍一下 SPU 的作用。
假设要用 JAX 写逻辑回归（SPU 不依赖 JAX，选择 JAX 因为简单），代码如下：
注】L15 使用了 JAX 的自动求导功能，对 loss function 进行了求导
对于上述程序，JAX 提供了 jit（全称 Just In Time，是编译技术的一种）方法，在不用任何算法改动的前提下将上述程序编译到 GPU 上，从而加速执行。
JAX 本质解决了两个问题
降低开发成本，在用户无感的情况下，利用 GPU/TPU 进行加速
降低学习成本，通过兼容 numpy API 并提供自动求到来吸引开发者
沿着这个思路，对安全计算做个类比
JAX 可以利用并行设备进行计算加速，我们可否用安全设备进行安全加固？
JAX 可以复用 numpy API，我们是否可以复用其他 AI 框架 API？
带着这两个问题，SPU 的核心 API 就是这样一个函数，将 AI 模型翻译到安全设备上执行
为了实现这个函数，我们需要两个子模块
SPU (Jit) Compiler: 将原生的 AI 程序翻译成 SPU 字节码
SPU VM：一个带安全语义的虚拟机，解释和执行 SPU 字节码
实际中实现稍微复杂一些
将 Tensorflow/PyTorch/JAX 翻译成 SPU 字节码本身是个复杂繁琐的工作
SPU 字节码使用的是密文类型系统，譬如没有 f32/f64 等，类型系统需要重新设计
SPU 后端是 MPC，本质上是个分布式系统，需要处理分布式系统的通信和协作
上述程序中，变量
x, y可能由不同的参与方提供，所以 IO 模块有些特别
细节会在后续部分进行简单介绍，在此不再复述。
2.SPU 的功能作用
介绍完 SPU 是什么，我们再来理一下为什么。
市面上的隐私计算框架有很多，比如 TFE，CrypTen，MP-SPDZ 等，为什么我们要重新造一套轮子呢？
领域之间的距离
安全机器学习是一个交叉领域，其实 AI 和安全之间有相当的距离。比如
安全开发者更关注基础算子，比如加减乘除的安全性
AI 开发者更关注高阶算子，比如 conv，tensordot
高阶算子和基础算子之间，有很大一段距离，譬如：
机器学习编译器处理的 lowering/tiling/fusing 等
运行时处理的的调度，并发等
无论是基于 AI 框架（TFE/CrypTen），还是从安全计算出发的框架（SPDZ），都有自己的问题。前者往往难部署，难做安全领域特定的优化。后者往往会需要写一些 Toy AI 框架，学习成本高。
SPU 试图缝合这两个领域之间的间隙，使得：
向上，SPU 原生对接 AI 框架（TF/JAX/PyTorch)，降低 AI 开发者的学习和开发成本。
向下，SPU 提供纯粹安全语义接口，只需要实现很少的安全协议（比如加乘与或）就能跑起来复杂的模型，让目光更聚焦安全本身。
算力和需求的距离
近些年，密态计算（MPC/HE）在算力上都巨大的进步，但是密态算力和 AI 的算法需求还是难以匹配。在算力无法匹配算法的时候，一个直观的想法就是 “明密文混合”，用来做安全和性能的 tradeoff。比如联邦学习，将算法的某一个子步骤使用安全计算实现，牺牲局部安全性以换取更高的性能。
隐语提供了非常自由的明密文混合编程范式，我们不限制明文的引擎，也不限制密文引擎，开发者可以用他自己熟悉的框架开发，然后标记其中的某一部分用明文引擎跑，另一部分用 SPU 跑。比如：
【注】图中 MPC Device 就是 SPU 实现的
作为对比，从安全和性能这种的角度，无论 TFE/CrypTen/SPDZ 等都很难进行这种 tradeoff。
理论和落地的距离
多方安全计算天然是一个分布式的系统，部署模式非常多样，比如：
论文中经常假设计算方和数据提供方分离（outsourcing）
真正进行业务落地时，数据提供方往往同时也是计算方（colocated）
在一个复杂的隐私计算网络里，计算方和数据方可以是任意组合的（hybrid）
如图，我们用三角形表示计算节点，圆形表示数据提供节点（不同颜色表示互不信任）
SPU 被设计成部署模式透明的，不用修改任何一行代码，你的模型都可以在上述任何一种部署场景上被安全且正确的执行。并且（相对于基于 AI 平台的隐私计算框架）SPU 运行时非常的轻量级，不需要 Python runtime，可以方便的进行部署和集成。
所以，Why SPU？
作为 AI 开发者，你不需要任何安全背景，就可以将你现有的模型安全的应用到多方数据上。
作为安全开发者，你不需要任何 AI 背景，仅仅实现安全计算的基本算子，就可以支持多种前端框架。
并且，你可以方便的部署和运维，在安全和性能之间折中，找到最佳的落地方案。
3.SPU 的架构
介绍完是什么和为什么之后，我们简单介绍一下 SPU 的架构。
SPU 上层对接了 XLA-HLO，然后利用 MLIR 将 HLO 翻译成 SPU IR，最后交给 SPU VM 进行解释执行，如图：
Workflow
基本工作流程是：
开发人员用自己熟悉的框架建模，然后用 AI compiler 将模型编译成 XLA IR。
SPU compiler 将 XLA IR 编译成 SPU IR（SPU 字节码）
参与方（Alice/Bob/Charlie) 将数据 infeed 给 SPU VM
SPU VM 执行字节码，接收输入，安全计算，并且产生输出
参与方协商将结果解密输出到某处
Why XLA？
这里对于 XLA 不熟悉的同学进行一个简单介绍，XLA 是一种针对特定领域的线性代数编译器，是 tensorflow 内部实现的一个子模块，使用编译器相关技术用来加速模型的执行。
我们可以将 SPU 理解成一个带安全语义的 Backend，对接 XLA 的理由是 Tensorflow/Jax/PyTorch 计算图最终都可以翻译到 XLA IR，只要 SPU 可以解释和执行 XLA IR，理论上就可以原生支持多种 AI 前端。所以理论上并不是选择了 XLA 编译器，而是选择了 XLA IR 作为 AI 和 MPC 的桥梁。
Arch
下面简单介绍一下 SPU 的编译和执行过程
前端，我们依赖 AI 前端将 python 代码翻译成 XLA IR
编译器，我们使用 MLIR 技术栈对 HLO 进行优化并翻译成 PPHLO（SPU 字节码）
运行时，我们逐渐将 Tensor ops 拆解，经过 SPU HAL (硬件抽象层，处理 fxp/int)，最终 dispatch 到协议层
协议层只需要实现 Ring or Field 上的基础运算即可
最终，通过编译时和运行时的层层翻译，SPU 将 AI 前端和 MPC 后端解耦，使得在 SPU 中扩展的任何安全协议都可以无感的支持多种前端。
Optimization
SPU 完全自主研发，所以我们可以针对安全计算的特点进行优化，比如
针对 MPC 延时敏感的计算类型进行并发调度
针对特定协议设计特殊的 VM 指令集
基于 C++ 语言，开发者可以得到 Low-level access，更有效的榨取性能
具体细节在此不再复述。
SPU 依然在一个高速迭代过程中，隐语期待更多的 AI 专家、编译专家、安全专家参与共建。
隐语社区：
隐语官网：https://www.secretflow.org.cn
👇欢迎关注：
公众号：隐语的小剧场
B 站：隐语 secretflow
邮箱：secretflow-contact@service.alipay.com
关注微信公众号：隐语的小剧场 2022-08-01 加入
隐语SecretFlow是蚂蚁自主研发的隐私计算开源框架，内置MPC、TEE、同态等多种密态计算虚拟设备供灵活选择。同时我们专注于隐私计算领域任何前沿技术、最新动态、行业资讯，隐语期待您的加入！
评论