深度解读隐语密态计算设备 SPU（Secretflow Processing Unit）
打开链接点亮社区Star，照亮技术的前进之路。每一个点赞，都是社区技术大佬前进的动力
Github 地址： https://github.com/secretflow
SPU 是_Secretflow Processing Unit_的简称，它作为隐语平台的密态计算单元，为隐语提供安全的计算服务。
1.SPU概念理解
密态计算单元这个概念听起来比较晦涩，我们用一个实际的例子介绍一下SPU的作用。
假设要用 JAX 写逻辑回归（SPU不依赖JAX，选择JAX因为简单），代码如下：
import jax
import 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 本质解决了两个问题
-
降低开发成本，在用户无感的情况下，利用GPU/TPU进行加速
-
降低学习成本，通过兼容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进行解释执行，如图：
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专家、编译专家、安全专家参与共建。
加入隐语社区，探索可信数据流通的未来！
隐语社区是一个面向融合可信数据空间、隐私计算、数据元件等多项数据流通利用基建技术设施在内的开源社区，致力于推动前沿技术探索、技术标准体系共建与产业应用场景共创，促进高质量数据资源流通利用与价值释放。
关注我们，获取最新资讯
Github仓库地址：
官网：
哔哩哔哩：
https://space.bilibili.com/2073575923
联系小助手：
SecretFlow04
开源可信隐私计算框架------ SecretFlow
支持MPC、FL、TEE等主流隐私计算技术，融合产学研生态共创能力，助力隐私计算更广泛应用到AI、数据分析等场景中，解决千行百业隐私保护和数据孤岛等应用痛点。
开源安全可信系统软件栈------星绽 Asterinas
聚焦安全可信技术软件栈的开源社区，面向通用计算和机密计算两大场景提供坚实的安全底座，为不断涌现的安全攸关和隐私敏感的应用保驾护航。
Github仓库地址：
官网：
本文由隐语社区统一发布，欢迎大家点 Star
