動機
Memory-safe PNG decoders now vastly outperform C PNG libraries（Rust の PNG デコーダが C の PNG ライブラリの性能を大幅に上回っているよ）という記事を読みました。性能向上の肝は DEFLATE のストリーミング展開と、自動ベクトル化による SIMD 命令の活用の二つであるとのことです。どういうコードを書いたら Rust で自動ベクトル化されるのか気になったので、簡単なコードを書いて試してみました。
環境
- Windows 10 WSL2 Debian12
- rustc 1.83.0
- Core i5 1345U
実装
二つの配列a, bの各要素の差|a[i]-b[i]|を計算して配列c[i]に保存します。画像処理なんかではよくあるパターンかと思います。
愚直に
まずは愚直にfor文で回して、各要素にアクセスしながら差を計算します。
src/lib.rs
pub fn diff(a: &[u32], b: &[u32], c: &mut [u32]) {
for i in 0..a.len() {
if a[i] > b[i] {
c[i] = a[i] - b[i];
} else {
c[i] = b[i] - a[i];
}
}
}
イテレータ
次にイテレータを使用します。zipでイテレータを結合してfor_eachで回します。
src/lib.rs
pub fn diff_iter(a: &[u32], b: &[u32], c: &mut [u32]) {
a.iter()
.zip(b.iter())
.zip(c.iter_mut())
.for_each(|((a, b), c)| {
if a > b {
*c = a - b;
} else {
*c = b - a;
}
});
}
アセンブリ
上と同様にイテレータを使用しますが、インラインアセンブリでsub命令を呼びます。（自動ベクトル化を無効にするためです）
src/lib.rs
pub fn diff_asm(a: &[u32], b: &[u32], c: &mut [u32]) {
a.iter()
.zip(b.iter())
.zip(c.iter_mut())
.for_each(|((a, b), c)| {
if a > b {
unsafe { asm!("sub {0:e}, {1:e}", inout(reg) *a => *c, in(reg) *b) };
} else {
unsafe { asm!("sub {0:e}, {1:e}", inout(reg) *b => *c, in(reg) *a) };
}
});
}
ベンチマーク
みんな大好きベンチマーク。Criterion で計測しました。
benches/my_bench.rs
use criterion::{Criterion, criterion_group, criterion_main};
use mytest::{diff, diff_asm, diff_iter, random};
const SIZE: usize = 1 << 24; // 16 MiB
fn bench(cr: &mut Criterion) {
let a = (0..SIZE).map(|_| random()).collect::<Vec<_>>();
let b = (0..SIZE).map(|_| random()).collect::<Vec<_>>();
let mut c = vec![0; SIZE];
let mut group = cr.benchmark_group("diff_bench");
group.bench_function("diff", |bm| bm.iter(|| diff(&a, &b, &mut c)));
group.bench_function("diff_iter", |bm| bm.iter(|| diff_iter(&a, &b, &mut c)));
group.bench_function("diff_asm", |bm| bm.iter(|| diff_asm(&a, &b, &mut c)));
group.finish();
}
criterion_group!(benches, bench);
criterion_main!(benches);
イテレータを使用した実装diff_iterが他よりも 6 倍くらい速いです。この実装だけ自動ベクトル化されているようです。
$ cargo bench
Compiling mytest v0.1.0 (/home/benki/rust/mytest)
Finished `bench` profile [optimized] target(s) in 4.32s
Running benches/my_bench.rs (target/release/deps/my_bench-6e8a4722e7cb98f8)
diff_bench/diff time: [56.597 ms 57.003 ms 57.451 ms]
diff_bench/diff_iter time: [9.3346 ms 9.4614 ms 9.5891 ms]
diff_bench/diff_asm time: [52.269 ms 52.632 ms 53.040 ms]
objdump でバイナリを覗いてみます（全部書くと長いので一部抜粋）。イテレータを使用した実装diff_iterでは SIMD 命令が使用されていますね。
$ objdump -d target/release/deps/my_bench-6e8a4722e7cb98f8
"fn diff()"
00000000000cf290 <_ZN6mytest4diff17he6146633e2ad21abE>:
（略）
cf2a5: 45 29 d3 sub %r10d,%r11d
cf2a8: 45 89 1c b8 mov %r11d,(%r8,%rdi,4)
cf2ac: 48 ff c7 inc %rdi
cf2af: 48 39 fe cmp %rdi,%rsi
cf2b2: 74 26 je cf2da <_ZN6mytest4diff17he6146633e2ad21abE+0x4a>
cf2b4: 48 39 f9 cmp %rdi,%rcx
cf2b7: 74 23 je cf2dc <_ZN6mytest4diff17he6146633e2ad21abE+0x4c>
cf2b9: 44 8b 14 b8 mov (%rax,%rdi,4),%r10d
cf2bd: 44 8b 1c ba mov (%rdx,%rdi,4),%r11d
cf2c1: 44 89 d3 mov %r10d,%ebx
cf2c4: 44 29 db sub %r11d,%ebx
cf2c7: 76 d7 jbe cf2a0 <_ZN6mytest4diff17he6146633e2ad21abE+0x10>
（略）
"fn diff_iter()"
00000000000cf310 <_ZN6mytest9diff_iter17h2d5bb75c37dc0c69E>:
（略）
cf36a: 66 0f fa eb psubd %xmm3,%xmm5
cf36e: 66 0f ef d8 pxor %xmm0,%xmm3
cf372: 66 0f ef c8 pxor %xmm0,%xmm1
cf376: 66 0f 66 cb pcmpgtd %xmm3,%xmm1
cf37a: 66 0f ef e9 pxor %xmm1,%xmm5
cf37e: 66 0f fa cd psubd %xmm5,%xmm1
cf382: 66 0f 6f da movdqa %xmm2,%xmm3
cf386: 66 0f fa dc psubd %xmm4,%xmm3
cf38a: 66 0f ef e0 pxor %xmm0,%xmm4
cf38e: 66 0f ef d0 pxor %xmm0,%xmm2
cf392: 66 0f 66 d4 pcmpgtd %xmm4,%xmm2
cf396: 66 0f ef da pxor %xmm2,%xmm3
cf39a: 66 0f fa d3 psubd %xmm3,%xmm2
cf39e: f3 41 0f 7f 0c b0 movdqu %xmm1,(%r8,%rsi,4)
（略）
"fn diff_asm()"
00000000000cf3e0 <_ZN6mytest8diff_asm17h1b3cd79d37687237E>:
（略）
cf400: 41 29 f1 sub %esi,%r9d
cf403: 44 89 ce mov %r9d,%esi
cf406: 41 89 34 80 mov %esi,(%r8,%rax,4)
cf40a: 48 ff c0 inc %rax
cf40d: 48 39 c1 cmp %rax,%rcx
cf410: 74 1b je cf42d <_ZN6mytest8diff_asm17h1b3cd79d37687237E+0x4d>
cf412: 44 8b 0c 87 mov (%rdi,%rax,4),%r9d
cf416: 8b 34 82 mov (%rdx,%rax,4),%esi
cf419: 41 39 f1 cmp %esi,%r9d
cf41c: 77 e2 ja cf400 <_ZN6mytest8diff_asm17h1b3cd79d37687237E+0x20>
cf41e: 44 29 ce sub %r9d,%esi
（略）
まとめ
今回の記事ではイテレータを使用した実装が自動ベクトル化されました。下手にインラインアセンブリ書くよりコンパイラに任せた方がいいのかもしれませんね（いろんなアーキテクチャに対応するのも楽だろうし）。いろいろ条件を変えてどんなときに自動ベクトル化されて、どんなときにされないのか調べていきたいと思います。おわり。
追記
上記の愚直にforで回すコードにabs_diffというメソッドを使ったら自動ベクトル化されました。最適化の奥は深い…
src/lib.rs
pub fn diff(a: &[u32], b: &[u32], c: &mut [u32]) {
for i in 0..a.len() {
- if a[i] > b[i] {
- c[i] = a[i] - b[i];
- } else {
- c[i] = b[i] - a[i];
- }
+ c[i] = a[i].abs_diff(b[i]);
}
}