AVX-512で実装した `tolower()` 関数

 (dotat.at)
1 ポイント 投稿者 GN⁺ 2024-07-30 | 1件のコメント | WhatsAppで共有
  • 文字列をコピーしながら ASCII 大文字を小文字に変換する処理を AVX-512-BW で64バイトずつ処理し、短い文字列でも SIMD の性能を引き出す実験
  • 実装の核心は、各バイトが 'A' 以上 'Z' 以下かを比較したうえで、その位置にだけ 'a' - 'A' を足すマスク演算 を行うこと
  • 短い文字列と長い文字列の末尾の残りは マスク付き load/store で処理し、SIMD コードでありがちな細切れ処理のコストを抑える
  • Clang 16、Debian 11、AMD Ryzen 9 7950X で約 1MiB のコピーを 1バイト〜1KiB チャンクで測定した結果、tolower64 は比較対象の中で一貫して高速な部類に入った
  • Zen 4 では AVX-512-BW が文字列処理に非常によく適していたが、ARM SVE と RISC-V Vector 拡張は詳しく直接検証できていない

AVX-512-BW で 64バイト tolower() を作る

  • 目標は、文字列をコピーしながら大文字の ASCII 文字を小文字に変換する tolower() カーネル を SIMD で実装すること
  • AVX-512-BW はバイトおよびワード単位の演算をサポートする拡張で、最近の AMD Zen プロセッサで利用できる
    • AVX-512 は複数の拡張に分かれており、サポート状況が複雑
    • Intel 側のサポートは特に一貫性がないと評価している
  • ARM SVE も文字列処理に適した バイト単位のマスク付き load/store を提供する
    • 最近の big-ARM Neoverse コア、たとえば Amazon Graviton で利用可能
    • Apple Silicon では利用できない
  • RISC-V Vector extension も ARM SVE に近いスタイルで、複数の小型シングルボードコンピュータで利用できる

tolower64() の動作方式

  • tolower64() は一度に 64バイト を処理する AVX-512 ベースのカーネル
  • まず 64 個のバイトが入るベクトルレジスタに基準値を詰める
    • 'A'
    • 'Z'
    • 'a' - 'A'
  • 入力文字ベクトル c を 'A'、'Z' と比較して、それぞれ 64ビットマスクを作る
    • c >= 'A' である位置
    • c <= 'Z' である位置
  • 2つのマスクを _kand_mask64() で結合し、大文字の位置 だけを示す is_upper マスクを作る
  • 最後に _mm512_mask_add_epi8() を適用する
    • is_upper が false のバイトは元の c を維持する
    • is_upper が true のバイトは c + ('a' - 'A') になる

長い文字列と短い文字列の処理

  • 長い文字列の大部分は通常の 非アラインベクトル load/store で処理する
    • _mm512_loadu_epi8()
    • tolower64()
    • _mm512_storeu_epi8()
  • 短い文字列と長い文字列の最後に残る断片には マスク付き非アライン load/store を使う
  • マスクは下位 len ビットだけが立つ形で作る
    • uint64_t len_bits = (~0ULL) >> (64 - len)
    • _cvtu64_mask64(len_bits) で SIMD マスクレジスタに載せる
  • _mm512_maskz_loadu_epi8() はマスクがオフの位置の宛先レジスタを 0 で埋める
  • _mm512_mask_storeu_epi8() はマスクがオンの位置だけを書き込む
  • この方式が小さな文字列断片を高速に処理する鍵になる

ベンチマーク条件と比較対象

  • ベンチマークは Clang 16、Debian 11、AMD Ryzen 9 7950X で実行
  • 測定対象は約 1MiB のコピー で、チャンク長は 1バイトから 1KiB まで変化させた
  • ソース文字列と宛先文字列のアラインメント差を反映するため、各文字列の間に数バイトを挟み、このバイト数は 1MiB の測定量には含めていない
  • Ryzen 9 7950X の L2 キャッシュはコアあたり 1MiB なので、各テスト実行は L3 キャッシュまで及ぶと予想される
  • 各関数はインライン化やコード移動の干渉を避けるため別々にコンパイルした
    • 実際のコードでは、インライン化を防ぐよりむしろ促進するほうがありそう

結果: tolower64 の滑らかな性能

  • ピンク色の tolower64 は全体としてテスト関数群の中でも 最速クラス に一貫して近い
    • 長さが 65バイトのとき、2本目のベクトルに移るところで少し落ちる
    • すぐに立ち上がり、深い性能の谷がなく、マスク付き load/store が短い文字列断片の処理に有効であることを示している
  • 緑色の copybytes64 は AVX-512 を似た形で使う memcpy バージョン
    • tolower64 より大幅に速いわけではない
    • 最新の Clang はこの関数の意味を認識して完全に書き換えるため、Clang 11 でコンパイルした
  • オレンジ色の copybytes1 はバイト単位の memcpy バージョン
    • Clang 11 でコンパイルした
    • 256バイト未満の文字列断片で Clang 11 の自動ベクトル化ヒューリスティックが相対的に良くないことを示している
  • 赤色の tolower は <ctype.h> の標準 tolower() を呼ぶベースラインで、非常に遅い
  • 紫色の tolower1 は Clang 16 でコンパイルしたバイト単位 tolower()
    • Clang 16 の自動ベクトル化は Clang 11 よりかなり改善している
    • 手書き版より遅く、はるかに複雑なコード を生成する
    • 短い文字列断片の処理が tolower64 ほど良くなく、性能グラフが鋭くギザギザする
  • 茶色の tolower8 は前回の記事の SWAR tolower()
    • Clang は自動ベクトル化を試みるが、関数が複雑なため結果は良くない
    • Clang 16 でコンパイルしたが、Clang 11 風の 256バイト性能崖が現れる
  • 青色の memcpy は glibc の memcpy を呼ぶ
    • 最初は速いが、copybytes64 の半分程度の速度まで落ちる区間がある
    • 原因は確認できていない

結論とコード

  • AVX-512-BW は文字列、とくに 短い文字列 を扱うのに非常によく適している
  • Zen 4 では非常に高速で、intrinsic 関数も比較的使いやすい
  • 最も目立つ特徴は 滑らかな性能
    • 自動ベクトル化が小さな文字列断片でスカラーコードへ切り替わる際に起こりがちな性能の谷がほとんど見られない
  • ARM SVE 対応機や RISC-V Vector extension 対応機に手軽にアクセスできないため、この2つの拡張は詳しく調べられていない
  • コードは Web サイトの git リポジトリ で見られる

関連記事

1件のコメント

 
GN⁺ 2024-07-30
Hacker News の意見

  • 「unsafe read beyond of death」トリックは、ハードウェアで許容されるとしても、Rust と LLVM のメモリモデルでは未定義動作と見なされる
    ほかの未定義動作と同じく、コンパイラは最適化の過程で「そんなことは起きない」と仮定できるため、予想外の結果になる可能性があり、回避するにはインラインアセンブリを使う必要がある
    https://github.com/ogxd/gxhash/issues/82

    • こういう場合にアセンブリ以外の選択肢があるとよい
      「割り当て範囲を超えた値は未指定要素として読み、ハードウェアが嫌がる場合にだけ未定義動作になる」といったロードはサポートが難しくなさそうで、内部的には該当するアセンブリ呼び出しのエイリアスでも十分
      さらに言えば、malloc、スタック、定数など、すべての割り当ての後ろに少なくとも 64 バイト程度の fault しないアドレスが保証されるとよいが、複数の構成要素の協力が必要なのでずっと複雑になる
      カスタムアロケータでは取るに足りないことだが、その場合はカスタムヒープ外のデータには SIMD コードを使いにくく、ごく小さなセグメンテーションフォルトの可能性に縛られることになる
      sanitizer や Valgrind も依然として有用。範囲外の値を未定義値として追跡でき、実際に使われたときにエラーを出せるため
    • ハードウェアレベルでも本当に正しいのか疑問
      マップされていないページや保護されたメモリを読むとどうなるのか気になるし、コードを見ていないので、アラインメント保証がそれを回避しているのかは分からない
    • 「未定義動作ならコンパイラがそんなことはないと仮定できる」という説明は間違っている
      未定義動作は C 標準の専門用語なので、一般化すること自体が妙だし、ANSI C はそのような仮定を明示的に許しておらず、ISO C もより開かれてはいるが、その仮定を具体的に正当化してはいない
      「UB = 起こり得ないと仮定」という説明は、かなり不誠実な恐怖あおりに近いと思う

  • 記事のすっきりしていて高性能なコードを見ると、AMD の AVX512 実装と Intel が予定している AVX10 がどう競争するのか気になる
    AVX10 の核心は Intel の P コア/E コア事情を解決することにあるように見えるが、AMD は状況に応じて Zen5 のフル幅実装や Zen4・Zen5 モバイルの 256 ビット 2 回処理方式を使いながら、API は滑らかに保つという、よりよいアプローチを取っているようだ
    記事の大きな性能向上もすべて Zen4 コアで出た結果であり、AVX512 には多くの利点があるため、Intel がこれを市場セグメンテーションで過度に制限し、汎用クライアントコードでの採用を事実上妨げた点はもどかしい

    • Intel が今後発売するすべての CPU に AVX10/256 を実際に入れるなら、最終的には普及性のために勝つだろう
      市場は CPU ごとのコードパス分岐を繰り返し拒んできたし、実際に重要な SIMD 実装は最小公倍数ならぬ最小共通分母である
      AVX10.1/256 と AVX512VL には共通部分集合があるので、十分な時間が経って大半の CPU が対応するようになれば、人々はその部分をターゲットにするようになるだろう
      AMD は AVX512 対応に更新されたいくつかのベンチマークアプリで楽な勝利を取り続けるだろうが、Intel が AVX10 計画を維持するなら、AMD も最終的には AVX512 互換性を保ちながら AVX10/256 を効率よくサポートするために、2 回処理の SIMD パイプを全面的に使う可能性が高い
      Intel は過去 10 年で悪い選択を多くしてきたが、命令セットで市場を分断したことは最悪の一つだった。最新イノベーションの推進力と関心を自ら殺したようなもので、幅そのものよりもマスク演算のような機能のほうがずっと重要なので、全ラインアップに AVX10/256 を入れてほしい
    • Zen 4 の AVX512 実装は 2 回処理（double-pumped）ではなく、技術記者はそう呼ぶのをやめるべき
      その言葉には特定の意味があり、実際の動作と合っていない
      Zen 4 は ZMM レジスタ演算を複数のマイクロ演算にデコードして空いている 256 ビットユニットにスケジューリングするだけで、512 ビットの全幅シャッフルは高価なエミュレーションを避けるための専用ハードウェアで特別扱いしている
      そのため、4 つの 256 ビット SIMD ユニットを持つ Zen 4 も強力な 2×512 ビットコアのように動作し、この実装は決して安価な方式ではなく、これまでのコンシューマー向けハードウェアでは最良の形である可能性が高い
    • Intel が E コアに 2 回処理 AVX512 を入れてこの問題を解決しない理由が分からない。あるいは、デスクトップには本来そうあるべきように、P コアだけの CPU を作ればよい
      すでに何年も直す時間があったし、AMD が対応しても市場シェアのために採用されない点が腹立たしく、AVX10 は残念ながら Intel が世界をさらに長く足止めすることになりそうだ
      デスクトップでは、より良いコア、より多くのコア、広い SIMD・float16・gather/scatter のような有用な機能を開く、よく標準化された命令セットを見たいし、AMD はかなりうまくやっている
      一方 Intel は、まともなコアの横に弱いコアを付け、弱いコアに合わせるためにまともなコアを制限し、何世代にもわたって同じコア数の CPU を出し、弱いコアでコア数が多いように見せかけ、有用な共通集合が生まれにくいほど多すぎる命令バリエーションを出し、自分たちが有望に見せていた命令対応さえ捨てている
      デスクトップメーカーとしての好みは、90 年代は Intel、2000 年代初頭は AMD、2000 年代後半と 2010 年代は Intel、今はまた AMD だった。Intel が相手を妨害すること以外で再び足場を得るために何をするのか気になるし、競争が続いてこそ一方が慢心しすぎずに済む

  • 面白半分に見るとよい資料: http://www.unicode.org/Public/3.1-Update1/CaseFolding-4.txt

    • こんなものもある: ドイツ語の ßを大文字に変換すると文字列長が変わる
      例えば "straße".upper() は 'STRASSE' になる
      また、ロケールを指定しないと、テュルク語圏の点なし i を大文字/小文字で往復変換する際に 'ı'.upper().lower() が 'i' になって壊れる
    • 幸い、このコードは DNS の作業から出てきたものなので ASCII 専用であり、そのような複雑さを扱う必要はない
      ASCII の大文字小文字を区別しないプロトコルはいくつもあり、多くのサーバーのホットパスで頻繁に現れる
    • 内部的に ID のように使う文字列と、人が入力したテキストは異なる
      前者は通常、8ビットエンコーディングの純粋な ASCII を使えばよいが、後者は複雑になる
      DNS アドレスが分かりやすい例で、技術的にはほぼすべての Unicode を含められるが、実際の DNS 解決では非常に制限された ASCII の部分集合に変換され、その解決過程は大文字小文字を区別しない
      もちろん、Unicode のあらゆる文字体系をサポートしつつ、識別子が大文字小文字を区別しないプログラミング言語のようなものもある。そういうものを扱っているなら、お悔やみ申し上げる
    • ドイツ語の maße が MASSE に変わる例に関連して、ドイツ語には大文字の Eszett である ẞ もある
      まだ広く普及しているわけではなく、対応フォントも少ないが、理論上は今や存在している

  • 記事の「マスク加算」の説明は間違っているのではないかと思う
    is_upper が false のときに加算し、true のときはそのままコピーすべきではないのか疑問

    • ああ、to_upper という変数名が逆で、to_lower と呼ぶべきだったことに後から気づいた
      紛らわしい点を指摘してくれてありがとう。記事とコードを直した
    • この演算は tolower である
      大文字の A は 0x40、小文字は 0x60 なので、0x20 を足すのは is_upper が true のときに起きるべき

  • こうした SWAR 最適化は、文字列が 8バイトアドレスにアラインされている場合にしか役に立たないことが多い
    アラインされていない文字列に SWAR アルゴリズムを適用すると、元のアルゴリズムより遅くなることがよくある
    先頭部分をアライン済みアドレスまで処理し、アラインされた本体を処理し、8バイト未満の末尾を処理するという3段階に分けると、命令数がさらに増える
    Go で utf8.IsValid がより速いという誤った主張と似た事例およびベンチマークがここにある: https://github.com/sugawarayuuta/charcoal/pull/1

    • AVX-512 と ARM SVE の マスク付き SIMD 演算は、その問題を解決するために出てきたもの
      メモリ演算は常にアラインされ、ベクトルサイズ全体で行いつつ、有効な要素だけにマスクをかけられる
      マスク付きベクトルメモリ演算がアラインされておらず、マップされていないページや保護されたページをまたいでいても、その lane がマスクでオフになっていれば fault は発生しない
      strlen() のように長さを事前に知らない演算のために、fault が発生する最初の要素の直前でベクトル長を縮める特殊な load 命令もある

  • マスク加算が格好よく見える。.NET intrinsic で AVX512 のマスクレジスタを直接操作できるとよいのだが、今は「認識されるイディオム」に頼るしかない
    GCC が生成した筆者の中核ループを uiCA（CQA/MAQAO）で Ice Lake 基準に分析すると約 32B/cycle となり、3GHz に換算すると、メモリのボトルネックがないと仮定してほぼ 96GiB/s になる。もちろん、この種のアルゴリズムではメモリアクセスが常にボトルネックになる
    ただし最適活用に非常に近いわけではなさそうで、Clang を使うと、よりうまく展開されたアンロール結果とより良い命令選択により 42.67B/cycle まで出る。L2 キャッシュでもそのスループットを維持するのは難しそうだが、中程度の長さの文字列の大文字・小文字変換が、画面の光が角膜に届くくらいの時間で終わるというのは興味深い
    数か月前に C# で似たような、UTF-8 内の ASCII 大文字・小文字変換を実装した: https://github.com/U8String/U8String/blob/main/Sources/U8Str...
    短い文字列が大半のコードベースを支配しているので、ベクトル化長未満向けのアンロール変換が重要で、switch はジャンプテーブルと branchless fall-through にコンパイルされる
    今は 256 ビットまでしか使っていないが、Zen 3 や 4 のように 256×4 SIMD ユニットしかない場合、すでに飽和するためだ。C 版と並べて比較した例はこちら: https://godbolt.org/z/eTGYhTPan
    AVX512 では vpternlogd で 3 命令の変換も可能そうで、AVX512 ハードウェアを使えたとき、.NET は 256 ビット幅 + AVX512VL でそのように最適化していたが、今は不思議なことに 512 ビット幅では再現できない
    switch dispatch 側の失敗した SWAR の試みも見えるはずだが、記事のライセンスが気になる。テストスイートを通るなら取り込んで使いたい

    • Clang と GCC は intrinsic の扱い方が異なり、とくに AVX-512 命令では、Clang のほうが GCC より Intel ガイドの指定 opcode やアルゴリズムから外れる可能性が高い
      2 つのコンパイラの構造を考えれば理解はできるが、結果が改善になることもあれば損になることもある
      数年前、両方でコンパイルされる必要がある高度にベクトル化されたプロジェクトをやっていて、C の参照版とともに、特定ターゲット向けのインラインアセンブリと .S ファイルをリポジトリに維持することになった
      Makefile がごちゃごちゃになり、テストスイートにベンチマークまで入れる必要があって保守負担が大きかったので、intrinsic を自動ベクトル化より優れた低レベル手段として使うことには非常に慎重であるべきだ、という結論に至った
      例: https://godbolt.org/z/T4Pjhrz5d では GCC の出力は予想どおりだったが、Clang の出力は意外で、実際により遅かった。ループで回すと uiCA 基準で GCC の 4 サイクルに対して 7 サイクルで、この関数が何十億回も実行される brute-force アルゴリズムの実アプリのベンチマークでもそれが現れた
      LLVM コードベースを覗いたところ、Clang 16 が内部リファクタリングのせいで一部のマスク付き AVX-512 命令をそもそも出せない可能性がある、という問題も見た覚えがある
    • 分析がとても有益
      可能な最高性能を狙ったわけではなく、最初は動くかどうかだけを見るつもりだったが、最初の試みがかなり良い結果になったのはおまけだった
      主な関心は ベクトルレジスタより短い文字列と、スループットグラフの谷をなくすことにある
      ブログ記事末尾のコードリンクをたどるとライセンス情報があり、BIND 用に元々書かれた MPL-2.0 部分を除けば 0BSD または MIT-0 だ
    • 大きなアセンブリの塊だけ見ても分かりにくいが、Clang は (x >= 'a' && x <= 'z') を (x - 'a') < ... 形式に書き直して命令を 1 つ減らしている
      奇妙な opcode エンコーディングのために、レジスタ load まで減る場合もある

  • swar が何か分からない

    • “SIMD Within A Register” の略
      通常、1 つのレジスタに複数の項目を packed して、明示的な SIMD 命令なしでも実質的に SIMD のように使う手法を指す
      たとえば 64 ビットレジスタに 31 ビットと 32 ビットの数値を入れ、carry 用に 1 ビットを残しておけば、64 ビット加算 1 回で 2 つの加算ができる
      ゲームでは RGB(A) 値を 32 ビット整数に packing してグラフィックスでこうしたトリックを使ってきたし、ScummVM にも 32 ビット値内の 16 ビット RGB ピクセル 2 つ、合計 6 成分を補間するコードがある: https://github.com/scummvm/scummvm/blob/master/graphics/scal...
    • レジスタ内の SIMD という意味

  • Unicode が登場して以降、大文字と小文字の概念は泥沼になった
    きちんとやるには多くのデータが必要
    ASCII の tolower の実行速度に、時間内に終わるかどうかがかかっている作業をしているなら、何か状況を変えて前提条件を変えたほうがよい

  • 以前、画像の周囲に 黒い枠を入れて、SIMD のバッファ越え読み取り問題を完全に回避したことがある
    非常にうまく動き、速度面で一部の OpenCV 実装に勝てたが、常にそのように入力を完全に制御できるわけではない

  • こういう形で試したのか気になる。自動ベクトル化の結果がかなりきれいに見える
    https://godbolt.org/z/1c5joKK5n

    • それは基本的に tolower1 と同じ。グラフ下の bullet を見ればよい