深度解读隐语密态计算设备 SPU

SPU 是 Secretflow Processing Unit 的简称，它作为隐语平台的密态计算单元，为隐语提供安全计算服务。

1.SPU 概念理解

密态计算单元这个概念听起来比较晦涩，我们用一个实际的例子介绍一下 SPU 的作用。假设要用 JAX 写逻辑回归（SPU 不依赖 JAX，选择 JAX 因为简单），代码如下：

import jaximport jax.numpy as jnp
def sigmoid(x):    return 1 / (1 + jnp.exp(-x))
def loss(x, y, w):    pred = sigmoid(jnp.dot(x, w))    label_prob = pred * y + (1 - pred) * (1 - y)    return -jnp.sum(jnp.log(label_prob))
def logit_regression(x, y, epochs=3, step_size=0.1):    w = jnp.zeros(x.shape[1])    for _ in range(epochs):        grad = jax.grad(loss, 2)(x, y, w)        w -= grad * step_size    return w
x, y = load_full_dataset()w = fit(x, y)

【注】L15 使用了 JAX 的自动求导功能，对 loss function 进行了求导

对于上述程序，JAX 提供了 jit（全称 Just In Time，是编译技术的一种）方法，在不用任何算法改动的前提下将上述程序编译到 GPU 上，从而加速执行。

jax.jit(logit_regression)(x, y)  # jax.jit将logit_regression翻译到GPU执行

JAX 本质解决了两个问题

1. 降低开发成本，在用户无感的情况下，利用 GPU/TPU 进行加速

2. 降低学习成本，通过兼容 numpy API 并提供自动求到来吸引开发者

沿着这个思路，对安全计算做个类比

• JAX 可以利用并行设备进行计算加速，我们可否用安全设备进行安全加固？

• JAX 可以复用 numpy API，我们是否可以复用其他 AI 框架 API？

带着这两个问题 ，SPU 的核心 API 就是这样一个函数，将 AI 模型翻译到安全设备上执行

spu.jit(logit_regression)(x, y)

为了实现这个函数，我们需要两个子模块

• 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 进行解释执行，如图：

工作流程

基本工作流程是：

1. 开发人员用自己熟悉的框架建模，然后用 AI compiler 将模型编译成 XLA IR。

2. SPU compiler 将 XLA IR 编译成 SPU IR（SPU 字节码）

3. 参与方（Alice/Bob/Charlie) 将数据 infeed 给 SPU VM

4. SPU VM 执行字节码，接收输入，安全计算，并且产生输出

5. 参与方协商将结果解密输出到某处

Why XLA？

这里对于 XLA 不熟悉的同学进行一个简单介绍，XLA 是一种针对特定领域的线性代数编译器， 是 tensorflow 内部实现的一个子模块，使用编译器相关技术用来加速模型的执行。

我们可以将 SPU 理解成一个带安全语义的 Backend，对接 XLA 的理由是 Tensorflow/Jax/PyTorch 计算图最终都可以翻译到 XLA IR，只要 SPU 可以解释和执行 XLA IR，理论上就可以原生支持多种 AI 前端。所以理论上并不是选择了 XLA 编译器，而是选择了 XLA IR 作为 AI 和 MPC 的桥梁。

编译和执行过程

• 前端，我们依赖 AI 前端将 python 代码翻译成 XLA IR

• 编译器，我们使用 MLIR 技术栈对 HLO 进行优化并翻译成 PPHLO（SPU 字节码）

• 运行时，我们逐渐将 Tensor ops 拆解，经过 SPU HAL(硬件抽象层，处理 fxp/int)，最终 dispatch 到协议层

• 协议层只需要实现 Ring or Field 上的基础运算即可

最终，通过编译时和运行时的层层翻译，SPU 将 AI 前端和 MPC 后端解耦，使得在 SPU 中扩展的任何安全协议都可以无感的支持多种前端。

优化

SPU 完全自主研发，所以我们可以针对安全计算的特点进行优化，比如

• 针对 MPC 延时敏感的计算类型进行并发调度

• 针对特定协议设计特殊的 VM 指令集

• 基于 C++语言，开发者可以得到 Low-level access，更有效的榨取性能

具体细节在此不再复述。SPU 依然在一个高速迭代过程中，隐语期待更多的 AI 专家、编译专家、安全专家参与共建。

获取隐语代码可访问以下地址：github.com/secretflow

了解更多关于隐语的信息可访问官网：www.secretflow.org.cn

扫码关注隐语微信公众号：隐语的小剧场，获取更多资讯！