// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
// Copyright (C) 2014 Navdeep Jaitly <ndjaitly@google.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#define EIGEN_TEST_NO_LONGDOUBLE
#define EIGEN_TEST_NO_COMPLEX

#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
#define EIGEN_USE_GPU

#include "main.h"
#include <contrib/Eigen/Tensor>

using Eigen::Tensor;
typedef Tensor<float, 1>::DimensionPair DimPair;

template <int DataLayout>
void test_gpu_contraction(int m_size, int k_size, int n_size) {
  Tensor<float, 2, DataLayout> t_left(m_size, k_size);
  Tensor<float, 2, DataLayout> t_right(k_size, n_size);
  Tensor<float, 2, DataLayout> t_result(m_size, n_size);
  Tensor<float, 2, DataLayout> t_result_gpu(m_size, n_size);
  Eigen::array<DimPair, 1> dims{DimPair(1, 0)};

  t_left.setRandom();
  t_right.setRandom();

  std::size_t t_left_bytes = t_left.size() * sizeof(float);
  std::size_t t_right_bytes = t_right.size() * sizeof(float);
  std::size_t t_result_bytes = t_result.size() * sizeof(float);

  float* d_t_left;
  float* d_t_right;
  float* d_t_result;

  gpuMalloc((void**)(&d_t_left), t_left_bytes);
  gpuMalloc((void**)(&d_t_right), t_right_bytes);
  gpuMalloc((void**)(&d_t_result), t_result_bytes);

  gpuMemcpy(d_t_left, t_left.data(), t_left_bytes, gpuMemcpyHostToDevice);
  gpuMemcpy(d_t_right, t_right.data(), t_right_bytes, gpuMemcpyHostToDevice);

  Eigen::GpuStreamDevice stream;
  Eigen::GpuDevice gpu_device(&stream);

  Eigen::TensorMap<Eigen::Tensor<float, 2, DataLayout>> gpu_t_left(d_t_left, Eigen::array<int, 2>{m_size, k_size});
  Eigen::TensorMap<Eigen::Tensor<float, 2, DataLayout>> gpu_t_right(d_t_right, Eigen::array<int, 2>{k_size, n_size});
  Eigen::TensorMap<Eigen::Tensor<float, 2, DataLayout>> gpu_t_result(d_t_result, Eigen::array<int, 2>{m_size, n_size});

  gpu_t_result.device(gpu_device) = gpu_t_left.contract(gpu_t_right, dims);
  t_result = t_left.contract(t_right, dims);

  gpuMemcpy(t_result_gpu.data(), d_t_result, t_result_bytes, gpuMemcpyDeviceToHost);
  for (DenseIndex i = 0; i < t_result.size(); i++) {
    // Absolute slack for accumulation-order differences between device and host contractions.
    if (numext::abs(t_result(i) - t_result_gpu(i)) < 1024 * NumTraits<float>::epsilon()) {
      continue;
    }
    if (Eigen::internal::isApprox(t_result(i), t_result_gpu(i), 1024 * NumTraits<float>::epsilon())) {
      continue;
    }
    std::cout << "mismatch detected at index " << i << ": " << t_result(i) << " vs " << t_result_gpu(i) << std::endl;
    assert(false);
  }

  gpuFree((void*)d_t_left);
  gpuFree((void*)d_t_right);
  gpuFree((void*)d_t_result);
}

template <int DataLayout, typename Scalar = float>
void test_gpu_contraction_zero_k(int m_size, int n_size) {
  const int k_size = 0;
  Tensor<Scalar, 2, DataLayout> t_left(m_size, k_size);
  Tensor<Scalar, 2, DataLayout> t_right(k_size, n_size);
  Tensor<Scalar, 2, DataLayout> t_result(m_size, n_size);
  Tensor<Scalar, 2, DataLayout> t_result_gpu(m_size, n_size);
  t_result_gpu.setConstant(Scalar(123));
  Eigen::array<DimPair, 1> dims{DimPair(1, 0)};

  std::size_t t_left_bytes = (std::max)(t_left.size(), DenseIndex(1)) * sizeof(Scalar);
  std::size_t t_right_bytes = (std::max)(t_right.size(), DenseIndex(1)) * sizeof(Scalar);
  std::size_t t_result_bytes = t_result.size() * sizeof(Scalar);

  Scalar* d_t_left;
  Scalar* d_t_right;
  Scalar* d_t_result;

  gpuMalloc((void**)(&d_t_left), t_left_bytes);
  gpuMalloc((void**)(&d_t_right), t_right_bytes);
  gpuMalloc((void**)(&d_t_result), t_result_bytes);

  // Prefill the device output with a nonzero value to verify that the GPU path actively zeroes the output buffer.
  gpuMemcpy(d_t_result, t_result_gpu.data(), t_result_bytes, gpuMemcpyHostToDevice);

  Eigen::GpuStreamDevice stream;
  Eigen::GpuDevice gpu_device(&stream);

  Eigen::TensorMap<Eigen::Tensor<Scalar, 2, DataLayout>> gpu_t_left(d_t_left, Eigen::array<int, 2>{m_size, k_size});
  Eigen::TensorMap<Eigen::Tensor<Scalar, 2, DataLayout>> gpu_t_right(d_t_right, Eigen::array<int, 2>{k_size, n_size});
  Eigen::TensorMap<Eigen::Tensor<Scalar, 2, DataLayout>> gpu_t_result(d_t_result, Eigen::array<int, 2>{m_size, n_size});

  gpu_t_result.device(gpu_device) = gpu_t_left.contract(gpu_t_right, dims);
  t_result = t_left.contract(t_right, dims);

  gpuMemcpy(t_result_gpu.data(), d_t_result, t_result_bytes, gpuMemcpyDeviceToHost);
  for (DenseIndex i = 0; i < t_result.size(); i++) {
    VERIFY_IS_EQUAL(t_result(i), Scalar(0));
    VERIFY_IS_EQUAL(t_result_gpu(i), Scalar(0));
  }

  gpuFree((void*)d_t_left);
  gpuFree((void*)d_t_right);
  gpuFree((void*)d_t_result);
}

template <int DataLayout>
void test_gpu_contraction_double(int m_size, int k_size, int n_size) {
  Tensor<double, 2, DataLayout> t_left(m_size, k_size);
  Tensor<double, 2, DataLayout> t_right(k_size, n_size);
  Tensor<double, 2, DataLayout> t_result(m_size, n_size);
  Tensor<double, 2, DataLayout> t_result_gpu(m_size, n_size);
  Eigen::array<Tensor<double, 1>::DimensionPair, 1> dims{Tensor<double, 1>::DimensionPair(1, 0)};

  t_left.setRandom();
  t_right.setRandom();

  std::size_t t_left_bytes = t_left.size() * sizeof(double);
  std::size_t t_right_bytes = t_right.size() * sizeof(double);
  std::size_t t_result_bytes = t_result.size() * sizeof(double);

  double* d_t_left;
  double* d_t_right;
  double* d_t_result;

  gpuMalloc((void**)(&d_t_left), t_left_bytes);
  gpuMalloc((void**)(&d_t_right), t_right_bytes);
  gpuMalloc((void**)(&d_t_result), t_result_bytes);

  gpuMemcpy(d_t_left, t_left.data(), t_left_bytes, gpuMemcpyHostToDevice);
  gpuMemcpy(d_t_right, t_right.data(), t_right_bytes, gpuMemcpyHostToDevice);

  Eigen::GpuStreamDevice stream;
  Eigen::GpuDevice gpu_device(&stream);

  Eigen::TensorMap<Eigen::Tensor<double, 2, DataLayout>> gpu_t_left(d_t_left, Eigen::array<int, 2>{m_size, k_size});
  Eigen::TensorMap<Eigen::Tensor<double, 2, DataLayout>> gpu_t_right(d_t_right, Eigen::array<int, 2>{k_size, n_size});
  Eigen::TensorMap<Eigen::Tensor<double, 2, DataLayout>> gpu_t_result(d_t_result, Eigen::array<int, 2>{m_size, n_size});

  gpu_t_result.device(gpu_device) = gpu_t_left.contract(gpu_t_right, dims);
  t_result = t_left.contract(t_right, dims);

  gpuMemcpy(t_result_gpu.data(), d_t_result, t_result_bytes, gpuMemcpyDeviceToHost);
  for (DenseIndex i = 0; i < t_result.size(); i++) {
    if (std::abs(t_result(i) - t_result_gpu(i)) < 524288 * NumTraits<double>::epsilon()) {
      continue;
    }
    if (Eigen::internal::isApprox(t_result(i), t_result_gpu(i), 524288 * NumTraits<double>::epsilon())) {
      continue;
    }
    std::cout << "mismatch detected at index " << i << ": " << t_result(i) << " vs " << t_result_gpu(i) << std::endl;
    assert(false);
  }

  gpuFree((void*)d_t_left);
  gpuFree((void*)d_t_right);
  gpuFree((void*)d_t_result);
}

template <int DataLayout>
void test_scalar(int m_size, int k_size, int n_size) {
  std::cout << "Testing for (" << m_size << "," << k_size << "," << n_size << ")" << std::endl;
  // with these dimensions, the output has 300 * 140 elements, which is
  // more than 30 * 1024, which is the number of threads in blocks on
  // a 15 SM GK110 GPU
  Tensor<float, 2, DataLayout> t_left(m_size, k_size);
  Tensor<float, 2, DataLayout> t_right(k_size, n_size);
  Tensor<float, 0, DataLayout> t_result;
  Tensor<float, 0, DataLayout> t_result_gpu;
  Eigen::array<DimPair, 2> dims{DimPair(0, 0), DimPair(1, 1)};

  t_left.setRandom();
  t_right.setRandom();

  std::size_t t_left_bytes = t_left.size() * sizeof(float);
  std::size_t t_right_bytes = t_right.size() * sizeof(float);
  std::size_t t_result_bytes = sizeof(float);

  float* d_t_left;
  float* d_t_right;
  float* d_t_result;

  gpuMalloc((void**)(&d_t_left), t_left_bytes);
  gpuMalloc((void**)(&d_t_right), t_right_bytes);
  gpuMalloc((void**)(&d_t_result), t_result_bytes);

  gpuMemcpy(d_t_left, t_left.data(), t_left_bytes, gpuMemcpyHostToDevice);
  gpuMemcpy(d_t_right, t_right.data(), t_right_bytes, gpuMemcpyHostToDevice);

  Eigen::GpuStreamDevice stream;
  Eigen::GpuDevice gpu_device(&stream);

  Eigen::TensorMap<Eigen::Tensor<float, 2, DataLayout>> gpu_t_left(d_t_left, m_size, k_size);
  Eigen::TensorMap<Eigen::Tensor<float, 2, DataLayout>> gpu_t_right(d_t_right, k_size, n_size);
  Eigen::TensorMap<Eigen::Tensor<float, 0, DataLayout>> gpu_t_result(d_t_result);

  gpu_t_result.device(gpu_device) = gpu_t_left.contract(gpu_t_right, dims);
  t_result = t_left.contract(t_right, dims);

  gpuMemcpy(t_result_gpu.data(), d_t_result, t_result_bytes, gpuMemcpyDeviceToHost);
  if (numext::abs(t_result() - t_result_gpu()) > 1024 * NumTraits<float>::epsilon() &&
      !Eigen::internal::isApprox(t_result(), t_result_gpu(), 1024 * NumTraits<float>::epsilon())) {
    std::cout << "mismatch detected: " << t_result() << " vs " << t_result_gpu() << std::endl;
    assert(false);
  }

  gpuFree((void*)d_t_left);
  gpuFree((void*)d_t_right);
  gpuFree((void*)d_t_result);
}

template <int DataLayout>
void test_gpu_contraction_m() {
  for (int k = 32; k < 256; k++) {
    test_gpu_contraction<DataLayout>(k, 128, 128);
  }
}

template <int DataLayout>
void test_gpu_contraction_k() {
  for (int k = 32; k < 256; k++) {
    test_gpu_contraction<DataLayout>(128, k, 128);
  }
}

template <int DataLayout>
void test_gpu_contraction_n() {
  for (int k = 32; k < 256; k++) {
    test_gpu_contraction<DataLayout>(128, 128, k);
  }
}

template <int DataLayout>
void test_gpu_contraction_sizes() {
  int m_sizes[3][5] = {{31, 39, 63, 64, 65}, {127, 129, 255, 257, 511}, {512, 513, 1023, 1024, 1025}};

  int n_sizes[3][5] = {{31, 39, 63, 64, 65}, {127, 129, 255, 257, 511}, {512, 513, 1023, 1024, 1025}};

  int k_sizes[3][6] = {{31, 39, 63, 64, 65, 95}, {96, 127, 129, 255, 257, 511}, {512, 513, 725, 1023, 1024, 1025}};

  // Some selection of specific cases.
  //  - m changes rows each iteration
  //  - n changes rows each 3 iterations
  //  - k changes rows each 9 iterations
  //  - within a row, advance once column each iteration
  const int m_cols = 5;
  const int n_cols = 5;
  const int k_cols = 6;
  int m_offset = 0;
  int n_offset = 1;
  int k_offset = 2;
  for (int i = 0; i < 3; ++i) {
    for (int j = 0; j < 3; ++j) {
      for (int l = 0; l < 3; ++l) {
        int m = m_sizes[l][m_offset];
        int n = n_sizes[j][n_offset];
        int k = k_sizes[i][k_offset];
        test_gpu_contraction<DataLayout>(m, n, k);
        n_offset = (n_offset + 1) % n_cols;
        k_offset = (k_offset + 1) % k_cols;
      }
      m_offset = (m_offset + 1) % m_cols;
      if (j < 2) {
        n_offset = (n_offset + n_cols - 3) % n_cols;  // Rewind 3.
      }
    }
    k_offset = (k_offset + 2 * k_cols - 9) % k_cols;  // Rewind 9.
  }
}

EIGEN_DECLARE_TEST(tensor_contract_gpu) {
  CALL_SUBTEST_1(test_gpu_contraction<ColMajor>(128, 128, 128));
  CALL_SUBTEST_1(test_gpu_contraction<RowMajor>(128, 128, 128));
  CALL_SUBTEST_1(test_gpu_contraction_double<ColMajor>(64, 64, 64));
  CALL_SUBTEST_1(test_gpu_contraction_double<RowMajor>(64, 64, 64));
  CALL_SUBTEST_1(test_gpu_contraction_double<ColMajor>(100, 100, 100));
  CALL_SUBTEST_1(test_gpu_contraction_double<RowMajor>(100, 100, 100));
  CALL_SUBTEST_1(test_gpu_contraction_double<ColMajor>(65, 129, 33));
  CALL_SUBTEST_1(test_gpu_contraction_double<RowMajor>(65, 129, 33));

  CALL_SUBTEST_1(test_scalar<ColMajor>(128, 128, 128));
  CALL_SUBTEST_1(test_scalar<RowMajor>(128, 128, 128));

  CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor>(10, 20)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor>(10, 20)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor>(64, 64)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor>(64, 64)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor, double>(10, 20)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor, double>(10, 20)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor, double>(64, 64)));
  CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor, double>(64, 64)));

  CALL_SUBTEST_2(test_gpu_contraction_m<ColMajor>());
  CALL_SUBTEST_3(test_gpu_contraction_m<RowMajor>());

  CALL_SUBTEST_4(test_gpu_contraction_k<ColMajor>());
  CALL_SUBTEST_5(test_gpu_contraction_k<RowMajor>());

  CALL_SUBTEST_6(test_gpu_contraction_n<ColMajor>());
  CALL_SUBTEST_7(test_gpu_contraction_n<RowMajor>());

#if !defined(EIGEN_USE_HIP)
  // disable these subtests for HIP
  CALL_SUBTEST_8(test_gpu_contraction_sizes<ColMajor>());
  CALL_SUBTEST_9(test_gpu_contraction_sizes<RowMajor>());
#endif
}