Cleanup TargetPtr and word size representations.

python/tvm/micro/base.py

@@ -77,7 +77,7 @@ def __init__(self, config):
         dev_funcs = tvm.micro.device.get_device_funcs(config['device_id'])
         self.toolchain_prefix = config['toolchain_prefix']
         self.mem_layout = config['mem_layout']
-        self.word_size = config['word_size']
+        self.word_size_bits = config['word_size_bits']
         self.thumb_mode = config['thumb_mode']
         self.use_device_timer = config['use_device_timer']
         self.comms_method = config['comms_method']
@@ -122,7 +122,7 @@ def __init__(self, config):
             self.mem_layout['workspace']['size'],
             self.mem_layout['stack'].get('start', 0),
             self.mem_layout['stack']['size'],
-            self.word_size,
+            self.word_size_bits,
             self.thumb_mode,
             self.use_device_timer,
             server_addr,

python/tvm/micro/device/arm/stm32f746xx.py

@@ -20,7 +20,7 @@
 
 DEVICE_ID = 'arm.stm32f746xx'
 TOOLCHAIN_PREFIX = 'arm-none-eabi-'
-WORD_SIZE = 4
+WORD_SIZE_BITS = 32
 #
 # [Device Memory Layout]
 #   RAM   (rwx) : START = 0x20000000, LENGTH = 320K
@@ -112,8 +112,8 @@ def generate_config(server_addr, server_port, section_constraints=None):
     return {
         'device_id': DEVICE_ID,
         'toolchain_prefix': TOOLCHAIN_PREFIX,
-        'mem_layout': gen_mem_layout(BASE_ADDR, AVAILABLE_MEM, WORD_SIZE, section_constraints),
-        'word_size': WORD_SIZE,
+        'mem_layout': gen_mem_layout(BASE_ADDR, AVAILABLE_MEM, WORD_SIZE_BITS, section_constraints),
+        'word_size_bits': WORD_SIZE_BITS,
         'thumb_mode': True,
         'use_device_timer': True,
         'comms_method': 'openocd',

python/tvm/micro/device/base.py

@@ -178,7 +178,7 @@ class MemConstraint(enum.Enum):
     WEIGHT = 1
 
 
-def gen_mem_layout(base_addr, available_mem, word_size, section_constraints):
+def gen_mem_layout(base_addr, available_mem, word_size_bits, section_constraints):
     """Template function to generate memory layout for devices.
 
     Parameters
@@ -189,12 +189,14 @@ def gen_mem_layout(base_addr, available_mem, word_size, section_constraints):
     available_mem: Number
         Available memory at base_addr, given in bytes.
 
-    word_size: Number
-        Number of bytes in one word on this device.
+    word_size_bits: Number
+        Number of bits in one word on this device.
 
     section_constraints: Optional[Dict[str, [Number, MemConstraint]]]
         maps section name to the quantity of available memory
     """
+    assert word_size_bits in (32, 64), 'only 32- or 64-bit devices are supported now'
+    word_size_bytes = word_size_bits // 8
     byte_sum = sum(x[0]
                    for x in section_constraints.values()
                    if x[1] == MemConstraint.ABSOLUTE_BYTES)
@@ -209,7 +211,7 @@ def gen_mem_layout(base_addr, available_mem, word_size, section_constraints):
     for section in DEVICE_SECTIONS:
         (val, cons_type) = section_constraints[section]
         if cons_type == MemConstraint.ABSOLUTE_BYTES:
-            assert val % word_size == 0, \
+            assert val % word_size_bytes == 0, \
                 f'constraint {val} for {section} section is not word-aligned'
             size = val
             res[section] = {
@@ -218,7 +220,7 @@ def gen_mem_layout(base_addr, available_mem, word_size, section_constraints):
             }
         else:
             size = int((val / weight_sum) * available_weight_mem)
-            size = (size // word_size) * word_size
+            size = (size // word_size_bytes) * word_size_bytes
             res[section] = {
                 'start': curr_addr,
                 'size': size,

python/tvm/micro/device/host.py

@@ -21,7 +21,7 @@
 
 DEVICE_ID = 'host'
 TOOLCHAIN_PREFIX = ''
-WORD_SIZE = 8 if sys.maxsize > 2**32 else 4
+WORD_SIZE_BITS = 64 if sys.maxsize > 2**32 else 32
 
 # we pretend we only have 320kb in the default case, so we can use `gen_mem_layout`
 DEFAULT_AVAILABLE_MEM = 3200000
@@ -89,7 +89,7 @@ def generate_config(available_mem=None, section_constraints=None):
         available_mem = DEFAULT_AVAILABLE_MEM
     if section_constraints is None:
         section_constraints = DEFAULT_SECTION_CONSTRAINTS
-    mem_layout = gen_mem_layout(0, available_mem, WORD_SIZE, section_constraints)
+    mem_layout = gen_mem_layout(0, available_mem, WORD_SIZE_BITS, section_constraints)
     # TODO the host emulated device is an outlier, since we don't know how what
     # its base address will be until we've created it in the C++. is there any
     # way to change the infrastructure around this so it's not so much of an
@@ -103,7 +103,7 @@ def generate_config(available_mem=None, section_constraints=None):
         'device_id': DEVICE_ID,
         'toolchain_prefix': TOOLCHAIN_PREFIX,
         'mem_layout': mem_layout,
-        'word_size': WORD_SIZE,
+        'word_size_bits': WORD_SIZE_BITS,
         'thumb_mode': False,
         'use_device_timer': False,
         'comms_method': 'host',

python/tvm/micro/device/riscv_spike.py

@@ -20,7 +20,7 @@
 
 DEVICE_ID = 'riscv_spike'
 TOOLCHAIN_PREFIX = 'riscv64-unknown-elf-'
-WORD_SIZE = 8
+WORD_SIZE_BITS = 64
 
 DEFAULT_SECTION_CONSTRAINTS = {
     'text': (18000, MemConstraint.ABSOLUTE_BYTES),
@@ -92,8 +92,8 @@ def generate_config(base_addr, available_mem, server_addr, server_port, section_
     return {
         'device_id': DEVICE_ID,
         'toolchain_prefix': TOOLCHAIN_PREFIX,
-        'mem_layout': gen_mem_layout(base_addr, available_mem, WORD_SIZE, section_constraints),
-        'word_size': WORD_SIZE,
+        'mem_layout': gen_mem_layout(base_addr, available_mem, WORD_SIZE_BITS, section_constraints),
+        'word_size_bits': WORD_SIZE_BITS,
         'thumb_mode': False,
         'use_device_timer': False,
         'comms_method': 'openocd',

src/runtime/micro/host_low_level_device.cc

@@ -43,14 +43,15 @@ class HostLowLevelDevice final : public LowLevelDevice {
    * \brief constructor to initialize on-host memory region to act as device
    * \param num_bytes size of the emulated on-device memory region
    */
-  explicit HostLowLevelDevice(size_t num_bytes, void** base_addr) : size_(num_bytes) {
+  explicit HostLowLevelDevice(size_t num_bytes, TargetPtr* base_addr) : size_(num_bytes) {
     size_t size_in_pages = (num_bytes + kPageSize - 1) / kPageSize;
     // TODO(weberlo): Set permissions per section (e.g., read-write perms for
     // the heap, execute perms for text, etc.).
     int mmap_prot = PROT_READ | PROT_WRITE | PROT_EXEC;
     int mmap_flags = MAP_ANONYMOUS | MAP_PRIVATE;
     base_addr_ = mmap(nullptr, size_in_pages * kPageSize, mmap_prot, mmap_flags, -1, 0);
-    *base_addr = base_addr_;
+    *base_addr = TargetPtr(TargetWordSize(sizeof(size_t) * 8),
+                           reinterpret_cast<uint64_t>(base_addr_));
   }
 
   /*!
@@ -83,9 +84,9 @@ class HostLowLevelDevice final : public LowLevelDevice {
   size_t size_;
 };
 
-const std::shared_ptr<LowLevelDevice> HostLowLevelDeviceCreate(size_t num_bytes, void** base_addr) {
-  std::shared_ptr<LowLevelDevice> lld =
-      std::make_shared<HostLowLevelDevice>(num_bytes, base_addr);
+const std::shared_ptr<LowLevelDevice> HostLowLevelDeviceCreate(size_t num_bytes,
+                                                               TargetPtr* base_addr) {
+  std::shared_ptr<LowLevelDevice> lld = std::make_shared<HostLowLevelDevice>(num_bytes, base_addr);
   return lld;
 }
 

src/runtime/micro/low_level_device.h

@@ -78,7 +78,8 @@ class LowLevelDevice {
  * \param num_bytes size of the memory region
  * \param base_addr pointer to write the host device's resulting base address into
  */
-const std::shared_ptr<LowLevelDevice> HostLowLevelDeviceCreate(size_t num_bytes, void** base_addr);
+const std::shared_ptr<LowLevelDevice> HostLowLevelDeviceCreate(size_t num_bytes,
+                                                               TargetPtr* base_addr);
 
 /*!
  * \brief connect to OpenOCD and create an OpenOCD low-level device

src/runtime/micro/micro_common.cc

@@ -51,7 +51,7 @@ const char* SectionToString(SectionKind section) {
 
 std::string RelocateBinarySections(
     const std::string& binary_path,
-    size_t word_size,
+    TargetWordSize word_size,
     TargetPtr text_start,
     TargetPtr rodata_start,
     TargetPtr data_start,
@@ -62,7 +62,7 @@ std::string RelocateBinarySections(
   CHECK(f != nullptr)
     << "Require tvm_callback_relocate_binary to exist in registry";
   std::string relocated_bin = (*f)(binary_path,
-                                   word_size,
+                                   word_size.bytes(),
                                    text_start.cast_to<uint64_t>(),
                                    rodata_start.cast_to<uint64_t>(),
                                    data_start.cast_to<uint64_t>(),
@@ -91,15 +91,15 @@ std::string ReadSection(const std::string& binary,
 size_t GetSectionSize(const std::string& binary_path,
                       SectionKind section,
                       const std::string& toolchain_prefix,
-                      size_t align) {
+                      TargetWordSize word_size) {
   CHECK(section == SectionKind::kText || section == SectionKind::kRodata ||
         section == SectionKind::kData || section == SectionKind::kBss)
       << "GetSectionSize requires section to be one of text, rodata, data, or bss.";
   const auto* f = Registry::Get("tvm_callback_get_section_size");
   CHECK(f != nullptr)
     << "Require tvm_callback_get_section_size to exist in registry";
   int size = (*f)(binary_path, SectionToString(section), toolchain_prefix);
-  return UpperAlignValue(size, align);
+  return UpperAlignValue(size, word_size.bytes());
 }
 
 }  // namespace runtime