「ZigはCよりも速いです。」をVector/SIMD最適化で検証する

「パフォーマンスについて言えば、ZigはCよりも速いです。」

この一文は Zig 言語の公式サイトの詳細な概要に実際に書いてある文章です。
その理由としていくつか挙げた上で最後にちらっと「Zig は SIMD ベクトル型を直接公開しており、移植性の高いベクトル化コードを簡単に記述することができます。」と書いてあります。今回の記事では画像処理っぽいお題を使って実際に SIMD 最適化される様子を検証してみます。

ベンチマーク環境

OS: Ubuntu20.04 on WSL2 on Windows11
CPU: AMD Ryzen 9 5900X 12-Core Processor

お題

一般的な画像データ構造の表現である、幅×高さ×4(RGBA)の長さを持つ8bit整数の配列を受け取り、in-placeでRGBの値を 10 加算して明るくするという関数を考えます。Aの値は元のままとします。

Zig で普通に書いてみる

brighten.zig

const std = @import("std");
const print = std.debug.print;
const microTimestamp = std.time.microTimestamp;

fn brighten(image: []u8) void {
    var idx: usize = 0;
    while (idx < image.len) : (idx += 4) {
        image[idx + 0] +|= 10;
        image[idx + 1] +|= 10;
        image[idx + 2] +|= 10;
    }
}

fn benchmark(comptime width: usize, comptime height: usize, comptime num_loop: usize, comptime func: fn (image: []u8) void, comptime name: []const u8) void {
    const size = width * height * 4;
    var image: [size]u8 = [_]u8{0} ** size;
    var cnt: usize = 0;
    const start = microTimestamp();
    while (cnt < num_loop) : (cnt += 1) {
        func(&image);
    }
    const stop = microTimestamp();
    const duration = @intToFloat(f64, stop - start);
    const call_duration = duration / @intToFloat(f64, num_loop);
    print("{s}: {d} times loop took {d}us us/call={d}us\n", .{ name, num_loop, duration, call_duration });
}

pub fn main() !void {
    const width = 1920;
    const height = 1080;
    const num_loop = 1000;
    benchmark(width, height, num_loop, brighten, "brighten");
}

brighten で見慣れない演算子を使ってますが、これは Zig ならではの飽和加算代入演算子です。こう書くだけでオーバーフローしても255に飽和させてくれます。

benchmark 関数を用意してベンチマークを取っています。画像処理関数を引数で取ることでこの後の改造に対応させています。

main ではベンチマークを 1920×1080の画像に対して1000回ループさせています。

ReleaseFastでビルドして実行します。

$ zig build-exe brighten.zig -O ReleaseFast && ./brighten
brighten: 1000 times loop took 1495369us us/call=1495.369us

FullHD画像1枚に対して1回あたりおよそ1.5msという結果が出ました。

@Vector を使う

@Vector(4, u8) を使ってRGBAを一気に処理するコードにしてみましょう。

brighten_simd.zig

const std = @import("std");
const print = std.debug.print;
const microTimestamp = std.time.microTimestamp;

fn brighten(image: []u8) void {
    var idx: usize = 0;
    while (idx < image.len) : (idx += 4) {
        image[idx + 0] +|= 10;
        image[idx + 1] +|= 10;
        image[idx + 2] +|= 10;
    }
}

fn brightenSIMD(image: []u8) void {
    var idx: usize = 0;
    const add_vec = @Vector(4, u8){ 10, 10, 10, 0 };
    while (idx < image.len) : (idx += 4) {
        const image_vec = @Vector(4, u8){ image[idx + 0], image[idx + 1], image[idx + 2], image[idx + 3] };
        const result_vec = image_vec +| add_vec;
        image[idx + 0] = result_vec[0];
        image[idx + 1] = result_vec[1];
        image[idx + 2] = result_vec[2];
        image[idx + 3] = result_vec[3]; // ここを消すと遅くなる
    }
}

fn benchmark(comptime width: usize, comptime height: usize, comptime num_loop: usize, comptime func: fn (image: []u8) void, comptime name: []const u8) void {
    const size = width * height * 4;
    var image: [size]u8 = [_]u8{0} ** size;
    var cnt: usize = 0;
    const start = microTimestamp();
    while (cnt < num_loop) : (cnt += 1) {
        func(&image);
    }
    const stop = microTimestamp();
    const duration = @intToFloat(f64, stop - start);
    const call_duration = duration / @intToFloat(f64, num_loop);
    print("{s}: {d} times loop took {d}us us/call={d}us\n", .{ name, num_loop, duration, call_duration });
}

pub fn main() !void {
    const width = 1920;
    const height = 1080;
    const num_loop = 1000;
    benchmark(width, height, num_loop, brighten, "brighten");
    benchmark(width, height, num_loop, brightenSIMD, "brightenSIMD");
}

$ zig build-exe brighten_simd.zig -O ReleaseFast && ./brighten_simd
brighten: 1000 times loop took 1492283us us/call=1492.283us
brightenSIMD: 1000 times loop took 462711us us/call=462.711us

3倍速くなりました。従来RGBそれぞれで飽和加算していたところがSIMDで一気に計算できるようになったので、3倍というのは妥当な数値です。

brightenSIMD() の中の「ここを消すと遅くなる」という行は、RGBAのAの値を代入しているところです。ここは値は変わらないのだから代入しなくても良いかと思うかもしれません。やってみましょう。

$ zig build-exe brighten_simd.zig -O ReleaseFast && ./brighten_simd
brighten: 1000 times loop took 1473778us us/call=1473.778us
brightenSIMD: 1000 times loop took 1467891us us/call=1467.891us

なんと @Vector 導入前と同じくらいまで遅くなってしまいました。代入のときにベクトル化演算の個数がズレてしまい、全体がベクトル化されなくなってしまったのでしょう。

もっと速く、かつポータブルにする

ベンチマークに使っている Ryzen 9 5900X の SIMD 演算器は 256bit 幅です。つまり、8bit整数なら32個まで同時に処理できるはずです。
また、CPUが異なれば SIMD 演算器の bit 幅も変わってきます。Zig ではどのように対処するかというと、std.simd.suggestVectorSize という関数があり、u8 や f32 などの型を与えるとSIMDで同時に処理出来る個数が Option で返ってきます SIMD を使わない方が良いときは null のようです。

brighten_simd2.zig

const std = @import("std");
const print = std.debug.print;
const microTimestamp = std.time.microTimestamp;

fn brighten(image: []u8) void {
    comptime var vector_size: ?usize = std.simd.suggestVectorSize(u8);
    if (vector_size) |size| {
        brightenSIMD(image, size);
    } else {
        brightenNaive(image);
    }
}

fn brightenNaive(image: []u8) void {
    var idx: usize = 0;
    while (idx < image.len) : (idx += 4) {
        image[idx + 0] +|= 10;
        image[idx + 1] +|= 10;
        image[idx + 2] +|= 10;
    }
}

fn brightenSIMD(image: []u8, comptime vector_size: usize) void {
    var idx: usize = 0;
    const repeat_count = @divExact(vector_size, 4);
    const add_vec: @Vector(vector_size, u8) = [_]u8{ 10, 10, 10, 0 } ** repeat_count;
    const add_vec4 = @Vector(4, u8){ 10, 10, 10, 0 };
    const loop_max = image.len -| vector_size;
    while (idx < loop_max) : (idx += vector_size) {
        const image_vec: @Vector(vector_size, u8) = image[idx..][0..vector_size].*;
        const result_vec = image_vec +| add_vec;
        inline for ([_]u8{0} ** vector_size) |_, i| {
            image[idx + i] = result_vec[i];
        }
    }
    while (idx + 3 < image.len) : (idx += 4) {
        const image_vec = @Vector(4, u8){ image[idx + 0], image[idx + 1], image[idx + 2], image[idx + 3] };
        const result_vec = image_vec +| add_vec4;
        image[idx + 0] = result_vec[0];
        image[idx + 1] = result_vec[1];
        image[idx + 2] = result_vec[2];
        image[idx + 3] = result_vec[3];
    }
}

fn benchmark(comptime width: usize, comptime height: usize, comptime num_loop: usize, comptime func: fn (image: []u8) void, comptime name: []const u8) void {
    const size = width * height * 4;
    var image: [size]u8 = [_]u8{0} ** size;
    var cnt: usize = 0;
    const start = microTimestamp();
    while (cnt < num_loop) : (cnt += 1) {
        func(&image);
    }
    const stop = microTimestamp();
    const duration = @intToFloat(f64, stop - start);
    const call_duration = duration / @intToFloat(f64, num_loop);
    print("{s}: {d} times loop took {d}us us/call={d}us\n", .{ name, num_loop, duration, call_duration });
}

pub fn main() !void {
    print("suggestVectorSize(u8):{?}\n", .{std.simd.suggestVectorSize(u8)});
    const width = 1920;
    const height = 1080;
    const num_loop = 1000;
    benchmark(width, height, num_loop, brighten, "brighten");
}

zig build-exe brighten_simd2.zig -O ReleaseFast && ./brighten_simd2 
suggestVectorSize(u8):32
brighten: 1000 times loop took 126154us us/call=126.154us

126マイクロ秒まで速くなりました。また、SIMD 演算器のビット幅にも依存しないポータブルなコードになっています。

（追記）Raspberry Piクロスコンパイル

$ zig build-exe brighten_simd2.zig -O ReleaseFast -target aarch64-linux

としてやれば aarch64 用のバイナリが生成できる。RasPi 400にこのバイナリを送り込んで実行すると

$ ./brighten_simd2
uggestVectorSize(u8):16
brightenNaive: 1000 times loop took 10296314us us/call=10296.314us
brighten: 1000 times loop took 6606317us us/call=6606.317us

RasPi環境でもSIMDが使われているっぽいことが確認できる。

まとめ

Zig 言語でベクトル化を表現する @Vector を使った画像処理関数を見てきました。コンパイラの最適化に頼るのでもなく、アーキテクチャ依存のインラインアセンブラを記述するのでもなく、自然な形でベクトル化を表現するコードが書けました。

このような形で SIMD に対応する機能を持っている言語を私は他に知りません。SIMD が使えるという一点だけを以てしても「Zig は C より速い」と言っても良いかもしれません。

今回は画像処理のお題でベクトル化しやすいデータ構造でしたが、任意の struct に対しても MultiArrayListを使えばフィールドが連続したメモリ領域に配置され、@Vector が SIMD 命令に展開されやすくなるはずです。