Arm backend: Test popular torch modules/ functions, adress issues (#9221)

- Adds tests for top 10 torch.nn.Modules, torch.* functions, and
torch.nn.functional.
- Add check to not partion ops with int64 input
- A few minor fixes

Signed-off-by: Erik Lundell <[email protected]>

Author: Erik-Lundell

backends/arm/_passes/arm_pass_manager.py

@@ -199,8 +199,8 @@ def transform_to_backend_pipeline(self, exported_program: ExportedProgram):
             )
 
     def transform_for_annotation_pipeline(self, graph_module: GraphModule):
-        self.add_pass(ScalarsToAttributePass())
         self.add_pass(ReplaceScalarWithTensorArgPass())
+        self.add_pass(ScalarsToAttributePass())
         self.add_pass(DecomposeLayerNormPass())
         self.add_pass(DecomposeVarPass())
         self.add_pass(DecomposeMeanDimPass())

backends/arm/_passes/arm_pass_utils.py

@@ -8,7 +8,7 @@
 # pyre-unsafe
 
 from inspect import isclass
-from typing import Optional
+from typing import Optional, Sequence
 
 import torch
 import torch.fx
@@ -149,7 +149,7 @@ def get_first_fake_tensor(node: torch.fx.Node) -> FakeTensor:
     If the node contains many fake tensors, return the first one.
     """
     if isinstance(
-        node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list)
+        node.meta["val"], (Sequence, torch.fx.immutable_collections.immutable_list)
     ):
         fake_tensor = node.meta["val"][0]
     else:

backends/arm/operator_support/tosa_supported_operators.py

@@ -18,7 +18,9 @@
     FuseQuantizedActivationPass,
 )
 from executorch.backends.arm.tosa_specification import Tosa_0_80, TosaSpecification
+from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export.graph_signature import InputKind
 from torch.fx.passes.operator_support import any_chain, chain, OperatorSupportBase
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
 
@@ -84,9 +86,10 @@ def get_registered_tosa_support_checks(
 
 def tosa_support_factory(
     tosa_spec: TosaSpecification,
+    exported_program: ExportedProgram,
     additional_checks: Optional[Sequence[OperatorSupportBase]] = None,
 ) -> OperatorSupportBase:
-    negative_checks: list[OperatorSupportBase] = []
+    negative_checks: list[OperatorSupportBase] = [CheckInt64Inputs(exported_program)]
     if not tosa_spec.support_float():
         negative_checks.append(NeedsDecompositionCheck())
         negative_checks.append(CheckProperQuantization())
@@ -247,6 +250,10 @@ def is_node_supported(
             exir_ops.edge.aten._log_softmax.default,
             exir_ops.edge.aten.var.correction,
             exir_ops.edge.aten.var.dim,
+            exir_ops.edge.aten.add.Scalar,
+            exir_ops.edge.aten.sub.Scalar,
+            exir_ops.edge.aten.mul.Scalar,
+            exir_ops.edge.aten.div.Scalar,
         ]
         return not needs_decomp
 
@@ -312,6 +319,8 @@ def is_node_supported(
             exir_ops.edge.aten.bmm.default,
             exir_ops.edge.aten.convolution.default,
             exir_ops.edge.aten.exp.default,
+            exir_ops.edge.aten.full.default,
+            exir_ops.edge.aten.full_like.default,
             exir_ops.edge.aten.hardtanh.default,
             exir_ops.edge.aten.linear.default,
             exir_ops.edge.aten.log.default,
@@ -371,3 +380,29 @@ def is_node_supported(
         if not output_quantized:
             return False
         return True
+
+
+class CheckInt64Inputs(OperatorSupportBase):
+
+    def __init__(self, exported_program: ExportedProgram):
+        self.input_names = [
+            spec.arg.name
+            for spec in exported_program.graph_signature.input_specs
+            if spec.kind == InputKind.USER_INPUT
+        ]
+        super().__init__()
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+
+        for input_node in node.all_input_nodes:
+            # We can cast constant placeholders AOT, not call_functions.
+            if (
+                input_node.name in self.input_names
+                or not input_node.op == "placeholder"
+            ):
+                tensor = get_first_fake_tensor(input_node)
+                if tensor.dtype == torch.int64:
+                    return False
+        return True

backends/arm/test/models/test_conformer.py

@@ -93,7 +93,7 @@ def test_conformer_tosa_BI(self):
             )
         )
 
-    @unittest.expectedFailure  # TODO(MLETORCH-635)
+    @conftest.expectedFailureOnFVP  # TODO(MLETORCH-635)
     def test_conformer_u55_BI(self):
         tester = (
             ArmTester(
@@ -115,7 +115,7 @@ def test_conformer_u55_BI(self):
                 inputs=get_test_inputs(self.dim, self.lengths, self.num_examples),
             )
 
-    @unittest.expectedFailure  # TODO(MLETORCH-635)
+    @conftest.expectedFailureOnFVP  # TODO(MLETORCH-635)
     def test_conformer_u85_BI(self):
         tester = (
             ArmTester(

backends/arm/test/models/test_nn_functional.py

@@ -0,0 +1,118 @@
+# Copyright 2025 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.
+
+"""
+Tests 10 popular torch.nn.functional not tested in other ways or training related
+- normalize
+- grid_sample
+- one_hot
+- softplus
+- cosine_similarity
+- unfold
+- elu
+- fold
+- affine_grid
+- max_pool1d
+- threshold
+"""
+from typing import Callable
+
+import torch
+from executorch.backends.arm.test.common import parametrize
+from executorch.backends.arm.test.tester.test_pipeline import (
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+
+def module_factory(function: Callable) -> torch.nn.Module:
+    class ModuleWrapper(torch.nn.Module):
+        def forward(self, *args):
+            return function(*args)
+
+    return ModuleWrapper()
+
+
+example_input = torch.rand(1, 6, 16, 16)
+
+module_tests = {
+    "normalize": (module_factory(torch.nn.functional.normalize), (example_input,)),
+    "grid_sample": (
+        module_factory(torch.nn.functional.grid_sample),
+        (torch.rand(1, 1, 4, 4), torch.rand(1, 5, 5, 2)),
+    ),
+    "one_hot": (
+        module_factory(torch.nn.functional.one_hot),
+        (torch.randint(0, 5, (2, 2, 5, 5)), 5),
+    ),
+    "softplus": (module_factory(torch.nn.functional.softplus), (example_input,)),
+    "cosine_similarity": (
+        module_factory(torch.nn.functional.cosine_similarity),
+        (example_input, example_input),
+    ),
+    "unfold": (
+        module_factory(torch.nn.functional.unfold),
+        (torch.randn(1, 3, 10, 12), (4, 5)),
+    ),
+    "elu": (module_factory(torch.nn.functional.elu), (example_input,)),
+    "fold": (
+        module_factory(torch.nn.functional.fold),
+        (torch.randn(1, 12, 12), (4, 5), (2, 2)),
+    ),
+    "affine_grid": (
+        module_factory(torch.nn.functional.affine_grid),
+        (torch.rand(1, 2, 3), (1, 2, 10, 10)),
+    ),
+    "max_pool1d": (
+        module_factory(torch.nn.functional.max_pool1d),
+        (torch.randn(20, 16, 50), 4),
+    ),
+    "threshold": (
+        module_factory(torch.nn.functional.threshold),
+        (example_input, 0.5, 0.1),
+    ),
+}
+
+input_t = tuple[torch.Tensor]
+
+
+@parametrize(
+    "test_data", module_tests, xfails={"max_pool1d": "ValueError: Invalid TOSA graph"}
+)
+def test_nn_functional_MI(test_data):
+    module, inputs = test_data
+    pipeline = TosaPipelineMI[input_t](
+        module, inputs, "", use_to_edge_transform_and_lower=True
+    )
+    pipeline.pop_stage("check.aten")
+    pipeline.pop_stage("check_count.exir")
+    try:
+        pipeline.run()
+    except RuntimeError as e:
+        if (
+            "Ran model with TosaReferenceModelDispatch but never ran TOSABackend delegate."
+            not in str(e)
+        ):
+            raise e
+
+
+@parametrize("test_data", module_tests)
+def test_nn_functional_BI(test_data):
+    module, inputs = test_data
+    pipeline = TosaPipelineBI[input_t](
+        module, inputs, "", use_to_edge_transform_and_lower=True
+    )
+    pipeline.pop_stage("check.aten")
+    pipeline.pop_stage("check_count.exir")
+    pipeline.pop_stage("check.quant_nodes")
+    pipeline.pop_stage("check_not.quant_nodes")
+    try:
+        pipeline.run()
+    except RuntimeError as e:
+        if (
+            "Ran model with TosaReferenceModelDispatch but never ran TOSABackend delegate."
+            not in str(e)
+        ):
+            raise e

backends/arm/test/models/test_nn_modules.py

@@ -0,0 +1,103 @@
+# Copyright 2025 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.
+
+"""
+Tests 10 popular nn modules not tested in other ways or training related.
+- Embedding
+- LeakyReLU
+- BatchNorm1d
+- AdaptiveAvgPool2d
+- ConvTranspose2d
+- GRU
+- GroupNorm
+- InstanceNorm2d
+- PReLU
+- Transformer
+"""
+
+import torch
+from executorch.backends.arm.test.common import parametrize
+from executorch.backends.arm.test.tester.test_pipeline import (
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+example_input = torch.rand(1, 6, 16, 16)
+
+module_tests = [
+    (torch.nn.Embedding(10, 10), (torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]]),)),
+    (torch.nn.LeakyReLU(), (example_input,)),
+    (torch.nn.BatchNorm1d(16), (torch.rand(6, 16, 16),)),
+    (torch.nn.AdaptiveAvgPool2d((12, 12)), (example_input,)),
+    (torch.nn.ConvTranspose2d(6, 3, 2), (example_input,)),
+    (torch.nn.GRU(10, 20, 2), (torch.randn(5, 3, 10), torch.randn(2, 3, 20))),
+    (torch.nn.GroupNorm(2, 6), (example_input,)),
+    (torch.nn.InstanceNorm2d(16), (example_input,)),
+    (torch.nn.PReLU(), (example_input,)),
+    (
+        torch.nn.Transformer(
+            d_model=64,
+            nhead=1,
+            num_encoder_layers=1,
+            num_decoder_layers=1,
+            dtype=torch.float32,
+        ),
+        (torch.rand((10, 32, 64)), torch.rand((20, 32, 64))),
+    ),
+]
+
+input_t = tuple[torch.Tensor]
+
+test_parameters = {str(test[0].__class__.__name__): test for test in module_tests}
+
+
+@parametrize(
+    "test_data",
+    test_parameters,
+    xfails={"Transformer": "Output 0 does not match reference output."},
+)
+def test_nn_Modules_MI(test_data):
+    module, inputs = test_data
+    pipeline = TosaPipelineMI[input_t](
+        module, inputs, "", use_to_edge_transform_and_lower=True
+    )
+    pipeline.pop_stage("check.aten")
+    pipeline.pop_stage("check_count.exir")
+    try:
+        pipeline.run()
+    except RuntimeError as e:
+        if (
+            "Ran model with TosaReferenceModelDispatch but never ran TOSABackend delegate."
+            not in str(e)
+        ):
+            raise e
+
+
+@parametrize(
+    "test_data",
+    test_parameters,
+    xfails={
+        "GRU": "RuntimeError: Node aten_linear_default with op <EdgeOpOverload: aten.linear[...]> was not decomposed or delegated.",
+        "PReLU": "RuntimeError: mul(): functions with out=... arguments don't support automatic differentiation, but one of the arguments requires grad.",
+        "Transformer": "RuntimeError: Expected out tensor to have dtype signed char, but got float",
+    },
+)
+def test_nn_Modules_BI(test_data):
+    module, inputs = test_data
+    pipeline = TosaPipelineBI[input_t](
+        module, inputs, "", use_to_edge_transform_and_lower=True
+    )
+    pipeline.pop_stage("check.aten")
+    pipeline.pop_stage("check_count.exir")
+    pipeline.pop_stage("check.quant_nodes")
+    pipeline.pop_stage("check_not.quant_nodes")
+    try:
+        pipeline.run()
+    except RuntimeError as e:
+        if (
+            "Ran model with TosaReferenceModelDispatch but never ran TOSABackend delegate."
+            not in str(e)
+        ):
+            raise e