[REFACTOR][FFI] Cleanup PackedFunc related redirection (#17923)

Author: tqchen

apps/cpp_rpc/rpc_env.cc

@@ -128,10 +128,10 @@ RPCEnv::RPCEnv(const std::string& wd) {
 
   ffi::Function::SetGlobal(
       "tvm.rpc.server.workpath",
-      ffi::Function::FromUnpacked([this](const std::string& path) { return this->GetPath(path); }));
+      ffi::Function::FromTyped([this](const std::string& path) { return this->GetPath(path); }));
 
   ffi::Function::SetGlobal("tvm.rpc.server.listdir",
-                           ffi::Function::FromUnpacked([this](const std::string& path) {
+                           ffi::Function::FromTyped([this](const std::string& path) {
                              std::string dir = this->GetPath(path);
                              std::ostringstream os;
                              for (auto d : ListDir(dir)) {
@@ -141,15 +141,15 @@ RPCEnv::RPCEnv(const std::string& wd) {
                            }));
 
   ffi::Function::SetGlobal("tvm.rpc.server.load_module",
-                           ffi::Function::FromUnpacked([this](const std::string& path) {
+                           ffi::Function::FromTyped([this](const std::string& path) {
                              std::string file_name = this->GetPath(path);
                              file_name = BuildSharedLibrary(file_name);
                              LOG(INFO) << "Load module from " << file_name << " ...";
                              return Module::LoadFromFile(file_name, "");
                            }));
 
   ffi::Function::SetGlobal("tvm.rpc.server.download_linked_module",
-                           ffi::Function::FromUnpacked([this](const std::string& path) {
+                           ffi::Function::FromTyped([this](const std::string& path) {
                              std::string file_name = this->GetPath(path);
                              file_name = BuildSharedLibrary(file_name);
                              std::string bin;

apps/hexagon_launcher/launcher_core.cc

@@ -137,27 +137,25 @@ Model::Model(tvm::runtime::Module executor, tvm::runtime::Module module, std::st
   run = get_module_func(model_executor, "run");
 }
 
-const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) {
+const tvm::ffi::Function get_runtime_func(const std::string& name) {
   if (auto pf = tvm::ffi::Function::GetGlobal(name)) {
     return *pf;
   }
-  return tvm::runtime::PackedFunc();
+  return tvm::ffi::Function();
 }
 
-const tvm::runtime::PackedFunc get_module_func(tvm::runtime::Module module,
-                                               const std::string& name) {
+const tvm::ffi::Function get_module_func(tvm::runtime::Module module, const std::string& name) {
   return module.GetFunction(name, false);
 }
 
 void reset_device_api() {
-  const tvm::runtime::PackedFunc api = get_runtime_func("device_api.hexagon");
+  const tvm::ffi::Function api = get_runtime_func("device_api.hexagon");
   tvm::ffi::Function::SetGlobal("device_api.cpu", api, true);
 }
 
 tvm::runtime::Module load_module(const std::string& file_name) {
-  static const tvm::runtime::PackedFunc loader =
-      get_runtime_func("runtime.module.loadfile_hexagon");
-  tvm::runtime::TVMRetValue rv = loader(file_name);
+  static const tvm::ffi::Function loader = get_runtime_func("runtime.module.loadfile_hexagon");
+  tvm::ffi::Any rv = loader(file_name);
   if (rv.type_code() == kTVMModuleHandle) {
     ICHECK_EQ(rv.type_code(), kTVMModuleHandle)
         << __func__ << ": loaded " << file_name << ", but did not get module handle";
@@ -180,26 +178,26 @@ tvm::runtime::Module create_graph_executor(const std::string& graph_json,
                                            tvm::runtime::Module graph_module, tvm::Device device) {
   std::string launcher_name = "tvm.graph_executor.create";
 
-  const tvm::runtime::PackedFunc create_executor = get_runtime_func(launcher_name);
+  const tvm::ffi::Function create_executor = get_runtime_func(launcher_name);
   uint64_t device_type = device.device_type;
   uint64_t device_id = device.device_id;
 
   if (graph_json.empty()) {
     LOG(ERROR) << __func__ << ": graph executor requires graph JSON";
     return tvm::runtime::Module();
   }
-  tvm::runtime::TVMRetValue rv = create_executor(graph_json, graph_module, device_type, device_id);
+  tvm::ffi::Any rv = create_executor(graph_json, graph_module, device_type, device_id);
   return rv.operator tvm::runtime::Module();
 }
 
 tvm::runtime::Module create_aot_executor(tvm::runtime::Module factory_module, tvm::Device device) {
-  tvm::runtime::PackedFunc list_modules = get_module_func(factory_module, "list_module_names");
+  tvm::ffi::Function list_modules = get_module_func(factory_module, "list_module_names");
   tvm::Array<tvm::String> module_names = list_modules();
   if (module_names.size() != 1) {
     LOG(WARNING) << __func__ << ": expecting single module, got: " << module_names << ", using "
                  << module_names[0];
   }
-  tvm::runtime::PackedFunc f = get_module_func(factory_module, module_names[0]);
+  tvm::ffi::Function f = get_module_func(factory_module, module_names[0]);
   if (f.get() == nullptr) {
     LOG(ERROR) << __func__ << ": failed to obtain function " << module_names[0];
     return tvm::runtime::Module();

apps/hexagon_launcher/launcher_core.h

@@ -90,7 +90,7 @@ struct Model {
   static tvm::Device device() { return tvm::Device{static_cast<DLDeviceType>(kDLHexagon), 0}; }
   static tvm::Device external() { return tvm::Device{static_cast<DLDeviceType>(kDLCPU), 0}; }
 
-  tvm::runtime::PackedFunc run;
+  tvm::ffi::Function run;
 };
 
 struct ExecutionSession {
@@ -123,9 +123,8 @@ void reset_device_api();
 
 tvm::runtime::Module load_module(const std::string& file_name);
 
-const tvm::runtime::PackedFunc get_runtime_func(const std::string& name);
-const tvm::runtime::PackedFunc get_module_func(tvm::runtime::Module module,
-                                               const std::string& name);
+const tvm::ffi::Function get_runtime_func(const std::string& name);
+const tvm::ffi::Function get_module_func(tvm::runtime::Module module, const std::string& name);
 
 tvm::runtime::Module create_aot_executor(tvm::runtime::Module factory_module, tvm::Device device);
 tvm::runtime::Module create_graph_executor(const std::string& graph_json,

apps/hexagon_launcher/launcher_hexagon.cc

@@ -47,15 +47,15 @@ static AEEResult error_too_small(const std::string& func_name, const std::string
 int __QAIC_HEADER(launcher_rpc_open)(const char* uri, remote_handle64* handle) {
   *handle = 0;  // Just use any value.
   reset_device_api();
-  static const tvm::runtime::PackedFunc acq_res =
+  static const tvm::ffi::Function acq_res =
       get_runtime_func("device_api.hexagon.acquire_resources");
   acq_res();
   return AEE_SUCCESS;
 }
 
 int __QAIC_HEADER(launcher_rpc_close)(remote_handle64 handle) {
   // Comment to stop clang-format from single-lining this function.
-  static const tvm::runtime::PackedFunc rel_res =
+  static const tvm::ffi::Function rel_res =
       get_runtime_func("device_api.hexagon.release_resources");
   rel_res();
   return AEE_SUCCESS;
@@ -104,8 +104,7 @@ AEEResult __QAIC_HEADER(launcher_rpc_get_num_inputs)(remote_handle64 handle, int
     return AEE_EBADSTATE;
   }
 
-  tvm::runtime::PackedFunc get_num_inputs =
-      get_module_func(TheModel->model_executor, "get_num_inputs");
+  tvm::ffi::Function get_num_inputs = get_module_func(TheModel->model_executor, "get_num_inputs");
   *num_inputs = get_num_inputs();
   return AEE_SUCCESS;
 }
@@ -140,7 +139,7 @@ AEEResult __QAIC_HEADER(launcher_rpc_set_input)(remote_handle64 handle, int inpu
 
   auto input = tvm::runtime::NDArray::FromDLPack(&managed);
 
-  tvm::runtime::PackedFunc set_input = get_module_func(TheModel->model_executor, "set_input");
+  tvm::ffi::Function set_input = get_module_func(TheModel->model_executor, "set_input");
   set_input(input_idx, input);
 
   return AEE_SUCCESS;
@@ -152,8 +151,7 @@ AEEResult __QAIC_HEADER(launcher_rpc_get_num_outputs)(remote_handle64 handle, in
     return AEE_EBADSTATE;
   }
 
-  tvm::runtime::PackedFunc get_num_outputs =
-      get_module_func(TheModel->model_executor, "get_num_outputs");
+  tvm::ffi::Function get_num_outputs = get_module_func(TheModel->model_executor, "get_num_outputs");
   *num_outputs = get_num_outputs();
   return AEE_SUCCESS;
 }
@@ -173,7 +171,7 @@ AEEResult __QAIC_HEADER(launcher_rpc_get_output)(remote_handle64 handle, int out
     return AEE_EBADPARM;
   }
 
-  tvm::runtime::PackedFunc get_output = get_module_func(TheModel->model_executor, "get_output");
+  tvm::ffi::Function get_output = get_module_func(TheModel->model_executor, "get_output");
   tvm::runtime::NDArray output = get_output(output_idx);
 
   std::vector<int64_t> shape_vec{output->shape, output->shape + output->ndim};

apps/ios_rpc/tvmrpc/RPCServer.mm

@@ -51,7 +51,7 @@
  *     2: need to write
  *     0: shutdown
  */
-using FEventHandler = PackedFunc;
+using FEventHandler = ffi::Function;
 
 /*!
  * \brief Create a server event handler.
@@ -68,7 +68,7 @@ FEventHandler CreateServerEventHandler(NSOutputStream* outputStream, std::string
       << "You are using tvm_runtime module built without RPC support. "
       << "Please rebuild it with USE_RPC flag.";
 
-  PackedFunc writer_func([outputStream](TVMArgs args, TVMRetValue* rv) {
+  ffi::Function writer_func([outputStream](ffi::PackedArgs args, ffi::Any* rv) {
     TVMByteArray* data = args[0].ptr<TVMByteArray>();
     int64_t nbytes = [outputStream write:reinterpret_cast<const uint8_t*>(data->data)
                                maxLength:data->size];

apps/ios_rpc/tvmrpc/TVMRuntime.mm

@@ -51,14 +51,15 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
 
 }  // namespace detail
 
-TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
-  static const std::string base_ = NSTemporaryDirectory().UTF8String;
-  const auto path = args[0].cast<std::string>();
-  *rv = base_ + "/" + path;
-});
+TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath")
+    .set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
+      static const std::string base_ = NSTemporaryDirectory().UTF8String;
+      const auto path = args[0].cast<std::string>();
+      *rv = base_ + "/" + path;
+    });
 
 TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module")
-    .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+    .set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
       auto name = args[0].cast<std::string>();
       std::string fmt = GetFileFormat(name, "");
       NSString* base;
@@ -109,7 +110,7 @@ void Init(const std::string& name) {
 
 // Add UnsignedDSOLoader plugin in global registry
 TVM_REGISTER_GLOBAL("runtime.module.loadfile_dylib_custom")
-    .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+    .set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
       auto n = make_object<UnsignedDSOLoader>();
       n->Init(args[0]);
       *rv = CreateModuleFromLibrary(n);

docs/arch/index.rst

@@ -198,7 +198,7 @@ tvm/runtime
 The runtime serves as the foundation of the TVM stack. It provides the mechanism to load and execute compiled artifacts.
 The runtime defines a stable standard set of C APIs to interface with frontend languages such as Python and Rust.
 
-`runtime::Object` is one of the primary data structures in TVM runtime besides the `runtime::PackedFunc`.
+`runtime::Object` is one of the primary data structures in TVM runtime besides the `ffi::Function`.
 It is a reference-counted base class with a type index to support runtime type checking and downcasting.
 The object system allows the developer to introduce new data structures to the runtime, such as Array, Map, and new IR data structures.
 

docs/arch/runtime.rst

@@ -54,7 +54,7 @@ The following code block provides an example in C++
 
     #include <tvm/runtime/packed_func.h>
 
-    void MyAdd(TVMArgs args, TVMRetValue* rv) {
+    void MyAdd(ffi::PackedArgs args, ffi::Any* rv) {
       // automatically convert arguments to desired type.
       int a = args[0].cast<int>();
       int b = args[1].cast<int>();
@@ -71,8 +71,8 @@ The following code block provides an example in C++
 In the above codeblock, we defined a PackedFunc MyAdd. It takes two arguments
 : ``args`` represents input arguments and ``rv`` represents return value.
 The function is type-erased, which means that the function signature does not restrict which input type to pass in or type to return.
-Under the hood, when we call a PackedFunc, it packs the input arguments to TVMArgs on stack,
-and gets the result back via TVMRetValue.
+Under the hood, when we call a PackedFunc, it packs the input arguments to ffi::PackedArgs on stack,
+and gets the result back via ffi::Any.
 
 Thanks to template tricks in C++, we can call a PackedFunc just like a normal function. Because of its type-erased nature, we can call a PackedFunc from dynamic languages like python, without additional glue code for each new type function created.
 The following example registers PackedFunc in C++ and calls from python.
@@ -91,7 +91,7 @@ The following example registers PackedFunc in C++ and calls from python.
     # prints 3
     print(myadd(1, 2))
 
-Most of the magic of PackedFunc lies in ``TVMArgs`` and ``TVMRetValue`` structure.
+Most of the magic of PackedFunc lies in ``ffi::PackedArgs`` and ``ffi::Any`` structure.
 We restrict a list of possible types which can be passed.
 Here are the common ones:
 
@@ -111,7 +111,7 @@ we can pass functions from python (as PackedFunc) to C++.
 .. code:: c
 
     TVM_REGISTER_GLOBAL("callhello")
-    .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+    .set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
       PackedFunc f = args[0];
       f("hello world");
     });