参考文献:
- [Har14] Harvey D. Faster arithmetic for number-theoretic transforms[J]. Journal of Symbolic Computation, 2014, 60: 113-119.
- [Sei18] Seiler G. Faster AVX2 optimized NTT multiplication for Ring-LWE lattice cryptography[J]. Cryptology ePrint Archive, 2018.
- [ZXZ+19] Zhou S, Xue H, Zhang D, et al. Preprocess-then-NTT technique and its applications to K yber and N ew H ope[C]//Information Security and Cryptology: 14th International Conference, Inscrypt 2018, Fuzhou, China, December 14-17, 2018, Revised Selected Papers 14. Springer International Publishing, 2019: 117-137.
- [ZLP21] Zhu Y, Liu Z, Pan Y. When NTT meets Karatsuba: preprocess-then-NTT technique revisited[C]//International Conference on Information and Communications Security. Cham: Springer International Publishing, 2021: 249-264.
- [CHK+21] Chung C M M, Hwang V, Kannwischer M J, et al. NTT multiplication for NTT-unfriendly rings: New speed records for Saber and NTRU on Cortex-M4 and AVX2[J]. IACR Transactions on Cryptographic Hardware and Embedded Systems, 2021: 159-188.
- [BHK+21] Becker H, Hwang V, Kannwischer M J, et al. Neon ntt: Faster dilithium, kyber, and saber on cortex-a72 and apple m1[J]. Cryptology ePrint Archive, 2021.
- [HLS+22] Hwang V, Liu J, Seiler G, et al. Verified NTT multiplications for NISTPQC KEM lattice finalists: Kyber, SABER, and NTRU[J]. IACR Transactions on Cryptographic Hardware and Embedded Systems, 2022: 718-750.
- [DL22] Duong-Ngoc P, Lee H. Configurable mixed-radix number theoretic transform architecture for lattice-based cryptography[J]. IEEE Access, 2022, 10: 12732-12741.
- [ZLH+23] Zhao Y, Liu X, Hu Y, et al. Design of an Efficient NTT/INTT Architecture with Low-Complex Memory Mapping Scheme[J]. IEEE Transactions on Circuits and Systems II: Express Briefs, 2023.
- 快速乘法技巧:Karatsuba, Toom, Good, Schonhage, Strassen, Nussbaumer
- Baby-Step Giant-Step & Homomorphic DFT
- ARM NEON技术 - 基础介绍
- Tensor Core 和 Cuda Core 的区别
- Nvidia Tensor Core-WMMA API 编程入门
文章目录
- 软件优化
- Harvey Butterfly
- Preprocess-then-NTT
- Improved PtNTT
- NTT-unfriendly rings
- 硬件优化
- Sei18
- Neon NTT
- Mixed-radix NTT
- TensorFHE
- Others
软件优化
Harvey Butterfly
在 Shoup’s NTL 中,radix-2 NTT 的蝴蝶实现如下:
它采用了 Barrett 算法的变体,Shoup’s modular multiplication:修改
W
′
≈
β
/
p
W' \approx \beta/p
W′≈β/p 为
W
≈
W
β
/
p
W \approx W\beta/p
W≈Wβ/p,于是 Barrett 取模算法就额外计算了与常数
W
W
W 的数乘运算。但是这个蝴蝶的 if-else
语句过多,容易使得 CPU 分支预测失败并导致回滚。
[Har14] 提出使用
Z
p
\mathbb Z_p
Zp 的冗余表示(
[
0
,
2
p
)
[0,2p)
[0,2p) 和
[
0
,
4
p
)
[0,4p)
[0,4p)),从而移除了一些 if-else
语句。正确性要求:Shoup 模乘算法
W
T
(
m
o
d
β
)
WT\pmod\beta
WT(modβ),只要求了
0
≤
T
<
β
0 \le T < \beta
0≤T<β,并不需要
T
<
p
T<p
T<p,因此只要维持
4
p
<
β
4p<\beta
4p<β 结果就是正确的。
GS 蝴蝶的实现:
CT 蝴蝶的实现:
另外,也可以使用 Montgomery 模乘(而非 Shoup’s Barrett 模乘)去实现蝴蝶,此时也可以继续采取冗余表示:
Preprocess-then-NTT
[ZXZ+19] 考虑了 Z q \mathbb Z_q Zq 不存在 ζ 2 n \zeta_{2n} ζ2n 的情况,并非采取 Incomplete NTT,而是先对多项式做一些预处理(其实就是 Nussbaumer 转换)
1-Round Preprocess-then-NTT(1PtNTT),给定
f
∈
Z
q
[
x
]
/
(
x
n
+
1
)
f \in \mathbb Z_q[x]/(x^n+1)
f∈Zq[x]/(xn+1),那么
ψ
:
Z
q
[
x
]
/
(
x
n
+
1
)
→
(
Z
q
[
y
]
/
(
y
n
/
2
+
1
)
)
[
x
]
/
(
x
2
−
y
)
f
e
v
e
n
(
x
2
)
+
x
⋅
f
o
d
d
(
x
2
)
↦
f
e
v
e
n
(
y
)
+
f
o
d
d
(
y
)
⋅
x
\begin{aligned} \psi: \mathbb Z_q[x]/(x^n+1) &\to (\mathbb Z_q[y]/(y^{n/2}+1))[x]/(x^2-y)\\ f_{even}(x^2)+x\cdot f_{odd}(x^2) &\mapsto f_{even}(y)+f_{odd}(y)\cdot x \end{aligned}
ψ:Zq[x]/(xn+1)feven(x2)+x⋅fodd(x2)→(Zq[y]/(yn/2+1))[x]/(x2−y)↦feven(y)+fodd(y)⋅x
此时,只需要
n
∣
q
−
1
n\mid q-1
n∣q−1(而非
2
n
∣
q
−
1
2n\mid q-1
2n∣q−1),那么两个系数
f
e
v
e
n
,
f
o
d
d
f_{even}, f_{odd}
feven,fodd 就可以完全 NTT,即
1
P
t
N
T
T
(
f
)
:
=
(
N
T
T
(
f
e
v
e
n
)
,
N
T
T
(
f
o
d
d
)
)
1PtNTT(f) := (NTT(f_{even}),\,\, NTT(f_{odd}))
1PtNTT(f):=(NTT(feven),NTT(fodd))
对于多项式乘法,就简单地采取 School 乘法即可。但是为了模
(
x
2
−
y
)
(x^2-y)
(x2−y) 方便,[ZXZ+19] 另外计算了
f
o
d
d
′
(
y
)
:
=
y
⋅
f
o
d
d
(
y
)
f_{odd}'(y):=y \cdot f_{odd}(y)
fodd′(y):=y⋅fodd(y) 以及它的 NTT 域,那么
h
e
v
e
n
=
f
e
v
e
n
⋅
g
e
v
e
n
+
f
o
d
d
⋅
g
o
d
d
′
h
o
d
d
=
f
e
v
e
n
⋅
g
o
d
d
+
f
o
d
d
⋅
g
e
v
e
n
\begin{aligned} h_{even} &= f_{even} \cdot g_{even} + f_{odd} \cdot g_{odd}'\\ h_{odd} &= f_{even} \cdot g_{odd} + f_{odd} \cdot g_{even} \end{aligned}
hevenhodd=feven⋅geven+fodd⋅godd′=feven⋅godd+fodd⋅geven
这一共需要计算
f
e
v
e
n
,
f
o
d
d
,
g
e
v
e
n
,
g
o
d
d
,
g
o
d
d
′
f_{even},f_{odd},g_{even},g_{odd},g_{odd}'
feven,fodd,geven,godd,godd′ 五个长度为
n
/
2
n/2
n/2 的 forward NTT,以及
h
e
v
e
n
,
h
o
d
d
h_{even},h_{odd}
heven,hodd 两个长度为
n
/
2
n/2
n/2 的 inverse NTT。计算复杂度为
7
n
/
2
log
n
+
2
n
7n/2\log n+2n
7n/2logn+2n
其实 y ∈ Z p [ y ] / ( y n / 2 + 1 ) y \in \mathbb Z_p[y]/(y^{n/2}+1) y∈Zp[y]/(yn/2+1) 的 NTT 域极其特殊,于是 g o d d ′ g_{odd}' godd′ 明明可以在 N T T ( g o d d ) NTT(g_{odd}) NTT(godd) 下直接计算出来,这个额外的 forward NTT 是不必要的。2-Round Preprocess-then-NTT(2PtNTT)的计算方法类似,就是采取了 x 4 = y x^4=y x4=y 的变换,此时只要求 n / 2 ∣ q − 1 n/2 \mid q-1 n/2∣q−1 即可。计算复杂度为 15 n / 4 log n + 4 n 15n/4\log n+4n 15n/4logn+4n
Improved PtNTT
[ZXZ+19] 实际上是采取了 “跨步” 转换。 [ZLP21] 采取 “聚合” 转换,它称之为 2-Part-Sepration,只需要
n
∣
q
−
1
n \mid q-1
n∣q−1(而非
2
n
∣
q
−
1
2n\mid q-1
2n∣q−1)
KaTeX parse error: Expected 'EOF', got '&' at position 31: …Z_q[x]/(x^n+1) &̲\to& \mathbb Z_…
采取 Karatsuba 算法,
f
↦
(
f
0
,
f
0
+
f
1
)
g
↦
(
g
0
,
g
0
+
g
1
)
u
:
=
f
1
g
1
h
=
f
0
g
0
⋅
(
1
−
y
)
+
(
f
0
+
f
1
)
(
g
0
+
g
1
)
⋅
x
+
u
⋅
(
y
2
−
y
)
=
(
f
0
g
0
−
u
)
+
(
(
f
0
+
f
1
)
(
g
0
+
g
1
)
−
f
0
g
0
−
u
)
⋅
y
\begin{aligned} f &\mapsto (f_0, f_0+f_1)\\ g &\mapsto (g_0, g_0+g_1)\\ u &:= f_1g_1\\ h &= f_0g_0 \cdot (1-y) + (f_0+f_1)(g_0+g_1) \cdot x + u \cdot (y^2-y)\\ &= (f_0g_0-u) + ((f_0+f_1)(g_0+g_1)-f_0g_0-u) \cdot y \end{aligned}
fguh↦(f0,f0+f1)↦(g0,g0+g1):=f1g1=f0g0⋅(1−y)+(f0+f1)(g0+g1)⋅x+u⋅(y2−y)=(f0g0−u)+((f0+f1)(g0+g1)−f0g0−u)⋅y
上述算法需要计算
f
0
,
f
1
,
g
0
,
g
1
f_0,f_1,g_0,g_1
f0,f1,g0,g1 四个长度为
n
/
2
n/2
n/2 的 forward NTT(应当是模
x
n
/
2
−
y
x^{n/2}-y
xn/2−y 的多项式,没法直接 NTT 啊!),以及
f
0
g
0
,
f
1
g
1
,
(
f
0
+
f
1
)
(
g
0
+
g
1
)
f_0g_0,f_1g_1,(f_0+f_1)(g_0+g_1)
f0g0,f1g1,(f0+f1)(g0+g1) 和
(
⋯
)
⋅
y
(\cdots)\cdot y
(⋯)⋅y 四个 point-wise mult,其中的
N
T
T
(
y
)
NTT(y)
NTT(y) 就只是常数而已。得到的
h
h
h 是长度
n
/
2
n/2
n/2 的向量(嗯?明显不正常啊),只需一次 inverse NTT 就可以恢复出
h
=
f
g
h=fg
h=fg
将它更加细分,
KaTeX parse error: Expected 'EOF', got '&' at position 31: …Z_q[x]/(x^n+1) &̲\to& \mathbb Z_…
此时的
f
f
f 被转换为
∑
i
f
i
(
x
)
⋅
y
i
\sum_i f_i(x) \cdot y^i
∑ifi(x)⋅yi,分成了
2
α
2^\alpha
2α 块。采取类似的乘法技巧,需要
2
α
+
1
2^{\alpha+1}
2α+1 次长度为
n
/
2
α
n/2^\alpha
n/2α 的 forward NTT,以及
2
2
α
+
2
α
+
1
−
4
2^{2\alpha}+2^{\alpha+1}-4
22α+2α+1−4 次的 point-wise mult,最终得到一个长度为
n
/
2
α
n/2^\alpha
n/2α 的结果(这是什么鬼!),执行一次 inverse NTT。[ZLP21] 说上述算法的复杂度为
5
n
log
n
+
O
(
n
)
5n\log n+O(n)
5nlogn+O(n),而原始 NTT 乘法的复杂度为
3
n
log
n
+
O
(
n
)
3n\log n+O(n)
3nlogn+O(n),因此减速因子是
5
/
3
5/3
5/3
[ZLP21] 另外还对 [ZXZ+19] 进行了优化,也就是不再计算 N T T ( g o d d ′ ) NTT(g_{odd}') NTT(godd′),而是使用 N T T ( y ) NTT(y) NTT(y) 计算乘积。多了一次 ponit-wise mult 的开销,但是减少了一次 forward NTT 运算。称其为:1-Round Improved-Preprocess-then-NTT(1IPtNTT),计算复杂度为 6 ⋅ n / 2 log ( n / 2 ) + 4 ⋅ n / 2 = 3 n log n − n 6\cdot n/2\log(n/2)+4\cdot n/2 = 3n\log n-n 6⋅n/2log(n/2)+4⋅n/2=3nlogn−n
另外,[ZLP21] 还将它扩展到更加细分,
α
\alpha
α-IPtNTT(其实就是 Nussbaumer 转换),
Z
q
[
x
]
/
(
x
n
+
1
)
≅
(
Z
q
[
y
]
/
(
y
n
/
2
α
+
1
)
)
[
x
]
/
(
x
2
α
−
y
)
\mathbb Z_q[x]/(x^n+1) \cong (\mathbb Z_q[y]/(y^{n/2^\alpha}+1))[x]/(x^{2^\alpha}-y)
Zq[x]/(xn+1)≅(Zq[y]/(yn/2α+1))[x]/(x2α−y)
然后只需
n
/
2
α
∣
q
−
1
n/2^\alpha \mid q-1
n/2α∣q−1,即可执行长度为
n
/
2
α
n/2^\alpha
n/2α 的完全 NTT,然后 [ZLP21] 采取 School 乘法,计算这个
(
m
o
d
x
2
α
−
y
)
\pmod{x^{2^\alpha}-y}
(modx2α−y) 的多项式乘法。计算复杂度为
3
n
log
n
+
(
3
⋅
2
α
−
2
−
3
α
+
1
/
2
)
⋅
n
3n\log n + (3 \cdot 2^{\alpha-2}-3\alpha+1/2)\cdot n
3nlogn+(3⋅2α−2−3α+1/2)⋅n,如果采取 Karatsuba 算法后面的线性项可以更小一些。
对于 α = 2 , 3 \alpha=2,3 α=2,3,达到最优的复杂度 3 n log n − 5 / 2 n 3n\log n-5/2n 3nlogn−5/2n,当 n = 1024 n=1024 n=1024 量级,甚至比原始的 NTT 算法的 3 n log n + O ( n ) 3n\log n + O(n) 3nlogn+O(n) 还要快不少(比率是 0.887 0.887 0.887)。换句话说,由于多项式的长度变短,蝴蝶层数减少,不完全的 NTT 乘法甚至可能会更快!
NTT-unfriendly rings
[CHK+21] 考虑了 PQC 中 NTT 不友好的 Saber、NTRU、LAC 方案的 NTT 加速实现。
- Saber 的代数结构 Z q [ x ] / ( x n + 1 ) \mathbb Z_q[x]/(x^n+1) Zq[x]/(xn+1),其中 q = 2 13 q=2^{13} q=213 不是素数,维度 n = 256 n=256 n=256
- NTRU 的代数结构有三个, Z 3 [ x ] / ( Φ n ( x ) ) \mathbb Z_3[x]/(\Phi_n(x)) Z3[x]/(Φn(x)), Z q [ x ] / ( Φ n ( x ) ) \mathbb Z_q[x]/(\Phi_n(x)) Zq[x]/(Φn(x)), Z q [ x ] / ( Φ 1 ( x ) ⋅ Φ n ( x ) ) \mathbb Z_q[x]/(\Phi_1(x)\cdot\Phi_n(x)) Zq[x]/(Φ1(x)⋅Φn(x)),其中的 n n n 是素数, q = 2 k q=2^k q=2k 不是素数
- LAC 的代数结构 Z q [ x ] / ( x n + 1 ) \mathbb Z_q[x]/(x^n+1) Zq[x]/(xn+1),其中 q = 251 q=251 q=251 是一种 min-split modulus,它使得 x n + 1 x^n+1 xn+1 仅能分解为两个长度 n / 2 n/2 n/2 的不可约因子
[CHK+21] 考虑的优化技术:Standard CT,Twisted GS,Negacyclic Convolutions, Incomplete NTTs, Good’s Trick,Mixed-Radix NTT,Multiple Moduli and Explicit CRT,
- 对 Saber 的优化:切换到很大的模数 q ′ q' q′(使得存在恰当的单位根),在 Z q ′ [ x ] / ( x n + 1 ) \mathbb Z_{q'}[x]/(x^n+1) Zq′[x]/(xn+1) 上执行不完全 NTT,最后计算 School 乘法。需要立即 InvNTT 并计算模约简,维持结果的正确性。
- 对 NTRU 的优化:切换到很大的维度 N N N(使得可以执行 NTT),切换到很大的模数 q ′ q' q′(使得存在恰当的单位根),在 Z q ′ [ x ] / ( x N + 1 ) \mathbb Z_{q'}[x]/(x^N+1) Zq′[x]/(xN+1) 上利用 Good 和 Mixed-radix 计算不完全 NTT,最后计算 School 乘法。需要立即 InvNTT 并计算模约简,维持结果的正确性。
- 对 LAC 的优化:切换到很大的模数 q ′ q' q′(使得存在恰当的单位根),在 Z q ′ [ x ] / ( x n + 1 ) \mathbb Z_{q'}[x]/(x^n+1) Zq′[x]/(xn+1) 上执行不完全 NTT,最后计算 School 乘法。需要立即 InvNTT 并计算模约简,维持结果的正确性。
采取 AVX2 实现上述的 NTT 乘法,考虑:快速模约简、层融合、延迟模约简、配置寄存器不相互依赖、不同 NTT 技巧的复杂度。
硬件优化
Sei18
[Sei18] 考虑了 Kyber 的 NTT 算法的 AVX2 实现。
首先是 Montgomery 模乘算法的修改:[Mon85] 采用了 q ′ = − q − 1 ( m o d β ) q'=-q^{-1}\pmod\beta q′=−q−1(modβ),计算无符号数的模乘,并保证输出结果是一个非负数。而 [Sei18] 采取了有符号数的变体,它最终的减法恰好消除了低位,没有进位,因此可以只计算高位。这就更加适合 AVX2,更密集的向量化。
其次是专用的模约简,对于 Kyber 采用的素数 q = 7681 q=7681 q=7681,它的二进制表示是稀疏的,
上述算法的输出范围是冗余的: − 2 15 + 4 q ≤ r < 2 15 − 3 q -2^{15}+4q \le r < 2^{15}-3q −215+4q≤r<215−3q,但是足够被用于加法/减法,将输入输出维持在单个 word 内。对于两个 words 的模约简,可以采用 Montgomery 模约简,常数 1 1 1 预计算为 β ( m o d q ) \beta \pmod q β(modq) 即可。
对于一般的素数 q q q,我们也希望只在单个 word 内完成模约简。采取 Barrett 算法:
它的输出范围是 0 ≤ r ≤ q 0 \le r \le q 0≤r≤q(对于 a ≡ 0 ( m o d q ) a\equiv 0\pmod q a≡0(modq) 会冗余)。另外,假如 step 1 采取了预计算 − v -v −v,并修改 step 4 成为 r = a + t r=a+t r=a+t,此时的输出范围是 − q ≤ r ≤ 0 -q \le r \le 0 −q≤r≤0。通过交错使用这两种 modes,可以维持模加的结果在 [ − q , q ] [-q,q] [−q,q] 范围内。
最后是 Lazy reduction:因为 Kyber 的模数满足 4 q < 2 15 = β / 2 4q<2^{15}=\beta/2 4q<215=β/2,因此加法结果可以累积起来,直到它溢出单个 word 之前,才执行一次 Barrett 模约简。在 NTT 中,我们采用了 Montgomery 模乘,它的结果范围是 − q < r ′ < q -q<r'<q −q<r′<q,因此每一层迭代,系数增长至多为 q q q,从而可以连续 3 3 3 层蝴蝶,累积但不溢出 β / 2 \beta/2 β/2,此时执行模约简依然可以得到正确结果。
Neon NTT
[BHK+21] 对比了 Montgomery 和 Barrett 的关系,提出了 Montgomery 模乘的类比:Barrett 模乘。不过,Shoup’s NTL 中其实已经采用了这种算法。
我们考虑四种 ”整数近似“ 函数:下取整 ⌊ z ⌋ \lfloor z \rfloor ⌊z⌋,上取整 ⌈ z ⌉ \lceil z \rceil ⌈z⌉,圆整 ⌊ z ⌉ \lfloor z \rceil ⌊z⌉,以及 “ 2 Z 2\mathbb Z 2Z-取值” ⌊ z ⌉ 2 : = 2 ⋅ ⌊ z / 2 ⌉ \lfloor z \rceil_2:= 2 \cdot \lfloor z/2 \rceil ⌊z⌉2:=2⋅⌊z/2⌉,这些函数可简记为 [ [ z ] ] [\![z]\!] [[z]],并且并不要求 [ [ z ] ] = z , ∀ z ∈ Z [\![z]\!]=z,\forall z\in \mathbb Z [[z]]=z,∀z∈Z
对于取模函数,可以采用上述的任意近似函数来定义,
z
(
m
o
d
[
[
⋅
]
]
N
)
:
=
z
−
N
⋅
[
[
z
N
]
]
z \pmod{^{[[\cdot]]} N} := z - N \cdot [\![\dfrac{z}{N}]\!]
z(mod[[⋅]]N):=z−N⋅[[Nz]]
- z ( m o d N ) z \pmod{N} z(modN),采用下取整的定义,范围 U N : = { 0 , 1 , ⋯ , N − 1 } U_N:=\{0,1,\cdots,N-1\} UN:={0,1,⋯,N−1},称为 canonical unsigned representative
- z ( m o d ± N ) z \pmod{^\pm N} z(mod±N),采用圆整的定义,范围 S N : = { − ⌊ N / 2 ⌋ , ⋯ , ⌊ ( N − 1 ) / 2 ⌋ } S_N:=\{-\lfloor N/2\rfloor ,\cdots,\lfloor (N-1)/2\rfloor\} SN:={−⌊N/2⌋,⋯,⌊(N−1)/2⌋},称为 canonical signed representative
- z ( m o d ⌊ ⋅ ⌉ 2 N ) z \pmod{^{\lfloor\cdot\rceil_2} N} z(mod⌊⋅⌉2N),采取 2 Z 2\mathbb Z 2Z-取值的定义,范围 { − N , ⋯ , N } \{-N,\cdots,N\} {−N,⋯,N},并且具有相同的奇偶性
我们首先给出 Barrett 和 Montgomery 的最基本描述:
根据这些整数近似函数的性质,可以计算出 Barrett 输出范围是 < 3 N / 2 <3N/2 <3N/2,假如继续约束 N < R / 3 N<R/3 N<R/3,那么输出结果 < R / 2 <R/2 <R/2,从而在 ( m o d R ) \pmod{R} (modR) 下的表示是唯一确定的。此时,就可以把 Barrett 的一些双精度运算简化为单精度运算,
对于 Montgomery,正如 [Sei18] 所说, m o n t + mont^+ mont+ 可以优化为单精度运算。但是 m o n t − mont^- mont− 出于进位的限制,无法这么优化。
两种 Montgomery 之间的关系:
Barrett 和 Montgomery 之间的关系:
类比着 Montgomery 模乘:
[BHK+21] 提出了 Barrett 模乘:
可以采取单精度指令的优化,只需要三条指令,
[BHK+21] 还继续考虑了 Armv8-A Neon vector instructions 提供的各种特殊指令,以优化 Barrett 和 Montgomery 的模约简、模乘的计算效率。
Mixed-radix NTT
[DL22] 考虑了 radix- 2 k 1 2^{k_1} 2k1 以及 radix- 2 k 2 2^{k_2} 2k2 的混合,给出了 FPGA 的实现。
对于一般的 radix-2 NTT 算法,在硬件上难以实现高吞吐量。因此他们将大的 NTT 拆解为若干小的 NTT,从而实现硬件的加速。
他们继续讨论了如何在 FPGA 上更好地实现这个算法。
TensorFHE
三种专用硬件:GPGPU(通用目的 GPU)、FPGA(可编程逻辑门阵列)、ASIC(专用集成电路)
[FWX+23] 给出了第一个使用 GPGPU 上的 TCU(Tensor Core Unit)加速的 FHE 实现。TCU 是最近的技术,它用于计算 4 × 4 4 \times 4 4×4 的矩阵乘法,速度比通用的 Cuda Core 并行度高得多。
- TCU 是一种专用于计算 multiply and accumulate (MAC) 运算的单元,不支持其他的运算
- TCU 仅支持低精度运算(至多 INT8 的整数,至多 FP16 的浮点数),返回值是 type-s32,但是只有低 16 比特是有效的
- TCU 是 warp 级别的,(二维)逐片运算;而 CUDA 是 thread 级别的,(一维)逐点运算
- 专用计算卡(比如 A100,好贵、已禁运)中包含远比 CUDA 多的 TCU,但是游戏卡中的 TCU 比 CUDA 少得多
[FWX+23] 的实现分为两层,
- API Layer:在 CPU 上运行,将用户的 FHE 操作自动分解为若干 basic kernel,并且确定批处理参数
- Kernel Layer:在 GPGPU 上运行,执行 7 7 7 种 basic kernel,包括 NTT(使用了 TCU 加速)、Hadamard Multiplication、Element-wise Addition/Subtract、Frobenius Map(槽置换)、Conjugate(共轭)、Basis Conversion(在不同 RNS 下切换)
他们测试了基本实现的各项性能,发现存在:Read-After-Write(RAW)数据依赖问题、线程间的资源竞争、GPU 占用率不高、TCU 不支持取模运算、TCU 仅支持低精度运算,一系列的问题。
因为 NTT 需要多层的蝴蝶迭代,每个蝴蝶都需要计算取模运算,不同层的蝴蝶有数据依赖关系,因此不适合在 TCU 上加速。[FWX+23] 简单地采取暴力计算,原本的 DFT 计算公式为:
∀
k
∈
[
N
]
,
A
k
=
(
∑
i
=
0
N
−
1
ζ
2
N
(
2
i
+
1
)
k
⋅
a
N
)
(
m
o
d
q
)
\forall k \in [N],\,\, A_k = \left( \sum_{i=0}^{N-1} \zeta_{2N}^{(2i+1)k} \cdot a_N \right) \pmod{q}\\
∀k∈[N],Ak=(i=0∑N−1ζ2N(2i+1)k⋅aN)(modq)
就是矩阵乘
A
N
=
W
N
×
a
N
A_N = W_{N} \times a_N
AN=WN×aN,我们采取 RNS 系统,使得环元素的系数是远小于
32
32
32 比特的整数,然后使用长度
64
64
64 比特的累加器,直到计算完矩阵乘法之后,才统一执行取模运算。
但是上述的矩阵
W
∈
Z
q
N
×
N
W \in \mathbb Z_q^{N \times N}
W∈ZqN×N 过大了,我们可以利用 Baby-Step Giant-Step(BSGS),将它转化为如下形式
A
N
1
×
N
2
=
(
(
W
1
×
a
N
1
×
N
2
)
⊙
W
2
)
×
W
3
T
(
m
o
d
q
)
\begin{aligned} A_{N_1 \times N_2} = \left( (W_1 \times a_{N_1 \times N_2}) \odot W_2 \right) \times W_3^T \pmod{q} \end{aligned}
AN1×N2=((W1×aN1×N2)⊙W2)×W3T(modq)
其中的
a
N
1
×
N
2
a_{N_1 \times N_2}
aN1×N2 是向量
a
N
,
N
=
N
1
N
2
a_{N}, N=N_1N_2
aN,N=N1N2 的二维化,三个矩阵分别为:
W
1
=
[
ζ
2
N
1
(
2
i
+
1
)
j
]
i
∈
[
N
1
]
,
j
∈
[
N
1
]
∈
Z
q
N
1
×
N
1
W
2
=
[
ζ
2
N
(
2
i
+
1
)
j
]
i
∈
[
N
1
]
,
j
∈
[
N
2
]
∈
Z
q
N
1
×
N
2
W
3
=
[
ζ
N
2
i
j
]
i
∈
[
N
2
]
,
j
∈
[
N
2
]
∈
Z
q
N
2
×
N
2
\begin{aligned} W_1 &= \begin{bmatrix} \zeta_{2N_1}^{(2i+1)j} \end{bmatrix}_{i \in [N_1],j \in [N_1]} \in \mathbb Z_q^{N_1 \times N_1}\\ W_2 &= \begin{bmatrix} \zeta_{2N}^{(2i+1)j} \end{bmatrix}_{i \in [N_1],j \in [N_2]} \in \mathbb Z_q^{N_1 \times N_2}\\ W_3 &= \begin{bmatrix} \zeta_{N_2}^{ij} \end{bmatrix}_{i \in [N_2],j \in [N_2]} \in \mathbb Z_q^{N_2 \times N_2} \end{aligned}
W1W2W3=[ζ2N1(2i+1)j]i∈[N1],j∈[N1]∈ZqN1×N1=[ζ2N(2i+1)j]i∈[N1],j∈[N2]∈ZqN1×N2=[ζN2ij]i∈[N2],j∈[N2]∈ZqN2×N2
这里的两个矩阵乘,以及一个阿达玛积,都是使用 GPGPU 暴力计算的。由于输入矩阵
a
N
1
×
N
2
,
W
1
,
W
3
a_{N_1 \times N_2}, W_1, W_3
aN1×N2,W1,W3 都是
32
32
32 比特的,因此可以将简单地按字节拆分为
4
4
4 个同样形状的矩阵,输入到 TCU 中计算 GEMM(利用
4
×
4
4 \times 4
4×4 基本矩阵乘搭建出来)。矩阵
W
2
W_2
W2 不必拆分,阿达玛积是直接在 CUDA 上运算的,这需要把 TCU 的计算结果融合为单个
32
32
32 比特矩阵。
此外,密文的 RNS 表示的索引是 ( L , N ) (L,N) (L,N),对于多个密文的同一个 L L L 的各种参数/资源是相同的。但是自然的索引 ( C , L , N ) (C,L,N) (C,L,N) 是以密文标号 C C C 主序的,导致了同一个 L L L 索引的多个密文数据是间断存储的。如果我们重排内存,按照 ( L , C , N ) (L,C,N) (L,C,N) 存储,那么就可以对于不同密文 C C C 的同一个 L L L 上的 NTT 域连续读取,它们共享了 NTT 参数,可以打包在一起计算,隐藏 I/O 延迟、资源复用。
最后 [FWX+23] 在 A100 上的实现,比 FPGA 表现的更好一些,但是依旧比 ASIC 慢得多。
Others
[HLS+22] 分别在 Intel AVX2 平台、ARM Cortex M4 平台,实现了 NTRU、Kyber、Saber 三种 KEM 方案,一共 6 6 6 个实现。他们使用汇编语言编写 NTT 算法,然后使用 CryptoLine 工具包(形式化语言,不依赖编程模型),半自动化地分析验证这些实现的正确性以及一些属性。
[ZLH+23] 优化了 High-radix NTT 的访存模式,提出了一种低复杂度的 cross-bank-write-back memory mapping scheme,通过时间延迟累积蝴蝶的结果,最后串行写回内存。最后,他们设计了 radix-4 NTT 的 FPGA 加速器。