/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/

#include "kernels/copy_blocks.mluh"
#include <map>
#include <vector>
#include "torch_ops_api.h"
#include "utils.h"

namespace tmo {
namespace torch_api {
void copy_blocks(const std::vector<torch::Tensor> &k_caches,
                 const std::vector<torch::Tensor> &v_caches,
                 const c10::Dict<int64_t, c10::List<int64_t>> &block_mapping_dict) {
  // Create block mapping array.
  std::vector<int32_t> block_mapping_vec;
  for (const auto &item : block_mapping_dict) {
    int64_t src_block_number = item.key();
    auto value_vec = item.value().vec();
    for (int64_t dst_block_number : value_vec) {
      block_mapping_vec.push_back(int32_t(src_block_number));
      block_mapping_vec.push_back(int32_t(dst_block_number));
    }
  }

  TORCH_CHECK(!k_caches.empty() && !block_mapping_vec.empty(),
              "k_caches and block_mapping_vec can not be empty.")
  int32_t num_layers = k_caches.size();
  for (auto i = 0; i < num_layers; i++) {
    TORCH_CHECK(k_caches[i].dim() == 4, "every layer k_cache must be 4d.")
  }
  // check same device and tensor type
  TORCH_CHECK(isMlu(k_caches[0]), "k_caches must on mlu.");
  TORCH_CHECK(k_caches[0].dtype() == torch::kInt8 || k_caches[0].dtype() == torch::kUInt8 ||
                  k_caches[0].dtype() == torch::kInt16 || k_caches[0].dtype() == torch::kInt32 ||
                  k_caches[0].dtype() == torch::kLong || k_caches[0].dtype() == torch::kFloat16 ||
                  k_caches[0].dtype() == torch::kFloat32 || k_caches[0].dtype() == torch::kBFloat16,
              "data type only supports torch::kInt8, torch::kUInt8, torch::kInt16, torch::kInt32, "
              "torch::kLong, torch::kFloat16, torch::kFloat32 and torch::kBFloat16");
  if (!v_caches.empty()) {
    TORCH_CHECK(k_caches.size() == v_caches.size(),
                "k_caches size must equal to "
                "v_caches size if v_caches is not none.")
    TORCH_CHECK(isMlu(v_caches[0]), "v_caches must on mlu.");
    for (auto i = 0; i < num_layers; i++) {
      TORCH_CHECK(k_caches[i].dtype() == v_caches[i].dtype(),
                  "the data type of k_caches and v_caches are not the same.")
      TORCH_CHECK(k_caches[i].dim() == v_caches[i].dim(),
                  "every layer k_cache dim must equal to v_cache dim.")
      // check shape
      TORCH_CHECK(k_caches[i][0].numel() == v_caches[0][0].numel(),
                  "the block_size of k_caches and v_caches are not the same.")
    }
  }

  const torch_mlu::mlu::MLUGuard device_guard(k_caches[0].device());
  auto queue = torch_mlu::getCurMLUStream();

  size_t block_size_bytes = k_caches[0][0].numel() * k_caches[0].element_size();
  std::vector<void *> new_key_caches, new_value_caches;
  for (auto i = 0; i < num_layers; ++i) {
    new_key_caches.push_back(k_caches[i].data_ptr());
    if (!v_caches.empty()) {
      new_value_caches.push_back(v_caches[i].data_ptr());
    }
  }
  TMO_KERNEL_CHECK_FATAL(invokeCopyBlocksKernel(queue, new_key_caches, new_value_caches,
                                                block_mapping_vec, block_size_bytes));
}

std::tuple<std::vector<at::Tensor>, std::vector<at::Tensor>> copy_blocks_out_of_place(
    const std::vector<at::Tensor> &k_caches,
    const std::vector<at::Tensor> &v_caches,
    const c10::Dict<int64_t, c10::List<int64_t>> &block_mapping) {
  copy_blocks(k_caches, v_caches, block_mapping);
  return std::make_tuple(k_caches, v_caches);
}

}  // namespace torch_api
}  // namespace tmo
