diff --git a/apps/cpp_rpc/rpc_env.cc b/apps/cpp_rpc/rpc_env.cc index fc19dd87b85e..607b2572c0ba 100644 --- a/apps/cpp_rpc/rpc_env.cc +++ b/apps/cpp_rpc/rpc_env.cc @@ -72,6 +72,13 @@ namespace runtime { */ void CleanDir(const std::string& dirname); +/*! + * \brief ListDir get the list of files in a directory + * \param dirname The root directory name + * \return vector Files in directory. + */ +std::vector ListDir(const std::string& dirname); + /*! * \brief buld a shared library if necessary * @@ -123,6 +130,15 @@ RPCEnv::RPCEnv(const std::string& wd) { *rv = this->GetPath(args[0]); }); + TVM_REGISTER_GLOBAL("tvm.rpc.server.listdir").set_body([this](TVMArgs args, TVMRetValue* rv) { + std::string dir = this->GetPath(args[0]); + std::ostringstream os; + for (auto d : ListDir(dir)) { + os << d << ","; + } + *rv = os.str(); + }); + TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module").set_body([this](TVMArgs args, TVMRetValue* rv) { std::string file_name = this->GetPath(args[0]); file_name = BuildSharedLibrary(file_name); diff --git a/python/tvm/rpc/client.py b/python/tvm/rpc/client.py index e984651c3eda..76a35ab4a73c 100644 --- a/python/tvm/rpc/client.py +++ b/python/tvm/rpc/client.py @@ -145,6 +145,23 @@ def remove(self, path): self._remote_funcs["remove"] = self.get_function("tvm.rpc.server.remove") self._remote_funcs["remove"](path) + def listdir(self, path): + """ls files from remote temp folder. + + Parameters + ---------- + path: str + The relative location to remote temp folder. + + Returns + ------- + dirs: str + The files in the given directory with split token ','. + """ + if "listdir" not in self._remote_funcs: + self._remote_funcs["listdir"] = self.get_function("tvm.rpc.server.listdir") + return self._remote_funcs["listdir"](path) + def load_module(self, path): """Load a remote module, the file need to be uploaded first.