glow/flow で, 正規分布のサンプリング(normal()
. https://pytorch.org/docs/stable/distributions.html )をよく使うので, CPU で実装がどうなっているか調べました.
ATen の DistributionTemplates.h に実装があります.
(Torch/ATen はいろいろ階層が深くてコード追うのがめんどうですね...)
16 個以上で並列化できるときと, scalar でコードパスが違います.
アルゴリズムはどちらも普通の Box-Muller 法(https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform)です.
16 並列生成版
16 個で sin, cos で半分づつ(8 こづつ)に分けて生成しています. AVX2 の場合は, 明示的な AVX instrinsics 関数を使って並列化しています(e.g. _mm_sincos256_ps など). SSE2 のコードパスが無いのは, C++ コンパイラが賢ければ自動で SSE2 命令で置き換えてくれる...はずだからでしょうか.
scalar 版
at::normal_distribution
で実装されています.
サンプルした値を cache しておいて, cache があればそれを使う, なければあらたに生成する, というアプローチをとっています. ちょっと実装が面倒ですね.
テスト
自前で normal distribution generator を実装したときなど, それが正しく正規分布しているか確かめる.
scipy に normaltest 関数があります.
その他統計テスト用関数
おまけ
もうひとつ乱数を生成して, それを元に cos を使うか sin を使うか判定に使うというのもあります.
GPU やアクセラレータで cache を持つのが面倒ときに使えますね.