Merge branch 'main' into ph-mypy-ops-enable-mypy-tests

Author: zingo

backends/arm/_passes/arm_pass_manager.py

@@ -159,7 +159,12 @@ def _transform(self, graph_module: GraphModule):
     def _tosa_pipeline(
         self, exported_program: ExportedProgram, graph_module: GraphModule
     ) -> GraphModule:
+        # Preprocessing passes
+
         self.add_pass(AnnotateOutputDimOrderPass())
+
+        # Node transformation passes (pre q/dq folding)
+
         self.add_pass(FuseQuantizedActivationPass())
         self.add_pass(RemoveGetItemPass())
         self.add_pass(ConvertToClampPass())
@@ -174,8 +179,19 @@ def _tosa_pipeline(
         self.add_pass(ConvertELUParamsPass())
         self.add_pass(ConvertSplitToSlicePass())
         self.add_pass(QuantizeOperatorArguments())
+
+        # Fold Q/DQ nodes, insert INT8/INT32 rescales.
+
         self.add_pass(FoldAndAnnotateQParamsPass(exported_program))  # type: ignore[call-arg]
         self.add_pass(FuseDuplicateUsersPass())
+        # TODO: DecomposeLinearPass should run after InsertRescaleInt32Pass or
+        # before FoldAndAnnotateQParamsPass but is unable to at the moment.
+        # Ticket: MLETORCH-1539
+        self.add_pass(DecomposeLinearPass())
+        self.add_pass(InsertRescaleInt32Pass())
+
+        # Node transformation passes (post q/dq folding)
+
         self.add_pass(DecomposeExpm1Pass())
         self.add_pass(DecomposeLogitPass())
         self.add_pass(DecomposeMaskedFill())
@@ -196,57 +212,67 @@ def _tosa_pipeline(
         self.add_pass(DecomposeSignPass())
         self.add_pass(DecomposeFloorDividePass())
         self.add_pass(DecomposeDivTensorModePass())
+        self.add_pass(DecomposeGeluPass())
+        self.add_pass(DecomposeAddSubAlphaPass())
+        self.add_pass(DecomposeGroupedConv())
+        self.add_pass(Conv1dUnsqueezePass())
+
+        # Scalars -> tensors, match tensor dtypes and ranks.
+
         self.add_pass(ReplaceScalarWithTensorByProfilePass())
+        self.add_pass(ConvertFullLikeToFullPass())
+        self.add_pass(MatchArgDtypePass())
+        self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
+        # TODO: Move DecomposeNotEqualPass to before or after this block of
+        # passes. Ticket: MLETORCH-1540
+        self.add_pass(DecomposeNotEqualPass())
+        self.add_pass(MatchArgRanksPass(exported_program))
+        self.add_pass(FuseConstantArgsPass(exported_program))
+
+        # Node transformation passes (post scalar-removal)
+
         self.add_pass(DecomposeRemainderPass())
         self.add_pass(DecomposeDivTensorModePass())
         self.add_pass(DecomposeEmbeddingPass())
         self.add_pass(FuseBatchnorm2DPass(exported_program))
         self.add_pass(ConvertMmToBmmPass())
         self.add_pass(DecomposeGluPass())
-        self.add_pass(DecomposeLinearPass())
         self.add_pass(DecomposeLeakyReLUPass())
-        self.add_pass(DecomposeNotEqualPass())
         self.add_pass(DecomposeDivPass())
-        self.add_pass(DecomposeAddSubAlphaPass())
         self.add_pass(DecomposeSoftmaxPass())
-        self.add_pass(DecomposeGeluPass())
-        self.add_pass(ConvertFullLikeToFullPass())
         self.add_pass(ConvertMinMaxPass())
         self.add_pass(ConvertAnyDefaultDimDimsPass())
-        self.add_pass(MatchArgDtypePass())
-        self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
-        self.add_pass(MatchArgRanksPass(exported_program))
         self.add_pass(DecomposeAdaptiveAvgPool2dPass())
         self.add_pass(DecomposeAvgPool2d())
         self.add_pass(
             DecorateFp32toInt32CastingPass()
         )  # Require that no new fp32->int32 is introduced after this pass
         self.add_pass(ComputeConstantOpsAOT(exported_program))
-
-        self.add_pass(DecomposeGroupedConv())
         self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(UnsqueezeBeforeRepeatPass())
         self.add_pass(DecomposeCumsumPass(exported_program))
-        self.add_pass(Conv1dUnsqueezePass())
         self.add_pass(DecomposeMaxPool2DPass())
         self.add_pass(SizeAdjustInputPass())
         self.add_pass(DecomposeSelectPass())
         self.add_pass(ConvertSqueezesToViewPass())
         self.add_pass(CastToInt32Pass())
         self.add_pass(BroadcastArgsPass())
-
         self.add_pass(ConvertPermuteSingletonToViewPass())
         self.add_pass(FuseViewCopyTransform())
-        self.add_pass(FuseConstantArgsPass(exported_program))
         self.add_pass(DecomposeConv2dWithInt16ActivationPass())
-        self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
+        self.add_pass(DecomposeSumPass())
         self.add_pass(InsertTableOpsPass(exported_program))
+
+        # Aten -> TOSA transformation passes
+
         self.add_pass(RewriteUpsamplePass())
         self.add_pass(RewriteConv2dPass(exported_program))
         self.add_pass(RewriteMatmulPass())
+
+        # Postprocessing/cleanup passes
+
+        self.add_pass(CastInt64BuffersToInt32Pass(exported_program))
         self.add_pass(FuseEqualPlaceholdersPass(exported_program))
-        self.add_pass(InsertRescaleInt32Pass())
-        self.add_pass(DecomposeSumPass())
         self.add_pass(ToTosaMemoryFormatPass(exported_program))
         self.add_pass(RemoveNoopPass())
         self.add_pass(InsertRescalePass())

backends/arm/_passes/decompose_linear_pass.py

@@ -12,6 +12,7 @@
     create_node,
     get_first_fake_tensor,
 )
+from executorch.backends.arm._passes.insert_rescales_pass import InsertRescaleInt32Pass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -26,7 +27,7 @@ class DecomposeLinearPass(ArmPass):
         output           = view(conv2d)
     """
 
-    _passes_required_after: Set[Type[ExportPass]] = set()
+    _passes_required_after: Set[Type[ExportPass]] = {InsertRescaleInt32Pass}
 
     def call(self, graph_module):
         for node in graph_module.graph.nodes:

backends/arm/_passes/match_arg_ranks_pass.py

@@ -57,6 +57,7 @@ def __init__(self, exported_program: ExportedProgram) -> None:
         exir_ops.edge.aten.lt.Tensor,
         exir_ops.edge.aten.le.Tensor,
         exir_ops.edge.aten.pow.Tensor_Tensor,
+        exir_ops.edge.aten.remainder.Tensor,
         exir_ops.edge.aten.where.self,
         exir_ops.edge.aten.bitwise_and.Tensor,
         exir_ops.edge.aten.bitwise_xor.Tensor,

backends/arm/operator_support/reduce_sum_support.py

@@ -2,7 +2,11 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+"""Declare operator support for ``aten.sum.dim_IntList`` in TOSA.
 
+Provide shape constraints for U55 subsets; otherwise allow reductions.
+
+"""
 from typing import cast
 
 import torch.fx as fx
@@ -16,14 +20,25 @@
 
 @register_tosa_support_check
 class SumSupported(SupportedTOSAOperatorCheck):
+    """Provide TOSA support check for sum over dimensions."""
+
     targets = [exir_ops.edge.aten.sum.dim_IntList]
 
     tosa_specs = [
         TosaSpecification.create_from_string("TOSA-1.0+INT"),
         TosaSpecification.create_from_string("TOSA-1.0+FP"),
     ]
 
-    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+    def is_node_tosa_supported(
+        self, node: fx.Node, tosa_spec: TosaSpecification
+    ) -> bool:
+        """Return True if the node is supported by TOSA.
+
+        On U55 subsets, enforce bounds on the reduced dimension and the products
+        of sizes before/after the reduction axis. On other targets, accept the
+        operation unconditionally.
+
+        """
         if not tosa_spec.is_U55_subset:
             return True
 

backends/arm/requirements-arm-models-test.txt

@@ -3,4 +3,4 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-diffusers[torch] == 0.33.1
+diffusers[torch] == 0.33.1

backends/arm/scripts/install_models_for_test.sh

@@ -6,3 +6,16 @@
 
 set -e
 pip install -r backends/arm/requirements-arm-models-test.txt
+
+# Install model gym repository
+git clone https://github.com/arm/neural-graphics-model-gym.git
+cd neural-graphics-model-gym
+# Remove model-converter installation from model-gym repository (to prevent overwriting executorch version)
+if [[ "$(uname)" == "Darwin" ]]; then
+    sed -i '' 's/^model-converter  = "ng_model_gym.bin.model_converter_launcher:main"/#&/' pyproject.toml
+else
+    sed -i 's/^model-converter  = "ng_model_gym.bin.model_converter_launcher:main"/#&/' pyproject.toml
+fi
+pip install . --no-deps
+cd ..
+rm -rf neural-graphics-model-gym

backends/arm/test/conftest.py

@@ -3,8 +3,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 import os
 import random
+import sys
 from typing import Any
 
 import pytest
@@ -27,6 +29,8 @@ def pytest_configure(config):
     if config.option.arm_run_tosa_version:
         pytest._test_options["tosa_version"] = config.option.arm_run_tosa_version
 
+    logging.basicConfig(stream=sys.stdout)
+
 
 def pytest_collection_modifyitems(config, items):
     pass

backends/arm/test/models/test_nss.py

@@ -0,0 +1,138 @@
+# 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.
+
+from typing import Tuple
+
+import pytest
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+from huggingface_hub import hf_hub_download
+
+from ng_model_gym.usecases.nss.model.model_blocks import AutoEncoderV1
+
+input_t = Tuple[torch.Tensor]  # Input x
+
+
+class NSS(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.auto_encoder = AutoEncoderV1()
+
+
+def nss() -> AutoEncoderV1:
+    """Get an instance of NSS with weights loaded."""
+
+    weights = hf_hub_download(
+        repo_id="Arm/neural-super-sampling", filename="nss_v0.1.0_fp32.pt"
+    )
+
+    nss_model = NSS()
+    nss_model.load_state_dict(
+        torch.load(weights, map_location=torch.device("cpu"), weights_only=True),
+        strict=False,
+    )
+    return nss_model.auto_encoder
+
+
+def example_inputs():
+    return (torch.randn((1, 12, 544, 960)),)
+
+
+def test_nss_tosa_FP():
+    pipeline = TosaPipelineFP[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.add_stage_after("export", pipeline.tester.dump_operator_distribution)
+    pipeline.run()
+
+
+def test_nss_tosa_INT():
+    pipeline = TosaPipelineINT[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+@pytest.mark.skip(reason="No support for aten_upsample_nearest2d_vec on U55")
+@common.XfailIfNoCorstone300
+def test_nss_u55_INT():
+    pipeline = EthosU55PipelineINT[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_ops=[],
+        exir_ops=[],
+        run_on_fvp=True,
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+@pytest.mark.skip(
+    reason="Fails at input memory allocation for input shape: [1, 12, 544, 960]"
+)
+@common.XfailIfNoCorstone320
+def test_nss_u85_INT():
+    pipeline = EthosU85PipelineINT[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_ops=[],
+        exir_ops=[],
+        run_on_fvp=True,
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+@pytest.mark.xfail(
+    reason="[MLETORCH-1430]: Double types are not supported in buffers in MSL"
+)
+@common.SkipIfNoModelConverter
+def test_nss_vgf_FP():
+    pipeline = VgfPipeline[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+FP",
+        use_to_edge_transform_and_lower=True,
+        run_on_vulkan_runtime=True,
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+def test_nss_vgf_INT():
+    pipeline = VgfPipeline[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+INT",
+        symmetric_io_quantization=True,
+        use_to_edge_transform_and_lower=True,
+        run_on_vulkan_runtime=True,
+    )
+    pipeline.run()
+
+
+ModelUnderTest = nss().eval()
+ModelInputs = example_inputs()