From 18fadc8dd07bbb7204f7c4060a95c34fc0cfc5e2 Mon Sep 17 00:00:00 2001
From: Giuseppe Rossini <giuseppe.rossini@arm.com>
Date: Tue, 18 May 2021 22:07:29 +0100
Subject: [PATCH] rebasing

Change-Id: I25e43d426f650e7ed9faa08023a4b59598c502b3
---
 tests/python/relay/aot/aot_test_utils.py | 27 ++++++++++++++++--------
 tests/python/relay/aot/test_crt_aot.py   | 10 ++++++++-
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/tests/python/relay/aot/aot_test_utils.py b/tests/python/relay/aot/aot_test_utils.py
index 24442ed9b49b2..cfa11d5d8e449 100644
--- a/tests/python/relay/aot/aot_test_utils.py
+++ b/tests/python/relay/aot/aot_test_utils.py
@@ -214,7 +214,13 @@ def extract_main_workspace_sizebytes(extract_dir):
 
 
 def compile_and_run(
-    mod, input_list, output_list, use_calculated_workspaces, params=None, workspace_byte_alignment=8, mod_name=None
+    mod,
+    input_list,
+    output_list,
+    use_calculated_workspaces,
+    params=None,
+    workspace_byte_alignment=8,
+    mod_name=None,
 ):
     """
     This method verifies the generated source
@@ -231,6 +237,7 @@ def compile_and_run(
 
     tmp_path = utils.tempdir()
     tmp_dir = tmp_path.temp_dir
+    tmp_dir = "."
 
     base_path = os.path.join(tmp_dir, "test")
     build_path = os.path.join(base_path, "build")
@@ -258,12 +265,18 @@ def compile_and_run(
             (f'{mangle_name(mod_name, "expected_output_data")}{i}'), output_list[i], build_path
         )
 
-    create_main("test.c", {mod_name: input_list}, {mod_name: output_list}, build_path, workspace_bytes)
+    create_main(
+        "test.c", {mod_name: input_list}, {mod_name: output_list}, build_path, workspace_bytes
+    )
 
     # Verify that compiles fine
     file_dir = os.path.dirname(os.path.abspath(__file__))
     makefile = os.path.join(file_dir, "aot_test.mk")
-    make_cmd = f"make -f {makefile} build_dir=" + build_path + f" TVM_ROOT={file_dir}/../../../.."
+    make_cmd = (
+        f"make CFLAGS='{cflags}' -f {makefile} build_dir="
+        + build_path
+        + f" TVM_ROOT={file_dir}/../../../.."
+    )
 
     compile_log_path = os.path.join(build_path, "test_compile.log")
     ret = subprocess_with_stdout_and_log(make_cmd, ".", compile_log_path, False)
@@ -317,16 +330,12 @@ def compile_and_run_multiple_models(mod_map, input_list_map, output_list_map, pa
                 (f'{mangle_name(mod_name,"expected_output_data")}{i}'), output_list[i], build_path
             )
 
-    create_main("test.c", input_list_map, output_list_map, build_path)
+    create_main("test.c", input_list_map, output_list_map, build_path, workspace_bytes=16384 * 1024)
 
     # Verify that compiles fine
     file_dir = os.path.dirname(os.path.abspath(__file__))
     makefile = os.path.join(file_dir, "aot_test.mk")
-    make_cmd = (
-        f"make CFLAGS='{cflags}' -f {makefile} build_dir="
-        + build_path
-        + f" TVM_ROOT={file_dir}/../../../.."
-    )
+    make_cmd = f"make -f {makefile} build_dir=" + build_path + f" TVM_ROOT={file_dir}/../../../.."
 
     compile_log_path = os.path.join(build_path, "test_compile.log")
     ret = subprocess_with_stdout_and_log(make_cmd, ".", compile_log_path, False)
diff --git a/tests/python/relay/aot/test_crt_aot.py b/tests/python/relay/aot/test_crt_aot.py
index 41be5632f66f6..ec04f56676ee6 100644
--- a/tests/python/relay/aot/test_crt_aot.py
+++ b/tests/python/relay/aot/test_crt_aot.py
@@ -380,7 +380,14 @@ def test_add_name_mangling_with_params():
     output_list = generate_ref_data(func, inputs, params)
 
     input_list = [y_in]
-    compile_and_run(func, input_list, output_list, params, "my_mod")
+    compile_and_run(
+        func,
+        input_list,
+        output_list,
+        use_calculated_workspaces=False,
+        params=params,
+        mod_name="my_mod",
+    )
 
 
 def test_multiple_models():
@@ -429,5 +436,6 @@ def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(8, 3, 5, 5),
 
     compile_and_run_multiple_models(mod_map, input_list_map, output_list_map, param_map)
 
+
 if __name__ == "__main__":
     pytest.main([__file__])