Merge pull request #1 from jermainewang/cpp

Cpp

include/dgl/runtime/registry.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/registry.h
+ * \brief This file defines the TVM global function registry.
+ *
+ *  The registered functions will be made available to front-end
+ *  as well as backend users.
+ *
+ *  The registry stores type-erased functions.
+ *  Each registered function is automatically exposed
+ *  to front-end language(e.g. python).
+ *
+ *  Front-end can also pass callbacks as PackedFunc, or register
+ *  then into the same global registry in C++.
+ *  The goal is to mix the front-end language and the TVM back-end.
+ *
+ * \code
+ *   // register the function as MyAPIFuncName
+ *   TVM_REGISTER_GLOBAL(MyAPIFuncName)
+ *   .set_body([](TVMArgs args, TVMRetValue* rv) {
+ *     // my code.
+ *   });
+ * \endcode
+ */
+#ifndef TVM_RUNTIME_REGISTRY_H_
+#define TVM_RUNTIME_REGISTRY_H_
+
+#include <string>
+#include <vector>
+#include "packed_func.h"
+
+namespace tvm {
+namespace runtime {
+
+/*! \brief Registry for global function */
+class Registry {
+ public:
+  /*!
+   * \brief set the body of the function to be f
+   * \param f The body of the function.
+   */
+  TVM_DLL Registry& set_body(PackedFunc f);  // NOLINT(*)
+  /*!
+   * \brief set the body of the function to be f
+   * \param f The body of the function.
+   */
+  Registry& set_body(PackedFunc::FType f) {  // NOLINT(*)
+    return set_body(PackedFunc(f));
+  }
+  /*!
+   * \brief set the body of the function to be TypedPackedFunc.
+   *
+   * \code
+   *
+   * TVM_REGISTER_API("addone")
+   * .set_body_typed<int(int)>([](int x) { return x + 1; });
+   *
+   * \endcode
+   *
+   * \param f The body of the function.
+   * \tparam FType the signature of the function.
+   * \tparam FLambda The type of f.
+   */
+  template<typename FType, typename FLambda>
+  Registry& set_body_typed(FLambda f) {
+    return set_body(TypedPackedFunc<FType>(f).packed());
+  }
+  /*!
+   * \brief Register a function with given name
+   * \param name The name of the function.
+   * \param override Whether allow oveeride existing function.
+   * \return Reference to theregistry.
+   */
+  TVM_DLL static Registry& Register(const std::string& name, bool override = false);  // NOLINT(*)
+  /*!
+   * \brief Erase global function from registry, if exist.
+   * \param name The name of the function.
+   * \return Whether function exist.
+   */
+  TVM_DLL static bool Remove(const std::string& name);
+  /*!
+   * \brief Get the global function by name.
+   * \param name The name of the function.
+   * \return pointer to the registered function,
+   *   nullptr if it does not exist.
+   */
+  TVM_DLL static const PackedFunc* Get(const std::string& name);  // NOLINT(*)
+  /*!
+   * \brief Get the names of currently registered global function.
+   * \return The names
+   */
+  TVM_DLL static std::vector<std::string> ListNames();
+
+  // Internal class.
+  struct Manager;
+
+ protected:
+  /*! \brief name of the function */
+  std::string name_;
+  /*! \brief internal packed function */
+  PackedFunc func_;
+  friend struct Manager;
+};
+
+/*! \brief helper macro to supress unused warning */
+#if defined(__GNUC__)
+#define TVM_ATTRIBUTE_UNUSED __attribute__((unused))
+#else
+#define TVM_ATTRIBUTE_UNUSED
+#endif
+
+#define TVM_STR_CONCAT_(__x, __y) __x##__y
+#define TVM_STR_CONCAT(__x, __y) TVM_STR_CONCAT_(__x, __y)
+
+#define TVM_FUNC_REG_VAR_DEF                                            \
+  static TVM_ATTRIBUTE_UNUSED ::tvm::runtime::Registry& __mk_ ## TVM
+
+#define TVM_TYPE_REG_VAR_DEF                                            \
+  static TVM_ATTRIBUTE_UNUSED ::tvm::runtime::ExtTypeVTable* __mk_ ## TVMT
+
+/*!
+ * \brief Register a function globally.
+ * \code
+ *   TVM_REGISTER_GLOBAL("MyPrint")
+ *   .set_body([](TVMArgs args, TVMRetValue* rv) {
+ *   });
+ * \endcode
+ */
+#define TVM_REGISTER_GLOBAL(OpName)                              \
+  TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) =            \
+      ::tvm::runtime::Registry::Register(OpName)
+
+/*!
+ * \brief Macro to register extension type.
+ *  This must be registered in a cc file
+ *  after the trait extension_class_info is defined.
+ */
+#define TVM_REGISTER_EXT_TYPE(T)                                 \
+  TVM_STR_CONCAT(TVM_TYPE_REG_VAR_DEF, __COUNTER__) =            \
+      ::tvm::runtime::ExtTypeVTable::Register_<T>()
+
+}  // namespace runtime
+}  // namespace tvm
+#endif  // TVM_RUNTIME_REGISTRY_H_

include/dgl/runtime/serializer.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/serializer.h
+ * \brief Serializer extension to support TVM data types
+ *  Include this file to enable serialization of DLDataType, DLContext
+ */
+#ifndef TVM_RUNTIME_SERIALIZER_H_
+#define TVM_RUNTIME_SERIALIZER_H_
+
+#include <dmlc/io.h>
+#include <dmlc/serializer.h>
+#include "c_runtime_api.h"
+#include "ndarray.h"
+
+namespace dmlc {
+namespace serializer {
+
+template<>
+struct Handler<DLDataType> {
+  inline static void Write(Stream *strm, const DLDataType& dtype) {
+    Handler<uint8_t>::Write(strm, dtype.code);
+    Handler<uint8_t>::Write(strm, dtype.bits);
+    Handler<uint16_t>::Write(strm, dtype.lanes);
+  }
+  inline static bool Read(Stream *strm, DLDataType* dtype) {
+    if (!Handler<uint8_t>::Read(strm, &(dtype->code))) return false;
+    if (!Handler<uint8_t>::Read(strm, &(dtype->bits))) return false;
+    if (!Handler<uint16_t>::Read(strm, &(dtype->lanes))) return false;
+    return true;
+  }
+};
+
+template<>
+struct Handler<DLContext> {
+  inline static void Write(Stream *strm, const DLContext& ctx) {
+    int32_t device_type = static_cast<int32_t>(ctx.device_type);
+    Handler<int32_t>::Write(strm, device_type);
+    Handler<int32_t>::Write(strm, ctx.device_id);
+  }
+  inline static bool Read(Stream *strm, DLContext* ctx) {
+    int32_t device_type = 0;
+    if (!Handler<int32_t>::Read(strm, &(device_type))) return false;
+    ctx->device_type = static_cast<DLDeviceType>(device_type);
+    if (!Handler<int32_t>::Read(strm, &(ctx->device_id))) return false;
+    return true;
+  }
+};
+
+}  // namespace serializer
+}  // namespace dmlc
+#endif  // TVM_RUNTIME_SERIALIZER_H_

include/dgl/runtime/threading_backend.h

+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file tvm/runtime/threading_backend.h
+ * \brief Utilities for manipulating thread pool threads.
+ */
+#ifndef TVM_RUNTIME_THREADING_BACKEND_H_
+#define TVM_RUNTIME_THREADING_BACKEND_H_
+
+#include <functional>
+#include <memory>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+namespace threading {
+
+/*!
+ * \brief A platform-agnostic abstraction for managing a collection of
+ *        thread pool threads.
+ */
+class ThreadGroup {
+ public:
+  class Impl;
+
+   /*!
+    * \brief Creates a collection of threads which run a provided function.
+    *
+    * \param num_workers The total number of worker threads in this group.
+             Includes main thread if `exclude_worker0 = true`
+    * \param worker_callback A callback which is run in its own thread.
+             Receives the worker_id as an argument.
+    * \param exclude_worker0 Whether to use the main thread as a worker.
+    *        If  `true`, worker0 will not be launched in a new thread and
+    *        `worker_callback` will only be called for values >= 1. This
+    *        allows use of the main thread as a worker.
+    */
+  ThreadGroup(int num_workers,
+              std::function<void(int)> worker_callback,
+              bool exclude_worker0 = false);
+  ~ThreadGroup();
+
+   /*!
+    * \brief Blocks until all non-main threads in the pool finish.
+    */
+  void Join();
+
+  enum AffinityMode : int {
+    kBig = 1,
+    kLittle = -1,
+  };
+
+  /*!
+   * \brief configure the CPU id affinity
+   *
+   * \param mode The preferred CPU type (1 = big, -1 = little).
+   * \param nthreads The number of threads to use (0 = use all).
+   * \param exclude_worker0 Whether to use the main thread as a worker.
+   *        If  `true`, worker0 will not be launched in a new thread and
+   *        `worker_callback` will only be called for values >= 1. This
+   *        allows use of the main thread as a worker.
+   *
+   * \return The number of workers to use.
+   */
+  int Configure(AffinityMode mode, int nthreads, bool exclude_worker0);
+
+ private:
+  Impl* impl_;
+};
+
+/*!
+ * \brief Platform-agnostic no-op.
+ */
+void Yield();
+
+/*!
+ * \return the maximum number of effective workers for this system.
+ */
+int MaxConcurrency();
+
+
+}  // namespace threading
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_THREADING_BACKEND_H_

include/dgl/runtime/util.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/util.h
+ * \brief Useful runtime util.
+ */
+#ifndef TVM_RUNTIME_UTIL_H_
+#define TVM_RUNTIME_UTIL_H_
+
+#include "c_runtime_api.h"
+
+namespace tvm {
+namespace runtime {
+
+/*!
+ * \brief Check whether type matches the given spec.
+ * \param t The type
+ * \param code The type code.
+ * \param bits The number of bits to be matched.
+ * \param lanes The number of lanes sin the type.
+ */
+inline bool TypeMatch(TVMType t, int code, int bits, int lanes = 1) {
+  return t.code == code && t.bits == bits && t.lanes == lanes;
+}
+}  // namespace runtime
+}  // namespace tvm
+// Forward declare the intrinsic id we need
+// in structure fetch to enable stackvm in runtime
+namespace tvm {
+namespace ir {
+namespace intrinsic {
+/*! \brief The kind of structure field info used in intrinsic */
+enum TVMStructFieldKind : int {
+  // array head address
+  kArrAddr,
+  kArrData,
+  kArrShape,
+  kArrStrides,
+  kArrNDim,
+  kArrTypeCode,
+  kArrTypeBits,
+  kArrTypeLanes,
+  kArrByteOffset,
+  kArrDeviceId,
+  kArrDeviceType,
+  kArrKindBound_,
+  // TVMValue field
+  kTVMValueContent,
+  kTVMValueKindBound_
+};
+}  // namespace intrinsic
+}  // namespace ir
+}  // namespace tvm
+#endif  // TVM_RUNTIME_UTIL_H_

python/dgl/__init__.py

-from .base import ALL
-
-from .batch import batch, unbatch
+from . import backend
+from . import data
+from . import function
+from . import generator
+from . import nn
 
-from .context import cpu, gpu
+from ._ffi.runtime_ctypes import TypeCode
+from ._ffi.function import register_func, get_global_func, list_global_func_names, extract_ext_funcs
+from ._ffi.base import DGLError, __version__
 
+from .base import ALL
+from .batched_graph import *
+from .generator import *
 from .graph import DGLGraph, __MSG__, __REPR__
-
 from .subgraph import DGLSubGraph
-
-from .generator import *

python/dgl/_ffi/README.md

+# C API and runtime
+
+Borrowed and adapted from TVM project.

python/dgl/_ffi/__init__.py

+"""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

+"""ctypes specific implementation of FFI"""

python/dgl/_ffi/_ctypes/function.py

+# 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 Number, Integral
+
+from ..base import _LIB, check_call
+from ..base import c_str, string_types
+from ..runtime_ctypes import TVMType, TVMByteArray, TVMContext
+from . import ndarray as _nd
+from .ndarray import NDArrayBase, _make_array
+from .types import TVMValue, TypeCode
+from .types import TVMPackedCFunc, TVMCFuncFinalizer
+from .types import RETURN_SWITCH, C_TO_PY_ARG_SWITCH, _wrap_arg_func
+
+FunctionHandle = ctypes.c_void_p
+ModuleHandle = ctypes.c_void_p
+TVMRetValueHandle = 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
+TVM_FREE_PYOBJ = TVMCFuncFinalizer(_ctypes_free_resource)
+ctypes.pythonapi.Py_IncRef(ctypes.py_object(TVM_FREE_PYOBJ))
+
+def convert_to_tvm_func(pyfunc):
+    """Convert a python function to TVM function
+
+    Parameters
+    ----------
+    pyfunc : python function
+        The python function to be converted.
+
+    Returns
+    -------
+    tvmfunc: tvm.nd.Function
+        The converted tvm 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.TVMAPISetLastError(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_tvm_args((rv,), temp_args)
+            if not isinstance(ret, TVMRetValueHandle):
+                ret = TVMRetValueHandle(ret)
+            check_call(_LIB.TVMCFuncSetReturn(ret, values, tcodes, ctypes.c_int(1)))
+            _ = temp_args
+            _ = rv
+        return 0
+
+    handle = FunctionHandle()
+    f = TVMPackedCFunc(cfun)
+    # NOTE: We will need to use python-api to increase ref count of the f
+    # TVM_FREE_PYOBJ will be called after it is no longer needed.
+    pyobj = ctypes.py_object(f)
+    ctypes.pythonapi.Py_IncRef(pyobj)
+    check_call(_LIB.TVMFuncCreateFromCFunc(
+        f, pyobj, TVM_FREE_PYOBJ, ctypes.byref(handle)))
+    return _CLASS_FUNCTION(handle, False)
+
+
+def _make_tvm_args(args, temp_args):
+    """Pack arguments into c args tvm call accept"""
+    num_args = len(args)
+    values = (TVMValue * 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, 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._TVM_COMPATS):
+            values[i].v_handle = ctypes.c_void_p(arg._tvm_handle)
+            type_codes[i] = arg.__class__._tvm_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, TVMType):
+            values[i].v_str = c_str(str(arg))
+            type_codes[i] = TypeCode.STR
+        elif isinstance(arg, TVMContext):
+            values[i].v_ctx = arg
+            type_codes[i] = TypeCode.TVM_CONTEXT
+        elif isinstance(arg, bytearray):
+            arr = TVMByteArray()
+            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_tvm_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.TVMFuncFree(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_tvm_args(args, temp_args)
+        ret_val = TVMValue()
+        ret_tcode = ctypes.c_int()
+        check_call(_LIB.TVMFuncCall(
+            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_tvm_args(args, temp_args)
+    ret_val = TVMValue()
+    ret_tcode = ctypes.c_int()
+    check_call(_LIB.TVMFuncCall(
+        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.NODE_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
+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

+# pylint: disable=invalid-name
+"""Runtime NDArray api"""
+from __future__ import absolute_import
+
+import ctypes
+from ..base import _LIB, check_call, c_str
+from ..runtime_ctypes import TVMArrayHandle
+from .types import RETURN_SWITCH, C_TO_PY_ARG_SWITCH, _wrap_arg_func, _return_handle
+
+
+TVMPyCapsuleDestructor = 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 = TVMArrayHandle()
+        check_call(_LIB.TVMArrayFromDLPack(ptr, ctypes.byref(handle)))
+        ctypes.pythonapi.PyCapsule_SetName(dltensor, _c_str_used_dltensor)
+        ctypes.pythonapi.PyCapsule_SetDestructor(dltensor, TVMPyCapsuleDestructor(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.TVMDLManagedTensorCallDeleter(ptr)
+        ctypes.pythonapi.PyCapsule_SetDestructor(pycapsule, TVMPyCapsuleDestructor(0))
+
+_c_dlpack_deleter = TVMPyCapsuleDestructor(_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 : TVMArrayHandle
+            the handle to the underlying C++ TVMArray
+        """
+        self.handle = handle
+        self.is_view = is_view
+
+    def __del__(self):
+        if not self.is_view and _LIB:
+            check_call(_LIB.TVMArrayFree(self.handle))
+
+    @property
+    def _tvm_handle(self):
+        return ctypes.cast(self.handle, ctypes.c_void_p).value
+
+    def to_dlpack(self):
+        """Produce an array from a DLPack Tensor without copying memory
+
+        Returns
+        -------
+        dlpack : DLPack tensor view of the array data
+        """
+        ptr = ctypes.c_void_p()
+        check_call(_LIB.TVMArrayToDLPack(self.handle, ctypes.byref(ptr)))
+        return ctypes.pythonapi.PyCapsule_New(ptr, _c_str_dltensor, _c_dlpack_deleter)
+
+
+def _make_array(handle, is_view):
+    handle = ctypes.cast(handle, TVMArrayHandle)
+    return _CLASS_NDARRAY(handle, is_view)
+
+_TVM_COMPATS = ()
+
+def _reg_extension(cls, fcreate):
+    global _TVM_COMPATS
+    _TVM_COMPATS += (cls,)
+    if fcreate:
+        fret = lambda x: fcreate(_return_handle(x))
+        RETURN_SWITCH[cls._tvm_tcode] = fret
+        C_TO_PY_ARG_SWITCH[cls._tvm_tcode] = _wrap_arg_func(fret, cls._tvm_tcode)
+
+
+_CLASS_NDARRAY = None
+
+def _set_class_ndarray(cls):
+    global _CLASS_NDARRAY
+    _CLASS_NDARRAY = cls

python/dgl/_ffi/_ctypes/types.py

+"""The C Types used in API."""
+# pylint: disable=invalid-name
+from __future__ import absolute_import as _abs
+
+import ctypes
+from ..base import py_str, check_call, _LIB
+from ..runtime_ctypes import TVMByteArray, TypeCode
+
+class TVMValue(ctypes.Union):
+    """TVMValue 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)]
+
+
+TVMPackedCFunc = ctypes.CFUNCTYPE(
+    ctypes.c_int,
+    ctypes.POINTER(TVMValue),
+    ctypes.POINTER(ctypes.c_int),
+    ctypes.c_int,
+    ctypes.c_void_p,
+    ctypes.c_void_p)
+
+
+TVMCFuncFinalizer = 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(TVMByteArray))[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.TVMCbArgToReturn(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
+}
+
+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
+}

python/dgl/_ffi/_cy2/__init__.py

+"""cython2 namespace"""

python/dgl/_ffi/_cy3/__init__.py

+"""cython3 namespace"""

python/dgl/_ffi/_cython/base.pxi

+from ..base import DGLError
+from libcpp.vector cimport vector
+from cpython.version cimport PY_MAJOR_VERSION
+from cpython cimport pycapsule
+from libc.stdint cimport int64_t, uint64_t, uint8_t, uint16_t
+import ctypes
+
+cdef enum TVMTypeCode:
+    kInt = 0
+    kUInt = 1
+    kFloat = 2
+    kHandle = 3
+    kNull = 4
+    kTVMType = 5
+    kTVMContext = 6
+    kArrayHandle = 7
+    kNodeHandle = 8
+    kModuleHandle = 9
+    kFuncHandle = 10
+    kStr = 11
+    kBytes = 12
+    kNDArrayContainer = 13
+    kExtBegin = 15
+
+cdef extern from "tvm/runtime/c_runtime_api.h":
+    ctypedef struct DLDataType:
+        uint8_t code
+        uint8_t bits
+        uint16_t lanes
+
+    ctypedef struct DLContext:
+        int device_type
+        int device_id
+
+    ctypedef struct DLTensor:
+        void* data
+        DLContext ctx
+        int ndim
+        DLDataType dtype
+        int64_t* shape
+        int64_t* strides
+        uint64_t byte_offset
+
+    ctypedef struct DLManagedTensor:
+        DLTensor dl_tensor
+        void* manager_ctx
+        void (*deleter)(DLManagedTensor* self)
+
+    ctypedef struct TVMValue:
+        int64_t v_int64
+        double v_float64
+        void* v_handle
+        const char* v_str
+        DLDataType v_type
+        DLContext v_ctx
+
+ctypedef int64_t tvm_index_t
+ctypedef DLTensor* DLTensorHandle
+ctypedef void* TVMStreamHandle
+ctypedef void* TVMRetValueHandle
+ctypedef void* TVMFunctionHandle
+ctypedef void* NodeHandle
+
+ctypedef int (*TVMPackedCFunc)(
+    TVMValue* args,
+    int* type_codes,
+    int num_args,
+    TVMRetValueHandle ret,
+    void* resource_handle)
+
+ctypedef void (*TVMPackedCFuncFinalizer)(void* resource_handle)
+
+cdef extern from "tvm/runtime/c_runtime_api.h":
+    void TVMAPISetLastError(const char* msg)
+    const char *TVMGetLastError()
+    int TVMFuncCall(TVMFunctionHandle func,
+                    TVMValue* arg_values,
+                    int* type_codes,
+                    int num_args,
+                    TVMValue* ret_val,
+                    int* ret_type_code)
+    int TVMFuncFree(TVMFunctionHandle func)
+    int TVMCFuncSetReturn(TVMRetValueHandle ret,
+                          TVMValue* value,
+                          int* type_code,
+                          int num_ret)
+    int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
+                               void* resource_handle,
+                               TVMPackedCFuncFinalizer fin,
+                               TVMFunctionHandle *out)
+    int TVMCbArgToReturn(TVMValue* value, int code)
+    int TVMArrayAlloc(tvm_index_t* shape,
+                      tvm_index_t ndim,
+                      DLDataType dtype,
+                      DLContext ctx,
+                      DLTensorHandle* out)
+    int TVMArrayFree(DLTensorHandle handle)
+    int TVMArrayCopyFromTo(DLTensorHandle src,
+                           DLTensorHandle to,
+                           TVMStreamHandle stream)
+    int TVMArrayFromDLPack(DLManagedTensor* arr_from,
+                           DLTensorHandle* out)
+    int TVMArrayToDLPack(DLTensorHandle arr_from,
+                         DLManagedTensor** out)
+    void TVMDLManagedTensorCallDeleter(DLManagedTensor* dltensor)
+
+cdef extern from "tvm/c_dsl_api.h":
+    int TVMNodeFree(NodeHandle handle)
+    int TVMNodeTypeKey2Index(const char* type_key,
+                             int* out_index)
+    int TVMNodeGetTypeIndex(NodeHandle handle,
+                            int* out_index)
+    int TVMNodeGetAttr(NodeHandle handle,
+                       const char* key,
+                       TVMValue* out_value,
+                       int* out_type_code,
+                       int* out_success)
+
+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(TVMGetLastError()))
+
+
+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
+    v_ptr = handle.value
+    return <void*>(v_ptr)

python/dgl/_ffi/_cython/core.pyx

+include "./base.pxi"
+include "./node.pxi"
+include "./function.pxi"
+include "./ndarray.pxi"

python/dgl/_ffi/_cython/function.pxi

+import ctypes
+import traceback
+from cpython cimport Py_INCREF, Py_DECREF
+from numbers import Number, Integral
+from ..base import string_types
+from ..node_generic import convert_to_node, NodeGeneric
+from ..runtime_ctypes import TVMType, TVMContext, TVMByteArray
+
+
+cdef void tvm_callback_finalize(void* fhandle):
+    local_pyfunc = <object>(fhandle)
+    Py_DECREF(local_pyfunc)
+
+cdef int tvm_callback(TVMValue* args,
+                      int* type_codes,
+                      int num_args,
+                      TVMRetValueHandle ret,
+                      void* fhandle) with gil:
+    cdef list pyargs
+    cdef TVMValue value
+    cdef int tcode
+    local_pyfunc = <object>(fhandle)
+    pyargs = []
+    for i in range(num_args):
+        value = args[i]
+        tcode = type_codes[i]
+        if (tcode == kNodeHandle or
+            tcode == kFuncHandle or
+            tcode == kModuleHandle or
+            tcode > kExtBegin):
+            CALL(TVMCbArgToReturn(&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()
+        TVMAPISetLastError(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(TVMCFuncSetReturn(ret, &value, &tcode, 1))
+    return 0
+
+
+def convert_to_tvm_func(object pyfunc):
+    """Convert a python function to TVM function
+
+    Parameters
+    ----------
+    pyfunc : python function
+        The python function to be converted.
+
+    Returns
+    -------
+    tvmfunc: tvm.Function
+        The converted tvm function.
+    """
+    cdef TVMFunctionHandle chandle
+    Py_INCREF(pyfunc)
+    CALL(TVMFuncCreateFromCFunc(tvm_callback,
+                                <void*>(pyfunc),
+                                tvm_callback_finalize,
+                                &chandle))
+    ret = _CLASS_FUNCTION(None, False)
+    (<FunctionBase>ret).chandle = chandle
+    return ret
+
+
+cdef inline int make_arg(object arg,
+                         TVMValue* value,
+                         int* tcode,
+                         list temp_args) except -1:
+    """Pack arguments into c args tvm call accept"""
+    cdef unsigned long long ptr
+    if isinstance(arg, NodeBase):
+        value[0].v_handle = (<NodeBase>arg).chandle
+        tcode[0] = kNodeHandle
+    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, _TVM_COMPATS):
+        ptr = arg._tvm_handle
+        value[0].v_handle = (<void*>ptr)
+        tcode[0] = arg.__class__._tvm_tcode
+    elif isinstance(arg, (int, long)):
+        value[0].v_int64 = arg
+        tcode[0] = kInt
+    elif isinstance(arg, float):
+        value[0].v_float64 = arg
+        tcode[0] = kFloat
+    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] = kFloat
+    elif isinstance(arg, TVMType):
+        tstr = c_str(str(arg))
+        value[0].v_str = tstr
+        tcode[0] = kStr
+        temp_args.append(tstr)
+    elif isinstance(arg, TVMContext):
+        value[0].v_ctx = (<DLContext*>(
+            <unsigned long long>ctypes.addressof(arg)))[0]
+        tcode[0] = kTVMContext
+    elif isinstance(arg, bytearray):
+        arr = TVMByteArray()
+        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, NodeGeneric)):
+        arg = convert_to_node(arg)
+        value[0].v_handle = (<NodeBase>arg).chandle
+        tcode[0] = kNodeHandle
+        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_tvm_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(TVMByteArray))[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(TVMValue value, int tcode):
+    """convert result to return value."""
+    if tcode == kNodeHandle:
+        return make_ret_node(value.v_handle)
+    elif tcode == kNull:
+        return None
+    elif tcode == kInt:
+        return value.v_int64
+    elif tcode == kFloat:
+        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 == kTVMContext:
+        return TVMContext(value.v_ctx.device_type, value.v_ctx.device_id)
+    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 _TVM_EXT_RET:
+        return _TVM_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,
+                          TVMValue* ret_val,
+                          int* ret_tcode) except -1:
+    cdef TVMValue[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)
+    CALL(TVMFuncCall(chandle, &values[0], &tcodes[0],
+                     nargs, ret_val, ret_tcode))
+    return 0
+
+cdef inline int FuncCall(void* chandle,
+                         tuple args,
+                         TVMValue* 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[TVMValue] 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)
+    CALL(TVMFuncCall(chandle, &values[0], &tcodes[0],
+                     nargs, ret_val, ret_tcode))
+    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 TVMValue 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 TVMFunctionHandle 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(TVMFuncFree(self.chandle))
+
+    def __call__(self, *args):
+        cdef TVMValue 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

+from ..runtime_ctypes import TVMArrayHandle
+
+cdef const char* _c_str_dltensor = "dltensor"
+cdef const char* _c_str_used_dltensor = "used_dltensor"
+
+
+cdef void _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)
+        TVMDLManagedTensorCallDeleter(dltensor)
+
+
+def _from_dlpack(object dltensor):
+    cdef DLManagedTensor* ptr
+    cdef DLTensorHandle chandle
+    if pycapsule.PyCapsule_IsValid(dltensor, _c_str_dltensor):
+        ptr = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(dltensor, _c_str_dltensor)
+        CALL(TVMArrayFromDLPack(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 DLTensor* 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 = <DLTensor*>(ptr)
+
+    property _tvm_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, TVMArrayHandle)
+
+        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(TVMArrayFree(self.chandle))
+
+    def to_dlpack(self):
+        """Produce an array from a DLPack Tensor without copying memory
+
+        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(TVMArrayToDLPack(self.chandle, &dltensor))
+        return pycapsule.PyCapsule_New(dltensor, _c_str_dltensor, _c_dlpack_deleter)
+
+
+cdef c_make_array(void* chandle, is_view):
+    ret = _CLASS_NDARRAY(None, is_view)
+    (<NDArrayBase>ret).chandle = <DLTensor*>chandle
+    return ret
+
+
+cdef _TVM_COMPATS = ()
+
+cdef _TVM_EXT_RET = {}
+
+def _reg_extension(cls, fcreate):
+    global _TVM_COMPATS
+    _TVM_COMPATS += (cls,)
+    if fcreate:
+        _TVM_EXT_RET[cls._tvm_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/node.pxi

+from ... import _api_internal
+from ..base import string_types
+from ..node_generic import _set_class_node_base
+
+"""Maps node type to its constructor"""
+NODE_TYPE = []
+
+def _register_node(int index, object cls):
+    """register node class"""
+    while len(NODE_TYPE) <= index:
+        NODE_TYPE.append(None)
+    NODE_TYPE[index] = cls
+
+
+cdef inline object make_ret_node(void* chandle):
+    global NODE_TYPE
+    cdef int tindex
+    cdef list node_type
+    cdef object cls
+    node_type = NODE_TYPE
+    CALL(TVMNodeGetTypeIndex(chandle, &tindex))
+    if tindex < len(node_type):
+        cls = node_type[tindex]
+        if cls is not None:
+            obj = cls.__new__(cls)
+        else:
+            obj = NodeBase.__new__(NodeBase)
+    else:
+        obj = NodeBase.__new__(NodeBase)
+    (<NodeBase>obj).chandle = chandle
+    return obj
+
+
+cdef class NodeBase:
+    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(TVMNodeFree(self.chandle))
+
+    def __getattr__(self, name):
+        cdef TVMValue ret_val
+        cdef int ret_type_code, ret_succ
+        CALL(TVMNodeGetAttr(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 Node object
+        instead of creating a new Node.
+        """
+        cdef void* chandle
+        ConstructorCall(
+            (<FunctionBase>fconstructor).chandle,
+            kNodeHandle, args, &chandle)
+        self.chandle = chandle
+
+_set_class_node_base(NodeBase)

python/dgl/_ffi/base.py

+# coding: utf-8
+# pylint: disable=invalid-name
+"""ctypes library and helper functions """
+from __future__ import absolute_import
+
+import sys
+import os
+import ctypes
+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
+
+
+def _load_lib():
+    """Load libary by searching possible path."""
+    lib_path = libinfo.find_lib_path()
+    lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL)
+    # DMatrix functions
+    lib.TVMGetLastError.restype = ctypes.c_char_p
+    return lib, os.path.basename(lib_path[0])
+
+# version number
+__version__ = libinfo.__version__
+# library instance of nnvm
+_LIB, _LIB_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.TVMGetLastError()))
+
+
+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)