Merge branch 'main' into add-q5k-mlx

Author: JaynouOliver

.ci/scripts/wheel/test_base.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -40,6 +41,23 @@ class ModelTest:
     backend: Backend
 
 
+def test_cmsis_nn_install():
+    import executorch.backends.cortex_m.library.cmsis_nn as cmsis_nn
+
+    buf_size = cmsis_nn.convolve_wrapper_buffer_size(
+        cmsis_nn.Backend.MVE,
+        cmsis_nn.DataType.A8W8,
+        input_nhwc=[1, 8, 8, 16],
+        filter_nhwc=[8, 3, 3, 16],
+        output_nhwc=[1, 6, 6, 8],
+        padding_hw=[0, 0],
+        stride_hw=[1, 1],
+        dilation_hw=[1, 1],
+    )
+
+    assert buf_size == 576
+
+
 def run_tests(model_tests: List[ModelTest]) -> None:
     # Test that we can import the portable_lib module - verifies RPATH is correct
     print("Testing portable_lib import...")

.ci/scripts/wheel/test_linux.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -38,6 +39,8 @@
         else:
             print("⚠ VulkanBackend not registered (expected for the default wheel)")
 
+        test_base.test_cmsis_nn_install()
+
     test_base.run_tests(
         model_tests=[
             test_base.ModelTest(

.ci/scripts/wheel/test_macos.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -9,6 +10,8 @@
 from examples.models import Backend, Model
 
 if __name__ == "__main__":
+    test_base.test_cmsis_nn_install()
+
     test_base.run_tests(
         model_tests=[
             test_base.ModelTest(

.ci/scripts/wheel/test_windows.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
 import platform
 from typing import List
 
+import test_base
+
 import torch
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.examples.models import Backend, Model, MODEL_NAME_TO_MODEL
@@ -74,6 +77,8 @@ def run_tests(model_tests: List[ModelTest]) -> None:
         else:
             print("⚠ VulkanBackend not registered (expected for the default wheel)")
 
+        test_base.test_cmsis_nn_install()
+
     run_tests(
         model_tests=[
             ModelTest(

CMakeLists.txt

@@ -755,7 +755,7 @@ if(EXECUTORCH_BUILD_XNNPACK)
   list(APPEND _executorch_backends xnnpack_backend)
 endif()
 
-if(EXECUTORCH_BUILD_CORTEX_M)
+if(EXECUTORCH_BUILD_CORTEX_M OR EXECUTORCH_BUILD_CMSIS_NN_PYBINDS)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cortex_m)
 endif()
 

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -92,6 +92,20 @@ def _is_tosa_dialect_op(target) -> bool:
             or "<EdgeOpOverload: tosa." in target_str
         )
 
+    @staticmethod
+    def _has_real_tosa_dialect_impl(target) -> bool:
+        schema = getattr(target, "_schema", None)
+        op_name = getattr(schema, "name", None)
+        if op_name is None:
+            return False
+
+        try:
+            return torch._C._dispatch_has_kernel_for_dispatch_key(
+                op_name, "CompositeExplicitAutograd"
+            )
+        except RuntimeError:
+            return False
+
     @staticmethod
     def _arg_contains_symbolic_shape(arg) -> bool:
         if isinstance(arg, torch.fx.Node):
@@ -197,19 +211,31 @@ def resolve_arg(arg, arg_index=None):
 
         return True
 
+    def maybe_delete(self, input_nodes_to_maybe_delete):
+        for input_node in input_nodes_to_maybe_delete:
+            if input_node.meta.get("is_input", False):
+                # Never delete submodule inputs, they need to match the parameters from the outer module.
+                continue
+            if len(input_node.users) == 0:
+                self._delete_constant_placeholder(input_node)
+
     def call(self, graph_module):
         modified = False
         input_nodes_to_maybe_delete = set()
         for node in graph_module.graph.nodes:
             if node.op != "call_function":
                 continue
-            # Don't fuse TOSA dialect ops as they do not have eager forward functions.
-            # Also don't fuse ops whose explicit args/kwargs include symbolic shape values.
-            if (
-                self._is_tosa_dialect_op(node.target)
-                or self._arg_contains_symbolic_shape(node.args)
-                or self._arg_contains_symbolic_shape(node.kwargs)
-            ):
+            if node.target == exir_ops.backend.tosa.RESCALE.default:
+                # Leave fusing of RESCALES to the compiler.
+                continue
+            if self._is_tosa_dialect_op(
+                node.target
+            ) and not self._has_real_tosa_dialect_impl(node.target):
+                continue
+            # Don't fuse ops whose explicit args/kwargs include symbolic shape values.
+            if self._arg_contains_symbolic_shape(
+                node.args
+            ) or self._arg_contains_symbolic_shape(node.kwargs):
                 continue
 
             input_nodes = node.all_input_nodes
@@ -241,10 +267,7 @@ def call(self, graph_module):
 
         if modified:
             graph_module.graph.eliminate_dead_code()
-            for input_node in input_nodes_to_maybe_delete:
-                if len(input_node.users) == 0:
-                    self._delete_constant_placeholder(input_node)
-
+            self.maybe_delete(input_nodes_to_maybe_delete)
             graph_module = super().call(graph_module).graph_module
 
         return PassResult(graph_module, modified)

backends/arm/operators/node_visitor.py

@@ -31,6 +31,7 @@
     TosaSpecMapping,
 )
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -246,3 +247,15 @@ def get_node_visitors(*args) -> Dict[str, NodeVisitor]:
         node_visitors[target] = visitor(*args)
 
     return node_visitors
+
+
+def get_node_visitor(target: str, tosa_spec: TosaSpecification):
+    # Ensure all operator modules are imported so visitors are registered.
+    import executorch.backends.arm.operators  # noqa: F401
+
+    node_visitor_tuples = _node_visitor_tuples.get(tosa_spec)
+    for target_name, node_visitor_cls in node_visitor_tuples:
+        if target_name == target:
+            return node_visitor_cls(tosa_spec)
+
+    raise ValueError(f"No {target} NodeVisitor registered for {tosa_spec}")

backends/arm/quantizer/quantization_annotator.py

@@ -859,6 +859,7 @@ def any_or_hardtanh_min_zero(n: Node):
         )
     elif node.target in (
         torch.ops.aten.cat.default,
+        torch.ops.aten.concat.default,
         torch.ops.aten.concatenate.default,
         torch.ops.aten.stack.default,
     ):

backends/arm/test/misc/tosa_dialect/test_tosa_gather.py

@@ -0,0 +1,44 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import executorch.backends.arm.tosa.dialect  # noqa: F401
+
+import torch
+from executorch.backends.arm.tosa.specification import (
+    TosaLoweringContext,
+    TosaSpecification,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+def test_gather_tosa_FP_fake() -> None:
+    values = torch.randn((1, 4, 3), dtype=torch.float32)
+    indices = torch.tensor([[0, 2]], dtype=torch.int32)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string("TOSA-1.0+FP")
+    ), FakeTensorMode() as mode:
+        output = exir_ops.backend.tosa.GATHER.default(
+            mode.from_tensor(values),
+            mode.from_tensor(indices),
+        )
+
+    assert output.dtype == values.dtype
+    assert tuple(output.shape) == (1, 2, 3)
+
+
+def test_gather_tosa_FP_real() -> None:
+    values = torch.tensor(
+        [[[1.0, 10.0], [2.0, 20.0], [3.0, 30.0], [4.0, 40.0]]],
+        dtype=torch.float32,
+    )
+    indices = torch.tensor([[3, 1]], dtype=torch.int32)
+
+    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
+        output = exir_ops.backend.tosa.GATHER.default(values, indices)
+
+    expected = values[:, indices[0], :]
+    torch.testing.assert_close(output, expected)

backends/arm/test/misc/tosa_dialect/test_tosa_identity.py

@@ -13,7 +13,7 @@
 from torch._subclasses.fake_tensor import FakeTensorMode
 
 
-def test_identity_tosa_FP() -> None:
+def test_identity_tosa_FP_fake() -> None:
     sample_input = torch.randn((1, 2, 3, 4), dtype=torch.float32)
 
     with TosaLoweringContext(
@@ -23,3 +23,12 @@ def test_identity_tosa_FP() -> None:
 
         assert output.dtype == sample_input.dtype
         assert tuple(output.shape) == tuple(sample_input.shape)
+
+
+def test_identity_tosa_FP_real() -> None:
+    sample_input = torch.randn((1, 2, 3, 4), dtype=torch.float32)
+
+    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
+        output = exir_ops.backend.tosa.IDENTITY.default(sample_input)
+
+    torch.testing.assert_close(output, sample_input)