Update on "[Executorch][SDPA] Refactor + Make quantized sdpa handle sequence at dim 1 or 2"

For quantized SDPA we want to evaluate performance impact of having seq at dim 1 as well as dim 2.
This diff refactors the code to enable this.

The same should be done also for float SDPA but left for future.

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

[ghstack-poisoned]

Author: kimishpatel

.ci/scripts/test_ios_ci.sh

@@ -7,7 +7,7 @@
 
 set -e
 
-APP_PATH="examples/demo-apps/apple_ios/ExecuTorchDemo/ExecuTorchDemo"
+APP_PATH="executorch-examples/apple/ExecuTorchDemo/ExecuTorchDemo"
 MODEL_NAME="mv3"
 SIMULATOR_NAME="executorch"
 
@@ -34,6 +34,10 @@ say() {
   echo -e "\033[1m\n\t** $1 **\n\033[0m"
 }
 
+say "Cloning the Demo App"
+
+git clone --depth 1 https://github.com/pytorch-labs/executorch-examples.git
+
 say "Installing CoreML Backend Requirements"
 
 ./backends/apple/coreml/scripts/install_requirements.sh

.github/workflows/android-release-artifacts.yml

@@ -7,6 +7,10 @@ on:
         description: Version name to be uploaded for AAR release
         required: false
         type: string
+      upload_to_maven:
+        description: Upload the AAR to maven staging repository
+        required: false
+        type: boolean
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
@@ -31,11 +35,14 @@ jobs:
   build-aar:
     name: build-aar
     needs: check-if-aar-exists
-    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    if: ${{ !github.event.pull_request.head.repo.fork }}
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@release/2.7
+    secrets: inherit
     permissions:
       id-token: write
       contents: read
     with:
+      secrets-env: EXECUTORCH_MAVEN_SIGNING_KEYID EXECUTORCH_MAVEN_SIGNING_PASSWORD EXECUTORCH_MAVEN_CENTRAL_PASSWORD EXECUTORCH_MAVEN_CENTRAL_USERNAME EXECUTORCH_MAVEN_SIGNING_GPG_KEY_CONTENTS
       runner: linux.2xlarge
       docker-image: executorch-ubuntu-22.04-clang12-android
       submodules: 'true'
@@ -52,6 +59,16 @@ jobs:
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool buck2
         export ARTIFACTS_DIR_NAME=artifacts-to-be-uploaded
 
+        mkdir -p ~/.gradle
+        touch ~/.gradle/gradle.properties
+        echo "signing.keyId=${SECRET_EXECUTORCH_MAVEN_SIGNING_KEYID}" >> ~/.gradle/gradle.properties
+        echo "signing.password=${SECRET_EXECUTORCH_MAVEN_SIGNING_PASSWORD}" >> ~/.gradle/gradle.properties
+        echo "mavenCentralUsername=${SECRET_EXECUTORCH_MAVEN_CENTRAL_USERNAME}" >> ~/.gradle/gradle.properties
+        echo "mavenCentralPassword=${SECRET_EXECUTORCH_MAVEN_CENTRAL_PASSWORD}" >> ~/.gradle/gradle.properties
+        echo "signing.secretKeyRingFile=/tmp/secring.gpg" >> ~/.gradle/gradle.properties
+
+        echo -n "$SECRET_EXECUTORCH_MAVEN_SIGNING_GPG_KEY_CONTENTS" | base64 -d > /tmp/secring.gpg
+
         # Build AAR Package
         mkdir aar-out
         export BUILD_AAR_DIR=aar-out
@@ -61,6 +78,12 @@ jobs:
 
         shasum -a 256 "${ARTIFACTS_DIR_NAME}/executorch.aar"
 
+        # Publish to maven staging
+        UPLOAD_TO_MAVEN="${{ inputs.upload_to_maven }}"
+        if [[ "$UPLOAD_TO_MAVEN" == "true" ]]; then
+          (cd aar-out; ANDROID_HOME="${ANDROID_SDK:-/opt/android/sdk}" ./gradlew :executorch_android:publishToMavenCentral)
+        fi
+
   upload-release-aar:
     name: upload-release-aar
     needs: build-aar

backends/arm/scripts/parse_test_names.py

@@ -0,0 +1,102 @@
+# 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.exir.dialects.edge.spec.utils import SAMPLE_INPUT
+
+# Add edge ops which we lower but which are not included in exir/dialects/edge/edge.yaml here.
+CUSTOM_EDGE_OPS = ["linspace.default", "eye.default"]
+ALL_EDGE_OPS = SAMPLE_INPUT.keys() | CUSTOM_EDGE_OPS
+
+# Add all targets and TOSA profiles we support here.
+TARGETS = {"tosa_BI", "tosa_MI", "u55_BI", "u85_BI"}
+
+
+def get_edge_ops():
+    """
+    Returns a set with edge_ops with names on the form to be used in unittests:
+    1. Names are in lowercase.
+    2. Overload is ignored if it is 'default', otherwise its appended with an underscore.
+    3. Overly verbose name are shortened by removing certain prefixes/suffixes.
+
+    Examples:
+        abs.default -> abs
+        split_copy.Tensor -> split_tensor
+    """
+    edge_ops = set()
+    for edge_name in ALL_EDGE_OPS:
+        op, overload = edge_name.split(".")
+
+        # Normalize names
+        op = op.lower()
+        op = op.removeprefix("_")
+        op = op.removesuffix("_copy")
+        op = op.removesuffix("_with_indices")
+        op = op.removesuffix("_no_training")
+        overload = overload.lower()
+
+        if overload == "default":
+            edge_ops.add(op)
+        else:
+            edge_ops.add(f"{op}_{overload}")
+
+    return edge_ops
+
+
+def parse_test_name(test_name: str, edge_ops: set[str]) -> tuple[str, str, bool]:
+    """
+    Parses a test name on the form
+        test_OP_TARGET_<not_delegated>_<any_other_info>
+    where OP must match a string in edge_ops and TARGET must match one string in TARGETS.
+    The "not_delegated" suffix indicates that the test tests that the op is not delegated.
+
+    Examples of valid names: "test_mm_u55_BI_not_delegated" or "test_add_scalar_tosa_MI_two_inputs".
+
+    Returns a tuple (OP, TARGET, IS_DELEGATED) if valid.
+    """
+    test_name = test_name.removeprefix("test_")
+    is_delegated = "not_delegated" not in test_name
+    assert (
+        "reject" not in test_name
+    ), f"Use 'not_delegated' instead of 'reject' in {test_name}"
+
+    op = "None"
+    target = "None"
+    for potential_target in TARGETS:
+        index = test_name.find(potential_target)
+        if index != -1:
+            op = test_name[: index - 1]
+            target = potential_target
+            break
+    # Special case for convolution
+    op = op.removesuffix("_1d")
+    op = op.removesuffix("_2d")
+
+    assert target != "None", f"{test_name} does not contain one of {TARGETS}"
+    assert (
+        op in edge_ops
+    ), f"Parsed unvalid OP from {test_name}, {op} does not exist in edge.yaml or CUSTOM_EDGE_OPS"
+
+    return op, target, is_delegated
+
+
+if __name__ == "__main__":
+    """Parses a list of test names given on the commandline."""
+    import sys
+
+    sys.tracebacklimit = 0  # Do not print stack trace
+
+    edge_ops = get_edge_ops()
+    exit_code = 0
+
+    for test_name in sys.argv[1:]:
+        try:
+            assert test_name[:5] == "test_", f"Unexpected input: {test_name}"
+            parse_test_name(test_name, edge_ops)
+        except AssertionError as e:
+            print(e)
+            exit_code = 1
+        else:
+            print(f"{test_name} OK")
+
+    sys.exit(exit_code)

backends/arm/scripts/pre-push

@@ -166,6 +166,44 @@ for COMMIT in ${COMMITS}; do
         fi
     fi
 
+    # Op test checks
+    op_test_files=$(echo $commit_files | grep -oE 'backends/arm/test/ops/\S+')
+    if [ "$op_test_files" ]; then
+
+        # TODO: These checks can be removed when all unittests are refactored.
+        if grep -icq "SkipIfNoCorstone" $op_test_files; then
+            echo -e "${ERROR} @SkipIfNoCorstone300/320 is deprecated;"\
+                "please use XfailIfNoCorstone300/320 instead." >&2
+            FAILED=1
+        fi
+
+        if grep -icq "conftest.expectedFailureOnFVP" $op_test_files; then
+            echo -e "${ERROR} @conftest.expectedFailureOnFVP is deprecated;"\
+                "please use XfailIfCorstone300/320 instead." >&2
+            FAILED=1
+        fi
+
+        if grep -icq "unittest.TestCase" $op_test_files; then
+            echo -e "${ERROR} Use of the Unittest test framework is deprecated;"\
+            "please use Pytest instead." >&2
+            FAILED=1
+        fi
+
+        if grep -icq "on_fvp(" $op_test_files; then
+            echo -e "${ERROR} All unittests should run on FVP if relevant,"\
+            "on_fvp suffix can be excluded." >&2
+            FAILED=1
+        fi
+
+        # Check that the tested op and target is parsed correctly from the test name
+        test_names=$(grep -h "def test_" $op_test_files | cut -d"(" -f1 | cut -d" " -f2)
+        python ./backends/arm/scripts/parse_test_names.py $test_names
+        if [ $? -ne 0 ]; then
+            echo -e "${ERROR} Failed op test name check." >&2
+            FAILED=1
+        fi
+    fi
+
     echo "" # Newline to visually separate commit processing
 done
 

backends/cadence/aot/quantizer/quantizer.py

@@ -43,7 +43,7 @@
 from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
 
 
-act_qspec_asym8u = QuantizationSpec(
+act_qspec_asym8s = QuantizationSpec(
     dtype=torch.int8,
     quant_min=-128,
     quant_max=127,
@@ -52,7 +52,7 @@
     observer_or_fake_quant_ctr=HistogramObserver.with_args(eps=2**-12),
 )
 
-wgt_qspec_asym8u = QuantizationSpec(
+wgt_qspec_asym8s = QuantizationSpec(
     dtype=torch.int8,
     quant_min=-128,
     quant_max=127,
@@ -61,7 +61,7 @@
     observer_or_fake_quant_ctr=MinMaxObserver,
 )
 
-wgt_qspec_asym8s = QuantizationSpec(
+wgt_qspec_sym8s = QuantizationSpec(
     dtype=torch.int8,
     quant_min=-128,
     quant_max=127,
@@ -72,17 +72,17 @@
 
 bias_qspec: Optional[QuantizationSpec] = None
 
-qconfig_A8uW8u = QuantizationConfig(
-    act_qspec_asym8u,
-    act_qspec_asym8u,
-    wgt_qspec_asym8u,
+qconfig_A8W8 = QuantizationConfig(
+    act_qspec_asym8s,
+    act_qspec_asym8s,
+    wgt_qspec_asym8s,
     None,
 )
 
-qconfig_A8uW8s = QuantizationConfig(
-    act_qspec_asym8u,
-    act_qspec_asym8u,
-    wgt_qspec_asym8s,
+qconfig_A8W8sym = QuantizationConfig(
+    act_qspec_asym8s,
+    act_qspec_asym8s,
+    wgt_qspec_sym8s,
     None,
 )
 
@@ -189,15 +189,15 @@ def get_supported_operators(cls) -> List[OperatorConfig]:
 
 def get_cadence_default_quantizers() -> List[Quantizer]:
     return [
-        CadenceAtenQuantizer(AddmmPattern(), qconfig_A8uW8u),
-        CadenceAtenQuantizer(BmmPattern(), qconfig_A8uW8u),
-        CadenceAtenQuantizer(Conv1dPattern(), qconfig_A8uW8s),
-        CadenceAtenQuantizer(Conv2dPattern(), qconfig_A8uW8s),
-        CadenceAtenQuantizer(LayerNormPattern(), qconfig_A8uW8u),
-        CadenceAtenQuantizer(LinearPattern(), qconfig_A8uW8u),
-        CadenceAtenQuantizer(MatmulPattern(), qconfig_A8uW8u),
-        CadenceAtenQuantizer(ReluPattern0(), qconfig_A8uW8u),
-        CadenceAtenQuantizer(ReluPattern1(), qconfig_A8uW8u),
+        CadenceAtenQuantizer(AddmmPattern(), qconfig_A8W8),
+        CadenceAtenQuantizer(BmmPattern(), qconfig_A8W8),
+        CadenceAtenQuantizer(Conv1dPattern(), qconfig_A8W8sym),
+        CadenceAtenQuantizer(Conv2dPattern(), qconfig_A8W8sym),
+        CadenceAtenQuantizer(LayerNormPattern(), qconfig_A8W8),
+        CadenceAtenQuantizer(LinearPattern(), qconfig_A8W8),
+        CadenceAtenQuantizer(MatmulPattern(), qconfig_A8W8),
+        CadenceAtenQuantizer(ReluPattern0(), qconfig_A8W8),
+        CadenceAtenQuantizer(ReluPattern1(), qconfig_A8W8),
     ]
 
 
@@ -244,6 +244,6 @@ class CadenceWakeWordQuantizer(CadenceQuantizer):
     def __init__(self, quantizers: Optional[list[Quantizer]] = None) -> None:
         if quantizers is None:
             quantizers = get_cadence_default_quantizers()
-        quantizers.append(CadenceAtenQuantizer(AddPattern(), qconfig_A8uW8u))
-        quantizers.append(CadenceAtenQuantizer(CatPattern(), qconfig_A8uW8u))
+        quantizers.append(CadenceAtenQuantizer(AddPattern(), qconfig_A8W8))
+        quantizers.append(CadenceAtenQuantizer(CatPattern(), qconfig_A8W8))
         super().__init__(quantizers)

backends/qualcomm/_passes/layout_transform.py

@@ -47,6 +47,7 @@ class LayoutTransform(ExportPass):
     layout_agnostic_ops = {
         exir_ops.edge.aten.abs.default,
         exir_ops.edge.aten.add.Tensor,
+        exir_ops.edge.aten.amax.default,
         exir_ops.edge.aten.bitwise_or.Tensor,
         exir_ops.edge.aten.bmm.default,
         exir_ops.edge.aten.bitwise_and.Tensor,

backends/qualcomm/builders/__init__.py

@@ -9,6 +9,7 @@
     op_abs,
     op_adaptive_avg_pool2d,
     op_add,
+    op_amax,
     op_and,
     op_arange,
     op_argmin,
@@ -95,6 +96,7 @@
     op_abs,
     op_adaptive_avg_pool2d,
     op_add,
+    op_amax,
     op_and,
     op_arange,
     op_argmin,