//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.4
//       Copyright (2021) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Siva Rajamanickam (srajama@sandia.gov)
//
// ************************************************************************
//@HEADER

//
// Created by Harvey, Evan on 7/13/21.
//

#ifndef KOKKOSKERNELS_KOKKOSBATCHED_KERNEL_HEADER_HPP
#define KOKKOSKERNELS_KOKKOSBATCHED_KERNEL_HEADER_HPP

#include "KokkosKernels_Error.hpp"

#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL)
#include <mkl.h>
#endif  // KOKKOSKERNELS_ENABLE_TPL_MKL

#if defined(KOKKOSKERNELS_ENABLE_TPL_ARMPL)
#include "armpl.h"
#endif  // KOKKOSKERNELS_ENABLE_TPL_ARMPL

#if defined(KOKKOSKERNELS_ENABLE_TPL_CUBLAS)
#include "cuda_runtime.h"
#include "cublas_v2.h"
#endif  // KOKKOSKERNELS_ENABLE_TPL_CUBLAS

#if defined(KOKKOSKERNELS_ENABLE_TPL_MAGMA)
#include <magma_v2.h>
#include <magma_batched.h>
#endif  // KOKKOSKERNELS_ENABLE_TPL_MAGMA

namespace KokkosBatched {

/// \brief Heuristic algorithm types. See BatchedKernelHandle for details.
namespace BaseHeuristicAlgos {
enum BASE_HEURISTIC_ALGOS : int { SQUARE = 0, TALL, WIDE, N };
}

/// \brief Tpl algorithm types. See BatchedKernelHandle for details.
namespace BaseTplAlgos {
enum BASE_TPL_ALGOS : int { ARMPL = BaseHeuristicAlgos::N, MKL, N };
}

/// \brief KokkosBatched algorithm types. See BatchedKernelHandle for details.
namespace BaseKokkosBatchedAlgos {
enum BASE_KOKKOS_BATCHED_ALGOS : int { KK_SERIAL = BaseTplAlgos::N, N };
}

#define N_BASE_ALGOS BaseKokkosBatchedAlgos::N
#define BASE_ALGO_STRS                                                      \
  "BaseHeuristicAlgos::SQUARE", "BaseHeuristicAlgos::TALL",                 \
      "BaseHeuristicAlgos::WIDE", "BaseTplAlgos::ARMPL", "BaseTplAlgosMKL", \
      "BaseKokkosBatchedAlgos::KK_SERIAL"

/// \brief TplParams abstracts underlying handle or execution queue type.
struct TplParams {
  union {
#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL)
    // queue mkl_queue;
    // TODO: Add queue header? Cannot find any declarations in intel-18, let
    // alone oneAPI 2021
#endif  // KOKKOSKERNELS_ENABLE_TPL_MKL

#if defined(KOKKOSKERNELS_ENABLE_TPL_ARMPL)
    armpl_int_t ninter = 1;
#endif  // KOKKOSKERNELS_ENABLE_TPL_ARMPL

#if defined(KOKKOSKERNELS_ENABLE_TPL_CUBLAS)
    cublasHandle_t *cublas_handle;
#endif  // KOKKOSKERNELS_ENABLE_TPL_CUBLAS

#if defined(KOKKOSKERNELS_ENABLE_TPL_MAGMA)
    magma_queue_t *magma_queue;
#endif  // KOKKOSKERNELS_ENABLE_TPL_MAGMA

    uint8_t empty_union;  // suppress empty union build warning
  };
};

// clang-format off
/// \brief Handle for selecting runtime behavior of the BatchedGemm interface.
///
/// \var kernelAlgoType  Specifies which algorithm to use for invocation (default, SQUARE).
///
///                    Specifies whether to select optimal invocations based on inputs and
///                    heuristics:
///                      SQUARE select invocations based on square matrix heuristics where M=N
///                      TALL   select invocations based on tall   matrix heuristics where M>N
///                      WIDE   select invocations based on wide   matrix heuristics where M<N
///                    Note: If the heuristics indicate SIMD views are required for optimal
///                          performance, notify the user that SIMD views are required for
///                          optimal performance.
///
///                    Specifies which cmake-enabled TPL algorithm to invoke:
///                      TPL algorithms currently SUPPORTED by all batched routines:
///
///                      TPL algorithms currently UNSUPPORTED by all batched routines:
///                        ARMPL    Invoke the ArmPL TPL interface
///                        MKL      Invoke the MKL TPL interface
///                        CUBLAS   Invoke the CuBLAS TPL interface
///                        MAGMA    Invoke the Magma TPL interface
///                    Note: See each routine's handle for details about which TPLs are supported.
///                    Note: Requires that input views for A, B, and C reside on either host
///                          or device depending on the TPL selected.
///                    Note: If the user selects a TPL, an error will be thrown if:
///                            1. The TPL is not enabled via cmake
///                            2. The input views do not reside on the host/device as needed
///
///                    Specifies which kokkos-kernels (KK) algorithm to invoke:
///                      KK algorithms currently SUPPORTED   by all KK batched routines:
///                        KK_SERIAL       Invoke SerialFUNC     via RangePolicy(BatchSz)
///                      KK algorithms currently UNSUPPORTED by all KK batched routines:
///                        KK_TEAM         Invoke TeamFUNC       via TeamPolicy(BatchSz)
///                        KK_TEAMVECTOR   Invoke TeamVectorFUNC via TeamPolicy(BatchSz)
///                        KK_SERIALSIMD   Invoke SerialFUNC     via TeamPolicy(BatchSz)
///                        KK_TEAMSIMD     Invoke TeamFUNC       via TeamPolicy(BatchSz)
///                        KK_SERIAL_RANK0 Invoke SerialFUNC     via RangePolicy(BatchSz*LHS_N*LHS_M)
///                        KK_SERIAL_SHMEM Invoke SerialFUNC     via TeamPolicy(BatchSz)
///                        KK_DBLBUF       Solve FUNC            via TeamPolicy(BatchSz*TILES)
///                                        Uses a tuned functor with tiling and double buffering
///                                        via shared memory and register buffers.
///                    Note: See each routine's handle for details about which KK algorithms are
///                          supported.
/// \var teamSz        Specifies the team size that will affect any KK algorithm which uses
///                    TeamPolicy (default, Kokkos::AUTO).
///                    Note: Only applied if useAlgo_type == KK_*
/// \var vecLen        Specifies the vector length that will affect any KK algorithm which
///                    uses TeamPolicy and Kokkos::ThreadVectorRange or Kokkos::TeamVectorRange
///                    (default, Kokkos::AUTO).
///                    Note: Only applied if useAlgo_type == KK_*
/// \var enabledDebug  toggle debug messages.
// clang-format on
class BatchedKernelHandle {
 public:
  int teamSz       = 0;
  int vecLen       = 0;
  bool enableDebug = false;

  BatchedKernelHandle(int kernelAlgoType = BaseHeuristicAlgos::SQUARE,
                      int teamSize = 0, int vecLength = 0)
      : teamSz(teamSize), vecLen(vecLength), _kernelAlgoType(kernelAlgoType) {
#if !defined(KOKKOSKERNELS_ENABLE_TPL_ARMPL) || ARMPL_BUILD < 1058
    if (_kernelAlgoType == BaseTplAlgos::ARMPL) {
      std::ostringstream os;
      os << "KokkosBatched::BatchedKernelHandle requires "
            "KOKKOSKERNELS_ENABLE_TPL_ARMPL and armpl version 21.0.0+"
         << std::endl;
      KokkosKernels::Impl::throw_runtime_exception(os.str());
    }
#endif  // !defined(KOKKOSKERNELS_ENABLE_TPL_ARMPL)
  };

  int get_kernel_algo_type() const { return _kernelAlgoType; }

  std::string get_kernel_algo_type_str() const {
    return algo_type_strs[_kernelAlgoType];
  }

  decltype(auto) get_tpl_params() const {
#if _kernelAlgoType == ARMPL && defined(KOKKOSKERNELS_ENABLE_TPL_ARMPL)
    return &_tplParamsSingleton.ninter;
#elif _kernelAlgoType == MKL && defined(KOKKOSKERNELS_ENABLE_TPL_MKL)
    return "BaseTplAlgos::MKL does not support any tpl parameters";
#else
    return "Unsupported kernelAlgoType:" + get_kernel_algo_type_str() + ".";
#endif
  }

  // clang-format off
  /// \var _kernelAlgoType Specifies which algorithm to use for invocation (default, SQUARE).
  /// \var _tplParams      a handle or queue specific to the TPL API.
  ///                      managed internally unless provided by user via
  ///                      constructor overload
  // clang-format on
 protected:
  // Define TPL params singleton as static class method variable
  // Only one instance of tplParamsGlobalStorage may exist per process.
  static TplParams &_get_tpl_params_singleton() {
    static TplParams tplParamsGlobalStorage;
    return tplParamsGlobalStorage;
  }

  int _kernelAlgoType            = BaseHeuristicAlgos::SQUARE;
  TplParams &_tplParamsSingleton = _get_tpl_params_singleton();
  bool _tplParamsSet             = false;

 private:
  const char *algo_type_strs[N_BASE_ALGOS] = {BASE_ALGO_STRS};
};

}  // namespace KokkosBatched

#endif  // KOKKOSKERNELS_KOKKOSBATCHED_KERNEL_HEADER_HPP
