// Benchmarks for diagonal operations.
//
// Tests diagonal extraction, diagonal-matrix product, and matrix-diagonal product.

#include <benchmark/benchmark.h>
#include <Eigen/Core>

using namespace Eigen;

// Extract diagonal from a square matrix and sum it.
template <typename Scalar>
static void BM_DiagonalExtract(benchmark::State& state) {
  const Index n = state.range(0);
  using Mat = Matrix<Scalar, Dynamic, Dynamic>;
  Mat A = Mat::Random(n, n);
  for (auto _ : state) {
    Scalar s = A.diagonal().sum();
    benchmark::DoNotOptimize(s);
  }
  state.SetBytesProcessed(state.iterations() * n * sizeof(Scalar));
}

// y = diag(d) * x  (diagonal matrix times vector).
template <typename Scalar>
static void BM_DiagonalTimesVector(benchmark::State& state) {
  const Index n = state.range(0);
  using Vec = Matrix<Scalar, Dynamic, 1>;
  Vec d = Vec::Random(n);
  Vec x = Vec::Random(n);
  Vec y(n);
  for (auto _ : state) {
    y = d.asDiagonal() * x;
    benchmark::DoNotOptimize(y.data());
  }
  state.SetBytesProcessed(state.iterations() * 3 * n * sizeof(Scalar));
}

// C = diag(d) * A  (diagonal matrix times dense matrix).
template <typename Scalar>
static void BM_DiagonalTimesMatrix(benchmark::State& state) {
  const Index n = state.range(0);
  using Mat = Matrix<Scalar, Dynamic, Dynamic>;
  using Vec = Matrix<Scalar, Dynamic, 1>;
  Vec d = Vec::Random(n);
  Mat A = Mat::Random(n, n);
  Mat C(n, n);
  for (auto _ : state) {
    C.noalias() = d.asDiagonal() * A;
    benchmark::DoNotOptimize(C.data());
  }
  state.SetBytesProcessed(state.iterations() * 2 * n * n * sizeof(Scalar));
}

// C = A * diag(d)  (dense matrix times diagonal matrix).
template <typename Scalar>
static void BM_MatrixTimesDiagonal(benchmark::State& state) {
  const Index n = state.range(0);
  using Mat = Matrix<Scalar, Dynamic, Dynamic>;
  using Vec = Matrix<Scalar, Dynamic, 1>;
  Vec d = Vec::Random(n);
  Mat A = Mat::Random(n, n);
  Mat C(n, n);
  for (auto _ : state) {
    C.noalias() = A * d.asDiagonal();
    benchmark::DoNotOptimize(C.data());
  }
  state.SetBytesProcessed(state.iterations() * 2 * n * n * sizeof(Scalar));
}

static void Sizes(::benchmark::Benchmark* b) {
  for (int n : {32, 64, 128, 256, 512, 1024}) b->Arg(n);
}

BENCHMARK(BM_DiagonalExtract<float>)->Apply(Sizes)->Name("DiagonalExtract_float");
BENCHMARK(BM_DiagonalExtract<double>)->Apply(Sizes)->Name("DiagonalExtract_double");
BENCHMARK(BM_DiagonalTimesVector<float>)->Apply(Sizes)->Name("DiagonalTimesVector_float");
BENCHMARK(BM_DiagonalTimesVector<double>)->Apply(Sizes)->Name("DiagonalTimesVector_double");
BENCHMARK(BM_DiagonalTimesMatrix<float>)->Apply(Sizes)->Name("DiagonalTimesMatrix_float");
BENCHMARK(BM_DiagonalTimesMatrix<double>)->Apply(Sizes)->Name("DiagonalTimesMatrix_double");
BENCHMARK(BM_MatrixTimesDiagonal<float>)->Apply(Sizes)->Name("MatrixTimesDiagonal_float");
BENCHMARK(BM_MatrixTimesDiagonal<double>)->Apply(Sizes)->Name("MatrixTimesDiagonal_double");