Rust の自動ベクトル化

# 動機

[Memory-safe PNG decoders now vastly outperform C PNG libraries](https://www.reddit.com/r/rust/comments/1ha7uyi/memorysafe_png_decoders_now_vastly_outperform_c/?rdt=59326)（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`文で回して、各要素にアクセスしながら差を計算します。

```rust: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`で回します。

```rust: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`命令を呼びます。（自動ベクトル化を無効にするためです）

```rust: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 で計測しました。

```rust: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 倍くらい速いです。この実装だけ自動ベクトル化されているようです。

```sh
$ 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 命令が使用されていますね。

```sh
$ 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`というメソッドを使ったら自動ベクトル化されました。最適化の奥は深い…:thinking:

```diff_rust: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]);
    }
}
```