Repository: neo-ai/neo-ai-dlr (language: python)

Issue:
Add documentation
It is convenient to host a rendered documentation for Python and C APIs. Use Sphinx, so that we can use ReadTheDocs hosting service.

EDIT. This also creates a rendered C++ documentation (using Doxygen)


Discussion / hints:

Fix patch:
diff --git a/3rdparty/treelite b/3rdparty/treelite
index 38964865d..fdf09b790 160000
--- a/3rdparty/treelite
+++ b/3rdparty/treelite
@@ -1,1 +1,1 @@
-Subproject commit 38964865d624048d1a432af2763ee1dcdbe3ce5d
+Subproject commit fdf09b790ae806dcc834ea242b90d20e9e3f32a5
diff --git a/3rdparty/tvm b/3rdparty/tvm
index 8b8e29c40..561ac4603 160000
--- a/3rdparty/tvm
+++ b/3rdparty/tvm
@@ -1,1 +1,1 @@
-Subproject commit 8b8e29c40f0c94cb29d00f4aa70e42fabc99ae26
+Subproject commit 561ac46035593e600d53b1ebdc3f8e180560f1cc
diff --git a/CMakeLists.txt b/CMakeLists.txt
index ad4b04805..97ab18ed9 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -17,7 +17,7 @@ if(ANDROID_BUILD)
     endif()
 endif(ANDROID_BUILD)
 
-project(dlr VERSION 1.9.0 LANGUAGES C CXX)
+project(dlr VERSION 1.10.0 LANGUAGES C CXX)
 
 message(STATUS "dlr version: ${dlr_VERSION}")
 
@@ -77,7 +77,14 @@ option(TEST_COVERAGE "C++ test coverage" OFF)
 # Compiler flags
 set(CMAKE_CXX_STANDARD 14)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
+# --exclude-libs is not available on Windows and macOS. As such, Windows and
+# Mac do not support the creation of multiple DLRModel instances (in Python) in
+# case model folders have their own libdlr.so.
+if (WIN32 OR APPLE)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
+else()
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--exclude-libs,ALL")
+endif()
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -funroll-loops")
 set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS}")
@@ -102,7 +109,6 @@ include_directories("${TVM_SRC}/src/runtime")
 include_directories("${DLPACK_SRC}/include")
 include_directories("${DMLC_CORE_SRC}/include")
 include_directories("${TREELITE_SRC}/include")
-include_directories("${TREELITE_SRC}/runtime/native/include")
 include_directories("${PROJECT_SOURCE_DIR}/include")
 include_directories("${JSON_SRC}")
 
@@ -141,7 +147,7 @@ endif(USE_OPENCL)
 
 if(USE_CUDA)
     message("USING CUDA")
-    find_cuda(${USE_CUDA})
+    find_cuda(${USE_CUDA} ${USE_CUDNN})
     if(NOT CUDA_FOUND)
         message(FATAL_ERROR "CUDA not found, please specify CUDA location with -DUSE_CUDA=/path/to/cuda/")
     endif(NOT CUDA_FOUND)
@@ -223,7 +229,7 @@ if(USE_TENSORRT)
     list(APPEND DLR_SRC ${RUNTIME_TENSORRT_SRCS})
 
     # Set defines
-    add_definitions(-DTVM_GRAPH_RUNTIME_TENSORRT)
+    add_definitions(-DTVM_GRAPH_EXECUTOR_TENSORRT)
 
     set(USE_TENSORRT OFF)
 endif()
@@ -233,7 +239,7 @@ if(WITH_HEXAGON)
 endif()
 
 if(ENABLE_DATATRANSFORM)
-    add_definitions(-DENABLE_DATATRANSFORM) 
+    add_definitions(-DENABLE_DATATRANSFORM)
 endif()
 
 if(AAR_BUILD)
@@ -259,23 +265,32 @@ if(USE_MKL)
 endif()
 
 if(USE_MLAS)
-    message(STATUS "Build with MLAS library")
-    if (NOT (USE_OPENMP STREQUAL "gnu" OR USE_OPENMP STREQUAL "intel"))
-        message(FATAL_ERROR "MLAS library must be built with USE_OPENMP=gnu or USE_OPENMP=intel")
-    endif()
-    file(GLOB RUNTIME_MLAS_SRCS ${TVM_SRC}/src/runtime/contrib/mlas/*.cc)
-    list(APPEND DLR_SRC ${RUNTIME_MLAS_SRCS})
-    list(APPEND DLR_LINKER_LIBS onnxruntime_mlas_static)
-    include_directories(${TVM_SRC}/3rdparty/mlas/inc)
+  message(STATUS "Build with MLAS library")
+  if (NOT (USE_OPENMP STREQUAL "gnu" OR USE_OPENMP STREQUAL "intel"))
+    message(FATAL_ERROR "MLAS library must be built with USE_OPENMP=gnu or USE_OPENMP=intel")
+  endif()
+  file(GLOB RUNTIME_MLAS_SRCS ${TVM_SRC}/src/runtime/contrib/mlas/*.cc)
+  list(APPEND DLR_SRC ${RUNTIME_MLAS_SRCS})
+  list(APPEND DLR_LINKER_LIBS onnxruntime_mlas_static)
+  include_directories(${TVM_SRC}/3rdparty/mlas/inc)
 endif()
 
 set(MAIN_EXEC "")
 FILE(GLOB MAIN_SRC src/*.cc)
 
+set(USE_LIBBACKTRACE OFF)
 add_subdirectory(${TVM_SRC} EXCLUDE_FROM_ALL)
+set(BUILD_SHARED_LIBS_SAVED "${BUILD_SHARED_LIBS}") # Save BUILD_SHARED_LIBS
+set(BUILD_STATIC_LIBS ON)
+set(BUILD_SHARED_LIBS OFF)
 add_subdirectory(${TREELITE_SRC} EXCLUDE_FROM_ALL)
+set(BUILD_SHARED_LIBS "${BUILD_SHARED_LIBS_SAVED}")  # Restore BUILD_SHARED_LIBS
+set(BUILD_STATIC_LIBS OFF)
 add_library(objdlr OBJECT ${DLR_SRC})
 
+target_compile_definitions(objdlr PUBLIC DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
+target_compile_definitions(objdlr PRIVATE TVM_USE_LIBBACKTRACE=0)
+
 #shared_library
 find_package(Threads)
 set(THREADS_PREFER_PTHREAD_FLAG TRUE)
@@ -333,6 +348,9 @@ install(EXPORT dlrTargets
   FILE dlrTargets.cmake
   DESTINATION lib/cmake/dlr
 )
+if(USE_MLAS)
+  install(TARGETS onnxruntime_mlas_static EXPORT dlrTargets)
+endif()
 
 include(CMakePackageConfigHelpers)
 # generate the config file that is includes the exports
@@ -373,7 +391,7 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
   file(GLOB TEST_SRCS tests/cpp/*.cc)
 
   if(NOT(ENABLE_DATATRANSFORM))
-    list(REMOVE_ITEM TEST_SRCS 
+    list(REMOVE_ITEM TEST_SRCS
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/cpp/dlr_data_transform_test.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/tests/cpp/dlr_pipeline_skl_xgb_test.cc)
   endif()
@@ -386,6 +404,8 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
     get_filename_component(__srcname ${__srcpath} NAME)
     string(REPLACE ".cc" "" __execname ${__srcname})
     add_executable(${__execname} ${__srcpath})
+    target_compile_definitions(${__execname} PUBLIC DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
+    target_compile_definitions(${__execname} PRIVATE TVM_USE_LIBBACKTRACE=0)
     target_link_libraries(${__execname} dlr gtest_main)
     set_output_directory(${__execname} ${CMAKE_BINARY_DIR})
     add_test(NAME ${__execname} COMMAND ${__execname})
@@ -397,7 +417,7 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
     get_filename_component(__srcname ${__srcpath} NAME)
     string(REPLACE ".cc" "" __execname ${__srcname})
     add_executable(${__execname} ${__srcpath})
-    target_link_libraries(${__execname} gtest_main dl) 
+    target_link_libraries(${__execname} gtest_main dl)
     set_output_directory(${__execname} ${CMAKE_BINARY_DIR})
     add_test(NAME ${__execname} COMMAND ${__execname})
     message(STATUS "Added dlsym Test: " ${__execname})
@@ -412,7 +432,7 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
     e35e82f3371bed37caa7ecece417f50876414077
   )
   endif()
-  
+
   set(STREET_IMAGE ${CMAKE_CURRENT_BINARY_DIR}/street_small.npy)
   if(NOT EXISTS ${STREET_IMAGE})
   download_file(
@@ -422,7 +442,7 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
     206d1a12646a5fe43d0877d38bd137c70a5adc3b
   )
   endif()
-  
+
   # Download compiled model for unit tests
   set(RESNET_MODEL ${CMAKE_CURRENT_BINARY_DIR}/resnet_v1_5_50)
   if(NOT IS_DIRECTORY ${RESNET_MODEL})
@@ -443,15 +463,15 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
   if(NOT IS_DIRECTORY ${XGBOOST_TEST_MODEL})
     file(MAKE_DIRECTORY ${XGBOOST_TEST_MODEL})
     download_file(
-      https://neo-ai-dlr-test-artifacts.s3-us-west-2.amazonaws.com/compiled-models/xgboost_test.tar.gz
-      /tmp/xgboost_test.tar.gz
+      https://neo-ai-dlr-test-artifacts.s3-us-west-2.amazonaws.com/compiled-models/release-1.10.0/xgboost-ml_m5.tar.gz
+      /tmp/xgboost-ml_m5.tar.gz
       SHA1
-      4c8ac5f20db6c7c6d674dd2ac9d9e14ce887281e
+      1e45bb9d6108d70ac4ff37855cf13061d61ef742
     )
     # this is OS-agnostic
-    execute_process(COMMAND ${CMAKE_COMMAND} -E tar -xf /tmp/xgboost_test.tar.gz
+    execute_process(COMMAND ${CMAKE_COMMAND} -E tar -xf /tmp/xgboost-ml_m5.tar.gz
                     WORKING_DIRECTORY ${XGBOOST_TEST_MODEL})
-    file(REMOVE /tmp/xgboost_test.tar.gz)
+    file(REMOVE /tmp/xgboost-ml_m5.tar.gz)
   endif()
 
   set(RELAYVM_MODEL ssd_mobilenet_v1_ppn_shared_box_predictor_300x300_coco14_sync_2018_07_03-LINUX_X86_64.tar.gz)
@@ -459,10 +479,10 @@ if(DLR_BUILD_TESTS AND NOT(AAR_BUILD))
   if(NOT IS_DIRECTORY ${RELAYVM_MODEL_DIR})
     file(MAKE_DIRECTORY ${RELAYVM_MODEL_DIR})
     download_file(
-      https://neo-ai-dlr-test-artifacts.s3-us-west-2.amazonaws.com/compiled-models/release-1.5.0/${RELAYVM_MODEL}
+      https://neo-ai-dlr-test-artifacts.s3-us-west-2.amazonaws.com/compiled-models/release-1.10.0/${RELAYVM_MODEL}
       /tmp/${RELAYVM_MODEL}
       SHA1
-      49ddd9e815c6cc14ef0e9a594c8c2d0d129e5e91
+      38111e6432d643122ebf6ed5493e415871dc3fa5
     )
     # this is OS-agnostic
     execute_process(COMMAND ${CMAKE_COMMAND} -E tar -xf /tmp/${RELAYVM_MODEL}
diff --git a/Jenkinsfile b/Jenkinsfile
index 953f42cc8..afc15583f 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -52,7 +52,7 @@ pipeline {
         unstash name: 'srcs'
         sh """
         tests/ci_build/git-clang-format.sh HEAD~1
-        tests/ci_build/git-clang-format.sh origin/main
+        tests/ci_build/git-clang-format.sh origin/$CHANGE_TARGET
         """
       }
     }
diff --git a/include/dlr.h b/include/dlr.h
index aa5dc3a5d..1bd48e4d6 100644
--- a/include/dlr.h
+++ b/include/dlr.h
@@ -19,7 +19,7 @@ extern "C" {  // Open extern "C" block
 /*! \brief major version */
 #define DLR_MAJOR 1
 /*! \brief minor version */
-#define DLR_MINOR 9
+#define DLR_MINOR 10
 /*! \brief patch version */
 #define DLR_PATCH 1
 /*! \brief DLR version */
@@ -198,7 +198,7 @@ int SetDLRInput(DLRModelHandle* handle, const char* name, const int64_t* shape,
 
 /*!
  * \brief Sets the input according the node name from existing DLTensor. Can only be
- *        used with TVM models (GraphRuntime and VMRuntime)
+ *        used with TVM models (GraphExecutor and VMRuntime)
  * \param handle The model handle returned from CreateDLRModel().
  * \param name The input node name.
  * \param tensor The input DLTensor.
@@ -209,7 +209,7 @@ int SetDLRInputTensor(DLRModelHandle* handle, const char* name, void* tensor);
 
 /*!
  * \brief Sets the input from existing DLTensor without copying data. Can only be
- *        used with TVM models (GraphRuntime). Input tensor device must match the device of the
+ *        used with TVM models (GraphExecutor). Input tensor device must match the device of the
  *        model, and data must be alligned to 128 bytes. GetDLRInput cannot be used for inputs set
  *        via SetDLRInputZeroCopy.
  * \param handle The model handle returned from CreateDLRModel().
@@ -289,7 +289,7 @@ int GetDLROutputPtr(DLRModelHandle* handle, int index, const void** out);
 
 /*!
  * \brief Gets the index-th output from the model and copies it into the given DLTensor.
- *        Can only be used with TVM models (GraphRuntime and VMRuntime)
+ *        Can only be used with TVM models (GraphExecutor and VMRuntime)
  * \param handle The model handle returned from CreateDLRModel().
  * \param index The index-th output.
  * \param tensor The pointer to an existing/allocated DLTensor to copy the output into.
@@ -300,7 +300,7 @@ int GetDLROutputTensor(DLRModelHandle* handle, int index, void* tensor);
 
 /*!
  * \brief Gets the index-th output from the model and sets the pointer to it.
- *        Can only be used with TVM models (GraphRuntime and VMRuntime)
+ *        Can only be used with TVM models (GraphExecutor and VMRuntime)
  * \param handle The model handle returned from CreateDLRModel().
  * \param index The index-th output.
  * \param tensor The pointer to an unallocated DLManagedTensor pointer, will be
diff --git a/include/dlr_common.h b/include/dlr_common.h
index 6da510d1b..1f406f24f 100644
--- a/include/dlr_common.h
+++ b/include/dlr_common.h
@@ -128,7 +128,7 @@ class DLR_DLL DLRModel {
   size_t num_inputs_ = 1;
   size_t num_weights_ = 0;
   size_t num_outputs_ = 1;
-  DLContext ctx_;
+  DLDevice dev_;
   std::vector<std::string> input_names_;
   std::vector<std::string> input_types_;
   std::vector<std::vector<int64_t>> input_shapes_;
@@ -136,7 +136,7 @@ class DLR_DLL DLRModel {
 
  public:
   nlohmann::json metadata_ = nullptr;
-  DLRModel(const DLContext& ctx, const DLRBackend& backend) : ctx_(ctx), backend_(backend) {}
+  DLRModel(const DLDevice& dev, const DLRBackend& backend) : dev_(dev), backend_(backend) {}
   virtual ~DLRModel() {}
 
   /* Input related functions */
diff --git a/include/dlr_data_transform.h b/include/dlr_data_transform.h
index 8a40b198e..489c35cff 100644
--- a/include/dlr_data_transform.h
+++ b/include/dlr_data_transform.h
@@ -23,7 +23,7 @@ class DLR_DLL Transformer {
   /*! \brief Helper function for TransformInput. Allocates NDArray to store
    * mapped input data. */
   virtual void InitNDArray(const nlohmann::json& input_json, const nlohmann::json& transform,
-                           DLDataType dtype, DLContext ctx,
+                           DLDataType dtype, DLDevice dev,
                            tvm::runtime::NDArray& input_array) const;
 };
 
@@ -79,7 +79,7 @@ class DLR_DLL DateTimeTransformer : public Transformer {
                     tvm::runtime::NDArray& input_array) const;
 
   void InitNDArray(const nlohmann::json& input_json, const nlohmann::json& transform,
-                   DLDataType dtype, DLContext ctx, tvm::runtime::NDArray& input_array) const;
+                   DLDataType dtype, DLDevice dev, tvm::runtime::NDArray& input_array) const;
 };
 
 class DLR_DLL TextTransformer : public Transformer {
@@ -90,7 +90,7 @@ class DLR_DLL TextTransformer : public Transformer {
                             tvm::runtime::NDArray& input_array) const override;
 
   virtual void InitNDArray(const nlohmann::json& input_json, const nlohmann::json& transform,
-                           DLDataType dtype, DLContext ctx,
+                           DLDataType dtype, DLDevice dev,
                            tvm::runtime::NDArray& input_array) const override;
 
   inline void SetIndex(int idx) const { column_idx_ = idx; };
@@ -150,7 +150,7 @@ class DLR_DLL DataTransform {
    * model input.
    */
   void TransformInput(const nlohmann::json& metadata, const int64_t* shape, const void* input,
-                      int dim, const std::vector<DLDataType>& dtypes, DLContext ctx,
+                      int dim, const std::vector<DLDataType>& dtypes, DLDevice dev,
                       std::vector<tvm::runtime::NDArray>* tvm_inputs) const;
 
   /*! \brief Transform integer output using CategoricalString output
diff --git a/include/dlr_hexagon/dlr_hexagon.h b/include/dlr_hexagon/dlr_hexagon.h
index 89929b8c9..5515352f1 100644
--- a/include/dlr_hexagon/dlr_hexagon.h
+++ b/include/dlr_hexagon/dlr_hexagon.h
@@ -46,7 +46,7 @@ class HexagonModel : public DLRModel {
  public:
   /*! \brief Load model files from given folder path.
    */
-  explicit HexagonModel(const std::vector<std::string>& files, const DLContext& ctx,
+  explicit HexagonModel(const std::vector<std::string>& files, const DLDevice& dev,
                         const int debug_level);
   ~HexagonModel();
 
diff --git a/include/dlr_pipeline.h b/include/dlr_pipeline.h
index 74475a5ba..f49a1dffb 100644
--- a/include/dlr_pipeline.h
+++ b/include/dlr_pipeline.h
@@ -1,7 +1,6 @@
 #ifndef DLR_PIPELINE_H_
 #define DLR_PIPELINE_H_
 
-#include <graph/graph_runtime.h>
 #include <tvm/runtime/memory.h>
 
 #include "dlr_common.h"
@@ -27,8 +26,8 @@ class DLR_DLL PipelineModel : public DLRModel {
  public:
   /*! \brief Load model files from given folder path.
    */
-  explicit PipelineModel(const std::vector<DLRModelPtr>& dlr_models, const DLContext& ctx)
-      : DLRModel(ctx, DLRBackend::kPIPELINE), dlr_models_(dlr_models) {
+  explicit PipelineModel(const std::vector<DLRModelPtr>& dlr_models, const DLDevice& dev)
+      : DLRModel(dev, DLRBackend::kPIPELINE), dlr_models_(dlr_models) {
     SetupPipelineModel();
   }
 
diff --git a/include/dlr_relayvm.h b/include/dlr_relayvm.h
index fde75cf11..a7f9462f5 100644
--- a/include/dlr_relayvm.h
+++ b/include/dlr_relayvm.h
@@ -2,7 +2,8 @@
 #define DLR_RELAYVM_H_
 
 #include <dlpack/dlpack.h>
-#include <tvm/runtime/container.h>
+#include <tvm/runtime/container/adt.h>
+#include <tvm/runtime/container/shape_tuple.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/object.h>
@@ -62,15 +63,15 @@ class DLR_DLL RelayVMModel : public DLRModel {
   DLDataType GetInputDLDataType(int index);
 
  public:
-  explicit RelayVMModel(const std::vector<std::string>& files, const DLContext& ctx)
-      : DLRModel(ctx, DLRBackend::kRELAYVM),
+  explicit RelayVMModel(const std::vector<std::string>& files, const DLDevice& dev)
+      : DLRModel(dev, DLRBackend::kRELAYVM),
         allocator_type_(tvm::runtime::vm::AllocatorType::kPooled) {
     SetupVMModule(files);
     FetchInputNodesData();
     FetchOutputNodesData();
   }
-  explicit RelayVMModel(std::vector<DLRModelElem> model_elems, const DLContext& ctx)
-      : DLRModel(ctx, DLRBackend::kRELAYVM),
+  explicit RelayVMModel(std::vector<DLRModelElem> model_elems, const DLDevice& dev)
+      : DLRModel(dev, DLRBackend::kRELAYVM),
         allocator_type_(tvm::runtime::vm::AllocatorType::kPooled) {
     SetupVMModule(model_elems);
     FetchInputNodesData();
diff --git a/include/dlr_treelite.h b/include/dlr_treelite.h
index ac416c8c0..f49cf4015 100644
--- a/include/dlr_treelite.h
+++ b/include/dlr_treelite.h
@@ -22,7 +22,7 @@ struct TreeliteInput {
   std::vector<size_t, DLRAllocator<size_t>> row_ptr;
   size_t num_row;
   size_t num_col;
-  CSRBatchHandle handle = nullptr;
+  DMatrixHandle handle = nullptr;
   ~TreeliteInput();
 };
 
@@ -58,8 +58,8 @@ class DLR_DLL TreeliteModel : public DLRModel {
  public:
   /*! \brief Load model files from given folder path.
    */
-  explicit TreeliteModel(const std::vector<std::string>& files, const DLContext& ctx)
-      : DLRModel(ctx, DLRBackend::kTREELITE) {
+  explicit TreeliteModel(const std::vector<std::string>& files, const DLDevice& dev)
+      : DLRModel(dev, DLRBackend::kTREELITE) {
     SetupTreeliteModule(files);
   }
 
diff --git a/include/dlr_tvm.h b/include/dlr_tvm.h
index 487be420e..12d5f4d26 100644
--- a/include/dlr_tvm.h
+++ b/include/dlr_tvm.h
@@ -1,7 +1,7 @@
 #ifndef DLR_TVM_H_
 #define DLR_TVM_H_
 
-#include <graph/graph_runtime.h>
+#include <graph_executor/graph_executor.h>
 #include <tvm/runtime/memory.h>
 #include <tvm/runtime/registry.h>
 
@@ -23,7 +23,7 @@ namespace dlr {
  */
 class DLR_DLL TVMModel : public DLRModel {
  private:
-  tvm::runtime::ObjectPtr<tvm::runtime::GraphRuntime> tvm_graph_runtime_;
+  tvm::runtime::ObjectPtr<tvm::runtime::GraphExecutor> tvm_graph_executor_;
   std::shared_ptr<tvm::runtime::Module> tvm_module_;
   std::vector<tvm::runtime::NDArray> inputs_;
   std::vector<tvm::runtime::NDArray> outputs_;
@@ -41,12 +41,12 @@ class DLR_DLL TVMModel : public DLRModel {
  public:
   /*! \brief Load model files from given folder path.
    */
-  explicit TVMModel(const std::vector<std::string>& files, const DLContext& ctx)
-      : DLRModel(ctx, DLRBackend::kTVM) {
+  explicit TVMModel(const std::vector<std::string>& files, const DLDevice& dev)
+      : DLRModel(dev, DLRBackend::kTVM) {
     SetupTVMModule(files);
   }
-  explicit TVMModel(std::vector<DLRModelElem> model_elems, const DLContext& ctx)
-      : DLRModel(ctx, DLRBackend::kTVM) {
+  explicit TVMModel(std::vector<DLRModelElem> model_elems, const DLDevice& dev)
+      : DLRModel(dev, DLRBackend::kTVM) {
     SetupTVMModule(model_elems);
   }
 
diff --git a/python/dlr/metadata.py b/python/dlr/metadata.py
index 75bcdef2b..93aacfbc8 100644
--- a/python/dlr/metadata.py
+++ b/python/dlr/metadata.py
@@ -1,3 +1,3 @@
 NAME = ['DLRModel']
 
-VERSION = "1.9.1"
+VERSION = "1.10.0"
diff --git a/src/dlr.cc b/src/dlr.cc
index 9ea18538b..b70bda28e 100644
--- a/src/dlr.cc
+++ b/src/dlr.cc
@@ -260,20 +260,19 @@ extern "C" int GetDLROutputByName(DLRModelHandle* handle, const char* name, void
   API_END();
 }
 
-
-DLRModelPtr CreateDLRModelPtr(const char* model_path, DLContext& ctx) {
+DLRModelPtr CreateDLRModelPtr(const char* model_path, DLDevice& dev) {
   std::vector<std::string> path_vec = dlr::MakePathVec(model_path);
   std::vector<std::string> files = FindFiles(path_vec);
   DLRBackend backend = dlr::GetBackend(files);
   if (backend == DLRBackend::kTVM) {
-    return std::make_shared<TVMModel>(files, ctx);
+    return std::make_shared<TVMModel>(files, dev);
   } else if (backend == DLRBackend::kRELAYVM) {
-    return std::make_shared<RelayVMModel>(files, ctx);
+    return std::make_shared<RelayVMModel>(files, dev);
   } else if (backend == DLRBackend::kTREELITE) {
-    return std::make_shared<TreeliteModel>(files, ctx);
+    return std::make_shared<TreeliteModel>(files, dev);
 #ifdef DLR_HEXAGON
   } else if (backend == DLRBackend::kHEXAGON) {
-    return std::make_shared<HexagonModel>(files, ctx, 1 /*debug_level*/);
+    return std::make_shared<HexagonModel>(files, dev, 1 /*debug_level*/);
 #endif  // DLR_HEXAGON
   } else {
     std::string err = "Unable to determine backend from path: '";
@@ -289,13 +288,13 @@ DLRModelPtr CreateDLRModelPtr(const char* model_path, DLContext& ctx) {
  */
 int CreateDLRModelFromHexagon(DLRModelHandle* handle, const char* model_path, int debug_level) {
   API_BEGIN();
-  // HexagonModel class does not use DLContext internally
-  DLContext ctx;
-  ctx.device_type = static_cast<DLDeviceType>(1);  // 1 - kDLCPU
-  ctx.device_id = 0;
+  // HexagonModel class does not use DLDevice internally
+  DLDevice dev;
+  dev.device_type = static_cast<DLDeviceType>(1);  // 1 - kDLCPU
+  dev.device_id = 0;
   std::vector<std::string> path_vec = dlr::MakePathVec(model_path);
   std::vector<std::string> files = FindFiles(path_vec);
-  DLRModel* model = new HexagonModel(files, ctx, debug_level);
+  DLRModel* model = new HexagonModel(files, dev, debug_level);
   *handle = model;
   API_END();
 }
@@ -306,9 +305,9 @@ int CreateDLRModelFromHexagon(DLRModelHandle* handle, const char* model_path, in
 extern "C" int CreateDLRModel(DLRModelHandle* handle, const char* model_path, int dev_type,
                               int dev_id) {
   API_BEGIN();
-  DLContext ctx;
-  ctx.device_type = static_cast<DLDeviceType>(dev_type);
-  ctx.device_id = dev_id;
+  DLDevice dev;
+  dev.device_type = static_cast<DLDeviceType>(dev_type);
+  dev.device_id = dev_id;
 
   std::vector<std::string> path_vec = dlr::MakePathVec(model_path);
   std::vector<std::string> files = FindFiles(path_vec);
@@ -317,14 +316,14 @@ extern "C" int CreateDLRModel(DLRModelHandle* handle, const char* model_path, in
   DLRModel* model;
   try {
     if (backend == DLRBackend::kTVM) {
-      model = new TVMModel(files, ctx);
+      model = new TVMModel(files, dev);
     } else if (backend == DLRBackend::kRELAYVM) {
-      model = new RelayVMModel(files, ctx);
+      model = new RelayVMModel(files, dev);
     } else if (backend == DLRBackend::kTREELITE) {
-      model = new TreeliteModel(files, ctx);
+      model = new TreeliteModel(files, dev);
 #ifdef DLR_HEXAGON
     } else if (backend == DLRBackend::kHEXAGON) {
-      model = new HexagonModel(files, ctx, 1 /*debug_level*/);
+      model = new HexagonModel(files, dev, 1 /*debug_level*/);
 #endif  // DLR_HEXAGON
     } else {
       std::string err = "Unable to determine backend from path: '";
@@ -344,9 +343,9 @@ extern "C" int CreateDLRModel(DLRModelHandle* handle, const char* model_path, in
 extern "C" int CreateDLRModelFromModelElem(DLRModelHandle* handle, const DLRModelElem* model_elems,
                                            size_t model_elems_size, int dev_type, int dev_id) {
   API_BEGIN();
-  DLContext ctx;
-  ctx.device_type = static_cast<DLDeviceType>(dev_type);
-  ctx.device_id = dev_id;
+  DLDevice dev;
+  dev.device_type = static_cast<DLDeviceType>(dev_type);
+  dev.device_id = dev_id;
 
   std::vector<DLRModelElem> model_elem_vec(model_elems, model_elems + model_elems_size);
 
@@ -354,9 +353,9 @@ extern "C" int CreateDLRModelFromModelElem(DLRModelHandle* handle, const DLRMode
   DLRModel* model;
   try {
     if (backend == DLRBackend::kTVM) {
-      model = new TVMModel(model_elem_vec, ctx);
+      model = new TVMModel(model_elem_vec, dev);
     } else if (backend == DLRBackend::kRELAYVM) {
-      model = new RelayVMModel(model_elem_vec, ctx);
+      model = new RelayVMModel(model_elem_vec, dev);
     } else {
       LOG(FATAL) << "Unsupported backend!";
       return -1;  // unreachable
@@ -375,20 +374,20 @@ extern "C" int CreateDLRModelFromModelElem(DLRModelHandle* handle, const DLRMode
 extern "C" int CreateDLRPipeline(DLRModelHandle* handle, int num_models, const char** model_paths,
                                  int dev_type, int dev_id) {
   API_BEGIN();
-  DLContext ctx;
-  ctx.device_type = static_cast<DLDeviceType>(dev_type);
-  ctx.device_id = dev_id;
+  DLDevice dev;
+  dev.device_type = static_cast<DLDeviceType>(dev_type);
+  dev.device_id = dev_id;
   std::vector<DLRModelPtr> dlr_models;
   for (int i = 0; i < num_models; i++) {
     try {
-      DLRModelPtr model_ptr = CreateDLRModelPtr(model_paths[i], ctx);
+      DLRModelPtr model_ptr = CreateDLRModelPtr(model_paths[i], dev);
       dlr_models.push_back(model_ptr);
     } catch (dmlc::Error& e) {
       LOG(ERROR) << e.what();
       return -1;
     }
   }
-  DLRModel* pipeline_model = new PipelineModel(dlr_models, ctx);
+  DLRModel* pipeline_model = new PipelineModel(dlr_models, dev);
   *handle = pipeline_model;
   API_END();
 }
diff --git a/src/dlr_common.cc b/src/dlr_common.cc
index 3021d5b22..1cca9290a 100644
--- a/src/dlr_common.cc
+++ b/src/dlr_common.cc
@@ -214,10 +214,10 @@ void DLRModel::ValidateDeviceTypeIfExists() {
     // Ignore missing metadata file or missing device type.
     return;
   }
-  if (device_type != 0 && ctx_.device_type != device_type) {
+  if (device_type != 0 && dev_.device_type != device_type) {
     std::string msg = "Compiled model requires device type \"";
     msg += GetStringFromDeviceType(device_type) + "\" but user gave \"";
-    msg += GetStringFromDeviceType(ctx_.device_type) + "\".";
+    msg += GetStringFromDeviceType(dev_.device_type) + "\".";
     throw dmlc::Error(msg);
   }
 }
@@ -239,7 +239,7 @@ DLDeviceType dlr::GetDeviceTypeFromString(const std::string& device_type_string)
   if (device_type_string == "cpu") {
     return DLDeviceType::kDLCPU;
   } else if (device_type_string == "gpu") {
-    return DLDeviceType::kDLGPU;
+    return DLDeviceType::kDLCUDA;
   } else if (device_type_string == "opencl") {
     return DLDeviceType::kDLOpenCL;
   }
@@ -249,7 +249,7 @@ DLDeviceType dlr::GetDeviceTypeFromString(const std::string& device_type_string)
 std::string dlr::GetStringFromDeviceType(DLDeviceType device_type) {
   if (device_type == DLDeviceType::kDLCPU) {
     return "cpu";
-  } else if (device_type == DLDeviceType::kDLGPU) {
+  } else if (device_type == DLDeviceType::kDLCUDA) {
     return "gpu";
   } else if (device_type == DLDeviceType::kDLOpenCL) {
     return "opencl";
diff --git a/src/dlr_data_transform.cc b/src/dlr_data_transform.cc
index e26b252fe..9dfd299b2 100644
--- a/src/dlr_data_transform.cc
+++ b/src/dlr_data_transform.cc
@@ -22,7 +22,7 @@ bool DataTransform::HasOutputTransform(const nlohmann::json& metadata, int index
 
 void DataTransform::TransformInput(const nlohmann::json& metadata, const int64_t* shape,
                                    const void* input, int dim,
-                                   const std::vector<DLDataType>& dtypes, DLContext ctx,
+                                   const std::vector<DLDataType>& dtypes, DLDevice dev,
                                    std::vector<tvm::runtime::NDArray>* tvm_inputs) const {
   nlohmann::json input_json = GetAsJson(shape, input, dim);
   const auto& transforms = metadata["DataTransform"]["Input"]["ColumnTransform"];
@@ -34,7 +34,7 @@ void DataTransform::TransformInput(const nlohmann::json& metadata, const int64_t
         << transformer_type << " is not a valid DataTransform type.";
     const Transformer* transformer = &(*it->second);
 
-    transformer->InitNDArray(input_json, transforms[i], dtypes[i], ctx, tvm_inputs->at(i));
+    transformer->InitNDArray(input_json, transforms[i], dtypes[i], dev, tvm_inputs->at(i));
     transformer->MapToNDArray(input_json, transforms[i], tvm_inputs->at(i));
   }
 }
@@ -55,18 +55,18 @@ nlohmann::json DataTransform::GetAsJson(const int64_t* shape, const void* input,
 }
 
 void Transformer::InitNDArray(const nlohmann::json& input_json, const nlohmann::json& transform,
-                              DLDataType dtype, DLContext ctx,
+                              DLDataType dtype, DLDevice dev,
                               tvm::runtime::NDArray& input_array) const {
   // Create NDArray for transformed input which will be passed to TVM.
   std::vector<int64_t> arr_shape = {static_cast<int64_t>(input_json.size()),
                                     static_cast<int64_t>(input_json[0].size())};
   CHECK(dtype.code == kDLFloat && dtype.bits == 32 && dtype.lanes == 1)
-      << "DataTransform CategoricalString is only supported for float32 "
-         "inputs.";
-  // Only allocate new buffer if not initialized or if shape or dtype has
-  // changed. Context will always match.
-  if (input_array == empty_ || input_array.Shape() != arr_shape) {
-    input_array = tvm::runtime::NDArray::Empty(arr_shape, dtype, ctx);
+      << "DataTransform CategoricalString is only supported for float32 inputs.";
+  // Only allocate new buffer if not initialized or if shape or dtype has changed. Context will
+  // always match.
+  if (input_array == empty_ || !std::equal(input_array.Shape().begin(), input_array.Shape().end(),
+                                           arr_shape.begin(), arr_shape.end())) {
+    input_array = tvm::runtime::NDArray::Empty(arr_shape, dtype, dev);
   }
 }
 
@@ -74,7 +74,7 @@ void FloatTransformer::MapToNDArray(const nlohmann::json& input_json,
                                     const nlohmann::json& transform,
                                     tvm::runtime::NDArray& input_array) const {
   DLTensor* input_tensor = const_cast<DLTensor*>(input_array.operator->());
-  CHECK_EQ(input_tensor->ctx.device_type, DLDeviceType::kDLCPU)
+  CHECK_EQ(input_tensor->device.device_type, DLDeviceType::kDLCPU)
       << "DataTransform is only supported for CPU.";
   float* data = static_cast<float*>(input_tensor->data);
   for (size_t r = 0; r < input_json.size(); ++r) {
@@ -99,10 +99,9 @@ void CategoricalStringTransformer::MapToNDArray(const nlohmann::json& input_json
                                                 tvm::runtime::NDArray& input_array) const {
   const nlohmann::json& mapping = transform["Map"];
   DLTensor* input_tensor = const_cast<DLTensor*>(input_array.operator->());
-  // Writing directly to the DLTensor will only work for CPU context. For other
-  // contexts, we would need to create an intermediate buffer on CPU and copy
-  // that to the context.
-  CHECK_EQ(input_tensor->ctx.device_type, DLDeviceType::kDLCPU)
+  // Writing directly to the DLTensor will only work for CPU context. For other contexts, we would
+  // need to create an intermediate buffer on CPU and copy that to the context.
+  CHECK_EQ(input_tensor->device.device_type, DLDeviceType::kDLCPU)
       << "DataTransform CategoricalString is only supported for CPU.";
   CHECK_EQ(input_json[0].size(), mapping.size())
       << "Input has " << input_json[0].size() << " columns, but model requires " << mapping.size();
@@ -134,7 +133,7 @@ void CategoricalStringTransformer::MapToNDArray(const nlohmann::json& input_json
 
 void DateTimeTransformer::InitNDArray(const nlohmann::json& input_json,
                                       const nlohmann::json& transform, DLDataType dtype,
-                                      DLContext ctx, tvm::runtime::NDArray& input_array) const {
+                                      DLDevice dev, tvm::runtime::NDArray& input_array) const {
   // Create NDArray for transformed input which will be passed to TVM. NUM_COL
   // fixed to 7
   auto date_col = transform["DateCol"].get<std::vector<int>>();
@@ -143,8 +142,9 @@ void DateTimeTransformer::InitNDArray(const nlohmann::json& input_json,
   CHECK(dtype.code == kDLFloat && dtype.bits == 32 && dtype.lanes == 1)
       << "DataTransform DateTimeTransformer is only supported for float32 "
          "inputs.";
-  if (input_array == empty_ || input_array.Shape() != arr_shape) {
-    input_array = tvm::runtime::NDArray::Empty(arr_shape, dtype, ctx);
+  if (input_array == empty_ || !std::equal(input_array.Shape().begin(), input_array.Shape().end(),
+                                           arr_shape.begin(), arr_shape.end())) {
+    input_array = tvm::runtime::NDArray::Empty(arr_shape, dtype, dev);
   }
 }
 
@@ -187,7 +187,7 @@ void DateTimeTransformer::MapToNDArray(const nlohmann::json& input_json,
                                        const nlohmann::json& transform,
                                        tvm::runtime::NDArray& input_array) const {
   DLTensor* input_tensor = const_cast<DLTensor*>(input_array.operator->());
-  CHECK_EQ(input_tensor->ctx.device_type, DLDeviceType::kDLCPU)
+  CHECK_EQ(input_tensor->device.device_type, DLDeviceType::kDLCPU)
       << "DataTransform DateTimeVectorizer is only supported for CPU.";
   float* data = static_cast<float*>(input_tensor->data);
 
@@ -250,7 +250,7 @@ TextTransformer::TextTransformer() {
 }
 
 void TextTransformer::InitNDArray(const nlohmann::json& input_json, const nlohmann::json& transform,
-                                  DLDataType dtype, DLContext ctx,
+                                  DLDataType dtype, DLDevice dev,
                                   tvm::runtime::NDArray& input_array) const {
   auto vocabularies = transform["Vocabularies"].get<std::vector<std::string>>();
   auto text_col = transform["TextCol"].get<int>();
@@ -268,8 +268,9 @@ void TextTransformer::InitNDArray(const nlohmann::json& input_json, const nlohma
 
   CHECK(dtype.code == kDLFloat && dtype.bits == 32 && dtype.lanes == 1)
       << "DataTransform TextTransformer is only supported for float32 inputs.";
-  if (input_array == empty_ || input_array.Shape() != arr_shape) {
-    input_array = tvm::runtime::NDArray::Empty(arr_shape, dtype, ctx);
+  if (input_array == empty_ || !std::equal(input_array.Shape().begin(), input_array.Shape().end(),
+                                           arr_shape.begin(), arr_shape.end())) {
+    input_array = tvm::runtime::NDArray::Empty(arr_shape, dtype, dev);
   }
 }
 
@@ -278,7 +279,7 @@ void TextTransformer::MapToNDArray(const nlohmann::json& input_json,
                                    tvm::runtime::NDArray& input_array) const {
   const nlohmann::json& vocab = transform["Vocabularies"];
   DLTensor* input_tensor = const_cast<DLTensor*>(input_array.operator->());
-  CHECK_EQ(input_tensor->ctx.device_type, DLDeviceType::kDLCPU)
+  CHECK_EQ(input_tensor->device.device_type, DLDeviceType::kDLCPU)
       << "DataTransform TfIdfVectorizer is only supported for CPU.";
 
   int id = col_to_id_->at(column_idx_);
@@ -319,7 +320,7 @@ void DataTransform::TransformOutput(const nlohmann::json& metadata, int index,
                                     const tvm::runtime::NDArray& output_array) {
   auto& transform = metadata["DataTransform"]["Output"][std::to_string(index)];
   const DLTensor* tensor = output_array.operator->();
-  CHECK_EQ(tensor->ctx.device_type, DLDeviceType::kDLCPU)
+  CHECK_EQ(tensor->device.device_type, DLDeviceType::kDLCPU)
       << "DataTransform CategoricalString is only supported for CPU.";
   CHECK(tensor->dtype.code == kDLInt && tensor->dtype.bits == 32 && tensor->dtype.lanes == 1)
       << "DataTransform CategoricalString is only supported for int32 outputs.";
diff --git a/src/dlr_hexagon/dlr_hexagon.cc b/src/dlr_hexagon/dlr_hexagon.cc
index 6947dd5f9..787195869 100644
--- a/src/dlr_hexagon/dlr_hexagon.cc
+++ b/src/dlr_hexagon/dlr_hexagon.cc
@@ -96,9 +96,9 @@ int HexagonModel::GetInputId(const char* name) {
 }
 
 // Constructor
-HexagonModel::HexagonModel(const std::vector<std::string>& files, const DLContext& ctx,
+HexagonModel::HexagonModel(const std::vector<std::string>& files, const DLDevice& dev,
                            const int debug_level)
-    : DLRModel(ctx, DLRBackend::kHEXAGON) {
+    : DLRModel(dev, DLRBackend::kHEXAGON) {
   const std::string model_so_file = GetHexagonModelFile(files);
   LOG(INFO) << "Model: " << model_so_file;
   const std::string model_folder = GetParentFolder(model_so_file);
diff --git a/src/dlr_relayvm.cc b/src/dlr_relayvm.cc
index aa60cbdca..20f61faa6 100644
--- a/src/dlr_relayvm.cc
+++ b/src/dlr_relayvm.cc
@@ -93,11 +93,11 @@ void RelayVMModel::SetupVMModule(const std::vector<DLRModelElem>& model_elems) {
   vm_module_ = std::make_shared<tvm::runtime::Module>(tvm::runtime::Module(vm));
 
   tvm::runtime::PackedFunc init = vm_module_->GetFunction("init");
-  if (ctx_.device_type == DLDeviceType::kDLCPU) {
-    init(static_cast<int>(ctx_.device_type), ctx_.device_id, static_cast<int>(allocator_type_));
+  if (dev_.device_type == DLDeviceType::kDLCPU) {
+    init(static_cast<int>(dev_.device_type), dev_.device_id, static_cast<int>(allocator_type_));
   } else {
     // CPU context also must be initialized because input/output data comes from CPU.
-    init(static_cast<int>(ctx_.device_type), ctx_.device_id, static_cast<int>(allocator_type_),
+    init(static_cast<int>(dev_.device_type), dev_.device_id, static_cast<int>(allocator_type_),
          static_cast<int>(DLDeviceType::kDLCPU), 0, static_cast<int>(allocator_type_));
   }
 }
@@ -192,7 +192,7 @@ void RelayVMModel::GetInput(const char* name, void* input) {
   auto in_array = inputs_[index];
   DLTensor input_tensor;
   input_tensor.data = input;
-  input_tensor.ctx = ctx_;
+  input_tensor.device = dev_;
   input_tensor.ndim = in_array->ndim;
   input_tensor.dtype = in_array->dtype;
   input_tensor.shape = in_array->shape;
@@ -286,7 +286,7 @@ void RelayVMModel::SetInput(const char* name, const int64_t* shape, const void*
     for (size_t i = 0; i < num_inputs_; ++i) {
       dtypes.emplace_back(GetInputDLDataType(i));
     }
-    data_transform_.TransformInput(metadata_, shape, input, dim, dtypes, ctx_, &inputs_);
+    data_transform_.TransformInput(metadata_, shape, input, dim, dtypes, dev_, &inputs_);
     return;
   }
 #endif
@@ -295,7 +295,7 @@ void RelayVMModel::SetInput(const char* name, const int64_t* shape, const void*
   DLDataType dtype = GetInputDLDataType(index);
   DLTensor input_tensor;
   input_tensor.data = const_cast<void*>(input);
-  input_tensor.ctx = DLContext{DLDeviceType::kDLCPU, 0};
+  input_tensor.device = DLDevice{DLDeviceType::kDLCPU, 0};
   input_tensor.ndim = dim;
   input_tensor.shape = const_cast<int64_t*>(shape);
   input_tensor.strides = nullptr;
@@ -305,9 +305,11 @@ void RelayVMModel::SetInput(const char* name, const int64_t* shape, const void*
 
   // Only allocate new buffer if not initialized or if shape or dtype has changed. Context will
   // always match.
-  if (inputs_[index] == empty_ || inputs_[index].Shape() != arr_shape ||
+  if (inputs_[index] == empty_ ||
+      !std::equal(inputs_[index].Shape().begin(), inputs_[index].Shape().end(), arr_shape.begin(),
+                  arr_shape.end()) ||
       !TypeEqual(inputs_[index].DataType(), dtype)) {
-    inputs_[index] = tvm::runtime::NDArray::Empty(arr_shape, dtype, ctx_);
+    inputs_[index] = tvm::runtime::NDArray::Empty(arr_shape, dtype, dev_);
   }
   inputs_[index].CopyFrom(&input_tensor);
 }
@@ -321,7 +323,7 @@ void RelayVMModel::SetInputTensor(const char* name, DLTensor* tensor) {
       dtypes.emplace_back(GetInputDLDataType(i));
     }
     data_transform_.TransformInput(metadata_, tensor->shape, tensor->data, tensor->ndim, dtypes,
-                                   ctx_, &inputs_);
+                                   dev_, &inputs_);
     return;
   }
 #endif
@@ -329,7 +331,7 @@ void RelayVMModel::SetInputTensor(const char* name, DLTensor* tensor) {
   int index = GetInputIndex(name);
   if (index > -1) {
     std::vector<int64_t> arr_shape(tensor->shape, tensor->shape + tensor->ndim);
-    tvm::runtime::NDArray input_arr = tvm::runtime::NDArray::Empty(arr_shape, tensor->dtype, ctx_);
+    tvm::runtime::NDArray input_arr = tvm::runtime::NDArray::Empty(arr_shape, tensor->dtype, dev_);
     input_arr.CopyFrom(tensor);
     inputs_[index] = input_arr;
   }
@@ -392,7 +394,7 @@ void RelayVMModel::GetOutput(int index, void* output) {
 
   DLTensor output_tensor;
   output_tensor.data = output;
-  output_tensor.ctx = DLContext{DLDeviceType::kDLCPU, 0};
+  output_tensor.device = DLDevice{DLDeviceType::kDLCPU, 0};
   output_tensor.ndim = out_array->ndim;
   output_tensor.dtype = out_array->dtype;
   output_tensor.shape = out_array->shape;
diff --git a/src/dlr_treelite.cc b/src/dlr_treelite.cc
index 015b32c24..b641db763 100644
--- a/src/dlr_treelite.cc
+++ b/src/dlr_treelite.cc
@@ -41,7 +41,7 @@ ModelPath dlr::SetTreelitePaths(const std::vector<std::string>& files) {
 }
 
 TreeliteInput::~TreeliteInput() {
-  if (handle) TreeliteDeleteSparseBatch(handle);
+  if (handle) TreeliteDMatrixFree(handle);
   handle = nullptr;
 }
 
@@ -62,22 +62,31 @@ void TreeliteModel::SetupTreeliteModule(const std::vector<std::string>& model_pa
       << TreeliteGetLastError();
   treelite_input_.reset(nullptr);
 
+  const char* output_type;
+  CHECK_EQ(TreelitePredictorQueryLeafOutputType(treelite_model_, &output_type), 0)
+      << TreeliteGetLastError();
+  CHECK_EQ(std::string(output_type), "float32")
+      << "Only float32 output types are supported, got " << output_type;
   size_t num_output_class;  // > 1 for multi-class classification; 1 otherwise
-  CHECK_EQ(TreelitePredictorQueryNumOutputGroup(treelite_model_, &num_output_class), 0)
+  CHECK_EQ(TreelitePredictorQueryNumClass(treelite_model_, &num_output_class), 0)
       << TreeliteGetLastError();
   treelite_output_buffer_size_ = num_output_class;
   treelite_output_.empty();
+
   // NOTE: second dimension of the output shape is smaller than num_output_class
   //       when a multi-class classifier outputs only the class prediction
-  //       (argmax) To detect this edge case, run TreelitePredictorPredictInst()
-  //       once.
-  std::vector<TreelitePredictorEntry> tmp_in(treelite_num_feature_);
-  std::vector<float> tmp_out(num_output_class);
-  CHECK_EQ(TreelitePredictorPredictInst(treelite_model_, tmp_in.data(), 0, tmp_out.data(),
-                                        &treelite_output_size_),
+  //       (argmax) To detect this edge case, run TreelitePredictorQueryResultSize()
+  DMatrixHandle tmp_matrix;
+  std::vector<float> tmp_in(treelite_num_feature_);
+  const float missing_value = 0.0f;
+  CHECK_EQ(TreeliteDMatrixCreateFromMat(tmp_in.data(), "float32", /*num_row=*/1,
+                                        treelite_num_feature_, &missing_value, &tmp_matrix),
            0)
       << TreeliteGetLastError();
+  CHECK_EQ(TreelitePredictorQueryResultSize(treelite_model_, tmp_matrix, &treelite_output_size_), 0)
+      << TreeliteGetLastError();
   CHECK_LE(treelite_output_size_, num_output_class) << "Precondition violated";
+
   UpdateInputShapes();
   has_sparse_input_ = false;
   if (!paths.metadata.empty() && !IsFileEmpty(paths.metadata)) {
@@ -166,9 +175,9 @@ void TreeliteModel::SetInput(const char* name, const int64_t* shape, const void*
 
   // Register CSR matrix with Treelite backend
   CHECK_EQ(
-      TreeliteAssembleSparseBatch(treelite_input_->data.data(), treelite_input_->col_ind.data(),
-                                  treelite_input_->row_ptr.data(), batch_size,
-                                  treelite_num_feature_, &treelite_input_->handle),
+      TreeliteDMatrixCreateFromCSR(treelite_input_->data.data(), "float32",
+                                   treelite_input_->col_ind.data(), treelite_input_->row_ptr.data(),
+                                   batch_size, treelite_num_feature_, &treelite_input_->handle),
       0)
       << TreeliteGetLastError();
   UpdateInputShapes();
@@ -214,8 +223,9 @@ void TreeliteModel::Run() {
   size_t out_result_size;
   CHECK(treelite_input_);
   treelite_output_.resize(treelite_input_->num_row * treelite_output_buffer_size_);
-  CHECK_EQ(TreelitePredictorPredictBatch(treelite_model_, treelite_input_->handle, 1, 0,
-                                         pred_margin, treelite_output_.data(), &out_result_size),
+  CHECK_EQ(TreelitePredictorPredictBatch(treelite_model_, treelite_input_->handle, 0, pred_margin,
+                                         (PredictorOutputHandle*)treelite_output_.data(),
+                                         &out_result_size),
            0)
       << TreeliteGetLastError();
 }
diff --git a/src/dlr_tvm.cc b/src/dlr_tvm.cc
index 6fcae8dfe..347c67ef8 100644
--- a/src/dlr_tvm.cc
+++ b/src/dlr_tvm.cc
@@ -96,22 +96,22 @@ void TVMModel::SetupTVMModule(const std::vector<DLRModelElem>& model_elems) {
   tvm::runtime::Module module;
   module = tvm::runtime::Module::LoadFromFile(model_lib_path);
 
-  tvm_graph_runtime_ = tvm::runtime::make_object<tvm::runtime::GraphRuntime>();
-  tvm_graph_runtime_->Init(graph_str, module, {ctx_}, nullptr);
+  tvm_graph_executor_ = tvm::runtime::make_object<tvm::runtime::GraphExecutor>();
+  tvm_graph_executor_->Init(graph_str, module, {dev_}, nullptr);
   dmlc::MemoryFixedSizeStream strm(const_cast<char*>(params_data), params_size);
-  tvm_graph_runtime_->LoadParams(&strm);
+  tvm_graph_executor_->LoadParams(&strm);
 
-  tvm_module_ = std::make_shared<tvm::runtime::Module>(tvm::runtime::Module(tvm_graph_runtime_));
+  tvm_module_ = std::make_shared<tvm::runtime::Module>(tvm::runtime::Module(tvm_graph_executor_));
 
   // Get list of weights.
-  weight_names_ = tvm_graph_runtime_->GetWeightNames();
+  weight_names_ = tvm_graph_executor_->GetWeightNames();
   num_weights_ = weight_names_.size();
   std::unordered_set<std::string> weight_names_set(weight_names_.begin(), weight_names_.end());
   // TVM inputs contains both inputs and weights.
-  const auto num_inputs_weights = tvm_graph_runtime_->NumInputs();
+  const auto num_inputs_weights = tvm_graph_executor_->NumInputs();
   // Filter out weights to get only inputs.
   for (int i = 0; i < num_inputs_weights; i++) {
-    auto name = tvm_graph_runtime_->GetInputName(i);
+    auto name = tvm_graph_executor_->GetInputName(i);
     if (weight_names_set.count(name) == 0) {
       input_names_.push_back(name);
     }
@@ -121,18 +121,18 @@ void TVMModel::SetupTVMModule(const std::vector<DLRModelElem>& model_elems) {
   inputs_.resize(num_inputs_);
   input_types_.resize(num_inputs_);
   for (int i = 0; i < num_inputs_; i++) {
-    inputs_[i] = tvm_graph_runtime_->GetInput(i);
-    input_types_[i] = tvm_graph_runtime_->GetInputType(i);
+    inputs_[i] = tvm_graph_executor_->GetInput(i);
+    input_types_[i] = tvm_graph_executor_->GetInputType(i);
   }
 
   // Get the number of output and reserve space to save output tensor
   // pointers.
-  num_outputs_ = tvm_graph_runtime_->NumOutputs();
+  num_outputs_ = tvm_graph_executor_->NumOutputs();
   outputs_.resize(num_outputs_);
   output_types_.resize(num_outputs_);
   for (int i = 0; i < num_outputs_; i++) {
-    outputs_[i] = tvm_graph_runtime_->GetOutput(i);
-    output_types_[i] = tvm_graph_runtime_->GetOutputType(i);
+    outputs_[i] = tvm_graph_executor_->GetOutput(i);
+    output_types_[i] = tvm_graph_executor_->GetOutputType(i);
   }
   UpdateInputShapes();
 }
@@ -141,7 +141,7 @@ void TVMModel::UpdateInputShapes() {
   input_shapes_.resize(num_inputs_);
   for (int i = 0; i < num_inputs_; i++) {
     std::vector<int64_t> input_shape;
-    tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(i);
+    tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(i);
     input_shape.assign(arr->shape, arr->shape + arr->ndim);
     input_shapes_[i] = input_shape;
   }
@@ -173,13 +173,13 @@ const char* TVMModel::GetInputType(int index) const {
 
 const int TVMModel::GetInputDim(int index) const {
   CHECK_LT(index, num_inputs_) << "Input index is out of range.";
-  tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(index);
+  tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(index);
   return arr->ndim;
 }
 
 const int64_t TVMModel::GetInputSize(int index) const {
   CHECK_LT(index, num_inputs_) << "Input index is out of range.";
-  tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(index);
+  tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(index);
   if (dlr::HasNegative(arr->shape, arr->ndim)) return -1;
   return std::accumulate(arr->shape, arr->shape + arr->ndim, 1, std::multiplies<int64_t>());
 }
@@ -197,16 +197,16 @@ void TVMModel::SetInput(const char* name, const int64_t* shape, const void* inpu
     for (size_t i = 0; i < num_inputs_; ++i) {
       dtypes.emplace_back(inputs_[i]->dtype);
     }
-    data_transform_.TransformInput(metadata_, shape, input, dim, dtypes, ctx_, &inputs_);
+    data_transform_.TransformInput(metadata_, shape, input, dim, dtypes, dev_, &inputs_);
     return;
   }
 #endif
 
   std::string str(name);
-  int index = tvm_graph_runtime_->GetInputIndex(str);
-  tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(index);
+  int index = tvm_graph_executor_->GetInputIndex(str);
+  tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(index);
   DLTensor input_tensor = *(arr.operator->());
-  input_tensor.ctx = DLContext{kDLCPU, 0};
+  input_tensor.device = DLDevice{kDLCPU, 0};
   input_tensor.data = const_cast<void*>(input);
   int64_t read_size = std::accumulate(shape, shape + dim, 1, std::multiplies<int64_t>());
   int64_t expected_size = std::accumulate(
@@ -226,44 +226,44 @@ void TVMModel::SetInputTensor(const char* name, DLTensor* tensor) {
       dtypes.emplace_back(inputs_[i]->dtype);
     }
     data_transform_.TransformInput(metadata_, tensor->shape, tensor->data, tensor->ndim, dtypes,
-                                   ctx_, &inputs_);
+                                   dev_, &inputs_);
     return;
   }
 #endif
 
   std::string str(name);
-  int index = tvm_graph_runtime_->GetInputIndex(str);
+  int index = tvm_graph_executor_->GetInputIndex(str);
   if (index > -1) {
-    tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(index);
+    tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(index);
     DLTensor input_tensor = *(arr.operator->());
     int64_t read_size =
         std::accumulate(tensor->shape, tensor->shape + tensor->ndim, 1, std::multiplies<int64_t>());
     int64_t expected_size = std::accumulate(
         input_tensor.shape, input_tensor.shape + input_tensor.ndim, 1, std::multiplies<int64_t>());
     CHECK_SHAPE("Mismatch found in input data size", read_size, expected_size);
-    tvm_graph_runtime_->SetInput(index, tensor);
+    tvm_graph_executor_->SetInput(index, tensor);
   }
 }
 
 void TVMModel::SetInputTensorZeroCopy(const char* name, DLTensor* tensor) {
   std::string str(name);
-  int index = tvm_graph_runtime_->GetInputIndex(str);
+  int index = tvm_graph_executor_->GetInputIndex(str);
   if (index == -1) return;
-  tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(index);
+  tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(index);
   const DLTensor* old_t = arr.operator->();
   CHECK_EQ(reinterpret_cast<size_t>(tensor->data) % 128, 0)
       << "Data must be alligned to 128 bits for SetDLRInputTensorZeroCopy.";
   CHECK_EQ(old_t->ndim, static_cast<size_t>(tensor->ndim))
       << "Model expected " << old_t->ndim << " dimensions, but input has " << tensor->ndim;
-  CHECK_EQ(old_t->ctx.device_type, tensor->ctx.device_type)
-      << "The input data must be on device \"" << GetStringFromDeviceType(old_t->ctx.device_type)
-      << "\", but user gave input on \"" << GetStringFromDeviceType(tensor->ctx.device_type)
+  CHECK_EQ(old_t->device.device_type, tensor->device.device_type)
+      << "The input data must be on device \"" << GetStringFromDeviceType(old_t->device.device_type)
+      << "\", but user gave input on \"" << GetStringFromDeviceType(tensor->device.device_type)
       << "\"";
-  CHECK_EQ(old_t->ctx.device_id, tensor->ctx.device_id);
+  CHECK_EQ(old_t->device.device_id, tensor->device.device_id);
   for (auto i = 0; i < tensor->ndim; ++i) {
     CHECK_EQ(old_t->shape[i], tensor->shape[i]);
   }
-  tvm_graph_runtime_->SetInputZeroCopy(index, tensor);
+  tvm_graph_executor_->SetInputZeroCopy(index, tensor);
 }
 
 void TVMModel::GetInput(const char* name, void* input) {
@@ -275,11 +275,11 @@ void TVMModel::GetInput(const char* name, void* input) {
 #endif
 
   std::string str(name);
-  int index = tvm_graph_runtime_->GetInputIndex(str);
-  tvm::runtime::NDArray arr = tvm_graph_runtime_->GetInput(index);
+  int index = tvm_graph_executor_->GetInputIndex(str);
+  tvm::runtime::NDArray arr = tvm_graph_executor_->GetInput(index);
   DLTensor input_tensor;
   input_tensor.data = input;
-  input_tensor.ctx = DLContext{kDLCPU, 0};
+  input_tensor.device = DLDevice{kDLCPU, 0};
   input_tensor.ndim = arr->ndim;
   input_tensor.dtype = arr->dtype;
   input_tensor.shape = arr->shape;
@@ -306,7 +306,7 @@ void TVMModel::GetOutput(int index, void* out) {
   }
 #endif
   DLTensor output_tensor = *outputs_[index].operator->();
-  output_tensor.ctx = DLContext{kDLCPU, 0};
+  output_tensor.device = DLDevice{kDLCPU, 0};
   output_tensor.data = out;
   tvm::runtime::PackedFunc get_output = tvm_module_->GetFunction("get_output");
   get_output(index, &output_tensor);
@@ -319,9 +319,9 @@ const void* TVMModel::GetOutputPtr(int index) const {
   }
 #endif
 
-  tvm::runtime::NDArray output = tvm_graph_runtime_->GetOutput(index);
+  tvm::runtime::NDArray output = tvm_graph_executor_->GetOutput(index);
   const DLTensor* tensor = output.operator->();
-  if (tensor->ctx.device_type == kDLCPU) {
+  if (tensor->device.device_type == kDLCPU) {
     return tensor->data;
   }
   throw dmlc::Error("GetOutputPtr is not supported for non-CPU device types");
@@ -332,7 +332,7 @@ void TVMModel::GetOutputManagedTensorPtr(int index, const DLManagedTensor** out)
   CHECK(!(HasMetadata() && data_transform_.HasOutputTransform(metadata_, index)))
       << "Output transforms are not supported with GetOutputManagedTensor.";
 #endif
-  tvm::runtime::NDArray output = tvm_graph_runtime_->GetOutput(index);
+  tvm::runtime::NDArray output = tvm_graph_executor_->GetOutput(index);
   *out = output.ToDLPack();
 }
 


Test patch:
diff --git a/tests/cpp/dlr_data_transform_test.cc b/tests/cpp/dlr_data_transform_test.cc
index 9442a3557..831f56f45 100644
--- a/tests/cpp/dlr_data_transform_test.cc
+++ b/tests/cpp/dlr_data_transform_test.cc
@@ -45,14 +45,14 @@ TEST(DLR, DataTransformCategoricalString) {
   std::vector<int64_t> shape = {static_cast<int64_t>(std::strlen(data))};
   // Model input
   std::vector<DLDataType> dtypes = {DLDataType{kDLFloat, 32, 1}};
-  DLContext ctx = DLContext{kDLCPU, 0};
+  DLDevice dev = DLDevice{kDLCPU, 0};
   std::vector<tvm::runtime::NDArray> transformed_data(1);
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
   // Test that same buffer is reused.
   const void* buffer = transformed_data[0]->data;
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
   EXPECT_EQ(buffer, transformed_data[0]->data);
 
   std::vector<float> expected_output = {0, 1, 2, -1, -1, -1};
@@ -88,10 +88,10 @@ TEST(DLR, DataTransformCategoricalStringNumericColumn) {
   std::vector<int64_t> shape = {static_cast<int64_t>(std::strlen(data))};
   // Model input
   std::vector<DLDataType> dtypes = {DLDataType{kDLFloat, 32, 1}};
-  DLContext ctx = DLContext{kDLCPU, 0};
+  DLDevice dev = DLDevice{kDLCPU, 0};
   std::vector<tvm::runtime::NDArray> transformed_data(1);
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
   const float kNan = std::numeric_limits<float>::quiet_NaN();
   const float kInf = std::numeric_limits<float>::infinity();
   std::vector<float> expected_output = {2.345, 7, 7, -9.7, kNan, -kInf, kNan, kInf};
@@ -131,10 +131,10 @@ TEST(DLR, DataTransformMultipleColumn) {
   std::vector<int64_t> shape = {static_cast<int64_t>(std::strlen(data))};
   // Model input
   std::vector<DLDataType> dtypes = {DLDataType{kDLFloat, 32, 1}, DLDataType{kDLFloat, 32, 1}};
-  DLContext ctx = DLContext{kDLCPU, 0};
+  DLDevice dev = DLDevice{kDLCPU, 0};
   std::vector<tvm::runtime::NDArray> transformed_data(2);
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
   const float kNan = std::numeric_limits<float>::quiet_NaN();
   const float kInf = std::numeric_limits<float>::infinity();
   std::vector<float> expected_output_float = {2.345, kNan, 7, 7};
@@ -175,10 +175,10 @@ TEST(DLR, DataTransformDateTime) {
   std::vector<int64_t> shape = {static_cast<int64_t>(std::strlen(data))};
 
   std::vector<DLDataType> dtypes = {DLDataType{kDLFloat, 32, 1}};
-  DLContext ctx = DLContext{kDLCPU, 0};
+  DLDevice dev = DLDevice{kDLCPU, 0};
   std::vector<tvm::runtime::NDArray> transformed_data(1);
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
 
   EXPECT_EQ(transformed_data[0]->ndim, 2);
   EXPECT_EQ(transformed_data[0]->shape[0], 1);
@@ -209,7 +209,7 @@ TEST(DLR, DataTransformDateTime) {
   shape = {static_cast<int64_t>(std::strlen(data))};
 
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
 
   EXPECT_EQ(transformed_data[0]->ndim, 2);
   EXPECT_EQ(transformed_data[0]->shape[0], 4);
@@ -234,15 +234,15 @@ TEST(DLR, DataTransformText) {
         "Input": {
           "ColumnTransform": [
             {
-              "Type": "Text", 
+              "Type": "Text",
               "Vocabularies": ["cats", "chase", "dogs", "eat", "hate", "like", "people", "rats"],
               "TextCol": 1
             },
             {
-              "Type": "Text", 
+              "Type": "Text",
               "Vocabularies": ["are", "mammals", "rats"],
               "TextCol": 2
-            } 
+            }
           ]
         }
       }
@@ -258,10 +258,10 @@ TEST(DLR, DataTransformText) {
 
   std::vector<int64_t> shape = {static_cast<int64_t>(std::strlen(data))};
   std::vector<DLDataType> dtypes = {DLDataType{kDLFloat, 32, 1}, DLDataType{kDLFloat, 32, 1}};
-  DLContext ctx = DLContext{kDLCPU, 0};
+  DLDevice dev = {kDLCPU, 0};
   std::vector<tvm::runtime::NDArray> transformed_data(2);
   EXPECT_NO_THROW(transform.TransformInput(metadata, shape.data(), const_cast<char*>(data),
-                                           shape.size(), dtypes, ctx, &transformed_data));
+                                           shape.size(), dtypes, dev, &transformed_data));
 
   EXPECT_EQ(transformed_data[0]->ndim, 2);
   EXPECT_EQ(transformed_data[0]->shape[0], 4);
@@ -286,11 +286,11 @@ TEST(DLR, DataTransformText) {
   }
 }
 
-TEST(DLR, RelayVMDataTransformInput) {
-  DLContext ctx = {kDLCPU, 0};
+TEST(DLR, DISABLED_RelayVMDataTransformInput) {
+  DLDevice dev = {kDLCPU, 0};
   std::vector<std::string> paths = {"./automl"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, ctx);
+  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, dev);
   EXPECT_EQ(model->GetNumInputs(), 1);
   EXPECT_STREQ(model->GetInputType(0), "json");
   EXPECT_STREQ(model->GetInputName(0), "input");
@@ -326,10 +326,10 @@ TEST(DLR, RelayVMDataTransformInput) {
 }
 
 TEST(DLR, RelayVMDataTransformOutput) {
-  DLContext ctx = {kDLCPU, 0};
+  DLDevice dev = {kDLCPU, 0};
   std::vector<std::string> paths = {"./inverselabel"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, ctx);
+  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, dev);
 
   int64_t size;
   int dim;
@@ -369,10 +369,10 @@ TEST(DLR, RelayVMDataTransformOutput) {
 }
 
 TEST(DLR, TVMDataTransformInput) {
-  DLContext ctx = {kDLCPU, 0};
+  DLDevice dev = {kDLCPU, 0};
   std::vector<std::string> paths = {"./automl_static"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::TVMModel* model = new dlr::TVMModel(files, ctx);
+  dlr::TVMModel* model = new dlr::TVMModel(files, dev);
   EXPECT_EQ(model->GetNumInputs(), 1);
   EXPECT_STREQ(model->GetInputType(0), "json");
   EXPECT_STREQ(model->GetInputName(0), "input");
@@ -408,10 +408,10 @@ TEST(DLR, TVMDataTransformInput) {
 }
 
 TEST(DLR, TVMDataTransformOutput) {
-  DLContext ctx = {kDLCPU, 0};
+  DLDevice dev = {kDLCPU, 0};
   std::vector<std::string> paths = {"./inverselabel_static"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::TVMModel* model = new dlr::TVMModel(files, ctx);
+  dlr::TVMModel* model = new dlr::TVMModel(files, dev);
 
   int64_t size;
   int dim;
diff --git a/tests/cpp/dlr_elem_test.cc b/tests/cpp/dlr_elem_test.cc
index 630aace7f..47efc872d 100644
--- a/tests/cpp/dlr_elem_test.cc
+++ b/tests/cpp/dlr_elem_test.cc
@@ -98,9 +98,9 @@ TEST(DLR, TestGetDLRWeightName) {
   auto model = GetDLRModel();
   const char* weight_name;
   EXPECT_EQ(GetDLRWeightName(&model, 0, &weight_name), 0);
-  EXPECT_STREQ(weight_name, "p45");
+  EXPECT_STREQ(weight_name, "p105");
   EXPECT_EQ(GetDLRWeightName(&model, 107, &weight_name), 0);
-  EXPECT_STREQ(weight_name, "p72");
+  EXPECT_STREQ(weight_name, "p76");
   DeleteDLRModel(&model);
 }
 
diff --git a/tests/cpp/dlr_pipeline_skl_xgb_test.cc b/tests/cpp/dlr_pipeline_skl_xgb_test.cc
index 0797dff2e..f572e768b 100644
--- a/tests/cpp/dlr_pipeline_skl_xgb_test.cc
+++ b/tests/cpp/dlr_pipeline_skl_xgb_test.cc
@@ -156,5 +156,7 @@ int main(int argc, char** argv) {
 #ifndef _WIN32
   testing::FLAGS_gtest_death_test_style = "threadsafe";
 #endif  // _WIN32
+  // temporarily disable tests (TODO: generate new artifacts)
+  testing::GTEST_FLAG(filter) = "-PipelineTest.*";
   return RUN_ALL_TESTS();
 }
diff --git a/tests/cpp/dlr_relayvm_elem_test.cc b/tests/cpp/dlr_relayvm_elem_test.cc
index 617ba4ae4..f1a26910f 100644
--- a/tests/cpp/dlr_relayvm_elem_test.cc
+++ b/tests/cpp/dlr_relayvm_elem_test.cc
@@ -20,7 +20,7 @@ class RelayVMElemTest : public ::testing::Test {
   std::vector<int8_t> img{std::vector<int8_t>(img_size)};
   const int device_type = kDLCPU;
   const int device_id = 0;
-  DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+  DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
   std::string ro_file = "./ssd_mobilenet_v1/code.ro";
   std::string so_file = "./ssd_mobilenet_v1/compiled.so";
   std::string meta_file = "./ssd_mobilenet_v1/compiled.meta";
@@ -35,7 +35,7 @@ class RelayVMElemTest : public ::testing::Test {
         {DLRModelElemType::RELAY_EXEC, nullptr, code_data.data(), code_data.size()},
         {DLRModelElemType::TVM_LIB, so_file.c_str(), nullptr, 0},
         {DLRModelElemType::NEO_METADATA, nullptr, meta_str.c_str(), 0}};
-    model = new dlr::RelayVMModel(model_elems, ctx);
+    model = new dlr::RelayVMModel(model_elems, dev);
   }
 
   ~RelayVMElemTest() { delete model; }
@@ -52,7 +52,7 @@ TEST_F(RelayVMElemTest, TestCreateModel_LibTvmIsPointer) {
   EXPECT_THROW(
       {
         try {
-          new dlr::RelayVMModel(model_elems, ctx);
+          new dlr::RelayVMModel(model_elems, dev);
         } catch (const dmlc::Error& e) {
           EXPECT_STREQ(e.what(),
                        "Invalid RelayVM model element TVM_LIB. TVM_LIB must be a file path.");
@@ -70,7 +70,7 @@ TEST_F(RelayVMElemTest, TestCreateModel_MetadataIsMissing) {
   EXPECT_THROW(
       {
         try {
-          new dlr::RelayVMModel(model_elems, ctx);
+          new dlr::RelayVMModel(model_elems, dev);
         } catch (const dmlc::Error& e) {
           EXPECT_STREQ(
               e.what(),
diff --git a/tests/cpp/dlr_relayvm_test.cc b/tests/cpp/dlr_relayvm_test.cc
index 920f3dd9d..4e1baac73 100644
--- a/tests/cpp/dlr_relayvm_test.cc
+++ b/tests/cpp/dlr_relayvm_test.cc
@@ -27,10 +27,10 @@ class RelayVMTest : public ::testing::Test {
   RelayVMTest() {
     const int device_type = kDLCPU;
     const int device_id = 0;
-    DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+    DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
     std::vector<std::string> paths = {"./ssd_mobilenet_v1"};
     std::vector<std::string> files = dlr::FindFiles(paths);
-    model = new dlr::RelayVMModel(files, ctx);
+    model = new dlr::RelayVMModel(files, dev);
   }
 
   ~RelayVMTest() { delete model; }
@@ -141,10 +141,10 @@ TEST_F(RelayVMTest, TestGetOutput) {
 }
 
 TEST(DLR, TestRelayVMAllocatorDefault) {
-  DLContext ctx = {static_cast<DLDeviceType>(kDLCPU), 0};
+  DLDevice dev = {static_cast<DLDeviceType>(kDLCPU), 0};
   std::vector<std::string> paths = {"./ssd_mobilenet_v1"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, ctx);
+  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, dev);
 
   EXPECT_EQ(model->GetAllocatorType(), tvm::runtime::vm::AllocatorType::kPooled);
 
@@ -153,10 +153,10 @@ TEST(DLR, TestRelayVMAllocatorDefault) {
 
 TEST(DLR, TestRelayVMAllocatorEnvVar) {
   EXPECT_EQ(SetEnv("DLR_RELAYVM_ALLOCATOR", "naive"), 0);
-  DLContext ctx = {static_cast<DLDeviceType>(kDLCPU), 0};
+  DLDevice dev = {static_cast<DLDeviceType>(kDLCPU), 0};
   std::vector<std::string> paths = {"./ssd_mobilenet_v1"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, ctx);
+  dlr::RelayVMModel* model = new dlr::RelayVMModel(files, dev);
 
   EXPECT_EQ(model->GetAllocatorType(), tvm::runtime::vm::AllocatorType::kNaive);
 
@@ -165,7 +165,7 @@ TEST(DLR, TestRelayVMAllocatorEnvVar) {
 }
 
 TEST(DLR, TestRelayVMAllocatorFromMetadata) {
-  DLContext ctx = {static_cast<DLDeviceType>(kDLCPU), 0};
+  DLDevice dev = {static_cast<DLDeviceType>(kDLCPU), 0};
   std::string ro_file = "./ssd_mobilenet_v1/code.ro";
   std::string so_file = "./ssd_mobilenet_v1/compiled.so";
   std::string meta_file = "./ssd_mobilenet_v1/compiled.meta";
@@ -177,7 +177,7 @@ TEST(DLR, TestRelayVMAllocatorFromMetadata) {
       {DLRModelElemType::RELAY_EXEC, ro_file.c_str(), nullptr, 0},
       {DLRModelElemType::TVM_LIB, so_file.c_str(), nullptr, 0},
       {DLRModelElemType::NEO_METADATA, nullptr, meta_str.c_str(), meta_str.size()}};
-  dlr::RelayVMModel* model = new dlr::RelayVMModel(model_elems, ctx);
+  dlr::RelayVMModel* model = new dlr::RelayVMModel(model_elems, dev);
   EXPECT_EQ(model->GetAllocatorType(), tvm::runtime::vm::AllocatorType::kNaive);
   delete model;
 }
diff --git a/tests/cpp/dlr_test.cc b/tests/cpp/dlr_test.cc
index bae1eedc7..04075aa86 100644
--- a/tests/cpp/dlr_test.cc
+++ b/tests/cpp/dlr_test.cc
@@ -57,9 +57,9 @@ TEST(DLR, TestGetDLRWeightName) {
   auto model = GetDLRModel();
   const char* weight_name;
   EXPECT_EQ(GetDLRWeightName(&model, 0, &weight_name), 0);
-  EXPECT_STREQ(weight_name, "p45");
+  EXPECT_STREQ(weight_name, "p105");
   EXPECT_EQ(GetDLRWeightName(&model, 107, &weight_name), 0);
-  EXPECT_STREQ(weight_name, "p72");
+  EXPECT_STREQ(weight_name, "p76");
   DeleteDLRModel(&model);
 }
 
@@ -278,7 +278,7 @@ TEST(DLR, TestCreateFromPaths_RelayVM) {
   EXPECT_EQ(SetDLRInputTensor(&model, input_name, &input), 0);
 
   // run inference
-  EXPECT_EQ(RunDLRModel(&model), 0);
+  EXPECT_EQ(RunDLRModel(&model), 0) << DLRGetLastError();
 
   // output 3
   float output3_d[100];
diff --git a/tests/cpp/dlr_treelite_test.cc b/tests/cpp/dlr_treelite_test.cc
index bf4dadfa7..6f0827364 100644
--- a/tests/cpp/dlr_treelite_test.cc
+++ b/tests/cpp/dlr_treelite_test.cc
@@ -30,10 +30,10 @@ class TreeliteTest : public ::testing::Test {
     // Instantiate model
     const int device_type = 1;
     const int device_id = 0;
-    DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+    DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
     std::vector<std::string> paths = {"./xgboost_test"};
     std::vector<std::string> files = dlr::FindFiles(paths);
-    model = new dlr::TreeliteModel(files, ctx);
+    model = new dlr::TreeliteModel(files, dev);
   }
 
   ~TreeliteTest() { delete model; }
diff --git a/tests/cpp/dlr_tvm_elem_test.cc b/tests/cpp/dlr_tvm_elem_test.cc
index 22ead3d6c..f85d19f6d 100644
--- a/tests/cpp/dlr_tvm_elem_test.cc
+++ b/tests/cpp/dlr_tvm_elem_test.cc
@@ -17,7 +17,7 @@ class TVMElemTest : public ::testing::Test {
   const std::string meta_file = "./resnet_v1_5_50/compiled.meta";
   int device_type = 1;
   int device_id = 0;
-  DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+  DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
 
   dlr::TVMModel* model;
 
@@ -36,7 +36,7 @@ class TVMElemTest : public ::testing::Test {
     img = LoadImageAndPreprocess("cat224-3.txt", img_size, batch_size);
 
     // Instantiate model
-    model = new dlr::TVMModel(model_elems, ctx);
+    model = new dlr::TVMModel(model_elems, dev);
   }
 
   ~TVMElemTest() { delete model; }
@@ -53,7 +53,7 @@ TEST_F(TVMElemTest, TestCreateModel_LibTvmIsPointer) {
   EXPECT_THROW(
       {
         try {
-          new dlr::TVMModel(model_elems, ctx);
+          new dlr::TVMModel(model_elems, dev);
         } catch (const dmlc::Error& e) {
           EXPECT_STREQ(e.what(), "Invalid TVM model element TVM_LIB. TVM_LIB must be a file path.");
           throw;
@@ -70,7 +70,7 @@ TEST_F(TVMElemTest, TestCreateModel_GraphIsMissing) {
   EXPECT_THROW(
       {
         try {
-          new dlr::TVMModel(model_elems, ctx);
+          new dlr::TVMModel(model_elems, dev);
         } catch (const dmlc::Error& e) {
           EXPECT_STREQ(e.what(),
                        "Invalid TVM model. Must have TVM_GRAPH, TVM_PARAMS and TVM_LIB elements");
diff --git a/tests/cpp/dlr_tvm_test.cc b/tests/cpp/dlr_tvm_test.cc
index de67e27d0..253abb3a3 100644
--- a/tests/cpp/dlr_tvm_test.cc
+++ b/tests/cpp/dlr_tvm_test.cc
@@ -26,9 +26,9 @@ class TVMTest : public ::testing::Test {
     int device_type = 1;
     int device_id = 0;
     std::vector<std::string> paths = {model_path};
-    DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+    DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
     std::vector<std::string> files = dlr::FindFiles(paths);
-    model = new dlr::TVMModel(files, ctx);
+    model = new dlr::TVMModel(files, dev);
   }
 
   ~TVMTest() { delete model; }
@@ -55,10 +55,10 @@ TEST_F(TVMTest, TestGetInputDim) { EXPECT_EQ(model->GetInputDim(0), 4); }
 TEST(TVM, TestTvmModelApisWithOutputMetadata) {
   const int device_type = 1;  // 1 - kDLCPU
   const int device_id = 0;
-  DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+  DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
   std::vector<std::string> paths = {"./resnet_v1_5_50"};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::TVMModel model = dlr::TVMModel(files, ctx);
+  dlr::TVMModel model = dlr::TVMModel(files, dev);
 
   EXPECT_EQ(model.GetNumInputs(), 1);
   EXPECT_EQ(model.GetNumOutputs(), 2);
@@ -120,10 +120,10 @@ TEST(TVM, TestTvmModelApisWithoutMetadata) {
 
   const int device_type = 1;  // 1 - kDLCPU
   const int device_id = 0;
-  DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+  DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
   std::vector<std::string> paths = {model_path};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::TVMModel model = dlr::TVMModel(files, ctx);
+  dlr::TVMModel model = dlr::TVMModel(files, dev);
 
   EXPECT_EQ(model.GetNumInputs(), 1);
   EXPECT_EQ(model.GetNumOutputs(), 2);
@@ -157,10 +157,10 @@ TEST(TVM, TestTvmModelApisWithoutOutputInMetadata) {
 
   const int device_type = 1;  // 1 - kDLCPU
   const int device_id = 0;
-  DLContext ctx = {static_cast<DLDeviceType>(device_type), device_id};
+  DLDevice dev = {static_cast<DLDeviceType>(device_type), device_id};
   std::vector<std::string> paths = {model_path};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::TVMModel model = dlr::TVMModel(files, ctx);
+  dlr::TVMModel model = dlr::TVMModel(files, dev);
 
   EXPECT_EQ(model.GetNumInputs(), 1);
   EXPECT_EQ(model.GetNumOutputs(), 2);
@@ -194,10 +194,10 @@ TEST(TVM, TestTvmModelGetDeviceTypeFromMetadata) {
       << "{\"Requirements\": { \"TargetDevice\": \"ML_C4\", \"TargetDeviceType\": \"cpu\"}}\n";
   mocked_metadata.close();
 
-  DLContext ctx = {DLDeviceType::kDLCPU, 0};
+  DLDevice dev = {DLDeviceType::kDLCPU, 0};
   std::vector<std::string> paths = {model_path};
   std::vector<std::string> files = dlr::FindFiles(paths);
-  dlr::TVMModel model = dlr::TVMModel(files, ctx);
+  dlr::TVMModel model = dlr::TVMModel(files, dev);
 
   EXPECT_EQ(model.GetNumInputs(), 1);
   EXPECT_EQ(model.GetNumOutputs(), 2);
@@ -210,9 +210,9 @@ TEST(TVM, TestTvmModelGetDeviceTypeFromMetadata) {
       << "{\"Requirements\": { \"TargetDevice\": \"ML_C4\", \"TargetDeviceType\": \"gpu\"}}\n";
   mocked_metadata.close();
 
-  auto throw_error = [&files, &ctx]() {
+  auto throw_error = [&files, &dev]() {
     try {
-      dlr::TVMModel(files, ctx);
+      dlr::TVMModel(files, dev);
     } catch (const dmlc::Error& e) {
       EXPECT_STREQ(e.what(), "Compiled model requires device type \"gpu\" but user gave \"cpu\".");
       throw;
diff --git a/tests/cpp/dlsym/dlr_dlsym_test.cc b/tests/cpp/dlsym/dlr_dlsym_test.cc
index f564b30a7..fe0ee4bb3 100644
--- a/tests/cpp/dlsym/dlr_dlsym_test.cc
+++ b/tests/cpp/dlsym/dlr_dlsym_test.cc
@@ -117,9 +117,9 @@ TEST(DLR, TestGetDLRWeightName) {
   auto model = GetDLRModel();
   const char* weight_name;
   EXPECT_EQ(GetDLRWeightName(&model, 0, &weight_name), 0);
-  EXPECT_STREQ(weight_name, "p45");
+  EXPECT_STREQ(weight_name, "p105");
   EXPECT_EQ(GetDLRWeightName(&model, 107, &weight_name), 0);
-  EXPECT_STREQ(weight_name, "p72");
+  EXPECT_STREQ(weight_name, "p76");
   DeleteDLRModel(&model);
 }
 
diff --git a/tests/cpp/test_utils.hpp b/tests/cpp/test_utils.hpp
index 2c8360a17..cd0771d8d 100644
--- a/tests/cpp/test_utils.hpp
+++ b/tests/cpp/test_utils.hpp
@@ -51,7 +51,7 @@ DLTensor GetInputDLTensor(int ndim, int64_t* shape, const char* filename) {
   for (int i = 0; i < ndim; i++) img_size *= shape[i];
 
   DLTensor dltensor;
-  dltensor.ctx = {kDLCPU, 0};
+  dltensor.device = {kDLCPU, 0};
   dltensor.ndim = ndim;
   dltensor.shape = (int64_t*)alligned_malloc(ndim * sizeof(int64_t), 128);
   dltensor.strides = 0;
@@ -86,7 +86,7 @@ DLTensor GetEmptyDLTensor(int ndim, int64_t* shape, uint8_t dtype, uint8_t bits)
   for (int i = 0; i < ndim; i++) size *= shape[i];
 
   DLTensor dltensor;
-  dltensor.ctx = {kDLCPU, 0};
+  dltensor.device = {kDLCPU, 0};
   dltensor.ndim = ndim;
   dltensor.shape = (int64_t*)malloc(dltensor.ndim * sizeof(int64_t));
   dltensor.strides = 0;
diff --git a/tests/python/integration/load_and_run_treelite_model.py b/tests/python/integration/load_and_run_treelite_model.py
index a52a4b02b..a07683aa3 100644
--- a/tests/python/integration/load_and_run_treelite_model.py
+++ b/tests/python/integration/load_and_run_treelite_model.py
@@ -12,16 +12,20 @@ def _sparse_to_dense(csr_matrix):
     return out
 
 def test_mnist():
-    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-mnist')
+    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-mnist-1.10.0')
     data_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost', 'mnist.libsvm')
     model = DLRModel(model_path, 'cpu', 0)
 
-    X, _ = load_svmlight_file(data_file, zero_based=True)
+    X, _ = load_svmlight_file(data_file)
     print('Testing inference on XGBoost MNIST...')
-    assert model.run(_sparse_to_dense(X))[0] == 7.0
+    res = model.run(X.toarray())[0]
+    # TODO investigate why output shape size in (1,10) instead of (1,1)
+    # the model uses multi:softmax objective which outputs one class with the maximum probability
+    # assert res.shape == (1,1)
+    assert res[0][0] == 7.0
 
 def test_iris():
-    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-iris')
+    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-iris-1.10.0')
     data_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost', 'iris.libsvm')
     model = DLRModel(model_path, 'cpu', 0)
 
@@ -31,7 +35,7 @@ def test_iris():
     assert np.allclose(model.run(_sparse_to_dense(X))[0], expected)
 
 def test_letor():
-    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-letor')
+    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost-letor-1.10.0')
     data_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xgboost', 'letor.libsvm')
     model = DLRModel(model_path, 'cpu', 0)
 
@@ -50,7 +54,7 @@ def test_letor():
 
 if __name__ == '__main__':
     arch = get_arch()
-    model_names = ['xgboost-mnist', 'xgboost-iris', 'xgboost-letor']
+    model_names = ['xgboost-mnist-1.10.0', 'xgboost-iris-1.10.0', 'xgboost-letor-1.10.0']
     for model_name in model_names:
         get_models(model_name, arch, kind='treelite')
     test_mnist()