Update on "[ET-VK] Allow aten.cat.default to handle any number of input tensors"

## Context

Previously, I updated the implementation of `aten.cat.default` in D76305343 (#11508) since the original implementation had a bug. The new implementation only supported up to 3 input tensors, but several models require the need for up to 6 input tensors.

This diff updates the capabilities of the `concat` op so that any arbitrary number of input tensors may be accepted.

## Changes

* Update implementation of the concat shader to be able to be called repeatedly, allowing support for any number of input tensors.

Differential Revision: [D79893084](https://our.internmc.facebook.com/intern/diff/D79893084/)

[ghstack-poisoned]

Author: SS-JIA

.github/workflows/trunk.yml

@@ -60,7 +60,7 @@ jobs:
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
       matrix:
-        model: [add]
+        model: [add, softmax, mv2]
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -72,6 +72,16 @@ jobs:
         MODEL_NAME=${{ matrix.model }}
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
+        if [[ ${{ matrix.model}} == "add" ]]; then
+          SIM_LIMIT_SEC=60
+        elif [[ ${{ matrix.model}} == "softmax" ]]; then
+          SIM_LIMIT_SEC=60
+        elif [[ ${{ matrix.model}} == "mv2" ]]; then
+          SIM_LIMIT_SEC=5000
+        else
+          echo "Failed unsupported model selection ${{ matrix.model }}"
+          exit 1
+        fi
 
         source .ci/scripts/utils.sh
         source .ci/scripts/zephyr-utils.sh
@@ -118,24 +128,22 @@ jobs:
             -C mps3_board.uart0.out_file='sim.out'  \
             -C cpu0.CFGITCMSZ=15 \
             -C cpu0.CFGDTCMSZ=15 \
-            --simlimit 120
+            --simlimit ${SIM_LIMIT_SEC}
 
         # Disable exit on error
         set +e
         # Report failure if any of the ouptut verification checks fail
-        # store 0 if found (failure), 1 if not (success)
         grep -qF "ERROR" sim.out
-        exit_status=$?
+        exit_status=$? #store 0 if found (failure), 1 if not (success)
         if [[ "$exit_status" -eq "0" ]]; then
             cat sim.out
             set -e
             exit 1
         fi
 
         # Report fail if simulation does not complete successfully
-        # store 0 if found (success), 1 if not (failure)
         grep -qF "SUCCESS: Program complete, exiting." sim.out
-        exit_status=$?
+        exit_status=$? #store 0 if found (success), 1 if not (failure)
         if [[ "$exit_status" -eq "1" ]]; then
             cat sim.out
             set -e

backends/apple/coreml/test/test_torch_ops.py

@@ -167,6 +167,10 @@ def test_dequantize_affine_c8w_embedding_b4w_linear(self):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs)
 
+    @unittest.skipIf(
+        not hasattr(torch.version, "git_version"),
+        "Enable in fbcode once D79658061 lands",
+    )
     def test_dequantize_codebook_linear(self):
         model, example_inputs = self._get_test_model()
         quantize_(
@@ -194,6 +198,10 @@ def test_dequantize_codebook_linear(self):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs)
 
+    @unittest.skipIf(
+        not hasattr(torch.version, "git_version"),
+        "Enable in fbcode once D79658061 lands",
+    )
     def test_dequantize_codebook_embedding(self):
         model, example_inputs = self._get_test_model()
         quantize_(

backends/arm/_passes/__init__.py

@@ -31,6 +31,7 @@
 from .decompose_atanh_pass import DecomposeAtanhPass  # noqa
 from .decompose_avg_pool2d import DecomposeAvgPool2d  # noqa
 from .decompose_batch_norm_no_stats import DecomposeBatchNormNoStatsPass  # noqa
+from .decompose_cosh_pass import DecomposeCoshPass  # noqa
 from .decompose_cosine_similarity_pass import DecomposeCosineSimilarityPass  # noqa
 from .decompose_div_pass import DecomposeDivPass  # noqa
 from .decompose_embedding_pass import DecomposeEmbeddingPass  # noqa  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -36,6 +36,7 @@
     DecomposeAtanPass,
     DecomposeAvgPool2d,
     DecomposeBatchNormNoStatsPass,
+    DecomposeCoshPass,
     DecomposeCosineSimilarityPass,
     DecomposeDivPass,
     DecomposeEmbeddingPass,
@@ -167,6 +168,7 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(DecomposeAcoshPass())
         self.add_pass(DecomposeAsinPass())
         self.add_pass(DecomposeAsinhPass())
+        self.add_pass(DecomposeCoshPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeAtanPass())
         self.add_pass(DecomposeAtanhPass())

backends/arm/_passes/decompose_cosh_pass.py

@@ -0,0 +1,48 @@
+# 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 executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+
+# For MI case
+edge_cosh = exir_ops.edge.aten.cosh.default
+
+
+class DecomposeCoshPass(ArmPass):
+    """
+    This pass replaces the cosh operator with a sequence of TOSA-equivalent operations that
+    compute the hyperbolic cosine using the formula:
+
+        cosh(x) = 0.5 * (e^x + e^(-x))
+
+    """
+
+    def call_operator(self, op, args, kwargs, meta, updated=False):
+        if op is not edge_cosh:
+            return super().call_operator(op, args, kwargs, meta, updated)
+
+        x = args
+
+        exp_op, mul_op, neg_op, add_op = (
+            exir_ops.edge.aten.exp.default,
+            exir_ops.edge.aten.mul.Scalar,
+            exir_ops.edge.aten.neg.default,
+            exir_ops.edge.aten.add.Tensor,
+        )
+
+        # exp1 = e^x
+        exp1 = super().call_operator(exp_op, x, {}, meta, updated=True)
+
+        # exp2 = e^(⁻x)
+        neg_x = super().call_operator(neg_op, x, {}, meta, updated=True)
+        exp2 = super().call_operator(exp_op, (neg_x,), {}, meta, updated=True)
+
+        # numer = exp1 + exp2
+        numer = super().call_operator(add_op, (exp1, exp2), {}, meta, updated=True)
+
+        # out = 0.5 * numer
+        out = super().call_operator(mul_op, (numer, 0.5), {}, meta, updated=True)
+
+        return out

backends/arm/_passes/insert_table_ops.py

@@ -59,6 +59,7 @@ class TableOps:
         exir_ops.edge.aten.acosh.default: torch.acosh,
         exir_ops.edge.aten.asin.default: torch.asin,
         exir_ops.edge.aten.asinh.default: torch.asinh,
+        exir_ops.edge.aten.cosh.default: torch.cosh,
     }
 
     # Targets that must be treated explicitly

backends/arm/operator_support/tosa_supported_operators.py

@@ -257,6 +257,7 @@ def is_node_supported(
             exir_ops.edge.aten.addmm.default,
             exir_ops.edge.aten.masked_fill.Scalar,
             exir_ops.edge.aten.asinh.default,
+            exir_ops.edge.aten.cosh.default,
         ]
 
         return supported

backends/arm/process_node.py

@@ -8,11 +8,12 @@
 from typing import Any, cast, Dict
 
 import numpy as np
+import serializer.tosa_serializer as ts
 import torch
 import torch.fx
 from executorch.backends.arm.operators.node_visitor import NodeVisitor
 from executorch.backends.arm.tosa_mapping import TosaArg
-from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification
+from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.arm.tosa_utils import getNodeArgs, tosa_shape
 from torch._export.utils import (
     get_buffer,
@@ -81,11 +82,6 @@ def process_inputs(
             "Is the original torch function supported?"
         ) from e
 
-    if isinstance(tosa_spec, Tosa_1_00):
-        import serializer.tosa_serializer as ts
-    else:
-        raise ValueError(f"Unsupported TOSA spec: {tosa_spec}")
-
     input_shape = tosa_arg.shape
     input_dim_order = tosa_arg.dim_order
     tensor = ts.TosaSerializerTensor(

backends/arm/quantizer/quantization_annotator.py

@@ -287,6 +287,7 @@ def _match_pattern(
     torch.ops.aten.asin.default,
     torch.ops.aten.atanh.default,
     torch.ops.aten.asinh.default,
+    torch.ops.aten.cosh.default,
 ]
 
 _one_to_one_shared_input_qspec = [

backends/arm/test/models/test_inception_v3_arm.py

@@ -0,0 +1,121 @@
+# 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 common
+import pytest
+
+import torch
+
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+from torchvision import models, transforms
+
+ic3 = models.inception_v3(weights=models.Inception_V3_Weights)
+ic3 = ic3.eval()
+
+# Normalization values referenced from here:
+# https://docs.pytorch.org/vision/main/models/generated/torchvision.models.quantization.inception_v3.html
+normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+model_inputs = (normalize(torch.rand(1, 3, 224, 224)),)
+input_t = Tuple[torch.Tensor]
+
+
+@pytest.mark.slow
+def test_ic3_tosa_FP():
+    pipeline = TosaPipelineFP[input_t](
+        ic3,
+        model_inputs,
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+@pytest.mark.slow
+def test_ic3_tosa_BI():
+    pipeline = TosaPipelineINT[input_t](
+        ic3,
+        model_inputs,
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+        atol=0.5,
+        qtol=1,
+    )
+    pipeline.run()
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="Takes too long to run on CI")
+@common.XfailIfNoCorstone300
+def test_ic3_u55_BI():
+    pipeline = EthosU55PipelineINT[input_t](
+        ic3,
+        model_inputs,
+        aten_ops=[],
+        exir_ops=[],
+        run_on_fvp=True,
+        use_to_edge_transform_and_lower=True,
+        atol=0.5,
+        qtol=1,
+    )
+    pipeline.run()
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="Takes too long to run on CI")
+@common.XfailIfNoCorstone320
+def test_ic3_u85_BI():
+    pipeline = EthosU85PipelineINT[input_t](
+        ic3,
+        model_inputs,
+        aten_ops=[],
+        exir_ops=[],
+        run_on_fvp=True,
+        use_to_edge_transform_and_lower=True,
+        atol=0.5,
+        qtol=1,
+    )
+    pipeline.run()
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="Takes too long to run on CI")
+@common.SkipIfNoModelConverter
+def test_ic3_vgf_FP():
+    pipeline = VgfPipeline[input_t](
+        ic3,
+        model_inputs,
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+FP",
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="Takes too long to run on CI")
+@common.SkipIfNoModelConverter
+def test_ic3_vgf_INT():
+    pipeline = VgfPipeline[input_t](
+        ic3,
+        model_inputs,
+        aten_op=[],
+        exir_op=[],
+        tosa_version="TOSA-1.0+INT",
+        use_to_edge_transform_and_lower=True,
+    )
+    pipeline.run()