[arrayfire 3.8.1]JITが改善されて速くなったようなので計測してみた

六花です。
約半年ぶりの投稿です。 もっと気軽に投稿した方がいいのかな。

ひと月ほど前にarrayfireのバージョンが3.8.1に上がっていたので、更新履歴(https://github.com/arrayfire/arrayfire/releases/tag/v3.8.1)を読んだところ、

CPU backend improvements – #3010 , #3138 , #3161
CUDA backend improvements – #3066 , #3091 , #3093 , #3125 , #3143 , #3161
OpenCL backend improvements – #3091 , #3068 , #3127 , #3010 , #3039 , #3138 , #3161
General(including JIT) performance improvements across backends – #3167

と、全体的な改良が行われたようなので、以前作成した3層NNを使って速度計測することにしました。

■環境
Microsoft Visual Studio Community 2019 Version 16.11.2
arrayfire 3.8.1 or 3.8.0
CUDA toolkit 11.6

■コード


#include <array>
#include <random>
#include <iostream>
#include <iomanip>

#include <arrayfire.h>



void print_mat(const af::array in)
{
    for (int i_row = 0; i_row < in.dims()[0]; ++i_row)
    {
        for (int i_col = 0; i_col < in.dims()[1]; ++i_col)
        {
            std::cout << std::setw(12) << af::sum<float>(in(i_row, i_col)) << "\t";
        }
        std::cout << std::endl;
    }
    std::cout << std::endl;
}

int main()
{
    af::setDevice(1);

    constexpr int size_input = 3;
    constexpr int size_hidden = 100;
    constexpr int size_output = 1;
    constexpr int size_data = 10000;

    // 乱数の準備
    std::random_device rd;
    std::array<std::seed_seq::result_type, std::mt19937_64::state_size> seed_data;
    std::generate(seed_data.begin(), seed_data.end(), ref(rd));
    std::seed_seq s_seq(seed_data.begin(), seed_data.end());
    std::mt19937_64 mt(s_seq);
    for (int i = 0; i < 100; ++i) { mt(); }
    auto dist_input = std::uniform_real_distribution<float>(-1.0, +1.0);

    float input_src[size_data][size_input];
    float output_src[size_data];

    // y = x1 + x2+ x3
    for (int i1 = 0; i1 < size_data; ++i1)
    {
        output_src[i1] = float(0);
        for (int i2 = 0; i2 < size_input; ++i2)
        {
            input_src[i1][i2] = dist_input(mt);
            output_src[i1] += input_src[i1][i2];
        }
    }

    // Arrayfire
    //-------------------
    {
        af::timer::start();

        {
            af::info();

            // af::constantの最初の引数はテンプレートになっているので、「0」にするとint型と判断されるのでバグの元
            af::array input = af::constant(0.0, size_input, size_data);
            af::array output = af::constant(0.0, size_output, size_data);
            for (int i1 = 0; i1 < size_data; ++i1)
            {
                for (int i2 = 0; i2 < size_input; ++i2)
                {
                    input(i2, i1) = input_src[i1][i2];
                }
                output(0, i1) = output_src[i1];
            }

            // z1 = W * x + B
            // z2 = sigmoid(z1)
            // y = V * z + C
            af::array x_data = af::constant(0.0, size_input, 1);
            af::array x_grad = af::constant(0.0, size_input, 1);
            af::array W_data = af::randn(size_hidden, size_input) * (1.0 / std::sqrt(size_input)); // Xavierの初期値
            af::array W_grad = af::constant(0.0, size_hidden, size_input);
            af::array B_data = af::constant(0.0, size_hidden, 1);
            af::array B_grad = af::constant(0.0, size_hidden, 1);
            af::array z1_data = af::constant(0.0, size_hidden, 1);
            af::array z1_grad = af::constant(0.0, size_hidden, 1);
            af::array z2_data = af::constant(0.0, size_hidden, 1);
            af::array z2_grad = af::constant(0.0, size_hidden, 1);
            af::array V_data = af::randn(size_output, size_hidden) * (1.0 / std::sqrt(size_hidden)); // Xavierの初期値
            af::array V_grad = af::constant(0.0, size_output, size_hidden);
            af::array C_data = af::constant(0.0, size_output, 1);
            af::array C_grad = af::constant(0.0, size_output, 1);
            af::array y_data = af::constant(0.0, size_output, 1);
            af::array y_grad = af::constant(0.0, size_output, 1);

            std::cout << "Arrayfire main loop" << std::endl;
            for (int i_epoch = 0; i_epoch < int(1e+5) / size_data; ++i_epoch)
            {
                // std::vectorでいうところのstd::shuffleの代わり
                // val_dataにはソート済みの行列、idx_dataにはval_dataの要素のソート前の行列の位置が入る
                // 乱数で生成した行列をソートすると、idx_dataに一意の数字が入っているshuffleされた配列が手に入る
                af::array idx_data;
                {
                    af::array val_data, src_data = af::randu(size_data, 1);
                    af::sort(val_data, idx_data, src_data);
                }

                for (int i_epoch_sub = 0; i_epoch_sub < size_data; ++i_epoch_sub)
                {
                    std::cout << i_epoch << "\t" << i_epoch_sub << "\r";

                    af::array idx_target = idx_data(i_epoch_sub);

                    // forward
                    x_data = input(af::span, idx_target);
                    z1_data = af::matmul(W_data, x_data) + B_data;
                    z2_data = af::sigmoid(z1_data);
                    y_data = af::matmul(V_data, z2_data) + C_data;

                    // diff
                    af::array diff = (y_data - output(af::span, idx_target));
                    y_grad = diff;

                    // backward
                    C_grad = y_grad;
                    V_grad = af::matmul(y_grad, z2_data.T());
                    z2_grad = af::matmul(V_data.T(), y_grad);
                    z1_grad = z2_grad * (af::sigmoid(z1_data) * (1.0 - af::sigmoid(z1_data)));
                    B_grad = z1_grad;
                    W_grad = af::matmul(z1_grad, x_data.T());
                    x_grad = matmul(W_data.T(), z1_grad);

                    // update
                    constexpr float alpha = 1e-2;
                    C_data -= alpha * C_grad;
                    V_data -= alpha * V_grad;
                    B_data -= alpha * B_grad;
                    W_data -= alpha * W_grad;

                    const float norm_diff = af::norm(diff);
                    static float min_diff = norm_diff;
                    if (norm_diff < min_diff)
                    {
                        min_diff = norm_diff;
                        std::cout << i_epoch << "\t" << i_epoch_sub << "\tnorm_diff :\t" << norm_diff << std::endl;
                    }
                }
            }
        }

        std::cout << std::endl << std::endl << "end" << std::endl;

        af::sync();
        std::cout << af::timer::stop() << std::endl;
    }

    while (true) {}
    return 0;
}

■結果
時間計測はarrayfire付属のaf::timerを使いました。
10回計測して、その平均を取ります。

3.8.0 GeForce RTX 3080 _ : 88.55873秒
3.8.1 GeForce GTX 1080 Ti : 76.22725秒
3.8.1 GeForce RTX 3080 _ : 73.27198秒

3080と1080 tiの違いがあまりなかったのはちょっと驚きましたが、計測だけにPCを使っているわけではなかったので何かしら負荷がかかっていたんだと思います。
少なくともバージョンを上げたら明確に処理速度が上がったのは実感できました。

■雑感
arrayfireはバージョン切り替えするのに再インストールするのが面倒くさかったです。
バージョン切り替えしやすい方法ないものかな。

おすすめ

コメントを残す

メールアドレスが公開されることはありません。