Rust Zero Overhead Abstraction

比较C/C++，Rust在保存内存可见性和底层操作才能的基础上，供给了更友爱的语法和类型体系支撑。为保障内存安全，Rust也给开发者做了重重约束。在此我罗列了一些细节，阐明如何在一切权与生命周期的约束下，充分利用Rust的零本钱笼统才能，写出最易读且高功能的代码。本文按以下条目安排：

功能评价方法

以下一切的功能评价均依据对汇编代码的剖析完结。我曾试图经过benchmark来获取比照，成果发现不同写法的差别可能仅是个位数条指令，benchmark本身的波动就足以掩盖这种差别了。因而终究挑选了直接剖析比照汇编代码的方法。经过如下操作装置和调用 cargo-show-asm 库就能够比较简略完结对汇编代码的剖析。

cargo install cargo-show-asm
cargo asm --rust --bin ntt ntt::main

一个简略的输出样例如下。 --rust 参数会将源码混合在汇编代码中输出。留意Rust在编译过程中做了许多的inline操作，因而你会看见许多标准库或第三方库中的代码。一起因为Rust在编译时会将编译期能确认的求值直接预核算好，嵌入到编译成果中，因而主张一切的输入都经过随机值产生，防止一些比较简略的核算逻辑被直接省掉。

一起关于重视的函数，能够添加 #[inline(never)]的标记，防止其被内联，方便剖析。

// /rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/core/src/ptr/mod.rs : 1178
	crate::intrinsics::read_via_copy(src)
mov esi, dword ptr [rsp + 4*rax + 152]
	// /rustc/8ede3aae28fe6e4d52b38157d7bfe0d3bceef225/library/core/src/ops/index_range.rs : 61
	self.start = unsafe { unchecked_add(value, 1) };
mov qword ptr [rsp + 968], rdx
	// /home/winkar/rntt/src/main.rs : 337
	!(fu >= lim || fu <= -lim || gu >= lim || gu <= -lim)
add esi, -50
xor ecx, ecx
cmp esi, -99

整型运算

Rust中不支撑的整型操作

比较C/C++，Rust的类型体系更为严厉。

无隐式类型转化： 1u32 + 2u64 这种操作在C/C++中会触发隐式的向上提升类型转化，但在Rust中需求显式对前者经过 as u64 做类型转化。

这个问题很简略处理，在一切需求的当地加上 as [target_type] 就能够。但在转写C代码时要留意一个细节，有时会呈现如下所示的代码：
```
let x: u64 = ...;
...
let y = x as u32 as u64;
```
这跟直接的 y = x as u64 是不同的，会将x在u32的数据规模内切断。因而不能简化成一次as。
制止无符号整型取负： let x = 3i32; let y = -x; 这样的操作在C/C++中是合法的，但在Rust中是不合法的。

无符号整型取负的本质，是求它的补码。因而我编写了如下的代码，能完结任意无符号整型取负的操作:
```
#[inline]
pub fn neg<T>(x: T) -> T
where
    T: Not<Output = T> + From<u8> + Add<T, Output = T>,
{
    (!x) + T::from(1u8)
}
```
也能够经过另一个办法： 0.wrapping_sub(x) 。它跟上面的代码是等价的，但运用的0需求依据x的类型做改换，比方x是u32时，就需求写成 0u32.wrapping_sub(x) 。wrapping_sub的本质是模当时数据类型上限（对u32来说便是2^32）的减法运算。尽管说是取模，但实践上在汇编层面便是一条一般的疏忽溢出的指令。

上述两种办法开优化编译成汇编后都是相同的，在u32类型上，终究编译成果如下：
```
shr rax, 32
neg eax
```
默许进行溢出检查：Rust在dev方法下，会对一切的整型核算进行溢出检查，并在溢出时panic。

在Falcon的核算代码里，许多时分是将整型核算的溢出当一个feature去用的。绕过这种检查有三个办法：
1. 运用wrapping核算办法
  
  上面的wrapping_sub便是其间一个比方，它不会形成溢出。相似的还有wrapping_add, wrapping_sub
2. 运用std::num::Wrapping类型
  
  将一切用到i32的当地封装成Wrapping能够到达相同的作用。这个操作是零运行时本钱的，但在源代码里会十分丑陋。
3. 关闭dev方法下的溢出检查
  
  实践上我现在便是这么做的。
```
[profile.dev]
overflow-checks = false
```

slice、数组与Vec的挑选

Rust中的slice &[T] 是以胖指针的方法完结的。其内部结构相似于下面的C结构体：

struct slice {
	void* ptr;
	size_t len;
}

假如将其作为参数传递，它实践上会仿制这样一个结构体。以x86_64为例，默许调用约定运用寄存器传递前两个参数，slice的长度和起始指针会被别离存放到两个寄存器中，作为两个参数传递。因而很简略发现，slice类型的巨细是指针的两倍。当然，这个4~8个字节的差异关于非热点函数调用来说影响很小，简直能够疏忽。

但slice在运用过程中有一个问题需求留意：slice有强制的运行期规模检查（release下依然存在），在超出slice规模时会触发panic。因而在运用时，假如实践规模巨细是确认的，简直总是主张转化成确认巨细的数组运用。这样编译器能够在编译期完结检查。

fn testSlice(x: &[i32]) {
    let mut s = 0;
    for i in 0..200 {
        s += x[i];
    }
    println!("{}", s);
}

fn testSlice(x: &[i32]) {
  let mut s = 0;
	let x:[i32;100] = x.try_into().unwrap();
	for i in 0..50 {
	    s += x[i];
	}
	println!("{}", s);
}

像上面这两段代码，Rust会为左面每次对slice的下标拜访x[i]刺进一个规模判断，对右边数组则不会。一切的判断在unwrap时已经完结了。

假如不确认巨细，无法在源代码中做一个确认的转化，也能够经过assert的办法来触发编译器的优化，到达相同的作用。简直总是主张经过assert或许try_into来帮助编译器优化后续代码。

fn testSlice(x: &[i32]) {
    let mut s = 0;
    assert!(x.len() > 200);
		// let x = x[..200]也能够起到相同的作用
    for i in 0..200 {
        s += x[i];
    }
    println!("{}", s);
}

此外，还有个有意思的case：假如做了转化，unwrap正常，但拜访的规模过大，会怎么样呢？

fn testSlice(x: &[i32]) {
    let mut s = 0;
    let x:[i32;100] = x.try_into().unwrap();
    for i in 0..200 {
        s += x[i];
    }
    println!("{}", s);
}

答案是：编译器生成的汇编中会完结前一百次加和，然后在下标大于100时直接panic。

Vec的拜访方法与slice相似，但其内存存储在堆上，且具有动态扩容的才能。其结构相似下面的格局

struct Vec {
	void* ptr;
  size_t len;
  size_t capacity;
}

比slice多了一个字段。因而主张假如没有动态扩容的需求，尽量传递slice而非vec。而假如有在堆上分配内存的需求，能够考虑Box<[u32]> ，它只占一个指针巨细。但留意因为Rust不供给safe的内存分配API，所以当其长度不确认时，你实践上需求经过vec.into_boxed_slice API来获取这个Box对象。当然，还是跟前面说的相同，这点传参的价值差别很小。

内存复用

在Falcon的完结中，为了防止分配内存，分配了很长的一段内存作为内存池，一切的中心成果都在这一段内存上存储。

在Rust中，这个写法就与一切权和类型规矩产生了抵触：

一起持有内存池中不同偏移的引证（比方两个内存池中存储的向量相加）

这个问题十分简略，运用split和split_mut API即可。其底层经过unsafe的API完结，但咱们无需关心其间细节——它能将一个slice拆成两个引证回来给咱们。所耗费的价值仅仅是一次额定的长度检查罢了。

而关于slice内部不同引证的操作（比方交流），Rust也供给了一系列封装好的办法，比方 slice::swap slice::copy_within 等等。它们能够帮助咱们完结一些在C/C++中很简略但在Rust中很困难的操作。
同一段内存以不同的类型引证（先以[u32]的方法拜访，再以[u64]的方法拜访）

对此，我给出了一些尝试性的处理方案。不过总的来说，这个问题在safe Rust中暂时可能是无解的：
1. 仅拜访不修正，[u64]→[u8]或[u8]→[u64]的景象
  
  不局限于u64，其它整数类型亦可，但转化的另一端必定得是u8（或有符号的i8）。
  
  将slice用struct做封装，[u8;8]→[u64]能够经过from_ne_bytes完结安全的转化，反之则能够经过to_ne_bytes完结。将一切加减乘除操作封装后，就能够完结像一般数字相同的简略运算拜访。并且因为Rust的优化，这一层封装简直是无价值的，仅在数据初始化时会多一次长度检查。
  
  缺陷：首先是此封装类型具有必定的传染性，需求在上下文做显式的转化。其次因为Rust不供给赋值的重载，所以这个写法在需求修正对应slice时会十分复杂。
2. unsafe transmute
  
  相似于C++中的reinterpret_cast，这应该是最简略也功能价值最低的处理方案。但它运用了unsafe API，让咱们的代码不纯洁了。
3. 仿制一份，核算完结后再仿制回去
  
  额定添加两轮仿制和一些内存分配（因为流程中许多时分需求的内存巨细不固定，vec处理更简略）的操作，功能耗费很高，但写起来简略，并且它是safe的。我终究挑选的便是这个方案。
4. 别离为不同类型开设不同的内存池
  
  我没完结过这个方案，它表面上看很美，但在整个过程中记载不同类型内存池的偏移并在层层调用中传递这个信息，可能会引发一些未预期的问题——并且这看起来很费事。

类型封装与操作符重载

Rust供给了struct Fpr(u64)这种方法的匿名成员界说。它为咱们供给了一个十分好用的零本钱笼统手法。

在Falcon的完结中，它许多运用了Fpr（以整型存储的浮点值）类型，其加减乘除均和一般整数不同。用上面的类型封装后，再去重载一切相关的操作符，完结成员函数，就能在调用特定运算操作的一起，代码像简略的整数运算相同整齐。

截取一个Falcon中的片段说明这个问题：

pub fn fpc_div(a_re: Fpr, a_im: Fpr, b_re: Fpr, b_im: Fpr) -> (Fpr, Fpr) {
    let m = b_re.sqr() + b_im.sqr();
    let m_inv: Fpr = m.inv();
    let b_re_scaled = b_re * m_inv;
    let b_im_scaled = b_im.neg() * m_inv;
    let d_re = a_re * b_re_scaled - a_im * b_im_scaled;
    let d_im = a_re * b_im_scaled + a_im * b_re_scaled;
    (d_re, d_im)
}

对应的C代码如下所示。能够发现用操作符重载和成员函数重写后，代码的复杂程度大大下降。当然，实践上我还能够把fpc这个类型也做相同的处理。如此笼统迭代之后，能够大大提高可读性。

#define FPC_DIV(d_re, d_im, a_re, a_im, b_re, b_im)   do { \
		fpr fpct_a_re, fpct_a_im; \
		fpr fpct_b_re, fpct_b_im; \
		fpr fpct_d_re, fpct_d_im; \
		fpr fpct_m; \
		fpct_a_re = (a_re); \
		fpct_a_im = (a_im); \
		fpct_b_re = (b_re); \
		fpct_b_im = (b_im); \
		fpct_m = fpr_add(fpr_sqr(fpct_b_re), fpr_sqr(fpct_b_im)); \
		fpct_m = fpr_inv(fpct_m); \
		fpct_b_re = fpr_mul(fpct_b_re, fpct_m); \
		fpct_b_im = fpr_mul(fpr_neg(fpct_b_im), fpct_m); \
		fpct_d_re = fpr_sub( \
			fpr_mul(fpct_a_re, fpct_b_re), \
			fpr_mul(fpct_a_im, fpct_b_im)); \
		fpct_d_im = fpr_add( \
			fpr_mul(fpct_a_re, fpct_b_im), \
			fpr_mul(fpct_a_im, fpct_b_re)); \
		(d_re) = fpct_d_re; \
		(d_im) = fpct_d_im; \
	} while (0)

这些笼统并不会引入额定的价值。该调函数调函数，该内联就内联，与原本代码的功能完全一致。

除了上述比方之外，多项式等类型也能够做相同的处理。不过也有类型无法做这样的处理：我一向想对环上的modq运算做相似的类型构建，但额定存储一个q作为成员价值似乎过高，不存储q的话又不太匹配通用的加减乘除trait，代码无法简化。

迭代器

Rust中供给了丰富的迭代器语义，一切slice都能够转化成迭代器Iter，在其基础上进行takewhile, map, fold等操作。因为Rust中默许不供给C sytle for句子，许多时分能够经过迭代器写出语义更明晰的循环。但在运用时要留意，迭代器作为一种笼统，会引入一些额定的价值。

例如 x.iter.step_by(p).take(q) 这样的句子，它表面上看可能跟 for (int i=0, a=x; i<q; i++, a+=p) 这样的循环是等价的，但实践上step_by和take都会刺进额定的判断句子。其间 step_by 会检查p是否为0（这只是单次检查，价值较小），而 take 会在每次迭代后检查是否超出长度约束（价值相对更大）。

防止这种检查的方法与上文相同：将x转化为固定长度的数组类型，或在前面刺进对x.len()的assert判断，帮助编译器优化迭代内部的长度校验。当咱们给编译器供给了充足信息时，它才能将迭代器优化为笼统价值最低的代码——Rust在默许的release级别上，会将简略的循环运算（迭代器也相同如此）做循环打开，并做SSE SIMD优化（下面有一个比方）。

迭代器API中的长度检查（step_by，take等）能够经过传入常量（比方 take(2) ）或许在循环之前对对应鸿沟做 assertion 来规避。

在迭代两个slice时，有个常用的迭代器函数是 zip。在数组上运用这个函数时，有个简略误用的点：

let x = [0u32;56];
let y = x;
for (&i, &j) in x.iter().zip(y) {
	//...
}

这个循环的写法是合法的，能经过编译，也能输出契合预期的成果。但它有一个小问题：会在迭代的时分对y进行一次不必要的memcpy。

原因很简略：咱们在zip(y)时直接传值，而数组作为一个完结了Copy trait的类型，不会被move，只会被值传递，因而这里就刺进了一次额定的memcpy。

要防止这次仿制也很简略：将迭代的代码改为 x.iter().zip(y.iter()) 即可。

别的留意一点：Rust关于迭代器的优化好于直接循环。

pub fn poly_add(a: &mut [Fpr], b: &[Fpr], logn: u32) {
    let n = 1usize << logn;
    assert!(a.len() >= n && b.len() >= n);
    // Compare to implementation: below,
    // iterator based code will not generate bound-check-and-panic code.
    // for u in 0..n {
    //     a[u] = a[u] + b[u];
    // }
    for (ax, bx) in a.iter_mut().zip(b.iter()).take(n) {
        *ax = *ax + *bx;
    }
}

能够参阅上面这组比照的代码，在语义上，毫无疑问咱们已经经过assertion约束了循环中绝不会呈现越界的情况，但实践上编译器依然会为循环生成鸿沟检查的代码。要防止这种情况，就能够运用下面的迭代器写法，尽管代码更难读了一些，但如此生成的代码功能更佳。

别的无论是哪一种写法，在启用了AVX（具体说明见SIMD节）之后，得到的中心汇编都如下：

		// /home/winkar/pqc-rust/src/falcon/fft.rs : 308
		*ax = *ax + *bx;
	vmovupd ymm0, ymmword ptr [rdi + 8*rsi]
	vmovupd ymm1, ymmword ptr [rdi + 8*rsi + 32]
	vmovupd ymm2, ymmword ptr [rdi + 8*rsi + 64]
	vmovupd ymm3, ymmword ptr [rdi + 8*rsi + 96]
		// /home/winkar/pqc-rust/src/falcon/fpr.rs : 443
		Fpr((x + y).to_bits())
	vaddpd ymm0, ymm0, ymmword ptr [rdx + 8*rsi]
	vaddpd ymm1, ymm1, ymmword ptr [rdx + 8*rsi + 32]
	vaddpd ymm2, ymm2, ymmword ptr [rdx + 8*rsi + 64]
	vaddpd ymm3, ymm3, ymmword ptr [rdx + 8*rsi + 96]
		// /home/winkar/pqc-rust/src/falcon/fft.rs : 308
		*ax = *ax + *bx;
	vmovupd ymmword ptr [rdi + 8*rsi], ymm0
	vmovupd ymmword ptr [rdi + 8*rsi + 32], ymm1
	vmovupd ymmword ptr [rdi + 8*rsi + 64], ymm2
	vmovupd ymmword ptr [rdi + 8*rsi + 96], ymm3

能够看出是进行了循环打开又用SIMD做了向量化的高度优化的代码——而咱们只需求写最原始的逻辑，循环打开和向量化的优化都由编译器主动完结。

依据宏和模板的代码复用

Rust中供给了泛型和宏的才能。尽管Rust的泛型自带了C++20才参加的concept支撑，但很可惜stable Rust到现在还不支撑const generic expr，导致咱们无法在密码算法中简略地经过泛型参数来完结不同安全等级的版别。

上述需求经过泛型尽管无法完结，但经过宏能够十分简略地做到。一个样例相似于此：

#[macro_export]
macro_rules! define_falcon_keypair {
    ($logn: expr, $sk_bytes: expr, $pk_bytes: expr) => {
        // Generate keypair, return (sk, pk)
        // ### Example
        // ```
        // # fn main() {
        // let (sk, pk) = keypair();
        // # }
        // ```
        pub fn keypair() -> (Vec<u8>, Vec<u8>) {
            let mut pk = [0u8; $pk_bytes];
            let mut sk = [0u8; $sk_bytes];
            let mut seed = [0u8; SEED_BYTES];
						// ...
				}
		}
}

但这个做法的缺陷也是显然的：目前Rust-Analyzer对宏的解析支撑不好，宏中无法进行代码跳转。因而仅主张在代码完结调试完结后再做宏封装。不然会十分影响开发功率。

别的当其间涉及到一些表达式较为复杂的常量参数，不适合直接作为单个参数传递时，也能够运用const function进行编译期求值：

pub const fn mkn(logn: u32) -> u32 {
    1 << logn
}
// 能够直接将mkn的eval成果作为常量运用，且mkn也依然能够正确作用于变量
let mut logn = 10;
const logn_2 = 9;
let buffer = [0u32; mkn(logn_2) as usize];
let n = mkn(logn);
// 上述句子都是合法的。

SIMD

SIMD在核算密集型任务中十分常用。因为我本身目前的运用场景，在此暂时不评论异构设备（比方GPU）上的SIMD，只评论针对CPU的优化。运用适宜的CPU指令，能够在一个指令周期中完结屡次加法、乘法核算，对核算密集型程序有较大的增益。

目前主流架构较新型号的CPU均供给了SIMD指令集，在x86/x86_64上有SSE*（Streaming SIMD Extension）和AVX*（Advanced Vector Extension）系列指令，在ARM上有NEON，SVE*（Scalable Vector Extension）指令集，别的在MIPS、RISC-V上也均有对应的SIMD指令集。这些指令集供给的原语和接口各有不同，需求在开发时做针对性的优化。

在正式评论SIMD代码的编写之前，咱们需求先知道如何在Rust中启用SIMD。依然以x86_64为例，能够经过 rustc --print=cfg | grep "target_feature” 检查默许启用的优化，在我的CPU上，能够看到默许只启用了SSE的优化，没有启用支撑更大向量的AVX指令集。

target_feature="fxsr"
target_feature="sse"
target_feature="sse2"

要启用AVX指令集，能够在编译之前设置RUSTFLAGS环境变量，export RUSTFLAGS="-C target-feature=+avx2,+fma"，如此便会在编译产物中主动引入AVX指令（fma指令是浮点乘加指令，是一个弥补，能够不开启）。

感兴趣的读者能够运用rustc --print=target-features检查编译器和当时CPU支撑的一切拓展。

在启用上述编译选项后，无论是否调用了专用的SIMD函数，编译器都会为适宜的句子生成AVX指令版的优化代码。

为便于编译器生成更适宜的SIMD代码，能够参阅这篇文章，揣摩如何写出编译器友爱的代码——最简略的思路是：对咱们的循环核算做打开，打开后的代码往往能更好地被向量化。

当然，相当多的时分，编译器并无法很好地主动对咱们的代码进行向量化，这时分就需求咱们手动调用对应的指令集进行运算。

这并不需求咱们去在Rust当中嵌入汇编指令——尽管这确实能够做到——只需求调用std::arch::x86_64（或许其它方针架构）中封装好的函数，比方_mm256_set1_pd，_mm256_mul_pd等等。即可用纯Rust代码完结向量化的核算。可惜的是，上述代码都是unsafe的。Rust unstable中供给了std::simd库，它目前以portable_simd的姓名在github上发布。用它能够完结safe的向量化核算。但这个库当时还无法运用。假如确实有需求，也能够运用 wide 替代。不过无论是 portable_simd还是wide，它们供给的都是通用的向量化API，可能无法完结一些只有特定架构上的SIMD指令才能完结的操作。

Rust零成本抽象指北

Rust Zero Overhead Abstraction

功能评价方法

整型运算

slice、数组与Vec的挑选

内存复用

类型封装与操作符重载

迭代器

依据宏和模板的代码复用

SIMD

作者信息

Rust零成本抽象指北

Rust Zero Overhead Abstraction

功能评价方法

整型运算

slice、数组与Vec的挑选

内存复用

类型封装与操作符重载

迭代器

依据宏和模板的代码复用

SIMD

相关文章

使用 Transformers 为多语种语音识别任务微调 Whisper 模型

三两下带你解决Android input问题

Cold Diffusion：高斯噪声，我一定非你不可吗？

Java开发高质量代码建议1：三元操作符的类型务必一致

作者信息