NXP backend: Remove IR optimization to remove dead branches. (#13574)

### Summary
This PR replaces an IR optimization that removes dead code from the
model, by an equivalent executorch call.

### Test plan
Unit test provided in `backends/nxp/tests/test_removing_dead_code.py`.


cc @digantdesai @JakeStevens @robert-kalmar

Author: MartinPavella

backends/nxp/backend/ir/tflite_optimizer/optimizations/eliminate_dead_branches.py

@@ -1,6 +1,6 @@
 #
 # Copyright 2023 Martin Pavella
-# Copyright 2024 NXP
+# Copyright 2024-2025 NXP
 #
 # License: MIT
 # See the LICENSE_MIT for more details.
@@ -14,9 +14,6 @@
 from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.combine_hard_sigmoid_and_mul_to_hard_swish import (
     CombineHardSigmoidAndMulIntoHardSwish,
 )
-from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.eliminate_dead_branches import (
-    EliminateDeadBranches,
-)
 from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.fuse_activation_functions import (
     FuseActivationFunctions,
 )
@@ -57,7 +54,6 @@ class Optimization(Enum):
     FUSE_PARALLEL_QUANTIZE_OPERATORS = 8
 
     REMOVE_UNUSED_TENSORS = 10
-    ELIMINATE_DEAD_BRANCHES = 11
     PERMUTE_FULLY_CONNECTED_WEIGHTS_AFTER_RESHAPE = 12
 
     MOVE_ACTIVATION_BEFORE_CONCAT = 15
@@ -115,9 +111,6 @@ def __init__(
             Optimization.REMOVE_UNUSED_TENSORS: RemoveUnusedTensorsAndBuffers(
                 builder, conversion_config
             ),
-            Optimization.ELIMINATE_DEAD_BRANCHES: EliminateDeadBranches(
-                builder, conversion_config
-            ),
             Optimization.PERMUTE_FULLY_CONNECTED_WEIGHTS_AFTER_RESHAPE: PermuteFullyConnectedWeightsAfterReshape(
                 builder, conversion_config
             ),

backends/nxp/backend/ir/tflite_optimizer/optimizer.py

@@ -7,6 +7,7 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
+
 from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import (
     NeutronAtenPassManager,
 )
@@ -242,8 +243,13 @@ def __init__(self):
     def transform_for_annotation(
         self, model: torch.fx.GraphModule
     ) -> torch.fx.GraphModule:
-        pass_runner = NeutronAtenPassManager()
-        return pass_runner(model).graph_module
+        model.graph.eliminate_dead_code()  # Remove dead code to simplify the graph for the passes.
+
+        model = NeutronAtenPassManager()(model).graph_module
+
+        model.graph.eliminate_dead_code()  # Remove dead code again, in case it was created by the passes.
+
+        return model
 
     def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         self._annotate_inputs(model)

backends/nxp/quantizer/neutron_quantizer.py

@@ -0,0 +1,60 @@
+# Copyright 2025 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import numpy as np
+import pytest
+import torch
+
+from executorch.backends.nxp.tests.executorch_pipeline import _quantize_model
+from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops
+
+
+@pytest.fixture(autouse=True)
+def reseed_model_per_test_run():
+    torch.manual_seed(42)
+    np.random.seed(23)
+
+
+class DeadCodeModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.eval()
+
+    def forward(self, x):
+        _ = torch.add(x, x)  # Dead code
+        return torch.mul(x, x)
+
+
+class TestRemovingDeadCode(unittest.TestCase):
+    __test__ = False  # Prevent interfering with PyTest tests
+
+    def test_removing_dead_code(self):
+        input_shape = (42,)
+        example_inputs = (torch.ones(input_shape),)
+        model = DeadCodeModule()
+
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
+
+        # Make sure the model contains the dead code.
+        assert graph_contains_any_of_ops(
+            exir_program_aten.module().graph, [torch.ops.aten.add.Tensor]
+        )
+
+        # The `NeutronQuantizer` should remove the dead code in the `transform_for_annotation()` method.
+        exir_program_aten_quant = _quantize_model(
+            exir_program_aten.module(), [example_inputs]
+        )
+
+        # Make sure the is no `add` operation in the graph anymore.
+        assert not any(
+            "add" in str(node.target) for node in exir_program_aten_quant.graph.nodes
+        )
+
+    @classmethod
+    def setUpClass(cls):
+        torch.manual_seed(23)
+        np.random.seed(23)