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/packed_func.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/packed_func.h
+ * \brief Type-erased function used across TVM API.
+ */
+#ifndef TVM_RUNTIME_PACKED_FUNC_H_
+#define TVM_RUNTIME_PACKED_FUNC_H_
+
+#include <dmlc/logging.h>
+#include <functional>
+#include <tuple>
+#include <vector>
+#include <string>
+#include <limits>
+#include <memory>
+#include <type_traits>
+#include "c_runtime_api.h"
+#include "module.h"
+#include "ndarray.h"
+
+namespace HalideIR {
+// Forward declare type for extensions
+// The header works fine without depending on this.
+struct Type;
+struct Expr;
+}
+
+// Whether use TVM runtime in header only mode.
+#ifndef TVM_RUNTIME_HEADER_ONLY
+#define TVM_RUNTIME_HEADER_ONLY 0
+#endif
+
+namespace tvm {
+// Forward declare NodeRef and Node for extensions.
+// This header works fine without depend on NodeRef
+// as long as it is not used.
+class Node;
+class NodeRef;
+
+namespace runtime {
+// forward declarations
+class TVMArgs;
+class TVMArgValue;
+class TVMRetValue;
+class TVMArgsSetter;
+
+/*!
+ * \brief Packed function is a type-erased function.
+ *  The arguments are passed by packed format.
+ *
+ *  This is an useful unified interface to call generated functions,
+ *  It is the unified function function type of TVM.
+ *  It corresponds to TVMFunctionHandle in C runtime API.
+ */
+class PackedFunc {
+ public:
+  /*!
+   * \brief The internal std::function
+   * \param args The arguments to the function.
+   * \param rv The return value.
+   *
+   * \code
+   *   // Example code on how to implemented FType
+   *   void MyPackedFunc(TVMArgs args, TVMRetValue* rv) {
+   *     // automatically convert arguments to desired type.
+   *     int a0 = args[0];
+   *     float a1 = args[1];
+   *     ...
+   *     // automatically assign values to rv
+   *     std::string my_return_value = "x";
+   *     *rv = my_return_value;
+   *   }
+   * \endcode
+   */
+  using FType = std::function<void (TVMArgs args, TVMRetValue* rv)>;
+  /*! \brief default constructor */
+  PackedFunc() {}
+  /*!
+   * \brief constructing a packed function from a std::function.
+   * \param body the internal container of packed function.
+   */
+  explicit PackedFunc(FType body) : body_(body) {}
+  /*!
+   * \brief Call packed function by directly passing in unpacked format.
+   * \param args Arguments to be passed.
+   * \tparam Args arguments to be passed.
+   *
+   * \code
+   *   // Example code on how to call packed function
+   *   void CallPacked(PackedFunc f) {
+   *     // call like normal functions by pass in arguments
+   *     // return value is automatically converted back
+   *     int rvalue = f(1, 2.0);
+   *   }
+   * \endcode
+   */
+  template<typename... Args>
+  inline TVMRetValue operator()(Args&& ...args) const;
+  /*!
+   * \brief Call the function in packed format.
+   * \param args The arguments
+   * \param rv The return value.
+   */
+  inline void CallPacked(TVMArgs args, TVMRetValue* rv) const;
+  /*! \return the internal body function */
+  inline FType body() const;
+  /*! \return Whether the packed function is nullptr */
+  bool operator==(std::nullptr_t null) const {
+    return body_ == nullptr;
+  }
+  /*! \return Whether the packed function is not nullptr */
+  bool operator!=(std::nullptr_t null) const {
+    return body_ != nullptr;
+  }
+
+ private:
+  /*! \brief internal container of packed function */
+  FType body_;
+};
+
+/*!
+ * \brief Please refer to \ref TypedPackedFuncAnchor "TypedPackedFunc<R(Args..)>"
+ */
+template<typename FType>
+class TypedPackedFunc;
+
+/*!
+ * \anchor TypedPackedFuncAnchor
+ * \brief A PackedFunc wrapper to provide typed function signature.
+ * It is backed by a PackedFunc internally.
+ *
+ * TypedPackedFunc enables compile time type checking.
+ * TypedPackedFunc works with the runtime system:
+ * - It can be passed as an argument of PackedFunc.
+ * - It can be assigned to TVMRetValue.
+ * - It can be directly converted to a type-erased PackedFunc.
+ *
+ * Developers should prefer TypedPackedFunc over PackedFunc in C++ code
+ * as it enables compile time checking.
+ * We can construct a TypedPackedFunc from a lambda function
+ * with the same signature.
+ *
+ * \code
+ *  // user defined lambda function.
+ *  auto addone = [](int x)->int {
+ *    return x + 1;
+ *  };
+ *  // We can directly convert
+ *  // lambda function to TypedPackedFunc
+ *  TypedPackedFunc<int(int)> ftyped(addone);
+ *  // invoke the function.
+ *  int y = ftyped(1);
+ *  // Can be directly converted to PackedFunc
+ *  PackedFunc packed = ftype;
+ * \endcode
+ * \tparam R The return value of the function.
+ * \tparam Args The argument signature of the function.
+ */
+template<typename R, typename ...Args>
+class TypedPackedFunc<R(Args...)> {
+ public:
+  /*! \brief short hand for this function type */
+  using TSelf = TypedPackedFunc<R(Args...)>;
+  /*! \brief default constructor */
+  TypedPackedFunc() {}
+  /*!
+   * \brief construct by wrap a PackedFunc
+   *
+   * Example usage:
+   * \code
+   * PackedFunc packed([](TVMArgs args, TVMRetValue *rv) {
+   *   int x = args[0];
+   *   *rv = x + 1;
+   *  });
+   * // construct from packed function
+   * TypedPackedFunc<int(int)> ftyped(packed);
+   * // call the typed version.
+   * CHECK_EQ(ftyped(1), 2);
+   * \endcode
+   *
+   * \param packed The packed function
+   */
+  inline explicit TypedPackedFunc(PackedFunc packed);
+  /*!
+   * \brief construct from a lambda function with the same signature.
+   *
+   * Example usage:
+   * \code
+   * auto typed_lambda = [](int x)->int { return x + 1; }
+   * // construct from packed function
+   * TypedPackedFunc<int(int)> ftyped(typed_lambda);
+   * // call the typed version.
+   * CHECK_EQ(ftyped(1), 2);
+   * \endcode
+   *
+   * \param typed_lambda typed lambda function.
+   * \tparam FLambda the type of the lambda function.
+   */
+  template<typename FLambda,
+           typename = typename std::enable_if<
+             std::is_convertible<FLambda,
+                                 std::function<R(Args...)>
+                                 >::value>::type>
+  explicit TypedPackedFunc(const FLambda& typed_lambda) {
+    this->AssignTypedLambda(typed_lambda);
+  }
+  /*!
+   * \brief copy assignment operator from typed lambda
+   *
+   * Example usage:
+   * \code
+   * // construct from packed function
+   * TypedPackedFunc<int(int)> ftyped;
+   * ftyped = [](int x) { return x + 1; }
+   * // call the typed version.
+   * CHECK_EQ(ftyped(1), 2);
+   * \endcode
+   *
+   * \param typed_lambda typed lambda function.
+   * \tparam FLambda the type of the lambda function.
+   * \returns reference to self.
+   */
+  template<typename FLambda,
+           typename = typename std::enable_if<
+             std::is_convertible<FLambda,
+                                 std::function<R(Args...)>
+                                 >::value>::type>
+  TSelf& operator=(FLambda typed_lambda) {  // NOLINT(*)
+    this->AssignTypedLambda(typed_lambda);
+    return *this;
+  }
+  /*!
+   * \brief copy assignment operator from PackedFunc.
+   * \param packed The packed function.
+   * \returns reference to self.
+   */
+  TSelf& operator=(PackedFunc packed) {
+    packed_ = packed;
+    return *this;
+  }
+  /*!
+   * \brief Invoke the operator.
+   * \param args The arguments
+   * \returns The return value.
+   */
+  inline R operator()(Args ...args) const;
+  /*!
+   * \brief convert to PackedFunc
+   * \return the internal PackedFunc
+   */
+  operator PackedFunc() const {
+    return packed();
+  }
+  /*!
+   * \return reference the internal PackedFunc
+   */
+  const PackedFunc& packed() const {
+    return packed_;
+  }
+
+ private:
+  friend class TVMRetValue;
+  /*! \brief The internal packed function */
+  PackedFunc packed_;
+  /*!
+   * \brief Assign the packed field using a typed lambda function.
+   *
+   * \param flambda The lambda function.
+   * \tparam FLambda The lambda function type.
+   * \note We capture the lambda when possible for maximum efficiency.
+   */
+  template<typename FLambda>
+  inline void AssignTypedLambda(FLambda flambda);
+};
+
+/*! \brief Arguments into TVM functions. */
+class TVMArgs {
+ public:
+  const TVMValue* values;
+  const int* type_codes;
+  int num_args;
+  /*!
+   * \brief constructor
+   * \param values The argument values
+   * \param type_codes The argument type codes
+   * \param num_args number of arguments.
+   */
+  TVMArgs(const TVMValue* values,
+          const int* type_codes,
+          int num_args)
+      : values(values),
+        type_codes(type_codes),
+        num_args(num_args) { }
+  /*! \return size of the arguments */
+  inline int size() const;
+  /*!
+   * \brief Get i-th argument
+   * \param i the index.
+   * \return the ith argument.
+   */
+  inline TVMArgValue operator[](int i) const;
+};
+
+/*!
+ * \brief Convert type code to its name
+ * \param type_code The type code .
+ * \return The name of type code.
+ */
+inline const char* TypeCode2Str(int type_code);
+
+/*!
+ * \brief convert a string to TVM type.
+ * \param s The string to be converted.
+ * \return The corresponding tvm type.
+ */
+inline TVMType String2TVMType(std::string s);
+
+/*!
+ * \brief convert a TVM type to string.
+ * \param t The type to be converted.
+ * \return The corresponding tvm type in string.
+ */
+inline std::string TVMType2String(TVMType t);
+
+// macro to check type code.
+#define TVM_CHECK_TYPE_CODE(CODE, T)                           \
+  CHECK_EQ(CODE, T) << " expected "                            \
+  << TypeCode2Str(T) << " but get " << TypeCode2Str(CODE)      \
+
+/*!
+ * \brief Type traits to mark if a class is tvm extension type.
+ *
+ * To enable extension type in C++ must be register () ed via marco.
+ * TVM_REGISTER_EXT_TYPE(TypeName) after defining this with this traits.
+ *
+ * Extension class can be passed and returned via PackedFunc in all tvm runtime.
+ * Internally extension class is stored as T*.
+ *
+ * \tparam T the typename
+ */
+template<typename T>
+struct extension_class_info {
+  static const int code = 0;
+};
+
+/*!
+ * \brief Runtime function table about extension type.
+ */
+class ExtTypeVTable {
+ public:
+  /*! \brief function to be called to delete a handle */
+  void (*destroy)(void* handle);
+  /*! \brief function to be called when clone a handle */
+  void* (*clone)(void* handle);
+  /*!
+   * \brief Register type
+   * \tparam T The type to be register.
+   * \return The registered vtable.
+   */
+  template <typename T>
+  static inline ExtTypeVTable* Register_();
+  /*!
+   * \brief Get a vtable based on type code.
+   * \param type_code The type code
+   * \return The registered vtable.
+   */
+  TVM_DLL static ExtTypeVTable* Get(int type_code);
+
+ private:
+  // Internal registration function.
+  TVM_DLL static ExtTypeVTable* RegisterInternal(int type_code, const ExtTypeVTable& vt);
+};
+
+/*!
+ * \brief Internal base class to
+ *  handle conversion to POD values.
+ */
+class TVMPODValue_ {
+ public:
+  operator double() const {
+    // Allow automatic conversion from int to float
+    // This avoids errors when user pass in int from
+    // the frontend while the API expects a float.
+    if (type_code_ == kDLInt) {
+      return static_cast<double>(value_.v_int64);
+    }
+    TVM_CHECK_TYPE_CODE(type_code_, kDLFloat);
+    return value_.v_float64;
+  }
+  operator int64_t() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
+    return value_.v_int64;
+  }
+  operator uint64_t() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
+    return value_.v_int64;
+  }
+  operator int() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
+    CHECK_LE(value_.v_int64,
+             std::numeric_limits<int>::max());
+    return static_cast<int>(value_.v_int64);
+  }
+  operator bool() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
+    return value_.v_int64 != 0;
+  }
+  operator void*() const {
+    if (type_code_ == kNull) return nullptr;
+    if (type_code_ == kArrayHandle) return value_.v_handle;
+    TVM_CHECK_TYPE_CODE(type_code_, kHandle);
+    return value_.v_handle;
+  }
+  operator DLTensor*() const {
+    if (type_code_ == kArrayHandle ||
+        type_code_ == kNDArrayContainer) {
+      return static_cast<DLTensor*>(value_.v_handle);
+    } else {
+      if (type_code_ == kNull) return nullptr;
+      LOG(FATAL) << "Expected "
+                 << "DLTensor* or NDArray but get "
+                 << TypeCode2Str(type_code_);
+      return nullptr;
+    }
+  }
+  operator NDArray() const {
+    if (type_code_ == kNull) return NDArray();
+    TVM_CHECK_TYPE_CODE(type_code_, kNDArrayContainer);
+    return NDArray(static_cast<NDArray::Container*>(value_.v_handle));
+  }
+  operator TVMContext() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kTVMContext);
+    return value_.v_ctx;
+  }
+  template<typename TExtension>
+  const TExtension& AsExtension() const {
+    CHECK_LT(type_code_, kExtEnd);
+    return static_cast<TExtension*>(value_.v_handle)[0];
+  }
+  int type_code() const {
+    return type_code_;
+  }
+  /*!
+   * \brief return handle as specific pointer type.
+   * \tparam T the data type.
+   * \return The pointer type.
+   */
+  template<typename T>
+  T* ptr() const {
+    return static_cast<T*>(value_.v_handle);
+  }
+
+ protected:
+  friend class TVMArgsSetter;
+  friend class TVMRetValue;
+  TVMPODValue_() : type_code_(kNull) {}
+  TVMPODValue_(TVMValue value, int type_code)
+      : value_(value), type_code_(type_code) {}
+
+  /*! \brief The value */
+  TVMValue value_;
+  /*! \brief the type code */
+  int type_code_;
+};
+
+/*!
+ * \brief A single argument value to PackedFunc.
+ *  Containing both type_code and TVMValue
+ *
+ *  Provides utilities to do type cast into other types.
+ */
+class TVMArgValue : public TVMPODValue_ {
+ public:
+  /*! \brief default constructor */
+  TVMArgValue() {}
+  /*!
+   * \brief constructor
+   * \param value of the function
+   * \param type_code The type code.
+   */
+  TVMArgValue(TVMValue value, int type_code)
+      : TVMPODValue_(value, type_code) {
+  }
+  // reuse converter from parent
+  using TVMPODValue_::operator double;
+  using TVMPODValue_::operator int64_t;
+  using TVMPODValue_::operator uint64_t;
+  using TVMPODValue_::operator int;
+  using TVMPODValue_::operator bool;
+  using TVMPODValue_::operator void*;
+  using TVMPODValue_::operator DLTensor*;
+  using TVMPODValue_::operator NDArray;
+  using TVMPODValue_::operator TVMContext;
+
+  // conversion operator.
+  operator std::string() const {
+    if (type_code_ == kTVMType) {
+      return TVMType2String(operator TVMType());
+    } else if (type_code_ == kBytes) {
+      TVMByteArray* arr = static_cast<TVMByteArray*>(value_.v_handle);
+      return std::string(arr->data, arr->size);
+    } else {
+      TVM_CHECK_TYPE_CODE(type_code_, kStr);
+      return std::string(value_.v_str);
+    }
+  }
+  operator TVMType() const {
+    if (type_code_ == kStr) {
+      return String2TVMType(operator std::string());
+    }
+    TVM_CHECK_TYPE_CODE(type_code_, kTVMType);
+    return value_.v_type;
+  }
+  operator PackedFunc() const {
+    if (type_code_ == kNull) return PackedFunc();
+    TVM_CHECK_TYPE_CODE(type_code_, kFuncHandle);
+    return *ptr<PackedFunc>();
+  }
+  template<typename FType>
+  operator TypedPackedFunc<FType>() const {
+    return TypedPackedFunc<FType>(operator PackedFunc());
+  }
+  operator Module() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kModuleHandle);
+    return *ptr<Module>();
+  }
+  const TVMValue& value() const {
+    return value_;
+  }
+  // Deferred extension handler.
+  template<typename TNodeRef>
+  inline TNodeRef AsNodeRef() const;
+  template<typename T,
+           typename = typename std::enable_if<
+             std::is_class<T>::value>::type>
+  inline operator T() const;
+  template<typename TNodeRef,
+           typename = typename std::enable_if<
+             std::is_class<TNodeRef>::value>::type>
+  inline bool IsNodeType() const;
+  inline operator HalideIR::Type() const;
+  inline operator HalideIR::Expr() const;
+  // get internal node ptr, if it is node
+  inline std::shared_ptr<Node>& node_sptr();
+};
+
+/*!
+ * \brief Return Value container,
+ *  Unlike TVMArgValue, which only holds reference and do not delete
+ *  the underlying container during destruction.
+ *
+ *  TVMRetValue holds value and will manage the underlying containers
+ *  when it stores a complicated data type.
+ */
+class TVMRetValue : public TVMPODValue_ {
+ public:
+  /*! \brief default constructor */
+  TVMRetValue() {}
+  /*!
+   * \brief move constructor from anoter return value.
+   * \param other The other return value.
+   */
+  TVMRetValue(TVMRetValue&& other)
+      : TVMPODValue_(other.value_, other.type_code_) {
+    other.value_.v_handle = nullptr;
+    other.type_code_ = kNull;
+  }
+  /*! \brief destructor */
+  ~TVMRetValue() {
+    this->Clear();
+  }
+  // reuse converter from parent
+  using TVMPODValue_::operator double;
+  using TVMPODValue_::operator int64_t;
+  using TVMPODValue_::operator uint64_t;
+  using TVMPODValue_::operator int;
+  using TVMPODValue_::operator bool;
+  using TVMPODValue_::operator void*;
+  using TVMPODValue_::operator DLTensor*;
+  using TVMPODValue_::operator TVMContext;
+  using TVMPODValue_::operator NDArray;
+  // Disable copy and assign from another value, but allow move.
+  TVMRetValue(const TVMRetValue& other) {
+    this->Assign(other);
+  }
+  // conversion operators
+  operator std::string() const {
+    if (type_code_ == kTVMType) {
+      return TVMType2String(operator TVMType());
+    } else if (type_code_ == kBytes) {
+      return *ptr<std::string>();
+    }
+    TVM_CHECK_TYPE_CODE(type_code_, kStr);
+    return *ptr<std::string>();
+  }
+  operator TVMType() const {
+    if (type_code_ == kStr) {
+      return String2TVMType(operator std::string());
+    }
+    TVM_CHECK_TYPE_CODE(type_code_, kTVMType);
+    return value_.v_type;
+  }
+  operator PackedFunc() const {
+    if (type_code_ == kNull) return PackedFunc();
+    TVM_CHECK_TYPE_CODE(type_code_, kFuncHandle);
+    return *ptr<PackedFunc>();
+  }
+  template<typename FType>
+  operator TypedPackedFunc<FType>() const {
+    return TypedPackedFunc<FType>(operator PackedFunc());
+  }
+  operator Module() const {
+    TVM_CHECK_TYPE_CODE(type_code_, kModuleHandle);
+    return *ptr<Module>();
+  }
+  // Assign operators
+  TVMRetValue& operator=(TVMRetValue&& other) {
+    this->Clear();
+    value_ = other.value_;
+    type_code_ = other.type_code_;
+    other.type_code_ = kNull;
+    return *this;
+  }
+  TVMRetValue& operator=(double value) {
+    this->SwitchToPOD(kDLFloat);
+    value_.v_float64 = value;
+    return *this;
+  }
+  TVMRetValue& operator=(std::nullptr_t value) {
+    this->SwitchToPOD(kNull);
+    value_.v_handle = value;
+    return *this;
+  }
+  TVMRetValue& operator=(void* value) {
+    this->SwitchToPOD(kHandle);
+    value_.v_handle = value;
+    return *this;
+  }
+  TVMRetValue& operator=(int64_t value) {
+    this->SwitchToPOD(kDLInt);
+    value_.v_int64 = value;
+    return *this;
+  }
+  TVMRetValue& operator=(int value) {
+    this->SwitchToPOD(kDLInt);
+    value_.v_int64 = value;
+    return *this;
+  }
+  TVMRetValue& operator=(TVMType t) {
+    this->SwitchToPOD(kTVMType);
+    value_.v_type = t;
+    return *this;
+  }
+  TVMRetValue& operator=(bool value) {
+    this->SwitchToPOD(kDLInt);
+    value_.v_int64 = value;
+    return *this;
+  }
+  TVMRetValue& operator=(std::string value) {
+    this->SwitchToClass(kStr, value);
+    return *this;
+  }
+  TVMRetValue& operator=(TVMByteArray value) {
+    this->SwitchToClass(kBytes, std::string(value.data, value.size));
+    return *this;
+  }
+  TVMRetValue& operator=(NDArray other) {
+    this->Clear();
+    type_code_ = kNDArrayContainer;
+    value_.v_handle = other.data_;
+    other.data_ = nullptr;
+    return *this;
+  }
+  TVMRetValue& operator=(PackedFunc f) {
+    this->SwitchToClass(kFuncHandle, f);
+    return *this;
+  }
+  template<typename FType>
+  TVMRetValue& operator=(const TypedPackedFunc<FType>& f) {
+    return operator=(f.packed());
+  }
+  TVMRetValue& operator=(Module m) {
+    this->SwitchToClass(kModuleHandle, m);
+    return *this;
+  }
+  TVMRetValue& operator=(const TVMRetValue& other) {  // NOLINT(*0
+    this->Assign(other);
+    return *this;
+  }
+  TVMRetValue& operator=(const TVMArgValue& other) {
+    this->Assign(other);
+    return *this;
+  }
+  template<typename T,
+           typename = typename std::enable_if<
+             extension_class_info<T>::code != 0>::type>
+  TVMRetValue& operator=(const T& other) {
+    this->SwitchToClass<T>(
+        extension_class_info<T>::code, other);
+    return *this;
+  }
+  /*!
+   * \brief Move the value back to front-end via C API.
+   *  This marks the current container as null.
+   *  The managed resources is moved to front-end and
+   *  the front end should take charge in managing them.
+   *
+   * \param ret_value The return value.
+   * \param ret_type_code The return type code.
+   */
+  void MoveToCHost(TVMValue* ret_value,
+                   int* ret_type_code) {
+    // cannot move str; need specially handle.
+    CHECK(type_code_ != kStr && type_code_ != kBytes);
+    *ret_value = value_;
+    *ret_type_code = type_code_;
+    type_code_ = kNull;
+  }
+  /*! \return The value field, if the data is POD */
+  const TVMValue& value() const {
+    CHECK(type_code_ != kNodeHandle &&
+          type_code_ != kFuncHandle &&
+          type_code_ != kModuleHandle &&
+          type_code_ != kStr) << "TVMRetValue.value can only be used for POD data";
+    return value_;
+  }
+  // NodeRef related extenstions: in tvm/packed_func_ext.h
+  template<typename T,
+           typename = typename std::enable_if<
+             std::is_class<T>::value>::type>
+  inline operator T() const;
+  template<typename TNodeRef>
+  inline TNodeRef AsNodeRef() const;
+  inline TVMRetValue& operator=(const NodeRef& other);
+  inline TVMRetValue& operator=(const std::shared_ptr<Node>& other);
+  // type related
+  inline operator HalideIR::Type() const;
+  inline TVMRetValue& operator=(const HalideIR::Type& other);
+
+ private:
+  template<typename T>
+  void Assign(const T& other) {
+    switch (other.type_code()) {
+      case kStr: {
+        SwitchToClass<std::string>(kStr, other);
+        break;
+      }
+      case kBytes: {
+        SwitchToClass<std::string>(kBytes, other);
+        break;
+      }
+      case kFuncHandle: {
+        SwitchToClass<PackedFunc>(kFuncHandle, other);
+        break;
+      }
+      case kModuleHandle: {
+        SwitchToClass<Module>(kModuleHandle, other);
+        break;
+      }
+      case kNDArrayContainer: {
+        *this = other.operator NDArray();
+        break;
+      }
+      case kNodeHandle: {
+        SwitchToClass<std::shared_ptr<Node> >(
+            kNodeHandle, *other.template ptr<std::shared_ptr<Node> >());
+        break;
+      }
+      default: {
+        if (other.type_code() < kExtBegin) {
+          SwitchToPOD(other.type_code());
+          value_ = other.value_;
+        } else {
+#if TVM_RUNTIME_HEADER_ONLY
+          LOG(FATAL) << "Header only mode do not support ext type";
+#else
+          this->Clear();
+          type_code_ = other.type_code();
+          value_.v_handle =
+              (*(ExtTypeVTable::Get(other.type_code())->clone))(
+                  other.value().v_handle);
+#endif
+        }
+        break;
+      }
+    }
+  }
+  // get the internal container.
+  void SwitchToPOD(int type_code) {
+    if (type_code_ != type_code) {
+      this->Clear();
+      type_code_ = type_code;
+    }
+  }
+  template<typename T>
+  void SwitchToClass(int type_code, T v) {
+    if (type_code_ != type_code) {
+      this->Clear();
+      type_code_ = type_code;
+      value_.v_handle = new T(v);
+    } else {
+      *static_cast<T*>(value_.v_handle) = v;
+    }
+  }
+  void Clear() {
+    if (type_code_ == kNull) return;
+    switch (type_code_) {
+      case kStr: delete ptr<std::string>(); break;
+      case kFuncHandle: delete ptr<PackedFunc>(); break;
+      case kModuleHandle: delete ptr<Module>(); break;
+      case kNodeHandle: delete ptr<std::shared_ptr<Node> >(); break;
+      case kNDArrayContainer: {
+        static_cast<NDArray::Container*>(value_.v_handle)->DecRef();
+        break;
+      }
+    }
+    if (type_code_ > kExtBegin) {
+#if TVM_RUNTIME_HEADER_ONLY
+          LOG(FATAL) << "Header only mode do not support ext type";
+#else
+      (*(ExtTypeVTable::Get(type_code_)->destroy))(value_.v_handle);
+#endif
+    }
+    type_code_ = kNull;
+  }
+};
+
+// implementation details
+inline const char* TypeCode2Str(int type_code) {
+  switch (type_code) {
+    case kDLInt: return "int";
+    case kDLUInt: return "uint";
+    case kDLFloat: return "float";
+    case kStr: return "str";
+    case kBytes: return "bytes";
+    case kHandle: return "handle";
+    case kNull: return "NULL";
+    case kNodeHandle: return "NodeHandle";
+    case kArrayHandle: return "ArrayHandle";
+    case kTVMType: return "TVMType";
+    case kTVMContext: return "TVMContext";
+    case kFuncHandle: return "FunctionHandle";
+    case kModuleHandle: return "ModuleHandle";
+    case kNDArrayContainer: return "NDArrayContainer";
+    default: LOG(FATAL) << "unknown type_code="
+                        << static_cast<int>(type_code); return "";
+  }
+}
+
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
+inline std::ostream& operator<<(std::ostream& os, TVMType t) {  // NOLINT(*)
+  os << TypeCode2Str(t.code);
+  if (t.code == kHandle) return os;
+  os << static_cast<int>(t.bits);
+  if (t.lanes != 1) {
+    os << 'x' << static_cast<int>(t.lanes);
+  }
+  return os;
+}
+#endif
+
+inline std::string TVMType2String(TVMType t) {
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
+  std::ostringstream os;
+  os << t;
+  return os.str();
+#else
+  std::string repr = "";
+  repr += TypeCode2Str(t.code);
+  if (t.code == kHandle) return repr;
+  repr += std::to_string(static_cast<int>(t.bits));
+  if (t.lanes != 1) {
+    repr += "x" + std::to_string(static_cast<int>(t.lanes));
+  }
+  return repr;
+#endif
+}
+
+inline TVMType String2TVMType(std::string s) {
+  TVMType t;
+  t.bits = 32; t.lanes = 1;
+  const char* scan;
+  if (s.substr(0, 3) == "int") {
+    t.code = kDLInt;  scan = s.c_str() + 3;
+  } else if (s.substr(0, 4) == "uint") {
+    t.code = kDLUInt; scan = s.c_str() + 4;
+  } else if (s.substr(0, 5) == "float") {
+    t.code = kDLFloat; scan = s.c_str() + 5;
+  } else if (s.substr(0, 6) == "handle") {
+    t.code = kHandle;
+    t.bits = 64;  // handle uses 64 bit by default.
+    scan = s.c_str() + 6;
+  } else {
+    scan = s.c_str();
+    LOG(FATAL) << "unknown type " << s;
+  }
+  char* xdelim;  // emulate sscanf("%ux%u", bits, lanes)
+  uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
+  if (bits != 0) t.bits = bits;
+  if (*xdelim == 'x') {
+    t.lanes = static_cast<uint16_t>(strtoul(xdelim + 1, nullptr, 10));
+  }
+  return t;
+}
+
+inline TVMArgValue TVMArgs::operator[](int i) const {
+  CHECK_LT(i, num_args)
+      << "not enough argument passed, "
+      << num_args << " passed"
+      << " but request arg[" << i << "].";
+  return TVMArgValue(values[i], type_codes[i]);
+}
+
+inline int TVMArgs::size() const {
+  return num_args;
+}
+
+inline void PackedFunc::CallPacked(TVMArgs args, TVMRetValue* rv) const {
+  body_(args, rv);
+}
+
+inline PackedFunc::FType PackedFunc::body() const {
+  return body_;
+}
+
+
+
+// internal namespace
+namespace detail {
+
+template<bool stop, std::size_t I, typename F>
+struct for_each_dispatcher {
+  template<typename T, typename ...Args>
+  static void run(const F& f, T&& value, Args&&... args) {  // NOLINT(*)
+    f(I, std::forward<T>(value));
+    for_each_dispatcher<sizeof...(Args) == 0, (I+1), F>
+        ::run(f, std::forward<Args>(args)...);
+  }
+};
+
+template<std::size_t I, typename F>
+struct for_each_dispatcher<true, I, F>  {
+  static void run(const F& f) {}  // NOLINT(*)
+};
+
+template<typename F, typename ...Args>
+inline void for_each(const F& f, Args&&... args) {  // NOLINT(*)
+  for_each_dispatcher<sizeof...(Args) == 0, 0, F>
+      ::run(f, std::forward<Args>(args)...);
+}
+}  // namespace detail
+
+/* \brief argument settter to PackedFunc */
+class TVMArgsSetter {
+ public:
+  TVMArgsSetter(TVMValue* values, int* type_codes)
+      : values_(values), type_codes_(type_codes) {}
+  // setters for POD types
+  template<typename T,
+           typename = typename std::enable_if<
+             std::is_integral<T>::value>::type>
+  void operator()(size_t i, T value) const {
+    values_[i].v_int64 = static_cast<int64_t>(value);
+    type_codes_[i] = kDLInt;
+  }
+  void operator()(size_t i, uint64_t value) const {
+    values_[i].v_int64 = static_cast<int64_t>(value);
+    CHECK_LE(value,
+             static_cast<uint64_t>(std::numeric_limits<int64_t>::max()));
+    type_codes_[i] = kDLInt;
+  }
+  void operator()(size_t i, double value) const {
+    values_[i].v_float64 = value;
+    type_codes_[i] = kDLFloat;
+  }
+  void operator()(size_t i, std::nullptr_t value) const {
+    values_[i].v_handle = value;
+    type_codes_[i] = kNull;
+  }
+  void operator()(size_t i, const TVMArgValue& value) const {
+    values_[i] = value.value_;
+    type_codes_[i] = value.type_code_;
+  }
+  void operator()(size_t i, void* value) const {
+    values_[i].v_handle = value;
+    type_codes_[i] = kHandle;
+  }
+  void operator()(size_t i, DLTensor* value) const {
+    values_[i].v_handle = value;
+    type_codes_[i] = kArrayHandle;
+  }
+  void operator()(size_t i, TVMContext value) const {
+    values_[i].v_ctx = value;
+    type_codes_[i] = kTVMContext;
+  }
+  void operator()(size_t i, TVMType value) const {
+    values_[i].v_type = value;
+    type_codes_[i] = kTVMType;
+  }
+  void operator()(size_t i, const char* value) const {
+    values_[i].v_str = value;
+    type_codes_[i] = kStr;
+  }
+  // setters for container type
+  // They must be reference(instead of const ref)
+  // to make sure they are alive in the tuple(instead of getting converted)
+  void operator()(size_t i, const std::string& value) const {  // NOLINT(*)
+    values_[i].v_str = value.c_str();
+    type_codes_[i] = kStr;
+  }
+  void operator()(size_t i, const TVMByteArray& value) const {  // NOLINT(*)
+    values_[i].v_handle = const_cast<TVMByteArray*>(&value);
+    type_codes_[i] = kBytes;
+  }
+  void operator()(size_t i, const PackedFunc& value) const {  // NOLINT(*)
+    values_[i].v_handle = const_cast<PackedFunc*>(&value);
+    type_codes_[i] = kFuncHandle;
+  }
+  template<typename FType>
+  void operator()(size_t i, const TypedPackedFunc<FType>& value) const {  // NOLINT(*)
+    operator()(i, value.packed());
+  }
+  void operator()(size_t i, const Module& value) const {  // NOLINT(*)
+    values_[i].v_handle = const_cast<Module*>(&value);
+    type_codes_[i] = kModuleHandle;
+  }
+  void operator()(size_t i, const NDArray& value) const {  // NOLINT(*)
+    values_[i].v_handle = value.data_;
+    type_codes_[i] = kNDArrayContainer;
+  }
+  void operator()(size_t i, const TVMRetValue& value) const {  // NOLINT(*)
+    if (value.type_code() == kStr) {
+      values_[i].v_str = value.ptr<std::string>()->c_str();
+      type_codes_[i] = kStr;
+    } else {
+      CHECK_NE(value.type_code(), kBytes) << "not handled.";
+      values_[i] = value.value_;
+      type_codes_[i] = value.type_code();
+    }
+  }
+  // extension
+  template<typename T,
+           typename = typename std::enable_if<
+             extension_class_info<T>::code != 0>::type>
+  inline void operator()(size_t i, const T& value) const;
+  // NodeRef related extenstions: in tvm/packed_func_ext.h
+  inline void operator()(size_t i, const NodeRef& other) const;  // NOLINT(*)
+  inline void operator()(size_t i, const HalideIR::Type& t) const;
+
+ private:
+  /*! \brief The values fields */
+  TVMValue* values_;
+  /*! \brief The type code fields */
+  int* type_codes_;
+};
+
+template<typename... Args>
+inline TVMRetValue PackedFunc::operator()(Args&& ...args) const {
+  const int kNumArgs = sizeof...(Args);
+  const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
+  TVMValue values[kArraySize];
+  int type_codes[kArraySize];
+  detail::for_each(TVMArgsSetter(values, type_codes),
+                   std::forward<Args>(args)...);
+  TVMRetValue rv;
+  body_(TVMArgs(values, type_codes, kNumArgs), &rv);
+  return rv;
+}
+
+namespace detail {
+template<typename R, int nleft, int index, typename F>
+struct unpack_call_dispatcher {
+  template<typename ...Args>
+  static void run(const F& f,
+                  const TVMArgs& args_pack,
+                  TVMRetValue* rv,
+                  Args&&... unpacked_args) {
+    unpack_call_dispatcher<R, nleft - 1, index + 1, F>
+        ::run(f, args_pack, rv,
+              std::forward<Args>(unpacked_args)...,
+              args_pack[index]);
+  }
+};
+
+template<typename R, int index, typename F>
+struct unpack_call_dispatcher<R, 0, index, F> {
+  template<typename ...Args>
+  static void run(const F& f,
+                  const TVMArgs& args_pack,
+                  TVMRetValue* rv,
+                  Args&&... unpacked_args) {
+    *rv = R(f(std::forward<Args>(unpacked_args)...));
+  }
+};
+
+template<int index, typename F>
+struct unpack_call_dispatcher<void, 0, index, F> {
+  template<typename ...Args>
+  static void run(const F& f,
+                  const TVMArgs& args_pack,
+                  TVMRetValue* rv,
+                  Args&&... unpacked_args) {
+    f(std::forward<Args>(unpacked_args)...);
+  }
+};
+
+template<typename R, int nargs, typename F>
+inline void unpack_call(const F& f, const TVMArgs& args, TVMRetValue* rv) {
+  unpack_call_dispatcher<R, nargs, 0, F>::run(f, args, rv);
+}
+
+template<typename R, typename ...Args>
+inline R call_packed(const PackedFunc& pf, Args&& ...args) {
+  return R(pf(std::forward<Args>(args)...));
+}
+
+template<typename R>
+struct typed_packed_call_dispatcher {
+  template<typename ...Args>
+  static inline R run(const PackedFunc& pf, Args&& ...args) {
+    return pf(std::forward<Args>(args)...);
+  }
+};
+
+template<>
+struct typed_packed_call_dispatcher<void> {
+  template<typename ...Args>
+  static inline void run(const PackedFunc& pf, Args&& ...args) {
+    pf(std::forward<Args>(args)...);
+  }
+};
+}  // namespace detail
+
+template<typename R, typename ...Args>
+TypedPackedFunc<R(Args...)>::TypedPackedFunc(PackedFunc packed)
+  : packed_(packed) {}
+
+template<typename R, typename ...Args>
+template<typename FType>
+inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda) {
+  packed_ = PackedFunc([flambda](const TVMArgs& args, TVMRetValue* rv) {
+      detail::unpack_call<R, sizeof...(Args)>(flambda, args, rv);
+    });
+}
+
+template<typename R, typename ...Args>
+inline R TypedPackedFunc<R(Args...)>::operator()(Args... args) const {
+  return detail::typed_packed_call_dispatcher<R>
+      ::run(packed_, std::forward<Args>(args)...);
+}
+
+// extension and node type handling
+namespace detail {
+template<typename T, typename TSrc, bool is_ext>
+struct TVMValueCast {
+  static T Apply(const TSrc* self) {
+    return self->template AsNodeRef<T>();
+  }
+};
+
+template<typename T, typename TSrc>
+struct TVMValueCast<T, TSrc, true> {
+  static T Apply(const TSrc* self) {
+    return self->template AsExtension<T>();
+  }
+};
+}  // namespace detail
+
+template<typename T, typename>
+inline TVMArgValue::operator T() const {
+  return detail::
+      TVMValueCast<T, TVMArgValue, extension_class_info<T>::code != 0>
+      ::Apply(this);
+}
+
+template<typename T, typename>
+inline TVMRetValue::operator T() const {
+  return detail::
+      TVMValueCast<T, TVMRetValue, extension_class_info<T>::code != 0>
+      ::Apply(this);
+}
+
+template<typename T, typename>
+inline void TVMArgsSetter::operator()(size_t i, const T& value) const {
+  static_assert(extension_class_info<T>::code != 0,
+                "Need to have extesion code");
+  type_codes_[i] = extension_class_info<T>::code;
+  values_[i].v_handle = const_cast<T*>(&value);
+}
+
+// extension type handling
+template<typename T>
+struct ExtTypeInfo {
+  static void destroy(void* handle) {
+    delete static_cast<T*>(handle);
+  }
+  static void* clone(void* handle) {
+    return new T(*static_cast<T*>(handle));
+  }
+};
+
+template<typename T>
+inline ExtTypeVTable* ExtTypeVTable::Register_() {
+  const int code = extension_class_info<T>::code;
+  static_assert(code != 0,
+                "require extension_class_info traits to be declared with non-zero code");
+  ExtTypeVTable vt;
+  vt.clone = ExtTypeInfo<T>::clone;
+  vt.destroy = ExtTypeInfo<T>::destroy;
+  return ExtTypeVTable::RegisterInternal(code, vt);
+}
+
+// Implement Module::GetFunction
+// Put implementation in this file so we have seen the PackedFunc
+inline PackedFunc Module::GetFunction(const std::string& name, bool query_imports) {
+  PackedFunc pf = node_->GetFunction(name, node_);
+  if (pf != nullptr) return pf;
+  if (query_imports) {
+    for (const Module& m : node_->imports_) {
+      pf = m.node_->GetFunction(name, m.node_);
+      if (pf != nullptr) return pf;
+    }
+  }
+  return pf;
+}
+}  // namespace runtime
+}  // namespace tvm
+#endif  // TVM_RUNTIME_PACKED_FUNC_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