|
| 1 | +# Copyright 2025 Arm Limited and/or its affiliates. |
| 2 | +# |
| 3 | +# This source code is licensed under the BSD-style license found in the |
| 4 | +# LICENSE file in the root directory of this source tree. |
| 5 | + |
| 6 | +# pyre-unsafe |
| 7 | + |
| 8 | +# |
| 9 | +# Main implementation of AoT flow to partition and preprocess for VGF target |
| 10 | +# backends. This flow converts via TOSA, to an encoding of TOSA known as VGF |
| 11 | +# this form is used where the final JIT compile is performed on target (in the |
| 12 | +# runtime delegate executorch::runtime::BackendInterface::init |
| 13 | +# |
| 14 | + |
| 15 | +import logging |
| 16 | +import os |
| 17 | +import subprocess |
| 18 | +import tempfile |
| 19 | +from typing import final, List |
| 20 | + |
| 21 | +from executorch.backends.arm.tosa_backend import ( |
| 22 | + arm_get_first_delegation_tag, |
| 23 | + TOSABackend, |
| 24 | +) |
| 25 | +from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult |
| 26 | +from executorch.exir.backend.compile_spec_schema import CompileSpec |
| 27 | +from torch.export.exported_program import ExportedProgram |
| 28 | + |
| 29 | +# debug functionality |
| 30 | +logger = logging.getLogger(__name__) |
| 31 | + |
| 32 | + |
| 33 | +@final |
| 34 | +class VgfBackend(BackendDetails): |
| 35 | + """ |
| 36 | + BackendDetails subclass for delegation to VGF compatible devices. This enables |
| 37 | + encapsulated TOSA on target device and JIT compilation on suitable platforms. |
| 38 | + """ |
| 39 | + |
| 40 | + @staticmethod |
| 41 | + def _compile_tosa_flatbuffer( |
| 42 | + tosa_flatbuffer: bytes, |
| 43 | + compile_spec: List[CompileSpec], |
| 44 | + tag_name: str = "", |
| 45 | + ) -> bytes: |
| 46 | + """ |
| 47 | + Static helper method to do the compilation of the TOSA flatbuffer |
| 48 | + representation to a target specific binary stream. |
| 49 | + """ |
| 50 | + compile_flags = [] |
| 51 | + artifact_path = None |
| 52 | + for spec in compile_spec: |
| 53 | + if spec.key == "compile_flags": |
| 54 | + compile_flags.append(spec.value.decode()) |
| 55 | + if spec.key == "debug_artifact_path": |
| 56 | + artifact_path = spec.value.decode() |
| 57 | + |
| 58 | + # Pass on the TOSA flatbuffer to the vgf compiler. |
| 59 | + binary = vgf_compile(tosa_flatbuffer, compile_flags, artifact_path, tag_name) |
| 60 | + return binary |
| 61 | + |
| 62 | + @staticmethod |
| 63 | + def preprocess( |
| 64 | + edge_program: ExportedProgram, |
| 65 | + compile_spec: List[CompileSpec], |
| 66 | + ) -> PreprocessResult: |
| 67 | + logger.info(f"{VgfBackend.__name__} preprocess") |
| 68 | + |
| 69 | + # deduce TOSA compile_spec from VGF compile spec. We get a new |
| 70 | + # compile spec list, containing only elements relevant for the |
| 71 | + # TOSABackend. |
| 72 | + tosa_compile_spec = TOSABackend.filter_tosa_compile_specs(compile_spec) |
| 73 | + |
| 74 | + # Backends doesn't allow inheritance, as stated in comments in exir/backend/backend_api.py |
| 75 | + # ('All backend implementation are final...'), so use composition instead. |
| 76 | + # preprocess returns the serialized TOSA flatbuffer in .processed_bytes, |
| 77 | + # which can be passed on to next compilation step. |
| 78 | + tosa_preprocess = TOSABackend.preprocess(edge_program, tosa_compile_spec) |
| 79 | + |
| 80 | + tag_name = arm_get_first_delegation_tag(edge_program.graph_module) |
| 81 | + |
| 82 | + binary = VgfBackend._compile_tosa_flatbuffer( |
| 83 | + tosa_preprocess.processed_bytes, compile_spec, tag_name |
| 84 | + ) |
| 85 | + |
| 86 | + return PreprocessResult(processed_bytes=binary) |
| 87 | + |
| 88 | + |
| 89 | +def vgf_compile( |
| 90 | + tosa_flatbuffer: bytes, |
| 91 | + compile_flags: List[str], |
| 92 | + artifact_path: str | None = None, |
| 93 | + tag_name: str = "", |
| 94 | +): |
| 95 | + with tempfile.TemporaryDirectory() as tmpdir: |
| 96 | + |
| 97 | + # We currently write out a flatbuffer as input to the converter |
| 98 | + tosaname = f"output_{tag_name}.tosa" |
| 99 | + tosa_path = os.path.join(tmpdir, tosaname) |
| 100 | + with open(tosa_path, "wb") as f: |
| 101 | + f.write(tosa_flatbuffer) |
| 102 | + |
| 103 | + additional_flags = " ".join(compile_flags) |
| 104 | + vgf_path = tosa_path + ".vgf" |
| 105 | + conversion_command = ( |
| 106 | + f"converter-backend {additional_flags} -i {tosa_path} -o {vgf_path}" |
| 107 | + ) |
| 108 | + try: |
| 109 | + subprocess.run( |
| 110 | + [conversion_command], shell=True, check=True, capture_output=True |
| 111 | + ) |
| 112 | + except subprocess.CalledProcessError as process_error: |
| 113 | + raise RuntimeError( |
| 114 | + f"Vgf compiler ('{conversion_command}') failed with error:\n \ |
| 115 | + {process_error.stderr.decode()}\n \ |
| 116 | + Stdout:\n{process_error.stdout.decode()}" |
| 117 | + ) |
| 118 | + |
| 119 | + if artifact_path is not None: |
| 120 | + logger.info(f"Emitting debug output to: {vgf_path=}") |
| 121 | + os.makedirs(artifact_path, exist_ok=True) |
| 122 | + cp = f"cp {vgf_path} {artifact_path}" |
| 123 | + subprocess.run(cp, shell=True, check=True, capture_output=False) |
| 124 | + |
| 125 | + vgf_bytes = open(vgf_path, "rb").read() |
| 126 | + return vgf_bytes |
0 commit comments