Added macro generation in MLF export

python/tvm/micro/model_library_format.py

@@ -47,15 +47,32 @@ class UnsupportedInModelLibraryFormatError(Exception):
 
 
 def generate_c_interface_header(
-    module_name, inputs, outputs, pools, io_pool_allocations, devices, workspace_size, include_path
+    module_name,
+    inputs,
+    outputs,
+    pools,
+    io_pool_allocations,
+    devices,
+    workspace_size,
+    include_path,
+    input_sizes,
+    output_sizes,
 ):
     """Generate C Interface header to be included in MLF"""
     mangled_name = to_c_variable_style(prefix_generated_name(module_name))
     metadata_header = os.path.join(include_path, f"{mangled_name}.h")
 
     interface_c_create = tvm._ffi.get_global_func("runtime.InterfaceCCreate")
     interface_c_module = interface_c_create(
-        module_name, inputs, outputs, pools, io_pool_allocations, devices, workspace_size
+        module_name,
+        inputs,
+        outputs,
+        pools,
+        io_pool_allocations,
+        devices,
+        workspace_size,
+        input_sizes,
+        output_sizes,
     )
 
     with open(metadata_header, "w") as header_file:
@@ -277,6 +294,37 @@ def _create_empty_entry(target_device_type):
             main_func_metadata.io_sizes[target]
         )
 
+        # Now, we also add the information about the size of each input and output of the main
+        # function (in bytes)
+        input_dict = {}
+        for input_param in main_func_metadata.relay_primfuncs[target].params:
+            if hasattr(input_param, "checked_type"):
+                input_dict[input_param.name_hint] = int(
+                    _shape_to_size(input_param.checked_type.shape, input_param.checked_type.dtype)
+                )
+            else:
+                # TODO: maybe fill checked_type here?
+                input_dict[input_param.name_hint] = 0
+        target_main_entries[int(target.kind.device_type)]["inputs"] = input_dict
+
+        output_dict = {}
+        # For output, we dont have the name of the output, so we enumerate them
+        if isinstance(main_func_metadata.relay_primfuncs[target].ret_type, tvm.ir.type.TupleType):
+            for i, output_type in enumerate(
+                main_func_metadata.relay_primfuncs[target].ret_type.fields
+            ):
+                if hasattr(output_type, "shape"):
+                    output_dict[i] = int(_shape_to_size(output_type.shape, output_type.dtype))
+                else:
+                    output_dict[i] = 0
+        else:
+            output_type = main_func_metadata.relay_primfuncs[target].ret_type
+            if hasattr(output_type, "shape"):
+                output_dict[0] = int(_shape_to_size(output_type.shape, output_type.dtype))
+            else:
+                output_dict[0] = 0
+        target_main_entries[int(target.kind.device_type)]["outputs"] = output_dict
+
     ret = {
         "operator_functions": func_entries,
         "main": list(target_main_entries.values()),
@@ -427,6 +475,18 @@ def _export_graph_model_library_format(
                     "workspace_size_bytes"
                 ]
             )
+            inputs_sizes = metadata["modules"][mod.libmod_name]["memory"]["functions"]["main"][0][
+                "inputs"
+            ]
+            # Here, we merge the output sizes with the actual output names
+            output_sizes = {}
+            for key in metadata["modules"][mod.libmod_name]["memory"]["functions"]["main"][0][
+                "outputs"
+            ].keys():
+                output_sizes[outputs[key]] = metadata["modules"][mod.libmod_name]["memory"][
+                    "functions"
+                ]["main"][0]["outputs"][key]
+
             generate_c_interface_header(
                 mod.libmod_name,
                 inputs,
@@ -436,6 +496,8 @@ def _export_graph_model_library_format(
                 devices,
                 workspace_size,
                 include_path,
+                inputs_sizes,
+                output_sizes,
             )
 
         is_aot = isinstance(mod, executor_factory.AOTExecutorFactoryModule)
@@ -459,7 +521,7 @@ class NonStaticShapeError(Exception):
 
 def _shape_to_size(shape, dtype):
     bits_per_item = int(
-        re.match(r"((float)|(int))(?P<width_bits>[0-9]+)", dtype).group("width_bits")
+        re.match(r"((float)|(int)|(uint))(?P<width_bits>[0-9]+)", dtype).group("width_bits")
     )
     assert bits_per_item is not None, f"don't know how to compute size of type {dtype}"
     total_bits = bits_per_item

src/target/source/interface_c.cc

@@ -46,20 +46,42 @@ class InterfaceCNode : public runtime::ModuleNode {
   InterfaceCNode(std::string module_name, Array<String> inputs, Array<String> outputs,
                  Array<tir::usmp::AllocatedPoolInfo> pools,
                  Map<String, tir::usmp::PoolAllocation> io_pool_allocations, Array<String> devices,
-                 int workspace_size)
+                 int workspace_size, Map<String, IntImm> input_sizes,
+                 Map<String, IntImm> output_sizes)
       : module_name_(module_name),
         inputs_(inputs),
         outputs_(outputs),
         devices_(devices),
         pools_(FilterExternalPools(pools)),
         io_pool_allocations_(io_pool_allocations),
-        workspace_size_(workspace_size) {}
+        workspace_size_(workspace_size),
+        input_sizes_(input_sizes),
+        output_sizes_(output_sizes) {}
   const char* type_key() const final { return "h"; }
 
   std::string GetSource(const std::string& format) final {
     std::stringstream code;
 
     EmitUpperHeaderGuard(code);
+
+    // Emit macros for input sizes
+    for (auto const& it : input_sizes_) {
+      std::string input_name = SanitizeName(it.first);
+      std::string input_macro_name = input_name + "_size";
+      int input_size = it.second->value;
+      EmitIntegerValueMacro(code, "Input tensor " + input_name + " size (in bytes)",
+                            input_macro_name, input_size);
+    }
+
+    // Emit macros for output sizes
+    for (auto const& it : output_sizes_) {
+      std::string output_name = SanitizeName(it.first);
+      std::string output_macro_name = output_name + "_size";
+      int output_size = it.second->value;
+      EmitIntegerValueMacro(code, "Output tensor " + output_name + " size (in bytes)",
+                            output_macro_name, output_size);
+    }
+
     EmitBrief(code, "Input tensor pointers");
     EmitStruct(code, "inputs", inputs_);
     EmitBrief(code, "Output tensor pointers");
@@ -277,14 +299,18 @@ class InterfaceCNode : public runtime::ModuleNode {
   Array<tir::usmp::AllocatedPoolInfo> pools_;
   Map<String, tir::usmp::PoolAllocation> io_pool_allocations_;
   int workspace_size_;
+  Map<String, IntImm> input_sizes_;
+  Map<String, IntImm> output_sizes_;
 };
 
 runtime::Module InterfaceCCreate(std::string module_name, Array<String> inputs,
                                  Array<String> outputs, Array<tir::usmp::AllocatedPoolInfo> pools,
                                  Map<String, tir::usmp::PoolAllocation> io_pool_allocations,
-                                 Array<String> devices, int workspace_size) {
+                                 Array<String> devices, int workspace_size,
+                                 Map<String, IntImm> input_sizes,
+                                 Map<String, IntImm> output_sizes) {
   auto n = make_object<InterfaceCNode>(module_name, inputs, outputs, pools, io_pool_allocations,
-                                       devices, workspace_size);
+                                       devices, workspace_size, input_sizes, output_sizes);
   return runtime::Module(n);
 }