Add a CLI component to enable rewrites

* Add flags for rewrite (--rewrite, --rewrite-start,
  --rewrite-end, --rewrite-target)
* Refactor CLI interfaces to accept tflite models with optimize for
  rewrite, keras models with optimize for clustering and pruning
* Refactor and move common.py and select.py out of the folder
  nn/tensorflow/optimizations
* Add file nn/rewrite/core/rewrite.py as placeholder
* Update/add unit tests
* Refactor OptimizeModel in ethos_u/data_collection.py
  for accepting tflite model case
* Extend the logic so that if "--rewrite" is specified, we don't add
  pruning to also accept TFLite models.
* Update README.md

Resolves: MLIA-750, MLIA-854, MLIA-865
Signed-off-by: Benjamin Klimczak <benjamin.klimczak@arm.com>
Change-Id: I67d85f71fa253d2bad4efe304ad8225970b9622c

9 files changed, 203 insertions, 28 deletions

diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py
index f3213c4..6765a53 100644
--- a/tests/test_cli_commands.py
+++ b/tests/test_cli_commands.py
@@ -3,6 +3,7 @@
 """Tests for cli.commands module."""
 from __future__ import annotations
 
+from contextlib import ExitStack as does_not_raise
 from pathlib import Path
 from typing import Any
 from unittest.mock import call
@@ -49,37 +50,148 @@ def test_performance_unknown_target(
 
 
 @pytest.mark.parametrize(
-    "target_profile, pruning, clustering, pruning_target, clustering_target",
+    "target_profile, pruning, clustering, pruning_target, clustering_target, "
+    "rewrite, rewrite_target, rewrite_start, rewrite_end, expected_error",
     [
-        ["ethos-u55-256", True, False, 0.5, None],
-        ["ethos-u65-512", False, True, 0.5, 32],
-        ["ethos-u55-256", True, True, 0.5, None],
-        ["ethos-u55-256", False, False, 0.5, None],
-        ["ethos-u55-256", False, True, "invalid", 32],
+        [
+            "ethos-u55-256",
+            True,
+            False,
+            0.5,
+            None,
+            False,
+            None,
+            "node_a",
+            "node_b",
+            does_not_raise(),
+        ],
+        [
+            "ethos-u55-256",
+            False,
+            False,
+            None,
+            None,
+            True,
+            "fully_connected",
+            "node_a",
+            "node_b",
+            does_not_raise(),
+        ],
+        [
+            "ethos-u55-256",
+            True,
+            False,
+            0.5,
+            None,
+            True,
+            "fully_connected",
+            "node_a",
+            "node_b",
+            pytest.raises(
+                Exception,
+                match=(r"Only 'rewrite' is supported for TensorFlow Lite files."),
+            ),
+        ],
+        [
+            "ethos-u65-512",
+            False,
+            True,
+            0.5,
+            32,
+            False,
+            None,
+            None,
+            None,
+            does_not_raise(),
+        ],
+        [
+            "ethos-u55-256",
+            False,
+            False,
+            0.5,
+            None,
+            True,
+            "random",
+            "node_x",
+            "node_y",
+            pytest.raises(
+                Exception,
+                match=(r"Currently only remove and fully_connected are supported."),
+            ),
+        ],
+        [
+            "ethos-u55-256",
+            False,
+            False,
+            0.5,
+            None,
+            True,
+            None,
+            "node_m",
+            "node_n",
+            pytest.raises(
+                Exception,
+                match=(
+                    r"To perform rewrite, rewrite-target, "
+                    r"rewrite-start and rewrite-end must be set."
+                ),
+            ),
+        ],
+        [
+            "ethos-u55-256",
+            False,
+            False,
+            "invalid",
+            None,
+            True,
+            "remove",
+            None,
+            "node_end",
+            pytest.raises(
+                Exception,
+                match=(
+                    r"To perform rewrite, rewrite-target, "
+                    r"rewrite-start and rewrite-end must be set."
+                ),
+            ),
+        ],
     ],
 )
-def test_opt_valid_optimization_target(
+def test_opt_valid_optimization_target(  # pylint: disable=too-many-arguments
     target_profile: str,
     sample_context: ExecutionContext,
     pruning: bool,
     clustering: bool,
     pruning_target: float | None,
     clustering_target: int | None,
+    rewrite: bool,
+    rewrite_target: str | None,
+    rewrite_start: str | None,
+    rewrite_end: str | None,
+    expected_error: Any,
     monkeypatch: pytest.MonkeyPatch,
     test_keras_model: Path,
+    test_tflite_model: Path,
 ) -> None:
     """Test that command should not fail with valid optimization targets."""
     mock_performance_estimation(monkeypatch)
 
-    optimize(
-        ctx=sample_context,
-        target_profile=target_profile,
-        model=str(test_keras_model),
-        pruning=pruning,
-        clustering=clustering,
-        pruning_target=pruning_target,
-        clustering_target=clustering_target,
-    )
+    model_type = test_tflite_model if rewrite else test_keras_model
+
+    with expected_error:
+        optimize(
+            ctx=sample_context,
+            target_profile=target_profile,
+            model=str(model_type),
+            pruning=pruning,
+            clustering=clustering,
+            pruning_target=pruning_target,
+            clustering_target=clustering_target,
+            rewrite=rewrite,
+            rewrite_target=rewrite_target,
+            rewrite_start=rewrite_start,
+            rewrite_end=rewrite_end,
+        )
 
 
 def mock_performance_estimation(monkeypatch: pytest.MonkeyPatch) -> None:
diff --git a/tests/test_cli_helpers.py b/tests/test_cli_helpers.py
index 6d19207..494ed89 100644
--- a/tests/test_cli_helpers.py
+++ b/tests/test_cli_helpers.py
@@ -10,7 +10,7 @@ import pytest
 
 from mlia.cli.helpers import CLIActionResolver
 from mlia.cli.helpers import copy_profile_file_to_output_dir
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.select import OptimizationSettings
 
 
 class TestCliActionResolver:
diff --git a/tests/test_cli_main.py b/tests/test_cli_main.py
index 2f89268..e415284 100644
--- a/tests/test_cli_main.py
+++ b/tests/test_cli_main.py
@@ -165,6 +165,11 @@ def wrap_mock_command(mock: MagicMock, command: Callable) -> Callable:
                 pruning_target=None,
                 clustering_target=None,
                 backend=None,
+                rewrite=False,
+                rewrite_target=None,
+                rewrite_start=None,
+                rewrite_end=None,
+                dataset=None,
             ),
         ],
         [
@@ -189,6 +194,11 @@ def wrap_mock_command(mock: MagicMock, command: Callable) -> Callable:
                 pruning_target=0.5,
                 clustering_target=32,
                 backend=None,
+                rewrite=False,
+                rewrite_target=None,
+                rewrite_start=None,
+                rewrite_end=None,
+                dataset=None,
             ),
         ],
         [
@@ -210,6 +220,11 @@ def wrap_mock_command(mock: MagicMock, command: Callable) -> Callable:
                 pruning_target=None,
                 clustering_target=None,
                 backend=["some_backend"],
+                rewrite=False,
+                rewrite_target=None,
+                rewrite_start=None,
+                rewrite_end=None,
+                dataset=None,
             ),
         ],
         [
diff --git a/tests/test_cli_options.py b/tests/test_cli_options.py
index c02ef89..75ace0b 100644
--- a/tests/test_cli_options.py
+++ b/tests/test_cli_options.py
@@ -30,6 +30,7 @@ from mlia.core.typing import OutputFormat
                     "optimization_type": "pruning",
                     "optimization_target": 0.5,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 }
             ],
         ],
@@ -44,6 +45,7 @@ from mlia.core.typing import OutputFormat
                     "optimization_type": "pruning",
                     "optimization_target": 0.5,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 }
             ],
         ],
@@ -58,6 +60,7 @@ from mlia.core.typing import OutputFormat
                     "optimization_type": "clustering",
                     "optimization_target": 32,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 }
             ],
         ],
@@ -72,11 +75,13 @@ from mlia.core.typing import OutputFormat
                     "optimization_type": "pruning",
                     "optimization_target": 0.5,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 },
                 {
                     "optimization_type": "clustering",
                     "optimization_target": 32,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 },
             ],
         ],
@@ -91,6 +96,7 @@ from mlia.core.typing import OutputFormat
                     "optimization_type": "pruning",
                     "optimization_target": 0.4,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 }
             ],
         ],
@@ -117,6 +123,7 @@ from mlia.core.typing import OutputFormat
                     "optimization_type": "clustering",
                     "optimization_target": 32.2,
                     "layers_to_optimize": None,
+                    "dataset": None,
                 }
             ],
         ],
diff --git a/tests/test_nn_tensorflow_optimizations_select.py b/tests/test_nn_select.py
index f5ba6f0..31628d2 100644
--- a/tests/test_nn_tensorflow_optimizations_select.py
+++ b/tests/test_nn_select.py
@@ -11,13 +11,13 @@ import pytest
 import tensorflow as tf
 
 from mlia.core.errors import ConfigurationError
+from mlia.nn.select import get_optimizer
+from mlia.nn.select import MultiStageOptimizer
+from mlia.nn.select import OptimizationSettings
 from mlia.nn.tensorflow.optimizations.clustering import Clusterer
 from mlia.nn.tensorflow.optimizations.clustering import ClusteringConfiguration
 from mlia.nn.tensorflow.optimizations.pruning import Pruner
 from mlia.nn.tensorflow.optimizations.pruning import PruningConfiguration
-from mlia.nn.tensorflow.optimizations.select import get_optimizer
-from mlia.nn.tensorflow.optimizations.select import MultiStageOptimizer
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_target_ethos_u_advice_generation.py b/tests/test_target_ethos_u_advice_generation.py
index 772fc56..ac4e5e9 100644
--- a/tests/test_target_ethos_u_advice_generation.py
+++ b/tests/test_target_ethos_u_advice_generation.py
@@ -12,7 +12,7 @@ from mlia.core.common import DataItem
 from mlia.core.context import ExecutionContext
 from mlia.core.helpers import ActionResolver
 from mlia.core.helpers import APIActionResolver
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.select import OptimizationSettings
 from mlia.target.ethos_u.advice_generation import EthosUAdviceProducer
 from mlia.target.ethos_u.advice_generation import EthosUStaticAdviceProducer
 from mlia.target.ethos_u.data_analysis import AllOperatorsSupportedOnNPU
diff --git a/tests/test_target_ethos_u_advisor.py b/tests/test_target_ethos_u_advisor.py
index 11aefc7..20131d2 100644
--- a/tests/test_target_ethos_u_advisor.py
+++ b/tests/test_target_ethos_u_advisor.py
@@ -1,7 +1,11 @@
 # SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
 """Tests for Ethos-U MLIA module."""
+from __future__ import annotations
+
+from contextlib import ExitStack as does_not_raise
 from pathlib import Path
+from typing import Any
 
 import pytest
 
@@ -16,16 +20,53 @@ def test_advisor_metadata() -> None:
     assert EthosUInferenceAdvisor.name() == "ethos_u_inference_advisor"
 
 
-def test_unsupported_advice_categories(tmp_path: Path, test_tflite_model: Path) -> None:
+@pytest.mark.parametrize(
+    "optimization_targets, expected_error",
+    [
+        [
+            [
+                {
+                    "optimization_type": "pruning",
+                    "optimization_target": 0.5,
+                    "layers_to_optimize": None,
+                }
+            ],
+            pytest.raises(
+                Exception,
+                match="Only 'rewrite' is supported for TensorFlow Lite files.",
+            ),
+        ],
+        [
+            [
+                {
+                    "optimization_type": "rewrite",
+                    "optimization_target": "fully_connected",
+                    "layers_to_optimize": [
+                        "MobileNet/avg_pool/AvgPool",
+                        "MobileNet/fc1/BiasAdd",
+                    ],
+                }
+            ],
+            does_not_raise(),
+        ],
+    ],
+)
+def test_unsupported_advice_categories(
+    tmp_path: Path,
+    test_tflite_model: Path,
+    optimization_targets: list[dict[str, Any]],
+    expected_error: Any,
+) -> None:
     """Test that advisor should throw an exception for unsupported categories."""
-    with pytest.raises(
-        Exception, match="Optimizations are not supported for TensorFlow Lite files."
-    ):
+    with expected_error:
         ctx = ExecutionContext(
             output_dir=tmp_path, advice_category={AdviceCategory.OPTIMIZATION}
         )
 
         advisor = configure_and_get_ethosu_advisor(
-            ctx, "ethos-u55-256", str(test_tflite_model)
+            ctx,
+            "ethos-u55-256",
+            str(test_tflite_model),
+            optimization_targets=optimization_targets,
         )
         advisor.configure(ctx)
diff --git a/tests/test_target_ethos_u_data_analysis.py b/tests/test_target_ethos_u_data_analysis.py
index 80f0603..713e8ef 100644
--- a/tests/test_target_ethos_u_data_analysis.py
+++ b/tests/test_target_ethos_u_data_analysis.py
@@ -12,7 +12,7 @@ from mlia.backend.vela.compat import Operator
 from mlia.backend.vela.compat import Operators
 from mlia.core.common import DataItem
 from mlia.core.data_analysis import Fact
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.select import OptimizationSettings
 from mlia.nn.tensorflow.tflite_compat import TFLiteCompatibilityInfo
 from mlia.nn.tensorflow.tflite_compat import TFLiteCompatibilityStatus
 from mlia.nn.tensorflow.tflite_compat import TFLiteConversionError
diff --git a/tests/test_target_ethos_u_data_collection.py b/tests/test_target_ethos_u_data_collection.py
index fd824ae..6244f8b 100644
--- a/tests/test_target_ethos_u_data_collection.py
+++ b/tests/test_target_ethos_u_data_collection.py
@@ -10,7 +10,7 @@ from mlia.backend.vela.compat import Operators
 from mlia.core.context import Context
 from mlia.core.data_collection import DataCollector
 from mlia.core.errors import FunctionalityNotSupportedError
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.select import OptimizationSettings
 from mlia.target.ethos_u.config import EthosUConfiguration
 from mlia.target.ethos_u.data_collection import EthosUOperatorCompatibility
 from mlia.target.ethos_u.data_collection import EthosUOptimizationPerformance