より高速なバイナリサーチ：コンパイル済みコードからメカニカルシンパシーへ

6倍高速なバイナリサーチ：コンパイル済みコードからメカニカルシンパシーへ by Itamar Turner-Trauring 最終更新日: 2026年7月15日、初版: 2026年7月11日 計算量の多いPythonコードを高速化するにはどうすればよいでしょうか？ 一般的で有用な出発点は、良いアルゴリズムを選択すること、コンパイル言語でPython拡張を作成すること、そして並列処理を追加して複数のCPUコアを利用することです。しかし、さらに速度が必要な場合はどうでしょうか？ 以下の実際の例を考えてみてください。これは、scikit-learnの勾配ヒストグラムブースティングアルゴリズムのステップの一つです。浮動小数点数の大きな配列があり、それらを0〜254の整数範囲に均等に割り当てたいとします。scikit-learnでは、浮動小数点値の全範囲を255個のバケットに分割し、バケット境界のソート済み配列を作成し、バイナリサーチを使用して各値に適切なバケットを選択します。このバイナリサーチはコンパイル言語で実装されており、複数のコアで並列実行できます。 最近、Quansightでの仕事の一環として、Paul Khuong氏の2つの投稿に触発されて、この実装を大幅に高速化しました。どのようにして？ コードがCPUと戦わないようにしたのです。この記事では、その高速化の過程を、簡単な例で説明します。その後、一連の追加最適化を実演し、最終バージョンは元のバージョンの6倍の速度で実行されます。多くの異なる低レベルハードウェアトピック（命令レベル並列性、分岐（誤）予測、メモリキャッシュ、SIMDなど）に触れることになるので、注意が必要です。これは単一の記事であり、可能なことの概要を簡単に紹介するだけで、詳細なチュートリアルとして機能することはできません。そのため、記事の最後にこれらのトピックについてさらに学習する方法について説明します。 出発点：標準的なバイナリサーチ 元のscikit-learnコードはCythonで実装されていましたが、この記事ではRustを使用します。以下は、境界の配列を与えられた場合にバケットを見つけるこのユースケースのために設計された、 pretty standardなバイナリサーチの実装です（NumPyのものをベースにしています）： ```rust use std::cmp::Ordering; /// Rustでは、通常の < 演算子で浮動小数点数を比較することはできません（NaNが比較結果を /// 一貫性のないものにするため）、そのためカスタム関数を実装します。 fn less_than(a: f64, b: &f64) -> bool { a.total_cmp(b) == Ordering::Less } /// バケット境界を与えられた場合、浮動小数点値をどのバケットに適合するかを見つけることで、 /// 整数値に変換します。 fn bucketize_classic_impl( arr: &[f64], boundaries: &[f64], ) -> Vec<usize> { // Vec（またはvector）はRustのPythonリストに相当します。 // ここでは、`arr.len()` の値を格納するのに十分なメモリを割り当てた空のVecを作成します： let mut result = Vec::with_capacity(arr.len()); for value in arr { // 標準的なバイナリサーチアルゴリズム： let mut min_idx = 0; let mut max_idx = boundaries.len(); while min_idx < max_idx { let middle = min_idx + ((max_idx - min_idx) / 2); if less_than(boundaries[middle], value) { min_idx = middle + 1; } else { max_idx = middle; } } // これはPythonの `a_list.append()` に相当します： result.push(min_idx); } // 結果を返します： result } ``` 完全を期すために、これをPythonにフックする方法を示します。これにより、NumPy配列を入力として受け取り、返します。これは定型的なコードなので、後続の関数では表示しません。RustとPyO3に慣れていない場合、または特に気にならない場合は、スキップしても構いません。記事の残りの部分とは関係ありません。コードを表示するにはここをクリックしてください。 ```rust use pyo3::prelude::*; use numpy::ndarray::Array; use numpy::{PyArray1, PyReadonlyArray1}; #[pyfunction] fn bucketize_classic<'py>( py: Python<'py>, // これら2つの引数は浮動小数点数の1次元配列です： arr: PyReadonlyArray1<f64>, boundaries: PyReadonlyArray1<f64>, ) -> PyResult<Bound<'py, PyArray1<usize>>> { let result = bucketize_classic_impl( arr.as_slice().unwrap(), boundaries.as_slice().unwrap(), ); // RustのベクトルをPythonに返すことができる1D NumPy配列に変換します： let result = PyArray1::from_owned_array(py, Array::from_vec(result)); Ok(result) } ``` 分岐予測の失敗はコードを遅くします このバイナリサーチの実装を高速化するにはどうすればよいでしょうか？ すでにスケーラブルなアルゴリズムとコンパイル言語を使用しています。並列処理はもちろん選択肢ですが、ここでは別のアプローチを使用します。それは、CPUがどのように機能するかをより良く理解する「メカニカルシンパシー」です。まず、最新のCPUが単一コア内でどのようにコードを並列実行するかを簡単にレビューします。Pythonコードの合理的なメンタルモデルは、コードが一度に1つの命令ずつ実行されるというものです。算術演算を2倍行えば、コードは2倍遅くなります。コンパイル言語に切り替えると、操作が1つまたは2つのCPU命令にマッピングされる場合があるため、そのメンタルモデルはもはや正しくありません。最新のCPUは、単一のCPUコアで複数の独立したCPU命令を同時に実行できる場合があり（「命令レベル並列性」）、実行速度が向上します。 ```rust fn two_adds(a: i64, b: i64, c: i64, d: i64) -> i64 { // あなたのCPUは、おそらくこれらの2つの加算を、単一コアで自動的に並列実行できます： let t1 = a + b; let t2 = c + d; // 結果を返します： t1 + t2 } ``` しかし、コード内のif/while/for式によって作成される分岐は問題を引き起こします。コードは一方の経路をたどるか、もう一方の経路をたどる可能性があります。2つの選択肢がある場合、CPUは将来のどの命令セットを並列実行しようと試みるべきでしょうか？ ```rust fn maybe_add( a: i64, b: i64, c: i64, d: i64, add: bool, ) -> i64 { let t1 = a + b; // `a + b` と並列に、これらの2つの分岐のどちらを実行すべきでしょうか？ let t2 = if add { c + d } else { c * d }; t1 + t2 } ``` 高速な実行を保証するために、CPUにはヒューリスティックにどの分岐を並列実行するかを選択する分岐予測器があります。推測が正しければ、コードは高速になります。推測が間違っていれば、CPUは最終的にそれに気づき、誤った作業を元に戻し、正しい分岐を実行します…これはコードが遅くなることを意味します。場合によっては、大幅に遅くなります。 上記のバイナリサーチアルゴリズムは、残念ながら入力データによっては非常に予測が困難です。バケット境界は、入力値がすべてのバケットに均等に分散されるように選択されていることを思い出してください。これは、バイナリサーチで左右どちらに進むかの選択が全く予測不可能であることを意味します。 ```rust // CPUはどちらのパスが取られるかを確実に推測できません： if less_than(boundaries[middle], value) { min_idx = middle + 1; } else { max_idx = middle; } ``` 同様に、反復回数も変動する可能性があります。 ```rust // このループは停止するまでに何回継続するでしょうか？ // 使用されている特定のデータによっては、知ることは不可能です。 while min_idx < max_idx { // ... } ``` この仮説を検証するために、CPUのハードウェアカウンター（py-perf-eventを介してPythonからアクセス可能）を使用して、実行中のコードが実行する分岐の数と、これらの分岐のうちいくつが誤って予測されたかを測定できます。使用する入力は次のとおりです。 ```python import numpy as np from numba import jit # 0から1の間の値： DATA = np.random.random(1_000_000) # バケット境界、均等に配置： BOUNDARIES = np.linspace(0.0, 1.0, 255)[1:-1] ``` そして、コードを実行した結果は次のとおりです。 ``` Code ➘ Elapsed µ-seconds ➘ Branch instructions ➘ Branch misprediction % bucketize_classic(DATA, BOUNDARIES) 45,870.2 26,997,038 16.6% ``` ➘ 低い数字ほど良い 16%の分岐が誤って予測されるのは良くありません。また、値あたりの分岐数もかなり多いです。DATAは1,000,000個の値があり、合計2700万回の分岐があるため、値あたり27回の分岐となります。 分岐なしの実行に切り替える 予測不可能な分岐の両方の原因を取り除きます。whileループの反復回数については、代わりに固定回数、つまりバケット数のlog2回反復することにします。値によっては、以前は1〜2回の反復でバケットが見つかっていた場合、少し多くの作業が発生するかもしれませんが、分岐予測の失敗を回避することによる速度の向上により、それだけの価値があります。if式の代わりに、現在のコードは一方の変数、もう一方の変数を設定していましたが、常にコードを置き換えます。