NTT 的各类优化：Harvey、PtNTT，Intel AVX2、ARM Neon、GPGPU

参考文献：

1. [Har14] Harvey D. Faster arithmetic for number-theoretic transforms[J]. Journal of Symbolic Computation, 2014, 60: 113-119.
2. [Sei18] Seiler G. Faster AVX2 optimized NTT multiplication for Ring-LWE lattice cryptography[J]. Cryptology ePrint Archive, 2018.
3. [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.
4. [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.
5. [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.
6. [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.
7. [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.
8. [DL22] Duong-Ngoc P, Lee H. Configurable mixed-radix number theoretic transform architecture for lattice-based cryptography[J]. IEEE Access, 2022, 10: 12732-12741.
9. [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.
10. [FWX+23] Fan S, Wang Z, Xu W, et al. Tensorfhe: Achieving practical computation on encrypted data using gpgpu[C]//2023 IEEE International Symposium on High-Performance Computer Architecture (HPCA). IEEE, 2023: 922-934.
11. 快速乘法技巧：Karatsuba, Toom, Good, Schonhage, Strassen, Nussbaumer
12. Baby-Step Giant-Step & Homomorphic DFT
13. ARM NEON技术 - 基础介绍
14. Tensor Core 和 Cuda Core 的区别
15. Nvidia Tensor Core-WMMA API 编程入门

软件优化

Harvey Butterfly

在 Shoup’s NTL 中，radix-2 NTT 的蝴蝶实现如下：

它采用了 Barrett 算法的变体，Shoup’s modular multiplication：修改 $W^{\prime }\approx \beta /p$ 为 $W\approx W\beta /p$，于是 Barrett 取模算法就额外计算了与常数 $W$ 的数乘运算。但是这个蝴蝶的 if-else 语句过多，容易使得 CPU 分支预测失败并导致回滚。

[Har14] 提出使用 ${\mathbb{Z}}_p$ 的冗余表示（$[0,2p)$ 和 $[0,4p)$），从而移除了一些 if-else 语句。正确性要求：Shoup 模乘算法 $WT(\mathrm{mod}\beta )$，只要求了 $0\le T<\beta$，并不需要 $T<p$，因此只要维持 $4p<\beta$ 结果就是正确的。

GS 蝴蝶的实现：

CT 蝴蝶的实现：

另外，也可以使用 Montgomery 模乘（而非 Shoup’s Barrett 模乘）去实现蝴蝶，此时也可以继续采取冗余表示：

Preprocess-then-NTT

[ZXZ+19] 考虑了 ${\mathbb{Z}}_q$ 不存在 ${\zeta }_{2n}$ 的情况，并非采取 Incomplete NTT，而是先对多项式做一些预处理（其实就是 Nussbaumer 转换）

1-Round Preprocess-then-NTT（1PtNTT），给定 $f\in {\mathbb{Z}}_q[x]/(x^n+1)$，那么
$$\begin{array}{rl}\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{array}$$
此时，只需要 $n\mid q-1$（而非 $2n\mid q-1$），那么两个系数 $f_{even},f_{odd}$ 就可以完全 NTT，即
$$1PtNTT(f):=(NTT(f_{even}),NTT(f_{odd}))$$
对于多项式乘法，就简单地采取 School 乘法即可。但是为了模 $(x^2-y)$ 方便，[ZXZ+19] 另外计算了 $f_{odd}^{\prime }(y):=y\cdot f_{odd}(y)$ 以及它的 NTT 域，那么
$$\begin{array}{rl}{h}_{even}& =f_{even}\cdot g_{even}+f_{odd}\cdot g_{odd}^{\prime }\\ h_{odd}& =f_{even}\cdot g_{odd}+f_{odd}\cdot g_{even}\end{array}$$
这一共需要计算 $f_{even},f_{odd},g_{even},g_{odd},g_{odd}^{\prime }$ 五个长度为 $n/2$ 的 forward NTT，以及 $h_{even},h_{odd}$ 两个长度为 $n/2$ 的 inverse NTT。计算复杂度为 $7n/2\log n+2n$

其实 $y\in {\mathbb{Z}}_p[y]/(y^{n/2}+1)$ 的 NTT 域极其特殊，于是 $g_{odd}^{\prime }$ 明明可以在 $NTT(g_{odd})$ 下直接计算出来，这个额外的 forward NTT 是不必要的。2-Round Preprocess-then-NTT（2PtNTT）的计算方法类似，就是采取了 $x^4=y$ 的变换，此时只要求 $n/2\mid q-1$ 即可。计算复杂度为 $15n/4\log n+4n$

Improved PtNTT

[ZXZ+19] 实际上是采取了 “跨步” 转换。 [ZLP21] 采取 “聚合” 转换，它称之为 2-Part-Sepration，只需要 $n\mid q-1$（而非 $2n\mid q-1$）
$$\begin{array}{crcl}\psi :& {\mathbb{Z}}_q[x]/(x^n+1)& \to & {\mathbb{Z}}_q[x][y]/(y^2+1,y-x^{n/2})\\ & f_0(x)+x^{n/2}\cdot f_1(x)& \mapsto & f_0(x)+f_1(x)\cdot y\end{array}$$

采取 Karatsuba 算法，
$$\begin{array}{rl}f& \mapsto (f_0,f_0+f_1)\\ g& \mapsto (g_0,g_0+g_1)\\ u& :=f_1{g}_1\\ h& =f_0{g}_0\cdot (1-y)+(f_0+f_1)(g_0+g_1)\cdot x+u\cdot (y^2-y)\\ & =(f_0{g}_0-u)+((f_0+f_1)(g_0+g_1)-f_0{g}_0-u)\cdot y\end{array}$$
上述算法需要计算 $f_0,f_1,g_0,g_1$ 四个长度为 $n/2$ 的 forward NTT（应当是模 $x^{n/2}-y$ 的多项式，没法直接 NTT 啊！），以及 $f_0{g}_0,f_1{g}_1,(f_0+f_1)(g_0+g_1)$ 和 $(\cdots )\cdot y$ 四个 point-wise mult，其中的 $NTT(y)$ 就只是常数而已。得到的 $h$ 是长度 $n/2$ 的向量（嗯？明显不正常啊），只需一次 inverse NTT 就可以恢复出 $h=fg$

将它更加细分，
$$\psi :{\mathbb{Z}}_q[x]/(x^n+1)\to {\mathbb{Z}}_q[x][y]/(y^{2^{\alpha }}+1,y-x^{n/2^{\alpha }})$$
此时的 $f$ 被转换为 ${\sum }_i{f}_i(x)\cdot y^i$，分成了 $2^{\alpha }$ 块。采取类似的乘法技巧，需要 $2^{\alpha +1}$ 次长度为 $n/2^{\alpha }$ 的 forward NTT，以及 $2^{2\alpha }+2^{\alpha +1}-4$ 次的 point-wise mult，最终得到一个长度为 $n/2^{\alpha }$ 的结果（这是什么鬼！），执行一次 inverse NTT。[ZLP21] 说上述算法的复杂度为 $5n\log n+O(n)$，而原始 NTT 乘法的复杂度为 $3n\log n+O(n)$，因此减速因子是 $5/3$

[ZLP21] 另外还对 [ZXZ+19] 进行了优化，也就是不再计算 $NTT(g_{odd}^{\prime })$，而是使用 $NTT(y)$ 计算乘积。多了一次 ponit-wise mult 的开销，但是减少了一次 forward NTT 运算。称其为：1-Round Improved-Preprocess-then-NTT（1IPtNTT），计算复杂度为 $6\cdot n/2\log (n/2)+4\cdot n/2=3n\log n-n$

另外，[ZLP21] 还将它扩展到更加细分，$\alpha$-IPtNTT（其实就是 Nussbaumer 转换），
$${\mathbb{Z}}_q[x]/(x^n+1)\cong ({\mathbb{Z}}_q[y]/(y^{n/2^{\alpha }}+1))[x]/(x^{2^{\alpha }}-y)$$
然后只需 $n/2^{\alpha }\mid q-1$，即可执行长度为 $n/2^{\alpha }$ 的完全 NTT，然后 [ZLP21] 采取 School 乘法，计算这个 $(\mathrm{mod}{x}^{2^{\alpha }}-y)$ 的多项式乘法。计算复杂度为 $3n\log n+(3\cdot 2^{\alpha -2}-3\alpha +1/2)\cdot n$，如果采取 Karatsuba 算法后面的线性项可以更小一些。

对于 $\alpha =2,3$，达到最优的复杂度 $3n\log n-5/2n$，当 $n=1024$ 量级，甚至比原始的 NTT 算法的 $3n\log n+O(n)$ 还要快不少（比率是 $0.887$）。换句话说，由于多项式的长度变短，蝴蝶层数减少，不完全的 NTT 乘法甚至可能会更快！

NTT-unfriendly rings

[CHK+21] 考虑了 PQC 中 NTT 不友好的 Saber、NTRU、LAC 方案的 NTT 加速实现。

• Saber 的代数结构 ${\mathbb{Z}}_q[x]/(x^n+1)$，其中 $q=2^{13}$ 不是素数，维度 $n=256$
• NTRU 的代数结构有三个，${\mathbb{Z}}_3[x]/({\Phi }_n(x))$，${\mathbb{Z}}_q[x]/({\Phi }_n(x))$，${\mathbb{Z}}_q[x]/({\Phi }_1(x)\cdot {\Phi }_n(x))$，其中的 $n$ 是素数，$q=2^k$ 不是素数
• LAC 的代数结构 ${\mathbb{Z}}_q[x]/(x^n+1)$，其中 $q=251$ 是一种 min-split modulus，它使得 $x^n+1$ 仅能分解为两个长度 $n/2$ 的不可约因子

[CHK+21] 考虑的优化技术：Standard CT，Twisted GS，Negacyclic Convolutions， Incomplete NTTs， Good’s Trick，Mixed-Radix NTT，Multiple Moduli and Explicit CRT，

1. 对 Saber 的优化：切换到很大的模数 $q^{\prime }$（使得存在恰当的单位根），在 ${\mathbb{Z}}_{q^{\prime }}[x]/(x^n+1)$ 上执行不完全 NTT，最后计算 School 乘法。需要立即 InvNTT 并计算模约简，维持结果的正确性。
2. 对 NTRU 的优化：切换到很大的维度 $N$（使得可以执行 NTT），切换到很大的模数 $q^{\prime }$（使得存在恰当的单位根），在 ${\mathbb{Z}}_{q^{\prime }}[x]/(x^N+1)$ 上利用 Good 和 Mixed-radix 计算不完全 NTT，最后计算 School 乘法。需要立即 InvNTT 并计算模约简，维持结果的正确性。
3. 对 LAC 的优化：切换到很大的模数 $q^{\prime }$（使得存在恰当的单位根），在 ${\mathbb{Z}}_{q^{\prime }}[x]/(x^n+1)$ 上执行不完全 NTT，最后计算 School 乘法。需要立即 InvNTT 并计算模约简，维持结果的正确性。

采取 AVX2 实现上述的 NTT 乘法，考虑：快速模约简、层融合、延迟模约简、配置寄存器不相互依赖、不同 NTT 技巧的复杂度。

硬件优化

Sei18

[Sei18] 考虑了 Kyber 的 NTT 算法的 AVX2 实现。

首先是 Montgomery 模乘算法的修改：[Mon85] 采用了 $q^{\prime }=-q^{-1}(\mathrm{mod}\beta )$，计算无符号数的模乘，并保证输出结果是一个非负数。而 [Sei18] 采取了有符号数的变体，它最终的减法恰好消除了低位，没有进位，因此可以只计算高位。这就更加适合 AVX2，更密集的向量化。

其次是专用的模约简，对于 Kyber 采用的素数 $q=7681$，它的二进制表示是稀疏的，

上述算法的输出范围是冗余的：$-2^{15}+4q\le r<2^{15}-3q$，但是足够被用于加法/减法，将输入输出维持在单个 word 内。对于两个 words 的模约简，可以采用 Montgomery 模约简，常数 $1$ 预计算为 $\beta (\mathrm{mod}q)$ 即可。

对于一般的素数 $q$，我们也希望只在单个 word 内完成模约简。采取 Barrett 算法：

它的输出范围是 $0\le r\le q$（对于 $a\equiv 0(\mathrm{mod}q)$ 会冗余）。另外，假如 step 1 采取了预计算 $-v$，并修改 step 4 成为 $r=a+t$，此时的输出范围是 $-q\le r\le 0$。通过交错使用这两种 modes，可以维持模加的结果在 $[-q,q]$ 范围内。

最后是 Lazy reduction：因为 Kyber 的模数满足 $4q<2^{15}=\beta /2$，因此加法结果可以累积起来，直到它溢出单个 word 之前，才执行一次 Barrett 模约简。在 NTT 中，我们采用了 Montgomery 模乘，它的结果范围是 $-q<r^{\prime }<q$，因此每一层迭代，系数增长至多为 $q$，从而可以连续 $3$ 层蝴蝶，累积但不溢出 $\beta /2$，此时执行模约简依然可以得到正确结果。

Neon NTT

[BHK+21] 对比了 Montgomery 和 Barrett 的关系，提出了 Montgomery 模乘的类比：Barrett 模乘。不过，Shoup’s NTL 中其实已经采用了这种算法。

我们考虑四种 ”整数近似“ 函数：下取整 $\lfloor z\rfloor$，上取整 $\lceil z\rceil$，圆整 $\lfloor z\rceil$，以及 “$2\mathbb{Z}$-取值” $\lfloor z{\rceil }_2:=2\cdot \lfloor z/2\rceil$，这些函数可简记为 $[[z]]$，并且并不要求 $[[z]]=z,\forall z\in \mathbb{Z}$

对于取模函数，可以采用上述的任意近似函数来定义，
$$z(\mathrm{mod}{}^{[[\cdot ]]}N):=z-N\cdot [[\frac{z}{N}]]$$

• $z(\mathrm{mod}N)$，采用下取整的定义，范围 U N : = { 0 , 1 , ⋯   , N − 1 } U_N:=\{0,1,\cdots,N-1\} UN​:={0,1,⋯,N−1}，称为 canonical unsigned representative
• $z(\mathrm{mod}{}^{\pm }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(\mathrm{mod}{}^{\lfloor \cdot {\rceil }_2}N)$，采取 $2\mathbb{Z}$-取值的定义，范围 { − N , ⋯   , N } \{-N,\cdots,N\} {−N,⋯,N}，并且具有相同的奇偶性

我们首先给出 Barrett 和 Montgomery 的最基本描述：

根据这些整数近似函数的性质，可以计算出 Barrett 输出范围是 $<3N/2$，假如继续约束 $N<R/3$，那么输出结果 $<R/2$，从而在 $(\mathrm{mod}R)$ 下的表示是唯一确定的。此时，就可以把 Barrett 的一些双精度运算简化为单精度运算，

对于 Montgomery，正如 [Sei18] 所说，$mon{t}^{+}$ 可以优化为单精度运算。但是 $mon{t}^{-}$ 出于进位的限制，无法这么优化。

两种 Montgomery 之间的关系：

Barrett 和 Montgomery 之间的关系：

类比着 Montgomery 模乘：

[BHK+21] 提出了 Barrett 模乘：

可以采取单精度指令的优化，只需要三条指令，

[BHK+21] 还继续考虑了 Armv8-A Neon vector instructions 提供的各种特殊指令，以优化 Barrett 和 Montgomery 的模约简、模乘的计算效率。

Mixed-radix NTT

[DL22] 考虑了 radix-$2^{k_1}$ 以及 radix-$2^{k_2}$ 的混合，给出了 FPGA 的实现。

对于一般的 radix-2 NTT 算法，在硬件上难以实现高吞吐量。因此他们将大的 NTT 拆解为若干小的 NTT，从而实现硬件的加速。

他们继续讨论了如何在 FPGA 上更好地实现这个算法。

TensorFHE

三种专用硬件：GPGPU（通用目的 GPU）、FPGA（可编程逻辑门阵列）、ASIC（专用集成电路）

[FWX+23] 给出了第一个使用 GPGPU 上的 TCU（Tensor Core Unit）加速的 FHE 实现。TCU 是最近的技术，它用于计算 $4\times 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] 的实现分为两层，

1. API Layer：在 CPU 上运行，将用户的 FHE 操作自动分解为若干 basic kernel，并且确定批处理参数
2. Kernel Layer：在 GPGPU 上运行，执行 $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 计算公式为：
$$\forall k\in [N],A_k=\left(\sum _{i=0}^{N-1}{\zeta }_{2N}^{(2i+1)k}\cdot a_N\right)(\mathrm{mod}q)$$
就是矩阵乘 $A_N=W_N\times a_N$，我们采取 RNS 系统，使得环元素的系数是远小于 $32$ 比特的整数，然后使用长度 $64$ 比特的累加器，直到计算完矩阵乘法之后，才统一执行取模运算。

但是上述的矩阵 $W\in {\mathbb{Z}}_q^{N\times N}$ 过大了，我们可以利用 Baby-Step Giant-Step（BSGS），将它转化为如下形式
$$\begin{array}{r}{A}_{N_1\times N_2}=\left((W_1\times a_{N_1\times N_2})\odot W_2\right)\times W_3^T(\mathrm{mod}q)\end{array}$$
其中的 $a_{N_1\times N_2}$ 是向量 $a_N,N=N_1{N}_2$ 的二维化，三个矩阵分别为：
$$\begin{array}{rl}{W}_1& ={\left[\begin{array}{c}{\zeta }_{2N_1}^{(2i+1)j}\end{array}\right]}_{i\in [N_1],j\in [N_1]}\in {\mathbb{Z}}_q^{N_1\times N_1}\\ W_2& ={\left[\begin{array}{c}{\zeta }_{2N}^{(2i+1)j}\end{array}\right]}_{i\in [N_1],j\in [N_2]}\in {\mathbb{Z}}_q^{N_1\times N_2}\\ W_3& ={\left[\begin{array}{c}{\zeta }_{N_2}^{ij}\end{array}\right]}_{i\in [N_2],j\in [N_2]}\in {\mathbb{Z}}_q^{N_2\times N_2}\end{array}$$
这里的两个矩阵乘，以及一个阿达玛积，都是使用 GPGPU 暴力计算的。由于输入矩阵 $a_{N_1\times N_2},W_1,W_3$ 都是 $32$ 比特的，因此可以将简单地按字节拆分为 $4$ 个同样形状的矩阵，输入到 TCU 中计算 GEMM（利用 $4\times 4$ 基本矩阵乘搭建出来）。矩阵 $W_2$ 不必拆分，阿达玛积是直接在 CUDA 上运算的，这需要把 TCU 的计算结果融合为单个 $32$ 比特矩阵。

此外，密文的 RNS 表示的索引是 $(L,N)$，对于多个密文的同一个 $L$ 的各种参数/资源是相同的。但是自然的索引 $(C,L,N)$ 是以密文标号 $C$ 主序的，导致了同一个 $L$ 索引的多个密文数据是间断存储的。如果我们重排内存，按照 $(L,C,N)$ 存储，那么就可以对于不同密文 $C$ 的同一个 $L$ 上的 NTT 域连续读取，它们共享了 NTT 参数，可以打包在一起计算，隐藏 I/O 延迟、资源复用。

最后 [FWX+23] 在 A100 上的实现，比 FPGA 表现的更好一些，但是依旧比 ASIC 慢得多。

Others

[HLS+22] 分别在 Intel AVX2 平台、ARM Cortex M4 平台，实现了 NTRU、Kyber、Saber 三种 KEM 方案，一共 $6$ 个实现。他们使用汇编语言编写 NTT 算法，然后使用 CryptoLine 工具包（形式化语言，不依赖编程模型），半自动化地分析验证这些实现的正确性以及一些属性。

[ZLH+23] 优化了 High-radix NTT 的访存模式，提出了一种低复杂度的 cross-bank-write-back memory mapping scheme，通过时间延迟累积蝴蝶的结果，最后串行写回内存。最后，他们设计了 radix-4 NTT 的 FPGA 加速器。