import ffi solution from TVM

.gitmodules

+[submodule "third_party/dmlc-core"]
+	path = third_party/dmlc-core
+	url = https://github.com/dmlc/dmlc-core.git
+[submodule "third_party/dlpack"]
+	path = third_party/dlpack
+	url = https://github.com/dmlc/dlpack.git

CMakeLists.txt

+########################################
+# Borrowed and adapted from TVM project
+########################################
+cmake_minimum_required(VERSION 3.2)
+project(dgl C CXX)
+
+# Utility functions
+include(cmake/util/Util.cmake)
+
+if(EXISTS ${CMAKE_CURRENT_BINARY_DIR}/config.cmake)
+  include(${CMAKE_CURRENT_BINARY_DIR}/config.cmake)
+else()
+  if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/config.cmake)
+    include(${CMAKE_CURRENT_SOURCE_DIR}/config.cmake)
+  endif()
+endif()
+
+# NOTE: do not modify this file to change option values.
+# You can create a config.cmake at build folder
+# and add set(OPTION VALUE) to override these build options.
+# Alernatively, use cmake -DOPTION=VALUE through command-line.
+#tvm_option(USE_CUDA "Build with CUDA" OFF)
+#tvm_option(USE_OPENCL "Build with OpenCL" OFF)
+#tvm_option(USE_VULKAN "Build with Vulkan" OFF)
+#tvm_option(USE_OPENGL "Build with OpenGL" OFF)
+#tvm_option(USE_METAL "Build with Metal" OFF)
+#tvm_option(USE_ROCM "Build with ROCM" OFF)
+#tvm_option(ROCM_PATH "The path to rocm" /opt/rocm)
+#tvm_option(USE_RPC "Build with RPC" ON)
+#tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" OFF)
+#tvm_option(USE_STACKVM_RUNTIME "Include stackvm into the runtime" OFF)
+#tvm_option(USE_GRAPH_RUNTIME "Build with tiny graph runtime" ON)
+#tvm_option(USE_GRAPH_RUNTIME_DEBUG "Build with tiny graph runtime debug mode" OFF)
+#tvm_option(USE_RTTI "Build with RTTI" ON)
+#tvm_option(USE_MSVC_MT "Build with MT" OFF)
+#tvm_option(INSTALL_DEV "Install compiler infrastructure" OFF)
+
+# Contrib library options
+#tvm_option(USE_BLAS "The blas library to be linked" none)
+#tvm_option(USE_MKL_PATH "MKL root path when use MKL blas" none)
+#tvm_option(USE_CUDNN "Build with cuDNN" OFF)
+#tvm_option(USE_CUBLAS "Build with cuBLAS" OFF)
+#tvm_option(USE_MIOPEN "Build with ROCM:MIOpen" OFF)
+#tvm_option(USE_ROCBLAS "Build with ROCM:RoCBLAS" OFF)
+#tvm_option(USE_SORT "Build with sort support" OFF)
+#tvm_option(USE_NNPACK "Build with nnpack support" OFF)
+#tvm_option(USE_RANDOM "Build with random support" OFF)
+
+# include directories
+include_directories("include")
+include_directories("third_party/dlpack/include")
+include_directories("third_party/dmlc-core/include")
+
+# initial variables
+set(DGL_LINKER_LIBS "")
+set(DGL_RUNTIME_LINKER_LIBS "")
+
+# Generic compilation options
+if(MSVC)
+  add_definitions(-DWIN32_LEAN_AND_MEAN)
+  add_definitions(-D_CRT_SECURE_NO_WARNINGS)
+  add_definitions(-D_SCL_SECURE_NO_WARNINGS)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /EHsc")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /MP")
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /bigobj")
+  if(USE_MSVC_MT)
+    foreach(flag_var
+        CMAKE_CXX_FLAGS CMAKE_CXX_FLAGS_DEBUG CMAKE_CXX_FLAGS_RELEASE
+        CMAKE_CXX_FLAGS_MINSIZEREL CMAKE_CXX_FLAGS_RELWITHDEBINFO)
+      if(${flag_var} MATCHES "/MD")
+        string(REGEX REPLACE "/MD" "/MT" ${flag_var} "${${flag_var}}")
+      endif(${flag_var} MATCHES "/MD")
+    endforeach(flag_var)
+  endif()
+else(MSVC)
+  include(CheckCXXCompilerFlag)
+  check_cxx_compiler_flag("-std=c++11"    SUPPORT_CXX11)
+  set(CMAKE_C_FLAGS "-O2 -Wall -fPIC ${CMAKE_C_FLAGS}")
+  set(CMAKE_CXX_FLAGS "-O2 -Wall -fPIC -std=c++11 ${CMAKE_CXX_FLAGS}")
+endif(MSVC)
+
+# add source group
+#FILE(GLOB_RECURSE GROUP_SOURCE "src/*.cc")
+#FILE(GLOB_RECURSE GROUP_INCLUDE "src/*.h" "include/*.h")
+#assign_source_group("Source" ${GROUP_SOURCE})
+#assign_source_group("Include" ${GROUP_INCLUDE})
+
+# Source file lists
+file(GLOB CORE_SRCS
+    src/graph/*.cc
+    )
+
+file(GLOB RUNTIME_SRCS src/runtime/*.cc)
+
+# Package runtime rules
+#if(NOT USE_RTTI)
+#  add_definitions(-DDMLC_ENABLE_RTTI=0)
+#endif()
+#
+#if(USE_RPC)
+#  message(STATUS "Build with RPC support...")
+#  file(GLOB RUNTIME_RPC_SRCS src/runtime/rpc/*.cc)
+#  list(APPEND RUNTIME_SRCS ${RUNTIME_RPC_SRCS})
+#endif(USE_RPC)
+#
+#file(GLOB STACKVM_RUNTIME_SRCS src/runtime/stackvm/*.cc)
+#file(GLOB STACKVM_CODEGEN_SRCS src/codegen/stackvm/*.cc)
+#list(APPEND COMPILER_SRCS ${STACKVM_CODEGEN_SRCS})
+#if(USE_STACKVM_RUNTIME)
+#  message(STATUS "Build with stackvm support in runtime...")
+#  list(APPEND RUNTIME_SRCS ${STACKVM_RUNTIME_SRCS})
+#else()
+#  list(APPEND COMPILER_SRCS ${STACKVM_RUNTIME_SRCS})
+#endif(USE_STACKVM_RUNTIME)
+#
+#if(USE_GRAPH_RUNTIME)
+#  message(STATUS "Build with Graph runtime support...")
+#  file(GLOB RUNTIME_GRAPH_SRCS src/runtime/graph/*.cc)
+#  list(APPEND RUNTIME_SRCS ${RUNTIME_GRAPH_SRCS})
+#
+#  if(USE_GRAPH_RUNTIME_DEBUG)
+#    set_source_files_properties(${RUNTIME_GRAPH_SRCS}
+#      PROPERTIES COMPILE_DEFINITIONS "TVM_GRAPH_RUNTIME_DEBUG")
+#  endif(USE_GRAPH_RUNTIME_DEBUG)
+#endif(USE_GRAPH_RUNTIME)
+
+# Module rules
+#include(cmake/modules/VTA.cmake)
+#include(cmake/modules/CUDA.cmake)
+#include(cmake/modules/OpenCL.cmake)
+#include(cmake/modules/OpenGL.cmake)
+#include(cmake/modules/Vulkan.cmake)
+#include(cmake/modules/Metal.cmake)
+#include(cmake/modules/ROCM.cmake)
+#include(cmake/modules/LLVM.cmake)
+#include(cmake/modules/contrib/BLAS.cmake)
+#include(cmake/modules/contrib/Random.cmake)
+#include(cmake/modules/contrib/Sort.cmake)
+#include(cmake/modules/contrib/NNPack.cmake)
+
+add_library(dgl SHARED ${CORE_SRCS} ${RUNTIME_SRCS})
+#add_library(dgl_runtime SHARED ${RUNTIME_SRCS})
+
+target_link_libraries(dgl ${DGL_LINKER_LIBS} ${DGL_RUNTIME_LINKER_LIBS})
+#target_link_libraries(dgl_runtime ${DGL_RUNTIME_LINKER_LIBS})
+
+# Related headers
+#target_include_directories(
+#  dgl
+#  PUBLIC "HalideIR/src"
+#  PUBLIC "topi/include")
+
+# Tests
+#set(TEST_EXECS "")
+#file(GLOB TEST_SRCS tests/cpp/*.cc)
+#find_library(GTEST_LIB gtest)
+
+#if(GTEST_LIB)
+#  foreach(__srcpath ${TEST_SRCS})
+#    get_filename_component(__srcname ${__srcpath} NAME)
+#    string(REPLACE ".cc" "" __execname ${__srcname})
+#    add_executable(${__execname} ${__srcpath})
+#    list(APPEND TEST_EXECS ${__execname})
+#    target_link_libraries(${__execname}
+#      tvm ${GTEST_LIB} pthread)
+#    set_target_properties(${__execname} PROPERTIES EXCLUDE_FROM_ALL 1)
+#    set_target_properties(${__execname} PROPERTIES EXCLUDE_FROM_DEFAULT_BUILD 1)
+#  endforeach()
+#  add_custom_target(cpptest DEPENDS ${TEST_EXECS})
+#endif()
+
+# Custom targets
+#add_custom_target(runtime DEPENDS tvm_runtime)
+
+# Installation rules
+install(TARGETS dgl DESTINATION lib${LIB_SUFFIX})
+#install(TARGETS dgl_runtime DESTINATION lib${LIB_SUFFIX})
+
+#if (INSTALL_DEV)
+#  install(
+#    DIRECTORY "include/." DESTINATION "include"
+#    FILES_MATCHING
+#    PATTERN "*.h"
+#  )
+#  install(
+#    DIRECTORY "topi/include/." DESTINATION "include"
+#    FILES_MATCHING
+#    PATTERN "*.h"
+#  )
+#  install(
+#    DIRECTORY "HalideIR/src/." DESTINATION "include/HalideIR"
+#    FILES_MATCHING
+#    PATTERN "*.h"
+#  )
+#  install(
+#    DIRECTORY "dlpack/include/." DESTINATION "include"
+#    FILES_MATCHING
+#    PATTERN "*.h"
+#    )
+#  install(
+#    DIRECTORY "nnvm/include/." DESTINATION "include"
+#    FILES_MATCHING
+#    PATTERN "*.h"
+#    )
+#else(INSTALL_DEV)
+#  install(
+#    DIRECTORY "include/tvm/runtime/." DESTINATION "include/tvm/runtime"
+#    FILES_MATCHING
+#    PATTERN "*.h"
+#    )
+#endif(INSTALL_DEV)
+
+# More target definitions
+#if(MSVC)
+#  target_compile_definitions(tvm PRIVATE -DHalide_EXPORTS)
+#  target_compile_definitions(tvm_runtime PRIVATE -DHalide_EXPORTS)
+#  target_compile_definitions(tvm PRIVATE -DTVM_EXPORTS)
+#  target_compile_definitions(tvm_runtime PRIVATE -DTVM_EXPORTS)
+#  target_compile_definitions(nnvm_compiler PRIVATE -DNNVM_EXPORTS)
+#endif()

cmake/util/Util.cmake

+########################################
+# Borrowed and adapted from TVM project
+########################################
+macro(__dgl_option variable description value)
+  if(NOT DEFINED ${variable})
+    set(${variable} ${value} CACHE STRING ${description})
+  endif()
+endmacro()
+
+#######################################################
+# An option that the user can select. Can accept condition to control when option is available for user.
+# Usage:
+#   dgl_option(<option_variable> "doc string" <initial value or boolean expression> [IF <condition>])
+macro(dgl_option variable description value)
+  set(__value ${value})
+  set(__condition "")
+  set(__varname "__value")
+  foreach(arg ${ARGN})
+    if(arg STREQUAL "IF" OR arg STREQUAL "if")
+      set(__varname "__condition")
+    else()
+      list(APPEND ${__varname} ${arg})
+    endif()
+  endforeach()
+  unset(__varname)
+  if("${__condition}" STREQUAL "")
+    set(__condition 2 GREATER 1)
+  endif()
+
+  if(${__condition})
+    if("${__value}" MATCHES ";")
+      if(${__value})
+        __dgl_option(${variable} "${description}" ON)
+      else()
+        __dgl_option(${variable} "${description}" OFF)
+      endif()
+    elseif(DEFINED ${__value})
+      if(${__value})
+        __dgl_option(${variable} "${description}" ON)
+      else()
+        __dgl_option(${variable} "${description}" OFF)
+      endif()
+    else()
+      __dgl_option(${variable} "${description}" "${__value}")
+    endif()
+  else()
+    unset(${variable} CACHE)
+  endif()
+endmacro()
+
+function(assign_source_group group)
+    foreach(_source IN ITEMS ${ARGN})
+        if (IS_ABSOLUTE "${_source}")
+            file(RELATIVE_PATH _source_rel "${CMAKE_CURRENT_SOURCE_DIR}" "${_source}")
+        else()
+            set(_source_rel "${_source}")
+        endif()
+        get_filename_component(_source_path "${_source_rel}" PATH)
+        string(REPLACE "/" "\\" _source_path_msvc "${_source_path}")
+        source_group("${group}\\${_source_path_msvc}" FILES "${_source}")
+    endforeach()
+endfunction(assign_source_group)

include/dgl/runtime/README.md

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

include/dgl/runtime/c_backend_api.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/c_backend_api.h
+ * \brief TVM runtime backend API.
+ *
+ *  The functions defined in this header are intended to be
+ *  used by compiled tvm operators, usually user do not need to use these
+ *  function directly.
+ */
+#ifndef TVM_RUNTIME_C_BACKEND_API_H_
+#define TVM_RUNTIME_C_BACKEND_API_H_
+
+#include "c_runtime_api.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Backend related functions.
+/*!
+ * \brief Backend function for modules to get function
+ *  from its environment mod_node (its imports and global function).
+ *  The user do should not call TVMFuncFree on func.
+ *
+ * \param mod_node The module handle.
+ * \param func_name The name of the function.
+ * \param out The result function.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendGetFuncFromEnv(void* mod_node,
+                                     const char* func_name,
+                                     TVMFunctionHandle *out);
+/*!
+ * \brief Backend function to register system-wide library symbol.
+ *
+ * \param name The name of the symbol
+ * \param ptr The symbol address.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendRegisterSystemLibSymbol(const char* name, void* ptr);
+
+/*!
+ * \brief Backend function to allocate temporal workspace.
+ *
+ * \note The result allocate spaced is ensured to be aligned to kTempAllocaAlignment.
+ *
+ * \param nbytes The size of the space requested.
+ * \param device_type The device type which the space will be allocated.
+ * \param device_id The device id which the space will be allocated.
+ * \param dtype_code_hint The type code of the array elements. Only used in
+ * certain backends such as OpenGL.
+ * \param dtype_bits_hint The type bits of the array elements. Only used in
+ * certain backends such as OpenGL.
+ * \return nullptr when error is thrown, a valid ptr if success
+ */
+TVM_DLL void* TVMBackendAllocWorkspace(int device_type,
+                                       int device_id,
+                                       uint64_t nbytes,
+                                       int dtype_code_hint,
+                                       int dtype_bits_hint);
+
+/*!
+ * \brief Backend function to free temporal workspace.
+ *
+ * \param ptr The result allocated space pointer.
+ * \param device_type The device type which the space will be allocated.
+ * \param device_id The device id which the space will be allocated.
+ * \return 0 when no error is thrown, -1 when failure happens
+ *
+ * \sa TVMBackendAllocWorkspace
+ */
+TVM_DLL int TVMBackendFreeWorkspace(int device_type,
+                                    int device_id,
+                                    void* ptr);
+
+/*!
+ * \brief Environment for TVM parallel task.
+ */
+typedef struct {
+  /*!
+   * \brief Auxiliary used for synchronization
+   */
+  void* sync_handle;
+  /*! \brief total amount of task */
+  int32_t num_task;
+} TVMParallelGroupEnv;
+
+/*!
+ * \brief The callback function to execute a parallel lambda
+ * \param task_id the task id of the function.
+ * \param penv The parallel environment backs the execution.
+ * \param cdata The supporting closure data.
+ */
+typedef int (*FTVMParallelLambda)(
+    int task_id, TVMParallelGroupEnv* penv, void* cdata);
+
+/*!
+ * \brief Backend function for running parallel jobs.
+ *
+ * \param flambda The parallel function to be launched.
+ * \param cdata The closure data.
+ * \param num_task Number of tasks to launch, can be 0, means launch
+ *           with all available threads.
+ *
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendParallelLaunch(FTVMParallelLambda flambda,
+                                     void* cdata,
+                                     int num_task);
+
+/*!
+ * \brief BSP barrrier between parallel threads
+ * \param task_id the task id of the function.
+ * \param penv The parallel environment backs the execution.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv);
+
+
+/*!
+ * \brief Simple static initialization fucntion.
+ *  Run f once and set handle to be not null.
+ *  This function is mainly used for test purpose.
+ *
+ * \param handle An global address to indicate f
+ * \param f The function to be ran
+ * \param cdata The closure data to pass to the function.
+ * \param nbytes Number of bytes in the closure data.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendRunOnce(void** handle,
+                              int (*f)(void*),
+                              void *cdata,
+                              int nbytes);
+
+#ifdef __cplusplus
+}  // TVM_EXTERN_C
+#endif
+#endif  // TVM_RUNTIME_C_BACKEND_API_H_

include/dgl/runtime/c_runtime_api.h

+/*!
+ *  Copyright (c) 2016 by Contributors
+ * \file tvm/runtime/c_runtime_api.h
+ * \brief TVM runtime library.
+ *
+ *  The philosophy of TVM project is to customize the compilation
+ *  stage to generate code that can used by other projects transparently.
+ *  So this is a minimum runtime code gluing, and some limited
+ *  memory management code to enable quick testing.
+ *
+ *  The runtime API is independent from TVM compilation stack and can
+ *  be linked via libtvm_runtime.
+ *
+ *  The common flow is:
+ *   - Use TVMFuncListGlobalNames to get global function name
+ *   - Use TVMFuncCall to call these functions.
+ */
+#ifndef TVM_RUNTIME_C_RUNTIME_API_H_
+#define TVM_RUNTIME_C_RUNTIME_API_H_
+
+// Macros to do weak linking
+#ifdef _MSC_VER
+#define TVM_WEAK __declspec(selectany)
+#else
+#define TVM_WEAK __attribute__((weak))
+#endif
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten/emscripten.h>
+#define TVM_DLL EMSCRIPTEN_KEEPALIVE
+#endif
+
+#ifndef TVM_DLL
+#ifdef _WIN32
+#ifdef TVM_EXPORTS
+#define TVM_DLL __declspec(dllexport)
+#else
+#define TVM_DLL __declspec(dllimport)
+#endif
+#else
+#define TVM_DLL
+#endif
+#endif
+
+// TVM version
+#define TVM_VERSION "0.5.dev"
+
+
+// TVM Runtime is DLPack compatible.
+#include <dlpack/dlpack.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+#include <stdint.h>
+#include <stddef.h>
+
+/*! \brief type of array index. */
+typedef int64_t tvm_index_t;
+
+/*! \brief Extension device types in TVM */
+typedef enum {
+  kDLAOCL = 5,
+  kDLSDAccel = 6,
+  kOpenGL = 11,
+  // Extension DRAM type, used for quickly test extension device
+  // The device api can differ depending on the xpu driver registered.
+  kExtDev = 12,
+  // AddExtraTVMType which is not in DLPack here
+} TVMDeviceExtType;
+
+/*!
+ * \brief The type code in TVMType
+ * \note TVMType is used in two places.
+ */
+typedef enum {
+  // The type code of other types are compatible with DLPack.
+  // The next few fields are extension types
+  // that is used by TVM API calls.
+  kHandle = 3U,
+  kNull = 4U,
+  kTVMType = 5U,
+  kTVMContext = 6U,
+  kArrayHandle = 7U,
+  kNodeHandle = 8U,
+  kModuleHandle = 9U,
+  kFuncHandle = 10U,
+  kStr = 11U,
+  kBytes = 12U,
+  kNDArrayContainer = 13U,
+  // Extension codes for other frameworks to integrate TVM PackedFunc.
+  // To make sure each framework's id do not conflict, use first and
+  // last sections to mark ranges.
+  // Open an issue at the repo if you need a section of code.
+  kExtBegin = 15U,
+  kNNVMFirst = 16U,
+  kNNVMLast = 20U,
+  // The following section of code is used for non-reserved types.
+  kExtReserveEnd = 64U,
+  kExtEnd = 128U
+} TVMTypeCode;
+
+/*!
+ * \brief The data type used in TVM Runtime.
+ *
+ *  Examples
+ *   - float: type_code = 2, bits = 32, lanes=1
+ *   - float4(vectorized 4 float): type_code = 2, bits = 32, lanes=4
+ *   - int8: type_code = 0, bits = 8, lanes=1
+ *
+ * \note Arguments TVM API function always takes bits=64 and lanes=1
+ */
+typedef DLDataType TVMType;
+
+/*!
+ * \brief The Device information, abstract away common device types.
+ */
+typedef DLContext TVMContext;
+
+/*!
+ * \brief The tensor array stucture to TVM API.
+ */
+typedef DLTensor TVMArray;
+
+/*! \brief the array handle */
+typedef TVMArray* TVMArrayHandle;
+
+/*!
+ * \brief Union type of values
+ *  being passed through API and function calls.
+ */
+typedef union {
+  int64_t v_int64;
+  double v_float64;
+  void* v_handle;
+  const char* v_str;
+  TVMType v_type;
+  TVMContext v_ctx;
+} TVMValue;
+
+/*!
+ * \brief Byte array type used to pass in byte array
+ *  When kBytes is used as data type.
+ */
+typedef struct {
+  const char* data;
+  size_t size;
+} TVMByteArray;
+
+/*! \brief Handle to TVM runtime modules. */
+typedef void* TVMModuleHandle;
+/*! \brief Handle to packed function handle. */
+typedef void* TVMFunctionHandle;
+/*! \brief Handle to hold return value. */
+typedef void* TVMRetValueHandle;
+/*!
+ * \brief The stream that is specific to device
+ * can be NULL, which indicates the default one.
+ */
+typedef void* TVMStreamHandle;
+
+/*!
+ * \brief Used for implementing C API function.
+ *  Set last error message before return.
+ * \param msg The error message to be set.
+ */
+TVM_DLL void TVMAPISetLastError(const char* msg);
+
+/*!
+ * \brief return str message of the last error
+ *  all function in this file will return 0 when success
+ *  and -1 when an error occured,
+ *  TVMGetLastError can be called to retrieve the error
+ *
+ *  this function is threadsafe and can be called by different thread
+ *  \return error info
+ */
+TVM_DLL const char *TVMGetLastError(void);
+/*!
+ * \brief Load module from file.
+ * \param file_name The file name to load the module from.
+ * \param format The format of the module.
+ * \param out The result module
+ *
+ * \return 0 when success, -1 when failure happens
+ * \note The resulting module do not contain import relation.
+ *  It can be reconstructed by TVMModImport.
+ */
+TVM_DLL int TVMModLoadFromFile(const char* file_name,
+                               const char* format,
+                               TVMModuleHandle* out);
+
+/*!
+ * \brief Add dep to mod's dependency.
+ *  This allows functions in this module to use modules.
+ *
+ * \param mod The module handle.
+ * \param dep The dependent module to be imported.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMModImport(TVMModuleHandle mod,
+                         TVMModuleHandle dep);
+
+/*!
+ * \brief Get function from the module.
+ * \param mod The module handle.
+ * \param func_name The name of the function.
+ * \param query_imports Whether to query imported modules
+ * \param out The result function, can be NULL if it is not available.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMModGetFunction(TVMModuleHandle mod,
+                              const char* func_name,
+                              int query_imports,
+                              TVMFunctionHandle *out);
+
+/*!
+ * \brief Free front-end extension type resource.
+ * \param handle The extension handle.
+ * \param type_code The type of of the extension type.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMExtTypeFree(void* handle, int type_code);
+
+/*!
+ * \brief Free the Module
+ * \param mod The module to be freed.
+ *
+ * \note This may not free up the module's resources.
+ *  If there is active TVMFunctionHandle uses the module
+ *  Or if this module is imported by another active module.
+ *
+ *  The all functions remains valid until TVMFuncFree is called.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMModFree(TVMModuleHandle mod);
+
+/*!
+ * \brief Free the function when it is no longer needed.
+ * \param func The function handle
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMFuncFree(TVMFunctionHandle func);
+
+/*!
+ * \brief Call a Packed TVM Function.
+ *
+ * \param func node handle of the function.
+ * \param arg_values The arguments
+ * \param type_codes The type codes of the arguments
+ * \param num_args Number of arguments.
+ *
+ * \param ret_val The return value.
+ * \param ret_type_code the type code of return value.
+ *
+ * \return 0 when success, -1 when failure happens
+ * \note TVM calls always exchanges with type bits=64, lanes=1
+ *
+ * \note API calls always exchanges with type bits=64, lanes=1
+ *   If API call returns container handles (e.g. FunctionHandle)
+ *   these handles should be managed by the front-end.
+ *   The front-end need to call free function (e.g. TVMFuncFree)
+ *   to free these handles.
+ */
+TVM_DLL int TVMFuncCall(TVMFunctionHandle func,
+                        TVMValue* arg_values,
+                        int* type_codes,
+                        int num_args,
+                        TVMValue* ret_val,
+                        int* ret_type_code);
+
+/*!
+ * \brief Set the return value of TVMPackedCFunc.
+ *
+ *  This function is called by TVMPackedCFunc to set the return value.
+ *  When this function is not called, the function returns null by default.
+ *
+ * \param ret The return value handle, pass by ret in TVMPackedCFunc
+ * \param value The value to be returned.
+ * \param type_code The type of the value to be returned.
+ * \param num_ret Number of return values, for now only 1 is supported.
+ */
+TVM_DLL int TVMCFuncSetReturn(TVMRetValueHandle ret,
+                              TVMValue* value,
+                              int* type_code,
+                              int num_ret);
+
+/*!
+ * \brief Inplace translate callback argument value to return value.
+ *  This is only needed for non-POD arguments.
+ *
+ * \param value The value to be translated.
+ * \param code The type code to be translated.
+ * \note This function will do a shallow copy when necessary.
+ *
+ * \return 0 when success, -1 when failure happens.
+ */
+TVM_DLL int TVMCbArgToReturn(TVMValue* value, int code);
+
+/*!
+ * \brief C type of packed function.
+ *
+ * \param args The arguments
+ * \param type_codes The type codes of the arguments
+ * \param num_args Number of arguments.
+ * \param ret The return value handle.
+ * \param resource_handle The handle additional resouce handle from fron-end.
+ * \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
+ * \sa TVMCFuncSetReturn
+ */
+typedef int (*TVMPackedCFunc)(
+    TVMValue* args,
+    int* type_codes,
+    int num_args,
+    TVMRetValueHandle ret,
+    void* resource_handle);
+
+/*!
+ * \brief C callback to free the resource handle in C packed function.
+ * \param resource_handle The handle additional resouce handle from fron-end.
+ */
+typedef void (*TVMPackedCFuncFinalizer)(void* resource_handle);
+
+/*!
+ * \brief Signature for extension function declarer.
+ *
+ *  TVM call this function to get the extension functions
+ *  The declarer will call register_func to register function and their name.
+ *
+ * \param register_func_handle The register function
+ * \return 0 if success, -1 if failure happens
+ */
+typedef int (*TVMExtensionFuncDeclarer)(TVMFunctionHandle register_func_handle);
+
+/*!
+ * \brief Wrap a TVMPackedCFunc to become a FunctionHandle.
+ *
+ * The resource_handle will be managed by TVM API, until the function is no longer used.
+ *
+ * \param func The packed C function.
+ * \param resource_handle The resource handle from front-end, can be NULL.
+ * \param fin The finalizer on resource handle when the FunctionHandle get freed, can be NULL
+ * \param out the result function handle.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
+                                   void* resource_handle,
+                                   TVMPackedCFuncFinalizer fin,
+                                   TVMFunctionHandle *out);
+
+/*!
+ * \brief Register the function to runtime's global table.
+ *
+ * The registered function then can be pulled by the backend by the name.
+ *
+ * \param name The name of the function.
+ * \param f The function to be registered.
+ * \param override Whether allow override already registered function.
+ */
+TVM_DLL int TVMFuncRegisterGlobal(
+    const char* name, TVMFunctionHandle f, int override);
+
+/*!
+ * \brief Get a global function.
+ *
+ * \param name The name of the function.
+ * \param out the result function pointer, NULL if it does not exist.
+ *
+ * \note The function handle of global function is managed by TVM runtime,
+ *  So TVMFuncFree is should not be called when it get deleted.
+ */
+TVM_DLL int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out);
+
+/*!
+ * \brief List all the globally registered function name
+ * \param out_size The number of functions
+ * \param out_array The array of function names.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMFuncListGlobalNames(int* out_size,
+                                   const char*** out_array);
+
+// Array related apis for quick proptyping
+/*!
+ * \brief Allocate a nd-array's memory,
+ *  including space of shape, of given spec.
+ *
+ * \param shape The shape of the array, the data content will be copied to out
+ * \param ndim The number of dimension of the array.
+ * \param dtype_code The type code of the dtype
+ * \param dtype_bits The number of bits of dtype
+ * \param dtype_lanes The number of lanes in the dtype.
+ * \param device_type The device type of context
+ * \param device_id The device id of context.
+ * \param out The output handle.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayAlloc(const tvm_index_t* shape,
+                          int ndim,
+                          int dtype_code,
+                          int dtype_bits,
+                          int dtype_lanes,
+                          int device_type,
+                          int device_id,
+                          TVMArrayHandle* out);
+
+/*!
+ * \brief Free the TVM Array.
+ * \param handle The array handle to be freed.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayFree(TVMArrayHandle handle);
+
+/*!
+ * \brief Copy array data from CPU byte array.
+ * \param handle The array handle.
+ * \param data the data pointer
+ * \param nbytes The number of bytes to copy.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayCopyFromBytes(TVMArrayHandle handle,
+                                  void* data,
+                                  size_t nbytes);
+
+/*!
+ * \brief Copy array data to CPU byte array.
+ * \param handle The array handle.
+ * \param data the data pointer
+ * \param nbytes The number of bytes to copy.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayCopyToBytes(TVMArrayHandle handle,
+                                void* data,
+                                size_t nbytes);
+
+/*!
+ * \brief Copy the array, both from and to must be valid during the copy.
+ * \param from The array to be copied from.
+ * \param to The target space.
+ * \param stream The stream where the copy happens, can be NULL.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayCopyFromTo(TVMArrayHandle from,
+                               TVMArrayHandle to,
+                               TVMStreamHandle stream);
+
+/*!
+ * \brief Produce an array from the DLManagedTensor that shares data memory
+ * with the DLManagedTensor.
+ * \param from The source DLManagedTensor.
+ * \param out The output array handle.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayFromDLPack(DLManagedTensor* from,
+                               TVMArrayHandle* out);
+
+/*!
+ * \brief Produce a DLMangedTensor from the array that shares data memory with
+ * the array.
+ * \param from The source array.
+ * \param out The DLManagedTensor handle.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMArrayToDLPack(TVMArrayHandle from,
+                             DLManagedTensor** out);
+
+/*!
+ * \brief Delete (free) a DLManagedTensor's data.
+ * \param dltensor Pointer to the DLManagedTensor.
+ */
+TVM_DLL void TVMDLManagedTensorCallDeleter(DLManagedTensor* dltensor);
+
+/*!
+ * \brief Create a new runtime stream.
+ *
+ * \param device_type The device type of context
+ * \param device_id The device id of context
+ * \param out The new stream handle
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMStreamCreate(int device_type, int device_id, TVMStreamHandle* out);
+
+/*!
+ * \brief Free a created stream handle.
+ *
+ * \param device_type The device type of context
+ * \param device_id The device id of context
+ * \param stream The stream to be freed
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMStreamFree(int device_type, int device_id, TVMStreamHandle stream);
+
+/*!
+ * \brief Set the runtime stream of current thread to be stream.
+ *  The subsequent calls to the same device_type
+ *  will use the setted stream handle.
+ *  The specific type of stream is runtime device dependent.
+ *
+ * \param device_type The device type of context
+ * \param device_id The device id of context.
+ * \param handle The stream handle.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMSetStream(int device_type, int device_id, TVMStreamHandle handle);
+
+/*!
+ * \brief Wait until all computations on stream completes.
+ *
+ * \param device_type The device type of context
+ * \param device_id The device id of context.
+ * \param stream The stream to be synchronized.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMSynchronize(int device_type, int device_id, TVMStreamHandle stream);
+
+/*!
+ * \brief Synchronize two streams of execution.
+ *
+ * \param device_type The device type of context
+ * \param device_id The device id of context
+ * \param src The source stream to synchronize.
+ * \param dst The destination stream to synchronize.
+ * \return 0 when success, -1 when failure happens
+ */
+TVM_DLL int TVMStreamStreamSynchronize(int device_type,
+                                       int device_id,
+                                       TVMStreamHandle src,
+                                       TVMStreamHandle dst);
+
+#ifdef __cplusplus
+}  // TVM_EXTERN_C
+#endif
+#endif  // TVM_RUNTIME_C_RUNTIME_API_H_

include/dgl/runtime/device_api.h

+/*!
+ *  Copyright (c) 2016 by Contributors
+ * \file tvm/runtime/device_api.h
+ * \brief Abstract device memory management API
+ */
+#ifndef TVM_RUNTIME_DEVICE_API_H_
+#define TVM_RUNTIME_DEVICE_API_H_
+
+#include <string>
+#include "packed_func.h"
+#include "c_runtime_api.h"
+
+namespace tvm {
+namespace runtime {
+/*!
+ * \brief the query type into GetAttr
+ */
+enum DeviceAttrKind : int {
+  kExist = 0,
+  kMaxThreadsPerBlock = 1,
+  kWarpSize = 2,
+  kMaxSharedMemoryPerBlock = 3,
+  kComputeVersion = 4,
+  kDeviceName = 5,
+  kMaxClockRate = 6,
+  kMultiProcessorCount = 7,
+  kMaxThreadDimensions = 8
+};
+
+/*! \brief Number of bytes each allocation must align to */
+constexpr int kAllocAlignment = 64;
+
+/*! \brief Number of bytes each allocation must align to in temporary allocation */
+constexpr int kTempAllocaAlignment = 64;
+
+/*! \brief Maximum size that can be allocated on stack */
+constexpr int kMaxStackAlloca = 1024;
+
+/*!
+ * \brief TVM Runtime Device API, abstracts the device
+ *  specific interface for memory management.
+ */
+class DeviceAPI {
+ public:
+  /*! \brief virtual destructor */
+  virtual ~DeviceAPI() {}
+  /*!
+   * \brief Set the environment device id to ctx
+   * \param ctx The context to be set.
+   */
+  virtual void SetDevice(TVMContext ctx) = 0;
+  /*!
+   * \brief Get attribute of specified device.
+   * \param ctx The device context
+   * \param kind The result kind
+   * \param rv The return value.
+   * \sa DeviceAttrKind
+   */
+  virtual void GetAttr(TVMContext ctx, DeviceAttrKind kind, TVMRetValue* rv) = 0;
+  /*!
+   * \brief Allocate a data space on device.
+   * \param ctx The device context to perform operation.
+   * \param nbytes The number of bytes in memory.
+   * \param alignment The alignment of the memory.
+   * \param type_hint The type of elements. Only needed by certain backends such
+   * as OpenGL, as nbytes & alignment are sufficient for most backends.
+   * \return The allocated device pointer.
+   */
+  virtual void* AllocDataSpace(TVMContext ctx,
+                               size_t nbytes,
+                               size_t alignment,
+                               TVMType type_hint) = 0;
+  /*!
+   * \brief Free a data space on device.
+   * \param ctx The device context to perform operation.
+   * \param ptr The data space.
+   */
+  virtual void FreeDataSpace(TVMContext ctx, void* ptr) = 0;
+  /*!
+   * \brief copy data from one place to another
+   * \param from The source array.
+   * \param from_offset The byte offeset in the from.
+   * \param to The target array.
+   * \param to_offset The byte offset in the to.
+   * \param num_bytes The size of the memory in bytes
+   * \param ctx_from The source context
+   * \param ctx_to The target context
+   * \param type_hint The type of elements, only neded by certain backends.
+   *                  can be useful for cross device endian converison.
+   * \param stream Optional stream object.
+   */
+  virtual void CopyDataFromTo(const void* from,
+                              size_t from_offset,
+                              void* to,
+                              size_t to_offset,
+                              size_t num_bytes,
+                              TVMContext ctx_from,
+                              TVMContext ctx_to,
+                              TVMType type_hint,
+                              TVMStreamHandle stream) = 0;
+    /*!
+   * \brief Create a new stream of execution.
+   *
+   * \param ctx The context of allocation.
+   */
+  TVM_DLL virtual TVMStreamHandle CreateStream(TVMContext ctx);
+
+  /*!
+   * \brief Free a stream of execution
+   *
+   * \param ctx The context of the stream
+   * \param stream The pointer to be freed.
+   */
+  TVM_DLL virtual void FreeStream(TVMContext ctx, TVMStreamHandle stream);
+
+  /*!
+   * \brief Synchronize the stream
+   * \param ctx The context to perform operation.
+   * \param stream The stream to be sync.
+   */
+  virtual void StreamSync(TVMContext ctx, TVMStreamHandle stream) = 0;
+  /*!
+   * \brief Set the stream
+   * \param ctx The context to set stream.
+   * \param stream The stream to be set.
+   */
+  virtual void SetStream(TVMContext ctx, TVMStreamHandle stream) {}
+  /*!
+   * \brief Synchronize 2 streams of execution.
+   *
+   * An event is created in event_src stream that the second then
+   * stream waits on.  Neither event_src or event_dst need to be of
+   * the same device ID as the context, but they must be of the same
+   * device type.
+   *
+   * \param ctx The context of the streams.
+   * \param event_src The source stream to synchronize.
+   * \param event_dst The destination stream to synchronize.
+   */
+  TVM_DLL virtual void SyncStreamFromTo(TVMContext ctx,
+                                        TVMStreamHandle event_src,
+                                        TVMStreamHandle event_dst);
+  /*!
+   * \brief Allocate temporal workspace for backend execution.
+   *
+   *  \note We have the following assumption about backend temporal
+   *   workspace allocation, and backend will optimize for such assumption:
+   *
+   *  - Only a few allocation will happen, and space will be released after use.
+   *  - The release order is usually in reverse order of allocate (stack style).
+   *  - Repeative pattern of same allocations over different runs.
+   *  - Workspace should not overlap between different threads(i.e. be threadlocal)
+   *
+   * \param ctx The context of allocation.
+   * \param nbytes The size to be allocated.
+   * \param type_hint The type of elements. Only needed by certain backends such
+   * as OpenGL, as nbytes is sufficient for most backends.
+   */
+  TVM_DLL virtual void* AllocWorkspace(TVMContext ctx,
+                                       size_t nbytes,
+                                       TVMType type_hint = {});
+  /*!
+   * \brief Free temporal workspace in backend execution.
+   *
+   * \param ctx The context of allocation.
+   * \param ptr The pointer to be freed.
+   */
+  TVM_DLL virtual void FreeWorkspace(TVMContext ctx, void* ptr);
+
+  /*!
+   * \brief Get device API base don context.
+   * \param ctx The context
+   * \param allow_missing Whether allow missing
+   * \return The corresponding device API.
+   */
+  TVM_DLL static DeviceAPI* Get(TVMContext ctx, bool allow_missing = false);
+};
+
+/*! \brief The device type bigger than this is RPC device */
+constexpr int kRPCSessMask = 128;
+}  // namespace runtime
+}  // namespace tvm
+#endif  // TVM_RUNTIME_DEVICE_API_H_

include/dgl/runtime/module.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/module.h
+ * \brief Runtime container of the functions generated by TVM,
+ *  This is used to support dynamically link, load and save
+ *  functions from different convention under unified API.
+ */
+#ifndef TVM_RUNTIME_MODULE_H_
+#define TVM_RUNTIME_MODULE_H_
+
+#include <dmlc/io.h>
+#include <memory>
+#include <vector>
+#include <string>
+#include <unordered_map>
+#include "c_runtime_api.h"
+
+namespace tvm {
+namespace runtime {
+
+// The internal container of module.
+class ModuleNode;
+class PackedFunc;
+
+/*!
+ * \brief Module container of TVM.
+ */
+class Module {
+ public:
+  Module() {}
+  // constructor from container.
+  explicit Module(std::shared_ptr<ModuleNode> n)
+      : node_(n) {}
+  /*!
+   * \brief Get packed function from current module by name.
+   *
+   * \param name The name of the function.
+   * \param query_imports Whether also query dependency modules.
+   * \return The result function.
+   *  This function will return PackedFunc(nullptr) if function do not exist.
+   * \note Implemented in packed_func.cc
+   */
+  inline PackedFunc GetFunction(const std::string& name, bool query_imports = false);
+  /*! \return internal container */
+  inline ModuleNode* operator->();
+  /*! \return internal container */
+  inline const ModuleNode* operator->() const;
+  // The following functions requires link with runtime.
+  /*!
+   * \brief Import another module into this module.
+   * \param other The module to be imported.
+   *
+   * \note Cyclic dependency is not allowed among modules,
+   *  An error will be thrown when cyclic dependency is detected.
+   */
+  TVM_DLL void Import(Module other);
+  /*!
+   * \brief Load a module from file.
+   * \param file_name The name of the host function module.
+   * \param format The format of the file.
+   * \note This function won't load the import relationship.
+   *  Re-create import relationship by calling Import.
+   */
+  TVM_DLL static Module LoadFromFile(const std::string& file_name,
+                                     const std::string& format = "");
+
+ private:
+  std::shared_ptr<ModuleNode> node_;
+};
+
+/*!
+ * \brief Base node container of module.
+ *  Do not create this directly, instead use Module.
+ */
+class ModuleNode {
+ public:
+  /*! \brief virtual destructor */
+  virtual ~ModuleNode() {}
+  /*! \return The module type key */
+  virtual const char* type_key() const = 0;
+  /*!
+   * \brief Get a PackedFunc from module.
+   *
+   *  The PackedFunc may not be fully initialized,
+   *  there might still be first time running overhead when
+   *  executing the function on certain devices.
+   *  For benchmarking, use prepare to eliminate
+   *
+   * \param name the name of the function.
+   * \param sptr_to_self The shared_ptr that points to this module node.
+   *
+   * \return PackedFunc(nullptr) when it is not available.
+   *
+   * \note The function will always remain valid.
+   *   If the function need resource from the module(e.g. late linking),
+   *   it should capture sptr_to_self.
+   */
+  virtual PackedFunc GetFunction(
+      const std::string& name,
+      const std::shared_ptr<ModuleNode>& sptr_to_self) = 0;
+  /*!
+   * \brief Save the module to file.
+   * \param file_name The file to be saved to.
+   * \param format The format of the file.
+   */
+  virtual void SaveToFile(const std::string& file_name,
+                          const std::string& format);
+  /*!
+   * \brief Save the module to binary stream.
+   * \param stream The binary stream to save to.
+   * \note It is recommended to implement this for device modules,
+   *   but not necessarily host modules.
+   *   We can use this to do AOT loading of bundled device functions.
+   */
+  TVM_DLL virtual void SaveToBinary(dmlc::Stream* stream);
+  /*!
+   * \brief Get the source code of module, when available.
+   * \param format Format of the source code, can be empty by default.
+   * \return Possible source code when available.
+   */
+  TVM_DLL virtual std::string GetSource(const std::string& format = "");
+  /*!
+   * \brief Get a function from current environment
+   *  The environment includes all the imports as well as Global functions.
+   *
+   * \param name name of the function.
+   * \return The corresponding function.
+   */
+  TVM_DLL const PackedFunc* GetFuncFromEnv(const std::string& name);
+  /*! \return The module it imports from */
+  const std::vector<Module>& imports() const {
+    return imports_;
+  }
+
+ protected:
+  friend class Module;
+  /*! \brief The modules this module depend on */
+  std::vector<Module> imports_;
+
+ private:
+  /*! \brief Cache used by GetImport */
+  std::unordered_map<std::string,
+                     std::unique_ptr<PackedFunc> > import_cache_;
+};
+
+/*! \brief namespace for constant symbols */
+namespace symbol {
+/*! \brief Global variable to store module context. */
+constexpr const char* tvm_module_ctx = "__tvm_module_ctx";
+/*! \brief Global variable to store device module blob */
+constexpr const char* tvm_dev_mblob = "__tvm_dev_mblob";
+/*! \brief Number of bytes of device module blob. */
+constexpr const char* tvm_dev_mblob_nbytes = "__tvm_dev_mblob_nbytes";
+/*! \brief global function to set device */
+constexpr const char* tvm_set_device = "__tvm_set_device";
+/*! \brief Auxiliary counter to global barrier. */
+constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
+/*! \brief Prepare the global barrier before kernels that uses global barrier. */
+constexpr const char* tvm_prepare_global_barrier = "__tvm_prepare_global_barrier";
+/*! \brief Placeholder for the module's entry function. */
+constexpr const char* tvm_module_main = "__tvm_main__";
+}  // namespace symbol
+
+// implementations of inline functions.
+inline ModuleNode* Module::operator->() {
+  return node_.get();
+}
+
+inline const ModuleNode* Module::operator->() const {
+  return node_.get();
+}
+
+}  // namespace runtime
+}  // namespace tvm
+
+#include "packed_func.h"
+#endif  // TVM_RUNTIME_MODULE_H_

include/dgl/runtime/ndarray.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/ndarray.h
+ * \brief Abstract device memory management API
+ */
+#ifndef TVM_RUNTIME_NDARRAY_H_
+#define TVM_RUNTIME_NDARRAY_H_
+
+#include <atomic>
+#include <vector>
+#include <utility>
+#include "c_runtime_api.h"
+#include "serializer.h"
+
+namespace tvm {
+namespace runtime {
+/*!
+ * \brief Managed NDArray.
+ *  The array is backed by reference counted blocks.
+ */
+class NDArray {
+ public:
+  // internal container type
+  struct Container;
+  /*! \brief default constructor */
+  NDArray() {}
+  /*!
+   * \brief cosntruct a NDArray that refers to data
+   * \param data The data this NDArray refers to
+   */
+  explicit inline NDArray(Container* data);
+  /*!
+   * \brief copy constructor
+   * \param other The value to be copied
+   */
+  inline NDArray(const NDArray& other);  // NOLINT(*)
+  /*!
+   * \brief move constructor
+   * \param other The value to be moved
+   */
+  NDArray(NDArray&& other) // NOLINT(*)
+      : data_(other.data_) {
+    other.data_ = nullptr;
+  }
+  /*! \brief destructor */
+  ~NDArray() {
+    this->reset();
+  }
+  /*!
+   * \brief Swap this array with another NDArray
+   * \param other The other NDArray
+   */
+  void swap(NDArray& other) {  // NOLINT(*)
+    std::swap(data_, other.data_);
+  }
+  /*!
+   * \brief copy assignmemt
+   * \param other The value to be assigned.
+   * \return reference to self.
+   */
+  NDArray& operator=(const NDArray& other) {  // NOLINT(*)
+    // copy-and-swap idiom
+    NDArray(other).swap(*this);  // NOLINT(*)
+    return *this;
+  }
+  /*!
+   * \brief move assignmemt
+   * \param other The value to be assigned.
+   * \return reference to self.
+   */
+  NDArray& operator=(NDArray&& other) {  // NOLINT(*)
+    // copy-and-swap idiom
+    NDArray(std::move(other)).swap(*this); // NOLINT(*)
+    return *this;
+  }
+  /*! \return If NDArray is defined */
+  bool defined() const {
+    return data_ != nullptr;
+  }
+  /*! \return If both NDArray reference the same container */
+  bool same_as(const NDArray& other) const {
+    return data_ == other.data_;
+  }
+  /*! \brief reset the content of NDArray to be nullptr */
+  inline void reset();
+  /*!
+   * \return the reference counter
+   * \note this number is approximate in multi-threaded setting.
+   */
+  inline int use_count() const;
+  /*! \return Pointer to content of DLTensor */
+  inline const DLTensor* operator->() const;
+  /*!
+   * \brief Copy data content from another array.
+   * \param other The source array to be copied from.
+   * \note The copy may happen asynchrously if it involves a GPU context.
+   *       TVMSynchronize is necessary.
+   */
+  inline void CopyFrom(DLTensor* other);
+  inline void CopyFrom(const NDArray& other);
+  /*!
+   * \brief Copy data content into another array.
+   * \param other The source array to be copied from.
+   * \note The copy may happen asynchrously if it involves a GPU context.
+   *       TVMSynchronize is necessary.
+   */
+  inline void CopyTo(DLTensor* other) const;
+  inline void CopyTo(const NDArray& other) const;
+  /*!
+   * \brief Copy the data to another context.
+   * \param ctx The target context.
+   * \return The array under another context.
+   */
+  inline NDArray CopyTo(const DLContext& ctx) const;
+  /*!
+   * \brief Load NDArray from stream
+   * \param stream The input data stream
+   * \return Whether load is successful
+   */
+  inline bool Load(dmlc::Stream* stream);
+  /*!
+   * \brief Save NDArray to stream
+   * \param stream The output data stream
+   */
+  inline void Save(dmlc::Stream* stream) const;
+  /*!
+   * \brief Create a NDArray that shares the data memory with the current one.
+   * \param shape The shape of the new array.
+   * \param dtype The data type of the new array.
+   * \note The memory size of new array must be smaller than the current one.
+   */
+  TVM_DLL NDArray CreateView(
+      std::vector<int64_t> shape, DLDataType dtype);
+  /*!
+   * \brief Create a reference view of NDArray that
+   *  represents as DLManagedTensor.
+   * \return A DLManagedTensor
+   */
+  TVM_DLL DLManagedTensor* ToDLPack() const;
+  /*!
+   * \brief Create an empty NDArray.
+   * \param shape The shape of the new array.
+   * \param dtype The data type of the new array.
+   * \param ctx The context of the Array.
+   * \return The created Array
+   */
+  TVM_DLL static NDArray Empty(std::vector<int64_t> shape,
+                               DLDataType dtype,
+                               DLContext ctx);
+  /*!
+   * \brief Create a NDArray backed by a dlpack tensor.
+   *
+   * This allows us to create a NDArray using the memory
+   * allocated by an external deep learning framework
+   * that is DLPack compatible.
+   *
+   * The memory is retained until the NDArray went out of scope.
+   * \param tensor The DLPack tensor to copy from.
+   * \return The created NDArray view.
+   */
+  TVM_DLL static NDArray FromDLPack(DLManagedTensor* tensor);
+  /*!
+   * \brief Function to copy data from one array to another.
+   * \param from The source array.
+   * \param to The target array.
+   * \param stream The stream used in copy.
+   */
+  TVM_DLL static void CopyFromTo(
+      DLTensor* from, DLTensor* to, TVMStreamHandle stream = nullptr);
+
+  // internal namespace
+  struct Internal;
+ private:
+  /*! \brief Internal Data content */
+  Container* data_{nullptr};
+  // enable internal functions
+  friend struct Internal;
+  friend class TVMRetValue;
+  friend class TVMArgsSetter;
+};
+
+/*!
+ * \brief Save a DLTensor to stream
+ * \param strm The outpu stream
+ * \param tensor The tensor to be saved.
+ */
+inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor);
+
+/*!
+ * \brief Reference counted Container object used to back NDArray.
+ *
+ *  This object is DLTensor compatible:
+ *    the pointer to the NDArrayContainer can be directly
+ *    interpreted as a DLTensor*
+ *
+ * \note: do not use this function directly, use NDArray.
+ */
+struct NDArray::Container {
+ public:
+  // NOTE: the first part of this structure is the same as
+  // DLManagedTensor, note that, however, the deleter
+  // is only called when the reference counter goes to 0
+  /*!
+   * \brief The corresponding dl_tensor field.
+   * \note it is important that the first field is DLTensor
+   *  So that this data structure is DLTensor compatible.
+   *  The head ptr of this struct can be viewed as DLTensor*.
+   */
+  DLTensor dl_tensor;
+  /*!
+   * \brief addtional context, reserved for recycling
+   * \note We can attach additional content here
+   *  which the current container depend on
+   *  (e.g. reference to original memory when creating views).
+   */
+  void* manager_ctx{nullptr};
+  /*!
+   * \brief Customized deleter
+   *
+   * \note The customized deleter is helpful to enable
+   *  different ways of memory allocator that are not
+   *  currently defined by the system.
+   */
+  void (*deleter)(Container* self) = nullptr;
+  /*! \brief default constructor */
+  Container() {
+    dl_tensor.data = nullptr;
+    dl_tensor.ndim = 0;
+    dl_tensor.shape = nullptr;
+    dl_tensor.strides = nullptr;
+    dl_tensor.byte_offset = 0;
+  }
+  /*! \brief developer function, increases reference counter */
+  void IncRef() {
+    ref_counter_.fetch_add(1, std::memory_order_relaxed);
+  }
+  /*! \brief developer function, decrease reference counter */
+  void DecRef() {
+    if (ref_counter_.fetch_sub(1, std::memory_order_release) == 1) {
+      std::atomic_thread_fence(std::memory_order_acquire);
+      if (this->deleter != nullptr) {
+        (*this->deleter)(this);
+      }
+    }
+  }
+
+ private:
+  friend class NDArray;
+  friend class RPCWrappedFunc;
+  /*!
+   * \brief The shape container,
+   *  can be used used for shape data.
+   */
+  std::vector<int64_t> shape_;
+  /*! \brief The internal array object */
+  std::atomic<int> ref_counter_{0};
+};
+
+// implementations of inline functions
+// the usages of functions are documented in place.
+inline NDArray::NDArray(Container* data)
+  : data_(data) {
+  data_->IncRef();
+}
+
+inline NDArray::NDArray(const NDArray& other)
+  : data_(other.data_) {
+  data_->IncRef();
+}
+
+inline void NDArray::reset() {
+  if (data_ != nullptr) {
+    data_->DecRef();
+    data_ = nullptr;
+  }
+}
+
+inline void NDArray::CopyFrom(DLTensor* other) {
+  CHECK(data_ != nullptr);
+  CopyFromTo(other, &(data_->dl_tensor));
+}
+
+inline void NDArray::CopyFrom(const NDArray& other) {
+  CHECK(data_ != nullptr);
+  CHECK(other.data_ != nullptr);
+  CopyFromTo(&(other.data_->dl_tensor), &(data_->dl_tensor));
+}
+
+inline void NDArray::CopyTo(DLTensor* other) const {
+  CHECK(data_ != nullptr);
+  CopyFromTo(&(data_->dl_tensor), other);
+}
+
+inline void NDArray::CopyTo(const NDArray& other) const {
+  CHECK(data_ != nullptr);
+  CHECK(other.data_ != nullptr);
+  CopyFromTo(&(data_->dl_tensor), &(other.data_->dl_tensor));
+}
+
+inline NDArray NDArray::CopyTo(const DLContext& ctx) const {
+  CHECK(data_ != nullptr);
+  const DLTensor* dptr = operator->();
+  NDArray ret = Empty(std::vector<int64_t>(dptr->shape, dptr->shape + dptr->ndim),
+                      dptr->dtype, ctx);
+  this->CopyTo(ret);
+  return ret;
+}
+
+inline int NDArray::use_count() const {
+  if (data_ == nullptr) return 0;
+  return data_->ref_counter_.load(std::memory_order_relaxed);
+}
+
+inline const DLTensor* NDArray::operator->() const {
+  return &(data_->dl_tensor);
+}
+
+/*! \brief Magic number for NDArray file */
+constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F;
+
+inline bool SaveDLTensor(dmlc::Stream* strm,
+                         DLTensor* tensor) {
+  uint64_t header = kTVMNDArrayMagic, reserved = 0;
+  strm->Write(header);
+  strm->Write(reserved);
+  // Always save data as CPU context
+  //
+  // Parameters that get serialized should be in CPU by default.
+  // So even the array's context is GPU, it will be stored as CPU array.
+  // This is used to prevent case when another user loads the parameters
+  // back on machine that do not have GPU or related context.
+  //
+  // We can always do array.CopyTo(target_ctx) to get a corresponding
+  // array in the target context.
+  DLContext cpu_ctx;
+  cpu_ctx.device_type = kDLCPU;
+  cpu_ctx.device_id = 0;
+  strm->Write(cpu_ctx);
+  strm->Write(tensor->ndim);
+  strm->Write(tensor->dtype);
+  int ndim = tensor->ndim;
+  strm->WriteArray(tensor->shape, ndim);
+  int type_bytes = tensor->dtype.bits / 8;
+  int64_t num_elems = 1;
+  for (int i = 0; i < ndim; ++i) {
+    num_elems *= tensor->shape[i];
+  }
+  int64_t data_byte_size = type_bytes * num_elems;
+  strm->Write(data_byte_size);
+
+  if (DMLC_IO_NO_ENDIAN_SWAP &&
+      tensor->ctx.device_type == kDLCPU &&
+      tensor->strides == nullptr &&
+      tensor->byte_offset == 0) {
+    // quick path
+    strm->Write(tensor->data, data_byte_size);
+  } else {
+    std::vector<uint8_t> bytes(data_byte_size);
+    CHECK_EQ(TVMArrayCopyToBytes(
+        tensor, dmlc::BeginPtr(bytes), data_byte_size), 0)
+        << TVMGetLastError();
+    if (!DMLC_IO_NO_ENDIAN_SWAP) {
+      dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
+    }
+    strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
+  }
+  return true;
+}
+
+inline void NDArray::Save(dmlc::Stream* strm) const {
+  SaveDLTensor(strm, const_cast<DLTensor*>(operator->()));
+}
+
+inline bool NDArray::Load(dmlc::Stream* strm) {
+  uint64_t header, reserved;
+  CHECK(strm->Read(&header))
+      << "Invalid DLTensor file format";
+  CHECK(strm->Read(&reserved))
+      << "Invalid DLTensor file format";
+  CHECK(header == kTVMNDArrayMagic)
+      << "Invalid DLTensor file format";
+  DLContext ctx;
+  int ndim;
+  DLDataType dtype;
+  CHECK(strm->Read(&ctx))
+      << "Invalid DLTensor file format";
+  CHECK(strm->Read(&ndim))
+      << "Invalid DLTensor file format";
+  CHECK(strm->Read(&dtype))
+      << "Invalid DLTensor file format";
+  CHECK_EQ(ctx.device_type, kDLCPU)
+      << "Invalid DLTensor context: can only save as CPU tensor";
+  std::vector<int64_t> shape(ndim);
+  if (ndim != 0) {
+    CHECK(strm->ReadArray(&shape[0], ndim))
+        << "Invalid DLTensor file format";
+  }
+  NDArray ret = NDArray::Empty(shape, dtype, ctx);
+  int64_t num_elems = 1;
+  int elem_bytes = (ret->dtype.bits + 7) / 8;
+  for (int i = 0; i < ret->ndim; ++i) {
+    num_elems *= ret->shape[i];
+  }
+  int64_t data_byte_size;
+  CHECK(strm->Read(&data_byte_size))
+      << "Invalid DLTensor file format";
+  CHECK(data_byte_size == num_elems * elem_bytes)
+      << "Invalid DLTensor file format";
+  CHECK(strm->Read(ret->data, data_byte_size))
+      << "Invalid DLTensor file format";
+  if (!DMLC_IO_NO_ENDIAN_SWAP) {
+    dmlc::ByteSwap(ret->data, elem_bytes, num_elems);
+  }
+  *this = ret;
+  return true;
+}
+
+}  // namespace runtime
+}  // namespace tvm
+#endif  // TVM_RUNTIME_NDARRAY_H_

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 . import backend
+from . import data
+from . import function
+from . import generator
+from . import nn
 
-from .batch import batch, unbatch
+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 .batch import batch, unbatch
 from .context import cpu, gpu
-
+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
+        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