python/mxnet/dlpack.py - master - githubmirror/incubator-mxnet

python/mxnet/dlpack.py
213 строк7 KB
Zhenghui Jin
[API] Add new dlpack API (#20546)
29 ноя 2021, 17:22
29 ноя 2021, 17:2245c7999
# 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