/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file tvm/runtime/tensor.h
 * \brief A device-independent managed Tensor abstraction.
 */
#ifndef TVM_RUNTIME_TENSOR_H_
#define TVM_RUNTIME_TENSOR_H_

#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/device_api.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>

#include <atomic>
#include <functional>
#include <utility>
#include <vector>

namespace tvm {
namespace runtime {

/*!
 * \brief Managed Tensor.
 *  The array is backed by reference counted blocks.
 */
class Tensor : public tvm::ffi::Tensor {
 public:
  Tensor() = default;
  /*!
   * \brief constructor.
   * \param data ffi::ObjectPtr to the data container.
   */
  explicit Tensor(ffi::ObjectPtr<ffi::TensorObj> data) : tvm::ffi::Tensor(data) {}
  explicit Tensor(ffi::UnsafeInit tag) : tvm::ffi::Tensor(tag) {}
  Tensor(ffi::Tensor&& other) : tvm::ffi::Tensor(std::move(other)) {}  // NOLINT(*)
  Tensor(const ffi::Tensor& other) : tvm::ffi::Tensor(other) {}        // NOLINT(*)

  ffi::ShapeView Shape() const { return this->shape(); }
  DLDataType DataType() const { return this->dtype(); }

  // DLPack handling
  static Tensor FromDLPack(DLManagedTensor* tensor) {
    return tvm::ffi::Tensor::FromDLPack(tensor, kAllocAlignment, true);
  }

  static Tensor FromDLPackVersioned(DLManagedTensorVersioned* tensor) {
    return tvm::ffi::Tensor::FromDLPackVersioned(tensor, kAllocAlignment, true);
  }
  inline const DLTensor* operator->() const { return this->get(); }
  /*!
   * \brief Copy data content from another array.
   * \param other The source array to be copied from.
   * \note The copy may happen asynchronously if it involves a GPU context.
   *       TVMSynchronize is necessary.
   */
  inline void CopyFrom(const DLTensor* other);
  inline void CopyFrom(const Tensor& other);
  /*!
   * \brief Copy data content from a byte buffer.
   * \param data The source bytes to be copied from.
   * \param nbytes The size of the buffer in bytes
   *        Must be equal to the size of the Tensor.
   * \note The copy always triggers a TVMSynchronize.
   */
  TVM_RUNTIME_DLL void CopyFromBytes(const void* data, size_t nbytes);
  /*!
   * \brief Copy data content into another array.
   * \param other The source array to be copied from.
   * \note The copy may happen asynchronously if it involves a GPU context.
   *       TVMSynchronize is necessary.
   */
  inline void CopyTo(DLTensor* other) const;
  inline void CopyTo(const Tensor& other) const;
  /*!
   * \brief Copy data content into another array.
   * \param data The source bytes to be copied from.
   * \param nbytes The size of the data buffer.
   *        Must be equal to the size of the Tensor.
   * \note The copy always triggers a TVMSynchronize.
   */
  TVM_RUNTIME_DLL void CopyToBytes(void* data, size_t nbytes) const;
  /*!
   * \brief Copy the data to another device.
   * \param dev The target device.
   * \param mem_scope The memory scope of the target array.
   * \return The array under another device.
   * \note The copy always triggers a TVMSynchronize.
   */
  TVM_RUNTIME_DLL Tensor CopyTo(const Device& dev,
                                ffi::Optional<ffi::String> mem_scope = std::nullopt) const;
  /*!
   * \brief Load Tensor from stream
   * \param stream The input data stream
   * \return Whether load is successful
   */
  inline bool Load(support::Stream* stream);
  /*!
   * \brief Save Tensor to stream
   * \param stream The output data stream
   */
  inline void Save(support::Stream* stream) const;

  /*!
   * \brief Create a Tensor that shares the data memory with the current one.
   *
   * \param shape The shape of the new array.
   *
   * \param dtype The data type of the new array.
   *
   * \param relative_byte_offset The offset of the output Tensor,
   *     relative to the current byte offset.
   *
   *     By default, the offset of the view is the same as the offset
   *     of the current array.
   *
   * \note The new array must not allow access of addresses which
   *       would be out of bounds in the current array.  If the new
   *       array is larger than the current array, or if the
   *       `relative_byte_offset` would place the end of the new array
   *       outside the bounds of the current array, this function will
   *       raise an exception.
   */
  TVM_RUNTIME_DLL Tensor CreateView(ffi::Shape shape, DLDataType dtype,
                                    uint64_t relative_byte_offset = 0) const;
  /*!
   * \brief Create an empty Tensor.
   * \param shape The shape of the new array.
   * \param dtype The data type of the new array.
   * \param dev The device of the array.
   * \param mem_scope The memory scope of the array.
   * \return The created Array
   */
  TVM_RUNTIME_DLL static Tensor Empty(ffi::Shape shape, DLDataType dtype, Device dev,
                                      ffi::Optional<ffi::String> mem_scope = std::nullopt);
  /*!
   * \brief Function to copy data from one array to another.
   * \param from The source array.
   * \param to The target array.
   * \param stream The stream used in copy.
   */
  TVM_RUNTIME_DLL static void CopyFromTo(const DLTensor* from, DLTensor* to,
                                         TVMStreamHandle stream = nullptr);

  /*!
   * \brief Function to copy data from one array to a byte buffer.
   * \param from The source array.
   * \param to The target byte buffer.
   * \param nbytes The size of the data buffer.
   * \param stream The stream used in copy.
   */
  TVM_RUNTIME_DLL static void CopyToBytes(const DLTensor* from, void* to, size_t nbytes,
                                          TVMStreamHandle stream = nullptr);

  /*!
   * \brief Function to copy data from one array to a byte buffer.
   * \param from The source array.
   * \param to The target byte buffer.
   * \param nbytes The size of the data buffer.
   * \param stream The stream used in copy.
   */
  TVM_RUNTIME_DLL static void CopyFromBytes(const DLTensor* to, void* from, size_t nbytes,
                                            TVMStreamHandle stream = nullptr);

  /*!
   * \brief Check if two tensors share the same underlying storage.
   *
   * This detects runtime storage aliasing (e.g. views from CreateView, etc.) but does
   * not imply either tensor was created by CreateView.
   *
   * \param a The first tensor.
   * \param b The second tensor.
   * \return True if the tensors share the same storage.
   */
  TVM_RUNTIME_DLL static bool IsStorageShared(const DLTensor* a, const DLTensor* b);

  /*!
   * \brief Tensor overload of IsStorageShared.
   * \param a The first tensor.
   * \param b The second tensor.
   * \return True if the tensors share the same storage.
   */
  static bool IsStorageShared(const Tensor& a, const Tensor& b);
};

/*!
 * \brief Save a DLTensor to stream
 * \param strm The output stream
 * \param tensor The tensor to be saved.
 */
inline bool SaveDLTensor(support::Stream* strm, const DLTensor* tensor);

inline void Tensor::CopyFrom(const DLTensor* other) {
  TVM_FFI_ICHECK(data_ != nullptr);
  CopyFromTo(other, get_mutable());
}

inline void Tensor::CopyFrom(const Tensor& other) {
  TVM_FFI_ICHECK(data_ != nullptr);
  TVM_FFI_ICHECK(other.data_ != nullptr);
  CopyFromTo(other.get_mutable(), get_mutable());
}

inline void Tensor::CopyTo(DLTensor* other) const {
  TVM_FFI_ICHECK(data_ != nullptr);
  CopyFromTo(get_mutable(), other);
}

inline void Tensor::CopyTo(const Tensor& other) const {
  TVM_FFI_ICHECK(data_ != nullptr);
  TVM_FFI_ICHECK(other.data_ != nullptr);
  CopyFromTo(get_mutable(), other.get_mutable());
}

/*! \brief Magic number for Tensor file */
constexpr uint64_t kTVMTensorMagic = 0xDD5E40F096B4A13F;

inline bool SaveDLTensor(support::Stream* strm, const DLTensor* tensor) {
  uint64_t header = kTVMTensorMagic, reserved = 0;
  strm->Write(header);
  strm->Write(reserved);
  // Always save data as CPU context
  //
  // Parameters that get serialized should be in CPU by default.
  // So even the array's context is GPU, it will be stored as CPU array.
  // This is used to prevent case when another user loads the parameters
  // back on machine that do not have GPU or related context.
  //
  // We can always do array.CopyTo(target_dev) to get a corresponding
  // array in the target context.
  Device cpu_dev;
  cpu_dev.device_type = kDLCPU;
  cpu_dev.device_id = 0;
  strm->Write(cpu_dev);
  strm->Write(tensor->ndim);
  strm->Write(tensor->dtype);
  int ndim = tensor->ndim;
  strm->WriteArray(tensor->shape, ndim);
  int type_bytes = (tensor->dtype.bits + 7) / 8;
  int64_t num_elems = 1;
  for (int i = 0; i < ndim; ++i) {
    num_elems *= tensor->shape[i];
  }
  int64_t data_byte_size = type_bytes * num_elems;
  strm->Write(data_byte_size);

  if (TVM_FFI_IO_NO_ENDIAN_SWAP && tensor->device.device_type == kDLCPU &&
      ffi::IsContiguous(*tensor) && tensor->byte_offset == 0) {
    // quick path
    strm->Write(tensor->data, data_byte_size);
  } else {
    std::vector<uint8_t> bytes(data_byte_size);
    Tensor::CopyToBytes(const_cast<DLTensor*>(tensor), bytes.data(), data_byte_size);
    if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
      ffi::ByteSwap(bytes.data(), type_bytes, num_elems);
    }
    strm->Write(bytes.data(), data_byte_size);
  }
  return true;
}

inline void Tensor::Save(support::Stream* strm) const { SaveDLTensor(strm, operator->()); }

inline bool Tensor::Load(support::Stream* strm) {
  uint64_t header, reserved;
  TVM_FFI_ICHECK(strm->Read(&header)) << "Invalid DLTensor file format";
  TVM_FFI_ICHECK(strm->Read(&reserved)) << "Invalid DLTensor file format";
  TVM_FFI_ICHECK(header == kTVMTensorMagic) << "Invalid DLTensor file format";
  Device dev;
  int ndim;
  DLDataType dtype;
  TVM_FFI_ICHECK(strm->Read(&dev)) << "Invalid DLTensor file format";
  TVM_FFI_ICHECK(strm->Read(&ndim)) << "Invalid DLTensor file format";
  TVM_FFI_ICHECK(strm->Read(&dtype)) << "Invalid DLTensor file format";
  TVM_FFI_ICHECK_EQ(dev.device_type, kDLCPU)
      << "Invalid DLTensor device: can only save as CPU tensor";
  std::vector<int64_t> shape(ndim);
  if (ndim != 0) {
    TVM_FFI_ICHECK(strm->ReadArray(&shape[0], ndim)) << "Invalid DLTensor file format";
  }
  Tensor ret = Tensor::Empty(ffi::Shape(shape), dtype, dev);
  int64_t num_elems = 1;
  int elem_bytes = (ret->dtype.bits + 7) / 8;
  for (int i = 0; i < ret->ndim; ++i) {
    num_elems *= ret->shape[i];
  }
  int64_t data_byte_size;
  TVM_FFI_ICHECK(strm->Read(&data_byte_size)) << "Invalid DLTensor file format";
  TVM_FFI_ICHECK(data_byte_size == num_elems * elem_bytes) << "Invalid DLTensor file format";
  auto read_ret = strm->Read(ret->data, data_byte_size);
  // Only check non-empty data
  if (ndim > 0 && shape[0] != 0) {
    TVM_FFI_ICHECK(read_ret) << "Invalid DLTensor file format";
  }
  if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
    ffi::ByteSwap(ret->data, elem_bytes, num_elems);
  }
  *this = ret;
  return true;
}

/*!
 * \brief Get the preferred host device from the input device.
 * - For CUDA and ROCm, CUDAHost and ROCMHost will be returned for pinned memory,
 * since pinned memory reduces copy overhead.
 * - For other devices, CPU is returned as a fallback.
 */
inline Device GetPreferredHostDevice(Device device) {
  if (device.device_type == DLDeviceType::kDLCUDA) {
    return Device{DLDeviceType::kDLCUDAHost, 0};
  } else if (device.device_type == DLDeviceType::kDLROCM) {
    return Device{DLDeviceType::kDLROCMHost, 0};
  } else {
    // Fallback to CPU.
    return Device{DLDeviceType::kDLCPU, 0};
  }
}

}  // namespace runtime
}  // namespace tvm

namespace std {
template <>
struct hash<tvm::Device> {
  std::size_t operator()(const tvm::Device& dev) const {
    return ((dev.device_id << 8) | dev.device_type);
  }
};

template <>
struct equal_to<tvm::Device> {
  bool operator()(const tvm::Device& lhs, const tvm::Device& rhs) const {
    return (lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id);
  }
};
}  // namespace std

#endif  // TVM_RUNTIME_TENSOR_H_