/*!
 *  Copyright (c) 2022 by Contributors
 * \file include/dgl/runtime/dlpack_convert.h
 * \brief Conversion between NDArray and DLPack.
 */
#ifndef DGL_RUNTIME_DLPACK_CONVERT_H_
#define DGL_RUNTIME_DLPACK_CONVERT_H_

#include "c_runtime_api.h"
#include "ndarray.h"

struct DLManagedTensor;

namespace dgl {
namespace runtime {

struct DLPackConvert {
  /*!
   * \brief Create a DGL NDArray from a DLPack tensor.
   *
   * This allows us to create a NDArray using the memory
   * allocated by an external deep learning framework
   * that is DLPack compatible.
   *
   * The memory is retained until the NDArray went out of scope.
   * \param tensor The DLPack tensor to copy from.
   * \return The created NDArray view.
   */
  static NDArray FromDLPack(DLManagedTensor* tensor);

  /*!
   * \brief Deleter for NDArray converted from DLPack.
   *
   * This is used from data which is passed from external DLPack(DLManagedTensor)
   * that are not allocated inside of DGL.
   * This enables us to create NDArray from memory allocated by other
   * frameworks that are DLPack compatible
   */
  static void DLPackDeleter(NDArray::Container* ptr);

  /*! \brief Convert a DGL NDArray to a DLPack tensor.
   *
   * \param from The DGL NDArray.
   * \return A DLPack tensor.
   */
  static DLManagedTensor* ToDLPack(const NDArray &from);
};

}  // namespace runtime
}  // namespace dgl

#ifdef __cplusplus
extern "C" {
#endif

/*!
 * \brief Delete (free) a DLManagedTensor's data.
 * \param dltensor Pointer to the DLManagedTensor.
 */
DGL_DLL void DGLDLManagedTensorCallDeleter(DLManagedTensor* dltensor);

/*!
 * \brief Produce an array from the DLManagedTensor that shares data memory
 * with the DLManagedTensor.
 * \param from The source DLManagedTensor.
 * \param out The output array handle.
 * \return 0 when success, -1 when failure happens
 */
DGL_DLL int DGLArrayFromDLPack(DLManagedTensor* from,
                               DGLArrayHandle* out);

/*!
 * \brief Produce a DLMangedTensor from the array that shares data memory with
 * the array.
 * \param from The source array.
 * \param out The DLManagedTensor handle.
 * \return 0 when success, -1 when failure happens
 */
DGL_DLL int DGLArrayToDLPack(DGLArrayHandle from, DLManagedTensor** out,
                             int alignment = 0);

#ifdef __cplusplus
}  // DGL_EXTERN_C
#endif
#endif  // DGL_RUNTIME_DLPACK_CONVERT_H_
