diff --git a/examples/arm/aot_arm_compiler.py b/examples/arm/aot_arm_compiler.py index 9a45195e58..b63d0b730b 100644 --- a/examples/arm/aot_arm_compiler.py +++ b/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) diff --git a/examples/arm/run.sh b/examples/arm/run.sh index 4a3f6dbf67..db46630a8a 100755 --- a/examples/arm/run.sh +++ b/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}" }