summaryrefslogtreecommitdiff
path: root/scripts/py/vela_configs.py
blob: b4af9fdddd7a45f5d2cd0f9843077357cbecf73a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
#!/usr/bin/env python3
#  SPDX-FileCopyrightText:  Copyright 2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
#  SPDX-License-Identifier: Apache-2.0
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
"""
Classes to represent NPU configurations for Vela
"""
import itertools
import typing
from dataclasses import dataclass

# The internal SRAM size for Corstone-300 implementation on MPS3 specified by AN552
# The internal SRAM size for Corstone-310 implementation on MPS3 specified by AN555
# is 4MB, but we are content with the 2MB specified below.
MPS3_MAX_SRAM_SZ = 2 * 1024 * 1024  # 2 MiB (2 banks of 1 MiB each)


@dataclass(frozen=True)
class NpuConfig:
    """
    Represents a Vela configuration for an NPU
    """
    name_prefix: str
    macs: int
    processor_id: str
    prefix_id: str
    memory_mode: str
    system_config: str
    arena_cache_size: int = 0

    @property
    def config_name(self) -> str:
        """
        Get the name of the configuration

        For example: "ethos-u55-128" would represent the Ethos-U55 NPU
        with a 128 MAC configuration.

        :return:    The NPU configuration name.
        """
        return f"{self.name_prefix}-{self.macs}"

    @property
    def config_id(self) -> str:
        """
        Get the configuration id as a string

        For example: "Y256" would represent the Ethos-U65 NPU
        with a 256 MAC configuration.

        :return:    The NPU configuration id.
        """
        return f"{self.prefix_id}{self.macs}"

    def overwrite_arena_cache_size(self, arena_cache_size):
        """
        Get a new NPU configuration with the specified
        arena cache size.

        By default, we use the `arena_cache_size` value in the
        `default_vela.ini` configuration file.

        :param  arena_cache_size:   The new arena cache size value.
        :return:                    A new NPU configuration with the new
                                    arena cache size value.
        """
        value = arena_cache_size

        if value == 0:
            value = MPS3_MAX_SRAM_SZ if self.memory_mode == "Shared_Sram" else None

        return NpuConfig(
            **{**self.__dict__, **{"arena_cache_size": value}}
        )


@dataclass(frozen=True)
class NpuConfigs:
    """
    Represents a collection of NPU configurations.
    """
    configs: typing.Dict[str, typing.Dict[int, NpuConfig]]

    @staticmethod
    def create(*configs: NpuConfig):
        """
        Create a new collection with the specified NPU configurations.

        :param configs: NPU configuration objects to add to the collection.
        :return:        A new collection of NPU configurations.
        """
        _configs = {}

        # Internal data structure of nested dictionaries based on
        # NPU name and MAC configuration, e.g.:
        #   _configs["ethos-u55"][128]

        for c in configs:
            if c.name_prefix not in _configs:
                _configs[c.name_prefix] = {}
            _configs[c.name_prefix][c.macs] = c
        return NpuConfigs(configs=_configs)

    def get(self, name_prefix: str, macs: typing.Union[int, str]) -> typing.Optional[NpuConfig]:
        """
        Get an NPU configuration by name prefix and MAC configuration.

        :param name_prefix: The name prefix, e.g. "ethos-u55".
        :param macs:        The MAC configuration, e.g. 128.
        :return:            The matching NPU configuration, or None if no such configuration
                            exists in the collection.
        """
        configs_for_name = self.configs.get(name_prefix)
        if not configs_for_name:
            return None
        return configs_for_name.get(int(macs))

    def get_by_name(self, name: str) -> typing.Optional[NpuConfig]:
        """
        Get an NPU configuration by name.

        :param name:    The NPU configuration name, e.g. "ethos-u55-128".
        :return:        The matching NPU configuration, or None if no such configuration
                        exists in the collection.
        """
        name_prefix, macs = name.rsplit("-", 1)
        return self.get(name_prefix, macs)

    @property
    def names(self):
        """
        Return a list of all NPU configuration names in the collection.

        :return:    The list of NPU configuration names.
        """
        return list(itertools.chain.from_iterable([
            [f"{c.name_prefix}-{c.macs}" for c in config.values()]
            for config in self.configs.values()
        ]))