| #!/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() |
| ])) |