Rustの画像処理でSIMDを使ってみる

Transcript

1. Rust の画像処理でSIMD を使ってみる 人工衛星の開発現場でLT 大会 〜地上と宇宙で活躍するRust 〜 2023/06/07 fmy
2. 自己紹介 fmy twitter: @fmy4159 仕事: 家族型ロボットを作っている at GROOVE X 趣味: GStreamer とかJetson を触っている Qiita:Rust でGStreamer Plugin を書く Qiita:Jetson のカメラドライバを書く
3. 今日の話 効率よく画像データを処理したい Rust で画像データ処理にSIMD を使ってみた
4. 背景: 効率よく画像データを処理したい ロボティクス周辺では画像を使うケースがよくある パフォーマンスが優先( 安全に動く+ ユーザー体験を良くする) 同時に省電力にもしなければならない
5. 小さな移動体は電力が厳しい 用途例 本体従量 容量 卓上移動ロボット 500g-2kg 15Wh 家庭屋内移動ロボット 2kg-4kg 85Wh 屋外移動ロボット 10kg-30kg 200Wh 以上 電気自動車 1680kg 62000Wh 今のロボットは30-90 分程度の稼働、同等以上の時間をかけて充電することが多い 90 分動くとして計算に使える電力としては5-20W ぐらい
6. SoC の計算能力 よく使われているのはRaspberryPi, Jetson, Snapdragon, Zynq クロックが低くシングルコアでの性能は低め GPU はハード依存が大きい -> 汎用部分はCPU の性能を活かしたい Arch: aarch64 Core: 2-8 Core Clock: 1GHz-2Ghz Mem: 4-16GB GPU: ものによる その他アクセラレータを持っている
7. 画像データの形式 イメージセンサの出力はMIPI-CSI2 でRAW(8,10,12,14bit) 後段のデータ処理のために1 画素を8bit 以上で表現する RGB のセンサーならBayer 形式(e.g.RGGB) で、ISP で現像してYUV などを得る
8. RAW10,RAW12 データ MIPI 規格でフォーマットが決まっている MIPI が8bit/16bit を一塊で扱うので、上位8bit と下位nbit を分離して転送している RAW10 RAW12
9. Jetson T_R16 Format Jetson のMIPI-CSI のパーサはソフトウェアで処理しやすいように、 8 もしくは16bit 長でフォーマットしている
10. やりたい操作 LttleEndian 、12bit の数字データが16bit 幅で左詰めされている 右に4bit シフトして12bit の数値データを得たい
11. ナイーブな実装 u8 単位でシフトして上書きする pub fn t16_to_raw16_u8(buf: &mut [u8]) { for i in 0..buf.len() / 2 { let i: usize = i * 2; buf[i] = buf[i] >> 4 | buf[i+1] << 4; buf[i+1] >>= 4; } }
12. ベンチマーク 1.2MPixel(Quad-VGA12, 1280×960) ぐらいのデータを使う CPU: AMD Ryzen 5 3400G (Base 3.7GHz, Boost 4.2GHz) Rust: 1.70.0(stable) criterion でベンチマーク format use u8 time: [1.1869 ms 1.1900 ms 1.1935 ms] カメラが5-60fps(200ms-16ms) で他の処理もあることを想定すると 1ms は十分早いとは言えない
13. 本題: Rust で画像データ処理にSIMD を使ってみた 画像処理をもっと早くしたい...!
14. アセンブリを見る cargo-asm を使う。 rustc --emit asm で出せるけどファイル単位 cargo-asm は関数を指定してアセンブリを出力できる cargo asm crate::path::to::foo
15. u8 アセンブリ .LBB6_2: lea rcx, [rax - 1] cmp rcx, rsi jae .LBB6_7 cmp rax, rsi jae .LBB6_6 # for ループ抜け movzx r9d, byte ptr [rdi + rax - 1] # 読み出し movzx ecx, byte ptr [rdi + rax] shr r9b, 4 mov edx, ecx shl dl, 4 or dl, r9b # 下位8bit mov byte ptr [rdi + rax - 1], dl # 書き込み shr cl, 4 # 上位8bit mov byte ptr [rdi + rax], cl # 書き込み add rax, 2 dec r8 jne .LBB6_2
16. 遅い理由の考察 1. shift 操作3 回 i. 16bit の単位で見たら全部右シフトで同じ操作にできる 2. mov 操作4 回 i. load とstore の2 回が理想 3. pixel 数x2 のループ回数 2. 64bit 単位で1 回にできるのでは? i. 64bit 単位ならループ回数が1/4 に減る
17. u64 実装 pub fn t16_to_raw16_u64(buf: &mut [u8]) { for i in 0..buf.len() / 8 { let i: usize = i * 8; let mut a = LittleEndian::read_u64(&buf[i..i + 8]); *x = (*x & 0xfff0fff0fff0fff0) >> 4; LittleEndian::write_u64(&mut buf[i..i + 8], a); } } byteorder を使ってu64 として操作する
18. ベンチマーク with u64 1/4 になった! loop 回数削減( pixel/2 -> pixel/8 ) の影響が大きい? format use u8 time: [1.1869 ms 1.1900 ms 1.1935 ms] format use u64 time: [210.27 µs 211.84 µs 213.57 µs]
19. u64 のアセンブリ シフト1 回、move も読み出しと書き込みの2 回 単純な操作は削減効果は少ないが、ループ展開は効くようだ .LBB2_2: cmp rax, rsi ja .LBB2_5 mov r8, qword ptr [rdi + rax - 8] shr r8, 4 and r8, rdx mov qword ptr [rdi + rax - 8], r8 add rax, 8 dec rcx jne .LBB2_2
20. u128 Rust にはu128 がプリミティブとして用意されている。 これはどうなる? pub fn t16_to_raw16_u128(buf: &mut [u8]) { for i in 0..buf.len() / 16 { let i: usize = i * 16; let mut a = LittleEndian::read_u128(&buf[i..i + 16]); a = (a & 0xfff0fff0fff0fff0fff0fff0fff0fff0) >> 4; LittleEndian::write_u128(&mut buf[i..i + 16], a); } }
21. ベンチマーク with u128 u16 も比較対象に追加: 命令が単純になった分高速化している u128 はu64 より早い format use u8 time: [1.1869 ms 1.1900 ms 1.1935 ms] format use u16 time: [706.00 µs 710.77 µs 716.34 µs] format use u64 time: [210.27 µs 211.84 µs 213.57 µs] format use u128 time: [159.99 µs 164.70 µs 169.74 µs]
22. U128 ASM 処理自体は64bit と同じだがループ回数が少なくなっている 高速化要因はループ展開でオーバーヘッド削減したから? .LBB3_2: cmp rax, rsi ja .LBB3_5 mov r8, qword ptr [rdi + rax - 16] mov r9, qword ptr [rdi + rax - 8] shr r8, 4 shr r9, 4 and r9, rdx and r8, rdx mov qword ptr [rdi + rax - 16], r8 mov qword ptr [rdi + rax - 8], r9 add rax, 16 dec rcx jne .LBB3_2
23. まだまだいけるはず 1 命令を複数データに適用するといえばSIMD Ryzen 5 3400G ではSSE2 による128bit 幅の操作、AVX2 による256bit 幅の操作が使える
24. SSE2 /// # Safety #[target_feature(enable = "sse2")] #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] pub unsafe fn t16_to_raw16_unroll_128(buf: &mut [u8]) { #[cfg(target_arch = "x86")] use std::arch::x86::*; #[cfg(target_arch = "x86_64")] use std::arch::x86_64::*; #[allow(overflowing_literals)] let shift4 = _mm_setr_epi16(4, 0, 0, 0, 0, 0, 0, 0); for i in 0..buf.len() / 16 { let i: usize = i * 16; let invec = _mm_loadu_si128(buf.as_ptr().add(i) as *const _); let shifted = _mm_srl_epi16(invec, shift4); // 論理右シフト _mm_storeu_si128(buf.as_mut_ptr().add(i) as *mut _, shifted); } }
25. AVX2 /// # Safety #[target_feature(enable = "avx2")] #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] pub unsafe fn t16_to_raw16_unroll_256(buf: &mut [u8]) { #[cfg(target_arch = "x86")] use std::arch::x86::*; #[cfg(target_arch = "x86_64")] use std::arch::x86_64::*; #[allow(overflowing_literals)] let shift4 = _mm_setr_epi16(4, 0, 0, 0, 0, 0, 0, 0); for i in 0..buf.len() / 32 { let i: usize = i * 32; let invec = _mm256_loadu_si256(buf.as_ptr().add(i) as *const _); let shifted = _mm256_srl_epi16(invec, shift4); // 論理右シフト _mm256_storeu_si256(buf.as_mut_ptr().add(i) as *mut _, shifted); } }
26. ベンチマーク with SIMD SSE2 はさらに半分ぐらいになった AVX2 は右シフト1 回程度の操作の場合ではあまり効果がない? format use u8 time: [1.1869 ms 1.1900 ms 1.1935 ms] format use u16 time: [706.00 µs 710.77 µs 716.34 µs] format use u64 time: [210.27 µs 211.84 µs 213.57 µs] format use u128 time: [159.99 µs 164.70 µs 169.74 µs] format use 128 SSE2 time: [95.795 µs 97.490 µs 99.380 µs] format use 256 AVX2 time: [92.268 µs 92.971 µs 93.704 µs]
27. SSE2 ASM ループ展開されている。これも高速化の一因か .LBB4_8: movdqu xmm0, xmmword ptr [rdx - 48] movdqu xmm1, xmmword ptr [rdx - 32] movdqu xmm2, xmmword ptr [rdx - 16] movdqu xmm3, xmmword ptr [rdx] psrlw xmm0, 4 movdqu xmmword ptr [rdx - 48], xmm0 psrlw xmm1, 4 movdqu xmmword ptr [rdx - 32], xmm1 psrlw xmm2, 4 movdqu xmmword ptr [rdx - 16], xmm2 add rcx, 4 psrlw xmm3, 4 movdqu xmmword ptr [rdx], xmm3 add rdx, 64 cmp rsi, rcx jne .LBB4_8
28. Jetson でも実験をする 手元にある Jetson Orin Nano 8GB 6-core Arm® Cortex®-A78AE v8.2 64-bit CPU 1.5MB L2 + 4MB L3 1.5Ghz シングルコアのクロック数は3400G の1/2 以下
29. aarch64 Neon ARM のSIMD 拡張命令としてNeon がある /// # Safety #[cfg(target_arch = "aarch64")] pub unsafe fn t16_to_raw16_unroll_128(buf: &mut [u8]) { use std::arch::aarch64::*; const SHIFT4_VEC: [i16; 8] = [-4, -4, -4, -4, -4, -4, -4, -4]; let shift4 = vld1q_s16(SHIFT4_VEC.as_ptr() as *const _); #[allow(clippy::never_loop)] for i in 0..buf.len() / 16 { let i: usize = i * 16; let invec = vld1q_u16(buf.as_ptr().add(i) as *const _); let res = vshlq_u16(invec, shift4); vst1q_u16(buf.as_mut_ptr().add(i) as *mut _, res); } }
30. ベンチマークその4 CPU: Cortex®-A78AE v8.2 64-bit 1.5GHz Rust: 1.70.0(stable) format use u8 time: [2.6863 ms 2.6927 ms 2.7000 ms] format use u16 time: [1.3942 ms 1.3950 ms 1.3964 ms] format use u64 time: [399.52 µs 402.12 µs 404.56 µs] format use u128 time: [280.36 µs 280.38 µs 280.41 µs] format use 128 neon time: [144.28 µs 144.79 µs 145.35 µs] クロックが遅いためかu8 では2.7ms もかかる u128 で1/10 になっている SIMD を使うとx86 の処理時間に肉薄する
31. aarch64 u128 x86_64 と同様にu64 に加えてループ展開されている .LBB3_2: add x8, x10, #16 cmp x8, x1 b.hi .LBB3_5 add x11, x0, x10 subs x9, x9, #1 ldp x10, x12, [x11] lsr x10, x10, #4 lsr x12, x12, #4 and x13, x10, #0xfff0fff0fff0fff and x12, x12, #0xfff0fff0fff0fff mov x10, x8 stp x13, x12, [x11] b.ne .LBB3_2
32. Neon ASM Neon 高速化している。SSE2 のようにループ展開はしてない。 ( なお、手動で展開したが結果はほぼ変わらず) .LBB4_2: ldr q1, [x0] subs x8, x8, #1 ushl v1.8h, v1.8h, v0.8h str q1, [x0], #16 b.ne .LBB4_2
33. Additional もっと強いCPU の場合はどうなる? CPU: Intel(R) Core(TM) i7-12700 4.7GHz SIMD 使えば早いというわけでもないらしい format use u8 time: [875.07 µs 876.88 µs 878.68 µs] format use u16 time: [435.49 µs 436.66 µs 437.85 µs] format use u64 time: [109.95 µs 110.36 µs 110.86 µs] format use u128 time: [77.175 µs 77.540 µs 77.984 µs] format use 128 SSE2 time: [83.290 µs 83.452 µs 83.639 µs] format use 256 AVX2 time: [86.874 µs 87.127 µs 87.402 µs]
34. まとめ cargo-asm とても便利 同一命令操作は最適化で性能を上げる余地はある 64bitCPU なら64bit 単位の処理で書くことで命令数を減らせる u128 を64bit 環境使った場合64bit 命令のループ展開になる SIMD は多くの場合には高速化が見込めそう データや状況によるので効果の計測を行うのが大事