text: # 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.
  
  # coding: utf-8
  # pylint: disable=protected-access
  # pylint: disable=import-error, no-name-in-module, undefined-variable
  
  """DLPack API of MXNet."""
  
  import ctypes
  import enum
  
  from mxnet.device import current_device
  from .base import _LIB, c_str, check_call, NDArrayHandle, mx_int
  
  DLPackHandle = ctypes.c_void_p
  
  PyCapsuleDestructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
  _c_str_dltensor = c_str('dltensor')
  _c_str_used_dltensor = c_str('used_dltensor')
  
  def _dlpack_deleter(pycapsule):
      pycapsule = ctypes.c_void_p(pycapsule)
      if ctypes.pythonapi.PyCapsule_IsValid(pycapsule, _c_str_dltensor):
          ptr = ctypes.c_void_p(
              ctypes.pythonapi.PyCapsule_GetPointer(pycapsule, _c_str_dltensor))
          check_call(_LIB.MXNDArrayCallDLPackDeleter(ptr))
  
  _c_dlpack_deleter = PyCapsuleDestructor(_dlpack_deleter)
  
  class DLDeviceType(enum.IntEnum):
      DLCPU = 1,
      DLGPU = 2,
      DLCPUPINNED = 3,
      DLOPENCL = 4,
      DLVULKAN = 7,
      DLMETAL = 8,
      DLVPI = 9,
      DLROCM = 10,
      DLEXTDEV = 12,
  
  
  class DLContext(ctypes.Structure):
      _fields_ = [("device_type", ctypes.c_int),
                  ("device_id", ctypes.c_int)]
  
  class DLDataType(ctypes.Structure):
      _fields_ = [("type_code", ctypes.c_uint8),
                  ("bits", ctypes.c_uint8),
                  ("lanes", ctypes.c_uint16)]
      TYPE_MAP = {
          "int32": (0, 32, 1),
          "int64": (0, 64, 1),
          "bool": (1, 1, 1),
          "uint8": (1, 8, 1),
          "uint32": (1, 32, 1),
          "uint64": (1, 64, 1),
          'float16': (2, 16, 1),
          "float32": (2, 32, 1),
          "float64": (2, 64, 1),
      }
  
  
  class DLTensor(ctypes.Structure):
      _fields_ = [("data", ctypes.c_void_p),
                  ("ctx", DLContext),
                  ("ndim", ctypes.c_int),
                  ("dtype", DLDataType),
                  ("shape", ctypes.POINTER(ctypes.c_int64)),
                  ("strides", ctypes.POINTER(ctypes.c_int64)),
                  ("byte_offset", ctypes.c_uint64)]
  
  class DLManagedTensor(ctypes.Structure):
      pass
  
  
  DeleterFunc = ctypes.CFUNCTYPE(None, ctypes.POINTER(DLManagedTensor))
  
  
  DLManagedTensor._fields_ = [("dl_tensor", DLTensor),           # pylint: disable=protected-access
                              ("manager_ctx", ctypes.c_void_p),
                              ("deleter", DeleterFunc)]
  
  @DeleterFunc
  def dl_managed_tensor_deleter(dl_managed_tensor_handle):
      void_p = dl_managed_tensor_handle.contents.manager_ctx
      pyobj = ctypes.cast(void_p, ctypes.py_object)
      ctypes.pythonapi.Py_DecRef(pyobj)
  
  def ndarray_from_dlpack(array_cls):
      """Returns a function that returns specified array_cls from dlpack.
  
      Returns
      -------
      fn : dlpack -> array_cls
      """
      def from_dlpack(dlpack):
          tp = type(dlpack)
          if tp.__module__ == "builtins" and tp.__name__ == "PyCapsule":
              dlpack = ctypes.py_object(dlpack)        
          elif hasattr(dlpack, "__dlpack__"):
              device, device_id = dlpack.__dlpack_device__()
              if device != DLDeviceType.DLGPU:
                  dlpack = ctypes.py_object(dlpack.__dlpack__())
              else:
                  s = mx_int()
                  check_call(_LIB.MXGetCurrentStream(
                      ctypes.c_int(device_id), ctypes.byref(s)))
                  dlpack = ctypes.py_object(dlpack.__dlpack__(stream=s.value))
          else:
              raise AttributeError("Required PyCapsule or object with __dlpack__")
          handle = NDArrayHandle()
          assert ctypes.pythonapi.PyCapsule_IsValid(dlpack, _c_str_dltensor), ValueError(
              'Invalid DLPack Tensor. DLTensor capsules can be consumed only once.')
          dlpack_handle = ctypes.c_void_p(ctypes.pythonapi.PyCapsule_GetPointer(dlpack, _c_str_dltensor))
          check_call(_LIB.MXNDArrayFromDLPack(dlpack_handle, False, ctypes.byref(handle)))
          # Rename PyCapsule (DLPack)
          ctypes.pythonapi.PyCapsule_SetName(dlpack, _c_str_used_dltensor)
          # delete the deleter of the old dlpack
          ctypes.pythonapi.PyCapsule_SetDestructor(dlpack, None)
          return array_cls(handle=handle)
      return from_dlpack
  
  
  def ndarray_to_dlpack_for_read():
      """Returns a function that returns dlpack for reading from mxnet array.
  
      Returns
      -------
      fn : tensor -> dlpack
      """
      def to_dlpack_for_read(data):
          data.wait_to_read()
          dlpack = DLPackHandle()
          check_call(_LIB.MXNDArrayToDLPack(data.handle, ctypes.byref(dlpack)))
          return ctypes.pythonapi.PyCapsule_New(dlpack, _c_str_dltensor, _c_dlpack_deleter)
      return to_dlpack_for_read
  
  def ndarray_to_dlpack_for_write():
      """Returns a function that returns dlpack for writing from mxnet array.
  
      Returns
      -------
      fn : tensor -> dlpack
      """
      def to_dlpack_for_write(data):
  
          check_call(_LIB.MXNDArrayWaitToWrite(data.handle))
          dlpack = DLPackHandle()
          check_call(_LIB.MXNDArrayToDLPack(data.handle, ctypes.byref(dlpack)))
          return ctypes.pythonapi.PyCapsule_New(dlpack, _c_str_dltensor, _c_dlpack_deleter)
      return to_dlpack_for_write
  
  def ndarray_from_numpy(array_cls, array_create_fn):
      """Returns a function that creates array_cls from numpy array.
  
      Returns
      -------
      fn : tensor -> dlpack
      """
      def from_numpy(ndarray, zero_copy=True):
          def _make_manager_ctx(obj):
              pyobj = ctypes.py_object(obj)
              void_p = ctypes.c_void_p.from_buffer(pyobj)
              ctypes.pythonapi.Py_IncRef(pyobj)
              return void_p
  
          def _make_dl_tensor(array):
              if str(array.dtype) not in DLDataType.TYPE_MAP:
                  raise ValueError(str(array.dtype) + " is not supported.")
              dl_tensor = DLTensor()
              dl_tensor.data = array.ctypes.data_as(ctypes.c_void_p)
              dl_tensor.ctx = DLContext(1, 0)
              dl_tensor.ndim = array.ndim
              dl_tensor.dtype = DLDataType.TYPE_MAP[str(array.dtype)]
              dl_tensor.shape = array.ctypes.shape_as(ctypes.c_int64)
              dl_tensor.strides = None
              dl_tensor.byte_offset = 0
              return dl_tensor
  
          def _make_dl_managed_tensor(array):
              c_obj = DLManagedTensor()
              c_obj.dl_tensor = _make_dl_tensor(array)
              c_obj.manager_ctx = _make_manager_ctx(array)
              c_obj.deleter = dl_managed_tensor_deleter
              return c_obj
  
          if not zero_copy:
              return array_create_fn(ndarray, dtype=ndarray.dtype)
  
          if not ndarray.flags['C_CONTIGUOUS']:
              raise ValueError("Only c-contiguous arrays are supported for zero-copy")
  
          ndarray.flags['WRITEABLE'] = False
          c_obj = _make_dl_managed_tensor(ndarray)
          handle = NDArrayHandle()
          check_call(_LIB.MXNDArrayFromDLPack(ctypes.byref(c_obj), True, ctypes.byref(handle)))
          return array_cls(handle=handle)
      return from_numpy
  