Benchmark column/row sums varmat vs. matvar

vector_sum.cpp

#include <benchmark/benchmark.h>
#include <stan/math.hpp>
#include <utility>

static void toss_me(benchmark::State& state) {
  using stan::math::var;
  Eigen::Matrix<double, -1, -1> x_vals = Eigen::MatrixXd::Random(256, 256);
  Eigen::Matrix<double, -1, -1> y_vals = Eigen::MatrixXd::Random(256, 256);
  using stan::math::var;
  using stan::math::sum;
  Eigen::Matrix<var, -1, -1> x = x_vals;
  Eigen::Matrix<var, -1, -1> y = y_vals;
  var lp = 0;
  lp -= sum((multiply(x, y) + x).eval());
  benchmark::DoNotOptimize(lp.vi_);
  for (auto _ : state) {
    lp.grad();
    benchmark::ClobberMemory();
    stan::math::set_zero_all_adjoints();
  }
  stan::math::recover_memory();
}

static void sum_columns_matvar(benchmark::State& state) {
  using stan::math::var;

  for (auto _ : state) {
    Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic> matrix =
      Eigen::MatrixXd::Random(state.range(0), state.range(0));
    Eigen::Matrix<stan::math::var, Eigen::Dynamic, 1> vector =
      Eigen::VectorXd::Random(state.range(0));

    auto start = std::chrono::high_resolution_clock::now();
    for(size_t i = 0; i < matrix.cols(); ++i)
      vector += matrix.col(i);
    stan::math::grad();
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed_seconds =
      std::chrono::duration_cast<std::chrono::duration<double>>(end - start);

    state.SetIterationTime(elapsed_seconds.count());
    stan::math::recover_memory();
    benchmark::ClobberMemory();
  }
}

static void sum_columns_varmat(benchmark::State& state) {
  using stan::math::var;

  for (auto _ : state) {
    stan::math::var_value<Eigen::MatrixXd> matrix =
      Eigen::MatrixXd::Random(state.range(0), state.range(0));
    stan::math::var_value<Eigen::VectorXd> vector =
      Eigen::VectorXd::Random(state.range(0));

    auto start = std::chrono::high_resolution_clock::now();
    for(size_t i = 0; i < matrix.cols(); ++i)
      vector = stan::math::add(vector, matrix.col(i));
    stan::math::grad();
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed_seconds =
      std::chrono::duration_cast<std::chrono::duration<double>>(end - start);

    state.SetIterationTime(elapsed_seconds.count());
    stan::math::recover_memory();
    benchmark::ClobberMemory();
  }
}

static void sum_rows_matvar(benchmark::State& state) {
  using stan::math::var;

  for (auto _ : state) {
    Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic> matrix =
      Eigen::MatrixXd::Random(state.range(0), state.range(0));
    Eigen::Matrix<stan::math::var, 1, Eigen::Dynamic> row_vector =
      Eigen::RowVectorXd::Random(state.range(0));

    auto start = std::chrono::high_resolution_clock::now();
    for(size_t i = 0; i < matrix.rows(); ++i)
      row_vector += matrix.row(i);
    stan::math::grad();
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed_seconds =
      std::chrono::duration_cast<std::chrono::duration<double>>(end - start);

    state.SetIterationTime(elapsed_seconds.count());
    stan::math::recover_memory();
    benchmark::ClobberMemory();
  }
}

static void sum_rows_varmat(benchmark::State& state) {
  using stan::math::var;

  for (auto _ : state) {
    stan::math::var_value<Eigen::MatrixXd> matrix =
      Eigen::MatrixXd::Random(state.range(0), state.range(0));
    stan::math::var_value<Eigen::RowVectorXd> row_vector =
      Eigen::RowVectorXd::Random(state.range(0));

    auto start = std::chrono::high_resolution_clock::now();
    for(size_t i = 0; i < matrix.rows(); ++i)
      row_vector = stan::math::add(row_vector, matrix.row(i));
    stan::math::grad();
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed_seconds =
      std::chrono::duration_cast<std::chrono::duration<double>>(end - start);

    state.SetIterationTime(elapsed_seconds.count());
    stan::math::recover_memory();
    benchmark::ClobberMemory();
  }
}

// The start and ending sizes for the benchmark                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
int start_val = 4;
int end_val = 1024;
BENCHMARK(toss_me);
BENCHMARK(sum_columns_matvar)->RangeMultiplier(2)->Range(start_val, end_val)->UseManualTime();
BENCHMARK(sum_columns_varmat)->RangeMultiplier(2)->Range(start_val, end_val)->UseManualTime();
BENCHMARK(sum_rows_matvar)->RangeMultiplier(2)->Range(start_val, end_val)->UseManualTime();
BENCHMARK(sum_rows_varmat)->RangeMultiplier(2)->Range(start_val, end_val)->UseManualTime();
BENCHMARK_MAIN();