CUDA libraryのexp(x)を読んだ

なんだか指数関数の実装ばかり見ているような気がしますが、GPGPU向けの実装はどうなっているんだろうなと思ったので見てみました。

Programming Guide :: CUDA Toolkit Documentationによれば、最大誤差1ULPを保証する実装になっているようです。 ULPの定義について詳しく語った論文を引用しているあたり、かなり丁寧なつくりになっていそうです。

中身を見てみたところ、11次の多項式近似をホーナー法で求め指数部をビット演算で操作する、というよくある感じの実装でした。

この方法でうまくいかない入力だった場合の例外処理が単精度浮動小数点数を使って実装されており面白かったです。 これは、倍精度浮動小数点数演算器をなるべく使わないような実装をしているためだと考えられます。 倍精度指数関数を計算する場合、倍精度浮動小数点数演算が大量に出てきますが、GPU上の倍精度浮動小数点数演算器は貴重なので、例外処理なんかに使いたくないということでしょう。

指数関数の入力は、大きく分けて以下の5つに分類されます*1

  • 小さすぎてexp(x)が0.0になるような入力(-Infを含む)
  • exp(x)が非正規化数になるような入力
  • exp(x)が正規化数になるような入力
  • 大きすぎてexp(x)が+Infになるような入力(+Infを含む)
  • NaN(exp(NaN)=NaN)

この場合分けを、入力倍精度浮動小数点数の上位32bitを単精度浮動小数点数とみなして絶対値をとってから特定の定数と比較することで行っています。

「倍精度浮動小数点数の上位32bitを単精度浮動小数点数とみなして」というのがいきなり意味不明ですが、これはGPUで倍精度浮動小数点数を持つときのフォーマットを利用した技です。 GPU上で倍精度浮動小数点数を持つときには、隣り合う*232bitレジスタ二つに入れているようなので、これらのレジスタのうち片方を命令オペランドに指定することで、簡単に上位32bitを単精度浮動小数点数として扱えるということです。

倍精度浮動小数点数と単精度浮動小数点数だと指数部のbit数が違う点が気になるかもしれませんが、比較の両辺が正(+0.0含む、NaNを含まない)であればただの辞書順比較と同じなので大丈夫です。

そんなわけで、「元の倍精度浮動小数点数入力の上位32bitを単精度浮動小数点数とみなした時、その絶対値が特定の定数より小さければ、必ずexp(x)が正規化数になる」ように、「特定の定数」が選ばれているようです。 実用上は大半の入力でこの条件が満たされるはずで、その場合にearly returnすれば追加の計算コストをほとんど払わずに済みます(この場合、単精度浮動小数点数の絶対値演算と比較演算がそれぞれ一回ずつで済みました)。

また、『元の倍精度浮動小数点数入力の上位32bitを単精度浮動小数点数とみなした時、その絶対値が別の特定の定数より大きい、またはその単精度浮動小数点数がNaNである*3場合、「小さすぎてexp(x)が0.0」「大きすぎてexp(x)が+Inf」「入力がNaNなのでexp(x)もNaN」のどれか』になるように「別の特定の定数」が選ばれているようです。 このようなケースは例外ケースの中では頻出なので、ここも最適化テクニックで高速化されていました。 具体的には、「exp(x)=x+(+Inf)を返す」という方法で「大きすぎてexp(x)が+Inf」と「入力がNaNなのでexp(x)がNaN」のケースをまとめて処理していました。どんな有限のxに対してもx+(+Inf)は+Infですし、xがNaNであればx+(+Inf)はNaNです*4

それ以外のパターン、つまり「exp(x)が非正規化数になるような入力x」に加えて、単精度浮動小数点数を使ったがために厳密な場合分けができずに漏れてしまった「exp(x)がかなり大きな正規化数になるような入力x」「exp(x)がかなり小さな正規化数になるような入力x」「exp(x)が最も小さな非正規化数よりもやや小さくて0.0になるような入力x」の場合は、指数部のビット演算を工夫した後、浮動小数点数演算で最終結果を算出していました。 これらのパターンが使われることはほとんどないですが、演算の量はそれなりに増えてしまうようです。

*1:実装アルゴリズムによっては、途中の計算結果がオーバーフローするが答えは正規化数になる、といったケースの例外処理が必要なこともあり得ます。

*2:論理レジスタ番号が2nと2n+1のレジスタという意味です。

*3:このような比較はunordered比較と言って一命令で実現可能です。

*4:xが-Infだとx+(+Inf)はNaNになりますが、まずxが負であれば0.0を返すというアルゴリズムなので、この計算結果が選ばれることはありません。