chore: import upstream snapshot with attribution

python/dgl/_ffi/README.md

@@ -0,0 +1,3 @@
+# C API and runtime
+
+Borrowed and adapted from TVM project. (commit: 2ce5277)

python/dgl/_ffi/__init__.py

@@ -0,0 +1,10 @@
+"""C interfacing code.
+
+This namespace contains everything that interacts with C code.
+Most C related object are ctypes compatible, which means
+they contains a handle field that is ctypes.c_void_p and can
+be used via ctypes function calls.
+
+Some performance critical functions are implemented by cython
+and have a ctypes fallback implementation.
+"""

python/dgl/_ffi/_ctypes/__init__.py

@@ -0,0 +1 @@
+"""ctypes specific implementation of FFI"""

python/dgl/_ffi/_ctypes/function.py

@@ -0,0 +1,298 @@
+# coding: utf-8
+# pylint: disable=invalid-name, protected-access, too-many-branches, global-statement
+"""Function configuration API."""
+from __future__ import absolute_import
+
+import ctypes
+import traceback
+from numbers import Integral, Number
+
+from ..base import _LIB, c_str, check_call, string_types
+from ..object_generic import convert_to_object, ObjectGeneric
+from ..runtime_ctypes import DGLByteArray, DGLContext, DGLDataType
+from . import ndarray as _nd, object as _object
+from .ndarray import _make_array, NDArrayBase
+from .object import ObjectBase
+from .types import (
+    _wrap_arg_func,
+    C_TO_PY_ARG_SWITCH,
+    DGLCFuncFinalizer,
+    DGLPackedCFunc,
+    DGLValue,
+    RETURN_SWITCH,
+    TypeCode,
+)
+
+FunctionHandle = ctypes.c_void_p
+ModuleHandle = ctypes.c_void_p
+DGLRetValueHandle = ctypes.c_void_p
+
+
+def _ctypes_free_resource(rhandle):
+    """callback to free resources when it it not needed."""
+    pyobj = ctypes.cast(rhandle, ctypes.py_object)
+    ctypes.pythonapi.Py_DecRef(pyobj)
+
+
+# Global callback that is always alive
+DGL_FREE_PYOBJ = DGLCFuncFinalizer(_ctypes_free_resource)
+ctypes.pythonapi.Py_IncRef(ctypes.py_object(DGL_FREE_PYOBJ))
+
+
+def convert_to_dgl_func(pyfunc):
+    """Convert a python function to DGL function
+
+    Parameters
+    ----------
+    pyfunc : python function
+        The python function to be converted.
+
+    Returns
+    -------
+    dglfunc: dgl.nd.Function
+        The converted dgl function.
+    """
+    local_pyfunc = pyfunc
+
+    def cfun(args, type_codes, num_args, ret, _):
+        """ctypes function"""
+        num_args = (
+            num_args.value if isinstance(num_args, ctypes.c_int) else num_args
+        )
+        pyargs = (
+            C_TO_PY_ARG_SWITCH[type_codes[i]](args[i]) for i in range(num_args)
+        )
+        # pylint: disable=broad-except
+        try:
+            rv = local_pyfunc(*pyargs)
+        except Exception:
+            msg = traceback.format_exc()
+            _LIB.DGLAPISetLastError(c_str(msg))
+            return -1
+
+        if rv is not None:
+            if isinstance(rv, tuple):
+                raise ValueError(
+                    "PackedFunction can only support one return value"
+                )
+            temp_args = []
+            values, tcodes, _ = _make_dgl_args((rv,), temp_args)
+            if not isinstance(ret, DGLRetValueHandle):
+                ret = DGLRetValueHandle(ret)
+            check_call(
+                _LIB.DGLCFuncSetReturn(ret, values, tcodes, ctypes.c_int(1))
+            )
+            _ = temp_args
+            _ = rv
+        return 0
+
+    handle = FunctionHandle()
+    f = DGLPackedCFunc(cfun)
+    # NOTE: We will need to use python-api to increase ref count of the f
+    # DGL_FREE_PYOBJ will be called after it is no longer needed.
+    pyobj = ctypes.py_object(f)
+    ctypes.pythonapi.Py_IncRef(pyobj)
+    check_call(
+        _LIB.DGLFuncCreateFromCFunc(
+            f, pyobj, DGL_FREE_PYOBJ, ctypes.byref(handle)
+        )
+    )
+    return _CLASS_FUNCTION(handle, False)
+
+
+def _make_dgl_args(args, temp_args):
+    """Pack arguments into c args dgl call accept.
+
+    temp_args is used to temporarily save the arguments so they will not be
+    freed during C API function call.
+    """
+    num_args = len(args)
+    values = (DGLValue * num_args)()
+    type_codes = (ctypes.c_int * num_args)()
+    for i, arg in enumerate(args):
+        if arg is None:
+            values[i].v_handle = None
+            type_codes[i] = TypeCode.NULL
+        elif isinstance(arg, ObjectBase):
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.OBJECT_HANDLE
+        elif isinstance(arg, (list, tuple, dict, ObjectGeneric)):
+            arg = convert_to_object(arg)
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.OBJECT_HANDLE
+            temp_args.append(arg)
+        elif isinstance(arg, NDArrayBase):
+            values[i].v_handle = ctypes.cast(arg.handle, ctypes.c_void_p)
+            type_codes[i] = (
+                TypeCode.NDARRAY_CONTAINER
+                if not arg.is_view
+                else TypeCode.ARRAY_HANDLE
+            )
+        elif isinstance(arg, _nd._DGL_COMPATS):
+            values[i].v_handle = ctypes.c_void_p(arg._dgl_handle)
+            type_codes[i] = arg.__class__._dgl_tcode
+        elif isinstance(arg, Integral):
+            values[i].v_int64 = arg
+            type_codes[i] = TypeCode.INT
+        elif isinstance(arg, Number):
+            values[i].v_float64 = arg
+            type_codes[i] = TypeCode.FLOAT
+        elif isinstance(arg, DGLDataType):
+            values[i].v_str = c_str(str(arg))
+            type_codes[i] = TypeCode.STR
+        elif isinstance(arg, DGLContext):
+            values[i].v_ctx = arg
+            type_codes[i] = TypeCode.DGL_CONTEXT
+        elif isinstance(arg, bytearray):
+            arr = DGLByteArray()
+            arr.data = ctypes.cast(
+                (ctypes.c_byte * len(arg)).from_buffer(arg),
+                ctypes.POINTER(ctypes.c_byte),
+            )
+            arr.size = len(arg)
+            values[i].v_handle = ctypes.c_void_p(ctypes.addressof(arr))
+            temp_args.append(arr)
+            type_codes[i] = TypeCode.BYTES
+        elif isinstance(arg, string_types):
+            values[i].v_str = c_str(arg)
+            type_codes[i] = TypeCode.STR
+        # NOTE(minjie): module is not used in DGL
+        # elif isinstance(arg, _CLASS_MODULE):
+        #    values[i].v_handle = arg.handle
+        #    type_codes[i] = TypeCode.MODULE_HANDLE
+        elif isinstance(arg, FunctionBase):
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.FUNC_HANDLE
+        elif isinstance(arg, ctypes.c_void_p):
+            values[i].v_handle = arg
+            type_codes[i] = TypeCode.HANDLE
+        elif callable(arg):
+            arg = convert_to_dgl_func(arg)
+            values[i].v_handle = arg.handle
+            type_codes[i] = TypeCode.FUNC_HANDLE
+            temp_args.append(arg)
+        else:
+            raise TypeError("Don't know how to handle type %s" % type(arg))
+    return values, type_codes, num_args
+
+
+class FunctionBase(object):
+    """Function base."""
+
+    __slots__ = ["handle", "is_global"]
+    # pylint: disable=no-member
+    def __init__(self, handle, is_global):
+        """Initialize the function with handle
+
+        Parameters
+        ----------
+        handle : FunctionHandle
+            the handle to the underlying function.
+
+        is_global : bool
+            Whether this is a global function in python
+        """
+        self.handle = handle
+        self.is_global = is_global
+
+    def __del__(self):
+        if not self.is_global and _LIB is not None:
+            check_call(_LIB.DGLFuncFree(self.handle))
+
+    def __call__(self, *args):
+        """Call the function with positional arguments
+
+        args : list
+           The positional arguments to the function call.
+        """
+        temp_args = []
+        values, tcodes, num_args = _make_dgl_args(args, temp_args)
+        ret_val = DGLValue()
+        ret_tcode = ctypes.c_int()
+        check_call(
+            _LIB.DGLFuncCall(
+                self.handle,
+                values,
+                tcodes,
+                ctypes.c_int(num_args),
+                ctypes.byref(ret_val),
+                ctypes.byref(ret_tcode),
+            )
+        )
+        _ = temp_args
+        _ = args
+        return RETURN_SWITCH[ret_tcode.value](ret_val)
+
+
+def __init_handle_by_constructor__(fconstructor, args):
+    """Initialize handle by constructor"""
+    temp_args = []
+    values, tcodes, num_args = _make_dgl_args(args, temp_args)
+    ret_val = DGLValue()
+    ret_tcode = ctypes.c_int()
+    check_call(
+        _LIB.DGLFuncCall(
+            fconstructor.handle,
+            values,
+            tcodes,
+            ctypes.c_int(num_args),
+            ctypes.byref(ret_val),
+            ctypes.byref(ret_tcode),
+        )
+    )
+    _ = temp_args
+    _ = args
+    assert ret_tcode.value == TypeCode.OBJECT_HANDLE
+    handle = ret_val.v_handle
+    return handle
+
+
+def _return_module(x):
+    """Return function"""
+    handle = x.v_handle
+    if not isinstance(handle, ModuleHandle):
+        handle = ModuleHandle(handle)
+    return _CLASS_MODULE(handle)
+
+
+def _handle_return_func(x):
+    """Return function"""
+    handle = x.v_handle
+    if not isinstance(handle, FunctionHandle):
+        handle = FunctionHandle(handle)
+    return _CLASS_FUNCTION(handle, False)
+
+
+# setup return handle for function type
+_object.__init_by_constructor__ = __init_handle_by_constructor__
+RETURN_SWITCH[TypeCode.FUNC_HANDLE] = _handle_return_func
+RETURN_SWITCH[TypeCode.MODULE_HANDLE] = _return_module
+RETURN_SWITCH[TypeCode.NDARRAY_CONTAINER] = lambda x: _make_array(
+    x.v_handle, False
+)
+C_TO_PY_ARG_SWITCH[TypeCode.FUNC_HANDLE] = _wrap_arg_func(
+    _handle_return_func, TypeCode.FUNC_HANDLE
+)
+C_TO_PY_ARG_SWITCH[TypeCode.MODULE_HANDLE] = _wrap_arg_func(
+    _return_module, TypeCode.MODULE_HANDLE
+)
+C_TO_PY_ARG_SWITCH[TypeCode.ARRAY_HANDLE] = lambda x: _make_array(
+    x.v_handle, True
+)
+C_TO_PY_ARG_SWITCH[TypeCode.NDARRAY_CONTAINER] = lambda x: _make_array(
+    x.v_handle, False
+)
+
+_CLASS_MODULE = None
+_CLASS_FUNCTION = None
+
+
+def _set_class_module(module_class):
+    """Initialize the module."""
+    global _CLASS_MODULE
+    _CLASS_MODULE = module_class
+
+
+def _set_class_function(func_class):
+    global _CLASS_FUNCTION
+    _CLASS_FUNCTION = func_class

python/dgl/_ffi/_ctypes/ndarray.py

@@ -0,0 +1,137 @@
+# pylint: disable=invalid-name
+"""Runtime NDArray api"""
+from __future__ import absolute_import
+
+import ctypes
+
+from ..base import _LIB, c_str, check_call
+from ..runtime_ctypes import DGLArrayHandle
+from .types import (
+    _return_handle,
+    _wrap_arg_func,
+    C_TO_PY_ARG_SWITCH,
+    RETURN_SWITCH,
+)
+
+DGLPyCapsuleDestructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+_c_str_dltensor = c_str("dltensor")
+_c_str_used_dltensor = c_str("used_dltensor")
+
+
+# used for PyCapsule manipulation
+if hasattr(ctypes, "pythonapi"):
+    ctypes.pythonapi.PyCapsule_GetName.restype = ctypes.c_char_p
+    ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
+    ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object
+
+
+def _from_dlpack(dltensor):
+    dltensor = ctypes.py_object(dltensor)
+    if ctypes.pythonapi.PyCapsule_IsValid(dltensor, _c_str_dltensor):
+        ptr = ctypes.pythonapi.PyCapsule_GetPointer(dltensor, _c_str_dltensor)
+        # XXX(minjie): The below cast should be unnecessary given the code to
+        #   set restype of PyCapsule calls. But weirdly, this does not
+        #   work out always.
+        ptr = ctypes.cast(ptr, ctypes.c_void_p)
+        handle = DGLArrayHandle()
+        check_call(_LIB.DGLArrayFromDLPack(ptr, ctypes.byref(handle)))
+        ctypes.pythonapi.PyCapsule_SetName(dltensor, _c_str_used_dltensor)
+        ctypes.pythonapi.PyCapsule_SetDestructor(
+            dltensor, DGLPyCapsuleDestructor(0)
+        )
+        return _make_array(handle, False)
+    raise ValueError(
+        "Expect a dltensor field, PyCapsule can only be consumed once"
+    )
+
+
+def _dlpack_deleter(pycapsule):
+    pycapsule = ctypes.cast(pycapsule, ctypes.py_object)
+    if ctypes.pythonapi.PyCapsule_IsValid(pycapsule, _c_str_dltensor):
+        ptr = ctypes.pythonapi.PyCapsule_GetPointer(pycapsule, _c_str_dltensor)
+        # XXX(minjie): The below cast should be unnecessary given the code to
+        #   set restype of PyCapsule calls. But weirdly, this does not
+        #   work out always.
+        ptr = ctypes.cast(ptr, ctypes.c_void_p)
+        _LIB.DGLDLManagedTensorCallDeleter(ptr)
+        ctypes.pythonapi.PyCapsule_SetDestructor(
+            pycapsule, DGLPyCapsuleDestructor(0)
+        )
+
+
+_c_dlpack_deleter = DGLPyCapsuleDestructor(_dlpack_deleter)
+
+
+class NDArrayBase(object):
+    """A simple Device/CPU Array object in runtime."""
+
+    __slots__ = ["handle", "is_view"]
+    # pylint: disable=no-member
+    def __init__(self, handle, is_view=False):
+        """Initialize the function with handle
+
+        Parameters
+        ----------
+        handle : DGLArrayHandle
+            the handle to the underlying C++ DGLArray
+        """
+        self.handle = handle
+        self.is_view = is_view
+
+    def __del__(self):
+        if not self.is_view and _LIB:
+            check_call(_LIB.DGLArrayFree(self.handle))
+
+    @property
+    def _dgl_handle(self):
+        return ctypes.cast(self.handle, ctypes.c_void_p).value
+
+    def to_dlpack(self, alignment=0):
+        """Produce an array from a DLPack Tensor without copying memory
+
+        Args
+        -------
+        alignment: int, default to be 0
+        Indicates the alignment requirement when converting to dlpack. Will copy to a
+        new tensor if the alignment requirement is not satisfied.
+        0 means no alignment requirement.
+
+
+        Returns
+        -------
+        dlpack : DLPack tensor view of the array data
+        """
+        ptr = ctypes.c_void_p()
+        check_call(
+            _LIB.DGLArrayToDLPack(self.handle, ctypes.byref(ptr), alignment)
+        )
+        return ctypes.pythonapi.PyCapsule_New(
+            ptr, _c_str_dltensor, _c_dlpack_deleter
+        )
+
+
+def _make_array(handle, is_view):
+    handle = ctypes.cast(handle, DGLArrayHandle)
+    return _CLASS_NDARRAY(handle, is_view)
+
+
+_DGL_COMPATS = ()
+
+
+def _reg_extension(cls, fcreate):
+    global _DGL_COMPATS
+    _DGL_COMPATS += (cls,)
+    if fcreate:
+        fret = lambda x: fcreate(_return_handle(x))
+        RETURN_SWITCH[cls._dgl_tcode] = fret
+        C_TO_PY_ARG_SWITCH[cls._dgl_tcode] = _wrap_arg_func(
+            fret, cls._dgl_tcode
+        )
+
+
+_CLASS_NDARRAY = None
+
+
+def _set_class_ndarray(cls):
+    global _CLASS_NDARRAY
+    _CLASS_NDARRAY = cls

python/dgl/_ffi/_ctypes/object.py

@@ -0,0 +1,109 @@
+"""ctypes object API."""
+from __future__ import absolute_import
+
+import ctypes
+
+from ..base import _LIB, c_str, check_call
+from ..object_generic import _set_class_object_base
+from .types import (
+    _wrap_arg_func,
+    C_TO_PY_ARG_SWITCH,
+    DGLValue,
+    RETURN_SWITCH,
+    TypeCode,
+)
+
+ObjectHandle = ctypes.c_void_p
+__init_by_constructor__ = None
+
+"""Maps object type to its constructor"""
+OBJECT_TYPE = {}
+
+
+def _register_object(index, cls):
+    """register object class in python"""
+    OBJECT_TYPE[index] = cls
+
+
+def _return_object(x):
+    """Construct a object object from the given DGLValue object"""
+    handle = x.v_handle
+    if not isinstance(handle, ObjectHandle):
+        handle = ObjectHandle(handle)
+    tindex = ctypes.c_int()
+    check_call(_LIB.DGLObjectGetTypeIndex(handle, ctypes.byref(tindex)))
+    cls = OBJECT_TYPE.get(tindex.value, ObjectBase)
+    # Avoid calling __init__ of cls, instead directly call __new__
+    # This allows child class to implement their own __init__
+    obj = cls.__new__(cls)
+    obj.handle = handle
+    return obj
+
+
+RETURN_SWITCH[TypeCode.OBJECT_HANDLE] = _return_object
+C_TO_PY_ARG_SWITCH[TypeCode.OBJECT_HANDLE] = _wrap_arg_func(
+    _return_object, TypeCode.OBJECT_HANDLE
+)
+
+
+class ObjectBase(object):
+    """Object base class"""
+
+    __slots__ = ["handle"]
+
+    # pylint: disable=no-member
+    def __del__(self):
+        if _LIB is not None and hasattr(self, "handle"):
+            check_call(_LIB.DGLObjectFree(self.handle))
+
+    def __getattr__(self, name):
+        if name == "handle":
+            raise AttributeError(
+                "'handle' is a reserved attribute name that should not be used"
+            )
+        ret_val = DGLValue()
+        ret_type_code = ctypes.c_int()
+        ret_success = ctypes.c_int()
+        check_call(
+            _LIB.DGLObjectGetAttr(
+                self.handle,
+                c_str(name),
+                ctypes.byref(ret_val),
+                ctypes.byref(ret_type_code),
+                ctypes.byref(ret_success),
+            )
+        )
+        if not ret_success.value:
+            raise AttributeError(
+                "'%s' object has no attribute '%s'" % (str(type(self)), name)
+            )
+        return RETURN_SWITCH[ret_type_code.value](ret_val)
+
+    def __init_handle_by_constructor__(self, fconstructor, *args):
+        """Initialize the handle by calling constructor function.
+
+        Parameters
+        ----------
+        fconstructor : Function
+            Constructor function.
+
+        args: list of objects
+            The arguments to the constructor
+
+        Note
+        ----
+        We have a special calling convention to call constructor functions.
+        So the return handle is directly set into the Object object
+        instead of creating a new Object.
+        """
+        # assign handle first to avoid error raising
+        self.handle = None
+        handle = __init_by_constructor__(
+            fconstructor, args
+        )  # pylint: disable=not-callable
+        if not isinstance(handle, ObjectHandle):
+            handle = ObjectHandle(handle)
+        self.handle = handle
+
+
+_set_class_object_base(ObjectBase)

python/dgl/_ffi/_ctypes/types.py

@@ -0,0 +1,91 @@
+"""The C Types used in API."""
+# pylint: disable=invalid-name
+from __future__ import absolute_import as _abs
+
+import ctypes
+
+from ..base import _LIB, check_call, py_str
+from ..runtime_ctypes import DGLByteArray, DGLContext, DGLDataType, TypeCode
+
+
+class DGLValue(ctypes.Union):
+    """DGLValue in C API"""
+
+    _fields_ = [
+        ("v_int64", ctypes.c_int64),
+        ("v_float64", ctypes.c_double),
+        ("v_handle", ctypes.c_void_p),
+        ("v_str", ctypes.c_char_p),
+        ("v_type", DGLDataType),
+        ("v_ctx", DGLContext),
+    ]
+
+
+DGLPackedCFunc = ctypes.CFUNCTYPE(
+    ctypes.c_int,
+    ctypes.POINTER(DGLValue),
+    ctypes.POINTER(ctypes.c_int),
+    ctypes.c_int,
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+)
+
+
+DGLCFuncFinalizer = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+
+
+def _return_handle(x):
+    """return handle"""
+    handle = x.v_handle
+    if not isinstance(handle, ctypes.c_void_p):
+        handle = ctypes.c_void_p(handle)
+    return handle
+
+
+def _return_bytes(x):
+    """return handle"""
+    handle = x.v_handle
+    if not isinstance(handle, ctypes.c_void_p):
+        handle = ctypes.c_void_p(handle)
+    arr = ctypes.cast(handle, ctypes.POINTER(DGLByteArray))[0]
+    size = arr.size
+    res = bytearray(size)
+    rptr = (ctypes.c_byte * size).from_buffer(res)
+    if not ctypes.memmove(rptr, arr.data, size):
+        raise RuntimeError("memmove failed")
+    return res
+
+
+def _wrap_arg_func(return_f, type_code):
+    tcode = ctypes.c_int(type_code)
+
+    def _wrap_func(x):
+        check_call(_LIB.DGLCbArgToReturn(ctypes.byref(x), tcode))
+        return return_f(x)
+
+    return _wrap_func
+
+
+RETURN_SWITCH = {
+    TypeCode.INT: lambda x: x.v_int64,
+    TypeCode.FLOAT: lambda x: x.v_float64,
+    TypeCode.HANDLE: _return_handle,
+    TypeCode.NULL: lambda x: None,
+    TypeCode.STR: lambda x: py_str(x.v_str),
+    TypeCode.BYTES: _return_bytes,
+    TypeCode.DGL_CONTEXT: lambda x: DGLContext(
+        x.v_ctx.device_type, x.v_ctx.device_id
+    ),
+}
+
+C_TO_PY_ARG_SWITCH = {
+    TypeCode.INT: lambda x: x.v_int64,
+    TypeCode.FLOAT: lambda x: x.v_float64,
+    TypeCode.HANDLE: _return_handle,
+    TypeCode.NULL: lambda x: None,
+    TypeCode.STR: lambda x: py_str(x.v_str),
+    TypeCode.BYTES: _return_bytes,
+    TypeCode.DGL_CONTEXT: lambda x: DGLContext(
+        x.v_ctx.device_type, x.v_ctx.device_id
+    ),
+}

python/dgl/_ffi/_cy2/__init__.py

@@ -0,0 +1 @@
+"""cython2 namespace"""

python/dgl/_ffi/_cy3/__init__.py

@@ -0,0 +1 @@
+"""cython3 namespace"""

python/dgl/_ffi/_cython/.gitignore

@@ -0,0 +1 @@
+*.cpp

python/dgl/_ffi/_cython/base.pxi

@@ -0,0 +1,169 @@
+from ..base import DGLError
+from libcpp.vector cimport vector
+from libcpp cimport bool
+from cpython.version cimport PY_MAJOR_VERSION
+from cpython cimport pycapsule
+from libc.stdint cimport int32_t, int64_t, uint64_t, uint8_t, uint16_t
+import ctypes
+
+cdef enum DGLObjectTypeCode:
+    kObjectInt = 0
+    kObjectUInt = 1
+    kObjectFloat = 2
+    kHandle = 3
+    kNull = 4
+    kDGLDataType = 5
+    kDGLContext = 6
+    kArrayHandle = 7
+    kObjectHandle = 8
+    kModuleHandle = 9
+    kFuncHandle = 10
+    kStr = 11
+    kBytes = 12
+    kNDArrayContainer = 13
+    kExtBegin = 15
+
+cdef extern from "dgl/runtime/c_runtime_api.h":
+    ctypedef struct DGLDataType:
+        uint8_t code
+        uint8_t bits
+        uint16_t lanes
+
+    ctypedef struct DGLContext:
+        int32_t device_type
+        int32_t device_id
+
+    ctypedef struct DGLArray:
+        void* data
+        DGLContext ctx
+        int32_t ndim
+        DGLDataType dtype
+        int64_t* shape
+        int64_t* strides
+        uint64_t byte_offset
+
+    ctypedef struct DLManagedTensor:
+        DGLArray dl_tensor
+        void* manager_ctx
+        void (*deleter)(DLManagedTensor* self)
+
+    ctypedef struct DGLValue:
+        int64_t v_int64
+        double v_float64
+        void* v_handle
+        const char* v_str
+        DGLDataType v_type
+        DGLContext v_ctx
+
+ctypedef int64_t dgl_index_t
+ctypedef DGLArray* DGLArrayHandle
+ctypedef void* DGLStreamHandle
+ctypedef void* DGLRetValueHandle
+ctypedef void* DGLFunctionHandle
+ctypedef void* ObjectHandle
+
+ctypedef int (*DGLPackedCFunc)(
+    DGLValue* args,
+    int* type_codes,
+    int num_args,
+    DGLRetValueHandle ret,
+    void* resource_handle)
+
+ctypedef void (*DGLPackedCFuncFinalizer)(void* resource_handle)
+
+cdef extern from "dgl/runtime/c_runtime_api.h":
+    void DGLAPISetLastError(const char* msg)
+    const char *DGLGetLastError()
+    int DGLFuncCall(DGLFunctionHandle func,
+                    DGLValue* arg_values,
+                    int* type_codes,
+                    int num_args,
+                    DGLValue* ret_val,
+                    int* ret_type_code) nogil
+    int DGLFuncFree(DGLFunctionHandle func)
+    int DGLCFuncSetReturn(DGLRetValueHandle ret,
+                          DGLValue* value,
+                          int* type_code,
+                          int num_ret)
+    int DGLFuncCreateFromCFunc(DGLPackedCFunc func,
+                               void* resource_handle,
+                               DGLPackedCFuncFinalizer fin,
+                               DGLFunctionHandle *out)
+    int DGLCbArgToReturn(DGLValue* value, int code)
+    int DGLArrayAlloc(dgl_index_t* shape,
+                      dgl_index_t ndim,
+                      DGLDataType dtype,
+                      DGLContext ctx,
+                      DGLArrayHandle* out)
+    int DGLArrayAllocSharedMem(const char *mem_name,
+                               const dgl_index_t *shape,
+                               int ndim,
+                               int dtype_code,
+                               int dtype_bits,
+                               int dtype_lanes,
+                               bool is_create,
+                               DGLArrayHandle* out)
+    int DGLArrayFree(DGLArrayHandle handle)
+    int DGLArrayCopyFromTo(DGLArrayHandle src,
+                           DGLArrayHandle to)
+
+cdef extern from "dgl/runtime/c_object_api.h":
+    int DGLObjectFree(ObjectHandle handle)
+    int DGLObjectTypeKey2Index(const char* type_key,
+                               int* out_index)
+    int DGLObjectGetTypeIndex(ObjectHandle handle,
+                              int* out_index)
+    int DGLObjectGetAttr(ObjectHandle handle,
+                         const char* key,
+                         DGLValue* out_value,
+                         int* out_type_code,
+                         int* out_success)
+
+cdef extern from "dgl/runtime/dlpack_convert.h":
+    int DGLArrayFromDLPack(DLManagedTensor* arr_from,
+                           DGLArrayHandle* out)
+    int DGLArrayToDLPack(DGLArrayHandle arr_from,
+                         DLManagedTensor** out,
+                         int alignment)
+    void DGLDLManagedTensorCallDeleter(DLManagedTensor* dltensor)
+
+cdef inline py_str(const char* x):
+    if PY_MAJOR_VERSION < 3:
+        return x
+    else:
+        return x.decode("utf-8")
+
+
+cdef inline c_str(pystr):
+    """Create ctypes char * from a python string
+    Parameters
+    ----------
+    string : string type
+        python string
+
+    Returns
+    -------
+    str : c_char_p
+        A char pointer that can be passed to C API
+    """
+    return pystr.encode("utf-8")
+
+
+cdef inline CALL(int ret):
+    if ret != 0:
+        raise DGLError(py_str(DGLGetLastError()))
+
+
+cdef inline object ctypes_handle(void* chandle):
+    """Cast C handle to ctypes handle."""
+    return ctypes.cast(<unsigned long long>chandle, ctypes.c_void_p)
+
+
+cdef inline void* c_handle(object handle):
+    """Cast C types handle to c handle."""
+    cdef unsigned long long v_ptr
+    if handle.value is None:
+        return NULL
+    else:
+        v_ptr = handle.value
+        return <void*>(v_ptr)

python/dgl/_ffi/_cython/core.pyx

@@ -0,0 +1,4 @@
+include "./base.pxi"
+include "./object.pxi"
+include "./function.pxi"
+include "./ndarray.pxi"

python/dgl/_ffi/_cython/function.pxi

@@ -0,0 +1,308 @@
+import ctypes
+import traceback
+from cpython cimport Py_INCREF, Py_DECREF
+from numbers import Number, Integral
+from ..base import string_types
+from ..object_generic import convert_to_object, ObjectGeneric
+from ..runtime_ctypes import DGLDataType as CTypesDGLDataType, \
+                             DGLContext as CTypesDGLContext, \
+                             DGLByteArray
+
+
+cdef void dgl_callback_finalize(void* fhandle):
+    local_pyfunc = <object>(fhandle)
+    Py_DECREF(local_pyfunc)
+
+cdef int dgl_callback(DGLValue* args,
+                      int* type_codes,
+                      int num_args,
+                      DGLRetValueHandle ret,
+                      void* fhandle) with gil:
+    cdef list pyargs
+    cdef DGLValue value
+    cdef int tcode
+    local_pyfunc = <object>(fhandle)
+    pyargs = []
+    for i in range(num_args):
+        value = args[i]
+        tcode = type_codes[i]
+        if (tcode == kObjectHandle or
+            tcode == kFuncHandle or
+            tcode == kModuleHandle or
+            tcode > kExtBegin):
+            CALL(DGLCbArgToReturn(&value, tcode))
+
+        if tcode != kArrayHandle:
+            pyargs.append(make_ret(value, tcode))
+        else:
+            pyargs.append(c_make_array(value.v_handle, True))
+    try:
+        rv = local_pyfunc(*pyargs)
+    except Exception:
+        msg = traceback.format_exc()
+        DGLAPISetLastError(c_str(msg))
+        return -1
+    if rv is not None:
+        if isinstance(rv, tuple):
+            raise ValueError("PackedFunction can only support one return value")
+        temp_args = []
+        make_arg(rv, &value, &tcode, temp_args)
+        CALL(DGLCFuncSetReturn(ret, &value, &tcode, 1))
+    return 0
+
+
+def convert_to_dgl_func(object pyfunc):
+    """Convert a python function to DGL function
+
+    Parameters
+    ----------
+    pyfunc : python function
+        The python function to be converted.
+
+    Returns
+    -------
+    dglfunc: dgl.Function
+        The converted dgl function.
+    """
+    cdef DGLFunctionHandle chandle
+    Py_INCREF(pyfunc)
+    CALL(DGLFuncCreateFromCFunc(dgl_callback,
+                                <void*>(pyfunc),
+                                dgl_callback_finalize,
+                                &chandle))
+    ret = _CLASS_FUNCTION(None, False)
+    (<FunctionBase>ret).chandle = chandle
+    return ret
+
+
+cdef inline int make_arg(object arg,
+                         DGLValue* value,
+                         int* tcode,
+                         list temp_args) except -1:
+    """Pack arguments into c args dgl call accept"""
+    cdef unsigned long long ptr
+    if isinstance(arg, ObjectBase):
+        value[0].v_handle = (<ObjectBase>arg).chandle
+        tcode[0] = kObjectHandle
+    elif isinstance(arg, NDArrayBase):
+        value[0].v_handle = (<NDArrayBase>arg).chandle
+        tcode[0] = (kNDArrayContainer if
+                    not (<NDArrayBase>arg).c_is_view else kArrayHandle)
+    elif isinstance(arg, _DGL_COMPATS):
+        ptr = arg._dgl_handle
+        value[0].v_handle = (<void*>ptr)
+        tcode[0] = arg.__class__._dgl_tcode
+    elif isinstance(arg, (int, long)):
+        value[0].v_int64 = arg
+        tcode[0] = kObjectInt
+    elif isinstance(arg, float):
+        value[0].v_float64 = arg
+        tcode[0] = kObjectFloat
+    elif isinstance(arg, str):
+        tstr = c_str(arg)
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif arg is None:
+        value[0].v_handle = NULL
+        tcode[0] = kNull
+    elif isinstance(arg, Number):
+        value[0].v_float64 = arg
+        tcode[0] = kObjectFloat
+    elif isinstance(arg, CTypesDGLDataType):
+        tstr = c_str(str(arg))
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif isinstance(arg, CTypesDGLContext):
+        value[0].v_ctx = (<DGLContext*>(
+            <unsigned long long>ctypes.addressof(arg)))[0]
+        tcode[0] = kDGLContext
+    elif isinstance(arg, bytearray):
+        arr = DGLByteArray()
+        arr.data = ctypes.cast(
+            (ctypes.c_byte * len(arg)).from_buffer(arg),
+            ctypes.POINTER(ctypes.c_byte))
+        arr.size = len(arg)
+        value[0].v_handle = <void*>(
+            <unsigned long long>ctypes.addressof(arr))
+        tcode[0] = kBytes
+        temp_args.append(arr)
+    elif isinstance(arg, string_types):
+        tstr = c_str(arg)
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif isinstance(arg, (list, tuple, dict, ObjectGeneric)):
+        arg = convert_to_object(arg)
+        value[0].v_handle = (<ObjectBase>arg).chandle
+        tcode[0] = kObjectHandle
+        temp_args.append(arg)
+    #elif isinstance(arg, _CLASS_MODULE):
+    #    value[0].v_handle = c_handle(arg.handle)
+    #    tcode[0] = kModuleHandle
+    elif isinstance(arg, FunctionBase):
+        value[0].v_handle = (<FunctionBase>arg).chandle
+        tcode[0] = kFuncHandle
+    elif isinstance(arg, ctypes.c_void_p):
+        value[0].v_handle = c_handle(arg)
+        tcode[0] = kHandle
+    elif callable(arg):
+        arg = convert_to_dgl_func(arg)
+        value[0].v_handle = (<FunctionBase>arg).chandle
+        tcode[0] = kFuncHandle
+        temp_args.append(arg)
+    else:
+        raise TypeError("Don't know how to handle type %s" % type(arg))
+    return 0
+
+cdef inline bytearray make_ret_bytes(void* chandle):
+    handle = ctypes_handle(chandle)
+    arr = ctypes.cast(handle, ctypes.POINTER(DGLByteArray))[0]
+    size = arr.size
+    res = bytearray(size)
+    rptr = (ctypes.c_byte * size).from_buffer(res)
+    if not ctypes.memmove(rptr, arr.data, size):
+        raise RuntimeError('memmove failed')
+    return res
+
+cdef inline object make_ret(DGLValue value, int tcode):
+    """convert result to return value."""
+    if tcode == kObjectHandle:
+        return make_ret_object(value.v_handle)
+    elif tcode == kNull:
+        return None
+    elif tcode == kObjectInt:
+        return value.v_int64
+    elif tcode == kObjectFloat:
+        return value.v_float64
+    elif tcode == kNDArrayContainer:
+        return c_make_array(value.v_handle, False)
+    elif tcode == kStr:
+        return py_str(value.v_str)
+    elif tcode == kBytes:
+        return make_ret_bytes(value.v_handle)
+    elif tcode == kHandle:
+        return ctypes_handle(value.v_handle)
+    elif tcode == kDGLContext:
+        return CTypesDGLContext(value.v_ctx.device_type, value.v_ctx.device_id)
+    # (minjie): class module are not used in DGL.
+    #elif tcode == kModuleHandle:
+    #    return _CLASS_MODULE(ctypes_handle(value.v_handle))
+    elif tcode == kFuncHandle:
+        fobj = _CLASS_FUNCTION(None, False)
+        (<FunctionBase>fobj).chandle = value.v_handle
+        return fobj
+    elif tcode in _DGL_EXT_RET:
+        return _DGL_EXT_RET[tcode](ctypes_handle(value.v_handle))
+
+    raise ValueError("Unhandled type code %d" % tcode)
+
+
+cdef inline int FuncCall3(void* chandle,
+                          tuple args,
+                          int nargs,
+                          DGLValue* ret_val,
+                          int* ret_tcode) except -1:
+    cdef DGLValue[3] values
+    cdef int[3] tcodes
+    nargs = len(args)
+    temp_args = []
+    for i in range(nargs):
+        make_arg(args[i], &values[i], &tcodes[i], temp_args)
+    with nogil:
+        ret = DGLFuncCall(chandle, &values[0], &tcodes[0],
+                          nargs, ret_val, ret_tcode)
+    if ret != 0:
+        raise DGLError(py_str(DGLGetLastError()))
+    return 0
+
+cdef inline int FuncCall(void* chandle,
+                         tuple args,
+                         DGLValue* ret_val,
+                         int* ret_tcode) except -1:
+    cdef int nargs
+    nargs = len(args)
+    if nargs <= 3:
+        FuncCall3(chandle, args, nargs, ret_val, ret_tcode)
+        return 0
+
+    cdef vector[DGLValue] values
+    cdef vector[int] tcodes
+    values.resize(max(nargs, 1))
+    tcodes.resize(max(nargs, 1))
+    temp_args = []
+    for i in range(nargs):
+        make_arg(args[i], &values[i], &tcodes[i], temp_args)
+    with nogil:
+        ret = DGLFuncCall(chandle, &values[0], &tcodes[0],
+                          nargs, ret_val, ret_tcode)
+    if ret != 0:
+        raise DGLError(py_str(DGLGetLastError()))
+    return 0
+
+
+cdef inline int ConstructorCall(void* constructor_handle,
+                                int type_code,
+                                tuple args,
+                                void** handle) except -1:
+    """Call contructor of a handle function"""
+    cdef DGLValue ret_val
+    cdef int ret_tcode
+    FuncCall(constructor_handle, args, &ret_val, &ret_tcode)
+    assert ret_tcode == type_code
+    handle[0] = ret_val.v_handle
+    return 0
+
+
+cdef class FunctionBase:
+    cdef DGLFunctionHandle chandle
+    cdef int is_global
+
+    cdef inline _set_handle(self, handle):
+        if handle is None:
+            self.chandle = NULL
+        else:
+            self.chandle = c_handle(handle)
+
+    property is_global:
+        def __get__(self):
+            return self.c_is_global != 0
+
+        def __set__(self, value):
+            self.c_is_global = value
+
+    property handle:
+        def __get__(self):
+            if self.chandle == NULL:
+                return None
+            else:
+                return ctypes.cast(<unsigned long long>self.chandle, ctypes.c_void_p)
+        def __set__(self, value):
+            self._set_handle(value)
+
+    def __init__(self, handle, is_global):
+        self._set_handle(handle)
+        self.c_is_global = is_global
+
+    def __dealloc__(self):
+        if self.is_global == 0:
+            CALL(DGLFuncFree(self.chandle))
+
+    def __call__(self, *args):
+        cdef DGLValue ret_val
+        cdef int ret_tcode
+        FuncCall(self.chandle, args, &ret_val, &ret_tcode)
+        return make_ret(ret_val, ret_tcode)
+
+_CLASS_FUNCTION = None
+_CLASS_MODULE = None
+
+def _set_class_module(module_class):
+    """Initialize the module."""
+    global _CLASS_MODULE
+    _CLASS_MODULE = module_class
+
+def _set_class_function(func_class):
+    global _CLASS_FUNCTION
+    _CLASS_FUNCTION = func_class

python/dgl/_ffi/_cython/ndarray.pxi

@@ -0,0 +1,110 @@
+from ..runtime_ctypes import DGLArrayHandle as PyDGLArrayHandle
+from cpython cimport PyCapsule_Destructor
+
+cdef const char* _c_str_dltensor = "dltensor"
+cdef const char* _c_str_used_dltensor = "used_dltensor"
+
+
+cdef _c_dlpack_deleter(object pycaps):
+    cdef DLManagedTensor* dltensor
+    if pycapsule.PyCapsule_IsValid(pycaps, _c_str_dltensor):
+        dltensor = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(pycaps, _c_str_dltensor)
+        DGLDLManagedTensorCallDeleter(dltensor)
+
+
+def _from_dlpack(object dltensor):
+    cdef DLManagedTensor* ptr
+    cdef DGLArrayHandle chandle
+    if pycapsule.PyCapsule_IsValid(dltensor, _c_str_dltensor):
+        ptr = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(dltensor, _c_str_dltensor)
+        CALL(DGLArrayFromDLPack(ptr, &chandle))
+        # set name and destructor to be empty
+        pycapsule.PyCapsule_SetDestructor(dltensor, NULL)
+        pycapsule.PyCapsule_SetName(dltensor, _c_str_used_dltensor)
+        return c_make_array(chandle, 0)
+    raise ValueError("Expect a dltensor field, pycapsule.PyCapsule can only be consumed once")
+
+
+cdef class NDArrayBase:
+    cdef DGLArray* chandle
+    cdef int c_is_view
+
+    cdef inline _set_handle(self, handle):
+        cdef unsigned long long ptr
+        if handle is None:
+            self.chandle = NULL
+        else:
+            ptr = ctypes.cast(handle, ctypes.c_void_p).value
+            self.chandle = <DGLArray*>(ptr)
+
+    property _dgl_handle:
+        def __get__(self):
+            return <unsigned long long>self.chandle
+
+    property handle:
+        def __get__(self):
+            if self.chandle == NULL:
+                return None
+            else:
+                return ctypes.cast(
+                    <unsigned long long>self.chandle, PyDGLArrayHandle)
+
+        def __set__(self, value):
+            self._set_handle(value)
+
+    def __init__(self, handle, is_view):
+        self._set_handle(handle)
+        self.c_is_view = is_view
+
+    def __dealloc__(self):
+        if self.c_is_view == 0:
+            CALL(DGLArrayFree(self.chandle))
+
+    def to_dlpack(self, alignment=0):
+        """Produce an array from a DLPack Tensor without copying memory
+
+        Args
+        -------
+        alignment: int, default to be 0
+        Indicates the alignment requirement when converting to dlpack. Will copy to a 
+        new tensor if the alignment requirement is not satisfied. 
+        0 means no alignment requirement.
+        
+        Returns
+        -------
+        dlpack : DLPack tensor view of the array data
+        """
+        cdef DLManagedTensor* dltensor
+        if self.c_is_view != 0:
+            raise ValueError("to_dlpack do not work with memory views")
+        CALL(DGLArrayToDLPack(self.chandle, &dltensor, alignment))
+        return pycapsule.PyCapsule_New(dltensor, _c_str_dltensor, <PyCapsule_Destructor>_c_dlpack_deleter)
+
+
+cdef c_make_array(void* chandle, is_view):
+    ret = _CLASS_NDARRAY(None, is_view)
+    (<NDArrayBase>ret).chandle = <DGLArray*>chandle
+    return ret
+
+
+cdef _DGL_COMPATS = ()
+
+cdef _DGL_EXT_RET = {}
+
+def _reg_extension(cls, fcreate):
+    global _DGL_COMPATS
+    _DGL_COMPATS += (cls,)
+    if fcreate:
+        _DGL_EXT_RET[cls._dgl_tcode] = fcreate
+
+
+def _make_array(handle, is_view):
+    cdef unsigned long long ptr
+    ptr = ctypes.cast(handle, ctypes.c_void_p).value
+    return c_make_array(<void*>ptr, is_view)
+
+cdef object _CLASS_NDARRAY = None
+
+def _set_class_ndarray(cls):
+    global _CLASS_NDARRAY
+    _CLASS_NDARRAY = cls

python/dgl/_ffi/_cython/object.pxi

@@ -0,0 +1,91 @@
+from ... import _api_internal
+from ..base import string_types
+from ..object_generic import _set_class_object_base
+
+"""Maps object type to its constructor"""
+OBJECT_TYPE = []
+
+def _register_object(int index, object cls):
+    """register object class"""
+    while len(OBJECT_TYPE) <= index:
+        OBJECT_TYPE.append(None)
+    OBJECT_TYPE[index] = cls
+
+
+cdef inline object make_ret_object(void* chandle):
+    global OBJECT_TYPE
+    cdef int tindex
+    cdef list object_type
+    cdef object cls
+    object_type = OBJECT_TYPE
+    CALL(DGLObjectGetTypeIndex(chandle, &tindex))
+    if tindex < len(object_type):
+        cls = object_type[tindex]
+        if cls is not None:
+            obj = cls.__new__(cls)
+        else:
+            obj = ObjectBase.__new__(ObjectBase)
+    else:
+        obj = ObjectBase.__new__(ObjectBase)
+    (<ObjectBase>obj).chandle = chandle
+    return obj
+
+
+cdef class ObjectBase:
+    cdef void* chandle
+
+    cdef _set_handle(self, handle):
+        cdef unsigned long long ptr
+        if handle is None:
+            self.chandle = NULL
+        else:
+            ptr = handle.value
+            self.chandle = <void*>(ptr)
+
+    property handle:
+        def __get__(self):
+            if self.chandle == NULL:
+                return None
+            else:
+                return ctypes_handle(self.chandle)
+
+        def __set__(self, value):
+            self._set_handle(value)
+
+    def __dealloc__(self):
+        CALL(DGLObjectFree(self.chandle))
+
+    def __getattr__(self, name):
+        cdef DGLValue ret_val
+        cdef int ret_type_code, ret_succ
+        CALL(DGLObjectGetAttr(self.chandle, c_str(name),
+                            &ret_val, &ret_type_code, &ret_succ))
+        if ret_succ == 0:
+            raise AttributeError(
+                "'%s' object has no attribute '%s'" % (type(self), name))
+        return make_ret(ret_val, ret_type_code)
+
+    def __init_handle_by_constructor__(self, fconstructor, *args):
+        """Initialize the handle by calling constructor function.
+
+        Parameters
+        ----------
+        fconstructor : Function
+            Constructor function.
+
+        args: list of objects
+            The arguments to the constructor
+
+        Note
+        ----
+        We have a special calling convention to call constructor functions.
+        So the return handle is directly set into the Object object
+        instead of creating a new Object.
+        """
+        cdef void* chandle
+        ConstructorCall(
+            (<FunctionBase>fconstructor).chandle,
+            kObjectHandle, args, &chandle)
+        self.chandle = chandle
+
+_set_class_object_base(ObjectBase)

python/dgl/_ffi/base.py

@@ -0,0 +1,155 @@
+# coding: utf-8
+# pylint: disable=invalid-name
+"""ctypes library and helper functions """
+from __future__ import absolute_import
+
+import ctypes
+import logging
+import os
+import sys
+
+import numpy as np
+
+from . import libinfo
+
+# ----------------------------
+# library loading
+# ----------------------------
+if sys.version_info[0] == 3:
+    string_types = (str,)
+    numeric_types = (float, int, np.float32, np.int32)
+    # this function is needed for python3
+    # to convert ctypes.char_p .value back to python str
+    py_str = lambda x: x.decode("utf-8")
+else:
+    string_types = (basestring,)
+    numeric_types = (float, int, long, np.float32, np.int32)
+    py_str = lambda x: x
+
+
+class DGLError(Exception):
+    """Error thrown by DGL function"""
+
+    pass  # pylint: disable=unnecessary-pass
+
+
+def _load_lib():
+    """Load libary by searching possible path."""
+    lib_path = libinfo.find_lib_path()
+    lib = ctypes.CDLL(lib_path[0])
+    dirname = os.path.dirname(lib_path[0])
+    basename = os.path.basename(lib_path[0])
+    # DMatrix functions
+    lib.DGLGetLastError.restype = ctypes.c_char_p
+    return lib, basename, dirname
+
+
+# version number
+__version__ = libinfo.__version__
+# library instance of nnvm
+_LIB, _LIB_NAME, _DIR_NAME = _load_lib()
+
+# The FFI mode of DGL
+_FFI_MODE = os.environ.get("DGL_FFI", "auto")
+
+# ----------------------------
+# helper function in ctypes.
+# ----------------------------
+def check_call(ret):
+    """Check the return value of C API call
+
+    This function will raise exception when error occurs.
+    Wrap every API call with this function
+
+    Parameters
+    ----------
+    ret : int
+        return value from API calls
+    """
+    if ret != 0:
+        raise DGLError(py_str(_LIB.DGLGetLastError()))
+
+
+def c_str(string):
+    """Create ctypes char * from a python string
+    Parameters
+    ----------
+    string : string type
+        python string
+
+    Returns
+    -------
+    str : c_char_p
+        A char pointer that can be passed to C API
+    """
+    return ctypes.c_char_p(string.encode("utf-8"))
+
+
+def c_array(ctype, values):
+    """Create ctypes array from a python array
+
+    Parameters
+    ----------
+    ctype : ctypes data type
+        data type of the array we want to convert to
+
+    values : tuple or list
+        data content
+
+    Returns
+    -------
+    out : ctypes array
+        Created ctypes array
+    """
+    return (ctype * len(values))(*values)
+
+
+def decorate(func, fwrapped):
+    """A wrapper call of decorator package, differs to call time
+
+    Parameters
+    ----------
+    func : function
+        The original function
+
+    fwrapped : function
+        The wrapped function
+    """
+    import decorator
+
+    return decorator.decorate(func, fwrapped)
+
+
+tensor_adapter_loaded = False
+
+
+def load_tensor_adapter(backend, version):
+    """Tell DGL to load a tensoradapter library for given backend and version.
+
+    Parameters
+    ----------
+    backend : str
+        The backend (currently ``pytorch``, ``mxnet`` or ``tensorflow``).
+    version : str
+        The version number of the backend.
+    """
+    global tensor_adapter_loaded
+    version = version.split("+")[0]
+    if sys.platform.startswith("linux"):
+        basename = "libtensoradapter_%s_%s.so" % (backend, version)
+    elif sys.platform.startswith("darwin"):
+        basename = "libtensoradapter_%s_%s.dylib" % (backend, version)
+    elif sys.platform.startswith("win"):
+        basename = "tensoradapter_%s_%s.dll" % (backend, version)
+    else:
+        raise NotImplementedError("Unsupported system: %s" % sys.platform)
+    path = os.path.join(_DIR_NAME, "tensoradapter", backend, basename)
+    tensor_adapter_loaded = _LIB.DGLLoadTensorAdapter(path.encode("utf-8")) == 0
+    if not tensor_adapter_loaded:
+        logger = logging.getLogger("dgl-core")
+        logger.debug("Memory optimization with PyTorch is not enabled.")
+
+
+def is_tensor_adaptor_enabled() -> bool:
+    """Check whether TensorAdaptor is enabled."""
+    return tensor_adapter_loaded

python/dgl/_ffi/capi.py

@@ -0,0 +1,4 @@
+"""Init all C APIs in the default namespace."""
+from .function import _init_api
+
+__all__ = _init_api("dgl.capi", __name__)

python/dgl/_ffi/function.py

@@ -0,0 +1,350 @@
+# pylint: disable=invalid-name, unused-import
+"""Function namespace."""
+from __future__ import absolute_import
+
+import ctypes
+import sys
+
+from .base import _FFI_MODE, _LIB, c_str, check_call, py_str, string_types
+
+IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
+
+try:
+    # pylint: disable=wrong-import-position
+    if _FFI_MODE == "ctypes":
+        raise ImportError()
+    if sys.version_info >= (3, 0):
+        from ._cy3.core import (
+            _set_class_function,
+            _set_class_module,
+            convert_to_dgl_func,
+            FunctionBase as _FunctionBase,
+        )
+    else:
+        from ._cy2.core import (
+            _set_class_function,
+            _set_class_module,
+            convert_to_dgl_func,
+            FunctionBase as _FunctionBase,
+        )
+except IMPORT_EXCEPT:
+    # pylint: disable=wrong-import-position
+    from ._ctypes.function import (
+        _set_class_function,
+        _set_class_module,
+        convert_to_dgl_func,
+        FunctionBase as _FunctionBase,
+    )
+
+FunctionHandle = ctypes.c_void_p
+
+
+class Function(_FunctionBase):
+    """The PackedFunc object.
+
+    Function plays an key role to bridge front and backend in DGL.
+    Function provide a type-erased interface, you can call function with positional arguments.
+
+    The compiled module returns Function.
+    DGL backend also registers and exposes its API as Functions.
+    For example, the developer function exposed in dgl.ir_pass are actually
+    C++ functions that are registered as PackedFunc
+
+    The following are list of common usage scenario of dgl.Function.
+
+    - Automatic exposure of C++ API into python
+    - To call PackedFunc from python side
+    - To call python callbacks to inspect results in generated code
+    - Bring python hook into C++ backend
+
+    See Also
+    --------
+    dgl.register_func: How to register global function.
+    dgl.get_global_func: How to get global function.
+    """
+
+    pass  # pylint: disable=unnecessary-pass
+
+
+class ModuleBase(object):
+    """Base class for module"""
+
+    __slots__ = ["handle", "_entry", "entry_name"]
+
+    def __init__(self, handle):
+        self.handle = handle
+        self._entry = None
+        self.entry_name = "__dgl_main__"
+
+    def __del__(self):
+        check_call(_LIB.DGLModFree(self.handle))
+
+    @property
+    def entry_func(self):
+        """Get the entry function
+
+        Returns
+        -------
+        f : Function
+            The entry function if exist
+        """
+        if self._entry:
+            return self._entry
+        self._entry = self.get_function(self.entry_name)
+        return self._entry
+
+    def get_function(self, name, query_imports=False):
+        """Get function from the module.
+
+        Parameters
+        ----------
+        name : str
+            The name of the function
+
+        query_imports : bool
+            Whether also query modules imported by this module.
+
+        Returns
+        -------
+        f : Function
+            The result function.
+        """
+        ret_handle = FunctionHandle()
+        check_call(
+            _LIB.DGLModGetFunction(
+                self.handle,
+                c_str(name),
+                ctypes.c_int(query_imports),
+                ctypes.byref(ret_handle),
+            )
+        )
+        if not ret_handle.value:
+            raise AttributeError("Module has no function '%s'" % name)
+        return Function(ret_handle, False)
+
+    def import_module(self, module):
+        """Add module to the import list of current one.
+
+        Parameters
+        ----------
+        module : Module
+            The other module.
+        """
+        check_call(_LIB.DGLModImport(self.handle, module.handle))
+
+    def __getitem__(self, name):
+        if not isinstance(name, string_types):
+            raise ValueError("Can only take string as function name")
+        return self.get_function(name)
+
+    def __call__(self, *args):
+        if self._entry:
+            return self._entry(*args)
+        f = self.entry_func
+        return f(*args)
+
+
+def register_func(func_name, f=None, override=False):
+    """Register global function
+
+    Parameters
+    ----------
+    func_name : str or function
+        The function name
+
+    f : function, optional
+        The function to be registered.
+
+    override: boolean optional
+        Whether override existing entry.
+
+    Returns
+    -------
+    fregister : function
+        Register function if f is not specified.
+
+    Examples
+    --------
+    The following code registers my_packed_func as global function.
+    Note that we simply get it back from global function table to invoke
+    it from python side. However, we can also invoke the same function
+    from C++ backend, or in the compiled DGL code.
+
+    .. code-block:: python
+
+      targs = (10, 10.0, "hello")
+      @dgl.register_func
+      def my_packed_func(*args):
+          assert(tuple(args) == targs)
+          return 10
+      # Get it out from global function table
+      f = dgl.get_global_func("my_packed_func")
+      assert isinstance(f, dgl.nd.Function)
+      y = f(*targs)
+      assert y == 10
+    """
+    if callable(func_name):
+        f = func_name
+        func_name = f.__name__
+
+    if not isinstance(func_name, str):
+        raise ValueError("expect string function name")
+
+    ioverride = ctypes.c_int(override)
+
+    def register(myf):
+        """internal register function"""
+        if not isinstance(myf, Function):
+            myf = convert_to_dgl_func(myf)
+        check_call(
+            _LIB.DGLFuncRegisterGlobal(c_str(func_name), myf.handle, ioverride)
+        )
+        return myf
+
+    if f:
+        return register(f)
+    return register
+
+
+def get_global_func(name, allow_missing=False):
+    """Get a global function by name
+
+    Parameters
+    ----------
+    name : str
+        The name of the global function
+
+    allow_missing : bool
+        Whether allow missing function or raise an error.
+
+    Returns
+    -------
+    func : dgl.Function
+        The function to be returned, None if function is missing.
+    """
+    handle = FunctionHandle()
+    check_call(_LIB.DGLFuncGetGlobal(c_str(name), ctypes.byref(handle)))
+    if handle.value:
+        return Function(handle, False)
+    else:
+        if allow_missing:
+            return None
+        else:
+            raise ValueError("Cannot find global function %s" % name)
+
+
+def list_global_func_names():
+    """Get list of global functions registered.
+
+    Returns
+    -------
+    names : list
+       List of global functions names.
+    """
+    plist = ctypes.POINTER(ctypes.c_char_p)()
+    size = ctypes.c_uint()
+
+    check_call(
+        _LIB.DGLFuncListGlobalNames(ctypes.byref(size), ctypes.byref(plist))
+    )
+    fnames = []
+    for i in range(size.value):
+        fnames.append(py_str(plist[i]))
+    return fnames
+
+
+def extract_ext_funcs(finit):
+    """
+    Extract the extension PackedFuncs from a C module.
+
+    Parameters
+    ----------
+    finit : ctypes function
+        a ctypes that takes signature of DGLExtensionDeclarer
+
+    Returns
+    -------
+    fdict : dict of str to Function
+        The extracted functions
+    """
+    fdict = {}
+
+    def _list(name, func):
+        fdict[name] = func
+
+    myf = convert_to_dgl_func(_list)
+    ret = finit(myf.handle)
+    _ = myf
+    if ret != 0:
+        raise RuntimeError("cannot initialize with %s" % finit)
+    return fdict
+
+
+def _get_api(f):
+    flocal = f
+    flocal.is_global = True
+    return flocal
+
+
+def _init_api(namespace, target_module_name=None):
+    """Initialize api for a given module name
+
+    namespace : str
+       The namespace of the source registry
+
+    target_module_name : str
+       The target module name if different from namespace
+    """
+    target_module_name = target_module_name if target_module_name else namespace
+    if namespace.startswith("dgl."):
+        return _init_api_prefix(target_module_name, namespace[4:])
+    else:
+        return _init_api_prefix(target_module_name, namespace)
+
+
+def _init_api_prefix(module_name, prefix):
+    module = sys.modules[module_name]
+    name_list = []
+
+    for name in list_global_func_names():
+        if name.startswith("_") and not name.startswith("_deprecate"):
+            # internal APIs are ignored
+            continue
+        name_split = name.rsplit(".", 1)
+        if name_split[0] != prefix:
+            continue
+
+        if len(name_split) == 1:
+            print('Warning: invalid API name "%s".' % name)
+            continue
+        fname = name_split[1]
+        target_module = module
+
+        f = get_global_func(name)
+        ff = _get_api(f)
+        ff.__name__ = fname
+        ff.__doc__ = "DGL PackedFunc %s. " % fname
+        setattr(target_module, ff.__name__, ff)
+        name_list.append(fname)
+
+    return name_list
+
+
+def _init_internal_api():
+    for name in list_global_func_names():
+        if not name.startswith("_") or name.startswith("_deprecate"):
+            # normal APIs are ignored
+            continue
+        target_module = sys.modules["dgl._api_internal"]
+        fname = name
+        if fname.find(".") != -1:
+            print('Warning: invalid API name "%s".' % fname)
+            continue
+        f = get_global_func(name)
+        ff = _get_api(f)
+        ff.__name__ = fname
+        ff.__doc__ = "DGL PackedFunc %s. " % fname
+        setattr(target_module, ff.__name__, ff)
+
+
+_set_class_function(Function)

python/dgl/_ffi/libinfo.py

@@ -0,0 +1,108 @@
+"""Library information."""
+from __future__ import absolute_import
+
+import os
+import pathlib
+import sys
+
+
+def find_lib_path(name=None, search_path=None, optional=False):
+    """Find dynamic library files.
+
+    Parameters
+    ----------
+    name : list of str
+        List of names to be found.
+
+    Returns
+    -------
+    lib_path : list(string)
+        List of all found path to the libraries
+    """
+    # See https://github.com/dmlc/tvm/issues/281 for some background.
+
+    # NB: This will either be the source directory (if DGL is run
+    # inplace) or the install directory (if DGL is installed).
+    # An installed DGL's curr_path will look something like:
+    #   $PREFIX/lib/python3.6/site-packages/dgl/_ffi
+    ffi_dir = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
+    source_dir = os.path.join(ffi_dir, "..", "..", "..")
+    install_lib_dir = os.path.join(ffi_dir, "..", "..", "..", "..")
+
+    dll_path = []
+
+    if os.environ.get("DGL_LIBRARY_PATH", None):
+        dll_path.append(os.environ["DGL_LIBRARY_PATH"])
+
+    if sys.platform.startswith("linux") and os.environ.get(
+        "LD_LIBRARY_PATH", None
+    ):
+        dll_path.extend(
+            [p.strip() for p in os.environ["LD_LIBRARY_PATH"].split(":")]
+        )
+    elif sys.platform.startswith("darwin") and os.environ.get(
+        "DYLD_LIBRARY_PATH", None
+    ):
+        dll_path.extend(
+            [p.strip() for p in os.environ["DYLD_LIBRARY_PATH"].split(":")]
+        )
+
+    # Pip lib directory
+    dll_path.append(os.path.join(ffi_dir, ".."))
+    # Default cmake build directory
+    dll_path.append(os.path.join(source_dir, "build"))
+    dll_path.append(os.path.join(source_dir, "build", "Release"))
+    # Default make build directory
+    dll_path.append(os.path.join(source_dir, "lib"))
+
+    dll_path.append(install_lib_dir)
+
+    if search_path is not None:
+        if isinstance(search_path, (list, tuple, set)):
+            dll_path = dll_path + list(search_path)
+        elif isinstance(search_path, str):
+            dll_path.append(search_path)
+        else:
+            raise ValueError(
+                "type(search_path)={} is invalid".format(type(search_path))
+            )
+    dll_path = [
+        str(x.absolute()) if isinstance(x, pathlib.Path) else os.path.abspath(x)
+        for x in dll_path
+    ]
+
+    if name is None:
+        if sys.platform.startswith("win32"):
+            name = ["libdgl.dll", "dgl.dll"]
+        elif sys.platform.startswith("darwin"):
+            name = "libdgl.dylib"
+        else:
+            name = "libdgl.so"
+
+    if isinstance(name, str):
+        name = [name]
+
+    lib_dll_path = []
+    for n in name:
+        lib_dll_path += [os.path.join(p, n) for p in dll_path]
+
+    lib_found = [p for p in lib_dll_path if os.path.isfile(p)]
+
+    if not lib_found:
+        message = (
+            "Cannot find the files.\n"
+            + "List of candidates:\n"
+            + str("\n".join(lib_dll_path))
+        )
+        if not optional:
+            raise RuntimeError(message)
+        return None
+
+    return lib_found
+
+
+# current version
+# We use the version of the incoming release for code
+# that is under development.
+# The following line is set by dgl/python/update_version.py
+__version__ = "2.5"

python/dgl/_ffi/ndarray.py

@@ -0,0 +1,448 @@
+# pylint: disable=invalid-name, unused-import
+"""Runtime NDArray api"""
+from __future__ import absolute_import
+
+import ctypes
+import sys
+
+import numpy as np
+
+from .base import _FFI_MODE, _LIB, c_array, c_str, check_call, string_types
+from .runtime_ctypes import (
+    dgl_shape_index_t,
+    DGLArray,
+    DGLArrayHandle,
+    DGLContext,
+    DGLDataType,
+    TypeCode,
+)
+
+IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
+
+try:
+    # pylint: disable=wrong-import-position
+    if _FFI_MODE == "ctypes":
+        raise ImportError()
+    if sys.version_info >= (3, 0):
+        from ._cy3.core import (
+            _from_dlpack,
+            _make_array,
+            _reg_extension,
+            _set_class_ndarray,
+            NDArrayBase as _NDArrayBase,
+        )
+    else:
+        from ._cy2.core import (
+            _from_dlpack,
+            _make_array,
+            _reg_extension,
+            _set_class_ndarray,
+            NDArrayBase as _NDArrayBase,
+        )
+except IMPORT_EXCEPT:
+    # pylint: disable=wrong-import-position
+    from ._ctypes.ndarray import (
+        _from_dlpack,
+        _make_array,
+        _reg_extension,
+        _set_class_ndarray,
+        NDArrayBase as _NDArrayBase,
+    )
+
+
+def context(dev_type, dev_id=0):
+    """Construct a DGL context with given device type and id.
+
+    Parameters
+    ----------
+    dev_type: int or str
+        The device type mask or name of the device.
+
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    ctx: DGLContext
+        The corresponding context.
+
+    Examples
+    --------
+    Context can be used to create reflection of context by
+    string representation of the device type.
+
+    .. code-block:: python
+
+      assert dgl.context("cpu", 1) == dgl.cpu(1)
+      assert dgl.context("gpu", 0) == dgl.gpu(0)
+      assert dgl.context("cuda", 0) == dgl.gpu(0)
+    """
+    if isinstance(dev_type, string_types):
+        dev_type = dev_type.split()[0]
+        if dev_type not in DGLContext.STR2MASK:
+            raise ValueError("Unknown device type %s" % dev_type)
+        dev_type = DGLContext.STR2MASK[dev_type]
+    return DGLContext(dev_type, dev_id)
+
+
+def numpyasarray(np_data):
+    """Return a DGLArray representation of a numpy array."""
+    data = np_data
+    assert data.flags["C_CONTIGUOUS"]
+    arr = DGLArray()
+    shape = c_array(dgl_shape_index_t, data.shape)
+    arr.data = data.ctypes.data_as(ctypes.c_void_p)
+    arr.shape = shape
+    arr.strides = None
+    arr.dtype = DGLDataType(np.dtype(data.dtype).name)
+    arr.ndim = data.ndim
+    # CPU device
+    arr.ctx = context(1, 0)
+    return arr, shape
+
+
+def empty(shape, dtype="float32", ctx=context(1, 0)):
+    """Create an empty array given shape and device
+
+    Parameters
+    ----------
+    shape : tuple of int
+        The shape of the array
+
+    dtype : type or str
+        The data type of the array.
+
+    ctx : DGLContext
+        The context of the array
+
+    Returns
+    -------
+    arr : dgl.nd.NDArray
+        The array dgl supported.
+    """
+    shape = c_array(dgl_shape_index_t, shape)
+    ndim = ctypes.c_int(len(shape))
+    handle = DGLArrayHandle()
+    dtype = DGLDataType(dtype)
+    check_call(
+        _LIB.DGLArrayAlloc(
+            shape,
+            ndim,
+            ctypes.c_int(dtype.type_code),
+            ctypes.c_int(dtype.bits),
+            ctypes.c_int(dtype.lanes),
+            ctx.device_type,
+            ctx.device_id,
+            ctypes.byref(handle),
+        )
+    )
+    return _make_array(handle, False)
+
+
+def empty_shared_mem(name, is_create, shape, dtype="float32"):
+    """Create an empty array with shared memory given shape and dtype
+
+    Parameters
+    ----------
+    name : string
+        The name of the shared memory. It's a file name in Unix.
+
+    is_create : bool
+        Whether to create the shared memory or use the one created by somewhere else.
+
+    shape : tuple of int
+        The shape of the array
+
+    dtype : type or str
+        The data type of the array.
+
+    Returns
+    -------
+    arr : dgl.nd.NDArray
+        The array dgl supported.
+    """
+    name = ctypes.c_char_p(name.encode("utf-8"))
+    shape = c_array(dgl_shape_index_t, shape)
+    ndim = ctypes.c_int(len(shape))
+    handle = DGLArrayHandle()
+    dtype = DGLDataType(dtype)
+    check_call(
+        _LIB.DGLArrayAllocSharedMem(
+            name,
+            shape,
+            ndim,
+            ctypes.c_int(dtype.type_code),
+            ctypes.c_int(dtype.bits),
+            ctypes.c_int(dtype.lanes),
+            is_create,
+            ctypes.byref(handle),
+        )
+    )
+    return _make_array(handle, False)
+
+
+def from_dlpack(dltensor):
+    """Produce an array from a DLPack tensor without memory copy.
+    Retrieves the underlying DLPack tensor's pointer to create an array from the
+    data. Removes the original DLPack tensor's destructor as now the array is
+    responsible for destruction.
+
+    Parameters
+    ----------
+    dltensor : DLPack tensor
+        Input DLManagedTensor, can only be consumed once.
+
+    Returns
+    -------
+    arr: dgl.nd.NDArray
+        The array view of the tensor data.
+    """
+    return _from_dlpack(dltensor)
+
+
+class NDArrayBase(_NDArrayBase):
+    """A simple Device/CPU Array object in runtime."""
+
+    @property
+    def shape(self):
+        """Shape of this array"""
+        return tuple(
+            self.handle.contents.shape[i]
+            for i in range(self.handle.contents.ndim)
+        )
+
+    @property
+    def dtype(self):
+        """Type of this array"""
+        return str(self.handle.contents.dtype)
+
+    @property
+    def ctx(self):
+        """context of this array"""
+        return self.handle.contents.ctx
+
+    @property
+    def context(self):
+        """context of this array"""
+        return self.ctx
+
+    def __hash__(self):
+        return ctypes.cast(self.handle, ctypes.c_void_p).value
+
+    def __eq__(self, other):
+        return self.same_as(other)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def same_as(self, other):
+        """Check object identity equality
+
+        Parameters
+        ----------
+        other : object
+            The other object to compare to
+
+        Returns
+        -------
+        same : bool
+            Whether other is same as self.
+        """
+        if not isinstance(other, NDArrayBase):
+            return False
+        return self.__hash__() == other.__hash__()
+
+    def __setitem__(self, in_slice, value):
+        """Set ndarray value"""
+        if (
+            not isinstance(in_slice, slice)
+            or in_slice.start is not None
+            or in_slice.stop is not None
+        ):
+            raise ValueError("Array only support set from numpy array")
+        if isinstance(value, NDArrayBase):
+            if value.handle is not self.handle:
+                value.copyto(self)
+        elif isinstance(value, (np.ndarray, np.generic)):
+            self.copyfrom(value)
+        else:
+            raise TypeError("type %s not supported" % str(type(value)))
+
+    def copyfrom(self, source_array):
+        """Perform a synchronized copy from the array.
+
+        Parameters
+        ----------
+        source_array : array_like
+            The data source we should like to copy from.
+
+        Returns
+        -------
+        arr : NDArray
+            Reference to self.
+        """
+        if isinstance(source_array, NDArrayBase):
+            source_array.copyto(self)
+            return self
+
+        if not isinstance(source_array, np.ndarray):
+            try:
+                source_array = np.asarray(source_array, dtype=self.dtype)
+            except:
+                raise TypeError(
+                    "array must be an array_like data,"
+                    + "type %s is not supported" % str(type(source_array))
+                )
+        t = DGLDataType(self.dtype)
+        shape, dtype = self.shape, self.dtype
+        if t.lanes > 1:
+            shape = shape + (t.lanes,)
+            t.lanes = 1
+            dtype = str(t)
+
+        if source_array.shape != shape:
+            raise ValueError(
+                "array shape do not match the shape of NDArray {0} vs {1}".format(
+                    source_array.shape, shape
+                )
+            )
+        source_array = np.ascontiguousarray(source_array, dtype=dtype)
+        assert source_array.flags["C_CONTIGUOUS"]
+        data = source_array.ctypes.data_as(ctypes.c_void_p)
+        nbytes = ctypes.c_size_t(
+            source_array.size * source_array.dtype.itemsize
+        )
+        check_call(_LIB.DGLArrayCopyFromBytes(self.handle, data, nbytes))
+        return self
+
+    def __repr__(self):
+        res = "dgl.{0}@{1}".format(self.asnumpy().__repr__(), self.context)
+        return res
+
+    def __str__(self):
+        return str(self.asnumpy())
+
+    def asnumpy(self):
+        """Convert this array to numpy array
+
+        Returns
+        -------
+        np_arr : numpy.ndarray
+            The corresponding numpy array.
+        """
+        t = DGLDataType(self.dtype)
+        shape, dtype = self.shape, self.dtype
+        if t.lanes > 1:
+            shape = shape + (t.lanes,)
+            t.lanes = 1
+            dtype = str(t)
+        np_arr = np.empty(shape, dtype=dtype)
+        assert np_arr.flags["C_CONTIGUOUS"]
+        data = np_arr.ctypes.data_as(ctypes.c_void_p)
+        nbytes = ctypes.c_size_t(np_arr.size * np_arr.dtype.itemsize)
+        check_call(_LIB.DGLArrayCopyToBytes(self.handle, data, nbytes))
+        return np_arr
+
+    def copyto(self, target):
+        """Copy array to target
+
+        Parameters
+        ----------
+        target : NDArray
+            The target array to be copied, must have same shape as this array.
+        """
+        if isinstance(target, DGLContext):
+            target = empty(self.shape, self.dtype, target)
+        if isinstance(target, NDArrayBase):
+            check_call(_LIB.DGLArrayCopyFromTo(self.handle, target.handle))
+        else:
+            raise ValueError("Unsupported target type %s" % str(type(target)))
+        return target
+
+    def pin_memory_(self):
+        """Pin host memory and map into GPU address space (in-place)"""
+        check_call(_LIB.DGLArrayPinData(self.handle))
+
+    def unpin_memory_(self):
+        """Unpin host memory pinned by pin_memory_()"""
+        check_call(_LIB.DGLArrayUnpinData(self.handle))
+
+    def record_stream(self, stream):
+        """Record the stream that is using this tensor.
+
+        Note
+        ----
+        This API is more for testing. Users should call ``record_stream``
+        on torch.Tensor or dgl.graph directly.
+
+        Parameters
+        ----------
+        stream : DGLStreamHandle
+        """
+        check_call(_LIB.DGLArrayRecordStream(self.handle, stream))
+
+
+def free_extension_handle(handle, type_code):
+    """Free c++ extension type handle
+
+    Parameters
+    ----------
+    handle : ctypes.c_void_p
+        The handle to the extension type.
+
+    type_code : int
+         The tyoe code
+    """
+    check_call(_LIB.DGLExtTypeFree(handle, ctypes.c_int(type_code)))
+
+
+def register_extension(cls, fcreate=None):
+    """Register a extension class to DGL.
+
+    After the class is registered, the class will be able
+    to directly pass as Function argument generated by DGL.
+
+    Parameters
+    ----------
+    cls : class
+        The class object to be registered as extension.
+
+    Note
+    ----
+    The registered class is requires one property: _dgl_handle and a class attribute _dgl_tcode.
+
+    - ```_dgl_handle``` returns integer represents the address of the handle.
+    - ```_dgl_tcode``` gives integer represents type code of the class.
+
+    Returns
+    -------
+    cls : class
+        The class being registered.
+
+    fcreate : function, optional
+        The creation function to create a class object given handle value.
+
+    Example
+    -------
+    The following code registers user defined class
+    MyTensor to be DLTensor compatible.
+
+    .. code-block:: python
+
+       @dgl.register_extension
+       class MyTensor(object):
+           _dgl_tcode = dgl.TypeCode.ARRAY_HANDLE
+
+           def __init__(self):
+               self.handle = _LIB.NewDLTensor()
+
+           @property
+           def _dgl_handle(self):
+               return self.handle.value
+    """
+    if fcreate and cls._dgl_tcode < TypeCode.EXT_BEGIN:
+        raise ValueError(
+            "Cannot register create when extension tcode is same as buildin"
+        )
+    _reg_extension(cls, fcreate)
+    return cls

python/dgl/_ffi/object.py

@@ -0,0 +1,117 @@
+"""Object namespace"""
+# pylint: disable=unused-import
+from __future__ import absolute_import
+
+import ctypes
+import sys
+
+from .. import _api_internal
+from .base import _FFI_MODE, _LIB, c_str, check_call, py_str
+from .object_generic import convert_to_object, ObjectGeneric
+
+# pylint: disable=invalid-name
+IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
+try:
+    # pylint: disable=wrong-import-position
+    if _FFI_MODE == "ctypes":
+        raise ImportError()
+    if sys.version_info >= (3, 0):
+        from ._cy3.core import _register_object, ObjectBase as _ObjectBase
+    else:
+        from ._cy2.core import _register_object, ObjectBase as _ObjectBase
+except IMPORT_EXCEPT:
+    # pylint: disable=wrong-import-position
+    from ._ctypes.object import _register_object, ObjectBase as _ObjectBase
+
+
+def _new_object(cls):
+    """Helper function for pickle"""
+    return cls.__new__(cls)
+
+
+class ObjectBase(_ObjectBase):
+    """ObjectBase is the base class of all DGL CAPI object.
+
+    The core attribute is ``handle``, which is a C raw pointer.  It must be initialized
+    via ``__init_handle_by_constructor__``.
+
+    Note that the same handle **CANNOT** be shared across multiple ObjectBase instances.
+    """
+
+    def __dir__(self):
+        plist = ctypes.POINTER(ctypes.c_char_p)()
+        size = ctypes.c_uint()
+        check_call(
+            _LIB.DGLObjectListAttrNames(
+                self.handle, ctypes.byref(size), ctypes.byref(plist)
+            )
+        )
+        names = []
+        for i in range(size.value):
+            names.append(py_str(plist[i]))
+        return names
+
+    def __hash__(self):
+        return _api_internal._raw_ptr(self)
+
+    def __eq__(self, other):
+        return self.same_as(other)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __reduce__(self):
+        cls = type(self)
+        return (_new_object, (cls,), self.__getstate__())
+
+    def __getstate__(self):
+        # TODO(minjie): TVM assumes that a Node (Object in DGL) can be serialized
+        #   to json. However, this is not true in DGL because DGL Object is meant
+        #   for runtime API, so it could contain binary data such as NDArray.
+        #   If this feature is required, please raise a RFC to DGL issue.
+        raise RuntimeError("__getstate__ is not supported for object type")
+
+    def __setstate__(self, state):
+        # pylint: disable=assigning-non-slot
+        # TODO(minjie): TVM assumes that a Node (Object in DGL) can be serialized
+        #   to json. However, this is not true in DGL because DGL Object is meant
+        #   for runtime API, so it could contain binary data such as NDArray.
+        #   If this feature is required, please raise a RFC to DGL issue.
+        raise RuntimeError("__setstate__ is not supported for object type")
+
+    def same_as(self, other):
+        """check object identity equality"""
+        if not isinstance(other, ObjectBase):
+            return False
+        return self.__hash__() == other.__hash__()
+
+
+def register_object(type_key=None):
+    """Decorator used to register object type
+
+    Examples
+    --------
+    >>> @register_object
+    >>> class MyObject:
+    >>> ... pass
+
+    Parameters
+    ----------
+    type_key : str or cls
+        The type key of the object
+    """
+    object_name = type_key if isinstance(type_key, str) else type_key.__name__
+
+    def register(cls):
+        """internal register function"""
+        tindex = ctypes.c_int()
+        ret = _LIB.DGLObjectTypeKey2Index(
+            c_str(object_name), ctypes.byref(tindex)
+        )
+        if ret == 0:
+            _register_object(tindex.value, cls)
+        return cls
+
+    if isinstance(type_key, str):
+        return register
+    return register(type_key)

python/dgl/_ffi/object_generic.py

@@ -0,0 +1,59 @@
+"""Common implementation of Object generic related logic"""
+# pylint: disable=unused-import
+from __future__ import absolute_import
+
+from numbers import Integral, Number
+
+from .. import _api_internal
+from .base import string_types
+
+# Object base class
+_CLASS_OBJECT_BASE = None
+
+
+def _set_class_object_base(cls):
+    global _CLASS_OBJECT_BASE
+    _CLASS_OBJECT_BASE = cls
+
+
+class ObjectGeneric(object):
+    """Base class for all classes that can be converted to object."""
+
+    def asobject(self):
+        """Convert value to object"""
+        raise NotImplementedError()
+
+
+def convert_to_object(value):
+    """Convert a python value to corresponding object type.
+
+    Parameters
+    ----------
+    value : str
+        The value to be inspected.
+
+    Returns
+    -------
+    object : Object
+        The corresponding object value.
+    """
+    if isinstance(value, _CLASS_OBJECT_BASE):
+        return value
+    if isinstance(value, (list, tuple)):
+        value = [convert_to_object(x) for x in value]
+        return _api_internal._List(*value)
+    if isinstance(value, dict):
+        vlist = []
+        for item in value.items():
+            if not isinstance(item[0], _CLASS_OBJECT_BASE) and not isinstance(
+                item[0], string_types
+            ):
+                raise ValueError(
+                    "key of map must already been a container type"
+                )
+            vlist.append(item[0])
+            vlist.append(convert_to_object(item[1]))
+        return _api_internal._Map(*vlist)
+    if isinstance(value, ObjectGeneric):
+        return value.asobject()
+    return _api_internal._Value(value)

python/dgl/_ffi/runtime_ctypes.py

@@ -0,0 +1,278 @@
+"""Common runtime ctypes."""
+# pylint: disable=invalid-name, super-init-not-called
+from __future__ import absolute_import
+
+import ctypes
+import json
+
+import numpy as np
+
+from .. import _api_internal
+from .base import _LIB, check_call
+
+dgl_shape_index_t = ctypes.c_int64
+
+
+class TypeCode(object):
+    """Type code used in API calls"""
+
+    INT = 0
+    UINT = 1
+    FLOAT = 2
+    HANDLE = 3
+    NULL = 4
+    DGL_DATA_TYPE = 5
+    DGL_CONTEXT = 6
+    ARRAY_HANDLE = 7
+    OBJECT_HANDLE = 8
+    MODULE_HANDLE = 9
+    FUNC_HANDLE = 10
+    STR = 11
+    BYTES = 12
+    NDARRAY_CONTAINER = 13
+    EXT_BEGIN = 15
+
+
+class DGLByteArray(ctypes.Structure):
+    """Temp data structure for byte array."""
+
+    _fields_ = [
+        ("data", ctypes.POINTER(ctypes.c_byte)),
+        ("size", ctypes.c_size_t),
+    ]
+
+
+class DGLDataType(ctypes.Structure):
+    """DGL datatype structure"""
+
+    _fields_ = [
+        ("type_code", ctypes.c_uint8),
+        ("bits", ctypes.c_uint8),
+        ("lanes", ctypes.c_uint16),
+    ]
+    CODE2STR = {0: "int", 1: "uint", 2: "float", 4: "handle"}
+    _cache = {}
+
+    def __new__(cls, type_str):
+        if type_str in cls._cache:
+            return cls._cache[type_str]
+
+        inst = super(DGLDataType, cls).__new__(DGLDataType)
+
+        if isinstance(type_str, np.dtype):
+            type_str = str(type_str)
+        arr = type_str.split("x")
+        head = arr[0]
+        inst.lanes = int(arr[1]) if len(arr) > 1 else 1
+        bits = 32
+
+        if head.startswith("int"):
+            inst.type_code = 0
+            head = head[3:]
+        elif head.startswith("uint"):
+            inst.type_code = 1
+            head = head[4:]
+        elif head.startswith("float"):
+            inst.type_code = 2
+            head = head[5:]
+        elif head.startswith("handle"):
+            inst.type_code = 4
+            bits = 64
+            head = ""
+        else:
+            raise ValueError("Do not know how to handle type %s" % type_str)
+        bits = int(head) if head else bits
+        inst.bits = bits
+
+        cls._cache[type_str] = inst
+        return inst
+
+    def __init__(self, type_str):
+        pass
+
+    def __repr__(self):
+        x = "%s%d" % (DGLDataType.CODE2STR[self.type_code], self.bits)
+        if self.lanes != 1:
+            x += "x%d" % self.lanes
+        return x
+
+    def __eq__(self, other):
+        return (
+            self.bits == other.bits
+            and self.type_code == other.type_code
+            and self.lanes == other.lanes
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+RPC_SESS_MASK = 128
+
+
+class DGLContext(ctypes.Structure):
+    """DGL context strucure."""
+
+    _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)]
+    MASK2STR = {
+        1: "cpu",
+        2: "gpu",
+        4: "opencl",
+        5: "aocl",
+        6: "sdaccel",
+        7: "vulkan",
+        8: "metal",
+        9: "vpi",
+        10: "rocm",
+        11: "opengl",
+        12: "ext_dev",
+    }
+    STR2MASK = {
+        "llvm": 1,
+        "stackvm": 1,
+        "cpu": 1,
+        "gpu": 2,
+        "cuda": 2,
+        "nvptx": 2,
+        "cl": 4,
+        "opencl": 4,
+        "aocl": 5,
+        "aocl_sw_emu": 5,
+        "sdaccel": 6,
+        "vulkan": 7,
+        "metal": 8,
+        "vpi": 9,
+        "rocm": 10,
+        "opengl": 11,
+        "ext_dev": 12,
+    }
+    _cache = {}
+
+    def __new__(cls, device_type, device_id):
+        if (device_type, device_id) in cls._cache:
+            return cls._cache[(device_type, device_id)]
+
+        inst = super(DGLContext, cls).__new__(DGLContext)
+
+        inst.device_type = device_type
+        inst.device_id = device_id
+
+        cls._cache[(device_type, device_id)] = inst
+        return inst
+
+    def __init__(self, device_type, device_id):
+        pass
+
+    @property
+    def exist(self):
+        """Whether this device exist."""
+        return (
+            _api_internal._GetDeviceAttr(self.device_type, self.device_id, 0)
+            != 0
+        )
+
+    @property
+    def max_threads_per_block(self):
+        """Maximum number of threads on each block."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 1)
+
+    @property
+    def warp_size(self):
+        """Number of threads that executes in concurrent."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 2)
+
+    @property
+    def max_shared_memory_per_block(self):
+        """Total amount of shared memory per block in bytes."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 3)
+
+    @property
+    def compute_version(self):
+        """Get compute verison number in string.
+
+        Currently used to get compute capability of CUDA device.
+
+        Returns
+        -------
+        version : str
+            The version string in `major.minor` format.
+        """
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 4)
+
+    @property
+    def device_name(self):
+        """Return the string name of device."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 5)
+
+    @property
+    def max_clock_rate(self):
+        """Return the max clock frequency of device."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 6)
+
+    @property
+    def multi_processor_count(self):
+        """Return the number of compute units of device."""
+        return _api_internal._GetDeviceAttr(self.device_type, self.device_id, 7)
+
+    @property
+    def max_thread_dimensions(self):
+        """Return the maximum size of each thread axis
+
+        Returns
+        -------
+        dims: List of int
+            The maximum length of threadIdx.x, threadIdx.y, threadIdx.z
+        """
+        return json.loads(
+            _api_internal._GetDeviceAttr(self.device_type, self.device_id, 8)
+        )
+
+    def sync(self):
+        """Synchronize until jobs finished at the context."""
+        check_call(_LIB.DGLSynchronize(self.device_type, self.device_id, None))
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, DGLContext)
+            and self.device_id == other.device_id
+            and self.device_type == other.device_type
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def __repr__(self):
+        if self.device_type >= RPC_SESS_MASK:
+            tbl_id = self.device_type / RPC_SESS_MASK - 1
+            dev_type = self.device_type % RPC_SESS_MASK
+            return "remote[%d]:%s(%d)" % (
+                tbl_id,
+                DGLContext.MASK2STR[dev_type],
+                self.device_id,
+            )
+        return "%s(%d)" % (
+            DGLContext.MASK2STR[self.device_type],
+            self.device_id,
+        )
+
+    def __hash__(self):
+        return hash((self.device_type, self.device_id))
+
+
+class DGLArray(ctypes.Structure):
+    """DGLValue in C API"""
+
+    _fields_ = [
+        ("data", ctypes.c_void_p),
+        ("ctx", DGLContext),
+        ("ndim", ctypes.c_int),
+        ("dtype", DGLDataType),
+        ("shape", ctypes.POINTER(dgl_shape_index_t)),
+        ("strides", ctypes.POINTER(dgl_shape_index_t)),
+        ("byte_offset", ctypes.c_uint64),
+    ]
+
+
+DGLArrayHandle = ctypes.POINTER(DGLArray)
+
+DGLStreamHandle = ctypes.c_void_p

python/dgl/_ffi/streams.py

@@ -0,0 +1,46 @@
+# pylint: disable=invalid-name, unused-import
+"""Runtime stream APIs which are mainly for internal test use only.
+For applications, please use PyTorch's stream management, of which DGL is aware.
+"""
+from __future__ import absolute_import
+
+import ctypes
+
+from .base import _FFI_MODE, _LIB, check_call
+from .runtime_ctypes import DGLStreamHandle
+
+
+def to_dgl_stream_handle(cuda_stream):
+    """Convert torch.cuda.Stream to DGL stream handle
+
+    Parameters
+    ----------
+    cuda_stream : torch.cuda.Stream.
+
+    Returns
+    -------
+    DGLStreamHandle
+        DGLStreamHandle of the input ``cuda_stream``.
+    """
+    return ctypes.c_void_p(cuda_stream.cuda_stream)
+
+
+def _dgl_get_stream(ctx):
+    """Get the current CUDA stream of the given DGL context.
+
+    Parameters
+    ----------
+    ctx : DGL context.
+
+    Returns
+    -------
+    DGLStreamHandle
+        DGLStreamHandle of the current CUDA stream.
+    """
+    current_cuda_stream = DGLStreamHandle()
+    check_call(
+        _LIB.DGLGetStream(
+            ctx.device_type, ctx.device_id, ctypes.byref(current_cuda_stream)
+        )
+    )
+    return current_cuda_stream