Update

[ghstack-poisoned]

Author: JCNTH

.github/workflows/doc-build.yml

@@ -104,9 +104,6 @@ jobs:
       repository: pytorch/executorch
       download-artifact: docs
       ref: gh-pages
-      # Publish on the host so actions/checkout's persisted credentials are
-      # available to git push.
-      run-with-docker: false
       timeout: 90
       script: |
         set -euo pipefail

.github/workflows/pull.yml

@@ -1427,7 +1427,7 @@ jobs:
       docker-image: ci-image:executorch-ubuntu-22.04-clang12
       submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 150
+      timeout: 120
       script: |
         set -eux
 

backends/mlx/passes.py

@@ -99,8 +99,49 @@ def call(self, exported_program) -> ExportedProgramPassResult:
         }
         if ops_to_inplace:
             reinplace_pass(exported_program, ops_to_inplace=ops_to_inplace)
+            self._resync_output_specs(exported_program)
         return ExportedProgramPassResult(exported_program, True)
 
+    @staticmethod
+    def _resync_output_specs(exported_program) -> None:
+        """Re-sync graph-signature output names after reinplace.
+
+        ``reinplace_pass`` rewrites an output-producing node (e.g. the final
+        ``exp`` -> ``exp_``) via ``replace_all_uses_with`` + erase, but does not
+        update ``graph_signature.output_specs``. Output order is preserved, so we
+        positionally re-sync each spec's argument name to the current output node
+        arg; otherwise ``ExportedProgram.validate()`` (run by the pass manager)
+        raises a SpecViolationError.
+
+        This positional pairing is only valid because reinplace does a 1:1
+        ``replace_all_uses_with`` + erase and never drops, adds, or reorders
+        outputs. We assert ``len(output_specs) == len(out_args)`` so that a
+        future change violating that invariant fails loudly here instead of
+        silently mis-pairing names (``zip`` would otherwise truncate). Specs
+        whose ``arg`` is not a named tensor (e.g. ``ConstantArgument``) carry no
+        ``name`` and are skipped by the ``getattr`` guard below.
+        """
+        out_node = next(
+            n for n in reversed(exported_program.graph.nodes) if n.op == "output"
+        )
+        out_args = out_node.args[0]
+        if not isinstance(out_args, (tuple, list)):
+            out_args = (out_args,)
+        output_specs = exported_program.graph_signature.output_specs
+        assert len(output_specs) == len(out_args), (
+            "reinplace changed graph output count: "
+            f"{len(output_specs)} output_specs vs {len(out_args)} output args. "
+            "Positional output-spec re-sync assumes a 1:1, order-preserving "
+            "rewrite."
+        )
+        for spec, arg in zip(output_specs, out_args):
+            if (
+                isinstance(arg, torch.fx.Node)
+                and getattr(spec.arg, "name", None) is not None
+                and spec.arg.name != arg.name
+            ):
+                spec.arg.name = arg.name
+
 
 @dataclass
 class RMSNormMatch(PatternMatch):

backends/nxp/backend/ir/converter/node_converter.py

@@ -6,7 +6,6 @@
 import logging
 import operator
 from abc import ABC, abstractmethod
-from math import prod
 from typing import Callable
 
 import torch
@@ -466,30 +465,29 @@ def uses_shape_broadcasting(node: Node) -> bool:
         )
 
     @staticmethod
-    def inputs_satisfy_broadcast_condition(node: Node) -> bool:
-        """Determine if given PyTorch fx Node has inputs that satisfy broadcasting conditions for Neutron or not.
+    def at_least_one_input_shape_matches_the_output_shape(node: Node) -> bool:
+        """Determine if given PyTorch fx Node uses at least one input shape broadcasting for it's input nodes or not.
 
         :param node: PyTorch fx Node with 'all_input_nodes' initialized.
-        :return: True, if at least one input has the same number of elements as the output node.
+        :return: True, if at least one input has the same shape as the output node.
                  False otherwise.
         """
 
         if node.all_input_nodes is None:
             logger.e(
                 logger.Code.INTERNAL_ERROR,
-                "node_converter.inputs_satisfy_broadcast_condition(): 'all_input_nodes' are None!",
+                "node_converter.at_least_one_input_shape_matches_the_output_shape(): 'all_input_nodes' are None!",
             )
 
         if len(node.all_input_nodes) == 0:
             logger.e(
                 logger.Code.INTERNAL_ERROR,
-                "node_converter.inputs_satisfy_broadcast_condition(): Operator has no inputs!",
+                "node_converter.at_least_one_input_shape_matches_the_output_shape(): Operator has no inputs!",
             )
 
         output_shape = node.meta["val"].shape
-        num_elements = prod(output_shape)
 
         return any(
-            prod(input_tensor.meta["val"].shape) == num_elements
+            input_tensor.meta["val"].shape == output_shape
             for input_tensor in node.all_input_nodes
         )

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

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
+        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
             return False
 
         supported_types = [torch.int8, torch.uint8]

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

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
+        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
             return False
 
         supported_types = [torch.int8, torch.uint8]

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

@@ -3,10 +3,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import torch
-
 from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT
-from executorch.backends.nxp.backend.edge_helper import input_rank
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
@@ -27,17 +24,38 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
+        node_shape = node.meta["val"].shape
+        rank = len(node_shape)
+
+        # According to Neutron spec., PReLU can be done only on 4D tensors
+        if rank != 4:
+            return False
+
+        ch_idx, h_idx, w_idx = PReLUConverter._get_channel_spatial_indices(node)
+        # According to Neutron spec., size of channels must be divisible by num_macs.
+        num_macs = neutron_target_spec.get_num_macs()
+        if node_shape[ch_idx] % num_macs != 0:
             return False
 
-        supported_types = [torch.int8, torch.uint8]
-        if not NodeConverter.uses_quantization_type_for_io(
-            node, supported_types, [0, 1], [0]
-        ):
+        # According to Neutron spec., height * width cannot be greater than a given constant.
+        if node_shape[w_idx] * node_shape[h_idx] > 4096:
             return False
 
         return True
 
+    @staticmethod
+    def _get_channel_spatial_indices(node: Node):
+        if node.meta[NXP_NODE_FORMAT].is_channels_first():
+            ch_idx = 1
+            h_idx = 2
+            w_idx = 3
+        else:
+            ch_idx = 3
+            h_idx = 1
+            w_idx = 2
+
+        return ch_idx, h_idx, w_idx
+
     @staticmethod
     def _is_supported_in_IR(
         node: Node,
@@ -47,12 +65,6 @@ def _is_supported_in_IR(
         if len(node.args) != 2:
             return False
 
-        if input_rank(node, 0) > 2 and tuple(
-            node.all_input_nodes[1].meta["val"].shape
-        ) != (1,):
-            if not node.meta[NXP_NODE_FORMAT].is_channels_first():
-                # Should never happen.
-                return False
         return True
 
     def convert(self, node: Node):

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

@@ -26,7 +26,7 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
+        if not NodeConverter.at_least_one_input_shape_matches_the_output_shape(node):
             return False
 
         supported_types = [torch.int8, torch.uint8]

backends/nxp/backend/node_format_inference.py

@@ -10,7 +10,6 @@
 
 from executorch.backends.nxp.backend.data_format import DataFormat, NXP_NODE_FORMAT
 from executorch.backends.nxp.backend.edge_helper import (
-    input_rank,
     is_channels_last_dim_order,
     try_get_arg,
 )
@@ -27,7 +26,6 @@
     MaxPool2DWithIndices,
     MeanDim,
     PermuteCopy,
-    Prelu,
     QuantizePerTensor,
     SumDimIntList,
     UpsampleBilinear2D,
@@ -56,11 +54,6 @@ class NodeFormatInference:
         UpsampleNearest2D: {"inputs": [0]},
     }
 
-    ops_conditionally_with_channels_first_nodes = {
-        Prelu: {"inputs": [0]},
-        torch.ops.aten.prelu.default: {"inputs": [0]},
-    }
-
     # A set of Edge Aten ops, which have the ability to change the format (for example - input nodes
     # are channels first but output is formatless).
     ops_that_can_change_tensor_format = {
@@ -71,8 +64,6 @@ class NodeFormatInference:
         SumDimIntList,
     }
 
-    prelu_targets = [Prelu, torch.ops.aten.prelu.default]
-
     _type_changed_during_last_run: bool
 
     # Mapping between Node and its ancestors (inputs)
@@ -169,13 +160,6 @@ def _infer_format_of_nodes(self, node: Node):
                     f"Node format inference for node type: {op_type} not found!"
                 )
 
-        elif node.target in self.prelu_targets:
-            num_parameters_shape = tuple(node.args[1].meta["val"].shape)
-            if input_rank(node, 0) > 2 and num_parameters_shape != (1,):
-                self._handle_node_which_uses_channels_first_format(node)
-            else:
-                self._handle_node_which_can_use_any_node_format(node)
-
         elif node.op != "call_function" or (
             hasattr(node, "target") and node.target in self._known_targets
         ):
@@ -264,14 +248,10 @@ def _handle_node_which_uses_channels_first_format(self, node: Node):
         Function for assigning format to nodes that require channels first input (Conv, MaxPool etc.)
         """
         op_type = self._get_node_op_type(node)
-        ops_using_channels_first_format = (
-            self.ops_with_channels_first_nodes
-            | self.ops_conditionally_with_channels_first_nodes
-        )
 
         for index, ancestor_node in enumerate(self._node_inputs[node]):
             # Go through input nodes and assign them correct format
-            if index in ops_using_channels_first_format[op_type]["inputs"]:
+            if index in self.ops_with_channels_first_nodes[op_type]["inputs"]:
                 self._assign_format_to_node(ancestor_node, DataFormat.CHANNELS_FIRST)
 
                 # We need to propagate channels first format up to already visited nodes

backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py

@@ -115,10 +115,6 @@ def test__basic_nsys_inference_qat(self, mocker, request):
                 [ModelInputSpec((1, 4, 8, 8)), ModelInputSpec((8, 8))],
                 id="2 inputs 4D + 2D.",
             ),
-            pytest.param(
-                [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
-                id="2 inputs 1D + 2D, num_elems of input == num_elems of output",
-            ),
         ],
     )
     def test__broadcast(self, mocker, request, input_spec):