Add New Ethos-U85 support to Arm aot_arm_compiler (pytorch#5345)

Summary:
Some changes to enable using the new compiler flow and backend for both Ethos-U55 and Ethos-U85.

This change means the run.sh flow uses the new compiler for Ethos-U55.

It enables Ethos-U85 to be targeted for compilation, but this feature isn't used in run.sh yet.

Pull Request resolved: pytorch#5345

Reviewed By: manuelcandales

Differential Revision: D62874758

Pulled By: digantdesai

fbshipit-source-id: fd9921b7dd656a0087aaffde5986f4bdc320742c

Author: robell

examples/arm/aot_arm_compiler.py

@@ -9,6 +9,7 @@
 
 import argparse
 import logging
+import os
 
 import torch
 
@@ -48,6 +49,19 @@ def get_model_and_inputs_from_name(model_name: str):
         model, example_inputs, _ = EagerModelFactory.create_model(
             *MODEL_NAME_TO_MODEL[model_name]
         )
+    # Case 3: Model is in an external python file loaded as a module.
+    #         ModelUnderTest should be a torch.nn.module instance
+    #         ModelInputs should be a tuple of inputs to the forward function
+    elif model_name.endswith(".py"):
+        import importlib.util
+
+        # load model's module and add it
+        spec = importlib.util.spec_from_file_location("tmp_model", model_name)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        model = module.ModelUnderTest
+        example_inputs = module.ModelInputs
+
     else:
         raise RuntimeError(
             f"Model '{model_name}' is not a valid name. Use --help for a list of available models."
@@ -133,7 +147,51 @@ def forward(self, x):
     "softmax": SoftmaxModule,
 }
 
-if __name__ == "__main__":
+targets = [
+    "ethos-u85-128",
+    "ethos-u55-128",
+    "TOSA",
+]
+
+
+def get_compile_spec(target: str, intermediates: bool) -> ArmCompileSpecBuilder:
+    spec_builder = None
+    if target == "TOSA":
+        spec_builder = (
+            ArmCompileSpecBuilder().tosa_compile_spec().set_permute_memory_format(True)
+        )
+    elif target == "ethos-u55-128":
+        spec_builder = (
+            ArmCompileSpecBuilder()
+            .ethosu_compile_spec(
+                "ethos-u55-128",
+                system_config="Ethos_U55_High_End_Embedded",
+                memory_mode="Shared_Sram",
+                extra_flags="--debug-force-regor --output-format=raw",
+            )
+            .set_permute_memory_format(args.model_name in MODEL_NAME_TO_MODEL.keys())
+            .set_quantize_io(True)
+        )
+    elif target == "ethos-u85-128":
+        spec_builder = (
+            ArmCompileSpecBuilder()
+            .ethosu_compile_spec(
+                "ethos-u85-128",
+                system_config="Ethos_U85_SYS_DRAM_Mid",
+                memory_mode="Shared_Sram",
+                extra_flags="--output-format=raw",
+            )
+            .set_permute_memory_format(True)
+            .set_quantize_io(True)
+        )
+
+    if intermediates is not None:
+        spec_builder.dump_intermediate_artifacts_to(args.intermediates)
+
+    return spec_builder.build()
+
+
+def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "-m",
@@ -149,6 +207,15 @@ def forward(self, x):
         default=False,
         help="Flag for producing ArmBackend delegated model",
     )
+    parser.add_argument(
+        "-t",
+        "--target",
+        action="store",
+        required=False,
+        default="ethos-u55-128",
+        choices=targets,
+        help=f"For ArmBackend delegated models, pick the target, and therefore the instruction set generated. valid targets are {targets}",
+    )
     parser.add_argument(
         "-q",
         "--quantize",
@@ -167,8 +234,26 @@ def forward(self, x):
     parser.add_argument(
         "--debug", action="store_true", help="Set the logging level to debug."
     )
-
+    parser.add_argument(
+        "-i",
+        "--intermediates",
+        action="store",
+        required=False,
+        help="Store intermediate output (like TOSA artefacts) somewhere.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        action="store",
+        required=False,
+        help="Location for outputs, if not the default of cwd.",
+    )
     args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = get_args()
 
     if args.debug:
         logging.basicConfig(level=logging.DEBUG, format=FORMAT, force=True)
@@ -191,7 +276,7 @@ def forward(self, x):
     ):
         raise RuntimeError(f"Model {args.model_name} cannot be delegated.")
 
-    # 1. pick model from one of the supported lists
+    # Pick model from one of the supported lists
     model, example_inputs = get_model_and_inputs_from_name(args.model_name)
     model = model.eval()
 
@@ -209,23 +294,18 @@ def forward(self, x):
             _check_ir_validity=False,
         ),
     )
+
+    # As we can target multiple output encodings from ArmBackend, one must
+    # be specified.
+    compile_spec = (
+        get_compile_spec(args.target, args.intermediates)
+        if args.delegate is True
+        else None
+    )
+
     logging.debug(f"Exported graph:\n{edge.exported_program().graph}")
     if args.delegate is True:
-        edge = edge.to_backend(
-            ArmPartitioner(
-                ArmCompileSpecBuilder()
-                .ethosu_compile_spec(
-                    "ethos-u55-128",
-                    system_config="Ethos_U55_High_End_Embedded",
-                    memory_mode="Shared_Sram",
-                )
-                .set_permute_memory_format(
-                    args.model_name in MODEL_NAME_TO_MODEL.keys()
-                )
-                .set_quantize_io(True)
-                .build()
-            )
-        )
+        edge = edge.to_backend(ArmPartitioner(compile_spec))
         logging.debug(f"Lowered graph:\n{edge.exported_program().graph}")
 
     try:
@@ -241,7 +321,12 @@ def forward(self, x):
         else:
             raise e
 
-    model_name = f"{args.model_name}" + (
-        "_arm_delegate" if args.delegate is True else ""
+    model_name = os.path.basename(os.path.splitext(args.model_name)[0])
+    output_name = f"{model_name}" + (
+        f"_arm_delegate_{args.target}" if args.delegate is True else ""
     )
-    save_pte_program(exec_prog, model_name)
+
+    if args.output is not None:
+        output_name = os.path.join(args.output, output_name)
+
+    save_pte_program(exec_prog, output_name)

examples/arm/run.sh

@@ -44,7 +44,9 @@ function generate_pte_file() {
 
     local model_filename=${model}.pte
     if [[ "${delegate}" == *"--delegate"* ]]; then
-        model_filename=${model}_arm_delegate.pte
+	# Name aligned with default aot_arm_compiler output - run.sh only supports
+	# running on Corstone-300 with Ethos-U55 FVP at the moment.
+        model_filename=${model}_arm_delegate_ethos-u55-128.pte
     fi
     cd $et_root_dir
 
@@ -56,7 +58,7 @@ function generate_pte_file() {
     SO_LIB=$(find cmake-out-aot-lib -name libquantized_ops_aot_lib.so)
 
     python3 -m examples.arm.aot_arm_compiler --model_name="${model}" ${delegate} --so_library="$SO_LIB" 1>&2
-    [[ -f ${pte_file} ]] || { echo "Failed to generate a pte file - ${pte_file}"; exit 1; }
+    [[ -f ${pte_file} ]] || { >&2 echo "Failed to generate a pte file - ${pte_file}"; exit 1; }
     echo "${pte_file}"
 }