massivelyを用いたモンテカルロ法による円周率計算

現在おれは、GPU並列アルゴリズムライブラリ「massively」をAIと一緒に作っており、その性能上のベンチマークとしてモンテカルロ法による円周率計算を設定している。説明不要と思うが、概要は以下である。

「乱数を2つ生成し、ダーツを投げる。それが円の中に入った本数によって円周率が計算出来る」

つまり、計算は3つに分解出来る。

乱数列を生成する（map）
それが円に入ってるか判定する（map）
円に入ってる本数を集計する（reduce）

さらにGPUの高速化として、1の結果をメモリに書き、2はそれを読んでから計算して結果をメモリに書き・・・とやってると帯域ネックとなって確実死ぬ（特におれが実験に使ってるのはiGPUなので死ぬ）ので、1と2をカーネル上の遅延計算とするようないわゆる「遅延イテレータ」の仕組みが必要になる。

というわけで概ね、

遅延イテレータの実装
効率的なreduceアルゴリズムの実装

が試されるというわけで、円周率計算は良いベンチマークといえるのである。

もう１つ、おれが円周率計算をベンチマークに選んだ理由はRayonをベースとしたCPU並列化をかなりハードに使った先行報告が存在したからである。しかもノートでやってるというのも、おれもAMDのiGPU（680M）を積んだノートでやろうとしているから公平性があると考えた。そこで今回はこいつが達成した2600億回のダーツ計算を「14秒」を超えることを目標にした。

→ 2600億回のモンテカルロを14秒で：Rust×rayon×SIMDで並列化してみた

以下が、massivelyを使って書いた円周率計算のコードである。massivelyを使ってGPUを使って計算するということが主眼なので、出来るだけ素直なコードを書くように努めた。「関数型的なコードを素直に書くとそれで速い」というのがmassivelyの売りだからである。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
fn main() -> Result<(), massively::Error> {
    let args = Args::parse();

    let n = args.n.get() as MIndex;
    let m = args.m.get() as MIndex;

    // Setup
    let exec = {
        // let device = cubecl::wgpu::WgpuDevice::Cpu;
        // let device = cubecl::wgpu::WgpuDevice::IntegratedGpu(0);
        // let device = cubecl::wgpu::WgpuDevice::DiscreteGpu(0);
        let device = cubecl::wgpu::WgpuDevice::DefaultDevice;
        // let device = cubecl::hip::device::AmdDevice::new(0);
        // let device = cubecl::cuda::CudaDevice::new(0);

        Executor::<cubecl::wgpu::WgpuRuntime>::new(device)
        // massively::Executor::<cubecl::hip::HipRuntime>::new(device)
        // massively::Executor::<cubecl::cuda::CudaRuntime>::new(device)
    };

    let mut sum_pi = 0.;

    for i in 0..n {
        let seed = i as u64 * 2;

        let x = massively::util::random::uniform_f32(0.0, 1.0, seed)?.take(m);
        let y = massively::util::random::uniform_f32(0.0, 1.0, seed+1)?.take(m);

        let hits = massively::lazy::transform(zip2(x,y), DetectHit);

        // Count the number of ones.
        let n_hits = massively::reduce(&exec, hits, 0_u32, CountHit)?;

        let pi = (n_hits as f64 / m as f64) * 4.;
        dbg!(pi);
        sum_pi += pi;
    }

    println!("pi={}", sum_pi / n as f64);
    Ok(())
}

実験では、１ループあたり4G（40億）個のダーツを投げて、それを65回繰り返している。結果としては複数回やって大体9-10秒と、CPU版をわずかに超えることが出来た。これでもGPT5.5の段階では超えることが出来ず、GPT5.6 Solを使って一日かけてライブラリのリワーク（ユーザのコードを型で表現し、それをもとにしてカーネルコードへのディスパッチをぎりぎりまで遅延するという再設計を行った）を行い、その後最適化するとようやく超えることが出来た。

現在のボトルネックはどこなのかというと、大体1ループあたり140msかかってる中、乱数生成120ms + 円内判定10ms + reduce10msくらいであり、乱数生成がネックとなってることがわかる。この乱数生成はPCG32というわりと高速＆高品質な実装を使っており、これ以上高速化するには乱数の品質を落とす他ないという結論になった。

現状計算ネックなので、現在の最新GPUを使えば10倍以上の性能を出すことは出来るだろうが、円周率を速く計算することが目的ではないので終わりにする。

→ akiradeveloper/pi-monte-carlo