[WIP][µTVM] Add OpenOCD Low-Level Device (RISC-V Support) (#3756)

python/tvm/contrib/binutil.py

@@ -75,8 +75,23 @@ def tvm_callback_get_section_size(binary_path, section_name, toolchain_prefix):
         entry_size = int(tokens[1])
         if entry_name in sections_to_sum:
             section_size += entry_size
-    return section_size
 
+    # NOTE: For some reason, the size of the BSS section on the RISC-V
+    # GCC is sometimes reported to be smaller than it is, so we need to adjust
+    # for this.
+    if "riscv" in toolchain_prefix and section_name == 'bss':
+        # TODO(weberlo): Figure out why 32 is the minimum constant that works.
+        #
+        # The current hypothesis is that the last symbols in the ".bss" and
+        # ".sbss" sections may have size zero, since the symbols in these
+        # sections are uninitialized and there's no address that follows that
+        # would enforce a particular size.
+        #
+        # If this is the case, then 32 just happens to be a safe amount of
+        # padding for most cases, but symbols can be arbitrarily large, so this
+        # isn't bulletproof.
+        return section_size + 32
+    return section_size
 
 @register_func("tvm_callback_relocate_binary")
 def tvm_callback_relocate_binary(
@@ -169,6 +184,7 @@ def tvm_callback_relocate_binary(
         msg = "linking error using ld:\n"
         msg += py_str(out)
         raise RuntimeError(msg)
+
     with open(rel_obj_path, "rb") as f:
         rel_bin = bytearray(f.read())
     return rel_bin

python/tvm/micro/base.py

@@ -29,7 +29,7 @@
 from .._ffi.function import _init_api
 from .._ffi.libinfo import find_include_path
 
-SUPPORTED_DEVICE_TYPES = ["host"]
+SUPPORTED_DEVICE_TYPES = ["host", "openocd"]
 
 class Session:
     """MicroTVM Device Session
@@ -50,15 +50,22 @@ class Session:
     .. code-block:: python
 
       c_mod = ...  # some module generated with "c" as the target
-      device_type = "host"
-      with tvm.micro.Session(device_type) as sess:
-          sess.create_micro_mod(c_mod)
+      device_type = "openocd"
+      toolchain_prefix = "riscv64-unknown-elf-"
+      with tvm.micro.Session(device_type,
+                             toolchain_prefix,
+                             base_addr=0x10010000,
+                             server_addr="127.0.0.1",
+                             port=6666):
+          c_mod.export_library(lib_obj_path, fcompile=tvm.micro.cross_compiler(toolchain_prefix))
+          micro_mod = tvm.module.load(lib_obj_path, "micro_dev")
     """
 
-    def __init__(self, device_type, toolchain_prefix):
+    def __init__(self, device_type, toolchain_prefix, **kwargs):
         if device_type not in SUPPORTED_DEVICE_TYPES:
             raise RuntimeError("unknown micro device type \"{}\"".format(device_type))
         self._check_system()
+        self._check_args(device_type, kwargs)
 
         # First, find and compile runtime library.
         runtime_src_path = os.path.join(_get_micro_device_dir(), "utvm_runtime.c")
@@ -67,7 +74,11 @@ def __init__(self, device_type, toolchain_prefix):
         create_micro_lib(
             runtime_obj_path, runtime_src_path, toolchain_prefix, include_dev_lib_header=False)
 
-        self.module = _CreateSession(device_type, runtime_obj_path, toolchain_prefix)
+        base_addr = kwargs.get("base_addr", 0)
+        server_addr = kwargs.get("server_addr", "")
+        port = kwargs.get("port", 0)
+        self.module = _CreateSession(
+            device_type, runtime_obj_path, toolchain_prefix, base_addr, server_addr, port)
         self._enter = self.module["enter"]
         self._exit = self.module["exit"]
 
@@ -83,6 +94,15 @@ def _check_system(self):
         if sys.maxsize <= 2**32:
             raise RuntimeError("microTVM is currently only supported on 64-bit platforms")
 
+    def _check_args(self, device_type, args):
+        """Check if the given configuration is valid."""
+        if device_type == "host":
+            pass
+        elif device_type == "openocd":
+            assert "base_addr" in args
+            assert "server_addr" in args
+            assert "port" in args
+
     def __enter__(self):
         self._enter()
 
@@ -181,7 +201,9 @@ def replace_suffix(s, new_suffix):
     options = ["-I" + path for path in find_include_path()]
     options += ["-I{}".format(_get_micro_device_dir())]
     options += ["-fno-stack-protector"]
-    if sys.maxsize > 2**32 and sys.platform.startswith("linux"):
+    # TODO(weberlo): Don't rely on the toolchain prefix to identify if this is the host
+    # device.
+    if toolchain_prefix == "" and sys.maxsize > 2**32 and sys.platform.startswith("linux"):
         # Only add this option if the host is a 64-bit Linux.
         options += ["-mcmodel=large"]
     compile_cmd = "{}gcc".format(toolchain_prefix)

src/runtime/micro/low_level_device.h

@@ -26,6 +26,7 @@
 #define TVM_RUNTIME_MICRO_LOW_LEVEL_DEVICE_H_
 
 #include <memory>
+#include <string>
 
 #include "micro_common.h"
 
@@ -66,9 +67,21 @@ class LowLevelDevice {
    */
   virtual void Execute(DevBaseOffset func_offset, DevBaseOffset breakpoint) = 0;
 
+  // TODO(weberlo): Should we just give the device the *entire* memory layout
+  // decided by the session?
+
+  /*!
+   * \brief sets the offset of the top of the stack section
+   * \param stack_top offset of the stack top
+   */
+  virtual void SetStackTop(DevBaseOffset stack_top) {
+    LOG(FATAL) << "unimplemented";
+  }
+
   /*!
    * \brief convert from base offset to absolute address
    * \param offset base offset
+   * \return absolute address
    */
   DevPtr ToDevPtr(DevBaseOffset offset) {
     return DevPtr(base_addr() + offset.value());
@@ -77,6 +90,7 @@ class LowLevelDevice {
   /*!
    * \brief convert from absolute address to base offset
    * \param ptr absolute address
+   * \return base offset
    */
   DevBaseOffset ToDevOffset(DevPtr ptr) {
     return DevBaseOffset(ptr.value() - base_addr());
@@ -102,6 +116,14 @@ class LowLevelDevice {
  */
 const std::shared_ptr<LowLevelDevice> HostLowLevelDeviceCreate(size_t num_bytes);
 
+/*!
+ * \brief connect to OpenOCD and create an OpenOCD low-level device
+ * \param port port of the OpenOCD server to connect to
+ */
+const std::shared_ptr<LowLevelDevice> OpenOCDLowLevelDeviceCreate(std::uintptr_t base_addr,
+                                                                  const std::string& addr,
+                                                                  int port);
+
 }  // namespace runtime
 }  // namespace tvm
 #endif  // TVM_RUNTIME_MICRO_LOW_LEVEL_DEVICE_H_

src/runtime/micro/micro_common.cc

@@ -39,21 +39,21 @@ namespace runtime {
 size_t GetDefaultSectionSize(SectionKind kind) {
   switch (kind) {
     case SectionKind::kText:
-      return 0xF0000;
+      return 0xF000;
     case SectionKind::kRodata:
       return 0xF000;
     case SectionKind::kData:
       return 0xF00;
     case SectionKind::kBss:
       return 0xF00;
     case SectionKind::kArgs:
-      return 0xF00000;
+      return 0xF0000;
     case SectionKind::kStack:
       return 0xF000;
     case SectionKind::kHeap:
-      return 0xF000000;
+      return 0xF00000;
     case SectionKind::kWorkspace:
-      return 0xF000000;
+      return 0xF0000;
     default:
       LOG(FATAL) << "invalid section " << static_cast<size_t>(kind);
       return 0;

src/runtime/micro/micro_session.cc

@@ -20,17 +20,6 @@
 /*!
  *  Copyright (c) 2019 by Contributors
  * \file micro_session.cc
- * \brief session to manage multiple micro modules
- *
- * Each session consists of an interaction with a *single* logical device.
- * Within that interaction, multiple TVM modules can be loaded on the logical
- * device.
- *
- * Multiple sessions can exist simultaneously, but there is only ever one
- * *active* session. The idea of an active session mainly has implications for
- * the frontend, in that one must make a session active in order to allocate
- * new TVM objects on it. Aside from that, previously allocated objects can be
- * used even if the session which they belong to is not currently active.
  */
 
 #include <dmlc/thread_local.h>
@@ -86,25 +75,38 @@ MicroSession::~MicroSession() {
   for (size_t i = 0; i < static_cast<size_t>(SectionKind::kNumKinds); i++) {
     section_allocators_[i] = nullptr;
   }
-
   low_level_device_ = nullptr;
 }
 
 void MicroSession::CreateSession(const std::string& device_type,
                                  const std::string& binary_path,
-                                 const std::string& toolchain_prefix) {
+                                 const std::string& toolchain_prefix,
+                                 std::uintptr_t base_addr,
+                                 const std::string& server_addr,
+                                 int port) {
   // TODO(weberlo): make device type enum
+  toolchain_prefix_ = toolchain_prefix;
   if (device_type == "host") {
     low_level_device_ = HostLowLevelDeviceCreate(memory_size_);
+  } else if (device_type == "openocd") {
+    // TODO(weberlo): We need a better way of configuring devices.
+    low_level_device_ = OpenOCDLowLevelDeviceCreate(base_addr, server_addr, port);
   } else {
     LOG(FATAL) << "unsupported micro low-level device";
   }
+
   SetRuntimeBinaryPath(binary_path);
   CHECK(!runtime_binary_path_.empty()) << "uTVM runtime not initialized";
   runtime_bin_info_ = LoadBinary(runtime_binary_path_, /* patch_dylib_pointers */ false);
   utvm_main_symbol_ = low_level_device()->ToDevOffset(runtime_symbol_map()["UTVMMain"]);
   utvm_done_symbol_ = low_level_device()->ToDevOffset(runtime_symbol_map()["UTVMDone"]);
 
+  if (device_type == "openocd") {
+    // Set OpenOCD device's stack pointer.
+    auto stack_section = GetAllocator(SectionKind::kStack);
+    low_level_device_->SetStackTop(stack_section->max_end_offset());
+  }
+
   // Patch workspace pointers to the start of the workspace section.
   DevBaseOffset workspace_start_offset = GetAllocator(SectionKind::kWorkspace)->start_offset();
   DevBaseOffset workspace_end_offset = GetAllocator(SectionKind::kWorkspace)->max_end_offset();
@@ -143,6 +145,7 @@ void MicroSession::PushToExecQueue(DevBaseOffset func, const TVMArgs& args) {
   };
   // Write the task.
   DevSymbolWrite(runtime_symbol_map(), "task", task);
+
   low_level_device()->Execute(utvm_main_symbol_, utvm_done_symbol_);
   // Check if there was an error during execution.  If so, log it.
   CheckDeviceError();
@@ -299,7 +302,7 @@ BinaryInfo MicroSession::LoadBinary(const std::string& binary_path, bool patch_d
 
 void MicroSession::PatchImplHole(const SymbolMap& symbol_map, const std::string& func_name) {
   void* runtime_impl_addr = runtime_symbol_map()[func_name].cast_to<void*>();
-  std::stringstream func_name_underscore;
+  std::ostringstream func_name_underscore;
   func_name_underscore << func_name << "_";
   DevSymbolWrite(symbol_map, func_name_underscore.str(), runtime_impl_addr);
 }
@@ -309,7 +312,7 @@ void MicroSession::SetRuntimeBinaryPath(std::string path) {
 }
 
 std::string MicroSession::ReadString(DevBaseOffset str_offset) {
-  std::stringstream result;
+  std::ostringstream result;
   const size_t buf_size = 256;
   std::vector<char> buf(buf_size, 0);
   size_t i = buf_size;
@@ -372,8 +375,12 @@ TVM_REGISTER_GLOBAL("micro._CreateSession")
     const std::string& device_type = args[0];
     const std::string& binary_path = args[1];
     const std::string& toolchain_prefix = args[2];
+    uint64_t base_addr = args[3];
+    const std::string& server_addr = args[4];
+    int port = args[5];
     std::shared_ptr<MicroSession> session = std::make_shared<MicroSession>();
-    session->CreateSession(device_type, binary_path, toolchain_prefix);
+    session->CreateSession(
+        device_type, binary_path, toolchain_prefix, base_addr, server_addr, port);
     *rv = Module(session);
     });
 

src/runtime/micro/micro_session.h

@@ -20,6 +20,17 @@
 /*!
  *  Copyright (c) 2019 by Contributors
  * \file micro_session.h
+ * \brief session to manage multiple micro modules
+ *
+ * Each session consists of an interaction with a *single* logical device.
+ * Within that interaction, multiple TVM modules can be loaded on the logical
+ * device.
+ *
+ * Multiple sessions can exist simultaneously, but there is only ever one
+ * *active* session. The idea of an active session mainly has implications for
+ * the frontend, in that one must make a session active in order to allocate
+ * new TVM objects on it. Aside from that, previously allocated objects can be
+ * used even if the session which they belong to is not currently active.
  */
 #ifndef TVM_RUNTIME_MICRO_MICRO_SESSION_H_
 #define TVM_RUNTIME_MICRO_MICRO_SESSION_H_
@@ -82,7 +93,10 @@ class MicroSession : public ModuleNode {
    */
   void CreateSession(const std::string& device_type,
                      const std::string& binary_path,
-                     const std::string& toolchain_prefix);
+                     const std::string& toolchain_prefix,
+                     std::uintptr_t base_addr,
+                     const std::string& server_addr,
+                     int port);
 
   /*!
    * \brief ends the session by destructing the low-level device and its allocators