Source code for afe.backends.ev.ev_checkers

#########################################################
# Copyright (C) 2022 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Ljubomir Papuga
#########################################################
from typing import Dict, Callable, List, Set
from afe.backends import Backend, ExprCheckInfo, BaseChecker
from afe.backends.backend_checker import Decision, Accept, Reject, Predicate

# The list of TVM expressions supported by EV
[docs] EV_SUPPORTED_REGISTRY = [ "layout_transform", "nn.batch_flatten", "reshape" ]
def _make_ev_supported_test(supported_operators: Set[str]) -> Predicate: """ Make a function that decides whether an operator is supported on the EV backend. :param supported_operators: TVM names of operators supported on the EV backend. Normally the names are the contents of EV_SUPPORTED_REGISTRY; however, some tests use different supported operators. Other code may mutate the list to change which operators are supported. :return: Predicate that decides whether an expression is supported on the EV backend. It reads supported_operators when it is called. Other code may mutate supported_operators to change the predicate's behavior. """ def do_ev_supported_test(args: ExprCheckInfo) -> Decision: """ Decide whether the TVM expression described by args can be assigned to the EV backend. """ if args.name in supported_operators: return Accept() else: return Reject([-1]) return do_ev_supported_test
[docs] class EVChecker(BaseChecker): """ Currently, no EV plugins are supported. When adding support for TVM operators for EV, developers should add operator name to EV_SUPPORTED_REGISTRY. This should include any supported composite operators. Composite operators will be decomposed if they are to be run on EV. :param _checkers_name: string :param _backend : a Backend class of EV """ _checkers_name: str = "EV Checker" _backend: Backend = Backend.EV _registry: Set[str] = set(EV_SUPPORTED_REGISTRY) # Mutable field _predicate: Predicate = _make_ev_supported_test(_registry) # Methods used in test cases to allow EVChecker to map certain operators to EV. # Should be used only from tests.network.utils.override_ev_checker decorator. # The user should take responsibility of handling the translation from RelayIR # expression to EVOp, if the operator is not handled in _TVM_EV_OPERATOR_CONVERSIONS_LIST. @classmethod
[docs] def add_operator_to_registry(cls, operator: str): cls._registry.add(operator)
@classmethod
[docs] def remove_operator_from_registry(cls, operator: str): cls._registry.discard(operator)