/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * \file micro_module.cc */ #include <tvm/runtime/registry.h> #include <tvm/runtime/c_runtime_api.h> #include <tvm/runtime/module.h> #include <unordered_map> #include <string> #include "micro_session.h" #include "low_level_device.h" #include "micro_common.h" #include "../pack_args.h" namespace tvm { namespace runtime { /*! * \brief module for uTVM micro devices */ class MicroModuleNode final : public ModuleNode { public: MicroModuleNode() {} ~MicroModuleNode() {} const char* type_key() const final { return "micro"; } PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final; /*! * \brief initializes module by establishing device connection and loads binary * \param binary_path path of the binary to be loaded */ void InitMicroModule(const std::string& binary_path) { session_ = MicroSession::Current(); symbol_map_ = session_->LoadBinary(binary_path, true).symbol_map; } private: SymbolMap symbol_map_; /*! \brief global session pointer */ ObjectPtr<MicroSession> session_; }; class MicroWrappedFunc { public: MicroWrappedFunc(ObjectPtr<MicroSession> session, DevPtr func_ptr) { session_ = session; func_ptr_ = func_ptr; } void operator()(TVMArgs args, TVMRetValue* rv) const { *rv = session_->PushToExecQueue(func_ptr_, args); } private: /*! \brief reference to the session for this function (to keep the session alive) */ ObjectPtr<MicroSession> session_; /*! \brief offset of the function to be called */ DevPtr func_ptr_; }; PackedFunc MicroModuleNode::GetFunction( const std::string& name, const ObjectPtr<Object>& sptr_to_self) { DevPtr func_ptr; if (name == tvm::runtime::symbol::tvm_module_main) { if (symbol_map_.HasSymbol(tvm::runtime::symbol::tvm_module_main)) { func_ptr = symbol_map_[tvm::runtime::symbol::tvm_module_main]; } else { func_ptr = symbol_map_["default_function"]; } } else { func_ptr = symbol_map_[name]; } MicroWrappedFunc f(session_, func_ptr); return PackedFunc(f); } // register loadfile function to load module from Python frontend TVM_REGISTER_GLOBAL("runtime.module.loadfile_micro_dev") .set_body([](TVMArgs args, TVMRetValue* rv) { auto n = make_object<MicroModuleNode>(); n->InitMicroModule(args[0]); *rv = runtime::Module(n); }); } // namespace runtime } // namespace tvm