arm_compile_spec.py

# Copyright 2023-2026 Arm Limited and/or its affiliates.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


#
# Main implementation of AoT flow to partition and preprocess for Arm target
# backends. Converts via TOSA as an intermediate form supported by AoT and
# JIT compiler flows.
#

import json
import warnings
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum

from executorch.backends.arm.common.pipeline_config import (
    ArmPassPipelineConfig,
    SoftmaxDecompositionConfig,
)
from executorch.backends.arm.tosa import TosaSpecification
from executorch.exir._warnings import deprecated

from executorch.exir.backend.compile_spec_schema import CompileSpec


@dataclass(init=False)
class ArmCompileSpec(ABC):
    class DebugMode(Enum):
        JSON = 1
        TOSA = 2

    tosa_spec: TosaSpecification
    compiler_flags: list[str] = field(default_factory=list)
    path_for_intermediates: str | None = None
    tosa_debug_mode: DebugMode | None = None
    tosa_dev_mode: bool | None = None

    _TOSA_SPEC_KEY = "tosa_spec"
    _COMPILE_FLAGS_KEY = "compile_flags"
    _OUTPUT_FORMAT_KEY = "output_format"
    _DEBUG_ARTIFACT_KEY = "debug_artifact_path"
    _DEBUG_MODE_KEY = "dump_debug_info"
    _OUTPUT_REORDER_KEY = "ouput_reorder_workaround"
    _TRANSFORM_PIPELINE_CONFIG_KEY = "transform_pipeline_config"
    _TOSA_DEV_MODE = "tosa_sw_dev_mode"

    def _set_compile_specs(
        self,
        tosa_spec: TosaSpecification,
        compiler_flags: list[str],
        path_for_intermediates: str | None = None,
        tosa_debug_mode: DebugMode | None = None,
        output_order_workaround: bool = False,
        pipeline_config: ArmPassPipelineConfig | None = None,
        tosa_dev_mode: bool | None = None,
    ):
        """Set all values of dataclass directly."""
        self.tosa_spec = tosa_spec
        self.compiler_flags = compiler_flags
        self.path_for_intermediates = path_for_intermediates
        self.tosa_debug_mode = tosa_debug_mode
        self._pipeline_config = pipeline_config
        self.output_order_workaround = output_order_workaround
        self.tosa_dev_mode = tosa_dev_mode
        if output_order_workaround:
            warnings.warn(
                "ArmCompileSpec(output_order_workaround=True) is deprecated and will be "
                "removed in v1.5; please remove this flag.",
                DeprecationWarning,
                stacklevel=2,
            )

    @classmethod
    def _from_list(cls, compile_specs: list[CompileSpec]):  # noqa: C901
        tosa_spec: TosaSpecification | None = None
        output_format: str | None = None
        compiler_flags: list[str] | None = None
        path_for_intermediates: str | None = None
        tosa_debug_mode: ArmCompileSpec.DebugMode | None = None
        output_order_workaround: bool = False
        pipeline_config: ArmPassPipelineConfig | None = None
        tosa_dev_mode: bool | None = None
        unknown_specs: dict[str, str] = {}
        for spec in compile_specs:
            key = spec.key
            val = (
                spec.value.decode()
                if isinstance(spec.value, (bytes, bytearray))
                else spec.value
            )
            if key == ArmCompileSpec._TOSA_SPEC_KEY:
                if tosa_spec is not None:
                    raise ValueError("More than one tosa_spec entry in compile spec.")
                tosa_spec = TosaSpecification.create_from_string(val)
            elif key == ArmCompileSpec._COMPILE_FLAGS_KEY:
                if compiler_flags is not None:
                    raise ValueError(
                        "More than one compiler flags entry in compile spec."
                    )
                compiler_flags = val.split(" ")
            elif key == ArmCompileSpec._OUTPUT_FORMAT_KEY:
                if output_format is not None:
                    raise ValueError(
                        "More than one output format entry in compile spec."
                    )
                output_format = val
            elif key == ArmCompileSpec._DEBUG_ARTIFACT_KEY:
                if path_for_intermediates is not None:
                    raise ValueError(
                        "More than one debug artifact path entry in compile spec."
                    )
                path_for_intermediates = val
            elif key == ArmCompileSpec._DEBUG_MODE_KEY:
                if tosa_debug_mode is not None:
                    raise ValueError(
                        "More than one tosa_debug_mode entry in compile spec."
                    )
                tosa_debug_mode = ArmCompileSpec.DebugMode[val]
            elif key == ArmCompileSpec._OUTPUT_REORDER_KEY:
                output_order_workaround = val  # type: ignore[assignment]
                if output_order_workaround:
                    warnings.warn(
                        "The 'output_order_workaround' compile spec entry is deprecated and will be removed in v1.5; please remove this entry.",
                        DeprecationWarning,
                        stacklevel=2,
                    )
            elif key == ArmCompileSpec._TRANSFORM_PIPELINE_CONFIG_KEY:
                if pipeline_config is not None:
                    raise ValueError(
                        "More than one transform pipeline entry in compile spec."
                    )
                pipeline_config = ArmPassPipelineConfig.from_dict(json.loads(val))
            elif key == ArmCompileSpec._TOSA_DEV_MODE:
                if tosa_dev_mode is not None:
                    raise ValueError(
                        "More than one tosa_sw_dev_mode entry in compile spec."
                    )
                raw = bytes(spec.value)
                if raw == b"\x01":
                    tosa_dev_mode = True
                elif raw == b"\x00":
                    tosa_dev_mode = False
                else:
                    raise ValueError(
                        f"Invalid tosa_sw_dev_mode byte value: {raw!r}, expected b'\\x00' or b'\\x01'."
                    )
            else:
                unknown_specs[key] = val

        if tosa_spec is None:
            raise ValueError("No tosa_spec in compile spec.")
        if output_format is None:
            raise ValueError("No output_format in compile spec.")
        if output_format != cls._get_output_format():
            raise ValueError(
                f"Incorrect output format '{output_format}' for {cls.__name__}, expected '{cls._get_output_format()}'"
            )
        if compiler_flags is None:
            compiler_flags = []

        # Create new object from class, but bypass __init__ and use _set_compile_specs instead.
        compile_spec = cls.__new__(cls)
        compile_spec._set_compile_specs(
            tosa_spec=tosa_spec,
            compiler_flags=compiler_flags,
            path_for_intermediates=path_for_intermediates,
            tosa_debug_mode=tosa_debug_mode,
            output_order_workaround=output_order_workaround,
            pipeline_config=pipeline_config,
            tosa_dev_mode=tosa_dev_mode,
        )
        cls._from_list_hook(compile_spec, unknown_specs)
        compile_spec._validate()
        return compile_spec

    @classmethod
    def _from_list_hook(cls, compile_spec, specs: dict[str, str]):  # noqa: B027
        """Allows subclasses to hook into parsing compile spec lists."""
        pass

    @abstractmethod
    def _validate(self):
        """Throws an error if the compile spec is not valid."""

    def _to_list(self):
        """Get the ArmCompileSpec in list form."""
        if not self.tosa_spec:
            raise ValueError("tosa_spec must be set before calling _to_list()")

        # Always supply a TOSA version
        compile_spec = [
            CompileSpec(ArmCompileSpec._TOSA_SPEC_KEY, str(self.tosa_spec).encode())
        ]

        # Add compile flags, these are backend specific, refer to the backend
        # documentation.
        if len(self.compiler_flags) > 0:
            compile_spec += [
                CompileSpec(
                    ArmCompileSpec._COMPILE_FLAGS_KEY,
                    " ".join(self.compiler_flags).encode(),
                ),
            ]

        # Add output format to identify kind of compile spec.
        compile_spec.append(
            CompileSpec(
                ArmCompileSpec._OUTPUT_FORMAT_KEY, self._get_output_format().encode()
            )
        )

        if self.path_for_intermediates is not None:
            compile_spec.append(
                CompileSpec(
                    ArmCompileSpec._DEBUG_ARTIFACT_KEY,
                    self.path_for_intermediates.encode(),
                )
            )

        if self.tosa_debug_mode is not None:
            if not self.path_for_intermediates:
                raise ValueError(
                    "dump_debug_info() must be used in conjunction with dump_intermediate_artifacts_to()"
                )

            compile_spec.append(
                CompileSpec(
                    ArmCompileSpec._DEBUG_MODE_KEY, self.tosa_debug_mode.name.encode()
                )
            )

        if not self.output_order_workaround:
            compile_spec.append(
                CompileSpec(
                    ArmCompileSpec._OUTPUT_REORDER_KEY,
                    bytes(self.output_order_workaround),
                )
            )

        if self._pipeline_config is not None and not self._pipeline_config.is_default():
            compile_spec.append(
                CompileSpec(
                    ArmCompileSpec._TRANSFORM_PIPELINE_CONFIG_KEY,
                    self._pipeline_config.serialize(),
                )
            )

        if self.tosa_dev_mode is not None:
            compile_spec.append(
                CompileSpec(
                    ArmCompileSpec._TOSA_DEV_MODE,
                    b"\x01" if self.tosa_dev_mode else b"\x00",
                )
            )

        return compile_spec

    def _get_pass_pipeline_config(self) -> ArmPassPipelineConfig:
        """Returns configuration that controls how the Arm pass pipeline should
        behave.

        Subclasses may override to tweak defaults for specific targets.

        """
        if self._pipeline_config is None:
            self._pipeline_config = self._create_default_pipeline_config()
        return self._pipeline_config

    def set_pass_pipeline_config(self, config: ArmPassPipelineConfig) -> None:
        """Sets the configuration that controls how the Arm pass pipeline should
        behave. Subclasses may override to tweak defaults for specific targets.

        Args:
            config: The custom ArmPassPipelineConfig to set.

        """
        self._pipeline_config = config

    def _create_default_pipeline_config(self) -> ArmPassPipelineConfig:
        config = ArmPassPipelineConfig()
        if self.tosa_spec.is_U55_subset:
            # Keep U55 on STABLE instead of the generic MASKED default:
            # MASKED also enables masked_fill decomposition, which lowers to
            # where/full_like and is not a good default fit for U55.
            config.softmax = SoftmaxDecompositionConfig.STABLE
        return config

    def _get_intermediate_path(self) -> str | None:
        """Gets the path used for dumping intermediate results such as tosa and
        pte.

        Returns:
            Path where intermediate results are saved.

        """
        return self.path_for_intermediates

    def dump_intermediate_artifacts_to(self, output_path: str | None):
        """Sets a path for dumping intermediate results during such as tosa and
        pte.

        Args:
            output_path: Path to dump intermediate results to.

        """
        self.path_for_intermediates = output_path
        return self

    def dump_debug_info(self, debug_mode: DebugMode | None):
        """Dump debugging information into the intermediates path.

        Args:
            debug_mode: The debug mode to use for dumping debug information.

        """
        self.tosa_debug_mode = debug_mode
        return self

    def _set_tosa_dev_mode(self, tosa_dev_mode: bool):
        """Sets whether to enable TOSA software development mode.

        Args:
            tosa_dev_mode: Boolean indicating whether to enable TOSA software development mode.

        """
        self.tosa_dev_mode = tosa_dev_mode
        return self

    @deprecated(
        "set_output_order_workaround() is deprecated and will be removed in v1.5; please remove this call."
    )
    def set_output_order_workaround(self, output_order_workaround: bool):
        """Sets whether to apply the output order workaround.

        Args:
            output_order_workaround: Boolean indicating whether to apply the workaround.

        """
        self.output_order_workaround = output_order_workaround
        return self

    @deprecated(
        "get_output_order_workaround() is deprecated and will be removed in v1.5; please remove this call."
    )
    def get_output_order_workaround(self) -> bool:
        """Gets whether the output order workaround is being applied."""
        return self.output_order_workaround

    @classmethod
    @abstractmethod
    def _get_output_format(cls) -> str:
        """Returns a constant string that is the output format of the class."""