[Fix] Keep update-to-date with upstream API change (#1209)

mlc-ai · Nov 6, 2023 · e2c99a8 · e2c99a8
1 parent a7f1183
commit e2c99a8
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/cpp/image_embed.cc b/cpp/image_embed.cc
@@ -9,10 +9,10 @@
 #include "image_embed.h"
 
 #include <picojson.h>
+#include <tvm/runtime/memory/memory_manager.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/registry.h>
-#include <tvm/runtime/relax_vm/memory_manager.h>
 
 #include <cctype>
 #include <chrono>
@@ -59,9 +59,9 @@ class LLMImage {
     ICHECK(fload_exec.defined()) << "TVM runtime cannot find vm_load_executable";
     vm_ = fload_exec();
     vm_->GetFunction("vm_initialization")(static_cast<int>(device_.device_type), device_.device_id,
-                                          static_cast<int>(relax_vm::AllocatorType::kPooled),
+                                          static_cast<int>(memory::AllocatorType::kPooled),
                                           static_cast<int>(kDLCPU), 0,
-                                          static_cast<int>(relax_vm::AllocatorType::kPooled));
+                                          static_cast<int>(memory::AllocatorType::kPooled));
 
     embed_func_ = vm_->GetFunction("embed");
 

diff --git a/cpp/llm_chat.cc b/cpp/llm_chat.cc
@@ -12,11 +12,11 @@
 #include <tokenizers_cpp.h>
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/disco/session.h>
+#include <tvm/runtime/memory/memory_manager.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/registry.h>
-#include <tvm/runtime/relax_vm/memory_manager.h>
 
 #include <cctype>
 #include <chrono>
@@ -159,8 +159,8 @@ struct FunctionTable {
       this->local_vm = fload_exec();
       this->local_vm->GetFunction("vm_initialization")(
           static_cast<int>(device.device_type), device.device_id,
-          static_cast<int>(relax_vm::AllocatorType::kPooled), static_cast<int>(kDLCPU), 0,
-          static_cast<int>(relax_vm::AllocatorType::kPooled));
+          static_cast<int>(memory::AllocatorType::kPooled), static_cast<int>(kDLCPU), 0,
+          static_cast<int>(memory::AllocatorType::kPooled));
       this->mod_get_func = [this](const std::string& name) -> PackedFunc {
         PackedFunc func = this->local_vm->GetFunction(name, false);
         return func;