chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,890 @@
+/**
+ *  Copyright (c) 2017-2022 by Contributors
+ * @file dgl/runtime/ndarray.h
+ * @brief Abstract device memory management API
+ */
+#ifndef DGL_RUNTIME_NDARRAY_H_
+#define DGL_RUNTIME_NDARRAY_H_
+
+#include <atomic>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "bfloat16.h"
+#include "c_runtime_api.h"
+#include "serializer.h"
+#include "shared_mem.h"
+
+#ifdef DGL_USE_CUDA
+#include <cuda_runtime.h>
+
+#define BF16_ENABLED (defined(CUDART_VERSION) && CUDART_VERSION >= 11000)
+
+#include <cuda_fp16.h>
+#if BF16_ENABLED
+#include <cuda_bf16.h>
+#endif  // BF16_ENABLED
+#endif  // DGL_USE_CUDA
+
+// forward declaration
+inline std::ostream& operator<<(std::ostream& os, DGLDataType t);
+
+namespace dgl {
+
+/**
+ * @brief Type traits that converts a C type to a DGLDataType.
+ *
+ * Usage:
+ * DGLDataTypeTraits<int>::dtype == dtype
+ */
+template <typename T>
+struct DGLDataTypeTraits {
+  static constexpr DGLDataType dtype{0, 0, 0};  // dummy
+};
+#define GEN_DGLDATATYPETRAITS_FOR(T, code, bits)       \
+  template <>                                          \
+  struct DGLDataTypeTraits<T> {                        \
+    static constexpr DGLDataType dtype{code, bits, 1}; \
+  }
+GEN_DGLDATATYPETRAITS_FOR(int8_t, kDGLInt, 8);
+GEN_DGLDATATYPETRAITS_FOR(uint8_t, kDGLUInt, 8);
+GEN_DGLDATATYPETRAITS_FOR(int16_t, kDGLInt, 16);
+GEN_DGLDATATYPETRAITS_FOR(int32_t, kDGLInt, 32);
+GEN_DGLDATATYPETRAITS_FOR(int64_t, kDGLInt, 64);
+// XXX(BarclayII) most DL frameworks do not support unsigned int and long
+// arrays, so I'm just converting uints to signed DTypes.
+GEN_DGLDATATYPETRAITS_FOR(uint32_t, kDGLInt, 32);
+GEN_DGLDATATYPETRAITS_FOR(uint64_t, kDGLInt, 64);
+#ifdef DGL_USE_CUDA
+GEN_DGLDATATYPETRAITS_FOR(__half, kDGLFloat, 16);
+#if BF16_ENABLED
+GEN_DGLDATATYPETRAITS_FOR(__nv_bfloat16, kDGLBfloat, 16);
+#endif  // BF16_ENABLED
+#endif  // DGL_USE_CUDA
+GEN_DGLDATATYPETRAITS_FOR(float, kDGLFloat, 32);
+GEN_DGLDATATYPETRAITS_FOR(double, kDGLFloat, 64);
+#undef GEN_DGLDATATYPETRAITS_FOR
+
+namespace runtime {
+
+/**
+ * @brief DLPack converter.
+ */
+struct DLPackConvert;
+
+/**
+ * @brief Managed NDArray.
+ *  The array is backed by reference counted blocks.
+ */
+class NDArray {
+ public:
+  // internal container type
+  struct Container;
+  /** @brief default constructor */
+  NDArray() {}
+  /**
+   * @brief cosntruct a NDArray that refers to data
+   * @param data The data this NDArray refers to
+   */
+  explicit inline NDArray(Container* data);
+  /**
+   * @brief copy constructor
+   * @param other The value to be copied
+   */
+  inline NDArray(const NDArray& other);  // NOLINT(*)
+  /**
+   * @brief move constructor
+   * @param other The value to be moved
+   */
+  NDArray(NDArray&& other)  // NOLINT(*)
+      : data_(other.data_) {
+    other.data_ = nullptr;
+  }
+  /** @brief destructor */
+  ~NDArray() { this->reset(); }
+  /**
+   * @brief Swap this array with another NDArray
+   * @param other The other NDArray
+   */
+  void swap(NDArray& other) {  // NOLINT(*)
+    std::swap(data_, other.data_);
+  }
+  /**
+   * @brief copy assignmemt
+   * @param other The value to be assigned.
+   * @return reference to self.
+   */
+  NDArray& operator=(const NDArray& other) {  // NOLINT(*)
+    // copy-and-swap idiom
+    NDArray(other).swap(*this);  // NOLINT(*)
+    return *this;
+  }
+  /**
+   * @brief move assignmemt
+   * @param other The value to be assigned.
+   * @return reference to self.
+   */
+  NDArray& operator=(NDArray&& other) {  // NOLINT(*)
+    // copy-and-swap idiom
+    NDArray(std::move(other)).swap(*this);  // NOLINT(*)
+    return *this;
+  }
+  /** @return If NDArray is defined */
+  bool defined() const { return data_ != nullptr; }
+  /** @return If both NDArray reference the same container */
+  bool same_as(const NDArray& other) const { return data_ == other.data_; }
+  /** @brief reset the content of NDArray to be nullptr */
+  inline void reset();
+  /**
+   * @return the reference counter
+   * @note this number is approximate in multi-threaded setting.
+   */
+  inline int use_count() const;
+  /** @return Pointer to content of DGLArray */
+  inline const DGLArray* operator->() const;
+  /** @return True if the ndarray is contiguous. */
+  bool IsContiguous() const;
+  /** @return the data pointer with type. */
+  template <typename T>
+  inline T* Ptr() const {
+    if (!defined())
+      return nullptr;
+    else
+      return static_cast<T*>(operator->()->data);
+  }
+
+  /**
+   * @brief Copy data content from/into another array.
+   * @param other The source array to be copied from.
+   * @note The copy runs on the dgl internal stream if it involves a GPU
+   * context.
+   */
+  inline void CopyFrom(DGLArray* other);
+  inline void CopyFrom(const NDArray& other);
+  inline void CopyTo(DGLArray* other) const;
+  inline void CopyTo(const NDArray& other) const;
+
+  /**
+   * @brief Copy the data to another context.
+   * @param ctx The target context.
+   * @return The array under another context.
+   */
+  inline NDArray CopyTo(const DGLContext& ctx) const;
+
+  /**
+   * @brief Return a new array with a copy of the content.
+   */
+  inline NDArray Clone() const;
+
+  /**
+   * @brief Return a copy of the current instance of NDArray in pinned
+   *     (page-locked) memory.
+   * @note This is an out-of-place method, which utilizes PyTorch's
+   *     CachingHostAllocator for allocating pinned memory and copying data
+   *     from the current NDAarray. As a result, PyTorch is responsible for
+   *     managing the lifecycle of the returned NDArray, including deciding
+   *     when to flush the data for reuse or call cudaFreeHost. The current
+   *     context must be kDGLCPU, otherwise, an error will be thrown.
+   */
+  inline NDArray PinMemory();
+
+  /**
+   * @brief In-place method to pin the current array by calling PinContainer
+   *        on the underlying NDArray:Container.
+   * @note This is an in-place method that flags the memory as page-locked by
+   *     utilizing cudaHostRegister at the underlying level to pin the current
+   *     instance of NDArray. The current context must be kDGLCPU, otherwise,
+   *     an error will be thrown.
+   */
+  inline void PinMemory_();
+
+  /**
+   * @brief In-place method to unpin the current array by calling UnpinContainer
+   *        on the underlying NDArray:Container.
+   * @note This is an in-place method. Behavior depends on the current context,
+   *       IsPinned: will be unpinned;
+   *       others: directly return.
+   */
+  inline void UnpinMemory_();
+
+  /**
+   * @brief Check if the array is pinned.
+   */
+  inline bool IsPinned() const;
+
+  /**
+   * @brief Record streams that are using the underlying tensor.
+   * @param stream The stream that is using the underlying tensor.
+   */
+  inline void RecordStream(DGLStreamHandle stream) const;
+
+  /**
+   * @brief Load NDArray from stream
+   * @param stream The input data stream
+   * @return Whether load is successful
+   */
+  bool Load(dmlc::Stream* stream);
+
+  /**
+   * @brief Save NDArray to stream
+   * @param stream The output data stream
+   */
+  void Save(dmlc::Stream* stream) const;
+
+  /**
+   * @brief Create a NDArray 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 offset The offset (in bytes) of the starting pointer.
+   * @note The memory size of new array must be smaller than the current one.
+   */
+  DGL_DLL NDArray
+  CreateView(std::vector<int64_t> shape, DGLDataType dtype, int64_t offset = 0);
+
+  /**
+   * @brief Create an empty NDArray.
+   * @param shape The shape of the new array.
+   * @param dtype The data type of the new array.
+   * @param ctx The context of the array.
+   * @return The created Array
+   */
+  DGL_DLL static NDArray Empty(
+      std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
+
+  /**
+   * @brief Create an empty NDArray in pinned memory.
+   * @param shape The shape of the new array.
+   * @param dtype The data type of the new array.
+   * @param ctx The context of the array.
+   * @return The created array.
+   */
+  DGL_DLL static NDArray PinnedEmpty(
+      std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
+
+  /**
+   * @brief Create an empty NDArray with shared memory.
+   * @param name The name of shared memory.
+   * @param shape The shape of the new array.
+   * @param dtype The data type of the new array.
+   * @param ctx The context of the array.
+   * @param is_create whether to create shared memory.
+   * @return The created Array
+   */
+  DGL_DLL static NDArray EmptyShared(
+      const std::string& name, std::vector<int64_t> shape, DGLDataType dtype,
+      DGLContext ctx, bool is_create);
+
+  /**
+   * @brief Get the size of the array in the number of bytes.
+   */
+  size_t GetSize() const;
+
+  /**
+   * @brief Get the number of elements in this array.
+   */
+  int64_t NumElements() const;
+
+  /**
+   * @brief Create a NDArray by copying from std::vector.
+   * @tparam T Type of vector data.  Determines the dtype of returned array.
+   */
+  template <typename T>
+  DGL_DLL static NDArray FromVector(
+      const std::vector<T>& vec, DGLContext ctx = DGLContext{kDGLCPU, 0});
+
+  /**
+   * @brief Create a NDArray from a raw pointer.
+   */
+  DGL_DLL static NDArray CreateFromRaw(
+      const std::vector<int64_t>& shape, DGLDataType dtype, DGLContext ctx,
+      void* raw, bool auto_free);
+
+  /**
+   * @brief Create a std::vector from a 1D NDArray.
+   * @tparam T Type of vector data.
+   * @note Type casting is NOT performed.  The caller has to make sure that the
+   * vector type matches the dtype of NDArray.
+   */
+  template <typename T>
+  std::vector<T> ToVector() const;
+
+  std::shared_ptr<SharedMemory> GetSharedMem() const;
+
+  /**
+   * @brief Function to copy data from one array to another.
+   * @param from The source array.
+   * @param to The target array.
+   * @param (optional) stream The stream used in copy.
+   */
+  DGL_DLL static void CopyFromTo(DGLArray* from, DGLArray* to);
+  DGL_DLL static void CopyFromTo(
+      DGLArray* from, DGLArray* to, DGLStreamHandle stream);
+
+  /**
+   * @brief Function to copy data between device and CPU while recording the
+   *     event.
+   * @param from The source array.
+   * @param to The target array.
+   * @param pytorch_ctx The context pointer from PyTorch's CachingHostAllocator.
+   * @note This function fuses data-copy and event recording to ensure
+   *     CachingHostAllocator works properly.
+   */
+  DGL_DLL static void RecordedCopyFromTo(
+      DGLArray* from, DGLArray* to, void* pytorch_ctx);
+
+  /**
+   * @brief Function to pin the DGLArray of a Container.
+   * @param ptr The container to be pinned.
+   * @note Data of the given array will be pinned inplace.
+   *       Behavior depends on the current context,
+   *       kDGLCPU: will be pinned;
+   *       IsPinned: directly return;
+   *       kDGLCUDA: invalid, will throw an error.
+   */
+  DGL_DLL static void PinContainer(Container* ptr);
+
+  /**
+   * @brief Function to unpin the DGLArray of a Container.
+   * @param ptr The container to be unpinned.
+   * @note Data of the given array will be unpinned inplace.
+   *       Behavior depends on the current context,
+   *       IsPinned: will be unpinned;
+   *       others: directly return.
+   */
+  DGL_DLL static void UnpinContainer(Container* ptr);
+
+  /**
+   * @brief Function check if the DGLArray of a Container is pinned.
+   * @param ptr The container to be checked.
+   * @return true if pinned.
+   */
+  DGL_DLL static bool IsContainerPinned(Container* ptr);
+
+  /**
+   * @brief Record streams that are using this tensor.
+   * @param ptr Pointer of the tensor to be recorded.
+   * @param stream The stream that is using this tensor.
+   */
+  DGL_DLL static void RecordStream(DGLArray* tensor, DGLStreamHandle stream);
+
+  // internal namespace
+  struct Internal {
+    // Default deleter for the container
+    static void DefaultDeleter(NDArray::Container* ptr);
+    // Local create function which allocates tensor metadata
+    // but does not allocate space for the data.
+    static NDArray Create(
+        std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
+    // Implementation of API function
+    static DGLArray* MoveAsDGLArray(NDArray arr);
+  };
+
+ private:
+  /** @brief Internal Data content */
+  Container* data_{nullptr};
+  // enable internal functions
+  friend struct Internal;
+  friend struct DLPackConvert;
+  friend class DGLRetValue;
+  friend class DGLArgsSetter;
+};
+
+/**
+ * @brief Save a DGLArray to stream
+ * @param strm The outpu stream
+ * @param tensor The tensor to be saved.
+ */
+inline bool SaveDGLArray(dmlc::Stream* strm, const DGLArray* tensor);
+
+/**
+ * @brief Reference counted Container object used to back NDArray.
+ *
+ *  This object is DGLArray compatible:
+ *    the pointer to the NDArrayContainer can be directly
+ *    interpreted as a DGLArray*
+ *
+ * @note: do not use this function directly, use NDArray.
+ */
+struct NDArray::Container {
+ public:
+  /** NOTE: the first part of this structure is the same as
+   * DLManagedTensor, note that, however, the deleter
+   * is only called when the reference counter goes to 0
+   */
+  /**
+   * @brief Tensor structure.
+   * @note it is important that the first field is DGLArray
+   *  So that this data structure is DGLArray compatible.
+   *  The head ptr of this struct can be viewed as DGLArray*.
+   */
+  DGLArray dl_tensor;
+  /**
+   * @brief addtional context, reserved for recycling
+   * @note We can attach additional content here
+   *  which the current container depend on
+   *  (e.g. reference to original memory when creating views).
+   */
+  void* manager_ctx{nullptr};
+  /**
+   * @brief Customized deleter
+   *
+   * @note The customized deleter is helpful to enable
+   *  different ways of memory allocator that are not
+   *  currently defined by the system.
+   */
+  void (*deleter)(Container* self) = nullptr;
+  /** @brief default constructor */
+  Container() {
+    dl_tensor.data = nullptr;
+    dl_tensor.ndim = 0;
+    dl_tensor.shape = nullptr;
+    dl_tensor.strides = nullptr;
+    dl_tensor.byte_offset = 0;
+  }
+  /** @brief pointer to shared memory */
+  std::shared_ptr<SharedMemory> mem;
+  /** @brief developer function, increases reference counter */
+  void IncRef() { ref_counter_.fetch_add(1, std::memory_order_relaxed); }
+  /** @brief developer function, decrease reference counter */
+  void DecRef() {
+    if (ref_counter_.fetch_sub(1, std::memory_order_release) == 1) {
+      std::atomic_thread_fence(std::memory_order_acquire);
+      if (this->deleter != nullptr) {
+        (*this->deleter)(this);
+      }
+    }
+  }
+
+ private:
+  friend struct DLPackConvert;
+  friend class NDArray;
+  friend class RPCWrappedFunc;
+  /**
+   * @brief The shape container,
+   *  can be used for shape data.
+   */
+  std::vector<int64_t> shape_;
+  /**
+   * @brief The stride container,
+   *  can be used for stride data.
+   */
+  std::vector<int64_t> stride_;
+  /** @brief The internal array object */
+  std::atomic<int> ref_counter_{0};
+
+  /** @brief Whether underlying dl_tensor is pinned by DGL. */
+  bool pinned_by_dgl_{false};
+
+  /** @brief Whether underlying dl_tensor is pinned by PyTorch
+   *    (CachingHostAllocator). */
+  bool pinned_by_pytorch_{false};
+
+  /** @brief The PyTorch storage ctx ptr if pinned_by_pytorch_ = True. */
+  void* pytorch_ctx_{nullptr};
+
+  /** @brief Pointer to the corresp. PyTorch deleter if pinned_by_pytorch_ =
+   *    True.
+   */
+  void* pytorch_raw_deleter_{nullptr};
+};
+
+// implementations of inline functions
+// the usages of functions are documented in place.
+inline NDArray::NDArray(Container* data) : data_(data) {
+  if (data_) data_->IncRef();
+}
+
+inline NDArray::NDArray(const NDArray& other) : data_(other.data_) {
+  if (data_) data_->IncRef();
+}
+
+inline void NDArray::reset() {
+  if (data_) {
+    data_->DecRef();
+    data_ = nullptr;
+  }
+}
+
+inline void NDArray::CopyFrom(DGLArray* other) {
+  CHECK(data_ != nullptr);
+  CopyFromTo(other, &(data_->dl_tensor));
+}
+
+inline void NDArray::CopyFrom(const NDArray& other) {
+  CHECK(other.data_ != nullptr);
+  // Copy between two devices
+  if (data_->dl_tensor.ctx.device_type !=
+      other.data_->dl_tensor.ctx.device_type) {
+    CHECK(data_ != nullptr);
+    auto to_ctx_type = data_->dl_tensor.ctx.device_type;
+    auto cpu_data = (to_ctx_type == kDGLCPU ? data_ : other.data_);
+    // Pinned by PyTorch
+    if (cpu_data->pinned_by_pytorch_) {
+      // To ensure correct behavior, the event must be recorded after
+      // cudaMemcpyAsync as long as the memory is pinned by PyTorch.
+      void* pytorch_ctx = cpu_data->pytorch_ctx_;
+      RecordedCopyFromTo(
+          &(other.data_->dl_tensor), &(data_->dl_tensor), pytorch_ctx);
+      return;
+    }
+  }
+  CopyFrom(&(other.data_->dl_tensor));
+}
+
+inline void NDArray::CopyTo(DGLArray* other) const {
+  CHECK(data_ != nullptr);
+  CopyFromTo(&(data_->dl_tensor), other);
+}
+
+inline void NDArray::CopyTo(const NDArray& other) const {
+  CHECK(other.data_ != nullptr);
+  // copy between two devices
+  if (data_->dl_tensor.ctx.device_type !=
+      other.data_->dl_tensor.ctx.device_type) {
+    CHECK(data_ != nullptr);
+    auto from_ctx_type = data_->dl_tensor.ctx.device_type;
+    auto cpu_data = (from_ctx_type == kDGLCPU ? data_ : other.data_);
+    // pinned by PyTorch
+    if (cpu_data->pinned_by_pytorch_) {
+      // To ensure correct behavior, the event must be recorded after
+      // cudaMemcpyAsync as long as the memory is pinned by PyTorch.
+      void* pytorch_ctx = cpu_data->pytorch_ctx_;
+      RecordedCopyFromTo(
+          &(data_->dl_tensor), &(other.data_->dl_tensor), pytorch_ctx);
+      return;
+    }
+  }
+  CopyTo(&(other.data_->dl_tensor));
+}
+
+inline NDArray NDArray::CopyTo(const DGLContext& ctx) const {
+  CHECK(data_ != nullptr);
+  const DGLArray* array = operator->();
+  NDArray ret = Empty(
+      std::vector<int64_t>(array->shape, array->shape + array->ndim),
+      array->dtype, ctx);
+  this->CopyTo(ret);
+  return ret;
+}
+
+inline NDArray NDArray::Clone() const {
+  CHECK(data_ != nullptr);
+  const DGLArray* array = operator->();
+  return this->CopyTo(array->ctx);
+}
+
+inline NDArray NDArray::PinMemory() {
+  CHECK(data_ != nullptr);
+  const DGLArray* array = operator->();
+  auto ctx = array->ctx;
+  NDArray ret = PinnedEmpty(
+      std::vector<int64_t>(array->shape, array->shape + array->ndim),
+      array->dtype, ctx);
+  this->CopyTo(ret);
+  return ret;
+}
+
+inline void NDArray::PinMemory_() {
+  CHECK(data_ != nullptr);
+  PinContainer(data_);
+}
+
+inline void NDArray::UnpinMemory_() {
+  CHECK(data_ != nullptr);
+  UnpinContainer(data_);
+}
+
+inline bool NDArray::IsPinned() const {
+  CHECK(data_ != nullptr);
+  return IsContainerPinned(data_);
+}
+
+inline void NDArray::RecordStream(DGLStreamHandle stream) const {
+  CHECK(data_ != nullptr);
+  RecordStream(&(data_->dl_tensor), stream);
+}
+
+inline int NDArray::use_count() const {
+  if (data_ == nullptr) return 0;
+  return data_->ref_counter_.load(std::memory_order_relaxed);
+}
+
+inline const DGLArray* NDArray::operator->() const {
+  return &(data_->dl_tensor);
+}
+
+/** @brief Magic number for NDArray file */
+constexpr uint64_t kDGLNDArrayMagic = 0xDD5E40F096B4A13F;
+
+inline bool SaveDGLArray(dmlc::Stream* strm, DGLArray* tensor) {
+  uint64_t header = kDGLNDArrayMagic, 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_ctx) to get a corresponding
+  // array in the target context.
+  DGLContext cpu_ctx;
+  cpu_ctx.device_type = kDGLCPU;
+  cpu_ctx.device_id = 0;
+  strm->Write(cpu_ctx);
+  strm->Write(tensor->ndim);
+  strm->Write(tensor->dtype);
+  int ndim = tensor->ndim;
+  strm->WriteArray(tensor->shape, ndim);
+  int type_bytes = tensor->dtype.bits / 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 (DMLC_IO_NO_ENDIAN_SWAP && tensor->ctx.device_type == kDGLCPU &&
+      tensor->strides == nullptr && tensor->byte_offset == 0) {
+    // quick path
+    strm->Write(tensor->data, data_byte_size);
+  } else {
+    std::vector<uint8_t> bytes(data_byte_size);
+    CHECK_EQ(
+        DGLArrayCopyToBytes(tensor, dmlc::BeginPtr(bytes), data_byte_size), 0)
+        << DGLGetLastError();
+    if (!DMLC_IO_NO_ENDIAN_SWAP) {
+      dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
+    }
+    strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
+  }
+  return true;
+}
+
+/**
+ * @brief Convert type code to its name
+ * @param type_code The type code .
+ * @return The name of type code.
+ */
+inline const char* TypeCode2Str(int type_code) {
+  switch (type_code) {
+    case kDGLInt:
+      return "int";
+    case kDGLUInt:
+      return "uint";
+    case kDGLFloat:
+      return "float";
+    case kStr:
+      return "str";
+    case kBytes:
+      return "bytes";
+    case kHandle:
+      return "handle";
+    case kNull:
+      return "NULL";
+    case kObjectHandle:
+      return "ObjectHandle";
+    case kArrayHandle:
+      return "ArrayHandle";
+    case kDGLDataType:
+      return "DGLDataType";
+    case kDGLContext:
+      return "DGLContext";
+    case kFuncHandle:
+      return "FunctionHandle";
+    case kModuleHandle:
+      return "ModuleHandle";
+    case kNDArrayContainer:
+      return "NDArrayContainer";
+    default:
+      LOG(FATAL) << "unknown type_code=" << static_cast<int>(type_code);
+      return "";
+  }
+}
+
+/**
+ * @brief Convert device type code to its name
+ * @param device_type The device type code.
+ * @return The name of the device.
+ */
+inline const char* DeviceTypeCode2Str(DGLDeviceType device_type) {
+  switch (device_type) {
+    case kDGLCPU:
+      return "cpu";
+    case kDGLCUDA:
+      return "cuda";
+    default:
+      LOG(FATAL) << "Unsupported device type code="
+                 << static_cast<int>(device_type);
+      return "";
+  }
+}
+
+/**
+ * @brief convert a string to DGL type.
+ * @param s The string to be converted.
+ * @return The corresponding dgl type.
+ */
+inline DGLDataType String2DGLDataType(std::string s) {
+  DGLDataType t;
+  t.bits = 32;
+  t.lanes = 1;
+  const char* scan;
+  if (s.substr(0, 3) == "int") {
+    t.code = kDGLInt;
+    scan = s.c_str() + 3;
+  } else if (s.substr(0, 4) == "uint") {
+    t.code = kDGLUInt;
+    scan = s.c_str() + 4;
+  } else if (s.substr(0, 5) == "float") {
+    t.code = kDGLFloat;
+    scan = s.c_str() + 5;
+  } else if (s.substr(0, 6) == "handle") {
+    t.code = kHandle;
+    t.bits = 64;  // handle uses 64 bit by default.
+    scan = s.c_str() + 6;
+  } else {
+    scan = s.c_str();
+    LOG(FATAL) << "unknown type " << s;
+  }
+  char* xdelim;  // emulate sscanf("%ux%u", bits, lanes)
+  uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
+  if (bits != 0) t.bits = bits;
+  if (*xdelim == 'x') {
+    t.lanes = static_cast<uint16_t>(strtoul(xdelim + 1, nullptr, 10));
+  }
+  return t;
+}
+
+/**
+ * @brief convert a DGL type to string.
+ * @param t The type to be converted.
+ * @return The corresponding dgl type in string.
+ */
+inline std::string DGLDataType2String(DGLDataType t) {
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
+  std::ostringstream os;
+  os << t;
+  return os.str();
+#else
+  std::string repr = "";
+  repr += TypeCode2Str(t.code);
+  if (t.code == kHandle) return repr;
+  repr += std::to_string(static_cast<int>(t.bits));
+  if (t.lanes != 1) {
+    repr += "x" + std::to_string(static_cast<int>(t.lanes));
+  }
+  return repr;
+#endif
+}
+
+// macro to check type code.
+#define DGL_CHECK_TYPE_CODE(CODE, T)                                  \
+  CHECK_EQ(CODE, T) << " expected " << TypeCode2Str(T) << " but get " \
+                    << TypeCode2Str(CODE)
+
+}  // namespace runtime
+}  // namespace dgl
+
+namespace dmlc {
+DMLC_DECLARE_TRAITS(has_saveload, dgl::runtime::NDArray, true);
+}  // namespace dmlc
+
+///////////////// Operator overloading for NDArray /////////////////
+dgl::runtime::NDArray operator+(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator-(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator*(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator/(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator%(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator+(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator-(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator*(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator/(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator%(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator+(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator-(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator*(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator/(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator%(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator-(const dgl::runtime::NDArray& array);
+
+dgl::runtime::NDArray operator>(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator<(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator>=(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator<=(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator==(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator!=(
+    const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator>(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator<(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator>=(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator<=(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator==(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator!=(const dgl::runtime::NDArray& a1, int64_t rhs);
+dgl::runtime::NDArray operator>(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator<(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator>=(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator<=(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator==(int64_t lhs, const dgl::runtime::NDArray& a2);
+dgl::runtime::NDArray operator!=(int64_t lhs, const dgl::runtime::NDArray& a2);
+
+std::ostream& operator<<(std::ostream& os, dgl::runtime::NDArray array);
+
+///////////////// Operator overloading for DGLDataType /////////////////
+
+/** @brief Check whether two data types are the same.*/
+inline bool operator==(const DGLDataType& ty1, const DGLDataType& ty2) {
+  return ty1.code == ty2.code && ty1.bits == ty2.bits && ty1.lanes == ty2.lanes;
+}
+
+/** @brief Check whether two data types are different.*/
+inline bool operator!=(const DGLDataType& ty1, const DGLDataType& ty2) {
+  return !(ty1 == ty2);
+}
+
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
+inline std::ostream& operator<<(std::ostream& os, DGLDataType t) {
+  os << dgl::runtime::TypeCode2Str(t.code);
+  if (t.code == kHandle) return os;
+  os << static_cast<int>(t.bits);
+  if (t.lanes != 1) {
+    os << 'x' << static_cast<int>(t.lanes);
+  }
+  return os;
+}
+#endif
+
+///////////////// Operator overloading for DGLContext /////////////////
+
+/** @brief Check whether two device contexts are the same.*/
+inline bool operator==(const DGLContext& ctx1, const DGLContext& ctx2) {
+  return ctx1.device_type == ctx2.device_type &&
+         ctx1.device_id == ctx2.device_id;
+}
+
+/** @brief Check whether two device contexts are different.*/
+inline bool operator!=(const DGLContext& ctx1, const DGLContext& ctx2) {
+  return !(ctx1 == ctx2);
+}
+
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
+inline std::ostream& operator<<(std::ostream& os, const DGLContext& ctx) {
+  return os << dgl::runtime::DeviceTypeCode2Str(ctx.device_type) << ":"
+            << ctx.device_id;
+}
+#endif
+
+#endif  // DGL_RUNTIME_NDARRAY_H_