// Benchmarks for Eigen Tensor contraction (generalized GEMM).
// Tests single-threaded (DefaultDevice) and multi-threaded (ThreadPoolDevice) variants.
// SPDX-FileCopyrightText: The Eigen Authors
// SPDX-License-Identifier: MPL-2.0

#define EIGEN_USE_THREADS

#include <benchmark/benchmark.h>
#include <contrib/Eigen/AutoDiff>
#include <contrib/Eigen/Tensor>
#include <Eigen/ThreadPool>

using namespace Eigen;

#ifndef SCALAR
#define SCALAR float
#endif

typedef SCALAR Scalar;

// --- DefaultDevice contraction (rank-2, equivalent to matrix multiply) ---
static void BM_Contraction(benchmark::State& state) {
  const int M = state.range(0);
  const int N = state.range(1);
  const int K = state.range(2);

  Tensor<Scalar, 2> A(M, K);
  Tensor<Scalar, 2> B(K, N);
  Tensor<Scalar, 2> C(M, N);
  A.setRandom();
  B.setRandom();

  using ContractDims = Tensor<Scalar, 2>::DimensionPair;
  Eigen::array<ContractDims, 1> contract_dims = {ContractDims(1, 0)};

  for (auto _ : state) {
    C = A.contract(B, contract_dims);
    benchmark::DoNotOptimize(C.data());
    benchmark::ClobberMemory();
  }
  state.counters["GFLOPS"] =
      benchmark::Counter(2.0 * M * N * K, benchmark::Counter::kIsIterationInvariantRate, benchmark::Counter::kIs1000);
}

// --- ThreadPoolDevice contraction ---
static void BM_Contraction_ThreadPool(benchmark::State& state) {
  const int M = state.range(0);
  const int N = state.range(1);
  const int K = state.range(2);
  const int threads = state.range(3);

  Tensor<Scalar, 2> A(M, K);
  Tensor<Scalar, 2> B(K, N);
  Tensor<Scalar, 2> C(M, N);
  A.setRandom();
  B.setRandom();

  ThreadPool tp(threads);
  ThreadPoolDevice dev(&tp, threads);

  using ContractDims = Tensor<Scalar, 2>::DimensionPair;
  Eigen::array<ContractDims, 1> contract_dims = {ContractDims(1, 0)};

  for (auto _ : state) {
    C.device(dev) = A.contract(B, contract_dims);
    benchmark::DoNotOptimize(C.data());
    benchmark::ClobberMemory();
  }
  state.counters["GFLOPS"] =
      benchmark::Counter(2.0 * M * N * K, benchmark::Counter::kIsIterationInvariantRate, benchmark::Counter::kIs1000);
  state.counters["threads"] = threads;
}

// --- Rank-3 batch contraction ---
// Contracts A(batch, M, K) with B(batch, K, N) over batch dim (0<->0)
// and K dim (2<->1), producing C(M, N). This sums over both the batch
// and inner dimensions: C(m, n) = sum_b sum_k A(b, m, k) * B(b, k, n).
static void BM_BatchContraction(benchmark::State& state) {
  const int batch = state.range(0);
  const int M = state.range(1);
  const int N = state.range(2);
  const int K = state.range(3);

  Tensor<Scalar, 3> A(batch, M, K);
  Tensor<Scalar, 3> B(batch, K, N);
  Tensor<Scalar, 2> C(M, N);
  A.setRandom();
  B.setRandom();

  using ContractDims = Tensor<Scalar, 3>::DimensionPair;
  Eigen::array<ContractDims, 2> contract_dims = {ContractDims(0, 0), ContractDims(2, 1)};

  for (auto _ : state) {
    C = A.contract(B, contract_dims);
    benchmark::DoNotOptimize(C.data());
    benchmark::ClobberMemory();
  }
  state.counters["GFLOPS"] = benchmark::Counter(2.0 * batch * M * N * K, benchmark::Counter::kIsIterationInvariantRate,
                                                benchmark::Counter::kIs1000);
}

// --- RowMajor contraction ---
static void BM_Contraction_RowMajor(benchmark::State& state) {
  const int M = state.range(0);
  const int N = state.range(1);
  const int K = state.range(2);

  Tensor<Scalar, 2, RowMajor> A(M, K);
  Tensor<Scalar, 2, RowMajor> B(K, N);
  Tensor<Scalar, 2, RowMajor> C(M, N);
  A.setRandom();
  B.setRandom();

  using ContractDims = Tensor<Scalar, 2, RowMajor>::DimensionPair;
  Eigen::array<ContractDims, 1> contract_dims = {ContractDims(1, 0)};

  for (auto _ : state) {
    C = A.contract(B, contract_dims);
    benchmark::DoNotOptimize(C.data());
    benchmark::ClobberMemory();
  }
  state.counters["GFLOPS"] =
      benchmark::Counter(2.0 * M * N * K, benchmark::Counter::kIsIterationInvariantRate, benchmark::Counter::kIs1000);
}

// --- Scalar type and layout variants, on DefaultDevice (threads == 0) or ThreadPoolDevice ---
// AutoDiffScalar requires initialization, so the contraction constructs its packed blocks and temporaries.
template <typename T>
static T make_scalar(float value, Index /*i*/) {
  return T(value);
}

template <>
AutoDiffScalar<VectorXf> make_scalar<AutoDiffScalar<VectorXf>>(float value, Index i) {
  const int num_derivatives = 4;
  return AutoDiffScalar<VectorXf>(value, num_derivatives, static_cast<int>(i % num_derivatives));
}

template <typename T, int Layout>
static void BM_ContractionTyped(benchmark::State& state) {
  const int M = state.range(0);
  const int N = state.range(1);
  const int K = state.range(2);
  const int threads = state.range(3);

  Tensor<float, 2, Layout> A_float(M, K);
  Tensor<float, 2, Layout> B_float(K, N);
  A_float.setRandom();
  B_float.setRandom();
  Tensor<T, 2, Layout> A(M, K);
  Tensor<T, 2, Layout> B(K, N);
  Tensor<T, 2, Layout> C(M, N);
  for (Index i = 0; i < A.size(); ++i) A.data()[i] = make_scalar<T>(A_float.data()[i], i);
  for (Index i = 0; i < B.size(); ++i) B.data()[i] = make_scalar<T>(B_float.data()[i], i);

  ThreadPool tp(threads > 0 ? threads : 1);
  ThreadPoolDevice dev(&tp, threads > 0 ? threads : 1);

  using ContractDims = typename Tensor<T, 2, Layout>::DimensionPair;
  Eigen::array<ContractDims, 1> contract_dims = {ContractDims(1, 0)};

  for (auto _ : state) {
    if (threads > 0) {
      C.device(dev) = A.contract(B, contract_dims);
    } else {
      C = A.contract(B, contract_dims);
    }
    benchmark::DoNotOptimize(C.data());
    benchmark::ClobberMemory();
  }
  state.counters["GFLOPS"] =
      benchmark::Counter(2.0 * M * N * K, benchmark::Counter::kIsIterationInvariantRate, benchmark::Counter::kIs1000);
  state.counters["threads"] = threads;
}

// clang-format off
#define CONTRACTION_SIZES \
  ->Args({32, 32, 32})->Args({64, 64, 64})->Args({128, 128, 128}) \
  ->Args({256, 256, 256})->Args({512, 512, 512})->Args({1024, 1024, 1024}) \
  ->Args({256, 256, 1024})->Args({1024, 64, 64})

#define CONTRACTION_THREADPOOL_SIZES \
  ->Args({64, 64, 64, 1})->Args({64, 64, 64, 2})->Args({64, 64, 64, 4}) \
  ->Args({64, 64, 64, 8})->Args({64, 64, 64, 16}) \
  ->Args({256, 256, 256, 1})->Args({256, 256, 256, 2})->Args({256, 256, 256, 4}) \
  ->Args({256, 256, 256, 8})->Args({256, 256, 256, 16}) \
  ->Args({512, 512, 512, 1})->Args({512, 512, 512, 2})->Args({512, 512, 512, 4}) \
  ->Args({512, 512, 512, 8})->Args({512, 512, 512, 16}) \
  ->Args({1024, 1024, 1024, 1})->Args({1024, 1024, 1024, 2})->Args({1024, 1024, 1024, 4}) \
  ->Args({1024, 1024, 1024, 8})->Args({1024, 1024, 1024, 16})

#define TYPED_SIZES \
  ->Args({64, 64, 64, 0})->Args({256, 256, 256, 0})->Args({256, 256, 256, 4}) \
  ->Args({2, 10, 30000, 4})

#define BATCH_SIZES \
  ->Args({1, 64, 64, 64})->Args({1, 256, 256, 256}) \
  ->Args({8, 64, 64, 64})->Args({8, 256, 256, 256}) \
  ->Args({32, 64, 64, 64})->Args({32, 256, 256, 256})
// clang-format on

BENCHMARK(BM_Contraction) CONTRACTION_SIZES;
BENCHMARK(BM_Contraction_RowMajor) CONTRACTION_SIZES;
BENCHMARK(BM_Contraction_ThreadPool) CONTRACTION_THREADPOOL_SIZES->UseRealTime();
BENCHMARK(BM_BatchContraction) BATCH_SIZES;
BENCHMARK_TEMPLATE(BM_ContractionTyped, float, ColMajor) TYPED_SIZES->UseRealTime();
BENCHMARK_TEMPLATE(BM_ContractionTyped, float, RowMajor) TYPED_SIZES->UseRealTime();
BENCHMARK_TEMPLATE(BM_ContractionTyped, AutoDiffScalar<VectorXf>, ColMajor) TYPED_SIZES->UseRealTime();
BENCHMARK_TEMPLATE(BM_ContractionTyped, AutoDiffScalar<VectorXf>, RowMajor) TYPED_SIZES->UseRealTime();