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

#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;
}