NXP backend: added aten.slice support

backends/nxp/backend/edge_program_converter.py

@@ -41,6 +41,7 @@
     exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter,  # noqa F405
     exir_ops.edge.aten.relu.default: ReLUConverter,  # noqa F405
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,  # noqa F405
+    exir_ops.edge.aten.slice_copy.Tensor: SliceTensorConverter,  # noqa F405
     exir_ops.edge.aten.sub.Tensor: SubTensorConverter,  # noqa F405
     exir_ops.edge.aten.tanh.default: TanhConverter,  # noqa F405
     exir_ops.edge.aten.view_copy.default: ViewCopyConverter,  # noqa F405

backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py

@@ -53,6 +53,9 @@
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.sigmoid_converter import (
     SigmoidConverter,
 )
+from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.slice_tensor_converter import (
+    SliceTensorConverter,
+)
 from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.softmax_converter import (
     SoftmaxConverter,
 )
@@ -90,4 +93,5 @@
     "HardTanhConverter",
     "SigmoidConverter",
     "TanhConverter",
+    "SliceTensorConverter",
 ]

backends/nxp/backend/ir/converter/node_converters/ops_converters/slice_tensor_converter.py

@@ -0,0 +1,128 @@
+# 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 numpy as np
+from backends.nxp.backend.edge_helper import input_tensor
+from backends.nxp.backend.ir.converter.conversion.common import OpsList
+from executorch.backends.nxp.backend.ir.converter.conversion import translator
+from executorch.backends.nxp.backend.ir.converter.conversion.common import (
+    node_uses_shape_broadcasting,
+)
+from executorch.backends.nxp.backend.ir.converter.node_converter import (
+    CustomDelegationOptions,
+    NodeConverter,
+)
+from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
+    slice_options,
+)
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
+from torch.fx import Node
+from torch.nn import Parameter
+
+
+class SliceTensorConverter(NodeConverter):
+    @staticmethod
+    def _is_supported_on_target(
+        node: Node,
+        neutron_target_spec: NeutronTargetSpec,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
+    ) -> bool:
+        if node_uses_shape_broadcasting(node):
+            # Shape broadcasting may require the addition of `Transpose` ops during conversion.
+            return False
+
+        # Provisional solution - slice conversion works for neutron software 2.2.1+
+        neutron_flavor = neutron_target_spec.neutron_target.__module__.split(".")[0]
+        if neutron_flavor != "neutron_converter_SDK_25_12":
+            return False
+
+        input_shape = input_tensor(node, 0).shape
+        dim = node.args[1]
+
+        # The rank of the dimension that we want to slice must be divisible by num_macs
+        num_macs = neutron_target_spec.get_num_macs()
+        return input_shape[dim] % num_macs == 0
+
+    @staticmethod
+    def _is_supported_in_IR(
+        node: Node,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
+    ) -> bool:
+        args = node.args
+        if len(args) != 4:
+            return False
+
+        _, start, end = SliceTensorConverter._get_clipped_slice_args(node)
+        if start >= end:
+            return False
+
+        return True
+
+    def _convert_to_slice(self, t_op, main_input, input_rank, dim, start, end) -> None:
+        # Prepare the TFLite parameters 'begin' and 'size' tensors
+        begin = [0] * input_rank  # By default, start the slice at 0
+        size = (
+            main_input.shape.vector.copy()
+        )  # By default, end the slice at the end of the dimension
+
+        size[dim] = max(end - start, 0)
+        begin[dim] = start
+
+        # We can slice only the channels dimension
+        # So we swap the sliced dimension with the channels dimension
+        begin[-1], begin[dim] = begin[dim], begin[-1]
+        size[-1], size[dim] = size[dim], size[-1]
+
+        # Create permutation for swapping
+        perm = list(range(0, input_rank))
+        perm[dim], perm[-1] = perm[-1], perm[dim]
+
+        begin_tensor = self.builder.create_tensor_for_data(
+            np.asarray(begin, np.int32), "begin"
+        )
+        size_tensor = self.builder.create_tensor_for_data(
+            np.asarray(size, np.int32), "size"
+        )
+
+        t_op.tmp_inputs = [main_input, begin_tensor, size_tensor]
+        t_op.builtin_options = slice_options.Slice()
+
+        ops = OpsList(middle_op=t_op)
+        # Insert forward and backward transpose
+        ops.add_pre(self.builder.create_transpose_operator_before(t_op, 0, perm))
+        ops.add_post(self.builder.create_transpose_operator_after(t_op, 0, perm))
+
+        self.builder.append_operators(ops.flatten())
+
+    Dim = Start = End = int
+
+    @staticmethod
+    def _get_clipped_slice_args(node: Node) -> (Dim, Start, End):
+        input_shape = input_tensor(node, 0).shape
+        _, dim, start, end = node.args
+        sliced_tensor_rank = input_shape[dim]
+
+        end = int(np.clip(end, 0, sliced_tensor_rank))
+        start = int(np.clip(start, 0, sliced_tensor_rank))
+
+        return dim, start, end
+
+    def convert(self, node: Node):
+        """Convert 'slice_tensor' operator to NeutronIR 'Slice'."""
+        self.assert_convertible(node)
+        t_op = self._create_tflite_op_with_io_tensors(node)
+        inputs = t_op.tmp_inputs[0]
+        rank = inputs.rank
+
+        dim, start, end = self._get_clipped_slice_args(node)
+
+        if t_op.tmp_inputs[0].tensor_format.is_channels_last():
+            dim = translator.create_channels_last_to_channels_first_permutation(
+                t_op.tmp_inputs[0].rank
+            )[dim]
+
+        self._convert_to_slice(t_op, inputs, rank, dim, start, end)

backends/nxp/backend/neutron_converter_manager.py

@@ -26,7 +26,10 @@ class NeutronConverterManager:
     contains NeutronGraph nodes.
     """
 
-    def __init__(self, neutron_converter_flavor: str = "SDK_25_09"):
+    def __init__(
+        self,
+        neutron_converter_flavor: str = "SDK_25_09",
+    ):
 
         neutron_converter_modules = [
             module.name
@@ -76,7 +79,9 @@ def convert(self, tflite_model: bytes, target: str) -> bytes:
         cctx = self.neutron_converter.CompilationContext()
         cctx.targetOpts = self.neutron_converter.getNeutronTarget(target)
         cctx.compilationOpts.minNumOpsPerGraph = 1
-        cctx.compilationOpts.excludeGraphPasses = "MergeTranspose"
+        cctx.compilationOpts.excludeGraphPasses = (
+            "HoistSliceAboveTranspose,MergeTranspose"
+        )
 
         logger = multiprocessing.log_to_stderr()
         logger.setLevel(logging.WARNING)

backends/nxp/neutron_partitioner.py

@@ -210,6 +210,7 @@ def tag_qdq_clusters(self, nodes: list[torch.fx.Node]):
     exir_ops.edge.aten.mm.default: MMConverter,  # noqa F405
     exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter,  # noqa F405
     exir_ops.edge.aten.relu.default: ReLUConverter,  # noqa F405
+    exir_ops.edge.aten.slice_copy.Tensor: SliceTensorConverter,  # noqa F405
     exir_ops.edge.aten._softmax.default: SoftmaxConverter,  # noqa F405
     exir_ops.edge.aten.sigmoid.default: SigmoidConverter,  # noqa F405
     exir_ops.edge.aten.sub.Tensor: SubTensorConverter,  # noqa F405

backends/nxp/quantizer/neutron_quantizer.py

@@ -38,6 +38,7 @@
     ReshapePattern,
     SharedSpecPattern,
     SigmoidPattern,
+    SliceTensorPattern,
     SoftMaxPattern,
     SubTensorPattern,
     TanhInPlacePattern,
@@ -214,6 +215,7 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec):
                 NeutronAtenQuantizer(ReluInPlacePattern(), static_qconfig),
                 NeutronAtenQuantizer(ReshapePattern(), static_qconfig),
                 NeutronAtenQuantizer(SigmoidPattern(), static_qconfig),
+                NeutronAtenQuantizer(SliceTensorPattern(), static_qconfig),
                 NeutronAtenQuantizer(SoftMaxPattern(), static_qconfig),
                 NeutronAtenQuantizer(SubTensorPattern(), static_qconfig),
                 NeutronAtenQuantizer(TanhPattern(), static_qconfig),

backends/nxp/quantizer/patterns.py

@@ -736,6 +736,15 @@ def partition_types(self):
         return [torch.ops.aten.view.default]
 
 
+class SliceTensorPattern(SharedSpecPattern):
+    """
+    Quantizer for Slice operator.
+    """
+
+    def partition_types(self):
+        return [torch.ops.aten.slice.Tensor]
+
+
 class SoftMaxPattern(QuantizationPattern):
     """
     Quantizer for Softmax operator.