数论变换（NTT）算法深度解析 —— 基于 Vitis-Tutorials 实现

1. 问题陈述

我们需要在有限域 \(\mathbb{F}_q\) 上高效计算长度为 \(N\) 的多项式向量的数论变换（NTT），其中：

• 模数 \(q = 3329\)（CRYSTALS-Kyber 后量子加密算法标准参数）
• 变换长度 \(N = 256 = 2^8\)
• 向量化维度 \(K = 128\)
• 支持原位变换和流水线化非原位变换两种实现方式

该问题的形式化定义为：给定多项式向量 \(\vec{f} = (f_0, f_1, \dots, f_{K-1})\)，其中每个 \(f_i = \sum_{j=0}^{N-1} f_{i,j} x^j \in \mathbb{F}_q[x]/(x^N + 1)\)，计算其 NTT 变换 \(\hat{\vec{f}} = (\hat{f}_0, \hat{f}_1, \dots, \hat{f}_{K-1})\)，满足：

2. 直觉

2.1 朴素方法的局限性

直接计算 NTT 需要 \(O(N^2)\) 次模乘运算，对于 \(N=256\) 而言虽然理论可行，但在硬件实现中存在两个致命问题：

1. 计算量过大：单多项式变换需要 \(256^2 = 65536\) 次模乘，无法满足后量子加密的实时性要求
2. 数据依赖复杂：原位变换引入的读写依赖会严重阻碍 FPGA 流水线的并行化

2.2 关键洞察

Vitis-Tutorials 中的实现基于以下核心洞察：

1. 分治策略：利用 Cooley-Tukey 算法将 NTT 分解为 \(\log_2 N = 8\) 个计算阶段（stage），每个阶段仅包含 \(O(N)\) 次操作
2. 旋转因子预计算：将所有单位根 \(\zeta^k\) 预存储到只读存储器（ROM）中，避免重复计算
3. 流水线化：将每个计算阶段封装为独立的硬件模块，通过乒乓缓冲区连接，实现任务级并行（TLP）
4. Montgomery 约减：使用预计算常数替代昂贵的模除法运算，将模乘延迟降低到单个时钟周期

3. 形式化定义

3.1 有限域约束

我们使用的有限域 \(\mathbb{F}_q\) 满足以下约束：

1. 素数模数 \(q = 3329 = 13 \times 2^8 + 1\)，因此 \(2^8 \mid q-1\)，保证存在 \(2^8\) 次本原单位根
2. 预计算常数 \(QINV = -q^{-1} \mod 2^{16} = 3327\)，用于 Montgomery 约减
3. 旋转因子数组 \(\text{zetas}[128]\) 满足 \(\text{zetas}[k] = \zeta^{\text{bitrev}_7(k)} \mod q\)，其中 \(\zeta = 17\) 是选定的本原单位根

3.2 Cooley-Tukey NTT 递推

对于长度为 \(N=2^m\) 的 NTT，我们使用迭代式 Cooley-Tukey 算法，其递推关系为：

3.3 Montgomery 约减

模乘运算 \(c = a \cdot b \mod q\) 使用 Montgomery 约减实现，其数学表达式为：

4. 算法

4.1 Version0：原位 NTT（基线实现）

伪代码

Algorithm: InPlaceNTT(p)
Input: Array p[0..N-1] in F_q
Output: NTT-transformed array p[0..N-1]
Constants: N=256, K=128, zetas[0..K-1]

k <- 1
for len from K down to 2 do
    len <- len >> 1
    start <- 0
    while start < 256 do
        zeta <- zetas[k]
        k <- k + 1
        for j from start to start + len - 1 do
            t <- fqmul(zeta, p[j + len])
            p[j + len] <- (p[j] - t) mod q
            p[j] <- (p[j] + t) mod q
        start <- j + len

执行流程图

实际实现代码

void ntt(uint16_t p[N]) {
    #pragma HLS INLINE OFF
    unsigned int len, start, j, k;
    int16_t t, zeta;

    k = 1;
    ntt_stage_outer:for(len = K; len >= 2; len >>= 1) {
        ntt_stage_loop:for(start = 0; start < 256; start = j + len) {
            zeta = zetas[k++];
            ntt_stage_inner:for(j = start; j < start + len; j++) {
                #pragma HLS PIPELINE
                t = fqmul(zeta, p[j + len]);
                p[j + len] = p[j] - t;
                p[j] = p[j] + t;
            }
        }
    }
}

4.2 Version3：流水线化 NTT（优化实现）

伪代码

Algorithm: PipelinedNTT(p_in, p_out)
Input: Array p_in[0..N-1] in F_q
Output: NTT-transformed array p_out[0..N-1]
Constants: N=256, K=128, zetas[0..K-1]
Intermediate Buffers: p_stage1[0..N-1], ..., p_stage6[0..N-1]

ntt_stage<0>(p_in, p_stage1)
ntt_stage<1>(p_stage1, p_stage2)
ntt_stage<2>(p_stage2, p_stage3)
ntt_stage<3>(p_stage3, p_stage4)
ntt_stage<4>(p_stage4, p_stage5)
ntt_stage<5>(p_stage5, p_stage6)
ntt_stage<6>(p_stage6, p_out)

Subroutine: ntt_stage<s>(p_in, p_out)
len <- K >> s
k <- 1 << s
start <- 0
while start < N do
    zeta <- zetas[k]
    k <- k + 1
    for j from start to start + len - 1 do
        t <- fqmul(zeta, p_in[j + len])
        p_out[j + len] <- (p_in[j] - t) mod q
        p_out[j] <- (p_in[j] + t) mod q
    start <- j + len

执行流程图

实际实现代码

template <int stage>
void ntt_stage(uint16_t p_in[N], uint16_t p_out[N]){
    unsigned int start=0, j=0;
    unsigned int len = K >> stage; 
    unsigned int k = 1 << stage; 

    ntt_stage_loop1: for(start = 0; start < N; start = j + len) {
        int16_t zeta = zetas[k]; k++;
        ntt_stage_loop2: for(j = start; j < start + len; j++) {
            #pragma HLS PIPELINE
            int16_t t = fqmul(zeta, p_in[j + len]);
            p_out[j + len] = p_in[j] - t;
            p_out[j] = p_in[j] + t;
        }
    }
}

void ntt(uint16_t p_in[N], uint16_t p_out[N]) { 
    uint16_t p_stage1[N];
    uint16_t p_stage2[N];
    uint16_t p_stage3[N];
    uint16_t p_stage4[N];
    uint16_t p_stage5[N];
    uint16_t p_stage6[N];
#pragma HLS DATAFLOW
    ntt_stage<0>(p_in,     p_stage1);
    ntt_stage<1>(p_stage1, p_stage2);
    ntt_stage<2>(p_stage2, p_stage3);
    ntt_stage<3>(p_stage3, p_stage4);
    ntt_stage<4>(p_stage4, p_stage5);
    ntt_stage<5>(p_stage5, p_stage6);
    ntt_stage<6>(p_stage6, p_out  );
}

5. 复杂度分析

5.1 时间复杂度

对于单个长度为 \(N=256\) 的多项式：

• 计算阶段数：\(m = \log_2 N = 8\)
• 每个阶段的蝴蝶运算数：\(N/2 = 128\)
• 每个蝴蝶运算的操作数：1 次 Montgomery 模乘，2 次模加/减

因此，总的时间复杂度为：

5.2 空间复杂度

• Version0（原位）：仅需输入数组 \(p[N]\) 和常数数组 \(\text{zetas}[K]\)，空间复杂度为 \(O(N)\)
• Version3（流水线）：需要输入数组、输出数组和 6 个中间缓冲区，总空间复杂度为 \(O(8N)\)，约 3KB 片上存储

5.3 硬件性能（来自 Vitis-Tutorials 实测数据）

6. 实现注意事项

6.1 与理论的差异

1. 简化模运算：实际实现省略了最终的模 \(q\) 约减，仅通过加减运算保证结果在 16 位整数范围内，依赖后续操作的隐式约减
2. 比特反转：理论上 NTT 需要前置或后置比特反转，但该实现通过调整旋转因子的访问顺序将其融合到计算阶段中
3. 常数展开：使用 C++ 模板参数将阶段索引 stage 转化为编译时常量，允许 HLS 工具为每个阶段生成专用硬件

6.2 工程权衡

1. 空间换时间：Version3 使用 6 个中间缓冲区（约 3KB BRAM）换取近 7 倍的吞吐量提升
2. 模板元编程：使用模板而非运行时参数会增加代码体积，但允许独立调度各个计算阶段
3. DATAFLOW 约束：所有中间缓冲区必须满足单生产者单消费者（SP-SC）约束，且函数必须无条件执行

6.3 Montgomery 约减实现

int16_t montgomery_reduce(int32_t a) {
#pragma HLS INLINE
    int16_t t1, t2;
    t1 = (int16_t)a*QINV;
    t2 = (a - (int32_t)t1*Q) >> 16;
    return t2;
}

注意这里的 (int16_t)a 是故意截断到低 16 位，符合 Montgomery 算法的要求。

7. 比较

7.1 与经典 Cooley-Tukey 算法的比较

7.2 与其他后量子 NTT 实现的比较

该实现专门针对 CRYSTALS-Kyber 算法优化，与通用 NTT 实现相比：

1. 固定参数：硬编码 \(q=3329\)、\(N=256\)，允许更激进的常量传播优化
2. 省略最终约减：利用 Kyber 算法的特性省略部分模运算
3. 向量化支持：顶层 polyvec_ntt 函数支持同时处理 128 个多项式，适合 Kyber 的密钥封装机制

7.3 Version0 与 Version3 的比较

Version0 是简洁的原位实现，适合资源受限的场景；Version3 是流水线化的非原位实现，通过 DATAFLOW 优化实现了任务级并行，在资源充足的 FPGA 上能提供更高的吞吐量。