Run AOT tests against reference system

This introduces an alternative way of running AOT tests using the reference system added in apache#8514. This gives us additional assurance that the AOT output runs successfully on embedded platforms in our core test suite.

I've also changed calculate_workspace_sizes to debug_workspace_sizes and default to False in most cases as it only needs to be True for a few cases to check theoutput with the debug flag - this was discovered trying to allocate 16MB in an embedded test 🙀

Co-authored-by: Grant Watson <[email protected]>

tests/python/relay/aot/aot_test_utils.py

@@ -39,6 +39,9 @@
 
 _LOG = logging.getLogger(__name__)
 
+AOT_SUCCESS_TOKEN = "AOT_TEST_SUCCESS"
+AOT_FAILURE_TOKEN = "AOT_TEST_FAILURE"
+
 
 class AOTTestModel(NamedTuple):
     """Class to describe a model under test
@@ -64,6 +67,38 @@ class AOTTestModel(NamedTuple):
     params: Optional[Dict[str, np.array]] = None
 
 
+class AOTTestRunner(NamedTuple):
+    """Class to describe a test runner for AOT code
+
+    Parameters
+    ----------
+    makefile: str
+        Premade Makefile to use from the AOT test folder
+    prologue: str
+        Code to prepend to the main function
+    includes: List[str]
+        Additional includes required to run the AOT test runner
+    parameters: Map[str, str]
+        Additional parameters to pass to the make command
+    """
+
+    makefile: str = "default"
+    prologue: str = ""
+    includes: List[str] = []
+    parameters: Dict[str, str] = {}
+
+
+AOT_DEFAULT_RUNNER = AOTTestRunner()
+AOT_CORSTONE300_RUNNER = AOTTestRunner(
+    makefile="corstone300",
+    prologue="""
+    uart_init();
+    """,
+    includes=["uart.h"],
+    parameters={"NPU_VARIANT": "256"},
+)
+
+
 def mangle_name(mod_name, name):
     mod_name = mangle_module_name(mod_name)
     return mod_name + "_" + name
@@ -114,17 +149,27 @@ def parametrize_aot_options(test):
 
     interface_api = ["packed", "c"]
     use_unpacked_api = [True, False]
-    use_calculated_workspaces = [True, False]
+    test_runner = [AOT_DEFAULT_RUNNER, AOT_CORSTONE300_RUNNER]
+
+    all_combinations = itertools.product(interface_api, use_unpacked_api, test_runner)
 
-    all_combinations = itertools.product(interface_api, use_unpacked_api, use_calculated_workspaces)
     # Filter out packed operators with c interface
     valid_combinations = filter(
-        lambda parameters: not (parameters[0] == "c" and parameters[1] == False),
+        lambda parameters: not (parameters[0] == "c" and not parameters[1]),
         all_combinations,
     )
 
+    # Only use reference system for C interface and unpacked API calls
+    valid_combinations = filter(
+        lambda parameters: not (
+            parameters[2] == AOT_CORSTONE300_RUNNER
+            and (parameters[0] == "packed" or not parameters[1])
+        ),
+        valid_combinations,
+    )
+
     return pytest.mark.parametrize(
-        ["interface_api", "use_unpacked_api", "use_calculated_workspaces"],
+        ["interface_api", "use_unpacked_api", "test_runner"],
         valid_combinations,
     )(test)
 
@@ -160,7 +205,7 @@ def subprocess_log_output(cmd, cwd, logfile):
     return proc.wait()
 
 
-def emit_main_prologue(main_file, workspace_bytes):
+def emit_main_prologue(main_file, custom_prologue, workspace_bytes):
     # Add TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES because of memory alignment.
     main_file.write(
         f"#define WORKSPACE_SIZE ({workspace_bytes} + TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES)\n"
@@ -185,6 +230,7 @@ def emit_main_prologue(main_file, workspace_bytes):
 int main(){\n
 """
     )
+    main_file.write(custom_prologue)
 
 
 def emit_main_data(main_file, input_map, output_list, mod_name):
@@ -297,11 +343,11 @@ def emit_main_compare(main_file, output_list, mod_name):
         main_file.write(f"for (int i = 0; i<{actual_data_name}{i}_len; i++){{\n")
         if is_float_dtype:
             main_file.write(
-                f'if (fabs({actual_data_name}{i}[i]-{expected_data_name}{i}[i]) > 0.001f){{\n\tprintf("ko\\n");\n\treturn -1;}}\n'
+                f'if (fabs({actual_data_name}{i}[i]-{expected_data_name}{i}[i]) > 0.001f){{\n\tprintf("{AOT_FAILURE_TOKEN}\\n");\n\treturn -1;}}\n'
             )
         else:
             main_file.write(
-                f'if ({actual_data_name}{i}[i]!={expected_data_name}{i}[i]){{\n\tprintf("ko\\n");\n\treturn -1;}}\n'
+                f'if ({actual_data_name}{i}[i]!={expected_data_name}{i}[i]){{\n\tprintf("{AOT_FAILURE_TOKEN}\\n");\n\treturn -1;}}\n'
             )
         main_file.write("}\n")
 
@@ -312,36 +358,40 @@ def emit_main_init_memory_manager(main_file):
 
 
 def emit_main_epilogue(main_file):
-    main_file.write('printf("ok\\n");')
+    main_file.write(f'printf("{AOT_SUCCESS_TOKEN}\\n");')
     main_file.write("return 0;")
     main_file.write("}\n")
 
 
-def emit_main_common_includes(main_file):
+def emit_main_common_includes(main_file, custom_includes):
     main_file.write("#include <stdio.h>\n")
     main_file.write("#include <math.h>\n")
     main_file.write('#include "tvm/runtime/c_runtime_api.h"\n')
     main_file.write('#include "tvm/runtime/crt/stack_allocator.h"\n')
+    for include in custom_includes:
+        main_file.write(f'#include "{include}"\n')
 
 
 def emit_main_micro_include(main_file, mod_name):
     main_file.write(f"#include <{mangle_module_name(mod_name)}.h>\n")
 
 
-def create_main(test_name, models, output_path, interface_api, workspace_bytes):
+def create_main(
+    test_name, models, output_path, custom_includes, custom_prologue, interface_api, workspace_bytes
+):
     file_path = pathlib.Path(f"{output_path}/" + test_name).resolve()
     # create header file
     raw_path = file_path.with_suffix(".c").resolve()
     with open(raw_path, "w") as main_file:
-        emit_main_common_includes(main_file)
+        emit_main_common_includes(main_file, custom_includes)
 
         if interface_api == "c":
             for model in models:
                 emit_main_micro_include(main_file, model.name)
-
-        emit_main_prologue(main_file, workspace_bytes)
         for model in models:
             emit_main_data(main_file, model.inputs, model.outputs, model.name)
+
+        emit_main_prologue(main_file, custom_prologue, workspace_bytes)
         emit_main_init_memory_manager(main_file)
 
         if interface_api == "c":
@@ -396,9 +446,10 @@ def extract_main_workspace_size_bytes(extract_dir):
 
 def compile_and_run(
     models: Union[List[AOTTestModel], AOTTestModel],
+    runner: AOTTestRunner,
     interface_api,
     use_unpacked_api,
-    use_calculated_workspaces,
+    debug_calculated_workspaces=False,
     workspace_byte_alignment=8,
     enable_op_fusion=True,
 ):
@@ -414,7 +465,7 @@ def compile_and_run(
         models = [models]
 
     # The calculated workspaces will not account for stack allocator tags used for debugging
-    if not use_calculated_workspaces:
+    if debug_calculated_workspaces:
         cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK "
 
     config = {"tir.disable_vectorize": True}
@@ -452,10 +503,7 @@ def compile_and_run(
         t = tarfile.open(tar_file)
         t.extractall(base_path)
 
-        if use_calculated_workspaces:
-            workspace_bytes += extract_main_workspace_size_bytes(base_path)
-        else:
-            workspace_bytes += 16384 * 1024
+        workspace_bytes += extract_main_workspace_size_bytes(base_path)
 
         for key in model.inputs:
             create_header_file(
@@ -480,31 +528,41 @@ def compile_and_run(
         "test.c",
         models,
         build_path,
+        runner.includes,
+        runner.prologue,
         interface_api,
         workspace_bytes,
     )
 
     # Verify that compiles fine
     file_dir = os.path.dirname(os.path.abspath(__file__))
     codegen_path = os.path.join(base_path, "codegen")
-    makefile = os.path.join(file_dir, "aot_test.mk")
-    make_cmd = (
-        f"make CFLAGS='{cflags}' -f {makefile} build_dir="
-        + build_path
+    makefile = os.path.join(file_dir, f"{runner.makefile}.mk")
+    custom_params = " ".join([f" {param}='{value}'" for param, value in runner.parameters.items()])
+    make_command = (
+        f"make -f {makefile} build_dir={build_path}"
+        + f" CFLAGS='{cflags}'"
         + f" TVM_ROOT={file_dir}/../../../.."
+        + f" AOT_TEST_ROOT={file_dir}"
         + f" CODEGEN_ROOT={codegen_path}"
         + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}"
+        + custom_params
     )
 
     compile_log_path = os.path.join(build_path, "test_compile.log")
-    ret = subprocess_log_output(make_cmd, ".", compile_log_path)
+    compile_command = f"{make_command} aot_test_runner"
+    ret = subprocess_log_output(compile_command, ".", compile_log_path)
     assert ret == 0
 
     # Verify that runs fine
     run_log_path = os.path.join(build_path, "test_run.log")
-    ret = subprocess_log_output("./aot_test_runner", build_path, run_log_path)
+    run_command = f"{make_command} run"
+    ret = subprocess_log_output(run_command, build_path, run_log_path)
     assert ret == 0
 
+    with open(run_log_path) as run_log:
+        assert AOT_SUCCESS_TOKEN in run_log.read()
+
 
 def generate_ref_data(mod, input_data, params=None, target="llvm"):
     """Generate reference data through executing the relay module"""