path: root/src/mlia/target/ethos_u/performance.py

# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Performance estimation."""
from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Union

import mlia.backend.vela.compiler as vela_comp
import mlia.backend.vela.performance as vela_perf
from mlia.backend.corstone import is_corstone_backend
from mlia.backend.corstone.performance import estimate_performance
from mlia.core.context import Context
from mlia.core.performance import PerformanceEstimator
from mlia.nn.select import OptimizationSettings
from mlia.nn.tensorflow.config import get_tflite_model
from mlia.nn.tensorflow.config import ModelConfiguration
from mlia.target.ethos_u.config import EthosUConfiguration
from mlia.target.registry import supported_backends
from mlia.utils.logging import log_action

logger = logging.getLogger(__name__)


@dataclass
class NPUCycles:
    """NPU cycles metrics."""

    npu_active_cycles: int
    npu_idle_cycles: int
    npu_total_cycles: int
    npu_axi0_rd_data_beat_received: int
    npu_axi0_wr_data_beat_written: int
    npu_axi1_rd_data_beat_received: int


BYTES_PER_KILOBYTE = 1024


class MemorySizeType(Enum):
    """Memory size type enumeration."""

    BYTES = 0
    KILOBYTES = 1


@dataclass
class MemoryUsage:
    """Memory usage metrics."""

    sram_memory_area_size: int | float
    dram_memory_area_size: int | float
    unknown_memory_area_size: int | float
    on_chip_flash_memory_area_size: int | float
    off_chip_flash_memory_area_size: int | float
    memory_size_type: MemorySizeType = MemorySizeType.BYTES

    _default_columns = [
        "SRAM used",
        "DRAM used",
        "Unknown memory used",
        "On chip flash used",
        "Off chip flash used",
    ]

    def in_kilobytes(self) -> MemoryUsage:
        """Return memory usage with values in kilobytes."""
        if self.memory_size_type == MemorySizeType.KILOBYTES:
            return self

        kilobytes = [
            value / BYTES_PER_KILOBYTE
            for value in [
                self.sram_memory_area_size,
                self.dram_memory_area_size,
                self.unknown_memory_area_size,
                self.on_chip_flash_memory_area_size,
                self.off_chip_flash_memory_area_size,
            ]
        ]

        return MemoryUsage(
            *kilobytes,  # type: ignore
            memory_size_type=MemorySizeType.KILOBYTES,
        )


@dataclass
class PerformanceMetrics:
    """Performance metrics."""

    target_config: EthosUConfiguration
    npu_cycles: NPUCycles | None
    memory_usage: MemoryUsage | None

    def in_kilobytes(self) -> PerformanceMetrics:
        """Return metrics with memory usage in KiB."""
        if self.memory_usage is None:
            return PerformanceMetrics(
                self.target_config, self.npu_cycles, self.memory_usage
            )

        return PerformanceMetrics(
            self.target_config, self.npu_cycles, self.memory_usage.in_kilobytes()
        )


@dataclass
class OptimizationPerformanceMetrics:
    """Optimization performance metrics."""

    original_perf_metrics: PerformanceMetrics
    optimizations_perf_metrics: list[
        tuple[list[OptimizationSettings], PerformanceMetrics]
    ]


class VelaPerformanceEstimator(
    PerformanceEstimator[Union[Path, ModelConfiguration], MemoryUsage]
):
    """Vela based performance estimator."""

    def __init__(self, context: Context, target_config: EthosUConfiguration) -> None:
        """Init Vela based performance estimator."""
        self.context = context
        self.target = target_config

    def estimate(self, model: Path | ModelConfiguration) -> MemoryUsage:
        """Estimate performance."""
        with log_action("Getting the memory usage metrics ..."):
            model_path = (
                Path(model.model_path)
                if isinstance(model, ModelConfiguration)
                else model
            )

            vela_perf_metrics = vela_perf.estimate_performance(
                model_path, self.target.compiler_options
            )

            return MemoryUsage(
                vela_perf_metrics.sram_memory_area_size,
                vela_perf_metrics.dram_memory_area_size,
                vela_perf_metrics.unknown_memory_area_size,
                vela_perf_metrics.on_chip_flash_memory_area_size,
                vela_perf_metrics.off_chip_flash_memory_area_size,
            )


class CorstonePerformanceEstimator(
    PerformanceEstimator[Union[Path, ModelConfiguration], NPUCycles]
):
    """Corstone-based performance estimator."""

    def __init__(
        self, context: Context, target_config: EthosUConfiguration, backend: str
    ) -> None:
        """Init Corstone-based performance estimator."""
        self.context = context
        self.target_config = target_config
        self.backend = backend

    def estimate(self, model: Path | ModelConfiguration) -> NPUCycles:
        """Estimate performance."""
        with log_action(f"Getting the performance metrics for '{self.backend}' ..."):
            logger.info(
                "WARNING: This task may require several minutes "
                "(press ctrl-c to interrupt)"
            )

            model_path = (
                Path(model.model_path)
                if isinstance(model, ModelConfiguration)
                else model
            )

            optimized_model_path = self.context.get_model_path(
                f"{model_path.stem}_vela.tflite"
            )

            vela_comp.optimize_model(
                model_path, self.target_config.compiler_options, optimized_model_path
            )

            corstone_perf_metrics = estimate_performance(
                self.target_config.target,
                self.target_config.mac,
                optimized_model_path,
                self.backend,
            )

            return NPUCycles(
                corstone_perf_metrics.npu_active_cycles,
                corstone_perf_metrics.npu_idle_cycles,
                corstone_perf_metrics.npu_total_cycles,
                corstone_perf_metrics.npu_axi0_rd_data_beat_received,
                corstone_perf_metrics.npu_axi0_wr_data_beat_written,
                corstone_perf_metrics.npu_axi1_rd_data_beat_received,
            )


class EthosUPerformanceEstimator(
    PerformanceEstimator[Union[Path, ModelConfiguration], PerformanceMetrics]
):
    """Ethos-U performance estimator."""

    def __init__(
        self,
        context: Context,
        target_config: EthosUConfiguration,
        backends: list[str] | None = None,
    ) -> None:
        """Init performance estimator."""
        self.context = context
        self.target_config = target_config
        if backends is None:
            backends = ["vela"]  # Only Vela is always available as default
        ethos_u_backends = supported_backends(target_config.target)
        for backend in backends:
            if backend != "vela" and backend not in ethos_u_backends:
                raise ValueError(
                    f"Unsupported backend '{backend}'. "
                    f"Only 'Vela' and {ethos_u_backends} "
                    "are supported."
                )
        self.backends = set(backends)

    def estimate(self, model: Path | ModelConfiguration) -> PerformanceMetrics:
        """Estimate performance."""
        model_path = (
            Path(model.model_path) if isinstance(model, ModelConfiguration) else model
        )

        tflite_model = get_tflite_model(model_path, self.context)

        memory_usage = None
        npu_cycles = None
        for backend in self.backends:
            if backend == "vela":
                vela_estimator = VelaPerformanceEstimator(
                    self.context, self.target_config
                )
                memory_usage = vela_estimator.estimate(tflite_model)
            elif is_corstone_backend(backend):
                corstone_estimator = CorstonePerformanceEstimator(
                    self.context, self.target_config, backend
                )
                npu_cycles = corstone_estimator.estimate(tflite_model)
            else:
                logger.warning(
                    "Backend '%s' is not supported for Ethos-U performance "
                    "estimation.",
                    backend,
                )

        return PerformanceMetrics(self.target_config, npu_cycles, memory_usage)