PythonとRustで比較

だいぶ前だけど、MojoでPythonが高速化できるという話があって、はーなるほどねぇ、ぐらいだったんだけど、ほとんどそのままな感じでRustにしたらどうなるんだっけ？というのをやってなかったのを急に思い出した。

#!/usr/bin/env python3.9
# -*- coding: utf-8 -*-

import time
import numpy as np
from numba import jit, prange

xmin = -1.75
xmax = 0.75
width = 4096
ymin = -1.25
ymax = 1.25
height = 4096
max_iter = 500


@jit(nopython=True, parallel=True)
def mandelbrot_kernel(c):
    z = c
    for i in range(max_iter):
        z = z * z + c
        if abs(z) > 2:
            return i
    return max_iter


@jit(nopython=True, parallel=True)
def compute_mandelbrot(image, xmin, xmax, ymin, ymax, width, height, max_iter):
    dx = (xmax - xmin) / width
    dy = (ymax - ymin) / height

    for j in prange(height):
        for i in range(width):
            y = ymin + j * dy
            x = xmin + i * dx
            image[j, i] = mandelbrot_kernel(complex(x, y))
    return image


image = np.zeros((height, width), dtype=np.int32)
t0 = time.time()
image = compute_mandelbrot(image, xmin, xmax, ymin, ymax, width, height, max_iter)
t1 = time.time()
print(f"elapsed time: {(t1-t0) * 1000:.3f} ms")

(mandelbrot) rifujita@macstudio mandelbrot % ./mandelbrot.py
elapsed time: 119282.660 ms

Pythonでnumbaを使っても119秒かかる。

use num_complex::Complex;
use rayon::prelude::*;
use std::time::Instant;

const XMIN: f64 = -1.75;
const XMAX: f64 = 0.75;
const WIDTH: usize = 4096;
const YMIN: f64 = -1.25;
const YMAX: f64 = 1.25;
const HEIGHT: usize = 4096;
const MAX_ITER: usize = 500;

fn mandelbrot_kernel(c: Complex<f64>) -> usize {
    let mut z = c;
    for i in 0..MAX_ITER {
        z = z * z + c;
        if z.norm() > 2.0 {
            return i;
        }
    }
    MAX_ITER
}

fn compute_mandelbrot(image: &mut Vec<Vec<usize>>) {
    let dx = (XMAX - XMIN) / WIDTH as f64;
    let dy = (YMAX - YMIN) / HEIGHT as f64;

    image.par_iter_mut().enumerate().for_each(|(j, row)| {
        for i in 0..WIDTH {
            let y = YMIN + j as f64 * dy;
            let x = XMIN + i as f64 * dx;
            row[i] = mandelbrot_kernel(Complex::new(x, y));
        }
    });
}

fn main() {
    let mut image = vec![vec![0; WIDTH]; HEIGHT];
    let start = Instant::now();
    compute_mandelbrot(&mut image);
    let duration = start.elapsed();
    println!("elapsed time: {:.3} ms", duration.as_secs_f64() * 1000.0);
}

rifujita@macstudio mandelbrot % cargo run --release
    Finished `release` profile [optimized] target(s) in 0.04s
     Running `target/release/mandelbrot`
elapsed time: 8239.660 ms

8.2秒だった。14.48倍ぐらいしか速くならなかった。もっと真面目にやらないとダメか。

関連