[SME][AOT] Add Fixed Virtual Platform (FVP) functional testing infrastructure

This commit adds the infrastructure required for testing compiled
functions that use SME. A more in depth discussion can be found here:
https://github.com/apache/tvm-rfcs/blob/main/rfcs/0107-scalable-matrix-extension-enablement.md#testing

Specifically, this commit adds: the installation of the AArch64
Architecture Envelope Model (AEM) Fixed Virtual Platform (FVP),
supporting files for compiling and running a graph on the FVP, sample
tests which can be removed once TVM can generate SME and some
enhancements to the AOT testing infrastructure so that TVM compiled
functions can be run on the FVP.

Change-Id: I60d39fc17b826a9f5c71991d86d3791de83a54d4

Author: lhutton1

docker/Dockerfile.ci_cpu

@@ -149,3 +149,8 @@ RUN bash /install/ubuntu_install_libxsmm.sh
 # ONNX and PyTorch
 COPY install/ubuntu_install_onnx.sh /install/ubuntu_install_onnx.sh
 RUN bash /install/ubuntu_install_onnx.sh
+
+# AArch64 Architecture Envelope Model (AEM)
+COPY install/ubuntu_install_aprofile_aem.sh /install
+RUN bash /install/ubuntu_install_aprofile_aem.sh
+ENV PATH $PATH:/opt/arm/fvp/Base_RevC_AEMvA_pkg/models/Linux64_GCC-9.3/:/opt/arm/gcc-aarch64-none-elf/bin

docker/install/ubuntu_install_aprofile_aem.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Install the AArch64 Architecture Envelope Model (AEM)
+
+set -e
+set -u
+set -o pipefail
+
+tmpdir=$(mktemp -d)
+
+cleanup()
+{
+    rm -rf "$tmpdir"
+}
+
+trap cleanup 0
+
+pushd "$tmpdir"
+
+# Install GCC toolchain
+gcc_install_dir="/opt/arm/gcc-aarch64-none-elf"
+gcc_url="https://developer.arm.com/-/media/Files/downloads/gnu/13.2.rel1/binrel/arm-gnu-toolchain-13.2.rel1-x86_64-aarch64-none-elf.tar.xz?rev=28d5199f6db34e5980aae1062e5a6703&hash=D87D4B558F0A2247B255BA15C32A94A9F354E6A8"
+gcc_sha="7fe7b8548258f079d6ce9be9144d2a10bd2bf93b551dafbf20fe7f2e44e014b8"
+gcc_tar="arm-gnu-toolchain-13.2.rel1-x86_64-aarch64-none-linux-gnu.tar.xz"
+mkdir -p $gcc_install_dir
+curl --retry 64 -sSL $gcc_url -o $gcc_tar
+echo "$gcc_sha $gcc_tar" | sha256sum --check
+tar -xf $gcc_tar -C $gcc_install_dir --strip-components=1
+
+# Download FVP
+fvp_dir="/opt/arm/fvp"
+fvp_url="https://developer.arm.com/-/media/Files/downloads/ecosystem-models/FVP_Base_RevC-2xAEMvA_11.24_11_Linux64.tgz"
+fvp_sha="0f132334834cbc66889a62dd72057c976d7c7dfcfeec21799e9c78fb2ce24720"
+curl --retry 64 -sSL $fvp_url -o fvp.tgz
+echo "$fvp_sha fvp.tgz" | sha256sum --check
+mkdir -p "$fvp_dir"
+tar -xzf fvp.tgz -C "$fvp_dir"
+rm -rf doc  # Remove some documentation bundled with the package

python/tvm/testing/aot.py

@@ -179,6 +179,15 @@ def _subprocess_check_log_output(cmd, cwd, logfile):
         raise RuntimeError(f"Subprocess failed: {cmd}\nstdout:\n{stdout}")
 
 
+def _get_entrypoint_suffix(target):
+    # LLVM modules don't use the same entrypoint suffix
+    # as C source generated modules.
+    if target.kind.name == "llvm":
+        return "__tvm_main__"
+    else:
+        return "run"
+
+
 def _mangle_name(mod_name, name):
     mod_name = mangle_module_name(mod_name)
     return mod_name + "_" + name
@@ -385,7 +394,14 @@ def _emit_main_fake_packed_values(main_file):
     )
 
 
-def _emit_main_packed_call(main_file, input_map, output_list, mod_name):
+def _emit_entry_function_forward_declaration(main_file, mod_name, entrypoint_suffix):
+    main_file.write(
+        f"int {_mangle_name(mod_name, entrypoint_suffix)}"
+        f"(TVMValue[], int32_t[], int32_t, void*, int32_t, void*);\n"
+    )
+
+
+def _emit_main_packed_call(main_file, input_map, output_list, mod_name, entrypoint_suffix):
     tensors_name = _mangle_name(mod_name, "tensors")
     values_name = _mangle_name(mod_name, "values")
     typeids_name = _mangle_name(mod_name, "typeids")
@@ -420,7 +436,8 @@ def fake_tensor(source, source_index, packed_index):
         fake_tensor(_mangle_name(mod_name, "outputs"), i, i + num_inputs)
 
     main_file.write(
-        f'{_mangle_name(mod_name, "run")}({values_name}, {typeids_name}, 0, NULL, 0, NULL);\n'
+        f"{_mangle_name(mod_name, entrypoint_suffix)}"
+        f"({values_name}, {typeids_name}, 0, NULL, 0, NULL);\n"
     )
     main_file.write("\n")
 
@@ -544,6 +561,15 @@ def _create_main(
             model = compiled_model.model
             _emit_main_data(main_file, model.inputs, model.outputs, model.name)
 
+        if interface_api == "packed":
+            for compiled_model in compiled_models:
+                entrypoint_suffix = _get_entrypoint_suffix(
+                    compiled_model.executor_factory.target[0]
+                )
+                _emit_entry_function_forward_declaration(
+                    main_file, compiled_model.model.name, entrypoint_suffix
+                )
+
         _emit_main_prologue(
             main_file,
             custom_prologue,
@@ -592,7 +618,12 @@ def _create_main(
             for compiled_model in compiled_models:
                 model = compiled_model.model
                 _emit_main_data_setup(main_file, model.inputs, model.outputs, model.name)
-                _emit_main_packed_call(main_file, model.inputs, model.outputs, model.name)
+                entrypoint_suffix = _get_entrypoint_suffix(
+                    compiled_model.executor_factory.target[0]
+                )
+                _emit_main_packed_call(
+                    main_file, model.inputs, model.outputs, model.name, entrypoint_suffix
+                )
 
         for compiled_model in compiled_models:
             model = compiled_model.model
@@ -665,14 +696,18 @@ def compile_models(
     workspace_memory_pools=None,
     constant_memory_pools=None,
     schedule_name: str = None,
+    runtime: tvm.relay.backend.Runtime = Runtime("crt"),
 ) -> List[AOTCompiledTestModel]:
     """
     This method generates runtime.Modules for the tests
     """
     if not isinstance(models, list):
         models = [models]
 
-    runtime = Runtime("crt")
+    assert (
+        runtime.name == "crt"
+    ), f"Currently only 'crt' is supported by the test framework, but got {runtime.name}"
+
     executor = Executor(
         "aot",
         {
@@ -835,10 +870,12 @@ def run_and_check_body(base_path):
         makefile_dir = os.path.join(file_dir, "../../../tests/python/relay/aot")
         codegen_path = os.path.join(base_path, "codegen")
         makefile = os.path.join(makefile_dir, f"{runner.makefile}.mk")
-        fvp_dir = "/opt/arm/FVP_Corstone_SSE-300/models/Linux64_GCC-6.4/"
-        # TODO(@grant-arm): Remove once ci_cpu docker image has been updated to FVP_Corstone_SSE
-        if not os.path.isdir(fvp_dir):
-            fvp_dir = "/opt/arm/FVP_Corstone_SSE-300_Ethos-U55/models/Linux64_GCC-6.4/"
+
+        if runner.makefile == "aprofile_aem":
+            fvp_dir = "/opt/arm/fvp/Base_RevC_AEMvA_pkg/models/Linux64_GCC-9.3/"
+        else:
+            fvp_dir = "/opt/arm/FVP_Corstone_SSE-300/models/Linux64_GCC-6.4/"
+
         custom_params = " ".join(
             [f" {param}='{value}'" for param, value in runner.parameters.items()]
         )
@@ -901,11 +938,28 @@ def compile_and_run(
     debug_last_error: bool = False,
     checker: Optional[Callable[[str], bool]] = None,
     print_output_on_mismatch: bool = False,
+    runtime: tvm.relay.backend.Runtime = Runtime("crt"),
 ) -> bool:
     """This is a wrapper API to compile and run models as test for AoT
 
     Parameters
     ----------
+    interface_api : str
+        The external calling convention interface API.
+
+        Examples: "c", "packed"
+
+    use_unpacked_api : bool
+        Whether or not to use type-erased API internally for the
+        operator calling convention.
+
+        Note: This feature can be useful for embedded targets
+        when space is at a premium.
+
+        Permitted values when interface API is:
+        > "c": True
+        > "packed": True/False
+
     test_dir : str
         This path will contain build, codegen, include directories.
 
@@ -935,6 +989,7 @@ def compile_and_run(
         use_runtime_executor=use_runtime_executor,
         target=target,
         schedule_name=schedule_name,
+        runtime=runtime,
     )
 
     return run_and_check(

tests/python/integration/test_arm_aprofile.py

@@ -16,13 +16,18 @@
 # under the License.
 """Tests for Arm(R) A-Profile Architecture."""
 import os
+import subprocess
+
 import numpy as np
 import pytest
+
 import tvm
 import tvm.testing
 from tvm import relay
 from tvm.relay.transform import ToMixedPrecision, FoldConstant
 from tvm.relay.build_module import bind_params_by_name
+from tvm.testing.aot import AOTTestModel, AOTTestRunner, generate_ref_data, compile_and_run
+from tvm.contrib import utils
 
 
 def get_mattr(dtype):
@@ -73,3 +78,96 @@ def test_conv2d(dtype):
     with tvm.transform.PassContext(opt_level=3):
         lib = tvm.relay.build(mod, target=target, params=params)
         lib.export_library(lib_path, cc="aarch64-linux-gnu-gcc")
+
+
+# AOT Test Runner using the AArch64 Architecture Envelope Model (AEM)
+# Fixed Virtual Platform (FVP) reference system.
+# See: https://developer.arm.com/Tools%20and%20Software/Fixed%20Virtual%20Platforms
+AOT_APROFILE_AEM_RUNNER = AOTTestRunner(
+    makefile="aprofile_aem",
+    pass_config={
+        "tir.usmp.enable": False,
+        "tir.disable_assert": True,  # AOT test infra creates 'fake' inputs that fail asserts
+    },
+)
+
+
+@tvm.testing.requires_x86
+def test_aem_simple_addition():
+    """Tests a simple addition running on the AArch64 AEM."""
+    inp = relay.var("data", shape=(1, 2, 4, 4))
+    add = relay.add(inp, relay.const(np.ones((1, 2, 4, 4))))
+    func = relay.Function([inp], add)
+    ir_mod = tvm.IRModule.from_expr(func)
+    ir_mod = tvm.relay.transform.InferType()(ir_mod)
+
+    main_func = ir_mod["main"]
+    shape_dict = {p.name_hint: p.checked_type.concrete_shape for p in main_func.params}
+    type_dict = {p.name_hint: p.checked_type.dtype for p in main_func.params}
+
+    input_data = np.random.uniform(size=shape_dict["data"]).astype(type_dict["data"])
+    params = {}
+    inputs = {"data": input_data}
+    ref_outputs = generate_ref_data(ir_mod, inputs, params)
+
+    compile_and_run(
+        AOTTestModel(module=ir_mod, inputs=inputs, outputs=ref_outputs, params=params),
+        target=tvm.target.Target("llvm -mtriple=aarch64-none-elf"),
+        runtime=tvm.relay.backend.Runtime("crt", {"system-lib": True}),
+        interface_api="packed",
+        use_unpacked_api=False,
+        runner=AOT_APROFILE_AEM_RUNNER,
+    )
+
+
+@tvm.testing.requires_x86
+def test_aem_asm_sme():
+    """
+    Tests SME assembly runs on the AArch64 AEM. This test is used as a simple
+    sanity check until the TVM schedules are able to produce SME.
+    """
+    c_code = """
+    #include <stdio.h>
+
+    int main(void) {
+        __asm volatile(
+            "smstart\\n"
+            "smstop\\n"
+        );
+        printf("EXITTHESIM\\n");
+        return 0;
+    }
+    """
+    runner = AOT_APROFILE_AEM_RUNNER
+
+    tmpdir = utils.tempdir()
+    build_path = os.path.join(tmpdir.path, "build")
+    os.makedirs(build_path, exist_ok=True)
+
+    with open(build_path + "/test.c", "w") as f:
+        f.write(c_code)
+
+    file_dir = os.path.dirname(os.path.abspath(__file__))
+    makefile_dir = os.path.join(file_dir, "../../../tests/python/relay/aot")
+    makefile = os.path.join(makefile_dir, f"{runner.makefile}.mk")
+
+    make_command = (
+        f"make -f {makefile} build_dir={build_path}"
+        + f" TVM_ROOT={file_dir}/../../.."
+        + f" AOT_TEST_ROOT={makefile_dir}"
+        + " FVP_DIR=/opt/arm/fvp/Base_RevC_AEMvA_pkg/models/Linux64_GCC-9.3/"
+    )
+
+    compile_command = f"{make_command} aot_test_runner"
+    popen = subprocess.Popen(compile_command, cwd=build_path, shell=True, stdout=subprocess.PIPE)
+    return_code = popen.wait()
+    assert not return_code, "Failed to compile"
+
+    run_command = f"{make_command} run"
+    popen = subprocess.Popen(run_command, cwd=build_path, shell=True, stdout=subprocess.PIPE)
+    return_code = popen.wait()
+    assert not return_code, "Failed to run"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()