AlphaDev がソートアルゴリズムを 70% 高速化! C言語ライブラリの作者がDeepMindの最新AIについて解説

数日前、DeepMind はソートアルゴリズムを 70% 直接的に高速化する AlphaDev をリリースしました。

この新しい AI システムは、チェスの名人 AlphaGo をベースにしています。

この研究は、元 Google 研究者のジャスティン・タニー氏の興味をそそりました。

「C ライブラリの作者として、私は常に可能な限り最高のものをキュレートする機会を探しています」と彼女は語った。

Justine が DeepMind のソートアルゴリズムを詳しく説明する方法を見てみましょう。

DeepMind ソートアルゴリズム

DeepMind はこの発見で当然の注目を集めましたが、残念ながら、AlphaDev の説明をもっとうまくできたはずです。

次に、DeepMind がリリースした、疑似アセンブリからアセンブリに変換された 3 つの項目の配列をソートするアセンブリ コードから始めましょう。

この関数を move37() と名付けたのは、DeepMind のブログ記事でこの関数が AlphaGo の驚異的な「37 手目」と比較されていたためです。

2016年の人間対機械の戦いで、AlphaGoは人間の直感に反する単純な肩の突進を打ち、伝説の囲碁選手であるイ・セドルを破った。

そこで、DeepMind コードを実行すると次のようになります。

しかし、私の意見ではこれは間違いです。

指定した配列は {3,1,2} でしたが、move37() によって {2,1,3} にソートされました。

2 が 1 より前に来るとは思えないので、DeepMind は私たちを騙しているに違いありません。 LLVM libcxx へのオープンソースの貢献についても見てみましょう。これにより、物事が明確になると思います。

したがって、move37() は実際にはソート関数ではなく、sort3() 関数の構成要素として使用することを目的としたソート カーネルです。

ほんの一瞬、私を大いに混乱させたので、この論文とブログ記事でこのことに言及してもらえればよかったと思います。以下は、欠落している swap 操作を含む、コードのより良いバージョンです。

彼らのコードがなぜ重要なのかを説明するために、このアルゴリズムが高レベルでどのように機能するかを考えてみましょう。私が最初に sort3() 問題を自分で解決しようとしたとき、次のような結論に至りました。

次に libcxx を調べたところ、同じことが行われていました。上記のコードの問題は、コンパイラがコードの最適化をあまりうまく行っていないことです。

上記のコードをコンパイルしようとすると、コンパイラが多数の分岐命令を挿入していることに気付くでしょう。これは、DeepMind が LLVM の貢献によって改善しようとしていることです。

しかし、これらの技術は理解するのが容易ではないことがよくあります。

私は実は、この一見無害そうなコードが気に入っています。目を細めて見てみると、DeepMind の最先端のアセンブリ コードと同じ基本的な考え方を共有するパターンがわかるからです。

この問題は本質的に 3 つの比較とスワップの操作に要約されるという考え方です。

上記のコードは、ネットワークをソートするためのこれまでの最先端のコードです。ここで、DeepMind の新たな発見が役に立ちます。彼らは、上記の mov 命令が不要な場合があることを発見しました。

上記のコードを実行してみると、削除された行があるかどうかに関係なく、100% 正しいことがわかります。

このコード行は何かを実行しているように見えますが、実際には何も実行していません。ですから、このようなことが何十年もコンピューター サイエンスによって無視されてきたとしても、私は驚きません。

AlphaDev がどのように機能するかもより明確になったはずです。

DeepMind は基本的に、アセンブリ コードを操作してランダムに何かを削除し、それが壊れるかどうかを確認できる AI を構築しました。

私は AlphaDev の知性を否定するためにこれを言っているのではありません。なぜなら、私も同じことをしていないと言ったら嘘になるからです。

上記のコードにはさらに 2 つの mov 命令がありますが、これらは削除できる可能性があります。 ARM64 命令セットを使用してこれを行うと、同様の問題に対してはるかに小さなコードが可能になります。

ここでは、一時変数を作成するための指示は必要ありません。

Arm は最近好調を維持しており、上記の例は同社の評判の証拠となると思います。

Arm はオープンソース分野でもトップクラスの企業の 1 つです。たとえば、彼らの MbedTLS ライブラリは、私がこれまでに見た中で最も過小評価されている gem の 1 つです。

作業を開始したとき、Arm コードを変更して x86 ハードウェアでより適切に動作するようにする計画を立てていました。

x86 上の OpenSSL と同じパフォーマンスを実現するために、これらすべての複雑なアセンブリ最適化を作成しました。

MbedTLS はシンプルで移植性があり、解読可能な C コードなので、Perl で生成されたアセンブリではない暗号ライブラリを必要とする人にとっては朗報です。

私はArmの人たちに自分がやっていることを伝えましたが、彼らはそれが混乱を招いているとは思っていませんでした。

いつか、DeepMind がやっていることと同じことをする時間を見つけて、その変更をアップストリーム化したいと考えています。 Arm の最適化されたライブラリも豊富で、その品質は二重変換のそれに匹敵するものがありません。

ここで C ライブラリが特に興味深いのは、オープンソース コミュニティが何十年もの間、90 年代初頭に Sun Microsystems によって作成された数学関数に依存して業務をこなしてきたためです。

Arm は、 pow(x,y) などのいくつかの関数を改善する方法を見つけました。これは数学における最も基本的な演算の 1 つであることを考慮すると、非常に強力なものです。

たとえば、Arm のソリューションを使用して x86 マシン上で純粋なソフトウェアで pow(x,y) を実装すると、Intel のネイティブ x87 命令よりも 5 倍高速になります。

幸運なことに、DeepMind もこのゲームに参加したので、私は彼らの libcxx diff を読み取り可能な C コードに翻訳する権限を得ました。

これは、論文やブログ投稿で見たいもう一つのことです。なぜなら、このコードには、コンパイラに分岐のない MOVcc 命令を生成させるために専門家が使用する標準的なトリックが含まれているからです。

Sort5() 関数を見たとき、DeepMind の研究の背後にある動機をより深く理解できたように感じました。

ARM64 で Sort5() 関数をコンパイルすると、コンパイラは 11 個のレジスタを処理する関数を生成します。数式について推論する場合、一度に 11 個の変数を作業記憶に保持できますか?

おそらくそうではないでしょう。これが、PartialSort3 のような優れたカーネル関数が非常に便利な理由です。

Sort3() と Sort5() は、従来のソート機能のビルディング ブロックとなることを意図しているため、それ自体がカーネルであることに留意してください。

ブログ投稿ではこのトピックについて取り上げていますが、実際に移植可能で実行可能なものを共有すると便利だと思いました。

上記のアルゴリズムは、新しく改良された libcxx が行っていることを表しています。基本的にはクイックソートですが、小さなスライスに再帰するときにソートカーネルと挿入ソートに切り替わります。libcxx では、ヒープソートでシュレッピングする追加の手順も実行されていると思います。これは少し遅いですが、敵がスタックを破壊するのを防ぎます。上記のアルゴリズムは、新しく改良された libcxx が行っていることを表しています。基本的にはクイックソートですが、小さなスライスに再帰するときにソートカーネルと挿入ソートに切り替わります。 libcxx の場合、ヒープソート内を移動するという追加の手順も実行されたと思います。これは少し遅くなりますが、攻撃者がスタックを破壊するのを防ぎます。

• この時点で、皆さんが疑問に思うのは、これを使用できるのか、これらのソート ネットワーク カーネルによってソートが実際に高速化されるのか、ということでしょう。答えはイエスでもありノーでもあります。昇順で long をソートしたいだけなら、上記のコードは C ライブラリで提供される標準の qsort() 関数よりも 2 倍高速になります。ただし、そのためにカーネルは必要ありません。これまでのところ、私のパーソナル コンピューター (Intel Core i9-12900KS 搭載) では、上記の関数は long を 255 MB/秒でソートすることがわかりました。ただし、ソート カーネルをコメント アウトすると、これらのソートカーネルは実際にソートを高速化するのでしょうか?はいとも言えますし、いいえとも言えます。長い文字列を昇順で並べ替えたいだけの場合、上記のコードは C ライブラリで提供される標準の qsort() 関数よりも約 2 倍高速になります。ただ、それを実行するためにカーネルは必要ありません。これまでのところ、私の PC (Intel Core i9-12900KS 搭載) では、上記の関数は 1 秒あたり約 255 MB の速度でソートすることがわかりました。しかし、ソートカーネルをコメントアウトすると、次のようになります。

私の longsort() 関数は 1 秒あたり 275 メガバイトで実行され、アルゴリズムを簡素化することで 7% のパフォーマンス向上が達成されました。

long の利点は、int のキーと値のペアを格納するのに十分な長さがあり、マップ エントリをすばやく並べ替えることができるという便利な機能があることです。

上記の関数は、わずか 181 バイトの x86-64 マシン コードにコンパイルされます。

DeepMind の sort3() はわずか 42 バイトなので、サイズをいくらか犠牲にしてパフォーマンス上の利点が得られることを期待していました。

これまでに私が見つけた次善のアルゴリズムは、代わりに基数ソートを使用することです。これにより、速度は 400 MB/秒になりますが、malloc() に依存することに加えて、763 バイトものバイナリ フットプリントが必要になります。したがって、これらのコアのパフォーマンスが向上すると良いと思います。

これは、DeepMind のアイデアに価値がないと言っているわけではありません。

昨年、DeepMind が (Google Brain として知られていた当時) ベクトル化されたクイック ソート ライブラリを私たちに提供してくれたことは注目に値すると思います。これにより、ソートにおいて誰にも挑戦されない優位性を獲得しました。

私のコンピューターでは、Vqsor ソートの長時間実行は 1155 MB/秒です。

これは、オープンソース コミュニティで最も人気のあるライブラリの 1 つである djbsor よりもわずかに優れていますが、int より多くのデータ型に一般化されることはありません。

どちらの実装も、ベクトル化されたソート ネットワークを通じて実現されます。ここがソートネットワーク技術が真価を発揮するところだと思います。

知的存在に関して言えば、AlphaDev が幼児でなければ、そうするだろうと思います。

第一原理から始めると、ベースライン命令セットだけをサポートするのは非常に困難です。待っていれば、AlphaDev が今後、より困難な課題に取り組む中で、素晴らしい成果が見られるようになると思います。

また、DeepMind がアルゴリズムを小さくしたという事実も気に入っています。これはあまり見られない特徴です。

サイズコーディングは私の好きな趣味の一つです。このブログでは、383 バイトのラムダ計算仮想マシンと 436 バイトのガベージ コレクションされた Lisp マシンを公開しました。

私は、Cospolitan C ライブラリで使用しているサイズ最適化テクニックについてもブログに書きました。

また、私は DeepMind の親会社も気に入っています。数週間前に Google が私に Open Source Peer Grant を授与してくれたからです。ソフトウェアを小さくしたいという私の情熱を彼らが共有してくれているのは素晴らしいことです。

ベクトル化されたクイックソートを改善するためにこれを使用するのを見るのは素晴らしいことです。

最後に、AI 企業が機械語のコードでマシンを作成するというアイデアは気に入っています。なぜそうしないのでしょうか?機械の本質は機械です。

ビルダーとして、私はこれが OpenAI が作り出している未来よりもはるかに実現可能性が低いと感じています。

彼らは、ゼロサム経済の中で地球上のあらゆる建設業者と競争する巨大な家父長制の機械を構築し、その後、世界中の利権追求者たちを誘惑して、政府の規制を通じてその機械をコントロールさせている。

私が最も好きなタスク（コーディングなど）をすべて自動化するという OpenAI の約束は、進歩だとは思えません。私が望んでいるのは、ソートカーネルの検出など、自分ではできないことができるマシンを制御できるようになることです。これは本当の進歩です。

私たちが削減できる組立ラインの一つ一つが、その夢の実現に向けた前向きな一歩になると思います。