/*************************************************************************
 * 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/swap_blocks.mluh"
#include "torch_ops_api.h"
#include "utils.h"

namespace tmo {
namespace torch_api {
void swap_blocks(torch::Tensor &dst,
                 const torch::Tensor &src,
                 const c10::Dict<int64_t, int64_t> &block_mapping_dict) {
  std::map<int64_t, int64_t> block_mapping;
  for (const auto &item : block_mapping_dict) {
    block_mapping[item.key()] = item.value();
  }

  // check shape
  TORCH_CHECK(src[0].numel() == dst[0].numel(), "the block_size of src and dst are not the same.")

  // check dtype
  TORCH_CHECK(src.dtype() == dst.dtype(), "the data type of src and dst are not the same.")
  TORCH_CHECK(src.dtype() == torch::kInt8 || src.dtype() == torch::kUInt8 ||
                  src.dtype() == torch::kInt16 || src.dtype() == torch::kInt32 ||
                  src.dtype() == torch::kLong || src.dtype() == torch::kFloat16 ||
                  src.dtype() == torch::kFloat32 || src.dtype() == torch::kBFloat16,
              "data type only supports torch::kInt8, torch::kUInt8, torch::kInt16, torch::kInt32, "
              "torch::kLong, torch::kFloat16, torch::kFloat32 and torch::kBFloat16");

  torch::Device src_device = src.device();
  torch::Device dst_device = dst.device();
  torch::Device mlu_device = src_device;
  cnrtMemTransDir_t memcpy_type;
  bool src_is_mlu = isMlu(src);
  bool dst_is_mlu = isMlu(dst);

  if (src_is_mlu && dst_is_mlu) {
    TORCH_CHECK(src_device.index() == dst_device.index(), "src and dst must be on the same MLU");
    memcpy_type = cnrtMemcpyDevToDev;
  } else if (src_is_mlu && dst_device.is_cpu()) {
    memcpy_type = cnrtMemcpyDevToHost;
  } else if (src_device.is_cpu() && dst_is_mlu) {
    memcpy_type = cnrtMemcpyHostToDev;
    mlu_device = dst_device;
  } else {
    TORCH_CHECK(false, "Invalid device combination");
  }

  const torch_mlu::mlu::MLUGuard device_guard(mlu_device);
  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  const int64_t block_size_in_bytes = src.element_size() * src[0].numel();
  TMO_KERNEL_CHECK_FATAL(invokeSwapBlocksKernel(handle, dst.data_ptr(), src.data_ptr(),
                                                block_size_in_bytes, memcpy_type, block_mapping));
}
}  // namespace torch_api
}  // namespace tmo
