[RPC] Expose module handle (#459)

* [RPC] Expose module handle * not include handle

[RPC] Expose module handle (#459)
* [RPC] Expose module handle * not include handle
d6007a24 · Tianqi Chen · GitHub · fe564d90 · d6007a24
Commit d6007a24 authored Sep 17, 2017 by Tianqi Chen Committed by GitHub Sep 17, 2017
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 0 deletions

src/runtime/rpc/rpc_module.cc
+8 -0

No files found.
--- a/src/runtime/rpc/rpc_module.cc
+++ b/src/runtime/rpc/rpc_module.cc
@@ -182,6 +182,14 @@ TVM_REGISTER_GLOBAL("contrib.rpc._ImportRemoteModule")
                             cmod->module_handle());
  });

+TVM_REGISTER_GLOBAL("contrib.rpc._ModuleHandle")
+.set_body([](TVMArgs args, TVMRetValue* rv) {
+    Module m = args[0];
+    std::string tkey = m->type_key();
+    CHECK_EQ(tkey, "rpc");
+    *rv = static_cast<RPCModuleNode*>(m.operator->())->module_handle();
+  });
+
 TVM_REGISTER_GLOBAL("contrib.rpc._SessTableIndex")
 .set_body([](TVMArgs args, TVMRetValue* rv) {
    Module m = args[0];