はじめに
Pythonは世界的にも人気のあるプログラミング言語ですが、実行速度については課題があります。Pythonの実行速度を高速化したい、という要求は根強く、これまでにも様々な処理系が開発されています。
この記事はPythonで書かれたコードを35000倍に高速化するにはどのような方法があるかについてまとめたものです。
この記事は：
- Pythonで書かれたアルゴリズムを35000倍に高速化する
- 事前コンパイル、並列化、SIMD演算を駆使する
- 最終的に44000倍まで高速化できた
なぜ35000倍？
2023年5月2日にModular社よりPythonの使いやすさとC言語の性能を兼ね備える新しいプログラミング言語、Mojoの開発について発表がありました。低レベルのハードウェア向けにコンパイル可能なこと、文法的にはPythonを踏襲しており、既存のPythonライブラリを利用可能であることが特徴です。
Mojoは2023年5月現在、まだ言語自体は公開されておらず、waitlistに登録することでMojo Playgroundと呼ばれるホスト型開発環境にアクセス可能になるようです。
エンジニア界隈で特に話題になっているのが、「Pythonより35000倍速い」というMojoのパフォーマンスです。公式サイトによれば、Pythonで1000秒かかる処理が0.03秒で計算できるということで、確かに35000倍になっているようです。
PythonとMojoの速度比較 (https://www.modular.com/mojo)
一般にPythonとCの実行速度は100倍程度に留まるため、35000倍というのはまさに桁違いの性能です。
しかし、ネイティブコードより速くなることが原理的にありえるのか？と疑問に感じ、どのような条件で35000倍になるかを調べたところ、公式サイトに以下の記述を見つけました。
Algorithm Mandelbrot
Instance AWS r7iz.metal-16xl Intel Xeon
どうやらマンデルブロ集合の計算を、マルチコアCPUで並列計算することで高速化しているようです。実行環境のAWS r7izはAWSのハイエンドインスタンスで、全コア最大ターボ周波数が 3.9 GHz の第 4 世代 Intel Xeon スケーラブルプロセッサを搭載したものです。r7iz.metal-16xlはvCPUが128個の最上位サーバで、まさに並列処理による高速化に最適です。
こういったマルチコア環境で並列化を駆使すれば、確かに35000倍の高速化ができそうです。
この記事では、Mojo以外の方法でもこのような高速化が実現できるかを調べてみたので、順番に紹介していきたいと思います。
Mandelbrotの描画を実装する
Mandelbrotの描画アルゴリズムは、フラクタル幾何学の分野で広く使用されるアルゴリズムの1つで、Mandelbrot集合を生成するために使用されます。
実際には複素数を反復的に評価することでMandelbrot集合を生成します。具体的には、与えられた複素数zについて、zの二乗に定数cを加えた結果を順次計算します。zがある値に収束するか、ある閾値を超えるまで、この計算を繰り返します。
Mandelbrot集合
Python
まずはMandelbrotアルゴリズムをPythonで実装していきます。
import time
xmin = -1.75
xmax = 0.75
width = 4096
ymin = -1.25
ymax = 1.25
height = 4096
max_iter = 500
def mandelbrot_kernel(c):
z = c
for i in range(max_iter):
z = z * z + c
# zが閾値を超えたら終了します
if abs(z) > 2:
return i
return max_iter
def compute_mandelbrot(image):
dx = (xmax - xmin) / width
dy = (ymax - ymin) / height
# 各ピクセルごとに複素数を計算します
for j in range(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 = [[0 for _ in range(width)] for _ in range(height)]
t0 = time.time()
image = compute_mandelbrot(image)
t1 = time.time()
print(f"elapsed time: {(t1-t0) * 1000:.3f} ms")
ここでx, yは複素数zの実部と虚部、width, heightは描画するイメージのピクセルサイズに相当します。max_iterは各ピクセルの計算の繰り返し数の上限値を表します。
こちらのコードを実際に実行してみると、計算が完了するまで283秒かかりました。これをベースラインとして高速化していきます。
$ python mandelbrot.py
elapsed time: 283104.711 ms
なお、実行環境はPython3.10.10 / Ubuntu20.04 / GCP c3-highcpu-176です。
PyPy
PyPyは、Python言語の高速実行環境の1つであり、JITコンパイラによってPythonコードを高速化します。PyPyは、一般的なCPython実行環境と互換性があり、Pythonの標準ライブラリやサードパーティライブラリと互換性があります。CPythonよりも高速なため、競技プログラミングの分野でも人気のようです。
PyPyは上記のPythonコードをそのまま実行できます。
結果は35秒でPythonより8.2倍高速になりました。
$ pypy3 mandelbrot.py
elapsed time: 34645.686 ms
Codon
Codonは高性能なPythonコンパイラです。実行時のオーバーヘッドなしにPythonコードをネイティブなマシンコードにコンパイルし、シングルスレッドで10-100倍以上の高速化が実現できます。Codonの開発はGithub上で行われており、2021年頃から現在まで様々な機能追加が行われています。
CodonはPythonコードを事前コンパイルして実行します。こちらも上記のPythonコードをそのまま実行できます。
結果は15秒で、18倍高速になりました。
$ codon run -release mandelbrot.py
elapsed time: 14979.343 ms
C++
しかし、まだ35000倍には程遠いようです。そこでPythonに見切りをつけて、C++でより高速化を図ります。C++の実装は次のようになります。
#include <chrono>
#include <cmath>
#include <iostream>
#include <vector>
constexpr double xmin = -1.75f;
constexpr double xmax = 0.75f;
constexpr int width = 4096;
constexpr double ymin = -1.25f;
constexpr double ymax = 1.25f;
constexpr int height = 4096;
constexpr int max_iter = 500;
class Complex {
public:
double real;
double imag;
Complex(double real, double imag) : real(real), imag(imag) {}
Complex operator+(const Complex& other) const {
return Complex(this->real + other.real, this->imag + other.imag);
}
Complex operator*(const Complex& other) const {
return Complex(this->real * other.real - this->imag * other.imag,
this->real * other.imag + this->imag * other.real);
}
double abs() const {
return std::sqrt(this->real * this->real + this->imag * this->imag);
}
};
int mandelbrot_kernel(const Complex& c) {
Complex z = c;
for (int i = 0; i < max_iter; i++) {
z = z * z + c;
if (z.abs() > 2) {
return i;
}
}
return max_iter;
}
std::vector<std::vector<int>> compute_mandelbrot(
std::vector<std::vector<int>>& image) {
double dx = (xmax - xmin) / width;
double dy = (ymax - ymin) / height;
for (int j = 0; j < height; j++) {
for (int i = 0; i < width; i++) {
double y = ymin + j * dy;
double x = xmin + i * dx;
image[j][i] = mandelbrot_kernel(Complex(x, y));
}
}
return image;
}
int main() {
std::vector<std::vector<int>> image(height, std::vector<int>(width));
auto t0 = std::chrono::high_resolution_clock::now();
image = compute_mandelbrot(image);
auto t1 = std::chrono::high_resolution_clock::now();
auto elapsed_time =
std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count();
std::cout << "elapsed time: " << double(elapsed_time) / 1000 << " ms"
<< std::endl;
return 0;
}
基本的な処理の流れはPythonと同様です。これを-Ofastでコンパイルして実行していきます。
実行速度は6.7秒、42倍の高速化です。どうもこのあたりがシングルスレッドによる限界のようです。
$ g++ -Ofast mandelbrot.cpp -o mandelbrot
$ ./mandelbrot
elapsed time: 6670.19 ms
C++ (並列化)
さて、いよいよ並列処理による高速化に入っていきます。
OpenMPを用いて並列処理を実装します。#pragma omp parallel forを並列化したい箇所に記載します。
#include <omp.h>
// ...（省略）... //
// この1行を追加します
#pragma omp parallel for collapse(2)
for (int j = 0; j < height; j++) {
for (int i = 0; i < width; i++) {
double y = ymin + j * dy;
double x = xmin + i * dx;
image[j][i] = mandelbrot_kernel(Complex(x, y));
}
}
collapseはOpenMPのオプションで、いくつの内部ループまで並列化するかを指定できます。（この場合は2重ループまで）
これを128スレッドで実行すると0.14秒、一気にPythonの2000倍まで速くなりました。
$ g++ -Ofast -fopenmp mandelbrot_omp.cpp -o mandelbrot
$ OMP_NUM_THREADS=128 ./mandelbrot
elapsed time: 139.868 ms
C++ (並列化 + SIMD)
128スレッドを使って並列処理しているのに、まだ足りないようです。そこで、CPUの性能を引き出すためSIMDを利用します。
SIMD（Single Instruction, Multiple Data）は、複数のデータに対して同じ命令を同時に実行することができるコンピュータアーキテクチャの一種です。これにより、データの並列処理を効率的に行うことができます。
現在のCPUには、SIMD演算をサポートする命令セットが多数用意されています。代表的なものとして、SSE（Streaming SIMD Extensions）、AVX（Advanced Vector Extensions）がありますが、今回はAVX2命令を使います。
次がSIMD演算を用いた実装です。非常にざっくりいうと__m256 にfloat型の変数を8個格納し、これらを一度に計算することで高速化しています。
#include <immintrin.h>
#include <omp.h>
#include <chrono>
#include <iostream>
#include <vector>
constexpr float xmin = -1.75f;
constexpr float xmax = 0.75f;
constexpr int width = 4096;
constexpr float ymin = -1.25f;
constexpr float ymax = 1.25f;
constexpr int height = 4096;
constexpr int max_iter = 500;
class ComplexSIMD {
public:
__m256 real;
__m256 imag;
ComplexSIMD(__m256 real, __m256 imag) : real(real), imag(imag) {}
ComplexSIMD operator+(const ComplexSIMD& other) const {
return ComplexSIMD(_mm256_add_ps(this->real, other.real),
_mm256_add_ps(this->imag, other.imag));
}
ComplexSIMD operator*(const ComplexSIMD& other) const {
__m256 real_part = _mm256_sub_ps(_mm256_mul_ps(this->real, other.real),
_mm256_mul_ps(this->imag, other.imag));
__m256 imag_part = _mm256_add_ps(_mm256_mul_ps(this->real, other.imag),
_mm256_mul_ps(this->imag, other.real));
return ComplexSIMD(real_part, imag_part);
}
__m256 abs() const {
return _mm256_sqrt_ps(_mm256_add_ps(_mm256_mul_ps(this->real, this->real),
_mm256_mul_ps(this->imag, this->imag)));
}
};
__m256i mandelbrot_kernel_simd(const ComplexSIMD& c) {
ComplexSIMD z = c;
__m256i iter_count = _mm256_setzero_si256();
for (int i = 0; i < max_iter; i++) {
z = z * z + c;
__m256 abs_z = z.abs();
__m256 mask = _mm256_cmp_ps(abs_z, _mm256_set1_ps(2.0f), _CMP_LE_OQ);
__m256i mask_int = _mm256_castps_si256(mask);
iter_count = _mm256_add_epi32(iter_count, mask_int);
if (_mm256_testz_si256(mask_int, mask_int)) {
break;
}
}
return iter_count;
}
std::vector<std::vector<int>> compute_mandelbrot_simd(
std::vector<std::vector<int>>& image) {
float dx = (xmax - xmin) / width;
float dy = (ymax - ymin) / height;
#pragma omp parallel for collapse(2)
for (int j = 0; j < height; j++) {
for (int i = 0; i < width; i += 8) {
float y = ymin + j * dy;
float x[8] = {xmin + i * dx, xmin + (i + 1) * dx,
xmin + (i + 2) * dx, xmin + (i + 3) * dx,
xmin + (i + 4) * dx, xmin + (i + 5) * dx,
xmin + (i + 6) * dx, xmin + (i + 7) * dx};
__m256 c_real = _mm256_loadu_ps(x);
__m256 c_imag = _mm256_set1_ps(y);
ComplexSIMD c(c_real, c_imag);
__m256i iter_counts = mandelbrot_kernel_simd(c);
int result[8];
_mm256_storeu_si256(reinterpret_cast<__m256i*>(result), iter_counts);
for (int k = 0; k < 8; k++) {
image[j][i + k] = result[k];
}
}
}
return image;
}
int main() {
std::vector<std::vector<int>> image(height, std::vector<int>(width));
auto t0 = std::chrono::high_resolution_clock::now();
image = compute_mandelbrot_simd(image);
auto t1 = std::chrono::high_resolution_clock::now();
auto elapsed_time =
std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count();
std::cout << "elapsed time: " << double(elapsed_time) / 1000 << " ms"
<< std::endl;
return 0;
}
SIMDによる実行速度は0.057秒、Pythonに比べて5200倍の速度になりました。
$ g++ -Ofast -fopenmp -march=native -mavx2 mandelbrot_simd.cpp -o mandelbrot
$ OMP_NUM_THREADS=128 ./mandelbrot
elapsed time: 56.808 ms
Pythonを遅くする
並列化やSIMDなどを駆使し最大限高速化してきましたが、目標の35000倍には届きませんでした。ここまで来たからには何としてでも35000倍を実現したいところです。
逆転の発想として、C++を高速化するのではなく、比較対象のPythonを遅くすればどうでしょう？
MojoのMandelbrot.ipynbをよく読むと、どうやら複素数の構造体を独自定義しているようです。
@register_passable("trivial")
struct Complex:
var real: F32
var imag: F32
fn __init__(real: F32, imag: F32) -> Self:
return Self {real: real, imag: imag}
fn __add__(lhs, rhs: Self) -> Self:
return Self(lhs.real + rhs.real, lhs.imag + rhs.imag)
fn __mul__(lhs, rhs: Self) -> Self:
return Self(
lhs.real * rhs.real - lhs.imag * rhs.imag,
lhs.real * rhs.imag + lhs.imag * rhs.real,
)
fn norm(self) -> F32:
return self.real * self.real + self.imag * self.imag
これは完全な推測ですが、Mojoと比較したPythonコードも、比較のためにこのようなComplexクラスを実装しているのではないかと想像されます。（比較元のPythonコードは公開されていないようです）
そこで、Python組み込みのcomplexではなく、独自にComplexクラスを実装してみます。Pythonは独自クラスのインスタンスの生成がかなり遅いため、実行速度に影響を与えると予想されます。
Python (Complexクラス実装)
class Complex(object):
real: float
imag: float
def __init__(self, real: float, imag: float):
self.real = real
self.imag = imag
def __add__(self, other):
return Complex(self.real + other.real, self.imag + other.imag)
def __mul__(self, other):
return Complex(
self.real * other.real - self.imag * other.imag,
self.real * other.imag + self.imag * other.real,
)
def __abs__(self):
return (self.real * self.real + self.imag * self.imag) ** 0.5
これを組み込みの complex の代わりに用います。
# image[j][i] = mandelbrot_kernel(complex(x, y))
image[j][i] = mandelbrot_kernel(Complex(x, y))
実行結果は、2501秒と、（予想通り）組み込みクラスを用いた実装と比較して約1/9の速度になりました。
$ python mandelbrot_complex.py
elapsed time: 2505104.201 ms
Pythonを遅くするのは本末転倒な気もしますが、とにかくこのPython実装（2505秒）と最速のC++の実行時間（0.057秒）を比較すると、最終的に44000倍の高速化が実現できたことになります。
結果
Complexクラスを実装したPythonコードをベースに再測定した最終的な結果がこちらです。なお、PyPyとCodonについてはそれぞれComplexクラスを実装したコードで再計測しました。
| 手法 | 時間 | vs Python |
|---|---|---|
| Python | 2505 秒 | x1 |
| Python (組み込みcomplex) | 283秒 | x9 |
| PyPy | 47 秒 | x53 |
| Codon | 7.9 秒 | x320 |
| Codon (並列化) | 0.34 秒 | x7300 |
| C++ | 6.7 秒 | x380 |
| C++ (並列化) | 0.14 秒 | x18000 |
| C++ (並列化 + SIMD) | 0.057 秒 | x44000 |
C++の並列化とSIMD演算を用いることで、Pythonと比べて大幅な高速化ができることがわかりました。Mojoはこのような苦労をせずとも35000倍の高速化ができるということで、非常にポテンシャルの高いプログラミング言語であるのは間違いなさそうです。
おわりに
この記事ではPythonより数万倍高速化するための方法について紹介しました。Mojoは「Pythonの35000倍速い」という部分がフォーカスされがちですが、速度に関しては（おそらく）他の言語でも同様の実行速度を実現することが可能かと思います。
個人的には、Mojoの本質はMLIR向けに設計・開発される最初のプログラミング言語であるという点です。MLIR（Multi-Level Intermediate Representation）は、コンパイラの中間表現として使用されることを目的として設計され、ハードウェアアーキテクチャ、プログラム言語、および最適化技術を組み合わせて、高速かつ効率的なコンパイルを実現することを目的としています。
現在の機械学習モデルの多くはPythonで学習され、速度が要求される推論時はC++など他の高速なプログラミング言語で使用されることがほとんどです。Mojo/MLIRを用いることで、学習からモデルの最適化、推論までをアーキテクチャによらず単一の処理系で一貫して実現できる可能性があります。Mojoは将来、機械学習モデルを扱う上で、最も重要なフレームワークの一つになるかもしれません。
最後に、筆者の所属するTuringでは、自動運転を実現するためのシステム開発を行っており、このような高速な処理をAIハードウェアに組み込む技術についても着目しています。Turingは完全自動運転EVの開発を目指すスタートアップで、機械学習から3Dアプリケーション、車体の制作に至るまで、ソフトウェア・ハードウェアを横断的なエンジニアリングに取り組んでいます。
もし興味のある方は私のTwitterのDMにてご連絡いただくか、採用ページで「スカウトを待つ」にぜひ登録をお願いします！
おまけ (75万倍高速化)
ちなみに、このような並列処理による高速化はGPUで計算するのが圧倒的に高速です。C++とCUDAを用いてMandelbrotを計算する例を紹介しておきます。
#include <cuda_runtime.h>
#include <iostream>
#include <vector>
#include <chrono>
constexpr float xmin = -1.75f;
constexpr float xmax = 0.75f;
constexpr int width = 4096;
constexpr float ymin = -1.25f;
constexpr float ymax = 1.25f;
constexpr int height = 4096;
constexpr int max_iter = 500;
__device__ int mandelbrot_kernel(float x, float y) {
float real = x;
float imag = y;
float real_sq = real * real;
float imag_sq = imag * imag;
int i;
for (i = 0; i < max_iter && real_sq + imag_sq <= 4.0f; ++i) {
imag = 2 * real * imag + y;
real = real_sq - imag_sq + x;
real_sq = real * real;
imag_sq = imag * imag;
}
return i;
}
__global__ void compute_mandelbrot(int *image) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
int j = blockIdx.y * blockDim.y + threadIdx.y;
if (i < width && j < height) {
float x = xmin + i * (xmax - xmin) / width;
float y = ymin + j * (ymax - ymin) / height;
image[j * width + i] = mandelbrot_kernel(x, y);
}
}
int main() {
int *image;
cudaMallocManaged(&image, width * height * sizeof(int));
dim3 blockSize(32, 32);
dim3 gridSize((width + blockSize.x - 1) / blockSize.x, (height + blockSize.y - 1) / blockSize.y);
auto t0 = std::chrono::high_resolution_clock::now();
compute_mandelbrot<<<gridSize, blockSize>>>(image);
cudaDeviceSynchronize();
auto t1 = std::chrono::high_resolution_clock::now();
auto elapsed_time =
std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count();
std::cout << "elapsed time: " << double(elapsed_time) / 1000 << " ms"
<< std::endl;
cudaFree(image);
return 0;
}
CUDAによる実行時間はRTX3090で0.003秒と、Pythonに比べて75万倍高速化されました。
$ nvcc -std=c++17 -arch=sm_86 mandelbrot.cu -o mandelbrot
$ ./mandelbrot
elapsed time: 3.324 ms
Discussion