[AOT] Added a test for detecting output size post MLF export

python/tvm/micro/model_library_format.py

@@ -485,6 +485,12 @@ def _export_graph_model_library_format(
                     "functions"
                 ]["main"][0]["outputs"][key]
 
+            input_name_to_size_map = {
+                name: property_map["size"] for name, property_map in inputs_sizes.items()
+            }
+            output_name_to_size_map = {
+                name: property_map["size"] for name, property_map in output_sizes.items()
+            }
             generate_c_interface_header(
                 mod.libmod_name,
                 inputs,
@@ -494,8 +500,8 @@ def _export_graph_model_library_format(
                 devices,
                 workspace_size,
                 include_path,
-                inputs_sizes,
-                output_sizes,
+                input_name_to_size_map,
+                output_name_to_size_map,
             )
 
         is_aot = isinstance(mod, executor_factory.AOTExecutorFactoryModule)

tests/python/relay/aot/test_crt_aot.py

@@ -225,6 +225,63 @@ def test_packed_global_variables():
             assert f"{func}_packed" not in tvmgen_names
 
 
+def test_io_size_definition():
+    """Check network IO size definitions in the codegen output."""
+    dtype = "float32"
+    ishape = (1, 32, 14, 14)
+    wshape = (32, 32, 3, 3)
+    interface_api = "c"
+    use_unpacked_api = True
+
+    data0 = relay.var("data", shape=ishape, dtype=dtype)
+    weight0 = relay.var("weight", shape=wshape, dtype=dtype)
+    out = relay.nn.conv2d(data0, weight0, kernel_size=(3, 3), padding=(1, 1), groups=1)
+    main_f = relay.Function([data0, weight0], out)
+    mod = tvm.IRModule()
+    mod["main"] = main_f
+    mod = transform.InferType()(mod)
+
+    i_data = np.random.uniform(0, 1, ishape).astype(dtype)
+    w_data = np.random.uniform(0, 1, wshape).astype(dtype)
+
+    inputs = OrderedDict([("data", i_data), ("weight", w_data)])
+
+    output_list = generate_ref_data(mod, inputs)
+    compiled_models_list = compile_models(
+        models=AOTTestModel(module=mod, inputs=inputs, outputs=output_list),
+        interface_api=interface_api,
+        use_unpacked_api=use_unpacked_api,
+        workspace_byte_alignment=8,
+        enable_op_fusion=True,
+        pass_config=AOT_DEFAULT_RUNNER.pass_config,
+        use_runtime_executor=True,
+        target=tvm.target.Target("c"),
+    )
+    dtype_itemsize = np.dtype(dtype).itemsize
+    ref_input_size = i_data.size * dtype_itemsize
+    ref_weight_size = w_data.size * dtype_itemsize
+    ref_output_size = output_list["output"].size * dtype_itemsize
+    compiled_model = compiled_models_list[0]
+
+    tmp_path = utils.tempdir()
+    base_path = tmp_path.temp_dir
+
+    model = compiled_model.model
+    tar_file = os.path.join(base_path, f"{model.name}.tar")
+    export_model_library_format(compiled_model.executor_factory, tar_file)
+    t = tarfile.open(tar_file)
+    t.extractall(base_path)
+
+    header_path = f"{base_path}/codegen/host/include/tvmgen_{model.name}.h"
+    with open(header_path, "r") as header:
+        contents = header.readlines()
+        contents = "".join(map(str, contents))
+        assert contents.count("_SIZE") == 4
+        assert f"TVMGEN_DEFAULT_DATA_SIZE {ref_input_size}" in contents
+        assert f"TVMGEN_DEFAULT_WEIGHT_SIZE {ref_weight_size}" in contents
+        assert f"TVMGEN_DEFAULT_OUTPUT_SIZE {ref_output_size}" in contents
+
+
 @parametrize_aot_options
 def test_concatenate(interface_api, use_unpacked_api, test_runner):
     """Tests compilation of concatenate"""