From e51c8ce569f7d0cd28606572f96260b93997b02c Mon Sep 17 00:00:00 2001
From: Dmitrii Agibov <dmitrii.agibov@arm.com>
Date: Thu, 21 Jul 2022 14:07:12 +0100
Subject: MLIA-550 Enable TOSA checker in the CLI

Change-Id: I9dfa9f17a93ac24be0aa9ad68e4a72b79aecfb9d
---
 src/mlia/cli/main.py    | 6 ++++--
 src/mlia/cli/options.py | 6 +++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

(limited to 'src/mlia')

diff --git a/src/mlia/cli/main.py b/src/mlia/cli/main.py
index 33fcdeb..f8fc00c 100644
--- a/src/mlia/cli/main.py
+++ b/src/mlia/cli/main.py
@@ -4,6 +4,7 @@
 import argparse
 import logging
 import sys
+from functools import partial
 from inspect import signature
 from pathlib import Path
 from typing import Dict
@@ -44,6 +45,7 @@ Supported targets:
 
  - Ethos-U55 <op compatibility, perf estimation, model opt>
  - Ethos-U65 <op compatibility, perf estimation, model opt>
+ - TOSA      <op compatibility>
 
 """.strip()
 
@@ -79,7 +81,7 @@ def get_commands() -> List[CommandInfo]:
             performance,
             ["perf"],
             [
-                add_target_options,
+                partial(add_target_options, profiles_to_skip=["tosa"]),
                 add_tflite_model_options,
                 add_output_options,
                 add_debug_options,
@@ -90,7 +92,7 @@ def get_commands() -> List[CommandInfo]:
             optimization,
             ["opt"],
             [
-                add_target_options,
+                partial(add_target_options, profiles_to_skip=["tosa"]),
                 add_keras_model_options,
                 add_multi_optimization_options,
                 add_output_options,
diff --git a/src/mlia/cli/options.py b/src/mlia/cli/options.py
index dc5cb73..29a0d89 100644
--- a/src/mlia/cli/options.py
+++ b/src/mlia/cli/options.py
@@ -16,9 +16,13 @@ from mlia.utils.filesystem import get_supported_profile_names
 from mlia.utils.types import is_number
 
 
-def add_target_options(parser: argparse.ArgumentParser) -> None:
+def add_target_options(
+    parser: argparse.ArgumentParser, profiles_to_skip: Optional[List[str]] = None
+) -> None:
     """Add target specific options."""
     target_profiles = get_supported_profile_names()
+    if profiles_to_skip:
+        target_profiles = [tp for tp in target_profiles if tp not in profiles_to_skip]
 
     default_target_profile = None
     default_help = ""
-- 
cgit v1.2.1