Add a16w8 per-op test for mean_dim (pytorch#19594) (pytorch#19594)

Summary:

Add int16 activation / int8 weight (a16w8) quantization tests for
`aten.mean.dim` on Ethos-U55 and Ethos-U85.

## Changes
- Add `a16w8_mean_test_parameters` dict with 11 test configurations
covering keepdim/no-keepdim, positive/negative dims, dim=None, and ranks
1-4
- Add `test_mean_dim_a16w8_u55_INT` using `EthosU55PipelineINT` with
`a16w8_quantization=True, symmetric_io_quantization=True`
- Add `test_mean_dim_a16w8_u85_INT` using `EthosU85PipelineINT` with
same kwargs
- Register `ops/test_mean_dim.py` in `fbcode/` and `xplat/`
`targets.bzl`

Differential Revision: D104532361

Author: christine-long-meta

backends/arm/test/ops/test_mean_dim.py

@@ -6,6 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 from typing import Callable
 
+import pytest
 import torch
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
@@ -393,3 +394,68 @@ def test_mean_tosa_INT(test_data):
         symmetric_io_quantization=True,
     )
     pipeline.run()
+
+
+a16w8_mean_test_parameters = {
+    "rank_2_keepdim": lambda: (torch.rand(7, 3), (0, 1), True),
+    "rank_3_keepdim": lambda: (torch.rand(5, 7, 3), (0, 1, 2), True),
+    "rand_23_keepdim": lambda: (torch.rand(1, 5, 7, 3), (2, 3), True),
+    "rand_0123_keepdim": lambda: (torch.rand(1, 5, 7, 3), (0, 1, 2, 3), True),
+    "rand_none_keepdim": lambda: (torch.rand(1, 5, 7, 3), None, True),
+    "rank_2": lambda: (torch.rand(5, 7), (-2, -1), False),
+    "rand_123": lambda: (torch.rand(1, 5, 7, 3), (-3, -2, -1), False),
+}
+
+a16w8_mean_test_parameters_u85_xfails = {
+    "rank_1_keepdim": lambda: (torch.rand(7), 0, True),
+    "rand_1_keepdim": lambda: (torch.rand(1, 5, 7, 3), (1), True),
+    "rank_1": lambda: (torch.rand(7), (-1), False),
+    "rand_3": lambda: (torch.rand(1, 5, 7, 3), (-1), False),
+}
+
+
+@common.parametrize(
+    "test_data", {**a16w8_mean_test_parameters, **a16w8_mean_test_parameters_u85_xfails}
+)
+@common.XfailIfNoCorstone300
+def test_mean_dim_a16w8_u55_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU55PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],
+        symmetric_io_quantization=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", a16w8_mean_test_parameters)
+@common.XfailIfNoCorstone320
+def test_mean_dim_a16w8_u85_INT(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU85PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],
+        symmetric_io_quantization=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", a16w8_mean_test_parameters_u85_xfails)
+@pytest.mark.xfail(
+    reason="U85 mean_dim a16w8 produces all-zero output for single-dim reductions"
+)
+@common.XfailIfNoCorstone320
+def test_mean_dim_a16w8_u85_INT_xfail(test_data):
+    test_data, dim, keep_dim = test_data()
+    pipeline = EthosU85PipelineINT[input_t](
+        MeanDim(dim, keep_dim),
+        (test_data,),
+        [],
+        symmetric_io_quantization=True,
+        a16w8_quantization=True,
+    )
+    pipeline.run()

backends/arm/test/targets.bzl

@@ -38,6 +38,7 @@ def define_arm_tests():
         "ops/test_to_copy.py",
         "ops/test_exp.py",
         "ops/test_reciprocal.py",
+        "ops/test_mean_dim.py",
     ]
 
     # Quantization