Halo中的快速傅里叶（逆）变换算法(I)FFT

1. FFT背景知识

FFT背景知识可参看博客十分简明易懂的FFT（快速傅里叶变换）。

2. Halo中的FFT代码实现

在4核8G ubuntu16.04服务器上运行：

cargo test test_fft -- --nocapture

test_fft函数中实现的是对两个999阶（1000个系数）多项式的乘法运算，在该函数内，分别进行了直接乘法运算naive_product和通过FFT实现的乘法运算multiply_polynomials。

2.1 a和b系数列表均扩展为$2^{exp}$

multiply_polynomials函数中会首先将两个多项式相乘后的系数总数扩展为$2^{exp}$，将a和b系数列表补零扩展为$2^{exp}$：

 	let degree_of_result = (a.len() - 1) + (b.len() - 1); //1998
    let coeffs_of_result = degree_of_result + 1; //1999

    // Compute the size of our evaluation domain
    let mut m = 1; //2048
    let mut exp = 0; //11
    while m < coeffs_of_result {
        m *= 2;
        exp += 1;

        // The pairing-friendly curve may not be able to support
        // large enough (radix2) evaluation domains.
        if exp >= F::S {
            panic!("polynomial too large");
        }
    }
	//将`a`和`b`系数列表补零扩展为$2^{exp}$
	// Extend the vectors with zeroes
    a.resize(m, F::zero());
    b.resize(m, F::zero());

2.2 获取$2^{exp}$-th primitive root of unity

F::ALPHA为$2^{32}$-th primitive root of unity，基于该值获取相应的$2^{exp}$-th primitive root of unity：

// Compute alpha, the 2^exp primitive root of unity
    let mut alpha = F::ALPHA;
    for _ in exp..F::S {
        alpha = alpha.square();
    }
   //alpha为$2^{exp}$-th primitive root of unity

2.3 分别对a和b系数列表做FFT

	////alpha为$2^{exp}$-th primitive root of unity，exp=11
	best_fft(&mut a, alpha, exp); 
    best_fft(&mut b, alpha, exp);

注意best_fft(&mut a, alpha, exp);返回的数组$a$中（对多项式$A(x)=a_0+a_{1}x+a_2{x}^2+...+a_n{x}^n$）依次$x$取$w_n^0,w_n^1,w_n^2,...,w_n^{(n-1)}$的值$a=[A(w_n^0),A(w_n^1),A(w_n^2),...,A(w_n^{(n-1)})]$。
也就是说，通过best_fft函数，可将系数表示的多项式转换为点值表示：$(w_n^0,A(w_n^0)),....,(w_n^{(n-1)},A(w_n^{(n-1)}))$。

best_fft中会针对exp与cpu核数的关系来决定调用串行方式serial_fft还是并行方式parallel_fft。

fn best_fft(a: &mut [F], omega: F, log_n: u32) {
    let cpus = num_cpus::get(); //4
    let log_cpus = log2_floor(cpus); //2

    if log_n <= log_cpus {
        serial_fft(a, omega, log_n);
    } else { //11>2
        parallel_fft(a, omega, log_n, log_cpus);
    }
}

2.3.1 并行FFT算法parallel_fft

// omega为$2^{exp}$-th primitive root of unity，exp=11, log_n=11, log_cpus=2
fn parallel_fft(a: &mut [F], omega: F, log_n: u32, log_cpus: u32) {
    assert!(log_n >= log_cpus);

    let num_cpus = 1 << log_cpus; //4
    let log_new_n = log_n - log_cpus; //11-2=9
    let mut tmp = vec![vec![F::zero(); 1 << log_new_n]; num_cpus]; //  [2^2行 * 2^9列] 矩阵。行数2^2=4为cpu核数。
    let new_omega = omega.pow(&[num_cpus as u64, 0, 0, 0]); // new_omega=power_mod(omega, num_cpus,p)，对应的new_omega为2^{exp-log_n}即2^9-th primitive root of unity。

    thread::scope(|scope| {
        let a = &*a;
		// 1）tmp为 [2^2行 * 2^9列] 矩阵，逐行操作。
        for (j, tmp) in tmp.iter_mut().enumerate() {
            scope.spawn(move |_| { //采用多线程方式，所有行同时处理。
                // Shuffle into a sub-FFT
                let omega_j = omega.pow(&[j as u64, 0, 0, 0]); // power_mod(omega,j,p)
                let omega_step = omega.pow(&[(j as u64) << log_new_n, 0, 0, 0]); //power_mod(omega, j*2^9,p)

                let mut elt = F::one(); //1
                for i in 0..(1 << log_new_n) { // 2）再逐列处理
                    for s in 0..num_cpus { // 3）再逐CPU处理
                        let idx = (i + (s << log_new_n)) % (1 << log_n);
                        //$idx$用于取系数列表`a`中的相应的值。
                        //s=0时，idx的取值范围为0~511；s=1时，idx:512~1023；s=2时，idx:1024~1535；s=3时，idx:1536~2047。
                        let mut t = a[idx]; 
                        t *= elt;
                        tmp[i] += t;
                        elt *= omega_step;
                    }
                    elt *= omega_j;
                }

                // Perform sub-FFT
                serial_fft(tmp, new_omega, log_new_n);
            });
        }
    })
    .unwrap();

    // Unshuffle
    let mask = (1 << log_cpus) - 1;
    for (idx, a) in a.iter_mut().enumerate() {
        *a = tmp[idx & mask][idx >> log_cpus];
    }
}

FFT算法的本质是将以系数表示的多项式转化为以点值表示。对于方程式：
$y=a_0+a_{1}x+a_2{x}^2+....+a_{n-1}{x}^{n-1}$
其中，$n=2^{exp}$，上例中$exp=11,n=2048$，omega【表示为$w_n^1$】为n-th root of unity（即满足$(w_n^1{)}^n=1$）。
转为$n$个互不相同的点值序列$(x_0,y_0),(x_1,y_1),(x_2,y_2),......,(x_{n-1},y_{n-1})$，其中$x_k=(w_n^1{)}^k=w_n^k$。上例中parallel_fft函数中tmp数组内，存储的即为所有的$y$值，即$tmp[k]=y_k$。

以4核（$p=4$）CPU为例，支持将2047阶（$n=2048，n/p=512$）多项式拆分为四个线程分别执行：
$$\begin{gathered} A(x)=a_0+a_{1}x+a_2{x}^2+...+a_{511}{x}^{511} \\ +x^{512}(a_{512}+a_{513}x+a_{514}{x}^2+...+a_{1023}{x}^{511}) \\ +x^{1024}(a_{1024}+a_{1025}x+a_{1026}{x}^2+...+a_{1535}{x}^{511}) \\ +x^{1536}(a_{1536}+a_{1537}x+a_{1538}{x}^2+...+a_{2047}{x}^{511}) \end{gathered}$$

逐列展开：
$A(x)=C_0(x^{512})+x{C}_1(x^{512})+x^2{C}_2(x^{512})+...+x^{511}{C}_{511}(x^{512})$
其中：
$C_0(x)=a_0+a_{512}x+a_{1024}{x}^2+a_{1536}{x}^3$
$C_1(x)=a_1+a_{513}x+a_{1025}{x}^2+a_{1537}{x}^3$
$.......$
$C_{511}(x)=a_{511}+a_{1023}x+a_{1535}{x}^2+a_{2047}{x}^3$

再设$k<\frac{n}{p}=\frac{2048}{4}=512$，把$w_n^k=w_{2048}^k$作为$x$值代入$A(x)$多项式，有：
$$\begin{gathered} A(w_{2048}^k)=C_0(w_{2048}^{512k})+w_{2048}^k{C}_1(w_{2048}^{512k})+...+w_{2048}^{511k}{C}_{511}(w_{2048}^{512k}) \\ =C_0(w_4^k)+w_{2048}^k{C}_1(w_4^k)+...+w_{2048}^{511k}{C}_{511}(w_4^k) \end{gathered}$$
因此，根据单位根的性质，只需计算分别取$k=0,1,2,3$时相应的$C_0(w_4^k),C_1(w_4^k),...,C_{511}(w_4^k)$值，即可很方便的计算取任意$k<512$时的$A(w_n^k)$的值。

在Halo代码中，$w_n^1$对应为omega，$w_4^k$对应为omega_step（let omega_step = omega.pow(&[(j as u64) << log_new_n, 0, 0, 0]); //power_mod(omega, j*2^9,p)），$w_n^k$对应为omega_j。

Halo代码中，tmp为$2^2\times 2^9$矩阵：
$tmp[0]=[C_0(w_4^0),w_{2048}^{(0\ast 1)}{C}_1(w_4^0),...,w_{2048}^{(0\ast 511)}{C}_{511}(w_4^0)]$
$tmp[1]=[C_0(w_4^1),w_{2048}^{(1\ast 1)}{C}_1(w_4^1),...,w_{2048}^{(1\ast 511)}{C}_{511}(w_4^1)]$
$tmp[2]=[C_0(w_4^2),w_{2048}^{(2\ast 1)}{C}_1(w_4^2),...,w_{2048}^{(2\ast 511)}{C}_{511}(w_4^2)]$
$tmp[3]=[C_0(w_4^3),w_{2048}^{(3\ast 1)}{C}_1(w_4^3),...,w_{2048}^{(3\ast 511)}{C}_{511}(w_4^3)]$

2.3.2 串行FFT算法serial_fft

	// Perform sub-FFT
 	serial_fft(tmp, new_omega, log_new_n);
 	//此处的tmp为tmp[j]第j行内容，new_omega为2^9-th primitive root of unity，log_new_n值为9。

以下代码段的作用是将多项式系数数组$[a_0,a_1,a_2,...,a_n]$按奇偶重新排列，数组的前半段为偶数系数，后半段为奇数系数，具体为：$[a_0,a_2,a_4,...,a_{n-2},a_1,a_3,a_5,...,a_{n-1}]$。

 	for k in 0..n {
        let rk = bitreverse(k, log_n);
        if k < rk {
            a.swap(rk as usize, k as usize);
        }
    }

以$A(x)=a_0+a_{1}x+a_2{x}^2+a_3{x}^3$为例，下面程序的演示效果如上图所示。

fn serial_fft(a: &mut [F], omega: F, log_n: u32) {
    fn bitreverse(mut n: u32, l: u32) -> u32 {
        let mut r = 0;
        for _ in 0..l {
            r = (r << 1) | (n & 1);
            n >>= 1;
        }
        r
    }

    let n = a.len() as u32;
    assert_eq!(n, 1 << log_n);

    for k in 0..n {
        let rk = bitreverse(k, log_n);
        if k < rk {
            a.swap(rk as usize, k as usize);
        }
    }
    

    let mut m = 1;
    for _ in 0..log_n {
        let w_m = omega.pow(&[u64::from(n / (2 * m)), 0, 0, 0]);

        let mut k = 0;
        while k < n {
            let mut w = F::one();
            for j in 0..m {
                let mut t = a[(k + j + m) as usize];
                t *= w;
                a[(k + j + m) as usize] = a[(k + j) as usize] - t;
                a[(k + j) as usize] += t;
                w *= w_m;
            }

            k += 2 * m;
        }

        m *= 2;
    }
}

2.4 点值表示的多项式乘法运算

  	best_fft(&mut a, alpha, exp); //由系数表示转换为点值表示
    best_fft(&mut b, alpha, exp);

    // Multiply pairwise。点值表示的多项式乘法运算
    let num_cpus = num_cpus::get();
    if a.len() > num_cpus {
        thread::scope(|scope| {
            let chunk = a.len() / num_cpus::get();

            for (a, b) in a.chunks_mut(chunk).zip(b.chunks(chunk)) {
                scope.spawn(move |_| {
                    for (a, b) in a.iter_mut().zip(b.iter()) {
                        *a *= *b;
                    }
                });
            }
        })
        .unwrap();
    } else {
        for (a, b) in a.iter_mut().zip(b.iter()) {
            *a *= *b;
        }
    }

2.5 傅里叶逆变换IFFT

IFFT的作用是将点值表示转换为系数表示。

	// Inverse FFT
    let alpha_inv = alpha.invert().unwrap();
    best_fft(&mut a, alpha_inv, exp);

    // Divide all elements by m = a.len()
    let minv = F::from_u64(m as u64).invert().unwrap();
    if a.len() > num_cpus {
        thread::scope(|scope| {
            let chunk = a.len() / num_cpus::get();

            for a in a.chunks_mut(chunk) {
                scope.spawn(move |_| {
                    for a in a.iter_mut() {
                        *a *= minv;
                    }
                });
            }
        })
        .unwrap();
    } else {
        for a in a.iter_mut() {
            *a *= minv;
        }
    }
	//为了递归调用，a数组的长度做了扩展补零。只截取乘积后相应的阶即可。	
    a.truncate(coeffs_of_result); 

2.6 补充资料

sage: root=155978335310571138272812138773814534618935470879470300630834867870977
....: 67520449
sage: p=417793508166910149535221561915641187328610041455049381805950185075960478
....: 43329
sage:  omega= power_mod(root,2^21,p)//21=32-exp, for exp=11.
sage: omega
35014335792849108923302692126549442116295992392289760687159465394416590439942
sage: hex(omega)
'4d696968d9c7e5b55e6a88fe57cbaa9e166872f777629c2cd200ba70d7cec606'
sage: new_omega=power_mod(omega,4,p)
sage: hex(new_omega)
'323286f8bcf390a8b2be7ef037eab34c127cf6fc0b1e0bc4866f1ba33bc1dc80'
//对应的new_omega为2^{exp-log_n}即2^9-th primitive root of unity。
sage: power_mod(new_omega, 2^9,p)
1
sage: R=2^256
sage: mod(R,p) //域内的one值表示，之所以用R来代替1值，是因为采用了montgomery_reduce表示，1*R=1， montgomery_reduce后R/R=1。
32233387603934165516526672625559670387547976374630687678267546992721033953278
sage: hex(3223338760393416551652667262555967038754797637463068767826754699272103
....: 3953278)
'4743736b947db12c8a7ab15117a98d9efc82c7cb9bfdb6facc000305fffffffe'

对用地

参考资料：
[1] 论文《Halo: Recursive Proof Composition without a Trusted Setup》
[2] https://electriccoin.co/zh/blog/halo-recursive-proof-composition-without-a-trusted-setup/
[3] https://github.com/ebfull/halo