current status, semi implemented micro session

Author: mutinifni

python/tvm/contrib/binutil.py

@@ -81,13 +81,13 @@ def tvm_relocate_binary(binary_name, text, data, bss):
 
 
 @register_func("tvm_read_binary_section")
-def tvm_read_binary_section(binary_name, section):
+def tvm_read_binary_section(binary, section):
     """Returns the contents of the specified section in the binary file
 
     Parameters
     ----------
-    binary_name : string 
-        name of the binary file
+    binary : bytearray
+        contents of the binary
 
     section : string
         type of section
@@ -98,10 +98,13 @@ def tvm_read_binary_section(binary_name, section):
         contents of the read section
     """
     tmp_dir = util.tempdir()
+    tmp_bin = tmp_dir.relpath("temp.bin")
     tmp_section = tmp_dir.relpath("tmp_section.bin")
+    with open(tmp_bin, "wb") as out_file:
+        out_file.write(bytes(binary))
     p1 = subprocess.Popen(["objcopy", "--dump-section",
                            "." + section + "=" + tmp_section,
-                           binary_name],
+                           tmp_bin],
                           stdout=subprocess.PIPE,
                           stderr=subprocess.STDOUT)
     (out, _) = p1.communicate()

src/runtime/micro/device/utvm_runtime.cc

@@ -12,8 +12,7 @@ UTVMTask task;
 void UTVMDone() {}
 
 // init stub
-int UTVMMain()
-{
+int UTVMMain() {
   task.func(task.args, task.arg_type_ids, *task.num_args);
   UTVMDone();
   return 0;

src/runtime/micro/low_level_device.h

@@ -3,8 +3,8 @@
  * \file low_level_device.h
  * \brief Abstract low-level micro device management
  */
-#ifndef TVM_RUNTIME_LOW_LEVEL_DEVICE_H_
-#define TVM_RUNTIME_LOW_LEVEL_DEVICE_H_
+#ifndef TVM_RUNTIME_MICRO_LOW_LEVEL_DEVICE_H_
+#define TVM_RUNTIME_MICRO_LOW_LEVEL_DEVICE_H_
 
 #include <cstddef>
 #include <memory>
@@ -64,6 +64,6 @@ const std::shared_ptr<LowLevelDevice> HostLowLevelDeviceCreate(size_t num_bytes)
  * \param port port of the OpenOCD server to connect to
  */
 const std::shared_ptr<LowLevelDevice> OpenOCDLowLevelDeviceCreate(int port);
-} // namespace runtime
-} // namespace tvm
-#endif // TVM_RUNTIME_LOW_LEVEL_DEVICE_H_
+}  // namespace runtime
+}  // namespace tvm
+#endif  // TVM_RUNTIME_MICRO_LOW_LEVEL_DEVICE_H_

src/runtime/micro/micro_common.cc

@@ -31,7 +31,7 @@ const char* SectionToString(SectionKind section) {
 
 void* GetSymbol(std::unordered_map<std::string, void*> symbol_map,
                 std::string name,
-                void* base_addr) {
+                const void* base_addr) {
   void* symbol_addr = symbol_map[name];
   return (void*)((uint8_t*) symbol_addr - (uint8_t*) base_addr);
 }
@@ -59,12 +59,12 @@ std::string RelocateBinarySections(std::string binary_name,
   return relocated_bin;
 }
 
-std::string ReadSection(std::string binary_name, SectionKind section) {
+std::string ReadSection(std::string binary, SectionKind section) {
   CHECK(section == kText || section == kData || section == kBss)
     << "ReadSection requires section to be one of text, data or bss.";
   const auto* f = Registry::Get("tvm_read_binary_section");
   CHECK(f != nullptr) << "Require tvm_read_binary_section to exist in registry";
-  std::string section_contents = (*f)(binary_name, SectionToString(section));
+  std::string section_contents = (*f)(binary, SectionToString(section));
   return section_contents;
 }
 

src/runtime/micro/micro_common.h

@@ -88,7 +88,7 @@ const char* SectionToString(SectionKind section);
  */
 void* GetSymbol(std::unordered_map<std::string, void*> symbol_map,
                 std::string name,
-                void* base_addr);
+                const void* base_addr);
 
 /*!
  * \brief links binary by repositioning section addresses
@@ -104,12 +104,12 @@ std::string RelocateBinarySections(std::string binary_name,
                                    void* bss);
 
 /*!
- * \brief reads section from binary file
- * \param binary_name input binary filename
+ * \brief reads section from binary
+ * \param binary input binary contents
  * \param section section type to be read
  * \return contents of the section
  */
-std::string ReadSection(std::string binary_name, SectionKind section);
+std::string ReadSection(std::string binary, SectionKind section);
 
 /*!
  * \brief finds size of the section in the binary

src/runtime/micro/micro_device_api.cc

@@ -29,13 +29,13 @@ class MicroDeviceAPI final : public DeviceAPI {
                        size_t alignment,
                        TVMType type_hint) final {
     // TODO: can make this a private member, but where to best init it?
-    MicroSession* session = MicroSession::Global();
+    std::shared_ptr<MicroSession> session = MicroSession::Global();
     void* alloc_ptr = session->AllocateInSection(kHeap, nbytes);
     return alloc_ptr;
   }
 
   void FreeDataSpace(TVMContext ctx, void* ptr) final {
-    MicroSession* session = MicroSession::Global();
+    std::shared_ptr<MicroSession> session = MicroSession::Global();
     session->FreeInSection(kHeap, ptr);
   }
 
@@ -48,7 +48,7 @@ class MicroDeviceAPI final : public DeviceAPI {
                       TVMContext ctx_to,
                       TVMType type_hint,
                       TVMStreamHandle stream) final {
-    MicroSession* session = MicroSession::Global();
+    std::shared_ptr<MicroSession> session = MicroSession::Global();
     uint8_t buffer[size];
     constexpr int micro_devtype = kDLMicroDev;
     std::tuple<int, int> type_from_to(ctx_from.device_type, ctx_to.device_type);
@@ -75,20 +75,20 @@ class MicroDeviceAPI final : public DeviceAPI {
     }
   }
 
-  // TODO: ignore this?
+  // TODO(): ignore this?
   void StreamSync(TVMContext ctx, TVMStreamHandle stream) final {
   }
 
   // TODO: what about ctx?
   void* AllocWorkspace(TVMContext ctx, size_t size, TVMType type_hint) final {
-    MicroSession* session = MicroSession::Global();
+    std::shared_ptr<MicroSession> session = MicroSession::Global();
     void* alloc_ptr = session->AllocateInSection(kWorkspace, size);
     return alloc_ptr;
   }
 
   // TODO: what about ctx?
   void FreeWorkspace(TVMContext ctx, void* data) final {
-    MicroSession* session = MicroSession::Global();
+    std::shared_ptr<MicroSession> session = MicroSession::Global();
     session->FreeInSection(kWorkspace, data);
   }
 
@@ -109,5 +109,5 @@ TVM_REGISTER_GLOBAL("device_api.micro_dev")
     DeviceAPI* ptr = MicroDeviceAPI::Global().get();
     *rv = static_cast<void*>(ptr);
   });
-} // namespace runtime
-} // namespace tvm
+}  // namespace runtime
+}  // namespace tvm

src/runtime/micro/micro_module.cc

@@ -11,6 +11,7 @@
 #include "micro_session.h"
 #include "low_level_device.h"
 #include "micro_common.h"
+#include "../pack_args.h"
 
 namespace tvm {
 namespace runtime {
@@ -32,14 +33,24 @@ class MicroModuleNode final : public ModuleNode {
    * \brief initializes module by establishing device connection and loads binary
    * \param binary name of the binary to be loaded
    */
-  void InitMicroModule(const std::string binary);
+  void InitMicroModule(const std::string binary) {
+    // TODO: if first MicroModule, then load init section in MicroSession 
+    session_ = MicroSession::Global();
+    // TODO: ensure low_level_device_ is initialized in MicroSession
+    lldevice_ = session_->low_level_device();
+    binary_ = binary;
+    LoadBinary();
+  }
 
   /*!
    * \brief runs selected function on the micro device
    * \param func name of the function to be run
+   * \param func_addr address of the function to be run
    * \param args type-erased arguments passed to the function
    */
-  void RunFunction(std::string func, TVMArgs args);
+  void RunFunction(std::string func, void* func_addr, TVMArgs args) {
+    session_->PushToExecQueue(func_addr, args);
+  }
 
  private:
   /*! \brief loaded module text start address */
@@ -49,7 +60,7 @@ class MicroModuleNode final : public ModuleNode {
   /*! \brief loaded module bss start address */
   void* bss_start_;
   /*! \brief size of module text section */
-  size_t code_size_;
+  size_t text_size_;
   /*! \brief size of module data section */
   size_t data_size_;
   /*! \brief size of module bss section */
@@ -61,7 +72,27 @@ class MicroModuleNode final : public ModuleNode {
   /*! \brief low-level device pointer */
   std::shared_ptr<LowLevelDevice> lldevice_;
   /*! \brief symbol map to addresses */
-  std::unordered_map<std::string, void*> symbol_map;
+  std::unordered_map<std::string, void*> symbol_map_;
+
+  void LoadBinary() {
+    text_size_ = GetSectionSize(binary_, kText);
+    data_size_ = GetSectionSize(binary_, kData);
+    bss_size_ = GetSectionSize(binary_, kBss);
+    text_start_ = session_->AllocateInSection(kText, text_size_);
+    data_start_ = session_->AllocateInSection(kData, data_size_);
+    bss_start_ = session_->AllocateInSection(kBss, bss_size_);
+    CHECK(text_start_ != nullptr && data_start_ != nullptr && bss_start_ != nullptr)
+      << "Not enough space to load module on device";
+    std::string relocated_bin = RelocateBinarySections(binary_, text_start_,
+                                                       data_start_, bss_start_);
+    std::string text_contents = ReadSection(relocated_bin, kText);
+    std::string data_contents = ReadSection(relocated_bin, kData);
+    std::string bss_contents = ReadSection(relocated_bin, kBss);
+    lldevice_->Write(text_start_, &text_contents[0], text_size_);
+    lldevice_->Write(data_start_, &data_contents[0], data_size_);
+    lldevice_->Write(bss_start_, &bss_contents[0], bss_size_);
+    symbol_map_ = GetSymbolMap(relocated_bin);
+  }
 };
 
 class MicroWrappedFunc {
@@ -74,7 +105,9 @@ class MicroWrappedFunc {
     func_addr_ = func_addr;
   }
 
-  void operator()(TVMArgs args, TVMRetValue* rv) const {
+  void operator()(TVMArgs args, TVMRetValue* rv, void** void_args) const {
+    // no return value yet, but may implement in the future
+    m_->RunFunction(func_name_, func_addr_, args);
   }
 
  private:
@@ -86,10 +119,20 @@ class MicroWrappedFunc {
   void* func_addr_;
 };
 
-// TODO: register module load function
+PackedFunc MicroModuleNode::GetFunction(
+    const std::string& name,
+    const std::shared_ptr<ModuleNode>& sptr_to_self) {
+  void* func_addr = GetSymbol(symbol_map_, name, lldevice_->base_addr());
+  MicroWrappedFunc f(this, name, func_addr);
+  return PackFuncVoidAddr(f, std::vector<TVMType>());
+}
+
 // register loadfile function to load module from Python frontend
 TVM_REGISTER_GLOBAL("module.loadfile_micro_dev")
 .set_body([](TVMArgs args, TVMRetValue* rv) {
+    std::shared_ptr<MicroModuleNode> n = std::make_shared<MicroModuleNode>();
+    n->InitMicroModule(args[0]);
+    *rv = runtime::Module(n);
   });
 }  // namespace runtime
 }  // namespace tvm