/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once

#include <glog/logging.h>
#include <limits>
#include <string>
#include <utility>
#include "paddle/common/layout.h"
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/api/lib/backend_set.h"
#include "paddle/phi/api/lib/data_type_set.h"
#include "paddle/phi/backends/all_context.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
#include "paddle/phi/core/selected_rows.h"
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"

// TODO(chenweihang): split Key, Kernel, Factory into diff files
#include "paddle/phi/core/kernel_factory.h"

namespace paddle {
namespace experimental {

namespace detail {
PADDLE_API BackendSet GetTensorBackendSet(const phi::TensorBase& t);
PADDLE_API std::size_t CountLeadingZeros(uint32_t val);
}  // namespace detail

PADDLE_API phi::DeviceContext* GetDeviceContextByBackend(phi::Backend backend);

enum class KernelType {
  DENSE_TENSOR_KERNEL,   // kernel for DenseTensor
  SELECTED_ROWS_KERNEL,  // kernel for SelectedRows
  SPARSE_COO_KERNEL,     // kernel for SparseCooTensor
  SPARSE_CSR_KERNEL      // kernel for SparseCsrTensor
};

// TODO(chenweihang): support DataLayout and DataType selected
struct KernelKeySet {
  BackendSet backend_set{Backend::UNDEFINED};
  phi::DataLayout layout{phi::DataLayout::UNDEFINED};
  DataType dtype{DataType::UNDEFINED};

  // TODO(chenweihang): iterate all kernel key for kernel selection
  phi::KernelKey GetHighestPriorityKernelKey() {
    uint32_t bitset_value = backend_set.bitset();
#ifdef PADDLE_WITH_CUSTOM_DEVICE
    if (backend_set.Has(Backend(4))) {
      return phi::KernelKey(Backend(4), layout, dtype);
    }
#endif
    std::size_t leading_zeros = detail::CountLeadingZeros(bitset_value);
    Backend selected_backend = static_cast<Backend>(32 - leading_zeros);
    VLOG(8) << "GetHighestPriorityKernelKey: selected_backend = "
            << selected_backend;

    return phi::KernelKey(selected_backend, layout, dtype);
  }
};

namespace detail {

template <typename Functor>
struct ArgsIterator {
  template <typename... Args>
  inline Functor& apply() {
    return self();
  }

  template <typename T, typename... Args>
  inline Functor& apply(T&& arg, Args&&... args) {
    self()(std::forward<T>(arg));
    if (self().short_circuit()) {
      return self();
    } else {
      return apply(std::forward<Args>(args)...);
    }
  }

  constexpr bool short_circuit() const { return false; }

 private:
  inline Functor& self() { return *static_cast<Functor*>(this); }
};

struct KernelKeyParser : ArgsIterator<KernelKeyParser> {
  KernelKeySet key_set;
  bool disable_gpudnn = false;
  // this dtype_set is used for cache multi-inputs dtype and used for
  // data_promote
  DataTypeSet dtype_set{DataType::UNDEFINED};

  inline void AssignKernelKeySet(const phi::TensorBase& tensor) {
    // assign Backend
    BackendSet tensor_backend_set = detail::GetTensorBackendSet(tensor);
    key_set.backend_set = key_set.backend_set | tensor_backend_set;
    // tensor's attribute use_gpudnn=False, explicitly disable gpudnn kernel
    if (tensor_backend_set ==
            BackendSet(paddle::experimental::get_accelerat_backend()) ||
        disable_gpudnn) {
      disable_gpudnn = true;
      key_set.backend_set = key_set.backend_set - BackendSet(Backend::GPUDNN);
      VLOG(8) << "Disable kernel backend: GPUDNN";
    }
    // assign DataLayout
    phi::DataLayout tensor_layout = tensor.layout();
    key_set.layout =
        tensor_layout > key_set.layout ? tensor_layout : key_set.layout;
    // assign DataType
    key_set.dtype = tensor.dtype();
    dtype_set = dtype_set | DataTypeSet(key_set.dtype);
    auto promote_result = PromoteTypes(dtype_set);
    if (promote_result != DataType::UNDEFINED) {
      key_set.dtype = promote_result;
      VLOG(8) << "promote kernel DataType:" << promote_result;
    }
  }

  void operator()(const Tensor& x) {
    const auto* tensor = x.impl().get();
    if (tensor) {
      AssignKernelKeySet(*tensor);
    }
  }

  void operator()(const std::vector<Tensor>& x) {
    if (!x.empty()) {
      const phi::TensorBase& tensor = *x.at(0).impl();
      AssignKernelKeySet(tensor);
    }
  }

  void operator()(const paddle::optional<Tensor>& x) {
    if (x) {
      const phi::TensorBase& tensor = *(x.get_ptr()->impl());
      AssignKernelKeySet(tensor);
    }
  }

  // skip other type args, these args don't used in kernel selection
  template <typename T>
  void operator()(const T& x) {
    // do nothing
  }
};

struct KernelTypeParser : ArgsIterator<KernelTypeParser> {
  KernelType kernel_type{KernelType::DENSE_TENSOR_KERNEL};

  // TODO(chenweihang): deal with multiple diff input Tensors
  // TODO(chenweihang): add global device guard method to set backend
  void operator()(const Tensor& x) {
    if (phi::SelectedRows::classof(x.impl().get())) {
      kernel_type = KernelType::SELECTED_ROWS_KERNEL;
    } else if (phi::SparseCooTensor::classof(x.impl().get())) {
      kernel_type = KernelType::SPARSE_COO_KERNEL;
    } else if (phi::SparseCsrTensor::classof(x.impl().get())) {
      kernel_type = KernelType::SPARSE_CSR_KERNEL;
    }
  }

  // skip other type args, these args don't used in kernel selection
  template <typename T>
  void operator()(const T& x) {
    // do nothing
  }
};

/* ------------------ for auto parallel ----------------------- */

struct DistTensorTypeParser : ArgsIterator<DistTensorTypeParser> {
  bool result = false;

  bool short_circuit() { return result; }

  void operator()(const Tensor& x) { result = x.is_dist_tensor(); }

  void operator()(const paddle::optional<Tensor>& x) {
    if (x) {
      result = x.get_ptr()->is_dist_tensor();
    }
  }

  void operator()(const std::vector<Tensor>& x) {
    if (!x.empty()) {
      for (auto& t : x) {
        result = result || t.is_dist_tensor();
        if (short_circuit()) break;
      }
    }
  }

  void operator()(const paddle::optional<std::vector<Tensor>>& x) {
    if (x && !x->empty()) {
      for (auto& t : *(x.get_ptr())) {
        result = result || t.is_dist_tensor();
        if (short_circuit()) break;
      }
    }
  }

  // skip other type args, these args don't used in kernel selection
  template <typename T>
  void operator()(const T& x) {
    // do nothing
  }
};

}  // namespace detail

template <typename... Args>
KernelKeySet ParseKernelKeyByInputArgs(const Args&... args) {
  return detail::KernelKeyParser().apply(args...).key_set;
}

template <typename... Args>
KernelType ParseKernelTypeByInputArgs(const Args&... args) {
  return detail::KernelTypeParser().apply(args...).kernel_type;
}

DataType ParseDataType(DataType dtype);
DataType ParseDataType(const Tensor& tensor);
DataType ParseDataType(const std::vector<Tensor>& tensors);
DataType ParseDataTypeWithInputOrder(DataType dtype, const Tensor& tensor);

PADDLE_API Backend ParseBackend(const Place& place);
Backend ParseBackend(const Tensor& tensor);
template <typename T, typename... Args>
Backend ParseBackend(T t, Args... args) {
  auto backend_set =
      BackendSet(ParseBackend(t)) | BackendSet(ParseBackend(args...));
  return static_cast<Backend>(32 -
                              detail::CountLeadingZeros(backend_set.bitset()));
}
Backend ParseBackendWithInputOrder(const Place& place, const Tensor& tensor);

PADDLE_API phi::DataLayout ParseLayout(phi::DataLayout layout);
phi::DataLayout ParseLayout(const Tensor& tensor);
phi::DataLayout ParseLayoutWithInputOrder(phi::DataLayout layout,
                                          const Tensor& tensor);

template <typename... Args>
bool AllInputsAreDistTensor(const Args&... args) {
  return detail::DistTensorTypeParser().apply(args...).result;
}

}  // namespace experimental
}  // namespace paddle